fix rpc unique addr

Signed-off-by: Superjomn <328693+Superjomn@users.noreply.github.com>

Author: Superjomn

tensorrt_llm/executor/rpc/__init__.py

@@ -1,10 +1,11 @@
 from .rpc_client import RPCClient
 from .rpc_common import (RPCCancelled, RPCError, RPCParams, RPCRequest,
-                         RPCResponse, RPCStreamingError, RPCTimeout)
+                         RPCResponse, RPCStreamingError, RPCTimeout,
+                         get_unique_ipc_addr)
 from .rpc_server import RPCServer, Server
 
 __all__ = [
     "RPCClient", "RPCServer", "Server", "RPCError", "RPCTimeout",
     "RPCCancelled", "RPCStreamingError", "RPCRequest", "RPCResponse",
-    "RPCParams"
+    "RPCParams", "get_unique_ipc_addr"
 ]

tensorrt_llm/executor/rpc/rpc_common.py

@@ -1,8 +1,20 @@
+import os
 import time
 from dataclasses import dataclass
 from typing import Any, Literal, NamedTuple, Optional
 
 
+def get_unique_ipc_addr() -> str:
+    """Generate a unique IPC address to avoid conflicts between tests.
+
+    Returns:
+        A unique IPC address string in the format: ipc:///tmp/rpc_test_{pid}_{timestamp}
+    """
+    pid = os.getpid()
+    timestamp = int(time.time() * 1000000)  # microseconds for better uniqueness
+    return f"ipc:///tmp/rpc_test_{pid}_{timestamp}"
+
+
 class RPCParams(NamedTuple):
     """ Parameters for RPC calls. """
 

tensorrt_llm/executor/rpc_proxy.py

@@ -1,7 +1,6 @@
 import asyncio
 import atexit
 import json
-import os
 import threading
 from typing import Optional
 
@@ -16,6 +15,7 @@
 from .request import GenerationRequest
 from .result import GenerationResult
 from .rpc import RPCClient
+from .rpc.rpc_common import get_unique_ipc_addr
 from .rpc_worker import RpcWorker
 from .utils import (ErrorResponse, create_mpi_comm_session,
                     get_spawn_proxy_process_env, is_llm_response)
