[None][feat] Support num_postprocess_workers > 0 under the Ray Orchestration

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,14 +1476,51 @@ 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],
         postproc_batches: Optional[List[List["PostprocWorker.Input"]]] = None,
         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.

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.

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:

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)

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):

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:

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