| 78 | |
| 79 | |
| 80 | class ProcessPoolExecutorSession(MpiSession): |
| 81 | # This process pool is introduced for better recoverable exceptions handling. |
| 82 | # It replaces MpiPoolExecutor for single-gpu case. |
| 83 | |
| 84 | def __init__(self, n_workers: int, **kwargs): |
| 85 | self.n_workers = n_workers |
| 86 | self.mpi_pool = ProcessPoolExecutor(max_workers=self.n_workers, |
| 87 | **kwargs) |
| 88 | |
| 89 | def submit(self, task: Callable, *args, |
| 90 | **kwargs) -> List[concurrent.futures.Future]: |
| 91 | return [ |
| 92 | self.mpi_pool.submit(task, *args, **kwargs) |
| 93 | for i in range(self.n_workers) |
| 94 | ] |
| 95 | |
| 96 | def submit_sync(self, task: Callable, *args, **kwargs) -> List[Any]: |
| 97 | futures = [ |
| 98 | self.mpi_pool.submit(task, *args, **kwargs) |
| 99 | for i in range(self.n_workers) |
| 100 | ] |
| 101 | return [future.result() for future in futures] |
| 102 | |
| 103 | def shutdown(self): |
| 104 | self.mpi_pool.shutdown(wait=True) |
| 105 | |
| 106 | |
| 107 | class ErrorResponse(NamedTuple): |