From 5c9a7ecff89e7921f67324f7cc9b66f9e46e0a38 Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Sun, 26 Jul 2026 10:33:52 -0700 Subject: [PATCH] [https://nvbugs/6480621][feat] disagg serving lifecycle fixes and orchestrator diagnostics MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three pull-forward fixes from the Fable disaggregated lifecycle design §13, plus a serve-layer orchestrator lifecycle emitter (WS1 serve layer). **Fix 1 — Stop retrying HTTP 4xx/5xx as network errors** (`openai_client.py`) `aiohttp.ClientResponseError` (raised for any HTTP error response) is a subclass of `aiohttp.ClientError`. Previously all `ClientError` exceptions were retried, which meant a GEN worker returning 503 (over capacity) or 422 (bad request) consumed the full retry budget against the same server instead of surfacing the error immediately. Now `ClientResponseError` is re-raised on the first attempt; only true TCP-level errors continue to be retried. **Fix 2 — Cancel surviving leg in gen-first non-streaming gather** (`openai_disagg_service.py`) In the `_send_disagg_request_gen_first` non-streaming path both legs were launched with `asyncio.gather`. If one failed, the surviving leg kept running to completion, holding a KV-receive session or transfer-block slot until its own timeout. The gather is now wrapped in a try/except that cancels all pending tasks on the first exception. **Fix 3 — Disconnect watcher for non-streaming disagg requests** (`openai_disagg_server.py`) The non-streaming disagg handler had no mechanism to detect client disconnect. `_run_with_disconnect_guard` races the handler coroutine against `_poll_disconnect` (polling `raw_request.is_disconnected()` at 0.5 s intervals) and cancels the handler on disconnect. The cancelled request is converted to an HTTP 499 response; `perf_metrics.DisaggPerfMetricsCollector` gains a `cancelled_requests` counter. Workers are not explicitly cancelled here (that requires WS4 abort RPCs — a later workstream). `tensorrt_llm/serve/disagg_lifecycle.py` adds a lightweight, opt-in per-process lifecycle event stream for the disagg orchestrator, enabled by `TRTLLM_DISAGG_ORCHESTRATOR_DIAGNOSTICS=1`. Events: `ctx_dispatch`, `ctx_complete`, `ctx_error`, `gen_dispatch`, `gen_complete`, `gen_error`, `gen_rejected` (HTTP 4xx/5xx), `client_disconnect`, `abort`. Each record is a JSON line on stdout with `(clock_id, seq, wall_ns, rid, elapsed_ms)` for cross-host correlation. `OrchestratorRequestTracer` (obtained via `DisaggOrchestratorLifecycle.tracer()`) provides paired dispatch/complete helpers that measure leg-level elapsed time from the orchestrator's perspective. Both ctx-first and gen-first flows are instrumented. Correlated with executor-layer events (`TRTLLM_DISAGG_TRANSFER_DIAGNOSTICS`) via `disagg_request_id`. - `test_disagg_openai_client.py`: verifies that `ClientResponseError` (5 status codes: 400, 422, 429, 500, 503) is not retried and that plain `ClientError` still is. - `test_openai_disagg_service.py`: two new tests verify that the surviving task is cancelled when CTX fails and when GEN fails in the gen-first non-streaming gather. - `test_disagg_lifecycle.py` (new): 16 tests covering disabled/enabled emitter, field presence, sequence increment, error truncation, tracer paired events, and `from_environment` with multiple env-var values. Signed-off-by: Chien-Chun Hung Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- tensorrt_llm/serve/disagg_lifecycle.py | 251 ++++++++++++++++++ tensorrt_llm/serve/openai_client.py | 8 + tensorrt_llm/serve/openai_disagg_server.py | 56 +++- tensorrt_llm/serve/openai_disagg_service.py | 89 ++++++- tensorrt_llm/serve/perf_metrics.py | 3 + .../disaggregated/test_disagg_lifecycle.py | 194 ++++++++++++++ .../test_disagg_openai_client.py | 49 ++++ .../test_openai_disagg_service.py | 76 ++++++ 8 files changed, 713 insertions(+), 13 deletions(-) create mode 100644 tensorrt_llm/serve/disagg_lifecycle.py create mode 100644 tests/unittest/disaggregated/test_disagg_lifecycle.py diff --git a/tensorrt_llm/serve/disagg_lifecycle.py b/tensorrt_llm/serve/disagg_lifecycle.py new file mode 100644 index 000000000000..4a3844d6cbce --- /dev/null +++ b/tensorrt_llm/serve/disagg_lifecycle.py @@ -0,0 +1,251 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Opt-in orchestrator-level lifecycle event stream for disaggregated serving. + +Each record is a single newline-terminated JSON line written to stdout. The +stream is *observational*: it does not coordinate request state across context +and generation workers. Enable with:: + + TRTLLM_DISAGG_ORCHESTRATOR_DIAGNOSTICS = 1 + +Correlated with executor-layer events (``TRTLLM_DISAGG_TRANSFER_DIAGNOSTICS``) +via ``disagg_request_id`` and, once WS3 attempt fencing lands, via +``attempt_id``. + +Timer semantics (Fable lifecycle design §5): + +* ``ctx_dispatch`` → ``ctx_complete``: measures full CTX leg including network + and prefill. The gap between ``ctx_complete`` and ``gen_dispatch`` is the + orchestrator's own scheduling overhead. +* ``gen_dispatch`` → ``gen_complete``: measures the GEN leg. Note that the + active-transfer clock should start at ``receiver_ready`` (WS2, not yet + implemented here), *not* at ``gen_dispatch`` — this field reflects the + orchestrator's wall-clock observation, which includes GEN queueing delay. +* On disconnect or abort the ``abort`` record carries the reason and the + elapsed wall time so consumers can distinguish queueing delay from transfer + stalls. + +Records may arrive out of order in high-concurrency deployments. Consumers +should order by ``(wall_ns, seq)`` within a single ``clock_id``. +""" + +from __future__ import annotations + +import json +import os +import secrets +import threading +import time +from enum import Enum +from typing import Any, Dict, Optional + +_ENV_VAR = "TRTLLM_DISAGG_ORCHESTRATOR_DIAGNOSTICS" +_LOG_PREFIX = "[DISAGG_DIAG][orchestrator]" +_SCHEMA_VERSION = 1 + + +class OrchestratorEvent(str, Enum): + """Lifecycle events emitted by the disagg orchestrator.""" + + CTX_DISPATCH = "ctx_dispatch" + CTX_COMPLETE = "ctx_complete" + CTX_ERROR = "ctx_error" + GEN_DISPATCH = "gen_dispatch" + GEN_COMPLETE = "gen_complete" + GEN_ERROR = "gen_error" + # GEN worker returned an HTTP 4xx/5xx — maps to GEN_REJECT in §4 + GEN_REJECTED = "gen_rejected" + # Client disconnect detected before both legs completed + CLIENT_DISCONNECT = "client_disconnect" + # Abort triggered by any leg error, timeout, or disconnect + ABORT = "abort" + + +class OrchestratorScheduleStyle(str, Enum): + CONTEXT_FIRST = "context_first" + GENERATION_FIRST = "generation_first" + + +class DisaggOrchestratorLifecycle: + """Lightweight per-process lifecycle emitter for the disagg orchestrator. + + Thread-safe; individual ``emit`` calls take a sequence number under a lock + but release it before formatting and printing so the hot path is not + serialised end-to-end. + """ + + def __init__(self, *, enabled: bool, node_id: Optional[str] = None): + self.enabled = bool(enabled) + self._node_id = str(node_id) if node_id is not None else "-" + self._clock_id = f"orch-{os.getpid()}-{secrets.token_hex(6)}" if enabled else "disabled" + self._seq = 0 + self._lock = threading.Lock() + + @classmethod + def from_environment(cls, *, node_id: Optional[str] = None) -> "DisaggOrchestratorLifecycle": + enabled = os.environ.get(_ENV_VAR, "0") not in ("0", "", "false", "False", "FALSE") + return cls(enabled=enabled, node_id=node_id) + + def tracer( + self, + *, + disagg_request_id: Optional[int] = None, + schedule_style: Optional[OrchestratorScheduleStyle] = None, + ) -> "OrchestratorRequestTracer": + return OrchestratorRequestTracer( + self, disagg_request_id=disagg_request_id, schedule_style=schedule_style + ) + + def _next_seq(self) -> int: + with self._lock: + self._seq += 1 + return self._seq + + def emit( + self, + event: OrchestratorEvent, + *, + disagg_request_id: Optional[int] = None, + ctx_server: Optional[str] = None, + gen_server: Optional[str] = None, + schedule_style: Optional[OrchestratorScheduleStyle] = None, + elapsed_ms: Optional[float] = None, + http_status: Optional[int] = None, + error: Optional[str] = None, + extra: Optional[Dict[str, Any]] = None, + ) -> None: + """Emit one lifecycle record. A no-op when disabled.""" + if not self.enabled: + return + seq = self._next_seq() + wall_ns = time.time_ns() + record: Dict[str, Any] = { + "v": _SCHEMA_VERSION, + "clock_id": self._clock_id, + "seq": seq, + "wall_ns": wall_ns, + "node": self._node_id, + "event": event.value, + } + if disagg_request_id is not None: + record["rid"] = disagg_request_id + if ctx_server is not None: + record["ctx"] = ctx_server + if gen_server is not None: + record["gen"] = gen_server + if schedule_style is not None: + record["sched"] = schedule_style.value + if elapsed_ms is not None: + record["elapsed_ms"] = round(elapsed_ms, 3) + if http_status is not None: + record["http_status"] = http_status + if error is not None: + record["error"] = error[:256] + if extra: + record.update(extra) + try: + print(f"{_LOG_PREFIX} {json.dumps(record, separators=(',', ':'))}", flush=True) + except Exception: + pass + + +class OrchestratorRequestTracer: + """Per-request helper that snapshots timestamps and emits paired events. + + Usage:: + + tracer = lifecycle.tracer(disagg_request_id=rid, ...) + tracer.ctx_dispatch(ctx_server) + ... + tracer.ctx_complete() + tracer.gen_dispatch(gen_server) + ... + tracer.gen_complete() + + All times are wall-clock milliseconds elapsed since the tracer was created. + """ + + def __init__( + self, + emitter: DisaggOrchestratorLifecycle, + *, + disagg_request_id: Optional[int] = None, + schedule_style: Optional[OrchestratorScheduleStyle] = None, + ): + self._emitter = emitter + self._rid = disagg_request_id + self._style = schedule_style + self._t0_ns = time.monotonic_ns() + self._ctx_dispatch_ns: Optional[int] = None + self._gen_dispatch_ns: Optional[int] = None + + def _elapsed_ms(self, since_ns: Optional[int] = None) -> float: + now = time.monotonic_ns() + ref = since_ns if since_ns is not None else self._t0_ns + return (now - ref) / 1e6 + + def _emit(self, event: OrchestratorEvent, **kw) -> None: + self._emitter.emit( + event, + disagg_request_id=self._rid, + schedule_style=self._style, + elapsed_ms=self._elapsed_ms(), + **kw, + ) + + def ctx_dispatch(self, server: str) -> None: + self._ctx_dispatch_ns = time.monotonic_ns() + self._emit(OrchestratorEvent.CTX_DISPATCH, ctx_server=server) + + def ctx_complete(self, ctx_server: str = "") -> None: + since = self._ctx_dispatch_ns + elapsed = self._elapsed_ms(since) if since is not None else None + self._emitter.emit( + OrchestratorEvent.CTX_COMPLETE, + disagg_request_id=self._rid, + schedule_style=self._style, + elapsed_ms=elapsed, + ctx_server=ctx_server or None, + ) + + def ctx_error( + self, error: str, ctx_server: str = "", http_status: Optional[int] = None + ) -> None: + self._emit( + OrchestratorEvent.CTX_ERROR, + ctx_server=ctx_server or None, + http_status=http_status, + error=error, + ) + + def gen_dispatch(self, server: str) -> None: + self._gen_dispatch_ns = time.monotonic_ns() + self._emit(OrchestratorEvent.GEN_DISPATCH, gen_server=server) + + def gen_complete(self, gen_server: str = "") -> None: + since = self._gen_dispatch_ns + elapsed = self._elapsed_ms(since) if since is not None else None + self._emitter.emit( + OrchestratorEvent.GEN_COMPLETE, + disagg_request_id=self._rid, + schedule_style=self._style, + elapsed_ms=elapsed, + gen_server=gen_server or None, + ) + + def gen_error( + self, error: str, gen_server: str = "", http_status: Optional[int] = None + ) -> None: + event = ( + OrchestratorEvent.GEN_REJECTED + if http_status is not None and http_status >= 400 + else OrchestratorEvent.GEN_ERROR + ) + self._emit(event, gen_server=gen_server or None, http_status=http_status, error=error) + + def client_disconnect(self) -> None: + self._emit(OrchestratorEvent.CLIENT_DISCONNECT) + + def abort(self, reason: str) -> None: + self._emit(OrchestratorEvent.ABORT, error=reason) diff --git a/tensorrt_llm/serve/openai_client.py b/tensorrt_llm/serve/openai_client.py index 26646a54344c..504de43301bf 100644 --- a/tensorrt_llm/serve/openai_client.py +++ b/tensorrt_llm/serve/openai_client.py @@ -269,6 +269,14 @@ async def _post_with_retry( ) break # break and skip retries if the whole response is processed without exception except (aiohttp.ClientError, OSError) as e: + # HTTP 4xx/5xx from a worker is a deterministic rejection — do + # not retry as a transient network error. In disaggregated + # serving a GEN worker rejection (e.g. 503 over capacity) must + # surface immediately so the orchestrator can reroute or fail + # fast rather than consuming retry budget against the same + # server. (Fable lifecycle design §13 pull-forward fix.) + if isinstance(e, aiohttp.ClientResponseError): + raise if lines_yielded > 0: logger.error( f"Client error to {url}: {e} - cannot retry since {lines_yielded} lines were yielded", diff --git a/tensorrt_llm/serve/openai_disagg_server.py b/tensorrt_llm/serve/openai_disagg_server.py index 5369c0cb6dbf..2bf0728c2296 100644 --- a/tensorrt_llm/serve/openai_disagg_server.py +++ b/tensorrt_llm/serve/openai_disagg_server.py @@ -57,6 +57,43 @@ code: f"\\x{code:02x}" for code in (*range(32), 127) } +_DISCONNECT_POLL_INTERVAL_SECS = 0.5 + + +async def _poll_disconnect(raw_req: Request) -> None: + """Return when the client has disconnected (detected via ASGI receive).""" + while not await raw_req.is_disconnected(): + await asyncio.sleep(_DISCONNECT_POLL_INTERVAL_SECS) + + +async def _run_with_disconnect_guard(coro, raw_req: Request): + """Await *coro* and cancel it when the HTTP client disconnects. + + Raises ``asyncio.CancelledError`` when the disconnect fires before the + coroutine completes so the caller can convert it to an HTTP 499 response. + Workers are not explicitly cancelled here (that requires WS4 abort RPCs); + this guard ensures the orchestrator stops consuming resources and releases + router load as soon as the client is gone. + """ + main_task = asyncio.ensure_future(coro) + disconnect_task = asyncio.ensure_future(_poll_disconnect(raw_req)) + try: + done, pending = await asyncio.wait( + {main_task, disconnect_task}, + return_when=asyncio.FIRST_COMPLETED, + ) + for t in pending: + t.cancel() + await asyncio.gather(*pending, return_exceptions=True) + if disconnect_task in done and main_task not in done: + logger.info("Disagg orchestrator: client disconnected, aborting request") + raise asyncio.CancelledError("client disconnected") + return main_task.result() + except asyncio.CancelledError: + main_task.cancel() + disconnect_task.cancel() + await asyncio.gather(main_task, disconnect_task, return_exceptions=True) + raise class RawRequestResponseHooks(ResponseHooks): def __init__(self, raw_req: Request, perf_metrics_collector: DisaggPerfMetricsCollector): @@ -287,12 +324,25 @@ async def wrapper(req: request_type, raw_req: Request) -> Response: raise HTTPException(status_code=400, detail=str(e)) from e self._extract_conversation_id(req, raw_req) hooks = RawRequestResponseHooks(raw_req, self._perf_metrics_collector) - response_or_generator = await entry_point(req, hooks) - self._perf_metrics_collector.total_responses.inc() if req.stream: - return StreamingResponse(content=response_or_generator, media_type="text/event-stream") + response_or_generator = await entry_point(req, hooks) + self._perf_metrics_collector.total_responses.inc() + return StreamingResponse( + content=response_or_generator, media_type="text/event-stream" + ) else: + # Non-streaming: race the disagg handler against client + # disconnect so an abandoned request does not hold a CTX KV + # pin or a GEN transfer-block slot after the client is gone. + # (Fable lifecycle design §13 pull-forward fix.) + response_or_generator = await _run_with_disconnect_guard( + entry_point(req, hooks), raw_req + ) + self._perf_metrics_collector.total_responses.inc() return JSONResponse(content=response_or_generator.model_dump()) + except asyncio.CancelledError: + self._perf_metrics_collector.cancelled_requests.inc() + raise HTTPException(status_code=499, detail="Client disconnected") except Exception as e: self._handle_exception(e) return wrapper diff --git a/tensorrt_llm/serve/openai_disagg_service.py b/tensorrt_llm/serve/openai_disagg_service.py index e2ee00257e98..4dee7fb56bfa 100644 --- a/tensorrt_llm/serve/openai_disagg_service.py +++ b/tensorrt_llm/serve/openai_disagg_service.py @@ -16,9 +16,16 @@ import os from typing import Callable, Optional +import aiohttp + from tensorrt_llm.llmapi.disagg_utils import ConditionalDisaggConfig, DisaggServerConfig, ServerRole from tensorrt_llm.logger import logger from tensorrt_llm.serve.disagg_coordinator import DisaggCoordinator +from tensorrt_llm.serve.disagg_lifecycle import ( + DisaggOrchestratorLifecycle, + OrchestratorEvent, + OrchestratorScheduleStyle, +) from tensorrt_llm.serve.openai_client import OpenAIClient from tensorrt_llm.serve.openai_protocol import ( ChatCompletionRequest, @@ -70,6 +77,11 @@ def __init__( self._gen_client = None self._schedule_style = DisaggScheduleStyle.CONTEXT_FIRST + # Opt-in orchestrator lifecycle emitter (TRTLLM_DISAGG_ORCHESTRATOR_DIAGNOSTICS=1). + self._lifecycle = DisaggOrchestratorLifecycle.from_environment( + node_id=str(config.node_id) if config.node_id is not None else None + ) + match self._config.schedule_style: case "generation_first": self._send_disagg_request = self._send_disagg_request_gen_first @@ -104,14 +116,22 @@ async def openai_completion( "Disaggregated server currently only supports single string prompt or list of integers in request" ) - return await self._send_disagg_request(request, hooks) + try: + return await self._send_disagg_request(request, hooks) + except asyncio.CancelledError: + self._lifecycle.emit(OrchestratorEvent.CLIENT_DISCONNECT) + raise async def openai_chat_completion( self, request: UCompletionRequest, hooks: Optional[ResponseHooks] = None ) -> UCompletionResponseOrGenerator: if not await self.is_ready(): raise RuntimeError("Cluster is not ready") - return await self._send_disagg_request(request, hooks) + try: + return await self._send_disagg_request(request, hooks) + except asyncio.CancelledError: + self._lifecycle.emit(OrchestratorEvent.CLIENT_DISCONNECT) + raise async def _send_disagg_request_ctx_first( self, request: UCompletionRequest, hooks: Optional[ResponseHooks] = None @@ -132,6 +152,10 @@ async def _send_disagg_request_ctx_first( need_ctx = need_ctx and not await self._check_gen_only_disagg(request) ctx_response = None gen_req = request + tracer = self._lifecycle.tracer( + disagg_request_id=disagg_request_id, + schedule_style=OrchestratorScheduleStyle.CONTEXT_FIRST, + ) if need_ctx: try: # Mark ctx-dispatch start: arrival->here is the pre-ctx wait in the @@ -143,9 +167,20 @@ async def _send_disagg_request_ctx_first( ctx_server, _ = await self._ctx_router.get_next_server( ctx_req, exclude_server=gen_server, req_id=disagg_request_id ) - ctx_response = await self._ctx_client.send_request( - ctx_req, server=ctx_server, hooks=hooks, req_id=disagg_request_id - ) + tracer.ctx_dispatch(ctx_server or "") + try: + ctx_response = await self._ctx_client.send_request( + ctx_req, server=ctx_server, hooks=hooks, req_id=disagg_request_id + ) + except Exception as exc: + if isinstance(exc, aiohttp.ClientResponseError): + tracer.ctx_error( + str(exc), ctx_server=ctx_server or "", http_status=exc.status + ) + else: + tracer.ctx_error(str(exc), ctx_server=ctx_server or "") + raise + tracer.ctx_complete(ctx_server=ctx_server or "") await self._verify_ctx_response(ctx_response) ctx_response_disagg_params = ctx_response.choices[0].disaggregated_params if ctx_response_disagg_params.disagg_request_id is not None: @@ -173,9 +208,19 @@ async def _send_disagg_request_ctx_first( gen_req, exclude_server=ctx_server, req_id=disagg_request_id ) gen_reservation_id = disagg_request_id - gen_response = await self._gen_client.send_request( - gen_req, server=gen_server, hooks=hooks, req_id=gen_reservation_id - ) + tracer.gen_dispatch(gen_server or "") + try: + gen_response = await self._gen_client.send_request( + gen_req, server=gen_server, hooks=hooks, req_id=gen_reservation_id + ) + except Exception as exc: + if isinstance(exc, aiohttp.ClientResponseError): + tracer.gen_error(str(exc), gen_server=gen_server or "", http_status=exc.status) + else: + tracer.gen_error(str(exc), gen_server=gen_server or "") + raise + if not request.stream: + tracer.gen_complete(gen_server=gen_server or "") return gen_response else: if gen_server: @@ -392,6 +437,10 @@ async def _send_disagg_request_gen_first( # Single-issuer disagg id (see _send_disagg_request_ctx_first): fetch from # the coordinator so fleet workers never mint colliding ids. disagg_request_id = await self._coordinator.get_disagg_request_id() + tracer = self._lifecycle.tracer( + disagg_request_id=disagg_request_id, + schedule_style=OrchestratorScheduleStyle.GENERATION_FIRST, + ) if need_ctx: # arrival->here = pre-ctx wait in the orchestrator/fleet. if hooks: @@ -417,6 +466,7 @@ async def _send_disagg_request_gen_first( # # Fix: eagerly start consuming the gen generator in a background # task so the HTTP POST fires, then pipe chunks through a queue. + tracer.gen_dispatch(gen_server or "") gen_response = await self._gen_client.send_request( gen_req, server=gen_server, hooks=hooks, req_id=disagg_request_id ) @@ -434,6 +484,7 @@ async def _consume_gen(): consume_task: asyncio.Task = asyncio.create_task(_consume_gen()) # Now send ctx request — gen server has received its request + tracer.ctx_dispatch(ctx_server or "") try: await self._ctx_client.send_request( ctx_req, @@ -441,13 +492,15 @@ async def _consume_gen(): hooks=hooks, req_id=disagg_request_id, ) - except Exception: + except Exception as exc: + tracer.ctx_error(str(exc), ctx_server=ctx_server or "") consume_task.cancel() try: await consume_task except (asyncio.CancelledError, Exception): pass raise + tracer.ctx_complete(ctx_server=ctx_server or "") async def _yield_from_queue(): try: @@ -470,6 +523,9 @@ async def _yield_from_queue(): else: # Non-streaming or no ctx needed: both HTTP POSTs fire eagerly # through generator consumption, so asyncio.gather works fine. + if need_ctx: + tracer.ctx_dispatch(ctx_server or "") + tracer.gen_dispatch(gen_server or "") tasks = [] if need_ctx: tasks.append( @@ -492,5 +548,18 @@ async def _yield_from_queue(): ) ) ) - responses = await asyncio.gather(*tasks) + # Cancel the surviving leg if the other fails. Without this, a + # failed CTX leg leaves the GEN worker holding a KV-receive session + # until its own timeout, consuming transfer budget needlessly. + # (Fable lifecycle design §13 pull-forward fix.) + try: + responses = await asyncio.gather(*tasks) + except Exception as exc: + tracer.abort(str(exc)) + for task in tasks: + if not task.done(): + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + raise + tracer.gen_complete(gen_server=gen_server or "") return responses[-1] diff --git a/tensorrt_llm/serve/perf_metrics.py b/tensorrt_llm/serve/perf_metrics.py index 4eaed580ba8f..a171f041ba57 100644 --- a/tensorrt_llm/serve/perf_metrics.py +++ b/tensorrt_llm/serve/perf_metrics.py @@ -135,6 +135,9 @@ def __getattr__( MetricsDefinition("http_exceptions", "Total number of HTTP exceptions", "counter"), MetricsDefinition("internal_errors", "Total number of internal errors", "counter"), MetricsDefinition("total_responses", "Total number of responses", "counter"), + MetricsDefinition( + "cancelled_requests", "Total number of requests cancelled by client disconnect", "counter" + ), MetricsDefinition( "queue_latency_seconds", "Histogram of latency from request arrival to being processed in seconds", diff --git a/tests/unittest/disaggregated/test_disagg_lifecycle.py b/tests/unittest/disaggregated/test_disagg_lifecycle.py new file mode 100644 index 000000000000..6b448f708e95 --- /dev/null +++ b/tests/unittest/disaggregated/test_disagg_lifecycle.py @@ -0,0 +1,194 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Unit tests for the serve-layer orchestrator lifecycle emitter.""" + +import json + +import pytest + +from tensorrt_llm.serve.disagg_lifecycle import ( + DisaggOrchestratorLifecycle, + OrchestratorEvent, + OrchestratorScheduleStyle, +) + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_emitter(enabled: bool = True, node_id: str = "n0") -> DisaggOrchestratorLifecycle: + return DisaggOrchestratorLifecycle(enabled=enabled, node_id=node_id) + + +def _capture(capsys): + """Return parsed JSON records from stdout since last capture.""" + out = capsys.readouterr().out + records = [] + for line in out.splitlines(): + line = line.strip() + if line.startswith("[DISAGG_DIAG][orchestrator]"): + payload = line[len("[DISAGG_DIAG][orchestrator]") :].strip() + records.append(json.loads(payload)) + return records + + +# --------------------------------------------------------------------------- +# Disabled emitter +# --------------------------------------------------------------------------- + + +class TestDisabledEmitter: + def test_disabled_produces_no_output(self, capsys): + em = _make_emitter(enabled=False) + em.emit(OrchestratorEvent.CTX_DISPATCH, disagg_request_id=1) + records = _capture(capsys) + assert records == [] + + def test_disabled_tracer_is_noop(self, capsys): + em = _make_emitter(enabled=False) + tracer = em.tracer(disagg_request_id=42) + tracer.ctx_dispatch("ctx:8000") + tracer.ctx_complete() + tracer.gen_dispatch("gen:9000") + tracer.gen_complete() + records = _capture(capsys) + assert records == [] + + +# --------------------------------------------------------------------------- +# Enabled emitter — basic emit +# --------------------------------------------------------------------------- + + +class TestEnabledEmitter: + def test_emit_contains_required_fields(self, capsys): + em = _make_emitter() + em.emit(OrchestratorEvent.CTX_DISPATCH, disagg_request_id=7, ctx_server="h:8") + records = _capture(capsys) + assert len(records) == 1 + r = records[0] + assert r["event"] == "ctx_dispatch" + assert r["rid"] == 7 + assert r["ctx"] == "h:8" + assert "wall_ns" in r + assert "seq" in r + assert r["v"] == 1 + + def test_sequence_increments(self, capsys): + em = _make_emitter() + em.emit(OrchestratorEvent.CTX_DISPATCH) + em.emit(OrchestratorEvent.GEN_DISPATCH) + records = _capture(capsys) + assert records[1]["seq"] == records[0]["seq"] + 1 + + def test_optional_fields_absent_when_none(self, capsys): + em = _make_emitter() + em.emit(OrchestratorEvent.ABORT, error="oops") + records = _capture(capsys) + r = records[0] + assert "ctx" not in r + assert "gen" not in r + assert r["error"] == "oops" + + def test_http_status_included(self, capsys): + em = _make_emitter() + em.emit(OrchestratorEvent.GEN_REJECTED, http_status=503, error="no capacity") + records = _capture(capsys) + r = records[0] + assert r["event"] == "gen_rejected" + assert r["http_status"] == 503 + + def test_error_truncated_to_256(self, capsys): + em = _make_emitter() + long_error = "x" * 512 + em.emit(OrchestratorEvent.ABORT, error=long_error) + records = _capture(capsys) + assert len(records[0]["error"]) == 256 + + +# --------------------------------------------------------------------------- +# Tracer — paired events +# --------------------------------------------------------------------------- + + +class TestTracer: + def test_ctx_dispatch_complete(self, capsys): + em = _make_emitter() + tracer = em.tracer( + disagg_request_id=99, schedule_style=OrchestratorScheduleStyle.CONTEXT_FIRST + ) + tracer.ctx_dispatch("ctx:7000") + tracer.ctx_complete(ctx_server="ctx:7000") + records = _capture(capsys) + assert len(records) == 2 + assert records[0]["event"] == "ctx_dispatch" + assert records[1]["event"] == "ctx_complete" + assert records[0]["rid"] == 99 + assert records[1]["sched"] == "context_first" + # ctx_complete elapsed_ms is measured from ctx_dispatch, not tracer start + assert records[1]["elapsed_ms"] is not None + assert records[1]["elapsed_ms"] >= 0 + + def test_gen_error_with_http_status(self, capsys): + em = _make_emitter() + tracer = em.tracer(disagg_request_id=5) + tracer.gen_error("worker unavailable", gen_server="gen:9000", http_status=503) + records = _capture(capsys) + assert len(records) == 1 + r = records[0] + assert r["event"] == "gen_rejected" + assert r["http_status"] == 503 + assert r["gen"] == "gen:9000" + + def test_gen_error_without_http_status(self, capsys): + em = _make_emitter() + tracer = em.tracer() + tracer.gen_error("transport error") + records = _capture(capsys) + assert records[0]["event"] == "gen_error" + assert "http_status" not in records[0] + + def test_abort(self, capsys): + em = _make_emitter() + tracer = em.tracer(disagg_request_id=1) + tracer.abort("client disconnected") + records = _capture(capsys) + assert records[0]["event"] == "abort" + assert "disconnected" in records[0]["error"] + + def test_client_disconnect(self, capsys): + em = _make_emitter() + tracer = em.tracer() + tracer.client_disconnect() + records = _capture(capsys) + assert records[0]["event"] == "client_disconnect" + + +# --------------------------------------------------------------------------- +# from_environment +# --------------------------------------------------------------------------- + + +class TestFromEnvironment: + def test_disabled_by_default(self, monkeypatch, capsys): + monkeypatch.delenv("TRTLLM_DISAGG_ORCHESTRATOR_DIAGNOSTICS", raising=False) + em = DisaggOrchestratorLifecycle.from_environment() + assert not em.enabled + em.emit(OrchestratorEvent.CTX_DISPATCH) + assert _capture(capsys) == [] + + @pytest.mark.parametrize("val", ["1", "true", "True", "TRUE"]) + def test_enabled_by_env(self, monkeypatch, val, capsys): + monkeypatch.setenv("TRTLLM_DISAGG_ORCHESTRATOR_DIAGNOSTICS", val) + em = DisaggOrchestratorLifecycle.from_environment() + assert em.enabled + em.emit(OrchestratorEvent.CTX_DISPATCH) + assert len(_capture(capsys)) == 1 + + @pytest.mark.parametrize("val", ["0", "", "false", "False", "FALSE"]) + def test_disabled_by_env(self, monkeypatch, val, capsys): + monkeypatch.setenv("TRTLLM_DISAGG_ORCHESTRATOR_DIAGNOSTICS", val) + em = DisaggOrchestratorLifecycle.from_environment() + assert not em.enabled diff --git a/tests/unittest/disaggregated/test_disagg_openai_client.py b/tests/unittest/disaggregated/test_disagg_openai_client.py index cead306ae3ff..6e9458cfac0e 100644 --- a/tests/unittest/disaggregated/test_disagg_openai_client.py +++ b/tests/unittest/disaggregated/test_disagg_openai_client.py @@ -277,6 +277,55 @@ async def test_invalid_request_type(self, openai_client): with pytest.raises(ValueError, match="Invalid request type"): await openai_client.send_request("invalid_request") + # ------------------------------------------------------------------ + # Fable lifecycle design §13 pull-forward fix 1: + # HTTP 4xx/5xx from a worker must NOT be retried as a network error. + # ------------------------------------------------------------------ + + @pytest.mark.asyncio + @pytest.mark.parametrize("status_code", [400, 422, 429, 500, 503]) + async def test_http_error_not_retried( + self, openai_client, completion_request, mock_session, mock_router, status_code + ): + """ClientResponseError (4xx/5xx) must raise immediately without retry.""" + request_info = Mock() + exc = aiohttp.ClientResponseError( + request_info, history=(), status=status_code, message="worker rejected" + ) + mock_session.post.side_effect = exc + + with pytest.raises(aiohttp.ClientResponseError) as exc_info: + await openai_client.send_request(completion_request) + + assert exc_info.value.status == status_code + # Must be exactly one attempt — no retry on HTTP errors. + assert mock_session.post.call_count == 1 + mock_router.finish_request.assert_called_once_with(completion_request) + + @pytest.mark.asyncio + async def test_network_error_still_retried( + self, openai_client, completion_request, mock_session, mock_router + ): + """Plain ClientError (TCP-level) continues to be retried as before.""" + mock_response = self.dummy_response() + mock_http_response = AsyncMock() + mock_http_response.headers = {"Content-Type": "application/json"} + mock_http_response.json = AsyncMock(return_value=mock_response.model_dump()) + mock_http_response.raise_for_status = Mock() + mock_http_response.__aenter__ = AsyncMock(return_value=mock_http_response) + mock_http_response.__aexit__ = AsyncMock() + + mock_session.post.side_effect = [ + aiohttp.ClientError("TCP reset"), + mock_http_response, + ] + + with patch("asyncio.sleep", new_callable=AsyncMock): + response = await openai_client.send_request(completion_request) + + assert isinstance(response, CompletionResponse) + assert mock_session.post.call_count == 2 + class TestHttpErrorBodyPreservation: """Test that HTTP 4xx/5xx errors include the response body (TRTLLM-11123).""" diff --git a/tests/unittest/disaggregated/test_openai_disagg_service.py b/tests/unittest/disaggregated/test_openai_disagg_service.py index 04defdd1cd37..9e0b93d5efd5 100644 --- a/tests/unittest/disaggregated/test_openai_disagg_service.py +++ b/tests/unittest/disaggregated/test_openai_disagg_service.py @@ -622,6 +622,82 @@ async def test_completed_response_with_null_disagg_request_id_passes(self): assert result is resp +@pytest.mark.asyncio +async def test_gen_first_nonstream_cancel_on_ctx_failure(monkeypatch): + """Fable §13 fix 2: when CTX fails, GEN must be cancelled. + + In gen-first non-streaming gather, if CTX raises the GEN task must be + cancelled rather than left running to completion. + """ + monkeypatch.delenv("TRTLLM_DISAGG_BENCHMARK_GEN_ONLY", raising=False) + service = _make_service("generation_first") + service._coordinator.get_disagg_request_id = AsyncMock(return_value=42) + service._ctx_router.get_next_server = AsyncMock(return_value=("ctx:9000", {})) + service._gen_router.get_next_server = AsyncMock(return_value=("gen:9001", {})) + + gen_task_cancelled = asyncio.Event() + + async def _ctx_raises(*_a, **_kw): + raise RuntimeError("ctx boom") + + async def _gen_hangs(*_a, **_kw): + try: + await asyncio.sleep(10) # would never finish naturally + except asyncio.CancelledError: + gen_task_cancelled.set() + raise + return _make_completion_response("gen", "stop", context_only=False) + + service._ctx_client = AsyncMock() + service._gen_client = AsyncMock() + service._ctx_client.send_request = AsyncMock(side_effect=_ctx_raises) + service._gen_client.send_request = AsyncMock(side_effect=_gen_hangs) + + request = CompletionRequest(model="test-model", prompt="hi", stream=False) + with pytest.raises(RuntimeError, match="ctx boom"): + await service._send_disagg_request_gen_first(request) + + assert gen_task_cancelled.is_set(), "GEN task was not cancelled after CTX failure" + + +@pytest.mark.asyncio +async def test_gen_first_nonstream_cancel_on_gen_failure(monkeypatch): + """When GEN fails in gen-first gather, CTX must be cancelled. + + In gen-first non-streaming mode, if GEN raises the CTX task must be + cancelled rather than left running. + """ + monkeypatch.delenv("TRTLLM_DISAGG_BENCHMARK_GEN_ONLY", raising=False) + service = _make_service("generation_first") + service._coordinator.get_disagg_request_id = AsyncMock(return_value=42) + service._ctx_router.get_next_server = AsyncMock(return_value=("ctx:9000", {})) + service._gen_router.get_next_server = AsyncMock(return_value=("gen:9001", {})) + + ctx_task_cancelled = asyncio.Event() + + async def _gen_raises(*_a, **_kw): + raise RuntimeError("gen boom") + + async def _ctx_hangs(*_a, **_kw): + try: + await asyncio.sleep(10) + except asyncio.CancelledError: + ctx_task_cancelled.set() + raise + return _make_completion_response("ctx", "length", context_only=True) + + service._ctx_client = AsyncMock() + service._gen_client = AsyncMock() + service._ctx_client.send_request = AsyncMock(side_effect=_ctx_hangs) + service._gen_client.send_request = AsyncMock(side_effect=_gen_raises) + + request = CompletionRequest(model="test-model", prompt="hi", stream=False) + with pytest.raises(RuntimeError, match="gen boom"): + await service._send_disagg_request_gen_first(request) + + assert ctx_task_cancelled.is_set(), "CTX task was not cancelled after GEN failure" + + class TestFirstGenLogProbsSerializeRoundtrip: """Roundtrip tests for _serialize/_deserialize_first_gen_log_probs."""