Preparation before running.
(self)
| 246 | ) |
| 247 | |
| 248 | async def prepare(self) -> None: |
| 249 | """Preparation before running.""" |
| 250 | try: |
| 251 | # make sure all rollout models are ready |
| 252 | self.models, self.auxiliary_models = await self.model_allocator.create_all_models() |
| 253 | self.logger.info("All models are ready.") |
| 254 | if not self.use_nccl_sync and self.model_type not in {"tinker", "external"}: |
| 255 | if self.config.mode == "serve": |
| 256 | # In serving mode, each engine will setup its own process group |
| 257 | await self.setup_model_level_weight_sync_group() |
| 258 | else: |
| 259 | master_address, master_port = await self.models[0].get_available_address_async( |
| 260 | random_port=True |
| 261 | ) |
| 262 | await self.setup_weight_sync_group(master_address, master_port) |
| 263 | |
| 264 | self.rollout_coordinator = RolloutCoordinator.get_actor(self.config) |
| 265 | await self.rollout_coordinator.prepare.remote() |
| 266 | self.logger.info("Rollout coordinator is ready.") |
| 267 | if self.config.explorer.eval_on_startup and self.explore_step_num == 0: |
| 268 | await self.eval() |
| 269 | |
| 270 | await self.synchronizer.set_explorer_status.remote(RunningStatus.RUNNING) |
| 271 | if self.sync_style == SyncStyle.FULLY_ASYNC: |
| 272 | self._async_watch_task = asyncio.create_task(self._watch_trainer_sync_signal()) |
| 273 | self.logger.info("Trainer sync watcher task started.") |
| 274 | self.logger.info("Explorer is ready.") |
| 275 | except Exception as e: |
| 276 | self.logger.error(f"Error during explorer preparation: {traceback.format_exc()}") |
| 277 | await self.shutdown() |
| 278 | raise e |
| 279 | |
| 280 | async def get_weight(self, name: str) -> torch.Tensor: |
| 281 | """Get the weight of the loaded model (For checkpoint weights update).""" |
nothing calls this directly
no test coverage detected