diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index c2166e7bb04b..a5a96a2ea8a3 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -417,6 +417,9 @@ common-files: &common_files | tensorrt_llm/scaffolding/worker.py | tensorrt_llm/scheduling_params.py | tensorrt_llm/serve/__init__.py | + tensorrt_llm/serve/_base/msgspec_route.py | + tensorrt_llm/serve/_base/sse.py | + tensorrt_llm/serve/api/anthropic/serving.py | tensorrt_llm/serve/chat_utils.py | tensorrt_llm/serve/cluster_storage.py | tensorrt_llm/serve/disagg_auto_scaling.py | @@ -1144,6 +1147,9 @@ legacy-files: &legacy_files | tensorrt_llm/scaffolding/worker.py | tensorrt_llm/scheduling_params.py | tensorrt_llm/serve/__init__.py | + tensorrt_llm/serve/_base/msgspec_route.py | + tensorrt_llm/serve/_base/sse.py | + tensorrt_llm/serve/api/anthropic/serving.py | tensorrt_llm/serve/chat_utils.py | tensorrt_llm/serve/cluster_storage.py | tensorrt_llm/serve/disagg_auto_scaling.py | @@ -1620,6 +1626,7 @@ repos: # Run mypy with no explicit paths so it type-checks the set declared in # pyproject.toml [tool.mypy] `files` (single source of truth). This # `files:` pattern only decides when the hook triggers — any change under - # the sampler source/test dirs. + # the sampler source/test dirs or the serving frontends. A PR that adds + # a typed module elsewhere extends this pattern in the same PR. pass_filenames: false - files: ^(tensorrt_llm/_torch/pyexecutor/sampler/|tests/unittest/_torch/sampler/) + files: ^(tensorrt_llm/_torch/pyexecutor/sampler/|tests/unittest/_torch/sampler/|tensorrt_llm/serve/|tensorrt_llm/commands/serve_cli/) diff --git a/jenkins/L0_MergeRequest.groovy b/jenkins/L0_MergeRequest.groovy index fed051af26b2..6d31bd8f7e27 100644 --- a/jenkins/L0_MergeRequest.groovy +++ b/jenkins/L0_MergeRequest.groovy @@ -1948,6 +1948,13 @@ def getMultiGpuFileChanged(pipeline, testFilter, globalVars) "tensorrt_llm/models/llama/", "tensorrt_llm/parameter.py", "tensorrt_llm/serve/__init__.py", + // _MsgspecRoute decodes the disaggregated frontend's msgpack bodies. + "tensorrt_llm/serve/_base/", + // The disaggregated /v1/messages runs through the Anthropic serving class. + "tensorrt_llm/serve/api/anthropic/api_router.py", + "tensorrt_llm/serve/api/anthropic/serving.py", + // /health is the disaggregated readiness probe. + "tensorrt_llm/serve/api/ops/", "tensorrt_llm/serve/cluster_storage.py", "tensorrt_llm/serve/disagg_auto_scaling.py", "tensorrt_llm/serve/extensions/", diff --git a/legacy-files.txt b/legacy-files.txt index 9850e7c810b4..3dd21ca755bf 100644 --- a/legacy-files.txt +++ b/legacy-files.txt @@ -409,6 +409,9 @@ tensorrt_llm/scaffolding/task_collection.py tensorrt_llm/scaffolding/worker.py tensorrt_llm/scheduling_params.py tensorrt_llm/serve/__init__.py +tensorrt_llm/serve/_base/msgspec_route.py +tensorrt_llm/serve/_base/sse.py +tensorrt_llm/serve/api/anthropic/serving.py tensorrt_llm/serve/chat_utils.py tensorrt_llm/serve/cluster_storage.py tensorrt_llm/serve/disagg_auto_scaling.py diff --git a/pyproject.toml b/pyproject.toml index 60e243d984b0..20a951c8b770 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -474,6 +474,9 @@ exclude = [ "tensorrt_llm/scaffolding/worker.py", "tensorrt_llm/scheduling_params.py", "tensorrt_llm/serve/__init__.py", + "tensorrt_llm/serve/_base/msgspec_route.py", + "tensorrt_llm/serve/_base/sse.py", + "tensorrt_llm/serve/api/anthropic/serving.py", "tensorrt_llm/serve/chat_utils.py", "tensorrt_llm/serve/cluster_storage.py", "tensorrt_llm/serve/disagg_auto_scaling.py", @@ -862,6 +865,13 @@ files = [ "tensorrt_llm/_torch/pyexecutor/sampler/token_ban.py", "tensorrt_llm/_torch/pyexecutor/sampler/ops/flashinfer.py", "tensorrt_llm/_torch/pyexecutor/sampler/ops/vanilla.py", + "tensorrt_llm/serve/_base/chat_backend.py", + "tensorrt_llm/serve/_base/errors.py", + "tensorrt_llm/serve/_base/request_context.py", + "tensorrt_llm/serve/_base/routing.py", + "tensorrt_llm/serve/api/anthropic/api_router.py", + "tensorrt_llm/serve/api/ops/system/api_router.py", + "tensorrt_llm/serve/api/ops/system/serving.py", "tests/unittest/_torch/sampler/test_torch_sampler.py", "tests/unittest/_torch/sampler/test_token_ban.py", "tests/unittest/_torch/sampler/test_beam_search.py", diff --git a/ruff-legacy-baseline.json b/ruff-legacy-baseline.json index 32f43a9f5486..eabd434cc0f9 100644 --- a/ruff-legacy-baseline.json +++ b/ruff-legacy-baseline.json @@ -528,6 +528,9 @@ "tensorrt_llm/scaffolding/task.py": { "F821": 1 }, + "tensorrt_llm/serve/_base/sse.py": { + "F821": 1 + }, "tensorrt_llm/serve/disagg_auto_scaling.py": { "D205": 2, "D212": 2 diff --git a/ruff-legacy.toml b/ruff-legacy.toml index dab55b4d0434..a091e2e35088 100644 --- a/ruff-legacy.toml +++ b/ruff-legacy.toml @@ -426,6 +426,9 @@ include = [ "tensorrt_llm/scaffolding/worker.py", "tensorrt_llm/scheduling_params.py", "tensorrt_llm/serve/__init__.py", + "tensorrt_llm/serve/_base/msgspec_route.py", + "tensorrt_llm/serve/_base/sse.py", + "tensorrt_llm/serve/api/anthropic/serving.py", "tensorrt_llm/serve/chat_utils.py", "tensorrt_llm/serve/cluster_storage.py", "tensorrt_llm/serve/disagg_auto_scaling.py", diff --git a/tensorrt_llm/commands/AGENTS.md b/tensorrt_llm/commands/AGENTS.md new file mode 100644 index 000000000000..35fcf993bc1a --- /dev/null +++ b/tensorrt_llm/commands/AGENTS.md @@ -0,0 +1,6 @@ + + +Rules for the trtllm-serve CLI (`commands/serve.py`) are in `tensorrt_llm/serve/AGENTS.md`. diff --git a/tensorrt_llm/commands/CLAUDE.md b/tensorrt_llm/commands/CLAUDE.md new file mode 100644 index 000000000000..40c6cf85d526 --- /dev/null +++ b/tensorrt_llm/commands/CLAUDE.md @@ -0,0 +1,2 @@ +# In ./CLAUDE.md +@AGENTS.md diff --git a/tensorrt_llm/serve/AGENTS.md b/tensorrt_llm/serve/AGENTS.md new file mode 100644 index 000000000000..6ff6ffeaf440 --- /dev/null +++ b/tensorrt_llm/serve/AGENTS.md @@ -0,0 +1,107 @@ + + +# 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 + +- New endpoint: a router and a serving class under `serve/api///` + (`api_router.py` + `serving.py`). Create the router with `make_router()` + (`serve/_base/routing.py`) and add a `RouterSpec` to `ROUTERS`. Never add a method or a + route to `OpenAIServer`. Check: `tests/unittest/api_stability/test_serve_routes.py`. +- Attach a router with `app.include_router(router)` and no `prefix=`; give the prefix to + `make_router(prefix=...)` instead. The composition root sets each serving object on + `app.state._serving_` before attaching; routers read it from + `request.app.state`, never from a closure. +- Anthropic Messages: `serve/api/anthropic/`. One `AnthropicMessagesServing` serves both + servers through a `ChatBackend`; the aggregated routes are in `api_router.py`, the + disaggregated frontend keeps its own `/v1/messages` route and calls the same class. + Check: `tests/unittest/serve/api/anthropic/test_messages_behavior.py`. +- 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. +- A move PR is redone on the latest `main` right before merge instead of being rebased, + and its PR body lists every moved definition with its old and new location, and every + definition the PR changes on purpose. +- 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. + +## Checklists + +- Endpoint: package under `serve/api/`; `RouterSpec` in `ROUTERS`; regenerate the route + snapshot (API review); multi-GPU list entry if it runs on the disaggregated path; CPU test + registered in `tests/integration/test_lists/test-db/l0_cpu.yml`. + +## 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 ""`. 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). diff --git a/tensorrt_llm/serve/CLAUDE.md b/tensorrt_llm/serve/CLAUDE.md new file mode 100644 index 000000000000..40c6cf85d526 --- /dev/null +++ b/tensorrt_llm/serve/CLAUDE.md @@ -0,0 +1,2 @@ +# In ./CLAUDE.md +@AGENTS.md diff --git a/tensorrt_llm/serve/_base/__init__.py b/tensorrt_llm/serve/_base/__init__.py new file mode 100644 index 000000000000..0256eaf43af3 --- /dev/null +++ b/tensorrt_llm/serve/_base/__init__.py @@ -0,0 +1,15 @@ +# 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. +"""Shared serving infrastructure with no dependency on endpoint packages.""" diff --git a/tensorrt_llm/serve/_base/chat_backend.py b/tensorrt_llm/serve/_base/chat_backend.py new file mode 100644 index 000000000000..06e90c8856d2 --- /dev/null +++ b/tensorrt_llm/serve/_base/chat_backend.py @@ -0,0 +1,64 @@ +# 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. +"""The chat implementation a wire-dialect serving class runs on.""" + +from __future__ import annotations + +from typing import AsyncIterator, Protocol, Union + +from tensorrt_llm.serve._base.errors import ServingError +from tensorrt_llm.serve._base.request_context import RequestContext +from tensorrt_llm.serve.openai_protocol import ChatCompletionRequest, ChatCompletionResponse + +# Non-streaming: the typed response. Streaming: the OpenAI SSE stream, as str +# frames (aggregated) or as the worker's raw bytes (disaggregated). +ChatResult = Union[ChatCompletionResponse, AsyncIterator[str], AsyncIterator[bytes]] + + +class ChatBackend(Protocol): + """The chat implementation a wire-dialect serving class runs on. + + The aggregated server implements it over its typed chat core, the + disaggregated frontend over its prepare, run and finish steps. Each side + keeps its own streamed model name and error-message extraction, which is + why both are methods of the backend and not of the serving class. + """ + + async def create_chat_completion( + self, request: ChatCompletionRequest, ctx: RequestContext + ) -> ChatResult: + """Run one chat completion. + + Non-streaming: the typed response. Streaming: the OpenAI SSE stream, as + str frames (aggregated) or the worker's raw bytes (disaggregated); one + stream never mixes the two. + + Raises: + ServingError: The request was rejected; the HTTP status travels + apart from the body. + EngineFatalError: The engine failed; SIGINT has been sent. + """ + ... + + def stream_model(self, request: ChatCompletionRequest) -> str: + """The model a translated stream announces when no chunk names one.""" + ... + + def error_message(self, error: ServingError) -> str: + """The client-facing message of an error. + + Each backend extracts it from the error body its own way. + """ + ... diff --git a/tensorrt_llm/serve/_base/errors.py b/tensorrt_llm/serve/_base/errors.py new file mode 100644 index 000000000000..e3fbd54b05a3 --- /dev/null +++ b/tensorrt_llm/serve/_base/errors.py @@ -0,0 +1,76 @@ +# 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. +"""Typed errors of the serving core. + +The core raises them; HTTP wrappers and wire dialects render them. ServingError keeps the HTTP +status apart from the error body, because ErrorResponse.code is not always a status (an +over-length prompt carries "context_length_exceeded"). +""" + +from http import HTTPStatus +from typing import Union + +from tensorrt_llm.executor.utils import ( + CONTEXT_LENGTH_EXCEEDED_CODE, + is_context_length_exceeded_message, +) +from tensorrt_llm.serve.openai_protocol import ErrorResponse + + +def error_body( + message: str, + err_type: str = "BadRequestError", + status_code: HTTPStatus = HTTPStatus.BAD_REQUEST, +) -> ErrorResponse: + """Build the OpenAI error body for a rejected request. + + ``code`` is the HTTP status, except for an over-length prompt, which is recognized by its + text because it crosses the executor IPC boundary as a plain string (RequestError), so no + exception type survives to here. + """ + code: Union[int, str] = status_code.value + if is_context_length_exceeded_message(message): + code = CONTEXT_LENGTH_EXCEEDED_CODE + return ErrorResponse(message=message, type=err_type, code=code) + + +class ServingError(Exception): + """A request the serving core rejects, with the HTTP status to answer it with. + + Attributes: + status: The HTTP status code. + body: The OpenAI error body; ``body.code`` is not necessarily ``status``. + """ + + def __init__(self, status: int, body: ErrorResponse) -> None: + super().__init__(body.message) + self.status = status + self.body = body + + +def serving_error( + message: str, + err_type: str = "BadRequestError", + status_code: HTTPStatus = HTTPStatus.BAD_REQUEST, +) -> ServingError: + """Return the ServingError for what OpenAIServer.create_error_response would render.""" + return ServingError(status_code.value, error_body(message, err_type, status_code)) + + +class EngineFatalError(Exception): + """The engine failed fatally; SIGINT has been sent and the server is shutting down. + + It carries no response; each caller decides what to answer. + """ diff --git a/tensorrt_llm/serve/_base/msgspec_route.py b/tensorrt_llm/serve/_base/msgspec_route.py new file mode 100644 index 000000000000..a7eaf33acc64 --- /dev/null +++ b/tensorrt_llm/serve/_base/msgspec_route.py @@ -0,0 +1,36 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Request and route classes that decode msgpack request bodies from the +disaggregated frontend.""" + +import json + +import msgspec + +from tensorrt_llm.serve.kv_hints import KvHintsRequest, KvHintsRoute + +_msgpack_decoder = msgspec.msgpack.Decoder() + + +class _MsgspecRequest(KvHintsRequest): + """Request that decodes X-TRTLLM-Msgpack bodies with msgspec.""" + + async def json(self): + if not hasattr(self, "_json_body"): + body = await self.body() + if not body: + self._json_body = {} + elif self.headers.get("x-trtllm-msgpack") == "1": + try: + self._json_body = _msgpack_decoder.decode(body) + except msgspec.DecodeError as e: + raise json.JSONDecodeError(f"msgpack: {e}", "", 0) from e + else: + self._json_body = json.loads(body) + return self.filter_kv_hints(self._json_body) + + +class _MsgspecRoute(KvHintsRoute): + """Parse msgpack/JSON and apply hint trust policy before model validation.""" + + request_type = _MsgspecRequest diff --git a/tensorrt_llm/serve/_base/request_context.py b/tensorrt_llm/serve/_base/request_context.py new file mode 100644 index 000000000000..2e2b5949267b --- /dev/null +++ b/tensorrt_llm/serve/_base/request_context.py @@ -0,0 +1,149 @@ +# 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. +"""What the chat core reads from an HTTP request, as one explicit object.""" + +from __future__ import annotations + +import json +from typing import Any, Dict, List, Mapping, Optional + +from starlette.datastructures import Address, Headers, State +from starlette.requests import Request + +_UNSET: Any = object() + + +class RequestContext: + """What the chat core reads from an HTTP request, made explicit. + + ``from_http()`` wraps the request FastAPI passed to a route and delegates to + it lazily, so the middleware and the core share one state + (``server_arrival_time``, ``server_first_token_time``, + ``perf_metrics_records``) and msgpack bodies keep decoding through the + request's own ``json()``. ``internal()`` is a request with no client behind + it, for Anthropic message batches. + + The properties ``headers``, ``state``, ``client``, ``json()`` and + ``is_disconnected()`` carry the names ``Request`` uses, so helpers that + still take a ``Request`` accept a context unchanged. The object is always + truthy, as every ``Request`` is. + """ + + def __init__( + self, + *, + request: Optional[Request] = None, + state: Optional[State] = None, + raw_json: Optional[Dict[str, Any]] = None, + ) -> None: + """Use ``from_http()`` or ``internal()``.""" + if request is None and state is None: + raise ValueError( + "RequestContext needs a request to delegate to or an internal state; " + "build it with RequestContext.from_http(request) or " + "RequestContext.internal(raw_json)" + ) + self._request = request + self._state = state + self._raw_json = raw_json + self._json_cache: Any = _UNSET + + @classmethod + def from_http(cls, request: Request) -> RequestContext: + """Wrap ``request``; nothing is read from it until the core asks.""" + return cls(request=request) + + @classmethod + def internal(cls, raw_json: Dict[str, Any]) -> RequestContext: + """A request without headers or a client, whose raw JSON is ``raw_json``. + + The state carries what the middleware would have seeded on a real + connection: an empty ``perf_metrics_records`` list, and + ``no_client_connection`` so that the disconnect watcher does not poll a + connection that can never drop. + """ + return cls( + state=State({"perf_metrics_records": [], "no_client_connection": True}), + raw_json=raw_json, + ) + + @property + def http_request(self) -> Optional[Request]: + """The wrapped HTTP request; ``None`` for an internal context.""" + return self._request + + @property + def headers(self) -> Mapping[str, str]: + if self._request is None: + return Headers() + return self._request.headers + + @property + def state(self) -> State: + if self._request is not None: + return self._request.state + assert self._state is not None # the constructor requires one of the two + return self._state + + @property + def client(self) -> Optional[Address]: + if self._request is None: + return None + return self._request.client + + async def json(self) -> Any: + """The request body as JSON. + + On the HTTP path this is the request's own ``json()`` (msgpack-aware on + workers, cached by the request). An internal context answers its raw + JSON after one encode-decode round trip, computed once, which is what a + body read from the wire looks like. + """ + if self._request is not None: + return await self._request.json() + if self._json_cache is _UNSET: + self._json_cache = json.loads(json.dumps(self._raw_json)) + return self._json_cache + + async def is_disconnected(self) -> bool: + """Whether the client went away; always ``False`` without a client.""" + if self._request is None: + return False + return await self._request.is_disconnected() + + # Typed views of the state, for code written against the context. + + @property + def no_client_connection(self) -> bool: + return bool(getattr(self.state, "no_client_connection", False)) + + @property + def server_arrival_time(self) -> Optional[float]: + value: Optional[float] = getattr(self.state, "server_arrival_time", None) + return value + + @property + def server_first_token_time(self) -> Optional[float]: + value: Optional[float] = getattr(self.state, "server_first_token_time", None) + return value + + @server_first_token_time.setter + def server_first_token_time(self, value: Optional[float]) -> None: + self.state.server_first_token_time = value + + @property + def perf_metrics_records(self) -> List[Dict[str, Any]]: + records: List[Dict[str, Any]] = self.state.perf_metrics_records + return records diff --git a/tensorrt_llm/serve/_base/routing.py b/tensorrt_llm/serve/_base/routing.py new file mode 100644 index 000000000000..79e9b0b4a0d9 --- /dev/null +++ b/tensorrt_llm/serve/_base/routing.py @@ -0,0 +1,115 @@ +# 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. +"""Router creation and role-based attachment for trtllm-serve.""" + +from __future__ import annotations + +import enum +import importlib +from dataclasses import dataclass +from typing import Any, Callable, Optional + +from fastapi import APIRouter, FastAPI + +from tensorrt_llm.llmapi.disagg_utils import ServerRole +from tensorrt_llm.serve._base.msgspec_route import _MsgspecRoute + + +class RoleGroup(enum.Enum): + LLM = "llm" + MM_ENCODER = "mm_encoder" + EMBEDDING = "embedding" + VISUAL_GEN = "visual_gen" + + +ALL_ROLE_GROUPS: frozenset[RoleGroup] = frozenset(RoleGroup) + + +def role_group_of(role: Optional[ServerRole]) -> RoleGroup: + """None, CONTEXT and GENERATION serve the LLM routes, as the dispatch in OpenAIServer does.""" + if role is ServerRole.MM_ENCODER: + return RoleGroup.MM_ENCODER + if role is ServerRole.EMBEDDING: + return RoleGroup.EMBEDDING + if role is ServerRole.VISUAL_GEN: + return RoleGroup.VISUAL_GEN + return RoleGroup.LLM + + +@dataclass(frozen=True) +class ServerCapabilities: + """What the composition root decides once at startup; routers are attached from it.""" + + role: Optional[ServerRole] + role_group: RoleGroup + # A harmony (gpt-oss) model serves chat through the harmony path, which + # the Anthropic Messages routes do not support; LLM role only. + harmony_model: bool = False + + @classmethod + def for_role(cls, role: Optional[ServerRole]) -> ServerCapabilities: + return cls(role=role, role_group=role_group_of(role)) + + +def _always(caps: ServerCapabilities) -> bool: + return True + + +def _anthropic_routes_enabled(caps: ServerCapabilities) -> bool: + return not caps.harmony_model + + +@dataclass(frozen=True) +class RouterSpec: + """One router module and the roles and capabilities that mount it.""" + + module: str # full dotted path of the api_router module + roles: frozenset[RoleGroup] + when: Callable[[ServerCapabilities], bool] = _always + + +# Attach order is registration order. Keep the route snapshot's conflict-relevant +# order: a router that registers GET /metrics must come before the Prometheus mount. +ROUTERS: tuple[RouterSpec, ...] = ( + RouterSpec("tensorrt_llm.serve.api.ops.system.api_router", roles=ALL_ROLE_GROUPS), + RouterSpec( + "tensorrt_llm.serve.api.anthropic.api_router", + roles=frozenset({RoleGroup.LLM}), + when=_anthropic_routes_enabled, + ), +) + + +def make_router(**kwargs: Any) -> APIRouter: + """The only way to create a trtllm-serve router: every route decodes msgpack bodies. + + The disaggregated frontend sends worker requests as msgpack; a route of a plain + APIRouter keeps APIRoute and rejects them. + """ + if "route_class" in kwargs: + raise TypeError( + "make_router() always uses _MsgspecRoute; do not pass route_class. " + "A plain APIRoute breaks msgpack request bodies from the disaggregated frontend." + ) + return APIRouter(route_class=_MsgspecRoute, **kwargs) + + +def attach_routers( + app: FastAPI, caps: ServerCapabilities, routers: tuple[RouterSpec, ...] = ROUTERS +) -> None: + """Import and attach only the routers this role needs.""" + for spec in routers: + if caps.role_group in spec.roles and spec.when(caps): + importlib.import_module(spec.module).attach_router(app) diff --git a/tensorrt_llm/serve/_base/sse.py b/tensorrt_llm/serve/_base/sse.py new file mode 100644 index 000000000000..ef935a266d10 --- /dev/null +++ b/tensorrt_llm/serve/_base/sse.py @@ -0,0 +1,43 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Server-sent-event generators shared by the chat serving paths.""" + +import traceback +from typing import AsyncGenerator + +from tensorrt_llm._utils import nvtx_mark +from tensorrt_llm.executor.postproc_worker import PostprocParams +from tensorrt_llm.llmapi.llm import RequestOutput +from tensorrt_llm.logger import logger + + +async def chat_stream_generator(promise: RequestOutput, + postproc_params: PostprocParams, *, server, + ctx) -> AsyncGenerator[str, None]: + try: + if not server.postproc_worker_enabled: + post_processor, args = postproc_params.post_processor, postproc_params.postproc_args + first_response = await anext(promise) + ctx.state.server_first_token_time = server._adjusted_steady_clock.now() + pp_results = first_response.outputs[ + 0]._postprocess_result if server.postproc_worker_enabled else post_processor( + first_response, args) + for pp_res in pp_results: + yield pp_res + # Making sure we can handle the situation where there is only one response + res = first_response + async for res in promise: + pp_results = res.outputs[ + 0]._postprocess_result if server.postproc_worker_enabled else post_processor( + res, args) + for pp_res in pp_results: + yield pp_res + yield "data: [DONE]\n\n" + await server._extract_metrics(res, ctx) + nvtx_mark("generation ends") + except Exception: + # Only failures are logged: a consumer that stops reading closes + # the stream with GeneratorExit, and a client that disconnects + # cancels it, and neither is an error of this stream. + logger.error(traceback.format_exc()) + raise diff --git a/tensorrt_llm/serve/api/__init__.py b/tensorrt_llm/serve/api/__init__.py new file mode 100644 index 000000000000..e5636921d165 --- /dev/null +++ b/tensorrt_llm/serve/api/__init__.py @@ -0,0 +1,15 @@ +# 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. +"""Endpoint packages. Do not import routers here: attach_routers imports them by role.""" diff --git a/tensorrt_llm/serve/api/anthropic/__init__.py b/tensorrt_llm/serve/api/anthropic/__init__.py new file mode 100644 index 000000000000..82f6904e8d05 --- /dev/null +++ b/tensorrt_llm/serve/api/anthropic/__init__.py @@ -0,0 +1,18 @@ +# 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. +"""Anthropic Messages wire dialect: protocol, adapter, batches, serving classes and router. + +Import submodules directly. +""" diff --git a/tensorrt_llm/serve/anthropic_adapter.py b/tensorrt_llm/serve/api/anthropic/adapter.py similarity index 99% rename from tensorrt_llm/serve/anthropic_adapter.py rename to tensorrt_llm/serve/api/anthropic/adapter.py index 267b77ce2d1d..74c5259f187f 100644 --- a/tensorrt_llm/serve/anthropic_adapter.py +++ b/tensorrt_llm/serve/api/anthropic/adapter.py @@ -37,7 +37,7 @@ from fastapi.responses import JSONResponse from tensorrt_llm.logger import logger -from tensorrt_llm.serve.anthropic_protocol import ( +from tensorrt_llm.serve.api.anthropic.protocol import ( AnthropicContentBlockDeltaEvent, AnthropicContentBlockStartEvent, AnthropicContentBlockStopEvent, diff --git a/tensorrt_llm/serve/api/anthropic/api_router.py b/tensorrt_llm/serve/api/anthropic/api_router.py new file mode 100644 index 000000000000..bf2b2432a812 --- /dev/null +++ b/tensorrt_llm/serve/api/anthropic/api_router.py @@ -0,0 +1,139 @@ +# 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. +"""The Anthropic Messages router of the aggregated server. + +``/v1/messages``, ``/v1/messages/count_tokens`` and the six Message Batches +routes. The disaggregated frontend keeps its own ``/v1/messages`` route and +calls the same serving class. +""" + +from __future__ import annotations + +from typing import Optional + +from fastapi import FastAPI, Query, Request, Response + +from tensorrt_llm.serve._base.request_context import RequestContext +from tensorrt_llm.serve._base.routing import make_router +from tensorrt_llm.serve.api.anthropic.protocol import ( + AnthropicCountTokensRequest, + AnthropicCreateBatchRequest, + AnthropicMessagesRequest, +) +from tensorrt_llm.serve.api.anthropic.serving import ( + AnthropicBatchesServing, + AnthropicCountTokensServing, + AnthropicMessagesServing, +) + +# The serving objects the composition root sets on app.state before attaching. +APP_STATE_KEYS = ( + "anthropic_serving_messages", + "anthropic_serving_count_tokens", + "anthropic_serving_batches", +) + +router = make_router() + + +# No app.state defaults: wiring bugs fail loudly. +def _messages(raw_request: Request) -> AnthropicMessagesServing: + serving: AnthropicMessagesServing = raw_request.app.state.anthropic_serving_messages + return serving + + +def _count_tokens(raw_request: Request) -> AnthropicCountTokensServing: + serving: AnthropicCountTokensServing = raw_request.app.state.anthropic_serving_count_tokens + return serving + + +def _batches(raw_request: Request) -> AnthropicBatchesServing: + serving: AnthropicBatchesServing = raw_request.app.state.anthropic_serving_batches + return serving + + +async def anthropic_messages(request: AnthropicMessagesRequest, raw_request: Request) -> Response: + """Serve the Anthropic Messages API on top of the chat pipeline.""" + return await _messages(raw_request).create_message( + request, RequestContext.from_http(raw_request) + ) + + +async def anthropic_count_tokens( + request: AnthropicCountTokensRequest, raw_request: Request +) -> Response: + """Report how many input tokens a Messages request would consume.""" + return await _count_tokens(raw_request).anthropic_count_tokens(request) + + +async def anthropic_create_batch( + request: AnthropicCreateBatchRequest, raw_request: Request +) -> Response: + return await _batches(raw_request).anthropic_create_batch(request) + + +async def anthropic_list_batches( + raw_request: Request, + limit: int = Query(20, ge=1, le=100), + after_id: Optional[str] = None, + before_id: Optional[str] = None, +) -> Response: + return await _batches(raw_request).anthropic_list_batches( + limit=limit, after_id=after_id, before_id=before_id + ) + + +async def anthropic_get_batch(batch_id: str, raw_request: Request) -> Response: + return await _batches(raw_request).anthropic_get_batch(batch_id) + + +async def anthropic_delete_batch(batch_id: str, raw_request: Request) -> Response: + return await _batches(raw_request).anthropic_delete_batch(batch_id) + + +async def anthropic_cancel_batch(batch_id: str, raw_request: Request) -> Response: + return await _batches(raw_request).anthropic_cancel_batch(batch_id) + + +async def anthropic_batch_results(batch_id: str, raw_request: Request) -> Response: + return await _batches(raw_request).anthropic_batch_results(batch_id) + + +# Registration order is match order. The literal batch paths come before the +# parameterized ones, as they always did. +router.add_api_route("/v1/messages", anthropic_messages, methods=["POST"]) +# Claude Code calls count_tokens before most turns to decide whether to +# compact its context; an unregistered route 404s every call. +router.add_api_route("/v1/messages/count_tokens", anthropic_count_tokens, methods=["POST"]) +router.add_api_route("/v1/messages/batches", anthropic_create_batch, methods=["POST"]) +router.add_api_route("/v1/messages/batches", anthropic_list_batches, methods=["GET"]) +router.add_api_route("/v1/messages/batches/{batch_id}", anthropic_get_batch, methods=["GET"]) +router.add_api_route("/v1/messages/batches/{batch_id}", anthropic_delete_batch, methods=["DELETE"]) +router.add_api_route( + "/v1/messages/batches/{batch_id}/cancel", anthropic_cancel_batch, methods=["POST"] +) +router.add_api_route( + "/v1/messages/batches/{batch_id}/results", anthropic_batch_results, methods=["GET"] +) + + +def attach_router(app: FastAPI) -> None: + for key in APP_STATE_KEYS: + if not hasattr(app.state, key): + raise RuntimeError( + f"app.state.{key} must be set before attaching " + "tensorrt_llm.serve.api.anthropic.api_router; the composition root sets it." + ) + app.include_router(router) diff --git a/tensorrt_llm/serve/anthropic_batches.py b/tensorrt_llm/serve/api/anthropic/batches.py similarity index 99% rename from tensorrt_llm/serve/anthropic_batches.py rename to tensorrt_llm/serve/api/anthropic/batches.py index 6d693df7c431..c6583367e91e 100644 --- a/tensorrt_llm/serve/anthropic_batches.py +++ b/tensorrt_llm/serve/api/anthropic/batches.py @@ -41,7 +41,7 @@ from typing import Any, Awaitable, Callable, Dict, List, Literal, Optional, Tuple from tensorrt_llm.logger import logger -from tensorrt_llm.serve.anthropic_protocol import ( +from tensorrt_llm.serve.api.anthropic.protocol import ( AnthropicBatchRequestCounts, AnthropicBatchRequestItem, AnthropicMessageBatch, diff --git a/tensorrt_llm/serve/anthropic_protocol.py b/tensorrt_llm/serve/api/anthropic/protocol.py similarity index 99% rename from tensorrt_llm/serve/anthropic_protocol.py rename to tensorrt_llm/serve/api/anthropic/protocol.py index 929f04d520e0..80eaacf94dac 100644 --- a/tensorrt_llm/serve/anthropic_protocol.py +++ b/tensorrt_llm/serve/api/anthropic/protocol.py @@ -344,7 +344,7 @@ class AnthropicCountTokensResponse(AnthropicBaseModel): AnthropicBatchResultType = Literal["succeeded", "errored", "canceled", "expired"] # Anthropic's documented ceiling. Enforced because every request in a batch is -# held in memory here (see anthropic_batches), so an unbounded batch is a way +# held in memory here (see batches), so an unbounded batch is a way # to exhaust the server's memory with a single call. ANTHROPIC_BATCH_MAX_REQUESTS = 100_000 diff --git a/tensorrt_llm/serve/api/anthropic/serving.py b/tensorrt_llm/serve/api/anthropic/serving.py new file mode 100644 index 000000000000..daeddfe2ed53 --- /dev/null +++ b/tensorrt_llm/serve/api/anthropic/serving.py @@ -0,0 +1,319 @@ +# 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. +"""Serving classes of the Anthropic Messages dialect. + +Messages, count_tokens and the Message Batches, served from values the +composition root injects; no class here reads a server attribute. +""" + +import asyncio +import json +import traceback +from typing import Any, Callable, Dict, Optional, Tuple + +from fastapi import Query +from fastapi.responses import JSONResponse, Response, StreamingResponse +from pydantic import ValidationError + +from tensorrt_llm.logger import logger +from tensorrt_llm.serve._base.chat_backend import ChatBackend +from tensorrt_llm.serve._base.errors import EngineFatalError, ServingError +from tensorrt_llm.serve._base.request_context import RequestContext +from tensorrt_llm.serve.api.anthropic.adapter import ( + AnthropicRequestError, AnthropicResponseError, anthropic_error_response, + convert_anthropic_count_tokens_request, convert_anthropic_request, + convert_chat_response, reframe_openai_stream) +from tensorrt_llm.serve.api.anthropic.batches import (AnthropicBatchStore, + BatchItemRunner, + BatchStoreFullError, + results_to_jsonl) +from tensorrt_llm.serve.api.anthropic.protocol import ( + AnthropicBatchDeleteResponse, AnthropicBatchList, + AnthropicCountTokensRequest, AnthropicCountTokensResponse, + AnthropicCreateBatchRequest, AnthropicMessagesRequest) +from tensorrt_llm.serve.openai_protocol import ChatCompletionResponse +from tensorrt_llm.serve.render import (HarmonyConversionError, RenderResources, + UnsupportedRenderError, render_chat, + render_chat_for_length) + + +def _error_type(status: int) -> str: + """The Anthropic error type both servers derive from an HTTP status.""" + return "invalid_request_error" if 400 <= status < 500 else "api_error" + + +class AnthropicMessagesServing: + """The Anthropic Messages orchestration, for any ChatBackend. + + One class serves the aggregated server (over its typed chat core), the + disaggregated frontend (over its prepare, run and finish steps) and the + aggregated Message Batches. What differs between the servers, the model a + stream without chunks announces and how an error's message is extracted, + comes from the backend. + """ + + def __init__(self, backend: ChatBackend) -> None: + self._backend = backend + + async def create_message(self, request: AnthropicMessagesRequest, + ctx: RequestContext) -> Response: + """Serve the Anthropic Messages API on top of the chat pipeline. + + The Anthropic request is translated into a ChatCompletionRequest, + handed to the chat backend unchanged (chat template, tool parser and + post-processing are reused verbatim), and the typed result is + translated back: JSON body for non-streaming, SSE reframing for + streaming. + """ + try: + chat_request = convert_anthropic_request(request) + except (AnthropicRequestError, ValidationError) as e: + return anthropic_error_response(str(e), "invalid_request_error", + 400) + try: + result = await self._backend.create_chat_completion( + chat_request, ctx) + except EngineFatalError: + return anthropic_error_response("Internal server error", + "api_error", 500) + except ServingError as e: + return anthropic_error_response(self._backend.error_message(e), + _error_type(e.status), e.status) + if not isinstance(result, ChatCompletionResponse): + return StreamingResponse( + content=reframe_openai_stream( + result, model=self._backend.stream_model(chat_request)), + media_type="text/event-stream", + ) + try: + anthropic_response = convert_chat_response(result) + except (AnthropicResponseError, ValidationError, json.JSONDecodeError): + logger.error("Invalid response from OpenAI chat pipeline:\n" + f"{traceback.format_exc()}") + return anthropic_error_response("Internal server error", + "api_error", 500) + return JSONResponse(content=anthropic_response.model_dump( + exclude_none=True)) + + async def run_batched_message( + self, + request: AnthropicMessagesRequest) -> Tuple[str, Dict[str, Any]]: + """Run one batched Messages request through the chat backend.""" + # Streaming makes no sense for a batch: the result is collected, not + # watched. Forcing it off here means a client that sets stream=true in + # params gets a usable message instead of an SSE body in its .jsonl. + request = request.model_copy(update={"stream": False}) + try: + chat_request = convert_anthropic_request(request) + except (AnthropicRequestError, ValidationError) as e: + return "errored", { + "type": "invalid_request_error", + "message": str(e) + } + # A batched request has no client connection behind it: the context + # carries the seeded state and the converted request as its raw JSON. + ctx = RequestContext.internal( + chat_request.model_dump(exclude_none=True)) + try: + result = await self._backend.create_chat_completion( + chat_request, ctx) + except EngineFatalError: + return "errored", { + "type": "api_error", + "message": "Internal server error" + } + except ServingError as e: + return "errored", { + "type": _error_type(e.status), + "message": self._backend.error_message(e) + } + try: + anthropic_response = convert_chat_response(result) + except (AnthropicResponseError, ValidationError, json.JSONDecodeError): + logger.error( + "Invalid response from OpenAI chat pipeline in batch:\n" + f"{traceback.format_exc()}") + return "errored", { + "type": "api_error", + "message": "Internal server error" + } + return "succeeded", anthropic_response.model_dump(exclude_none=True) + + +class AnthropicCountTokensServing: + """Counts the input tokens of a Messages request over the rendered prompt.""" + + def __init__(self, *, render_resources: Callable[[], + RenderResources]) -> None: + # Called per request: the resources are read at call time, as the + # chat route reads them. + self._render_resources = render_resources + + async def anthropic_count_tokens( + self, request: AnthropicCountTokensRequest) -> Response: + """Report how many input tokens a Messages request would consume. + + The count is taken over the rendered prompt rather than over the raw + messages, so system blocks, tool definitions and the thinking prefix + are all included - those are exactly what make a client's own estimate + drift from the server's. + """ + try: + chat_request = convert_anthropic_count_tokens_request(request) + except AnthropicRequestError as e: + return anthropic_error_response(str(e), "invalid_request_error", + 400) + except ValidationError as e: + return anthropic_error_response(str(e), "invalid_request_error", + 400) + + try: + + # Both the render and the encode are synchronous and both scale + # with prompt size. Claude Code calls this before most turns with + # the whole conversation attached, so running them inline would + # block the event loop -- and therefore every other in-flight + # request on this server -- for a full template render plus a + # tokenizer pass. The chat path avoids this the same way. + def _count() -> int: + # The count is over the prompt tokens the chat route would + # execute: same pipeline, same template, same tokenization. + resources = self._render_resources() + try: + rendered = render_chat(chat_request, + resources, + tokenize=True) + except (UnsupportedRenderError, HarmonyConversionError): + # Media, an own-processor model or an unconvertible Harmony + # request: count the text, tools and template without the media. + rendered = render_chat_for_length(chat_request, resources) + return len(rendered.token_ids) + + input_tokens = await asyncio.to_thread(_count) + except Exception: + logger.error("count_tokens failed to render the prompt:\n" + f"{traceback.format_exc()}") + return anthropic_error_response("Internal server error", + "api_error", 500) + + return JSONResponse(content=AnthropicCountTokensResponse( + input_tokens=input_tokens).model_dump()) + + +class AnthropicBatchesServing: + """The Message Batches routes of the aggregated server.""" + + def __init__(self, *, runner: BatchItemRunner) -> None: + # The batch methods create the store on first use, with this runner. + self._run_batched_message = runner + self._anthropic_batch_store: Optional[AnthropicBatchStore] = None + + def _batch_store(self) -> AnthropicBatchStore: + """Created on first use so servers that never batch pay nothing.""" + if getattr(self, "_anthropic_batch_store", None) is None: + self._anthropic_batch_store = AnthropicBatchStore( + runner=self._run_batched_message) + return self._anthropic_batch_store + + async def anthropic_create_batch( + self, request: AnthropicCreateBatchRequest) -> Response: + try: + batch = self._batch_store().create(request.requests) + except BatchStoreFullError as e: + # 429, not 400: the request is valid and an identical one will + # succeed later, so the client should back off rather than edit it. + return anthropic_error_response(str(e), "rate_limit_error", 429) + except ValueError as e: + return anthropic_error_response(str(e), "invalid_request_error", + 400) + return JSONResponse(content=batch.model_dump()) + + async def anthropic_list_batches( + self, + # Validated here rather than clamped in the store: clamping turned + # limit=0 into 1 instead of a 400, and bounded the page against the + # retention limit, which is an unrelated quantity. + limit: int = Query(20, ge=1, le=100), + after_id: Optional[str] = None, + before_id: Optional[str] = None) -> Response: + try: + batches, has_more = self._batch_store().list(limit=limit, + after_id=after_id, + before_id=before_id) + except LookupError as e: + # Saying so beats an empty page: an empty page is indistinguishable + # from a finished walk, so the client stops early having silently + # skipped every batch past the cursor. + return anthropic_error_response( + f"Unknown pagination cursor {e.args[0]!r}. The batch may have " + "been deleted or expired; restart the listing without a cursor.", + "invalid_request_error", 400) + listing = AnthropicBatchList( + data=batches, + # Reported honestly rather than hardcoded false: a client that + # asked for N and received N cannot otherwise tell a full page from + # the end of the list, and silently stops early. It pages on with + # ?after_id=. + has_more=has_more, + first_id=batches[0].id if batches else None, + last_id=batches[-1].id if batches else None, + ) + return JSONResponse(content=listing.model_dump()) + + async def anthropic_get_batch(self, batch_id: str) -> Response: + batch = self._batch_store().get(batch_id) + if batch is None: + return self._no_such_batch(batch_id) + return JSONResponse(content=batch.model_dump()) + + async def anthropic_cancel_batch(self, batch_id: str) -> Response: + batch = self._batch_store().cancel(batch_id) + if batch is None: + return self._no_such_batch(batch_id) + return JSONResponse(content=batch.model_dump()) + + async def anthropic_delete_batch(self, batch_id: str) -> Response: + try: + deleted = self._batch_store().delete(batch_id) + except ValueError as e: + return anthropic_error_response(str(e), "invalid_request_error", + 400) + if not deleted: + return self._no_such_batch(batch_id) + return JSONResponse(content=AnthropicBatchDeleteResponse( + id=batch_id).model_dump()) + + async def anthropic_batch_results(self, batch_id: str) -> Response: + store = self._batch_store() + results = store.results(batch_id) + if results is None: + if not store.has(batch_id): + return self._no_such_batch(batch_id) + # Known batch, but not finished. 404 would say "no such batch", + # which is wrong and would stop a client from polling. + return anthropic_error_response( + f"Message Batch {batch_id} is still processing; results are " + "available once processing_status is 'ended'", + "invalid_request_error", 400) + return Response(content=results_to_jsonl(results), + media_type="application/x-ndjson") + + @staticmethod + def _no_such_batch(batch_id: str) -> Response: + # Batches are held in memory only, so an id from before a restart lands + # here. Say so, rather than leaving the caller to wonder. + return anthropic_error_response( + f"Message Batch {batch_id} not found. Batches are stored in memory " + "and do not survive a server restart.", "not_found_error", 404) diff --git a/tensorrt_llm/serve/api/ops/__init__.py b/tensorrt_llm/serve/api/ops/__init__.py new file mode 100644 index 000000000000..1c18ed58c923 --- /dev/null +++ b/tensorrt_llm/serve/api/ops/__init__.py @@ -0,0 +1,15 @@ +# 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. +"""Operational endpoints (health, version, model listing).""" diff --git a/tensorrt_llm/serve/api/ops/system/__init__.py b/tensorrt_llm/serve/api/ops/system/__init__.py new file mode 100644 index 000000000000..18515e1ab013 --- /dev/null +++ b/tensorrt_llm/serve/api/ops/system/__init__.py @@ -0,0 +1,15 @@ +# 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. +"""/health, /version and /v1/models, served for every server role.""" diff --git a/tensorrt_llm/serve/api/ops/system/api_router.py b/tensorrt_llm/serve/api/ops/system/api_router.py new file mode 100644 index 000000000000..2bece7f2519f --- /dev/null +++ b/tensorrt_llm/serve/api/ops/system/api_router.py @@ -0,0 +1,60 @@ +# 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. +"""The system router: /health, /version and /v1/models for every server role.""" + +from __future__ import annotations + +from fastapi import FastAPI, Request, Response +from fastapi.responses import JSONResponse + +from tensorrt_llm.serve._base.routing import make_router +from tensorrt_llm.serve.api.ops.system.serving import SystemServing + +APP_STATE_KEY = "ops_serving_system" +_UNAVAILABLE = "LLM is unavailable. Please check the server logs for more details." + +router = make_router() + + +def _serving(raw_request: Request) -> SystemServing: + # No app.state default: wiring bugs fail loudly. + serving: SystemServing = raw_request.app.state.ops_serving_system + return serving + + +@router.get("/health") +async def health(raw_request: Request) -> Response: + if _serving(raw_request).is_healthy(): + return Response(status_code=200) + return Response(status_code=503, content=_UNAVAILABLE) + + +@router.get("/version") +async def version(raw_request: Request) -> JSONResponse: + return JSONResponse(content=_serving(raw_request).version()) + + +@router.get("/v1/models") +async def get_model(raw_request: Request) -> JSONResponse: + return JSONResponse(content=_serving(raw_request).models().model_dump()) + + +def attach_router(app: FastAPI) -> None: + if not hasattr(app.state, APP_STATE_KEY): + raise RuntimeError( + f"app.state.{APP_STATE_KEY} (a SystemServing) must be set before attaching " + "tensorrt_llm.serve.api.ops.system; the composition root sets it." + ) + app.include_router(router) diff --git a/tensorrt_llm/serve/api/ops/system/serving.py b/tensorrt_llm/serve/api/ops/system/serving.py new file mode 100644 index 000000000000..c6a204580df5 --- /dev/null +++ b/tensorrt_llm/serve/api/ops/system/serving.py @@ -0,0 +1,44 @@ +# 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. +"""/health, /version and /v1/models for every server role.""" + +from __future__ import annotations + +from typing import Callable + +from tensorrt_llm.serve.openai_protocol import ModelCard, ModelList +from tensorrt_llm.version import __version__ as VERSION + + +class SystemServing: + """Answers the system endpoints from values the composition root injects. + + The health probe stays on the composition root (it reads executor state and + may raise SIGINT); this class only calls it. + """ + + def __init__(self, *, model: str, health_probe: Callable[[], bool]) -> None: + self._model = model + self._health_probe = health_probe + + def is_healthy(self) -> bool: + return self._health_probe() + + def version(self) -> dict[str, str]: + return {"version": VERSION} + + def models(self) -> ModelList: + # Built per request: ModelCard.created is the request time, as today. + return ModelList(data=[ModelCard(id=self._model)]) diff --git a/tensorrt_llm/serve/openai_disagg_server.py b/tensorrt_llm/serve/openai_disagg_server.py index d8ebec1a6ce7..60f34f5f4ac7 100644 --- a/tensorrt_llm/serve/openai_disagg_server.py +++ b/tensorrt_llm/serve/openai_disagg_server.py @@ -22,14 +22,13 @@ import traceback import uuid from contextlib import asynccontextmanager -from typing import Callable, Optional +from typing import Callable, Optional, Union import aiohttp import uvicorn from fastapi import FastAPI, HTTPException, Request from fastapi.exceptions import RequestValidationError from fastapi.responses import JSONResponse, Response, StreamingResponse -from pydantic import ValidationError from tensorrt_llm.executor import CppExecutorError from tensorrt_llm.executor.executor import CppExecutorError @@ -41,15 +40,13 @@ from tensorrt_llm.llmapi.disagg_utils import (DisaggServerConfig, MetadataServerConfig, ServerRole) from tensorrt_llm.logger import logger +from tensorrt_llm.serve._base.errors import EngineFatalError, ServingError +from tensorrt_llm.serve._base.request_context import RequestContext from tensorrt_llm.serve._telemetry import create_uvicorn_server -from tensorrt_llm.serve.anthropic_adapter import (AnthropicRequestError, - AnthropicResponseError, - anthropic_error_response, - convert_anthropic_request, - convert_chat_response, - reframe_openai_stream) -from tensorrt_llm.serve.anthropic_protocol import (AnthropicCountTokensRequest, - AnthropicMessagesRequest) +from tensorrt_llm.serve.api.anthropic.adapter import anthropic_error_response +from tensorrt_llm.serve.api.anthropic.protocol import ( + AnthropicCountTokensRequest, AnthropicMessagesRequest) +from tensorrt_llm.serve.api.anthropic.serving import AnthropicMessagesServing from tensorrt_llm.serve.cluster_storage import ( HttpClusterStorageServer, create_cluster_storage, validate_http_cluster_storage_scope) @@ -62,14 +59,15 @@ from tensorrt_llm.serve.openai_disagg_service import ( OpenAIDisaggregatedService, ResponseHooks) from tensorrt_llm.serve.openai_protocol import ( - ChatCompletionRequest, ChatCompletionResponse, CompletionRequest, - ConversationParams, ErrorResponse, ResponsesRequest, UCompletionRequest, - UCompletionResponse, ensure_request_chat_template_allowed) + ChatCompletionRequest, CompletionRequest, ConversationParams, ErrorResponse, + ResponsesRequest, UCompletionRequest, UCompletionResponse, + ensure_request_chat_template_allowed) from tensorrt_llm.serve.perf_metrics import (DisaggPerfMetricsCollector, PerfMetricsJsonlWriter, 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) @@ -213,7 +211,72 @@ def _upstream_error_message(error: aiohttp.ClientResponseError) -> str: return f"context worker rejected the request with status {error.status}" +def _envelope_message(response: Response) -> str: + """The message of an error envelope, read the way the disaggregated /v1/messages did.""" + try: + payload = json.loads(response.body) + message = (payload.get("message") or payload.get("detail") + or payload.get("error") or json.dumps(payload)) + except (json.JSONDecodeError, AttributeError, TypeError): + message = "Internal server error" + return str(message) + + +def _disagg_serving_error(status: int, message: str) -> ServingError: + """A disaggregated pipeline failure as a ServingError; the message is final.""" + code: Union[int, str] = status + if is_context_length_exceeded_message(message): + code = CONTEXT_LENGTH_EXCEEDED_CODE + return ServingError(status, ErrorResponse( + message=message, + type="BadRequestError" if 400 <= status < 500 else "InternalServerError", + code=code)) + + +class _DisaggChatBackend: + """ChatBackend over the prepare, run and finish steps (no serialize).""" + + def __init__(self, server: "OpenAIDisaggServer") -> None: + self._server = server + + async def create_chat_completion(self, request: ChatCompletionRequest, + ctx: RequestContext): + try: + # The service is read per request, so tests may replace it. + outcome = await self._server._call_entry_point( + self._server._service.openai_chat_completion, request, + ctx.http_request) + except HTTPException as e: + raise _disagg_serving_error(e.status_code, str(e.detail)) from e + if outcome is None: + # Engine failure: _handle_exception has sent SIGINT. + raise EngineFatalError() + if isinstance(outcome, Response): + # The envelope _handle_exception returns for a worker error that + # carries a machine-readable code (context_length_exceeded). + raise _disagg_serving_error(outcome.status_code, + _envelope_message(outcome)) + return outcome + + def stream_model(self, request: ChatCompletionRequest) -> str: + # The client's name, as the disaggregated route always announced it. + return request.model + + def error_message(self, error: ServingError) -> str: + return error.body.message + + 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, @@ -270,6 +333,11 @@ def __init__(self, self._service = OpenAIDisaggregatedService( self._config, self._coordinator, self._create_client, req_timeout_secs=self._req_timeout_secs) + # The Anthropic Messages orchestration is shared with the aggregated + # server; here it runs over the steps of _wrap_entry_point, without + # the serialization the OpenAI routes need. + self._anthropic_messages_serving = AnthropicMessagesServing( + _DisaggChatBackend(self)) try: otlp_cfg = config.otlp_config @@ -426,93 +494,115 @@ 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 - async def anthropic_messages(self, request: AnthropicMessagesRequest, - raw_request: Request) -> Response: - """Serve Anthropic Messages through the disaggregated chat pipeline.""" - try: - chat_request = convert_anthropic_request(request) - except (AnthropicRequestError, ValidationError) as e: - return anthropic_error_response(str(e), "invalid_request_error", - 400) + @tracing.trace_span("disaggregated_request") + async def _call_entry_point(self, entry_point: Callable, + req: UCompletionRequest, raw_req: Request): + """Prepare, run and finish one request without serializing it. + Returns the service's result (a typed response or the worker stream) + or, on failure, what _finish_failed_request returns (an error + envelope, or None after an engine failure); its HTTPException + propagates. + """ + hooks: Optional[RawRequestResponseHooks] = None try: - openai_response = await self._wrap_entry_point( - self._service.openai_chat_completion)(chat_request, raw_request) - except HTTPException as e: - error_type = ("invalid_request_error" - if 400 <= e.status_code < 500 else "api_error") - return anthropic_error_response(str(e.detail), error_type, - e.status_code) - - if isinstance(openai_response, StreamingResponse): - return StreamingResponse( - content=reframe_openai_stream(openai_response.body_iterator, - model=request.model), - media_type="text/event-stream", - ) + hooks = self._prepare_request(req, raw_req) + result = await self._run_entry_point(entry_point, req, hooks, + raw_req) + self._finish_request(req, hooks) + return result + except Exception as e: + return self._finish_failed_request(e, req, hooks) - status = getattr(openai_response, "status_code", 500) - if status != 200: - try: - payload = json.loads(openai_response.body) - message = (payload.get("message") or payload.get("detail") - or payload.get("error") or json.dumps(payload)) - except (json.JSONDecodeError, AttributeError, TypeError): - message = "Internal server error" - error_type = ("invalid_request_error" - if 400 <= status < 500 else "api_error") - return anthropic_error_response(str(message), error_type, status) + 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: - chat_response = ChatCompletionResponse( - **json.loads(openai_response.body)) - anthropic_response = convert_chat_response(chat_response) - except (AnthropicResponseError, ValidationError, json.JSONDecodeError): - logger.error( - "Invalid response from OpenAI chat pipeline:\n" - f"{traceback.format_exc()}") - return anthropic_error_response("Internal server error", - "api_error", 500) - return JSONResponse(content=anthropic_response.model_dump( - exclude_none=True)) + 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.""" + return await self._anthropic_messages_serving.create_message( + request, RequestContext.from_http(raw_request)) async def get_model(self) -> Response: """List the served model, as a context worker reports it.""" diff --git a/tensorrt_llm/serve/openai_disagg_service.py b/tensorrt_llm/serve/openai_disagg_service.py index c62d62b31f8e..bc9e02945756 100644 --- a/tensorrt_llm/serve/openai_disagg_service.py +++ b/tensorrt_llm/serve/openai_disagg_service.py @@ -23,7 +23,7 @@ from tensorrt_llm._utils import get_global_steady_clock_now_in_seconds from tensorrt_llm.llmapi.disagg_utils import ConditionalDisaggConfig, DisaggServerConfig, ServerRole from tensorrt_llm.logger import logger -from tensorrt_llm.serve.anthropic_protocol import ( +from tensorrt_llm.serve.api.anthropic.protocol import ( AnthropicCountTokensRequest, AnthropicCountTokensResponse, ) diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py index eb79cdb18397..242995bab904 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -5,6 +5,7 @@ import array import asyncio import base64 +import dataclasses import functools import json import os @@ -22,12 +23,11 @@ from datetime import datetime from http import HTTPStatus from pathlib import Path -from typing import (TYPE_CHECKING, Annotated, Any, AsyncGenerator, - AsyncIterator, Dict, List, Optional, Tuple, Union) +from typing import (TYPE_CHECKING, Annotated, Any, AsyncIterator, Dict, List, + Optional, Union) -import msgspec import uvicorn -from fastapi import Body, Depends, FastAPI, HTTPException, Query, Request +from fastapi import Body, Depends, FastAPI, HTTPException, Request from fastapi.exceptions import RequestValidationError from fastapi.responses import (FileResponse, JSONResponse, Response, StreamingResponse) @@ -43,8 +43,6 @@ from tensorrt_llm.executor.postproc_worker import PostprocArgs, PostprocParams from tensorrt_llm.executor.request import DEFAULT_REQUEST_PRIORITY from tensorrt_llm.executor.rpc import RPCError -from tensorrt_llm.executor.utils import (CONTEXT_LENGTH_EXCEEDED_CODE, - is_context_length_exceeded_message) from tensorrt_llm.inputs import prompt_inputs from tensorrt_llm.inputs.data import TokensPrompt from tensorrt_llm.inputs.media_io import (BaseMediaIO, @@ -67,20 +65,19 @@ from tensorrt_llm.runtime.kv_cache_hash import \ get_effective_kv_cache_event_hash_algo from tensorrt_llm.sampling_params import GuidedDecodingParams, SamplingParams +from tensorrt_llm.serve._base.errors import (EngineFatalError, ServingError, + error_body, serving_error) +from tensorrt_llm.serve._base.msgspec_route import _MsgspecRoute +from tensorrt_llm.serve._base.request_context import RequestContext +from tensorrt_llm.serve._base.routing import (RoleGroup, ServerCapabilities, + attach_routers, role_group_of) +from tensorrt_llm.serve._base.sse import chat_stream_generator from tensorrt_llm.serve._perf_metrics_config import resolve_perf_metrics_config -from tensorrt_llm.serve.anthropic_adapter import ( - AnthropicRequestError, AnthropicResponseError, anthropic_error_response, - convert_anthropic_count_tokens_request, convert_anthropic_request, - convert_chat_response, reframe_openai_stream) -from tensorrt_llm.serve.anthropic_batches import (AnthropicBatchStore, - BatchStoreFullError, - results_to_jsonl) -from tensorrt_llm.serve.anthropic_protocol import (AnthropicBatchDeleteResponse, - AnthropicBatchList, - AnthropicCountTokensRequest, - AnthropicCountTokensResponse, - AnthropicCreateBatchRequest, - AnthropicMessagesRequest) +from tensorrt_llm.serve.api.anthropic.adapter import anthropic_error_response +from tensorrt_llm.serve.api.anthropic.serving import ( + AnthropicBatchesServing, AnthropicCountTokensServing, + AnthropicMessagesServing) +from tensorrt_llm.serve.api.ops.system.serving import SystemServing from tensorrt_llm.serve.chat_tokenization import tokenize_harmony_chat_request from tensorrt_llm.serve.chat_utils import (load_chat_template, parse_chat_messages_coroutines) @@ -104,14 +101,13 @@ ChatCompletionResponse, ChatCompletionResponseChoice, ChatCompletionToolsParam, ChatMessage, CompletionRequest, CompletionResponse, CompletionResponseChoice, EmbeddingRequest, - EmbeddingResponse, EmbeddingResponseData, EmbeddingUsageInfo, ErrorResponse, + EmbeddingResponse, EmbeddingResponseData, EmbeddingUsageInfo, ImageEditRequest, ImageGenerationRequest, ImageGenerationResponse, - ImageObject, MemoryUpdateRequest, ModelCard, ModelList, PromptTokensDetails, - ResponseFormat, ResponsesRequest, ResponsesResponse, - RuntimeMemoryUpdateRequest, StartProfileRequest, TokenizeRequest, - TokenizeResponse, UpdateWeightsRequest, UsageInfo, - ensure_request_chat_template_allowed, to_llm_conversation_params, - to_llm_disaggregated_params) + ImageObject, MemoryUpdateRequest, PromptTokensDetails, ResponseFormat, + ResponsesRequest, ResponsesResponse, RuntimeMemoryUpdateRequest, + StartProfileRequest, TokenizeRequest, TokenizeResponse, + UpdateWeightsRequest, UsageInfo, ensure_request_chat_template_allowed, + to_llm_conversation_params, to_llm_disaggregated_params) from tensorrt_llm.serve.openai_video_routes import _VideoRoutesMixin from tensorrt_llm.serve.perf_metrics import (PerfMetricsJsonlWriter, PerfMetricsMiddleware, @@ -123,11 +119,8 @@ chat_stream_post_processor, completion_response_post_processor, completion_stream_post_processor, responses_api_post_processor, responses_api_streaming_post_processor) -from tensorrt_llm.serve.render import (HarmonyConversionError, RenderResources, - UnsupportedRenderError, - mount_render_endpoints, - prepare_chat_request, render_chat, - render_chat_for_length, +from tensorrt_llm.serve.render import (RenderResources, mount_render_endpoints, + prepare_chat_request, render_conversation) from tensorrt_llm.serve.responses_utils import (ClosingStreamingResponse, ConversationHistoryStore, @@ -153,11 +146,10 @@ from tensorrt_llm.version import __version__ as VERSION from .._utils import (AdjustedSteadyClock, - get_global_steady_clock_now_in_seconds, nvtx_mark, + get_global_steady_clock_now_in_seconds, set_prometheus_multiproc_dir) from ._telemetry import create_uvicorn_server from .harmony_adapter import HarmonyAdapter, get_harmony_adapter -from .kv_hints import KvHintsRequest, KvHintsRoute if TYPE_CHECKING: from tensorrt_llm.visual_gen import VisualGen @@ -179,33 +171,6 @@ def _is_visual_gen_instance(obj) -> bool: # yapf: enable -_msgpack_decoder = msgspec.msgpack.Decoder() - - -class _MsgspecRequest(KvHintsRequest): - """Request that decodes X-TRTLLM-Msgpack bodies with msgspec.""" - - async def json(self): - if not hasattr(self, "_json_body"): - body = await self.body() - if not body: - self._json_body = {} - elif self.headers.get("x-trtllm-msgpack") == "1": - try: - self._json_body = _msgpack_decoder.decode(body) - except msgspec.DecodeError as e: - raise json.JSONDecodeError(f"msgpack: {e}", "", 0) from e - else: - self._json_body = json.loads(body) - return self.filter_kv_hints(self._json_body) - - -class _MsgspecRoute(KvHintsRoute): - """Parse msgpack/JSON and apply hint trust policy before model validation.""" - - request_type = _MsgspecRequest - - TIMEOUT_KEEP_ALIVE = 5 # seconds. _warned_unresolvable_thinking = False @@ -571,6 +536,30 @@ def resolve_spec_decode_num_spec_tokens(args: Any) -> Optional[int]: return getattr(spec_config, "max_draft_len", None) +class _AggregatedChatBackend: + """ChatBackend over OpenAIServer's typed chat core. + + Transitional: deleted when the chat core moves out of OpenAIServer into a + serving class that implements stream_model and error_message itself. + """ + + def __init__(self, server: "OpenAIServer") -> None: + self._server = server + + async def create_chat_completion(self, request: ChatCompletionRequest, + ctx: RequestContext): + return await self._server.create_chat_completion(request, ctx) + + def stream_model(self, request: ChatCompletionRequest) -> str: + # The served name, as the aggregated route always announced it. + return self._server.model + + def error_message(self, error: ServingError) -> str: + # What the rejected chat route's body said: its message, else the + # whole body. + return error.body.message or json.dumps(error.body.model_dump()) + + class OpenAIServer(_VideoRoutesMixin): @staticmethod @@ -907,6 +896,35 @@ async def validation_exception_handler(request, exc): len(detail) - 4096) return JSONResponse(status_code=400, content={"error": detail}) + # /health, /version and /v1/models for every role. The health + # probe stays here, with its private accesses, until the public + # health API exists. + self.app.state.ops_serving_system = SystemServing( + model=self.model, health_probe=self._health_probe) + if role_group_of(self.server_role) is RoleGroup.LLM: + # An attached frontend of a multi-frontend server forwards the + # launcher-owned routes, the Message Batches routes among them. + # They register before the routers so that they keep winning + # over the local handlers: FastAPI matches in registration order. + self._register_launcher_forwarding_routes() + messages_serving = AnthropicMessagesServing( + _AggregatedChatBackend(self)) + self.app.state.anthropic_serving_messages = messages_serving + # Transitional: RenderResources per request until it is built once. + self.app.state.anthropic_serving_count_tokens = ( + AnthropicCountTokensServing( + render_resources=lambda: RenderResources.from_server(self))) + self.app.state.anthropic_serving_batches = AnthropicBatchesServing( + runner=messages_serving.run_batched_message) + attach_routers( + self.app, + dataclasses.replace(ServerCapabilities.for_role(self.server_role), + harmony_model=(not self._is_visual_gen) + and self.use_harmony)) + # Runs before the role dispatch for now. It moves after the dispatch + # once the serving objects attach_router checks for are built per + # role. + if self.server_role is ServerRole.VISUAL_GEN: assert self._is_visual_gen, \ "generator must be a VisualGen for VISUAL_GEN server" @@ -923,7 +941,6 @@ async def validation_exception_handler(request, exc): self._init_embedding_batcher() self.register_embedding_routes() else: - self._register_launcher_forwarding_routes() self.register_routes() if self._collect_perf_metrics: @@ -1148,15 +1165,7 @@ def create_error_response( message: str, err_type: str = "BadRequestError", status_code: HTTPStatus = HTTPStatus.BAD_REQUEST) -> Response: - # Over-length prompts are detected from the message text because they - # cross the executor IPC boundary as plain strings (RequestError), so - # no exception type survives to here. - code: Union[int, str] = status_code.value - if is_context_length_exceeded_message(message): - code = CONTEXT_LENGTH_EXCEEDED_CODE - error_response = ErrorResponse(message=message, - type=err_type, - code=code) + error_response = error_body(message, err_type, status_code) return JSONResponse(content=error_response.model_dump(), status_code=status_code.value) @@ -1223,12 +1232,9 @@ def _check_health(self) -> bool: return True def register_routes(self): - self.app.add_api_route("/health", self.health, methods=["GET"]) self.app.add_api_route("/health_generate", self.health_generate, methods=["GET"]) - self.app.add_api_route("/version", self.version, methods=["GET"]) - self.app.add_api_route("/v1/models", self.get_model, methods=["GET"]) self.app.add_api_route("/v1/data_transceiver_state", self.data_transceiver_state, methods=["GET"]) @@ -1285,38 +1291,6 @@ def register_routes(self): "/v1/chat/completions", self.openai_chat if not self.use_harmony else self.chat_harmony, methods=["POST"]) - # Anthropic Messages API adapter (e.g. Claude Code as a client). - # Not supported together with the harmony chat path (gpt-oss). - if not self.use_harmony: - self.app.add_api_route("/v1/messages", - self.anthropic_messages, - methods=["POST"]) - # Claude Code calls this before most turns to decide whether to - # compact its context. Leaving it unregistered 404s every call, - # which the client cannot act on. - self.app.add_api_route("/v1/messages/count_tokens", - self.anthropic_count_tokens, - methods=["POST"]) - # Message Batches. Registered before /v1/messages/{...} would be, - # so the literal paths win; FastAPI matches in registration order. - self.app.add_api_route("/v1/messages/batches", - self.anthropic_create_batch, - methods=["POST"]) - self.app.add_api_route("/v1/messages/batches", - self.anthropic_list_batches, - methods=["GET"]) - self.app.add_api_route("/v1/messages/batches/{batch_id}", - self.anthropic_get_batch, - methods=["GET"]) - self.app.add_api_route("/v1/messages/batches/{batch_id}", - self.anthropic_delete_batch, - methods=["DELETE"]) - self.app.add_api_route("/v1/messages/batches/{batch_id}/cancel", - self.anthropic_cancel_batch, - methods=["POST"]) - self.app.add_api_route("/v1/messages/batches/{batch_id}/results", - self.anthropic_batch_results, - methods=["GET"]) self.app.add_api_route("/v1/responses", self.openai_responses, methods=["POST"]) @@ -1373,9 +1347,6 @@ def mount_metrics(self): self.app.routes.append(metrics_route) def register_mm_encoder_routes(self): - self.app.add_api_route("/health", self.health, methods=["GET"]) - self.app.add_api_route("/version", self.version, methods=["GET"]) - self.app.add_api_route("/v1/models", self.get_model, methods=["GET"]) # TODO: the metrics endpoint only reports iteration stats, not the runtime stats for now self.app.add_api_route("/metrics", self.get_iteration_stats, @@ -1480,9 +1451,6 @@ def encode_fn(token_ids_batch): ) def register_embedding_routes(self): - self.app.add_api_route("/health", self.health, methods=["GET"]) - self.app.add_api_route("/version", self.version, methods=["GET"]) - self.app.add_api_route("/v1/models", self.get_model, methods=["GET"]) self.app.add_api_route("/v1/embeddings", self.openai_embedding, methods=["POST"]) @@ -1587,10 +1555,6 @@ async def openai_embedding(self, request: EmbeddingRequest, def register_visual_gen_routes(self): """Register routes for diffusion model serving.""" - # Health and info endpoints - self.app.add_api_route("/health", self.health, methods=["GET"]) - self.app.add_api_route("/version", self.version, methods=["GET"]) - self.app.add_api_route("/v1/models", self.get_model, methods=["GET"]) self.app.add_api_route("/metrics", self.get_iteration_stats, methods=["GET"]) @@ -1707,9 +1671,9 @@ def _on_launcher_lost(self, error: BaseException) -> None: _record_generator_termination(self.generator) signal.raise_signal(signal.SIGINT) - async def health(self) -> Response: + def _health_probe(self) -> bool: if self._check_health(): - return Response(status_code=200) + return True else: # If the engine has a fatal error, trigger server shutdown so the # pod doesn't linger as a zombie that accepts connections but never @@ -1727,11 +1691,7 @@ async def health(self) -> Response: f"server shutdown: {executor._fatal_error}") _record_generator_termination(self.generator) signal.raise_signal(signal.SIGINT) - return Response( - status_code=503, - content= - "LLM is unavailable. Please check the server logs for more details." - ) + return False async def health_generate(self, raw_request: Request) -> Response: """Health check that performs a minimal generation.""" @@ -1782,14 +1742,6 @@ async def health_generate(self, raw_request: Request) -> Response: return Response(status_code=500, content=f"Generation health check failed: {str(e)}") - async def version(self) -> JSONResponse: - ver = {"version": VERSION} - return JSONResponse(content=ver) - - async def get_model(self) -> JSONResponse: - model_list = ModelList(data=[ModelCard(id=self.model)]) - return JSONResponse(content=model_list.model_dump()) - async def get_iteration_stats(self) -> JSONResponse: # Visual-gen deployments use a separate iteration-stats producer # (see ``DiffusionRemoteClient._iter_stats``). The LLM PyExecutor @@ -2056,6 +2008,43 @@ def _record_iteration_stat(self, llm_stat: dict) -> None: async def openai_chat(self, request: ChatCompletionRequest, raw_request: Request) -> Response: + try: + result = await self.create_chat_completion( + request, RequestContext.from_http(raw_request)) + if isinstance(result, ChatCompletionResponse): + return JSONResponse(content=result.model_dump()) + return StreamingResponse(content=result, + media_type="text/event-stream") + except ServingError as e: + return JSONResponse(content=e.body.model_dump(), + status_code=e.status) + except EngineFatalError: + # Answered as HTTP 200 with body null (FastAPI's rendering of + # None); the in-process callers map that to their own errors. + # Answering 5xx would be a behavior change of its own. + return None + except ValueError as e: + # JSONResponse encodes in its constructor (allow_nan=False), so a + # response that cannot be encoded fails here and keeps its 400. + return self.create_error_response(str(e)) + except Exception as e: + logger.error(traceback.format_exc()) + return self.create_error_response(str(e)) + + async def create_chat_completion( + self, request: ChatCompletionRequest, ctx: RequestContext + ) -> Union[ChatCompletionResponse, AsyncIterator[str]]: + """Run one chat completion and return a typed result. + + Returns the ChatCompletionResponse for a non-streaming request, and + for a streaming one the SSE stream as strings, ending with + "data: [DONE]". + + Raises: + ServingError: The request is rejected; the HTTP status travels + apart from the body. + EngineFatalError: The engine failed; SIGINT has been sent. + """ def get_role() -> str: if request.add_generation_prompt: @@ -2064,38 +2053,6 @@ def get_role() -> str: role = request.messages[-1]["role"] return role - async def chat_stream_generator( - promise: RequestOutput, - postproc_params: PostprocParams) -> AsyncGenerator[str, None]: - try: - if not self.postproc_worker_enabled: - post_processor, args = postproc_params.post_processor, postproc_params.postproc_args - first_response = await anext(promise) - raw_request.state.server_first_token_time = self._adjusted_steady_clock.now( - ) - pp_results = first_response.outputs[ - 0]._postprocess_result if self.postproc_worker_enabled else post_processor( - first_response, args) - for pp_res in pp_results: - yield pp_res - # Making sure we can handle the situation where there is only one response - res = first_response - async for res in promise: - pp_results = res.outputs[ - 0]._postprocess_result if self.postproc_worker_enabled else post_processor( - res, args) - for pp_res in pp_results: - yield pp_res - yield "data: [DONE]\n\n" - await self._extract_metrics(res, raw_request) - nvtx_mark("generation ends") - except Exception: - # Only failures are logged: a consumer that stops reading closes - # the stream with GeneratorExit, and a client that disconnects - # cancels it, and neither is an error of this stream. - logger.error(traceback.format_exc()) - raise - try: # Extension preprocessing, the required-tool check and tool dumping # are the first stage of the shared prompt-preparation pipeline. @@ -2151,7 +2108,7 @@ async def chat_stream_generator( if isinstance(request.tool_choice, ChatCompletionNamedToolChoiceParam): if not request.tools: - return self.create_error_response( + raise serving_error( message=("tool_choice with a named function requires " "`tools` to be set."), err_type="BadRequestError", @@ -2169,7 +2126,7 @@ async def chat_stream_generator( # testing ``prompt_token_ids`` alone lets a base64 request # through to the pre-tokenized branch, where the begin # prefix is never appended. - return self.create_error_response( + raise serving_error( message=("tool_choice with a named function is not " "supported alongside `prompt_token_ids`; pass " "`messages` instead."), @@ -2184,7 +2141,7 @@ async def chat_stream_generator( # instead. Honouring only one of them silently would either # ignore an explicit instruction or emit a tool call whose # arguments are unconstrained text, so reject instead. - return self.create_error_response( + raise serving_error( message=("tool_choice with a named function cannot be " "combined with `response_format` or other " "guided-decoding settings."), @@ -2207,14 +2164,13 @@ async def chat_stream_generator( # ends up reported as ``function.arguments``. Fail loudly # instead of returning a tool call whose arguments are # unconstrained text. - return self.create_error_response( - message=( - "tool_choice with a named function requires the " - "server to be started with a guided decoding " - "backend (set `guided_decoding_backend` to " - "'xgrammar' or 'llguidance')."), - err_type="BadRequestError", - status_code=HTTPStatus.BAD_REQUEST) + raise serving_error(message=( + "tool_choice with a named function requires the " + "server to be started with a guided decoding " + "backend (set `guided_decoding_backend` to " + "'xgrammar' or 'llguidance')."), + err_type="BadRequestError", + status_code=HTTPStatus.BAD_REQUEST) forced_tool_name = request.tool_choice.function.name reasoning_parser_name = self.generator.args.reasoning_parser @@ -2267,7 +2223,7 @@ async def chat_stream_generator( request.tools, self.tool_parser, forced_tool_name)) except ValueError as e: - return self.create_error_response( + raise serving_error( message=str(e), err_type="BadRequestError", status_code=HTTPStatus.BAD_REQUEST) @@ -2282,7 +2238,7 @@ async def chat_stream_generator( sampling_params.guided_decoding = strict_guided elif forced_tool_name is not None: # tools were provided but the server has no tool_parser configured. - return self.create_error_response( + raise serving_error( message=("tool_choice with a named function requires the " "server to be started with --tool_parser."), err_type="BadRequestError", @@ -2304,7 +2260,7 @@ async def chat_stream_generator( # calls in the model output. postproc_args.tools = (postproc_args.tools or []) + dynamic_tool_params - self._validate_internal_disagg_request(request, raw_request) + self._validate_internal_disagg_request(request, ctx) disaggregated_params = to_llm_disaggregated_params( request.disaggregated_params) @@ -2336,7 +2292,7 @@ async def chat_stream_generator( ) except ValidationError: # ValidatorIterator rejects extra fields; fall back to raw JSON. - raw_body = await raw_request.json() + raw_body = await ctx.json() raw_messages = raw_body.get("messages", []) conversation, mm_coroutines, mm_placeholder_counts, mm_item_order = parse_chat_messages_coroutines( raw_messages, @@ -2457,15 +2413,14 @@ async def chat_stream_generator( postproc_args=postproc_args, ) - trace_headers = (None if raw_request is None else - tracing.extract_trace_headers(raw_request.headers)) + trace_headers = (None if ctx is None else + tracing.extract_trace_headers(ctx.headers)) resolve_request_conversation_id( - request, None if raw_request is None else raw_request.headers) + request, None if ctx is None else ctx.headers) conversation_params = to_llm_conversation_params( request.conversation_params) - scheduling_params = self._get_scheduling_params( - request, raw_request) + scheduling_params = self._get_scheduling_params(request, ctx) generate_inputs = prompt preprocess_fn = getattr(self.generator, "preprocess", None) @@ -2492,19 +2447,19 @@ async def chat_stream_generator( priority=request.priority if request.priority is not None else DEFAULT_REQUEST_PRIORITY, ) - asyncio.create_task(self.await_disconnected(raw_request, promise)) + asyncio.create_task(self.await_disconnected(ctx, promise)) if not self.postproc_worker_enabled: postproc_args.tokenizer = self.tokenizer postproc_args.num_prompt_tokens = len(promise.prompt_token_ids) if request.stream: - response_generator = chat_stream_generator( - promise, postproc_params) - return StreamingResponse(content=response_generator, - media_type="text/event-stream") + return chat_stream_generator(promise, + postproc_params, + server=self, + ctx=ctx) else: response = await self._create_chat_response( - promise, postproc_params, raw_request, disaggregated_params) + promise, postproc_params, ctx, disaggregated_params) # Context-only: optionally return prompt_token_ids as a base64 # int32 buffer (one string) so the disagg orchestrator relays it # without materializing the int list on its event loop. The encode @@ -2517,313 +2472,20 @@ async def chat_stream_generator( np.asarray(response.prompt_token_ids, dtype=np.int32).tobytes()).decode("ascii") response.prompt_token_ids = None - return JSONResponse(content=response.model_dump()) - except CppExecutorError: + return response + except ServingError: + raise + except CppExecutorError as e: logger.error(traceback.format_exc()) # If internal executor error is raised, shutdown the server _record_generator_termination(self.generator) signal.raise_signal(signal.SIGINT) + raise EngineFatalError() from e except ValueError as e: - return self.create_error_response(str(e)) + raise serving_error(str(e)) from e except Exception as e: logger.error(traceback.format_exc()) - return self.create_error_response(str(e)) - - async def anthropic_messages(self, request: AnthropicMessagesRequest, - raw_request: Request) -> Response: - """Serve the Anthropic Messages API on top of the openai_chat path. - - The Anthropic request is translated into a ChatCompletionRequest, - handed to ``openai_chat`` unchanged (chat template, tool parser and - post-processing are reused verbatim), and the OpenAI-shaped result is - translated back: JSON body for non-streaming, SSE reframing for - streaming. - """ - try: - chat_request = convert_anthropic_request(request) - except AnthropicRequestError as e: - return anthropic_error_response(str(e), "invalid_request_error", - 400) - except ValidationError as e: - return anthropic_error_response(str(e), "invalid_request_error", - 400) - - response = await self.openai_chat(chat_request, raw_request) - if response is None: - return anthropic_error_response("Internal server error", - "api_error", 500) - - if isinstance(response, StreamingResponse): - return StreamingResponse( - content=reframe_openai_stream(response.body_iterator, - model=self.model), - media_type="text/event-stream", - ) - - status = getattr(response, "status_code", 500) - if status != 200: - try: - payload = json.loads(response.body) - message = payload.get("message") or json.dumps(payload) - except (json.JSONDecodeError, AttributeError): - message = "Internal server error" - err_type = ("invalid_request_error" - if 400 <= status < 500 else "api_error") - return anthropic_error_response(message, err_type, status) - - try: - chat_response = ChatCompletionResponse(**json.loads(response.body)) - anthropic_response = convert_chat_response(chat_response) - except (AnthropicResponseError, ValidationError, json.JSONDecodeError): - logger.error("Invalid response from OpenAI chat pipeline:\n" - f"{traceback.format_exc()}") - return anthropic_error_response("Internal server error", - "api_error", 500) - return JSONResponse(content=anthropic_response.model_dump( - exclude_none=True)) - - async def anthropic_count_tokens( - self, request: AnthropicCountTokensRequest) -> Response: - """Report how many input tokens a Messages request would consume. - - The count is taken over the rendered prompt rather than over the raw - messages, so system blocks, tool definitions and the thinking prefix - are all included - those are exactly what make a client's own estimate - drift from the server's. - """ - try: - chat_request = convert_anthropic_count_tokens_request(request) - except AnthropicRequestError as e: - return anthropic_error_response(str(e), "invalid_request_error", - 400) - except ValidationError as e: - return anthropic_error_response(str(e), "invalid_request_error", - 400) - - try: - - # Both the render and the encode are synchronous and both scale - # with prompt size. Claude Code calls this before most turns with - # the whole conversation attached, so running them inline would - # block the event loop -- and therefore every other in-flight - # request on this server -- for a full template render plus a - # tokenizer pass. The chat path avoids this the same way. - def _count() -> int: - # The count is over the prompt tokens the chat route would - # execute: same pipeline, same template, same tokenization. - resources = RenderResources.from_server(self) - try: - rendered = render_chat(chat_request, - resources, - tokenize=True) - except (UnsupportedRenderError, HarmonyConversionError): - # Media, an own-processor model or an unconvertible Harmony - # request: count the text, tools and template without the media. - rendered = render_chat_for_length(chat_request, resources) - return len(rendered.token_ids) - - input_tokens = await asyncio.to_thread(_count) - except Exception: - logger.error("count_tokens failed to render the prompt:\n" - f"{traceback.format_exc()}") - return anthropic_error_response("Internal server error", - "api_error", 500) - - return JSONResponse(content=AnthropicCountTokensResponse( - input_tokens=input_tokens).model_dump()) - - # -- Message Batches --------------------------------------------------- - - @staticmethod - def _synthetic_request(body: Dict[str, Any]) -> Request: - """A Request for work with no client connection behind it. - - The chat path reads ``raw_request.state`` and may re-read the body via - ``.json()``. A batched request has no live connection, so rather than - duck-type a stub that drifts as that path evolves, build a real Request - over a synthetic ASGI scope - anything the chat path can legitimately - do to a Request keeps working. - """ - payload = json.dumps(body).encode() - - async def receive() -> Dict[str, Any]: - return {"type": "http.request", "body": payload, "more_body": False} - - return Request( - { - "type": "http", - "http_version": "1.1", - "method": "POST", - "scheme": "http", - "path": "/v1/messages", - "raw_path": b"/v1/messages", - "query_string": b"", - "root_path": "", - "headers": [(b"content-type", b"application/json")], - "client": None, - "server": None, - # PerfMetricsMiddleware seeds state["perf_metrics_records"] on a - # real connection (see perf_metrics.py). A batched request never - # passes through the middleware, so seed it here: with - # return_perf_metrics enabled, _extract_metrics appends to this - # list unconditionally and every batched request would otherwise - # die with AttributeError and come back "errored". - # No socket sits behind this request, so await_disconnected - # keys off no_client_connection to avoid polling forever on a - # connection that can never drop. - "state": { - "perf_metrics_records": [], - "no_client_connection": True - }, - }, - receive, - ) - - def _batch_store(self) -> AnthropicBatchStore: - """Created on first use so servers that never batch pay nothing.""" - if getattr(self, "_anthropic_batch_store", None) is None: - self._anthropic_batch_store = AnthropicBatchStore( - runner=self._run_batched_message) - return self._anthropic_batch_store - - async def _run_batched_message( - self, - request: AnthropicMessagesRequest) -> Tuple[str, Dict[str, Any]]: - """Run one batched Messages request through the ordinary chat path.""" - # Streaming makes no sense for a batch: the result is collected, not - # watched. Forcing it off here means a client that sets stream=true in - # params gets a usable message instead of an SSE body in its .jsonl. - request = request.model_copy(update={"stream": False}) - try: - chat_request = convert_anthropic_request(request) - except (AnthropicRequestError, ValidationError) as e: - return "errored", { - "type": "invalid_request_error", - "message": str(e) - } - - raw_request = self._synthetic_request( - chat_request.model_dump(exclude_none=True)) - response = await self.openai_chat(chat_request, raw_request) - if response is None: - return "errored", { - "type": "api_error", - "message": "Internal server error" - } - - status = getattr(response, "status_code", 500) - if status != 200: - try: - payload = json.loads(response.body) - message = payload.get("message") or json.dumps(payload) - except (json.JSONDecodeError, AttributeError): - message = "Internal server error" - err_type = "invalid_request_error" if 400 <= status < 500 else "api_error" - return "errored", {"type": err_type, "message": message} - - try: - chat_response = ChatCompletionResponse(**json.loads(response.body)) - anthropic_response = convert_chat_response(chat_response) - except (AnthropicResponseError, ValidationError, json.JSONDecodeError): - logger.error( - "Invalid response from OpenAI chat pipeline in batch:\n" - f"{traceback.format_exc()}") - return "errored", { - "type": "api_error", - "message": "Internal server error" - } - return "succeeded", anthropic_response.model_dump(exclude_none=True) - - async def anthropic_create_batch( - self, request: AnthropicCreateBatchRequest) -> Response: - try: - batch = self._batch_store().create(request.requests) - except BatchStoreFullError as e: - # 429, not 400: the request is valid and an identical one will - # succeed later, so the client should back off rather than edit it. - return anthropic_error_response(str(e), "rate_limit_error", 429) - except ValueError as e: - return anthropic_error_response(str(e), "invalid_request_error", - 400) - return JSONResponse(content=batch.model_dump()) - - async def anthropic_list_batches( - self, - # Validated here rather than clamped in the store: clamping turned - # limit=0 into 1 instead of a 400, and bounded the page against the - # retention limit, which is an unrelated quantity. - limit: int = Query(20, ge=1, le=100), - after_id: Optional[str] = None, - before_id: Optional[str] = None) -> Response: - try: - batches, has_more = self._batch_store().list(limit=limit, - after_id=after_id, - before_id=before_id) - except LookupError as e: - # Saying so beats an empty page: an empty page is indistinguishable - # from a finished walk, so the client stops early having silently - # skipped every batch past the cursor. - return anthropic_error_response( - f"Unknown pagination cursor {e.args[0]!r}. The batch may have " - "been deleted or expired; restart the listing without a cursor.", - "invalid_request_error", 400) - listing = AnthropicBatchList( - data=batches, - # Reported honestly rather than hardcoded false: a client that - # asked for N and received N cannot otherwise tell a full page from - # the end of the list, and silently stops early. It pages on with - # ?after_id=. - has_more=has_more, - first_id=batches[0].id if batches else None, - last_id=batches[-1].id if batches else None, - ) - return JSONResponse(content=listing.model_dump()) - - async def anthropic_get_batch(self, batch_id: str) -> Response: - batch = self._batch_store().get(batch_id) - if batch is None: - return self._no_such_batch(batch_id) - return JSONResponse(content=batch.model_dump()) - - async def anthropic_cancel_batch(self, batch_id: str) -> Response: - batch = self._batch_store().cancel(batch_id) - if batch is None: - return self._no_such_batch(batch_id) - return JSONResponse(content=batch.model_dump()) - - async def anthropic_delete_batch(self, batch_id: str) -> Response: - try: - deleted = self._batch_store().delete(batch_id) - except ValueError as e: - return anthropic_error_response(str(e), "invalid_request_error", - 400) - if not deleted: - return self._no_such_batch(batch_id) - return JSONResponse(content=AnthropicBatchDeleteResponse( - id=batch_id).model_dump()) - - async def anthropic_batch_results(self, batch_id: str) -> Response: - store = self._batch_store() - results = store.results(batch_id) - if results is None: - if not store.has(batch_id): - return self._no_such_batch(batch_id) - # Known batch, but not finished. 404 would say "no such batch", - # which is wrong and would stop a client from polling. - return anthropic_error_response( - f"Message Batch {batch_id} is still processing; results are " - "available once processing_status is 'ended'", - "invalid_request_error", 400) - return Response(content=results_to_jsonl(results), - media_type="application/x-ndjson") - - @staticmethod - def _no_such_batch(batch_id: str) -> Response: - # Batches are held in memory only, so an id from before a restart lands - # here. Say so, rather than leaving the caller to wonder. - return anthropic_error_response( - f"Message Batch {batch_id} not found. Batches are stored in memory " - "and do not survive a server restart.", "not_found_error", 404) + raise serving_error(str(e)) from e async def openai_mm_encoder(self, request: ChatCompletionRequest, raw_request: Request) -> Response: diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index 41e0cabd17ea..af02b6fa2746 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -341,6 +341,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 diff --git a/tests/test_common/serve/__init__.py b/tests/test_common/serve/__init__.py new file mode 100644 index 000000000000..467079831e16 --- /dev/null +++ b/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. diff --git a/tests/test_common/serve/app_factory.py b/tests/test_common/serve/app_factory.py new file mode 100644 index 000000000000..1e958e6b44f7 --- /dev/null +++ b/tests/test_common/serve/app_factory.py @@ -0,0 +1,366 @@ +# 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. +"""Factory that builds every trtllm-serve app on CPU through the real constructors. + +``isolated_serve_env`` is the only entry point: building a server touches +process-global state (the default Prometheus registry, ``BaseMediaIO``'s +executor, environment variables, the prometheus multiprocess directory, the +telemetry session), and the context manager both prepares and restores all of +it. Use it instead of ``object.__new__(OpenAIServer)``: + + with isolated_serve_env(tmp_path) as env: + server = env.build_openai_server(AggConfig()) + client = env.client(server.app) + ... + +or the ``serve_env`` fixture from ``tests/unittest/serve/conftest.py``. + +Inside an environment, change the state it saves through the environment +(``AggConfig.env``, ``record_terminations()``), not through ``monkeypatch``: +pytest may undo a monkeypatch change after the environment's exit, which then +overwrites the value the environment restored. +""" + +import contextlib +import json +import os +from contextlib import contextmanager +from dataclasses import dataclass +from pathlib import Path +from types import SimpleNamespace +from typing import Iterator, List, Optional +from unittest import mock + +from fastapi import FastAPI +from fastapi.testclient import TestClient +from prometheus_client import REGISTRY + +import tensorrt_llm._utils +from tensorrt_llm.inputs.media_io import BaseMediaIO +from tensorrt_llm.llmapi.disagg_utils import ( + CtxGenServerConfig, + DisaggClusterConfig, + DisaggServerConfig, + RouterConfig, + ServerRole, +) +from tensorrt_llm.serve.coordinator_server import CoordinatorServer +from tensorrt_llm.serve.disagg_coordinator import DisaggCoordinatorService +from tensorrt_llm.serve.openai_disagg_server import OpenAIDisaggServer +from tensorrt_llm.serve.openai_server import OpenAIServer +from tensorrt_llm.usage import usage_lib +from test_common.serve.fakes import ( + FakeAsyncLLM, + FakeEncodeOnlyLLM, + FakeLLM, + FakeLlmArgs, + FakeMultimodalEncoder, +) + +__all__ = [ + "AggConfig", + "CoordinatorConfig", + "DisaggConfig", + "ServeTestEnv", + "isolated_serve_env", +] + +_MISSING = object() + +# Environment variables that change route registration or request handling; +# removed on entry so every configuration starts from the same baseline and +# sets what it needs through AggConfig.env. +_CLEARED_ENV_VARS = ( + "ENABLE_METRICS", + "TRTLLM_ENABLE_RENDER_ENDPOINTS", + "DISABLE_HARMONY_ADAPTER", + "TRTLLM_RESPONSES_API_DISABLE_STORE", + "TRTLLM_KVCACHE_TIME_OUTPUT_PATH", +) + + +@dataclass(frozen=True) +class AggConfig: + """One ``OpenAIServer`` configuration, built through the real constructor.""" + + role: Optional[ServerRole] = None + model_type: str = "llama" # "gpt_oss" makes use_harmony true + governor_queue: bool = False # FakeExecutor(resource_governor_queue=object()) + rl_control: bool = False # FakeAsyncLLM + enable_rl_control_endpoints + runtime_control: bool = False # sleep_config + enable_runtime_control_endpoints + return_perf_metrics: bool = False + perf_metrics_output_dir: bool = False # sets args.perf_metrics_output_dir=/perf + env: tuple = () # (name, value) pairs applied before construction + model_name: str = "fake-model" # OpenAIServer(model=...): the /v1/models id + tool_parser: Optional[str] = None + reasoning_parser: Optional[str] = None + num_postprocess_workers: int = 0 # > 0 takes the postprocess-worker path + # Attached frontend of a multi-frontend server: registers the launcher + # forwarding routes ahead of the local handlers, which they shadow. + multi_frontend_attached: bool = False + + +@dataclass(frozen=True) +class DisaggConfig: + """One ``OpenAIDisaggServer`` configuration.""" + + fleet: bool = False # delegate to a remote coordinator (CoordinatorClient) + http_cluster_storage: bool = False + return_perf_metrics: bool = False + + +@dataclass(frozen=True) +class CoordinatorConfig: + """One ``CoordinatorServer`` configuration (the fleet-mode coordinator app).""" + + http_cluster_storage: bool = False + + +def _make_disagg_server_config( + *, return_perf_metrics: bool = False, http_cluster_storage: bool = False +) -> DisaggServerConfig: + # Loopback addresses on closed ports: construction is network-free, and + # nothing in these tests runs the disaggregated lifespan. + return DisaggServerConfig( + server_configs=[ + CtxGenServerConfig(type="ctx", hostname="127.0.0.1", port=9), + CtxGenServerConfig(type="gen", hostname="127.0.0.1", port=10), + ], + hostname="127.0.0.1", + port=0, + ctx_router_config=RouterConfig(server_role=ServerRole.CONTEXT), + gen_router_config=RouterConfig(server_role=ServerRole.GENERATION), + node_id=0, + return_perf_metrics=return_perf_metrics, + disagg_cluster_config=( + DisaggClusterConfig(cluster_uri="http://127.0.0.1:9", cluster_name="test") + if http_cluster_storage + else None + ), + ) + + +class _StubTelemetrySession: + """Collects termination observations for ``ServeTestEnv.record_terminations``.""" + + def __init__(self, sink: list) -> None: + self.owner_pid = os.getpid() + self._sink = sink + + def record_termination_observation(self, outcome, lifecycle_phase=None) -> None: + self._sink.append(outcome) + + +class ServeTestEnv: + """Handle returned by ``isolated_serve_env``; builds servers and clients.""" + + def __init__(self, tmp_dir: Path) -> None: + self.tmp_dir = Path(tmp_dir) + self._servers: List[OpenAIServer] = [] + self._entered_clients: List[TestClient] = [] + self._plain_clients: List[TestClient] = [] + + # -- builders --------------------------------------------------------- + + def _write_fake_checkpoint(self, model_type: str) -> Path: + """A directory with only a config.json, so load_pretrained_config works. + + No tokenizer files are written: they would make AutoProcessor succeed + and change ``server.processor`` from its fake-generator value (None). + """ + model_dir = self.tmp_dir / f"model-{model_type}" + model_dir.mkdir(parents=True, exist_ok=True) + architectures = { + "llama": ["LlamaForCausalLM"], + "gpt_oss": ["GptOssForCausalLM"], + }[model_type] + (model_dir / "config.json").write_text( + json.dumps({"model_type": model_type, "architectures": architectures}) + ) + return model_dir + + def build_openai_server(self, config: AggConfig) -> OpenAIServer: + from test_common.serve.fakes import FakeExecutor # local: keeps the import graph flat + + for name, value in config.env: + os.environ[name] = value + + args = FakeLlmArgs( + model=str(self._write_fake_checkpoint(config.model_type)), + return_perf_metrics=config.return_perf_metrics, + perf_metrics_output_dir=( + str(self.tmp_dir / "perf") if config.perf_metrics_output_dir else None + ), + sleep_config=SimpleNamespace() if config.runtime_control else None, + encode_only=config.role is ServerRole.EMBEDDING, + reasoning_parser=config.reasoning_parser, + num_postprocess_workers=config.num_postprocess_workers, + ) + executor = FakeExecutor(resource_governor_queue=object() if config.governor_queue else None) + if config.role is ServerRole.VISUAL_GEN: + # Lazy import: only VisualGen tests pay the visual_gen import cost, + # and importing it is what lets _is_visual_gen_instance say yes. + from test_common.serve.fakes_visual_gen import FakeVisualGen + + generator = FakeVisualGen(model=config.model_name) + elif config.role is ServerRole.MM_ENCODER: + generator = FakeMultimodalEncoder(args=args, executor=executor) + elif config.role is ServerRole.EMBEDDING: + generator = FakeEncodeOnlyLLM(args=args) + elif config.rl_control: + generator = FakeAsyncLLM(args=args, executor=executor) + else: + generator = FakeLLM(args=args, executor=executor) + + multi_frontend_serving = None + if config.multi_frontend_attached: + from tensorrt_llm.serve.multi_frontend import MultiFrontendServing + + ipc_dir = self.tmp_dir / "multi_frontend_ipc" + ipc_dir.mkdir(parents=True, exist_ok=True) + multi_frontend_serving = MultiFrontendServing( + launcher_uds=str(ipc_dir / "launcher.sock"), is_launcher=False + ) + + server = OpenAIServer( + generator=generator, + model=config.model_name, + tool_parser=config.tool_parser, + server_role=config.role, + metadata_server_cfg=None, + enable_rl_control_endpoints=config.rl_control, + rl_control_api_key="test-key" if config.rl_control else None, + enable_runtime_control_endpoints=config.runtime_control, + runtime_control_api_key="test-key" if config.runtime_control else None, + multi_frontend_serving=multi_frontend_serving, + ) + self._servers.append(server) + return server + + def build_disagg_server(self, config: DisaggConfig) -> OpenAIDisaggServer: + # At most one per environment: each registers the disaggregated server + # counters in the default Prometheus registry, and a second one fails + # with "Duplicated timeseries". + server_config = _make_disagg_server_config( + return_perf_metrics=config.return_perf_metrics, + http_cluster_storage=config.http_cluster_storage, + ) + return OpenAIDisaggServer( + server_config, + coordinator_url="http://127.0.0.1:9" if config.fleet else None, + ) + + def build_coordinator_server(self, config: CoordinatorConfig) -> CoordinatorServer: + server_config = _make_disagg_server_config(http_cluster_storage=config.http_cluster_storage) + coordinator = DisaggCoordinatorService(server_config, lambda *args, **kwargs: None) + return CoordinatorServer(coordinator) + + # -- request plumbing ------------------------------------------------- + + def client(self, app: FastAPI, *, lifespan: bool = True) -> TestClient: + """A TestClient closed at environment exit; runs the lifespan by default.""" + client = TestClient(app) + if lifespan: + client.__enter__() + self._entered_clients.append(client) + else: + self._plain_clients.append(client) + return client + + def record_terminations(self) -> list: + """Route termination observations of this process into the returned list. + + Installs a stub telemetry session; the environment's exit restores the + session that existed before entry. + """ + sink: list = [] + usage_lib._SESSION = _StubTelemetrySession(sink) + return sink + + # -- teardown helpers (called by isolated_serve_env) ------------------- + + def _close(self) -> None: + for client in self._entered_clients: + client.__exit__(None, None, None) + for client in self._plain_clients: + client.close() + self._entered_clients.clear() + self._plain_clients.clear() + for server in self._servers: + for pool_name in ( + "_input_proc_executor", + "_media_load_executor", + "_metrics_stats_executor", + ): + pool = getattr(server, pool_name, None) + if pool is not None: + pool.shutdown(wait=False, cancel_futures=True) + self._servers.clear() + + +@contextmanager +def isolated_serve_env(tmp_dir: Path) -> Iterator[ServeTestEnv]: + """Process-global state that building trtllm-serve apps touches, restored on exit. + + The only supported forms are ``with isolated_serve_env(tmp_path) as env:`` + and the ``serve_env`` fixture. The context-manager object has no builder + methods, so ``isolated_serve_env(tmp_path).build_openai_server(...)`` fails + with ``AttributeError`` instead of leaving global state behind. + + Build one configuration per environment when configurations differ in + environment variables or perf metrics. + """ + tmp_dir = Path(tmp_dir) + env_patcher = mock.patch.dict(os.environ) + env_patcher.start() + try: + os.environ["HF_HUB_OFFLINE"] = "1" + os.environ["TRANSFORMERS_OFFLINE"] = "1" + prometheus_dir = tmp_dir / "prometheus" + prometheus_dir.mkdir(parents=True, exist_ok=True) + os.environ["PROMETHEUS_MULTIPROC_DIR"] = str(prometheus_dir) + os.environ["TRTLLM_MEDIA_STORAGE_PATH"] = str(tmp_dir / "media") + for name in _CLEARED_ENV_VARS: + os.environ.pop(name, None) + + saved_media_executor = BaseMediaIO._executor + # The module global exists only after the first + # set_prometheus_multiproc_dir() call; save it with a sentinel. + saved_prometheus_dir = getattr(tensorrt_llm._utils, "prometheus_multiproc_dir", _MISSING) + # An unhealthy /health records a termination observation on the + # process-wide telemetry session, which keeps only its first + # observation; detach any session so tests cannot poison it. + saved_session = usage_lib._SESSION + usage_lib._SESSION = None + collectors_before = set(REGISTRY._collector_to_names.keys()) + + env = ServeTestEnv(tmp_dir) + try: + yield env + finally: + env._close() + for collector in list(REGISTRY._collector_to_names.keys()): + if collector not in collectors_before: + with contextlib.suppress(KeyError): + REGISTRY.unregister(collector) + BaseMediaIO._executor = saved_media_executor + if saved_prometheus_dir is _MISSING: + if hasattr(tensorrt_llm._utils, "prometheus_multiproc_dir"): + delattr(tensorrt_llm._utils, "prometheus_multiproc_dir") + else: + tensorrt_llm._utils.prometheus_multiproc_dir = saved_prometheus_dir + usage_lib._SESSION = saved_session + finally: + env_patcher.stop() diff --git a/tests/test_common/serve/cpu_tests.py b/tests/test_common/serve/cpu_tests.py new file mode 100644 index 000000000000..44575b7c24f1 --- /dev/null +++ b/tests/test_common/serve/cpu_tests.py @@ -0,0 +1,183 @@ +# 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. +"""The one command that runs every serve CPU test, as the CPU CI stage runs them. + + LLM_MODELS_ROOT=... PYTHONPATH=tests python3 -m test_common.serve.cpu_tests [--list] [pytest args...] + +Reads the CPU test-db list, keeps the entries whose files are serve CPU tests, +and runs one pytest process per kept entry from ``tests/`` with +``tests/unittest`` on ``PYTHONPATH`` — one process per entry, as the CPU stage +does — so process-global state does not leak between entries. Extra arguments go +to every process. Exits with the first non-zero pytest exit code, or 0. +""" + +import argparse +import os +import subprocess +import sys +from dataclasses import dataclass +from pathlib import Path +from typing import List, Optional, Sequence, Tuple + +import yaml + +__all__ = ["Entry", "select_entries"] + +REPO_ROOT = Path(__file__).resolve().parents[3] +CPU_LIST = REPO_ROOT / "tests" / "integration" / "test_lists" / "test-db" / "l0_cpu.yml" + +_MODULE_MARKERS = ( + "tensorrt_llm.serve", + "tensorrt_llm.commands.serve", + "from tensorrt_llm.commands import serve", + "test_common.serve", +) + + +@dataclass(frozen=True) +class Entry: + """One kept test-db entry with the targets it runs.""" + + entry: str # the test-db line, e.g. "unittest/serve" + targets: Tuple[str, ...] # pytest targets relative to tests/ + + +def _iter_test_entries(data) -> List[str]: + """Every string test entry of a test-db document, in file order.""" + entries: List[str] = [] + + def _walk(node) -> None: + if isinstance(node, str): + entries.append(node) + elif isinstance(node, list): + for item in node: + _walk(item) + elif isinstance(node, dict): + for value in node.values(): + _walk(value) + + _walk(data) + return entries + + +def _is_serve_cpu_test(path: Path) -> bool: + """Whether one test file is a serve CPU test. + + The file must contain the literal ``pytest.mark.cpu_only`` (the CPU stage + collects nothing from it otherwise) and either live under + ``tests/unittest/serve/`` or mention a serve module or ``test_common.serve`` + in its text. The path condition catches files that name no module, such as + checks over the test tree itself. + """ + try: + text = path.read_text(errors="replace") + except OSError: + return False + if "pytest.mark.cpu_only" not in text: + return False + serve_dir = REPO_ROOT / "tests" / "unittest" / "serve" + if serve_dir in path.parents: + return True + return any(marker in text for marker in _MODULE_MARKERS) + + +def select_entries(cpu_list: Path = CPU_LIST) -> Tuple[List[Entry], List[str]]: + """(kept entries with targets, notes about skipped entries).""" + data = yaml.safe_load(cpu_list.read_text()) + tests_root = REPO_ROOT / "tests" + kept: List[Entry] = [] + notes: List[str] = [] + for raw_entry in _iter_test_entries(data): + entry = raw_entry.strip() + if not entry.startswith("unittest/"): + continue + target_part = entry.split("::")[0].split(" ")[0] + if entry != target_part and "::" not in entry: + # Entries with -k, -m, TIMEOUT or ISOLATION run with options this + # runner does not reproduce; run them through the CI stage instead. + notes.append(f"skipped (has options): {entry}") + continue + if "::" in entry and " " in entry: + notes.append(f"skipped (has options): {entry}") + continue + target = entry.split("::")[0].rstrip("/") + path = tests_root / target + if "::" in entry: + # A node-id entry runs only that node, never the whole file. + if path.is_file() and _is_serve_cpu_test(path): + kept.append(Entry(entry=raw_entry, targets=(entry,))) + continue + if path.is_file(): + if _is_serve_cpu_test(path): + kept.append(Entry(entry=raw_entry, targets=(target,))) + continue + if path.is_dir(): + selected = tuple( + str(candidate.relative_to(tests_root)) + for candidate in sorted(path.rglob("test_*.py")) + if _is_serve_cpu_test(candidate) + ) + if selected: + kept.append(Entry(entry=raw_entry, targets=selected)) + return kept, notes + + +def main(argv: Optional[Sequence[str]] = None) -> int: + parser = argparse.ArgumentParser( + prog="python3 -m test_common.serve.cpu_tests", + description="Run every serve CPU test, one pytest process per test-db entry.", + ) + parser.add_argument( + "--list", action="store_true", help="print the kept entries and their targets" + ) + args, pytest_args = parser.parse_known_args(argv) + + kept, notes = select_entries() + if args.list: + for note in notes: + print(note) + for entry in kept: + print(entry.entry) + for target in entry.targets: + print(f" {target}") + return 0 + + if not os.environ.get("LLM_MODELS_ROOT"): + print( + "LLM_MODELS_ROOT is not set. Some selected tests read model files through " + "autouse fixtures and would error in setup; set LLM_MODELS_ROOT and rerun.", + file=sys.stderr, + ) + return 2 + + env = dict(os.environ) + unittest_dir = str(REPO_ROOT / "tests" / "unittest") + env["PYTHONPATH"] = os.pathsep.join( + part for part in (unittest_dir, env.get("PYTHONPATH", "")) if part + ) + first_failure = 0 + for note in notes: + print(note) + for entry in kept: + command = [sys.executable, "-m", "pytest", "-m", "cpu_only", *entry.targets, *pytest_args] + print(f"[{entry.entry}] {' '.join(command)}") + result = subprocess.run(command, cwd=REPO_ROOT / "tests", env=env) + if result.returncode != 0 and first_failure == 0: + first_failure = result.returncode + return first_failure + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/test_common/serve/fakes.py b/tests/test_common/serve/fakes.py new file mode 100644 index 000000000000..2318776cd4c8 --- /dev/null +++ b/tests/test_common/serve/fakes.py @@ -0,0 +1,493 @@ +# 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. +"""Fake generators with scripted generation for CPU tests of trtllm-serve. + +The fakes subclass the real generator classes, so every ``isinstance`` check in +``tensorrt_llm/serve`` passes without patching module attributes, but they never +run the real constructors. ``generate_async`` records every call and replays a +queued :class:`ScriptedOutput` through the same ``RequestOutput`` surface the +request handlers consume, so handler tests run through production code on CPU. + +The fakes fail loudly on purpose: a handler that reads a generator attribute the +fakes lack raises ``AttributeError``, which is the signal to extend the fakes +here (with a self-test in ``tests/unittest/serve/test_serve_test_helpers.py``), +not to catch it in the test. + +Not scripted: sampling, logits processors, guided decoding, logprobs, perf +metrics, multimodal data. A test that needs one of those adds it to these fakes +with a self-test, not as a private fake in its own test file. +""" + +import asyncio +import collections +import itertools +from dataclasses import dataclass, field +from types import SimpleNamespace +from typing import Any, AsyncIterator, Optional, Sequence, Union + +from tokenizers import Tokenizer +from tokenizers.models import WordLevel +from tokenizers.pre_tokenizers import Whitespace +from transformers import PreTrainedTokenizerFast + +from tensorrt_llm._torch.async_llm import AsyncLLM +from tensorrt_llm.executor.result import GenerationResultBase +from tensorrt_llm.inputs.registry import DefaultInputProcessor +from tensorrt_llm.llmapi import MultimodalEncoder +from tensorrt_llm.llmapi.llm import LLM, PreprocessedInputs +from tensorrt_llm.llmapi.tokenizer import TransformersTokenizer +from tensorrt_llm.sampling_params import SamplingParams + +__all__ = [ + "FAKE_CHAT_TEMPLATE", + "FakeAsyncLLM", + "FakeEncodeOnlyLLM", + "FakeEncoderExecutor", + "FakeExecutor", + "FakeLLM", + "FakeLlmArgs", + "FakeMultimodalEncoder", + "GenerateCall", + "PreprocessCall", + "ScriptedOutput", + "ScriptedRequestOutput", + "ScriptedStep", + "UserProbeMiddleware", + "make_fake_tokenizer", +] + +# Renders each message as <|role|>content and appends <|assistant|> when +# add_generation_prompt is true, so chat requests render offline through the +# production template machinery. +FAKE_CHAT_TEMPLATE = ( + "{% for message in messages %}" + "<|{{ message['role'] }}|>{{ message['content'] }}" + "{% endfor %}" + "{% if add_generation_prompt %}<|assistant|>{% endif %}" +) + + +def make_fake_tokenizer() -> TransformersTokenizer: + """An in-memory tokenizer with a working chat template, built offline. + + Gives the handlers ``.tokenizer``, ``encode``, ``decode`` and a chat + template without reading any file. The vocabulary only needs to exist; + unknown words map to ````. + """ + vocab = { + "": 0, + "": 1, + "<|system|>": 2, + "<|user|>": 3, + "<|assistant|>": 4, + "hi": 5, + "hello": 6, + "ok": 7, + } + backend = Tokenizer(WordLevel(vocab, unk_token="")) + backend.pre_tokenizer = Whitespace() + hf_tokenizer = PreTrainedTokenizerFast( + tokenizer_object=backend, + unk_token="", + eos_token="", + chat_template=FAKE_CHAT_TEMPLATE, + ) + return TransformersTokenizer(hf_tokenizer) + + +@dataclass +class FakeLlmArgs: + """The LLM args fields ``OpenAIServer`` reads at construction, in its lifespan and in its handlers. + + Defaults are the ``TorchLlmArgs`` defaults. A new direct read of a field + missing here fails loudly with ``AttributeError``; add the field. + """ + + model: str # path of the fake checkpoint directory + backend: str = "pytorch" + dtype: str = "auto" + trust_remote_code: bool = False + checkpoint_format: Optional[str] = None + num_postprocess_workers: int = 0 + post_processor_hook: Optional[str] = None + return_perf_metrics: bool = False + perf_metrics_output_dir: Optional[str] = None + per_request_spec_decode_stats: bool = False + speculative_config: Any = None + decoding_config: Any = None + prometheus_metrics_config: Any = None + quant_config: Any = None + max_seq_len: Optional[int] = None + max_batch_size: Optional[int] = 2048 # get_server_info reads it + parallel_config: Any = field( + default_factory=lambda: SimpleNamespace( + tp_size=1, pp_size=1, cp_size=1, world_size=1, moe_ep_size=None + ) + ) + kv_cache_config: Any = None + cache_transceiver_config: Any = None + enable_energy_metrics: bool = False + enable_iter_perf_stats: bool = False + iter_stats_max_iterations: Optional[int] = None + encode_only: bool = False + sleep_config: Any = None + num_serve_frontends: int = 1 + max_beam_width: int = 1 + reasoning_parser: Optional[str] = None # chat, completions, Responses + gather_generation_logits: bool = False # to_sampling_params in chat and completions + guided_decoding_backend: Optional[str] = None # the required-tool grammar check + + +class FakeExecutor: + """Stands in for the executor where ``OpenAIServer`` reads it.""" + + def __init__(self, *, resource_governor_queue: object = None) -> None: + # register_routes reads it to decide on the resource-governor routes. + self.resource_governor_queue = resource_governor_queue + self.healthy = True + self._fatal_error: Optional[BaseException] = None # read by the /health handler + self.doing_shutdown = False # read by the /health handler + + @property + def fatal_error(self) -> Optional[BaseException]: + """Public name of ``_fatal_error``; nothing reads it yet.""" + return self._fatal_error + + def check_health(self) -> bool: + # Like the real executor: healthy unless a fatal error was recorded. + return self.healthy and self._fatal_error is None + + def shutdown(self) -> None: + self.doing_shutdown = True + + +class FakeEncoderExecutor: + """Stands in for the encoder executor of an encode-only LLM.""" + + def __init__(self) -> None: + self.model_engine = SimpleNamespace(batch_size=8, max_num_tokens=8192, max_seq_len=2048) + + def shutdown(self) -> None: + pass + + +@dataclass(frozen=True) +class ScriptedStep: + """One engine response: what it appends to the outputs.""" + + text: str = "" + token_ids: tuple = () + finish_reason: Optional[str] = None # set on the last step: "stop", "length", ... + stop_reason: Optional[Union[int, str]] = None + index: Optional[int] = None # output index; None applies the step to every output + + +@dataclass(frozen=True) +class ScriptedOutput: + """What one ``generate_async`` call produces.""" + + steps: tuple = (ScriptedStep(text="ok", token_ids=(7,), finish_reason="stop"),) + raise_on_call: Optional[BaseException] = None # raised after the call is recorded + raise_after_steps: Optional[BaseException] = None # raised once the steps are replayed + hold_after_steps: bool = False # wait until release(), abort() or cancellation + prompt_token_ids: Optional[tuple] = None # default: derived from the inputs + + +@dataclass +class GenerateCall: + """One recorded ``generate_async`` call.""" + + inputs: Any # exactly what the handler passed + kwargs: dict # every other argument by name, including sampling_params + result: Optional["ScriptedRequestOutput"] # None when raise_on_call fired + + @property + def sampling_params(self) -> Optional[SamplingParams]: + return self.kwargs.get("sampling_params") + + +@dataclass +class PreprocessCall: + """One recorded ``preprocess`` call.""" + + inputs: Any + sampling_params: Any + disaggregated_params: Any + + +def _derive_prompt_fields(inputs: Any, tokenizer: TransformersTokenizer) -> tuple: + """(prompt token ids, prompt text or None) the way the engine would see them.""" + if isinstance(inputs, PreprocessedInputs): + return list(inputs.prompt_token_ids), None + if isinstance(inputs, dict): + if inputs.get("prompt_token_ids") is not None: # TokensPrompt + return list(inputs["prompt_token_ids"]), None + if inputs.get("prompt") is not None: # TextPrompt + return list(tokenizer.encode(inputs["prompt"])), inputs["prompt"] + if isinstance(inputs, str): + return list(tokenizer.encode(inputs)), inputs + if isinstance(inputs, Sequence) and all(isinstance(t, int) for t in inputs): + return list(inputs), None + raise TypeError(f"cannot derive prompt token ids from inputs of type {type(inputs)!r}") + + +class ScriptedRequestOutput(GenerationResultBase): + """Replays a :class:`ScriptedOutput` with the ``RequestOutput`` surface handlers use.""" + + def __init__( + self, + *, + request_id: int, + script: ScriptedOutput, + inputs: Any, + sampling_params: Optional[SamplingParams], + tokenizer: TransformersTokenizer, + postproc_params: Any = None, + ) -> None: + GenerationResultBase.__init__( + self, + request_id, + sampling_params if sampling_params is not None else SamplingParams(), + postproc_params=postproc_params, + ) + self._script = script + self._fake_tokenizer = tokenizer + if script.prompt_token_ids is not None: + self.prompt_token_ids = list(script.prompt_token_ids) + try: + _, self.prompt = _derive_prompt_fields(inputs, tokenizer) + except TypeError: + self.prompt = None + else: + self.prompt_token_ids, self.prompt = _derive_prompt_fields(inputs, tokenizer) + self._pending = collections.deque(script.steps) + self._release_event = asyncio.Event() + self._after_steps_raised = False + + @property + def request_id(self) -> int: + return self.id + + @property + def finished(self) -> bool: + return self._done + + def release(self) -> None: + """End a ``hold_after_steps`` wait normally.""" + self._release_event.set() + + def abort(self) -> None: + # An aborted engine request ends with finish_reason "cancelled". + for output in self._outputs: + output.finish_reason = "cancelled" + super().abort() + self._done = True + self._release_event.set() + + def _replay_step(self, step: ScriptedStep) -> None: + targets = self._outputs if step.index is None else [self._outputs[step.index]] + for output in targets: + # Update the diff markers the way the result handler does, so + # text_diff / token_ids_diff work for streaming post-processors. + output._last_text_len = len(output.text) + output._last_token_ids_len = len(output.token_ids) + output._last_logprobs_len = len(output.logprobs) + output.text += step.text + output.token_ids = list(output.token_ids) + list(step.token_ids) + if step.finish_reason is not None: + output.finish_reason = step.finish_reason + output.stop_reason = step.stop_reason + if not self._pending and not self._script.hold_after_steps: + self._done = True + if self.postproc_params is not None: + # As the postprocess worker does: fill _postprocess_result after + # each response, with the tokenizer bound on the args. + args = self.postproc_params.postproc_args + args.tokenizer = self._fake_tokenizer + self._outputs[0]._postprocess_result = self.postproc_params.post_processor(self, args) + + def _finish_after_steps(self) -> None: + if self._script.raise_after_steps is not None and not self._after_steps_raised: + self._after_steps_raised = True + self._done = True + raise self._script.raise_after_steps + + async def aresult(self) -> "ScriptedRequestOutput": + while self._pending: + self._replay_step(self._pending.popleft()) + self._finish_after_steps() + if self._script.hold_after_steps and not self._done: + await self._release_event.wait() + self._done = True + return self + + def __await__(self): + return self.aresult().__await__() + + def __aiter__(self) -> "ScriptedRequestOutput": + return self + + async def __anext__(self) -> "ScriptedRequestOutput": + if self._pending: + self._replay_step(self._pending.popleft()) + return self + self._finish_after_steps() + if self._script.hold_after_steps and not self._done: + await self._release_event.wait() + self._done = True + raise StopAsyncIteration + + +class _FakeGeneratorMixin: + """Shared state and scripted generation of every fake generator.""" + + def _setup_fake( + self, + *, + args: FakeLlmArgs, + executor: Optional[FakeExecutor], + encode_only: bool = False, + ) -> None: + self.args = args + self._hf_model_dir = args.model # read by OpenAIServer._init_llm + self._tokenizer = make_fake_tokenizer() + # BaseLLM.tokenizer returns input_processor.tokenizer; completions call + # the input processor for text prompts. The real DefaultInputProcessor + # needs only a tokenizer. + self.input_processor = DefaultInputProcessor( + model_path=args.model, config=None, tokenizer=self._tokenizer + ) + self._executor = executor # BaseLLM._check_health and /health read it + self._encode_only = encode_only # BaseLLM._check_health reads it + self._llm_id = "fake-llm" + self.generate_calls: list = [] + self.preprocess_calls: list = [] + self._scripts: collections.deque = collections.deque() + self._request_ids = itertools.count(1) + + @property + def llm_id(self) -> str: + return self._llm_id + + def script(self, *outputs: ScriptedOutput) -> None: + """Queue outputs for the next ``generate_async`` calls.""" + self._scripts.extend(outputs) + + def preprocess( + self, + inputs: Any, + sampling_params: Any = None, + disaggregated_params: Any = None, + ) -> PreprocessedInputs: + # Replaces the inherited BaseLLM.preprocess, which needs + # args.generation_config and a model. + self.preprocess_calls.append(PreprocessCall(inputs, sampling_params, disaggregated_params)) + token_ids, _ = _derive_prompt_fields(inputs, self._tokenizer) + return PreprocessedInputs(prompt_token_ids=token_ids) + + def generate_async( + self, + inputs: Any, + sampling_params: Optional[SamplingParams] = None, + **kwargs: Any, + ) -> ScriptedRequestOutput: + # Record first, so a scripted failure still leaves the call inspectable. + call = GenerateCall( + inputs=inputs, + kwargs={"sampling_params": sampling_params, **kwargs}, + result=None, + ) + self.generate_calls.append(call) + script = self._scripts.popleft() if self._scripts else ScriptedOutput() + if script.raise_on_call is not None: + raise script.raise_on_call + postproc_params = kwargs.get("_postproc_params") + promise = ScriptedRequestOutput( + request_id=next(self._request_ids), + script=script, + inputs=inputs, + sampling_params=sampling_params, + tokenizer=self._tokenizer, + postproc_params=postproc_params, + ) + if postproc_params is not None: + # As GenerationExecutor.generate_async does before handing the + # request to the postprocess worker. + postproc_params.postproc_args.num_prompt_tokens = len(promise.prompt_token_ids) + call.result = promise + return promise + + async def get_stats_async(self, timeout: Optional[float] = None) -> AsyncIterator[dict]: + return + yield # pragma: no cover - makes this an async generator + + def get_data_transceiver_state(self) -> None: + return None + + def shutdown(self) -> None: + # Overrides BaseLLM.shutdown, which BaseLLM.__del__ calls. + for name in ("_executor", "_encoder_executor"): + executor = getattr(self, name, None) + if executor is not None: + executor.shutdown() + + +class FakeLLM(_FakeGeneratorMixin, LLM): + """LLM generator fake: the default, CONTEXT, GENERATION and runtime-control roles.""" + + def __init__(self, *, args: FakeLlmArgs, executor: Optional[FakeExecutor] = None) -> None: + self._setup_fake(args=args, executor=executor or FakeExecutor()) + + +class FakeAsyncLLM(_FakeGeneratorMixin, AsyncLLM): + """AsyncLLM generator fake: the RL-control role gate requires AsyncLLM.""" + + def __init__(self, *, args: FakeLlmArgs, executor: Optional[FakeExecutor] = None) -> None: + self._setup_fake(args=args, executor=executor or FakeExecutor()) + + +class FakeMultimodalEncoder(_FakeGeneratorMixin, MultimodalEncoder): + """MultimodalEncoder generator fake: the MM_ENCODER role gate.""" + + def __init__(self, *, args: FakeLlmArgs, executor: Optional[FakeExecutor] = None) -> None: + self._setup_fake(args=args, executor=executor or FakeExecutor()) + + +class FakeEncodeOnlyLLM(FakeLLM): + """Encode-only LLM fake for the EMBEDDING role. + + A production encode-only LLM has no generation executor: ``_build_model`` + returns before creating one and ``_check_health`` reads only + ``_encoder_executor`` in encode-only mode. + """ + + def __init__(self, *, args: FakeLlmArgs) -> None: + self._setup_fake(args=args, executor=None, encode_only=True) + self._encoder_executor = FakeEncoderExecutor() + + def encode(self, token_ids_batch: Any) -> Any: + raise NotImplementedError( + "FakeEncodeOnlyLLM does not encode; script it when a test needs it." + ) + + +class UserProbeMiddleware: + """Pure ASGI no-op, used to check where ``--middleware`` lands in the stack.""" + + def __init__(self, app: Any) -> None: + self.app = app + + async def __call__(self, scope: Any, receive: Any, send: Any) -> None: + await self.app(scope, receive, send) diff --git a/tests/test_common/serve/fakes_visual_gen.py b/tests/test_common/serve/fakes_visual_gen.py new file mode 100644 index 000000000000..c1241fb6df77 --- /dev/null +++ b/tests/test_common/serve/fakes_visual_gen.py @@ -0,0 +1,44 @@ +# 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. +"""Fake VisualGen generator for CPU tests of the VISUAL_GEN serving role. + +Kept in its own module so that only tests that build a VisualGen app import the +``tensorrt_llm.visual_gen`` tree: ``OpenAIServer`` detects VisualGen generators +by probing ``sys.modules`` for ``tensorrt_llm.visual_gen.visual_gen``, so +importing this module is what makes ``_is_visual_gen_instance`` able to say yes. +""" + +import itertools +from typing import Any + +from tensorrt_llm.visual_gen.args import VisualGenArgs +from tensorrt_llm.visual_gen.visual_gen import VisualGen + +__all__ = ["FakeVisualGen"] + + +class FakeVisualGen(VisualGen): + """VisualGen generator fake; subclassing keeps ``isinstance`` checks true.""" + + def __init__(self, model: str = "fake-visual-gen") -> None: + # Never run VisualGen.__init__ (it loads a pipeline). VisualGen.shutdown, + # which __del__ calls, returns early when executor is None. + self.model = model + self.args = VisualGenArgs() + self.executor = None + self._req_counter = itertools.count() + + def generate_async(self, *args: Any, **kwargs: Any) -> Any: + raise NotImplementedError("FakeVisualGen does not generate.") diff --git a/tests/test_common/serve/layering.py b/tests/test_common/serve/layering.py new file mode 100644 index 000000000000..aba66f10df46 --- /dev/null +++ b/tests/test_common/serve/layering.py @@ -0,0 +1,222 @@ +# 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. +"""Import-direction contracts between serving layers, read from the source tree. + +The import graph is built with ``ast`` from the repository files, not the +installed package, and needs no new dependency. Each package a PR creates adds +its :class:`Contract` to ``tests/unittest/serve/test_layering.py``. +""" + +import ast +from dataclasses import dataclass, field +from pathlib import Path +from typing import FrozenSet, List, Optional, Sequence, Tuple + +__all__ = [ + "Contract", + "ImportEdge", + "build_import_graph", + "check_contracts", + "stale_allowances", +] + + +@dataclass(frozen=True) +class ImportEdge: + importer: str # dotted module of the importing file + imported: str # absolute dotted module (relative imports resolved) + names: Tuple[str, ...] # names of a from-import; () for "import x" + lineno: int + in_function: bool # inside a def or class body (runs at call time) + type_checking: bool # under "if TYPE_CHECKING:" + + +@dataclass(frozen=True) +class Contract: + name: str + importers: Tuple[str, ...] # module prefixes the rule applies to + forbidden: Tuple[str, ...] # imported-module prefixes that are not allowed + allowed: FrozenSet[Tuple[str, str, str]] = field(default_factory=frozenset) + include_type_checking: bool = False + fix: str = "" # printed on failure: which contract, where the code goes + + +def _matches_prefix(module: str, prefix: str) -> bool: + return module == prefix or module.startswith(prefix + ".") + + +class _ImportCollector(ast.NodeVisitor): + def __init__(self, module: str, package: str) -> None: + self._module = module + self._package = package + self._function_depth = 0 + self._type_checking_depth = 0 + self.edges: List[ImportEdge] = [] + + def _add(self, imported: Optional[str], names: Tuple[str, ...], lineno: int) -> None: + if imported is None: + return + self.edges.append( + ImportEdge( + importer=self._module, + imported=imported, + names=names, + lineno=lineno, + in_function=self._function_depth > 0, + type_checking=self._type_checking_depth > 0, + ) + ) + + def visit_Import(self, node: ast.Import) -> None: + for alias in node.names: + self._add(alias.name, (), node.lineno) + + def visit_ImportFrom(self, node: ast.ImportFrom) -> None: + imported = _resolve_relative(self._package, node.level, node.module) + self._add(imported, tuple(alias.name for alias in node.names), node.lineno) + + def visit_Call(self, node: ast.Call) -> None: + # importlib.import_module("") calls are edges with names=(). + target = node.func + is_import_module = ( + isinstance(target, ast.Attribute) + and target.attr == "import_module" + and isinstance(target.value, ast.Name) + and target.value.id == "importlib" + ) or (isinstance(target, ast.Name) and target.id == "import_module") + if is_import_module and node.args and isinstance(node.args[0], ast.Constant): + value = node.args[0].value + if isinstance(value, str): + self._add(value, (), node.lineno) + self.generic_visit(node) + + def visit_If(self, node: ast.If) -> None: + if _is_type_checking_test(node.test): + self._type_checking_depth += 1 + for child in node.body: + self.visit(child) + self._type_checking_depth -= 1 + for child in node.orelse: + self.visit(child) + return + self.generic_visit(node) + + def _visit_scope(self, node) -> None: + self._function_depth += 1 + self.generic_visit(node) + self._function_depth -= 1 + + def visit_FunctionDef(self, node: ast.FunctionDef) -> None: + self._visit_scope(node) + + def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> None: + self._visit_scope(node) + + def visit_ClassDef(self, node: ast.ClassDef) -> None: + self._visit_scope(node) + + +def _is_type_checking_test(test: ast.expr) -> bool: + if isinstance(test, ast.Name): + return test.id == "TYPE_CHECKING" + if isinstance(test, ast.Attribute): + return test.attr == "TYPE_CHECKING" + return False + + +def _resolve_relative(package: str, level: int, module: Optional[str]) -> Optional[str]: + if level == 0: + return module + parts = package.split(".") if package else [] + if level - 1 > len(parts): + return None + base = parts[: len(parts) - (level - 1)] + if module: + base.extend(module.split(".")) + return ".".join(base) if base else None + + +def build_import_graph(package_dir: Path, package: str) -> List[ImportEdge]: + """Every import edge of every ``*.py`` under ``package_dir``. + + Module-level and function-level imports, relative imports resolved, and + ``importlib.import_module("")`` calls recorded as edges with + ``names=()``. + """ + edges: List[ImportEdge] = [] + for path in sorted(package_dir.rglob("*.py")): + relative = path.relative_to(package_dir) + parts = [package, *relative.parts] + if parts[-1] == "__init__.py": + parts = parts[:-1] + else: + parts[-1] = parts[-1][: -len(".py")] + module = ".".join(parts) + file_package = module if path.name == "__init__.py" else module.rpartition(".")[0] + try: + tree = ast.parse(path.read_text(errors="replace")) + except SyntaxError: + continue + collector = _ImportCollector(module, file_package) + collector.visit(tree) + edges.extend(collector.edges) + return edges + + +def _edge_triples(edge: ImportEdge) -> List[Tuple[str, str, str]]: + names = edge.names or ("",) + return [(edge.importer, edge.imported, name) for name in names] + + +def check_contracts(edges: Sequence[ImportEdge], contracts: Sequence[Contract]) -> List[str]: + """Violations of each contract, as messages that name the contract and the fix. + + A from-import is matched as ``.``, so ``from pkg import sub`` + and ``from .. import sub`` count as imports of the submodule ``pkg.sub``. + """ + problems = [] + for contract in contracts: + for edge in edges: + if edge.type_checking and not contract.include_type_checking: + continue + if not any(_matches_prefix(edge.importer, prefix) for prefix in contract.importers): + continue + for triple in _edge_triples(edge): + importer, imported, name = triple + target = f"{imported}.{name}" if name else imported + if not any(_matches_prefix(target, prefix) for prefix in contract.forbidden): + continue + if triple in contract.allowed: + continue + problems.append( + f"contract '{contract.name}': {importer} (line {edge.lineno}) imports " + f"`{target}`. {contract.fix}" + ) + return problems + + +def stale_allowances(edges: Sequence[ImportEdge], contracts: Sequence[Contract]) -> List[str]: + """Baseline allowances that match no current import; the baseline only shrinks.""" + live = {triple for edge in edges for triple in _edge_triples(edge)} + problems = [] + for contract in contracts: + for triple in sorted(contract.allowed): + if triple not in live: + importer, imported, name = triple + problems.append( + f"contract '{contract.name}': the allowance ({importer}, {imported}, " + f"{name!r}) matches no import anymore; delete it, the baseline only shrinks." + ) + return problems diff --git a/tests/test_common/serve/middleware.py b/tests/test_common/serve/middleware.py new file mode 100644 index 000000000000..2cedec9a1741 --- /dev/null +++ b/tests/test_common/serve/middleware.py @@ -0,0 +1,42 @@ +# 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. +"""Middleware-order snapshot of a trtllm-serve app.""" + +from typing import Dict, List + +__all__ = ["middleware_snapshot"] + + +def middleware_snapshot(app) -> List[Dict]: + """One dict per middleware in ``app.user_middleware``, outermost first. + + Starlette's ``add_middleware`` inserts at position 0, so the last one added + is outermost and ``app.user_middleware`` already lists outermost first. For + a class whose module starts with ``tensorrt_llm.`` the keyword names are + recorded (plus ``expose_headers`` for ``PerfMetricsMiddleware``); for any + other class only the class name is, because third-party keyword names may + change between versions. + """ + snapshot = [] + for middleware in app.user_middleware: + cls = middleware.cls + if cls.__module__.startswith("tensorrt_llm."): + entry = {"class": cls.__name__, "kwargs": sorted(middleware.kwargs)} + if cls.__name__ == "PerfMetricsMiddleware": + entry["expose_headers"] = middleware.kwargs["expose_headers"] + else: + entry = {"class": cls.__name__} + snapshot.append(entry) + return snapshot diff --git a/tests/test_common/serve/routes.py b/tests/test_common/serve/routes.py new file mode 100644 index 000000000000..b4dc02823074 --- /dev/null +++ b/tests/test_common/serve/routes.py @@ -0,0 +1,494 @@ +# 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. +"""Route snapshot of every trtllm-serve app configuration. + +A route record is one string: ``" "``: + +- ````: the route's methods, sorted and comma-joined (``GET,HEAD`` for + FastAPI's documentation routes); ``MOUNT`` for a ``Mount``. +- ````: the effective path from :func:`effective_routes`. +- ````: the route's bare class name, on purpose: moving the route class + to another module must not change the snapshot. +- When two routes in one configuration share `` ``, both records + get `` #`` appended so their order can be compared. Names are not + recorded otherwise, so moving a handler into a router does not change the + snapshot. + +Regenerate the reference with:: + + PYTHONPATH=tests python3 -m test_common.serve.routes --write + +run from the repository root with the checkout's build importable. +""" + +import argparse +import sys +import tempfile +from collections import Counter +from dataclasses import dataclass +from pathlib import Path +from typing import Dict, List, Literal, Optional, Sequence, Tuple, Union + +import yaml +from fastapi import FastAPI + +from tensorrt_llm.llmapi.disagg_utils import ServerRole +from test_common.serve.app_factory import ( + AggConfig, + CoordinatorConfig, + DisaggConfig, + isolated_serve_env, +) + +__all__ = [ + "CONFIGS", + "VISUAL_GEN_CONFIGS", + "SnapshotConfig", + "compare", + "effective_routes", + "read_reference", + "reference_path", + "route_record", + "snapshot_config", + "snapshot_routes", + "write_reference", +] + +REPO_ROOT = Path(__file__).resolve().parents[3] + +REFERENCE_HEADER = """\ +# Routes of every trtllm-serve app configuration. Generated, do not edit by hand: +# PYTHONPATH=tests python3 -m test_common.serve.routes --write +# Each route is " ", in registration order. "#" is +# added only when two routes share methods and path. The checker compares the order +# only for routes that can match the same request. Changing this file changes the +# HTTP surface of trtllm-serve; the API review committee approves it. +""" + + +def reference_path(repo_root: Path = REPO_ROOT) -> Path: + return ( + repo_root + / "tests" + / "unittest" + / "api_stability" + / "references" + / "trtllm_serve_routes.yaml" + ) + + +# --------------------------------------------------------------------------- +# Walking routes +# --------------------------------------------------------------------------- + + +def effective_routes(app_or_routes, prefix: str = "") -> list: + """(effective path, leaf route) in match order, on FastAPI 0.136.x and 0.139.x. + + FastAPI 0.139 keeps one ``_IncludedRouter`` object per ``include_router`` + call in ``app.routes``; it has no ``path`` attribute, so reading ``.path`` + off ``app.routes`` raises once any router is included. This walks into such + objects and applies their include prefix. On 0.136.x included routes are + already flat, so the recursion never runs and the result equals + ``app.routes``. The duck-typed check below uses two attribute names of a + private FastAPI class; if a later FastAPI renames them, + ``test_effective_routes_flattens_included_routers`` fails first, naming this + function. + """ + routes = app_or_routes.routes if isinstance(app_or_routes, FastAPI) else app_or_routes + out = [] + for route in routes: + included = getattr(route, "original_router", None) + context = getattr(route, "include_context", None) + if included is not None and context is not None: # FastAPI 0.139 _IncludedRouter + out.extend(effective_routes(included.routes, prefix + context.prefix)) + continue + # Mounts keep their path and are not descended into. + out.append((prefix + getattr(route, "path", ""), route)) + return out + + +def route_record(path: str, route) -> str: + """``" "`` for one route, without the ``#`` suffix.""" + methods = getattr(route, "methods", None) + methods_text = ",".join(sorted(methods)) if methods else "MOUNT" + return f"{methods_text} {path} {type(route).__name__}" + + +def snapshot_routes(app: FastAPI) -> List[str]: + """Route records of ``app`` in registration order, disambiguated by name.""" + pairs = effective_routes(app) + records = [route_record(path, route) for path, route in pairs] + key_counts = Counter(record.rsplit(" ", 1)[0] for record in records) + out = [] + for record, (path, route) in zip(records, pairs): + methods_and_path = record.rsplit(" ", 1)[0] + if key_counts[methods_and_path] > 1: + record = f"{record} #{getattr(route, 'name', '')}" + out.append(record) + return out + + +# --------------------------------------------------------------------------- +# Configuration matrix +# --------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class SnapshotConfig: + name: str + kind: Literal["openai_server", "disagg_server", "coordinator_server"] + build: Union[AggConfig, DisaggConfig, CoordinatorConfig] + + +CONFIGS: Tuple[SnapshotConfig, ...] = ( + SnapshotConfig("llm", "openai_server", AggConfig()), + SnapshotConfig("llm_context", "openai_server", AggConfig(role=ServerRole.CONTEXT)), + SnapshotConfig("llm_generation", "openai_server", AggConfig(role=ServerRole.GENERATION)), + SnapshotConfig("llm_gpt_oss", "openai_server", AggConfig(model_type="gpt_oss")), + SnapshotConfig( + "llm_gpt_oss_harmony_disabled", + "openai_server", + AggConfig(model_type="gpt_oss", env=(("DISABLE_HARMONY_ADAPTER", "1"),)), + ), + SnapshotConfig("llm_governor", "openai_server", AggConfig(governor_queue=True)), + SnapshotConfig("llm_rl_control", "openai_server", AggConfig(rl_control=True)), + SnapshotConfig("llm_runtime_control", "openai_server", AggConfig(runtime_control=True)), + SnapshotConfig( + "llm_rl_and_runtime_control", + "openai_server", + AggConfig(rl_control=True, runtime_control=True), + ), + SnapshotConfig("llm_perf_metrics", "openai_server", AggConfig(return_perf_metrics=True)), + SnapshotConfig( + "llm_perf_metrics_enable_metrics_env", + "openai_server", + AggConfig(return_perf_metrics=True, env=(("ENABLE_METRICS", "true"),)), + ), + SnapshotConfig( + "llm_perf_metrics_output_dir_only", + "openai_server", + AggConfig(perf_metrics_output_dir=True), + ), + SnapshotConfig( + "llm_render_endpoints", + "openai_server", + AggConfig(env=(("TRTLLM_ENABLE_RENDER_ENDPOINTS", "1"),)), + ), + # An attached frontend of a multi-frontend server registers the launcher + # forwarding routes before the local table; same method and path, so + # registration order decides which handler serves them. + SnapshotConfig( + "llm_multi_frontend_attached", + "openai_server", + AggConfig(multi_frontend_attached=True), + ), + SnapshotConfig("mm_encoder", "openai_server", AggConfig(role=ServerRole.MM_ENCODER)), + SnapshotConfig( + "mm_encoder_rl_control", + "openai_server", + AggConfig(role=ServerRole.MM_ENCODER, rl_control=True), + ), + SnapshotConfig("embedding", "openai_server", AggConfig(role=ServerRole.EMBEDDING)), + SnapshotConfig( + "embedding_runtime_control", + "openai_server", + AggConfig(role=ServerRole.EMBEDDING, runtime_control=True), + ), + SnapshotConfig("disagg_local", "disagg_server", DisaggConfig()), + SnapshotConfig( + "disagg_local_perf_metrics", "disagg_server", DisaggConfig(return_perf_metrics=True) + ), + SnapshotConfig( + "disagg_local_http_cluster_storage", + "disagg_server", + DisaggConfig(http_cluster_storage=True), + ), + SnapshotConfig("disagg_fleet", "disagg_server", DisaggConfig(fleet=True)), + SnapshotConfig( + "disagg_fleet_perf_metrics", + "disagg_server", + DisaggConfig(fleet=True, return_perf_metrics=True), + ), + SnapshotConfig( + "disagg_fleet_http_cluster_storage", + "disagg_server", + DisaggConfig(fleet=True, http_cluster_storage=True), + ), + SnapshotConfig("coordinator", "coordinator_server", CoordinatorConfig()), + SnapshotConfig( + "coordinator_http_cluster_storage", + "coordinator_server", + CoordinatorConfig(http_cluster_storage=True), + ), +) + +VISUAL_GEN_CONFIGS: Tuple[SnapshotConfig, ...] = ( + SnapshotConfig( + "visual_gen", + "openai_server", + AggConfig(role=ServerRole.VISUAL_GEN, model_name="fake-visual-gen"), + ), +) + + +def snapshot_config(config: SnapshotConfig, tmp_dir: Path) -> Dict: + """``{"kind": ..., "routes": [...]}`` or ``{"kind": ..., "error": ": "}``.""" + with isolated_serve_env(tmp_dir) as env: + try: + if config.kind == "openai_server": + app = env.build_openai_server(config.build).app + elif config.kind == "disagg_server": + app = env.build_disagg_server(config.build).app + else: + app = env.build_coordinator_server(config.build).app + except ValueError as error: + return {"kind": config.kind, "error": f"{type(error).__name__}: {error}"} + return {"kind": config.kind, "routes": snapshot_routes(app)} + + +# --------------------------------------------------------------------------- +# Comparison +# --------------------------------------------------------------------------- + + +def _parse_record(record: str) -> Tuple[str, str, str, str]: + """(methods, path, class, name) of one record; name is "" when absent.""" + parts = record.split(" ") + methods, path, cls = parts[0], parts[1], parts[2] + name = parts[3][1:] if len(parts) > 3 and parts[3].startswith("#") else "" + return methods, path, cls, name + + +def _literal_prefix(path: str) -> str: + """The path up to its first ``{``.""" + brace = path.find("{") + return path if brace < 0 else path[:brace] + + +def _paths_can_match_same_request(path_a: str, path_b: str) -> bool: + segments_a = path_a.split("/") + segments_b = path_b.split("/") + for i in range(max(len(segments_a), len(segments_b))): + a = segments_a[i] if i < len(segments_a) else None + b = segments_b[i] if i < len(segments_b) else None + # A {name:path} parameter matches the rest of the path. + if a is not None and "{" in a and a.rstrip("}").endswith(":path"): + return True + if b is not None and "{" in b and b.rstrip("}").endswith(":path"): + return True + if a is None or b is None: + # Different segment counts with no path parameter: no overlap. + return False + if "{" in a or "{" in b: + continue # a parameter matches any literal + if a != b: + return False + return True + + +def _records_conflict(record_a: str, record_b: str) -> bool: + """Whether two records can match the same request (over-approximated). + + Two records conflict when (a) their methods intersect or either is a + mount, and (b) their path templates can match one request path. + Over-approximating conflicts only locks more pairs, which is safe. + """ + methods_a, path_a, _, _ = _parse_record(record_a) + methods_b, path_b, _, _ = _parse_record(record_b) + a_is_mount = methods_a == "MOUNT" + b_is_mount = methods_b == "MOUNT" + if not a_is_mount and not b_is_mount: + if not set(methods_a.split(",")) & set(methods_b.split(",")): + return False + return _paths_can_match_same_request(path_a, path_b) + # A mount at P conflicts with every route whose literal prefix starts + # with P, or that P starts with. + if a_is_mount and b_is_mount: + return path_a.startswith(path_b) or path_b.startswith(path_a) + mount_path = path_a if a_is_mount else path_b + other_prefix = _literal_prefix(path_b if a_is_mount else path_a) + return other_prefix.startswith(mount_path) or mount_path.startswith(other_prefix) + + +def compare(expected: Sequence[str], actual: Sequence[str]) -> List[str]: + """Human-readable differences between the reference and the live records.""" + differences: List[str] = [] + expected_counts = Counter(expected) + actual_counts = Counter(actual) + removed = list((expected_counts - actual_counts).elements()) + added = list((actual_counts - expected_counts).elements()) + + # A removed and an added record with the same methods and path but another + # class are one "route class changed" line. + def _methods_path(record: str) -> str: + methods, path, _, _ = _parse_record(record) + return f"{methods} {path}" + + remaining_added = list(added) + class_changes = [] + for removed_record in list(removed): + for added_record in remaining_added: + if _methods_path(added_record) == _methods_path(removed_record): + old_cls = _parse_record(removed_record)[2] + new_cls = _parse_record(added_record)[2] + why = ( + " (a plain APIRouter breaks msgpack bodies from the disaggregated frontend)" + if old_cls == "_MsgspecRoute" + else "" + ) + class_changes.append( + f"route class changed: {_methods_path(removed_record)}: " + f"{old_cls} -> {new_cls}{why}" + ) + removed.remove(removed_record) + remaining_added.remove(added_record) + break + for record in remaining_added: + differences.append(f"added: {record}") + for record in removed: + differences.append(f"removed: {record}") + differences.extend(class_changes) + + # Conflict-relevant order: for every conflicting pair a-before-b in + # expected, a must come before b in actual, when both are present. For + # repeated identical records, the k-th occurrence pairs with the k-th. + actual_positions: Dict[str, List[int]] = {} + for position, record in enumerate(actual): + actual_positions.setdefault(record, []).append(position) + expected_occurrence: List[Optional[int]] = [] + seen: Counter = Counter() + for record in expected: + positions = actual_positions.get(record, []) + occurrence = seen[record] + seen[record] += 1 + expected_occurrence.append(positions[occurrence] if occurrence < len(positions) else None) + for i, record_a in enumerate(expected): + position_a = expected_occurrence[i] + if position_a is None: + continue + for j in range(i + 1, len(expected)): + record_b = expected[j] + position_b = expected_occurrence[j] + if position_b is None or record_a == record_b: + continue + if _records_conflict(record_a, record_b) and position_a > position_b: + differences.append( + f"order changed: {record_a} must come before {record_b} " + "(FastAPI matches in registration order)" + ) + return differences + + +def format_failure(name: str, differences: Sequence[str]) -> str: + lines = "\n".join(f" {difference}" for difference in differences) + return ( + f"Route snapshot of configuration '{name}' differs from " + "tests/unittest/api_stability/references/trtllm_serve_routes.yaml:\n" + f"{lines}\n" + "If the change is intended, regenerate with " + "`PYTHONPATH=tests python3 -m test_common.serve.routes --write`, commit the YAML, " + "and request review from @NVIDIA/trt-llm-noncommitted-api-review-committee " + "(owner of tests/unittest/api_stability/). New endpoints are a router and a " + "serving class under tensorrt_llm/serve/api///: create the " + "router with make_router() and add a RouterSpec to ROUTERS " + "(tensorrt_llm/serve/_base/routing.py). Do not add routes or methods to " + "OpenAIServer; see tensorrt_llm/serve/AGENTS.md." + ) + + +# --------------------------------------------------------------------------- +# Reference file +# --------------------------------------------------------------------------- + + +def build_reference_data(tmp_root: Path) -> Dict: + configs: Dict[str, Dict] = {} + for index, config in enumerate(CONFIGS + VISUAL_GEN_CONFIGS): + configs[config.name] = snapshot_config(config, tmp_root / f"cfg{index}") + return {"version": 1, "configs": configs} + + +def write_reference(data: Dict, path: Optional[Path] = None) -> Path: + path = path or reference_path() + body = yaml.safe_dump(data, sort_keys=False, width=4096) + path.write_text(REFERENCE_HEADER + body) + return path + + +def read_reference(path: Optional[Path] = None) -> Dict: + path = path or reference_path() + return yaml.safe_load(path.read_text()) + + +def main(argv: Optional[Sequence[str]] = None) -> int: + parser = argparse.ArgumentParser( + prog="python3 -m test_common.serve.routes", + description="Check, rewrite or print the trtllm-serve route reference.", + ) + parser.add_argument( + "--write", action="store_true", help="rebuild every configuration and rewrite the reference" + ) + parser.add_argument( + "--print", + dest="print_name", + metavar="CONFIG", + help="print one configuration's live records", + ) + args = parser.parse_args(argv) + + by_name = {config.name: config for config in CONFIGS + VISUAL_GEN_CONFIGS} + with tempfile.TemporaryDirectory() as tmp: + tmp_root = Path(tmp) + if args.print_name: + if args.print_name not in by_name: + print(f"unknown configuration '{args.print_name}'; one of: {sorted(by_name)}") + return 2 + snapshot = snapshot_config(by_name[args.print_name], tmp_root / "print") + for line in snapshot.get("routes", [snapshot.get("error", "")]): + print(line) + return 0 + data = build_reference_data(tmp_root) + if args.write: + path = write_reference(data) + print(f"wrote {path}") + return 0 + reference = read_reference() + exit_code = 0 + for name, live in data["configs"].items(): + recorded = reference["configs"].get(name) + if recorded is None: + print(f"configuration '{name}' missing from the reference") + exit_code = 1 + continue + if recorded.get("error") != live.get("error"): + print( + f"configuration '{name}': error changed: " + f"{recorded.get('error')!r} -> {live.get('error')!r}" + ) + exit_code = 1 + differences = compare(recorded.get("routes", []), live.get("routes", [])) + if differences: + print(format_failure(name, differences)) + exit_code = 1 + stale = set(reference["configs"]) - set(data["configs"]) + if stale: + print(f"configurations only in the reference: {sorted(stale)}") + exit_code = 1 + return exit_code + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/test_common/serve/shims.py b/tests/test_common/serve/shims.py new file mode 100644 index 000000000000..98ae1a1cfa04 --- /dev/null +++ b/tests/test_common/serve/shims.py @@ -0,0 +1,376 @@ +# 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. +"""Shim identity, FutureWarning and old-path-import checks for moved serve modules. + +A module move keeps the old path only when code outside this repository imports +it. Such a kept shim re-exports the public names (each the same object as at the +new path), warns once with ``FutureWarning``, and code inside the repository +must import the new path. These checks hold all three properties; the PR that +creates a shim adds its :class:`ShimRow` rows to the tables in +``tests/unittest/serve/test_shim_identity.py``. +""" + +import ast +import importlib +import re +import sys +import warnings +from dataclasses import dataclass +from pathlib import Path +from typing import Dict, List, Optional, Sequence, Tuple + +__all__ = [ + "SCAN_ROOTS", + "ShimRow", + "check_identity", + "check_warnings", + "find_old_path_imports", +] + +SCAN_ROOTS = ("tensorrt_llm", "examples", "scripts", "triton_backend", "docs", "setup.py") + + +@dataclass(frozen=True) +class ShimRow: + """One re-exported name of one kept old path.""" + + old_module: str # "tensorrt_llm.serve.openai_disagg_server" + name: str # "OpenAIDisaggServer"; public names only, except in interface rows + new_module: str # "tensorrt_llm.serve.disagg.openai_disagg_server" + retire_in: str # release that removes the shim; "not retired" for an interface row + ban_old_path: bool = True # False: repo code may keep importing name from old_module + warns: bool = True # False: the re-export is the module's own interface, not a shim + + +# --------------------------------------------------------------------------- +# Identity +# --------------------------------------------------------------------------- + + +def check_identity(rows: Sequence[ShimRow]) -> List[str]: + """Each old-path name must be the same object as the new-path name. + + Identity only, never ``__module__``: some re-exported names are dicts or + constants. + """ + problems = [] + for row in rows: + old_module = importlib.import_module(row.old_module) + new_module = importlib.import_module(row.new_module) + try: + old_object = getattr(old_module, row.name) + except AttributeError: + problems.append( + f"`{row.old_module}.{row.name}` no longer exists: the re-export was removed, " + "probably by autoflake or ruff during pre-commit, which drop an import that " + f"`__all__` does not list. Restore `from {row.new_module} import {row.name}` in " + f"{row.old_module.replace('.', '/')}.py and list the name in `__all__` (a pure " + "shim), or restore its row in `_LAZY_NAMES` (`commands/serve.py`)." + ) + continue + new_object = getattr(new_module, row.name) + if old_object is not new_object: + problems.append( + f"`{row.old_module}.{row.name}` is not the same object as " + f"`{row.new_module}.{row.name}`. The old module must re-export the object, " + "not define a copy." + ) + return problems + + +# --------------------------------------------------------------------------- +# FutureWarning +# --------------------------------------------------------------------------- + +_IMPORT_WARNING = re.compile( + r"^(?P[A-Za-z0-9_.]+) has moved to (?P[A-Za-z0-9_.]+) " + r"and will be removed in (?P\S+?)\.$" +) +_READ_WARNING = re.compile( + r"^(?P[A-Za-z0-9_.]+)\.(?P[A-Za-z0-9_]+) has moved to " + r"(?P[A-Za-z0-9_.]+) and will be removed in (?P\S+?)\.$" +) + + +def _matching_warnings(caught, prefix: str) -> List[str]: + """Messages of the FutureWarnings that start with ``prefix``. + + Filtering by message, as ``tests/unittest/_torch/peft/test_shim_peft.py`` + does, also drops the warning of a parent shim package that a fresh import of + one of its submodules emits. + """ + return [ + str(record.message) + for record in caught + if issubclass(record.category, FutureWarning) and str(record.message).startswith(prefix) + ] + + +def _no_warning_message(old_module: str, where: str, name: str = "") -> str: + what = f"reading `{old_module}.{name}`" if where == "read" else f"importing `{old_module}`" + return ( + f"{what} emitted no FutureWarning. A kept shim warns once with " + "warnings.warn(..., FutureWarning, stacklevel=2): the TRT-LLM logger defaults to level " + "error, so a logged warning would not reach callers outside the repository. Restore " + "that call." + ) + + +def check_warnings(rows: Sequence[ShimRow]) -> List[str]: + """Each kept shim warns once with the expected FutureWarning text. + + A pure shim (no ``_LAZY_NAMES``) warns exactly once at import with + `` has moved to and will be removed in .``, + where ```` is the row's new module or a package that contains it, + and never on reads. A ``_LAZY_NAMES`` module warns never at import and + exactly once on the first read of each lazy name; a name it imports eagerly + never warns. Rows with ``warns=False`` (interface rows) are skipped. + + Imports each old module fresh (a module body runs once per process) and + puts any previously imported module object back in ``sys.modules`` and on + its parent package, so later tests keep the objects they imported. + """ + problems: List[str] = [] + by_module: Dict[str, List[ShimRow]] = {} + for row in rows: + if row.warns: + by_module.setdefault(row.old_module, []).append(row) + + for module_name, module_rows in by_module.items(): + saved = sys.modules.get(module_name) + if saved is not None: + del sys.modules[module_name] + try: + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + module = importlib.import_module(module_name) + import_messages = _matching_warnings(caught, f"{module_name} has moved to ") + lazy_names = getattr(module, "_LAZY_NAMES", None) + + if lazy_names is None: + problems.extend( + _check_pure_shim_import(module_name, module_rows, import_messages) + ) + elif import_messages: + problems.append( + f"importing `{module_name}` emitted a FutureWarning; trtllm-serve " + "imports commands/serve.py on every start, so it may warn only when a " + "moved name is read." + ) + + for row in module_rows: + is_lazy_name = lazy_names is not None and row.name in lazy_names + for read_index in (1, 2): + before = len(caught) + getattr(module, row.name) + read_messages = _matching_warnings( + caught[before:], f"{module_name}.{row.name} has moved to " + ) + if is_lazy_name and read_index == 1: + problems.extend(_check_lazy_first_read(row, read_messages)) + elif read_messages: + problems.append( + f"the FutureWarning of `{module_name}` says " + f"`{read_messages[0]}`; expected no warning on this read of " + f"`{row.name}`." + ) + finally: + if saved is not None: + sys.modules[module_name] = saved + parent_name, _, tail = module_name.rpartition(".") + if parent_name and parent_name in sys.modules: + setattr(sys.modules[parent_name], tail, saved) + return problems + + +def _check_pure_shim_import( + module_name: str, module_rows: Sequence[ShimRow], import_messages: List[str] +) -> List[str]: + if not import_messages: + return [_no_warning_message(module_name, "import")] + problems = [] + if len(import_messages) > 1: + problems.append( + f"importing `{module_name}` emitted {len(import_messages)} FutureWarnings; a pure " + "shim warns exactly once, at import." + ) + message = import_messages[0] + match = _IMPORT_WARNING.match(message) + retire_in = module_rows[0].retire_in + expected = f"{module_name} has moved to and will be removed in {retire_in}." + if match is None or match.group("release") != retire_in: + problems.append( + f"the FutureWarning of `{module_name}` says `{message}`; expected `{expected}`: it " + "names the new location and the removal release." + ) + return problems + location = match.group("location") + for row in module_rows: + if not (row.new_module == location or row.new_module.startswith(location + ".")): + problems.append( + f"the FutureWarning of `{module_name}` says `{message}`; expected a location " + f"that contains `{row.new_module}`: it names the new location and the removal " + "release." + ) + return problems + + +def _check_lazy_first_read(row: ShimRow, read_messages: List[str]) -> List[str]: + if not read_messages: + return [_no_warning_message(row.old_module, "read", row.name)] + expected = ( + f"{row.old_module}.{row.name} has moved to {row.new_module} " + f"and will be removed in {row.retire_in}." + ) + if read_messages != [expected]: + return [ + f"the FutureWarning of `{row.old_module}` says `{read_messages[0]}`; expected " + f"`{expected}`: it names the new location and the removal release." + ] + return [] + + +# --------------------------------------------------------------------------- +# Old-path import ban +# --------------------------------------------------------------------------- + + +def _module_name_of(path: Path, repo_root: Path) -> str: + relative = path.relative_to(repo_root) + parts = list(relative.parts) + if parts[-1] == "__init__.py": + parts = parts[:-1] + else: + parts[-1] = parts[-1][: -len(".py")] + return ".".join(parts) + + +def _package_of(path: Path, repo_root: Path) -> str: + module = _module_name_of(path, repo_root) + if path.name == "__init__.py": + return module + return module.rpartition(".")[0] + + +def _resolve_relative(package: str, level: int, module: Optional[str]) -> Optional[str]: + if level == 0: + return module + parts = package.split(".") if package else [] + if level - 1 > len(parts): + return None + base = parts[: len(parts) - (level - 1)] + if module: + base.append(module) + return ".".join(base) if base else None + + +def find_old_path_imports( + rows: Sequence[ShimRow], repo_root: Path, roots: Sequence[str] = SCAN_ROOTS +) -> List[str]: + """Imports of banned names from kept old paths, anywhere under ``roots``. + + Flags ``from import `` (absolute, or relative resolved + against the file's package) and ``.`` attribute access where + ```` is bound to the old module in that file. Importing an old module + that still has code is allowed; only the banned names are flagged. The old + module's own file is exempt, and ``tests/`` is not scanned. + """ + banned: Dict[str, set] = {} + for row in rows: + if row.ban_old_path: + banned.setdefault(row.old_module, set()).add(row.name) + if not banned: + return [] + + retire_by_module = {row.old_module: row.retire_in for row in rows} + new_by_module_name: Dict[Tuple[str, str], str] = { + (row.old_module, row.name): row.new_module for row in rows + } + problems: List[str] = [] + files: List[Path] = [] + for root in roots: + root_path = repo_root / root + if root_path.is_file(): + files.append(root_path) + elif root_path.is_dir(): + files.extend(sorted(root_path.rglob("*.py"))) + + for path in files: + text = path.read_text(errors="replace") + # Cheap prefilter: parse only files that mention an old module. + if not any(module.rsplit(".", 1)[-1] in text for module in banned): + continue + try: + tree = ast.parse(text) + except SyntaxError: + continue + file_module = _module_name_of(path, repo_root) + exempt_modules = {module for module in banned if file_module == module} + package = _package_of(path, repo_root) + + aliases: Dict[str, str] = {} # local name -> old module it is bound to + hits: List[Tuple[int, str, str]] = [] # (line, old module, name) + + for node in ast.walk(tree): + if isinstance(node, ast.Import): + for alias in node.names: + if alias.name in banned and alias.name not in exempt_modules and alias.asname: + aliases[alias.asname] = alias.name + elif isinstance(node, ast.ImportFrom): + resolved = _resolve_relative(package, node.level, node.module) + if resolved is None: + continue + for alias in node.names: + if resolved in banned and resolved not in exempt_modules: + if alias.name in banned[resolved]: + hits.append((node.lineno, resolved, alias.name)) + # `from import ` binds the module. + candidate = f"{resolved}.{alias.name}" + if candidate in banned and candidate not in exempt_modules: + aliases[alias.asname or alias.name] = candidate + + for node in ast.walk(tree): + if not isinstance(node, ast.Attribute): + continue + dotted = _dotted_name(node) + if dotted is None: + continue + prefix, _, attribute = dotted.rpartition(".") + bound = aliases.get(prefix) + if bound is None and prefix in banned and prefix not in exempt_modules: + bound = prefix # the dotted `import ` used as `.` + if bound is not None and attribute in banned.get(bound, ()): + hits.append((node.lineno, bound, attribute)) + + for lineno, module, name in sorted(set(hits)): + problems.append( + f"{path.relative_to(repo_root)}:{lineno} imports `{name}` from the old path " + f"`{module}`; import it from `{new_by_module_name[(module, name)]}`. The old " + "path is kept only for callers outside the repository until " + f"{retire_by_module[module]}." + ) + return problems + + +def _dotted_name(node: ast.Attribute) -> Optional[str]: + parts = [node.attr] + value = node.value + while isinstance(value, ast.Attribute): + parts.append(value.attr) + value = value.value + if isinstance(value, ast.Name): + parts.append(value.id) + return ".".join(reversed(parts)) + return None diff --git a/tests/unittest/api_stability/references/trtllm_serve_routes.yaml b/tests/unittest/api_stability/references/trtllm_serve_routes.yaml new file mode 100644 index 000000000000..523342d4afae --- /dev/null +++ b/tests/unittest/api_stability/references/trtllm_serve_routes.yaml @@ -0,0 +1,663 @@ +# Routes of every trtllm-serve app configuration. Generated, do not edit by hand: +# PYTHONPATH=tests python3 -m test_common.serve.routes --write +# Each route is " ", in registration order. "#" is +# added only when two routes share methods and path. The checker compares the order +# only for routes that can match the same request. Changing this file changes the +# HTTP surface of trtllm-serve; the API review committee approves it. +version: 1 +configs: + llm: + kind: openai_server + routes: + - GET,HEAD /openapi.json Route + - GET,HEAD /docs Route + - GET,HEAD /docs/oauth2-redirect Route + - GET,HEAD /redoc Route + - GET /health _MsgspecRoute + - GET /health_generate _MsgspecRoute + - GET /version _MsgspecRoute + - GET /v1/models _MsgspecRoute + - GET /v1/data_transceiver_state _MsgspecRoute + - GET /metrics _MsgspecRoute + - GET /energy_metrics _MsgspecRoute + - GET /steady_clock_offset _MsgspecRoute + - POST /steady_clock_offset _MsgspecRoute + - POST /kv_cache_events _MsgspecRoute + - POST /v1/completions _MsgspecRoute + - POST /v1/chat/completions _MsgspecRoute + - POST /v1/messages _MsgspecRoute + - POST /v1/messages/count_tokens _MsgspecRoute + - POST /v1/messages/batches _MsgspecRoute + - GET /v1/messages/batches _MsgspecRoute + - GET /v1/messages/batches/{batch_id} _MsgspecRoute + - DELETE /v1/messages/batches/{batch_id} _MsgspecRoute + - POST /v1/messages/batches/{batch_id}/cancel _MsgspecRoute + - GET /v1/messages/batches/{batch_id}/results _MsgspecRoute + - POST /v1/responses _MsgspecRoute + - GET /v1/responses/{response_id} _MsgspecRoute + - DELETE /v1/responses/{response_id} _MsgspecRoute + - POST /_internal/tokenize _MsgspecRoute + - POST /start_profile _MsgspecRoute + - POST /stop_profile _MsgspecRoute + - GET /server_info _MsgspecRoute + llm_context: + kind: openai_server + routes: + - GET,HEAD /openapi.json Route + - GET,HEAD /docs Route + - GET,HEAD /docs/oauth2-redirect Route + - GET,HEAD /redoc Route + - GET /health _MsgspecRoute + - GET /health_generate _MsgspecRoute + - GET /version _MsgspecRoute + - GET /v1/models _MsgspecRoute + - GET /v1/data_transceiver_state _MsgspecRoute + - GET /metrics _MsgspecRoute + - GET /energy_metrics _MsgspecRoute + - GET /steady_clock_offset _MsgspecRoute + - POST /steady_clock_offset _MsgspecRoute + - POST /kv_cache_events _MsgspecRoute + - POST /v1/completions _MsgspecRoute + - POST /v1/chat/completions _MsgspecRoute + - POST /v1/messages _MsgspecRoute + - POST /v1/messages/count_tokens _MsgspecRoute + - POST /v1/messages/batches _MsgspecRoute + - GET /v1/messages/batches _MsgspecRoute + - GET /v1/messages/batches/{batch_id} _MsgspecRoute + - DELETE /v1/messages/batches/{batch_id} _MsgspecRoute + - POST /v1/messages/batches/{batch_id}/cancel _MsgspecRoute + - GET /v1/messages/batches/{batch_id}/results _MsgspecRoute + - POST /v1/responses _MsgspecRoute + - GET /v1/responses/{response_id} _MsgspecRoute + - DELETE /v1/responses/{response_id} _MsgspecRoute + - POST /_internal/tokenize _MsgspecRoute + - POST /start_profile _MsgspecRoute + - POST /stop_profile _MsgspecRoute + - GET /server_info _MsgspecRoute + llm_generation: + kind: openai_server + routes: + - GET,HEAD /openapi.json Route + - GET,HEAD /docs Route + - GET,HEAD /docs/oauth2-redirect Route + - GET,HEAD /redoc Route + - GET /health _MsgspecRoute + - GET /health_generate _MsgspecRoute + - GET /version _MsgspecRoute + - GET /v1/models _MsgspecRoute + - GET /v1/data_transceiver_state _MsgspecRoute + - GET /metrics _MsgspecRoute + - GET /energy_metrics _MsgspecRoute + - GET /steady_clock_offset _MsgspecRoute + - POST /steady_clock_offset _MsgspecRoute + - POST /kv_cache_events _MsgspecRoute + - POST /v1/completions _MsgspecRoute + - POST /v1/chat/completions _MsgspecRoute + - POST /v1/messages _MsgspecRoute + - POST /v1/messages/count_tokens _MsgspecRoute + - POST /v1/messages/batches _MsgspecRoute + - GET /v1/messages/batches _MsgspecRoute + - GET /v1/messages/batches/{batch_id} _MsgspecRoute + - DELETE /v1/messages/batches/{batch_id} _MsgspecRoute + - POST /v1/messages/batches/{batch_id}/cancel _MsgspecRoute + - GET /v1/messages/batches/{batch_id}/results _MsgspecRoute + - POST /v1/responses _MsgspecRoute + - GET /v1/responses/{response_id} _MsgspecRoute + - DELETE /v1/responses/{response_id} _MsgspecRoute + - POST /_internal/tokenize _MsgspecRoute + - POST /start_profile _MsgspecRoute + - POST /stop_profile _MsgspecRoute + - GET /server_info _MsgspecRoute + llm_gpt_oss: + kind: openai_server + routes: + - GET,HEAD /openapi.json Route + - GET,HEAD /docs Route + - GET,HEAD /docs/oauth2-redirect Route + - GET,HEAD /redoc Route + - GET /health _MsgspecRoute + - GET /health_generate _MsgspecRoute + - GET /version _MsgspecRoute + - GET /v1/models _MsgspecRoute + - GET /v1/data_transceiver_state _MsgspecRoute + - GET /metrics _MsgspecRoute + - GET /energy_metrics _MsgspecRoute + - GET /steady_clock_offset _MsgspecRoute + - POST /steady_clock_offset _MsgspecRoute + - POST /kv_cache_events _MsgspecRoute + - POST /v1/completions _MsgspecRoute + - POST /v1/chat/completions _MsgspecRoute + - POST /v1/responses _MsgspecRoute + - GET /v1/responses/{response_id} _MsgspecRoute + - DELETE /v1/responses/{response_id} _MsgspecRoute + - POST /_internal/tokenize _MsgspecRoute + - POST /start_profile _MsgspecRoute + - POST /stop_profile _MsgspecRoute + - GET /server_info _MsgspecRoute + llm_gpt_oss_harmony_disabled: + kind: openai_server + routes: + - GET,HEAD /openapi.json Route + - GET,HEAD /docs Route + - GET,HEAD /docs/oauth2-redirect Route + - GET,HEAD /redoc Route + - GET /health _MsgspecRoute + - GET /health_generate _MsgspecRoute + - GET /version _MsgspecRoute + - GET /v1/models _MsgspecRoute + - GET /v1/data_transceiver_state _MsgspecRoute + - GET /metrics _MsgspecRoute + - GET /energy_metrics _MsgspecRoute + - GET /steady_clock_offset _MsgspecRoute + - POST /steady_clock_offset _MsgspecRoute + - POST /kv_cache_events _MsgspecRoute + - POST /v1/completions _MsgspecRoute + - POST /v1/chat/completions _MsgspecRoute + - POST /v1/messages _MsgspecRoute + - POST /v1/messages/count_tokens _MsgspecRoute + - POST /v1/messages/batches _MsgspecRoute + - GET /v1/messages/batches _MsgspecRoute + - GET /v1/messages/batches/{batch_id} _MsgspecRoute + - DELETE /v1/messages/batches/{batch_id} _MsgspecRoute + - POST /v1/messages/batches/{batch_id}/cancel _MsgspecRoute + - GET /v1/messages/batches/{batch_id}/results _MsgspecRoute + - POST /v1/responses _MsgspecRoute + - GET /v1/responses/{response_id} _MsgspecRoute + - DELETE /v1/responses/{response_id} _MsgspecRoute + - POST /_internal/tokenize _MsgspecRoute + - POST /start_profile _MsgspecRoute + - POST /stop_profile _MsgspecRoute + - GET /server_info _MsgspecRoute + llm_governor: + kind: openai_server + routes: + - GET,HEAD /openapi.json Route + - GET,HEAD /docs Route + - GET,HEAD /docs/oauth2-redirect Route + - GET,HEAD /redoc Route + - GET /health _MsgspecRoute + - GET /health_generate _MsgspecRoute + - GET /version _MsgspecRoute + - GET /v1/models _MsgspecRoute + - GET /v1/data_transceiver_state _MsgspecRoute + - GET /metrics _MsgspecRoute + - GET /energy_metrics _MsgspecRoute + - GET /steady_clock_offset _MsgspecRoute + - POST /steady_clock_offset _MsgspecRoute + - POST /kv_cache_events _MsgspecRoute + - POST /_resource_governor/truncate _MsgspecRoute + - POST /_control/kv_cache/truncate _MsgspecRoute + - POST /_control/kv_cache/truncate_tokens _MsgspecRoute + - POST /v1/completions _MsgspecRoute + - POST /v1/chat/completions _MsgspecRoute + - POST /v1/messages _MsgspecRoute + - POST /v1/messages/count_tokens _MsgspecRoute + - POST /v1/messages/batches _MsgspecRoute + - GET /v1/messages/batches _MsgspecRoute + - GET /v1/messages/batches/{batch_id} _MsgspecRoute + - DELETE /v1/messages/batches/{batch_id} _MsgspecRoute + - POST /v1/messages/batches/{batch_id}/cancel _MsgspecRoute + - GET /v1/messages/batches/{batch_id}/results _MsgspecRoute + - POST /v1/responses _MsgspecRoute + - GET /v1/responses/{response_id} _MsgspecRoute + - DELETE /v1/responses/{response_id} _MsgspecRoute + - POST /_internal/tokenize _MsgspecRoute + - POST /start_profile _MsgspecRoute + - POST /stop_profile _MsgspecRoute + - GET /server_info _MsgspecRoute + llm_rl_control: + kind: openai_server + routes: + - GET,HEAD /openapi.json Route + - GET,HEAD /docs Route + - GET,HEAD /docs/oauth2-redirect Route + - GET,HEAD /redoc Route + - GET /health _MsgspecRoute + - GET /health_generate _MsgspecRoute + - GET /version _MsgspecRoute + - GET /v1/models _MsgspecRoute + - GET /v1/data_transceiver_state _MsgspecRoute + - GET /metrics _MsgspecRoute + - GET /energy_metrics _MsgspecRoute + - GET /steady_clock_offset _MsgspecRoute + - POST /steady_clock_offset _MsgspecRoute + - POST /kv_cache_events _MsgspecRoute + - POST /v1/completions _MsgspecRoute + - POST /v1/chat/completions _MsgspecRoute + - POST /v1/messages _MsgspecRoute + - POST /v1/messages/count_tokens _MsgspecRoute + - POST /v1/messages/batches _MsgspecRoute + - GET /v1/messages/batches _MsgspecRoute + - GET /v1/messages/batches/{batch_id} _MsgspecRoute + - DELETE /v1/messages/batches/{batch_id} _MsgspecRoute + - POST /v1/messages/batches/{batch_id}/cancel _MsgspecRoute + - GET /v1/messages/batches/{batch_id}/results _MsgspecRoute + - POST /v1/responses _MsgspecRoute + - GET /v1/responses/{response_id} _MsgspecRoute + - DELETE /v1/responses/{response_id} _MsgspecRoute + - POST /_internal/tokenize _MsgspecRoute + - POST /start_profile _MsgspecRoute + - POST /stop_profile _MsgspecRoute + - POST /release_memory _MsgspecRoute + - POST /resume_memory _MsgspecRoute + - POST /update_weights _MsgspecRoute + - GET /server_info _MsgspecRoute + llm_runtime_control: + kind: openai_server + routes: + - GET,HEAD /openapi.json Route + - GET,HEAD /docs Route + - GET,HEAD /docs/oauth2-redirect Route + - GET,HEAD /redoc Route + - GET /health _MsgspecRoute + - GET /health_generate _MsgspecRoute + - GET /version _MsgspecRoute + - GET /v1/models _MsgspecRoute + - GET /v1/data_transceiver_state _MsgspecRoute + - GET /metrics _MsgspecRoute + - GET /energy_metrics _MsgspecRoute + - GET /steady_clock_offset _MsgspecRoute + - POST /steady_clock_offset _MsgspecRoute + - POST /kv_cache_events _MsgspecRoute + - POST /v1/completions _MsgspecRoute + - POST /v1/chat/completions _MsgspecRoute + - POST /v1/messages _MsgspecRoute + - POST /v1/messages/count_tokens _MsgspecRoute + - POST /v1/messages/batches _MsgspecRoute + - GET /v1/messages/batches _MsgspecRoute + - GET /v1/messages/batches/{batch_id} _MsgspecRoute + - DELETE /v1/messages/batches/{batch_id} _MsgspecRoute + - POST /v1/messages/batches/{batch_id}/cancel _MsgspecRoute + - GET /v1/messages/batches/{batch_id}/results _MsgspecRoute + - POST /v1/responses _MsgspecRoute + - GET /v1/responses/{response_id} _MsgspecRoute + - DELETE /v1/responses/{response_id} _MsgspecRoute + - POST /_internal/tokenize _MsgspecRoute + - POST /start_profile _MsgspecRoute + - POST /stop_profile _MsgspecRoute + - POST /release_memory _MsgspecRoute + - POST /resume_memory _MsgspecRoute + - GET /memory_status _MsgspecRoute + - GET /server_info _MsgspecRoute + llm_rl_and_runtime_control: + kind: openai_server + error: 'ValueError: Legacy RL control and runtime control endpoints cannot both be enabled.' + llm_perf_metrics: + kind: openai_server + routes: + - GET,HEAD /openapi.json Route + - GET,HEAD /docs Route + - GET,HEAD /docs/oauth2-redirect Route + - GET,HEAD /redoc Route + - GET /health _MsgspecRoute + - GET /health_generate _MsgspecRoute + - GET /version _MsgspecRoute + - GET /v1/models _MsgspecRoute + - GET /v1/data_transceiver_state _MsgspecRoute + - GET /metrics _MsgspecRoute + - GET /energy_metrics _MsgspecRoute + - GET /steady_clock_offset _MsgspecRoute + - POST /steady_clock_offset _MsgspecRoute + - POST /kv_cache_events _MsgspecRoute + - POST /v1/completions _MsgspecRoute + - POST /v1/chat/completions _MsgspecRoute + - POST /v1/messages _MsgspecRoute + - POST /v1/messages/count_tokens _MsgspecRoute + - POST /v1/messages/batches _MsgspecRoute + - GET /v1/messages/batches _MsgspecRoute + - GET /v1/messages/batches/{batch_id} _MsgspecRoute + - DELETE /v1/messages/batches/{batch_id} _MsgspecRoute + - POST /v1/messages/batches/{batch_id}/cancel _MsgspecRoute + - GET /v1/messages/batches/{batch_id}/results _MsgspecRoute + - POST /v1/responses _MsgspecRoute + - GET /v1/responses/{response_id} _MsgspecRoute + - DELETE /v1/responses/{response_id} _MsgspecRoute + - POST /_internal/tokenize _MsgspecRoute + - POST /start_profile _MsgspecRoute + - POST /stop_profile _MsgspecRoute + - GET /server_info _MsgspecRoute + - MOUNT /prometheus/metrics Mount + llm_perf_metrics_enable_metrics_env: + kind: openai_server + routes: + - GET,HEAD /openapi.json Route + - GET,HEAD /docs Route + - GET,HEAD /docs/oauth2-redirect Route + - GET,HEAD /redoc Route + - GET /health _MsgspecRoute + - GET /health_generate _MsgspecRoute + - GET /version _MsgspecRoute + - GET /v1/models _MsgspecRoute + - GET /v1/data_transceiver_state _MsgspecRoute + - 'GET /metrics _MsgspecRoute #get_iteration_stats' + - GET /energy_metrics _MsgspecRoute + - GET /steady_clock_offset _MsgspecRoute + - POST /steady_clock_offset _MsgspecRoute + - POST /kv_cache_events _MsgspecRoute + - POST /v1/completions _MsgspecRoute + - POST /v1/chat/completions _MsgspecRoute + - POST /v1/messages _MsgspecRoute + - POST /v1/messages/count_tokens _MsgspecRoute + - POST /v1/messages/batches _MsgspecRoute + - GET /v1/messages/batches _MsgspecRoute + - GET /v1/messages/batches/{batch_id} _MsgspecRoute + - DELETE /v1/messages/batches/{batch_id} _MsgspecRoute + - POST /v1/messages/batches/{batch_id}/cancel _MsgspecRoute + - GET /v1/messages/batches/{batch_id}/results _MsgspecRoute + - POST /v1/responses _MsgspecRoute + - GET /v1/responses/{response_id} _MsgspecRoute + - DELETE /v1/responses/{response_id} _MsgspecRoute + - POST /_internal/tokenize _MsgspecRoute + - POST /start_profile _MsgspecRoute + - POST /stop_profile _MsgspecRoute + - GET /server_info _MsgspecRoute + - 'GET /metrics _MsgspecRoute #metrics' + - MOUNT /prometheus/metrics Mount + llm_perf_metrics_output_dir_only: + kind: openai_server + routes: + - GET,HEAD /openapi.json Route + - GET,HEAD /docs Route + - GET,HEAD /docs/oauth2-redirect Route + - GET,HEAD /redoc Route + - GET /health _MsgspecRoute + - GET /health_generate _MsgspecRoute + - GET /version _MsgspecRoute + - GET /v1/models _MsgspecRoute + - GET /v1/data_transceiver_state _MsgspecRoute + - GET /metrics _MsgspecRoute + - GET /energy_metrics _MsgspecRoute + - GET /steady_clock_offset _MsgspecRoute + - POST /steady_clock_offset _MsgspecRoute + - POST /kv_cache_events _MsgspecRoute + - POST /v1/completions _MsgspecRoute + - POST /v1/chat/completions _MsgspecRoute + - POST /v1/messages _MsgspecRoute + - POST /v1/messages/count_tokens _MsgspecRoute + - POST /v1/messages/batches _MsgspecRoute + - GET /v1/messages/batches _MsgspecRoute + - GET /v1/messages/batches/{batch_id} _MsgspecRoute + - DELETE /v1/messages/batches/{batch_id} _MsgspecRoute + - POST /v1/messages/batches/{batch_id}/cancel _MsgspecRoute + - GET /v1/messages/batches/{batch_id}/results _MsgspecRoute + - POST /v1/responses _MsgspecRoute + - GET /v1/responses/{response_id} _MsgspecRoute + - DELETE /v1/responses/{response_id} _MsgspecRoute + - POST /_internal/tokenize _MsgspecRoute + - POST /start_profile _MsgspecRoute + - POST /stop_profile _MsgspecRoute + - GET /server_info _MsgspecRoute + - MOUNT /prometheus/metrics Mount + llm_render_endpoints: + kind: openai_server + routes: + - GET,HEAD /openapi.json Route + - GET,HEAD /docs Route + - GET,HEAD /docs/oauth2-redirect Route + - GET,HEAD /redoc Route + - GET /health _MsgspecRoute + - GET /health_generate _MsgspecRoute + - GET /version _MsgspecRoute + - GET /v1/models _MsgspecRoute + - GET /v1/data_transceiver_state _MsgspecRoute + - GET /metrics _MsgspecRoute + - GET /energy_metrics _MsgspecRoute + - GET /steady_clock_offset _MsgspecRoute + - POST /steady_clock_offset _MsgspecRoute + - POST /kv_cache_events _MsgspecRoute + - POST /v1/completions _MsgspecRoute + - POST /v1/chat/completions _MsgspecRoute + - POST /v1/messages _MsgspecRoute + - POST /v1/messages/count_tokens _MsgspecRoute + - POST /v1/messages/batches _MsgspecRoute + - GET /v1/messages/batches _MsgspecRoute + - GET /v1/messages/batches/{batch_id} _MsgspecRoute + - DELETE /v1/messages/batches/{batch_id} _MsgspecRoute + - POST /v1/messages/batches/{batch_id}/cancel _MsgspecRoute + - GET /v1/messages/batches/{batch_id}/results _MsgspecRoute + - POST /v1/responses _MsgspecRoute + - GET /v1/responses/{response_id} _MsgspecRoute + - DELETE /v1/responses/{response_id} _MsgspecRoute + - POST /_internal/tokenize _MsgspecRoute + - POST /start_profile _MsgspecRoute + - POST /stop_profile _MsgspecRoute + - GET /server_info _MsgspecRoute + - POST /v1/chat/completions/render _MsgspecRoute + - POST /v1/completions/render _MsgspecRoute + llm_multi_frontend_attached: + kind: openai_server + routes: + - GET,HEAD /openapi.json Route + - GET,HEAD /docs Route + - GET,HEAD /docs/oauth2-redirect Route + - GET,HEAD /redoc Route + - 'GET /metrics _MsgspecRoute #forward' + - 'POST /kv_cache_events _MsgspecRoute #forward' + - 'POST /start_profile _MsgspecRoute #forward' + - 'POST /stop_profile _MsgspecRoute #forward' + - 'POST /steady_clock_offset _MsgspecRoute #forward' + - 'GET /v1/messages/batches _MsgspecRoute #forward' + - 'POST /v1/messages/batches _MsgspecRoute #forward' + - GET /v1/messages/batches/{rest:path} _MsgspecRoute + - POST /v1/messages/batches/{rest:path} _MsgspecRoute + - DELETE /v1/messages/batches/{rest:path} _MsgspecRoute + - GET /health _MsgspecRoute + - GET /health_generate _MsgspecRoute + - GET /version _MsgspecRoute + - GET /v1/models _MsgspecRoute + - GET /v1/data_transceiver_state _MsgspecRoute + - 'GET /metrics _MsgspecRoute #get_iteration_stats' + - GET /energy_metrics _MsgspecRoute + - GET /steady_clock_offset _MsgspecRoute + - 'POST /steady_clock_offset _MsgspecRoute #set_steady_clock_offset' + - 'POST /kv_cache_events _MsgspecRoute #get_kv_cache_events' + - POST /v1/completions _MsgspecRoute + - POST /v1/chat/completions _MsgspecRoute + - POST /v1/messages _MsgspecRoute + - POST /v1/messages/count_tokens _MsgspecRoute + - 'POST /v1/messages/batches _MsgspecRoute #anthropic_create_batch' + - 'GET /v1/messages/batches _MsgspecRoute #anthropic_list_batches' + - GET /v1/messages/batches/{batch_id} _MsgspecRoute + - DELETE /v1/messages/batches/{batch_id} _MsgspecRoute + - POST /v1/messages/batches/{batch_id}/cancel _MsgspecRoute + - GET /v1/messages/batches/{batch_id}/results _MsgspecRoute + - POST /v1/responses _MsgspecRoute + - GET /v1/responses/{response_id} _MsgspecRoute + - DELETE /v1/responses/{response_id} _MsgspecRoute + - POST /_internal/tokenize _MsgspecRoute + - 'POST /start_profile _MsgspecRoute #start_profile' + - 'POST /stop_profile _MsgspecRoute #stop_profile' + - GET /server_info _MsgspecRoute + mm_encoder: + kind: openai_server + routes: + - GET,HEAD /openapi.json Route + - GET,HEAD /docs Route + - GET,HEAD /docs/oauth2-redirect Route + - GET,HEAD /redoc Route + - GET /health _MsgspecRoute + - GET /version _MsgspecRoute + - GET /v1/models _MsgspecRoute + - GET /metrics _MsgspecRoute + - POST /v1/chat/completions _MsgspecRoute + mm_encoder_rl_control: + kind: openai_server + error: 'ValueError: RL control endpoints require AsyncLLM' + embedding: + kind: openai_server + routes: + - GET,HEAD /openapi.json Route + - GET,HEAD /docs Route + - GET,HEAD /docs/oauth2-redirect Route + - GET,HEAD /redoc Route + - GET /health _MsgspecRoute + - GET /version _MsgspecRoute + - GET /v1/models _MsgspecRoute + - POST /v1/embeddings _MsgspecRoute + embedding_runtime_control: + kind: openai_server + error: 'ValueError: Runtime control endpoints are not supported for embedding servers' + disagg_local: + kind: disagg_server + routes: + - GET,HEAD /openapi.json Route + - GET,HEAD /docs Route + - GET,HEAD /docs/oauth2-redirect Route + - GET,HEAD /redoc Route + - POST /v1/completions KvHintsRoute + - POST /v1/chat/completions KvHintsRoute + - POST /v1/messages KvHintsRoute + - POST /v1/messages/count_tokens KvHintsRoute + - POST /v1/responses KvHintsRoute + - GET /v1/models KvHintsRoute + - GET /health KvHintsRoute + - GET /cluster_info KvHintsRoute + - GET /version KvHintsRoute + - MOUNT /prometheus/metrics Mount + disagg_local_perf_metrics: + kind: disagg_server + routes: + - GET,HEAD /openapi.json Route + - GET,HEAD /docs Route + - GET,HEAD /docs/oauth2-redirect Route + - GET,HEAD /redoc Route + - POST /v1/completions KvHintsRoute + - POST /v1/chat/completions KvHintsRoute + - POST /v1/messages KvHintsRoute + - POST /v1/messages/count_tokens KvHintsRoute + - POST /v1/responses KvHintsRoute + - GET /v1/models KvHintsRoute + - GET /health KvHintsRoute + - GET /cluster_info KvHintsRoute + - GET /version KvHintsRoute + - MOUNT /prometheus/metrics Mount + disagg_local_http_cluster_storage: + kind: disagg_server + routes: + - GET,HEAD /openapi.json Route + - GET,HEAD /docs Route + - GET,HEAD /docs/oauth2-redirect Route + - GET,HEAD /redoc Route + - POST /v1/completions KvHintsRoute + - POST /v1/chat/completions KvHintsRoute + - POST /v1/messages KvHintsRoute + - POST /v1/messages/count_tokens KvHintsRoute + - POST /v1/responses KvHintsRoute + - GET /v1/models KvHintsRoute + - GET /health KvHintsRoute + - GET /cluster_info KvHintsRoute + - GET /version KvHintsRoute + - MOUNT /prometheus/metrics Mount + - POST /set KvHintsRoute + - GET /get KvHintsRoute + - DELETE /delete KvHintsRoute + - GET /expire KvHintsRoute + - GET /get_prefix KvHintsRoute + disagg_fleet: + kind: disagg_server + routes: + - GET,HEAD /openapi.json Route + - GET,HEAD /docs Route + - GET,HEAD /docs/oauth2-redirect Route + - GET,HEAD /redoc Route + - POST /v1/completions KvHintsRoute + - POST /v1/chat/completions KvHintsRoute + - POST /v1/messages KvHintsRoute + - POST /v1/messages/count_tokens KvHintsRoute + - POST /v1/responses KvHintsRoute + - GET /v1/models KvHintsRoute + - GET /health KvHintsRoute + - GET /cluster_info KvHintsRoute + - GET /version KvHintsRoute + - MOUNT /prometheus/metrics Mount + disagg_fleet_perf_metrics: + kind: disagg_server + routes: + - GET,HEAD /openapi.json Route + - GET,HEAD /docs Route + - GET,HEAD /docs/oauth2-redirect Route + - GET,HEAD /redoc Route + - POST /v1/completions KvHintsRoute + - POST /v1/chat/completions KvHintsRoute + - POST /v1/messages KvHintsRoute + - POST /v1/messages/count_tokens KvHintsRoute + - POST /v1/responses KvHintsRoute + - GET /v1/models KvHintsRoute + - GET /health KvHintsRoute + - GET /cluster_info KvHintsRoute + - GET /version KvHintsRoute + - MOUNT /prometheus/metrics Mount + disagg_fleet_http_cluster_storage: + kind: disagg_server + routes: + - GET,HEAD /openapi.json Route + - GET,HEAD /docs Route + - GET,HEAD /docs/oauth2-redirect Route + - GET,HEAD /redoc Route + - POST /v1/completions KvHintsRoute + - POST /v1/chat/completions KvHintsRoute + - POST /v1/messages KvHintsRoute + - POST /v1/messages/count_tokens KvHintsRoute + - POST /v1/responses KvHintsRoute + - GET /v1/models KvHintsRoute + - GET /health KvHintsRoute + - GET /cluster_info KvHintsRoute + - GET /version KvHintsRoute + - MOUNT /prometheus/metrics Mount + - POST /set KvHintsRoute + - GET /get KvHintsRoute + - DELETE /delete KvHintsRoute + - GET /expire KvHintsRoute + - GET /get_prefix KvHintsRoute + coordinator: + kind: coordinator_server + routes: + - GET,HEAD /openapi.json Route + - GET,HEAD /docs Route + - GET,HEAD /docs/oauth2-redirect Route + - GET,HEAD /redoc Route + - POST /select APIRoute + - POST /finish APIRoute + - POST /renew APIRoute + - GET /cluster_info APIRoute + - GET /health APIRoute + - GET /version APIRoute + coordinator_http_cluster_storage: + kind: coordinator_server + routes: + - GET,HEAD /openapi.json Route + - GET,HEAD /docs Route + - GET,HEAD /docs/oauth2-redirect Route + - GET,HEAD /redoc Route + - POST /select APIRoute + - POST /finish APIRoute + - POST /renew APIRoute + - GET /cluster_info APIRoute + - GET /health APIRoute + - GET /version APIRoute + - POST /set APIRoute + - GET /get APIRoute + - DELETE /delete APIRoute + - GET /expire APIRoute + - GET /get_prefix APIRoute + visual_gen: + kind: openai_server + routes: + - GET,HEAD /openapi.json Route + - GET,HEAD /docs Route + - GET,HEAD /docs/oauth2-redirect Route + - GET,HEAD /redoc Route + - GET /health _MsgspecRoute + - GET /version _MsgspecRoute + - GET /v1/models _MsgspecRoute + - GET /metrics _MsgspecRoute + - POST /v1/images/generations _MsgspecRoute + - POST /v1/images/edits _MsgspecRoute + - GET /v1/images/{image_id}/content _MsgspecRoute + - POST /v1/videos _MsgspecRoute + - POST /v1/videos/sync _MsgspecRoute + - POST /v1/videos/generations _MsgspecRoute + - GET /v1/videos _MsgspecRoute + - GET /v1/videos/{video_id} _MsgspecRoute + - GET /v1/videos/{video_id}/content _MsgspecRoute + - DELETE /v1/videos/{video_id} _MsgspecRoute diff --git a/tests/unittest/api_stability/test_serve_routes.py b/tests/unittest/api_stability/test_serve_routes.py new file mode 100644 index 000000000000..1392c1428ec3 --- /dev/null +++ b/tests/unittest/api_stability/test_serve_routes.py @@ -0,0 +1,251 @@ +# 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. +"""Route snapshot of every trtllm-serve configuration: a new api-stability gate. + +Each configuration's live routes must equal the generated reference +``references/trtllm_serve_routes.yaml``. Regenerating the reference cannot hide +the invariants at the bottom of this file: they are checked on the live apps. +The VisualGen configuration is checked by +``tests/unittest/serve/api/visual_gen/test_routes_visual_gen.py`` so that it +runs on VisualGen-only pull requests. +""" + +import pytest +from test_common.serve.routes import ( + CONFIGS, + VISUAL_GEN_CONFIGS, + compare, + format_failure, + read_reference, + snapshot_config, +) + +pytestmark = pytest.mark.cpu_only + + +@pytest.fixture(scope="module") +def reference(): + return read_reference() + + +@pytest.fixture(scope="module") +def live_snapshots(tmp_path_factory): + root = tmp_path_factory.mktemp("serve_routes") + return {config.name: snapshot_config(config, root / config.name) for config in CONFIGS} + + +def test_reference_configuration_names(reference): + expected = [config.name for config in CONFIGS + VISUAL_GEN_CONFIGS] + assert list(reference["configs"].keys()) == expected, ( + "The configuration names in references/trtllm_serve_routes.yaml must equal the " + "names in CONFIGS plus VISUAL_GEN_CONFIGS of tests/test_common/serve/routes.py, in " + "that order. Regenerate with `PYTHONPATH=tests python3 -m test_common.serve.routes " + "--write`." + ) + + +@pytest.mark.parametrize("name", [config.name for config in CONFIGS]) +def test_routes_match_reference(name, reference, live_snapshots): + recorded = reference["configs"][name] + live = live_snapshots[name] + assert recorded.get("error") == live.get("error"), ( + f"Configuration '{name}' now fails with {live.get('error')!r} instead of " + f"{recorded.get('error')!r}. If the change is intended, regenerate with " + "`PYTHONPATH=tests python3 -m test_common.serve.routes --write` and request review " + "from @NVIDIA/trt-llm-noncommitted-api-review-committee." + ) + differences = compare(recorded.get("routes", []), live.get("routes", [])) + assert not differences, format_failure(name, differences) + + +# --------------------------------------------------------------------------- +# Invariants checked on the live apps, independent of the reference. +# Regenerating the reference cannot hide these. +# --------------------------------------------------------------------------- + +_DOC_ROUTES = {"/openapi.json", "/docs", "/docs/oauth2-redirect", "/redoc"} + + +def _routes_of(live_snapshots, name): + return live_snapshots[name]["routes"] + + +def _parse(record): + parts = record.split(" ") + return parts[0], parts[1], parts[2] + + +def test_openai_server_routes_are_msgspec(live_snapshots): + for name, snapshot in live_snapshots.items(): + if snapshot.get("error") or snapshot["kind"] != "openai_server": + continue + for record in snapshot["routes"]: + methods, path, cls = _parse(record) + if path in _DOC_ROUTES: + assert cls == "Route", record + elif methods == "MOUNT": + assert cls == "Mount", record + else: + assert cls == "_MsgspecRoute", ( + f"configuration '{name}': {record}: every API route of the aggregated " + "server must keep the msgpack route class, or workers reject the " + "disaggregated frontend's msgpack bodies with 400/422 while plain JSON " + "still works." + ) + + +_API_ROUTE_CLASSES = { + "disagg_server": ( + "KvHintsRoute", + "every API route of the disaggregated frontend must keep KvHintsRoute: it strips " + "kv_hints from request bodies unless the server trusts its ingress " + "(TRTLLM_TRUST_KV_HINTS=1), and a plain APIRoute skips that policy.", + ), + "coordinator_server": ( + "APIRoute", + "the coordinator app registers FastAPI's plain routes; it takes neither the " + "disaggregated frontend's KvHintsRoute nor the aggregated server's msgpack route " + "class.", + ), +} + + +def test_disagg_and_coordinator_route_classes(live_snapshots): + for name, snapshot in live_snapshots.items(): + if snapshot.get("error") or snapshot["kind"] == "openai_server": + continue + expected, reason = _API_ROUTE_CLASSES[snapshot["kind"]] + for record in snapshot["routes"]: + methods, path, cls = _parse(record) + if path in _DOC_ROUTES: + assert cls == "Route", record + elif methods == "MOUNT": + assert cls == "Mount", record + else: + assert cls == expected, f"configuration '{name}': {record}: {reason}" + + +def test_control_route_groups(live_snapshots): + def paths(name): + return {_parse(record)[1] for record in _routes_of(live_snapshots, name)} + + rl = paths("llm_rl_control") + assert {"/release_memory", "/resume_memory", "/update_weights"} <= rl + assert "/memory_status" not in rl, ( + "the RL control group has no /memory_status; the runtime control group owns it" + ) + runtime = paths("llm_runtime_control") + assert {"/release_memory", "/resume_memory", "/memory_status"} <= runtime + assert "/update_weights" not in runtime, ( + "the runtime control group has no /update_weights; the RL control group owns it" + ) + control_paths = {"/release_memory", "/resume_memory", "/update_weights", "/memory_status"} + for name in ("llm", "llm_context", "llm_generation", "mm_encoder", "embedding"): + assert not control_paths & paths(name), ( + f"configuration '{name}' must register no memory-control route: those routes " + "exist only behind their construction-time switches." + ) + + +def test_metrics_presence_per_role(live_snapshots): + embedding_paths = {_parse(r)[1] for r in _routes_of(live_snapshots, "embedding")} + assert "/metrics" not in embedding_paths, ( + "the embedding server registers no /metrics today; keep the asymmetry until a " + "change deliberately adds it." + ) + mm_paths = {_parse(r)[1] for r in _routes_of(live_snapshots, "mm_encoder")} + assert "/metrics" in mm_paths + + +def test_launcher_forwarding_shadows_local_routes(live_snapshots): + """An attached frontend's forwarding routes precede the local handlers. + + Launcher-owned routes are registered before the local table with the same + method and path (or a path parameter that matches the same requests), and + FastAPI matches in registration order: flipping the order silently serves + launcher-owned state from the local process instead. + """ + records = _routes_of(live_snapshots, "llm_multi_frontend_attached") + base = _routes_of(live_snapshots, "llm") + assert len(records) == len(base) + 10, ( + "the attached-frontend configuration registers the ten launcher " + "forwarding routes on top of the default table" + ) + forward_positions = [i for i, r in enumerate(records) if r.endswith("#forward")] + assert len(forward_positions) == 7, records + for position in forward_positions: + methods_path = " ".join(records[position].split(" ")[:2]) + shadowed = [ + j + for j, record in enumerate(records) + if j != position and " ".join(record.split(" ")[:2]) == methods_path + ] + assert shadowed and all(position < j for j in shadowed), ( + f"{records[position]} must stay registered before the local handler of {methods_path}" + ) + rest_positions = [i for i, r in enumerate(records) if "{rest:path}" in _parse(r)[1]] + batch_id_positions = [i for i, r in enumerate(records) if "{batch_id}" in _parse(r)[1]] + assert rest_positions and batch_id_positions + assert max(rest_positions) < min(batch_id_positions), ( + "the batches forwarding routes (a path parameter that matches every " + "batch id) must stay registered before the local batch-id handlers" + ) + + +def test_metrics_shadowing_order(live_snapshots): + records = _routes_of(live_snapshots, "llm_perf_metrics_enable_metrics_env") + metrics = [r for r in records if _parse(r)[1] == "/metrics" and _parse(r)[0] == "GET"] + assert len(metrics) == 2, metrics + assert metrics[0].endswith("#get_iteration_stats") and metrics[1].endswith("#metrics"), ( + "the JSON /metrics route must stay registered before the Prometheus-text /metrics " + "route: FastAPI matches in registration order, so swapping them silently changes " + "the /metrics payload when ENABLE_METRICS is set." + ) + + +# The only configurations that fail at construction, with their exact errors. +_ERROR_CONFIGURATIONS = { + "llm_rl_and_runtime_control": ( + "ValueError: Legacy RL control and runtime control endpoints cannot both be enabled." + ), + "mm_encoder_rl_control": "ValueError: RL control endpoints require AsyncLLM", + "embedding_runtime_control": ( + "ValueError: Runtime control endpoints are not supported for embedding servers" + ), +} + + +def test_valid_configurations_build(live_snapshots): + failing = { + name: snapshot.get("error") + for name, snapshot in live_snapshots.items() + if name not in _ERROR_CONFIGURATIONS and "routes" not in snapshot + } + assert not failing, ( + f"these configurations must build but fail at construction: {failing}. Only the " + "configurations in _ERROR_CONFIGURATIONS may fail, and a regenerated reference that " + "records the error does not change that." + ) + + +def test_error_configurations(tmp_path): + by_name = {config.name: config for config in CONFIGS} + for name, message in _ERROR_CONFIGURATIONS.items(): + snapshot = snapshot_config(by_name[name], tmp_path / name) + assert snapshot.get("error") == message, ( + f"configuration '{name}' must fail at construction with exactly " + f"{message!r}, got {snapshot.get('error')!r}: the conflicting route groups are " + "rejected before any route is registered." + ) diff --git a/tests/unittest/disaggregated/test_coordinator_e2e.py b/tests/unittest/disaggregated/test_coordinator_e2e.py index 9b058001c2c8..530015a7f2c1 100644 --- a/tests/unittest/disaggregated/test_coordinator_e2e.py +++ b/tests/unittest/disaggregated/test_coordinator_e2e.py @@ -54,11 +54,11 @@ ServerRole, ) from tensorrt_llm.logger import logger +from tensorrt_llm.serve._base.msgspec_route import _MsgspecRequest from tensorrt_llm.serve.coordinator_server import CoordinatorServer from tensorrt_llm.serve.disagg_coordinator import DisaggCoordinatorService from tensorrt_llm.serve.openai_client import OpenAIHttpClient from tensorrt_llm.serve.openai_disagg_server import OpenAIDisaggServer -from tensorrt_llm.serve.openai_server import _MsgspecRequest GEN_TEXT = "HELLO_FROM_GEN" diff --git a/tests/unittest/disaggregated/test_disagg_entry_point.py b/tests/unittest/disaggregated/test_disagg_entry_point.py new file mode 100644 index 000000000000..e239f2451ff0 --- /dev/null +++ b/tests/unittest/disaggregated/test_disagg_entry_point.py @@ -0,0 +1,457 @@ +# 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. +"""Behavior of the disaggregated server's entry-point wrapper, pinned. + +``OpenAIDisaggServer._wrap_entry_point`` turns a service entry point into the +FastAPI handler of the OpenAI routes: it counts the request, applies the +chat-template permission, resolves the conversation identity from the headers, +builds the per-request hooks, runs the entry point, counts the response, +serializes it, and sends every failure through ``_handle_exception`` while the +exception is being handled. These tests pin that behavior - the per-request +counters, the status codes, the serialization and the single trace span - so +the wrapper's internals can be reorganized without changing what a client or +a metrics scrape observes. +""" + +import asyncio +import contextlib +import json +import signal +from types import SimpleNamespace +from unittest.mock import AsyncMock, Mock + +import aiohttp +import pytest +from fastapi import HTTPException +from fastapi.responses import JSONResponse, StreamingResponse +from openai.types.responses import ResponseFormatTextJSONSchemaConfig, ResponseTextConfig +from opentelemetry.trace import SpanKind +from starlette.requests import Request +from test_common.serve.app_factory import DisaggConfig, isolated_serve_env + +import tensorrt_llm.serve.openai_disagg_server as disagg_server_module +from tensorrt_llm.executor import CppExecutorError +from tensorrt_llm.executor.utils import context_length_exceeded_message +from tensorrt_llm.llmapi import tracing +from tensorrt_llm.serve.openai_disagg_server import RawRequestResponseHooks +from tensorrt_llm.serve.openai_protocol import ( + ChatCompletionRequest, + ChatCompletionResponse, + ChatCompletionResponseChoice, + ChatMessage, + ResponsesRequest, + ResponsesResponse, + UsageInfo, +) +from tensorrt_llm.serve.openai_server import OpenAIServer + +pytestmark = pytest.mark.cpu_only + +COUNTER_NAMES = ( + "total_requests", + "stream_requests", + "nonstream_requests", + "validation_exceptions", + "http_exceptions", + "internal_errors", + "total_responses", +) + + +def _mock_collector() -> SimpleNamespace: + """The seven server counters and the queue-latency histogram, as mocks.""" + return SimpleNamespace( + **{name: SimpleNamespace(inc=Mock()) for name in COUNTER_NAMES}, + queue_latency_seconds=SimpleNamespace(observe=Mock()), + ) + + +def _server_with_counters(env): + server = env.build_disagg_server(DisaggConfig()) + counters = _mock_collector() + server._perf_metrics_collector = counters + return server, counters + + +def _assert_counts(counters: SimpleNamespace, case: str, **expected: int) -> None: + for name in COUNTER_NAMES: + actual = getattr(counters, name).inc.call_count + assert actual == expected.get(name, 0), ( + f"the {name} counter changed for {case}; _wrap_entry_point and its " + "steps must keep the request accounting exactly" + ) + + +def _raw_request(*headers: tuple) -> Request: + return Request( + { + "type": "http", + "headers": list(headers), + "state": {"server_arrival_time": 1.0, "perf_metrics_records": []}, + } + ) + + +def _chat_request(**overrides) -> ChatCompletionRequest: + payload = {"model": "m", "messages": [{"role": "user", "content": "hi"}]} + payload.update(overrides) + return ChatCompletionRequest(**payload) + + +def _chat_response() -> ChatCompletionResponse: + return ChatCompletionResponse( + id="chatcmpl-entry-point-test", + model="m", + choices=[ + ChatCompletionResponseChoice( + index=0, + message=ChatMessage(role="assistant", content="hello"), + finish_reason="stop", + ) + ], + usage=UsageInfo(prompt_tokens=3, completion_tokens=2, total_tokens=5), + ) + + +def test_non_streaming_request_is_counted_and_serialized(tmp_path): + with isolated_serve_env(tmp_path) as env: + server, counters = _server_with_counters(env) + chat_response = _chat_response() + entry_point = AsyncMock(return_value=chat_response) + wrapper = server._wrap_entry_point(entry_point, ChatCompletionRequest) + req = _chat_request() + raw_req = _raw_request() + + response = asyncio.run(wrapper(req, raw_req)) + + assert isinstance(response, JSONResponse) + assert json.loads(bytes(response.body)) == chat_response.model_dump(by_alias=True) + entry_point.assert_awaited_once() + called_req, called_hooks = entry_point.await_args.args + assert called_req is req + assert isinstance(called_hooks, RawRequestResponseHooks) + assert called_hooks.raw_req is raw_req + _assert_counts( + counters, + "a non-streaming chat request", + total_requests=1, + nonstream_requests=1, + total_responses=1, + ) + + +def test_streaming_request_returns_the_generator_unchanged(tmp_path): + with isolated_serve_env(tmp_path) as env: + server, counters = _server_with_counters(env) + + async def generated(): + yield b"data: [DONE]\n\n" + + generator = generated() + entry_point = AsyncMock(return_value=generator) + wrapper = server._wrap_entry_point(entry_point, ChatCompletionRequest) + + response = asyncio.run(wrapper(_chat_request(stream=True), _raw_request())) + + assert isinstance(response, StreamingResponse) + assert response.media_type == "text/event-stream" + assert response.body_iterator is generator + _assert_counts( + counters, + "a streaming chat request", + total_requests=1, + stream_requests=1, + total_responses=1, + ) + + +def test_forbidden_chat_template_is_a_400_before_any_worker(tmp_path): + with isolated_serve_env(tmp_path) as env: + server, counters = _server_with_counters(env) + entry_point = AsyncMock() + wrapper = server._wrap_entry_point(entry_point, ChatCompletionRequest) + + with pytest.raises(HTTPException) as exc_info: + asyncio.run(wrapper(_chat_request(chat_template="x"), _raw_request())) + + assert exc_info.value.status_code == 400 + entry_point.assert_not_awaited() + _assert_counts( + counters, + "a request with a forbidden chat template", + total_requests=1, + nonstream_requests=1, + http_exceptions=1, + ) + + +def test_conversation_id_comes_from_the_headers(tmp_path): + with isolated_serve_env(tmp_path) as env: + server, _ = _server_with_counters(env) + entry_point = AsyncMock(return_value=_chat_response()) + wrapper = server._wrap_entry_point(entry_point, ChatCompletionRequest) + + asyncio.run(wrapper(_chat_request(), _raw_request((b"x-session-id", b"s")))) + + (called_req, _), _ = entry_point.await_args + assert called_req.conversation_params is not None + assert called_req.conversation_params.conversation_id == "s" + + +def test_upstream_error_keeps_its_status(tmp_path): + with isolated_serve_env(tmp_path) as env: + server, counters = _server_with_counters(env) + entry_point = AsyncMock( + side_effect=aiohttp.ClientResponseError( + request_info=Mock(), history=(), status=404, message='{"message": "nope"}' + ) + ) + wrapper = server._wrap_entry_point(entry_point, ChatCompletionRequest) + + with pytest.raises(HTTPException) as exc_info: + asyncio.run(wrapper(_chat_request(), _raw_request())) + + assert exc_info.value.status_code == 404 + assert exc_info.value.detail == "nope" + _assert_counts( + counters, + "an upstream worker error", + total_requests=1, + nonstream_requests=1, + http_exceptions=1, + ) + + +def test_context_length_error_returns_the_worker_envelope(tmp_path): + with isolated_serve_env(tmp_path) as env: + server, counters = _server_with_counters(env) + worker_body = OpenAIServer.create_error_response( + context_length_exceeded_message(2048, 4096) + ).body.decode() + # The error the way openai_client raises one: post_json/_send_request + # use message=f"{reason}: {body[:2048]}" with the worker's response text. + entry_point = AsyncMock( + side_effect=aiohttp.ClientResponseError( + request_info=Mock(), + history=(), + status=400, + message=f"Bad Request: {worker_body}", + ) + ) + wrapper = server._wrap_entry_point(entry_point, ChatCompletionRequest) + + response = asyncio.run(wrapper(_chat_request(), _raw_request())) + + assert isinstance(response, JSONResponse) + assert response.status_code == 400 + assert json.loads(bytes(response.body))["code"] == "context_length_exceeded" + _assert_counts( + counters, + "a context-length rejection", + total_requests=1, + nonstream_requests=1, + http_exceptions=1, + ) + + +def test_engine_fatal_error_returns_none_after_sigint(tmp_path, monkeypatch): + with isolated_serve_env(tmp_path) as env: + server, counters = _server_with_counters(env) + recorder = Mock() + monkeypatch.setattr(signal, "raise_signal", recorder) + monkeypatch.setattr(disagg_server_module, "record_termination_observation", Mock()) + entry_point = AsyncMock(side_effect=CppExecutorError("x")) + wrapper = server._wrap_entry_point(entry_point, ChatCompletionRequest) + + response = asyncio.run(wrapper(_chat_request(), _raw_request())) + + assert response is None + recorder.assert_called_once_with(signal.SIGINT) + _assert_counts( + counters, + "an engine fatal error", + total_requests=1, + nonstream_requests=1, + ) + + +def test_unexpected_error_is_a_500(tmp_path): + with isolated_serve_env(tmp_path) as env: + server, counters = _server_with_counters(env) + entry_point = AsyncMock(side_effect=RuntimeError("boom")) + wrapper = server._wrap_entry_point(entry_point, ChatCompletionRequest) + + with pytest.raises(HTTPException) as exc_info: + asyncio.run(wrapper(_chat_request(), _raw_request())) + + assert exc_info.value.status_code == 500 + assert exc_info.value.detail == "Internal server error boom" + _assert_counts( + counters, + "an unexpected error", + total_requests=1, + nonstream_requests=1, + internal_errors=1, + ) + + +def test_serialization_failure_goes_to_the_error_path(tmp_path): + with isolated_serve_env(tmp_path) as env: + server, counters = _server_with_counters(env) + + class _ExplodingResponse: + def model_dump(self, **kwargs): + raise RuntimeError("cannot serialize") + + entry_point = AsyncMock(return_value=_ExplodingResponse()) + wrapper = server._wrap_entry_point(entry_point, ChatCompletionRequest) + + with pytest.raises(HTTPException) as exc_info: + asyncio.run(wrapper(_chat_request(), _raw_request())) + + assert exc_info.value.status_code == 500 + _assert_counts( + counters, + "a response that fails to serialize", + total_requests=1, + nonstream_requests=1, + total_responses=1, + internal_errors=1, + ) + + +class _FakeTracer: + def __init__(self): + self.spans = [] + + @contextlib.contextmanager + def start_as_current_span(self, name, kind=None, context=None): + self.spans.append((name, kind)) + yield SimpleNamespace(set_status=Mock(), record_exception=Mock()) + + +def test_one_trace_span_per_request(tmp_path, monkeypatch): + with isolated_serve_env(tmp_path) as env: + server, _ = _server_with_counters(env) + fake = _FakeTracer() + monkeypatch.setattr(tracing, "global_otlp_tracer", lambda: fake) + entry_point = AsyncMock(return_value=_chat_response()) + wrapper = server._wrap_entry_point(entry_point, ChatCompletionRequest) + + asyncio.run(wrapper(_chat_request(), _raw_request())) + assert fake.spans == [("disaggregated_request", SpanKind.SERVER)] + + with pytest.raises(HTTPException): + asyncio.run(wrapper(_chat_request(chat_template="x"), _raw_request())) + assert fake.spans == [("disaggregated_request", SpanKind.SERVER)] * 2 + + +def test_anthropic_messages_still_counts_like_chat(tmp_path): + with isolated_serve_env(tmp_path) as env: + server, counters = _server_with_counters(env) + server._service = SimpleNamespace( + openai_chat_completion=AsyncMock(return_value=_chat_response()) + ) + client = env.client(server.app, lifespan=False) + + response = client.post( + "/v1/messages", + json={ + "model": "m", + "max_tokens": 16, + "messages": [{"role": "user", "content": "hi"}], + }, + ) + + assert response.status_code == 200 + assert response.json()["content"][0]["text"] == "hello" + _assert_counts( + counters, + "an Anthropic Messages request", + total_requests=1, + nonstream_requests=1, + total_responses=1, + ) + + +def test_non_streaming_responses_request_serializes_by_alias(tmp_path): + with isolated_serve_env(tmp_path) as env: + server, counters = _server_with_counters(env) + responses_response = ResponsesResponse( + model="m", + output=[], + parallel_tool_calls=False, + temperature=1.0, + tool_choice="auto", + tools=[], + top_p=1.0, + background=False, + service_tier="auto", + status="completed", + text=ResponseTextConfig( + format=ResponseFormatTextJSONSchemaConfig( + name="shape", schema={"type": "object"}, type="json_schema" + ) + ), + top_logprobs=0, + truncation="disabled", + ) + entry_point = AsyncMock(return_value=responses_response) + wrapper = server._wrap_entry_point(entry_point, ResponsesRequest) + + response = asyncio.run(wrapper(ResponsesRequest(model="m", input="hi"), _raw_request())) + + assert isinstance(response, JSONResponse) + body = json.loads(bytes(response.body)) + # by_alias serialization: text.format.schema goes out under its wire + # name, not the Python field name schema_. + assert body["text"]["format"]["schema"] == {"type": "object"} + assert "schema_" not in body["text"]["format"] + assert body == responses_response.model_dump(by_alias=True) + _assert_counts( + counters, + "a non-streaming Responses request", + total_requests=1, + nonstream_requests=1, + total_responses=1, + ) + + +def test_streamed_responses_request_is_guarded(tmp_path): + with isolated_serve_env(tmp_path) as env: + server, counters = _server_with_counters(env) + + async def generated(): + yield b"data: {}\n\n" + + generator = generated() + entry_point = AsyncMock(return_value=generator) + wrapper = server._wrap_entry_point(entry_point, ResponsesRequest) + req = ResponsesRequest(model="m", input="hi", stream=True) + + response = asyncio.run(wrapper(req, _raw_request())) + + assert isinstance(response, StreamingResponse) + assert response.media_type == "text/event-stream" + # A relayed Responses stream is wrapped so that an early failure ends + # with an `error` event; the raw generator must not go out as-is. + assert response.body_iterator is not generator + _assert_counts( + counters, + "a streamed Responses request", + total_requests=1, + stream_requests=1, + total_responses=1, + ) diff --git a/tests/unittest/disaggregated/test_disagg_server_contract.py b/tests/unittest/disaggregated/test_disagg_server_contract.py new file mode 100644 index 000000000000..28a7121f2f78 --- /dev/null +++ b/tests/unittest/disaggregated/test_disagg_server_contract.py @@ -0,0 +1,347 @@ +# 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. +"""What a subclass of ``OpenAIDisaggServer`` outside this repository relies on. + +NeMo-RL's ``OpenAIDisaggServerAdaptor`` subclasses ``OpenAIDisaggServer``. It +constructs the server as ``(config, req_timeout_secs=..., coordinator_url=None)`` +plus one keyword argument of its own, adds a pure-ASGI middleware to +``self.app`` right after ``super().__init__()`` (so the app must exist with +every route registered when ``__init__`` returns), replaces ``__call__(host, +port, sockets=None)`` to keep a handle on its own uvicorn server, overrides +``_wrap_entry_point`` (swapping the request type and re-wrapping the returned +handler) and ``_handle_exception`` (answering a 4xx upstream error itself, +counting it on ``_perf_metrics_collector.http_exceptions``, delegating the +rest to ``super()``). These tests build a subclass of the same shape and pin +that surface. A failing test here means the change needs the NeMo-RL owners' +agreement first; see the ``OpenAIDisaggServer`` class docstring. +""" + +import asyncio +import importlib +import inspect +import json +from types import SimpleNamespace +from unittest.mock import AsyncMock, Mock + +import aiohttp +import pytest +from fastapi import FastAPI, HTTPException, Request, Response +from fastapi.responses import JSONResponse +from test_common.serve.app_factory import isolated_serve_env +from test_common.serve.routes import effective_routes + +from tensorrt_llm.llmapi.disagg_utils import ( + CtxGenServerConfig, + DisaggServerConfig, + RouterConfig, + ServerRole, +) +from tensorrt_llm.serve.openai_protocol import ( + ChatCompletionRequest, + ChatCompletionResponse, + ChatCompletionResponseChoice, + ChatMessage, + CompletionRequest, + UCompletionRequest, + UsageInfo, +) + +pytestmark = pytest.mark.cpu_only + +# The import paths of the class. A move that keeps the old path as a shim adds +# the new path here and keeps the old one, so the contract holds through both. +CONTRACT_MODULES = ("tensorrt_llm.serve.openai_disagg_server",) + + +@pytest.fixture(params=CONTRACT_MODULES) +def server_cls(request): + return importlib.import_module(request.param).OpenAIDisaggServer + + +def _contract(item: str) -> str: + return ( + f"{item}: NeMo-RL subclasses OpenAIDisaggServer and overrides this; " + "agree the change with the NeMo-RL owners first " + "(see the OpenAIDisaggServer docstring)" + ) + + +class _HeaderMiddleware: + """Pure ASGI, the way NeMo-RL installs its request adaptor.""" + + def __init__(self, app): + self.app = app + + async def __call__(self, scope, receive, send): + if scope["type"] != "http": + await self.app(scope, receive, send) + return + + async def send_with_header(message): + if message["type"] == "http.response.start": + message.setdefault("headers", []).append((b"x-contract", b"1")) + await send(message) + + await self.app(scope, receive, send_with_header) + + +def _contract_subclass(server_cls): + """A subclass shaped like NeMo-RL's ``OpenAIDisaggServerAdaptor``.""" + + class _Adaptor(server_cls): + def __init__(self, *args, **kwargs): + self.marker = kwargs.pop("marker") + self.wrapped = [] + super().__init__(*args, **kwargs) + self.app.add_middleware(_HeaderMiddleware) + + async def __call__(self, host, port, sockets=None): + raise AssertionError("the contract tests never serve") + + def _handle_exception(self, exception): + if isinstance(exception, aiohttp.ClientResponseError) and 400 <= exception.status < 500: + self._perf_metrics_collector.http_exceptions.inc() + raise HTTPException(status_code=exception.status, detail=exception.message) + return super()._handle_exception(exception) + + def _wrap_entry_point(self, entry_point, request_type=UCompletionRequest): + self.wrapped.append((entry_point, request_type)) + inner = super()._wrap_entry_point(entry_point, request_type) + + async def wrapper(req: request_type, raw_req: Request) -> Response: + response = await inner(req, raw_req) + if isinstance(response, JSONResponse): + response.headers["x-wrapped"] = "1" + return response + + return wrapper + + return _Adaptor + + +def _disagg_server_config() -> DisaggServerConfig: + """Two workers on closed loopback ports: construction is network-free.""" + return DisaggServerConfig( + server_configs=[ + CtxGenServerConfig(type="ctx", hostname="127.0.0.1", port=9), + CtxGenServerConfig(type="gen", hostname="127.0.0.1", port=10), + ], + hostname="127.0.0.1", + port=0, + ctx_router_config=RouterConfig(server_role=ServerRole.CONTEXT), + gen_router_config=RouterConfig(server_role=ServerRole.GENERATION), + node_id=0, + ) + + +def _build_adaptor(server_cls): + """Construct the subclass with NeMo-RL's call shape.""" + config = _disagg_server_config() + return _contract_subclass(server_cls)( + config, req_timeout_secs=7, coordinator_url=None, marker=1 + ) + + +def _raw_request() -> Request: + return Request( + { + "type": "http", + "headers": [], + "state": {"server_arrival_time": 1.0, "perf_metrics_records": []}, + } + ) + + +def _run_through_subclass(coroutine): + """Await a wrapped request built by the subclass. + + A TypeError here means the base changed how it calls what the subclass + overrides: NeMo-RL's _handle_exception takes exactly one positional + argument (the exception) and its wrapper takes (req, raw_req). + """ + try: + return asyncio.run(coroutine) + except TypeError as error: + raise AssertionError( + _contract( + "_handle_exception called with exactly one positional argument (the " + "exception) and the wrapper called with (req, raw_req)" + ) + ) from error + + +def _chat_response() -> ChatCompletionResponse: + return ChatCompletionResponse( + id="chatcmpl-contract-test", + model="m", + choices=[ + ChatCompletionResponseChoice( + index=0, + message=ChatMessage(role="assistant", content="hello"), + finish_reason="stop", + ) + ], + usage=UsageInfo(prompt_tokens=3, completion_tokens=2, total_tokens=5), + ) + + +def test_constructor_takes_config_timeout_and_coordinator_url(tmp_path, server_cls): + config = _disagg_server_config() + try: + inspect.signature(server_cls).bind(config, req_timeout_secs=7, coordinator_url=None) + except TypeError as error: + raise AssertionError( + _contract( + "the constructor call shape (config, req_timeout_secs=..., coordinator_url=None)" + ) + ) from error + with isolated_serve_env(tmp_path): + server = _build_adaptor(server_cls) + assert server.marker == 1, _contract( + "a subclass keyword argument popped before super().__init__" + ) + + +def test_app_and_routes_exist_when_init_returns(tmp_path, server_cls): + with isolated_serve_env(tmp_path): + server = _build_adaptor(server_cls) + assert isinstance(server.app, FastAPI), _contract("self.app set when __init__ returns") + post_routes = { + path: route + for path, route in effective_routes(server.app) + if "POST" in (getattr(route, "methods", None) or ()) + } + for path in ("/v1/completions", "/v1/chat/completions"): + assert path in post_routes, _contract( + f"the {path} route registered when __init__ returns" + ) + assert post_routes[path].endpoint.__module__ == __name__, _contract( + f"the {path} endpoint is the wrapper the subclass returned" + ) + wrapped_types = [request_type for _, request_type in server.wrapped] + assert any(t is CompletionRequest for t in wrapped_types), _contract( + "_wrap_entry_point called with CompletionRequest, compared by identity" + ) + assert any(t is ChatCompletionRequest for t in wrapped_types), _contract( + "_wrap_entry_point called with ChatCompletionRequest, compared by identity" + ) + subclass_endpoints = [ + route + for _, route in effective_routes(server.app) + if getattr(route, "endpoint", None) is not None + and route.endpoint.__module__ == __name__ + ] + assert len(server.wrapped) == len(subclass_endpoints), _contract( + "every route built through _wrap_entry_point goes through the override" + ) + + +def test_middleware_added_after_init_runs(tmp_path, server_cls): + with isolated_serve_env(tmp_path) as env: + server = _build_adaptor(server_cls) + response = env.client(server.app, lifespan=False).get("/version") + assert response.status_code == 200, _contract( + "a request served by self.app after construction" + ) + assert response.headers.get("x-contract") == "1", _contract( + "middleware added to self.app right after super().__init__ runs" + ) + + +def test_wrapper_result_reaches_the_client(tmp_path, server_cls): + with isolated_serve_env(tmp_path): + server = _build_adaptor(server_cls) + chat_response = _chat_response() + wrapper = server._wrap_entry_point( + AsyncMock(return_value=chat_response), ChatCompletionRequest + ) + req = ChatCompletionRequest(model="m", messages=[{"role": "user", "content": "hi"}]) + + response = _run_through_subclass(wrapper(req, _raw_request())) + + assert isinstance(response, JSONResponse), _contract( + "the base wrapper's non-streaming result, re-wrappable by a subclass" + ) + assert response.headers.get("x-wrapped") == "1", _contract( + "the subclass wrapper around the base wrapper's result" + ) + assert json.loads(bytes(response.body)) == chat_response.model_dump(by_alias=True), ( + _contract("the response body the base wrapper serializes") + ) + + +def test_every_failure_reaches_the_overridden_handle_exception(tmp_path, server_cls): + with isolated_serve_env(tmp_path): + server = _build_adaptor(server_cls) + assert callable(server._perf_metrics_collector.http_exceptions.inc), _contract( + "_perf_metrics_collector.http_exceptions.inc" + ) + server._perf_metrics_collector = SimpleNamespace( + http_exceptions=SimpleNamespace(inc=Mock()), + internal_errors=SimpleNamespace(inc=Mock()), + queue_latency_seconds=SimpleNamespace(observe=Mock()), + total_requests=SimpleNamespace(inc=Mock()), + stream_requests=SimpleNamespace(inc=Mock()), + nonstream_requests=SimpleNamespace(inc=Mock()), + total_responses=SimpleNamespace(inc=Mock()), + ) + req = ChatCompletionRequest(model="m", messages=[{"role": "user", "content": "hi"}]) + + upstream = AsyncMock( + side_effect=aiohttp.ClientResponseError( + request_info=Mock(), history=(), status=413, message="too long" + ) + ) + with pytest.raises(HTTPException) as exc_info: + _run_through_subclass( + server._wrap_entry_point(upstream, ChatCompletionRequest)(req, _raw_request()) + ) + assert exc_info.value.status_code == 413, _contract( + "a 4xx upstream error answered by the overridden _handle_exception" + ) + assert exc_info.value.detail == "too long", _contract( + "the overridden _handle_exception's own error text, not the base's" + ) + assert server._perf_metrics_collector.http_exceptions.inc.call_count == 1, _contract( + "the override counting on _perf_metrics_collector.http_exceptions" + ) + + unexpected = AsyncMock(side_effect=RuntimeError("boom")) + with pytest.raises(HTTPException) as exc_info: + _run_through_subclass( + server._wrap_entry_point(unexpected, ChatCompletionRequest)(req, _raw_request()) + ) + assert exc_info.value.status_code == 500, _contract( + "errors the override delegates reach super()._handle_exception" + ) + assert exc_info.value.detail == "Internal server error boom", _contract( + "super()._handle_exception's 500 for a delegated error" + ) + assert server._perf_metrics_collector.internal_errors.inc.call_count == 1, _contract( + "super()._handle_exception counting internal errors" + ) + + +def test_call_serves_host_port_sockets(server_cls): + assert inspect.iscoroutinefunction(server_cls.__call__), _contract( + "__call__ as the coroutine that serves" + ) + parameters = inspect.signature(server_cls.__call__).parameters + assert list(parameters) == ["self", "host", "port", "sockets"], _contract( + "__call__(self, host, port, sockets)" + ) + assert parameters["sockets"].default is None, _contract( + "the sockets argument of __call__ defaulting to None" + ) diff --git a/tests/unittest/disaggregated/test_msgspec_route.py b/tests/unittest/disaggregated/test_msgspec_route.py index e9c5e0269c08..f8a988d59e9f 100644 --- a/tests/unittest/disaggregated/test_msgspec_route.py +++ b/tests/unittest/disaggregated/test_msgspec_route.py @@ -19,8 +19,8 @@ from fastapi.testclient import TestClient from pydantic import BaseModel +from tensorrt_llm.serve._base.msgspec_route import _MsgspecRequest, _MsgspecRoute from tensorrt_llm.serve.openai_client import MSGPACK_HEADERS, _msgpack_encoder -from tensorrt_llm.serve.openai_server import _MsgspecRequest, _MsgspecRoute PAYLOAD = { "model": "m", diff --git a/tests/unittest/grpc/openengine/test_metrics.py b/tests/unittest/grpc/openengine/test_metrics.py index 3e97cc47d4b5..df3662fbeb30 100644 --- a/tests/unittest/grpc/openengine/test_metrics.py +++ b/tests/unittest/grpc/openengine/test_metrics.py @@ -15,6 +15,7 @@ from urllib.request import urlopen import pytest +from starlette.routing import Match pytestmark = [pytest.mark.cpu_only, pytest.mark.threadleak(enabled=False)] @@ -620,10 +621,16 @@ def _http_enablement() -> None: server = OpenAIServer(generator, "test-model", None, None, None) try: enabled = requested or output_dir is not None - assert ( - any(route.path == "/prometheus/metrics" for route in server.app.routes) - == enabled - ) + # Match a request instead of reading .path off app.routes: FastAPI + # 0.139 keeps included routers there as objects without a path. + scope = { + "type": "http", + "method": "GET", + "path": "/prometheus/metrics", + "path_params": {}, + } + exposed = any(route.matches(scope)[0] == Match.FULL for route in server.app.routes) + assert exposed == enabled assert server._expose_perf_metrics is requested assert server._collect_perf_metrics is enabled assert args.return_perf_metrics is enabled diff --git a/tests/unittest/llmapi/apps/test_anthropic_adapter.py b/tests/unittest/llmapi/apps/test_anthropic_adapter.py index c06404d993dc..7c53b736ca20 100644 --- a/tests/unittest/llmapi/apps/test_anthropic_adapter.py +++ b/tests/unittest/llmapi/apps/test_anthropic_adapter.py @@ -26,8 +26,8 @@ import pytest from pydantic import ValidationError -import tensorrt_llm.serve.anthropic_adapter as adapter -from tensorrt_llm.serve.anthropic_adapter import ( +import tensorrt_llm.serve.api.anthropic.adapter as adapter +from tensorrt_llm.serve.api.anthropic.adapter import ( AnthropicRequestError, AnthropicResponseError, AnthropicStreamReframer, @@ -37,7 +37,7 @@ map_stop_reason, reframe_openai_stream, ) -from tensorrt_llm.serve.anthropic_protocol import AnthropicMessagesRequest +from tensorrt_llm.serve.api.anthropic.protocol import AnthropicMessagesRequest from tensorrt_llm.serve.openai_protocol import ( ChatCompletionResponse, ChatCompletionResponseChoice, @@ -1758,7 +1758,7 @@ def test_every_known_block_type_selects_its_own_model(): catch-all -- where it would be reported as unsupported even though the server models it. Nothing else would notice. """ - from tensorrt_llm.serve.anthropic_protocol import ( + from tensorrt_llm.serve.api.anthropic.protocol import ( _CONTENT_BLOCK_MODELS, AnthropicMessage, AnthropicUnknownBlock, @@ -1858,7 +1858,7 @@ def test_cached_tokens_exceeding_prompt_tokens_is_logged(monkeypatch): tensorrt_llm.logger does not propagate to the stdlib root, so caplog sees nothing and the test would pass without the warning existing. """ - import tensorrt_llm.serve.anthropic_adapter as adapter + import tensorrt_llm.serve.api.anthropic.adapter as adapter warnings = [] monkeypatch.setattr(adapter.logger, "warning", warnings.append) diff --git a/tests/unittest/llmapi/apps/test_anthropic_batches.py b/tests/unittest/llmapi/apps/test_anthropic_batches.py index 3c3a69ba69f9..a8a16b11327a 100644 --- a/tests/unittest/llmapi/apps/test_anthropic_batches.py +++ b/tests/unittest/llmapi/apps/test_anthropic_batches.py @@ -23,12 +23,12 @@ import pytest -from tensorrt_llm.serve.anthropic_batches import ( +from tensorrt_llm.serve.api.anthropic.batches import ( AnthropicBatchStore, BatchStoreFullError, results_to_jsonl, ) -from tensorrt_llm.serve.anthropic_protocol import ( +from tensorrt_llm.serve.api.anthropic.protocol import ( AnthropicBatchRequestItem, AnthropicCreateBatchRequest, ) diff --git a/tests/unittest/llmapi/apps/test_anthropic_routes.py b/tests/unittest/llmapi/apps/test_anthropic_routes.py index b7e6cf62a3d6..f292c347e9ed 100644 --- a/tests/unittest/llmapi/apps/test_anthropic_routes.py +++ b/tests/unittest/llmapi/apps/test_anthropic_routes.py @@ -12,7 +12,13 @@ # 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. -"""Route-level tests for the Anthropic Messages compatibility handlers.""" +"""Route-level tests for the Anthropic Messages compatibility handlers. + +The servers are built through the real constructors with the fake generator +from ``tests/test_common/serve/``, so every request goes through the real app: +its validation handler, its routers and the serving classes. Only the chat core +(aggregated) or the service entry point (disaggregated) is stubbed. +""" import asyncio import json @@ -22,12 +28,17 @@ import aiohttp import pytest -from fastapi import FastAPI -from fastapi.responses import JSONResponse, StreamingResponse -from fastapi.testclient import TestClient +from fastapi.responses import StreamingResponse +from test_common.serve.app_factory import AggConfig, DisaggConfig, isolated_serve_env +from test_common.serve.routes import effective_routes from tensorrt_llm.inputs.chat_template_guard import UnusedChatTemplateKwargsError -from tensorrt_llm.serve import openai_server as openai_server_module +from tensorrt_llm.serve._base import sse as sse_module +from tensorrt_llm.serve.api.anthropic.protocol import AnthropicCreateBatchRequest +from tensorrt_llm.serve.api.anthropic.serving import ( + AnthropicBatchesServing, + AnthropicCountTokensServing, +) from tensorrt_llm.serve.openai_disagg_server import OpenAIDisaggServer from tensorrt_llm.serve.openai_protocol import ( ChatCompletionRequest, @@ -40,6 +51,7 @@ UsageInfo, ) from tensorrt_llm.serve.openai_server import OpenAIServer +from tensorrt_llm.serve.render import RenderResources from tensorrt_llm.serve.render import chat as render_chat_module # These tests are CPU-only (no GPU, engine or sockets) and run in the @@ -47,6 +59,15 @@ pytestmark = pytest.mark.cpu_only MODEL = "test-model" +COUNTER_NAMES = ( + "total_requests", + "stream_requests", + "nonstream_requests", + "validation_exceptions", + "http_exceptions", + "internal_errors", + "total_responses", +) def _request(**overrides): @@ -82,14 +103,8 @@ def _chat_response(*, tool_arguments=None): ) -def _json_chat_response(*, tool_arguments=None, status_code=200): - return JSONResponse( - content=_chat_response(tool_arguments=tool_arguments).model_dump(), - status_code=status_code, - ) - - -def _streaming_chat_response(): +def _streaming_chat_result(server_kind): + """The chat pipeline's stream: SSE strings on the aggregated server, the worker's bytes on the disaggregated one.""" chunks = [ ChatCompletionStreamResponse( model=MODEL, @@ -111,40 +126,46 @@ def _streaming_chat_response(): usage=UsageInfo(prompt_tokens=3, completion_tokens=2, total_tokens=5), ), ] + frames = [f"data: {chunk.model_dump_json()}\n\n" for chunk in chunks] + ["data: [DONE]\n\n"] async def source(): - for chunk in chunks: - yield f"data: {chunk.model_dump_json()}\n\n" - yield "data: [DONE]\n\n" + for frame in frames: + yield frame if server_kind == "standard" else frame.encode() - return StreamingResponse(source(), media_type="text/event-stream") + return source() -def _make_route_client(server_kind, openai_response): - app = FastAPI() - # Body validation fails before the route function runs, so without the - # server's own handler these tests would see FastAPI's 422 instead of the - # 400 + Anthropic envelope a client actually gets. - _anthropic_validation_handler(app) +def _counters(): + return SimpleNamespace( + **{name: SimpleNamespace(inc=Mock()) for name in COUNTER_NAMES}, + queue_latency_seconds=SimpleNamespace(observe=Mock()), + ) + + +def _route_client(env, server_kind, result): + """A client over the real app of one server, its chat pipeline stubbed to answer ``result``. + + Returns the client and the stub; the stub's first positional argument is + the chat request the Anthropic request was converted to. + """ + core = AsyncMock(return_value=result) if server_kind == "standard": - server = object.__new__(OpenAIServer) - server.model = MODEL - backend = AsyncMock(return_value=openai_response) - server.openai_chat = backend - else: - server = object.__new__(OpenAIDisaggServer) - server._service = SimpleNamespace(openai_chat_completion=object()) - backend = AsyncMock(return_value=openai_response) - server._wrap_entry_point = Mock(return_value=backend) - app.add_api_route("/v1/messages", server.anthropic_messages, methods=["POST"]) - return TestClient(app), backend + server = env.build_openai_server(AggConfig(model_name=MODEL)) + server.create_chat_completion = core + return env.client(server.app), core + server = env.build_disagg_server(DisaggConfig()) + server._perf_metrics_collector = _counters() + server._service = SimpleNamespace(openai_chat_completion=core) + # The disaggregated lifespan connects to workers; these tests never run it. + return env.client(server.app, lifespan=False), core @pytest.mark.parametrize("server_kind", ["standard", "disagg"]) -def test_messages_route_converts_nonstream_response(server_kind): - client, backend = _make_route_client(server_kind, _json_chat_response()) +def test_messages_route_converts_nonstream_response(server_kind, tmp_path): + with isolated_serve_env(tmp_path) as env: + client, core = _route_client(env, server_kind, _chat_response()) - response = client.post("/v1/messages", json=_request()) + response = client.post("/v1/messages", json=_request()) assert response.status_code == 200 assert response.json() | {"id": "ignored"} == { @@ -156,47 +177,49 @@ def test_messages_route_converts_nonstream_response(server_kind): "stop_reason": "end_turn", "usage": {"input_tokens": 3, "output_tokens": 2}, } - chat_request = backend.await_args.args[0] + chat_request = core.await_args.args[0] assert chat_request.model == MODEL assert chat_request.max_completion_tokens == 64 assert not chat_request.stream @pytest.mark.parametrize("server_kind", ["standard", "disagg"]) -def test_messages_route_accepts_thinking_display(server_kind): +def test_messages_route_accepts_thinking_display(server_kind, tmp_path): """The thinking config Claude Code sends is served, not answered with a 400.""" - client, backend = _make_route_client(server_kind, _json_chat_response()) + with isolated_serve_env(tmp_path) as env: + client, core = _route_client(env, server_kind, _chat_response()) - response = client.post( - "/v1/messages", - json=_request(max_tokens=4096, thinking={"type": "adaptive", "display": "updates"}), - ) + response = client.post( + "/v1/messages", + json=_request(max_tokens=4096, thinking={"type": "adaptive", "display": "updates"}), + ) assert response.status_code == 200 - assert backend.await_args.args[0].chat_template_kwargs["enable_thinking"] is True + assert core.await_args.args[0].chat_template_kwargs["enable_thinking"] is True @pytest.mark.parametrize("server_kind", ["standard", "disagg"]) -def test_messages_route_rejects_anthropic_server_tools(server_kind): - client, backend = _make_route_client(server_kind, _json_chat_response()) - - response = client.post( - "/v1/messages", - json=_request( - tools=[ - { - "name": "web_search", - "type": "web_search_20250305", - } - ] - ), - ) +def test_messages_route_rejects_anthropic_server_tools(server_kind, tmp_path): + with isolated_serve_env(tmp_path) as env: + client, core = _route_client(env, server_kind, _chat_response()) + + response = client.post( + "/v1/messages", + json=_request( + tools=[ + { + "name": "web_search", + "type": "web_search_20250305", + } + ] + ), + ) assert response.status_code == 400 assert response.json()["type"] == "error" assert response.json()["error"]["type"] == "invalid_request_error" assert "server tool" in response.json()["error"]["message"] - backend.assert_not_awaited() + core.assert_not_awaited() @pytest.mark.parametrize( @@ -214,7 +237,7 @@ def test_messages_route_rejects_anthropic_server_tools(server_kind): ), ], ) -def test_malformed_known_block_is_a_400_against_its_own_type(block, tag, field): +def test_malformed_known_block_is_a_400_against_its_own_type(block, tag, field, tmp_path): """A known type with bad fields must fail as that type, not as the catch-all. The union is tagged on `type`, so the block is checked against the model the @@ -223,14 +246,17 @@ def test_malformed_known_block_is_a_400_against_its_own_type(block, tag, field): the adapter's dispatch on `type`, and raise AttributeError on an attribute the block never had -- a 500 for a plain client error. - Errors carry `loc` as a tuple, so the tag and field appear adjacent. + Errors carry `loc` as a tuple, so the tag and field appear adjacent. The + 400 and the Anthropic envelope come from the server's own validation + handler, which the real app carries. """ - client, backend = _make_route_client("standard", _json_chat_response()) + with isolated_serve_env(tmp_path) as env: + client, core = _route_client(env, "standard", _chat_response()) - response = client.post( - "/v1/messages", - json=_request(messages=[{"role": "user", "content": [block]}]), - ) + response = client.post( + "/v1/messages", + json=_request(messages=[{"role": "user", "content": [block]}]), + ) assert response.status_code == 400 assert response.json()["error"]["type"] == "invalid_request_error" @@ -238,33 +264,37 @@ def test_malformed_known_block_is_a_400_against_its_own_type(block, tag, field): assert f"'{tag}', '{field}'" in message # The catch-all must not have absorbed it; that is the regression itself. assert "'unknown'," not in message - backend.assert_not_awaited() + core.assert_not_awaited() -def test_unknown_block_type_still_reaches_the_adapter(): +def test_unknown_block_type_still_reaches_the_adapter(tmp_path): """The catch-all exists for types this server does not model; keep it working. Tagging the union must not close it: an unmodelled type has to validate and be named by the adapter, rather than returning a wall of failed members. """ - client, backend = _make_route_client("standard", _json_chat_response()) + with isolated_serve_env(tmp_path) as env: + client, core = _route_client(env, "standard", _chat_response()) - response = client.post( - "/v1/messages", - json=_request(messages=[{"role": "user", "content": [{"type": "document", "source": {}}]}]), - ) + response = client.post( + "/v1/messages", + json=_request( + messages=[{"role": "user", "content": [{"type": "document", "source": {}}]}] + ), + ) assert response.status_code == 400 assert response.json()["error"]["type"] == "invalid_request_error" assert "document" in response.json()["error"]["message"] - backend.assert_not_awaited() + core.assert_not_awaited() @pytest.mark.parametrize("server_kind", ["standard", "disagg"]) -def test_messages_route_hides_invalid_generated_tool_arguments(server_kind): - client, _ = _make_route_client(server_kind, _json_chat_response(tool_arguments="{not-json")) +def test_messages_route_hides_invalid_generated_tool_arguments(server_kind, tmp_path): + with isolated_serve_env(tmp_path) as env: + client, _ = _route_client(env, server_kind, _chat_response(tool_arguments="{not-json")) - response = client.post("/v1/messages", json=_request()) + response = client.post("/v1/messages", json=_request()) assert response.status_code == 500 assert response.json() == { @@ -274,10 +304,11 @@ def test_messages_route_hides_invalid_generated_tool_arguments(server_kind): @pytest.mark.parametrize("server_kind", ["standard", "disagg"]) -def test_messages_route_reframes_streaming_response(server_kind): - client, backend = _make_route_client(server_kind, _streaming_chat_response()) +def test_messages_route_reframes_streaming_response(server_kind, tmp_path): + with isolated_serve_env(tmp_path) as env: + client, core = _route_client(env, server_kind, _streaming_chat_result(server_kind)) - response = client.post("/v1/messages", json=_request(stream=True)) + response = client.post("/v1/messages", json=_request(stream=True)) assert response.status_code == 200 assert response.headers["content-type"].startswith("text/event-stream") @@ -285,7 +316,7 @@ def test_messages_route_reframes_streaming_response(server_kind): assert "event: content_block_delta\n" in response.text assert '"text":"hello"' in response.text assert response.text.rstrip().endswith('event: message_stop\ndata: {"type":"message_stop"}') - assert backend.await_args.args[0].stream + assert core.await_args.args[0].stream class _StreamedGeneration: @@ -371,8 +402,9 @@ async def render(**kwargs): @pytest.fixture def server_logger(monkeypatch): + """The logger of the chat SSE generator, which lives in serve/_base/sse.py.""" log = Mock() - monkeypatch.setattr(openai_server_module, "logger", log) + monkeypatch.setattr(sse_module, "logger", log) return log @@ -450,61 +482,19 @@ async def run(): assert "engine failed" in server_logger.error.call_args.args[0] -@pytest.fixture -def prometheus_multiproc_dir(tmp_path, monkeypatch): - """PROMETHEUS_MULTIPROC_DIR for one test. +def _paths(app) -> set: + return {path for path, _ in effective_routes(app)} - prometheus_client picks and caches its value class the first time a metric - is built with the variable set, so the class is restored as well as the - environment; otherwise later metrics in the process fail without a - directory. - """ - import prometheus_client.values - - monkeypatch.setenv("PROMETHEUS_MULTIPROC_DIR", str(tmp_path)) - saved_value_class = prometheus_client.values.ValueClass - try: - yield tmp_path - finally: - prometheus_client.values.ValueClass = saved_value_class - - -def test_standard_and_disagg_register_messages_route(prometheus_multiproc_dir): - # register_routes mounts a prometheus multiprocess collector, which needs a - # real PROMETHEUS_MULTIPROC_DIR; object.__new__ skips the startup that - # creates one. - - standard = object.__new__(OpenAIServer) - standard.app = FastAPI() - standard.generator = SimpleNamespace( - _executor=SimpleNamespace(resource_governor_queue=None), - args=SimpleNamespace(return_perf_metrics=False), - ) - standard.use_harmony = False - # register_routes() reads these to decide whether to mount the optional - # control endpoints. Building the server with object.__new__ skips - # __init__, so anything register_routes() touches has to be supplied. - standard._enable_runtime_control_endpoints = False - standard._enable_rl_control_endpoints = False - standard.register_routes() - - disagg = object.__new__(OpenAIDisaggServer) - disagg.app = FastAPI() - disagg._service = SimpleNamespace( - openai_completion=AsyncMock(), - openai_chat_completion=AsyncMock(), - anthropic_count_tokens=AsyncMock(), - openai_responses=AsyncMock(), - ) - disagg._perf_metrics_collector = SimpleNamespace(get_perf_metrics=AsyncMock()) - disagg._disagg_cluster_storage = None - # register_routes reads this; the server is built with object.__new__ here, - # so every attribute it touches has to be stubbed explicitly. - disagg._coordinator = None - disagg.register_routes() - - for server in (standard, disagg): - paths = {route.path for route in server.app.routes} + +def test_standard_and_disagg_register_messages_route(tmp_path): + # One server per environment: both register collectors in the default + # Prometheus registry, which the environment restores on exit. + with isolated_serve_env(tmp_path / "standard") as env: + standard_paths = _paths(env.build_openai_server(AggConfig()).app) + with isolated_serve_env(tmp_path / "disagg") as env: + disagg_paths = _paths(env.build_disagg_server(DisaggConfig()).app) + + for paths in (standard_paths, disagg_paths): assert "/v1/messages" in paths # Claude Code calls count_tokens before most turns to size its context. @@ -512,42 +502,22 @@ def test_standard_and_disagg_register_messages_route(prometheus_multiproc_dir): # so registration is asserted separately from /v1/messages -- and on both # servers, because a client pointed at a disaggregated deployment cannot # tell the difference and must not lose the endpoint. - for server in (standard, disagg): - paths = {route.path for route in server.app.routes} + for paths in (standard_paths, disagg_paths): assert "/v1/messages/count_tokens" in paths - disagg_paths = {route.path for route in disagg.app.routes} assert "/v1/responses" in disagg_paths assert "/v1/models" in disagg_paths -def _batch_client(runner=None): - """A TestClient over just the batch routes, with a stubbed item runner.""" - from tensorrt_llm.serve.anthropic_batches import AnthropicBatchStore +def _batch_client(env, runner=None): + """A client over the real app, with the batch serving's item runner stubbed.""" async def ok(request): return "succeeded", {"type": "message", "role": "assistant", "content": []} - app = FastAPI() - _anthropic_validation_handler(app) - server = object.__new__(OpenAIServer) - server.model = MODEL - server._anthropic_batch_store = AnthropicBatchStore(runner=runner or ok) - app.add_api_route("/v1/messages/batches", server.anthropic_create_batch, methods=["POST"]) - app.add_api_route("/v1/messages/batches", server.anthropic_list_batches, methods=["GET"]) - app.add_api_route( - "/v1/messages/batches/{batch_id}", server.anthropic_get_batch, methods=["GET"] - ) - app.add_api_route( - "/v1/messages/batches/{batch_id}", server.anthropic_delete_batch, methods=["DELETE"] - ) - app.add_api_route( - "/v1/messages/batches/{batch_id}/cancel", server.anthropic_cancel_batch, methods=["POST"] - ) - app.add_api_route( - "/v1/messages/batches/{batch_id}/results", server.anthropic_batch_results, methods=["GET"] - ) - return TestClient(app) + server = env.build_openai_server(AggConfig(model_name=MODEL)) + server.app.state.anthropic_serving_batches = AnthropicBatchesServing(runner=runner or ok) + return env.client(server.app) def _batch_body(*custom_ids): @@ -566,22 +536,10 @@ def _batch_body(*custom_ids): } -def test_batch_routes_are_registered_on_the_standard_server(): - standard = object.__new__(OpenAIServer) - standard.app = FastAPI() - standard.generator = SimpleNamespace( - _executor=SimpleNamespace(resource_governor_queue=None), - args=SimpleNamespace(return_perf_metrics=False), - ) - standard.use_harmony = False - # register_routes() reads these to decide whether to mount the optional - # control endpoints. Building the server with object.__new__ skips - # __init__, so anything register_routes() touches has to be supplied. - standard._enable_runtime_control_endpoints = False - standard._enable_rl_control_endpoints = False - standard.register_routes() - - paths = {route.path for route in standard.app.routes} +def test_batch_routes_are_registered_on_the_standard_server(tmp_path): + with isolated_serve_env(tmp_path) as env: + paths = _paths(env.build_openai_server(AggConfig()).app) + for path in ( "/v1/messages/batches", "/v1/messages/batches/{batch_id}", @@ -591,9 +549,10 @@ def test_batch_routes_are_registered_on_the_standard_server(): assert path in paths, f"{path} not registered" -def test_create_batch_returns_an_in_progress_batch(): - client = _batch_client() - response = client.post("/v1/messages/batches", json=_batch_body("a", "b")) +def test_create_batch_returns_an_in_progress_batch(tmp_path): + with isolated_serve_env(tmp_path) as env: + client = _batch_client(env) + response = client.post("/v1/messages/batches", json=_batch_body("a", "b")) assert response.status_code == 200 body = response.json() @@ -603,18 +562,20 @@ def test_create_batch_returns_an_in_progress_batch(): assert body["request_counts"]["processing"] + body["request_counts"]["succeeded"] == 2 -def test_create_batch_rejects_duplicate_custom_ids(): - client = _batch_client() - response = client.post("/v1/messages/batches", json=_batch_body("same", "same")) +def test_create_batch_rejects_duplicate_custom_ids(tmp_path): + with isolated_serve_env(tmp_path) as env: + client = _batch_client(env) + response = client.post("/v1/messages/batches", json=_batch_body("same", "same")) assert response.status_code == 400 assert response.json()["error"]["type"] == "invalid_request_error" -def test_unknown_batch_says_it_is_not_found_and_explains_why(): +def test_unknown_batch_says_it_is_not_found_and_explains_why(tmp_path): """The in-memory limitation is the likeliest cause, so name it.""" - client = _batch_client() - response = client.get("/v1/messages/batches/msgbatch_nope") + with isolated_serve_env(tmp_path) as env: + client = _batch_client(env) + response = client.get("/v1/messages/batches/msgbatch_nope") assert response.status_code == 404 error = response.json()["error"] @@ -625,13 +586,10 @@ def test_unknown_batch_says_it_is_not_found_and_explains_why(): def test_results_of_an_unfinished_batch_are_400_not_404(): """404 would tell a polling client the batch never existed. - Driven directly rather than through TestClient: each TestClient call runs - in its own portal, so the batch's background task is not guaranteed to - still be in flight by the time the second request arrives, and the batch - would have already ended. + Driven on the serving class directly rather than through a client, so the + batch's background task is still in flight when the results are asked for. """ - from tensorrt_llm.serve.anthropic_batches import AnthropicBatchStore - from tensorrt_llm.serve.anthropic_protocol import AnthropicBatchRequestItem + from tensorrt_llm.serve.api.anthropic.protocol import AnthropicBatchRequestItem async def scenario(): blocked = asyncio.Event() @@ -640,8 +598,7 @@ async def never_finishes(request): await blocked.wait() return "succeeded", {"type": "message"} - server = object.__new__(OpenAIServer) - server._anthropic_batch_store = AnthropicBatchStore(runner=never_finishes) + serving = AnthropicBatchesServing(runner=never_finishes) item = AnthropicBatchRequestItem( custom_id="a", params={ @@ -650,15 +607,15 @@ async def never_finishes(request): "messages": [{"role": "user", "content": "hi"}], }, ) - batch = server._anthropic_batch_store.create([item]) + batch = serving._batch_store().create([item]) await asyncio.sleep(0) # let the worker start and block - response = await server.anthropic_batch_results(batch.id) + response = await serving.anthropic_batch_results(batch.id) assert response.status_code == 400 assert "still processing" in json.loads(response.body)["error"]["message"] # An id that was never issued is a genuine 404, and must stay distinct. - missing = await server.anthropic_batch_results("msgbatch_nope") + missing = await serving.anthropic_batch_results("msgbatch_nope") assert missing.status_code == 404 blocked.set() @@ -673,8 +630,8 @@ def test_interrupted_batch_still_reports_counts_that_sum(): end the batch with fewer results than requests - and request_counts is the number clients use to decide whether every request was accounted for. """ - from tensorrt_llm.serve.anthropic_batches import AnthropicBatchStore - from tensorrt_llm.serve.anthropic_protocol import AnthropicBatchRequestItem + from tensorrt_llm.serve.api.anthropic.batches import AnthropicBatchStore + from tensorrt_llm.serve.api.anthropic.protocol import AnthropicBatchRequestItem async def scenario(): blocked = asyncio.Event() @@ -714,29 +671,32 @@ async def never_finishes(request): asyncio.run(scenario()) -def test_finished_batch_serves_ndjson_results(): - client = _batch_client() - created = client.post("/v1/messages/batches", json=_batch_body("a", "b")).json() +def test_finished_batch_serves_ndjson_results(tmp_path): + with isolated_serve_env(tmp_path) as env: + client = _batch_client(env) + created = client.post("/v1/messages/batches", json=_batch_body("a", "b")).json() + + for _ in range(200): + if ( + client.get(f"/v1/messages/batches/{created['id']}").json()["processing_status"] + == "ended" + ): + break + time.sleep(0.01) - for _ in range(200): - if ( - client.get(f"/v1/messages/batches/{created['id']}").json()["processing_status"] - == "ended" - ): - break - time.sleep(0.01) + response = client.get(f"/v1/messages/batches/{created['id']}/results") - response = client.get(f"/v1/messages/batches/{created['id']}/results") assert response.status_code == 200 assert response.headers["content-type"].startswith("application/x-ndjson") lines = [json.loads(line) for line in response.text.splitlines()] assert {line["custom_id"] for line in lines} == {"a", "b"} -def test_list_batches_returns_a_page(): - client = _batch_client() - client.post("/v1/messages/batches", json=_batch_body("a")) - body = client.get("/v1/messages/batches").json() +def test_list_batches_returns_a_page(tmp_path): + with isolated_serve_env(tmp_path) as env: + client = _batch_client(env) + client.post("/v1/messages/batches", json=_batch_body("a")) + body = client.get("/v1/messages/batches").json() assert isinstance(body["data"], list) and body["data"] assert body["first_id"] == body["data"][0]["id"] @@ -840,125 +800,104 @@ def encode(self, text, add_special_tokens=True, **kwargs): return text.split() -def _anthropic_validation_handler(app): - """Mirror the server's handler so boundary rejections keep the envelope. +def _count_tokens_client(env, resources): + """A client over the real app, counting over ``resources`` (what the renderer reads off a server). - Model-level validation fails before the route function runs, so without - this a bare test app answers with FastAPI's 422 shape while the real - server answers 400 with the Anthropic envelope -- the test would pin the - fixture's behaviour instead of the product's. + Transitional: the resources get built directly once the server assembly + stops reading them off the server per request. """ - from fastapi.exceptions import RequestValidationError - - from tensorrt_llm.serve.anthropic_adapter import anthropic_error_response - - @app.exception_handler(RequestValidationError) - async def _handler(request, exc): # noqa: ANN001 - if request.url.path.startswith("/v1/messages"): - return anthropic_error_response(str(exc), "invalid_request_error", 400) - raise exc + server = env.build_openai_server(AggConfig(model_name=MODEL)) + server.app.state.anthropic_serving_count_tokens = AnthropicCountTokensServing( + render_resources=lambda: RenderResources.from_server(resources) + ) + return env.client(server.app) -def _count_tokens_client(): - app = FastAPI() - _anthropic_validation_handler(app) - server = object.__new__(OpenAIServer) - server.model = MODEL - server.tokenizer = _FakeTokenizer() - app.add_api_route("/v1/messages/count_tokens", server.anthropic_count_tokens, methods=["POST"]) - return TestClient(app) +def _count_resources(chat_template=None): + return SimpleNamespace(tokenizer=_FakeTokenizer(), chat_template=chat_template) -def test_count_tokens_route_returns_a_count(): - client = _count_tokens_client() +def test_count_tokens_route_returns_a_count(tmp_path): + with isolated_serve_env(tmp_path) as env: + client = _count_tokens_client(env, _count_resources()) - response = client.post( - "/v1/messages/count_tokens", - json={ - "model": MODEL, - "messages": [{"role": "user", "content": "one two three"}], - }, - ) + response = client.post( + "/v1/messages/count_tokens", + json={ + "model": MODEL, + "messages": [{"role": "user", "content": "one two three"}], + }, + ) assert response.status_code == 200 assert response.json() == {"input_tokens": 3} -def test_count_tokens_counts_the_rendered_prompt_not_the_raw_messages(): +def test_count_tokens_counts_the_rendered_prompt_not_the_raw_messages(tmp_path): """Tools inflate the real prompt, so they have to inflate the count.""" - client = _count_tokens_client() body = { "model": MODEL, "messages": [{"role": "user", "content": "one two three"}], } - - bare = client.post("/v1/messages/count_tokens", json=body).json() - with_tools = client.post( - "/v1/messages/count_tokens", - json={ - **body, - "tools": [ - { - "name": "get_weather", - "input_schema": {"type": "object"}, - } - ], - }, - ).json() + with isolated_serve_env(tmp_path) as env: + client = _count_tokens_client(env, _count_resources()) + + bare = client.post("/v1/messages/count_tokens", json=body).json() + with_tools = client.post( + "/v1/messages/count_tokens", + json={ + **body, + "tools": [ + { + "name": "get_weather", + "input_schema": {"type": "object"}, + } + ], + }, + ).json() assert with_tools["input_tokens"] > bare["input_tokens"] -def test_count_tokens_honours_the_server_chat_template(): +def test_count_tokens_honours_the_server_chat_template(tmp_path): """A server started with --chat_template must count the prompt it renders. - The renderer only consults request.chat_template; openai_chat resolves + The renderer only consults request.chat_template; the chat route resolves `request.chat_template or self.chat_template`. Without the same resolution here, count_tokens reports a prompt the server never builds - the exact drift this endpoint exists to remove. """ - app = FastAPI() - server = object.__new__(OpenAIServer) - server.model = MODEL - server.tokenizer = _FakeTokenizer() # A template that measurably inflates the rendered prompt. - server.chat_template = "SERVER-TEMPLATE-MARKER {{ messages }}" - app.add_api_route("/v1/messages/count_tokens", server.anthropic_count_tokens, methods=["POST"]) - client = TestClient(app) - - with_template = client.post( - "/v1/messages/count_tokens", - json={ - "model": MODEL, - "messages": [{"role": "user", "content": "one two three"}], - }, - ).json() + resources = _count_resources(chat_template="SERVER-TEMPLATE-MARKER {{ messages }}") + body = { + "model": MODEL, + "messages": [{"role": "user", "content": "one two three"}], + } + with isolated_serve_env(tmp_path) as env: + client = _count_tokens_client(env, resources) - server.chat_template = None - without = client.post( - "/v1/messages/count_tokens", - json={ - "model": MODEL, - "messages": [{"role": "user", "content": "one two three"}], - }, - ).json() + with_template = client.post("/v1/messages/count_tokens", json=body).json() + + resources.chat_template = None + without = client.post("/v1/messages/count_tokens", json=body).json() assert with_template["input_tokens"] != without["input_tokens"], ( "server chat_template had no effect on the count, so it was ignored" ) -def test_count_tokens_rejects_a_malformed_request(): - client = _count_tokens_client() +def test_count_tokens_rejects_a_malformed_request(tmp_path): + with isolated_serve_env(tmp_path) as env: + client = _count_tokens_client(env, _count_resources()) - response = client.post( - "/v1/messages/count_tokens", - json={ - "model": MODEL, - "messages": [{"role": "user", "content": "hi"}], - "thinking": {"type": "bogus"}, - }, - ) + response = client.post( + "/v1/messages/count_tokens", + json={ + "model": MODEL, + "messages": [{"role": "user", "content": "hi"}], + "thinking": {"type": "bogus"}, + }, + ) assert response.status_code == 400 assert response.json()["error"]["type"] == "invalid_request_error" @@ -983,7 +922,7 @@ def _disagg_service(servers, post_json=None): def test_count_tokens_forwards_to_a_context_worker(): """The disagg server has no tokenizer, so the count must come from a worker.""" - from tensorrt_llm.serve.anthropic_protocol import ( + from tensorrt_llm.serve.api.anthropic.protocol import ( AnthropicCountTokensRequest, AnthropicCountTokensResponse, ) @@ -1004,7 +943,7 @@ def test_count_tokens_forwards_to_a_context_worker(): def test_count_tokens_spreads_over_context_workers(): """Round-robin: one worker must not absorb every count.""" - from tensorrt_llm.serve.anthropic_protocol import ( + from tensorrt_llm.serve.api.anthropic.protocol import ( AnthropicCountTokensRequest, AnthropicCountTokensResponse, ) @@ -1029,7 +968,7 @@ def test_count_tokens_without_context_workers_is_an_error_not_a_zero(): Returning 0 would read as an empty prompt and silently skew a client's context budgeting. """ - from tensorrt_llm.serve.anthropic_protocol import AnthropicCountTokensRequest + from tensorrt_llm.serve.api.anthropic.protocol import AnthropicCountTokensRequest service = _disagg_service([]) request = AnthropicCountTokensRequest(model=MODEL, messages=[{"role": "user", "content": "hi"}]) @@ -1085,7 +1024,7 @@ def test_post_json_sends_msgpack_the_worker_can_decode(): """ import msgspec - from tensorrt_llm.serve.anthropic_protocol import ( + from tensorrt_llm.serve.api.anthropic.protocol import ( AnthropicCountTokensRequest, AnthropicCountTokensResponse, ) @@ -1117,7 +1056,7 @@ def test_post_json_raises_on_an_error_status(): """ import aiohttp - from tensorrt_llm.serve.anthropic_protocol import ( + from tensorrt_llm.serve.api.anthropic.protocol import ( AnthropicCountTokensRequest, AnthropicCountTokensResponse, ) @@ -1190,29 +1129,28 @@ def test_create_batch_at_capacity_is_429_not_400(): instead of surfacing a hard failure to the user. This pins the status mapping only. Whether the store actually refuses at - capacity is store policy, covered in test_anthropic_batches.py -- driving a - real store to its limit through TestClient is not possible anyway, since - each request runs in its own event loop and a batch's task does not - survive between them. + capacity is store policy, covered in test_anthropic_batches.py. Driven on + the serving class directly: the mapping needs no server. """ - from tensorrt_llm.serve.anthropic_batches import BatchStoreFullError + from tensorrt_llm.serve.api.anthropic.batches import BatchStoreFullError - app = FastAPI() - server = object.__new__(OpenAIServer) - server.model = MODEL + async def ok(request): + return "succeeded", {"type": "message"} def full(_requests): raise BatchStoreFullError( "the server is holding 100 unfinished batches (limit 100); retry once one of them ends" ) - server._anthropic_batch_store = SimpleNamespace(create=full) - app.add_api_route("/v1/messages/batches", server.anthropic_create_batch, methods=["POST"]) + serving = AnthropicBatchesServing(runner=ok) + serving._anthropic_batch_store = SimpleNamespace(create=full) - response = TestClient(app).post("/v1/messages/batches", json=_batch_body("a")) + response = asyncio.run( + serving.anthropic_create_batch(AnthropicCreateBatchRequest(**_batch_body("a"))) + ) assert response.status_code == 429 - body = response.json() + body = json.loads(response.body) assert body["type"] == "error" assert body["error"]["type"] == "rate_limit_error" # The message has to say what to do about it; an "invalid request" shape @@ -1220,46 +1158,17 @@ def full(_requests): assert "retry" in body["error"]["message"].lower() -def test_await_disconnected_returns_immediately_for_a_synthetic_request(): - """A batched request has no socket, so the watchdog must not poll on it. - - is_disconnected() can never become true for a synthesised request, so - without an explicit opt-out this coroutine loops at 1Hz for the life of - the process -- one live task per batched request, each holding its Request - and RequestOutput. - """ - server = object.__new__(OpenAIServer) - raw_request = OpenAIServer._synthetic_request( - { - "model": MODEL, - "max_tokens": 8, - "messages": [{"role": "user", "content": "hi"}], - } - ) - promise = SimpleNamespace(finished=False, request_id="req", abort=Mock()) - - async def run(): - # A real socket-backed request would never let this return; a bounded - # wait keeps the failure a clear timeout rather than a hung suite. - await asyncio.wait_for( - OpenAIServer.await_disconnected(server, raw_request, promise), - timeout=5.0, - ) - - asyncio.run(run()) - # Returning early must not be mistaken for a disconnect. - promise.abort.assert_not_called() - - @pytest.mark.parametrize("limit", [0, -1, 101]) -def test_list_batches_rejects_out_of_range_limit(limit): +def test_list_batches_rejects_out_of_range_limit(limit, tmp_path): """Clamping hid both a client error and a surprising page size. limit=0 became 1 rather than an error, and the page was bounded by the retention limit -- how many batches are kept, which has nothing to do with - how many fit on a page. + how many fit on a page. The 400 and the envelope come from the server's own + validation handler. """ - response = _batch_client().get(f"/v1/messages/batches?limit={limit}") + with isolated_serve_env(tmp_path) as env: + response = _batch_client(env).get(f"/v1/messages/batches?limit={limit}") assert response.status_code == 400, response.text assert response.json()["error"]["type"] == "invalid_request_error" diff --git a/tests/unittest/llmapi/apps/test_chat_template_kwargs_guard.py b/tests/unittest/llmapi/apps/test_chat_template_kwargs_guard.py index a3de0428a00a..34a55bb1990d 100644 --- a/tests/unittest/llmapi/apps/test_chat_template_kwargs_guard.py +++ b/tests/unittest/llmapi/apps/test_chat_template_kwargs_guard.py @@ -373,8 +373,8 @@ class TestServerInjectedControls: @staticmethod def _anthropic_chat_request(): - from tensorrt_llm.serve.anthropic_adapter import convert_anthropic_request - from tensorrt_llm.serve.anthropic_protocol import AnthropicMessagesRequest + from tensorrt_llm.serve.api.anthropic.adapter import convert_anthropic_request + from tensorrt_llm.serve.api.anthropic.protocol import AnthropicMessagesRequest return convert_anthropic_request( AnthropicMessagesRequest( diff --git a/tests/unittest/llmapi/apps/test_kimi_serve_extensions.py b/tests/unittest/llmapi/apps/test_kimi_serve_extensions.py index f2cc94d81a50..156d03085b13 100644 --- a/tests/unittest/llmapi/apps/test_kimi_serve_extensions.py +++ b/tests/unittest/llmapi/apps/test_kimi_serve_extensions.py @@ -506,8 +506,8 @@ class TestKimiAnthropicThinkingControls: """Anthropic ``thinking`` / ``output_config.effort`` reach the K3 template kwargs.""" def convert(self, **fields) -> ChatCompletionRequest: - from tensorrt_llm.serve.anthropic_adapter import convert_anthropic_request - from tensorrt_llm.serve.anthropic_protocol import AnthropicMessagesRequest + from tensorrt_llm.serve.api.anthropic.adapter import convert_anthropic_request + from tensorrt_llm.serve.api.anthropic.protocol import AnthropicMessagesRequest req = convert_anthropic_request( AnthropicMessagesRequest(model="m", max_tokens=4096, messages=[USER_MSG], **fields) diff --git a/tests/unittest/llmapi/apps/test_openai_chat_multimodal_prompt.py b/tests/unittest/llmapi/apps/test_openai_chat_multimodal_prompt.py index d6ed4425262e..fd575c107a98 100644 --- a/tests/unittest/llmapi/apps/test_openai_chat_multimodal_prompt.py +++ b/tests/unittest/llmapi/apps/test_openai_chat_multimodal_prompt.py @@ -29,6 +29,12 @@ from fastapi.testclient import TestClient from PIL import Image +from tensorrt_llm.serve.openai_protocol import ( + ChatCompletionResponse, + ChatCompletionResponseChoice, + ChatMessage, + UsageInfo, +) from tensorrt_llm.serve.openai_server import OpenAIServer from .openai_server_stub import make_stub_openai_server @@ -112,8 +118,20 @@ def generate_async(*, inputs, **kwargs): tmp_path, "test-vlm", _StubModelConfig(), generate_async, tokenizer=tokenizer ) server.await_disconnected = AsyncMock() + # The chat handler serializes a ChatCompletionResponse and streams anything + # else, so the stubbed response is a real one. server._create_chat_response = AsyncMock( - return_value=SimpleNamespace(model_dump=lambda: {"ok": True}) + return_value=ChatCompletionResponse( + model="test-vlm", + choices=[ + ChatCompletionResponseChoice( + index=0, + message=ChatMessage(role="assistant", content="ok"), + finish_reason="stop", + ) + ], + usage=UsageInfo(prompt_tokens=2, completion_tokens=1, total_tokens=3), + ) ) app = FastAPI() app.add_api_route("/v1/chat/completions", server.openai_chat, methods=["POST"]) diff --git a/tests/unittest/llmapi/apps/test_render_callers.py b/tests/unittest/llmapi/apps/test_render_callers.py index 752b8794b5cf..fd2d92dd6939 100644 --- a/tests/unittest/llmapi/apps/test_render_callers.py +++ b/tests/unittest/llmapi/apps/test_render_callers.py @@ -664,22 +664,25 @@ def _client( from fastapi import FastAPI from fastapi.testclient import TestClient - from tensorrt_llm.serve.openai_server import OpenAIServer + from tensorrt_llm.serve.api.anthropic.serving import AnthropicCountTokensServing - app = FastAPI() - server = object.__new__(OpenAIServer) - server.model = "m" - server.tokenizer = tokenizer - server.chat_template = chat_template + # What RenderResources.from_server reads off a server, as a namespace. + # Transitional: the resources get built directly once the server + # assembly stops reading them off the server per request. + ns = SimpleNamespace(model="m", tokenizer=tokenizer, chat_template=chat_template) if use_harmony: - server.use_harmony = True - server.model_config = SimpleNamespace(model_type="gpt_oss") + ns.use_harmony = True + ns.model_config = SimpleNamespace(model_type="gpt_oss") elif model_type is not None: - server.model_config = SimpleNamespace(model_type=model_type) + ns.model_config = SimpleNamespace(model_type=model_type) if input_processor is not None: - server.generator = SimpleNamespace(input_processor=input_processor) + ns.generator = SimpleNamespace(input_processor=input_processor) + serving = AnthropicCountTokensServing( + render_resources=lambda: RenderResources.from_server(ns) + ) + app = FastAPI() app.add_api_route( - "/v1/messages/count_tokens", server.anthropic_count_tokens, methods=["POST"] + "/v1/messages/count_tokens", serving.anthropic_count_tokens, methods=["POST"] ) return TestClient(app) diff --git a/tests/unittest/llmapi/apps/test_responses_usage.py b/tests/unittest/llmapi/apps/test_responses_usage.py index 7b36f70a5d00..789b53eaa421 100644 --- a/tests/unittest/llmapi/apps/test_responses_usage.py +++ b/tests/unittest/llmapi/apps/test_responses_usage.py @@ -30,6 +30,7 @@ from tensorrt_llm.executor.result import CompletionOutput from tensorrt_llm.inputs.tokenization import tokenize_prompt +from tensorrt_llm.serve.api.anthropic.api_router import anthropic_messages from tensorrt_llm.serve.openai_disagg_service import _ctx_usage_info from tensorrt_llm.serve.openai_protocol import ( PromptTokensDetails, @@ -380,7 +381,10 @@ def generate_async(*, inputs, disaggregated_params=None, **kwargs): server._extract_metrics = AsyncMock() app = FastAPI() app.add_api_route("/v1/chat/completions", server.openai_chat, methods=["POST"]) - app.add_api_route("/v1/messages", server.anthropic_messages, methods=["POST"]) + # The Messages route function reads its serving object from app.state; the + # server's constructor built it over the same chat core the chat route runs. + app.state.anthropic_serving_messages = server.app.state.anthropic_serving_messages + app.add_api_route("/v1/messages", anthropic_messages, methods=["POST"]) app.add_api_route("/v1/responses", server.openai_responses, methods=["POST"]) return TestClient(app), prompts diff --git a/tests/unittest/llmapi/test_kv_hints.py b/tests/unittest/llmapi/test_kv_hints.py index 218e7482d8a6..1e2d4d51b827 100644 --- a/tests/unittest/llmapi/test_kv_hints.py +++ b/tests/unittest/llmapi/test_kv_hints.py @@ -662,11 +662,15 @@ def test_http_handlers_forward_hints(request_kind: str, with_hints: bool) -> Non ) handler = server.chat_harmony if request_kind == "harmony" else server.openai_chat + # The request FastAPI would pass, reduced to what the handlers read before + # the engine boundary. The chat handler wraps it in a RequestContext, which + # needs an object to delegate to. + raw_request = SimpleNamespace(state=SimpleNamespace(), headers={}) with patch( "tensorrt_llm.serve.openai_server.responses_api_request_preprocess", new=AsyncMock(return_value=([1, 2], SamplingParams(max_tokens=4))), ): - response = asyncio.run(handler(request, None)) + response = asyncio.run(handler(request, raw_request)) server.generator.generate_async.assert_called_once() assert server.generator.generate_async.call_args.kwargs["kv_hints"] == hints assert b"captured at engine boundary" in response.body @@ -739,9 +743,9 @@ def test_http_ingress_policy_before_nested_validation( from fastapi import FastAPI from fastapi.testclient import TestClient + from tensorrt_llm.serve._base.msgspec_route import _MsgspecRoute from tensorrt_llm.serve.kv_hints import KvHintsRoute from tensorrt_llm.serve.openai_protocol import CompletionRequest, ResponsesRequest - from tensorrt_llm.serve.openai_server import _MsgspecRoute hints = _wire() if hint_kind == "unknown": @@ -832,9 +836,9 @@ def test_ignored_hints_do_not_relax_http_request_validation( from fastapi import FastAPI from fastapi.testclient import TestClient + from tensorrt_llm.serve._base.msgspec_route import _MsgspecRoute from tensorrt_llm.serve.kv_hints import KvHintsRoute from tensorrt_llm.serve.openai_protocol import CompletionRequest - from tensorrt_llm.serve.openai_server import _MsgspecRoute app = FastAPI() app.state.trust_kv_hints = True diff --git a/tests/unittest/llmapi/test_serve_multi_frontend_helpers.py b/tests/unittest/llmapi/test_serve_multi_frontend_helpers.py index 0a6f3b9055bb..6bb5d3e5fbc1 100644 --- a/tests/unittest/llmapi/test_serve_multi_frontend_helpers.py +++ b/tests/unittest/llmapi/test_serve_multi_frontend_helpers.py @@ -35,7 +35,14 @@ from tensorrt_llm._utils import AdjustedSteadyClock from tensorrt_llm.llmapi.disagg_utils import ServerRole from tensorrt_llm.serve import _frontend_processes, openai_server +from tensorrt_llm.serve._base.routing import ServerCapabilities, attach_routers from tensorrt_llm.serve._frontend_processes import _watch_attached_frontends +from tensorrt_llm.serve.api.anthropic.serving import ( + AnthropicBatchesServing, + AnthropicCountTokensServing, + AnthropicMessagesServing, +) +from tensorrt_llm.serve.api.ops.system.serving import SystemServing from tensorrt_llm.serve.multi_frontend import ( FORWARDED_ROUTES, AttachedFrontendWatchdog, @@ -111,8 +118,18 @@ async def lifespan(_: FastAPI): server._enable_runtime_control_endpoints = False server.resource_governor = None server._launcher_forwarder = forwarder - # The order OpenAIServer.__init__ registers them in. + # The order OpenAIServer.__init__ registers them in: the launcher + # forwarding routes first, then the routers (system, Anthropic), then the + # local table. The routes are enumerated and forwarded, never served + # locally, so the Anthropic serving objects get unused collaborators. + server.app.state.ops_serving_system = SystemServing(model="m", health_probe=lambda: True) + server.app.state.anthropic_serving_messages = AnthropicMessagesServing(backend=None) + server.app.state.anthropic_serving_count_tokens = AnthropicCountTokensServing( + render_resources=lambda: None + ) + server.app.state.anthropic_serving_batches = AnthropicBatchesServing(runner=None) server._register_launcher_forwarding_routes() + attach_routers(server.app, ServerCapabilities.for_role(None)) server.register_routes() return server diff --git a/tests/unittest/llmapi/test_serve_multi_frontend_route_gate.py b/tests/unittest/llmapi/test_serve_multi_frontend_route_gate.py index e477c020d314..b682e4f6b9e4 100644 --- a/tests/unittest/llmapi/test_serve_multi_frontend_route_gate.py +++ b/tests/unittest/llmapi/test_serve_multi_frontend_route_gate.py @@ -25,7 +25,15 @@ import pytest from fastapi import FastAPI from starlette.routing import Match - +from test_common.serve.routes import effective_routes + +from tensorrt_llm.serve._base.routing import ServerCapabilities, attach_routers +from tensorrt_llm.serve.api.anthropic.serving import ( + AnthropicBatchesServing, + AnthropicCountTokensServing, + AnthropicMessagesServing, +) +from tensorrt_llm.serve.api.ops.system.serving import SystemServing from tensorrt_llm.serve.multi_frontend import FORWARDED_ROUTES from tensorrt_llm.serve.openai_server import OpenAIServer @@ -73,10 +81,23 @@ def _llm_server_routes() -> list[tuple[str, str]]: server._enable_rl_control_endpoints = True server._enable_runtime_control_endpoints = True server.resource_governor = None + # As OpenAIServer.__init__ does: the routers (the system router for every + # role, the Anthropic router for the LLM role) are attached before the + # role's own table. The routes are enumerated, never called, so the + # serving objects get collaborators that are never used. + server.app.state.ops_serving_system = SystemServing(model="m", health_probe=lambda: True) + server.app.state.anthropic_serving_messages = AnthropicMessagesServing(backend=None) + server.app.state.anthropic_serving_count_tokens = AnthropicCountTokensServing( + render_resources=lambda: None + ) + server.app.state.anthropic_serving_batches = AnthropicBatchesServing(runner=None) + attach_routers(server.app, ServerCapabilities.for_role(None)) server.register_routes() + # effective_routes walks included routers, which FastAPI 0.139 keeps as + # objects without .path or .methods in app.routes. return sorted( - (method, route.path) - for route in server.app.routes + (method, path) + for path, route in effective_routes(server.app) if getattr(route, "methods", None) for method in route.methods if method != "HEAD" diff --git a/tests/unittest/serve/__init__.py b/tests/unittest/serve/__init__.py new file mode 100644 index 000000000000..467079831e16 --- /dev/null +++ b/tests/unittest/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. diff --git a/tests/unittest/serve/_base/__init__.py b/tests/unittest/serve/_base/__init__.py new file mode 100644 index 000000000000..467079831e16 --- /dev/null +++ b/tests/unittest/serve/_base/__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. diff --git a/tests/unittest/serve/_base/test_errors.py b/tests/unittest/serve/_base/test_errors.py new file mode 100644 index 000000000000..21280bbd21fc --- /dev/null +++ b/tests/unittest/serve/_base/test_errors.py @@ -0,0 +1,70 @@ +# 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. +"""The typed errors of the serving core (``tensorrt_llm/serve/_base/errors.py``). + +``error_body`` is the one place that builds an OpenAI error body, so a +``ServingError`` carried out of the chat core and an error response rendered +by ``OpenAIServer.create_error_response`` can never differ. +""" + +import json +from http import HTTPStatus + +import pytest + +from tensorrt_llm.executor.utils import context_length_exceeded_message +from tensorrt_llm.serve._base.errors import ServingError, error_body, serving_error +from tensorrt_llm.serve.openai_server import OpenAIServer + +pytestmark = pytest.mark.cpu_only + +CONTEXT_LENGTH_MESSAGE = context_length_exceeded_message(2048, 4096) + + +@pytest.mark.parametrize( + "message, err_type, status", + [ + ("bad", "BadRequestError", HTTPStatus.BAD_REQUEST), + ("gone", "NotFoundError", HTTPStatus.NOT_FOUND), + (CONTEXT_LENGTH_MESSAGE, "BadRequestError", HTTPStatus.BAD_REQUEST), + ], + ids=["bad-request", "not-found", "context-length"], +) +def test_error_body_matches_create_error_response(message, err_type, status): + rendered = OpenAIServer.create_error_response( + message=message, err_type=err_type, status_code=status + ) + body = error_body(message, err_type, status) + assert body.model_dump() == json.loads(rendered.body) + assert rendered.status_code == status.value + + error = serving_error(message, err_type, status) + assert error.status == status.value + assert error.body == body + assert str(error) == message + + if message is CONTEXT_LENGTH_MESSAGE: + # The code names the condition; the HTTP status stays a 400. + assert body.code == "context_length_exceeded" + assert error.status == 400 + else: + assert body.code == status.value + + +def test_serving_error_is_not_a_value_error(): + assert not issubclass(ServingError, ValueError), ( + "ServingError must not be a ValueError: openai_chat's core catches ValueError " + "and would re-wrap it" + ) diff --git a/tests/unittest/serve/_base/test_request_context.py b/tests/unittest/serve/_base/test_request_context.py new file mode 100644 index 000000000000..dfecf34ab828 --- /dev/null +++ b/tests/unittest/serve/_base/test_request_context.py @@ -0,0 +1,171 @@ +# 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. +"""``RequestContext``: what the chat core reads from a request, on the HTTP and the internal path.""" + +import asyncio +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +from starlette.requests import Request +from test_common.serve.app_factory import AggConfig + +from tensorrt_llm.serve._base.msgspec_route import _MsgspecRequest +from tensorrt_llm.serve._base.request_context import RequestContext +from tensorrt_llm.serve.openai_client import MSGPACK_HEADERS, _msgpack_encoder + +pytestmark = pytest.mark.cpu_only + +BODY = {"model": "m", "messages": [{"role": "user", "content": "hi"}]} + + +def _http_request(body: bytes = b"", headers=(), state=None) -> Request: + async def receive(): + return {"type": "http.request", "body": body, "more_body": False} + + scope = { + "type": "http", + "method": "POST", + "path": "/v1/messages", + "headers": [(k.lower().encode(), v.encode()) for k, v in headers], + "state": {"perf_metrics_records": []} if state is None else state, + } + return Request(scope, receive) + + +# -- from_http: delegation to the request FastAPI passed ---------------------- + + +def test_from_http_shares_the_request_state(): + request = _http_request() + ctx = RequestContext.from_http(request) + + assert ctx.http_request is request + assert ctx.state is request.state + assert ctx.perf_metrics_records is request.state.perf_metrics_records + # A write through the context is a write on the request's state, which the + # perf-metrics middleware reads when the response goes out. + ctx.server_first_token_time = 2.0 + assert request.state.server_first_token_time == 2.0 + assert ctx.server_first_token_time == 2.0 + assert bool(ctx), "the chat core tests `if raw_request and ...`; a context is truthy" + + +def test_from_http_reads_headers_and_client_from_the_request(): + request = _http_request(headers=(("X-Conversation-Id", "c1"),)) + ctx = RequestContext.from_http(request) + + assert ctx.headers["x-conversation-id"] == "c1" + assert ctx.client is None # the scope above names no client + assert ctx.server_arrival_time is None + assert ctx.no_client_connection is False + + +def test_from_http_reads_nothing_at_construction(): + # Tests hand the chat core a bare namespace as the request; wrapping it + # must not touch any attribute. + ctx = RequestContext.from_http(SimpleNamespace()) + assert ctx.http_request is not None + + +def test_from_http_keeps_the_msgpack_decoding_of_the_route_request(): + """A worker's request decodes msgpack bodies; the context must reuse it, not rebuild one.""" + body = _msgpack_encoder.encode(BODY) + request = _MsgspecRequest( + _http_request(body=body, headers=tuple(MSGPACK_HEADERS.items())).scope, + _http_request(body=body).receive, + ) + ctx = RequestContext.from_http(request) + + assert asyncio.run(ctx.json()) == BODY + + +def test_from_http_is_disconnected_asks_the_request(): + request = _http_request() + disconnected = {"value": False} + + async def is_disconnected(): + return disconnected["value"] + + request.is_disconnected = is_disconnected # type: ignore[method-assign] + ctx = RequestContext.from_http(request) + + assert asyncio.run(ctx.is_disconnected()) is False + disconnected["value"] = True + assert asyncio.run(ctx.is_disconnected()) is True + + +# -- internal: a request with no client behind it ------------------------------ + + +def test_internal_context_has_no_client_and_seeds_the_state(): + ctx = RequestContext.internal(BODY) + + assert ctx.http_request is None + assert dict(ctx.headers) == {} + assert ctx.client is None + assert ctx.no_client_connection is True + assert ctx.perf_metrics_records == [] + assert ctx.server_arrival_time is None + assert ctx.server_first_token_time is None + assert asyncio.run(ctx.json()) == BODY + assert asyncio.run(ctx.is_disconnected()) is False + assert bool(ctx) + + +def test_internal_json_is_what_a_wire_body_decodes_to_and_is_computed_once(): + # A body read from the wire has been through JSON: tuples become lists, + # non-string keys become strings. The internal context answers the same. + ctx = RequestContext.internal({"a": (1, 2), 3: "x"}) + + first = asyncio.run(ctx.json()) + assert first == {"a": [1, 2], "3": "x"} + assert asyncio.run(ctx.json()) is first + + +def test_internal_state_writes_stay_in_the_context(): + ctx = RequestContext.internal(BODY) + ctx.server_first_token_time = 5.0 + ctx.perf_metrics_records.append({"k": "v"}) + + assert ctx.state.server_first_token_time == 5.0 + assert ctx.perf_metrics_records == [{"k": "v"}] + + +def test_a_context_needs_a_request_or_an_internal_state(): + with pytest.raises(ValueError, match="from_http"): + RequestContext() + + +def test_await_disconnected_returns_immediately_for_an_internal_context(serve_env): + """A batched request has no socket, so the watchdog must not poll on it. + + is_disconnected() can never become true without a client, so without the + no_client_connection opt-out this coroutine would loop at 1Hz for the life + of the process, one live task per batched request. + """ + server = serve_env.build_openai_server(AggConfig()) + promise = SimpleNamespace(finished=False, request_id="req", abort=Mock()) + + async def run(): + # A socket-backed request would never let this return; a bounded wait + # keeps the failure a clear timeout rather than a hung suite. + await asyncio.wait_for( + server.await_disconnected(RequestContext.internal(BODY), promise), timeout=5.0 + ) + + asyncio.run(run()) + # Returning early must not be mistaken for a disconnect. + promise.abort.assert_not_called() diff --git a/tests/unittest/serve/_base/test_routing.py b/tests/unittest/serve/_base/test_routing.py new file mode 100644 index 000000000000..37fbd7fec2c8 --- /dev/null +++ b/tests/unittest/serve/_base/test_routing.py @@ -0,0 +1,264 @@ +# 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. +"""``make_router()``, ``attach_routers()`` and the role mapping of ``serve/_base/routing.py``.""" + +import ast +import importlib +import sys +import textwrap +from pathlib import Path + +import pytest +from fastapi import Depends, FastAPI, Request +from fastapi.routing import APIRoute +from fastapi.testclient import TestClient +from test_common.serve.app_factory import AggConfig +from test_common.serve.routes import effective_routes + +from tensorrt_llm.llmapi.disagg_utils import ServerRole +from tensorrt_llm.serve._base.msgspec_route import _MsgspecRoute +from tensorrt_llm.serve._base.routing import ( + ROUTERS, + RoleGroup, + RouterSpec, + ServerCapabilities, + attach_routers, + make_router, + role_group_of, +) +from tensorrt_llm.serve.openai_client import MSGPACK_HEADERS, _msgpack_encoder + +pytestmark = pytest.mark.cpu_only + + +def test_make_router_uses_msgspec_route(): + router = make_router() + assert router.route_class is _MsgspecRoute + + @router.get("/x") + async def x(): + return {} + + assert len(router.routes) == 1 + assert isinstance(router.routes[0], _MsgspecRoute) + + +def test_make_router_rejects_route_class(): + with pytest.raises(TypeError, match="_MsgspecRoute"): + make_router(route_class=APIRoute) + + +def test_role_group_of_every_role(): + expected = { + None: RoleGroup.LLM, + ServerRole.CONTEXT: RoleGroup.LLM, + ServerRole.GENERATION: RoleGroup.LLM, + ServerRole.MM_ENCODER: RoleGroup.MM_ENCODER, + ServerRole.EMBEDDING: RoleGroup.EMBEDDING, + ServerRole.VISUAL_GEN: RoleGroup.VISUAL_GEN, + } + # Every ServerRole member must have a row here: a new role must decide its + # role group (and its routes) explicitly. + assert set(expected) - {None} == set(ServerRole) + for role, group in expected.items(): + assert role_group_of(role) is group + caps = ServerCapabilities.for_role(role) + assert caps.role is role + assert caps.role_group is group + + +_ROUTER_MODULE_TEMPLATE = """ +from fastapi import FastAPI + +from tensorrt_llm.serve._base.routing import make_router + +router = make_router() + + +@router.get("{path}") +async def handler(): + return {{}} + + +def attach_router(app: FastAPI) -> None: + app.include_router(router) +""" + + +def _write_router_package(tmp_path, package: str, modules: dict) -> None: + pkg = tmp_path / package + pkg.mkdir() + (pkg / "__init__.py").write_text("") + for name, path in modules.items(): + (pkg / f"{name}.py").write_text(_ROUTER_MODULE_TEMPLATE.format(path=path)) + + +def test_attach_routers_imports_only_matching_specs(tmp_path, monkeypatch): + _write_router_package( + tmp_path, + "attachpkg", + {"llm_router": "/llm-only", "mm_router": "/mm-only", "gated_router": "/gated"}, + ) + monkeypatch.syspath_prepend(str(tmp_path)) + module_names = ("attachpkg.llm_router", "attachpkg.mm_router", "attachpkg.gated_router") + specs = ( + RouterSpec("attachpkg.llm_router", roles=frozenset({RoleGroup.LLM})), + RouterSpec("attachpkg.mm_router", roles=frozenset({RoleGroup.MM_ENCODER})), + RouterSpec( + "attachpkg.gated_router", + roles=frozenset({RoleGroup.LLM}), + when=lambda caps: False, + ), + ) + try: + app = FastAPI() + attach_routers(app, ServerCapabilities.for_role(None), routers=specs) + assert "attachpkg.llm_router" in sys.modules + assert "attachpkg.mm_router" not in sys.modules, "role mismatch must not be imported" + assert "attachpkg.gated_router" not in sys.modules, "a false `when` must not be imported" + paths = [path for path, _ in effective_routes(app)] + assert "/llm-only" in paths + assert "/mm-only" not in paths + assert "/gated" not in paths + finally: + for name in ("attachpkg", *module_names): + sys.modules.pop(name, None) + + +def test_router_routes_decode_msgpack(): + router = make_router() + + @router.post("/echo") + async def echo(raw_request: Request): + return {"body": await raw_request.json()} + + app = FastAPI() + app.include_router(router) + body = _msgpack_encoder.encode({"prompt": [1, 2, 3]}) + with TestClient(app) as client: + response = client.post("/echo", content=body, headers=MSGPACK_HEADERS) + assert response.status_code == 200 + assert response.json() == {"body": {"prompt": [1, 2, 3]}} + + +def _route_path(raw_request: Request) -> str: + return raw_request.scope["route"].path + + +def test_include_keeps_route_path(): + # A path prefix goes to make_router(prefix=...), never include_router(prefix=...): + # FastAPI 0.139 drops include prefixes from scope["route"].path. + router = make_router(prefix="/p") + + @router.get("/x") + async def x(route_path: str = Depends(_route_path)): + return {"route_path": route_path} + + app = FastAPI() + app.include_router(router) + with TestClient(app) as client: + response = client.get("/p/x") + assert response.status_code == 200 + assert response.json() == {"route_path": "/p/x"} + + +def test_every_router_spec_module_exposes_attach_router(): + # ROUTERS holds module paths as strings, so a renamed router module would + # break attachment only at server startup; import each one here instead. + assert ROUTERS, "ROUTERS must list at least the system router" + for spec in ROUTERS: + module = importlib.import_module(spec.module) + assert callable(getattr(module, "attach_router", None)), ( + f"{spec.module} must define attach_router(app)" + ) + tree = ast.parse(Path(module.__file__).read_text()) + for node in ast.walk(tree): + is_include_call = ( + isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and node.func.attr == "include_router" + ) + if is_include_call: + assert all(keyword.arg != "prefix" for keyword in node.keywords), ( + f"{spec.module} passes prefix= to include_router; give the prefix to " + "make_router(prefix=...) instead: FastAPI 0.139 drops include prefixes " + "from scope['route'].path." + ) + + +_ECHO_ROUTER_MODULE = textwrap.dedent( + ''' + """A new endpoint is a router plus a serving class: the router reads its + serving object from app.state; the composition root sets it and attaches.""" + + from fastapi import FastAPI, Request + + from tensorrt_llm.serve._base.routing import make_router + + router = make_router() + + + @router.post("/v1/test-echo") + async def echo(raw_request: Request): + serving = raw_request.app.state.test_serving_echo + return serving.render(await raw_request.json()) + + + def attach_router(app: FastAPI) -> None: + if not hasattr(app.state, "test_serving_echo"): + raise RuntimeError("app.state.test_serving_echo must be set before attaching") + app.include_router(router) + ''' +) + + +class _EchoServing: + def render(self, body): + return {"echo": body} + + +# The TestClient portal thread and the fake server's lazy worker-pool threads +# are alive when the leak check runs; the serve_env fixture closes them in its +# teardown (precedent: tests/unittest/serve/api/ops/system/test_system_endpoints.py). +@pytest.mark.threadleak(enabled=False) +def test_new_endpoint_lands_as_router_and_serving_class(tmp_path, monkeypatch, serve_env): + pkg = tmp_path / "newendpointpkg" + pkg.mkdir() + (pkg / "__init__.py").write_text("") + (pkg / "api_router.py").write_text(_ECHO_ROUTER_MODULE) + monkeypatch.syspath_prepend(str(tmp_path)) + spec = RouterSpec("newendpointpkg.api_router", roles=frozenset({RoleGroup.LLM})) + body = _msgpack_encoder.encode({"prompt": "hi"}) + try: + server = serve_env.build_openai_server(AggConfig()) + server.app.state.test_serving_echo = _EchoServing() + attach_routers(server.app, ServerCapabilities.for_role(None), routers=(spec,)) + client = serve_env.client(server.app) + response = client.post("/v1/test-echo", content=body, headers=MSGPACK_HEADERS) + assert response.status_code == 200 + assert response.json() == {"echo": {"prompt": "hi"}} + + # The same spec on a role outside its role set: not imported again, not + # attached, and the route is absent. + mm_server = serve_env.build_openai_server(AggConfig(role=ServerRole.MM_ENCODER)) + attach_routers( + mm_server.app, ServerCapabilities.for_role(ServerRole.MM_ENCODER), routers=(spec,) + ) + mm_client = serve_env.client(mm_server.app) + mm_response = mm_client.post("/v1/test-echo", content=body, headers=MSGPACK_HEADERS) + assert mm_response.status_code == 404 + finally: + for name in ("newendpointpkg", "newendpointpkg.api_router"): + sys.modules.pop(name, None) diff --git a/tests/unittest/serve/_base/test_sse.py b/tests/unittest/serve/_base/test_sse.py new file mode 100644 index 000000000000..517d93d383a4 --- /dev/null +++ b/tests/unittest/serve/_base/test_sse.py @@ -0,0 +1,147 @@ +# 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. +"""The chat SSE generator (``tensorrt_llm/serve/_base/sse.py``), driven with fakes. + +The server and the engine promise are local stand-ins: the generator reads +three things from the server (``postproc_worker_enabled``, the adjusted clock +and ``_extract_metrics``) and iterates the promise, and nothing else. +""" + +import asyncio +import logging +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest + +from tensorrt_llm.executor.postproc_worker import PostprocParams +from tensorrt_llm.logger import logger +from tensorrt_llm.serve._base.sse import chat_stream_generator + +pytestmark = pytest.mark.cpu_only + + +class _Promise: + """An async iterator of engine results, as a streaming ``RequestOutput`` is.""" + + def __init__(self, results, *, then=None): + self._results = list(results) + self._then = then + + def __aiter__(self): + return self + + async def __anext__(self): + if self._results: + return self._results.pop(0) + if self._then is not None: + raise self._then + raise StopAsyncIteration + + +def _result(frames): + """One engine result whose post-processing worker output is ``frames``.""" + return SimpleNamespace(outputs=[SimpleNamespace(_postprocess_result=frames)]) + + +def _server(*, postproc_worker_enabled): + return SimpleNamespace( + postproc_worker_enabled=postproc_worker_enabled, + _adjusted_steady_clock=SimpleNamespace(now=lambda: 7.0), + _extract_metrics=AsyncMock(), + ) + + +def _raw_request(): + return SimpleNamespace(state=SimpleNamespace()) + + +async def _collect(stream): + return [frame async for frame in stream] + + +@pytest.fixture +def error_records(): + """Messages of the ERROR records the TRT-LLM logger emits during the test. + + The library logger does not propagate, so ``caplog`` would see nothing; a + handler is attached to it instead. + """ + records = [] + + class Capture(logging.Handler): + def emit(self, record): + records.append(record.getMessage()) + + handler = Capture(level=logging.ERROR) + logger._logger.addHandler(handler) + try: + yield records + finally: + logger._logger.removeHandler(handler) + + +def test_frames_then_done_then_metrics(error_records): + results = [_result(["data: a\n\n"]), _result(["data: b\n\n"])] + server = _server(postproc_worker_enabled=True) + raw_request = _raw_request() + # With the post-processing worker on, the generator never reads the params. + params = PostprocParams(post_processor=None, postproc_args=None) + + frames = asyncio.run( + _collect(chat_stream_generator(_Promise(results), params, server=server, ctx=raw_request)) + ) + + assert frames == ["data: a\n\n", "data: b\n\n", "data: [DONE]\n\n"] + assert raw_request.state.server_first_token_time == 7.0 + server._extract_metrics.assert_awaited_once_with(results[-1], raw_request) + assert error_records == [] + + +def test_in_process_post_processor_is_used(): + results = [_result(None), _result(None)] + server = _server(postproc_worker_enabled=False) + raw_request = _raw_request() + seen = [] + + def post_processor(result, args): + seen.append((result, args)) + return [f"data: {len(seen)}\n\n"] + + params = PostprocParams(post_processor=post_processor, postproc_args="the args") + + frames = asyncio.run( + _collect(chat_stream_generator(_Promise(results), params, server=server, ctx=raw_request)) + ) + + assert frames == ["data: 1\n\n", "data: 2\n\n", "data: [DONE]\n\n"] + assert seen == [(results[0], "the args"), (results[1], "the args")] + assert raw_request.state.server_first_token_time == 7.0 + server._extract_metrics.assert_awaited_once_with(results[-1], raw_request) + + +def test_a_failure_is_logged_and_raised(error_records): + server = _server(postproc_worker_enabled=True) + promise = _Promise([_result(["data: a\n\n"])], then=RuntimeError("boom")) + params = PostprocParams(post_processor=None, postproc_args=None) + + with pytest.raises(RuntimeError, match="boom"): + asyncio.run( + _collect(chat_stream_generator(promise, params, server=server, ctx=_raw_request())) + ) + + assert len(error_records) == 1, error_records + assert "RuntimeError: boom" in error_records[0] + server._extract_metrics.assert_not_awaited() diff --git a/tests/unittest/serve/api/__init__.py b/tests/unittest/serve/api/__init__.py new file mode 100644 index 000000000000..467079831e16 --- /dev/null +++ b/tests/unittest/serve/api/__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. diff --git a/tests/unittest/serve/api/anthropic/__init__.py b/tests/unittest/serve/api/anthropic/__init__.py new file mode 100644 index 000000000000..467079831e16 --- /dev/null +++ b/tests/unittest/serve/api/anthropic/__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. diff --git a/tests/unittest/serve/api/anthropic/test_api_router.py b/tests/unittest/serve/api/anthropic/test_api_router.py new file mode 100644 index 000000000000..b7f654ce7e36 --- /dev/null +++ b/tests/unittest/serve/api/anthropic/test_api_router.py @@ -0,0 +1,140 @@ +# 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. +"""The Anthropic router: attachment contract, route provenance per role, msgpack bodies.""" + +from types import SimpleNamespace + +import pytest +from fastapi import FastAPI +from test_common.serve.app_factory import AggConfig +from test_common.serve.routes import effective_routes + +from tensorrt_llm.llmapi.disagg_utils import ServerRole +from tensorrt_llm.serve._base.msgspec_route import _MsgspecRoute +from tensorrt_llm.serve.api.anthropic.api_router import attach_router +from tensorrt_llm.serve.api.anthropic.serving import ( + AnthropicBatchesServing, + AnthropicCountTokensServing, + AnthropicMessagesServing, +) +from tensorrt_llm.serve.openai_client import MSGPACK_HEADERS, _msgpack_encoder +from tensorrt_llm.serve.render import RenderResources + +pytestmark = pytest.mark.cpu_only + +ANTHROPIC_ROUTER_MODULE = "tensorrt_llm.serve.api.anthropic.api_router" +ANTHROPIC_ROUTES = ( + ("POST", "/v1/messages"), + ("POST", "/v1/messages/count_tokens"), + ("POST", "/v1/messages/batches"), + ("GET", "/v1/messages/batches"), + ("GET", "/v1/messages/batches/{batch_id}"), + ("DELETE", "/v1/messages/batches/{batch_id}"), + ("POST", "/v1/messages/batches/{batch_id}/cancel"), + ("GET", "/v1/messages/batches/{batch_id}/results"), +) + + +def _anthropic_routes(app: FastAPI) -> dict: + """(method, path) -> [routes] for every Anthropic path the app registers.""" + found = {key: [] for key in ANTHROPIC_ROUTES} + for path, route in effective_routes(app): + for method in getattr(route, "methods", None) or (): + if (method, path) in found: + found[(method, path)].append(route) + return found + + +def test_attach_requires_every_serving_object(): + # The composition root sets the three serving objects before attaching, so + # a wiring bug fails at startup, not with an AttributeError per request. + with pytest.raises(RuntimeError, match="anthropic_serving_messages"): + attach_router(FastAPI()) + + +def test_serving_classes_require_their_collaborators(): + # No defaults: a serving class built without its collaborators fails at + # construction instead of falling back silently. + with pytest.raises(TypeError): + AnthropicMessagesServing() + with pytest.raises(TypeError): + AnthropicCountTokensServing() + with pytest.raises(TypeError): + AnthropicBatchesServing() + + +@pytest.mark.parametrize( + "role", [None, ServerRole.CONTEXT, ServerRole.GENERATION], ids=["llm", "context", "generation"] +) +def test_the_eight_routes_come_from_the_anthropic_router(role, serve_env): + server = serve_env.build_openai_server(AggConfig(role=role)) + + for key, routes in _anthropic_routes(server.app).items(): + assert len(routes) == 1, f"{key} must be registered exactly once, found {len(routes)}" + route = routes[0] + assert isinstance(route, _MsgspecRoute), ( + f"{key} must be a _MsgspecRoute; a plain APIRoute breaks msgpack request bodies " + "from the disaggregated frontend" + ) + assert route.endpoint.__module__ == ANTHROPIC_ROUTER_MODULE + + +def test_a_harmony_model_has_no_anthropic_routes(serve_env): + server = serve_env.build_openai_server(AggConfig(model_type="gpt_oss")) + assert all(not routes for routes in _anthropic_routes(server.app).values()) + + +@pytest.mark.parametrize("role", [ServerRole.MM_ENCODER, ServerRole.EMBEDDING]) +def test_other_roles_have_no_anthropic_routes(role, serve_env): + server = serve_env.build_openai_server(AggConfig(role=role)) + assert all(not routes for routes in _anthropic_routes(server.app).values()) + + +class _WordTokenizer: + """Renders the conversation as the words of its messages and counts them.""" + + def get_chat_template(self, chat_template=None, tools=None): + return chat_template or "WORDS" + + def apply_chat_template(self, messages=None, **kwargs): + if messages is None: + messages = kwargs.pop("conversation") + return " ".join(str(m.get("content", "")) for m in messages) + + def encode(self, text, add_special_tokens=True, **kwargs): + return text.split() + + +# The TestClient portal thread and the fake server's lazy worker-pool threads +# are alive when the leak check runs; the serve_env fixture closes them in its +# teardown (precedent: tests/unittest/serve/api/ops/system/test_system_endpoints.py). +@pytest.mark.threadleak(enabled=False) +def test_count_tokens_decodes_the_disaggregated_frontends_msgpack_body(serve_env): + """The disaggregated frontend forwards count_tokens as msgpack to a context worker.""" + server = serve_env.build_openai_server(AggConfig()) + resources = SimpleNamespace(tokenizer=_WordTokenizer(), chat_template=None) + server.app.state.anthropic_serving_count_tokens = AnthropicCountTokensServing( + render_resources=lambda: RenderResources.from_server(resources) + ) + body = _msgpack_encoder.encode( + {"model": "m", "messages": [{"role": "user", "content": "one two three"}]} + ) + + response = serve_env.client(server.app).post( + "/v1/messages/count_tokens", content=body, headers=MSGPACK_HEADERS + ) + + assert response.status_code == 200, response.text + assert response.json() == {"input_tokens": 3} diff --git a/tests/unittest/serve/api/anthropic/test_messages_behavior.py b/tests/unittest/serve/api/anthropic/test_messages_behavior.py new file mode 100644 index 000000000000..ce384da56005 --- /dev/null +++ b/tests/unittest/serve/api/anthropic/test_messages_behavior.py @@ -0,0 +1,630 @@ +# 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. +"""What the aggregated and the disaggregated server answer on the Anthropic routes, pinned. + +Both servers serve ``POST /v1/messages`` by translating the request into a chat +completion, running it through their own chat pipeline and translating the +answer back; the aggregated server also runs Message Batches over the same +path. The two pipelines answer some cases differently today (the model name of +a stream without chunks, the message of an unexpected error) and are meant to +keep doing so. These tests pin each server's own answer, case by case, so that +a change to how the Anthropic routes reach the chat pipelines can show it kept +every one of them. + +Only HTTP and the engine seams are used. The aggregated chat path is stubbed +where it meets the engine (``generate_async`` and the response builder +``_create_chat_response``); the disaggregated one where it meets the service +(``openai_chat_completion``). The served model is ``server-model`` and every +request names ``client-model``, so an answer that echoes the wrong one is +caught. +""" + +import json +import logging +import signal +import time +from types import SimpleNamespace +from unittest.mock import AsyncMock, Mock + +import aiohttp +import pytest +from test_common.serve.app_factory import AggConfig, DisaggConfig + +import tensorrt_llm.serve.openai_disagg_server as disagg_server_module +import tensorrt_llm.serve.openai_server as openai_server_module +from tensorrt_llm.executor import CppExecutorError +from tensorrt_llm.executor.utils import context_length_exceeded_message +from tensorrt_llm.inputs.chat_template_guard import UnusedChatTemplateKwargsError +from tensorrt_llm.logger import logger as trtllm_logger +from tensorrt_llm.serve.openai_protocol import ( + ChatCompletionResponse, + ChatCompletionResponseChoice, + ChatCompletionStreamResponse, + ChatMessage, + FunctionCall, + ToolCall, + UsageInfo, +) +from tensorrt_llm.serve.openai_server import OpenAIServer + +# threadleak is off: the TestClient's portal thread and the fake server's +# worker pools are still alive when the leak check runs; the serve_env fixture +# closes them in its teardown, which runs after that check. +pytestmark = [pytest.mark.cpu_only, pytest.mark.threadleak(enabled=False)] + +SIDES = ["aggregated", "disaggregated"] +SERVER_MODEL = "server-model" +CLIENT_MODEL = "client-model" +# Two chunks of streamed text. The accent gives the disaggregated byte stream a +# multi-byte character to be cut inside. +STREAMED_CHUNKS = ("hél", "lo") +STREAMED_TEXT = "".join(STREAMED_CHUNKS) +COUNTER_NAMES = ( + "total_requests", + "stream_requests", + "nonstream_requests", + "validation_exceptions", + "http_exceptions", + "internal_errors", + "total_responses", +) + + +def _changed(case: str, side: str, expected, got) -> str: + return ( + f"Anthropic {case} on the {side} server changed: expected {expected!r}, got {got!r}. " + "Each server must keep its own answer; the aggregated and disaggregated servers " + "differ on purpose." + ) + + +# -- requests and engine answers ----------------------------------------------- + + +def _request(text: str = "hello", **overrides) -> dict: + payload = { + "model": CLIENT_MODEL, + "max_tokens": 32, + "messages": [{"role": "user", "content": text}], + } + payload.update(overrides) + return payload + + +def _chat_response(*, tool_arguments=None) -> ChatCompletionResponse: + """The chat pipeline's non-streaming answer: the text "hello", optionally with a tool call.""" + tool_calls = [] + if tool_arguments is not None: + tool_calls = [ + ToolCall( + id="call_1", + function=FunctionCall(name="get_weather", arguments=tool_arguments), + ) + ] + return ChatCompletionResponse( + id="chatcmpl-messages-behavior", + model=SERVER_MODEL, + choices=[ + ChatCompletionResponseChoice( + index=0, + message=ChatMessage(role="assistant", content="hello", tool_calls=tool_calls), + finish_reason="tool_calls" if tool_calls else "stop", + ) + ], + usage=UsageInfo(prompt_tokens=3, completion_tokens=2, total_tokens=5), + ) + + +def _stream_frames(*texts: str) -> list: + """The chat pipeline's SSE frames for a stream of text chunks, without the final [DONE].""" + chunks = [ + ChatCompletionStreamResponse( + model=SERVER_MODEL, choices=[{"index": 0, "delta": {"content": text}}] + ) + for text in texts + ] + chunks.append( + ChatCompletionStreamResponse( + model=SERVER_MODEL, + choices=[{"index": 0, "delta": {}, "finish_reason": "stop"}], + usage=UsageInfo(prompt_tokens=3, completion_tokens=2, total_tokens=5), + ) + ) + return [f"data: {chunk.model_dump_json()}\n\n" for chunk in chunks] + + +def _promise() -> SimpleNamespace: + """The engine's handle for a non-streaming request; the response builder is stubbed.""" + return SimpleNamespace(prompt_token_ids=[1, 2, 3]) + + +class _StreamedGeneration: + """A streamed generation whose steps are already post-processed SSE frames.""" + + prompt_token_ids = [1, 2, 3] + + def __init__(self, steps): + self._steps = list(steps) + + def __aiter__(self): + return self + + async def __anext__(self): + if not self._steps: + raise StopAsyncIteration + frames = self._steps.pop(0) + return SimpleNamespace( + finished=not self._steps, + outputs=[SimpleNamespace(_postprocess_result=frames)], + ) + + +def _relayed_bytes(frames: list): + """The worker's SSE bytes as the disaggregated service relays them. + + Delivered in two reads, cut inside the multi-byte character of the streamed + text, as network reads need not align with characters or lines. + """ + payload = "".join(frames).encode("utf-8") + cut = payload.index("é".encode("utf-8")) + 1 + + async def relay(): + yield payload[:cut] + yield payload[cut:] + + return relay() + + +def _done_only(): + """A relayed stream finished by the context worker alone: only the [DONE] sentinel.""" + + async def relay(): + yield b"data: [DONE]\n\n" + + return relay() + + +def _upstream_error(worker_body: str, status: int) -> aiohttp.ClientResponseError: + """A worker rejection the way the disaggregated client raises it: reason, then the body.""" + return aiohttp.ClientResponseError( + request_info=Mock(), history=(), status=status, message=f"Bad Request: {worker_body}" + ) + + +# -- servers ------------------------------------------------------------------- + + +def _aggregated_server(env, *, generate_async, chat_response=None) -> OpenAIServer: + """The aggregated server with its chat path stubbed where it meets the engine. + + Post-processing runs in worker mode, so a streamed generation's steps are + the SSE frames themselves. The engine's preprocessing and the disconnect + watcher are not part of any answer and are left out. + """ + server = env.build_openai_server(AggConfig(model_name=SERVER_MODEL, num_postprocess_workers=1)) + server.generator.preprocess = None + server.generator.generate_async = generate_async + server.await_disconnected = AsyncMock() + if chat_response is not None: + server._create_chat_response = AsyncMock(return_value=chat_response) + return server + + +def _disaggregated_server(env, openai_chat_completion): + """The disaggregated server with the service stubbed and its request counters mocked.""" + server = env.build_disagg_server(DisaggConfig()) + counters = SimpleNamespace( + **{name: SimpleNamespace(inc=Mock()) for name in COUNTER_NAMES}, + queue_latency_seconds=SimpleNamespace(observe=Mock()), + ) + server._perf_metrics_collector = counters + server._service = SimpleNamespace(openai_chat_completion=openai_chat_completion) + return server, counters + + +def _post_messages(env, side: str, server, payload: dict): + # The disaggregated lifespan connects to workers; these tests never run it. + client = env.client(server.app, lifespan=side == "aggregated") + return client.post("/v1/messages", json=payload) + + +def _assert_error(response, *, case: str, side: str, status: int, error_type: str, message: str): + body = response.json() + got = (response.status_code, body.get("type"), body.get("error")) + expected = (status, "error", {"type": error_type, "message": message}) + assert got == expected, _changed(case, side, expected, got) + + +def _assert_counts(counters, case: str, **expected: int) -> None: + for name in COUNTER_NAMES: + got = getattr(counters, name).inc.call_count + assert got == expected.get(name, 0), _changed( + f"{case} ({name} counter)", "disaggregated", expected.get(name, 0), got + ) + + +def _events(body: str) -> list: + """The (event name, decoded data) pairs of an Anthropic SSE body.""" + events = [] + for frame in body.strip().split("\n\n"): + name, data = None, None + for line in frame.splitlines(): + if line.startswith("event: "): + name = line[len("event: ") :] + elif line.startswith("data: "): + data = json.loads(line[len("data: ") :]) + events.append((name, data)) + return events + + +def _streamed_text(events: list) -> str: + return "".join( + data["delta"]["text"] + for name, data in events + if name == "content_block_delta" and data["delta"]["type"] == "text_delta" + ) + + +@pytest.fixture +def error_records(): + """The ERROR records the TensorRT-LLM logger emits during the test.""" + records = [] + + class Capture(logging.Handler): + def emit(self, record): + records.append(record) + + handler = Capture(level=logging.ERROR) + trtllm_logger._logger.addHandler(handler) + try: + yield records + finally: + trtllm_logger._logger.removeHandler(handler) + + +@pytest.fixture +def engine_failure(monkeypatch): + """What an engine failure does to the process, recorded instead of done. + + Both servers answer an engine failure by recording the termination and + raising SIGINT; the recorder is patched in the module whose code raises it. + """ + recorded = SimpleNamespace(signals=[], aggregated=Mock(), disaggregated=Mock()) + monkeypatch.setattr(signal, "raise_signal", recorded.signals.append) + monkeypatch.setattr(openai_server_module, "_record_generator_termination", recorded.aggregated) + monkeypatch.setattr( + disagg_server_module, "record_termination_observation", recorded.disaggregated + ) + return recorded + + +# -- /v1/messages -------------------------------------------------------------- + + +@pytest.mark.parametrize("side", SIDES) +def test_non_streaming_success_is_the_translated_chat_response(side, serve_env): + if side == "aggregated": + server = _aggregated_server( + serve_env, generate_async=Mock(return_value=_promise()), chat_response=_chat_response() + ) + counters = None + else: + server, counters = _disaggregated_server( + serve_env, AsyncMock(return_value=_chat_response()) + ) + + response = _post_messages(serve_env, side, server, _request()) + + body = response.json() + got = (response.status_code, body.get("content"), body.get("model")) + expected = (200, [{"type": "text", "text": "hello"}], SERVER_MODEL) + assert got == expected, _changed("non-streaming success", side, expected, got) + if counters is not None: + _assert_counts( + counters, + "non-streaming success", + total_requests=1, + nonstream_requests=1, + total_responses=1, + ) + + +@pytest.mark.parametrize("side", SIDES) +def test_streaming_success_reports_the_chunks_model_and_text(side, serve_env): + frames = _stream_frames(*STREAMED_CHUNKS) + if side == "aggregated": + generation = _StreamedGeneration([[frame] for frame in frames]) + server = _aggregated_server(serve_env, generate_async=Mock(return_value=generation)) + else: + relayed = _relayed_bytes(frames + ["data: [DONE]\n\n"]) + server, _ = _disaggregated_server(serve_env, AsyncMock(return_value=relayed)) + + response = _post_messages(serve_env, side, server, _request(stream=True)) + + assert response.status_code == 200, _changed( + "streaming success", side, 200, (response.status_code, response.text) + ) + assert response.headers["content-type"].startswith("text/event-stream") + events = _events(response.text) + got = (events[0][0], events[0][1]["message"]["model"], _streamed_text(events), events[-1][0]) + expected = ("message_start", SERVER_MODEL, STREAMED_TEXT, "message_stop") + assert got == expected, _changed("streaming success", side, expected, got) + + +@pytest.mark.parametrize( + "side, expected_model", + [("aggregated", SERVER_MODEL), ("disaggregated", CLIENT_MODEL)], +) +def test_stream_without_chunks_names_each_servers_own_model(side, expected_model, serve_env): + """With no chunk to take the model from, each server falls back to its own choice.""" + if side == "aggregated": + generation = _StreamedGeneration([[]]) + server = _aggregated_server(serve_env, generate_async=Mock(return_value=generation)) + else: + server, _ = _disaggregated_server(serve_env, AsyncMock(return_value=_done_only())) + + response = _post_messages(serve_env, side, server, _request(stream=True)) + + assert response.status_code == 200, _changed( + "stream without chunks", side, 200, (response.status_code, response.text) + ) + events = _events(response.text) + got = (events[0][0], events[0][1]["message"]["model"], events[-1][0]) + expected = ("message_start", expected_model, "message_stop") + assert got == expected, _changed("stream without chunks", side, expected, got) + + +@pytest.mark.parametrize("side", SIDES) +def test_rejected_request_is_a_400_with_the_rejection(side, serve_env): + if side == "aggregated": + server = _aggregated_server(serve_env, generate_async=Mock(side_effect=ValueError("bad"))) + else: + server, _ = _disaggregated_server( + serve_env, AsyncMock(side_effect=UnusedChatTemplateKwargsError("bad")) + ) + + response = _post_messages(serve_env, side, server, _request()) + + _assert_error( + response, + case="rejected request", + side=side, + status=400, + error_type="invalid_request_error", + message="bad", + ) + + +@pytest.mark.parametrize("side", SIDES) +def test_context_length_error_keeps_its_message(side, serve_env): + message = context_length_exceeded_message(2048, 4096) + if side == "aggregated": + server = _aggregated_server(serve_env, generate_async=Mock(side_effect=ValueError(message))) + else: + worker_body = OpenAIServer.create_error_response(message).body.decode() + server, _ = _disaggregated_server( + serve_env, AsyncMock(side_effect=_upstream_error(worker_body, status=400)) + ) + + response = _post_messages(serve_env, side, server, _request()) + + _assert_error( + response, + case="context length", + side=side, + status=400, + error_type="invalid_request_error", + message=message, + ) + + +def test_upstream_error_keeps_the_worker_status_and_message(serve_env): + error = aiohttp.ClientResponseError( + request_info=Mock(), + history=(), + status=429, + message='Too Many Requests: {"error": {"message": "slow down"}}', + ) + server, counters = _disaggregated_server(serve_env, AsyncMock(side_effect=error)) + + response = _post_messages(serve_env, "disaggregated", server, _request()) + + _assert_error( + response, + case="upstream error", + side="disaggregated", + status=429, + error_type="invalid_request_error", + message="slow down", + ) + _assert_counts( + counters, "upstream error", total_requests=1, nonstream_requests=1, http_exceptions=1 + ) + + +@pytest.mark.parametrize( + "side, status, error_type, message", + [ + ("aggregated", 400, "invalid_request_error", "boom"), + ("disaggregated", 500, "api_error", "Internal server error boom"), + ], +) +def test_unexpected_error_is_answered_each_servers_way( + side, status, error_type, message, serve_env +): + if side == "aggregated": + server = _aggregated_server( + serve_env, generate_async=Mock(side_effect=RuntimeError("boom")) + ) + counters = None + else: + server, counters = _disaggregated_server( + serve_env, AsyncMock(side_effect=RuntimeError("boom")) + ) + + response = _post_messages(serve_env, side, server, _request()) + + _assert_error( + response, + case="unexpected error", + side=side, + status=status, + error_type=error_type, + message=message, + ) + if counters is not None: + _assert_counts( + counters, "unexpected error", total_requests=1, nonstream_requests=1, internal_errors=1 + ) + + +@pytest.mark.parametrize("side", SIDES) +def test_engine_failure_is_a_500_after_one_sigint(side, serve_env, engine_failure): + if side == "aggregated": + server = _aggregated_server( + serve_env, generate_async=Mock(side_effect=CppExecutorError("dead")) + ) + recorder = engine_failure.aggregated + else: + server, _ = _disaggregated_server( + serve_env, AsyncMock(side_effect=CppExecutorError("dead")) + ) + recorder = engine_failure.disaggregated + + response = _post_messages(serve_env, side, server, _request()) + + _assert_error( + response, + case="engine failure", + side=side, + status=500, + error_type="api_error", + message="Internal server error", + ) + got = (engine_failure.signals, recorder.call_count) + expected = ([signal.SIGINT], 1) + assert got == expected, _changed( + "engine failure (SIGINT, termination record)", side, expected, got + ) + + +@pytest.mark.parametrize("side", SIDES) +def test_unconvertible_chat_response_is_a_500_with_one_error_log(side, serve_env, error_records): + response_with_bad_tool_call = _chat_response(tool_arguments="{not-json") + if side == "aggregated": + server = _aggregated_server( + serve_env, + generate_async=Mock(return_value=_promise()), + chat_response=response_with_bad_tool_call, + ) + else: + server, _ = _disaggregated_server( + serve_env, AsyncMock(return_value=response_with_bad_tool_call) + ) + + response = _post_messages(serve_env, side, server, _request()) + + _assert_error( + response, + case="unconvertible response", + side=side, + status=500, + error_type="api_error", + message="Internal server error", + ) + messages = [record.getMessage() for record in error_records] + assert len(messages) == 1 and "Invalid response from OpenAI chat pipeline" in messages[0], ( + _changed( + "unconvertible response (error log)", + side, + "one record naming the invalid response", + messages, + ) + ) + + +# -- /v1/messages/batches (aggregated only) -------------------------------------- + + +def test_batch_items_report_each_outcome(serve_env, engine_failure): + def generate_async(inputs, **kwargs): + prompt = inputs.get("prompt") or "" + if "reject me" in prompt: + raise ValueError("bad") + if "kill the engine" in prompt: + raise CppExecutorError("dead") + return _promise() + + server = _aggregated_server( + serve_env, generate_async=generate_async, chat_response=_chat_response() + ) + client = serve_env.client(server.app) + items = [ + {"custom_id": custom_id, "params": _request(text)} + for custom_id, text in ( + ("ok", "hello"), + ("rejected", "reject me"), + ("engine", "kill the engine"), + ) + ] + + created = client.post("/v1/messages/batches", json={"requests": items}) + assert created.status_code == 200, created.text + batch_id = created.json()["id"] + deadline = time.monotonic() + 30 + while True: + status = client.get(f"/v1/messages/batches/{batch_id}").json()["processing_status"] + if status == "ended": + break + assert time.monotonic() < deadline, f"batch {batch_id} is still {status} after 30 s" + time.sleep(0.01) + results = client.get(f"/v1/messages/batches/{batch_id}/results") + + assert results.status_code == 200, results.text + by_id = { + line["custom_id"]: line["result"] for line in map(json.loads, results.text.splitlines()) + } + got = { + "ok": ( + by_id["ok"]["type"], + by_id["ok"]["message"]["content"], + by_id["ok"]["message"]["model"], + ), + "rejected": by_id["rejected"], + "engine": by_id["engine"], + } + expected = { + "ok": ("succeeded", [{"type": "text", "text": "hello"}], SERVER_MODEL), + "rejected": { + "type": "errored", + "error": { + "type": "error", + "error": {"type": "invalid_request_error", "message": "bad"}, + "request_id": None, + }, + }, + "engine": { + "type": "errored", + "error": { + "type": "error", + "error": {"type": "api_error", "message": "Internal server error"}, + "request_id": None, + }, + }, + } + assert got == expected, _changed("batch item outcomes", "aggregated", expected, got) + assert engine_failure.signals == [signal.SIGINT], _changed( + "batch engine failure (SIGINT)", "aggregated", [signal.SIGINT], engine_failure.signals + ) diff --git a/tests/unittest/serve/api/anthropic/test_serving.py b/tests/unittest/serve/api/anthropic/test_serving.py new file mode 100644 index 000000000000..bb50d5772344 --- /dev/null +++ b/tests/unittest/serve/api/anthropic/test_serving.py @@ -0,0 +1,355 @@ +# 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. +"""``AnthropicMessagesServing`` over a fake ``ChatBackend``: what reaches the answer and how.""" + +import asyncio +import json +import logging + +import pytest +from fastapi.responses import JSONResponse, StreamingResponse + +from tensorrt_llm.logger import logger as trtllm_logger +from tensorrt_llm.serve._base.errors import EngineFatalError, ServingError +from tensorrt_llm.serve._base.request_context import RequestContext +from tensorrt_llm.serve.api.anthropic import serving as serving_module +from tensorrt_llm.serve.api.anthropic.protocol import AnthropicMessagesRequest +from tensorrt_llm.serve.api.anthropic.serving import AnthropicMessagesServing, _error_type +from tensorrt_llm.serve.openai_protocol import ( + ChatCompletionResponse, + ChatCompletionResponseChoice, + ChatCompletionStreamResponse, + ChatMessage, + ErrorResponse, + FunctionCall, + ToolCall, + UsageInfo, +) + +pytestmark = pytest.mark.cpu_only + +CLIENT_MODEL = "client-model" +RESPONSE_MODEL = "response-model" + + +class _Backend: + """A ChatBackend that answers with a fixed result or raises a fixed error.""" + + def __init__(self, *, result=None, error=None, model="x-model", message="x-message"): + self.result = result + self.error = error + self.model = model + self.message = message + self.calls = [] + + async def create_chat_completion(self, request, ctx): + self.calls.append((request, ctx)) + if self.error is not None: + raise self.error + return self.result + + def stream_model(self, request): + return self.model + + def error_message(self, error): + return self.message + + +def _request(**overrides) -> AnthropicMessagesRequest: + payload = { + "model": CLIENT_MODEL, + "max_tokens": 32, + "messages": [{"role": "user", "content": "hello"}], + } + payload.update(overrides) + return AnthropicMessagesRequest(**payload) + + +def _chat_response(*, tool_arguments=None) -> ChatCompletionResponse: + tool_calls = [] + if tool_arguments is not None: + tool_calls = [ + ToolCall( + id="call_1", function=FunctionCall(name="get_weather", arguments=tool_arguments) + ) + ] + return ChatCompletionResponse( + id="chatcmpl-serving-test", + model=RESPONSE_MODEL, + choices=[ + ChatCompletionResponseChoice( + index=0, + message=ChatMessage(role="assistant", content="hello", tool_calls=tool_calls), + finish_reason="tool_calls" if tool_calls else "stop", + ) + ], + usage=UsageInfo(prompt_tokens=3, completion_tokens=2, total_tokens=5), + ) + + +def _frames(*texts: str) -> list: + chunks = [ + ChatCompletionStreamResponse( + model=RESPONSE_MODEL, choices=[{"index": 0, "delta": {"content": text}}] + ) + for text in texts + ] + return [f"data: {chunk.model_dump_json()}\n\n" for chunk in chunks] + ["data: [DONE]\n\n"] + + +async def _stream(frames: list): + for frame in frames: + yield frame + + +def _serving_error(status: int) -> ServingError: + return ServingError(status, ErrorResponse(message="body message", type="t", code=status)) + + +def _create(backend: _Backend, request: AnthropicMessagesRequest): + serving = AnthropicMessagesServing(backend) + return asyncio.run(serving.create_message(request, RequestContext.internal({}))) + + +def _body(response: JSONResponse) -> dict: + return json.loads(bytes(response.body)) + + +def _events(response: StreamingResponse) -> list: + async def collect(): + return [frame async for frame in response.body_iterator] + + events = [] + for frame in "".join(asyncio.run(collect())).strip().split("\n\n"): + name, data = None, None + for line in frame.splitlines(): + if line.startswith("event: "): + name = line[len("event: ") :] + elif line.startswith("data: "): + data = json.loads(line[len("data: ") :]) + events.append((name, data)) + return events + + +@pytest.fixture +def error_records(): + """Messages of the ERROR records the TensorRT-LLM logger emits (it does not propagate to caplog).""" + records = [] + + class Capture(logging.Handler): + def emit(self, record): + records.append(record.getMessage()) + + handler = Capture(level=logging.ERROR) + trtllm_logger._logger.addHandler(handler) + try: + yield records + finally: + trtllm_logger._logger.removeHandler(handler) + + +# -- create_message -------------------------------------------------------------- + + +def test_non_streaming_result_is_translated_without_a_round_trip(): + backend = _Backend(result=_chat_response()) + + response = _create(backend, _request()) + + assert isinstance(response, JSONResponse) and response.status_code == 200 + body = _body(response) + assert body["content"] == [{"type": "text", "text": "hello"}] + assert body["model"] == RESPONSE_MODEL + assert body["usage"] == {"input_tokens": 3, "output_tokens": 2} + (chat_request, ctx), *_ = backend.calls + assert chat_request.model == CLIENT_MODEL and chat_request.stream is False + assert isinstance(ctx, RequestContext) + + +def test_streaming_result_is_reframed_as_event_stream_from_the_backends_iterator(monkeypatch): + stream = _stream(_frames("he", "llo")) + backend = _Backend(result=stream) + seen = {} + + async def fake_reframe(openai_sse, model): + seen["iterator"] = openai_sse + seen["model"] = model + yield "event: ping\ndata: {}\n\n" + + monkeypatch.setattr(serving_module, "reframe_openai_stream", fake_reframe) + + response = _create(backend, _request(stream=True)) + + assert isinstance(response, StreamingResponse) + assert response.media_type == "text/event-stream" + assert _events(response) == [("ping", {})] + # The backend's stream object goes to the reframer as it is: no copy, no + # wrapper, so the adapter reads the worker's bytes or the core's frames. + assert seen["iterator"] is stream + assert seen["model"] == "x-model" + + +def test_stream_without_chunks_names_the_backends_model(): + backend = _Backend(result=_stream(["data: [DONE]\n\n"]), model="x-model") + + response = _create(backend, _request(stream=True)) + + events = _events(response) + assert events[0][0] == "message_start" + assert events[0][1]["message"]["model"] == "x-model" + assert events[-1][0] == "message_stop" + + +def test_stream_with_chunks_adopts_the_chunks_model(): + backend = _Backend(result=_stream(_frames("he", "llo")), model="x-model") + + response = _create(backend, _request(stream=True)) + + events = _events(response) + assert events[0][1]["message"]["model"] == RESPONSE_MODEL + text = "".join( + data["delta"]["text"] + for name, data in events + if name == "content_block_delta" and data["delta"]["type"] == "text_delta" + ) + assert text == "hello" + + +@pytest.mark.parametrize( + "status, expected_type", + [(400, "invalid_request_error"), (429, "invalid_request_error"), (503, "api_error")], +) +def test_serving_error_is_rendered_with_the_backends_message_and_the_status(status, expected_type): + backend = _Backend(error=_serving_error(status), message="x-message") + + response = _create(backend, _request()) + + assert response.status_code == status + assert _body(response) == { + "type": "error", + "error": {"type": expected_type, "message": "x-message"}, + } + + +def test_engine_failure_is_a_500_api_error(): + backend = _Backend(error=EngineFatalError()) + + response = _create(backend, _request()) + + assert response.status_code == 500 + assert _body(response)["error"] == {"type": "api_error", "message": "Internal server error"} + + +def test_request_conversion_failure_is_a_400_before_the_backend_runs(): + backend = _Backend(result=_chat_response()) + + response = _create( + backend, _request(tools=[{"name": "web_search", "type": "web_search_20250305"}]) + ) + + assert response.status_code == 400 + assert _body(response)["error"]["type"] == "invalid_request_error" + assert "server tool" in _body(response)["error"]["message"] + assert backend.calls == [] + + +def test_unconvertible_chat_response_is_a_500_with_one_error_log(error_records): + backend = _Backend(result=_chat_response(tool_arguments="{not-json")) + + response = _create(backend, _request()) + + assert response.status_code == 500 + assert _body(response)["error"] == {"type": "api_error", "message": "Internal server error"} + assert len(error_records) == 1 + assert "Invalid response from OpenAI chat pipeline" in error_records[0] + + +# -- run_batched_message ---------------------------------------------------------- + + +def _run_batched(backend: _Backend, request: AnthropicMessagesRequest): + serving = AnthropicMessagesServing(backend) + return asyncio.run(serving.run_batched_message(request)) + + +def test_batched_message_runs_an_internal_context_without_streaming(): + backend = _Backend(result=_chat_response()) + + outcome, payload = _run_batched(backend, _request(stream=True)) + + assert outcome == "succeeded" + assert payload["content"] == [{"type": "text", "text": "hello"}] + (chat_request, ctx), *_ = backend.calls + # Streaming makes no sense for a batch: the result is collected, not watched. + assert chat_request.stream is False + assert isinstance(ctx, RequestContext) + assert ctx.http_request is None + assert dict(ctx.headers) == {} + assert ctx.no_client_connection is True + assert asyncio.run(ctx.json()) == chat_request.model_dump(exclude_none=True) + + +@pytest.mark.parametrize( + "error, expected", + [ + (_serving_error(400), {"type": "invalid_request_error", "message": "x-message"}), + (_serving_error(503), {"type": "api_error", "message": "x-message"}), + (EngineFatalError(), {"type": "api_error", "message": "Internal server error"}), + ], + ids=["rejected", "upstream_5xx", "engine_failure"], +) +def test_batched_message_reports_errors_as_errored_items(error, expected): + backend = _Backend(error=error, message="x-message") + + assert _run_batched(backend, _request()) == ("errored", expected) + + +def test_batched_message_conversion_failure_does_not_reach_the_backend(): + backend = _Backend(result=_chat_response()) + + outcome, payload = _run_batched( + backend, _request(tools=[{"name": "web_search", "type": "web_search_20250305"}]) + ) + + assert outcome == "errored" + assert payload["type"] == "invalid_request_error" + assert backend.calls == [] + + +def test_batched_unconvertible_response_is_an_errored_item_with_one_log(error_records): + backend = _Backend(result=_chat_response(tool_arguments="{not-json")) + + assert _run_batched(backend, _request()) == ( + "errored", + {"type": "api_error", "message": "Internal server error"}, + ) + assert len(error_records) == 1 + assert "Invalid response from OpenAI chat pipeline in batch" in error_records[0] + + +# -- _error_type ---------------------------------------------------------------- + + +@pytest.mark.parametrize( + "status, expected", + [ + (399, "api_error"), + (400, "invalid_request_error"), + (499, "invalid_request_error"), + (500, "api_error"), + ], +) +def test_error_type_boundaries(status, expected): + assert _error_type(status) == expected diff --git a/tests/unittest/serve/api/ops/__init__.py b/tests/unittest/serve/api/ops/__init__.py new file mode 100644 index 000000000000..467079831e16 --- /dev/null +++ b/tests/unittest/serve/api/ops/__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. diff --git a/tests/unittest/serve/api/ops/system/__init__.py b/tests/unittest/serve/api/ops/system/__init__.py new file mode 100644 index 000000000000..467079831e16 --- /dev/null +++ b/tests/unittest/serve/api/ops/system/__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. diff --git a/tests/unittest/serve/api/ops/system/test_system_endpoints.py b/tests/unittest/serve/api/ops/system/test_system_endpoints.py new file mode 100644 index 000000000000..12a88b166e90 --- /dev/null +++ b/tests/unittest/serve/api/ops/system/test_system_endpoints.py @@ -0,0 +1,132 @@ +# 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. +"""/health, /version and /v1/models behavior, through the production handlers. + +These pin what the system endpoints do today, per serving role, so a later +change can show it changed nothing. An unhealthy engine makes ``/health`` raise +``SIGINT`` (after recording a termination observation) so the process does not +linger as a zombie; the signal is patched out here before any unhealthy request. +""" + +import signal + +import pytest +from test_common.serve.app_factory import AggConfig + +import tensorrt_llm.version +from tensorrt_llm.llmapi.disagg_utils import ServerRole + +# threadleak is off: these tests create TestClients and run requests inside the +# test body, so the client's portal thread and the fake server's worker-pool +# threads are still alive when the leak check runs; the serve_env fixture closes +# them in its teardown, which runs after that check (precedent: +# tests/unittest/llmapi/apps/test_serve_lazy_import.py). +pytestmark = [pytest.mark.cpu_only, pytest.mark.threadleak(enabled=False)] + +UNAVAILABLE_BODY = "LLM is unavailable. Please check the server logs for more details." + + +@pytest.fixture +def sigint_recorder(monkeypatch): + calls = [] + monkeypatch.setattr(signal, "raise_signal", calls.append) + return calls + + +@pytest.mark.parametrize( + "role", [None, ServerRole.MM_ENCODER, ServerRole.EMBEDDING], ids=["llm", "mm", "embedding"] +) +def test_health_returns_200_with_empty_body(role, serve_env): + server = serve_env.build_openai_server(AggConfig(role=role)) + response = serve_env.client(server.app).get("/health") + assert response.status_code == 200 + assert response.content == b"" + + +@pytest.mark.parametrize("role", [None, ServerRole.MM_ENCODER], ids=["llm", "mm"]) +def test_unhealthy_executor_returns_503_without_shutdown(role, serve_env, sigint_recorder): + server = serve_env.build_openai_server(AggConfig(role=role)) + server.generator._executor.healthy = False + response = serve_env.client(server.app).get("/health") + assert response.status_code == 503 + assert response.text == UNAVAILABLE_BODY + assert sigint_recorder == [], "an unhealthy executor without a fatal error must not SIGINT" + + +@pytest.mark.parametrize("role", [None, ServerRole.MM_ENCODER], ids=["llm", "mm"]) +def test_fatal_error_triggers_sigint_and_termination_record(role, serve_env, sigint_recorder): + server = serve_env.build_openai_server(AggConfig(role=role)) + client = serve_env.client(server.app) + seen = serve_env.record_terminations() + server.generator._executor._fatal_error = RuntimeError("boom") + server.generator._executor.doing_shutdown = False + response = client.get("/health") + assert response.status_code == 503 + assert sigint_recorder == [signal.SIGINT] + assert len(seen) == 1 + assert seen[0].termination_kind == "exception" + assert seen[0].component == "llm" + assert seen[0].reporting_source == "self" + + +@pytest.mark.parametrize("role", [None, ServerRole.MM_ENCODER], ids=["llm", "mm"]) +def test_fatal_error_during_shutdown_signals_once(role, serve_env, sigint_recorder): + server = serve_env.build_openai_server(AggConfig(role=role)) + client = serve_env.client(server.app) + seen = serve_env.record_terminations() + server.generator._executor._fatal_error = RuntimeError("boom") + server.generator._executor.doing_shutdown = True + response = client.get("/health") + assert response.status_code == 503 + assert sigint_recorder == [], "a teardown already in progress must not stack signals" + assert seen == [] + + +def test_embedding_health_depends_on_batcher_and_encoder(serve_env, sigint_recorder): + server = serve_env.build_openai_server(AggConfig(role=ServerRole.EMBEDDING)) + # Without the lifespan the encode batcher never started, so every + # /v1/embeddings request would hang; /health reports unhealthy. + without_lifespan = serve_env.client(server.app, lifespan=False) + assert without_lifespan.get("/health").status_code == 503 + with_lifespan = serve_env.client(server.app) + assert with_lifespan.get("/health").status_code == 200 + # Encode-only generators have no generation executor; health reads the + # encoder executor instead. + server.generator._encoder_executor = None + assert with_lifespan.get("/health").status_code == 503 + assert sigint_recorder == [], ( + "the embedding health path finds no generation executor, so it must not SIGINT" + ) + + +def test_version_reports_package_version(serve_env): + server = serve_env.build_openai_server(AggConfig()) + response = serve_env.client(server.app).get("/version") + assert response.status_code == 200 + assert response.json() == {"version": tensorrt_llm.version.__version__} + + +def test_models_lists_the_served_model(serve_env): + server = serve_env.build_openai_server(AggConfig()) + response = serve_env.client(server.app).get("/v1/models") + assert response.status_code == 200 + body = response.json() + assert body["object"] == "list" + assert len(body["data"]) == 1 + entry = body["data"][0] + assert entry["id"] == "fake-model" + assert entry["object"] == "model" + assert entry["owned_by"] == "tensorrt_llm" + assert isinstance(entry["created"], int) # the request time; never compare the value diff --git a/tests/unittest/serve/api/ops/system/test_system_router.py b/tests/unittest/serve/api/ops/system/test_system_router.py new file mode 100644 index 000000000000..36b680d1e7c2 --- /dev/null +++ b/tests/unittest/serve/api/ops/system/test_system_router.py @@ -0,0 +1,65 @@ +# 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. +"""The ops.system router: attachment contract and route provenance per role.""" + +import pytest +from fastapi import FastAPI +from test_common.serve.app_factory import AggConfig +from test_common.serve.routes import effective_routes + +from tensorrt_llm.llmapi.disagg_utils import ServerRole +from tensorrt_llm.serve._base.msgspec_route import _MsgspecRoute +from tensorrt_llm.serve.api.ops.system.api_router import attach_router +from tensorrt_llm.serve.api.ops.system.serving import SystemServing + +pytestmark = pytest.mark.cpu_only + +SYSTEM_PATHS = ("/health", "/version", "/v1/models") +SYSTEM_ROUTER_MODULE = "tensorrt_llm.serve.api.ops.system.api_router" + + +def test_attach_requires_serving_state(): + # The composition root must set the serving object before attaching, so a + # wiring bug fails at startup, not with an AttributeError per request. + with pytest.raises(RuntimeError, match="ops_serving_system"): + attach_router(FastAPI()) + + +def test_system_serving_requires_arguments(): + # No defaults: a serving class built without its collaborators must fail at + # construction, not fall back silently. + with pytest.raises(TypeError): + SystemServing() + + +@pytest.mark.parametrize( + "role", + [None, ServerRole.MM_ENCODER, ServerRole.EMBEDDING], + ids=["llm", "mm_encoder", "embedding"], +) +def test_routes_come_from_ops_system(role, serve_env): + server = serve_env.build_openai_server(AggConfig(role=role)) + by_path = {path: [] for path in SYSTEM_PATHS} + for path, route in effective_routes(server.app): + if path in by_path: + by_path[path].append(route) + for path, routes in by_path.items(): + assert len(routes) == 1, f"{path} must be registered exactly once, found {len(routes)}" + route = routes[0] + assert isinstance(route, _MsgspecRoute), ( + f"{path} must be a _MsgspecRoute; a plain APIRoute breaks msgpack request " + "bodies from the disaggregated frontend" + ) + assert route.endpoint.__module__ == SYSTEM_ROUTER_MODULE diff --git a/tests/unittest/serve/api/visual_gen/__init__.py b/tests/unittest/serve/api/visual_gen/__init__.py new file mode 100644 index 000000000000..467079831e16 --- /dev/null +++ b/tests/unittest/serve/api/visual_gen/__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. diff --git a/tests/unittest/serve/api/visual_gen/test_ops_system_router_visual_gen.py b/tests/unittest/serve/api/visual_gen/test_ops_system_router_visual_gen.py new file mode 100644 index 000000000000..488a5e621d05 --- /dev/null +++ b/tests/unittest/serve/api/visual_gen/test_ops_system_router_visual_gen.py @@ -0,0 +1,48 @@ +# 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. +"""Route provenance of the system endpoints for the VisualGen serving role. + +Lives under a ``visual_gen/`` path so that VisualGen-only pull requests, whose +CI keeps only test entries containing ``visual_gen/``, still run it. +""" + +import pytest +from test_common.serve.app_factory import isolated_serve_env +from test_common.serve.routes import VISUAL_GEN_CONFIGS, effective_routes + +from tensorrt_llm.serve._base.msgspec_route import _MsgspecRoute + +pytestmark = pytest.mark.cpu_only + +SYSTEM_PATHS = ("/health", "/version", "/v1/models") +SYSTEM_ROUTER_MODULE = "tensorrt_llm.serve.api.ops.system.api_router" + + +@pytest.mark.parametrize("config", VISUAL_GEN_CONFIGS, ids=lambda config: config.name) +def test_routes_come_from_ops_system_visual_gen(config, tmp_path): + with isolated_serve_env(tmp_path) as env: + server = env.build_openai_server(config.build) + by_path = {path: [] for path in SYSTEM_PATHS} + for path, route in effective_routes(server.app): + if path in by_path: + by_path[path].append(route) + for path, routes in by_path.items(): + assert len(routes) == 1, f"{path} must be registered exactly once, found {len(routes)}" + route = routes[0] + assert isinstance(route, _MsgspecRoute), ( + f"{path} must be a _MsgspecRoute; a plain APIRoute breaks msgpack request " + "bodies from the disaggregated frontend" + ) + assert route.endpoint.__module__ == SYSTEM_ROUTER_MODULE diff --git a/tests/unittest/serve/api/visual_gen/test_ops_system_visual_gen.py b/tests/unittest/serve/api/visual_gen/test_ops_system_visual_gen.py new file mode 100644 index 000000000000..205391381c13 --- /dev/null +++ b/tests/unittest/serve/api/visual_gen/test_ops_system_visual_gen.py @@ -0,0 +1,52 @@ +# 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. +"""/health, /version and /v1/models of the VisualGen serving role.""" + +import pytest +from test_common.serve.app_factory import AggConfig, isolated_serve_env + +import tensorrt_llm.version +from tensorrt_llm.llmapi.disagg_utils import ServerRole + +pytestmark = pytest.mark.cpu_only + + +@pytest.fixture +def visual_gen_client(tmp_path): + with isolated_serve_env(tmp_path) as env: + server = env.build_openai_server( + AggConfig(role=ServerRole.VISUAL_GEN, model_name="fake-visual-gen") + ) + yield env.client(server.app) + + +def test_health_always_succeeds_today(visual_gen_client): + # The VisualGen generator exposes no health probe, so /health answers 200 + # unconditionally today; a change to that is a behavior change. + assert visual_gen_client.get("/health").status_code == 200 + + +def test_version(visual_gen_client): + response = visual_gen_client.get("/version") + assert response.status_code == 200 + assert response.json() == {"version": tensorrt_llm.version.__version__} + + +def test_models_lists_the_served_model(visual_gen_client): + response = visual_gen_client.get("/v1/models") + assert response.status_code == 200 + body = response.json() + assert body["object"] == "list" + assert [entry["id"] for entry in body["data"]] == ["fake-visual-gen"] diff --git a/tests/unittest/serve/api/visual_gen/test_routes_visual_gen.py b/tests/unittest/serve/api/visual_gen/test_routes_visual_gen.py new file mode 100644 index 000000000000..c20344b6eea7 --- /dev/null +++ b/tests/unittest/serve/api/visual_gen/test_routes_visual_gen.py @@ -0,0 +1,56 @@ +# 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. +"""Route snapshot of the VisualGen serving configuration. + +Lives under a ``visual_gen/`` path so that VisualGen-only pull requests, whose +CI keeps only test entries containing ``visual_gen/``, still run it. +""" + +import pytest +from test_common.serve.routes import ( + VISUAL_GEN_CONFIGS, + compare, + format_failure, + read_reference, + snapshot_config, +) + +pytestmark = pytest.mark.cpu_only + +_DOC_ROUTES = {"/openapi.json", "/docs", "/docs/oauth2-redirect", "/redoc"} + + +@pytest.mark.parametrize("config", VISUAL_GEN_CONFIGS, ids=lambda config: config.name) +def test_visual_gen_routes_match_reference(config, tmp_path): + reference = read_reference() + recorded = reference["configs"][config.name] + live = snapshot_config(config, tmp_path) + assert recorded.get("error") == live.get("error") + differences = compare(recorded.get("routes", []), live.get("routes", [])) + assert not differences, format_failure(config.name, differences) + + +@pytest.mark.parametrize("config", VISUAL_GEN_CONFIGS, ids=lambda config: config.name) +def test_visual_gen_routes_keep_msgpack_route_class(config, tmp_path): + live = snapshot_config(config, tmp_path) + for record in live["routes"]: + methods, path, cls = record.split(" ")[:3] + if path in _DOC_ROUTES: + assert cls == "Route", record + else: + assert cls == "_MsgspecRoute", ( + f"{record}: every VisualGen API route must keep the msgpack route class, " + "like every other aggregated-server route." + ) diff --git a/tests/unittest/serve/api/visual_gen/test_shim_identity_visual_gen.py b/tests/unittest/serve/api/visual_gen/test_shim_identity_visual_gen.py new file mode 100644 index 000000000000..1e2d2d57bbe0 --- /dev/null +++ b/tests/unittest/serve/api/visual_gen/test_shim_identity_visual_gen.py @@ -0,0 +1,46 @@ +# 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. +"""Identity rows of VisualGen names re-exported from moved serve modules. + +``VISUAL_GEN_SHIM_ROWS`` starts empty. No kept shim has a VisualGen path, so +this table only ever receives interface rows (``warns=False``, +``ban_old_path=False``): names a package facade keeps re-exporting as its own +interface, checked for identity only. Lives under a ``visual_gen/`` path so that +VisualGen-only pull requests still run it; the machinery's self-tests are in +``tests/unittest/serve/test_shim_identity.py``. +""" + +from pathlib import Path + +import pytest +from test_common.serve.shims import check_identity, check_warnings, find_old_path_imports + +pytestmark = pytest.mark.cpu_only + +REPO_ROOT = Path(__file__).resolve().parents[5] + +VISUAL_GEN_SHIM_ROWS: tuple = () + + +def test_visual_gen_shim_identity(): + assert check_identity(VISUAL_GEN_SHIM_ROWS) == [] + + +def test_visual_gen_shim_warnings(): + assert check_warnings(VISUAL_GEN_SHIM_ROWS) == [] + + +def test_visual_gen_no_old_path_imports(): + assert find_old_path_imports(VISUAL_GEN_SHIM_ROWS, REPO_ROOT) == [] diff --git a/tests/unittest/serve/conftest.py b/tests/unittest/serve/conftest.py new file mode 100644 index 000000000000..cda9b5cc0c2e --- /dev/null +++ b/tests/unittest/serve/conftest.py @@ -0,0 +1,26 @@ +# 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. +import pytest + + +@pytest.fixture +def serve_env(tmp_path): + """A :class:`test_common.serve.app_factory.ServeTestEnv` per test.""" + # Imported here, not at module level: this conftest also governs checks + # that need no tensorrt_llm build (the ratchet, the agent-docs test). + from test_common.serve.app_factory import isolated_serve_env + + with isolated_serve_env(tmp_path) as env: + yield env diff --git a/tests/unittest/serve/test_agents_md.py b/tests/unittest/serve/test_agents_md.py new file mode 100644 index 000000000000..b95dac68ff16 --- /dev/null +++ b/tests/unittest/serve/test_agents_md.py @@ -0,0 +1,132 @@ +# 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. +"""The serve agent docs name only paths and commands that exist. + +Rules written for agents rot silently when the file or command they name moves; +this test keeps every inline-code path and command in the serve ``AGENTS.md`` +files real, and bans line-number anchors, which drift without failing anything. +""" + +import re +from pathlib import Path + +import pytest + +pytestmark = pytest.mark.cpu_only + +REPO_ROOT = Path(__file__).resolve().parents[3] + +DOCS = ("tensorrt_llm/serve/AGENTS.md", "tensorrt_llm/commands/AGENTS.md") + +PATH_ROOTS = ( + "tensorrt_llm/", + "serve/", + "commands/", + "tests/", + "scripts/", + "jenkins/", + ".github/", + "docs/", + "examples/", +) +ROOT_FILES = ("AGENTS.md", "CLAUDE.md", "legacy-files.txt", "CODING_GUIDELINES.md") +COMMAND_STARTS = ("pytest", "python", "python3", "PYTHONPATH=", "LLM_MODELS_ROOT=") + +_INLINE_CODE = re.compile(r"`([^`]+)`") +_LINE_ANCHOR = re.compile(r"\.(py|md|ya?ml|groovy|txt)\s*:\s*\d") +_BARE_ANCHOR_SPAN = re.compile(r"`:\d") + + +def _spans(doc: Path): + return _INLINE_CODE.findall(doc.read_text()) + + +def _resolve(path_text: str) -> Path: + if path_text.startswith(("serve/", "commands/")): + return REPO_ROOT / "tensorrt_llm" / path_text + return REPO_ROOT / path_text + + +def _is_template(token: str) -> bool: + return any(symbol in token for symbol in "<>*") + + +def _missing(doc: str, token: str) -> str: + return ( + f"{doc} names `{token}`, which does not exist. Update the rule in the PR that " + "moved or deleted it." + ) + + +def _check_command(doc: str, span: str, problems: list) -> None: + tokens = span.split(" ") + pythonpath_roots = [] + for token in tokens: + if token.startswith("PYTHONPATH="): + pythonpath_roots = token.split("=", 1)[1].split(":") + for index, token in enumerate(tokens): + if _is_template(token): + continue + if index > 0 and tokens[index - 1] == "-m" and "/" not in token and tokens[0] != "pytest": + # `python3 -m `: resolve to a file under the PYTHONPATH= + # roots or the repository root. + candidates = [] + for root in [*pythonpath_roots, "."]: + base = REPO_ROOT / root / Path(*token.split(".")) + candidates.extend((base.with_suffix(".py"), base / "__init__.py")) + if not any(candidate.exists() for candidate in candidates): + problems.append(_missing(doc, token)) + continue + if token.startswith(PATH_ROOTS): + target = token.split("::")[0] + if not _resolve(target).exists(): + problems.append(_missing(doc, token)) + + +@pytest.mark.parametrize("doc", DOCS) +def test_paths_and_commands_are_real(doc): + doc_path = REPO_ROOT / doc + problems = [] + for span in _spans(doc_path): + if " " not in span and (span.startswith(PATH_ROOTS) or span in ROOT_FILES): + if _is_template(span): + continue + target = _resolve(span.split("::")[0]) + if not target.exists(): + problems.append(_missing(doc, span)) + elif span.startswith(COMMAND_STARTS): + _check_command(doc, span, problems) + assert not problems, "\n".join(problems) + + +@pytest.mark.parametrize("doc", DOCS) +def test_no_line_number_anchors(doc): + text = (REPO_ROOT / doc).read_text() + assert not _LINE_ANCHOR.search(text) and not _BARE_ANCHOR_SPAN.search(text), ( + f"{doc} contains a line-number anchor. Name files and symbols, never line numbers: " + "anchors drift silently when the file changes." + ) + + +@pytest.mark.parametrize("doc", DOCS) +def test_sibling_claude_md(doc): + doc_path = REPO_ROOT / doc + sibling = doc_path.with_name("CLAUDE.md") + assert sibling.exists(), f"{doc} needs a sibling CLAUDE.md pointer file" + lines = [line.strip() for line in sibling.read_text().splitlines() if line.strip()] + assert lines == ["# In ./CLAUDE.md", "@AGENTS.md"], ( + f"{sibling.relative_to(REPO_ROOT)} must hold exactly the repo's two-line pointer " + "form: `# In ./CLAUDE.md` and `@AGENTS.md`." + ) diff --git a/tests/unittest/serve/test_layering.py b/tests/unittest/serve/test_layering.py new file mode 100644 index 000000000000..e09971963ba0 --- /dev/null +++ b/tests/unittest/serve/test_layering.py @@ -0,0 +1,361 @@ +# 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. +"""Import directions between serving layers, written as tests. + +Each pull request that creates a package under ``tensorrt_llm/serve`` adds its +contract here. A contract forbids imports by module prefix; a baseline allowance +records the imports that existed when the contract was written and only shrinks. +""" + +import ast +import textwrap +from pathlib import Path + +import pytest +from test_common.serve.layering import ( + Contract, + ImportEdge, + build_import_graph, + check_contracts, + stale_allowances, +) + +pytestmark = pytest.mark.cpu_only + +REPO_ROOT = Path(__file__).resolve().parents[3] + +CONTRACTS = ( + Contract( + name="render core has no HTTP or executor dependency", + importers=( + "tensorrt_llm.serve.render.chat", + "tensorrt_llm.serve.render.resources", + "tensorrt_llm.serve.render.fingerprint", + "tensorrt_llm.serve.render.prompt_types", + ), + forbidden=("fastapi", "starlette", "tensorrt_llm.executor"), + fix=( + "Keep HTTP code in serve/render/api_router.py or _http.py and pass engine data " + "in through RenderResources." + ), + ), + Contract( + name="gRPC frontends do not reach into serve", + importers=("tensorrt_llm.grpc",), + forbidden=("tensorrt_llm.serve",), + allowed=frozenset( + { + ( + "tensorrt_llm.grpc.openengine.server", + "tensorrt_llm.serve._frontend_processes", + name, + ) + for name in ( + "FrontendStartupCancelled", + "_init_multi_frontend_mode", + "_signal_frontend_ready", + "_spawn_attached_frontends", + "_terminate_attached_frontends", + ) + } + | { + ( + "tensorrt_llm.grpc.openengine.server", + "tensorrt_llm.serve._perf_metrics_config", + "resolve_perf_metrics_config", + ) + } + ), + fix=( + "gRPC frontends use serving code only through public entry points; add a public " + "API in tensorrt_llm/llmapi instead. The baseline only shrinks: delete an " + "allowance when its import goes away." + ), + ), + Contract( + name="serve/_base/ does not import serve/api/", + importers=("tensorrt_llm.serve._base",), + forbidden=("tensorrt_llm.serve.api",), + fix=( + "serve/_base/ is shared infrastructure and must not depend on endpoint " + "packages; pass what it needs in from the composition root." + ), + ), + Contract( + name="the Anthropic package does not import a server", + importers=("tensorrt_llm.serve.api.anthropic",), + forbidden=( + "tensorrt_llm.serve.openai_server", + "tensorrt_llm.serve.openai_disagg_server", + ), + fix=( + "serve/api/anthropic/ serves both servers; take the chat backend and other " + "dependencies through constructors and never import a server module." + ), + ), + Contract( + name="serve/api/ and serve/_base/ do not import the composition root", + importers=("tensorrt_llm.serve.api", "tensorrt_llm.serve._base"), + forbidden=("tensorrt_llm.serve.openai_server", "tensorrt_llm.serve.app"), + fix=( + "Endpoint packages and serve/_base/ must not import the composition root " + "(serve/openai_server.py or serve/app.py); the root imports them and injects " + "what they need. Move the shared helper to serve/_base/ instead." + ), + ), +) + + +@pytest.fixture(scope="module") +def edges(): + return build_import_graph(REPO_ROOT / "tensorrt_llm", "tensorrt_llm") + + +def test_contracts_hold(edges): + assert check_contracts(edges, CONTRACTS) == [] + + +def test_no_stale_allowances(edges): + assert stale_allowances(edges, CONTRACTS) == [] + + +_API_PREFIX = "tensorrt_llm.serve.api" + + +def _api_package_of(module: str): + """The endpoint package of a module under ``tensorrt_llm.serve.api``. + + The package is the first two components below ``api`` for ``openai`` and + ``ops`` (``openai.chat``, ``ops.system``) and the first one otherwise + (``anthropic``, ``visual_gen``). + """ + below = module.split(".")[3:] + if below and below[0] in ("openai", "ops"): + return tuple(below[:2]) + return tuple(below[:1]) + + +def _under_api(module: str) -> bool: + return module == _API_PREFIX or module.startswith(_API_PREFIX + ".") + + +def _cross_package_imports(edges) -> list: + """Imports between endpoint packages, matched per imported name as check_contracts does. + + ``from tensorrt_llm.serve import api`` names the api root package and + ``from tensorrt_llm.serve.api.ops import system`` names ``ops.system``: the + package is read off ``.``, not off the ``from`` module alone. + """ + problems = [] + for edge in edges: + if edge.type_checking or not _under_api(edge.importer): + continue + for name in edge.names or ("",): + target = f"{edge.imported}.{name}" if name else edge.imported + if not _under_api(target): + continue + if _api_package_of(target) != _api_package_of(edge.importer): + problems.append( + f"serve/api/ packages do not import each other: {edge.importer} " + f"(line {edge.lineno}) imports `{target}`. Share code through serve/_base/ " + "or inject it from the composition root." + ) + return problems + + +def test_api_packages_do_not_import_each_other(edges): + assert _cross_package_imports(edges) == [] + + +def test_cross_package_classification_reads_the_imported_name(): + def edge(importer, imported, names): + return ImportEdge( + importer=importer, + imported=imported, + names=names, + lineno=1, + in_function=False, + type_checking=False, + ) + + system = "tensorrt_llm.serve.api.ops.system.api_router" + # Parent-package forms resolve to the submodule they name: the same package. + assert _cross_package_imports([edge(system, "tensorrt_llm.serve.api.ops", ("system",))]) == [] + assert ( + _cross_package_imports([edge(system, "tensorrt_llm.serve.api.ops.system", ("serving",))]) + == [] + ) + # The api root, a sibling package and another dialect are cross-package, whatever the form. + assert len(_cross_package_imports([edge(system, "tensorrt_llm.serve", ("api",))])) == 1 + assert ( + len(_cross_package_imports([edge(system, "tensorrt_llm.serve.api.ops", ("metrics",))])) == 1 + ) + anthropic = "tensorrt_llm.serve.api.anthropic.serving" + assert len(_cross_package_imports([edge(system, anthropic, ("AnthropicServing",))])) == 1 + # Importers outside serve/api/ are not classified here. + root = "tensorrt_llm.serve.openai_server" + assert ( + _cross_package_imports( + [edge(root, "tensorrt_llm.serve.api.ops.system.serving", ("SystemServing",))] + ) + == [] + ) + + +def test_api_init_imports_nothing(): + tree = ast.parse((REPO_ROOT / "tensorrt_llm" / "serve" / "api" / "__init__.py").read_text()) + imports = [node for node in ast.walk(tree) if isinstance(node, (ast.Import, ast.ImportFrom))] + assert imports == [], ( + "serve/api/__init__.py imports nothing: attach_routers imports routers by role; an " + "import here would load them for every role and break the VisualGen CI selection." + ) + + +# --------------------------------------------------------------------------- +# Self-test on a synthetic package: the graph must see every import form. +# --------------------------------------------------------------------------- + + +@pytest.fixture +def synthetic_package(tmp_path): + pkg = tmp_path / "layerpkg" + pkg.mkdir() + (pkg / "__init__.py").write_text("") + (pkg / "inner.py").write_text("VALUE = 1\n") + (pkg / "top.py").write_text( + textwrap.dedent( + """ + from typing import TYPE_CHECKING + + import os + from json import dumps, loads + + from .inner import VALUE + + if TYPE_CHECKING: + import socket + + def late(): + import importlib + + importlib.import_module("collections.abc") + from ..sibling import thing # noqa: F401 + return VALUE + """ + ) + ) + return pkg + + +def test_import_graph_records_every_form(synthetic_package): + edges = build_import_graph(synthetic_package, "layerpkg") + by_imported = {edge.imported: edge for edge in edges} + + assert by_imported["os"].names == () + assert not by_imported["os"].in_function and not by_imported["os"].type_checking + assert by_imported["json"].names == ("dumps", "loads") + assert by_imported["layerpkg.inner"].names == ("VALUE",) # relative import resolved + assert by_imported["socket"].type_checking + assert by_imported["importlib"].in_function + assert by_imported["collections.abc"].names == () # importlib.import_module literal + assert by_imported["collections.abc"].in_function + assert by_imported["sibling"].names == ("thing",) # level-2 relative from layerpkg.top + + +def test_contract_violation_and_allowance(synthetic_package): + edges = build_import_graph(synthetic_package, "layerpkg") + contract = Contract( + name="top does not import json", + importers=("layerpkg.top",), + forbidden=("json",), + fix="Use the serializer helper instead.", + ) + problems = check_contracts(edges, [contract]) + assert len(problems) == 2 # one per from-import name + assert "contract 'top does not import json'" in problems[0] + assert "Use the serializer helper instead." in problems[0] + + allowed = Contract( + name="top does not import json", + importers=("layerpkg.top",), + forbidden=("json",), + allowed=frozenset({("layerpkg.top", "json", "dumps"), ("layerpkg.top", "json", "loads")}), + fix="Use the serializer helper instead.", + ) + assert check_contracts(edges, [allowed]) == [] + assert stale_allowances(edges, [allowed]) == [] + + stale = Contract( + name="top does not import json", + importers=("layerpkg.top",), + forbidden=("json",), + allowed=frozenset({("layerpkg.top", "json", "gone")}), + fix="Use the serializer helper instead.", + ) + messages = stale_allowances(edges, [stale]) + assert len(messages) == 1 and "the baseline only shrinks" in messages[0] + + +def test_parent_package_import_of_a_forbidden_module(synthetic_package): + # Both forms record the edge against the parent package, with the forbidden + # submodule only as an imported name. + (synthetic_package / "via_parent.py").write_text( + textwrap.dedent( + """ + from layerpkg import inner, top + from . import inner as relative_inner + """ + ) + ) + edges = build_import_graph(synthetic_package, "layerpkg") + contract = Contract( + name="no inner", + importers=("layerpkg.via_parent",), + forbidden=("layerpkg.inner",), + fix="Import VALUE through layerpkg.top.", + ) + problems = check_contracts(edges, [contract]) + assert len(problems) == 2 # `top` from the same statement is not flagged + assert "(line 2) imports `layerpkg.inner`." in problems[0] # from layerpkg import inner + assert "(line 3) imports `layerpkg.inner`." in problems[1] # from . import inner + + allowed = Contract( + name="no inner", + importers=("layerpkg.via_parent",), + forbidden=("layerpkg.inner",), + allowed=frozenset({("layerpkg.via_parent", "layerpkg", "inner")}), + fix="Import VALUE through layerpkg.top.", + ) + assert check_contracts(edges, [allowed]) == [] + assert stale_allowances(edges, [allowed]) == [] + + +def test_type_checking_imports_are_skipped_by_default(synthetic_package): + edges = build_import_graph(synthetic_package, "layerpkg") + contract = Contract( + name="no socket", + importers=("layerpkg",), + forbidden=("socket",), + fix="unused", + ) + assert check_contracts(edges, [contract]) == [] + including = Contract( + name="no socket", + importers=("layerpkg",), + forbidden=("socket",), + include_type_checking=True, + fix="unused", + ) + assert len(check_contracts(edges, [including])) == 1 diff --git a/tests/unittest/serve/test_middleware_order.py b/tests/unittest/serve/test_middleware_order.py new file mode 100644 index 000000000000..aec17dd4f79d --- /dev/null +++ b/tests/unittest/serve/test_middleware_order.py @@ -0,0 +1,133 @@ +# 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. +"""Middleware-order snapshot of the aggregated and disaggregated apps. + +Starlette's ``add_middleware`` inserts at position 0, so the last one added is +outermost. ``ServerArrivalTimeMiddleware`` assigns ``scope['state'] = {}`` and +``PerfMetricsMiddleware`` only ``setdefault``s into it, so their relative order +is load-bearing: swapping them wipes every perf-metrics record. VisualGen is not +listed: it takes the default code path without perf metrics, and no VisualGen +path can change middleware registration. +""" + +import pytest +from test_common.serve.app_factory import AggConfig, CoordinatorConfig, DisaggConfig +from test_common.serve.middleware import middleware_snapshot + +from tensorrt_llm.llmapi.disagg_utils import ServerRole + +pytestmark = pytest.mark.cpu_only + +_ARRIVAL_WITH_CLOCK = {"class": "ServerArrivalTimeMiddleware", "kwargs": ["adjusted_clock"]} +_ARRIVAL_PLAIN = {"class": "ServerArrivalTimeMiddleware", "kwargs": []} + + +def _perf(kwargs, expose_headers): + return { + "class": "PerfMetricsMiddleware", + "kwargs": sorted(kwargs), + "expose_headers": expose_headers, + } + + +# Outermost first, matching middleware_snapshot. +EXPECTED = { + "llm": (AggConfig(), [_ARRIVAL_WITH_CLOCK]), + "mm_encoder": (AggConfig(role=ServerRole.MM_ENCODER), [_ARRIVAL_WITH_CLOCK]), + "embedding": (AggConfig(role=ServerRole.EMBEDDING), [_ARRIVAL_WITH_CLOCK]), + "llm_perf_metrics": ( + AggConfig(return_perf_metrics=True), + [ + _ARRIVAL_WITH_CLOCK, + _perf(["adjusted_clock", "expose_headers", "writer"], True), + ], + ), + "llm_perf_metrics_output_dir_only": ( + AggConfig(perf_metrics_output_dir=True), + [ + _ARRIVAL_WITH_CLOCK, + _perf(["adjusted_clock", "expose_headers", "writer"], False), + ], + ), + "llm_perf_metrics_enable_metrics_env": ( + AggConfig(return_perf_metrics=True, env=(("ENABLE_METRICS", "true"),)), + [ + _ARRIVAL_WITH_CLOCK, + _perf(["adjusted_clock", "expose_headers", "writer"], True), + {"class": "PrometheusInstrumentatorMiddleware"}, + ], + ), + "disagg_local": (DisaggConfig(), [_ARRIVAL_PLAIN]), + "disagg_fleet": (DisaggConfig(fleet=True), [_ARRIVAL_PLAIN]), + "disagg_local_perf_metrics": ( + DisaggConfig(return_perf_metrics=True), + [_ARRIVAL_PLAIN, _perf(["expose_headers", "writer"], True)], + ), + "disagg_fleet_perf_metrics": ( + DisaggConfig(fleet=True, return_perf_metrics=True), + [_ARRIVAL_PLAIN, _perf(["expose_headers", "writer"], True)], + ), + "coordinator": (CoordinatorConfig(), []), +} + + +def _build(serve_env, config): + if isinstance(config, AggConfig): + return serve_env.build_openai_server(config).app + if isinstance(config, DisaggConfig): + return serve_env.build_disagg_server(config).app + return serve_env.build_coordinator_server(config).app + + +def _failure(name, expected, got): + return ( + f"Middleware stack of '{name}' changed: expected {expected}, got {got}. " + "Starlette's add_middleware inserts at position 0, so the last one added is " + "outermost. ServerArrivalTimeMiddleware assigns scope['state'] = {} and must stay " + "outside PerfMetricsMiddleware, or every perf-metrics record is wiped. Agg passes " + "adjusted_clock to both, disagg to neither. If the change is intended, update " + "EXPECTED in tests/unittest/serve/test_middleware_order.py and say why in the PR " + "body." + ) + + +@pytest.mark.parametrize("name", sorted(EXPECTED)) +def test_middleware_order(name, serve_env): + config, expected = EXPECTED[name] + snapshot = middleware_snapshot(_build(serve_env, config)) + assert snapshot == expected, _failure(name, expected, snapshot) + + +@pytest.mark.parametrize("name", sorted(EXPECTED)) +def test_arrival_time_stays_outside_perf_metrics(name, serve_env): + """When both are present, ServerArrivalTimeMiddleware is outermost.""" + config, _ = EXPECTED[name] + classes = [entry["class"] for entry in middleware_snapshot(_build(serve_env, config))] + if "PerfMetricsMiddleware" not in classes or "ServerArrivalTimeMiddleware" not in classes: + return + assert classes.index("ServerArrivalTimeMiddleware") < classes.index("PerfMetricsMiddleware"), ( + _failure(name, "ServerArrivalTimeMiddleware before PerfMetricsMiddleware", classes) + ) + + +def test_user_middleware_is_outermost(serve_env): + """``--middleware`` classes are applied after construction, so they end up outermost.""" + from tensorrt_llm.commands.serve import _apply_fastapi_middlewares + + server = serve_env.build_openai_server(AggConfig()) + _apply_fastapi_middlewares(server.app, ["test_common.serve.fakes.UserProbeMiddleware"]) + snapshot = middleware_snapshot(server.app) + expected = [{"class": "UserProbeMiddleware"}, _ARRIVAL_WITH_CLOCK] + assert snapshot == expected, _failure("llm_user_middleware", expected, snapshot) diff --git a/tests/unittest/serve/test_no_object_new.py b/tests/unittest/serve/test_no_object_new.py new file mode 100644 index 000000000000..99fa1ed463aa --- /dev/null +++ b/tests/unittest/serve/test_no_object_new.py @@ -0,0 +1,290 @@ +# 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. +"""Ratchet against new ``object.__new__()`` test sites. + +Tests that build a server with ``object.__new__`` bypass the constructor and pin +whatever attribute set the handler happens to read; they break or silently +mistest when handlers move. New tests build servers with ``isolated_serve_env`` +from ``tests/test_common/serve/app_factory.py`` instead. The existing sites stay +until the change that moves each handler rewrites them, so this test holds a +per-file baseline of sites that may only shrink. A site is named by its +enclosing function or class and the server class it builds, so a site added in +another function fails even when one is removed from the same file. + +The scan counts AST call nodes, not text, so a docstring or comment that +mentions ``object.__new__(OpenAIServer)`` does not count. A name bound by +``from ... import OpenAIServer as Server`` counts as ``OpenAIServer``. +""" + +import ast +from collections import Counter +from pathlib import Path + +import pytest + +pytestmark = pytest.mark.cpu_only + +REPO_ROOT = Path(__file__).resolve().parents[3] +TESTS_ROOT = REPO_ROOT / "tests" + +SERVER_CLASSES = ("OpenAIServer", "OpenAIDisaggServer", "CoordinatorServer") + +# Per-file sites that existed when this ratchet landed, as +# " " (see find_new_sites), repeated +# when one scope builds the class twice. The baseline only shrinks: rewrite a +# site on the factory, delete its entry. +BASELINE = { + "tests/unittest/disaggregated/test_coordinator_worker.py": [ + "test_coordinator_keep_alive_timeout_is_passed_to_uvicorn CoordinatorServer", + "test_coordinator_tcp_listener_uses_the_prebound_socket CoordinatorServer", + ], + "tests/unittest/disaggregated/test_disagg_internal_auth.py": [ + "test_aggregated_worker_ignores_affinity_header OpenAIServer", + "test_worker_affinity_survives_wire_roundtrip OpenAIServer", + "test_worker_rejects_protected_fields_without_cache_transceiver_config OpenAIServer", + ], + "tests/unittest/disaggregated/test_disagg_openai_client.py": [ + "test_worker_clock_calibration_uses_global_clock OpenAIServer", + ], + "tests/unittest/disaggregated/test_disagg_responses_api.py": [ + "test_a_context_worker_with_postprocessing_workers_returns_the_prompt OpenAIServer", + "test_a_models_failure_is_an_error_response OpenAIDisaggServer", + ], + "tests/unittest/disaggregated/test_openai_disagg_server.py": [ + "test_create_client_does_not_register_with_server_metrics_collector OpenAIDisaggServer", + "test_http_cluster_storage_request_is_proxied_to_coordinator OpenAIDisaggServer", + "test_server_keep_alive_timeout_is_passed_to_uvicorn OpenAIDisaggServer", + ], + "tests/unittest/disaggregated/test_openai_disagg_service.py": [ + "test_create_chat_response_sets_prompt_token_ids_for_context_only OpenAIServer", + "test_subagent_affinity_http_worker_adp_placement.make_worker OpenAIServer", + ], + "tests/unittest/disaggregated/test_openai_server_info.py": [ + "_make_server OpenAIServer", + ], + "tests/unittest/disaggregated/test_openai_server_stream_errors.py": [ + "_bare_server OpenAIServer", + ], + "tests/unittest/disaggregated/test_paired_request_cancellation.py": [ + "test_http_failure_reaches_paired_worker_abort OpenAIDisaggServer", + "test_http_failure_reaches_paired_worker_abort OpenAIServer", + ], + "tests/unittest/disaggregated/test_router.py": [ + "test_chat_harmony_preserves_original_tool_conversion_error OpenAIServer", + "test_gpt_oss_router_tokens_match_chat_harmony_server_input OpenAIServer", + ], + "tests/unittest/executor/test_spec_decode_stats_payload.py": [ + "TestServerOptIn._server OpenAIServer", + ], + "tests/unittest/llmapi/apps/_test_trtllm_serve_profile.py": [ + "server_with_mock_generator OpenAIServer", + ], + "tests/unittest/llmapi/apps/test_anthropic_routes.py": [ + "_chat_stream OpenAIServer", + "_handle OpenAIDisaggServer", + ], + "tests/unittest/llmapi/apps/test_context_length_error.py": [ + "_disagg_server OpenAIDisaggServer", + ], + "tests/unittest/llmapi/apps/test_openai_server_iteration_stats.py": [ + "_make_server OpenAIServer", + ], + "tests/unittest/llmapi/apps/test_render_callers.py": [ + "TestServerInfoFingerprint.test_the_first_call_hashes_off_the_event_loop OpenAIServer", + ], + "tests/unittest/llmapi/apps/test_responses_store_replay.py": [ + "_server_up_to_preprocess OpenAIServer", + ], + "tests/unittest/llmapi/apps/test_responses_streaming_events.py": [ + "_server_streaming OpenAIServer", + ], + "tests/unittest/llmapi/apps/test_serving_extensions.py": [ + "TestOpenAIChatIntegration.chat_client OpenAIServer", + ], + "tests/unittest/llmapi/apps/test_tool_parsers.py": [ + "TestConfigureParserSpecialTokenDecoding.test_the_responses_route_applies_it OpenAIServer", + ], + "tests/unittest/llmapi/test_kv_hints.py": [ + "test_http_handlers_forward_hints OpenAIServer", + ], + "tests/unittest/llmapi/test_llm.py": [ + "test_openai_completion_list_prompt_stream_reuses_stream_metadata.run_request OpenAIServer", + "test_startup_metrics_by_server_info OpenAIServer", + ], + "tests/unittest/llmapi/test_rl_control_auth.py": [ + "_make_server OpenAIServer", + ], + "tests/unittest/llmapi/test_runtime_control.py": [ + "_make_server OpenAIServer", + ], + "tests/unittest/llmapi/test_serve_multi_frontend_helpers.py": [ + "_attached_server OpenAIServer", + "test_on_launcher_lost_fails_fast_then_shuts_down OpenAIServer", + "test_only_the_launcher_registers_the_server OpenAIServer", + "test_set_steady_clock_offset OpenAIServer", + ], + "tests/unittest/llmapi/test_serve_multi_frontend_route_gate.py": [ + "_llm_server_routes OpenAIServer", + ], +} + + +def _class_name(node) -> str: + if isinstance(node, ast.Name): + return node.id + if isinstance(node, ast.Attribute): + return node.attr + return "" + + +def _server_class_names(tree: ast.AST) -> dict[str, str]: + """Local name -> server class, with the names bound by ``import ... as``.""" + names = {server_class: server_class for server_class in SERVER_CLASSES} + for node in ast.walk(tree): + if isinstance(node, ast.ImportFrom): + for alias in node.names: + if alias.name in SERVER_CLASSES and alias.asname: + names[alias.asname] = alias.name + return names + + +def _built_server_class(node: ast.AST, names: dict[str, str]) -> str: + """The server class that a ``__new__`` call node builds, or "" for any other node.""" + if not isinstance(node, ast.Call) or not isinstance(node.func, ast.Attribute): + return "" + if node.func.attr != "__new__": + return "" + owner = node.func.value + if _class_name(owner) == "object": + if not node.args: + return "" + owner = node.args[0] + if isinstance(owner, ast.Name): + return names.get(owner.id, "") + if isinstance(owner, ast.Attribute) and owner.attr in SERVER_CLASSES: + return owner.attr + return "" + + +def find_new_sites(tree: ast.AST) -> list[str]: + """``" "`` per ``__new__`` call that builds one. + + The scope is the dotted chain of enclosing ``def``/``class`` names, or + ```` at module level; the list is sorted. + """ + names = _server_class_names(tree) + sites = [] + stack = [(tree, "")] + while stack: + node, scope = stack.pop() + server_class = _built_server_class(node, names) + if server_class: + sites.append(f"{scope or ''} {server_class}") + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)): + scope = f"{scope}.{node.name}" if scope else node.name + stack.extend((child, scope) for child in ast.iter_child_nodes(node)) + return sorted(sites) + + +def scan_tests(tests_root: Path = TESTS_ROOT) -> dict[str, list[str]]: + sites_by_file = {} + for path in sorted(tests_root.rglob("*.py")): + try: + tree = ast.parse(path.read_text(errors="replace")) + except SyntaxError: + continue + sites = find_new_sites(tree) + if sites: + sites_by_file[str(path.relative_to(REPO_ROOT))] = sites + return sites_by_file + + +def ratchet_problems( + sites_by_file: dict[str, list[str]], baseline: dict[str, list[str]] +) -> list[str]: + """Sites missing from ``baseline``, then baseline entries with no site left, per file.""" + problems = [] + for file in sorted(set(sites_by_file) | set(baseline)): + live = Counter(sites_by_file.get(file, [])) + recorded = Counter(baseline.get(file, [])) + for site in sorted((live - recorded).elements()): + problems.append( + f"{file}: '{site}' builds a server with object.__new__ and is not in the " + "baseline; use `isolated_serve_env` from " + "tests/test_common/serve/app_factory.py instead; see " + "tensorrt_llm/serve/AGENTS.md." + ) + for site in sorted((recorded - live).elements()): + problems.append( + f"{file}: '{site}' no longer builds a server with object.__new__; delete it " + f"from BASELINE[{file!r}] in tests/unittest/serve/test_no_object_new.py; the " + "baseline only shrinks." + ) + return problems + + +def test_no_new_object_new_sites(): + problems = ratchet_problems(scan_tests(), BASELINE) + assert not problems, "\n".join(problems) + + +# --------------------------------------------------------------------------- +# Self-test on synthetic sources: the scan counts calls, not text. +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "source,expected", + [ + ('"""object.__new__(OpenAIServer) in a docstring"""', 0), + ("# object.__new__(OpenAIServer) in a comment\nx = 1", 0), + ("server = object.__new__(OpenAIServer)", 1), + ("server = object.__new__(mod.OpenAIDisaggServer)", 1), + ("server = CoordinatorServer.__new__(CoordinatorServer)", 1), + ("thing = object.__new__(SomethingElse)", 0), + ("from m import OpenAIServer as Server\nserver = object.__new__(Server)", 1), + ("from m import OpenAIDisaggServer as S\nserver = S.__new__(S)", 1), + ("from m import SomethingElse as Server\nthing = object.__new__(Server)", 0), + ], +) +def test_scan_counts_call_nodes(source, expected): + assert len(find_new_sites(ast.parse(source))) == expected + + +def test_scan_names_sites_by_scope_and_class(): + source = ( + "from tensorrt_llm.serve.openai_server import OpenAIServer as Server\n" + "SERVER = object.__new__(OpenAIServer)\n" + "class TestServer:\n" + " def test_build(self):\n" + " return object.__new__(Server)\n" + ) + assert find_new_sites(ast.parse(source)) == [ + " OpenAIServer", + "TestServer.test_build OpenAIServer", + ] + + +def test_ratchet_flags_a_replaced_site(): + # Removing one site and adding one in another function keeps the file's + # total; the ratchet still reports both. + file = "tests/unittest/test_example.py" + before = find_new_sites(ast.parse("def test_old():\n object.__new__(OpenAIServer)\n")) + after = find_new_sites(ast.parse("def test_new():\n object.__new__(OpenAIServer)\n")) + assert ratchet_problems({file: after}, {file: after}) == [] + problems = ratchet_problems({file: after}, {file: before}) + assert len(problems) == 2, problems + assert "'test_new OpenAIServer' builds a server" in problems[0] + assert "'test_old OpenAIServer' no longer builds a server" in problems[1] diff --git a/tests/unittest/serve/test_openai_server_chat_core.py b/tests/unittest/serve/test_openai_server_chat_core.py new file mode 100644 index 000000000000..da96de12d1c0 --- /dev/null +++ b/tests/unittest/serve/test_openai_server_chat_core.py @@ -0,0 +1,353 @@ +# 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. +"""The chat completion handler's answers, pinned through the production code. + +The HTTP-level tests here run on ``OpenAIServer.openai_chat`` before and after +it is split into a typed core and an HTTP wrapper, so the split can show that +it changed no answer: a rejected request keeps its 400 body byte for byte, an +engine failure keeps today's HTTP 200 with body ``null`` (and +``/health_generate`` keeps turning that into a 500), a response that cannot be +JSON-encoded keeps its 400, and a stream keeps its frames and media type. + +Servers are built through the real constructor with the fake generator from +``tests/test_common/serve/``; only instance attributes are stubbed. +""" + +import asyncio +import json +import logging +import signal +from http import HTTPStatus +from types import SimpleNamespace +from unittest.mock import AsyncMock, Mock + +import pytest +from fastapi.responses import Response +from test_common.serve.app_factory import AggConfig +from test_common.serve.fakes import ScriptedOutput + +from tensorrt_llm.executor import CppExecutorError +from tensorrt_llm.logger import logger +from tensorrt_llm.serve import openai_server +from tensorrt_llm.serve._base.errors import EngineFatalError, ServingError +from tensorrt_llm.serve._base.request_context import RequestContext +from tensorrt_llm.serve.openai_protocol import ( + ChatCompletionRequest, + ChatCompletionResponse, + ChatCompletionResponseChoice, + ChatMessage, + UsageInfo, +) +from tensorrt_llm.serve.openai_server import OpenAIServer + +# threadleak is off: the tests create TestClients inside the test body, and the +# serve_env fixture closes them in its teardown, which runs after the leak check +# (precedent: tests/unittest/serve/api/ops/system/test_system_endpoints.py). +pytestmark = [pytest.mark.cpu_only, pytest.mark.threadleak(enabled=False)] + +MODEL = "fake-model" +KEEP = ( + "openai_chat must keep this answer (status and body); see the branches of " + "the HTTP wrapper openai_chat" +) +# The first rejection the handler itself makes (the request model already +# rejects a named tool_choice without tools): a named tool_choice on a +# pre-tokenized prompt. +NAMED_TOOL_ON_TOKEN_IDS_MESSAGE = ( + "tool_choice with a named function is not supported alongside " + "`prompt_token_ids`; pass `messages` instead." +) +_TOOL = { + "type": "function", + "function": {"name": "get_weather", "parameters": {"type": "object", "properties": {}}}, +} + + +def _chat_body(**overrides): + """A chat request whose prompt is pre-tokenized, so no rendering happens.""" + body = { + "model": MODEL, + "messages": [{"role": "user", "content": "hi"}], + "prompt_token_ids": [1, 2, 3], + } + body.update(overrides) + return body + + +def _named_tool_body(): + return _chat_body( + tools=[_TOOL], + tool_choice={"type": "function", "function": {"name": "get_weather"}}, + ) + + +def _chat_response(**choice_fields) -> ChatCompletionResponse: + return ChatCompletionResponse( + id="chatcmpl-chat-core-test", + created=1700000000, + model=MODEL, + choices=[ + ChatCompletionResponseChoice( + index=0, + message=ChatMessage(role="assistant", content="ok"), + finish_reason="stop", + **choice_fields, + ) + ], + usage=UsageInfo(prompt_tokens=3, completion_tokens=1, total_tokens=4), + ) + + +def _raw_request(): + """What the handler reads from a request when it is called directly.""" + return SimpleNamespace(state=SimpleNamespace(), headers={}) + + +class _StreamedGeneration: + """Engine output already post-processed into SSE frames. + + This is what the post-processing worker delivers: each step carries its + frames in ``outputs[0]._postprocess_result``. + """ + + def __init__(self, steps, *, then=None): + self._steps = list(steps) + self._then = then + + def __aiter__(self): + return self + + async def __anext__(self): + if self._steps: + frames = self._steps.pop(0) + return SimpleNamespace( + finished=not self._steps, + outputs=[SimpleNamespace(_postprocess_result=frames)], + candidate_metrics=None, + metrics_dict=None, + ) + if self._then is not None: + raise self._then + raise StopAsyncIteration + + +@pytest.fixture +def engine_failure_recorder(monkeypatch): + """Record the two side effects of an engine failure instead of running them.""" + recorder = SimpleNamespace(signals=[], terminations=[]) + monkeypatch.setattr(signal, "raise_signal", recorder.signals.append) + monkeypatch.setattr( + openai_server, "_record_generator_termination", recorder.terminations.append + ) + return recorder + + +def _build(serve_env, config=None) -> OpenAIServer: + server = serve_env.build_openai_server(config or AggConfig()) + # The disconnect poll needs a client socket; none of these tests has one. + server.await_disconnected = AsyncMock() + return server + + +# --------------------------------------------------------------------------- +# HTTP answers that the split must keep +# --------------------------------------------------------------------------- + + +def test_http_wrapper_keeps_every_answer(serve_env, engine_failure_recorder): + server = _build(serve_env, AggConfig(num_postprocess_workers=1)) + client = serve_env.client(server.app) + + # A request the handler rejects: 400 with today's body bytes. + expected = OpenAIServer.create_error_response( + message=NAMED_TOOL_ON_TOKEN_IDS_MESSAGE, + err_type="BadRequestError", + status_code=HTTPStatus.BAD_REQUEST, + ) + response = client.post("/v1/chat/completions", json=_named_tool_body()) + assert response.status_code == 400, f"rejected request: {KEEP}" + assert response.content == expected.body, f"rejected request: {KEEP}" + + # An engine failure: HTTP 200 with body null, after SIGINT. + server.generator.script(ScriptedOutput(raise_on_call=CppExecutorError("dead"))) + response = client.post("/v1/chat/completions", json=_chat_body()) + assert response.status_code == 200, f"engine failure: {KEEP}" + assert response.content == b"null", f"engine failure: {KEEP}" + assert engine_failure_recorder.signals == [signal.SIGINT] + assert len(engine_failure_recorder.terminations) == 1 + + # A non-streaming success: the typed response, serialized. + chat_response = _chat_response() + server._create_chat_response = AsyncMock(return_value=chat_response) + response = client.post("/v1/chat/completions", json=_chat_body()) + assert response.status_code == 200, f"non-streaming success: {KEEP}" + assert response.json() == chat_response.model_dump(), f"non-streaming success: {KEEP}" + + # A streaming success: the frames, then [DONE], as text/event-stream. + server.generator.generate_async = Mock( + return_value=_StreamedGeneration([["data: a\n\n"], ["data: b\n\n"]]) + ) + response = client.post("/v1/chat/completions", json=_chat_body(stream=True)) + assert response.status_code == 200, f"streaming success: {KEEP}" + assert response.headers["content-type"].startswith("text/event-stream"), ( + f"streaming success: {KEEP}" + ) + assert response.text == "data: a\n\ndata: b\n\ndata: [DONE]\n\n", f"streaming success: {KEEP}" + + +def test_unencodable_response_is_still_a_400(serve_env, error_records): + server = _build(serve_env) + server._create_chat_response = AsyncMock( + return_value=_chat_response(avg_decoded_tokens_per_iter=float("nan")) + ) + response = serve_env.client(server.app).post("/v1/chat/completions", json=_chat_body()) + assert response.status_code == 400, f"unencodable response: {KEEP}" + assert response.json()["message"].startswith( + "Out of range float values are not JSON compliant" + ), f"unencodable response: {KEEP}" + # The wrapper answers this from its own ValueError branch, without the + # traceback the generic handler would log. + assert error_records == [], f"unencodable response: {KEEP}" + + +def test_health_generate_still_fails_with_500_on_engine_failure(serve_env, engine_failure_recorder): + server = _build(serve_env) + + server.generator.script(ScriptedOutput(raise_on_call=CppExecutorError("dead"))) + response = asyncio.run(server.health_generate(_raw_request())) + assert response.status_code == 500, f"/health_generate on an engine failure: {KEEP}" + assert engine_failure_recorder.signals == [signal.SIGINT] + + server._create_chat_response = AsyncMock(return_value=_chat_response()) + response = asyncio.run(server.health_generate(_raw_request())) + assert response.status_code == 200, f"/health_generate on a success: {KEEP}" + assert response.body == b"Generation health check OK" + + +# --------------------------------------------------------------------------- +# The typed core: create_chat_completion +# --------------------------------------------------------------------------- + + +@pytest.fixture +def error_records(): + """Messages of the ERROR records the TRT-LLM logger emits during the test. + + The library logger does not propagate, so ``caplog`` would see nothing; a + handler is attached to it instead. The positive control that the handler + sees records is ``test_unexpected_error_is_a_serving_error_with_one_traceback``. + """ + records = [] + + class Capture(logging.Handler): + def emit(self, record): + records.append(record.getMessage()) + + handler = Capture(level=logging.ERROR) + logger._logger.addHandler(handler) + try: + yield records + finally: + logger._logger.removeHandler(handler) + + +def _request(**overrides) -> ChatCompletionRequest: + return ChatCompletionRequest(**_chat_body(**overrides)) + + +def _run_core(server, request): + return asyncio.run( + server.create_chat_completion(request, RequestContext.from_http(_raw_request())) + ) + + +def test_typed_result_for_a_non_streaming_request(serve_env): + server = _build(serve_env) + chat_response = _chat_response() + server._create_chat_response = AsyncMock(return_value=chat_response) + + result = _run_core(server, _request()) + + assert result is chat_response + assert not isinstance(result, Response) + + +def test_stream_is_returned_as_sse_strings(serve_env): + server = _build(serve_env, AggConfig(num_postprocess_workers=1)) + server.generator.generate_async = Mock( + return_value=_StreamedGeneration([["data: a\n\n"], ["data: b\n\n"]]) + ) + + async def run(): + result = await server.create_chat_completion( + _request(stream=True), RequestContext.from_http(_raw_request()) + ) + assert not isinstance(result, Response) + return [frame async for frame in result] + + frames = asyncio.run(run()) + + assert all(isinstance(frame, str) for frame in frames) + assert frames == ["data: a\n\n", "data: b\n\n", "data: [DONE]\n\n"] + + +def test_named_tool_choice_on_token_ids_raises_serving_error(serve_env, error_records): + server = _build(serve_env) + expected = OpenAIServer.create_error_response( + message=NAMED_TOOL_ON_TOKEN_IDS_MESSAGE, err_type="BadRequestError" + ) + + with pytest.raises(ServingError) as raised: + _run_core(server, ChatCompletionRequest(**_named_tool_body())) + + assert raised.value.status == 400 + assert raised.value.body.model_dump() == json.loads(expected.body) + assert error_records == [], "a rejected request is not an error of the server" + + +def test_value_error_is_a_serving_error_without_a_log(serve_env, error_records): + server = _build(serve_env) + server.generator.script(ScriptedOutput(raise_on_call=ValueError("bad"))) + + with pytest.raises(ServingError) as raised: + _run_core(server, _request()) + + assert raised.value.status == 400 + assert raised.value.body.message == "bad" + assert error_records == [] + + +def test_unexpected_error_is_a_serving_error_with_one_traceback(serve_env, error_records): + server = _build(serve_env) + server.generator.script(ScriptedOutput(raise_on_call=RuntimeError("boom"))) + + with pytest.raises(ServingError) as raised: + _run_core(server, _request()) + + assert raised.value.status == 400 + assert raised.value.body.message == "boom" + assert len(error_records) == 1, error_records + assert "RuntimeError: boom" in error_records[0] + + +def test_engine_failure_raises_engine_fatal_error_after_sigint(serve_env, engine_failure_recorder): + server = _build(serve_env) + server.generator.script(ScriptedOutput(raise_on_call=CppExecutorError("dead"))) + + with pytest.raises(EngineFatalError): + _run_core(server, _request()) + + assert engine_failure_recorder.signals == [signal.SIGINT] + assert len(engine_failure_recorder.terminations) == 1 diff --git a/tests/unittest/serve/test_serve_test_helpers.py b/tests/unittest/serve/test_serve_test_helpers.py new file mode 100644 index 000000000000..aa70512d9986 --- /dev/null +++ b/tests/unittest/serve/test_serve_test_helpers.py @@ -0,0 +1,404 @@ +# 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. +"""Self-tests of the serve test seam: the helpers other serve tests build on.""" + +import asyncio +import json +import os +from pathlib import Path + +import pytest +from fastapi import APIRouter, FastAPI +from prometheus_client import REGISTRY +from test_common.serve.app_factory import AggConfig, isolated_serve_env +from test_common.serve.cpu_tests import select_entries +from test_common.serve.fakes import ( + FakeLLM, + FakeLlmArgs, + ScriptedOutput, + ScriptedStep, + UserProbeMiddleware, +) +from test_common.serve.middleware import middleware_snapshot +from test_common.serve.routes import _records_conflict, compare, effective_routes + +import tensorrt_llm._utils +from tensorrt_llm.inputs.media_io import BaseMediaIO +from tensorrt_llm.sampling_params import SamplingParams +from tensorrt_llm.usage import usage_lib + +# threadleak is off: the scripted-generation tests create TestClients and run +# requests inside the test body, so the client's portal thread and the fake +# server's worker-pool threads are still alive when the leak check runs; the +# serve_env fixture closes them in its teardown, which runs after that check +# (precedent: tests/unittest/llmapi/apps/test_serve_lazy_import.py). +pytestmark = [pytest.mark.cpu_only, pytest.mark.threadleak(enabled=False)] + +REPO_ROOT = Path(__file__).resolve().parents[3] +_MISSING = object() + + +# --------------------------------------------------------------------------- +# effective_routes +# --------------------------------------------------------------------------- + + +def test_effective_routes_flattens_included_routers(): + app = FastAPI() + app.add_api_route("/direct", lambda: None, methods=["GET"]) + router = APIRouter() + router.add_api_route("/from_router", lambda: None, methods=["POST"]) + app.include_router(router) + prefixed = APIRouter() + prefixed.add_api_route("/suffix", lambda: None, methods=["GET"]) + app.include_router(prefixed, prefix="/pre") + paths = [path for path, _ in effective_routes(app)] + assert "/direct" in paths + assert "/from_router" in paths, ( + "effective_routes must walk into included routers; FastAPI 0.139 keeps them as " + "objects without a path attribute (duck-typed on original_router and " + "include_context in test_common/serve/routes.py)." + ) + assert "/pre/suffix" in paths, "an include_router prefix must be applied to the leaf paths" + + +def test_effective_routes_matches_app_routes_without_includes(): + app = FastAPI() + app.add_api_route("/only", lambda: None, methods=["GET"]) + assert [route for _, route in effective_routes(app)] == list(app.routes) + + +# --------------------------------------------------------------------------- +# compare and the conflict rule +# --------------------------------------------------------------------------- + + +def test_compare_reports_each_difference_kind(): + expected = [ + "GET /a _MsgspecRoute", + "GET /metrics _MsgspecRoute #get_iteration_stats", + "GET /metrics _MsgspecRoute #metrics", + "POST /v1/chat/completions _MsgspecRoute", + ] + actual = [ + "GET /metrics _MsgspecRoute #metrics", + "GET /metrics _MsgspecRoute #get_iteration_stats", + "POST /v1/chat/completions APIRoute", + "POST /new _MsgspecRoute", + ] + differences = compare(expected, actual) + assert any(d.startswith("added:") and "/new" in d for d in differences) + assert any(d.startswith("removed:") and "GET /a" in d for d in differences) + assert any( + "route class changed" in d + and "_MsgspecRoute -> APIRoute" in d + and "(a plain APIRouter breaks msgpack bodies from the disaggregated frontend)" in d + for d in differences + ) + assert any("order changed" in d and "#get_iteration_stats" in d for d in differences) + + +def test_compare_ignores_order_of_non_conflicting_routes(): + expected = ["GET /a APIRoute", "GET /b APIRoute"] + assert compare(expected, ["GET /b APIRoute", "GET /a APIRoute"]) == [] + assert compare(expected, list(expected)) == [] + + +# The first six rows use records of the real apps, the rest are synthetic. +@pytest.mark.parametrize( + "record_a,record_b,conflict", + [ + ( + "GET /metrics _MsgspecRoute #get_iteration_stats", + "GET /metrics _MsgspecRoute #metrics", + True, + ), + ( + "GET /v1/responses/{response_id} _MsgspecRoute", + "DELETE /v1/responses/{response_id} _MsgspecRoute", + False, + ), + ( + "GET /v1/messages/batches _MsgspecRoute", + "GET /v1/messages/batches/{batch_id} _MsgspecRoute", + False, + ), + ( + "GET /v1/images/{image_id}/content _MsgspecRoute", + "GET /v1/videos/{video_id}/content _MsgspecRoute", + False, + ), + ("MOUNT /prometheus/metrics Mount", "GET /metrics _MsgspecRoute", False), + ("GET,HEAD /docs Route", "GET,HEAD /docs/oauth2-redirect Route", False), + ("GET /v1/models _MsgspecRoute", "GET /v1/{name} APIRoute", True), + ("GET,HEAD /openapi.json Route", "GET /openapi.json APIRoute", True), + ("GET /files/{rest:path} APIRoute", "GET /files/a/b APIRoute", True), + ("MOUNT /prometheus/metrics Mount", "POST /prometheus/metrics/x APIRoute", True), + ], +) +def test_conflict_rule(record_a, record_b, conflict): + assert _records_conflict(record_a, record_b) is conflict + + +# --------------------------------------------------------------------------- +# middleware_snapshot format +# --------------------------------------------------------------------------- + + +def test_middleware_snapshot_format(): + from tensorrt_llm.serve.responses_utils import ServerArrivalTimeMiddleware + + app = FastAPI() + app.add_middleware(ServerArrivalTimeMiddleware, adjusted_clock=None) + app.add_middleware(UserProbeMiddleware) # added last, so outermost + snapshot = middleware_snapshot(app) + assert snapshot == [ + {"class": "UserProbeMiddleware"}, # not a tensorrt_llm class: name only + {"class": "ServerArrivalTimeMiddleware", "kwargs": ["adjusted_clock"]}, + ] + + +# --------------------------------------------------------------------------- +# isolated_serve_env +# --------------------------------------------------------------------------- + + +def test_isolated_serve_env_restores_global_state(tmp_path): + environ_before = dict(os.environ) + media_executor_before = BaseMediaIO._executor + prometheus_dir_before = getattr(tensorrt_llm._utils, "prometheus_multiproc_dir", _MISSING) + session_before = usage_lib._SESSION + collectors_before = set(REGISTRY._collector_to_names.keys()) + + with isolated_serve_env(tmp_path) as env: + env.build_openai_server(AggConfig(return_perf_metrics=True)) + os.environ["A2_RESTORE_SENTINEL"] = "1" + sink = env.record_terminations() + assert usage_lib._SESSION is not None and sink == [] + + assert dict(os.environ) == environ_before + assert BaseMediaIO._executor is media_executor_before + prometheus_dir_after = getattr(tensorrt_llm._utils, "prometheus_multiproc_dir", _MISSING) + if prometheus_dir_before is _MISSING: + assert prometheus_dir_after is _MISSING, ( + "isolated_serve_env must delete tensorrt_llm._utils.prometheus_multiproc_dir " + "when it did not exist before entry" + ) + else: + assert prometheus_dir_after is prometheus_dir_before + assert usage_lib._SESSION is session_before + assert set(REGISTRY._collector_to_names.keys()) == collectors_before, ( + "isolated_serve_env must unregister every Prometheus collector the build added" + ) + + +def test_isolated_serve_env_rejects_use_without_with(tmp_path): + with pytest.raises(AttributeError): + isolated_serve_env(tmp_path).build_openai_server(AggConfig()) + + +# --------------------------------------------------------------------------- +# cpu_tests selection +# --------------------------------------------------------------------------- + + +def test_cpu_tests_selection(): + kept, _ = select_entries() + targets = [target for entry in kept for target in entry.targets] + expected_new_files = [ + "unittest/api_stability/test_serve_routes.py", + "unittest/serve/test_agents_md.py", + "unittest/serve/test_layering.py", + "unittest/serve/test_middleware_order.py", + "unittest/serve/test_no_object_new.py", + "unittest/serve/test_serve_test_helpers.py", + "unittest/serve/test_shim_identity.py", + "unittest/serve/api/ops/system/test_system_endpoints.py", + "unittest/serve/api/visual_gen/test_ops_system_visual_gen.py", + "unittest/serve/api/visual_gen/test_routes_visual_gen.py", + "unittest/serve/api/visual_gen/test_shim_identity_visual_gen.py", + ] + for target in expected_new_files: + assert target in targets, f"the serve CPU runner must select {target}" + assert "unittest/llmapi/apps/test_anthropic_routes.py" in targets + for target in targets: + file_part = target.split("::")[0] + text = (REPO_ROOT / "tests" / file_part).read_text(errors="replace") + assert "pytest.mark.cpu_only" in text, ( + f"{file_part} is selected but carries no cpu_only marker; the CPU stage would " + "collect nothing from it" + ) + + +def test_cpu_tests_selection_keeps_node_ids(tmp_path): + # A node-id entry stays a node id: only that node runs, never the whole file. + node = "unittest/serve/test_serve_test_helpers.py::test_cpu_tests_selection" + cpu_list = tmp_path / "l0_cpu.yml" + cpu_list.write_text(f"l0_cpu:\n- tests:\n - {node}\n") + kept, notes = select_entries(cpu_list) + assert [entry.targets for entry in kept] == [(node,)], kept + assert not notes, notes + + +# --------------------------------------------------------------------------- +# Scripted generation through the production handlers +# --------------------------------------------------------------------------- + +_CHAT_BODY = { + "model": "fake-model", + "messages": [{"role": "user", "content": "hi"}], + "cache_salt": "salted", +} + + +def test_scripted_chat_non_streaming(serve_env): + server = serve_env.build_openai_server(AggConfig()) + client = serve_env.client(server.app) + response = client.post("/v1/chat/completions", json=_CHAT_BODY) + assert response.status_code == 200, response.text + body = response.json() + assert body["choices"][0]["message"]["content"] == "ok" # the default scripted text + generator = server.generator + assert len(generator.preprocess_calls) == 1 + call = generator.generate_calls[0] + assert isinstance(call.sampling_params, SamplingParams) + assert call.kwargs["streaming"] is False + assert call.kwargs["cache_salt"] == "salted" + + +def test_scripted_chat_streaming(serve_env): + server = serve_env.build_openai_server(AggConfig()) + server.generator.script( + ScriptedOutput( + steps=( + ScriptedStep(text="he", token_ids=(6,)), + ScriptedStep(text="llo", token_ids=(6,), finish_reason="stop"), + ) + ) + ) + client = serve_env.client(server.app) + with client.stream( + "POST", "/v1/chat/completions", json={**_CHAT_BODY, "stream": True} + ) as response: + assert response.status_code == 200 + lines = [line for line in response.iter_lines() if line.startswith("data: ")] + assert lines[-1] == "data: [DONE]" + deltas = [json.loads(line[len("data: ") :]) for line in lines[:-1]] + text = "".join( + choice["delta"].get("content") or "" for delta in deltas for choice in delta["choices"] + ) + assert text == "hello" + assert server.generator.generate_calls[0].kwargs["streaming"] is True + + +def test_scripted_completion_through_real_input_processor(serve_env): + server = serve_env.build_openai_server(AggConfig()) + server.generator.script( + ScriptedOutput(steps=(ScriptedStep(text="done", token_ids=(7,), finish_reason="stop"),)) + ) + client = serve_env.client(server.app) + response = client.post("/v1/completions", json={"model": "fake-model", "prompt": "hello hi"}) + assert response.status_code == 200, response.text + assert response.json()["choices"][0]["text"] == "done" + call = server.generator.generate_calls[0] + # A text prompt goes through the real DefaultInputProcessor, which + # tokenizes with the fake tokenizer: "hello hi" -> ids 6 and 5. + assert call.inputs == {"prompt_token_ids": [6, 5]} + + +def test_postprocess_worker_path_answers_through_postprocess_result(serve_env): + server = serve_env.build_openai_server(AggConfig(num_postprocess_workers=1)) + client = serve_env.client(server.app) + response = client.post("/v1/chat/completions", json=_CHAT_BODY) + assert response.status_code == 200, response.text + assert response.json()["choices"][0]["message"]["content"] == "ok" + promise = server.generator.generate_calls[0].result + assert promise.outputs[0]._postprocess_result is not None + with client.stream( + "POST", "/v1/chat/completions", json={**_CHAT_BODY, "stream": True} + ) as streaming_response: + lines = [line for line in streaming_response.iter_lines() if line.startswith("data: ")] + assert lines[-1] == "data: [DONE]" + assert any('"content"' in line for line in lines[:-1]) + + +def test_raise_on_call_returns_400_and_records_the_call(serve_env): + server = serve_env.build_openai_server(AggConfig()) + server.generator.script(ScriptedOutput(raise_on_call=RuntimeError("scripted failure"))) + client = serve_env.client(server.app) + response = client.post("/v1/chat/completions", json=_CHAT_BODY) + assert response.status_code == 400 + assert response.json()["message"] == "scripted failure" + call = server.generator.generate_calls[0] + assert call.result is None # recorded before the scripted exception fired + assert isinstance(call.sampling_params, SamplingParams) + + +def test_raise_after_steps_ends_a_stream_with_the_scripted_error(serve_env): + server = serve_env.build_openai_server(AggConfig()) + server.generator.script( + ScriptedOutput( + steps=(ScriptedStep(text="partial", token_ids=(5,)),), + raise_after_steps=RuntimeError("stream exploded"), + ) + ) + client = serve_env.client(server.app) + received = [] + with pytest.raises(RuntimeError, match="stream exploded"): + with client.stream( + "POST", "/v1/chat/completions", json={**_CHAT_BODY, "stream": True} + ) as response: + for line in response.iter_lines(): + received.append(line) + assert any("partial" in line for line in received), ( + "the first scripted chunk must reach the client before the scripted error" + ) + + +def test_hold_after_steps_release_and_cancellation(): + async def scenario(): + llm = FakeLLM(args=FakeLlmArgs(model="unused")) + llm.script( + ScriptedOutput( + steps=(ScriptedStep(text="first", token_ids=(5,)),), hold_after_steps=True + ) + ) + promise = llm.generate_async({"prompt_token_ids": [1, 2]}, SamplingParams()) + first = await promise.__anext__() + assert first.outputs[0].text == "first" + assert not promise.finished + + async def consume_rest(): + async for _ in promise: + pass + + consumer = asyncio.create_task(consume_rest()) + await asyncio.sleep(0) # the consumer reaches the hold + promise.release() + await consumer # release ends the wait with a normal end of stream + assert promise.finished + + llm.script( + ScriptedOutput(steps=(ScriptedStep(text="x", token_ids=(5,)),), hold_after_steps=True) + ) + held = llm.generate_async({"prompt_token_ids": [1]}, SamplingParams()) + await held.__anext__() + hanging = asyncio.create_task(held.__anext__()) + await asyncio.sleep(0) + hanging.cancel() # a client disconnect cancels the consuming task + with pytest.raises(asyncio.CancelledError): + await hanging + + asyncio.run(scenario()) diff --git a/tests/unittest/serve/test_shim_identity.py b/tests/unittest/serve/test_shim_identity.py new file mode 100644 index 000000000000..2843e2bb1353 --- /dev/null +++ b/tests/unittest/serve/test_shim_identity.py @@ -0,0 +1,245 @@ +# 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. +"""Identity, FutureWarning and old-path-import checks of the kept serve shims. + +``SHIM_ROWS`` starts empty: only the pull requests that keep an old import path +as a shim (old paths that code outside the repository imports) add rows here. +Each row then checks that the old name is the same object as the new one, that +the shim warns with the expected ``FutureWarning``, and that nothing inside the +repository imports the old path. The self-tests below prove the three checks +can fail, so the empty table is not vacuous. +""" + +import sys +import textwrap +from pathlib import Path + +import pytest +from test_common.serve.shims import ShimRow, check_identity, check_warnings, find_old_path_imports + +pytestmark = pytest.mark.cpu_only + +REPO_ROOT = Path(__file__).resolve().parents[3] + +# Rows are added only by the pull requests that create the kept shims. +SHIM_ROWS: tuple = () + + +def test_shim_identity(): + assert check_identity(SHIM_ROWS) == [] + + +def test_shim_warnings(): + assert check_warnings(SHIM_ROWS) == [] + + +def test_no_old_path_imports(): + assert find_old_path_imports(SHIM_ROWS, REPO_ROOT) == [] + + +# --------------------------------------------------------------------------- +# Self-test on a synthetic package: each check must be able to fail. +# --------------------------------------------------------------------------- + + +@pytest.fixture +def shim_package(tmp_path, monkeypatch): + """A synthetic package with good, silent, copying, lazy and interface shims.""" + pkg = "shim_selftest_pkg" + root = tmp_path / pkg + root.mkdir() + (root / "__init__.py").write_text("") + (root / "new.py").write_text( + textwrap.dedent( + """ + class X: + pass + + TABLE = {"a": 1} + + def main(): + pass + """ + ) + ) + (root / "good.py").write_text( + textwrap.dedent( + f""" + import warnings + + from {pkg}.new import TABLE, X + + warnings.warn( + "{pkg}.good has moved to {pkg}.new and will be removed in 1.6.0.", + FutureWarning, + stacklevel=2, + ) + __all__ = ["TABLE", "X"] + """ + ) + ) + (root / "silent.py").write_text(f"from {pkg}.new import X\n__all__ = ['X']\n") + (root / "bad.py").write_text("class X:\n pass\n__all__ = ['X']\n") + (root / "lazy.py").write_text( + textwrap.dedent( + f""" + import warnings + + from {pkg}.new import main + + _LAZY_NAMES = {{"X": "{pkg}.new"}} + + def __getattr__(name): + if name in _LAZY_NAMES: + warnings.warn( + f"{pkg}.lazy.{{name}} has moved to {pkg}.new " + "and will be removed in 1.6.0.", + FutureWarning, + stacklevel=2, + ) + import importlib + + value = getattr(importlib.import_module(_LAZY_NAMES[name]), name) + globals()[name] = value # warn once: later reads skip __getattr__ + return value + raise AttributeError(name) + """ + ) + ) + (root / "iface.py").write_text( + textwrap.dedent( + f""" + def __getattr__(name): + if name == "X": + from {pkg}.new import X + + return X + raise AttributeError(name) + """ + ) + ) + sub = root / "sub" + sub.mkdir() + (sub / "__init__.py").write_text( + textwrap.dedent( + f""" + import warnings + + from {pkg}.new import X + + warnings.warn( + "{pkg}.sub has moved to {pkg}.new and will be removed in 1.6.0.", + FutureWarning, + stacklevel=2, + ) + __all__ = ["X"] + """ + ) + ) + (sub / "mod.py").write_text( + textwrap.dedent( + f""" + import warnings + + from {pkg}.new import X + + warnings.warn( + "{pkg}.sub.mod has moved to {pkg}.new and will be removed in 1.6.0.", + FutureWarning, + stacklevel=2, + ) + __all__ = ["X"] + """ + ) + ) + (root / "importer.py").write_text(f"from {pkg}.good import X\n") + (root / "module_user.py").write_text( + f"from {pkg} import good\n\n\ndef use():\n return good.X\n" + ) + (root / "other_name_user.py").write_text( + f"from {pkg} import good\n\n\ndef use():\n return good.TABLE\n" + ) + monkeypatch.syspath_prepend(str(tmp_path)) + yield pkg, tmp_path + for name in list(sys.modules): + if name == pkg or name.startswith(pkg + "."): + del sys.modules[name] + + +def _row(pkg, module, name, **kwargs): + retire_in = kwargs.pop("retire_in", "1.6.0") + return ShimRow(f"{pkg}.{module}", name, f"{pkg}.new", retire_in=retire_in, **kwargs) + + +def test_identity_check_passes_and_fails(shim_package): + pkg, _ = shim_package + good_rows = [_row(pkg, "good", "X"), _row(pkg, "good", "TABLE")] + assert check_identity(good_rows) == [] + problems = check_identity([_row(pkg, "bad", "X")]) + assert len(problems) == 1 and "not the same object" in problems[0] + problems = check_identity([_row(pkg, "good", "GONE")]) + assert len(problems) == 1 and "no longer exists" in problems[0] + + +def test_warning_check_passes_for_shim_styles(shim_package): + pkg, _ = shim_package + rows = [ + _row(pkg, "good", "X"), + _row(pkg, "good", "TABLE"), + _row(pkg, "lazy", "X"), # warns on the first read only + _row(pkg, "lazy", "main"), # eager name: never warns + _row(pkg, "iface", "X", warns=False, ban_old_path=False, retire_in="not retired"), + # A fresh import of a shim submodule also imports (and warns for) its + # parent shim package; the message filter must not count that warning. + _row(pkg, "sub.mod", "X"), + ] + assert f"{pkg}.sub" not in sys.modules + assert check_warnings(rows) == [] + + +def test_warning_check_reports_silent_and_wrong_release(shim_package): + pkg, _ = shim_package + problems = check_warnings([_row(pkg, "silent", "X")]) + assert len(problems) == 1 and "emitted no FutureWarning" in problems[0] + wrong_release = ShimRow(f"{pkg}.good", "X", f"{pkg}.new", retire_in="9.9.9") + problems = check_warnings([wrong_release]) + assert len(problems) == 1 and "names the new location and the removal release" in problems[0] + + +def test_warning_check_restores_sys_modules(shim_package): + pkg, _ = shim_package + import importlib + + before = importlib.import_module(f"{pkg}.good") + assert check_warnings([_row(pkg, "good", "X")]) == [] + assert sys.modules[f"{pkg}.good"] is before + assert getattr(sys.modules[pkg], "good") is before + + +def test_old_path_import_scan(shim_package): + pkg, tmp_path = shim_package + rows = [_row(pkg, "good", "X")] + problems = find_old_path_imports(rows, tmp_path, roots=(pkg,)) + importing_lines = [p for p in problems if "importer.py" in p] + module_attribute_lines = [p for p in problems if "module_user.py" in p] + other_name_lines = [p for p in problems if "other_name_user.py" in p] + assert len(importing_lines) == 1, problems + assert f"import it from `{pkg}.new`" in importing_lines[0] + # `from import good` + `good.X` loads the shim as surely as a from-import. + assert len(module_attribute_lines) == 1, problems + # Reading another name of the module is not banned by this row. + assert other_name_lines == [], problems + # The shim's own file is exempt. + assert not any("good.py" in p for p in problems), problems