TensorRT-LLMs/tests/unittest/llmapi/_run_mpi_comm_task.py
Kaiyu Xie 2631f21089
Update (#2978)
Signed-off-by: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com>
2025-03-23 16:39:35 +08:00

21 lines
469 B
Python

import os
import subprocess # nosec B404
import time
from tensorrt_llm.llmapi.mpi_session import RemoteMpiCommSessionClient
from tensorrt_llm.llmapi.utils import print_colored
def main():
tasks = [0]
client = RemoteMpiCommSessionClient(
os.environ['TLLM_SPAWN_PROXY_PROCESS_IPC_ADDR'])
for task in tasks:
client.submit(print_colored, f"{task}\n", "green")
time.sleep(10)
client.shutdown()
if __name__ == "__main__":
main()