[JunyiXu-nv]

Stacked on #21406 ([TRTLLM-17261][test] Add route, middleware-order and shim-identity snapshots and a CPU test seam for trtllm-serve): the first commits below belong to that PR and disappear from this diff when it merges; this PR's own change is the last 2 commits.

@coderabbitai summary

Description

Every OpenAI request on the disaggregated frontend, and every Anthropic /v1/messages request, goes through one closure, OpenAIDisaggServer._wrap_entry_point (tensorrt_llm/serve/openai_disagg_server.py). Upcoming disaggregated-serving work (client-disconnect watching, in-stream errors for streamed Responses, request tracing) all lands in that closure; splitting it first means each of those changes lands in a step of its own instead of growing one closure. This PR splits the closure, with no behavior change, into a composition of five methods:

Step	Owns	Raises or returns
_prepare_request	total_requests, stream_requests / nonstream_requests; chat-template permission; conversation identity (mutates the request); RawRequestResponseHooks (queue-latency histogram, perf records)	HTTPException(400) for a forbidden template
_run_entry_point	the call into OpenAIDisaggregatedService	whatever the service raises (aiohttp.ClientResponseError, CppExecutorError, UnusedChatTemplateKwargsError, ...)
_finish_request	success: total_responses	—
_finish_failed_request	failure: _handle_exception, so http_exceptions or internal_errors, the SIGINT on CppExecutorError, the context-length envelope	raises HTTPException, or returns JSONResponse / None
_serialize_response	StreamingResponse (with the Responses stream guard) / JSONResponse(content=...model_dump(by_alias=True)) for the OpenAI routes	serialization errors go to the failure path, as today
_wrap_entry_point	composition, the disaggregated_request trace span, FastAPI's request model	—

Where in-flight changes to this closure land after the split:

