(rank, world_size, num_split)
| 249 | |
| 250 | |
| 251 | def run_worker(rank, world_size, num_split): |
| 252 | os.environ['MASTER_ADDR'] = 'localhost' |
| 253 | os.environ['MASTER_PORT'] = '29500' |
| 254 | |
| 255 | # Higher timeout is added to accommodate for kernel compilation time in case of ROCm. |
| 256 | options = rpc.TensorPipeRpcBackendOptions(num_worker_threads=256, rpc_timeout=300) |
| 257 | |
| 258 | if rank == 0: |
| 259 | rpc.init_rpc( |
| 260 | "master", |
| 261 | rank=rank, |
| 262 | world_size=world_size, |
| 263 | rpc_backend_options=options |
| 264 | ) |
| 265 | run_master(num_split) |
| 266 | else: |
| 267 | rpc.init_rpc( |
| 268 | f"worker{rank}", |
| 269 | rank=rank, |
| 270 | world_size=world_size, |
| 271 | rpc_backend_options=options |
| 272 | ) |
| 273 | pass |
| 274 | |
| 275 | # block until all rpcs finish |
| 276 | rpc.shutdown() |
| 277 | |
| 278 | |
| 279 | if __name__=="__main__": |
nothing calls this directly
no test coverage detected