(
self,
config: Config,
rollout_model_id: int,
auxiliary_model_ids: Optional[List[int]] = None,
runner_id: Optional[int] = None,
)
| 46 | """A Ray remote actor to run the workflow and generate experiences.""" |
| 47 | |
| 48 | def __init__( |
| 49 | self, |
| 50 | config: Config, |
| 51 | rollout_model_id: int, |
| 52 | auxiliary_model_ids: Optional[List[int]] = None, |
| 53 | runner_id: Optional[int] = None, |
| 54 | ) -> None: |
| 55 | self.name = f"{config.explorer.name}_runner_{runner_id}" |
| 56 | self.logger = get_logger(self.name, in_ray_actor=True) |
| 57 | self.config = config |
| 58 | allocator = Allocator(config.explorer) |
| 59 | self.model_wrapper: ModelWrapper = allocator.get_model( |
| 60 | config.explorer.rollout_model, "rollout", rollout_model_id |
| 61 | ) |
| 62 | self.auxiliary_model_wrappers: List[ModelWrapper] = [ |
| 63 | allocator.get_model( |
| 64 | config.explorer.auxiliary_models[index], f"auxiliary_{index}", auxiliary_model_id |
| 65 | ) |
| 66 | for index, auxiliary_model_id in enumerate(auxiliary_model_ids or []) |
| 67 | ] |
| 68 | self.workflow_instance: Workflow = None |
| 69 | self.runner_id = runner_id |
| 70 | self.runner_state = { |
| 71 | "workflow_id": None, |
| 72 | "model_version": None, |
| 73 | "begin_time": 0, |
| 74 | "terminate_time": 0, |
| 75 | } |
| 76 | self.concurrent_mode = config.explorer.concurrent_mode |
| 77 | if self.concurrent_mode == "sequential": |
| 78 | self.concurrent_run_fn = self._sequential_run |
| 79 | elif self.concurrent_mode == "asynchronous": |
| 80 | self.concurrent_run_fn = self._asynchronous_run |
| 81 | elif self.concurrent_mode == "multi-threading": |
| 82 | self.concurrent_run_fn = self._multi_threading_run |
| 83 | else: |
| 84 | self.logger.warning( |
| 85 | f"Unknown concurrent_mode {self.concurrent_mode}, defaulting to sequential." |
| 86 | ) |
| 87 | self.concurrent_run_fn = self._sequential_run |
| 88 | self.logger.info( |
| 89 | f"WorkflowRunner [{self.name}]({self.concurrent_mode}) initialized:\n" |
| 90 | f" > rollout model: {self.config.explorer.rollout_model.model_path}\n" |
| 91 | f" > auxiliary models: {[aux_model_config.model_path for aux_model_config in self.config.explorer.auxiliary_models]}" |
| 92 | ) |
| 93 | |
| 94 | async def prepare(self) -> None: |
| 95 | """Prepare the runner.""" |
no test coverage detected