Setup both explorer and trainer. For the explorer, a step contains `batch_size * sync_interval` number of rollout tasks. For the trainer, it has to consume all experiences generated by the explorer in the latest step. The specific number of experiences may vary for different al
(config: Config)
| 221 | |
| 222 | |
| 223 | def both(config: Config) -> StageStatus: |
| 224 | """Setup both explorer and trainer. |
| 225 | |
| 226 | For the explorer, a step contains `batch_size * sync_interval` number |
| 227 | of rollout tasks. |
| 228 | |
| 229 | For the trainer, it has to consume all experiences generated by the explorer in |
| 230 | the latest step. The specific number of experiences may vary for different |
| 231 | algorithms and tasks. |
| 232 | """ |
| 233 | from trinity.common.constants import SyncMethod |
| 234 | from trinity.explorer.explorer import Explorer |
| 235 | from trinity.manager.synchronizer import Synchronizer |
| 236 | from trinity.trainer.trainer import Trainer |
| 237 | |
| 238 | explorer = Explorer.get_actor(config) |
| 239 | trainer = Trainer.get_actor(config) |
| 240 | started_at = time.perf_counter() |
| 241 | try: |
| 242 | ray.get([explorer.__ray_ready__.remote(), trainer.__ray_ready__.remote()]) |
| 243 | ray.get( |
| 244 | [ |
| 245 | explorer.prepare.remote(), |
| 246 | trainer.prepare.remote(), |
| 247 | ] |
| 248 | ) |
| 249 | # Set up NCCL weight sync group between Trainer and Explorer. |
| 250 | # This must happen after both sides are prepared (Trainer has model |
| 251 | # meta cached, Explorer has rollout models created) and before the |
| 252 | # first weight sync. |
| 253 | if config.synchronizer.sync_method == SyncMethod.NCCL: |
| 254 | synchronizer = Synchronizer.get_actor(namespace=config.ray_namespace) |
| 255 | ray.get(synchronizer.coordinate_weight_sync_setup.remote()) |
| 256 | ray.get( |
| 257 | [ |
| 258 | explorer.sync_weight.remote(), |
| 259 | trainer.sync_weight.remote(), |
| 260 | ] |
| 261 | ) |
| 262 | ready_ref, wait_ref = ray.wait( |
| 263 | [ |
| 264 | explorer.explore.remote(), |
| 265 | trainer.train.remote(), |
| 266 | ], |
| 267 | num_returns=1, |
| 268 | ) |
| 269 | |
| 270 | ready = ray.get(ready_ref[0]) |
| 271 | if ready == config.trainer.name: |
| 272 | logger.info( |
| 273 | "===========================================================\n" |
| 274 | "> Launcher detected that the `Trainer` process has finished.\n" |
| 275 | "> Stopping the explorer process immediately.\n" |
| 276 | "===========================================================" |
| 277 | ) |
| 278 | ray.wait(wait_ref, timeout=5) |
| 279 | elif ready == config.explorer.name: |
| 280 | logger.info( |