(
self,
worker_kwargs: dict,
model_world_size: int = 1,
mpi_session: Optional[MpiSession] = None,
*,
worker_cls: type = GenerationExecutorWorker,
postproc_worker_config: Optional[PostprocWorkerConfig] = None,
is_llm_executor: Optional[bool] = None,
)
| 39 | READY_SIGNAL = b"READY" |
| 40 | |
| 41 | def __init__( |
| 42 | self, |
| 43 | worker_kwargs: dict, |
| 44 | model_world_size: int = 1, |
| 45 | mpi_session: Optional[MpiSession] = None, |
| 46 | *, |
| 47 | worker_cls: type = GenerationExecutorWorker, |
| 48 | postproc_worker_config: Optional[PostprocWorkerConfig] = None, |
| 49 | is_llm_executor: Optional[bool] = None, |
| 50 | ) -> None: |
| 51 | postproc_worker_config = postproc_worker_config or PostprocWorkerConfig( |
| 52 | ) |
| 53 | super().__init__( |
| 54 | num_postprocess_workers=postproc_worker_config. |
| 55 | num_postprocess_workers, |
| 56 | postprocess_tokenizer_dir=postproc_worker_config. |
| 57 | postprocess_tokenizer_dir, |
| 58 | is_llm_executor=is_llm_executor, |
| 59 | ) |
| 60 | |
| 61 | self.workers_started = False |
| 62 | self.worker_cls = worker_cls |
| 63 | |
| 64 | mpi_process_pre_spawned: bool = get_spawn_proxy_process_env() |
| 65 | |
| 66 | if mpi_session is None: |
| 67 | if mpi_process_pre_spawned: |
| 68 | logger_debug('create comm session ...\n', "yellow") |
| 69 | self.mpi_session = create_mpi_comm_session(model_world_size) |
| 70 | else: |
| 71 | logger_debug('create pool session ...\n', "yellow") |
| 72 | self.mpi_session = MpiPoolSession(n_workers=model_world_size) |
| 73 | else: |
| 74 | logger_debug('using external mpi session ...\n', "yellow") |
| 75 | self.mpi_session = mpi_session |
| 76 | |
| 77 | if isinstance(self.mpi_session, |
| 78 | (MpiCommSession, RemoteMpiCommSessionClient)): |
| 79 | print_colored( |
| 80 | f"rank {mpi_rank()} using MpiCommSession to bind to external MPI processes\n", |
| 81 | "yellow") |
| 82 | else: |
| 83 | print_colored( |
| 84 | f"rank {mpi_rank()} using MpiPoolSession to spawn MPI processes\n", |
| 85 | "yellow") |
| 86 | |
| 87 | self._results: Dict[int, GenerationResult] = {} |
| 88 | |
| 89 | self.model_world_size = model_world_size |
| 90 | |
| 91 | self.garbage_collection_gen0_threshold = worker_kwargs[ |
| 92 | "llm_args"].garbage_collection_gen0_threshold if worker_kwargs.get( |
| 93 | "llm_args", None) is not None else None |
| 94 | |
| 95 | # Generate RPC address and key for stats RPC |
| 96 | self.rpc_addr = get_unique_ipc_addr() |
| 97 | self.hmac_key = os.urandom(32) |
| 98 |
nothing calls this directly
no test coverage detected