mirror of
https://github.com/NVIDIA/TensorRT-LLM.git
synced 2026-01-23 20:23:08 +08:00
13 lines
428 B
Python
13 lines
428 B
Python
#!/usr/bin/env python3
|
|
import logging
|
|
|
|
from mpi4py.futures import MPICommExecutor
|
|
from mpi4py.MPI import COMM_WORLD
|
|
|
|
# For multi-node MPI, the worker nodes should launch MPICommExecutor to accept tasks sent from rank0
|
|
with MPICommExecutor(COMM_WORLD) as executor:
|
|
if executor is not None:
|
|
raise RuntimeError(f"rank{COMM_WORLD.rank} should not have executor")
|
|
|
|
logging.warning(f"worker rank{COMM_WORLD.rank} quited")
|