(
n_workers: int)
| 46 | |
| 47 | |
| 48 | def create_mpi_comm_session( |
| 49 | n_workers: int) -> RemoteMpiCommSessionClient | MpiPoolSession: |
| 50 | assert mpi_rank( |
| 51 | ) == 0, f"create_mpi_comm_session must be called by rank 0, but it was called by rank {mpi_rank()}" |
| 52 | if get_spawn_proxy_process_env(): |
| 53 | assert get_spawn_proxy_process_ipc_addr_env( |
| 54 | ), f"{LlmLauncherEnvs.TLLM_SPAWN_PROXY_PROCESS_IPC_ADDR} is not set." |
| 55 | logger_debug( |
| 56 | f"Using RemoteMpiPoolSessionClient to bind to external MPI processes at {get_spawn_proxy_process_ipc_addr_env()}\n", |
| 57 | "yellow") |
| 58 | hmac_key = get_spawn_proxy_process_ipc_hmac_key_env() |
| 59 | return RemoteMpiCommSessionClient( |
| 60 | addr=get_spawn_proxy_process_ipc_addr_env(), hmac_key=hmac_key) |
| 61 | else: |
| 62 | logger_debug( |
| 63 | f"Using MpiCommSession to bind to external MPI processes\n", |
| 64 | "yellow") |
| 65 | return MpiCommSession(n_workers=n_workers) |
| 66 | |
| 67 | |
| 68 | def has_event_loop() -> bool: |
no test coverage detected