[TRTLLM-8922][fix] Address review comments on gen-first disagg PR

- Fix all([]) returning True when no non-gen-first requests exist
- Catch RPCError instead of broad Exception in proxy
- Log only ctx_request_id instead of full disaggregated_params
- Cancel background consume task when stream generator stops early
- Fix parametrize indentation lint in test
- Include overlap_gen_first* in overlap client test validation

Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>

Author: reasonsolo

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -2728,12 +2728,14 @@ def _check_disagg_gen_transfer_status(self):
             req.is_disagg_generation_transmission_in_progress
             for req in self.active_requests
         ])
-        need_check_one = all([
+        non_gen_first_reqs = [
+            req for req in self.active_requests
+            if req.py_disaggregated_params and req.py_disaggregated_params.
+            schedule_style != DisaggScheduleStyle.GENERATION_FIRST
+        ]
+        need_check_one = bool(non_gen_first_reqs) and all(
             req.is_disagg_generation_transmission_in_progress
-            for req in self.active_requests
-            if req.py_disaggregated_params \
-                and req.py_disaggregated_params.schedule_style != DisaggScheduleStyle.GENERATION_FIRST
-        ])
+            for req in non_gen_first_reqs)
 
         if need_check:
             at_least_num = 1 if need_check_one else 0
@@ -2943,12 +2945,14 @@ def _recv_disagg_gen_cache(self, new_gen_reqs):
                 if req.state == LlmRequestState.DISAGG_GENERATION_TRANS_IN_PROGRESS:
                     req.py_kv_transfer_start_time = time.time()
 
-        block_transfer = all([
+        non_gen_first_active = [
+            req for req in self.active_requests
+            if req.py_disaggregated_params and req.py_disaggregated_params.
+            schedule_style != DisaggScheduleStyle.GENERATION_FIRST
+        ]
+        block_transfer = bool(non_gen_first_active) and all(
             req.is_disagg_generation_transmission_in_progress
-            and req.py_disaggregated_params.schedule_style
-            != DisaggScheduleStyle.GENERATION_FIRST
-            for req in self.active_requests
-        ])
+            for req in non_gen_first_active)
         self._check_disagg_gen_cache_transfer_status(1 if block_transfer else 0)
 
         return

tensorrt_llm/executor/proxy.py

@@ -24,7 +24,7 @@
 from .request import CancellingRequest, GenerationRequest
 from .result import GenerationResult, IterationResult
 from .rpc import RPCClient
-from .rpc.rpc_common import get_unique_ipc_addr
+from .rpc.rpc_common import RPCError, get_unique_ipc_addr
 from .utils import (ErrorResponse, WorkerCommIpcAddrs, create_mpi_comm_session,
                     get_spawn_proxy_process_env, is_llm_response,
                     print_alive_threads)
@@ -396,7 +396,7 @@ def get_disaggregated_params(self) -> dict:
         try:
             params = self.rpc_client.get_disaggregated_params().remote()
             return params if isinstance(params, dict) else {}
-        except Exception as e:
+        except RPCError as e:
             logger.warning(f"Error fetching disaggregated params via RPC: {e}")
             return {}
 

tensorrt_llm/serve/openai_client.py

@@ -127,7 +127,9 @@ async def _send_request(
         if server is None:
             server, _ = await self._router.get_next_server(request)
         url = f"http://{server}/{endpoint}"
-        logger.debug(f"Sending {self._role} request {request.disaggregated_params} to {url}")
+        logger.debug(
+            f"Sending {self._role} request {request.disaggregated_params.ctx_request_id} to {url}"
+        )
         try:
             self._metrics_collector.total_requests.inc()
             resp_generator = self._post_with_retry(server, url, request, hooks)

tensorrt_llm/serve/openai_disagg_service.py

@@ -405,19 +405,27 @@ async def _consume_gen():
                     await queue.put(e)
                 await queue.put(None)  # sentinel
 
-            consume_task = asyncio.create_task(_consume_gen())  # noqa: F841 prevent GC
+            consume_task: asyncio.Task = asyncio.create_task(_consume_gen())
 
             # Now send ctx request — gen server has received its request
             await self._ctx_client.send_request(ctx_req, server=ctx_server, hooks=hooks)
 
             async def _yield_from_queue():
-                while True:
-                    item = await queue.get()
-                    if item is None:
-                        break
-                    if isinstance(item, Exception):
-                        raise item
-                    yield item
+                try:
+                    while True:
+                        item = await queue.get()
+                        if item is None:
+                            break
+                        if isinstance(item, Exception):
+                            raise item
+                        yield item
+                finally:
+                    if not consume_task.done():
+                        consume_task.cancel()
+                    try:
+                        await consume_task
+                    except asyncio.CancelledError:
+                        pass
 
             return _yield_from_queue()
         else:

tests/integration/defs/accuracy/test_disaggregated_serving.py

@@ -1616,12 +1616,13 @@ def test_auto_dtype_with_helix(self, comms_medium, cuda_graph_config,
     @pytest.mark.parametrize(
         "gen_tp_pp", [(1, 1), (1, 2), (2, 1), (2, 2)],
         ids=["gen_tp1pp1", "gen_tp1pp2", "gen_tp2pp1", "gen_tp2pp2"])
-    @pytest.mark.parametrize("ctx_tp_pp", [(1, 1), (1, 2), (2, 1), (2, 2),
-                                           (1, 4)],
-                             ids=[
-                                 "ctx_tp1pp1", "ctx_tp1pp2", "ctx_tp2pp1",
-                                 "ctx_tp2pp2", "ctx_tp1pp4"
-                             ])
+    @pytest.mark.parametrize(
+        "ctx_tp_pp",
+        [(1, 1), (1, 2), (2, 1), (2, 2), (1, 4)],
+        ids=[
+            "ctx_tp1pp1", "ctx_tp1pp2", "ctx_tp2pp1", "ctx_tp2pp2", "ctx_tp1pp4"
+        ],
+    )
     def test_gen_first(self, ctx_tp_pp, gen_tp_pp):
         ctx_tp, ctx_pp = ctx_tp_pp
         gen_tp, gen_pp = gen_tp_pp

tests/integration/defs/disaggregated/test_disaggregated.py

@@ -284,7 +284,7 @@ def get_client_test_set(test_desc):
                              verify_streaming_completion=True,
                              verify_chat=False,
                              verify_streaming_chat=False)
-    if test_desc in ("overlap", "trtllm_sampler"):
+    if test_desc.startswith("overlap") or test_desc == "trtllm_sampler":
         return ClientTestSet(completion=True,
                              completion_streaming=True,
                              chat=True,