| 181 | |
| 182 | |
| 183 | class MpiCommSession(MpiSession): |
| 184 | |
| 185 | def __init__(self, comm=None, n_workers: int = 1): |
| 186 | self.comm = comm |
| 187 | self.n_workers = n_workers |
| 188 | self.thread_pool: Optional[ThreadPoolExecutor] = None |
| 189 | self.mpi_pool: Optional[MPIPoolExecutor] = None |
| 190 | self.owns_mpi_pool = False # Track if this instance owns the mpi_pool |
| 191 | |
| 192 | if n_workers <= 0: |
| 193 | raise ValueError( |
| 194 | f'n_workers must be non-negative, but got {n_workers}') |
| 195 | |
| 196 | if ENABLE_MULTI_DEVICE: |
| 197 | if not self.comm: |
| 198 | self.comm = mpi4py.MPI.COMM_WORLD |
| 199 | |
| 200 | if self.comm.Get_rank() != 0: |
| 201 | raise RuntimeError( |
| 202 | f'only rank 0 can start multi-node session, got {self.comm.Get_rank()}' |
| 203 | ) |
| 204 | |
| 205 | if self.comm.Get_size() != n_workers: |
| 206 | raise ValueError( |
| 207 | f'n_workers must be equal to the number of processes in MPI, got {n_workers} vs {get_mpi_world_size()}' |
| 208 | ) |
| 209 | |
| 210 | self._start_mpi_pool() |
| 211 | |
| 212 | def get_comm(self): |
| 213 | return self.comm |
| 214 | |
| 215 | def submit(self, task: Callable[..., T], *args, |
| 216 | **kwargs) -> List[Future[T]]: |
| 217 | ''' Submit a task to MPI workers. |
| 218 | |
| 219 | Args: |
| 220 | task: The task to be submitted. |
| 221 | args: Positional arguments for the task. |
| 222 | kwargs: Keyword arguments for the task. |
| 223 | ''' |
| 224 | assert self.mpi_pool is not None, 'MPI session not started' |
| 225 | worker_futures = [ |
| 226 | self.mpi_pool.submit(task, *args, **kwargs) |
| 227 | for i in range(self.n_workers - 1) |
| 228 | ] |
| 229 | |
| 230 | rank0_future = self.thread_pool.submit(task, *args, **kwargs) |
| 231 | return [rank0_future] + worker_futures |
| 232 | |
| 233 | def submit_sync(self, task: Callable[..., T], *args, **kwargs) -> List[T]: |
| 234 | futures = self.submit(task, *args, **kwargs) |
| 235 | return [future.result() for future in futures] |
| 236 | |
| 237 | def shutdown(self, wait=True): |
| 238 | # Only shutdown the mpi_pool if this instance created it |
| 239 | # For shared global mpi_pool, we don't shut it down |
| 240 | if self.mpi_pool is not None and self.owns_mpi_pool: |
no outgoing calls