[TRTLLM-17262][refactor] Split the disaggregated server's entry-point wrapper into prepare, run, finish and serialize steps

tensorrt_llm/commands/AGENTS.md

@@ -0,0 +1,6 @@
+<!--
+SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+SPDX-License-Identifier: Apache-2.0
+-->
+
+Rules for the trtllm-serve CLI (`commands/serve.py`) are in `tensorrt_llm/serve/AGENTS.md`.

tensorrt_llm/commands/CLAUDE.md

@@ -0,0 +1,2 @@
+# In ./CLAUDE.md
+@AGENTS.md

tensorrt_llm/serve/AGENTS.md

@@ -0,0 +1,86 @@
+<!--
+SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+SPDX-License-Identifier: Apache-2.0
+-->
+
+# AGENTS.md - tensorrt_llm/serve
+
+Scope: `tensorrt_llm/serve/` and everything under it. Supplements the repo-root `AGENTS.md`
+and `tests/AGENTS.md`; `tensorrt_llm/commands/AGENTS.md` points here for the serve CLI.
+Where a check enforces a rule, the rule names it; a failing check says what to change.
+Name files and symbols here, never line numbers.
+
+## Where code goes
+
+- Model-specific behavior: a hook on `ServingExtension` (`serve/serving_extensions.py`,
+  `serve/extensions/`). Do not branch on a model name in a handler.
+- Prompt rendering: `serve/render/` only. Handlers do not call `apply_chat_template`.
+- Request and response models: a new field needs a row in
+  `tests/unittest/api_stability/references/trtllm_serve_api.yaml`.
+  Check: `tests/unittest/api_stability/test_serve_api.py`.
+- Routes: every server configuration's routes (path, method, route class, and the order of
+  routes that can match the same request) are pinned in
+  `tests/unittest/api_stability/references/trtllm_serve_routes.yaml`. A route change
+  regenerates it and needs the API review committee.
+  Check: `tests/unittest/api_stability/test_serve_routes.py`.
+- Middleware order matters: `add_middleware` inserts at the front, so the last one added is
+  outermost. Do not reorder. Check: `tests/unittest/serve/test_middleware_order.py`.
+- Import directions between layers are tests; a PR that creates a package adds its contract.
+  Check: `tests/unittest/serve/test_layering.py`.
+- VisualGen serve tests live under a path that contains `visual_gen/`
+  (`tests/unittest/serve/api/visual_gen/`), or CI selection misses them.
+
+## Moving code
+
+- A move PR only moves code: no reformatting, no logic change.
+- In the same PR update the multi-GPU file list (`getMultiGpuFileChanged` in
+  `jenkins/L0_MergeRequest.groovy`), `.github/CODEOWNERS` and `legacy-files.txt`. A new file
+  takes the formatting group of the file its code comes from.
+- Old import paths: a move deletes the old module (or the moved names) without a shim and
+  rewrites every import of it in the repository, tests included; module paths written as
+  strings (patch targets, subprocess command lines, `_LAZY_ATTRS` in `serve/__init__.py`) are
+  fixed by hand. The PR body lists "old path → new path" for the release notes. Only five old
+  paths keep a shim, because code outside the repository imports them:
+  `tensorrt_llm/llmapi/reasoning_parser.py`, `serve/tool_parser/`,
+  `serve/openai_disagg_server.py`, `serve/openai_server.py` and `commands/serve.py`. A shim
+  re-exports public names only, lists them in `__all__` (or `_LAZY_NAMES` in
+  `commands/serve.py`), warns once with `FutureWarning`, and has its rows in
+  `tests/unittest/serve/test_shim_identity.py`.
+- `OpenAIDisaggServer` (`serve/openai_disagg_server.py`) is subclassed outside the repository
+  (NeMo-RL). Keep its constructor `(config, req_timeout_secs=..., coordinator_url=None)`,
+  `self.app` with every route registered when `__init__` returns,
+  `__call__(host, port, sockets=None)`, `_wrap_entry_point`, `_handle_exception` and
+  `_perf_metrics_collector.http_exceptions`, or agree the change with the NeMo-RL owners
+  first. Check: `tests/unittest/disaggregated/test_disagg_server_contract.py`.
+- `mock.patch("old.module.name")` does not follow moved code. Grep the tests for every moved
+  name, including `monkeypatch.setattr(module, "name", ...)` and dotted paths split across lines.
+
+## Tests
+
+- Register every new test file in a CI list in the same PR (`tests/AGENTS.md`). Serve tests
+  that need no GPU go in `tests/integration/test_lists/test-db/l0_cpu.yml` and must contain
+  `pytest.mark.cpu_only`, or the CPU stage collects nothing from them.
+- New serve tests mirror the package: `tests/unittest/serve/` for `tensorrt_llm/serve/`.
+- Build servers with the factory and fake generators in `tests/test_common/serve/`
+  (`isolated_serve_env`), never `object.__new__(OpenAIServer)` or
+  `object.__new__(OpenAIDisaggServer)`. Check: `tests/unittest/serve/test_no_object_new.py`.
+- List an app's routes with `effective_routes()` from `tests/test_common/serve/routes.py`, not
+  `app.routes`: FastAPI 0.139 keeps included routers as objects without `.path`.
+- The harmony tests guard whole files with try-imports. After renaming anything they import,
+  run them with `pytest -rs` and check that nothing is skipped as "not importable": a green run
+  can mean every test was skipped.
+
+## Commands
+
+- All serve CPU tests (the runner exits when `LLM_MODELS_ROOT` is unset):
+  `LLM_MODELS_ROOT=/path/to/models PYTHONPATH=tests python3 -m test_common.serve.cpu_tests`
+- Route snapshot: `pytest -m cpu_only tests/unittest/api_stability/test_serve_routes.py tests/unittest/serve/api/visual_gen/test_routes_visual_gen.py`
+- Regenerate the route snapshot: `PYTHONPATH=tests python3 -m test_common.serve.routes --write`
+- API gates: `pytest tests/unittest/api_stability/test_serve_api.py tests/unittest/api_stability/test_serve_cli.py`
+- Extensions and parsers: `pytest tests/unittest/llmapi/apps/test_serving_extensions.py tests/unittest/llmapi/apps/test_tool_parsers.py`
+- Disaggregated router and coordinator (CPU): `pytest tests/unittest/disaggregated -m cpu_only -q`
+- Which CI stage runs a test: `python scripts/test_to_stage_mapping.py --tests "<test>"`. It
+  matches the text of list entries and does not expand directories: a file listed through its
+  directory maps only when you pass the directory entry, e.g. `--tests "unittest/serve"` for
+  anything under `tests/unittest/serve/`.
+- Tests that load model files need `LLM_MODELS_ROOT` (the harmony tests' autouse fixture does).

tensorrt_llm/serve/CLAUDE.md

@@ -0,0 +1,2 @@
+# In ./CLAUDE.md
+@AGENTS.md

tensorrt_llm/serve/openai_disagg_server.py

@@ -68,6 +68,7 @@
                                              PerfMetricsMiddleware,
                                              combine_disagg_metrics)
 from tensorrt_llm.serve.responses_utils import (ServerArrivalTimeMiddleware,
+                                                UCompletionResponseOrGenerator,
                                                 get_steady_clock_now_in_seconds,
                                                 guard_responses_stream,
                                                 stream_error_event)
