[TRTLLM-8922][chore] Remove hang detector from py_cache_transceiver_mp test

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
@@ -2874,7 +2876,6 @@ def _prepare_disagg_gen_transmission_complete(self, scheduled_batch):
                 req.decoding_iter = 1
                 req.py_decoding_iter = 1
                 req.py_kv_transfer_start_time = None
-                req.decoding_iter = 1
                 first_gen_tokens = req.context_phase_params.first_gen_tokens
                 ctx_draft_tokens = req.context_phase_params.draft_tokens
                 req.py_draft_tokens = [] if ctx_draft_tokens is None else ctx_draft_tokens
@@ -2944,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/commands/serve.py

@@ -24,8 +24,8 @@
 from tensorrt_llm.executor.utils import LlmLauncherEnvs
 from tensorrt_llm.inputs.multimodal import MultimodalServerConfig
 from tensorrt_llm.llmapi import (BuildConfig, CapacitySchedulerPolicy,
-                                 DisaggScheduleStyle, DynamicBatchConfig,
-                                 KvCacheConfig, SchedulerConfig, VisualGen)
+                                 DynamicBatchConfig, KvCacheConfig,
+                                 SchedulerConfig, VisualGen)
 from tensorrt_llm.llmapi.disagg_utils import (DisaggClusterConfig,
                                               MetadataServerConfig, ServerRole,
                                               extract_disagg_cluster_config,
@@ -1011,7 +1011,8 @@ def serve_encoder(model: str, host: str, port: int, log_level: str,
               help="The logging level.")
 @click.option("-s",
               "--schedule_style",
-              type=None,
+              type=click.Choice(["context_first", "generation_first"],
+                                case_sensitive=False),
               default=None,
               help="The schedule style for the disaggregated server.")
 @click.option(
@@ -1041,13 +1042,6 @@ def disaggregated(
 
     disagg_cfg = parse_disagg_config_file(config_file)
     if schedule_style:
-        valid_styles = [
-            key.lower() for key in DisaggScheduleStyle.__members__.keys()
-        ]
-        if schedule_style not in valid_styles:
-            raise ValueError(
-                f"Invalid schedule style: {schedule_style}, options: {valid_styles}"
-            )
         disagg_cfg.schedule_style = schedule_style
     with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
         try:

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/llmapi/disagg_utils.py

@@ -178,7 +178,7 @@ def extract_disagg_cfg(hostname: str = 'localhost',
                                 conditional_disagg_config, otlp_config,
                                 max_retries, perf_metrics_max_requests,
                                 disagg_cluster_config)
-    if node_id:
+    if node_id is not None:
         config.node_id = node_id
     if schedule_style:
         config.schedule_style = schedule_style

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,40 +405,43 @@ async def _consume_gen():
                     await queue.put(e)
                 await queue.put(None)  # sentinel
 
-            asyncio.create_task(_consume_gen())
+            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:
             # Non-streaming or no ctx needed: both HTTP POSTs fire eagerly
             # through generator consumption, so asyncio.gather works fine.
             tasks = []
             if need_ctx:
-                async def request_ctx():
-                    response = await self._ctx_client.send_request(
-                        ctx_req, server=ctx_server, hooks=hooks
+                tasks.append(
+                    asyncio.create_task(
+                        self._ctx_client.send_request(ctx_req, server=ctx_server, hooks=hooks)
                     )
-                    return response
-
-                tasks.append(asyncio.create_task(request_ctx()))
-
-            async def request_gen():
-                response = await self._gen_client.send_request(
-                    gen_req, server=gen_server, hooks=hooks
                 )
-                return response
-
-            tasks.append(asyncio.create_task(request_gen()))
+            tasks.append(
+                asyncio.create_task(
+                    self._gen_client.send_request(gen_req, server=gen_server, hooks=hooks)
+                )
+            )
             responses = await asyncio.gather(*tasks)
             return responses[-1]

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_configs/disagg_config_overlap_gen_first.yaml

@@ -0,0 +1,39 @@
+model: TinyLlama/TinyLlama-1.1B-Chat-v1.0
+hostname: localhost
+port: 8000
+backend: "pytorch"
+cuda_graph_config: null
+free_gpu_memory_fraction: 0.2
+context_servers:
+  num_instances: 1
+  max_batch_size: 8
+  max_num_tokens: 3000
+  max_seq_len: 4096
+  tensor_parallel_size: 1
+  pipeline_parallel_size: 1
+  kv_cache_config:
+    enable_block_reuse: False
+    free_gpu_memory_fraction: 0.2
+    enable_partial_reuse: False
+
+  cache_transceiver_config:
+    backend: DEFAULT
+    transceiver_runtime: PYTHON
+  urls:
+      - "localhost:8001"
+generation_servers:
+  num_instances: 1
+  tensor_parallel_size: 1
+  pipeline_parallel_size: 1
+  max_batch_size: 256
+  max_num_tokens: 4096
+  max_seq_len: 4096
+  kv_cache_config:
+    enable_block_reuse: False
+    free_gpu_memory_fraction: 0.2
+    enable_partial_reuse: False
+  cache_transceiver_config:
+    backend: DEFAULT
+    transceiver_runtime: PYTHON
+  urls:
+      - "localhost:8002"

tests/integration/defs/disaggregated/test_configs/disagg_config_overlap_gen_first_pp4.yaml

@@ -0,0 +1,39 @@
+model: TinyLlama/TinyLlama-1.1B-Chat-v1.0
+hostname: localhost
+port: 8000
+backend: "pytorch"
+cuda_graph_config: null
+free_gpu_memory_fraction: 0.2
+context_servers:
+  num_instances: 1
+  max_batch_size: 8
+  max_num_tokens: 3000
+  max_seq_len: 4096
+  tensor_parallel_size: 1
+  pipeline_parallel_size: 4
+  kv_cache_config:
+    enable_block_reuse: False
+    free_gpu_memory_fraction: 0.2
+    enable_partial_reuse: False
+
+  cache_transceiver_config:
+    backend: DEFAULT
+    transceiver_runtime: PYTHON
+  urls:
+      - "localhost:8001"
+generation_servers:
+  num_instances: 1
+  tensor_parallel_size: 1
+  pipeline_parallel_size: 1
+  max_batch_size: 256
+  max_num_tokens: 4096
+  max_seq_len: 4096
+  kv_cache_config:
+    enable_block_reuse: False
+    free_gpu_memory_fraction: 0.2
+    enable_partial_reuse: False
+  cache_transceiver_config:
+    backend: DEFAULT
+    transceiver_runtime: PYTHON
+  urls:
+      - "localhost:8002"

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,