mirror of
https://github.com/NVIDIA/TensorRT-LLM.git
synced 2026-01-14 06:27:45 +08:00
75057cd036
* Update TensorRT-LLM --------- Co-authored-by: Puneesh Khanna <puneesh.khanna@tii.ae> Co-authored-by: Ethan Zhang <26497102+ethnzhng@users.noreply.github.com>
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")
|