@@ -212,6 +213,16 @@ def _upstream_error_message(error: aiohttp.ClientResponseError) -> str:
 
 
 class OpenAIDisaggServer:
+    """OpenAI- and Anthropic-compatible frontend of disaggregated serving.
+
+    Extension surface used outside this repository (NeMo-RL subclasses this
+    class): the constructor (config, req_timeout_secs=...,
+    coordinator_url=None) with self.app built and every route registered when
+    it returns; __call__(host, port, sockets=None); _wrap_entry_point;
+    _handle_exception; and _perf_metrics_collector.http_exceptions. Agree a
+    change to any of them with the NeMo-RL owners first; the contract test is
+    tests/unittest/disaggregated/test_disagg_server_contract.py.
+    """
     def __init__(self,
                  config: DisaggServerConfig,
                  req_timeout_secs: int = 180,
@@ -422,44 +433,90 @@ def _wrap_entry_point(self, entry_point: Callable, request_type: type = UComplet
         # annotation with request_type (as openai_server.py does).
         @tracing.trace_span("disaggregated_request")
         async def wrapper(req: request_type, raw_req: Request) -> Response:
+            hooks: Optional[RawRequestResponseHooks] = None
             try:
-                self._perf_metrics_collector.total_requests.inc()
-                if req.stream:
-                    self._perf_metrics_collector.stream_requests.inc()
-                else:
-                    self._perf_metrics_collector.nonstream_requests.inc()
-                try:
-                    ensure_request_chat_template_allowed(
-                        req, self._allow_request_chat_template)
-                except ValueError as e:
-                    raise HTTPException(status_code=400, detail=str(e)) from e
-                self._extract_conversation_id(
-                    req, raw_req,
-                    self._config.conversation_affinity_header_for_subagents)
-                hooks = RawRequestResponseHooks(
-                    raw_req, self._perf_metrics_collector.queue_latency_seconds,
-                    self._collect_perf_metrics)
-                response_or_generator = await entry_point(req, hooks)
-                self._perf_metrics_collector.total_responses.inc()
-                if req.stream:
-                    stream = response_or_generator
-                    if isinstance(req, ResponsesRequest):
-                        # A relayed Responses stream that fails early ends with
-                        # an `error` event; other protocols end on a sentinel.
-                        stream = guard_responses_stream(stream,
-                                                        stream_error_event)
-                    return StreamingResponse(
-                        content=stream,
-                        media_type="text/event-stream")
-                # by_alias: e.g. text.format.schema goes out under its wire name.
-                return JSONResponse(
-                    content=response_or_generator.model_dump(by_alias=True))
+                hooks = self._prepare_request(req, raw_req)
+                response_or_generator = await self._run_entry_point(
+                    entry_point, req, hooks, raw_req)
+                self._finish_request(req, hooks)
+                return self._serialize_response(req, response_or_generator)
             except Exception as e:
                 # Usually raises; returns a Response for worker errors that
                 # carry a machine-readable code (context_length_exceeded).
-                return self._handle_exception(e)
+                return self._finish_failed_request(e, req, hooks)
         return wrapper
 
+    def _prepare_request(self, req: UCompletionRequest,
+                         raw_req: Request) -> RawRequestResponseHooks:
+        """Account for a request and check it before any worker is involved.
+
+        Counts the request, applies the chat-template permission (HTTPException
+        400), resolves the conversation identity from the headers and builds
+        the per-request hooks.
+        """
+        self._perf_metrics_collector.total_requests.inc()
+        if req.stream:
+            self._perf_metrics_collector.stream_requests.inc()
+        else:
+            self._perf_metrics_collector.nonstream_requests.inc()
+        try:
+            ensure_request_chat_template_allowed(
+                req, self._allow_request_chat_template)
+        except ValueError as e:
+            raise HTTPException(status_code=400, detail=str(e)) from e
+        self._extract_conversation_id(
+            req, raw_req,
+            self._config.conversation_affinity_header_for_subagents)
+        return RawRequestResponseHooks(
+            raw_req, self._perf_metrics_collector.queue_latency_seconds,
+            self._collect_perf_metrics)
+
+    async def _run_entry_point(self, entry_point: Callable,
+                               req: UCompletionRequest,
+                               hooks: RawRequestResponseHooks,
+                               raw_req: Request) -> UCompletionResponseOrGenerator:
+        """Run the service entry point for one request.
+
+        A typed response for non-streaming requests, an async generator of the
+        worker's raw SSE bytes for streaming ones.
+        """
+        return await entry_point(req, hooks)
+
+    def _finish_request(self, req: UCompletionRequest,
+                        hooks: RawRequestResponseHooks) -> None:
+        """Account for a request the service answered."""
+        self._perf_metrics_collector.total_responses.inc()
+
+    def _finish_failed_request(
+            self, error: Exception, req: UCompletionRequest,
+            hooks: Optional[RawRequestResponseHooks]) -> Optional[Response]:
+        """Answer a failed request through _handle_exception.
+
+        Must be called while ``error`` is being handled: _handle_exception logs
+        with traceback.format_exc(). Raises HTTPException in most cases, returns
+        a JSONResponse for worker errors that carry a machine-readable code, and
+        returns None after an engine fatal error (SIGINT already sent).
+        """
+        return self._handle_exception(error)
+
+    def _serialize_response(
+            self, req: UCompletionRequest,
+            response_or_generator: UCompletionResponseOrGenerator) -> Response:
+        """The HTTP response of the OpenAI routes."""
+        if req.stream:
+            stream = response_or_generator
+            if isinstance(req, ResponsesRequest):
+                # A relayed Responses stream that fails early ends with
+                # an `error` event; other protocols end on a sentinel.
+                stream = guard_responses_stream(stream,
+                                                stream_error_event)
+            return StreamingResponse(
+                content=stream,
+                media_type="text/event-stream")
+        # by_alias: e.g. text.format.schema goes out under its wire name.
+        return JSONResponse(
+            content=response_or_generator.model_dump(by_alias=True))
+
     async def anthropic_messages(self, request: AnthropicMessagesRequest,
                                  raw_request: Request) -> Response:
         """Serve Anthropic Messages through the disaggregated chat pipeline."""

tests/integration/test_lists/test-db/l0_cpu.yml

@@ -334,6 +334,11 @@ l0_cpu:
   - unittest/quantization/test_mode.py
   - unittest/scaffolding
   - unittest/scripts
+  - unittest/serve
+  # VisualGen serve tests are also collected by the entry above. This entry keeps
+  # them for VisualGen-only PRs, where CBTS keeps only entries containing visual_gen/
+  # (with the trailing slash).
+  - unittest/serve/api/visual_gen/
   - unittest/tools
   - unittest/usage
   - unittest/utils/test_logger.py

tests/test_common/serve/__init__.py

@@ -0,0 +1,14 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.