diff --git a/tensorrt_llm/executor/base_worker.py b/tensorrt_llm/executor/base_worker.py index 098d4674603d..131363011759 100644 --- a/tensorrt_llm/executor/base_worker.py +++ b/tensorrt_llm/executor/base_worker.py @@ -255,9 +255,20 @@ def set_result_queue(self, queue): assert self.frontend_result_queues is None self.result_queue = queue - def set_postproc_queues(self, queues: List["IpcQueue"]): - """ Set the IPC queues for feeding post-processing processes. """ - assert self.result_queue is None + def set_postproc_queues(self, + queues: List["IpcQueue"], + *, + coexist_with_result_queue: bool = False): + """ Set the IPC queues for feeding post-processing processes. + + coexist_with_result_queue: the classic proxy gives each PostprocWorker + its own push lane to the frontend, so a result_queue must not exist + there. Under RPC/Ray orchestration the finished PostprocWorker.Output + records are collected back INTO the result queue (the single RPC + response stream), so both queues legitimately coexist. + """ + if not coexist_with_result_queue: + assert self.result_queue is None assert self.frontend_result_queues is None self.postproc_queues = queues @@ -1465,6 +1476,35 @@ def _get_logprobs(worker, return logprobs_result +def _send_rsp_to_postproc( + worker, response: Union[tllm.Response, ResponseWrapper, ErrorResponse], + postproc_batches: Optional[List[List["PostprocWorker.Input"]]]): + """Shard a raw response to the postproc workers (batched or direct put).""" + sampling_params, postproc_params, disaggregated_params = ( + _get_params_for_first_rsp(worker, response.client_id)) + inp = PostprocWorker.Input( + response, + # sampling_params is necessary for creating fake GenerationResult + # instances in the postproc processes. They are for incremental + # detokenize. They should be transmitted only once for each + # Request. + sampling_params=sampling_params, + postproc_params=postproc_params, + disaggregated_params=disaggregated_params, + streaming=worker._results.get(response.client_id, None)._streaming) + + # Group the responses into buckets for the postprocessing steps. + # Bucketing is used instead of random dispatching because the + # incremental detokenization during postprocessing relies on the + # prior CompletionOutput of a given request. + pid = response.client_id % worker.postproc_config.num_postprocess_workers + + if postproc_batches is None: + worker.postproc_queues[pid].put(inp) + else: + postproc_batches[pid].append(inp) + + def _send_rsp( worker, response: Union[tllm.Response, ResponseWrapper, ErrorResponse], @@ -1472,7 +1512,15 @@ def _send_rsp( rsp_batch: Optional[List[tllm.Response]] = None): # if postproc_batches is set, append to batch instead of putting to IpcQueue - if worker.frontend_result_queues is not None: + # Postprocess parallelism takes priority over the direct result routes: + # under RPC/Ray orchestration the worker holds a result_queue (the RPC + # response stream feed) AND postproc input queues at the same time, and + # raw responses must go to the postproc workers first — their finished + # Output records re-enter the result_queue via the collector thread + # (see RpcWorkerMixin.init_postproc_workers). + if postproc_batches is not None: + _send_rsp_to_postproc(worker, response, postproc_batches) + elif worker.frontend_result_queues is not None: # Route to the origin frontend's result lane; None/out-of-range ids # fall back to lane 0 (see frontend_lane_index). if rsp_batch is not None: @@ -1487,29 +1535,7 @@ def _send_rsp( else: worker.result_queue.put(response) else: - sampling_params, postproc_params, disaggregated_params = ( - _get_params_for_first_rsp(worker, response.client_id)) - inp = PostprocWorker.Input( - response, - # sampling_params is necessary for creating fake GenerationResult - # instances in the postproc processes. They are for incremental - # detokenize. They should be transmitted only once for each - # Request. - sampling_params=sampling_params, - postproc_params=postproc_params, - disaggregated_params=disaggregated_params, - streaming=worker._results.get(response.client_id, None)._streaming) - - pid = response.client_id % worker.postproc_config.num_postprocess_workers - - if not postproc_batches: - # Group the responses into buckets for the postprocessing steps. - # Bucketing is used instead of random dispatching because the - # incremental detokenization during postprocessing relies on the - # prior CompletionOutput of a given request. - worker.postproc_queues[pid].put(inp) - else: - postproc_batches[pid].append(inp) + _send_rsp_to_postproc(worker, response, None) # Eliminate the finished GenerationRequest instances timely, which may # take considerable memory. diff --git a/tensorrt_llm/executor/ray/executor.py b/tensorrt_llm/executor/ray/executor.py index cd7368fe8dd9..89096308a806 100644 --- a/tensorrt_llm/executor/ray/executor.py +++ b/tensorrt_llm/executor/ray/executor.py @@ -562,8 +562,7 @@ def _get_default(tp_size): # path 2 return _get_default(tp_size) - @property - def enable_postprocess_parallel(self) -> bool: - ret = super().enable_postprocess_parallel - assert ret == False, "Postprocess parallel is not supported in RayExecutor" - return ret + # Postprocess parallelism is supported: the rank-0 RayGPUWorker spawns a + # local PostprocWorker pool and feeds finished Output records back into + # the RPC response stream (see RpcWorkerMixin.init_postproc_workers), so + # the base-class property applies unchanged. diff --git a/tensorrt_llm/executor/ray/gpu_worker.py b/tensorrt_llm/executor/ray/gpu_worker.py index cd5165fbba0a..f05220ab777d 100644 --- a/tensorrt_llm/executor/ray/gpu_worker.py +++ b/tensorrt_llm/executor/ray/gpu_worker.py @@ -261,6 +261,11 @@ def __init__( raise RuntimeError( "RPC mode enabled but no rpc_addr provided to RayGPUWorker") self.init_rpc_worker(self.global_rank, rpc_addr, hmac_key) + # Postprocess parallelism: spawn the local PostprocWorker pool on the + # response-producing rank. Outputs re-enter _response_queue via the + # collector thread, so the RPC stream needs no protocol change. + if self.global_rank == 0 and self.postproc_config.enabled: + self.init_postproc_workers() self.start_rpc_server() def setup_engine(self): @@ -329,6 +334,9 @@ def shutdown(self): if hasattr(self, 'shutdown_event'): self.shutdown_event.set() + if getattr(self, '_postproc_pool', None) is not None: + self.shutdown_postproc_workers() + if hasattr(self, 'rpc_server') and self.rpc_server is not None: logger.info(f"[Rank {self.global_rank}] Shutting down RPC server") try: diff --git a/tensorrt_llm/executor/result.py b/tensorrt_llm/executor/result.py index c459d1faa1b3..cfa17257b39d 100644 --- a/tensorrt_llm/executor/result.py +++ b/tensorrt_llm/executor/result.py @@ -154,6 +154,23 @@ class CompletionOutput: # the result of result_handler passed to postprocess workers _postprocess_result: Any = None + def __getstate__(self): + # _incremental_states holds a tokenizers.DecodeStream (a Rust object, + # not picklable) used for process-local incremental detokenization. + # Receivers never need it, so drop it when this output crosses a + # process boundary — e.g. a PostprocWorker streaming a raw + # CompletionOutput (no postproc_params) back over IPC. + # slots=True dataclass: object.__getstate__ returns + # (dict_or_None, slots_dict). + state = object.__getstate__(self) + if isinstance(state, tuple) and len(state) == 2: + d, slots = state + if slots and '_incremental_states' in slots: + slots = dict(slots) + slots['_incremental_states'] = None + return (d, slots) + return state + @property def length(self) -> int: return len(self.token_ids) diff --git a/tensorrt_llm/executor/rpc_proxy_mixin.py b/tensorrt_llm/executor/rpc_proxy_mixin.py index ecbb86e25ead..b3779b2ecbae 100644 --- a/tensorrt_llm/executor/rpc_proxy_mixin.py +++ b/tensorrt_llm/executor/rpc_proxy_mixin.py @@ -8,6 +8,7 @@ from ..llmapi.tracer import global_tracer from ..llmapi.utils import _SyncQueue from ..logger import logger +from .postproc_worker import PostprocWorker from .request import GenerationRequest from .result import GenerationResult from .rpc import RPCClient @@ -121,8 +122,16 @@ def process_res(res: list): else: queue.put(r) - if (is_llm_response(r) and r.result.is_final) or isinstance(r, ErrorResponse): - self._results.pop(client_id) + # PostprocWorker.Output records arrive on the same stream when + # postprocess parallelism is enabled (worker-side pool); they + # carry their own is_final, mirroring the classic proxy's + # dispatch_result_task handling. + if ( + (is_llm_response(r) and r.result.is_final) + or isinstance(r, ErrorResponse) + or (isinstance(r, PostprocWorker.Output) and r.is_final) + ): + self._results.pop(client_id, None) # Handle the case where responses might not be a list of lists if responses and not isinstance(responses[0], list): diff --git a/tensorrt_llm/executor/rpc_worker_mixin.py b/tensorrt_llm/executor/rpc_worker_mixin.py index 1b3ff9b06f41..be2bd1b4fd4f 100644 --- a/tensorrt_llm/executor/rpc_worker_mixin.py +++ b/tensorrt_llm/executor/rpc_worker_mixin.py @@ -53,6 +53,87 @@ def init_rpc_worker(self, rank: int, rpc_addr: Optional[str], hmac_key: bytes): self.rpc_server = None self.rpc_addr = rpc_addr + self._postproc_pool = None + self._postproc_input_queues = None + self._postproc_collector = None + self._postproc_collector_thread = None + + def init_postproc_workers(self): + """Spawn local PostprocWorker processes feeding the RPC response stream. + + The classic (MPI proxy) path gives each PostprocWorker a dedicated + push lane straight to the frontend process. Under RPC/Ray + orchestration no such lane exists: every record must travel the single + RPC response stream. Instead of teaching PostprocWorker a new + transport, its push pipe is pointed at a local collector socket whose + consumer thread enqueues the already-final ``PostprocWorker.Output`` + batches into ``_response_queue`` — the same queue ``fetch_responses`` + drains — so the stream, the proxy demux, and the client stay on the + code they already run for the classic postproc path. + + Call on the response-producing rank only (rank 0), after + ``init_rpc_worker``. + """ + import threading + from concurrent.futures import ProcessPoolExecutor + + import zmq + + from .ipc import IpcQueue + from .postproc_worker import PostprocWorker, postproc_worker_main + + num = self.postproc_config.num_postprocess_workers + assert num > 0 + + self._postproc_input_queues = [ + IpcQueue(is_server=True, name=f"rpc_worker_postproc_input_{i}") for i in range(num) + ] + self._postproc_collector = IpcQueue( + is_server=True, socket_type=zmq.PULL, name="rpc_worker_postproc_collector" + ) + # Both the result_queue (RPC stream feed) and the postproc input + # queues are live on this worker — see set_postproc_queues docstring. + self.set_postproc_queues(self._postproc_input_queues, coexist_with_result_queue=True) + + # fork (default), matching the classic path. spawn is NOT usable + # here: the spawn bootstrap re-imports the Ray worker's __main__, + # which deadlocks inside a Ray actor (verified empirically). + self._postproc_pool = ProcessPoolExecutor(max_workers=num) + for i in range(num): + self._postproc_pool.submit( + postproc_worker_main, + self._postproc_input_queues[i].address, + [self._postproc_collector.address], + self.postproc_config.postprocess_tokenizer_dir, + PostprocWorker.default_record_creator, + self.postproc_config.post_processor_hook, + ) + + def _collect_postproc_outputs(): + while not self.shutdown_event.is_set(): + batch = self._postproc_collector.get() + if batch is None: + break + # fetch_responses drains _response_queue batch-wise; Output + # batches ride the RPC stream exactly like final responses. + self._response_queue.put(batch) + + self._postproc_collector_thread = threading.Thread( + target=_collect_postproc_outputs, name="rpc_worker_postproc_collector", daemon=True + ) + self._postproc_collector_thread.start() + + def shutdown_postproc_workers(self): + """Best-effort teardown of the local postproc pool (idempotent).""" + if self._postproc_input_queues: + for q in self._postproc_input_queues: + try: + q.put(None) # PostprocWorker mainloop exits on None + except Exception: + pass + if self._postproc_pool is not None: + self._postproc_pool.shutdown(wait=False) + self._postproc_pool = None def start_rpc_server(self): if self.rank == 0: diff --git a/tests/unittest/_torch/ray_orchestrator/single_gpu/test_postproc_workers.py b/tests/unittest/_torch/ray_orchestrator/single_gpu/test_postproc_workers.py new file mode 100644 index 000000000000..5ead6295813b --- /dev/null +++ b/tests/unittest/_torch/ray_orchestrator/single_gpu/test_postproc_workers.py @@ -0,0 +1,94 @@ +"""Postprocess parallelism under the Ray orchestrator. + +The rank-0 RayGPUWorker spawns a local PostprocWorker pool; finished +``PostprocWorker.Output`` records re-enter the RPC response stream via the +collector thread (RpcWorkerMixin.init_postproc_workers). These tests pin the +two user-visible contracts: + + 1. workers>0 produces token-for-token identical text to the inline + (workers=0) path, sync and streaming; + 2. the stream/proxy demux finalizes Output records (no leaked results, + no hang on completion). +""" + +import asyncio +import os + +from utils.llm_data import llm_models_root + +from tensorrt_llm import LLM +from tensorrt_llm.llmapi import KvCacheConfig, SamplingParams + +PROMPTS = [ + "Hello, my name is", + "The president of the United States is", + "The capital of France is", + "The future of AI is", +] + + +def _model_path() -> str: + override = os.environ.get("POSTPROC_TEST_MODEL") + if override: + return override + return str(llm_models_root() / "llama-models-v2/TinyLlama-1.1B-Chat-v1.0") + + +def _make_llm(num_postprocess_workers: int) -> LLM: + model = _model_path() + extra = {} + if num_postprocess_workers > 0: + extra = dict( + num_postprocess_workers=num_postprocess_workers, postprocess_tokenizer_dir=model + ) + return LLM( + model=model, + kv_cache_config=KvCacheConfig(enable_block_reuse=False, max_tokens=16384), + **extra, + ) + + +def _generate_sync(num_postprocess_workers: int) -> list[str]: + sampling_params = SamplingParams(temperature=0, max_tokens=32) + with _make_llm(num_postprocess_workers) as llm: + outputs = llm.generate(PROMPTS, sampling_params) + return [output.outputs[0].text for output in outputs] + + +def test_postproc_workers_match_inline(): + """Greedy outputs with workers=2 must equal the inline workers=0 path.""" + baseline = _generate_sync(0) + assert all(text for text in baseline) + with_postproc = _generate_sync(2) + assert with_postproc == baseline + + +def test_streaming_control_no_postproc(): + """Control: streaming with workers=0 must work (isolates any pre-existing + ray-streaming issue from the postproc path).""" + _streaming_body(num_postprocess_workers=0) + + +def test_postproc_workers_streaming(): + """Streaming rides the same Output records; final text must match sync.""" + _streaming_body(num_postprocess_workers=2) + + +def _streaming_body(num_postprocess_workers: int): + sampling_params = SamplingParams(temperature=0, max_tokens=32) + + with _make_llm(num_postprocess_workers) as llm: + sync_texts = [output.outputs[0].text for output in llm.generate(PROMPTS, sampling_params)] + + async def collect(prompt: str) -> str: + final = None + async for output in llm.generate_async(prompt, sampling_params, streaming=True): + final = output + return final.outputs[0].text + + async def run_all() -> list[str]: + return await asyncio.gather(*[collect(p) for p in PROMPTS]) + + streamed_texts = asyncio.run(run_all()) + + assert streamed_texts == sync_texts