• "Responses agent-client fixes and Responses on disaggregated serving" ([https://nvbugs/6781271][fix] Responses API: agent-client fixes and disaggregated serving #19883, merged) added the Responses stream guard and model_dump(by_alias=True); both moved into _serialize_response unchanged.
• "Responses: fail over-length streamed requests in-stream with context_length_exceeded" ([None][fix] Responses: fail over-length streamed requests in-stream with context_length_exceeded #21060, open) edits the failure branch: after this split it re-applies in _finish_failed_request, which already receives the request.
• A fix that cancels a disconnected client's request end to end from this wrapper goes into _run_entry_point (the disconnect watch; the method already receives the raw request) and _finish_failed_request (the cancellation / 499 branch). "Cancel paired GEN on CTX failure or cancellation" ([TRTLLM-17083][feat] Cancel paired GEN on CTX failure or cancellation #20491, merged) landed in the client/service layer and does not touch this wrapper.
• Request-trace collection starts in _prepare_request; stream early-start work such as "Close the Mooncake store feature gap" ([None][feat] Close the Mooncake store feature gap: async loads, attention DP, speculative prefill and KV connector support in KVCacheManagerV2 #21195, draft) lands in _run_entry_point / _serialize_response.
• "Recover disaggregated requests after timeouts" ([https://nvbugs/6312828][fix] Recover disaggregated requests after timeouts #16402, open) rewrites this closure on its pre-Responses shape: its deadline-watched await entry_point(...) becomes the body of _run_entry_point, its StreamingResponse cleanup wrapper goes into _serialize_response, and its except asyncio.CancelledError branch goes into the wrapper ahead of except Exception, with the cancellation handling in _finish_failed_request. It is the only open PR that adds a disconnect/deadline watch here, so it has to be re-applied onto these steps.
• "Disaggregated Anthropic Messages: start the stream and ping while the context phase runs" ([None][fix] Disaggregated Anthropic Messages: start the stream and ping while the context phase runs #21388, open) and "Anthropic Messages: answer an over-length prompt" ([None][fix] Anthropic Messages: answer an over-length prompt with "prompt is too long", before any stream #21107, open) edit anthropic_messages, the wrapper's caller, which this PR does not change.

Some step parameters are not read yet (raw_req of _run_entry_point, req and hooks of the finish steps): they are the inputs the known upcoming changes to each step need, so those changes will not have to alter the signatures again.

OpenAIDisaggServer is subclassed outside this repository: NVIDIA NeMo-RL's OpenAIDisaggServerAdaptor constructs it as (config, req_timeout_secs=..., coordinator_url=None), adds middleware to self.app right after __init__, and overrides __call__(host, port, sockets=None), _wrap_entry_point and _handle_exception (counting on _perf_metrics_collector.http_exceptions). This PR keeps that whole surface and pins it three ways: a contract test (tests/unittest/disaggregated/test_disagg_server_contract.py, which builds a subclass of the same shape), the new OpenAIDisaggServer class docstring, and a tensorrt_llm/serve/AGENTS.md bullet. Changing any item of that surface needs the NeMo-RL owners' agreement first; the tests' failure messages say so.

Part of the trtllm-serve refactor (epic TRTLLM-17056, task TRTLLM-17262). Builds on the test seam of #21406 (the server factory the new tests build the disaggregated server with).

Behavior

No behavior change. The wrapper keeps the same trace span, the same try block and the same statement order: the counters run before the template check, total_responses before serialization, and a serialization failure still reaches _handle_exception, which is called while the exception is being handled (its traceback logging needs that). The commit structure separates the pin from the change: the first of the two commits adds the characterization tests, the contract test, the class docstring and the AGENTS.md bullet against the unsplit closure; the second performs the split. Both test files pass identically on both commits.

CI bookkeeping

• Multi-GPU table: none (the file is already listed; openai_disagg_server.py changes, so the multi-GPU stages run and need the ci: full pre-merge approved label)
• CODEOWNERS: none
• legacy-files.txt: none (the five new methods live in the already-listed file; the two test files are new code, formatted by the default hooks)
• Test lists: none to edit — both new test files are CPU-only (pytest.mark.cpu_only) and collected by the existing unittest/disaggregated entry of tests/integration/test_lists/test-db/l0_cpu.yml

Test Coverage

New, both collected by the unittest/disaggregated entry of l0_cpu.yml and run on the tests-only commit and on the final commit with identical results:

• tests/unittest/disaggregated/test_disagg_entry_point.py (13 tests): per-request counters for non-streaming, streaming, Anthropic and Responses requests; chat-template 400 before any worker; conversation identity from the headers; upstream error status kept; the context-length worker envelope; None after an engine fatal error (SIGINT recorded); unexpected errors as 500; serialization failures reaching the failure path; exactly one disaggregated_request span per request; the streamed Responses relay guarded; text.format.schema serialized under its wire name (by_alias).
• tests/unittest/disaggregated/test_disagg_server_contract.py (6 tests): the NeMo-RL subclass surface above, exercised through a subclass of the same shape.

Commands run locally (CPU devel container, no GPU):

• pytest -m cpu_only tests/unittest/disaggregated/test_disagg_entry_point.py tests/unittest/disaggregated/test_disagg_server_contract.py -q — 19 passed on the tests-only commit and on the final commit.
• pytest -m cpu_only tests/unittest/disaggregated -q — 1415 passed, 7 failed on this branch; the same run on the base commit gives 1402 passed, 9 failed, and every failure on this branch also fails on the base (4 TestEtcdClusterStorage cases that spawn a real etcd on a port shared with other test processes on the box, 3 test_router.py cases that need LLM_MODELS_ROOT, which that run did not set; they pass 4/4 when rerun with it); no new failure. (8 TestHttpClusterStorage errors in one run were port-18000 collisions with another test process on the shared box; the class passes 8/8 when rerun alone.)
• LLM_MODELS_ROOT=... PYTHONPATH=tests python3 -m test_common.serve.cpu_tests (every serve CPU test, one pytest process per test-db entry, 66 entries) — all entries green except the 4 pre-existing TestEtcdClusterStorage cases above and unittest/grpc/openengine/test_frontend_lifecycle.py::test_frontend_group_cleanup[launcher-death], a process-lifecycle test unrelated to this change (no gRPC file is touched) that passes 8/8 when rerun alone. The route, middleware-order, shim-identity, layering and object.__new__ ratchet checks are among the green entries.
• pytest -m cpu_only tests/unittest/llmapi/apps/test_anthropic_routes.py tests/unittest/llmapi/apps/test_context_length_error.py -q — 58 passed.
• pytest tests/unittest/api_stability -q — 98 passed.
• pytest -m cpu_only tests/unittest/serve/test_agents_md.py -q — 6 passed (the new AGENTS.md bullet names only real paths).
• python scripts/legacy_utils.py lint-precommit tensorrt_llm/serve/openai_disagg_server.py — no new violations.
• pre-commit run --files <changed files> — clean on the second run.
• Negative check (local only, reverted): making the failure step call _handle_exception through the class instead of self fails test_every_failure_reaches_the_overridden_handle_exception with the NeMo-RL message.

GPU coverage is left to CI: the single-GPU disaggregated block of l0_a10.yml (in particular test_disaggregated_single_gpu, test_disaggregated_perf_metrics, test_disaggregated_chat_completion_tool_calls) and the multi-GPU disaggregated stages; openai_disagg_server.py is on the multi-GPU trigger list, so those stages need the ci: full pre-merge approved label.

PR Checklist

Please review the following before submitting your PR:

• PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

• PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

• Test cases are provided for new code paths (see test instructions)

• If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.

• Any new dependencies have been scanned for license and vulnerabilities

• CODEOWNERS updated if ownership changes

• Documentation updated as needed

• Update tava architecture diagram if there is a significant design change in PR.

• The reviewers assigned automatically/manually are appropriate for the PR.

• Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

🤖 Generated with Claude Code