@@ -45,7 +45,7 @@ def __init__(
             kv_connector_config: the kv cache connector config
         """
         GenerationExecutorRpcProxy.INSTANCE_COUNTER += 1
-        self.rpc_addr = self.gen_uniq_rpc_addr()
+        self.rpc_addr = get_unique_ipc_addr()
         self.rpc_client = RPCClient(self.rpc_addr)
 
         postproc_worker_config = postproc_worker_config or PostprocWorkerConfig(
@@ -368,8 +368,3 @@ def _create_mpi_session(self, model_world_size: int,
         else:
             print_colored_debug('using external mpi session ...\n', "yellow")
             self.mpi_session = mpi_session
-
-    @staticmethod
-    def gen_uniq_rpc_addr() -> str:
-        process_id = os.getpid()
-        return f"ipc:///tmp/rpc-proxy-{process_id}-{GenerationExecutorRpcProxy.INSTANCE_COUNTER}"

tests/unittest/executor/test_rpc.py

@@ -5,6 +5,7 @@
 
 from tensorrt_llm.executor.rpc import (RPCCancelled, RPCClient, RPCError,
                                        RPCServer, RPCStreamingError, RPCTimeout)
+from tensorrt_llm.executor.rpc.rpc_common import get_unique_ipc_addr
 
 
 class RpcServerWrapper(RPCServer):
@@ -31,7 +32,8 @@ class App:
             def hello(self):
                 print("hello")
 
-        with RpcServerWrapper(App(), addr="ipc:///tmp/rpc_test") as server:
+        addr = get_unique_ipc_addr()
+        with RpcServerWrapper(App(), addr=addr) as server:
             pass
 
     def test_remote_call_without_arg(self):
@@ -42,8 +44,9 @@ def hello(self):
                 print("hello")
                 return "world"
 
-        with RpcServerWrapper(App(), addr="ipc:///tmp/rpc_test") as server:
-            with RPCClient("ipc:///tmp/rpc_test") as client:
+        addr = get_unique_ipc_addr()
+        with RpcServerWrapper(App(), addr=addr) as server:
+            with RPCClient(addr) as client:
                 ret = client.hello().remote()  # sync call
                 assert ret == "world"
 
@@ -55,8 +58,9 @@ def hello(self, name: str, location: str):
                 print("hello")
                 return f"hello {name} from {location}"
 
-        with RpcServerWrapper(App(), addr="ipc:///tmp/rpc_test") as server:
-            with RPCClient("ipc:///tmp/rpc_test") as client:
+        addr = get_unique_ipc_addr()
+        with RpcServerWrapper(App(), addr=addr) as server:
+            with RPCClient(addr) as client:
                 ret = client.hello("app", "Marvel").remote()
                 assert ret == "hello app from Marvel"
 
@@ -68,8 +72,9 @@ def hello(self, name: str, location: str):
                 print("hello")
                 return f"hello {name} from {location}"
 
-        with RpcServerWrapper(App(), addr="ipc:///tmp/rpc_test") as server:
-            with RPCClient("ipc:///tmp/rpc_test") as client:
+        addr = get_unique_ipc_addr()
+        with RpcServerWrapper(App(), addr=addr) as server:
+            with RPCClient(addr) as client:
                 ret = client.hello(name="app", location="Marvel").remote()
                 assert ret == "hello app from Marvel"
 
@@ -81,8 +86,9 @@ def hello(self, name: str, location: str):
                 print("hello")
                 return f"hello {name} from {location}"
 
-        with RpcServerWrapper(App(), addr="ipc:///tmp/rpc_test") as server:
-            with RPCClient("ipc:///tmp/rpc_test") as client:
+        addr = get_unique_ipc_addr()
+        with RpcServerWrapper(App(), addr=addr) as server:
+            with RPCClient(addr) as client:
                 ret = client.hello(name="app", location="Marvel").remote()
                 assert ret == "hello app from Marvel"
 
@@ -91,8 +97,9 @@ def test_rpc_server_address(self):
         class App:
             pass
 
-        with RpcServerWrapper(App(), addr="ipc:///tmp/rpc_test") as server:
-            assert server.address == "ipc:///tmp/rpc_test"
+        addr = get_unique_ipc_addr()
+        with RpcServerWrapper(App(), addr=addr) as server:
+            assert server.address == addr
 
     def test_rpc_with_error(self):
 
@@ -101,9 +108,9 @@ class App:
             def hello(self):
                 raise ValueError("hello")
 
-        with RpcServerWrapper(App(),
-                              addr="ipc:///tmp/rpc_test_error") as server:
-            with RPCClient("ipc:///tmp/rpc_test_error") as client:
+        addr = get_unique_ipc_addr()
+        with RpcServerWrapper(App(), addr=addr) as server:
+            with RPCClient(addr) as client:
                 with pytest.raises(RPCError):
                     client.hello().remote()
 
@@ -123,9 +130,9 @@ def send_task(self) -> None:
             def get_task_submitted(self) -> bool:
                 return self.task_submitted
 
-        with RpcServerWrapper(App(),
-                              addr="ipc:///tmp/rpc_test_no_wait") as server:
-            with RPCClient("ipc:///tmp/rpc_test_no_wait") as client:
+        addr = get_unique_ipc_addr()
+        with RpcServerWrapper(App(), addr=addr) as server:
+            with RPCClient(addr) as client:
                 client.send_task().remote(need_response=False)
                 time.sleep(
                     0.1
@@ -152,11 +159,12 @@ def divide_by_zero(self):
             def custom_exception(self):
                 raise TestRpcError.CustomError("Custom error occurred")
 
+        addr = get_unique_ipc_addr()
         with RPCServer(App()) as server:
-            server.bind("ipc:///tmp/rpc_test_error")
+            server.bind(addr)
             server.start()
             time.sleep(0.1)
-            with RPCClient("ipc:///tmp/rpc_test_error") as client:
+            with RPCClient(addr) as client:
                 # Test ValueError handling
                 with pytest.raises(RPCError) as exc_info:
                     client.hello().remote()
@@ -196,7 +204,7 @@ def task(self):
                 time.sleep(10)
                 return True
 
-        addr = "ipc:///tmp/rpc_test_cancelled"
+        addr = get_unique_ipc_addr()
 
         server = RPCServer(
             App(),
@@ -218,7 +226,6 @@ def task(self):
 
         client.close()
 
-    @pytest.mark.skip(reason="https://nvbugs/5579234")
     def test_timeout_error(self):
         """Test that requests that exceed timeout are handled with proper error."""
 
@@ -229,13 +236,12 @@ def slow_method(self):
                 time.sleep(2.0)
                 return "completed"
 
-        with RpcServerWrapper(App(),
-                              addr="ipc:///tmp/rpc_test_timeout") as server:
+        addr = get_unique_ipc_addr()
+        with RpcServerWrapper(App(), addr=addr) as server:
             time.sleep(0.1)
 
             # Create client with short timeout
-            with RPCClient("ipc:///tmp/rpc_test_timeout",
-                           timeout=0.5) as client:
+            with RPCClient(addr, timeout=0.5) as client:
                 with pytest.raises(RPCError) as exc_info:
                     client.slow_method().remote(timeout=0.5)
 
@@ -252,11 +258,11 @@ class App:
             def existing_method(self):
                 return "exists"
 
-        with RpcServerWrapper(App(),
-                              addr="ipc:///tmp/rpc_test_not_found") as server:
+        addr = get_unique_ipc_addr()
+        with RpcServerWrapper(App(), addr=addr) as server:
             time.sleep(0.1)
 
-            with RPCClient("ipc:///tmp/rpc_test_not_found") as client:
+            with RPCClient(addr) as client:
                 with pytest.raises(RPCError) as exc_info:
                     client.non_existent_method().remote()
 
@@ -272,11 +278,12 @@ class App:
         def hello(self):
             return "world"
 
+    addr = get_unique_ipc_addr()
     with RPCServer(App()) as server:
-        server.bind("ipc:///tmp/rpc_test_shutdown")
+        server.bind(addr)
         server.start()
         time.sleep(0.1)
-        with RPCClient("ipc:///tmp/rpc_test_shutdown") as client:
+        with RPCClient(addr) as client:
             ret = client.hello().remote()
             assert ret == "world"
 
@@ -298,11 +305,12 @@ def send_task(self) -> None:
             time.sleep(0.001)
             return None
 
+    addr = get_unique_ipc_addr()
     with RPCServer(App(), num_workers=10) as server:
-        server.bind("ipc:///tmp/rpc_test_no_wait")
+        server.bind(addr)
         server.start()
         time.sleep(0.1)
-        with RPCClient("ipc:///tmp/rpc_test_no_wait") as client:
+        with RPCClient(addr) as client:
             time_start = time.time()
             for i in range(100):
                 client.send_task().remote(need_response=False)
@@ -329,7 +337,7 @@ def cal(self, n: int):
             return n * 2
 
     with RPCServer(App(), async_run_task=async_run_task) as server:
-        address = "ipc:///tmp/rpc_test" if use_ipc_addr else "tcp://127.0.0.1:*"
+        address = get_unique_ipc_addr() if use_ipc_addr else "tcp://127.0.0.1:*"
 
         server.bind(address)
         server.start()
@@ -428,10 +436,10 @@ class App:
             def quick_task(self, task_id: int):
                 return f"quick_task_{task_id}"
 
-        with RpcServerWrapper(App(),
-                              addr="ipc:///tmp/rpc_test_shutdown") as server:
+        addr = get_unique_ipc_addr()
+        with RpcServerWrapper(App(), addr=addr) as server:
             time.sleep(0.1)
-            with RPCClient("ipc:///tmp/rpc_test_shutdown") as client:
+            with RPCClient(addr) as client:
                 client.quick_task(1).remote()
 
                 # repeated shutdown should not raise an error
@@ -446,12 +454,13 @@ def foo(self, delay: int):
                 time.sleep(delay)
                 return "foo"
 
+        addr = get_unique_ipc_addr()
         server = RPCServer(App())
-        server.bind("ipc:///tmp/rpc_test_shutdown")
+        server.bind(addr)
         server.start()
 
         time.sleep(0.1)
-        with RPCClient("ipc:///tmp/rpc_test_shutdown") as client:
+        with RPCClient(addr) as client:
             # This task should be continued after server shutdown
             res = client.foo(10).remote_future(timeout=12)
 
@@ -639,21 +648,20 @@ def nested_function():
             yield nested_function
 
     def test_unpickleable_error(self):
-        with RpcServerWrapper(
-                self.App(), addr="ipc:///tmp/rpc_test_pickle_error") as server:
-            with RPCClient("ipc:///tmp/rpc_test_pickle_error") as client:
+        addr = get_unique_ipc_addr()
+        with RpcServerWrapper(self.App(), addr=addr) as server:
+            with RPCClient(addr) as client:
                 with pytest.raises(RPCError) as exc_info:
                     client.unpickleable_return().remote()
 
                 assert "Failed to pickle response" in str(exc_info.value)
 
     @pytest.mark.asyncio
     async def test_unpickleable_streaming_error(self):
-        with RpcServerWrapper(self.App(),
-                              addr="ipc:///tmp/rpc_test_pickle_error_streaming",
+        addr = get_unique_ipc_addr()
+        with RpcServerWrapper(self.App(), addr=addr,
                               async_run_task=True) as server:
-            with RPCClient(
-                    "ipc:///tmp/rpc_test_pickle_error_streaming") as client:
+            with RPCClient(addr) as client:
                 with pytest.raises(RPCStreamingError) as exc_info:
                     async for _ in client.unpickleable_streaming_return(
                     ).remote_streaming():

tests/unittest/executor/test_rpc_proxy.py

@@ -55,7 +55,6 @@ def create_proxy(self, tp_size: int):
 
         return proxy
 
-    @pytest.mark.skip(reason="https://nvbugs/5579234")
     @pytest.mark.parametrize("num_reqs", [1, 10])
     def test_tp1(self, num_reqs):
         tokenizer = TransformersTokenizer.from_pretrained(model_path)

tests/unittest/executor/test_rpc_worker.py

@@ -10,7 +10,7 @@
 
 from tensorrt_llm.executor.request import GenerationRequest
 from tensorrt_llm.executor.rpc import RPCClient
-from tensorrt_llm.executor.rpc_proxy import GenerationExecutorRpcProxy
+from tensorrt_llm.executor.rpc.rpc_common import get_unique_ipc_addr
 from tensorrt_llm.executor.rpc_worker import RpcWorker
 from tensorrt_llm.llmapi.mpi_session import MpiPoolSession
 from tensorrt_llm.sampling_params import SamplingParams
@@ -42,7 +42,7 @@ def teardown_method(self):
         self.client.close()
 
     def create_worker_pool(self):
-        addr = GenerationExecutorRpcProxy.gen_uniq_rpc_addr()
+        addr = get_unique_ipc_addr()
         mp_context = multiprocessing.get_context(
             'spawn')  # spawn for CUDA context
         pool = ProcessPoolExecutor(max_workers=1, mp_context=mp_context)
@@ -211,7 +211,7 @@ def teardown_method(self):
 
     def create_worker_session(self):
         session = MpiPoolSession(n_workers=2)
-        addr = GenerationExecutorRpcProxy.gen_uniq_rpc_addr()
+        addr = get_unique_ipc_addr()
         futures = session.submit(RpcWorker.main_task,
                                  engine=model_path,
                                  rpc_addr=addr,