kanshan/TensorRT-LLMs
1
0
Fork 0
mirror of https://github.com/NVIDIA/TensorRT-LLM.git synced 2026-01-13 22:18:36 +08:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

[TRTLLM-9144][fix] enhance RPC robustness (#8711)

Browse Source 
Signed-off-by: Superjomn <328693+Superjomn@users.noreply.github.com>
Signed-off-by: Erin Ho <14718778+hchings@users.noreply.github.com>
Signed-off-by: Yan Chunwei <328693+Superjomn@users.noreply.github.com>
Co-authored-by: Erin Ho <14718778+hchings@users.noreply.github.com>
This commit is contained in:
Yan Chunwei 2025-12-02 21:37:59 +08:00 committed by GitHub
parent 21e3dc11d8
commit b86256eb54
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
25 changed files with 2562 additions and 1323 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
tensorrt_llm/_utils.py
View File

@ -524,13 +524,6 @@ def mpi_disabled() -> bool:
return os.environ.get("TLLM_DISABLE_MPI") == "1"
def ray_use_rpc() -> bool:
"""True if TLLM_RAY_USE_RPC is set to "1", False otherwise.
# TODO: deprecate this once Ray is fully moved to use RPC client/server.
"""
return os.environ.get("TLLM_RAY_USE_RPC") == "1"
def mpi_rank():
if mpi_disabled():
try:

3
tensorrt_llm/executor/executor.py
View File

@ -103,9 +103,6 @@ class GenerationExecutor(ABC):
self._iter_kv_events_result: IterationResult | None = None
self._iter_stats_result: IterationResult | None = None
def use_ray_queue(self) -> bool:
return False
@abstractmethod
def submit(self, request: GenerationRequest) -> GenerationResult:
pass

99
tensorrt_llm/executor/ipc.py
View File

@ -3,6 +3,7 @@ import hashlib
import hmac
import os
import pickle # nosec B403
import threading
import time
import traceback
from queue import Queue
@ -65,6 +66,13 @@ class ZeroMqQueue:
self.hmac_key = address[1] if address is not None else None
self.use_hmac_encryption = use_hmac_encryption
self._setup_lock = threading.Lock()
# Thread safety debugging
self._zmq_thread_id = None
self._zmq_debug_enabled = os.environ.get('TLLM_LLMAPI_ZMQ_DEBUG',
'0') != '0'
# Check HMAC key condition
if self.use_hmac_encryption and not self.is_server and self.hmac_key is None:
raise ValueError(
@ -93,18 +101,44 @@ class ZeroMqQueue:
self.address = (self.address_endpoint, self.hmac_key)
def setup_lazily(self):
# Early return if setup is already done
if self._setup_done:
return
self._setup_done = True
if not self.is_server:
with self._setup_lock:
if self._setup_done:
return
self._setup_done = True
if not self.is_server:
logger_debug(
f"Client [{self.name}] connecting to {self.address_endpoint} in {self.socket_type_str[self.socket_type]}\n",
"green")
self.socket.connect(self.address_endpoint)
self.poller = zmq.Poller()
self.poller.register(self.socket, zmq.POLLIN)
def _check_thread_safety(self):
"""Check if the current thread is the same as the thread that first used the socket."""
if not self._zmq_debug_enabled:
return
current_thread_id = threading.get_ident()
if self._zmq_thread_id is None:
# First call - capture the thread ID
self._zmq_thread_id = current_thread_id
logger_debug(
f"Client [{self.name}] connecting to {self.address_endpoint} in {self.socket_type_str[self.socket_type]}\n",
"green")
self.socket.connect(self.address_endpoint)
self.poller = zmq.Poller()
self.poller.register(self.socket, zmq.POLLIN)
f"ZMQ socket [{self.name}] initialized on thread {current_thread_id}",
"cyan")
elif self._zmq_thread_id != current_thread_id:
# Thread mismatch - raise error
raise RuntimeError(
f"ZMQ thread safety violation detected in [{self.name}]: "
f"Socket created on thread {self._zmq_thread_id}, "
f"but accessed from thread {current_thread_id}. "
f"ZMQ sockets are not thread-safe!")
def poll(self, timeout: int) -> bool:
"""
@ -112,6 +146,7 @@ class ZeroMqQueue:
timeout (int): Timeout in seconds
"""
self.setup_lazily()
self._check_thread_safety()
events = dict(self.poller.poll(timeout=timeout * 1000))
if self.socket in events and events[self.socket] == zmq.POLLIN:
@ -121,6 +156,7 @@ class ZeroMqQueue:
def put(self, obj: Any):
self.setup_lazily()
self._check_thread_safety()
with nvtx_range_debug("send", color="blue", category="IPC"):
if self.use_hmac_encryption or self.socket_type == zmq.ROUTER:
# Need manual serialization for encryption or ROUTER multipart
@ -148,6 +184,7 @@ class ZeroMqQueue:
assert retry >= 0 and retry <= 10, "Retry must be between 0 and 10, adjust the wait_time if needed"
self.setup_lazily()
self._check_thread_safety()
with nvtx_range_debug("send", color="blue", category="IPC"):
data = self._prepare_data(obj)
@ -162,6 +199,7 @@ class ZeroMqQueue:
async def put_async(self, obj: Any):
self.setup_lazily()
self._check_thread_safety()
try:
if self.use_hmac_encryption or self.socket_type == zmq.ROUTER:
# Need manual serialization for encryption or ROUTER multipart
@ -182,6 +220,7 @@ class ZeroMqQueue:
async def put_async_noblock(self, obj: Any):
self.setup_lazily()
self._check_thread_safety()
try:
if self.use_hmac_encryption:
data = pickle.dumps(obj) # nosec B301
@ -196,14 +235,55 @@ class ZeroMqQueue:
def get(self) -> Any:
self.setup_lazily()
self._check_thread_safety()
return self._recv_data()
async def get_async(self) -> Any:
self.setup_lazily()
self._check_thread_safety()
return await self._recv_data_async()
async def get_async_noblock(self, timeout: float = 0.5) -> Any:
return await asyncio.wait_for(self.get_async(), timeout)
"""Get data with timeout using polling to avoid message drops.
This method uses ZMQ's NOBLOCK flag with polling instead of asyncio.wait_for
to prevent cancelling recv operations which can cause message drops.
Args:
timeout: Timeout in seconds
Returns:
The received object
Raises:
asyncio.TimeoutError: If timeout is reached without receiving data
"""
self.setup_lazily()
self._check_thread_safety()
# Use polling loop instead of asyncio.wait_for to avoid cancelling recv
# which can cause message drops
deadline = asyncio.get_event_loop().time() + timeout
while True:
try:
# Try non-blocking receive
if self.socket_type == zmq.ROUTER:
identity, data = await self.socket.recv_multipart(
flags=zmq.NOBLOCK)
self._last_identity = identity
return self._parse_data(data)
else:
if self.use_hmac_encryption:
data = await self.socket.recv(flags=zmq.NOBLOCK)
return self._parse_data(data)
else:
return await self.socket.recv_pyobj(flags=zmq.NOBLOCK)
except zmq.Again:
# No message available yet
if asyncio.get_event_loop().time() >= deadline:
raise asyncio.TimeoutError()
# Short sleep to avoid busy-waiting
await asyncio.sleep(0.01)
def close(self):
if self.socket:
@ -311,6 +391,7 @@ class ZeroMqQueue:
raise ValueError(
"notify_with_retry is only supported for DEALER socket for now")
self._check_thread_safety()
retry_count = 0
while retry_count < max_retries:

159
tensorrt_llm/executor/ray_executor.py
View File

@ -13,7 +13,7 @@ from ray.util.placement_group import (PlacementGroup,
placement_group)
from tensorrt_llm._ray_utils import unwrap_ray_errors
from tensorrt_llm._utils import get_free_port, nvtx_range_debug, ray_use_rpc
from tensorrt_llm._utils import get_free_port, nvtx_range_debug
from tensorrt_llm.logger import logger
from ..llmapi.utils import logger_debug
@ -21,8 +21,8 @@ from .executor import GenerationExecutor
from .postproc_worker import PostprocWorkerConfig
from .ray_gpu_worker import RayGPUWorker, RayWorkerWrapper
from .request import GenerationRequest
from .result import GenerationResult, RayAsyncQueue, RaySyncQueue
from .rpc_proxy import RpcExecutorMixin
from .result import GenerationResult
from .rpc_proxy_mixin import RpcExecutorMixin
__all__ = [
"RayExecutor",
@ -76,38 +76,18 @@ class RayExecutor(RpcExecutorMixin, GenerationExecutor):
self.tp_size = tp_size
self.master_address = ray.util.get_node_ip_address()
self.master_port = get_free_port()
self.use_rpc = ray_use_rpc()
worker_kwargs = dict(**worker_kwargs,
postproc_worker_config=postproc_worker_config,
is_llm_executor=is_llm_executor)
if self.use_rpc:
self.init_rpc_executor()
worker_kwargs['rpc_addr'] = self.rpc_addr
self.create_workers(RayGPUWorker, worker_kwargs)
self.setup_engine_remote()
self.setup_mainloop(tasks=[self._fetch_responses_loop_async],
thread_name="ray_executor_main_loop")
logger.info(f"Connecting to RPC server at {self.rpc_addr}")
else:
self.response_queue = RayAsyncQueue.options(runtime_env={
"env_vars": {
"TLLM_DISABLE_MPI": "1"
}
}).remote()
self.response_sync_queue = RaySyncQueue.options(runtime_env={
"env_vars": {
"TLLM_DISABLE_MPI": "1"
}
}).remote()
self.async_response_queue_weakref = self.create_actor_weak_ref(
self.response_queue)
self.sync_response_queue_weakref = self.create_actor_weak_ref(
self.response_sync_queue)
self.response_queue.warmup.remote()
self.response_sync_queue.warmup.remote()
self.create_workers(RayGPUWorker, worker_kwargs)
self.init_rpc_executor()
worker_kwargs['rpc_addr'] = self.rpc_addr
self.create_workers(RayGPUWorker, worker_kwargs)
self.setup_engine_remote()
self.setup_mainloop(tasks=[self._fetch_responses_loop_async],
thread_name="ray_executor_main_loop")
logger.info(f"Connecting to RPC server at {self.rpc_addr}")
except Exception as e:
self.shutdown()
@ -192,37 +172,21 @@ class RayExecutor(RpcExecutorMixin, GenerationExecutor):
def submit(self, request: "GenerationRequest") -> "GenerationResult":
"""
Low-level API to the executor. Return a "future" GenerationResult
which can be waited.
Forwards the request to the workers through RPC or Ray queues depending on mode.
which can be waited. Forwards the request to the workers through RPC.
"""
request.set_id(self._get_next_client_id())
logprob_params = self._get_logprob_params(request)
if self.use_rpc:
with nvtx_range_debug("rpc_submit"):
self.rpc_client.submit(request).remote(need_response=False)
with nvtx_range_debug("rpc_submit"):
self.rpc_client.submit(request).remote(need_response=False)
result = GenerationResult(
request,
background_error_handler=self._handle_background_error,
executor=self,
disaggregated_params=request.disaggregated_params,
logprob_params=logprob_params)
self._results[request.id] = result
else:
result = GenerationResult(
request,
background_error_handler=self._handle_background_error,
executor=self,
disaggregated_params=request.disaggregated_params,
logprob_params=logprob_params)
with nvtx_range_debug("request_queue.put"):
self.call_all_ray_workers("enqueue_request",
leader_only=True,
request=request,
async_call=True,
result_wait_queue=result.queue)
result = GenerationResult(
request,
background_error_handler=self._handle_background_error,
executor=self,
disaggregated_params=request.disaggregated_params,
logprob_params=logprob_params)
self._results[request.id] = result
return result
@ -238,9 +202,6 @@ class RayExecutor(RpcExecutorMixin, GenerationExecutor):
async_call=False)
return sorted(gpu_ids)
def use_ray_queue(self) -> bool:
return not self.use_rpc
def abort_request(self, request_id: int) -> None:
self.call_all_ray_workers("abort_request",
leader_only=True,
@ -253,54 +214,40 @@ class RayExecutor(RpcExecutorMixin, GenerationExecutor):
if hasattr(self, '_shutdown_event'):
self._shutdown_event.set()
mode_str = "RPC mode" if self.use_rpc else "Ray queue mode"
logger_debug(f"Shutting down RayExecutor ({mode_str})", color="yellow")
logger_debug(f"Shutting down RayExecutor", color="yellow")
if self.use_rpc:
if hasattr(self, 'main_loop') and self.main_loop and hasattr(
self, 'main_loop_task_obj') and self.main_loop_task_obj:
logger_debug("Cancelling main loop task.", color="yellow")
try:
self.main_loop.call_soon_threadsafe(
self.main_loop_task_obj.cancel)
except Exception as e:
logger_debug(f"Error cancelling main loop task: {e}",
color="yellow")
if hasattr(self, 'main_loop') and self.main_loop and hasattr(
self, 'main_loop_task_obj') and self.main_loop_task_obj:
logger_debug("Cancelling main loop task.", color="yellow")
try:
self.main_loop.call_soon_threadsafe(
self.main_loop_task_obj.cancel)
except Exception as e:
logger_debug(f"Error cancelling main loop task: {e}",
color="yellow")
if hasattr(self, 'main_loop_thread'):
self.main_loop_thread.join()
if hasattr(self, 'main_loop_thread'):
self.main_loop_thread.join()
# Then, shutdown the workers
if hasattr(self, 'workers') and self.workers is not None:
try:
logger_debug("Shutting down RPC remote", color="yellow")
shutdown_refs = [
worker.shutdown.remote() for worker in self.workers
]
# Add timeout to prevent indefinite hanging
ray.get(shutdown_refs, timeout=30.0)
except ray.exceptions.GetTimeoutError:
logger.warning(
"Timeout waiting for workers to shutdown after 30 seconds"
)
except Exception as e:
logger.warning(f"Error shutting down RPC remote: {e}")
# Then, shutdown the workers
if hasattr(self, 'workers') and self.workers is not None:
try:
shutdown_refs = [
worker.shutdown.remote() for worker in self.workers
]
# Add timeout to prevent indefinite hanging
ray.get(shutdown_refs, timeout=30.0)
except ray.exceptions.GetTimeoutError:
logger.warning(
"Timeout waiting for workers to shutdown after 30 seconds")
except Exception as e:
logger.warning(f"Error shutting down: {e}")
if hasattr(self, 'rpc_client') and self.rpc_client is not None:
try:
self.rpc_client.close()
except Exception as e:
# Suppress errors during RPC client shutdown
# These can occur if the client is already closed or if there are
# pending operations that get cancelled during cleanup
logger_debug(
f"Suppressed error during RPC client close: {e}")
else:
# Release actors
self.response_queue = None
self.response_sync_queue = None
self.async_response_queue_weakref = None
self.sync_response_queue_weakref = None
if hasattr(self, 'rpc_client') and self.rpc_client is not None:
try:
self.rpc_client.close()
except Exception as e:
logger_debug(f"Suppressed error during RPC client close: {e}")
self.workers = None
if hasattr(self,
@ -387,9 +334,3 @@ class RayExecutor(RpcExecutorMixin, GenerationExecutor):
ret = super().enable_postprocess_parallel
assert ret == False, "Postprocess parallel is not supported in RayExecutor"
return ret
@staticmethod
def create_actor_weak_ref(actor_handle: ray.actor.ActorHandle):
state, _, _ = actor_handle._serialization_helper()
return ray.actor.ActorHandle._deserialization_helper(state,
weak_ref=True)

16
tensorrt_llm/executor/ray_gpu_worker.py
View File

@ -12,7 +12,6 @@ from tensorrt_llm._torch.utils import get_device_uuid
from tensorrt_llm._torch.virtual_memory import (materialize_with_tag,
release_with_tag,
verify_sleep_wakeup_tags)
from tensorrt_llm._utils import ray_use_rpc
from ..bindings import executor as tllm
from ..builder import Engine
@ -23,7 +22,7 @@ from .base_worker import BaseWorker
from .postproc_worker import PostprocWorkerConfig
from .request import GenerationRequest
from .result import GenerationResult
from .rpc_worker import RpcWorkerMixin
from .rpc_worker_mixin import RpcWorkerMixin
__all__ = [
"RayGPUWorker",
@ -189,14 +188,11 @@ class RayGPUWorker(RpcWorkerMixin, BaseWorker):
if self.global_rank > 1:
logger.set_rank(self.global_rank)
if ray_use_rpc():
if rpc_addr is None:
raise RuntimeError(
"RPC mode enabled but no rpc_addr provided to RayGPUWorker")
self.init_rpc_worker(self.global_rank, rpc_addr)
self.start_rpc_server()
else:
self.setup_engine()
if rpc_addr is None:
raise RuntimeError(
"RPC mode enabled but no rpc_addr provided to RayGPUWorker")
self.init_rpc_worker(self.global_rank, rpc_addr)
self.start_rpc_server()
def setup_engine(self):
if torch.distributed.is_initialized(

146
tensorrt_llm/executor/result.py
View File

@ -1,6 +1,5 @@
import asyncio
import json
import threading
import time
import weakref
from dataclasses import dataclass, field
@ -15,12 +14,11 @@ import torch.nn.functional as F
from tensorrt_llm.llmapi import tracing
try:
import ray
pass
except ModuleNotFoundError:
from tensorrt_llm import ray_stub as ray
pass
from .._ray_utils import unwrap_ray_errors
from .._utils import mpi_disabled, nvtx_range_debug, ray_use_rpc
from .._utils import nvtx_range_debug
from ..bindings import executor as tllm
from ..disaggregated_params import DisaggregatedParams
from ..llmapi.tracer import global_tracer
@ -160,104 +158,12 @@ class CompletionOutput:
return self.logprobs[self._last_logprobs_len:]
def warmup_tensorrt_llm():
import tensorrt_llm
print("Warmup by importing tensorrt_llm with version",
tensorrt_llm.version.__version__)
@ray.remote(max_concurrency=1000000, num_cpus=2)
class RayAsyncQueue:
"""Ray actor for async response handling."""
def __init__(self):
self.data = {}
self.event_map = {}
self.warmup_done = False
def register(self, key: int):
assert key not in self.event_map, f"Key {key} already registered"
self.event_map[key] = asyncio.Event()
def unregister(self, key: int):
if key in self.event_map:
del self.event_map[key]
if key in self.data:
del self.data[key]
def warmup(self):
if self.warmup_done:
return
warmup_tensorrt_llm()
self.warmup_done = True
def put_response(self, key: int, item: Any):
assert key in self.event_map, f"Key {key} not registered"
self.data[key] = item
self.event_map[key].set()
async def get_async(self, key: int):
assert key in self.event_map, f"Key {key} not registered"
await self.event_map[key].wait()
self.event_map[key].clear()
ret = self.data[key]
del self.data[key]
return ret
SYNC_QUEUE_MAX_CONCURRENCY = 2
@ray.remote(max_concurrency=SYNC_QUEUE_MAX_CONCURRENCY,
num_cpus=SYNC_QUEUE_MAX_CONCURRENCY)
class RaySyncQueue:
"""Ray actor for sync response handling."""
def __init__(self):
self.data = {}
self.event_map = {}
self.semaphore = threading.Semaphore(SYNC_QUEUE_MAX_CONCURRENCY - 1)
self.warmup_done = False
def register(self, key: int):
assert key not in self.event_map, f"Key {key} already registered"
self.event_map[key] = threading.Event()
self.event_map[key]
def unregister(self, key: int):
if key in self.event_map:
del self.event_map[key]
if key in self.data:
del self.data[key]
def warmup(self):
if self.warmup_done:
return
warmup_tensorrt_llm()
self.warmup_done = True
def put_response(self, key: int, item: Any):
self.data[key] = item
self.event_map[key].set()
def get(self, key: int):
with self.semaphore:
self.event_map[key].wait()
self.event_map[key].clear()
ret = self.data[key]
del self.data[key]
return ret
class GenerationResultBase:
''' This holds the core logic of the GenerationResult class. '''
def __init__(self,
id: int,
sampling_params: SamplingParams,
ray_queue: Optional[RayAsyncQueue] = None,
background_error_handler: Optional[Callable] = None,
postproc_params: "Optional[PostprocParams]" = None):
self.id = id
@ -275,22 +181,12 @@ class GenerationResultBase:
# torch backend will use trtllm sampler in beam search mode, but it does not support return logprobs incrementally
self.use_trtllm_sampler = sampling_params.use_beam_search and sampling_params.best_of > 1
if ray_queue is not None and not ray_use_rpc():
if has_event_loop():
self.aqueue = ray_queue
self.queue = self.aqueue
else:
self.queue = ray_queue
self.aqueue = None
with unwrap_ray_errors():
ray.get(self.queue.register.remote(id))
if has_event_loop():
self.aqueue = AsyncQueue()
self.queue = self.aqueue.sync_q
else:
if has_event_loop():
self.aqueue = AsyncQueue()
self.queue = self.aqueue.sync_q
else:
self.queue = Queue()
self.aqueue = None
self.queue = Queue()
self.aqueue = None
# In Sampling mode, the Executor runtime will return best_of sequences
# in total, which the LLM API will select the n-best sequences among
@ -557,12 +453,6 @@ class GenerationResultBase:
else:
raise ValueError(f"Unknown response type: {response}")
if self._done and mpi_disabled() and not ray_use_rpc():
assert hasattr(
self.queue, "unregister"
), "Ray path should be activated for unregistering the Ray queue."
self.queue.unregister.remote(self.id)
def record_stats(self,
output: CompletionOutput,
stats: Optional[dict[str, float]] = None) -> None:
@ -787,15 +677,9 @@ class GenerationResult(GenerationResultBase):
disaggregated_params: Optional[DisaggregatedParams] = None,
logprob_params: Optional[LogprobParams] = None,
) -> None:
use_async_queue = has_event_loop()
shared_queue = None
if executor and executor.use_ray_queue() and not ray_use_rpc():
shared_queue = executor.async_response_queue_weakref if use_async_queue else executor.sync_response_queue_weakref
super().__init__(
generation_request.id,
generation_request.sampling_params,
shared_queue,
background_error_handler,
postproc_params=generation_request.postproc_params,
)
@ -854,22 +738,12 @@ class GenerationResult(GenerationResultBase):
return response
def _result_step(self, timeout: Optional[float] = None):
if mpi_disabled() and not ray_use_rpc():
with unwrap_ray_errors():
response = ray.get(self.queue.get.remote(self.request_id))
response = self._handle_ray_response(response)
else:
response = self.queue.get()
response = self.queue.get()
self._handle_response(response)
async def _aresult_step(self):
assert self.aqueue is not None, "The asyncio event loop was not present during initialization, so async operations are not available."
if mpi_disabled() and not ray_use_rpc():
response = await self.aqueue.get_async.remote(self.request_id)
response = self._handle_ray_response(response)
else:
response = await self.aqueue.get()
response = await self.aqueue.get()
global_tracer().log_instant("result_step.get")
self._handle_response(response)

359
tensorrt_llm/executor/rpc/rpc_client.py
View File

@ -1,8 +1,10 @@
import asyncio
import concurrent.futures
import os
import threading
import time
import uuid
from datetime import datetime
from typing import Any, AsyncIterator, Dict, Optional
import zmq
@ -94,14 +96,26 @@ class RPCClient:
'''
self._address = address
self._timeout = timeout
# Check if PAIR mode is enabled via environment variable
use_pair_mode = os.environ.get('TLLM_LLMAPI_ZMQ_PAIR', '0') != '0'
socket_type = zmq.PAIR if use_pair_mode else zmq.DEALER
if use_pair_mode:
logger_debug(
"[client] Using zmq.PAIR socket type for RPC communication")
self._client_socket = ZeroMqQueue(address=(address, hmac_key),
is_server=False,
is_async=True,
use_hmac_encryption=False,
socket_type=zmq.DEALER)
socket_type=socket_type,
name="rpc_client")
self._pending_futures = {}
# map request_id to the queue for streaming responses
self._streaming_queues: Dict[str, AsyncQueue] = {}
self._streaming_queues_lock = threading.RLock(
) # Protect cross-thread access
self._reader_task = None
self._executor = concurrent.futures.ThreadPoolExecutor(
max_workers=num_workers, thread_name_prefix="rpc_client_worker")
@ -111,8 +125,25 @@ class RPCClient:
self._loop = None
self._loop_thread = None
self._reader_asyncio_task = None # Track the asyncio task for proper cancellation
self._loop_lock = threading.Lock(
) # Lock to protect event loop initialization
logger_debug(f"RPC Client initialized. Connected to {self._address}")
# Eagerly create the background event loop so that all subsequent
# RPC calls (sync or streaming) can assume it exists. This removes a
# race between the first streaming call (which previously created the
# loop lazily) and immediate fire-and-forget calls like `submit()`.
self._ensure_event_loop()
# Force ZeroMqQueue client connection during initialization
# This ensures the socket is connected before any RPC operations
self._client_socket.setup_lazily()
# Start the response reader eagerly to avoid race conditions with streaming
# This ensures the reader is processing responses before any RPC calls
self._start_response_reader_eagerly()
logger_debug(
f"[client] RPC Client initialized. Connected to {self._address}")
def shutdown_server(self):
"""Shutdown the server."""
@ -130,14 +161,19 @@ class RPCClient:
return
self._closed = True
logger_debug("RPC Client closing")
logger_debug("[client] RPC Client closing")
# Cancel the reader task first to avoid socket closure errors
# Notify any active streaming consumers so they don't hang waiting for
# further data. This must be done *before* shutting down the event
# loop/executor because they may depend on the loop to complete.
self._broadcast_streaming_error(RPCCancelled("RPC client closed"))
# 1. Cancel the reader task
if self._reader_task and not self._reader_task.done():
if self._loop and self._loop.is_running(
) and self._reader_asyncio_task:
try:
# Cancel the asyncio task in its event loop
async def cancel_reader_task():
if self._reader_asyncio_task and not self._reader_asyncio_task.done(
):
@ -145,35 +181,36 @@ class RPCClient:
try:
await self._reader_asyncio_task
except asyncio.CancelledError:
pass # Expected
pass
cancel_future = asyncio.run_coroutine_threadsafe(
cancel_reader_task(), self._loop)
cancel_future.result(timeout=2.0)
logger_debug("Reader task cancelled successfully")
logger_debug("[client] Reader task cancelled successfully")
except concurrent.futures.TimeoutError:
logger.warning("Reader task did not exit gracefully")
except Exception as e:
logger_debug(f"Reader task cleanup: {e}")
self._reader_task = None
self._reader_asyncio_task = None
logger_debug(f"[client] Reader task cleanup: {e}")
# Now close the socket after reader has stopped
if self._client_socket:
self._client_socket.close()
self._client_socket = None
# Stop the event loop
# 2. Stop the event loop
if self._loop and self._loop.is_running():
self._loop.call_soon_threadsafe(self._loop.stop)
# 3. Join the event loop thread
if self._loop_thread:
self._loop_thread.join(timeout=2.0)
self._loop_thread = None
if self._loop_thread.is_alive():
logger.warning("Event loop thread did not exit gracefully")
# 4. Shutdown the executor
if self._executor:
self._executor.shutdown(wait=True)
logger_debug("RPC Client closed")
# 5. Close the socket
if self._client_socket:
self._client_socket.close()
logger_debug("[client] RPC Client closed")
def _handle_streaming_response(self, response: RPCResponse):
"""Handle a streaming response by putting it in the appropriate queue.
@ -185,19 +222,25 @@ class RPCClient:
'start', 'data', 'end', 'error'
], f"Invalid stream status: {response.stream_status}"
queue = self._streaming_queues.get(response.request_id)
with self._streaming_queues_lock:
queue = self._streaming_queues.get(response.request_id)
if queue:
logger_debug(
f"[client] [{datetime.now().isoformat()}] Found streaming queue for response: request_id={response.request_id}, "
f"status={response.stream_status}")
if queue:
# put to the sync queue, as the current event loop is
# different from the one in call_async or call_streaming
assert isinstance(queue, AsyncQueue)
if enable_llmapi_debug() or logger.level == 'debug':
logger_debug(
f"RPC Client putting response to AsyncQueue: status={response.stream_status}, request_id={response.request_id}"
f"[client] RPC Client putting response to AsyncQueue: status={response.stream_status}, request_id={response.request_id}"
)
queue.sync_q.put(response)
# Clean up if stream ended
if response.stream_status in ['end', 'error']:
self._streaming_queues.pop(response.request_id, None)
with self._streaming_queues_lock:
self._streaming_queues.pop(response.request_id, None)
def _handle_regular_response(self, response: RPCResponse):
"""Handle a regular (non-streaming) response by setting the future result.
@ -223,24 +266,25 @@ class RPCClient:
# This is expected in high-load scenarios, just log and continue
if enable_llmapi_debug() or logger.level == 'debug':
logger_debug(
f"Future already done for request_id: {response.request_id}, skipping"
f"[client] Future already done for request_id: {response.request_id}, skipping"
)
if enable_llmapi_debug() or logger.level == 'debug':
if response.error is None:
logger_debug(
f"Setting result for request_id: {response.request_id}"
f"[client] Setting result for request_id: {response.request_id}"
)
else:
logger_debug(
f"Setting exception for request_id: {response.request_id}, error: {response.error}"
f"[client] Setting exception for request_id: {response.request_id}, error: {response.error}"
)
target_loop.call_soon_threadsafe(safe_set_result)
else:
if enable_llmapi_debug() or logger.level == 'debug':
logger_debug(
f"No future found for request_id: {response.request_id}")
f"[client] No future found for request_id: {response.request_id}"
)
self._pending_futures.pop(response.request_id, None)
@ -267,9 +311,8 @@ class RPCClient:
target_loop.call_soon_threadsafe(safe_set_exception)
# Also signal error to streaming queues
for queue in self._streaming_queues.values():
await queue.put(RPCResponse("", None, exception, False, 0, 'error'))
# Propagate to streaming queues via common helper
self._broadcast_streaming_error(exception)
async def _wait_for_response(self) -> RPCResponse:
"""Wait for a response from the socket.
@ -277,21 +320,49 @@ class RPCClient:
Returns:
RPCResponse from the server
"""
# Directly await the socket - cancellation will be handled by task cancellation
return await self._client_socket.get_async()
# Use timeout-based recv to handle cancellation gracefully
# This prevents the CancelledError from being logged as an exception
while True:
try:
# Short timeout allows periodic checks for cancellation
return await self._client_socket.get_async_noblock(timeout=2)
except asyncio.TimeoutError:
# Check if we should exit due to cancellation
if self._closed or (self._reader_asyncio_task
and self._reader_asyncio_task.cancelled()):
raise asyncio.CancelledError("Reader task cancelled")
# Otherwise continue polling
continue
async def _response_reader(self):
"""Task to read responses from the socket and set results on futures."""
logger_debug("Response reader started")
logger_debug("[client] Response reader started")
self._reader_asyncio_task = asyncio.current_task()
try:
# Add initial delay to ensure socket is fully connected
# This helps prevent race conditions during initialization
await asyncio.sleep(0.1)
logger_debug("[client] Response reader ready to process messages")
with customized_gc_thresholds(10000):
while True:
last_alive_log = time.time()
while not self._closed:
# Periodic alive logging for debugging
if time.time() - last_alive_log > 5.0:
logger_debug(
"[client] Response reader is alive and waiting for responses"
)
last_alive_log = time.time()
with nvtx_range_debug("response_reader",
color="cyan",
category="RPC"):
try:
response = await self._wait_for_response()
logger_debug(
f"[client] [{datetime.now().isoformat()}] Received response: {response}"
)
nvtx_mark_debug(
f"RPC.response.{'streaming' if response.is_streaming else 'sync'}",
@ -302,7 +373,7 @@ class RPCClient:
# This avoids holding GIL for f-string evaluation when debug is disabled
if enable_llmapi_debug() or logger.level == 'debug':
logger_debug(
f"RPC Client received response: request_id={response.request_id}, "
f"[client] [{datetime.now().isoformat()}] RPC Client received response: request_id={response.request_id}, "
f"is_streaming={response.is_streaming}, "
f"pending_futures={len(self._pending_futures)}"
)
@ -315,31 +386,58 @@ class RPCClient:
else:
self._handle_regular_response(response)
except asyncio.CancelledError:
# Re-raise cancellation to exit cleanly
raise
except Exception as e:
await self._handle_reader_exception(e)
break
# Log the error but continue reading unless it's a critical error
logger.error(f"Error processing response: {e}",
exc_info=True)
# For critical errors, propagate and exit
if isinstance(e, (ConnectionError, zmq.ZMQError)):
await self._handle_reader_exception(e)
break
# For other errors, try to continue
await asyncio.sleep(
0.1
) # Brief pause to avoid tight loop on repeated errors
except asyncio.CancelledError:
logger_debug("Response reader cancelled")
logger_debug("[client] Response reader cancelled")
finally:
logger_debug("Response reader exiting gracefully")
logger_debug("[client] Response reader exiting gracefully")
self._reader_task = None
self._reader_asyncio_task = None
def _start_response_reader_lazily(self):
if self._reader_task is None or self._reader_task.done():
# Ensure we have a persistent background loop
self._ensure_event_loop()
def _start_response_reader_eagerly(self):
"""Start the response reader immediately during initialization.
# Wrapper to track the asyncio task
async def run_reader():
self._reader_asyncio_task = asyncio.current_task()
await self._response_reader()
This ensures the reader is ready before any RPC calls are made,
preventing race conditions with streaming responses.
"""
if self._reader_task is not None and not self._reader_task.done():
return # Already running
# Start the reader task on the persistent loop
future = asyncio.run_coroutine_threadsafe(run_reader(), self._loop)
# Store the concurrent.futures.Future
self._reader_task = future
try:
if self._loop and self._loop.is_running():
future = asyncio.run_coroutine_threadsafe(
self._response_reader(), self._loop)
self._reader_task = future
# Wait a bit to ensure the reader is actually processing
time.sleep(0.2)
logger_debug(
"[client] Response reader started eagerly during initialization"
)
else:
raise RuntimeError(
"Event loop not running when trying to start response reader"
)
except Exception as e:
logger.error(f"Failed to start response reader eagerly: {e}")
self._reader_task = None
raise
async def _call_async(self, method_name, *args, **kwargs):
"""Async version of RPC call.
@ -353,26 +451,28 @@ class RPCClient:
The result of the remote method call
"""
if enable_llmapi_debug() or logger.level == 'debug':
logger_debug(f"RPC client calling method: {method_name}")
logger_debug(
f"[client] [{datetime.now().isoformat()}] RPC client calling method: {method_name}"
)
nvtx_mark_debug(f"RPC.async.{method_name}",
color="yellow",
category="RPC")
if self._server_stopped:
raise RPCCancelled("Server is shutting down, request cancelled")
self._start_response_reader_lazily()
rpc_params = kwargs.pop("__rpc_params", RPCParams())
need_response = rpc_params.need_response
timeout = rpc_params.timeout if rpc_params.timeout is not None else self._timeout
request_id = uuid.uuid4().hex
request = RPCRequest(request_id,
method_name,
args,
kwargs,
need_response,
method_name=method_name,
args=args,
kwargs=kwargs,
need_response=need_response,
timeout=timeout)
await self._client_socket.put_async(request)
logger_debug(f"[client] RPC Client sent request: {request}")
if not need_response:
return None
@ -403,44 +503,35 @@ class RPCClient:
self._pending_futures.pop(request_id, None)
def _ensure_event_loop(self):
"""Ensure we have a running event loop in a background thread."""
if self._loop is None or not self._loop.is_running():
self._loop = asyncio.new_event_loop()
"""Create and start the background event loop.
# TODO: WAR. Remove after RPC shutdown is fixed.
def custom_exception_handler(loop, context):
exception = context.get('exception')
message = context.get('message', '')
This is called once during initialization to create the dedicated
event loop for all socket I/O operations.
"""
if self._loop is not None:
return # Already created
if isinstance(exception,
asyncio.CancelledError) or "pending" in message:
logger.debug(f"Suppressed error during shutdown: {message}")
return
self._loop = asyncio.new_event_loop()
loop.default_exception_handler(context)
def run_loop():
asyncio.set_event_loop(self._loop)
self._loop.run_forever()
self._loop.set_exception_handler(custom_exception_handler)
self._loop_thread = threading.Thread(target=run_loop,
daemon=True,
name="rpc_client_loop")
self._loop_thread.start()
def run_loop():
asyncio.set_event_loop(self._loop)
self._loop.run_forever()
self._loop_thread = threading.Thread(target=run_loop,
daemon=True,
name="rpc_client_loop")
self._loop_thread.start()
# Give the loop a moment to start
time.sleep(0.1)
# Wait briefly to ensure the loop is running before returning
time.sleep(0.2)
def _call_sync(self, method_name, *args, **kwargs):
"""Synchronous version of RPC call."""
if enable_llmapi_debug() or logger.level == 'debug':
logger_debug(f"RPC Client calling method: {method_name}")
logger_debug(f"[client] RPC Client calling method: {method_name}")
nvtx_mark_debug(f"RPC.sync.{method_name}",
color="green",
category="RPC")
self._ensure_event_loop()
future = asyncio.run_coroutine_threadsafe(
self._call_async(method_name, *args, **kwargs), self._loop)
result = future.result()
@ -462,7 +553,6 @@ class RPCClient:
nvtx_mark_debug(f"RPC.future.{name}", color="blue", category="RPC")
def _async_to_sync():
self._ensure_event_loop()
future = asyncio.run_coroutine_threadsafe(
self._call_async(name, *args, **kwargs), self._loop)
return future.result()
@ -474,20 +564,14 @@ class RPCClient:
"""
Call a remote async generator method and get streaming results.
Args:
name: Method name to call
*args: Positional arguments
**kwargs: Keyword arguments
Yields:
Results from the remote async generator
Implementation note: The outgoing request is sent on the RPCClient’s
private event-loop to obey the single-loop rule. The returned items
are yielded in the caller’s loop via AsyncQueue, which is thread-safe.
"""
nvtx_mark_debug(f"RPC.streaming.{name}", color="red", category="RPC")
if self._server_stopped:
raise RPCCancelled("Server is shutting down, request cancelled")
self._start_response_reader_lazily()
rpc_params = kwargs.pop("__rpc_params", RPCParams())
timeout = rpc_params.timeout if rpc_params.timeout is not None else self._timeout
@ -495,21 +579,47 @@ class RPCClient:
# Use AsyncQueue to ensure proper cross-thread communication
queue = AsyncQueue()
# Recreate sync_q with the current running loop for proper cross-thread communication
# This ensures the background _response_reader thread can properly notify this event loop
queue._sync_q = _SyncQueue(queue, asyncio.get_running_loop())
self._streaming_queues[request_id] = queue
# Register queue with lock to ensure thread-safe access
with self._streaming_queues_lock:
self._streaming_queues[request_id] = queue
#logger_debug(f"[{datetime.now().isoformat()}] Registered streaming queue for request_id={request_id}")
# Build the RPCRequest object here – it's pickle-able and small – but
# *do not* touch the ZeroMQ socket from this (caller) event-loop.
request = RPCRequest(request_id,
method_name=name,
args=args,
kwargs=kwargs,
need_response=True,
timeout=timeout,
is_streaming=True)
# Send the request on the RPCClient's dedicated loop to guarantee that
# **all** socket I/O happens from exactly one thread / loop.
async def _send_streaming_request(req: RPCRequest):
"""Private helper executed in the client loop to put the request."""
logger_debug(
f"[client] [{datetime.now().isoformat()}] Sending streaming request: {req.method_name}, request_id={req.request_id}"
)
await self._client_socket.put_async(req)
logger_debug(
f"[client][{datetime.now().isoformat()}] Streaming request sent successfully: {req.method_name}, request_id={req.request_id}"
)
send_future = asyncio.run_coroutine_threadsafe(
_send_streaming_request(request), self._loop)
# Wait until the request is actually on the wire before entering the
# user-visible streaming loop. We wrap the concurrent.futures.Future so
# we can await it in the caller's asyncio context.
await asyncio.wrap_future(send_future)
try:
# Send streaming request
request = RPCRequest(request_id,
name,
args,
kwargs,
need_response=True,
timeout=timeout,
is_streaming=True)
await self._client_socket.put_async(request)
logger_debug(
f"[client] [{datetime.now().isoformat()}] Starting to read streaming responses for request_id={request_id}"
)
# Read streaming responses
while True:
if timeout is None:
@ -520,7 +630,7 @@ class RPCClient:
if enable_llmapi_debug() or logger.level == 'debug':
logger_debug(
f"RPC Client _call_streaming received [{response.stream_status}] response",
f"[client] [{datetime.now().isoformat()}] RPC Client _call_streaming received [{response.stream_status}] response",
color="green")
if response.stream_status == 'start':
@ -543,7 +653,39 @@ class RPCClient:
f"Streaming request '{name}' timed out after {timeout}s")
finally:
# Clean up
self._streaming_queues.pop(request_id, None)
with self._streaming_queues_lock:
self._streaming_queues.pop(request_id, None)
def _broadcast_streaming_error(self, exc: Exception):
"""Send an error response to all pending streaming queues so that
any coroutines blocked in _call_streaming can exit promptly.
Args:
exc: The exception instance to propagate downstream.
"""
# Iterate over a copy because callbacks may mutate the dict
with self._streaming_queues_lock:
streaming_items = list(self._streaming_queues.items())
for request_id, queue in streaming_items:
if not isinstance(queue, AsyncQueue):
continue
try:
# Use the underlying sync_q to ensure cross-thread delivery
queue.sync_q.put(
RPCResponse(
request_id,
result=None,
error=exc,
is_streaming=True,
chunk_index=0,
stream_status='error',
))
except Exception as e:
# Best-effort; log and continue
if enable_llmapi_debug() or logger.level == 'debug':
logger_debug(
f"[client] [{datetime.now().isoformat()}] Failed to broadcast streaming error for {request_id}: {e}"
)
def get_server_attr(self, name: str):
""" Get the attribute of the RPC server.
@ -561,7 +703,9 @@ class RPCClient:
client.method(args).remote_future()
async for x in client.method(args).remote_streaming()
"""
logger_debug(f"RPC Client getting attribute: {name}")
logger_debug(
f"[client] [{datetime.now().isoformat()}] RPC Client getting attribute: {name}"
)
def method_caller(*args, **kwargs):
return RemoteCall(self, name, *args, **kwargs)
@ -573,6 +717,3 @@ class RPCClient:
def __exit__(self, exc_type, exc_value, traceback):
self.close()
def __del__(self):
self.close()

9
tensorrt_llm/executor/rpc/rpc_common.py
View File

@ -2,7 +2,7 @@ import os
import tempfile
import time
import uuid
from dataclasses import dataclass
from dataclasses import KW_ONLY, dataclass
from typing import Any, Literal, NamedTuple, Optional
@ -66,6 +66,7 @@ class RPCStreamingError(RPCError):
@dataclass
class RPCRequest:
request_id: str
_: KW_ONLY
method_name: str
args: tuple
kwargs: dict
@ -81,10 +82,12 @@ class RPCRequest:
self.creation_timestamp = time.time()
class RPCResponse(NamedTuple):
@dataclass
class RPCResponse:
request_id: str
_: KW_ONLY
result: Any
error: Optional[RPCError] = None
is_streaming: bool = False # True if more responses coming
sequence_number: int = 0 # For ordering streaming responses
chunk_index: int = 0 # For ordering streaming responses
stream_status: Literal['start', 'data', 'end', 'error'] = 'data'

616
tensorrt_llm/executor/rpc/rpc_server.py
View File

@ -1,19 +1,19 @@
import asyncio
import inspect
import queue
import os
import threading
import time
import traceback
from concurrent.futures import ThreadPoolExecutor
from typing import Optional
from typing import Any, Callable, Dict, List, Optional
import zmq
from ...llmapi.utils import ManagedThread, logger_debug
from ...llmapi.utils import logger_debug
from ...logger import logger
from ..ipc import ZeroMqQueue
from .rpc_common import (RPCError, RPCRequest, RPCResponse, RPCStreamingError,
RPCTimeout)
from .rpc_common import (RPCCancelled, RPCError, RPCRequest, RPCResponse,
RPCStreamingError, RPCTimeout)
class RPCServer:
@ -22,11 +22,11 @@ class RPCServer:
"""
def __init__(self,
instance,
hmac_key=None,
instance: Any,
hmac_key: Optional[bytes] = None,
num_workers: int = 4,
timeout: float = 0.5,
async_run_task: bool = False):
async_run_task: bool = False) -> None:
"""
Initializes the server with an instance.
@ -37,7 +37,8 @@ class RPCServer:
timeout (int): Timeout for RPC calls.
async_run_task (bool): Whether to run the task asynchronously.
NOTE: make num_workers larger if there are some streaming tasks runs infinitely.
NOTE: make num_workers larger or the remote() and remote_future() may
be blocked by the thread pool.
"""
self._instance = instance
self._hmac_key = hmac_key
@ -46,42 +47,48 @@ class RPCServer:
self._timeout = timeout
self._client_socket = None
# set the stop event to True, and all the workers will exit
self._stop_event = threading.Event()
# Asyncio components
self._loop: Optional[asyncio.AbstractEventLoop] = None
self._main_task: Optional[asyncio.Task] = None
self._worker_tasks: List[asyncio.Task] = []
self._shutdown_event: Optional[asyncio.Event] = None
self._server_thread: Optional[threading.Thread] = None
self._stop_event: threading.Event = threading.Event(
) # for thread-safe shutdown
self._num_pending_requests = 0
self._functions = {
self._functions: Dict[str, Callable[..., Any]] = {
# Some built-in methods for RPC server
"_rpc_shutdown": lambda: self.shutdown(is_remote_call=True),
"_rpc_get_attr": lambda name: self.get_attr(name),
}
self._dispatcher_thread: Optional[ManagedThread] = None
if async_run_task:
self._executor = ThreadPoolExecutor(
max_workers=num_workers, thread_name_prefix="rpc_server_worker")
else:
self._executor = None
self._queue = None
# Automatically register the instance
self.register_instance(instance)
logger_debug(f"RPC Server initialized with {num_workers} workers.",
color="green")
logger_debug(
f"[server] RPCServer initialized with {num_workers} workers.",
color="green")
@property
def address(self) -> str:
assert self._client_socket is not None, "Client socket is not bound"
return self._client_socket.address[0]
def __enter__(self):
def __enter__(self) -> 'RPCServer':
return self
def __exit__(self, exc_type, exc_value, traceback):
def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
self.shutdown()
def bind(self, address="tcp://*:5555"):
def bind(self, address: str = "tcp://*:5555") -> None:
"""
Bind the server to the specified address.
@ -89,14 +96,24 @@ class RPCServer:
address (str): The ZMQ address to bind the client-facing socket.
"""
self._address = address
# Check if PAIR mode is enabled via environment variable
use_pair_mode = os.environ.get('TLLM_LLMAPI_ZMQ_PAIR', '0') != '0'
socket_type = zmq.PAIR if use_pair_mode else zmq.ROUTER
if use_pair_mode:
logger_debug(
"[server] Using zmq.PAIR socket type for RPC communication")
self._client_socket = ZeroMqQueue(address=(address, self._hmac_key),
is_server=True,
is_async=True,
use_hmac_encryption=False,
socket_type=zmq.ROUTER)
logger.info(f"RPC Server bound to {self._address}")
socket_type=socket_type,
name="rpc_server")
logger.info(f"RPCServer is bound to {self._address}")
def shutdown(self, is_remote_call: bool = False):
def shutdown(self, is_remote_call: bool = False) -> None:
"""Internal method to trigger server shutdown.
Args:
@ -110,105 +127,154 @@ class RPCServer:
return
logger_debug(
"RPC Server shutdown signal received. Terminating server...")
# Set the stop event to True, this will trigger the dispatcher routine and
# the worker routine to prepare for exit, like stopping accepting new requests,
# and continue to process the pending requests.
self._stop_event.set()
# The worker routine should process the pending requests
logger_debug(
f"RPC Server shutdown: {self._num_pending_requests} pending requests"
"[server] RPCServer is shutting down. Terminating server immediately..."
)
while self._num_pending_requests > 0:
time.sleep(0.01)
logger_debug(f"RPC Server shutdown finished pending requests")
# Set the stop event to True, this will trigger immediate shutdown
self._stop_event.set()
# Log pending requests that will be cancelled
logger_debug(
f"[server] RPCServer is shutting down: {self._num_pending_requests} pending requests will be cancelled"
)
# Signal asyncio shutdown event if available
if self._shutdown_event and self._loop:
self._loop.call_soon_threadsafe(self._shutdown_event.set)
if not is_remote_call:
# Block the thread until shutdown is finished
# 1. Wait for the dispatcher thread to exit, so that no new requests are accepted
logger_debug(f"RPC Server dispatcher thread joining")
if self._dispatcher_thread:
self._dispatcher_thread.join()
self._dispatcher_thread = None
logger_debug(f"RPC Server dispatcher thread joined")
# 1. Cancel the main task gracefully which will trigger proper cleanup
if self._main_task and not self._main_task.done():
self._loop.call_soon_threadsafe(self._main_task.cancel)
# 2. Wait for the executor to exit, it will wait for the pending requests to be processed
# 2. Wait for the server thread to exit (this will wait for proper cleanup)
if self._server_thread and self._server_thread.is_alive():
logger_debug(
"[server] RPCServer is waiting for server thread to exit")
self._server_thread.join()
self._server_thread = None
logger_debug("[server] RPCServer thread joined")
# 3. Shutdown the executor immediately without waiting for tasks
if self._executor:
self._executor.shutdown(wait=True)
self._executor.shutdown(wait=False)
self._executor = None
# 3. (Optionally) Close the client socket, this doesn't affect
# anything since zmq client will not timeout even if the target is not available
# 4. Close the client socket
if self._client_socket:
self._client_socket.close()
else:
# if the shutdown is called by a remote call, this method itself will
# be executed in a executor thread, so we cannot join the dispatcher thread as
# the dispatcher thread is awaiting for the shutdown result.
# be executed in a executor thread, so we cannot join the server thread
logger_debug(
f"RPC Server to shutdown: {self._num_pending_requests} pending requests"
f"[server] RPC Server shutdown initiated: {self._num_pending_requests} pending requests will be cancelled"
)
while self._num_pending_requests > 0:
time.sleep(0.01)
logger_debug(f"RPC Server shutdown finished pending requests")
logger_debug("[server] RPCServer is shutdown successfully",
color="yellow")
def register_function(self, func, name=None):
"""Exposes a single function to clients."""
def register_function(self,
func: Callable[..., Any],
name: Optional[str] = None) -> None:
"""Exposes a single function to clients.
Args:
func: The function to register.
name: The name of the function. If not provided, the name of the function will be used.
"""
fname = name or func.__name__
if fname in self._functions:
logger.warning(
f"Function '{fname}' is already registered. Overwriting.")
self._functions[fname] = func
logger_debug(f"Registered function: {fname}")
logger_debug(f"[server] Registered function: {fname}")
def register_instance(self, instance):
"""Exposes all public methods of a class instance."""
def register_instance(self, instance: Any) -> None:
"""Exposes all public methods of a class instance.
Args:
instance: The instance to register.
"""
logger_debug(
f"Registering instance of class: {instance.__class__.__name__}")
f"[server] Registering instance of class: {instance.__class__.__name__}"
)
for name in dir(instance):
if not name.startswith('_'):
attr = getattr(instance, name)
if callable(attr):
self.register_function(attr, name)
def get_attr(self, name: str):
def get_attr(self, name: str) -> Any:
""" Get the attribute of the RPC server.
This is mainly used for testing. """
Args:
name: The name of the attribute to get.
"""
return getattr(self, name)
async def _dispatcher_routine(self, stop_event: threading.Event):
assert self._client_socket is not None, "Client socket is not bound"
assert self._queue is not None, "RPC queue is not initialized"
async def _drain_pending_requests(self) -> None:
"""Drain any remaining requests from the socket and send cancellation responses."""
if self._client_socket is None:
return
# Once shutdown, the dispatcher will exit first, and the workers will
# continue to process the pending requests.
while not stop_event.is_set():
logger_debug("[server] Draining pending requests after shutdown")
drained_count = 0
# Give a short window to drain any in-flight requests
end_time = asyncio.get_event_loop().time() + 2
while asyncio.get_event_loop().time() < end_time:
try:
req: RPCRequest = await self._client_socket.get_async_noblock(
timeout=0.5)
logger_debug(f"RPC dispatcher got request: {req}")
req: RPCRequest = await asyncio.wait_for(
self._client_socket.get_async_noblock(), timeout=2)
drained_count += 1
logger_debug(f"[server] Draining request after shutdown: {req}")
# Send cancellation response
await self._send_error_response(
req,
RPCCancelled("Server is shutting down, request cancelled"))
except asyncio.TimeoutError:
await asyncio.sleep(0)
continue
# No more requests to drain
break
except Exception as e:
logger.error(f"RPC dispatcher caught an exception: {e}")
logger.error(traceback.format_exc())
continue
logger.debug(f"Error draining request: {e}")
break
await self._queue.put(req) # type: ignore
if drained_count > 0:
logger_debug(
f"[server] Drained {drained_count} requests after shutdown")
# shutdown methods depend on _num_pending_requests, so
# they should not be counted
if req.method_name not in ["_rpc_shutdown", "shutdown"]:
self._num_pending_requests += 1
logger_debug(
f"Dispatcher received request {req}, pending: {self._num_pending_requests}"
)
async def _run_server(self) -> None:
"""Main server loop that handles incoming requests directly."""
assert self._client_socket is not None, "Client socket is not bound"
logger_debug("[server] RPC Server main loop started")
# Create worker tasks
for i in range(self._num_workers):
task = asyncio.create_task(self._process_requests())
self._worker_tasks.append(task)
try:
# Wait for all worker tasks to complete
await asyncio.gather(*self._worker_tasks)
except asyncio.CancelledError:
logger_debug("[server] RPC Server main loop cancelled")
# Cancel all worker tasks
for task in self._worker_tasks:
if not task.done():
task.cancel()
# Wait for all tasks to finish cancellation
await asyncio.gather(*self._worker_tasks, return_exceptions=True)
except Exception as e:
logger.error(f"RPC Server main loop error: {e}")
logger.error(traceback.format_exc())
finally:
logger_debug("[server] RPC Server main loop exiting")
# TODO optimization: resolve the sequential scheduling for the remote calls
# Suppose tons of submit remote call block the FIFO queue, and the later get_stats remote calls may be blocked
@ -218,80 +284,130 @@ class RPCServer:
# - get_stats() - 1, remote_call -> dedicated queue -> dedicated routine/pool
# - submit() - 3 -> dedicated queue -> dedicated routine/pool
# TODO potential optimization: for submit(), batch the ad-hoc requests in an interval like 5ms, reduce the IPC count
async def _worker_routine(self, stop_event: threading.Event):
"""The routine executed by each worker thread."""
assert self._client_socket is not None, "Client socket is not bound"
assert self._queue is not None, "RPC queue is not initialized"
async def _send_error_response(self, req: RPCRequest,
error: Exception) -> None:
"""Send an error response for a request."""
if not req.need_response:
return
while (not stop_event.is_set()) or self._num_pending_requests > 0:
if req.is_streaming:
await self._client_socket.put_async(
RPCResponse(
req.request_id,
result=None,
error=error,
is_streaming=
True, # Important: mark as streaming so it gets routed correctly
stream_status='error'))
logger_debug(
f"[server] Sent error response for request {req.request_id}",
color="green")
else:
await self._client_socket.put_async(
RPCResponse(req.request_id, result=None, error=error))
logger_debug(
f"[server] Sent error response for request {req.request_id}",
color="green")
async def _handle_shutdown_request(self, req: RPCRequest) -> bool:
"""Handle a request during shutdown. Returns True if handled."""
if not self._shutdown_event.is_set():
return False
# Allow shutdown methods to proceed
if req.method_name in ["_rpc_shutdown", "shutdown"]:
return False
# Send cancellation error for all other requests
await self._send_error_response(
req, RPCCancelled("Server is shutting down, request cancelled"))
# Decrement pending count
self._num_pending_requests -= 1
return True
async def _process_requests(self) -> None:
"""Process incoming requests directly from the socket."""
assert self._client_socket is not None, "Client socket is not bound"
while not self._shutdown_event.is_set():
try:
#logger_debug(f"[server] Worker waiting for request", color="green")
# Read request directly from socket with timeout
req: RPCRequest = await asyncio.wait_for(
self._queue.get(), # type: ignore
timeout=self._timeout)
self._client_socket.get_async_noblock(), timeout=2)
logger_debug(f"[server] Worker got request: {req}",
color="green")
except asyncio.TimeoutError:
await asyncio.sleep(0)
continue
except asyncio.CancelledError:
logger_debug("[server] RPC worker cancelled")
break
except Exception as e:
if self._shutdown_event.is_set():
break
logger.error(f"RPC worker caught an exception: {e}")
logger.error(traceback.format_exc())
continue
# check if the method name is in the functions
# shutdown methods depend on _num_pending_requests, so
# they should not be counted
if req.method_name not in ["_rpc_shutdown", "shutdown"]:
self._num_pending_requests += 1
logger_debug(
f"[server] Worker received request {req}, pending: {self._num_pending_requests}"
)
# Check if we should cancel due to shutdown
if await self._handle_shutdown_request(req):
continue
# Check if the method exists
if req.method_name not in self._functions:
logger.error(
f"Method '{req.method_name}' not found in RPC server.")
self._num_pending_requests -= 1
if not req.need_response:
continue
if req.is_streaming:
await self._client_socket.put_async(
RPCResponse(
req.request_id,
None,
RPCStreamingError(
f"Method '{req.method_name}' not found in RPC server.",
traceback=traceback.format_exc()),
stream_status='error'))
else:
response = RPCResponse(
req.request_id,
None,
RPCError(
f"Method '{req.method_name}' not found in RPC server.",
traceback=traceback.format_exc()),
)
await self._client_socket.put_async(response)
error = RPCStreamingError if req.is_streaming else RPCError
await self._send_error_response(
req,
error(
f"Method '{req.method_name}' not found in RPC server.",
traceback=traceback.format_exc()))
continue
func = self._functions[req.method_name]
# Final shutdown check before processing
if await self._handle_shutdown_request(req):
continue
# Process the request
if req.is_streaming:
if inspect.isasyncgenfunction(func):
await self._process_streaming_request(req)
else:
# Non-streaming function called with streaming flag
response = RPCResponse(
req.request_id,
None,
await self._send_error_response(
req,
RPCStreamingError(
f"Method '{req.method_name}' is not a streaming function."
),
# need to redirect the error to the client's streaming queue
is_streaming=True,
stream_status='error',
)
await self._client_socket.put_async(response)
))
else:
# Process regular request
response = await self._process_request(req)
# Some tasks don't need response, e.g. submit_request or shutdown
# Send response if needed
if req.need_response and response is not None:
logger_debug(
f"RPC Server sending response for request {req}, pending: {self._num_pending_requests}"
f"[server] RPC Server sending response for request {req}, pending: {self._num_pending_requests}"
)
if await self._send_response(req, response):
logger_debug(
f"RPC Server sent response for request {req}")
f"[server] RPC Server sent response for request {req}"
)
# Only decrement if this request was counted in the first place
# Decrement pending count
if req.method_name not in ["_rpc_shutdown", "shutdown"]:
self._num_pending_requests -= 1
@ -315,7 +431,7 @@ class RPCServer:
if pending_time > 0.1: # Only log if significant pending time
method_type = "streaming " if is_streaming else ""
logger_debug(
f"RPC Server adjusted timeout for {method_type}{req.method_name}: "
f"[server] RPC Server adjusted timeout for {method_type}{req.method_name}: "
f"original={req.timeout}s, pending={pending_time:.3f}s, adjusted={adjusted_timeout:.3f}s"
)
return adjusted_timeout
@ -331,7 +447,7 @@ class RPCServer:
if inspect.iscoroutinefunction(func):
# Execute async function directly in event loop, no need to run in executor due to the GIL
logger_debug(
f"RPC Server running async task {req.method_name} in dispatcher"
f"[server] RPC Server running async task {req.method_name} in dispatcher"
)
result = await asyncio.wait_for(func(*req.args, **req.kwargs),
timeout=adjusted_timeout)
@ -343,31 +459,34 @@ class RPCServer:
return func(*req.args, **req.kwargs)
logger_debug(
f"RPC Server running async task {req.method_name} in worker"
f"[server] RPC Server running async task {req.method_name} in worker"
)
# TODO: let num worker control the pool size
result = await asyncio.wait_for(loop.run_in_executor(
self._executor, call_with_kwargs),
timeout=adjusted_timeout)
logger_debug(f"RPC Server returned result for request {req}")
response = RPCResponse(req.request_id, result)
response = RPCResponse(req.request_id, result=result)
except asyncio.TimeoutError:
response = RPCResponse(
req.request_id, None,
RPCTimeout(
req.request_id,
result=None,
error=RPCTimeout(
f"Method '{req.method_name}' timed out after {req.timeout} seconds",
traceback=traceback.format_exc()))
except Exception as e:
response = RPCResponse(
req.request_id, None,
RPCError(str(e), cause=e, traceback=traceback.format_exc()))
response = RPCResponse(req.request_id,
result=None,
error=RPCError(
str(e),
cause=e,
traceback=traceback.format_exc()))
return response
async def _process_streaming_request(self, req: RPCRequest):
async def _process_streaming_request(self, req: RPCRequest) -> None:
"""Process a streaming request by sending multiple responses."""
func = self._functions[req.method_name]
@ -375,43 +494,66 @@ class RPCServer:
await self._client_socket.put_async(
RPCResponse(
req.request_id,
None,
RPCStreamingError(
result=None,
error=RPCStreamingError(
f"Method '{req.method_name}' is not an async generator.",
traceback=traceback.format_exc()),
# need to redirect the error to the client's streaming queue
is_streaming=True,
stream_status='error'))
return
sequence_number = 0
chunk_index = 0
# Calculate adjusted timeout based on pending overhead
adjusted_timeout = self._calculate_adjusted_timeout(req,
is_streaming=True)
adjusted_timeout: float = self._calculate_adjusted_timeout(
req, is_streaming=True)
try:
logger_debug(f"RPC Server running streaming task {req.method_name}")
logger_debug(
f"[server] RPC Server running streaming task {req.method_name}")
# Send start signal
await self._client_socket.put_async(
RPCResponse(req.request_id, None, None, True, sequence_number,
'start'))
sequence_number += 1
RPCResponse(req.request_id,
result=None,
error=None,
is_streaming=True,
chunk_index=chunk_index,
stream_status='start'))
logger_debug(
f"[server] Sent start signal for request {req.request_id}",
color="green")
chunk_index += 1
# Apply timeout to the entire streaming operation if specified
if adjusted_timeout is not None and adjusted_timeout > 0:
# Create a task for the async generator with timeout
async def stream_with_timeout():
nonlocal sequence_number
nonlocal chunk_index
async for result in func(*req.args, **req.kwargs):
if result is None or result == []:
# Skip None values or empty list to save bandwidth
# TODO[Superjomn]: add a flag to control this behavior
continue
# Check if shutdown was triggered
if self._shutdown_event.is_set():
raise RPCCancelled(
"Server is shutting down, streaming cancelled")
logger_debug(
f"RPC Server got data and ready to send result {result}"
f"[server] RPC Server got data and ready to send result {result}"
)
response = RPCResponse(req.request_id, result, None,
True, sequence_number, 'data')
response = RPCResponse(req.request_id,
result=result,
error=None,
is_streaming=True,
chunk_index=chunk_index,
stream_status='data')
if not await self._send_response(req, response):
# Stop streaming after a pickle error
return
sequence_number += 1
logger_debug(
f"[server] Sent response for request {req.request_id}",
color="green")
chunk_index += 1
# Use wait_for for timeout handling
await asyncio.wait_for(stream_with_timeout(),
@ -419,35 +561,71 @@ class RPCServer:
else:
# No timeout specified, stream normally
async for result in func(*req.args, **req.kwargs):
if result is None or result == []:
continue # Skip None values or empty list
# Check if shutdown was triggered
if self._shutdown_event.is_set():
raise RPCCancelled(
"Server is shutting down, streaming cancelled")
logger_debug(
f"RPC Server got data and ready to send result {result}"
f"[server] RPC Server got data and ready to send result {result}"
)
response = RPCResponse(req.request_id, result, None, True,
sequence_number, 'data')
response = RPCResponse(req.request_id,
result=result,
error=None,
is_streaming=True,
chunk_index=chunk_index,
stream_status='data')
if not await self._send_response(req, response):
# Stop streaming after a pickle error
return
sequence_number += 1
chunk_index += 1
# Send end signal
await self._client_socket.put_async(
RPCResponse(req.request_id, None, None, True, sequence_number,
'end'))
RPCResponse(req.request_id,
result=None,
error=None,
is_streaming=True,
chunk_index=chunk_index,
stream_status='end'))
logger_debug(
f"[server] Sent end signal for request {req.request_id}",
color="green")
except RPCCancelled as e:
# Server is shutting down, send cancelled error
await self._client_socket.put_async(
RPCResponse(req.request_id,
result=None,
error=e,
is_streaming=True,
chunk_index=chunk_index,
stream_status='error'))
logger_debug(
f"[server] Sent error signal for request {req.request_id}",
color="green")
except asyncio.TimeoutError:
await self._client_socket.put_async(
RPCResponse(
req.request_id, None,
RPCTimeout(
req.request_id,
result=None,
error=RPCTimeout(
f"Streaming method '{req.method_name}' timed out",
traceback=traceback.format_exc()), True,
sequence_number, 'error'))
traceback=traceback.format_exc()),
is_streaming=True,
chunk_index=chunk_index,
stream_status='error'))
except Exception as e:
response = RPCResponse(
req.request_id, None,
RPCStreamingError(str(e), traceback=traceback.format_exc()),
True, sequence_number, 'error')
req.request_id,
result=None,
error=RPCStreamingError(str(e),
traceback=traceback.format_exc()),
is_streaming=True,
chunk_index=chunk_index,
stream_status='error')
await self._send_response(req, response)
async def _send_response(self, req: RPCRequest,
@ -455,6 +633,8 @@ class RPCServer:
"""Safely sends a response, handling pickle errors."""
try:
await self._client_socket.put_async(response)
logger_debug(f"[server] Sent response for request {req.request_id}",
color="green")
return True
except Exception as e:
logger.error(
@ -462,30 +642,35 @@ class RPCServer:
error_msg = f"Failed to pickle response: {e}"
if req.is_streaming:
error_cls = RPCStreamingError
# For streaming, we also need sequence number. The original response has it.
sequence_number = response.sequence_number if response else None
chunk_index = response.chunk_index if response else None
error_response = RPCResponse(
req.request_id,
None,
error_cls(error_msg, traceback=traceback.format_exc()),
result=None,
error=error_cls(error_msg,
traceback=traceback.format_exc()),
is_streaming=True,
sequence_number=sequence_number,
chunk_index=chunk_index,
stream_status='error')
else:
error_cls = RPCError
error_response = RPCResponse(
req.request_id, None,
error_cls(error_msg, traceback=traceback.format_exc()))
req.request_id,
result=None,
error=error_cls(error_msg,
traceback=traceback.format_exc()))
try:
await self._client_socket.put_async(error_response)
logger_debug(
f"[server] Sent error response for request {req.request_id}",
color="green")
except Exception as e_inner:
logger.error(
f"Failed to send error response for request {req.request_id}: {e_inner}"
)
return False
def start(self):
def start(self) -> None:
"""Binds sockets, starts workers, and begins proxying messages."""
if self._client_socket is None:
raise RuntimeError(
@ -495,24 +680,73 @@ class RPCServer:
self._client_socket.setup_lazily()
logger.info(f"RPC Server started and listening on {self._address}")
async def tasks():
self._queue = asyncio.Queue()
await asyncio.gather(
self._dispatcher_routine(self._stop_event), *[
self._worker_routine(self._stop_event)
for i in range(self._num_workers)
])
# Create and configure the event loop
self._loop = asyncio.new_event_loop()
def loop() -> bool:
asyncio.run(tasks())
return True # ManagedThread
self._shutdown_event = asyncio.Event()
error_queue = queue.Queue()
self._dispatcher_thread = ManagedThread(task=loop,
stop_event=self._stop_event,
name="rpc_dispatcher_thread",
error_queue=error_queue)
self._dispatcher_thread.start()
async def run_server():
"""Run the server until shutdown."""
try:
await self._run_server()
except asyncio.CancelledError:
logger_debug("[server] Server task cancelled")
except Exception as e:
logger.error(f"Server error: {e}")
logger.error(traceback.format_exc())
finally:
# Cancel all worker tasks
for task in self._worker_tasks:
if not task.done():
task.cancel()
# Wait for all tasks to complete
if self._worker_tasks:
await asyncio.gather(*self._worker_tasks,
return_exceptions=True)
# Drain any remaining requests and send cancellation responses
await self._drain_pending_requests()
logger_debug("[server] All server tasks completed")
self._main_task = self._loop.create_task(run_server())
def run_loop():
asyncio.set_event_loop(self._loop)
try:
self._loop.run_until_complete(self._main_task)
except RuntimeError as e:
# This can happen if the event loop is stopped while futures are pending
error_str = str(e)
if "Event loop stopped before Future completed" in error_str:
# This is expected during shutdown - ignore it
logger.debug(
f"[server] Expected shutdown error: {error_str}")
else:
# This is an unexpected RuntimeError - log full details
import traceback
logger.error(f"Event loop error: {error_str}")
logger.error(f"Traceback: {traceback.format_exc()}")
except Exception as e:
logger.error(f"Event loop error: {e}")
finally:
# Clean up any remaining tasks
pending = asyncio.all_tasks(self._loop)
for task in pending:
task.cancel()
if pending:
try:
self._loop.run_until_complete(
asyncio.gather(*pending, return_exceptions=True))
except RuntimeError:
# Event loop might already be closed
pass
self._loop.close()
self._server_thread = threading.Thread(target=run_loop,
name="rpc_server_thread",
daemon=True)
self._server_thread.start()
logger.info("RPC Server has started.")

283
tensorrt_llm/executor/rpc_proxy.py
View File

@ -1,277 +1,14 @@
import asyncio
import atexit
import json
import threading
from typing import Callable, List, Optional
from typing import Optional
from .._utils import nvtx_range_debug
from ..llmapi.mpi_session import MpiPoolSession, MpiSession
from ..llmapi.tracer import global_tracer
from ..llmapi.utils import AsyncQueue, _SyncQueue, logger_debug
from ..llmapi.utils import logger_debug
from ..logger import logger
from .executor import GenerationExecutor
from .postproc_worker import PostprocWorkerConfig
from .request import GenerationRequest
from .result import GenerationResult
from .rpc import RPCClient
from .rpc.rpc_common import get_unique_ipc_addr
from .rpc_proxy_mixin import RpcExecutorMixin
from .rpc_worker import RpcWorker
from .utils import (ErrorResponse, create_mpi_comm_session,
get_spawn_proxy_process_env, is_llm_response)
class RpcExecutorMixin:
"""Mixin for executors that use RPC client for hot path communication.
Provides:
- RPC client initialization
- Response handling loop
- Main loop thread management
- Shutdown logic for RPC components
The inheriting class should call init_rpc_executor() to set up RPC client.
"""
def init_rpc_executor(self):
self.rpc_addr = get_unique_ipc_addr()
self.rpc_client = RPCClient(self.rpc_addr)
self._results = {}
self._shutdown_event = threading.Event()
self.main_loop_task_obj = None
self.main_loop = None
self.main_loop_thread = None
def setup_mainloop(self,
tasks: Optional[List[Callable]] = None,
thread_name: str = "rpc_proxy_main_loop"):
"""Setup main loop thread with custom async tasks.
Args:
tasks: List of async coroutine functions to run.
thread_name: Name for the main loop thread
"""
if tasks is None:
tasks = [
self._fetch_responses_loop_async,
self._fetch_stats_loop_async,
]
# Only add kv_cache_events loop if it's enabled
if self._iter_kv_events_result:
tasks.append(self._fetch_kv_cache_events_loop_async)
async def main_loop_task():
await asyncio.gather(*[task() for task in tasks])
def _run_main_loop_task():
"""Local method to run the main loop task."""
self.main_loop = asyncio.new_event_loop()
asyncio.set_event_loop(self.main_loop)
self.main_loop_task_obj = self.main_loop.create_task(
main_loop_task())
try:
self.main_loop.run_until_complete(self.main_loop_task_obj)
except asyncio.CancelledError:
pass # Task cancellation is expected during shutdown
finally:
self.main_loop.close()
self.main_loop_thread = threading.Thread(target=_run_main_loop_task,
daemon=True,
name=thread_name)
self.main_loop_thread.start()
atexit.register(self.shutdown)
def submit(self, request: GenerationRequest) -> GenerationResult:
request.set_id(self._get_next_client_id())
logprob_params = self._get_logprob_params(request)
# submit is a fire-and-forget operation, don't need to wait for response
with nvtx_range_debug("RPCExecutor.submit",
color="green",
category="Proxy"):
self.rpc_client.submit(request).remote(need_response=False)
result = GenerationResult(
request,
background_error_handler=self._handle_background_error,
executor=self,
disaggregated_params=request.disaggregated_params,
logprob_params=logprob_params)
self._results[request.id] = result
return result
def handle_responses(self, responses: list[GenerationResult]) -> bool:
async_queues = []
event_loop = None
def process_res(res: list):
for r in res:
client_id = r.client_id
nonlocal event_loop
nonlocal async_queues
if client_id not in self._results:
logger.warning(
f"Received response for unknown client_id: {client_id}")
continue
queue = self._results[client_id].queue
if isinstance(queue, _SyncQueue):
queue.put_nowait(r)
async_queues.append(queue)
# all the loops are identical
event_loop = event_loop or queue.loop
else:
queue.put(r)
if (is_llm_response(r) and r.result.is_final) or isinstance(
r, ErrorResponse):
self._results.pop(client_id)
# Handle the case where responses might not be a list of lists
if responses and not isinstance(responses[0], list):
# If responses is a flat list, wrap it
responses = [responses]
for res in responses:
global_tracer().log_instant("RPC.get")
process_res(res)
if async_queues:
_SyncQueue.notify_many(event_loop, async_queues)
def handle_stats(self, stats):
"""Handle stats received from RPC worker and put them into the stats result queue.
Args:
stats: Statistics data from the RPC worker (can be dict, str, or list)
"""
self._handle_iteration_data(stats, self._iter_stats_result, "stats")
def handle_kv_cache_events(self, events):
"""Handle KV cache events received from RPC worker and put them into the events result queue.
Args:
events: KV cache events data from the RPC worker (can be dict, str, or list)
"""
self._handle_iteration_data(events, self._iter_kv_events_result,
"kv_cache_events")
async def _generic_fetch_loop_async(self, fetch_method_name: str,
handler_method: Callable,
method_name: str):
"""Generic method for fetching data in a loop from RPC worker.
Args:
fetch_method_name: Name of the RPC client method to call
handler_method: The handler method to call with the fetched data
method_name: Name of the method for logging
"""
try:
fetch_method = getattr(self.rpc_client, fetch_method_name)
async for data in fetch_method().remote_streaming():
if self._shutdown_event.is_set():
return
handler_method(data)
except asyncio.CancelledError:
logger.debug(f"{method_name} task cancelled")
except Exception as e:
logger.error(f"Error in {method_name}: {e}")
raise
async def _fetch_responses_loop_async(self):
await self._generic_fetch_loop_async(
fetch_method_name="fetch_responses_loop_async",
handler_method=self.handle_responses,
method_name="_fetch_responses_loop_async")
async def _fetch_stats_loop_async(self):
await self._generic_fetch_loop_async(
fetch_method_name="fetch_stats_loop_async",
handler_method=self.handle_stats,
method_name="_fetch_stats_loop_async")
async def _fetch_kv_cache_events_loop_async(self):
await self._generic_fetch_loop_async(
fetch_method_name="fetch_kv_cache_events_loop_async",
handler_method=self.handle_kv_cache_events,
method_name="_fetch_kv_cache_events_loop_async")
def _handle_iteration_data(self, data, result_singleton, data_type: str):
"""Generic method to handle iteration data received from RPC worker.
Args:
data: Data from the RPC worker (can be dict, str, or list)
result_singleton: The iteration result singleton to put data into
data_type: Type of data for logging (e.g., "stats", "kv_cache_events")
"""
# Make sure we have initialized the iteration results
self._maybe_initialize_iteration_results()
if not result_singleton:
logger.debug(
f"Skipping {data_type} handling while result_singleton=None")
return
# Get the queue from the result singleton
queue = result_singleton.queue
async_queues = []
# Clear old data if queue is full (similar to _iteration_result_task)
while queue.full():
queue.get()
try:
# Handle different types of data
if isinstance(data, str):
# Already JSON serialized
data_json = data
elif isinstance(data, list):
# Skip empty lists to avoid putting nothing in the queue
if not data:
logger.debug(
f"rpc_proxy.py: Skipping empty {data_type} list")
return
# Handle list of data (multiple iterations)
for item in data:
if isinstance(item, str):
item_json = item
else:
item_json = json.dumps(item)
if isinstance(queue, _SyncQueue):
queue.put_nowait(item_json)
async_queues.append(queue)
else:
queue.put(item_json)
if async_queues:
_SyncQueue.notify_many(queue.loop, async_queues)
return
else:
# Convert dict/other to JSON string as expected by IterationResult
data_json = json.dumps(data)
if isinstance(queue, _SyncQueue):
queue.put_nowait(data_json)
async_queues.append(queue)
else:
queue.put(data_json)
if async_queues:
_SyncQueue.notify_many(queue.loop, async_queues)
except AsyncQueue.EventLoopShutdownError:
# This happens when the event loop is already closed
logger.debug(
f"rpc_proxy.py: EventLoopShutdownError in handle_{data_type}")
except Exception as e:
logger.error(f"rpc_proxy.py: Error in handle_{data_type}: {e}")
raise e
from .utils import create_mpi_comm_session, get_spawn_proxy_process_env
class GenerationExecutorRpcProxy(RpcExecutorMixin, GenerationExecutor):
@ -350,7 +87,7 @@ class GenerationExecutorRpcProxy(RpcExecutorMixin, GenerationExecutor):
def shutdown_remote(self):
logger_debug(f"Shutting down rpc remote", color="yellow")
self.rpc_client.shutdown().remote()
self.rpc_client.shutdown().remote(need_response=False)
def abort_request(self, request_id: int) -> None:
return self.rpc_client.abort_request(request_id).remote()
@ -380,7 +117,9 @@ class GenerationExecutorRpcProxy(RpcExecutorMixin, GenerationExecutor):
# (e.g., during garbage collection in that thread)
if self.main_loop_thread and threading.current_thread(
) != self.main_loop_thread:
self.main_loop_thread.join()
self.main_loop_thread.join(timeout=2.0)
if self.main_loop_thread.is_alive():
logger.warning("Main loop thread did not exit gracefully")
# 3. shutdown the mpi session, this should wait until all the PyExecutor
# processes are shutdown
@ -403,11 +142,11 @@ class GenerationExecutorRpcProxy(RpcExecutorMixin, GenerationExecutor):
mpi_process_pre_spawned: bool = get_spawn_proxy_process_env()
if mpi_session is None:
if mpi_process_pre_spawned:
logger_debug('create comm session ...\n', "yellow")
logger_debug('[proxy] create comm session ...\n', "yellow")
self.mpi_session = create_mpi_comm_session(model_world_size)
else:
logger_debug('create pool session ...\n', "yellow")
logger_debug('[proxy] create pool session ...\n', "yellow")
self.mpi_session = MpiPoolSession(n_workers=model_world_size)
else:
logger_debug('using external mpi session ...\n', "yellow")
logger_debug('[proxy] using external mpi session ...\n', "yellow")
self.mpi_session = mpi_session

264
tensorrt_llm/executor/rpc_proxy_mixin.py Normal file
View File

@ -0,0 +1,264 @@
import asyncio
import atexit
import json
import threading
from typing import Callable, List, Optional
from .._utils import nvtx_range_debug
from ..llmapi.tracer import global_tracer
from ..llmapi.utils import AsyncQueue, _SyncQueue
from ..logger import logger
from .request import GenerationRequest
from .result import GenerationResult
from .rpc import RPCClient
from .rpc.rpc_common import get_unique_ipc_addr
from .utils import ErrorResponse, is_llm_response
class RpcExecutorMixin:
"""Mixin for executors that use RPC client for hot path communication.
Provides:
- RPC client initialization
- Response handling loop
- Main loop thread management
- Shutdown logic for RPC components
The inheriting class should call init_rpc_executor() to set up RPC client.
"""
def init_rpc_executor(self):
self.rpc_addr = get_unique_ipc_addr()
self.rpc_client = RPCClient(self.rpc_addr)
self._results = {}
self._shutdown_event = threading.Event()
self.main_loop_task_obj = None
self.main_loop = None
self.main_loop_thread = None
def setup_mainloop(
self, tasks: Optional[List[Callable]] = None, thread_name: str = "rpc_proxy_main_loop"
):
"""Setup main loop thread with custom async tasks.
Args:
tasks: List of async coroutine functions to run.
thread_name: Name for the main loop thread
"""
if tasks is None:
tasks = [
self._fetch_responses_loop_async,
self._fetch_stats_loop_async,
]
# Only add kv_cache_events loop if it's enabled
if self._iter_kv_events_result:
tasks.append(self._fetch_kv_cache_events_loop_async)
async def main_loop_task():
await asyncio.gather(*[task() for task in tasks])
def _run_main_loop_task():
"""Local method to run the main loop task."""
self.main_loop = asyncio.new_event_loop()
asyncio.set_event_loop(self.main_loop)
self.main_loop_task_obj = self.main_loop.create_task(main_loop_task())
try:
self.main_loop.run_until_complete(self.main_loop_task_obj)
except asyncio.CancelledError:
pass # Task cancellation is expected during shutdown
finally:
self.main_loop.close()
self.main_loop_thread = threading.Thread(
target=_run_main_loop_task, daemon=True, name=thread_name
)
self.main_loop_thread.start()
atexit.register(self.shutdown)
def submit(self, request: GenerationRequest) -> GenerationResult:
request.set_id(self._get_next_client_id())
logprob_params = self._get_logprob_params(request)
# submit is a fire-and-forget operation, don't need to wait for response
with nvtx_range_debug("RPCExecutor.submit", color="green", category="Proxy"):
self.rpc_client.submit(request).remote(need_response=False)
result = GenerationResult(
request,
background_error_handler=self._handle_background_error,
executor=self,
disaggregated_params=request.disaggregated_params,
logprob_params=logprob_params,
)
self._results[request.id] = result
return result
def handle_responses(self, responses: list[GenerationResult]) -> bool:
async_queues = []
event_loop = None
def process_res(res: list):
for r in res:
client_id = r.client_id
nonlocal event_loop
nonlocal async_queues
if client_id not in self._results:
logger.warning(f"Received response for unknown client_id: {client_id}")
continue
queue = self._results[client_id].queue
if isinstance(queue, _SyncQueue):
queue.put_nowait(r)
async_queues.append(queue)
# all the loops are identical
event_loop = event_loop or queue.loop
else:
queue.put(r)
if (is_llm_response(r) and r.result.is_final) or isinstance(r, ErrorResponse):
self._results.pop(client_id)
# Handle the case where responses might not be a list of lists
if responses and not isinstance(responses[0], list):
# If responses is a flat list, wrap it
responses = [responses]
for res in responses:
global_tracer().log_instant("RPC.get")
process_res(res)
if async_queues:
_SyncQueue.notify_many(event_loop, async_queues)
def handle_stats(self, stats):
"""Handle stats received from RPC worker and put them into the stats result queue.
Args:
stats: Statistics data from the RPC worker (can be dict, str, or list)
"""
self._handle_iteration_data(stats, self._iter_stats_result, "stats")
def handle_kv_cache_events(self, events):
"""Handle KV cache events received from RPC worker and put them into the events result queue.
Args:
events: KV cache events data from the RPC worker (can be dict, str, or list)
"""
self._handle_iteration_data(events, self._iter_kv_events_result, "kv_cache_events")
async def _generic_fetch_loop_async(
self, fetch_method_name: str, handler_method: Callable, method_name: str
):
"""Generic method for fetching data in a loop from RPC worker.
Args:
fetch_method_name: Name of the RPC client method to call
handler_method: The handler method to call with the fetched data
method_name: Name of the method for logging
"""
try:
fetch_method = getattr(self.rpc_client, fetch_method_name)
async for data in fetch_method().remote_streaming():
if self._shutdown_event.is_set():
return
handler_method(data)
except asyncio.CancelledError:
logger.debug(f"{method_name} task cancelled")
except Exception as e:
logger.error(f"Error in {method_name}: {e}")
raise
async def _fetch_responses_loop_async(self):
await self._generic_fetch_loop_async(
fetch_method_name="fetch_responses_loop_async",
handler_method=self.handle_responses,
method_name="_fetch_responses_loop_async",
)
async def _fetch_stats_loop_async(self):
await self._generic_fetch_loop_async(
fetch_method_name="fetch_stats_loop_async",
handler_method=self.handle_stats,
method_name="_fetch_stats_loop_async",
)
async def _fetch_kv_cache_events_loop_async(self):
await self._generic_fetch_loop_async(
fetch_method_name="fetch_kv_cache_events_loop_async",
handler_method=self.handle_kv_cache_events,
method_name="_fetch_kv_cache_events_loop_async",
)
def _handle_iteration_data(self, data, result_singleton, data_type: str):
"""Generic method to handle iteration data received from RPC worker.
Args:
data: Data from the RPC worker (can be dict, str, or list)
result_singleton: The iteration result singleton to put data into
data_type: Type of data for logging (e.g., "stats", "kv_cache_events")
"""
# Make sure we have initialized the iteration results
self._maybe_initialize_iteration_results()
if not result_singleton:
logger.debug(f"Skipping {data_type} handling while result_singleton=None")
return
# Get the queue from the result singleton
queue = result_singleton.queue
async_queues = []
# Clear old data if queue is full (similar to _iteration_result_task)
while queue.full():
queue.get()
try:
# Handle different types of data
if isinstance(data, str):
# Already JSON serialized
data_json = data
elif isinstance(data, list):
# Skip empty lists to avoid putting nothing in the queue
if not data:
logger.debug(f"rpc_proxy.py: Skipping empty {data_type} list")
return
# Handle list of data (multiple iterations)
for item in data:
if isinstance(item, str):
item_json = item
else:
item_json = json.dumps(item)
if isinstance(queue, _SyncQueue):
queue.put_nowait(item_json)
async_queues.append(queue)
else:
queue.put(item_json)
if async_queues:
_SyncQueue.notify_many(queue.loop, async_queues)
return
else:
# Convert dict/other to JSON string as expected by IterationResult
data_json = json.dumps(data)
if isinstance(queue, _SyncQueue):
queue.put_nowait(data_json)
async_queues.append(queue)
else:
queue.put(data_json)
if async_queues:
_SyncQueue.notify_many(queue.loop, async_queues)
except AsyncQueue.EventLoopShutdownError:
# This happens when the event loop is already closed
logger.debug(f"rpc_proxy.py: EventLoopShutdownError in handle_{data_type}")
except Exception as e:
logger.error(f"rpc_proxy.py: Error in handle_{data_type}: {e}")
raise e

211
tensorrt_llm/executor/rpc_worker.py
View File

@ -1,15 +1,14 @@
import asyncio
from pathlib import Path
from queue import Queue
from threading import Event
from typing import AsyncGenerator, Optional, Union
from typing import Optional, Union
import nvtx
from tensorrt_llm._utils import mpi_comm
from tensorrt_llm.llmapi.utils import enable_llm_debug, logger_debug
from .._utils import mpi_rank, nvtx_range_debug
from .._utils import mpi_rank
from ..bindings import executor as tllm
from ..builder import Engine
from ..llmapi.llm_args import BaseLlmArgs
@ -18,152 +17,8 @@ from ..logger import set_level
from ..sampling_params import BatchedLogitsProcessor
from .base_worker import BaseWorker
from .postproc_worker import PostprocWorkerConfig
from .request import GenerationRequest
from .rpc import RPCServer
class RpcWorkerMixin:
"""Mixin for workers that serve RPC requests.
Provides:
- RPC server initialization
- Response queue management
- Async response fetching methods
- Shutdown logic for RPC components
The inheriting class should call init_rpc_worker() in its __init__.
"""
# Number of RPC server workers
NUM_WORKERS = 6
def init_rpc_worker(self, rank: int, rpc_addr: Optional[str]):
if rpc_addr is None:
raise RuntimeError(
"RPC mode enabled but no rpc_addr provided to worker")
self.rank = rank
self.shutdown_event = Event()
self._response_queue = Queue()
self.set_result_queue(self._response_queue)
self.rpc_server = None
self.rpc_addr = rpc_addr
def start_rpc_server(self):
if self.rank == 0:
self.rpc_server = RPCServer(self,
num_workers=RpcWorkerMixin.NUM_WORKERS)
self.rpc_server.bind(self.rpc_addr)
self.rpc_server.start()
def submit(self, request: GenerationRequest):
""" Submits a request to the worker. """
with nvtx_range_debug("RpcWorker.submit",
color="blue",
category="Worker"):
super().submit(request)
def fetch_responses(self, timeout: Optional[float] = None) -> list:
"""Fetch responses from the response queue (blocking)."""
logger_debug(f"RpcWorker {self.rank} is fetching responses",
color="yellow")
with nvtx_range_debug("RpcWorker.fetch_responses",
color="orange",
category="Worker"):
# NOTE: This is a blocking call, it will wait for the responses to be available.
responses = super().await_responses(timeout)
self._await_response_helper.responses_handler(responses)
qsize = self._response_queue.qsize()
logger_debug(f"RpcWorker returning {qsize} responses", color="yellow")
all_responses = []
for _ in range(qsize):
# The queue contains batches of responses, so extend the list
all_responses.extend(self._response_queue.get())
return all_responses
async def fetch_responses_async(self,
timeout: Optional[float] = None) -> list:
"""Async version of fetch_responses using asyncio.to_thread."""
# A really async version of fetch_responses
logger_debug(f"RpcWorker {self.rank} is fetching responses async",
color="yellow")
# First, await any pending responses without blocking the event loop
responses = await asyncio.to_thread(self.fetch_responses,
timeout=timeout)
return responses
async def fetch_responses_loop_async(self) -> AsyncGenerator[list, None]:
while not self.shutdown_event.is_set():
responses = await self.fetch_responses_async()
if responses: # Only yield if there are actual responses
logger_debug(
f"RpcWorker {self.rank} is yielding responses: {responses}",
color="yellow")
yield responses # batching the responses to opt IPC performance
else:
# Small delay to prevent busy waiting when no responses
await asyncio.sleep(0)
logger_debug(
f"RpcWorker {self.rank} quitting fetch_responses_loop_async",
color="yellow")
async def fetch_stats_async(self, timeout: Optional[float] = None) -> list:
"""Async version of fetch_stats using asyncio.to_thread."""
return await asyncio.to_thread(self.fetch_stats)
async def fetch_kv_cache_events_async(self,
timeout: Optional[float] = None
) -> list:
"""Async version of fetch_kv_cache_events using asyncio.to_thread."""
return await asyncio.to_thread(self.fetch_kv_cache_events)
async def fetch_stats_loop_async(
self,
timeout: Optional[float] = None) -> AsyncGenerator[list, None]:
async for data in self._generic_fetch_loop_async(
fetch_method=self.fetch_stats_async,
serializer=self._stats_serializer,
method_name="fetch_stats_loop_async",
timeout=timeout):
yield data
async def fetch_kv_cache_events_loop_async(
self,
timeout: Optional[float] = None) -> AsyncGenerator[list, None]:
async for data in self._generic_fetch_loop_async(
fetch_method=self.fetch_kv_cache_events_async,
serializer=self._kv_cache_events_serializer,
method_name="fetch_kv_cache_events_loop_async",
timeout=timeout):
yield data
async def _generic_fetch_loop_async(
self,
fetch_method,
serializer,
method_name: str,
timeout: Optional[float] = None) -> AsyncGenerator[list, None]:
"""Generic method for fetching data in a loop.
Args:
fetch_method: The async method to call for fetching data
serializer: The serializer function to apply to each item
method_name: Name of the method for logging
timeout: Optional timeout between fetches
"""
while not self.shutdown_event.is_set():
timeout = timeout or 0.1
await asyncio.sleep(timeout)
data = await fetch_method()
# Always yield data, even if empty, to prevent the client looks like hanging
# TODO: Remove the empty data to reduce the IPC overhead
yield [serializer(item) for item in data]
logger_debug(f"RpcWorker {self.rank} quitting {method_name}",
color="yellow")
from .rpc_worker_mixin import RpcWorkerMixin
class RpcWorker(RpcWorkerMixin, BaseWorker):
@ -179,6 +34,18 @@ class RpcWorker(RpcWorkerMixin, BaseWorker):
- `shutdown`: Shutdown the worker.
"""
# Default number of RPC server workers
# Increased to handle concurrent requests and prevent thread pool exhaustion
# Need enough workers for: submit requests + fetch_responses + other operations
# Can be overridden via constructor parameter
DEFAULT_NUM_WORKERS = 32
# Default timeout for fetch_responses in seconds
# This is a short timeout to prevent blocking the event loop while still allowing
# responses to be fetched efficiently. The value is tuned to balance responsiveness
# and CPU usage. Can be overridden via constructor parameter.
DEFAULT_FETCH_TIMEOUT = 0.1
def __init__(
self,
engine: Union[Path, Engine],
@ -189,18 +56,26 @@ class RpcWorker(RpcWorkerMixin, BaseWorker):
hf_model_dir: Optional[Path] = None,
tokenizer: Optional[TokenizerBase] = None,
llm_args: Optional[BaseLlmArgs] = None,
num_workers: Optional[int] = None,
fetch_timeout: Optional[float] = None,
) -> None:
super().__init__(
engine=engine,
executor_config=executor_config,
is_llm_executor=is_llm_executor,
llm_args=llm_args,
batched_logits_processor=batched_logits_processor,
postproc_worker_config=postproc_worker_config,
is_llm_executor=is_llm_executor,
hf_model_dir=hf_model_dir,
tokenizer=tokenizer,
llm_args=llm_args,
)
# Configure number of RPC workers
self.num_workers = num_workers if num_workers is not None else self.DEFAULT_NUM_WORKERS
# Configure fetch timeout
self._fetch_timeout = fetch_timeout if fetch_timeout is not None else self.DEFAULT_FETCH_TIMEOUT
# Extract garbage_collection_gen0_threshold from llm_args if available
self.garbage_collection_gen0_threshold = (
llm_args.garbage_collection_gen0_threshold if llm_args is not None
@ -211,6 +86,10 @@ class RpcWorker(RpcWorkerMixin, BaseWorker):
self._response_queue = Queue()
self.set_result_queue(self._response_queue)
# Note: We don't create a persistent ThreadPoolExecutor anymore
# to avoid thread leaks. Instead, we use asyncio.to_thread() which
# manages threads internally.
def setup_engine(self):
# Force all the ranks to wait here, and start creating the executor simultaneously.
# Only call barrier if we have multiple ranks to avoid hanging in single-process tests
@ -219,6 +98,14 @@ class RpcWorker(RpcWorkerMixin, BaseWorker):
super().setup_engine()
def shutdown(self):
logger_debug(f"[worker] RpcWorker #{mpi_rank()} is shutting down",
color="yellow")
self.shutdown_event.set()
super().shutdown()
logger_debug(f"[worker] RpcWorker #{mpi_rank()} is shutdown",
color="yellow")
def start(self):
pass
@ -257,32 +144,30 @@ class RpcWorker(RpcWorkerMixin, BaseWorker):
# The non-leader worker will setup the engine immediately.
# The leader worker will wait for the RPC call to propagate the
# potential error.
logger_debug(f"Worker {mpi_rank()} is setting up the engine",
color="yellow")
logger_debug(
f"[worker] Worker {mpi_rank()} is setting up the engine",
color="yellow")
worker.setup_engine()
else:
logger_debug(f"Worker {mpi_rank()} is creating the RPC service",
color="yellow")
logger_debug(
f"[worker] Worker {mpi_rank()} is creating the RPC service with {worker.num_workers} workers",
color="yellow")
# Step 2: Create the RPC service, it will expose all the APIs of the worker as remote call to the client
# Set num_workers to larger than 1 since there are some streaming tasks runs infinitely, such as await_responses_async.
rpc_server = RPCServer(worker, num_workers=RpcWorker.NUM_WORKERS)
rpc_server = RPCServer(worker, num_workers=worker.num_workers)
rpc_server.bind(rpc_addr)
rpc_server.start()
logger_debug(f"[worker] RPC server {mpi_rank()} is started",
color="yellow")
# Step 3: Wait for the worker to shutdown
logger_debug(
f"Worker {mpi_rank()} is waiting for the worker to shutdown")
f"[worker] Worker {mpi_rank()} is waiting for shutdown event",
color="yellow")
worker.shutdown_event.wait()
rpc_server.shutdown()
def shutdown(self):
logger_debug(f"RPC worker {mpi_rank()} is shutting down",
color="yellow")
self.shutdown_event.set()
super().shutdown()
logger_debug(f"RPC worker {mpi_rank()} is shutdown", color="yellow")
def __enter__(self):
return self

151
tensorrt_llm/executor/rpc_worker_mixin.py Normal file
View File

@ -0,0 +1,151 @@
import asyncio
from queue import Queue
from threading import Event
from typing import AsyncGenerator, Optional
from .._utils import nvtx_range_debug
from ..llmapi.utils import logger_debug
from .request import GenerationRequest
from .rpc import RPCServer
class RpcWorkerMixin:
"""Mixin for workers that serve RPC requests.
Provides:
- RPC server initialization
- Response queue management
- Async response fetching methods
- Shutdown logic for RPC components
The inheriting class should call init_rpc_worker() in its __init__.
"""
# Default number of RPC server workers
# This can be overridden by setting num_workers in the inheriting class
NUM_WORKERS = 6
def init_rpc_worker(self, rank: int, rpc_addr: Optional[str]):
if rpc_addr is None:
raise RuntimeError("RPC mode enabled but no rpc_addr provided to worker")
self.rank = rank
self.shutdown_event = Event()
self._response_queue = Queue()
self.set_result_queue(self._response_queue)
self.rpc_server = None
self.rpc_addr = rpc_addr
def start_rpc_server(self):
if self.rank == 0:
# Use num_workers if set on the instance, otherwise use class default
num_workers = getattr(self, "num_workers", RpcWorkerMixin.NUM_WORKERS)
self.rpc_server = RPCServer(self, num_workers=num_workers)
self.rpc_server.bind(self.rpc_addr)
self.rpc_server.start()
def submit(self, request: GenerationRequest):
"""Submits a request to the worker."""
with nvtx_range_debug("RpcWorker.submit", color="blue", category="Worker"):
logger_debug(f"[worker] Submitting request {request.id}", color="green")
super().submit(request)
logger_debug(f"[worker] Submitted request {request.id}", color="green")
def fetch_responses(self, timeout: Optional[float] = None) -> list:
"""Fetch responses from the response queue (blocking)."""
logger_debug(f"[worker] RpcWorker {self.rank} is fetching responses", color="yellow")
with nvtx_range_debug("RpcWorker.fetch_responses", color="orange", category="Worker"):
# NOTE: This is a blocking call, it will wait for the responses to be available.
# Use the configured fetch timeout if no timeout is provided
actual_timeout = (
timeout if timeout is not None else getattr(self, "_fetch_timeout", 0.1)
)
responses = super().await_responses(timeout=actual_timeout)
self._await_response_helper.responses_handler(responses)
logger_debug(f"[worker] Fetched {len(responses)} responses", color="green")
qsize = self._response_queue.qsize()
logger_debug(f"[worker] RpcWorker returning {qsize} responses", color="yellow")
all_responses = []
for _ in range(qsize):
# The queue contains batches of responses, so extend the list
all_responses.extend(self._response_queue.get())
return all_responses
async def fetch_responses_async(self, timeout: Optional[float] = None) -> list:
"""Async version of fetch_responses using asyncio.to_thread."""
# Use asyncio.to_thread to avoid blocking the event loop
# This is similar to fetch_stats_async and fetch_kv_cache_events_async
responses = await asyncio.to_thread(self.fetch_responses, timeout=timeout)
return responses
async def fetch_responses_loop_async(self) -> AsyncGenerator[list, None]:
"""Stream responses in a loop until shutdown."""
while not self.shutdown_event.is_set():
responses = await self.fetch_responses_async()
if responses: # Only yield if there are actual responses
logger_debug(
f"[worker] RpcWorker {self.rank} is yielding responses: {responses}",
color="yellow",
)
yield responses # batching the responses to opt IPC performance
else:
# Small delay to prevent busy waiting when no responses
await asyncio.sleep(0)
logger_debug(
f"[worker] RpcWorker {self.rank} quitting fetch_responses_loop_async", color="yellow"
)
async def fetch_stats_async(self, timeout: Optional[float] = None) -> list:
"""Async version of fetch_stats using asyncio.to_thread."""
return await asyncio.to_thread(self.fetch_stats)
async def fetch_kv_cache_events_async(self, timeout: Optional[float] = None) -> list:
"""Async version of fetch_kv_cache_events using asyncio.to_thread."""
return await asyncio.to_thread(self.fetch_kv_cache_events)
async def fetch_stats_loop_async(
self, timeout: Optional[float] = None
) -> AsyncGenerator[list, None]:
"""Stream stats in a loop until shutdown."""
async for data in self._generic_fetch_loop_async(
fetch_method=self.fetch_stats_async,
serializer=self._stats_serializer,
method_name="fetch_stats_loop_async",
timeout=timeout,
):
yield data
async def fetch_kv_cache_events_loop_async(
self, timeout: Optional[float] = None
) -> AsyncGenerator[list, None]:
"""Stream KV cache events in a loop until shutdown."""
async for data in self._generic_fetch_loop_async(
fetch_method=self.fetch_kv_cache_events_async,
serializer=self._kv_cache_events_serializer,
method_name="fetch_kv_cache_events_loop_async",
timeout=timeout,
):
yield data
async def _generic_fetch_loop_async(
self, fetch_method, serializer, method_name: str, timeout: Optional[float] = None
) -> AsyncGenerator[list, None]:
"""Generic method for fetching data in a loop.
Args:
fetch_method: The async method to call for fetching data
serializer: The serializer function to apply to each item
method_name: Name of the method for logging
timeout: Optional timeout between fetches
"""
while not self.shutdown_event.is_set():
timeout = timeout or 0.1
await asyncio.sleep(timeout)
data = await fetch_method()
# Always yield data, even if empty, to prevent the client looks like hanging
# TODO: Remove the empty data to reduce the IPC overhead
yield [serializer(item) for item in data]
logger_debug(f"[worker] RpcWorker {self.rank} quitting {method_name}", color="yellow")

2
tensorrt_llm/llmapi/utils.py
View File

@ -56,6 +56,7 @@ def print_colored(message,
bold_red="\x1b[31;1m",
bold_green="\033[1;32m",
green="\033[0;32m",
cyan="\033[0;36m",
)
reset = "\x1b[0m"
@ -113,6 +114,7 @@ def logger_debug(message,
location) > 50 else location
print_colored(f"{timestamp} [{cur_dualname}]", "bold_green", writer)
print_colored(f" {message}\n", color, writer)
writer.flush()
else:
# Fallback to logger.debug
logger.debug(message)

9
tests/integration/defs/examples/test_ray.py
View File

@ -12,11 +12,7 @@ def ray_example_root(llm_root):
return example_root
@pytest.mark.parametrize("use_rpc", [True, False], ids=["rpc", "no_rpc"])
def test_llm_inference_async_ray(ray_example_root, llm_venv, monkeypatch,
use_rpc):
if use_rpc:
monkeypatch.setenv("TLLM_RAY_USE_RPC", "1")
def test_llm_inference_async_ray(ray_example_root, llm_venv):
script_path = os.path.join(ray_example_root, "llm_inference_async_ray.py")
model_path = f"{llm_models_root()}/llama-models-v2/TinyLlama-1.1B-Chat-v1.0"
venv_check_call(llm_venv, [script_path, "--model", model_path])
@ -60,6 +56,9 @@ def test_llm_inference_distributed_ray(ray_example_root, llm_venv, tp_size,
@pytest.mark.skip_less_device(2)
@pytest.mark.parametrize("tp_size", [1, 2], ids=["tp1", "tp2"])
def test_ray_disaggregated_serving(ray_example_root, llm_venv, tp_size):
if tp_size == 1:
pytest.skip("https://nvbugs/5682551")
if get_device_count() < tp_size * 2:
pytest.skip(f"Need {tp_size * 2} GPUs.")

2
tests/integration/test_lists/test-db/l0_h100.yml
View File

@ -140,7 +140,7 @@ l0_h100:
- unittest/_torch/executor
- unittest/_torch/ray_orchestrator/single_gpu
- unittest/llmapi/test_llm_pytorch.py
- examples/test_ray.py::test_llm_inference_async_ray[no_rpc]
- examples/test_ray.py::test_llm_inference_async_ray
- condition:
ranges:
system_gpu_count:

7
tests/integration/test_lists/waives.txt
View File

@ -303,13 +303,11 @@ full:L40S/accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_
accuracy/test_llm_api.py::TestMixtral8x7BInstruct::test_awq_tp2 SKIP (https://nvbugs/5598847)
unittest/executor/test_rpc.py SKIP (https://nvbugs/5596365)
unittest/executor/test_rpc_worker.py SKIP (https://nvbugs/5583261)
test_e2e.py::test_ptp_quickstart_multimodal_multiturn[gemma-3-27b-it-gemma/gemma-3-27b-it] SKIP (https://nvbugs/5568836)
unittest/llmapi/test_llm_pytorch.py::test_llm_capture_request_error SKIP (https://nvbugs/5599176)
examples/test_phi.py::test_phi_fp8_with_bf16_lora[Phi-3.5-MoE-instruct] SKIP (https://nvbugs/5465143)
unittest/llmapi/test_llm_multi_gpu_pytorch.py::test_llm_rpc_tp2 SKIP (https://nvbugs/5594753)
unittest/llmapi/test_llm_pytorch.py::test_llm_rpc SKIP (https://nvbugs/5594753)
unittest/llmapi/test_llm_pytorch.py::test_llm_rpc_streaming SKIP (https://nvbugs/5594753)
unittest/llmapi/test_memory_profiling.py SKIP (https://nvbugs/5580781)
unittest/executor/test_rpc_proxy.py SKIP (https://nvbugs/5605741)
triton_server/test_triton.py::test_llava[llava] SKIP (https://nvbugs/5547414)
full:RTX/accuracy/test_llm_api_pytorch.py::TestGemma3_1BInstruct::test_auto_dtype SKIP (https://nvbugs/5569696)
accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_4gpus[ep4-cutlass-auto] SKIP (https://nvbugs/5596343)
accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_4gpus[tp4-cutlass-auto] SKIP (https://nvbugs/5596343)
@ -317,7 +315,6 @@ examples/test_phi.py::test_llm_phi_lora_1gpu[Phi-3-mini-4k-instruct-ru-lora-Phi-
accuracy/test_disaggregated_serving.py::TestGemma3_1BInstruct::test_auto_dtype[False] SKIP (https://nvbugs/5569696)
accuracy/test_disaggregated_serving.py::TestGemma3_1BInstruct::test_auto_dtype[True] SKIP (https://nvbugs/5569696)
test_e2e.py::test_trtllm_serve_multimodal_example SKIP (https://nvbugs/5596377)
unittest/llmapi/test_llm_multi_gpu_pytorch.py::test_llm_rpc_streaming_tp2 SKIP (https://nvbugs/5594753)
triton_server/test_triton.py::test_cpp_unit_tests[cpp-unit-tests] SKIP (https://nvbugs/5619359)
triton_server/test_triton_rcca.py::test_rcca_bug_4934893[Temperature:0.5-TOP_P:0.95-TOP_K:10-False-1---False-True-False-0-2048-enableDecoupleMode-inflight_fused_batching-disableTrtOverlap--max_utilization---1-1-1-False-ensemble] SKIP (https://nvbugs/5619369)
accuracy/test_disaggregated_serving.py::TestQwen3_30B_A3B::test_mixed_ctx_gen_model[ctxpp2gentp2] SKIP (https://nvbugs/5582258)

1
tests/unittest/_torch/ray_orchestrator/multi_gpu/test_executor.py
View File

@ -31,7 +31,6 @@ def test_bundle_indices(monkeypatch):
"""Placement via bundle indices"""
monkeypatch.setenv("RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES", "1")
monkeypatch.setenv("TLLM_RAY_USE_RPC", "1")
pg = None
try:

24
tests/unittest/executor/test_base_worker.py
View File

@ -23,6 +23,30 @@ default_model_name = "llama-models-v2/TinyLlama-1.1B-Chat-v1.0"
model_path = llm_models_root() / default_model_name
def create_fake_executor_config(engine_path, tp_size: int = 1):
"""Create TorchLlmArgs and executor_config for testing.
Args:
engine_path: Path to the model
tp_size: Tensor parallel size
Returns:
Tuple of (llm_args, executor_config)
"""
llm_args = TorchLlmArgs(
model=engine_path,
tensor_parallel_size=tp_size,
backend='pytorch',
enable_iter_perf_stats=True,
max_seq_len=2048, # Set reasonable max sequence length
max_batch_size=8, # Set reasonable batch size for tests
max_num_tokens=2048, # Set reasonable max tokens
)
# executor_config is not needed for PyTorch backend
executor_config = None
return llm_args, executor_config
class FakeWorker(BaseWorker):
def __init__(self, engine: str, tp_size: int = 1):

754
tests/unittest/executor/test_ipc.py Normal file
View File

@ -0,0 +1,754 @@
import asyncio
import time
from threading import Thread
import pytest
import zmq
from tensorrt_llm.executor.ipc import ZeroMqQueue
class TestIpcBasics:
"""Test basic synchronous IPC operations."""
def test_pair_socket_with_hmac(self):
"""Test PAIR socket with HMAC encryption."""
# Create server
server = ZeroMqQueue(
address=None,
socket_type=zmq.PAIR,
is_server=True,
is_async=False,
name="test_server",
use_hmac_encryption=True,
)
# Create client with server's address
client = ZeroMqQueue(
address=server.address,
socket_type=zmq.PAIR,
is_server=False,
is_async=False,
name="test_client",
use_hmac_encryption=True,
)
try:
# Test basic send/receive
test_data = {"message": "hello", "value": 42}
client.put(test_data)
received = server.get()
assert received == test_data
# Test reverse direction
response = {"status": "ok", "result": 100}
server.put(response)
received = client.get()
assert received == response
finally:
client.close()
server.close()
def test_pair_socket_without_hmac(self):
"""Test PAIR socket without HMAC encryption."""
# Create server without HMAC
server = ZeroMqQueue(
address=None,
socket_type=zmq.PAIR,
is_server=True,
is_async=False,
name="test_server_no_hmac",
use_hmac_encryption=False,
)
# Create client
client = ZeroMqQueue(
address=(server.address[0], None),
socket_type=zmq.PAIR,
is_server=False,
is_async=False,
name="test_client_no_hmac",
use_hmac_encryption=False,
)
try:
# Test send/receive
test_data = {"message": "hello without encryption", "numbers": [1, 2, 3]}
client.put(test_data)
received = server.get()
assert received == test_data
finally:
client.close()
server.close()
def test_poll_timeout(self):
"""Test poll timeout behavior."""
server = ZeroMqQueue(
address=None,
socket_type=zmq.PAIR,
is_server=True,
is_async=False,
name="test_poll_server",
use_hmac_encryption=False,
)
try:
# Poll should timeout when no data available
start = time.time()
result = server.poll(timeout=1)
elapsed = time.time() - start
assert result is False
assert elapsed >= 1.0
assert elapsed < 1.5 # Allow some margin
finally:
server.close()
def test_poll_with_data(self):
"""Test poll returns True when data is available."""
server = ZeroMqQueue(
address=None,
socket_type=zmq.PAIR,
is_server=True,
is_async=False,
name="test_poll_data_server",
use_hmac_encryption=False,
)
client = ZeroMqQueue(
address=(server.address[0], None),
socket_type=zmq.PAIR,
is_server=False,
is_async=False,
name="test_poll_data_client",
use_hmac_encryption=False,
)
try:
# Send data in background
def send_data():
time.sleep(0.1) # Small delay
client.put({"data": "test"})
thread = Thread(target=send_data)
thread.start()
# Poll should return True
result = server.poll(timeout=2)
assert result is True
# Verify data
received = server.get()
assert received == {"data": "test"}
thread.join()
finally:
client.close()
server.close()
def test_router_socket_with_hmac(self):
"""Test ROUTER socket with HMAC encryption and identity tracking."""
# Create ROUTER server
server = ZeroMqQueue(
address=None,
socket_type=zmq.ROUTER,
is_server=True,
is_async=False,
name="test_router_server",
use_hmac_encryption=True,
)
# Create DEALER client
client = ZeroMqQueue(
address=server.address,
socket_type=zmq.DEALER,
is_server=False,
is_async=False,
name="test_dealer_client",
use_hmac_encryption=True,
)
try:
# Client sends request
request = {"action": "process", "data": [1, 2, 3]}
client.put(request)
# Server receives and tracks identity
received = server.get()
assert received == request
# Server sends response (using stored identity)
response = {"status": "done", "result": 6}
server.put(response)
# Client receives response
received = client.get()
assert received == response
finally:
client.close()
server.close()
def test_dealer_notify_with_retry(self):
"""Test DEALER socket notify_with_retry mechanism."""
# Create ROUTER server
server = ZeroMqQueue(
address=None,
socket_type=zmq.ROUTER,
is_server=True,
is_async=False,
name="test_router_ack_server",
use_hmac_encryption=False,
)
# Create DEALER client
client = ZeroMqQueue(
address=(server.address[0], None),
socket_type=zmq.DEALER,
is_server=False,
is_async=False,
name="test_dealer_ack_client",
use_hmac_encryption=False,
)
try:
# Server thread that acknowledges messages
def server_ack():
msg = server.get()
assert msg == {"notify": "test"}
# Send ACK
server.put({"ack": True})
thread = Thread(target=server_ack)
thread.start()
# Client sends with retry
result = client.notify_with_retry({"notify": "test"}, max_retries=3, timeout=1)
assert result is True
thread.join()
finally:
client.close()
server.close()
def test_dealer_notify_with_retry_timeout(self):
"""Test DEALER socket notify_with_retry timeout behavior."""
# Create ROUTER server (but don't respond)
server = ZeroMqQueue(
address=None,
socket_type=zmq.ROUTER,
is_server=True,
is_async=False,
name="test_router_no_ack_server",
use_hmac_encryption=False,
)
# Create DEALER client
client = ZeroMqQueue(
address=(server.address[0], None),
socket_type=zmq.DEALER,
is_server=False,
is_async=False,
name="test_dealer_no_ack_client",
use_hmac_encryption=False,
)
try:
# Client sends but server doesn't acknowledge
result = client.notify_with_retry({"notify": "test"}, max_retries=2, timeout=0.5)
assert result is False
finally:
client.close()
server.close()
def test_hmac_key_generation(self):
"""Test that server generates HMAC key when encryption is enabled."""
server = ZeroMqQueue(
address=None,
socket_type=zmq.PAIR,
is_server=True,
is_async=False,
name="test_hmac_gen",
use_hmac_encryption=True,
)
try:
# Server should have generated an HMAC key
assert server.hmac_key is not None
assert len(server.hmac_key) == 32
finally:
server.close()
def test_hmac_validation_error_client_no_key(self):
"""Test that client without HMAC key raises ValueError when encryption enabled."""
with pytest.raises(ValueError, match="Client must receive HMAC key"):
ZeroMqQueue(
address=("tcp://127.0.0.1:5555", None), # No HMAC key
socket_type=zmq.PAIR,
is_server=False,
is_async=False,
name="test_client_no_key",
use_hmac_encryption=True, # But encryption enabled
)
def test_hmac_validation_error_key_when_disabled(self):
"""Test that providing HMAC key when encryption disabled raises ValueError."""
with pytest.raises(ValueError, match="should not receive HMAC key"):
ZeroMqQueue(
address=("tcp://127.0.0.1:5555", b"some_key"), # Has key
socket_type=zmq.PAIR,
is_server=False,
is_async=False,
name="test_client_key_disabled",
use_hmac_encryption=False, # But encryption disabled
)
def test_put_noblock_retry(self):
"""Test put_noblock with retry mechanism."""
server = ZeroMqQueue(
address=None,
socket_type=zmq.PAIR,
is_server=True,
is_async=False,
name="test_noblock_server",
use_hmac_encryption=False,
)
client = ZeroMqQueue(
address=(server.address[0], None),
socket_type=zmq.PAIR,
is_server=False,
is_async=False,
name="test_noblock_client",
use_hmac_encryption=False,
)
try:
# Send with put_noblock
test_data = {"nonblocking": True, "value": 123}
client.put_noblock(test_data, retry=3, wait_time=0.001)
# Should be able to receive
received = server.get()
assert received == test_data
finally:
client.close()
server.close()
class TestIpcAsyncBasics:
"""Test asynchronous IPC operations."""
@pytest.mark.asyncio
async def test_async_pair_with_hmac(self):
"""Test async PAIR socket with HMAC encryption."""
# Create async server
server = ZeroMqQueue(
address=None,
socket_type=zmq.PAIR,
is_server=True,
is_async=True,
name="async_server",
use_hmac_encryption=True,
)
# Create async client
client = ZeroMqQueue(
address=server.address,
socket_type=zmq.PAIR,
is_server=False,
is_async=True,
name="async_client",
use_hmac_encryption=True,
)
try:
# Test async send/receive
test_data = {"async": True, "value": 999}
await client.put_async(test_data)
received = await server.get_async()
assert received == test_data
# Test reverse direction
response = {"status": "async_ok"}
await server.put_async(response)
received = await client.get_async()
assert received == response
finally:
client.close()
server.close()
@pytest.mark.asyncio
async def test_async_pair_without_hmac(self):
"""Test async PAIR socket without HMAC encryption."""
server = ZeroMqQueue(
address=None,
socket_type=zmq.PAIR,
is_server=True,
is_async=True,
name="async_server_no_hmac",
use_hmac_encryption=False,
)
client = ZeroMqQueue(
address=(server.address[0], None),
socket_type=zmq.PAIR,
is_server=False,
is_async=True,
name="async_client_no_hmac",
use_hmac_encryption=False,
)
try:
# Test async operations
test_data = {"no_encryption": True, "items": [1, 2, 3, 4, 5]}
await client.put_async(test_data)
received = await server.get_async()
assert received == test_data
finally:
client.close()
server.close()
@pytest.mark.asyncio
async def test_async_router_with_identity(self):
"""Test async ROUTER socket with identity handling."""
server = ZeroMqQueue(
address=None,
socket_type=zmq.ROUTER,
is_server=True,
is_async=True,
name="async_router_server",
use_hmac_encryption=True,
)
client = ZeroMqQueue(
address=server.address,
socket_type=zmq.DEALER,
is_server=False,
is_async=True,
name="async_dealer_client",
use_hmac_encryption=True,
)
try:
# Client sends async request
request = {"async_request": "process"}
await client.put_async(request)
# Server receives with identity
received = await server.get_async()
assert received == request
# Server replies
response = {"async_response": "completed"}
await server.put_async(response)
# Client receives
received = await client.get_async()
assert received == response
finally:
client.close()
server.close()
@pytest.mark.asyncio
async def test_get_async_noblock_timeout(self):
"""Test get_async_noblock timeout expiration."""
server = ZeroMqQueue(
address=None,
socket_type=zmq.PAIR,
is_server=True,
is_async=True,
name="async_timeout_server",
use_hmac_encryption=False,
)
try:
# Should timeout when no data available
with pytest.raises(asyncio.TimeoutError):
await server.get_async_noblock(timeout=0.5)
finally:
server.close()
@pytest.mark.asyncio
async def test_get_async_noblock_success(self):
"""Test get_async_noblock successful receive before timeout."""
server = ZeroMqQueue(
address=None,
socket_type=zmq.PAIR,
is_server=True,
is_async=True,
name="async_noblock_server",
use_hmac_encryption=False,
)
client = ZeroMqQueue(
address=(server.address[0], None),
socket_type=zmq.PAIR,
is_server=False,
is_async=True,
name="async_noblock_client",
use_hmac_encryption=False,
)
try:
# Send data in background
async def send_delayed():
await asyncio.sleep(0.1)
await client.put_async({"delayed": True})
send_task = asyncio.create_task(send_delayed())
# Should receive before timeout
received = await server.get_async_noblock(timeout=2.0)
assert received == {"delayed": True}
await send_task
finally:
client.close()
server.close()
@pytest.mark.asyncio
async def test_put_async_noblock(self):
"""Test put_async_noblock with NOBLOCK flag."""
server = ZeroMqQueue(
address=None,
socket_type=zmq.PAIR,
is_server=True,
is_async=True,
name="async_put_noblock_server",
use_hmac_encryption=False,
)
client = ZeroMqQueue(
address=(server.address[0], None),
socket_type=zmq.PAIR,
is_server=False,
is_async=True,
name="async_put_noblock_client",
use_hmac_encryption=False,
)
try:
# Send with noblock
test_data = {"noblock_async": True}
await client.put_async_noblock(test_data)
# Should be able to receive
received = await server.get_async()
assert received == test_data
finally:
client.close()
server.close()
class TestIpcPressureTest:
"""Test performance and load handling."""
def test_high_frequency_small_messages(self):
"""Test sending many small messages rapidly."""
server = ZeroMqQueue(
address=None,
socket_type=zmq.PAIR,
is_server=True,
is_async=False,
name="pressure_server",
use_hmac_encryption=False,
)
client = ZeroMqQueue(
address=(server.address[0], None),
socket_type=zmq.PAIR,
is_server=False,
is_async=False,
name="pressure_client",
use_hmac_encryption=False,
)
num_messages = 10000
try:
# Send many small messages
def sender():
for i in range(num_messages):
client.put({"id": i, "data": f"msg_{i}"})
# Receive in parallel
def receiver():
received_count = 0
for i in range(num_messages):
msg = server.get()
assert msg["id"] == i
assert msg["data"] == f"msg_{i}"
received_count += 1
return received_count
send_thread = Thread(target=sender)
start_time = time.time()
send_thread.start()
count = receiver()
send_thread.join()
elapsed = time.time() - start_time
# Verify all messages received
assert count == num_messages
print(
f"\nHigh frequency test: {num_messages} messages in {elapsed:.2f}s "
f"({num_messages / elapsed:.0f} msg/s)"
)
finally:
client.close()
server.close()
def test_large_message_size(self):
"""Test sending large messages with HMAC encryption."""
server = ZeroMqQueue(
address=None,
socket_type=zmq.PAIR,
is_server=True,
is_async=False,
name="large_msg_server",
use_hmac_encryption=True,
)
client = ZeroMqQueue(
address=server.address,
socket_type=zmq.PAIR,
is_server=False,
is_async=False,
name="large_msg_client",
use_hmac_encryption=True,
)
num_messages = 100
message_size = 1024 * 1024 # 1 MB
try:
start_time = time.time()
for i in range(num_messages):
# Create large message (1 MB of data)
large_data = {"id": i, "payload": "x" * message_size}
client.put(large_data)
received = server.get()
assert received["id"] == i
assert len(received["payload"]) == message_size
elapsed = time.time() - start_time
total_mb = (num_messages * message_size) / (1024 * 1024)
print(
f"\nLarge message test: {num_messages} x 1MB messages in {elapsed:.2f}s "
f"({total_mb / elapsed:.1f} MB/s)"
)
finally:
client.close()
server.close()
@pytest.mark.asyncio
async def test_concurrent_async_access(self):
"""Test multiple async coroutines sending/receiving simultaneously."""
server = ZeroMqQueue(
address=None,
socket_type=zmq.PAIR,
is_server=True,
is_async=True,
name="concurrent_server",
use_hmac_encryption=False,
)
client = ZeroMqQueue(
address=(server.address[0], None),
socket_type=zmq.PAIR,
is_server=False,
is_async=True,
name="concurrent_client",
use_hmac_encryption=False,
)
num_messages = 1000
try:
# Sender coroutine
async def sender():
for i in range(num_messages):
await client.put_async({"id": i, "data": f"concurrent_{i}"})
if i % 100 == 0:
await asyncio.sleep(0.001) # Small yield
# Receiver coroutine
async def receiver():
received_ids = set()
for _ in range(num_messages):
msg = await server.get_async()
received_ids.add(msg["id"])
return received_ids
# Run concurrently
start_time = time.time()
sender_task = asyncio.create_task(sender())
receiver_task = asyncio.create_task(receiver())
received_ids = await receiver_task
await sender_task
elapsed = time.time() - start_time
# Verify all messages received
assert len(received_ids) == num_messages
assert received_ids == set(range(num_messages))
print(f"\nConcurrent async test: {num_messages} messages in {elapsed:.2f}s")
finally:
client.close()
server.close()
def test_router_socket_multiple_requests(self):
"""Test ROUTER socket handling multiple sequential requests."""
server = ZeroMqQueue(
address=None,
socket_type=zmq.ROUTER,
is_server=True,
is_async=False,
name="router_load_server",
use_hmac_encryption=False,
)
client = ZeroMqQueue(
address=(server.address[0], None),
socket_type=zmq.DEALER,
is_server=False,
is_async=False,
name="dealer_load_client",
use_hmac_encryption=False,
)
num_requests = 1000
try:
start_time = time.time()
for i in range(num_requests):
# Client sends request
client.put({"request_id": i, "action": "process"})
# Server receives
request = server.get()
assert request["request_id"] == i
# Server responds
server.put({"request_id": i, "result": i * 2})
# Client receives response
response = client.get()
assert response["request_id"] == i
assert response["result"] == i * 2
elapsed = time.time() - start_time
print(
f"\nROUTER socket test: {num_requests} round-trips in {elapsed:.2f}s "
f"({num_requests / elapsed:.0f} req/s)"
)
finally:
client.close()
server.close()

467
tests/unittest/executor/test_rpc.py
View File

@ -1,4 +1,5 @@
import asyncio
import threading
import time
import pytest
@ -9,10 +10,11 @@ from tensorrt_llm.executor.rpc.rpc_common import get_unique_ipc_addr
class RpcServerWrapper(RPCServer):
""" A helper class to wrap the RPCServer and manage its lifecycle. """
def __init__(self, *args, addr: str, **kwargs):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.addr = addr
self.addr = get_unique_ipc_addr()
def __enter__(self):
self.bind(self.addr)
@ -24,6 +26,7 @@ class RpcServerWrapper(RPCServer):
class TestRpcBasics:
""" Test the basic functionality of the RPC server and client. """
def test_rpc_server_basics(self):
@ -32,8 +35,7 @@ class TestRpcBasics:
def hello(self):
print("hello")
addr = get_unique_ipc_addr()
with RpcServerWrapper(App(), addr=addr) as server:
with RpcServerWrapper(App()) as server:
pass
def test_remote_call_without_arg(self):
@ -44,9 +46,8 @@ class TestRpcBasics:
print("hello")
return "world"
addr = get_unique_ipc_addr()
with RpcServerWrapper(App(), addr=addr) as server:
with RPCClient(addr) as client:
with RpcServerWrapper(App()) as server:
with RPCClient(server.addr) as client:
ret = client.hello().remote() # sync call
assert ret == "world"
@ -58,9 +59,8 @@ class TestRpcBasics:
print("hello")
return f"hello {name} from {location}"
addr = get_unique_ipc_addr()
with RpcServerWrapper(App(), addr=addr) as server:
with RPCClient(addr) as client:
with RpcServerWrapper(App()) as server:
with RPCClient(server.addr) as client:
ret = client.hello("app", "Marvel").remote()
assert ret == "hello app from Marvel"
@ -72,9 +72,8 @@ class TestRpcBasics:
print("hello")
return f"hello {name} from {location}"
addr = get_unique_ipc_addr()
with RpcServerWrapper(App(), addr=addr) as server:
with RPCClient(addr) as client:
with RpcServerWrapper(App()) as server:
with RPCClient(server.addr) as client:
ret = client.hello(name="app", location="Marvel").remote()
assert ret == "hello app from Marvel"
@ -86,9 +85,8 @@ class TestRpcBasics:
print("hello")
return f"hello {name} from {location}"
addr = get_unique_ipc_addr()
with RpcServerWrapper(App(), addr=addr) as server:
with RPCClient(addr) as client:
with RpcServerWrapper(App()) as server:
with RPCClient(server.addr) as client:
ret = client.hello(name="app", location="Marvel").remote()
assert ret == "hello app from Marvel"
@ -97,9 +95,8 @@ class TestRpcBasics:
class App:
pass
addr = get_unique_ipc_addr()
with RpcServerWrapper(App(), addr=addr) as server:
assert server.address == addr
with RpcServerWrapper(App()) as server:
assert server.address == server.addr
def test_rpc_with_error(self):
@ -108,9 +105,8 @@ class TestRpcBasics:
def hello(self):
raise ValueError("hello")
addr = get_unique_ipc_addr()
with RpcServerWrapper(App(), addr=addr) as server:
with RPCClient(addr) as client:
with RpcServerWrapper(App()) as server:
with RPCClient(server.addr) as client:
with pytest.raises(RPCError):
client.hello().remote()
@ -130,9 +126,8 @@ class TestRpcBasics:
def get_task_submitted(self) -> bool:
return self.task_submitted
addr = get_unique_ipc_addr()
with RpcServerWrapper(App(), addr=addr) as server:
with RPCClient(addr) as client:
with RpcServerWrapper(App()) as server:
with RPCClient(server.addr) as client:
client.send_task().remote(need_response=False)
time.sleep(
0.1
@ -140,6 +135,90 @@ class TestRpcBasics:
assert client.get_task_submitted().remote()
class TestRpcCorrectness:
""" Test the correctness of the RPC framework with various large tasks. """
class App:
def incremental_task(self, v: int):
return v + 1
async def incremental_task_async(self, v: int):
return v + 1
async def streaming_task(self, n: int):
for i in range(n):
yield i
def test_incremental_task(self, num_tasks: int = 10000):
with RpcServerWrapper(TestRpcCorrectness.App()) as server:
with RPCClient(server.addr) as client:
for i in range(num_tasks): # a large number of tasks
result = client.incremental_task(i).remote()
if i % 1000 == 0:
print(f"incremental_task {i} done")
assert result == i + 1, f"result {result} != {i + 1}"
def test_incremental_task_async(self, num_tasks: int = 10000):
with RpcServerWrapper(TestRpcCorrectness.App()) as server:
with RPCClient(server.addr) as client:
async def test_incremental_task_async():
for i in range(num_tasks): # a large number of tasks
result = await client.incremental_task_async(
i).remote_async()
if i % 1000 == 0:
print(f"incremental_task_async {i} done")
assert result == i + 1, f"result {result} != {i + 1}"
asyncio.run(test_incremental_task_async())
@pytest.mark.skip(reason="This test is flaky, need to fix it")
def test_incremental_task_future(self):
with RpcServerWrapper(TestRpcCorrectness.App()) as server:
# Create client with more workers to handle concurrent futures
with RPCClient(server.addr, num_workers=16) as client:
# Process in smaller batches to avoid overwhelming the system
batch_size = 50
total_tasks = 1000 # Reduced from 10000 for stability
for batch_start in range(0, total_tasks, batch_size):
batch_end = min(batch_start + batch_size, total_tasks)
futures = []
# Create futures for this batch
for i in range(batch_start, batch_end):
futures.append(
client.incremental_task(i).remote_future())
# Wait for all futures in this batch to complete
for idx, future in enumerate(futures):
no = batch_start + idx
if no % 100 == 0:
print(f"incremental_task_future {no} done")
assert future.result(
) == no + 1, f"result {future.result()} != {no + 1}"
def test_incremental_task_streaming(self):
with RpcServerWrapper(TestRpcCorrectness.App()) as server:
with RPCClient(server.addr) as client:
async def test_streaming_task():
results = []
no = 0
async for result in client.streaming_task(
10000).remote_streaming():
results.append(result)
if no % 1000 == 0:
print(f"streaming_task {no} done")
no += 1
assert results == [
i for i in range(10000)
], f"results {results} != {[i for i in range(10000)]}"
asyncio.run(test_streaming_task())
class TestRpcError:
class CustomError(Exception):
@ -214,18 +293,21 @@ class TestRpcError:
server.start()
time.sleep(0.1)
with RPCClient(addr) as client:
client = RPCClient(addr)
try:
client.shutdown_server()
pending_futures = [client.task().remote_future() for _ in range(10)]
for future in pending_futures:
with pytest.raises(RPCCancelled):
future.result()
finally:
# Ensure proper cleanup
client.close()
# Wait for background threads to exit
time.sleep(1.0)
time.sleep(5)
client.close()
@pytest.mark.skip(reason="This test is flaky, need to fix it")
def test_timeout_error(self):
"""Test that requests that exceed timeout are handled with proper error."""
@ -236,12 +318,11 @@ class TestRpcError:
time.sleep(2.0)
return "completed"
addr = get_unique_ipc_addr()
with RpcServerWrapper(App(), addr=addr) as server:
with RpcServerWrapper(App()) as server:
time.sleep(0.1)
# Create client with short timeout
with RPCClient(addr, timeout=0.5) as client:
with RPCClient(server.addr, timeout=0.5) as client:
with pytest.raises(RPCError) as exc_info:
client.slow_method().remote(timeout=0.5)
@ -258,11 +339,10 @@ class TestRpcError:
def existing_method(self):
return "exists"
addr = get_unique_ipc_addr()
with RpcServerWrapper(App(), addr=addr) as server:
with RpcServerWrapper(App()) as server:
time.sleep(0.1)
with RPCClient(addr) as client:
with RPCClient(server.addr) as client:
with pytest.raises(RPCError) as exc_info:
client.non_existent_method().remote()
@ -279,19 +359,22 @@ def test_rpc_shutdown_server():
return "world"
addr = get_unique_ipc_addr()
with RPCServer(App()) as server:
server.bind(addr)
server.start()
time.sleep(0.1)
server = RPCServer(App())
server.bind(addr)
server.start()
time.sleep(0.1)
try:
with RPCClient(addr) as client:
ret = client.hello().remote()
assert ret == "world"
client.shutdown_server()
time.sleep(5) # the server dispatcher thread need some time to quit
finally:
# Wait for the server dispatcher thread to quit
time.sleep(1.0)
@pytest.mark.skip(reason="This test is flaky, need to fix it")
def test_rpc_without_response_performance():
# At any circumstances, the RPC call without response should be faster than the one with response
class App:
@ -367,9 +450,8 @@ class TestRpcTimeout:
def setup_method(self, method):
"""Setup RPC server and client for timeout tests."""
# Use unique address based on the test parameter to avoid socket conflicts
test_name = method.__name__
self.address = f"ipc:///tmp/rpc_test_timeout_{test_name}_{id(self)}"
# Use unique address to avoid socket conflicts
self.address = get_unique_ipc_addr()
self.server = RPCServer(self.App())
self.server.bind(self.address)
self.server.start()
@ -378,10 +460,14 @@ class TestRpcTimeout:
def teardown_method(self):
"""Shutdown server and close client."""
self.client.close()
self.server.shutdown()
# Add a small delay to ensure the socket is fully released before the next test
time.sleep(0.5)
# Shutdown server first to stop accepting new requests
if hasattr(self, 'server') and self.server:
self.server.shutdown()
# Then close client to clean up connections
if hasattr(self, 'client') and self.client:
self.client.close()
# Wait longer to ensure all background threads exit completely
time.sleep(1.0)
def run_sync_timeout_test(self):
with pytest.raises(RPCTimeout) as exc_info:
@ -436,16 +522,16 @@ class TestRpcShutdown:
def quick_task(self, task_id: int):
return f"quick_task_{task_id}"
addr = get_unique_ipc_addr()
with RpcServerWrapper(App(), addr=addr) as server:
with RpcServerWrapper(App()) as server:
time.sleep(0.1)
with RPCClient(addr) as client:
with RPCClient(server.addr) as client:
client.quick_task(1).remote()
# repeated shutdown should not raise an error
for i in range(10):
server.shutdown()
@pytest.mark.skip(reason="This test is flaky, need to fix it")
def test_submit_request_after_server_shutdown(self):
class App:
@ -461,13 +547,15 @@ class TestRpcShutdown:
time.sleep(0.1)
with RPCClient(addr) as client:
# This task should be continued after server shutdown
# This task should be cancelled when server shuts down
res = client.foo(10).remote_future(timeout=12)
# The shutdown will block until all pending requests are finished
# The shutdown will now immediately cancel pending requests
server.shutdown()
assert res.result() == "foo"
# Verify the request was cancelled
with pytest.raises(RPCCancelled):
res.result()
class TestApp:
@ -483,14 +571,12 @@ class TestApp:
async def async_multiply(self, x: int, y: int) -> int:
"""Async method."""
await asyncio.sleep(0.01)
self.call_count += 1
return x * y
async def streaming_range(self, n: int):
"""Streaming generator."""
for i in range(n):
await asyncio.sleep(0.01)
yield i
async def streaming_error(self, n: int):
@ -501,11 +587,35 @@ class TestApp:
yield i
async def streaming_timeout(self, delay: float):
"""Streaming generator with configurable delay."""
"""Streaming generator with configurable delay for timeout testing."""
for i in range(10):
await asyncio.sleep(delay)
yield i
async def streaming_forever(self):
"""Streaming generator that never ends, used for cancellation testing."""
i = 0
while True:
await asyncio.sleep(0.1)
yield i
i += 1
@pytest.mark.asyncio
async def test_streaming_task_cancelled():
# Test the streaming task cancelled when the server is shutdown
# This emulates the RpcWorker.fetch_responses_loop_async behavior
app = TestApp()
with RpcServerWrapper(app, num_workers=2, async_run_task=True) as server:
with RPCClient(server.address) as client:
iter = client.streaming_forever().remote_streaming()
# Only get the first 3 values
for i in range(3):
v = await iter.__anext__()
print(f"value {i}: {v}")
# The server should be shutdown while the task is not finished
class TestRpcAsync:
# Use setup_method/teardown_method for pytest class-based setup/teardown
@ -648,9 +758,8 @@ class TestResponsePickleError:
yield nested_function
def test_unpickleable_error(self):
addr = get_unique_ipc_addr()
with RpcServerWrapper(self.App(), addr=addr) as server:
with RPCClient(addr) as client:
with RpcServerWrapper(self.App()) as server:
with RPCClient(server.addr) as client:
with pytest.raises(RPCError) as exc_info:
client.unpickleable_return().remote()
@ -658,13 +767,241 @@ class TestResponsePickleError:
@pytest.mark.asyncio
async def test_unpickleable_streaming_error(self):
addr = get_unique_ipc_addr()
with RpcServerWrapper(self.App(), addr=addr,
async_run_task=True) as server:
with RPCClient(addr) as client:
with RpcServerWrapper(self.App(), async_run_task=True) as server:
with RPCClient(server.addr) as client:
with pytest.raises(RPCStreamingError) as exc_info:
async for _ in client.unpickleable_streaming_return(
).remote_streaming():
pass
assert "Failed to pickle response" in str(exc_info.value)
class TestRpcRobustness:
class App:
LARGE_RESPONSE_SIZE = 1024 * 1024 * 10 # 10MB
def remote_with_large_response(self):
return b"a" * self.LARGE_RESPONSE_SIZE
async def streaming_with_large_response(self):
for i in range(1000):
yield b"a" * self.LARGE_RESPONSE_SIZE
async def get_streaming(self):
for i in range(1000):
yield i
def test_remote_with_large_response(self):
with RpcServerWrapper(self.App()) as server:
with RPCClient(server.addr) as client:
for i in range(100):
result = client.remote_with_large_response().remote()
assert result == b"a" * self.App.LARGE_RESPONSE_SIZE
@pytest.mark.asyncio
async def test_streaming_with_large_response(self):
with RpcServerWrapper(self.App()) as server:
with RPCClient(server.addr) as client:
async for result in client.streaming_with_large_response(
).remote_streaming():
assert result == b"a" * self.App.LARGE_RESPONSE_SIZE
def test_threaded_streaming(self):
"""Test that get_streaming can be safely called from multiple threads."""
# All the async remote calls will be submitted to the RPCClient._loop, let
# it handle the concurrent requests. Once the response arrives, it will
# be processed by the RPCClient._loop, and dispatch to the corresponding
# task via the dedicated AsyncQueue.
num_threads = 100
items_per_stream = 100
# Use shorter stream for faster test
class TestApp:
async def get_streaming(self):
for i in range(items_per_stream):
yield i
with RpcServerWrapper(TestApp(), async_run_task=True) as server:
errors = []
results = [None] * num_threads
def stream_consumer(thread_id: int):
"""Function to be executed in each thread."""
print(f"Thread {thread_id} started")
try:
# Each thread creates its own client connection
with RPCClient(server.addr) as client:
collected = []
async def consume_stream():
async for value in client.get_streaming(
).remote_streaming():
collected.append(value)
# Run the async streaming call in this thread
asyncio.run(consume_stream())
# Verify we got all expected values
expected = list(range(items_per_stream))
if collected != expected:
errors.append(
f"Thread {thread_id}: Expected {expected}, got {collected}"
)
else:
results[thread_id] = collected
except Exception as e:
errors.append(
f"Thread {thread_id}: {type(e).__name__}: {str(e)}")
# Create and start multiple threads
threads = []
for i in range(num_threads):
thread = threading.Thread(target=stream_consumer, args=(i, ))
threads.append(thread)
thread.start()
# Wait for all threads to complete
for thread in threads:
thread.join(timeout=30) # 30 second timeout per thread
# Check for any errors
if errors:
error_msg = "\n".join(errors)
pytest.fail(
f"Thread safety test failed with errors:\n{error_msg}")
# Verify all threads completed successfully
for i, result in enumerate(results):
assert result is not None, f"Thread {i} did not complete successfully"
assert len(
result
) == items_per_stream, f"Thread {i} got {len(result)} items, expected {items_per_stream}"
def test_threaded_remote_call(self):
"""Test that regular remote calls can be safely made from multiple threads."""
# Each thread will make multiple synchronous remote calls
# This tests if RPCClient can handle concurrent requests from different threads
num_threads = 100
calls_per_thread = 100
class TestApp:
def __init__(self):
self.call_count = 0
self.lock = threading.Lock()
def increment(self, v):
with self.lock:
self.call_count += 1
threading.get_ident()
return v + 1
app = TestApp()
with RpcServerWrapper(app) as server:
errors = []
results = [None] * num_threads
client = RPCClient(server.addr)
def remote_caller(thread_id: int):
"""Function to be executed in each thread."""
print(f"Thread {thread_id} started")
try:
thread_results = []
for i in range(calls_per_thread):
result = client.increment(i).remote()
expected = i + 1
if result != expected:
errors.append(
f"Thread {thread_id}, call {i}: Expected {expected}, got {result}"
)
thread_results.append(result)
results[thread_id] = thread_results
except Exception as e:
errors.append(
f"Thread {thread_id}: {type(e).__name__}: {str(e)}")
finally:
print(f"Thread {thread_id} completed")
# Create and start multiple threads
threads = []
for i in range(num_threads):
thread = threading.Thread(target=remote_caller,
args=(i, ),
daemon=True)
threads.append(thread)
thread.start()
# Wait for all threads to complete
for thread in threads:
thread.join(timeout=30) # 30 second timeout per thread
client.close()
# Check for any errors
if errors:
error_msg = "\n".join(errors)
pytest.fail(
f"Thread safety test failed with errors:\n{error_msg}")
# Verify all threads completed successfully
for i, result in enumerate(results):
assert result is not None, f"Thread {i} did not complete successfully"
assert len(
result
) == calls_per_thread, f"Thread {i} made {len(result)} calls, expected {calls_per_thread}"
# Verify total call count
expected_total_calls = num_threads * calls_per_thread
assert app.call_count == expected_total_calls, \
f"Expected {expected_total_calls} total calls, but got {app.call_count}"
def test_repeated_creation_and_destruction(self, num_calls: int = 100):
"""Test robustness of repeated RPCServer/RPCClient creation and destruction.
This test ensures there are no resource leaks, socket exhaustion, or other
issues when repeatedly creating and destroying server/client pairs.
"""
class TestApp:
def __init__(self):
self.counter = 0
def increment(self, value: int) -> int:
self.counter += 1
return value + 1
def get_counter(self) -> int:
return self.counter
for i in range(num_calls):
# Create app, server, and client
# RpcServerWrapper automatically generates unique addresses
app = TestApp()
with RpcServerWrapper(app) as server:
with RPCClient(server.addr) as client:
# Perform a few remote calls to verify functionality
result1 = client.increment(10).remote()
assert result1 == 11, f"Iteration {i}: Expected 11, got {result1}"
result2 = client.increment(20).remote()
assert result2 == 21, f"Iteration {i}: Expected 21, got {result2}"
counter = client.get_counter().remote()
assert counter == 2, f"Iteration {i}: Expected counter=2, got {counter}"
if i % 10 == 0:
print(
f"Iteration {i}/{num_calls} completed successfully")
print(f"All {num_calls} iterations completed successfully")

22
tests/unittest/executor/test_rpc_proxy.py
View File

@ -7,8 +7,8 @@ from test_base_worker import create_fake_executor_config
from tensorrt_llm.executor.rpc_proxy import GenerationExecutorRpcProxy
from tensorrt_llm.llmapi.llm_args import KvCacheConfig
from tensorrt_llm.llmapi.mpi_session import MpiPoolSession
from tensorrt_llm.llmapi.tokenizer import TransformersTokenizer
from tensorrt_llm.llmapi.utils import logger_debug
from tensorrt_llm.sampling_params import SamplingParams
# isort: off
@ -31,9 +31,9 @@ class TestRpcProxy:
llm_args.kv_cache_config = KvCacheConfig(
event_buffer_max_size=1000, # Enable event buffer
enable_block_reuse=True, # Required for KV cache events
free_gpu_memory_fraction=0.6,
)
mpi_session = MpiPoolSession(n_workers=tp_size)
proxy = GenerationExecutorRpcProxy(
worker_kwargs={
"engine": model_path,
@ -43,7 +43,6 @@ class TestRpcProxy:
"hf_model_dir": model_path,
},
model_world_size=tp_size,
mpi_session=mpi_session,
is_llm_executor=True, # Enable stats collection
)
@ -55,8 +54,7 @@ class TestRpcProxy:
return proxy
@pytest.mark.skip(reason="https://nvbugs/5579234")
@pytest.mark.parametrize("num_reqs", [1, 10])
@pytest.mark.parametrize("num_reqs", [1, 5, 10])
def test_tp1(self, num_reqs):
tokenizer = TransformersTokenizer.from_pretrained(model_path)
prompt = "A B C D"
@ -64,19 +62,21 @@ class TestRpcProxy:
max_tokens = 8
with self.create_proxy(tp_size=1) as proxy:
logger_debug(f"[Test] Proxy created", color="green")
sampling_params = SamplingParams(max_tokens=max_tokens)
for _ in range(num_reqs):
logger_debug(f"[Test] Generating {_}th", color="green")
result = proxy.generate(prompt_token_ids, sampling_params)
print(f"get result: {result}")
assert similar(tokenizer.decode(result.outputs[0].token_ids),
'E F G H I J K L')
logger_debug(f"req {_} get result: {result}", color="green")
stats = proxy.get_stats(timeout=2)
assert stats
#stats = proxy.get_stats(timeout=2)
#assert stats
kv_cache_events = proxy.get_kv_events(timeout=2)
#kv_cache_events = proxy.get_kv_events(timeout=2)
# KV cache events may be empty if no cache operations occurred
assert isinstance(kv_cache_events, list)
#assert isinstance(kv_cache_events, list)
@pytest.mark.parametrize("num_reqs", [1, 10])
@skip_single_gpu
@ -97,4 +97,4 @@ class TestRpcProxy:
if __name__ == "__main__":
TestRpcProxy().test_tp1(1)
TestRpcProxy().test_tp1(20)

268
tests/unittest/executor/test_rpc_worker.py
View File

@ -1,24 +1,16 @@
import asyncio
import multiprocessing
import os
import sys
import time
from concurrent.futures import ProcessPoolExecutor
import pytest
from tensorrt_llm.executor.request import GenerationRequest
from tensorrt_llm.executor.rpc import RPCClient
from tensorrt_llm.executor.rpc.rpc_common import get_unique_ipc_addr
from tensorrt_llm.executor.rpc_worker import RpcWorker
from tensorrt_llm.llmapi.llm_args import TorchLlmArgs
from tensorrt_llm.llmapi.mpi_session import MpiPoolSession
from tensorrt_llm.llmapi.llm_args import KvCacheConfig, TorchLlmArgs
from tensorrt_llm.sampling_params import SamplingParams
# isort: off
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + "/..")
from utils.llm_data import llm_models_root
from utils.util import skip_single_gpu
# isort: on
model_path = llm_models_root() / "llama-models-v2/TinyLlama-1.1B-Chat-v1.0"
@ -33,230 +25,62 @@ class TestRpcWorkerTP1:
tensor_parallel_size=1,
backend='pytorch',
enable_iter_perf_stats=True,
kv_cache_config=KvCacheConfig(free_gpu_memory_fraction=0.5, ),
)
self.pool, self.addr = self.create_worker_pool()
self.client = self.create_rpc_client(self.addr)
self.client.setup_engine().remote()
print(f"Worker setup engine done")
time.sleep(10)
def teardown_method(self):
self.client.shutdown().remote()
self.pool.shutdown()
self.client.close()
def create_worker_pool(self):
addr = get_unique_ipc_addr()
mp_context = multiprocessing.get_context(
'spawn') # spawn for CUDA context
pool = ProcessPoolExecutor(max_workers=1, mp_context=mp_context)
pool.submit(
RpcWorker.main_task,
# Create RpcWorker instance
self.worker = RpcWorker(
engine=model_path,
rpc_addr=addr,
llm_args=self.llm_args,
hf_model_dir=model_path,
)
return pool, addr
def create_rpc_client(self, addr: str):
client = RPCClient(addr)
return client
def test_create_shutdown(self):
pass
def test_fetch_responses_sync(self):
# Wait a bit to ensure engine is ready
time.sleep(1)
print(f"start to submit")
self.client.submit(
GenerationRequest(prompt_token_ids=[3, 4, 5],
sampling_params=SamplingParams(
max_tokens=5)), ).remote(need_response=False)
print(f"submit done")
time.sleep(3)
results = []
# Fetch responses
results.extend(self.client.fetch_responses().remote())
assert len(results) == 1
@pytest.mark.skip(reason="https://nvbugs/5583261")
def test_fetch_responses_streaming_sync(self):
self.client.submit(
GenerationRequest(prompt_token_ids=[3, 4, 5],
sampling_params=SamplingParams(max_tokens=5),
streaming=True), ).remote(need_response=False)
results = []
for i in range(10):
res = self.client.fetch_responses().remote(timeout=1.0)
results.extend(res)
print(f"fetch_responses {i} result: {results}")
# If we've received enough results, break early
if len(results) >= 5:
break
assert 0 < len(results) <= 5
@pytest.mark.skip(reason="https://nvbugs/5583261")
@pytest.mark.asyncio
@pytest.mark.parametrize("req_count", [10])
async def test_main_loop_async(self, req_count: int):
await asyncio.sleep(1)
async def process_request_streaming():
for i in range(req_count):
ret = self.client.submit(
GenerationRequest(
prompt_token_ids=[3, 4, 5],
sampling_params=SamplingParams(max_tokens=5),
streaming=True), ).remote(need_response=False)
assert ret is None
print("submit result: ", ret)
# NOTE: known issue, the responses should be fetched before shutdown,
# or the shutdown will hang.
results = []
responses_per_client = {}
expected_responses_per_client = 5 # max_tokens=5
print(f"start to fetch_responses_async")
no = 0
async for result in self.client.fetch_responses_loop_async(
).remote_streaming():
if result: # result is already a list of lists
print(
f"fetch_responses_async batch {no}, received {len(result)} sub-batches"
)
for batch in result:
if isinstance(batch, list):
print(f" Sub-batch has {len(batch)} responses")
results.extend(batch)
# Track responses per client
for response in batch:
client_id = response.client_id
if client_id not in responses_per_client:
responses_per_client[client_id] = 0
responses_per_client[client_id] += 1
else:
# Single response
results.append(batch)
client_id = batch.client_id
if client_id not in responses_per_client:
responses_per_client[client_id] = 0
responses_per_client[client_id] += 1
no += 1
# Check if all clients have received their expected responses
completed_clients = sum(
1 for count in responses_per_client.values()
if count >= expected_responses_per_client)
print(f"Responses per client: {responses_per_client}")
print(f"Completed clients: {completed_clients}/{req_count}")
# Break when we've received all expected responses
if completed_clients >= req_count:
print(
f"All {completed_clients} clients completed after {no} batches"
)
break
# Safety break to prevent infinite loop
if no >= req_count * 20: # Much higher limit as safety
print(f"Safety break after {no} batches")
break
print(f"Received {no} batches of streaming responses")
print(f"Total responses received: {len(results)}")
print(f"Final responses per client: {responses_per_client}")
assert results
assert len(responses_per_client) >= req_count
await process_request_streaming()
@pytest.mark.skip(reason="https://nvbugs/5583261")
@pytest.mark.asyncio
async def test_fetch_stats_loop_async(self):
await asyncio.sleep(1)
results = []
max_batches = 5
async def consume_stats():
async for stats in self.client.fetch_stats_loop_async(
).remote_streaming():
results.append(stats)
assert not stats # empty stats
if len(results) >= max_batches:
break
await asyncio.wait_for(consume_stats(), timeout=5)
assert len(results) == max_batches
assert all(not stats for stats in results)
class TestRpcWorkerTP2:
def setup_method(self):
self.llm_args = TorchLlmArgs(
model=model_path,
tensor_parallel_size=2,
backend='pytorch',
enable_iter_perf_stats=True,
)
self.session, self.addr, self.futures = self.create_worker_session()
self.client = self.create_rpc_client(self.addr)
self.client.setup_engine().remote()
time.sleep(10)
# Initialize the engine
self.worker.setup_engine()
def teardown_method(self):
self.client.shutdown().remote()
self.session.shutdown()
self.client.close()
# Clean up the worker
self.worker.shutdown()
def create_worker_session(self):
session = MpiPoolSession(n_workers=2)
addr = get_unique_ipc_addr()
futures = session.submit(RpcWorker.main_task,
engine=model_path,
rpc_addr=addr,
llm_args=self.llm_args,
hf_model_dir=model_path,
model_world_size=2)
return session, addr, futures
def test_fetch_responses_async(self):
"""Test that fetch_responses_async can be called and returns a list."""
# Submit a request first
sampling_params = SamplingParams(max_tokens=10)
request = GenerationRequest(prompt_token_ids=[3, 4, 5],
sampling_params=sampling_params)
self.worker.submit(request)
def create_rpc_client(self, addr: str):
return RPCClient(addr)
# Sleep a bit to let the request start processing
time.sleep(0.5)
@skip_single_gpu
@pytest.mark.gpu2
@pytest.mark.skip(reason="https://nvbugs/5583261")
def test_create_shutdown(self):
# Invoke setup_engine in rank 0, and that will unblock all the ranks to
# invoke setup_engine simultaneously.
pass
# Fetch responses with a timeout to prevent hanging
responses = asyncio.run(self.worker.fetch_responses_async(timeout=1.0))
assert isinstance(responses, list)
@skip_single_gpu
@pytest.mark.gpu2
@pytest.mark.skip(reason="https://nvbugs/5583261")
def test_fetch_responses_sync(self):
# Wait a bit to ensure engine is ready
time.sleep(1)
def test_fetch_stats_async(self):
"""Test that fetch_stats_async can be called and returns a list."""
# Submit a request first to generate some stats
sampling_params = SamplingParams(max_tokens=10)
request = GenerationRequest(prompt_token_ids=[3, 4, 5],
sampling_params=sampling_params)
self.worker.submit(request)
self.client.submit(
GenerationRequest(prompt_token_ids=[3, 4, 5],
sampling_params=SamplingParams(
max_tokens=5)), ).remote(need_response=False)
# Sleep a bit to let the request start processing
time.sleep(0.5)
# Wait for generation to complete
time.sleep(3)
# Fetch stats
stats = asyncio.run(self.worker.fetch_stats_async())
assert isinstance(stats, list)
results = []
# Fetch responses with timeout
results.extend(self.client.fetch_responses().remote(timeout=5))
assert len(results) == 1
def test_fetch_kv_cache_events_async(self):
"""Test that fetch_kv_cache_events_async can be called and returns a list."""
# Submit a request first to generate some kv cache events
sampling_params = SamplingParams(max_tokens=10)
request = GenerationRequest(prompt_token_ids=[3, 4, 5],
sampling_params=sampling_params)
self.worker.submit(request)
# Sleep a bit to let the request start processing
time.sleep(0.5)
# Fetch kv cache events
events = asyncio.run(self.worker.fetch_kv_cache_events_async())
assert isinstance(events, list)

1
tests/unittest/llmapi/test_llm_multi_gpu_pytorch.py
View File

@ -47,6 +47,7 @@ def test_llama_7b_lora_tp2():
@pytest.mark.gpu2
@pytest.mark.skip(reason="https://nvbugs/5682551")
def test_llama_7b_multi_lora_tp2():
# For LoRA checkpoints without finetuned embedding and lm_head, we can either:
# (1) specify lora_target_modules, or

5
tests/unittest/llmapi/test_llm_pytorch.py
View File

@ -366,6 +366,7 @@ def _check_llama_7b_multi_lora_evict_load_new_adapters(
@skip_gpu_memory_less_than_40gb
@skip_ray # https://nvbugs/5682551
def test_llama_7b_multi_lora_evict_and_reload_lora_gpu_cache():
"""Test eviction and re-loading a previously evicted adapter from the LoRA GPU cache, within a single
llm.generate call, that's repeated twice.
@ -460,6 +461,7 @@ def test_llama_7b_peft_cache_config_affects_peft_cache_size():
cuda_graph_config=None)
@skip_ray # https://nvbugs/5682551
@skip_gpu_memory_less_than_40gb
def test_llama_7b_lora_config_overrides_peft_cache_config():
"""Tests that cache size args in lora_config LLM arg override the cache size
@ -938,7 +940,8 @@ class TestLlmError:
@skip_ray
def test_llm_rpc():
@pytest.mark.parametrize("num_requests", [1, 5, 10])
def test_llm_rpc(num_requests: int):
# TODO: remove the with-statement when shutdown hang issue is fixed
with LLM(model=llama_model_path,
kv_cache_config=global_kvcache_config,