From b291cca0c4ee996340250e68e3cac5cc24705315 Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Wed, 28 Oct 2026 02:34:11 +0000 Subject: [PATCH 01/26] [TRTLLM-17261][test] Add route, middleware-order and shim-identity snapshots and a CPU test seam for trtllm-serve Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- tensorrt_llm/commands/AGENTS.md | 6 + tensorrt_llm/commands/CLAUDE.md | 2 + tensorrt_llm/serve/AGENTS.md | 76 +++ tensorrt_llm/serve/CLAUDE.md | 2 + .../integration/test_lists/test-db/l0_cpu.yml | 5 + tests/test_common/serve/__init__.py | 14 + tests/test_common/serve/app_factory.py | 352 +++++++++++++ tests/test_common/serve/cpu_tests.py | 183 +++++++ tests/test_common/serve/fakes.py | 493 ++++++++++++++++++ tests/test_common/serve/fakes_visual_gen.py | 44 ++ tests/test_common/serve/layering.py | 218 ++++++++ tests/test_common/serve/middleware.py | 42 ++ tests/test_common/serve/routes.py | 477 +++++++++++++++++ tests/test_common/serve/shims.py | 376 +++++++++++++ .../api_stability/test_serve_routes.py | 182 +++++++ tests/unittest/serve/__init__.py | 14 + tests/unittest/serve/api/__init__.py | 14 + tests/unittest/serve/api/ops/__init__.py | 14 + .../unittest/serve/api/ops/system/__init__.py | 14 + .../api/ops/system/test_system_endpoints.py | 132 +++++ .../unittest/serve/api/visual_gen/__init__.py | 14 + .../visual_gen/test_ops_system_visual_gen.py | 50 ++ .../api/visual_gen/test_routes_visual_gen.py | 56 ++ .../test_shim_identity_visual_gen.py | 46 ++ tests/unittest/serve/conftest.py | 26 + tests/unittest/serve/test_agents_md.py | 132 +++++ tests/unittest/serve/test_layering.py | 193 +++++++ tests/unittest/serve/test_middleware_order.py | 133 +++++ tests/unittest/serve/test_no_object_new.py | 146 ++++++ .../unittest/serve/test_serve_test_helpers.py | 393 ++++++++++++++ tests/unittest/serve/test_shim_identity.py | 245 +++++++++ 31 files changed, 4094 insertions(+) create mode 100644 tensorrt_llm/commands/AGENTS.md create mode 100644 tensorrt_llm/commands/CLAUDE.md create mode 100644 tensorrt_llm/serve/AGENTS.md create mode 100644 tensorrt_llm/serve/CLAUDE.md create mode 100644 tests/test_common/serve/__init__.py create mode 100644 tests/test_common/serve/app_factory.py create mode 100644 tests/test_common/serve/cpu_tests.py create mode 100644 tests/test_common/serve/fakes.py create mode 100644 tests/test_common/serve/fakes_visual_gen.py create mode 100644 tests/test_common/serve/layering.py create mode 100644 tests/test_common/serve/middleware.py create mode 100644 tests/test_common/serve/routes.py create mode 100644 tests/test_common/serve/shims.py create mode 100644 tests/unittest/api_stability/test_serve_routes.py create mode 100644 tests/unittest/serve/__init__.py create mode 100644 tests/unittest/serve/api/__init__.py create mode 100644 tests/unittest/serve/api/ops/__init__.py create mode 100644 tests/unittest/serve/api/ops/system/__init__.py create mode 100644 tests/unittest/serve/api/ops/system/test_system_endpoints.py create mode 100644 tests/unittest/serve/api/visual_gen/__init__.py create mode 100644 tests/unittest/serve/api/visual_gen/test_ops_system_visual_gen.py create mode 100644 tests/unittest/serve/api/visual_gen/test_routes_visual_gen.py create mode 100644 tests/unittest/serve/api/visual_gen/test_shim_identity_visual_gen.py create mode 100644 tests/unittest/serve/conftest.py create mode 100644 tests/unittest/serve/test_agents_md.py create mode 100644 tests/unittest/serve/test_layering.py create mode 100644 tests/unittest/serve/test_middleware_order.py create mode 100644 tests/unittest/serve/test_no_object_new.py create mode 100644 tests/unittest/serve/test_serve_test_helpers.py create mode 100644 tests/unittest/serve/test_shim_identity.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..fa82b67fa4c6 --- /dev/null +++ b/tensorrt_llm/serve/AGENTS.md @@ -0,0 +1,76 @@ + + +# AGENTS.md - tensorrt_llm/serve + +Scope: `tensorrt_llm/serve/` and everything under it. Supplements the repo-root `AGENTS.md` +and `tests/AGENTS.md`; `tensorrt_llm/commands/AGENTS.md` points here for the serve CLI. +Where a check enforces a rule, the rule names it; a failing check says what to change. +Name files and symbols here, never line numbers. + +## Where code goes + +- Model-specific behavior: a hook on `ServingExtension` (`serve/serving_extensions.py`, + `serve/extensions/`). Do not branch on a model name in a handler. +- Prompt rendering: `serve/render/` only. Handlers do not call `apply_chat_template`. +- Request and response models: a new field needs a row in + `tests/unittest/api_stability/references/trtllm_serve_api.yaml`. + Check: `tests/unittest/api_stability/test_serve_api.py`. +- Routes: every server configuration's routes (path, method, route class, and the order of + routes that can match the same request) are pinned in + `tests/unittest/api_stability/references/trtllm_serve_routes.yaml`. A route change + regenerates it and needs the API review committee. + Check: `tests/unittest/api_stability/test_serve_routes.py`. +- Middleware order matters: `add_middleware` inserts at the front, so the last one added is + outermost. Do not reorder. Check: `tests/unittest/serve/test_middleware_order.py`. +- Import directions between layers are tests; a PR that creates a package adds its contract. + Check: `tests/unittest/serve/test_layering.py`. +- VisualGen serve tests live under a path that contains `visual_gen/` + (`tests/unittest/serve/api/visual_gen/`), or CI selection misses them. + +## Moving code + +- A move PR only moves code: no reformatting, no logic change. +- In the same PR update the multi-GPU file list (`getMultiGpuFileChanged` in + `jenkins/L0_MergeRequest.groovy`), `.github/CODEOWNERS` and `legacy-files.txt`. A new file + takes the formatting group of the file its code comes from. +- Old import paths: a move deletes the old module (or the moved names) without a shim and + rewrites every import of it in the repository, tests included; module paths written as + strings (patch targets, subprocess command lines, `_LAZY_ATTRS` in `serve/__init__.py`) are + fixed by hand. The PR body lists "old path → new path" for the release notes. Only five old + paths keep a shim, because code outside the repository imports them: + `tensorrt_llm/llmapi/reasoning_parser.py`, `serve/tool_parser/`, + `serve/openai_disagg_server.py`, `serve/openai_server.py` and `commands/serve.py`. A shim + re-exports public names only, lists them in `__all__` (or `_LAZY_NAMES` in + `commands/serve.py`), warns once with `FutureWarning`, and has its rows in + `tests/unittest/serve/test_shim_identity.py`. +- `mock.patch("old.module.name")` does not follow moved code. Grep the tests for every moved + name, including `monkeypatch.setattr(module, "name", ...)` and dotted paths split across lines. + +## Tests + +- Register every new test file in a CI list in the same PR (`tests/AGENTS.md`). Serve tests + that need no GPU go in `tests/integration/test_lists/test-db/l0_cpu.yml` and must contain + `pytest.mark.cpu_only`, or the CPU stage collects nothing from them. +- New serve tests mirror the package: `tests/unittest/serve/` for `tensorrt_llm/serve/`. +- Build servers with the factory and fake generators in `tests/test_common/serve/` + (`isolated_serve_env`), never `object.__new__(OpenAIServer)` or + `object.__new__(OpenAIDisaggServer)`. Check: `tests/unittest/serve/test_no_object_new.py`. +- List an app's routes with `effective_routes()` from `tests/test_common/serve/routes.py`, not + `app.routes`: FastAPI 0.139 keeps included routers as objects without `.path`. +- The harmony tests guard whole files with try-imports. After renaming anything they import, + run them with `pytest -rs` and check that nothing is skipped as "not importable": a green run + can mean every test was skipped. + +## Commands + +- All serve CPU tests: `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 ""` +- 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/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..aeac2ed4cfb2 --- /dev/null +++ b/tests/test_common/serve/app_factory.py @@ -0,0 +1,352 @@ +# 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 + + +@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) + + 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, + ) + 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..20472ac446da --- /dev/null +++ b/tests/test_common/serve/layering.py @@ -0,0 +1,218 @@ +# 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.""" + 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 + if not any(_matches_prefix(edge.imported, prefix) for prefix in contract.forbidden): + continue + for triple in _edge_triples(edge): + if triple in contract.allowed: + continue + importer, imported, name = triple + what = f"`{imported}.{name}`" if name else f"`{imported}`" + problems.append( + f"contract '{contract.name}': {importer} (line {edge.lineno}) imports " + f"{what}. {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..84ae3d719887 --- /dev/null +++ b/tests/test_common/serve/routes.py @@ -0,0 +1,477 @@ +# 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"),)), + ), + 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] + class_changes.append( + f"route class changed: {_methods_path(removed_record)}: {old_cls} -> {new_cls}" + ) + 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/). 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/test_serve_routes.py b/tests/unittest/api_stability/test_serve_routes.py new file mode 100644 index 000000000000..78d253183a15 --- /dev/null +++ b/tests/unittest/api_stability/test_serve_routes.py @@ -0,0 +1,182 @@ +# 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." + ) + + +def test_disagg_and_coordinator_routes_are_plain(live_snapshots): + for name, snapshot in live_snapshots.items(): + if snapshot.get("error") or snapshot["kind"] == "openai_server": + continue + for record in snapshot["routes"]: + cls = _parse(record)[2] + assert cls in ("APIRoute", "Route", "Mount"), ( + f"configuration '{name}': {record}: the disaggregated frontend and the " + "coordinator app register plain routes; they do not take the aggregated " + "server's msgpack route class." + ) + + +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_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." + ) + + +def test_error_configurations(tmp_path): + expected = { + "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" + ), + } + by_name = {config.name: config for config in CONFIGS} + for name, message in expected.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/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/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/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/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_visual_gen.py b/tests/unittest/serve/api/visual_gen/test_ops_system_visual_gen.py new file mode 100644 index 000000000000..d9b55a9b2624 --- /dev/null +++ b/tests/unittest/serve/api/visual_gen/test_ops_system_visual_gen.py @@ -0,0 +1,50 @@ +# 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): + body = visual_gen_client.get("/v1/models").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..42327c1202d3 --- /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=") + +_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..c547dd6d80c8 --- /dev/null +++ b/tests/unittest/serve/test_layering.py @@ -0,0 +1,193 @@ +# 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 textwrap +from pathlib import Path + +import pytest +from test_common.serve.layering import ( + Contract, + 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", + ) + ), + 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." + ), + ), +) + + +@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) == [] + + +# --------------------------------------------------------------------------- +# 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_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..9c72a1cfb62b --- /dev/null +++ b/tests/unittest/serve/test_no_object_new.py @@ -0,0 +1,146 @@ +# 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 that may only shrink. + +The scan counts AST call nodes, not text, so a docstring or comment that +mentions ``object.__new__(OpenAIServer)`` does not count. +""" + +import ast +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 counts of the sites that existed when this ratchet landed. +# The baseline only shrinks: rewrite a site on the factory, lower its count. +BASELINE = { + "tests/unittest/disaggregated/test_coordinator_worker.py": 2, + "tests/unittest/disaggregated/test_disagg_internal_auth.py": 3, + "tests/unittest/disaggregated/test_disagg_openai_client.py": 1, + "tests/unittest/disaggregated/test_disagg_responses_api.py": 2, + "tests/unittest/disaggregated/test_openai_disagg_server.py": 3, + "tests/unittest/disaggregated/test_openai_disagg_service.py": 2, + "tests/unittest/disaggregated/test_openai_server_info.py": 1, + "tests/unittest/disaggregated/test_openai_server_stream_errors.py": 1, + "tests/unittest/disaggregated/test_router.py": 2, + "tests/unittest/executor/test_spec_decode_stats_payload.py": 1, + "tests/unittest/llmapi/apps/_test_trtllm_serve_profile.py": 1, + "tests/unittest/llmapi/apps/test_anthropic_routes.py": 13, + "tests/unittest/llmapi/apps/test_context_length_error.py": 1, + "tests/unittest/llmapi/apps/test_harmony_parsing.py": 1, + "tests/unittest/llmapi/apps/test_openai_server_iteration_stats.py": 1, + "tests/unittest/llmapi/apps/test_render_callers.py": 2, + "tests/unittest/llmapi/apps/test_responses_store_replay.py": 1, + "tests/unittest/llmapi/apps/test_responses_streaming_events.py": 1, + "tests/unittest/llmapi/apps/test_serving_extensions.py": 1, + "tests/unittest/llmapi/test_llm.py": 2, + "tests/unittest/llmapi/test_rl_control_auth.py": 1, + "tests/unittest/llmapi/test_runtime_control.py": 1, +} + + +def _class_name(node) -> str: + if isinstance(node, ast.Name): + return node.id + if isinstance(node, ast.Attribute): + return node.attr + return "" + + +def count_new_sites(tree: ast.AST) -> int: + """Call nodes that build a server class with ``__new__``.""" + sites = 0 + for node in ast.walk(tree): + if not isinstance(node, ast.Call) or not isinstance(node.func, ast.Attribute): + continue + if node.func.attr != "__new__": + continue + owner = _class_name(node.func.value) + if owner == "object": + if node.args and _class_name(node.args[0]) in SERVER_CLASSES: + sites += 1 + elif owner in SERVER_CLASSES: + sites += 1 + return sites + + +def scan_tests(tests_root: Path = TESTS_ROOT): + counts = {} + for path in sorted(tests_root.rglob("*.py")): + try: + tree = ast.parse(path.read_text(errors="replace")) + except SyntaxError: + continue + sites = count_new_sites(tree) + if sites: + counts[str(path.relative_to(REPO_ROOT))] = sites + return counts + + +def test_no_new_object_new_sites(): + counts = scan_tests() + problems = [] + for file, sites in sorted(counts.items()): + baseline = BASELINE.get(file, 0) + if sites > baseline: + problems.append( + f"{file} builds a server with object.__new__ ({sites} sites, baseline " + f"{baseline}); use `isolated_serve_env` from " + "tests/test_common/serve/app_factory.py instead; see " + "tensorrt_llm/serve/AGENTS.md." + ) + for file, baseline in sorted(BASELINE.items()): + sites = counts.get(file, 0) + if sites < baseline: + problems.append( + f"{file} now has {sites} object.__new__ sites; lower BASELINE[{file!r}] to " + f"{sites} in tests/unittest/serve/test_no_object_new.py; the baseline only " + "shrinks." + ) + 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), + ], +) +def test_scan_counts_call_nodes(source, expected): + assert count_new_sites(ast.parse(source)) == expected 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..8d7a32fa98f5 --- /dev/null +++ b/tests/unittest/serve/test_serve_test_helpers.py @@ -0,0 +1,393 @@ +# 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 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 + # A node-id entry stays a node id: only that node runs, never the whole file. + harmony = "unittest/llmapi/apps/test_harmony_parsing.py::TestStripIncompleteMessagesReporting" + assert harmony in targets + assert "unittest/llmapi/apps/test_harmony_parsing.py" not 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" + ) + + +# --------------------------------------------------------------------------- +# 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 From c9311abec8bfcd4f9b9642c196640385e28e3238 Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Sun, 1 Nov 2026 15:58:05 +0000 Subject: [PATCH 02/26] [TRTLLM-17261][test] Record the trtllm-serve route reference and test baselines at the merge base The route reference is generated from the live apps; the object.__new__ baseline and the gRPC import allowance record what exists on main today and only shrink. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- .../references/trtllm_serve_routes.yaml | 619 ++++++++++++++++++ tests/unittest/serve/test_layering.py | 33 +- tests/unittest/serve/test_no_object_new.py | 4 + 3 files changed, 644 insertions(+), 12 deletions(-) create mode 100644 tests/unittest/api_stability/references/trtllm_serve_routes.yaml 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..a744e75b8686 --- /dev/null +++ b/tests/unittest/api_stability/references/trtllm_serve_routes.yaml @@ -0,0 +1,619 @@ +# 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 + 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 APIRoute + - POST /v1/chat/completions APIRoute + - POST /v1/messages APIRoute + - POST /v1/messages/count_tokens APIRoute + - POST /v1/responses APIRoute + - GET /v1/models APIRoute + - GET /health APIRoute + - GET /cluster_info APIRoute + - GET /version APIRoute + - 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 APIRoute + - POST /v1/chat/completions APIRoute + - POST /v1/messages APIRoute + - POST /v1/messages/count_tokens APIRoute + - POST /v1/responses APIRoute + - GET /v1/models APIRoute + - GET /health APIRoute + - GET /cluster_info APIRoute + - GET /version APIRoute + - 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 APIRoute + - POST /v1/chat/completions APIRoute + - POST /v1/messages APIRoute + - POST /v1/messages/count_tokens APIRoute + - POST /v1/responses APIRoute + - GET /v1/models APIRoute + - GET /health APIRoute + - GET /cluster_info APIRoute + - GET /version APIRoute + - MOUNT /prometheus/metrics Mount + - POST /set APIRoute + - GET /get APIRoute + - DELETE /delete APIRoute + - GET /expire APIRoute + - GET /get_prefix APIRoute + 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 APIRoute + - POST /v1/chat/completions APIRoute + - POST /v1/messages APIRoute + - POST /v1/messages/count_tokens APIRoute + - POST /v1/responses APIRoute + - GET /v1/models APIRoute + - GET /health APIRoute + - GET /cluster_info APIRoute + - GET /version APIRoute + - 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 APIRoute + - POST /v1/chat/completions APIRoute + - POST /v1/messages APIRoute + - POST /v1/messages/count_tokens APIRoute + - POST /v1/responses APIRoute + - GET /v1/models APIRoute + - GET /health APIRoute + - GET /cluster_info APIRoute + - GET /version APIRoute + - 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 APIRoute + - POST /v1/chat/completions APIRoute + - POST /v1/messages APIRoute + - POST /v1/messages/count_tokens APIRoute + - POST /v1/responses APIRoute + - GET /v1/models APIRoute + - GET /health APIRoute + - GET /cluster_info APIRoute + - GET /version APIRoute + - MOUNT /prometheus/metrics Mount + - POST /set APIRoute + - GET /get APIRoute + - DELETE /delete APIRoute + - GET /expire APIRoute + - GET /get_prefix APIRoute + 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/serve/test_layering.py b/tests/unittest/serve/test_layering.py index c547dd6d80c8..1b3c44427000 100644 --- a/tests/unittest/serve/test_layering.py +++ b/tests/unittest/serve/test_layering.py @@ -54,18 +54,27 @@ 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._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 " diff --git a/tests/unittest/serve/test_no_object_new.py b/tests/unittest/serve/test_no_object_new.py index 9c72a1cfb62b..a883a28c496a 100644 --- a/tests/unittest/serve/test_no_object_new.py +++ b/tests/unittest/serve/test_no_object_new.py @@ -48,6 +48,7 @@ "tests/unittest/disaggregated/test_openai_disagg_service.py": 2, "tests/unittest/disaggregated/test_openai_server_info.py": 1, "tests/unittest/disaggregated/test_openai_server_stream_errors.py": 1, + "tests/unittest/disaggregated/test_paired_request_cancellation.py": 2, "tests/unittest/disaggregated/test_router.py": 2, "tests/unittest/executor/test_spec_decode_stats_payload.py": 1, "tests/unittest/llmapi/apps/_test_trtllm_serve_profile.py": 1, @@ -59,9 +60,12 @@ "tests/unittest/llmapi/apps/test_responses_store_replay.py": 1, "tests/unittest/llmapi/apps/test_responses_streaming_events.py": 1, "tests/unittest/llmapi/apps/test_serving_extensions.py": 1, + "tests/unittest/llmapi/apps/test_tool_parsers.py": 1, "tests/unittest/llmapi/test_llm.py": 2, "tests/unittest/llmapi/test_rl_control_auth.py": 1, "tests/unittest/llmapi/test_runtime_control.py": 1, + "tests/unittest/llmapi/test_serve_multi_frontend_helpers.py": 4, + "tests/unittest/llmapi/test_serve_multi_frontend_route_gate.py": 1, } From 81b035082b4157aeaefe2f65f5f5a0e06bc7200c Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Sun, 1 Nov 2026 16:56:07 +0000 Subject: [PATCH 03/26] [TRTLLM-17261][test] Cover the attached-frontend route registration in the snapshot matrix An attached frontend of a multi-frontend server registers ten launcher forwarding routes before the local table; seven share a method and path with local handlers and three use a path parameter that matches every batch id, so registration order decides which handler serves them. The new configuration records all of it, a live invariant pins the forwarding-before-local order independently of the reference, and the route-class failure message now says why the msgpack route class matters. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- tests/test_common/serve/app_factory.py | 14 ++++++ tests/test_common/serve/routes.py | 16 ++++++- .../references/trtllm_serve_routes.yaml | 44 +++++++++++++++++++ .../api_stability/test_serve_routes.py | 35 +++++++++++++++ .../unittest/serve/test_serve_test_helpers.py | 7 ++- 5 files changed, 114 insertions(+), 2 deletions(-) diff --git a/tests/test_common/serve/app_factory.py b/tests/test_common/serve/app_factory.py index aeac2ed4cfb2..1e958e6b44f7 100644 --- a/tests/test_common/serve/app_factory.py +++ b/tests/test_common/serve/app_factory.py @@ -107,6 +107,9 @@ class AggConfig: 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) @@ -221,6 +224,16 @@ def build_openai_server(self, config: AggConfig) -> OpenAIServer: 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, @@ -231,6 +244,7 @@ def build_openai_server(self, config: AggConfig) -> OpenAIServer: 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 diff --git a/tests/test_common/serve/routes.py b/tests/test_common/serve/routes.py index 84ae3d719887..0c983b8022a6 100644 --- a/tests/test_common/serve/routes.py +++ b/tests/test_common/serve/routes.py @@ -187,6 +187,14 @@ class SnapshotConfig: "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", @@ -336,8 +344,14 @@ def _methods_path(record: str) -> str: 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)}: {old_cls} -> {new_cls}" + f"route class changed: {_methods_path(removed_record)}: " + f"{old_cls} -> {new_cls}{why}" ) removed.remove(removed_record) remaining_added.remove(added_record) diff --git a/tests/unittest/api_stability/references/trtllm_serve_routes.yaml b/tests/unittest/api_stability/references/trtllm_serve_routes.yaml index a744e75b8686..60741bbdb96d 100644 --- a/tests/unittest/api_stability/references/trtllm_serve_routes.yaml +++ b/tests/unittest/api_stability/references/trtllm_serve_routes.yaml @@ -424,6 +424,50 @@ configs: - 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: diff --git a/tests/unittest/api_stability/test_serve_routes.py b/tests/unittest/api_stability/test_serve_routes.py index 78d253183a15..2216f623eb43 100644 --- a/tests/unittest/api_stability/test_serve_routes.py +++ b/tests/unittest/api_stability/test_serve_routes.py @@ -151,6 +151,41 @@ def test_metrics_presence_per_role(live_snapshots): 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"] diff --git a/tests/unittest/serve/test_serve_test_helpers.py b/tests/unittest/serve/test_serve_test_helpers.py index 8d7a32fa98f5..f97435cc6efa 100644 --- a/tests/unittest/serve/test_serve_test_helpers.py +++ b/tests/unittest/serve/test_serve_test_helpers.py @@ -101,7 +101,12 @@ def test_compare_reports_each_difference_kind(): 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 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) From 2eb184ae58b4344f034a428a5365ef60ead0dc80 Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Mon, 2 Nov 2026 10:33:02 +0000 Subject: [PATCH 04/26] [None][test] Address review of the serve snapshot and layering tests - AGENTS.md: the CPU test command sets LLM_MODELS_ROOT, which the runner requires, and the stage-mapping note says the mapper matches list-entry text, so a file listed through its directory maps by the directory. - check_contracts matches each from-import as ., so 'from pkg import sub' and 'from .. import sub' count as imports of the submodule; new test for both forms and for the allowance entry. - The VisualGen /v1/models test asserts the 200 status. - The AGENTS.md command check also recognizes commands that start with LLM_MODELS_ROOT=. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- tensorrt_llm/serve/AGENTS.md | 8 +++-- tests/test_common/serve/layering.py | 16 +++++---- .../visual_gen/test_ops_system_visual_gen.py | 4 ++- tests/unittest/serve/test_agents_md.py | 2 +- tests/unittest/serve/test_layering.py | 34 +++++++++++++++++++ 5 files changed, 54 insertions(+), 10 deletions(-) diff --git a/tensorrt_llm/serve/AGENTS.md b/tensorrt_llm/serve/AGENTS.md index fa82b67fa4c6..bdf1fe8c23b6 100644 --- a/tensorrt_llm/serve/AGENTS.md +++ b/tensorrt_llm/serve/AGENTS.md @@ -66,11 +66,15 @@ Name files and symbols here, never line numbers. ## Commands -- All serve CPU tests: `PYTHONPATH=tests python3 -m test_common.serve.cpu_tests` +- 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 ""` +- 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/tests/test_common/serve/layering.py b/tests/test_common/serve/layering.py index 20472ac446da..aba66f10df46 100644 --- a/tests/test_common/serve/layering.py +++ b/tests/test_common/serve/layering.py @@ -181,7 +181,11 @@ def _edge_triples(edge: ImportEdge) -> List[Tuple[str, str, str]]: def check_contracts(edges: Sequence[ImportEdge], contracts: Sequence[Contract]) -> List[str]: - """Violations of each contract, as messages that name the contract and the fix.""" + """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: @@ -189,16 +193,16 @@ def check_contracts(edges: Sequence[ImportEdge], contracts: Sequence[Contract]) continue if not any(_matches_prefix(edge.importer, prefix) for prefix in contract.importers): continue - if not any(_matches_prefix(edge.imported, prefix) for prefix in contract.forbidden): - 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 - importer, imported, name = triple - what = f"`{imported}.{name}`" if name else f"`{imported}`" problems.append( f"contract '{contract.name}': {importer} (line {edge.lineno}) imports " - f"{what}. {contract.fix}" + f"`{target}`. {contract.fix}" ) return problems 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 index d9b55a9b2624..205391381c13 100644 --- 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 @@ -45,6 +45,8 @@ def test_version(visual_gen_client): def test_models_lists_the_served_model(visual_gen_client): - body = visual_gen_client.get("/v1/models").json() + 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/test_agents_md.py b/tests/unittest/serve/test_agents_md.py index 42327c1202d3..b95dac68ff16 100644 --- a/tests/unittest/serve/test_agents_md.py +++ b/tests/unittest/serve/test_agents_md.py @@ -42,7 +42,7 @@ "examples/", ) ROOT_FILES = ("AGENTS.md", "CLAUDE.md", "legacy-files.txt", "CODING_GUIDELINES.md") -COMMAND_STARTS = ("pytest", "python", "python3", "PYTHONPATH=") +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") diff --git a/tests/unittest/serve/test_layering.py b/tests/unittest/serve/test_layering.py index 1b3c44427000..6c4fb30bccc8 100644 --- a/tests/unittest/serve/test_layering.py +++ b/tests/unittest/serve/test_layering.py @@ -183,6 +183,40 @@ def test_contract_violation_and_allowance(synthetic_package): 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( From 5138e15fe0e9e5fb53d94f7f56471ca40c8719cc Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Wed, 4 Nov 2026 01:35:30 +0000 Subject: [PATCH 05/26] [TRTLLM-17261][test] Refresh the serve route reference and baselines on the current main The disaggregated frontend now registers its routes with KvHintsRoute (#20269), which strips kv_hints from request bodies unless the server trusts its ingress. Regenerate the route reference: the six disaggregated configurations record KvHintsRoute where they recorded APIRoute; their route sets and order are unchanged. The route-class invariant, which a regenerated reference cannot hide, now requires KvHintsRoute on every API route of the disaggregated frontend and APIRoute on the coordinator app. Dropping the KV-hint policy from the frontend fails it even after a regeneration, and both apps still reject the aggregated server's msgpack route class. The object.__new__ baseline records the sites on main at the merge base: test_kv_hints.py builds one server that way (#20269), and test_harmony_parsing.py no longer builds any (#21176). The harmony parsing tests are now registered as a whole file (#21411), so the node-id entry that the CPU-selection test used as its example is gone. The node-id check now runs select_entries on a one-entry list of its own, so it no longer depends on which node-id entries l0_cpu.yml holds. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- .../references/trtllm_serve_routes.yaml | 128 +++++++++--------- .../api_stability/test_serve_routes.py | 32 ++++- tests/unittest/serve/test_no_object_new.py | 2 +- .../unittest/serve/test_serve_test_helpers.py | 14 +- 4 files changed, 100 insertions(+), 76 deletions(-) diff --git a/tests/unittest/api_stability/references/trtllm_serve_routes.yaml b/tests/unittest/api_stability/references/trtllm_serve_routes.yaml index 60741bbdb96d..523342d4afae 100644 --- a/tests/unittest/api_stability/references/trtllm_serve_routes.yaml +++ b/tests/unittest/api_stability/references/trtllm_serve_routes.yaml @@ -504,15 +504,15 @@ configs: - GET,HEAD /docs Route - GET,HEAD /docs/oauth2-redirect Route - GET,HEAD /redoc Route - - POST /v1/completions APIRoute - - POST /v1/chat/completions APIRoute - - POST /v1/messages APIRoute - - POST /v1/messages/count_tokens APIRoute - - POST /v1/responses APIRoute - - GET /v1/models APIRoute - - GET /health APIRoute - - GET /cluster_info APIRoute - - GET /version APIRoute + - 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 @@ -521,15 +521,15 @@ configs: - GET,HEAD /docs Route - GET,HEAD /docs/oauth2-redirect Route - GET,HEAD /redoc Route - - POST /v1/completions APIRoute - - POST /v1/chat/completions APIRoute - - POST /v1/messages APIRoute - - POST /v1/messages/count_tokens APIRoute - - POST /v1/responses APIRoute - - GET /v1/models APIRoute - - GET /health APIRoute - - GET /cluster_info APIRoute - - GET /version APIRoute + - 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 @@ -538,21 +538,21 @@ configs: - GET,HEAD /docs Route - GET,HEAD /docs/oauth2-redirect Route - GET,HEAD /redoc Route - - POST /v1/completions APIRoute - - POST /v1/chat/completions APIRoute - - POST /v1/messages APIRoute - - POST /v1/messages/count_tokens APIRoute - - POST /v1/responses APIRoute - - GET /v1/models APIRoute - - GET /health APIRoute - - GET /cluster_info APIRoute - - GET /version APIRoute + - 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 APIRoute - - GET /get APIRoute - - DELETE /delete APIRoute - - GET /expire APIRoute - - GET /get_prefix APIRoute + - POST /set KvHintsRoute + - GET /get KvHintsRoute + - DELETE /delete KvHintsRoute + - GET /expire KvHintsRoute + - GET /get_prefix KvHintsRoute disagg_fleet: kind: disagg_server routes: @@ -560,15 +560,15 @@ configs: - GET,HEAD /docs Route - GET,HEAD /docs/oauth2-redirect Route - GET,HEAD /redoc Route - - POST /v1/completions APIRoute - - POST /v1/chat/completions APIRoute - - POST /v1/messages APIRoute - - POST /v1/messages/count_tokens APIRoute - - POST /v1/responses APIRoute - - GET /v1/models APIRoute - - GET /health APIRoute - - GET /cluster_info APIRoute - - GET /version APIRoute + - 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 @@ -577,15 +577,15 @@ configs: - GET,HEAD /docs Route - GET,HEAD /docs/oauth2-redirect Route - GET,HEAD /redoc Route - - POST /v1/completions APIRoute - - POST /v1/chat/completions APIRoute - - POST /v1/messages APIRoute - - POST /v1/messages/count_tokens APIRoute - - POST /v1/responses APIRoute - - GET /v1/models APIRoute - - GET /health APIRoute - - GET /cluster_info APIRoute - - GET /version APIRoute + - 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 @@ -594,21 +594,21 @@ configs: - GET,HEAD /docs Route - GET,HEAD /docs/oauth2-redirect Route - GET,HEAD /redoc Route - - POST /v1/completions APIRoute - - POST /v1/chat/completions APIRoute - - POST /v1/messages APIRoute - - POST /v1/messages/count_tokens APIRoute - - POST /v1/responses APIRoute - - GET /v1/models APIRoute - - GET /health APIRoute - - GET /cluster_info APIRoute - - GET /version APIRoute + - 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 APIRoute - - GET /get APIRoute - - DELETE /delete APIRoute - - GET /expire APIRoute - - GET /get_prefix APIRoute + - POST /set KvHintsRoute + - GET /get KvHintsRoute + - DELETE /delete KvHintsRoute + - GET /expire KvHintsRoute + - GET /get_prefix KvHintsRoute coordinator: kind: coordinator_server routes: diff --git a/tests/unittest/api_stability/test_serve_routes.py b/tests/unittest/api_stability/test_serve_routes.py index 2216f623eb43..a2b8737889c8 100644 --- a/tests/unittest/api_stability/test_serve_routes.py +++ b/tests/unittest/api_stability/test_serve_routes.py @@ -106,17 +106,35 @@ def test_openai_server_routes_are_msgspec(live_snapshots): ) -def test_disagg_and_coordinator_routes_are_plain(live_snapshots): +_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"]: - cls = _parse(record)[2] - assert cls in ("APIRoute", "Route", "Mount"), ( - f"configuration '{name}': {record}: the disaggregated frontend and the " - "coordinator app register plain routes; they do not take the aggregated " - "server's msgpack route class." - ) + 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): diff --git a/tests/unittest/serve/test_no_object_new.py b/tests/unittest/serve/test_no_object_new.py index a883a28c496a..3c5e0628f53e 100644 --- a/tests/unittest/serve/test_no_object_new.py +++ b/tests/unittest/serve/test_no_object_new.py @@ -54,13 +54,13 @@ "tests/unittest/llmapi/apps/_test_trtllm_serve_profile.py": 1, "tests/unittest/llmapi/apps/test_anthropic_routes.py": 13, "tests/unittest/llmapi/apps/test_context_length_error.py": 1, - "tests/unittest/llmapi/apps/test_harmony_parsing.py": 1, "tests/unittest/llmapi/apps/test_openai_server_iteration_stats.py": 1, "tests/unittest/llmapi/apps/test_render_callers.py": 2, "tests/unittest/llmapi/apps/test_responses_store_replay.py": 1, "tests/unittest/llmapi/apps/test_responses_streaming_events.py": 1, "tests/unittest/llmapi/apps/test_serving_extensions.py": 1, "tests/unittest/llmapi/apps/test_tool_parsers.py": 1, + "tests/unittest/llmapi/test_kv_hints.py": 1, "tests/unittest/llmapi/test_llm.py": 2, "tests/unittest/llmapi/test_rl_control_auth.py": 1, "tests/unittest/llmapi/test_runtime_control.py": 1, diff --git a/tests/unittest/serve/test_serve_test_helpers.py b/tests/unittest/serve/test_serve_test_helpers.py index f97435cc6efa..aa70512d9986 100644 --- a/tests/unittest/serve/test_serve_test_helpers.py +++ b/tests/unittest/serve/test_serve_test_helpers.py @@ -233,10 +233,6 @@ def test_cpu_tests_selection(): 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 - # A node-id entry stays a node id: only that node runs, never the whole file. - harmony = "unittest/llmapi/apps/test_harmony_parsing.py::TestStripIncompleteMessagesReporting" - assert harmony in targets - assert "unittest/llmapi/apps/test_harmony_parsing.py" not in targets for target in targets: file_part = target.split("::")[0] text = (REPO_ROOT / "tests" / file_part).read_text(errors="replace") @@ -246,6 +242,16 @@ def test_cpu_tests_selection(): ) +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 # --------------------------------------------------------------------------- From 045a69eae058c80c866a0cd5ac7a9c592426b737 Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Wed, 4 Nov 2026 07:43:16 +0000 Subject: [PATCH 06/26] [None][test] Make the route snapshot and the object.__new__ ratchet fail on silent regressions The route snapshot now requires every configuration other than the three invalid ones to build, so a configuration that starts failing at construction cannot pass through a regenerated reference. The object.__new__ ratchet resolves 'from ... import as ' aliases and records each baseline site as its enclosing function or class plus the server class, so swapping one site for another in the same file no longer passes. The baseline keeps the same 52 sites in the same 26 files. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- .../api_stability/test_serve_routes.py | 36 ++- tests/unittest/serve/test_no_object_new.py | 278 ++++++++++++++---- 2 files changed, 241 insertions(+), 73 deletions(-) diff --git a/tests/unittest/api_stability/test_serve_routes.py b/tests/unittest/api_stability/test_serve_routes.py index a2b8737889c8..1392c1428ec3 100644 --- a/tests/unittest/api_stability/test_serve_routes.py +++ b/tests/unittest/api_stability/test_serve_routes.py @@ -215,18 +215,34 @@ def test_metrics_shadowing_order(live_snapshots): ) -def test_error_configurations(tmp_path): - expected = { - "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" - ), +# 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 expected.items(): + 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 " diff --git a/tests/unittest/serve/test_no_object_new.py b/tests/unittest/serve/test_no_object_new.py index 3c5e0628f53e..0aa59460caf8 100644 --- a/tests/unittest/serve/test_no_object_new.py +++ b/tests/unittest/serve/test_no_object_new.py @@ -19,13 +19,17 @@ 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 that may only shrink. +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. +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 @@ -37,35 +41,115 @@ SERVER_CLASSES = ("OpenAIServer", "OpenAIDisaggServer", "CoordinatorServer") -# Per-file counts of the sites that existed when this ratchet landed. -# The baseline only shrinks: rewrite a site on the factory, lower its count. +# 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": 2, - "tests/unittest/disaggregated/test_disagg_internal_auth.py": 3, - "tests/unittest/disaggregated/test_disagg_openai_client.py": 1, - "tests/unittest/disaggregated/test_disagg_responses_api.py": 2, - "tests/unittest/disaggregated/test_openai_disagg_server.py": 3, - "tests/unittest/disaggregated/test_openai_disagg_service.py": 2, - "tests/unittest/disaggregated/test_openai_server_info.py": 1, - "tests/unittest/disaggregated/test_openai_server_stream_errors.py": 1, - "tests/unittest/disaggregated/test_paired_request_cancellation.py": 2, - "tests/unittest/disaggregated/test_router.py": 2, - "tests/unittest/executor/test_spec_decode_stats_payload.py": 1, - "tests/unittest/llmapi/apps/_test_trtllm_serve_profile.py": 1, - "tests/unittest/llmapi/apps/test_anthropic_routes.py": 13, - "tests/unittest/llmapi/apps/test_context_length_error.py": 1, - "tests/unittest/llmapi/apps/test_openai_server_iteration_stats.py": 1, - "tests/unittest/llmapi/apps/test_render_callers.py": 2, - "tests/unittest/llmapi/apps/test_responses_store_replay.py": 1, - "tests/unittest/llmapi/apps/test_responses_streaming_events.py": 1, - "tests/unittest/llmapi/apps/test_serving_extensions.py": 1, - "tests/unittest/llmapi/apps/test_tool_parsers.py": 1, - "tests/unittest/llmapi/test_kv_hints.py": 1, - "tests/unittest/llmapi/test_llm.py": 2, - "tests/unittest/llmapi/test_rl_control_auth.py": 1, - "tests/unittest/llmapi/test_runtime_control.py": 1, - "tests/unittest/llmapi/test_serve_multi_frontend_helpers.py": 4, - "tests/unittest/llmapi/test_serve_multi_frontend_route_gate.py": 1, + "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": [ + "_batch_client OpenAIServer", + "_chat_stream OpenAIServer", + "_count_tokens_client OpenAIServer", + "_handle OpenAIDisaggServer", + "_make_route_client OpenAIDisaggServer", + "_make_route_client OpenAIServer", + "test_await_disconnected_returns_immediately_for_a_synthetic_request OpenAIServer", + "test_batch_routes_are_registered_on_the_standard_server OpenAIServer", + "test_count_tokens_honours_the_server_chat_template OpenAIServer", + "test_create_batch_at_capacity_is_429_not_400 OpenAIServer", + "test_results_of_an_unfinished_batch_are_400_not_404.scenario OpenAIServer", + "test_standard_and_disagg_register_messages_route OpenAIDisaggServer", + "test_standard_and_disagg_register_messages_route OpenAIServer", + ], + "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": [ + "TestCountTokens._client OpenAIServer", + "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", + ], } @@ -77,56 +161,94 @@ def _class_name(node) -> str: return "" -def count_new_sites(tree: ast.AST) -> int: - """Call nodes that build a server class with ``__new__``.""" - sites = 0 +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 not isinstance(node, ast.Call) or not isinstance(node.func, ast.Attribute): - continue - if node.func.attr != "__new__": - continue - owner = _class_name(node.func.value) - if owner == "object": - if node.args and _class_name(node.args[0]) in SERVER_CLASSES: - sites += 1 - elif owner in SERVER_CLASSES: - sites += 1 - return sites + 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. -def scan_tests(tests_root: Path = TESTS_ROOT): - counts = {} + 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 = count_new_sites(tree) + sites = find_new_sites(tree) if sites: - counts[str(path.relative_to(REPO_ROOT))] = sites - return counts + sites_by_file[str(path.relative_to(REPO_ROOT))] = sites + return sites_by_file -def test_no_new_object_new_sites(): - counts = scan_tests() +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, sites in sorted(counts.items()): - baseline = BASELINE.get(file, 0) - if sites > baseline: + 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} builds a server with object.__new__ ({sites} sites, baseline " - f"{baseline}); use `isolated_serve_env` from " + 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 file, baseline in sorted(BASELINE.items()): - sites = counts.get(file, 0) - if sites < baseline: + for site in sorted((recorded - live).elements()): problems.append( - f"{file} now has {sites} object.__new__ sites; lower BASELINE[{file!r}] to " - f"{sites} in tests/unittest/serve/test_no_object_new.py; the baseline only " - "shrinks." + 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) @@ -144,7 +266,37 @@ def test_no_new_object_new_sites(): ("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 count_new_sites(ast.parse(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] From 1a3a615318d7ed7355a6184fb445577a54ceb7f5 Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Mon, 2 Nov 2026 09:50:26 +0000 Subject: [PATCH 07/26] [TRTLLM-17262][refactor] Shared middleware package for the aggregated and disaggregated servers Create tensorrt_llm/serve/middleware/ and move ServerArrivalTimeMiddleware (from serve/responses_utils.py), PerfMetricsMiddleware (from serve/perf_metrics.py) and _apply_fastapi_middlewares (from commands/serve.py) into it, verbatim. Both servers now install their built-in middleware through serve/middleware/register.py::register_middlewares, which reproduces each app's stack exactly: the aggregated server passes adjusted_clock to both middlewares, the disaggregated server to neither. User --middleware entries are still applied after construction and stay outermost. The old import paths are removed with no compatibility shim; no code outside the repository imports them. commands/serve.py keeps an import of _apply_fastapi_middlewares for its own two launcher calls. CI bookkeeping in the same change: the multi-GPU trigger table gains tensorrt_llm/serve/middleware/; legacy-files.txt registers middleware/arrival_time.py and middleware/user_middleware.py with the configs gen-configs regenerates; ruff-legacy-baseline.json carries the moved docstring's one D212; mypy checks the new register.py. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- .pre-commit-config.yaml | 4 + jenkins/L0_MergeRequest.groovy | 1 + legacy-files.txt | 2 + pyproject.toml | 3 + ruff-legacy-baseline.json | 3 + ruff-legacy.toml | 2 + tensorrt_llm/commands/serve.py | 23 +-- tensorrt_llm/serve/AGENTS.md | 5 +- tensorrt_llm/serve/middleware/__init__.py | 15 ++ tensorrt_llm/serve/middleware/arrival_time.py | 42 ++++ tensorrt_llm/serve/middleware/perf_metrics.py | 100 +++++++++ tensorrt_llm/serve/middleware/register.py | 72 +++++++ .../serve/middleware/user_middleware.py | 38 ++++ tensorrt_llm/serve/openai_disagg_server.py | 15 +- tensorrt_llm/serve/openai_server.py | 17 +- tensorrt_llm/serve/perf_metrics.py | 72 ------- tensorrt_llm/serve/responses_utils.py | 35 ---- .../test_disagg_openai_client.py | 5 +- .../llmapi/apps/test_request_metrics.py | 2 +- tests/unittest/llmapi/test_config_database.py | 2 +- tests/unittest/serve/middleware/__init__.py | 14 ++ .../serve/middleware/test_register.py | 193 ++++++++++++++++++ tests/unittest/serve/test_layering.py | 14 ++ tests/unittest/serve/test_middleware_order.py | 2 +- .../unittest/serve/test_serve_test_helpers.py | 2 +- 25 files changed, 528 insertions(+), 155 deletions(-) create mode 100644 tensorrt_llm/serve/middleware/__init__.py create mode 100644 tensorrt_llm/serve/middleware/arrival_time.py create mode 100644 tensorrt_llm/serve/middleware/perf_metrics.py create mode 100644 tensorrt_llm/serve/middleware/register.py create mode 100644 tensorrt_llm/serve/middleware/user_middleware.py create mode 100644 tests/unittest/serve/middleware/__init__.py create mode 100644 tests/unittest/serve/middleware/test_register.py diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index c2166e7bb04b..b1d4513dd66d 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -422,6 +422,8 @@ common-files: &common_files | tensorrt_llm/serve/disagg_auto_scaling.py | tensorrt_llm/serve/harmony_adapter.py | tensorrt_llm/serve/metadata_server.py | + tensorrt_llm/serve/middleware/arrival_time.py | + tensorrt_llm/serve/middleware/user_middleware.py | tensorrt_llm/serve/openai_disagg_server.py | tensorrt_llm/serve/openai_protocol.py | tensorrt_llm/serve/openai_server.py | @@ -1149,6 +1151,8 @@ legacy-files: &legacy_files | tensorrt_llm/serve/disagg_auto_scaling.py | tensorrt_llm/serve/harmony_adapter.py | tensorrt_llm/serve/metadata_server.py | + tensorrt_llm/serve/middleware/arrival_time.py | + tensorrt_llm/serve/middleware/user_middleware.py | tensorrt_llm/serve/openai_disagg_server.py | tensorrt_llm/serve/openai_protocol.py | tensorrt_llm/serve/openai_server.py | diff --git a/jenkins/L0_MergeRequest.groovy b/jenkins/L0_MergeRequest.groovy index fed051af26b2..d949a16d90ff 100644 --- a/jenkins/L0_MergeRequest.groovy +++ b/jenkins/L0_MergeRequest.groovy @@ -1952,6 +1952,7 @@ def getMultiGpuFileChanged(pipeline, testFilter, globalVars) "tensorrt_llm/serve/disagg_auto_scaling.py", "tensorrt_llm/serve/extensions/", "tensorrt_llm/serve/metadata_server.py", + "tensorrt_llm/serve/middleware/", "tensorrt_llm/serve/openai_client.py", "tensorrt_llm/serve/openai_disagg_server.py", "tensorrt_llm/serve/openai_disagg_service.py", diff --git a/legacy-files.txt b/legacy-files.txt index 9850e7c810b4..33c46057a53b 100644 --- a/legacy-files.txt +++ b/legacy-files.txt @@ -414,6 +414,8 @@ tensorrt_llm/serve/cluster_storage.py tensorrt_llm/serve/disagg_auto_scaling.py tensorrt_llm/serve/harmony_adapter.py tensorrt_llm/serve/metadata_server.py +tensorrt_llm/serve/middleware/arrival_time.py +tensorrt_llm/serve/middleware/user_middleware.py tensorrt_llm/serve/openai_disagg_server.py tensorrt_llm/serve/openai_protocol.py tensorrt_llm/serve/openai_server.py diff --git a/pyproject.toml b/pyproject.toml index 60e243d984b0..133a66f71076 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -479,6 +479,8 @@ exclude = [ "tensorrt_llm/serve/disagg_auto_scaling.py", "tensorrt_llm/serve/harmony_adapter.py", "tensorrt_llm/serve/metadata_server.py", + "tensorrt_llm/serve/middleware/arrival_time.py", + "tensorrt_llm/serve/middleware/user_middleware.py", "tensorrt_llm/serve/openai_disagg_server.py", "tensorrt_llm/serve/openai_protocol.py", "tensorrt_llm/serve/openai_server.py", @@ -862,6 +864,7 @@ 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/middleware/register.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..efa2253b8cb6 100644 --- a/ruff-legacy-baseline.json +++ b/ruff-legacy-baseline.json @@ -539,6 +539,9 @@ "D212": 19, "D415": 1 }, + "tensorrt_llm/serve/middleware/arrival_time.py": { + "D212": 1 + }, "tensorrt_llm/serve/openai_disagg_server.py": { "D210": 1, "F811": 1 diff --git a/ruff-legacy.toml b/ruff-legacy.toml index dab55b4d0434..524064c8b7a4 100644 --- a/ruff-legacy.toml +++ b/ruff-legacy.toml @@ -431,6 +431,8 @@ include = [ "tensorrt_llm/serve/disagg_auto_scaling.py", "tensorrt_llm/serve/harmony_adapter.py", "tensorrt_llm/serve/metadata_server.py", + "tensorrt_llm/serve/middleware/arrival_time.py", + "tensorrt_llm/serve/middleware/user_middleware.py", "tensorrt_llm/serve/openai_disagg_server.py", "tensorrt_llm/serve/openai_protocol.py", "tensorrt_llm/serve/openai_server.py", diff --git a/tensorrt_llm/commands/serve.py b/tensorrt_llm/commands/serve.py index e8234ac09a2c..8cb97c58abf2 100644 --- a/tensorrt_llm/commands/serve.py +++ b/tensorrt_llm/commands/serve.py @@ -16,8 +16,6 @@ import atexit import contextlib import gc -import importlib -import inspect import json import os import secrets @@ -74,6 +72,8 @@ _frontend_affinity_sockets, _init_multi_frontend_mode, _signal_frontend_ready, _spawn_attached_frontends, _terminate_attached_frontends, _watch_attached_frontends) +from tensorrt_llm.serve.middleware.user_middleware import \ + _apply_fastapi_middlewares from tensorrt_llm.serve.multi_frontend import (LAUNCHER_UDS_ENV, LAUNCHER_UDS_NAME, MultiFrontendServing) @@ -158,25 +158,6 @@ def _pop_bool_config_option(config: dict[str, Any], key: str) -> bool: return validate_config_bool(config.pop(key, False), key) -def _apply_fastapi_middlewares(app, middlewares: Sequence[str]) -> None: - """Import and register middleware objects on a FastAPI app.""" - for middleware in middlewares: - try: - module_path, object_name = middleware.rsplit(".", 1) - except ValueError as e: - raise ValueError(f"Invalid middleware import path '{middleware}'. " - "Expected format: ..") from e - - imported = getattr(importlib.import_module(module_path), object_name) - if inspect.isclass(imported): - app.add_middleware(imported) - elif inspect.iscoroutinefunction(imported): - app.middleware("http")(imported) - else: - raise ValueError(f"Invalid middleware {middleware}. " - "Must be a class or an async function.") - - def is_non_default_or_required(param_name, value, explicit_cli_keys): """ Check if a parameter should be explicitly included in llm_args. diff --git a/tensorrt_llm/serve/AGENTS.md b/tensorrt_llm/serve/AGENTS.md index bdf1fe8c23b6..5e0de500ef43 100644 --- a/tensorrt_llm/serve/AGENTS.md +++ b/tensorrt_llm/serve/AGENTS.md @@ -23,8 +23,9 @@ Name files and symbols here, never line numbers. `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`. +- Middleware: register built-in middleware in `serve/middleware/register.py`, which both apps + call. 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/` diff --git a/tensorrt_llm/serve/middleware/__init__.py b/tensorrt_llm/serve/middleware/__init__.py new file mode 100644 index 000000000000..1fd42a075078 --- /dev/null +++ b/tensorrt_llm/serve/middleware/__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. +"""ASGI middleware of the trtllm-serve apps. Register built-in middleware in register.py.""" diff --git a/tensorrt_llm/serve/middleware/arrival_time.py b/tensorrt_llm/serve/middleware/arrival_time.py new file mode 100644 index 000000000000..a0290227e69b --- /dev/null +++ b/tensorrt_llm/serve/middleware/arrival_time.py @@ -0,0 +1,42 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""ASGI middleware that stamps each HTTP request with its arrival time.""" + +from typing import Optional + +from tensorrt_llm._utils import AdjustedSteadyClock + + +class ServerArrivalTimeMiddleware: + """ + Custom ASGI middleware to track server arrival time. + + We implement this as a pure ASGI middleware instead of using FastAPI's + @app.middleware("http") decorator because the decorator internally uses + BaseHTTPMiddleware, which wraps the ASGI `receive` callable. This wrapping + breaks Request.is_disconnected() functionality - the wrapped receive doesn't + properly forward http.disconnect events while the middleware is waiting in + call_next(), preventing detection of client disconnections during long-running + non-streaming requests. + + By implementing pure ASGI middleware, we pass through the original receive/send + callables unchanged, preserving the ability to detect client disconnections. + + See: https://github.com/encode/starlette/discussions/2094 + """ + + def __init__(self, + app, + adjusted_clock: Optional[AdjustedSteadyClock] = None): + self.app = app + self._adjusted_clock = adjusted_clock or AdjustedSteadyClock() + + async def __call__(self, scope, receive, send): + if scope["type"] == "http": + # Add arrival time to scope + scope["state"] = {} + scope["state"]["server_arrival_time"] = self._adjusted_clock.now() + + # Pass through the original receive/send - no wrapping! + await self.app(scope, receive, send) diff --git a/tensorrt_llm/serve/middleware/perf_metrics.py b/tensorrt_llm/serve/middleware/perf_metrics.py new file mode 100644 index 000000000000..bd901be1d482 --- /dev/null +++ b/tensorrt_llm/serve/middleware/perf_metrics.py @@ -0,0 +1,100 @@ +# 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. +"""ASGI middleware that returns per-request perf metrics and hands completed records to the JSONL writer.""" + +from typing import Any, Dict, Optional + +from tensorrt_llm._utils import AdjustedSteadyClock +from tensorrt_llm.serve.perf_metrics import ( + _CLOCK_SYNC_HEADER_BYTES, + _RETURN_METRICS_HEADER_BYTES, + PerfMetricsJsonlWriter, + _limit_metrics_headers, + build_clock_sync_header, + build_metrics_headers, + build_metrics_sse_event, +) + + +class PerfMetricsMiddleware: + """Expose request metrics and optionally persist completed records.""" + + def __init__( + self, + app: Any, + expose_headers: bool, + writer: Optional[PerfMetricsJsonlWriter] = None, + adjusted_clock: Optional[AdjustedSteadyClock] = None, + ): + self._app = app + self._expose_headers = expose_headers + self._writer = writer + self._adjusted_clock = adjusted_clock + + async def __call__(self, scope: Dict[str, Any], receive: Any, send: Any) -> None: + if scope["type"] != "http": + await self._app(scope, receive, send) + return + is_stream = False + metrics_headers = None + records = scope.setdefault("state", {}).setdefault("perf_metrics_records", []) + return_metrics = self._expose_headers and any( + name.lower() == _RETURN_METRICS_HEADER_BYTES and value.strip() == b"1" + for name, value in scope.get("headers", []) + ) + + async def send_metrics(message: Dict[str, Any]) -> None: + nonlocal is_stream, metrics_headers + if message["type"] == "http.response.start": + headers = list(message.get("headers", [])) + is_stream = any( + key.lower() == b"content-type" and b"text/event-stream" in value.lower() + for key, value in headers + ) + if records and return_metrics and not is_stream: + metrics_headers = build_metrics_headers(records) + public_headers = _limit_metrics_headers(metrics_headers) + headers.extend( + (name.encode(), value.encode()) for name, value in public_headers.items() + ) + if return_metrics and self._adjusted_clock is not None: + receive_time = scope.get("state", {}).get("server_arrival_time") + if receive_time is not None: + clock_sync = build_clock_sync_header( + receive_time, + self._adjusted_clock.now(), + ) + headers.append((_CLOCK_SYNC_HEADER_BYTES, clock_sync.encode())) + message["headers"] = headers + + elif message["type"] == "http.response.body" and not message.get("more_body", False): + if records: + if metrics_headers is None: + metrics_headers = build_metrics_headers(records) + if self._writer is not None: + for record in records: + self._writer.submit(record) + if return_metrics and is_stream: + message["body"] = message.get("body", b"") + build_metrics_sse_event( + metrics_headers + ) + try: + await send(message) + except OSError: + pass + return + await send(message) + + await self._app(scope, receive, send_metrics) diff --git a/tensorrt_llm/serve/middleware/register.py b/tensorrt_llm/serve/middleware/register.py new file mode 100644 index 000000000000..c6cdbfaa9344 --- /dev/null +++ b/tensorrt_llm/serve/middleware/register.py @@ -0,0 +1,72 @@ +# 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. +"""Built-in middleware of the aggregated and disaggregated trtllm-serve apps. + +Starlette's ``add_middleware`` inserts at the front of the stack, so the middleware added +last runs first. Both apps call ``register_middlewares``; do not add middleware elsewhere. +""" + +from typing import Optional + +from fastapi import FastAPI + +from tensorrt_llm._utils import AdjustedSteadyClock +from tensorrt_llm.serve.middleware.arrival_time import ServerArrivalTimeMiddleware +from tensorrt_llm.serve.middleware.perf_metrics import PerfMetricsMiddleware +from tensorrt_llm.serve.perf_metrics import PerfMetricsJsonlWriter + + +def register_middlewares( + app: FastAPI, + *, + collect_perf_metrics: bool, + expose_perf_metrics: bool, + perf_metrics_writer: Optional[PerfMetricsJsonlWriter], + adjusted_clock: Optional[AdjustedSteadyClock] = None, +) -> None: + """Install the built-in middleware of a trtllm-serve app. + + The stack, outermost first, is ServerArrivalTimeMiddleware, then PerfMetricsMiddleware + when perf metrics are collected. ``--middleware`` entries are added after the server is + built (serve/middleware/user_middleware.py), so they end up outside both. + + Args: + app: The app of OpenAIServer or OpenAIDisaggServer. + collect_perf_metrics: Install PerfMetricsMiddleware (metrics returned to clients or + written to JSONL). + expose_perf_metrics: Let PerfMetricsMiddleware return metrics in response headers and + in the SSE metrics event. + perf_metrics_writer: Writer for completed records. + adjusted_clock: The aggregated server's clock, mapped to the disaggregated server's + clock domain. The disaggregated server passes None, and then neither middleware + receives the argument. + """ + if adjusted_clock is None: + if collect_perf_metrics: + app.add_middleware( + PerfMetricsMiddleware, + expose_headers=expose_perf_metrics, + writer=perf_metrics_writer, + ) + app.add_middleware(ServerArrivalTimeMiddleware) + return + if collect_perf_metrics: + app.add_middleware( + PerfMetricsMiddleware, + expose_headers=expose_perf_metrics, + writer=perf_metrics_writer, + adjusted_clock=adjusted_clock, + ) + app.add_middleware(ServerArrivalTimeMiddleware, adjusted_clock=adjusted_clock) diff --git a/tensorrt_llm/serve/middleware/user_middleware.py b/tensorrt_llm/serve/middleware/user_middleware.py new file mode 100644 index 000000000000..36f743b88a8a --- /dev/null +++ b/tensorrt_llm/serve/middleware/user_middleware.py @@ -0,0 +1,38 @@ +# 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 that users add with ``trtllm-serve --middleware``.""" + +import importlib +import inspect +from typing import Sequence + + +def _apply_fastapi_middlewares(app, middlewares: Sequence[str]) -> None: + """Import and register middleware objects on a FastAPI app.""" + for middleware in middlewares: + try: + module_path, object_name = middleware.rsplit(".", 1) + except ValueError as e: + raise ValueError(f"Invalid middleware import path '{middleware}'. " + "Expected format: ..") from e + + imported = getattr(importlib.import_module(module_path), object_name) + if inspect.isclass(imported): + app.add_middleware(imported) + elif inspect.iscoroutinefunction(imported): + app.middleware("http")(imported) + else: + raise ValueError(f"Invalid middleware {middleware}. " + "Must be a class or an async function.") diff --git a/tensorrt_llm/serve/openai_disagg_server.py b/tensorrt_llm/serve/openai_disagg_server.py index d8ebec1a6ce7..90e1848cd72c 100644 --- a/tensorrt_llm/serve/openai_disagg_server.py +++ b/tensorrt_llm/serve/openai_disagg_server.py @@ -58,6 +58,7 @@ from tensorrt_llm.serve.disagg_coordinator import (CoordinatorClient, DisaggCoordinatorService) from tensorrt_llm.serve.kv_hints import KvHintsRoute +from tensorrt_llm.serve.middleware.register import register_middlewares from tensorrt_llm.serve.openai_client import OpenAIClient, OpenAIHttpClient from tensorrt_llm.serve.openai_disagg_service import ( OpenAIDisaggregatedService, ResponseHooks) @@ -67,10 +68,8 @@ 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, - get_steady_clock_now_in_seconds, +from tensorrt_llm.serve.responses_utils import (get_steady_clock_now_in_seconds, guard_responses_stream, stream_error_event) from tensorrt_llm.serve.router import Router @@ -295,12 +294,10 @@ async def lifespan(app) -> None: self.app.state.trust_kv_hints = os.getenv("TRTLLM_TRUST_KV_HINTS", "0") == "1" self.app.router.route_class = KvHintsRoute - if self._collect_perf_metrics: - self.app.add_middleware( - PerfMetricsMiddleware, - expose_headers=self._expose_perf_metrics, - writer=self._perf_metrics_writer) - self.app.add_middleware(ServerArrivalTimeMiddleware) + register_middlewares(self.app, + collect_perf_metrics=self._collect_perf_metrics, + expose_perf_metrics=self._expose_perf_metrics, + perf_metrics_writer=self._perf_metrics_writer) # Log request-body validation failures so a client/server schema mismatch # shows up server-side. Throttled (first, then every 1000th) to avoid diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py index eb79cdb18397..93ce610ca9a2 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -93,6 +93,7 @@ from tensorrt_llm.serve.encode_batcher import (EncodeBatcher, InputTooLongError, QueueFullError) from tensorrt_llm.serve.metadata_server import create_metadata_server +from tensorrt_llm.serve.middleware.register import register_middlewares from tensorrt_llm.serve.multi_frontend import (FORWARDED_ROUTES, AttachedFrontendWatchdog, LauncherForwarder, @@ -114,7 +115,6 @@ to_llm_disaggregated_params) from tensorrt_llm.serve.openai_video_routes import _VideoRoutesMixin from tensorrt_llm.serve.perf_metrics import (PerfMetricsJsonlWriter, - PerfMetricsMiddleware, build_request_metrics_record) from tensorrt_llm.serve.postprocess_handlers import ( ChatCompletionPostprocArgs, ChatPostprocArgs, CompletionPostprocArgs, @@ -131,8 +131,7 @@ render_conversation) from tensorrt_llm.serve.responses_utils import (ClosingStreamingResponse, ConversationHistoryStore, - ResponsesStreamingProcessor, - ServerArrivalTimeMiddleware) + ResponsesStreamingProcessor) from tensorrt_llm.serve.responses_utils import \ create_response as responses_api_create_response from tensorrt_llm.serve.responses_utils import (guard_responses_stream, @@ -926,13 +925,11 @@ async def validation_exception_handler(request, exc): self._register_launcher_forwarding_routes() self.register_routes() - if self._collect_perf_metrics: - self.app.add_middleware(PerfMetricsMiddleware, - expose_headers=self._expose_perf_metrics, - writer=self._perf_metrics_writer, - adjusted_clock=self._adjusted_steady_clock) - self.app.add_middleware(ServerArrivalTimeMiddleware, - adjusted_clock=self._adjusted_steady_clock) + register_middlewares(self.app, + collect_perf_metrics=self._collect_perf_metrics, + expose_perf_metrics=self._expose_perf_metrics, + perf_metrics_writer=self._perf_metrics_writer, + adjusted_clock=self._adjusted_steady_clock) def _init_visual_gen(self): self.processor = None diff --git a/tensorrt_llm/serve/perf_metrics.py b/tensorrt_llm/serve/perf_metrics.py index 5c6c7b2715b3..371e9ae02f56 100644 --- a/tensorrt_llm/serve/perf_metrics.py +++ b/tensorrt_llm/serve/perf_metrics.py @@ -786,75 +786,3 @@ def build_metrics_sse_event(headers: Dict[str, str]) -> bytes: return b"" payload = json.dumps(headers, separators=(",", ":")) return (f"event: {SSE_METRICS_EVENT}\ndata: {payload}\n\n").encode() - - -class PerfMetricsMiddleware: - """Expose request metrics and optionally persist completed records.""" - - def __init__( - self, - app: Any, - expose_headers: bool, - writer: Optional[PerfMetricsJsonlWriter] = None, - adjusted_clock: Optional[AdjustedSteadyClock] = None, - ): - self._app = app - self._expose_headers = expose_headers - self._writer = writer - self._adjusted_clock = adjusted_clock - - async def __call__(self, scope: Dict[str, Any], receive: Any, send: Any) -> None: - if scope["type"] != "http": - await self._app(scope, receive, send) - return - is_stream = False - metrics_headers = None - records = scope.setdefault("state", {}).setdefault("perf_metrics_records", []) - return_metrics = self._expose_headers and any( - name.lower() == _RETURN_METRICS_HEADER_BYTES and value.strip() == b"1" - for name, value in scope.get("headers", []) - ) - - async def send_metrics(message: Dict[str, Any]) -> None: - nonlocal is_stream, metrics_headers - if message["type"] == "http.response.start": - headers = list(message.get("headers", [])) - is_stream = any( - key.lower() == b"content-type" and b"text/event-stream" in value.lower() - for key, value in headers - ) - if records and return_metrics and not is_stream: - metrics_headers = build_metrics_headers(records) - public_headers = _limit_metrics_headers(metrics_headers) - headers.extend( - (name.encode(), value.encode()) for name, value in public_headers.items() - ) - if return_metrics and self._adjusted_clock is not None: - receive_time = scope.get("state", {}).get("server_arrival_time") - if receive_time is not None: - clock_sync = build_clock_sync_header( - receive_time, - self._adjusted_clock.now(), - ) - headers.append((_CLOCK_SYNC_HEADER_BYTES, clock_sync.encode())) - message["headers"] = headers - - elif message["type"] == "http.response.body" and not message.get("more_body", False): - if records: - if metrics_headers is None: - metrics_headers = build_metrics_headers(records) - if self._writer is not None: - for record in records: - self._writer.submit(record) - if return_metrics and is_stream: - message["body"] = message.get("body", b"") + build_metrics_sse_event( - metrics_headers - ) - try: - await send(message) - except OSError: - pass - return - await send(message) - - await self._app(scope, receive, send_metrics) diff --git a/tensorrt_llm/serve/responses_utils.py b/tensorrt_llm/serve/responses_utils.py index e139cfe0938c..0173bcb34431 100644 --- a/tensorrt_llm/serve/responses_utils.py +++ b/tensorrt_llm/serve/responses_utils.py @@ -56,7 +56,6 @@ from tensorrt_llm._utils import \ get_steady_clock_now_in_seconds # noqa: F401 (re-export) -from tensorrt_llm._utils import AdjustedSteadyClock from tensorrt_llm.executor import (EngineDeadError, GenerationResult, RequestError) from tensorrt_llm.inputs import PromptInputs, prompt_inputs @@ -3488,40 +3487,6 @@ def until_keepalive() -> Optional[float]: await items.aclose() -class ServerArrivalTimeMiddleware: - """ - Custom ASGI middleware to track server arrival time. - - We implement this as a pure ASGI middleware instead of using FastAPI's - @app.middleware("http") decorator because the decorator internally uses - BaseHTTPMiddleware, which wraps the ASGI `receive` callable. This wrapping - breaks Request.is_disconnected() functionality - the wrapped receive doesn't - properly forward http.disconnect events while the middleware is waiting in - call_next(), preventing detection of client disconnections during long-running - non-streaming requests. - - By implementing pure ASGI middleware, we pass through the original receive/send - callables unchanged, preserving the ability to detect client disconnections. - - See: https://github.com/encode/starlette/discussions/2094 - """ - - def __init__(self, - app, - adjusted_clock: Optional[AdjustedSteadyClock] = None): - self.app = app - self._adjusted_clock = adjusted_clock or AdjustedSteadyClock() - - async def __call__(self, scope, receive, send): - if scope["type"] == "http": - # Add arrival time to scope - scope["state"] = {} - scope["state"]["server_arrival_time"] = self._adjusted_clock.now() - - # Pass through the original receive/send - no wrapping! - await self.app(scope, receive, send) - - class PeriodicLatencyLogger: """Periodically log latency percentiles for a named coordinator API. diff --git a/tests/unittest/disaggregated/test_disagg_openai_client.py b/tests/unittest/disaggregated/test_disagg_openai_client.py index 61519e07bb3b..a008ef463cdc 100644 --- a/tests/unittest/disaggregated/test_disagg_openai_client.py +++ b/tests/unittest/disaggregated/test_disagg_openai_client.py @@ -36,6 +36,8 @@ validate_internal_disagg_request, validate_subagent_affinity, ) +from tensorrt_llm.serve.middleware.arrival_time import ServerArrivalTimeMiddleware +from tensorrt_llm.serve.middleware.perf_metrics import PerfMetricsMiddleware from tensorrt_llm.serve.openai_client import OpenAIHttpClient, _encode_request_body from tensorrt_llm.serve.openai_protocol import ( CompletionRequest, @@ -53,10 +55,9 @@ SERVER_TIMING_HEADER, SSE_METRICS_EVENT, START_END_TIME_HEADER, - PerfMetricsMiddleware, adjusted_clock_from_headers, ) -from tensorrt_llm.serve.responses_utils import ResponseHooks, ServerArrivalTimeMiddleware +from tensorrt_llm.serve.responses_utils import ResponseHooks from tensorrt_llm.serve.router import Router pytestmark = pytest.mark.cpu_only diff --git a/tests/unittest/llmapi/apps/test_request_metrics.py b/tests/unittest/llmapi/apps/test_request_metrics.py index cffed81b75e1..addbf0e2605a 100644 --- a/tests/unittest/llmapi/apps/test_request_metrics.py +++ b/tests/unittest/llmapi/apps/test_request_metrics.py @@ -21,6 +21,7 @@ from tensorrt_llm.logger import logger from tensorrt_llm.serve import perf_metrics as perf_metrics_module +from tensorrt_llm.serve.middleware.perf_metrics import PerfMetricsMiddleware from tensorrt_llm.serve.perf_metrics import ( CTX_CHUNK_METRICS_HEADER, RETURN_METRICS_HEADER, @@ -29,7 +30,6 @@ START_END_TIME_HEADER, STEP_METRICS_HEADER, PerfMetricsJsonlWriter, - PerfMetricsMiddleware, _jsonl_record, build_metrics_headers, build_metrics_record_from_headers, diff --git a/tests/unittest/llmapi/test_config_database.py b/tests/unittest/llmapi/test_config_database.py index 33dc9011af9f..75a21fa1c98a 100644 --- a/tests/unittest/llmapi/test_config_database.py +++ b/tests/unittest/llmapi/test_config_database.py @@ -25,9 +25,9 @@ from starlette.middleware.base import BaseHTTPMiddleware from tensorrt_llm._torch.models.modeling_utils import get_registered_model_class -from tensorrt_llm.commands.serve import _apply_fastapi_middlewares from tensorrt_llm.commands.serve import main as serve_main from tensorrt_llm.llmapi import llm_args as llm_args_module +from tensorrt_llm.serve.middleware.user_middleware import _apply_fastapi_middlewares from .yaml_validation_harness import ( assert_no_default_valued_leaves, diff --git a/tests/unittest/serve/middleware/__init__.py b/tests/unittest/serve/middleware/__init__.py new file mode 100644 index 000000000000..467079831e16 --- /dev/null +++ b/tests/unittest/serve/middleware/__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/middleware/test_register.py b/tests/unittest/serve/middleware/test_register.py new file mode 100644 index 000000000000..03aa768ec2c1 --- /dev/null +++ b/tests/unittest/serve/middleware/test_register.py @@ -0,0 +1,193 @@ +# 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. +"""``register_middlewares`` reproduces each app's middleware stack exactly. + +The aggregated server passes ``adjusted_clock`` to both middlewares, the +disaggregated server to neither; ``--middleware`` entries are applied after the +server is built and stay outermost. +""" + +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +from fastapi import FastAPI, Request +from starlette.testclient import TestClient +from test_common.serve.fakes import UserProbeMiddleware + +from tensorrt_llm.serve.middleware.arrival_time import ServerArrivalTimeMiddleware +from tensorrt_llm.serve.middleware.perf_metrics import PerfMetricsMiddleware +from tensorrt_llm.serve.middleware.register import register_middlewares +from tensorrt_llm.serve.middleware.user_middleware import _apply_fastapi_middlewares +from tensorrt_llm.serve.perf_metrics import RETURN_METRICS_HEADER, SERVER_TIMING_HEADER + +pytestmark = pytest.mark.cpu_only + + +def _record(status="complete"): + # Copy of the record helper of tests/unittest/llmapi/apps/test_request_metrics.py: + # arrival 1.0 and first_scheduled 1.01 give Server-Timing server_queue;dur=10.000000. + return { + "schema_version": 1, + "request_id": "42", + "status": status, + "phases": { + "server": { + "timing_metrics": { + "arrival_time": 1.0, + "first_scheduled_time": 1.01, + "first_token_time": 1.02, + "last_token_time": 1.05, + "kv_cache_transfer_start": None, + "kv_cache_transfer_end": None, + }, + "time_breakdown_metrics": { + "step_metrics": [ + { + "iter": 3, + "forward_start_time": 2.0, + "forward_end_time": 2.002, + "sample_start_time": 2.002, + "sample_end_time": 2.003, + "gpu_forward_time": 1.5, + "gpu_sample_time": 0.5, + } + ], + "ctx_chunk_metrics": [ + { + "forward_start_time": 1.0, + "forward_end_time": 1.004, + "sample_start_time": 1.004, + "sample_end_time": 1.005, + "gpu_forward_time": 3.0, + "gpu_sample_time": 0.25, + } + ], + }, + } + }, + } + + +def test_aggregated_stack_is_unchanged(): + app = FastAPI() + w = SimpleNamespace(submit=Mock()) + c = SimpleNamespace(now=lambda: 123.25) + register_middlewares( + app, + collect_perf_metrics=True, + expose_perf_metrics=True, + perf_metrics_writer=w, + adjusted_clock=c, + ) + assert [(m.cls, m.kwargs) for m in app.user_middleware] == [ + (ServerArrivalTimeMiddleware, {"adjusted_clock": c}), + (PerfMetricsMiddleware, {"expose_headers": True, "writer": w, "adjusted_clock": c}), + ], ( + "register_middlewares changed the aggregated stack; it must equal the two " + "add_middleware calls OpenAIServer made before, in the same order and with the same " + "keyword arguments" + ) + + +def test_disaggregated_stack_passes_no_clock(): + app = FastAPI() + w = SimpleNamespace(submit=Mock()) + register_middlewares( + app, + collect_perf_metrics=True, + expose_perf_metrics=True, + perf_metrics_writer=w, + ) + assert [(m.cls, m.kwargs) for m in app.user_middleware] == [ + (ServerArrivalTimeMiddleware, {}), + (PerfMetricsMiddleware, {"expose_headers": True, "writer": w}), + ], "the disaggregated app passes no adjusted_clock to either middleware" + + +def test_no_perf_metrics_installs_only_the_arrival_time_middleware(): + app = FastAPI() + register_middlewares( + app, + collect_perf_metrics=False, + expose_perf_metrics=False, + perf_metrics_writer=None, + ) + assert [m.cls for m in app.user_middleware] == [ServerArrivalTimeMiddleware] + + +def test_user_middleware_is_outermost(): + app = FastAPI() + w = SimpleNamespace(submit=Mock()) + c = SimpleNamespace(now=lambda: 123.25) + register_middlewares( + app, + collect_perf_metrics=True, + expose_perf_metrics=True, + perf_metrics_writer=w, + adjusted_clock=c, + ) + _apply_fastapi_middlewares(app, ["test_common.serve.fakes.UserProbeMiddleware"]) + assert [m.cls for m in app.user_middleware] == [ + UserProbeMiddleware, + ServerArrivalTimeMiddleware, + PerfMetricsMiddleware, + ], "--middleware entries must wrap the built-in middleware" + + +def test_records_reach_the_writer_and_headers(): + app = FastAPI() + writer = SimpleNamespace(submit=Mock()) + record = _record() + register_middlewares( + app, + collect_perf_metrics=True, + expose_perf_metrics=True, + perf_metrics_writer=writer, + ) + + @app.get("/probe") + async def probe(request: Request) -> dict: + request.state.perf_metrics_records.append(record) + return {} + + with TestClient(app) as client: + response = client.get("/probe", headers={RETURN_METRICS_HEADER: "1"}) + assert response.status_code == 200 + assert "server_queue;dur=10.000000" in response.headers.get(SERVER_TIMING_HEADER, ""), ( + "PerfMetricsMiddleware no longer sees the records handlers append to " + "request.state.perf_metrics_records; check that both middlewares share scope['state'] " + "and that their order is unchanged" + ) + writer.submit.assert_called_once_with(record) + + +def test_arrival_time_is_set(): + app = FastAPI() + register_middlewares( + app, + collect_perf_metrics=False, + expose_perf_metrics=False, + perf_metrics_writer=None, + adjusted_clock=SimpleNamespace(now=lambda: 123.25), + ) + + @app.get("/probe") + async def probe(request: Request) -> dict: + return {"server_arrival_time": request.state.server_arrival_time} + + with TestClient(app) as client: + response = client.get("/probe") + assert response.json() == {"server_arrival_time": 123.25} diff --git a/tests/unittest/serve/test_layering.py b/tests/unittest/serve/test_layering.py index 6c4fb30bccc8..009b50e2fb4c 100644 --- a/tests/unittest/serve/test_layering.py +++ b/tests/unittest/serve/test_layering.py @@ -82,6 +82,20 @@ "allowance when its import goes away." ), ), + Contract( + name="serve/middleware/ depends on no server, endpoint or Responses module", + importers=("tensorrt_llm.serve.middleware",), + forbidden=( + "tensorrt_llm.serve.api", + "tensorrt_llm.serve.openai_server", + "tensorrt_llm.serve.openai_disagg_server", + "tensorrt_llm.serve.responses_utils", + ), + fix=( + "Middleware is shared by both apps: pass what it needs into register_middlewares " + "from the composition root instead of importing a server or an endpoint package." + ), + ), ) diff --git a/tests/unittest/serve/test_middleware_order.py b/tests/unittest/serve/test_middleware_order.py index aec17dd4f79d..aa0fd889a095 100644 --- a/tests/unittest/serve/test_middleware_order.py +++ b/tests/unittest/serve/test_middleware_order.py @@ -124,7 +124,7 @@ def test_arrival_time_stays_outside_perf_metrics(name, serve_env): 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 + from tensorrt_llm.serve.middleware.user_middleware import _apply_fastapi_middlewares server = serve_env.build_openai_server(AggConfig()) _apply_fastapi_middlewares(server.app, ["test_common.serve.fakes.UserProbeMiddleware"]) diff --git a/tests/unittest/serve/test_serve_test_helpers.py b/tests/unittest/serve/test_serve_test_helpers.py index aa70512d9986..8e4e688e6c51 100644 --- a/tests/unittest/serve/test_serve_test_helpers.py +++ b/tests/unittest/serve/test_serve_test_helpers.py @@ -158,7 +158,7 @@ def test_conflict_rule(record_a, record_b, conflict): def test_middleware_snapshot_format(): - from tensorrt_llm.serve.responses_utils import ServerArrivalTimeMiddleware + from tensorrt_llm.serve.middleware.arrival_time import ServerArrivalTimeMiddleware app = FastAPI() app.add_middleware(ServerArrivalTimeMiddleware, adjusted_clock=None) From a0aa0dfcb75bb27ba1f77ded71b7bdb12a78ea04 Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Mon, 2 Nov 2026 09:54:19 +0000 Subject: [PATCH 08/26] [TRTLLM-17262][refactor] ServerArrivalTimeMiddleware keeps an existing request state Replace the assignment scope["state"] = {} with scope.setdefault("state", {}) in ServerArrivalTimeMiddleware.__call__. State already present when a request reaches this middleware - written by a user --middleware entry, or copied in by the ASGI server from a lifespan that yields a dict - is now kept, and server_arrival_time is added to it. Before, that state was replaced by an empty dict. Neither trtllm-serve app yields lifespan state today and the built-in middleware order is unchanged, so no built-in path changes behavior. Two tests pin the new behavior (state from an outer middleware and from a lifespan survive); both fail on the previous commit. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- tensorrt_llm/serve/middleware/arrival_time.py | 2 +- .../serve/middleware/test_register.py | 69 +++++++++++++++++++ tests/unittest/serve/test_middleware_order.py | 12 ++-- 3 files changed, 76 insertions(+), 7 deletions(-) diff --git a/tensorrt_llm/serve/middleware/arrival_time.py b/tensorrt_llm/serve/middleware/arrival_time.py index a0290227e69b..088b8f33623e 100644 --- a/tensorrt_llm/serve/middleware/arrival_time.py +++ b/tensorrt_llm/serve/middleware/arrival_time.py @@ -35,7 +35,7 @@ def __init__(self, async def __call__(self, scope, receive, send): if scope["type"] == "http": # Add arrival time to scope - scope["state"] = {} + scope.setdefault("state", {}) scope["state"]["server_arrival_time"] = self._adjusted_clock.now() # Pass through the original receive/send - no wrapping! diff --git a/tests/unittest/serve/middleware/test_register.py b/tests/unittest/serve/middleware/test_register.py index 03aa768ec2c1..35be18163d35 100644 --- a/tests/unittest/serve/middleware/test_register.py +++ b/tests/unittest/serve/middleware/test_register.py @@ -191,3 +191,72 @@ async def probe(request: Request) -> dict: with TestClient(app) as client: response = client.get("/probe") assert response.json() == {"server_arrival_time": 123.25} + + +_SETDEFAULT_FAILURE = ( + "ServerArrivalTimeMiddleware replaced scope['state']; it must use setdefault, so state " + "set by outer middleware or the lifespan survives" +) + + +class _OuterStateMiddleware: + """Pure ASGI middleware that writes into the request state before the built-ins run.""" + + def __init__(self, app): + self.app = app + + async def __call__(self, scope, receive, send): + if scope["type"] == "http": + scope.setdefault("state", {})["outer"] = 1 + await self.app(scope, receive, send) + + +def test_state_from_an_outer_middleware_survives(): + app = FastAPI() + register_middlewares( + app, + collect_perf_metrics=False, + expose_perf_metrics=False, + perf_metrics_writer=None, + ) + app.add_middleware(_OuterStateMiddleware) + + @app.get("/probe") + async def probe(request: Request) -> dict: + return { + "outer": request.state.outer, + "has_arrival_time": hasattr(request.state, "server_arrival_time"), + } + + with TestClient(app) as client: + response = client.get("/probe") + assert response.status_code == 200 and response.json() == { + "outer": 1, + "has_arrival_time": True, + }, _SETDEFAULT_FAILURE + + +def test_lifespan_state_survives(): + from contextlib import asynccontextmanager + + @asynccontextmanager + async def lifespan(app): + yield {"from_lifespan": 1} + + app = FastAPI(lifespan=lifespan) + register_middlewares( + app, + collect_perf_metrics=False, + expose_perf_metrics=False, + perf_metrics_writer=None, + ) + + @app.get("/probe") + async def probe(request: Request) -> dict: + return {"from_lifespan": request.state.from_lifespan} + + with TestClient(app) as client: + response = client.get("/probe") + assert response.status_code == 200 and response.json() == {"from_lifespan": 1}, ( + _SETDEFAULT_FAILURE + ) diff --git a/tests/unittest/serve/test_middleware_order.py b/tests/unittest/serve/test_middleware_order.py index aa0fd889a095..1d4419c87961 100644 --- a/tests/unittest/serve/test_middleware_order.py +++ b/tests/unittest/serve/test_middleware_order.py @@ -15,10 +15,10 @@ """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 +outermost. ``ServerArrivalTimeMiddleware`` and ``PerfMetricsMiddleware`` both +``setdefault`` into ``scope['state']``; the order stays load-bearing because the +arrival time must be stamped before ``PerfMetricsMiddleware`` runs. VisualGen is +not listed: it takes the default code path without perf metrics, and no VisualGen path can change middleware registration. """ @@ -95,8 +95,8 @@ 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 " + "outermost. ServerArrivalTimeMiddleware must stay outside PerfMetricsMiddleware, so " + "the arrival time is stamped before PerfMetricsMiddleware runs. 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." From 68344ed711c3021d9cab017cccee69c87c665031 Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Mon, 2 Nov 2026 05:38:33 +0000 Subject: [PATCH 09/26] [TRTLLM-17261][refactor] Move the msgpack route classes to tensorrt_llm/serve/_base/msgspec_route.py Generated move: _msgpack_decoder, _MsgspecRequest and _MsgspecRoute leave tensorrt_llm/serve/openai_server.py byte for byte; the module imports _MsgspecRoute back for its app router and keeps no re-export (the three names are private). The two disaggregated test importers now import the new module. OpenAIServer.health, .version and .get_model are removed here and replaced by the system router added in the following commits. tensorrt_llm/serve/_base/msgspec_route.py joins legacy-files.txt (its code comes from a legacy-formatted file), and the generated lint configs are regenerated with scripts/legacy_utils.py gen-configs. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- .pre-commit-config.yaml | 2 + legacy-files.txt | 1 + pyproject.toml | 1 + ruff-legacy.toml | 1 + tensorrt_llm/serve/_base/msgspec_route.py | 34 +++++++++ tensorrt_llm/serve/openai_server.py | 75 ++----------------- .../disaggregated/test_coordinator_e2e.py | 2 +- .../disaggregated/test_msgspec_route.py | 2 +- tests/unittest/llmapi/test_kv_hints.py | 4 +- 9 files changed, 49 insertions(+), 73 deletions(-) create mode 100644 tensorrt_llm/serve/_base/msgspec_route.py diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index c2166e7bb04b..148d141826d9 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -417,6 +417,7 @@ 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/chat_utils.py | tensorrt_llm/serve/cluster_storage.py | tensorrt_llm/serve/disagg_auto_scaling.py | @@ -1144,6 +1145,7 @@ 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/chat_utils.py | tensorrt_llm/serve/cluster_storage.py | tensorrt_llm/serve/disagg_auto_scaling.py | diff --git a/legacy-files.txt b/legacy-files.txt index 9850e7c810b4..e6dd3e54e62c 100644 --- a/legacy-files.txt +++ b/legacy-files.txt @@ -409,6 +409,7 @@ 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/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..7bc96a689072 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -474,6 +474,7 @@ 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/chat_utils.py", "tensorrt_llm/serve/cluster_storage.py", "tensorrt_llm/serve/disagg_auto_scaling.py", diff --git a/ruff-legacy.toml b/ruff-legacy.toml index dab55b4d0434..1bd2bbfab0e5 100644 --- a/ruff-legacy.toml +++ b/ruff-legacy.toml @@ -426,6 +426,7 @@ 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/chat_utils.py", "tensorrt_llm/serve/cluster_storage.py", "tensorrt_llm/serve/disagg_auto_scaling.py", diff --git a/tensorrt_llm/serve/_base/msgspec_route.py b/tensorrt_llm/serve/_base/msgspec_route.py new file mode 100644 index 000000000000..24b4e955ef85 --- /dev/null +++ b/tensorrt_llm/serve/_base/msgspec_route.py @@ -0,0 +1,34 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +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/openai_server.py b/tensorrt_llm/serve/openai_server.py index eb79cdb18397..df7c6476fff7 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -25,7 +25,6 @@ from typing import (TYPE_CHECKING, Annotated, Any, AsyncGenerator, AsyncIterator, Dict, List, Optional, Tuple, Union) -import msgspec import uvicorn from fastapi import Body, Depends, FastAPI, HTTPException, Query, Request from fastapi.exceptions import RequestValidationError @@ -67,6 +66,7 @@ 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.msgspec_route import _MsgspecRoute from tensorrt_llm.serve._perf_metrics_config import resolve_perf_metrics_config from tensorrt_llm.serve.anthropic_adapter import ( AnthropicRequestError, AnthropicResponseError, anthropic_error_response, @@ -106,12 +106,11 @@ CompletionResponse, CompletionResponseChoice, EmbeddingRequest, EmbeddingResponse, EmbeddingResponseData, EmbeddingUsageInfo, ErrorResponse, 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, @@ -157,7 +156,6 @@ 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 +177,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 @@ -1707,32 +1678,6 @@ def _on_launcher_lost(self, error: BaseException) -> None: _record_generator_termination(self.generator) signal.raise_signal(signal.SIGINT) - async def health(self) -> Response: - if self._check_health(): - return Response(status_code=200) - 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 - # produces tokens. Uses SIGINT (not SIGTERM) to match the - # existing CppExecutorError handlers and let uvicorn perform - # its graceful shutdown sequence. The doing_shutdown guard - # ensures we only send SIGINT once — repeated health probes - # won't stack signals during an in-progress teardown. - executor = getattr(self.generator, '_executor', None) - if executor is not None and getattr(executor, '_fatal_error', - None) is not None: - if not getattr(executor, 'doing_shutdown', True): - logger.error( - "Health check detected fatal engine error, initiating " - 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." - ) - async def health_generate(self, raw_request: Request) -> Response: """Health check that performs a minimal generation.""" extra_args = {} @@ -1782,14 +1727,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 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_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/llmapi/test_kv_hints.py b/tests/unittest/llmapi/test_kv_hints.py index 218e7482d8a6..e73a6097d2d8 100644 --- a/tests/unittest/llmapi/test_kv_hints.py +++ b/tests/unittest/llmapi/test_kv_hints.py @@ -739,9 +739,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 +832,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 From 9f4a48ec95e578a25f60886ab476f19601ba0d4a Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Mon, 2 Nov 2026 05:53:54 +0000 Subject: [PATCH 10/26] [TRTLLM-17261][refactor] Add the trtllm-serve router skeleton and serve the system endpoints through it New structure: tensorrt_llm/serve/_base/ (shared serving infrastructure) and tensorrt_llm/serve/api/ (endpoint packages, one router plus one serving class each; the package __init__ imports nothing so routers load lazily, by role). serve/_base/routing.py provides make_router(), which always uses the msgpack-decoding route class, and attach_routers(), which imports and attaches only the routers whose RouterSpec matches the server role. The first endpoint package, serve/api/ops/system/, serves /health, /version and /v1/models for every role. OpenAIServer builds a SystemServing with the injected _health_probe (today's probe body, SIGINT side effect included) and attaches the routers before the role dispatch; the 12 add_api_route lines for the three paths leave the four role tables. Same routes, methods, route classes, bodies and conflict-relevant order: the recorded route reference is unchanged. The layering tests gain the new packages' contracts, the route-snapshot failure message points new endpoints at the router pattern, and the serve agent docs describe where a new endpoint goes. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- tensorrt_llm/serve/AGENTS.md | 17 ++ tensorrt_llm/serve/_base/__init__.py | 15 + tensorrt_llm/serve/_base/msgspec_route.py | 2 + tensorrt_llm/serve/_base/routing.py | 103 +++++++ tensorrt_llm/serve/api/__init__.py | 15 + tensorrt_llm/serve/api/ops/__init__.py | 15 + tensorrt_llm/serve/api/ops/system/__init__.py | 15 + .../serve/api/ops/system/api_router.py | 60 ++++ tensorrt_llm/serve/api/ops/system/serving.py | 44 +++ tensorrt_llm/serve/openai_server.py | 47 +++- tests/test_common/serve/routes.py | 5 +- tests/unittest/serve/_base/__init__.py | 14 + tests/unittest/serve/_base/test_routing.py | 264 ++++++++++++++++++ .../api/ops/system/test_system_router.py | 65 +++++ .../test_ops_system_router_visual_gen.py | 48 ++++ tests/unittest/serve/test_layering.py | 63 +++++ 16 files changed, 778 insertions(+), 14 deletions(-) create mode 100644 tensorrt_llm/serve/_base/__init__.py create mode 100644 tensorrt_llm/serve/_base/routing.py create mode 100644 tensorrt_llm/serve/api/__init__.py create mode 100644 tensorrt_llm/serve/api/ops/__init__.py create mode 100644 tensorrt_llm/serve/api/ops/system/__init__.py create mode 100644 tensorrt_llm/serve/api/ops/system/api_router.py create mode 100644 tensorrt_llm/serve/api/ops/system/serving.py create mode 100644 tests/unittest/serve/_base/__init__.py create mode 100644 tests/unittest/serve/_base/test_routing.py create mode 100644 tests/unittest/serve/api/ops/system/test_system_router.py create mode 100644 tests/unittest/serve/api/visual_gen/test_ops_system_router_visual_gen.py diff --git a/tensorrt_llm/serve/AGENTS.md b/tensorrt_llm/serve/AGENTS.md index bdf1fe8c23b6..1e47ff2651b7 100644 --- a/tensorrt_llm/serve/AGENTS.md +++ b/tensorrt_llm/serve/AGENTS.md @@ -12,6 +12,14 @@ 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. - 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`. @@ -33,6 +41,9 @@ Name files and symbols here, never line numbers. ## 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. @@ -64,6 +75,12 @@ Name files and symbols here, never line numbers. 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): 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/msgspec_route.py b/tensorrt_llm/serve/_base/msgspec_route.py index 24b4e955ef85..a7eaf33acc64 100644 --- a/tensorrt_llm/serve/_base/msgspec_route.py +++ b/tensorrt_llm/serve/_base/msgspec_route.py @@ -1,5 +1,7 @@ # 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 diff --git a/tensorrt_llm/serve/_base/routing.py b/tensorrt_llm/serve/_base/routing.py new file mode 100644 index 000000000000..3f461a029282 --- /dev/null +++ b/tensorrt_llm/serve/_base/routing.py @@ -0,0 +1,103 @@ +# 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 + + @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 + + +@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), +) + + +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/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/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_server.py b/tensorrt_llm/serve/openai_server.py index df7c6476fff7..92f2d9478baf 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -67,6 +67,7 @@ get_effective_kv_cache_event_hash_algo from tensorrt_llm.sampling_params import GuidedDecodingParams, SamplingParams from tensorrt_llm.serve._base.msgspec_route import _MsgspecRoute +from tensorrt_llm.serve._base.routing import ServerCapabilities, attach_routers from tensorrt_llm.serve._perf_metrics_config import resolve_perf_metrics_config from tensorrt_llm.serve.anthropic_adapter import ( AnthropicRequestError, AnthropicResponseError, anthropic_error_response, @@ -81,6 +82,7 @@ AnthropicCountTokensResponse, AnthropicCreateBatchRequest, AnthropicMessagesRequest) +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) @@ -878,6 +880,16 @@ 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) + attach_routers(self.app, ServerCapabilities.for_role(self.server_role)) + # 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" @@ -1194,12 +1206,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"]) @@ -1344,9 +1353,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, @@ -1451,9 +1457,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"]) @@ -1558,10 +1561,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"]) @@ -1678,6 +1677,28 @@ def _on_launcher_lost(self, error: BaseException) -> None: _record_generator_termination(self.generator) signal.raise_signal(signal.SIGINT) + def _health_probe(self) -> bool: + if self._check_health(): + 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 + # produces tokens. Uses SIGINT (not SIGTERM) to match the + # existing CppExecutorError handlers and let uvicorn perform + # its graceful shutdown sequence. The doing_shutdown guard + # ensures we only send SIGINT once — repeated health probes + # won't stack signals during an in-progress teardown. + executor = getattr(self.generator, '_executor', None) + if executor is not None and getattr(executor, '_fatal_error', + None) is not None: + if not getattr(executor, 'doing_shutdown', True): + logger.error( + "Health check detected fatal engine error, initiating " + f"server shutdown: {executor._fatal_error}") + _record_generator_termination(self.generator) + signal.raise_signal(signal.SIGINT) + return False + async def health_generate(self, raw_request: Request) -> Response: """Health check that performs a minimal generation.""" extra_args = {} diff --git a/tests/test_common/serve/routes.py b/tests/test_common/serve/routes.py index 0c983b8022a6..b4dc02823074 100644 --- a/tests/test_common/serve/routes.py +++ b/tests/test_common/serve/routes.py @@ -401,7 +401,10 @@ def format_failure(name: str, differences: Sequence[str]) -> str: "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/). Do not add routes or methods to " + "(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." ) 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_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/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/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/test_layering.py b/tests/unittest/serve/test_layering.py index 6c4fb30bccc8..cefd3e87d818 100644 --- a/tests/unittest/serve/test_layering.py +++ b/tests/unittest/serve/test_layering.py @@ -19,6 +19,7 @@ records the imports that existed when the contract was written and only shrinks. """ +import ast import textwrap from pathlib import Path @@ -82,6 +83,25 @@ "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="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." + ), + ), ) @@ -98,6 +118,49 @@ 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 test_api_packages_do_not_import_each_other(edges): + problems = [ + f"serve/api/ packages do not import each other: {edge.importer} (line {edge.lineno}) " + f"imports `{edge.imported}`. Share code through serve/_base/ or inject it from the " + "composition root." + for edge in edges + if _under_api(edge.importer) + and _under_api(edge.imported) + and not edge.type_checking + and _api_package_of(edge.importer) != _api_package_of(edge.imported) + ] + assert problems == [] + + +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. # --------------------------------------------------------------------------- From caf3a29974e1799d30172b1ac9f4588d7105aba2 Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Mon, 2 Nov 2026 06:02:39 +0000 Subject: [PATCH 11/26] [TRTLLM-17261][infra] Multi-GPU triggers and type checking for the new serve packages tensorrt_llm/serve/_base/ and tensorrt_llm/serve/api/ops/ join the multi-GPU trigger list: the msgpack route class decodes the disaggregated frontend's worker requests, and /health is the disaggregated readiness probe, so changes under these directories keep running the multi-GPU stages that exercise that path. The three new modules join the strict mypy file set, and the type-check pre-commit hook now also triggers on changes under tensorrt_llm/serve/ and tensorrt_llm/commands/serve_cli/, so edits to typed serving modules get local type checking; the hook still checks only the declared file set. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- .pre-commit-config.yaml | 5 +++-- jenkins/L0_MergeRequest.groovy | 4 ++++ pyproject.toml | 3 +++ 3 files changed, 10 insertions(+), 2 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 148d141826d9..4aa5b0f82bcd 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1622,6 +1622,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..541ec1cb7626 100644 --- a/jenkins/L0_MergeRequest.groovy +++ b/jenkins/L0_MergeRequest.groovy @@ -1948,6 +1948,10 @@ 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/", + // /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/pyproject.toml b/pyproject.toml index 7bc96a689072..d17ac01fee08 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -863,6 +863,9 @@ 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/routing.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", From 4a6325da005f235bd045ec21d783312c2c17f712 Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Mon, 2 Nov 2026 14:52:33 +0000 Subject: [PATCH 12/26] [TRTLLM-17261][test] Register the system router in the multi-frontend tests that build routes by hand Two multi-frontend tests build an OpenAIServer without its constructor and call register_routes() themselves. Since /health, /version and /v1/models are attached through the system router by the constructor, those tests no longer saw the three routes: one asserted that /version stays local on an attached frontend (404), the other reported the three per-frontend-safe entries as removed. Both helpers now set the serving object and attach the routers before the hand registration, in the constructor's order, and the route enumeration walks included routers, which FastAPI 0.139 keeps as objects without a path in app.routes. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- .../llmapi/test_serve_multi_frontend_helpers.py | 8 +++++++- .../llmapi/test_serve_multi_frontend_route_gate.py | 13 +++++++++++-- 2 files changed, 18 insertions(+), 3 deletions(-) diff --git a/tests/unittest/llmapi/test_serve_multi_frontend_helpers.py b/tests/unittest/llmapi/test_serve_multi_frontend_helpers.py index 0a6f3b9055bb..add864bb829d 100644 --- a/tests/unittest/llmapi/test_serve_multi_frontend_helpers.py +++ b/tests/unittest/llmapi/test_serve_multi_frontend_helpers.py @@ -35,7 +35,9 @@ 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.ops.system.serving import SystemServing from tensorrt_llm.serve.multi_frontend import ( FORWARDED_ROUTES, AttachedFrontendWatchdog, @@ -111,7 +113,11 @@ 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 system router + # (/health, /version, /v1/models) first, then the launcher forwarding + # routes, then the local table. + server.app.state.ops_serving_system = SystemServing(model="m", health_probe=lambda: True) + attach_routers(server.app, ServerCapabilities.for_role(None)) server._register_launcher_forwarding_routes() 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..f419d5458da0 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,10 @@ 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.ops.system.serving import SystemServing from tensorrt_llm.serve.multi_frontend import FORWARDED_ROUTES from tensorrt_llm.serve.openai_server import OpenAIServer @@ -73,10 +76,16 @@ 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 system router (/health, /version, + # /v1/models) is attached for every role before the role's own table. + server.app.state.ops_serving_system = SystemServing(model="m", health_probe=lambda: True) + 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" From 7e2950968a462863e9ea1a4ae33b08e590c2db3b Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Mon, 2 Nov 2026 15:10:33 +0000 Subject: [PATCH 13/26] [TRTLLM-17261][test] Match a request instead of reading route paths in the OpenEngine metrics test The HTTP-enablement case builds a real OpenAIServer through its constructor and looked for /prometheus/metrics by reading .path off every entry of app.routes. Since the constructor attaches the system router, FastAPI 0.139 keeps an included-router object there that has no .path, and the child process failed with AttributeError. The test now matches a GET /prometheus/metrics request against the routes with the public route-matching API, which works on both ends of the pinned FastAPI range. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- tests/unittest/grpc/openengine/test_metrics.py | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) 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 From d3aacdd27485b3a9da6fd8fab4e3c3d241e50af1 Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Mon, 2 Nov 2026 16:01:01 +0000 Subject: [PATCH 14/26] [TRTLLM-17261][test] Classify endpoint-package imports by the imported name The check that endpoint packages under tensorrt_llm/serve/api do not import each other read the package off the from-module alone, so `from tensorrt_llm.serve import api` inside an endpoint package passed and `from tensorrt_llm.serve.api.ops import system` inside ops/system failed. It now reads the package off . per imported name, as the layering contracts do, and a synthetic-edge test pins both forms. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- tests/unittest/serve/test_layering.py | 72 +++++++++++++++++++++++---- 1 file changed, 61 insertions(+), 11 deletions(-) diff --git a/tests/unittest/serve/test_layering.py b/tests/unittest/serve/test_layering.py index cefd3e87d818..357dcd4b02fc 100644 --- a/tests/unittest/serve/test_layering.py +++ b/tests/unittest/serve/test_layering.py @@ -26,6 +26,7 @@ import pytest from test_common.serve.layering import ( Contract, + ImportEdge, build_import_graph, check_contracts, stale_allowances, @@ -138,18 +139,67 @@ 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): - problems = [ - f"serve/api/ packages do not import each other: {edge.importer} (line {edge.lineno}) " - f"imports `{edge.imported}`. Share code through serve/_base/ or inject it from the " - "composition root." - for edge in edges - if _under_api(edge.importer) - and _under_api(edge.imported) - and not edge.type_checking - and _api_package_of(edge.importer) != _api_package_of(edge.imported) - ] - assert problems == [] + 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(): From 543322fc974f21b8370768b3c5b1f56ab34c11a4 Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Tue, 3 Nov 2026 00:32:20 +0000 Subject: [PATCH 15/26] [TRTLLM-17263][test] Pin the chat completion handler's HTTP answers before splitting it Three tests run OpenAIServer.openai_chat through a server built with the real constructor and the fake generator, and pin what it answers today: a request the handler rejects keeps its 400 body byte for byte, an engine failure keeps 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. They pass on the unsplit handler and must still pass once it is split into a typed core and an HTTP wrapper. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- .../serve/test_openai_server_chat_core.py | 227 ++++++++++++++++++ 1 file changed, 227 insertions(+) create mode 100644 tests/unittest/serve/test_openai_server_chat_core.py 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..3c9c782bb0fc --- /dev/null +++ b/tests/unittest/serve/test_openai_server_chat_core.py @@ -0,0 +1,227 @@ +# 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 signal +from http import HTTPStatus +from types import SimpleNamespace +from unittest.mock import AsyncMock, Mock + +import pytest +from test_common.serve.app_factory import AggConfig +from test_common.serve.fakes import ScriptedOutput + +from tensorrt_llm.executor import CppExecutorError +from tensorrt_llm.serve import openai_server +from tensorrt_llm.serve.openai_protocol import ( + 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): + 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}" + + +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 json.loads(json.dumps(response.body.decode())) == "Generation health check OK" From 6174ee4b71effb4c81b8626549e6d72114ce4bae Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Tue, 3 Nov 2026 00:39:56 +0000 Subject: [PATCH 16/26] [TRTLLM-17263][refactor] Move the chat SSE generator to tensorrt_llm/serve/_base/sse.py Generated move: the chat_stream_generator closure of OpenAIServer.openai_chat is lifted to module level in tensorrt_llm/serve/_base/sse.py unchanged, apart from its two free variables, which become the keyword-only parameters `server` (was `self`) and `raw_request`. openai_chat imports it back and passes server=self, raw_request=raw_request at its one call. The generator keeps its `except Exception:` clause and its two `_postprocess_result` reads. tensorrt_llm/serve/_base/sse.py joins legacy-files.txt (its code comes from a legacy-formatted file) and the generated lint configs are regenerated with scripts/legacy_utils.py gen-configs; ruff-legacy-baseline.json gains the module's one recorded F821 (the `anext` builtin ruff does not know), which openai_server.py keeps counting until its baseline is regenerated. The Anthropic stream tests patched the logger of openai_server to check what the generator logs; they now patch the logger of serve/_base/sse.py, where the generator lives. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- .pre-commit-config.yaml | 2 + legacy-files.txt | 1 + pyproject.toml | 1 + ruff-legacy-baseline.json | 3 ++ ruff-legacy.toml | 1 + tensorrt_llm/serve/_base/sse.py | 44 +++++++++++++++++++ tensorrt_llm/serve/openai_server.py | 44 ++++--------------- .../llmapi/apps/test_anthropic_routes.py | 5 ++- 8 files changed, 63 insertions(+), 38 deletions(-) create mode 100644 tensorrt_llm/serve/_base/sse.py diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 4aa5b0f82bcd..aeabd348c165 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -418,6 +418,7 @@ common-files: &common_files | 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/chat_utils.py | tensorrt_llm/serve/cluster_storage.py | tensorrt_llm/serve/disagg_auto_scaling.py | @@ -1146,6 +1147,7 @@ legacy-files: &legacy_files | 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/chat_utils.py | tensorrt_llm/serve/cluster_storage.py | tensorrt_llm/serve/disagg_auto_scaling.py | diff --git a/legacy-files.txt b/legacy-files.txt index e6dd3e54e62c..c1903c47e3f8 100644 --- a/legacy-files.txt +++ b/legacy-files.txt @@ -410,6 +410,7 @@ 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/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 d17ac01fee08..aa271dd8ff02 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -475,6 +475,7 @@ exclude = [ "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/chat_utils.py", "tensorrt_llm/serve/cluster_storage.py", "tensorrt_llm/serve/disagg_auto_scaling.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 1bd2bbfab0e5..4ccd76cfbd7c 100644 --- a/ruff-legacy.toml +++ b/ruff-legacy.toml @@ -427,6 +427,7 @@ include = [ "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/chat_utils.py", "tensorrt_llm/serve/cluster_storage.py", "tensorrt_llm/serve/disagg_auto_scaling.py", diff --git a/tensorrt_llm/serve/_base/sse.py b/tensorrt_llm/serve/_base/sse.py new file mode 100644 index 000000000000..6b91472ef45d --- /dev/null +++ b/tensorrt_llm/serve/_base/sse.py @@ -0,0 +1,44 @@ +# 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, + raw_request) -> 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) + raw_request.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, 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 diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py index 92f2d9478baf..c706284f957d 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -22,8 +22,8 @@ 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, Tuple, Union) import uvicorn from fastapi import Body, Depends, FastAPI, HTTPException, Query, Request @@ -68,6 +68,7 @@ from tensorrt_llm.sampling_params import GuidedDecodingParams, SamplingParams from tensorrt_llm.serve._base.msgspec_route import _MsgspecRoute from tensorrt_llm.serve._base.routing import ServerCapabilities, attach_routers +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, @@ -154,7 +155,7 @@ 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 @@ -2022,38 +2023,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. @@ -2457,7 +2426,10 @@ async def chat_stream_generator( if request.stream: response_generator = chat_stream_generator( - promise, postproc_params) + promise, + postproc_params, + server=self, + raw_request=raw_request) return StreamingResponse(content=response_generator, media_type="text/event-stream") else: diff --git a/tests/unittest/llmapi/apps/test_anthropic_routes.py b/tests/unittest/llmapi/apps/test_anthropic_routes.py index b7e6cf62a3d6..82620eb57553 100644 --- a/tests/unittest/llmapi/apps/test_anthropic_routes.py +++ b/tests/unittest/llmapi/apps/test_anthropic_routes.py @@ -27,7 +27,7 @@ from fastapi.testclient import TestClient 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.openai_disagg_server import OpenAIDisaggServer from tensorrt_llm.serve.openai_protocol import ( ChatCompletionRequest, @@ -371,8 +371,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 From bd6b3d734eb85b1349927b1e1b8ace131aa8e1b0 Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Tue, 3 Nov 2026 00:59:25 +0000 Subject: [PATCH 17/26] [TRTLLM-17263][refactor] Split openai_chat into a typed create_chat_completion and an HTTP wrapper OpenAIServer.create_chat_completion is the former body of openai_chat with a typed contract: it returns the ChatCompletionResponse for a non-streaming request and the SSE stream as strings for a streaming one, and it raises ServingError (the HTTP status apart from the OpenAI error body) where it used to return an error response, or EngineFatalError after the three side effects of an engine failure. Nothing else in the body changes: the six early rejections raise where they returned, the ValueError and the generic handler raise what they rendered, and the ServingError handler comes first so an early rejection is neither logged nor re-wrapped. openai_chat keeps its name, signature and route and becomes the wrapper that renders the typed result: JSONResponse for the response, a streaming response for the stream, the error body with its status for ServingError, and HTTP 200 with body null after an engine failure, as before. Its own ValueError branch keeps the 400 for a response that cannot be JSON-encoded, which used to fail inside the core. Every caller (the route, /health_generate, the Anthropic messages and batch paths) is unchanged. tensorrt_llm/serve/_base/errors.py holds error_body, ServingError, serving_error and EngineFatalError; create_error_response builds its body with error_body, so a ServingError body and a rendered error can never differ (the over-length-prompt code lives in one place). Tests: the typed core and its errors through a server built with the real constructor; the error builders against create_error_response; the lifted SSE generator with fakes. The three answer-pinning tests added before the split pass unchanged. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- tensorrt_llm/serve/_base/errors.py | 76 +++++++++ tensorrt_llm/serve/openai_server.py | 99 +++++++----- tests/unittest/serve/_base/test_errors.py | 70 ++++++++ tests/unittest/serve/_base/test_sse.py | 153 ++++++++++++++++++ .../serve/test_openai_server_chat_core.py | 120 +++++++++++++- 5 files changed, 481 insertions(+), 37 deletions(-) create mode 100644 tensorrt_llm/serve/_base/errors.py create mode 100644 tests/unittest/serve/_base/test_errors.py create mode 100644 tests/unittest/serve/_base/test_sse.py 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/openai_server.py b/tensorrt_llm/serve/openai_server.py index c706284f957d..982861aa69fe 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -42,8 +42,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, @@ -66,6 +64,8 @@ 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.routing import ServerCapabilities, attach_routers from tensorrt_llm.serve._base.sse import chat_stream_generator @@ -107,7 +107,7 @@ ChatCompletionResponse, ChatCompletionResponseChoice, ChatCompletionToolsParam, ChatMessage, CompletionRequest, CompletionResponse, CompletionResponseChoice, EmbeddingRequest, - EmbeddingResponse, EmbeddingResponseData, EmbeddingUsageInfo, ErrorResponse, + EmbeddingResponse, EmbeddingResponseData, EmbeddingUsageInfo, ImageEditRequest, ImageGenerationRequest, ImageGenerationResponse, ImageObject, MemoryUpdateRequest, PromptTokensDetails, ResponseFormat, ResponsesRequest, ResponsesResponse, RuntimeMemoryUpdateRequest, @@ -1132,15 +1132,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) @@ -2015,6 +2007,42 @@ 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, 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, raw_request: Request + ) -> 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: @@ -2078,7 +2106,7 @@ def get_role() -> str: 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", @@ -2096,7 +2124,7 @@ def get_role() -> str: # 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."), @@ -2111,7 +2139,7 @@ def get_role() -> str: # 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."), @@ -2134,14 +2162,13 @@ def get_role() -> str: # 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 @@ -2194,7 +2221,7 @@ def get_role() -> str: 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) @@ -2209,7 +2236,7 @@ def get_role() -> str: 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", @@ -2425,13 +2452,10 @@ def get_role() -> str: postproc_args.num_prompt_tokens = len(promise.prompt_token_ids) if request.stream: - response_generator = chat_stream_generator( - promise, - postproc_params, - server=self, - raw_request=raw_request) - return StreamingResponse(content=response_generator, - media_type="text/event-stream") + return chat_stream_generator(promise, + postproc_params, + server=self, + raw_request=raw_request) else: response = await self._create_chat_response( promise, postproc_params, raw_request, disaggregated_params) @@ -2447,17 +2471,20 @@ def get_role() -> str: 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)) + raise serving_error(str(e)) from e async def anthropic_messages(self, request: AnthropicMessagesRequest, raw_request: Request) -> Response: 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_sse.py b/tests/unittest/serve/_base/test_sse.py new file mode 100644 index 000000000000..b62aa6f53cb1 --- /dev/null +++ b/tests/unittest/serve/_base/test_sse.py @@ -0,0 +1,153 @@ +# 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, raw_request=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, raw_request=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, raw_request=_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/test_openai_server_chat_core.py b/tests/unittest/serve/test_openai_server_chat_core.py index 3c9c782bb0fc..1af6161999df 100644 --- a/tests/unittest/serve/test_openai_server_chat_core.py +++ b/tests/unittest/serve/test_openai_server_chat_core.py @@ -27,18 +27,23 @@ 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.openai_protocol import ( + ChatCompletionRequest, ChatCompletionResponse, ChatCompletionResponseChoice, ChatMessage, @@ -224,4 +229,117 @@ def test_health_generate_still_fails_with_500_on_engine_failure(serve_env, engin 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 json.loads(json.dumps(response.body.decode())) == "Generation health check OK" + 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, _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), _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 From 936cc5489204157d0305bb0377038ae66b624fff Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Tue, 3 Nov 2026 01:11:10 +0000 Subject: [PATCH 18/26] [TRTLLM-17263][infra] Type-check tensorrt_llm/serve/_base/errors.py The new error module is written from scratch, so it joins the mypy file list in pyproject.toml; scripts/run_mypy.sh passes with it (the strict global settings apply). tensorrt_llm/serve/_base/sse.py stays out: it is moved legacy code and joins type checking with its package later. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- pyproject.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/pyproject.toml b/pyproject.toml index aa271dd8ff02..86a343c97bc3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -864,6 +864,7 @@ 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/errors.py", "tensorrt_llm/serve/_base/routing.py", "tensorrt_llm/serve/api/ops/system/api_router.py", "tensorrt_llm/serve/api/ops/system/serving.py", From 974e4bdf2c36dff37c8988fc232a38b915b37c6d Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Tue, 3 Nov 2026 02:48:21 +0000 Subject: [PATCH 19/26] [TRTLLM-17263][test] Stub the chat response as a ChatCompletionResponse in the multimodal prompt test The chat handler now serializes a ChatCompletionResponse and streams any other result of its core, so a stand-in for _create_chat_response that only has model_dump() is no longer rendered as JSON. The test stubs a real ChatCompletionResponse instead; what it checks (the rendered prompt, the decoded image and chat_template_applied reaching the engine) is unchanged. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- .../test_openai_chat_multimodal_prompt.py | 20 ++++++++++++++++++- 1 file changed, 19 insertions(+), 1 deletion(-) 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"]) From d8c56b57a505fde637f26c789a78fa36fea53bf1 Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Tue, 3 Nov 2026 04:54:11 +0000 Subject: [PATCH 20/26] [TRTLLM-17263][test] Pin that an unencodable chat response is answered 400 without a traceback The HTTP wrapper answers a response that cannot be JSON-encoded from its own ValueError branch, which logs nothing; the generic handler would render the same 400 but log a traceback. The answer-pinning test now also asserts that no ERROR record is emitted, so dropping that branch fails the test. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- tests/unittest/serve/test_openai_server_chat_core.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/unittest/serve/test_openai_server_chat_core.py b/tests/unittest/serve/test_openai_server_chat_core.py index 1af6161999df..3ed4257650a8 100644 --- a/tests/unittest/serve/test_openai_server_chat_core.py +++ b/tests/unittest/serve/test_openai_server_chat_core.py @@ -206,7 +206,7 @@ def test_http_wrapper_keeps_every_answer(serve_env, engine_failure_recorder): 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): +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")) @@ -216,6 +216,9 @@ def test_unencodable_response_is_still_a_400(serve_env): 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): From 633a3325ffad13b447a203b8935b9fc655309585 Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Tue, 3 Nov 2026 10:34:55 +0000 Subject: [PATCH 21/26] [TRTLLM-17269][refactor] Add build_http_app and serve_uvicorn, the shared assembly of the trtllm-serve apps tensorrt_llm/serve/_base/app_builder.py holds the two functions both servers build on instead of a shared base class. build_http_app() assembles a FastAPI app in one fixed order: the lifespan, the optional route class before any route exists, the exception handlers, the routers in order, and the built-in middleware last so it stays outermost. serve_uvicorn() runs the app through the telemetry-aware uvicorn server, with an optional task that runs next to the server and receives it. The module is typed (added to the mypy file list) and imports only FastAPI, Starlette, uvicorn, the standard library and serve/_telemetry.py. Its tests pin the assembly order, the route class, the middleware stack order, the lifespan and exception-handler pass-through, every uvicorn argument, and the import boundary. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- pyproject.toml | 1 + tensorrt_llm/serve/_base/app_builder.py | 121 +++++++ .../unittest/serve/_base/test_app_builder.py | 297 ++++++++++++++++++ 3 files changed, 419 insertions(+) create mode 100644 tensorrt_llm/serve/_base/app_builder.py create mode 100644 tests/unittest/serve/_base/test_app_builder.py diff --git a/pyproject.toml b/pyproject.toml index fa2c9c96ce57..08b8de321a2f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -865,6 +865,7 @@ 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/app_builder.py", "tensorrt_llm/serve/_base/routing.py", "tensorrt_llm/serve/api/ops/system/api_router.py", "tensorrt_llm/serve/api/ops/system/serving.py", diff --git a/tensorrt_llm/serve/_base/app_builder.py b/tensorrt_llm/serve/_base/app_builder.py new file mode 100644 index 000000000000..bfb0564c93d8 --- /dev/null +++ b/tensorrt_llm/serve/_base/app_builder.py @@ -0,0 +1,121 @@ +# 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 assembly of the trtllm-serve HTTP apps (shared functions, no base class). + +``build_http_app()`` fixes the order in which an app is put together; ``serve_uvicorn()`` +runs it. Each server passes its own lifespan, route class, error handlers, routes and +middleware, so the differences between the aggregated server and the disaggregated +frontend stay explicit at the call site. + +This module imports only FastAPI, Starlette, uvicorn, the standard library and +``serve/_telemetry.py``: it must not import ``serve/api/`` or either server module. +""" + +import asyncio +import socket +from typing import Any, Callable, Coroutine, List, Optional, Sequence, Type, Union + +import uvicorn +from fastapi import FastAPI +from fastapi.routing import APIRoute +from starlette.types import Lifespan + +from tensorrt_llm.serve._telemetry import create_uvicorn_server + +__all__ = ["AppInstaller", "build_http_app", "serve_uvicorn"] + +# A callable that installs something (handlers, routes, middleware) into an app. +AppInstaller = Callable[[FastAPI], None] + + +def build_http_app( + *, + lifespan: Lifespan[FastAPI], + routers: Sequence[AppInstaller], + dialect_error_renderers: AppInstaller, + middleware: AppInstaller, + route_class: Optional[Type[APIRoute]] = None, +) -> FastAPI: + """Build a FastAPI app in the order trtllm-serve relies on. + + 1. ``FastAPI(lifespan=lifespan)``. + 2. ``route_class``, when given, on ``app.router``, before any route exists, so every + route registered with ``app.add_api_route`` uses it. The aggregated server passes + the msgpack-decoding route class and the disaggregated frontend the KV-hint policy + route class; ``None`` keeps FastAPI's plain ``APIRoute``. + 3. ``dialect_error_renderers(app)``: the exception handlers. + 4. Each of ``routers``, in order. Route registration order is load-bearing (of two + ``GET /metrics`` routes the first registered wins), and a router may add middleware + of its own (the Prometheus instrumentator does). + 5. ``middleware(app)`` last, so its middleware ends up outside any middleware a router added + (``add_middleware`` inserts at the front of the stack). + + User middleware from ``--middleware`` is not part of this: the CLI adds it to the + returned app after construction, outermost. + + Args: + lifespan: The app's lifespan context manager factory. + routers: Callables that register routes on the app, run in order. + dialect_error_renderers: Callable that registers the app's exception handlers. + middleware: Callable that registers the app's built-in middleware. + route_class: Route class for ``app.router``; ``None`` keeps FastAPI's default. + + Returns: + The assembled app, ready for ``serve_uvicorn``. + """ + app = FastAPI(lifespan=lifespan) + if route_class is not None: + app.router.route_class = route_class + dialect_error_renderers(app) + for install_routes in routers: + install_routes(app) + middleware(app) + return app + + +async def serve_uvicorn( + app: FastAPI, + *, + host: str, + port: int, + sockets: Optional[List[socket.socket]] = None, + log_level: Union[str, int], + timeout_keep_alive: int, + after_start: Optional[Callable[[uvicorn.Server], Coroutine[Any, Any, None]]] = None, +) -> None: + """Run ``app`` on uvicorn through ``create_uvicorn_server`` (the telemetry-aware server). + + Args: + app: The app to serve. + host: Bind host, passed to ``uvicorn.Config``. + port: Bind port, passed to ``uvicorn.Config``. + sockets: Pre-bound sockets to serve on instead of binding ``host:port``. + log_level: uvicorn's log level, passed through unchanged. + timeout_keep_alive: HTTP keep-alive timeout in seconds. + after_start: When given, runs as a task next to the server. It receives the + ``uvicorn.Server`` so it can wait for ``server.started`` or set + ``server.should_exit``. The task is not kept. + """ + config = uvicorn.Config( + app, + host=host, + port=port, + log_level=log_level, + timeout_keep_alive=timeout_keep_alive, + ) + server = create_uvicorn_server(config) + if after_start is not None: + asyncio.create_task(after_start(server)) + await server.serve(sockets=sockets) diff --git a/tests/unittest/serve/_base/test_app_builder.py b/tests/unittest/serve/_base/test_app_builder.py new file mode 100644 index 000000000000..a2f6cca8c70e --- /dev/null +++ b/tests/unittest/serve/_base/test_app_builder.py @@ -0,0 +1,297 @@ +# 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. +"""``build_http_app()`` and ``serve_uvicorn()`` of ``serve/_base/app_builder.py``. + +Both trtllm-serve apps are assembled by ``build_http_app`` in one fixed order (route +class, exception handlers, routers, middleware) and run by ``serve_uvicorn``. The tests +pin that order and the pass-through of every uvicorn argument with small callables +written here; the route and middleware-order snapshots pin what the real servers get. +""" + +import ast +import asyncio +import socket +import sys +from contextlib import asynccontextmanager +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +from fastapi import FastAPI +from fastapi.routing import APIRoute +from fastapi.testclient import TestClient + +from tensorrt_llm.serve._base import app_builder +from tensorrt_llm.serve._base.app_builder import build_http_app, serve_uvicorn + +pytestmark = pytest.mark.cpu_only + +_ORDER_FAILURE = "build_http_app order changed; the middleware and /metrics order depend on it" + + +@asynccontextmanager +async def _lifespan(app): + yield + + +async def _endpoint(): + return {} + + +def _nothing(app): + pass + + +class _CustomRoute(APIRoute): + """Stands in for the msgpack-decoding route class the aggregated server uses.""" + + +class _PureAsgiMiddleware: + def __init__(self, app): + self.app = app + + async def __call__(self, scope, receive, send): + await self.app(scope, receive, send) + + +class _MiddlewareA(_PureAsgiMiddleware): + pass + + +class _MiddlewareB(_PureAsgiMiddleware): + pass + + +class _MiddlewareC(_PureAsgiMiddleware): + pass + + +def _route(app, path): + # Routes added with app.add_api_route sit directly on app.router on both + # ends of the supported FastAPI range (only included routers are wrapped). + (route,) = [r for r in app.router.routes if getattr(r, "path", None) == path] + return route + + +def _build(**overrides): + kwargs = dict( + lifespan=_lifespan, + routers=(), + dialect_error_renderers=_nothing, + middleware=_nothing, + ) + kwargs.update(overrides) + return build_http_app(**kwargs) + + +def test_route_class_applies_to_routes_added_by_routers(): + def add_probe(app): + app.add_api_route("/probe", _endpoint, methods=["GET"]) + + custom = _build(routers=(add_probe,), route_class=_CustomRoute) + assert type(_route(custom, "/probe")) is _CustomRoute, ( + "build_http_app must set route_class before any router runs" + ) + + plain = _build(routers=(add_probe,)) + assert type(_route(plain, "/probe")) is APIRoute, ( + "without route_class, build_http_app must leave FastAPI's APIRoute in place: " + "the disaggregated frontend must not decode msgpack bodies on its public port" + ) + + +def test_assembly_order(): + calls = [] + apps_seen = [] + + def recorder(name): + def install(app): + calls.append(name) + apps_seen.append(app) + + return install + + app = _build( + routers=(recorder("router-1"), recorder("router-2")), + dialect_error_renderers=recorder("error-handlers"), + middleware=recorder("middleware"), + ) + assert calls == ["error-handlers", "router-1", "router-2", "middleware"], _ORDER_FAILURE + assert all(seen is app for seen in apps_seen), ( + "every installer must receive the app build_http_app returns" + ) + + +def test_middleware_installed_after_router_middleware(): + def router_with_middleware(app): + app.add_api_route("/probe", _endpoint, methods=["GET"]) + app.add_middleware(_MiddlewareA) + + def built_in_middleware(app): + app.add_middleware(_MiddlewareB) + app.add_middleware(_MiddlewareC) + + app = _build(routers=(router_with_middleware,), middleware=built_in_middleware) + assert [m.cls for m in app.user_middleware] == [_MiddlewareC, _MiddlewareB, _MiddlewareA], ( + "middleware(app) must run after the routers so its middleware is outermost" + ) + + +def test_lifespan_and_exception_handlers_are_the_given_ones(): + events = [] + + @asynccontextmanager + async def lifespan(app): + events.append("enter") + yield + events.append("exit") + + class _Boom(Exception): + pass + + async def boom(): + raise _Boom() + + async def render_boom(request, exc): + from fastapi.responses import JSONResponse + + return JSONResponse(status_code=418, content={"error": "rendered"}) + + def add_boom(app): + app.add_api_route("/boom", boom, methods=["GET"]) + + app = _build( + lifespan=lifespan, + routers=(add_boom,), + dialect_error_renderers=lambda app: app.add_exception_handler(_Boom, render_boom), + ) + with TestClient(app) as client: + assert events == ["enter"], "build_http_app must pass lifespan to FastAPI unchanged" + response = client.get("/boom") + assert events == ["enter", "exit"] + assert (response.status_code, response.json()) == (418, {"error": "rendered"}), ( + "the handler registered by dialect_error_renderers must render the exception" + ) + + +def _fake_uvicorn(monkeypatch, serve): + """Patch create_uvicorn_server where serve_uvicorn looks it up; return (fake server, configs).""" + configs = [] + server = SimpleNamespace(serve=serve, started=True, should_exit=False) + + def create(config): + configs.append(config) + return server + + monkeypatch.setattr(app_builder, "create_uvicorn_server", create) + return server, configs + + +def _count_create_task(monkeypatch): + """Count asyncio.create_task calls made by serve_uvicorn, without changing them.""" + create_task = Mock(side_effect=asyncio.create_task) + monkeypatch.setattr(app_builder, "asyncio", SimpleNamespace(create_task=create_task)) + return create_task + + +def test_serve_uvicorn_config_and_after_start(monkeypatch): + app = FastAPI() + started = asyncio.Event() + received = [] + serve_calls = [] + + async def after_start(server): + received.append(server) + started.set() + + async def serve(sockets=None): + serve_calls.append(sockets) + # The server keeps running while after_start works next to it. + await asyncio.wait_for(started.wait(), timeout=5) + + server, configs = _fake_uvicorn(monkeypatch, serve) + create_task = _count_create_task(monkeypatch) + with socket.socket() as sock: + asyncio.run( + serve_uvicorn( + app, + host="127.0.0.1", + port=8123, + sockets=[sock], + log_level="warning", + timeout_keep_alive=7, + after_start=after_start, + ) + ) + failure = ( + "serve_uvicorn must pass host/port/log_level/timeout_keep_alive/sockets through " + "and start after_start" + ) + (config,) = configs + assert config.app is app, failure + assert (config.host, config.port) == ("127.0.0.1", 8123), failure + assert config.log_level == "warning", failure + assert config.timeout_keep_alive == 7, failure + assert serve_calls == [[sock]], failure + assert received == [server], failure + assert create_task.call_count == 1, failure + + +def test_serve_uvicorn_passes_integer_log_level_unchanged(monkeypatch): + async def serve(sockets=None): + pass + + _, configs = _fake_uvicorn(monkeypatch, serve) + asyncio.run(serve_uvicorn(FastAPI(), host="h", port=1, log_level=30, timeout_keep_alive=5)) + assert configs[0].log_level == 30, ( + "serve_uvicorn must hand uvicorn whatever log level the server passes, unchanged" + ) + + +def test_no_after_start_task_when_none(monkeypatch): + serve_calls = [] + + async def serve(sockets=None): + serve_calls.append(sockets) + + _fake_uvicorn(monkeypatch, serve) + create_task = _count_create_task(monkeypatch) + asyncio.run(serve_uvicorn(FastAPI(), host="h", port=1, log_level="info", timeout_keep_alive=5)) + assert serve_calls == [None] + create_task.assert_not_called() + + +def test_app_builder_imports_only_its_allowed_dependencies(): + """``serve/_base/app_builder.py`` depends on nothing above ``serve/_base/``.""" + source = Path(app_builder.__file__).read_text() + imported = set() + for node in ast.walk(ast.parse(source)): + if isinstance(node, ast.Import): + imported.update(alias.name for alias in node.names) + elif isinstance(node, ast.ImportFrom): + imported.add(node.module or "") + allowed_roots = {"fastapi", "starlette", "uvicorn"} | set(sys.stdlib_module_names) + disallowed = sorted( + module + for module in imported + if module != "tensorrt_llm.serve._telemetry" and module.split(".")[0] not in allowed_roots + ) + assert disallowed == [], ( + f"serve/_base/app_builder.py imports {disallowed}; it may import only FastAPI, " + "Starlette, uvicorn, the standard library and serve/_telemetry.py. It must not " + "import serve/api/ or either server module: pass what it needs in from the server's " + "__init__ as a callable." + ) From 24f9585e8c7a657efb2e5c527461f3cacf770568 Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Tue, 3 Nov 2026 10:35:54 +0000 Subject: [PATCH 22/26] [TRTLLM-17269][refactor] Build the aggregated and disaggregated trtllm-serve apps with build_http_app and serve_uvicorn OpenAIServer and OpenAIDisaggServer construct their FastAPI app through build_http_app() and run uvicorn through serve_uvicorn(). Every difference between the two apps stays an explicit argument at the call site: the msgpack-decoding route class on the aggregated app and the KV-hint policy route class (KvHintsRoute) on the disaggregated frontend, both of which strip kv_hints from untrusted request bodies; the adjusted clock for the aggregated middleware and none for the disaggregated one; "info" and the 5 s keep-alive for the aggregated server, logger.level and the configured keep-alive for the disaggregated one; the worker-registration task after start only on the aggregated server. Both middleware installers read the server's perf-metrics attributes when they run, after the routes, as the aggregated server did before. Both servers keep reading TRTLLM_TRUST_KV_HINTS into app.state.trust_kv_hints, now right after the build; nothing reads the flag before the first request. No behavior change. The validation handlers keep their bodies and are registered with app.add_exception_handler instead of the decorator; the aggregated server's role dispatch goes unchanged into _attach_role_routes(app), and its post-start closure becomes the method _register_after_serving(server). On the disaggregated app the two built-in middlewares are now added after the routes instead of before; nothing in its route registration adds middleware, so the resulting stack is the same and the middleware-order snapshot is unchanged. Tests: tests/unittest/serve/test_server_uvicorn_settings.py builds both servers through the test factory and pins what each hands to uvicorn (log level, keep-alive, sockets, the post-start task only on the aggregated server) and what build_http_app must carry for each app: its route class on every API route and the KV-hint trust flag. The unit test that patched uvicorn through openai_disagg_server now patches it where serve_uvicorn looks it up and also checks the disaggregated log level. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- tensorrt_llm/serve/openai_disagg_server.py | 46 +++--- tensorrt_llm/serve/openai_server.py | 81 +++++----- .../test_openai_disagg_server.py | 11 +- .../serve/test_server_uvicorn_settings.py | 147 ++++++++++++++++++ 4 files changed, 225 insertions(+), 60 deletions(-) create mode 100644 tests/unittest/serve/test_server_uvicorn_settings.py diff --git a/tensorrt_llm/serve/openai_disagg_server.py b/tensorrt_llm/serve/openai_disagg_server.py index 90e1848cd72c..6808f38d579f 100644 --- a/tensorrt_llm/serve/openai_disagg_server.py +++ b/tensorrt_llm/serve/openai_disagg_server.py @@ -25,7 +25,6 @@ from typing import Callable, Optional import aiohttp -import uvicorn from fastapi import FastAPI, HTTPException, Request from fastapi.exceptions import RequestValidationError from fastapi.responses import JSONResponse, Response, StreamingResponse @@ -41,7 +40,7 @@ from tensorrt_llm.llmapi.disagg_utils import (DisaggServerConfig, MetadataServerConfig, ServerRole) from tensorrt_llm.logger import logger -from tensorrt_llm.serve._telemetry import create_uvicorn_server +from tensorrt_llm.serve._base.app_builder import build_http_app, serve_uvicorn from tensorrt_llm.serve.anthropic_adapter import (AnthropicRequestError, AnthropicResponseError, anthropic_error_response, @@ -290,20 +289,10 @@ async def lifespan(app) -> None: await self._service.teardown() await self._perf_metrics_writer.close() - self.app = FastAPI(lifespan=lifespan) - self.app.state.trust_kv_hints = os.getenv("TRTLLM_TRUST_KV_HINTS", "0") == "1" - self.app.router.route_class = KvHintsRoute - - register_middlewares(self.app, - collect_perf_metrics=self._collect_perf_metrics, - expose_perf_metrics=self._expose_perf_metrics, - perf_metrics_writer=self._perf_metrics_writer) - # Log request-body validation failures so a client/server schema mismatch # shows up server-side. Throttled (first, then every 1000th) to avoid # flooding the event loop when every request fails identically. self._val_err_n = 0 - @self.app.exception_handler(RequestValidationError) async def validation_exception_handler(request: Request, exc): self._perf_metrics_collector.validation_exceptions.inc() # Anthropic clients parse the error envelope, so every /v1/messages @@ -327,6 +316,26 @@ async def validation_exception_handler(request: Request, exc): f"(n={self._val_err_n}): {brief}") return JSONResponse(status_code=400, content={"error": str(exc)}) + # KvHintsRoute: the frontend's public port strips kv_hints from request + # bodies unless this process trusts its ingress (TRTLLM_TRUST_KV_HINTS=1); + # it does not decode msgpack bodies. The middleware takes no clock. + self.app = build_http_app( + lifespan=lifespan, + route_class=KvHintsRoute, + dialect_error_renderers=lambda app: app.add_exception_handler( + RequestValidationError, validation_exception_handler), + routers=(self._attach_routes,), + middleware=lambda app: register_middlewares( + app, + collect_perf_metrics=self._collect_perf_metrics, + expose_perf_metrics=self._expose_perf_metrics, + perf_metrics_writer=self._perf_metrics_writer), + ) + self.app.state.trust_kv_hints = os.getenv("TRTLLM_TRUST_KV_HINTS", "0") == "1" + + def _attach_routes(self, app: FastAPI) -> None: + # register_routes adds its routes through self.app. + self.app = app self.register_routes() def _create_client(self, router: Router, role: ServerRole, max_retries: int = 1) -> OpenAIClient: @@ -616,13 +625,12 @@ async def version(self) -> JSONResponse: return JSONResponse(content={"version": VERSION}) async def __call__(self, host: str, port: int, sockets: list[socket.socket] | None = None): - keep_alive_timeout = self._config.server_keep_alive_timeout - config = uvicorn.Config(self.app, - host=host, - port=port, - log_level=logger.level, - timeout_keep_alive=keep_alive_timeout) - await create_uvicorn_server(config).serve(sockets=sockets) + await serve_uvicorn(self.app, + host=host, + port=port, + sockets=sockets, + log_level=logger.level, + timeout_keep_alive=self._config.server_keep_alive_timeout) async def _sync_server_clock(self, server: str): """ Sync the ctx/gen server's steady clock with the disagg-server's steady clock (in case NTP service is not running). diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py index 026011c19bb3..bbb31c1f5658 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -66,6 +66,7 @@ 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.app_builder import build_http_app, serve_uvicorn from tensorrt_llm.serve._base.msgspec_route import _MsgspecRoute from tensorrt_llm.serve._base.routing import ServerCapabilities, attach_routers from tensorrt_llm.serve._perf_metrics_config import resolve_perf_metrics_config @@ -155,7 +156,6 @@ from .._utils import (AdjustedSteadyClock, get_global_steady_clock_now_in_seconds, nvtx_mark, set_prometheus_multiproc_dir) -from ._telemetry import create_uvicorn_server from .harmony_adapter import HarmonyAdapter, get_harmony_adapter if TYPE_CHECKING: @@ -849,12 +849,6 @@ async def lifespan(app: FastAPI): self._metrics_stats_executor.shutdown(wait=False, cancel_futures=True) - self.app = FastAPI(lifespan=lifespan) - self.app.state.trust_kv_hints = os.getenv("TRTLLM_TRUST_KV_HINTS", - "0") == "1" - self.app.router.route_class = _MsgspecRoute - - @self.app.exception_handler(RequestValidationError) async def validation_exception_handler(request, exc): if request.url.path.startswith("/v1/messages"): # Anthropic clients expect a 400 with the Anthropic error @@ -879,12 +873,31 @@ async def validation_exception_handler(request, exc): len(detail) - 4096) return JSONResponse(status_code=400, content={"error": detail}) + self.app = build_http_app( + lifespan=lifespan, + route_class=_MsgspecRoute, + dialect_error_renderers=lambda app: app.add_exception_handler( + RequestValidationError, validation_exception_handler), + routers=(self._attach_role_routes, ), + middleware=lambda app: register_middlewares( + app, + collect_perf_metrics=self._collect_perf_metrics, + expose_perf_metrics=self._expose_perf_metrics, + perf_metrics_writer=self._perf_metrics_writer, + adjusted_clock=self._adjusted_steady_clock), + ) + self.app.state.trust_kv_hints = os.getenv("TRTLLM_TRUST_KV_HINTS", + "0") == "1" + + def _attach_role_routes(self, app: FastAPI) -> None: + # The register_* methods below add their routes through self.app. + self.app = app # /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( + app.state.ops_serving_system = SystemServing( model=self.model, health_probe=self._health_probe) - attach_routers(self.app, ServerCapabilities.for_role(self.server_role)) + attach_routers(app, ServerCapabilities.for_role(self.server_role)) # Runs before the role dispatch for now. It moves after the dispatch # once the serving objects attach_router checks for are built per # role. @@ -908,12 +921,6 @@ async def validation_exception_handler(request, exc): self._register_launcher_forwarding_routes() self.register_routes() - register_middlewares(self.app, - collect_perf_metrics=self._collect_perf_metrics, - expose_perf_metrics=self._expose_perf_metrics, - perf_metrics_writer=self._perf_metrics_writer, - adjusted_clock=self._adjusted_steady_clock) - def _init_visual_gen(self): self.processor = None self.model_config = None @@ -4208,26 +4215,24 @@ async def __call__(self, self.binding_addr = f"http://{host}:{port}" self.host = host self.port = port - config = uvicorn.Config(self.app, - host=host, - port=port, - log_level="info", - timeout_keep_alive=TIMEOUT_KEEP_ALIVE) - server = create_uvicorn_server(config) - - async def _register_after_serving(): - with _StartupTimer("http_server_start"): - while not server.started: - await asyncio.sleep(0.1) - if self.disagg_cluster_worker: - try: - with _StartupTimer( - f"service_registration/{self.disagg_cluster_worker.worker_info.worker_id}" - ): - await self.disagg_cluster_worker.register_worker() - except Exception as e: - logger.error(f"Worker registration failed: {e}") - server.should_exit = True - - asyncio.create_task(_register_after_serving()) - await server.serve(sockets=sockets) + await serve_uvicorn(self.app, + host=host, + port=port, + sockets=sockets, + log_level="info", + timeout_keep_alive=TIMEOUT_KEEP_ALIVE, + after_start=self._register_after_serving) + + async def _register_after_serving(self, server: uvicorn.Server) -> None: + with _StartupTimer("http_server_start"): + while not server.started: + await asyncio.sleep(0.1) + if self.disagg_cluster_worker: + try: + with _StartupTimer( + f"service_registration/{self.disagg_cluster_worker.worker_info.worker_id}" + ): + await self.disagg_cluster_worker.register_worker() + except Exception as e: + logger.error(f"Worker registration failed: {e}") + server.should_exit = True diff --git a/tests/unittest/disaggregated/test_openai_disagg_server.py b/tests/unittest/disaggregated/test_openai_disagg_server.py index 4984a1ad749e..980bd75a1057 100644 --- a/tests/unittest/disaggregated/test_openai_disagg_server.py +++ b/tests/unittest/disaggregated/test_openai_disagg_server.py @@ -19,7 +19,8 @@ from starlette.datastructures import Headers from tensorrt_llm.llmapi.disagg_utils import ServerRole, extract_disagg_cfg -from tensorrt_llm.serve import openai_disagg_server +from tensorrt_llm.logger import logger +from tensorrt_llm.serve._base import app_builder from tensorrt_llm.serve.openai_disagg_server import OpenAIDisaggServer from tensorrt_llm.serve.openai_protocol import ( CompletionRequest, @@ -55,12 +56,16 @@ async def test_server_keep_alive_timeout_is_passed_to_uvicorn( config_factory = Mock(return_value=uvicorn_config) uvicorn_server = SimpleNamespace(serve=AsyncMock()) server_factory = Mock(return_value=uvicorn_server) - monkeypatch.setattr(openai_disagg_server.uvicorn, "Config", config_factory) - monkeypatch.setattr(openai_disagg_server.uvicorn, "Server", server_factory) + # Both servers run uvicorn through serve_uvicorn, which looks the names up + # on the uvicorn module imported by app_builder. + monkeypatch.setattr(app_builder.uvicorn, "Config", config_factory) + monkeypatch.setattr(app_builder.uvicorn, "Server", server_factory) await server(host="localhost", port=8000) assert config_factory.call_args.kwargs["timeout_keep_alive"] == expected_timeout + # The disaggregated frontend logs at the TRT-LLM logger's level, not at "info". + assert config_factory.call_args.kwargs["log_level"] == logger.level server_factory.assert_called_once_with(uvicorn_config) uvicorn_server.serve.assert_awaited_once_with(sockets=None) diff --git a/tests/unittest/serve/test_server_uvicorn_settings.py b/tests/unittest/serve/test_server_uvicorn_settings.py new file mode 100644 index 000000000000..7cb7207005f4 --- /dev/null +++ b/tests/unittest/serve/test_server_uvicorn_settings.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 uvicorn settings each server hands to ``serve_uvicorn`` when it is called. + +The two servers differ on purpose: the aggregated server logs at ``"info"``, keeps +connections alive for ``TIMEOUT_KEEP_ALIVE`` seconds and schedules the worker +registration next to the server; the disaggregated frontend logs at the TRT-LLM +logger's level, uses its configured keep-alive and schedules nothing. Both servers are +built through the real constructors; ``create_uvicorn_server`` is replaced where +``serve_uvicorn`` looks it up, so no port is bound. + +The last test pins what ``build_http_app`` must carry for each app: the aggregated +server registers its API routes through the msgpack-decoding ``_MsgspecRoute`` and +the disaggregated frontend through ``KvHintsRoute``, and both read +``TRTLLM_TRUST_KV_HINTS`` into ``app.state.trust_kv_hints``. +""" + +import asyncio +import os +import socket +from types import SimpleNamespace + +import pytest +from fastapi.routing import APIRoute +from test_common.serve.app_factory import AggConfig, DisaggConfig + +from tensorrt_llm.logger import logger +from tensorrt_llm.serve._base import app_builder +from tensorrt_llm.serve._base.msgspec_route import _MsgspecRoute +from tensorrt_llm.serve.kv_hints import KvHintsRoute +from tensorrt_llm.serve.openai_server import TIMEOUT_KEEP_ALIVE + +pytestmark = pytest.mark.cpu_only + + +class _FakeUvicorn: + """Records the ``uvicorn.Config`` and the tasks ``serve_uvicorn`` creates.""" + + def __init__(self, monkeypatch): + self.configs = [] + self.serve_sockets = [] + self.task_names = [] + self._tasks = [] + self.server = SimpleNamespace(serve=self._serve, started=True, should_exit=False) + monkeypatch.setattr(app_builder, "create_uvicorn_server", self._create) + monkeypatch.setattr(app_builder, "asyncio", SimpleNamespace(create_task=self._create_task)) + + def _create(self, config): + self.configs.append(config) + return self.server + + def _create_task(self, coroutine): + self.task_names.append(coroutine.cr_code.co_name) + task = asyncio.create_task(coroutine) + self._tasks.append(task) + return task + + async def _serve(self, sockets=None): + self.serve_sockets.append(sockets) + # A task started next to the server runs while the server serves. + if self._tasks: + await asyncio.gather(*self._tasks) + + +def test_aggregated_server_uvicorn_settings(serve_env, monkeypatch): + server = serve_env.build_openai_server(AggConfig()) + fake = _FakeUvicorn(monkeypatch) + with socket.socket() as sock: + asyncio.run(server("127.0.0.1", 8123, sockets=[sock])) + assert fake.serve_sockets == [[sock]] + (config,) = fake.configs + assert config.app is server.app + assert (config.host, config.port) == ("127.0.0.1", 8123) + assert config.log_level == "info", "the aggregated server runs uvicorn at log level info" + assert config.timeout_keep_alive == TIMEOUT_KEEP_ALIVE, ( + "the aggregated server keeps connections alive for TIMEOUT_KEEP_ALIVE seconds" + ) + assert fake.task_names == ["_register_after_serving"], ( + "the aggregated server schedules _register_after_serving next to the uvicorn server" + ) + assert (server.binding_addr, server.host, server.port) == ( + "http://127.0.0.1:8123", + "127.0.0.1", + 8123, + ) + + +def test_disaggregated_server_uvicorn_settings(serve_env, monkeypatch): + server = serve_env.build_disagg_server(DisaggConfig()) + fake = _FakeUvicorn(monkeypatch) + asyncio.run(server("127.0.0.1", 8123)) + assert fake.serve_sockets == [None] + (config,) = fake.configs + assert config.app is server.app + assert (config.host, config.port) == ("127.0.0.1", 8123) + assert config.log_level == logger.level, ( + "the disaggregated frontend runs uvicorn at the TRT-LLM logger's level, not at info" + ) + assert config.timeout_keep_alive == server._config.server_keep_alive_timeout, ( + "the disaggregated frontend uses the configured server_keep_alive_timeout" + ) + assert fake.task_names == [], "the disaggregated frontend schedules no post-start task" + + +@pytest.mark.parametrize("trust", [False, True], ids=["untrusted-ingress", "trusted-ingress"]) +def test_route_class_and_kv_hint_policy_survive_the_assembly(serve_env, trust): + """Each app keeps its route class and its KV-hint trust flag through build_http_app. + + The aggregated server decodes msgpack bodies through ``_MsgspecRoute`` (which also + applies the KV-hint policy); the disaggregated frontend applies the policy through + ``KvHintsRoute``. A plain ``APIRoute`` on either app would let ``kv_hints`` through + from an untrusted ingress. ``TRTLLM_TRUST_KV_HINTS=1`` turns the trust on; the + environment variable is restored at the serve environment's exit. + """ + if trust: + os.environ["TRTLLM_TRUST_KV_HINTS"] = "1" + else: + os.environ.pop("TRTLLM_TRUST_KV_HINTS", None) + aggregated = serve_env.build_openai_server(AggConfig()).app + disaggregated = serve_env.build_disagg_server(DisaggConfig()).app + assert aggregated.router.route_class is _MsgspecRoute + assert disaggregated.router.route_class is KvHintsRoute + for name, app, route_class in ( + ("aggregated", aggregated, _MsgspecRoute), + ("disaggregated", disaggregated, KvHintsRoute), + ): + api_routes = [route for route in app.routes if isinstance(route, APIRoute)] + assert api_routes, f"the {name} app registers API routes" + assert {type(route) for route in api_routes} == {route_class}, ( + f"every API route of the {name} app must be a {route_class.__name__}: " + f"{sorted({type(route).__name__ for route in api_routes})}" + ) + assert app.state.trust_kv_hints is trust, ( + f"the {name} app reads TRTLLM_TRUST_KV_HINTS into app.state.trust_kv_hints" + ) From 0cd435421f025384e1df87b88cbfc714727a3816 Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Tue, 3 Nov 2026 17:00:14 +0000 Subject: [PATCH 23/26] [TRTLLM-17269][test] Pin the validation-error answers of both trtllm-serve apps before rendering them per dialect Posts invalid bodies to servers built through the real constructors and checks status, body and side effects: the aggregated server's OpenAI shape with its 4096-character cap, the Anthropic envelope, the VisualGen 422 envelope and the metrics call; the disaggregated frontend's uncapped OpenAI shape, Anthropic envelope, counter and throttled log with escaped control characters. The validation text is read from the exception FastAPI raised, so the test does not re-implement the rendering. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- .../serve/exception_handling/__init__.py | 14 + .../test_server_validation_errors.py | 278 ++++++++++++++++++ 2 files changed, 292 insertions(+) create mode 100644 tests/unittest/serve/exception_handling/__init__.py create mode 100644 tests/unittest/serve/exception_handling/test_server_validation_errors.py diff --git a/tests/unittest/serve/exception_handling/__init__.py b/tests/unittest/serve/exception_handling/__init__.py new file mode 100644 index 000000000000..467079831e16 --- /dev/null +++ b/tests/unittest/serve/exception_handling/__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/exception_handling/test_server_validation_errors.py b/tests/unittest/serve/exception_handling/test_server_validation_errors.py new file mode 100644 index 000000000000..e8141a83e31a --- /dev/null +++ b/tests/unittest/serve/exception_handling/test_server_validation_errors.py @@ -0,0 +1,278 @@ +# 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 validation-error answers of both trtllm-serve apps, pinned on the real servers. + +Each test posts an invalid body to a server built through the real constructor and +checks status, body and side effects (the metrics call, the counter, the throttled log) +against today's answers. The validation text is read from the exception FastAPI raised +for that very request, so no test re-implements how a server renders it. + +The two servers differ on purpose and the tests use values that show it: the aggregated +server caps the OpenAI-shaped text at 4096 characters and counts a rejected request on +its metrics collector; the disaggregated frontend keeps the whole text, counts on its +own Prometheus counter and logs the first and every 1000th OpenAI-shaped rejection. +""" + +import asyncio +import logging +from unittest.mock import MagicMock, Mock + +import fastapi.routing +import pytest +from fastapi import Request +from fastapi.exceptions import RequestValidationError +from test_common.serve.app_factory import AggConfig, DisaggConfig + +from tensorrt_llm.llmapi.disagg_utils import ServerRole +from tensorrt_llm.logger import logger + +# threadleak is off: TestClients are created in the test body and closed by the serve_env +# fixture's teardown, which runs after the leak check. +pytestmark = [pytest.mark.cpu_only, pytest.mark.threadleak(enabled=False)] + +# Rejected by every request model: the model name is not a string, and the batch request's +# list of requests is not a list. +INVALID_BODY = {"model": 1.5, "requests": "not-a-list"} +# A body whose echo in the validation text exceeds the aggregated server's cap. +LONG_INVALID_BODY = {"model": "m" * 10_000} +AGG_DETAIL_LIMIT = 4096 +KEEP = "today's answer must be kept (status, body and side effects)" + + +def _openai_body(text): + return {"error": text} + + +def _anthropic_body(text): + return {"type": "error", "error": {"type": "invalid_request_error", "message": text}} + + +def _absent_collector(): + """A metrics collector the server treats as not configured (falsy), yet observable.""" + collector = MagicMock() + collector.__bool__.return_value = False + return collector + + +@pytest.fixture +def raised_validation_errors(monkeypatch): + """The RequestValidationError objects FastAPI raised, in order. + + FastAPI raises the error from its routing module; a recording subclass installed there + makes ``str(exc)`` of the very request available, so bodies are compared to the text the + server saw, not to a re-implementation of it. The app's handler, registered for the base + class, still receives the subclass. + """ + raised = [] + + class Recording(RequestValidationError): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + raised.append(self) + + monkeypatch.setattr(fastapi.routing, "RequestValidationError", Recording) + return raised + + +@pytest.fixture +def warning_records(monkeypatch): + """Messages the TRT-LLM logger emits at WARNING or above during the test. + + The library logger does not propagate, so ``caplog`` sees nothing; a handler is attached + instead, and the logger is lowered to ``warning`` for the test because its default level + drops warnings before any handler. + """ + records = [] + + class Capture(logging.Handler): + def emit(self, record): + records.append(record.getMessage()) + + handler = Capture(level=logging.WARNING) + previous_level = logger.level + monkeypatch.setattr(logger, "_set_from_env", False) + logger.set_level("warning") + logger._logger.addHandler(handler) + try: + yield records + finally: + logger._logger.removeHandler(handler) + logger.set_level(previous_level) + + +def _validation_warnings(records): + # The TRT-LLM logger prefixes each message with its own tags; match the text, not the start. + return [message for message in records if "[validation]" in message] + + +def _counter_value(counter) -> float: + return next( + sample.value for sample in counter.collect()[0].samples if sample.name.endswith("_total") + ) + + +# --------------------------------------------------------------------------- +# Aggregated server +# --------------------------------------------------------------------------- + +AGG_LLM_ROUTES = [ + ("/v1/chat/completions", _openai_body), + ("/v1/completions", _openai_body), + ("/v1/responses", _openai_body), + ("/v1/messages", _anthropic_body), + ("/v1/messages/count_tokens", _anthropic_body), + ("/v1/messages/batches", _anthropic_body), +] + + +@pytest.mark.parametrize("metrics", ["collector", "none"]) +@pytest.mark.parametrize("path, body_of", AGG_LLM_ROUTES, ids=[path for path, _ in AGG_LLM_ROUTES]) +def test_agg_llm_bodies_and_metrics(serve_env, raised_validation_errors, path, body_of, metrics): + server = serve_env.build_openai_server(AggConfig()) + collector = Mock() if metrics == "collector" else _absent_collector() + server.metrics_collector = collector + client = serve_env.client(server.app) + + response = client.post(path, json=INVALID_BODY) + + text = str(raised_validation_errors[-1]) + assert (response.status_code, response.json()) == (400, body_of(text)), f"{path}: {KEEP}" + if metrics == "collector": + collector.log_request_error.assert_called_once_with(http_code=400) + else: + collector.log_request_error.assert_not_called() + + +def test_agg_caps_only_the_openai_validation_text(serve_env, raised_validation_errors): + server = serve_env.build_openai_server(AggConfig()) + client = serve_env.client(server.app) + + response = client.post("/v1/chat/completions", json=LONG_INVALID_BODY) + text = str(raised_validation_errors[-1]) + assert len(text) > AGG_DETAIL_LIMIT, "the long body must make the validation text long" + cut = text[:AGG_DETAIL_LIMIT] + " ... [truncated %d bytes]" % (len(text) - AGG_DETAIL_LIMIT) + assert (response.status_code, response.json()) == (400, _openai_body(cut)), ( + f"OpenAI cap: {KEEP}" + ) + + response = client.post("/v1/messages", json=LONG_INVALID_BODY) + text = str(raised_validation_errors[-1]) + assert len(text) > AGG_DETAIL_LIMIT + assert (response.status_code, response.json()) == (400, _anthropic_body(text)), ( + f"the Anthropic envelope carries the whole text: {KEEP}" + ) + + +def test_agg_visual_gen_body(serve_env): + server = serve_env.build_openai_server(AggConfig(role=ServerRole.VISUAL_GEN)) + server.metrics_collector = Mock() + client = serve_env.client(server.app, lifespan=False) + + response = client.post("/v1/images/generations", json={"prompt": 123}) + assert (response.status_code, response.json()) == ( + 422, + { + "object": "error", + "message": "prompt: Input should be a valid string", + "type": "BadRequestError", + "param": None, + "code": 422, + }, + ), f"VisualGen envelope: {KEEP}" + + response = client.post("/v1/images/generations", json={"prompt": "a cat", "bogus": 1}) + assert response.status_code == 422 + assert response.json()["message"] == ( + "Unknown request field 'bogus'. Pass model-specific parameters via 'extra_params' instead." + ), f"VisualGen unknown-field message: {KEEP}" + + server.metrics_collector.log_request_error.assert_not_called() + + +# --------------------------------------------------------------------------- +# Disaggregated frontend +# --------------------------------------------------------------------------- + + +def test_disagg_bodies_counter_and_throttled_log( + serve_env, raised_validation_errors, warning_records +): + server = serve_env.build_disagg_server(DisaggConfig()) + client = serve_env.client(server.app, lifespan=False) + counter = server._perf_metrics_collector.validation_exceptions + rejected_before = _counter_value(counter) + + # An OpenAI route: the whole text (no cap), counted, logged as the first rejection with + # the compact error list rather than the echoed input. + response = client.post("/v1/chat/completions", json=LONG_INVALID_BODY) + text = str(raised_validation_errors[-1]) + assert len(text) > AGG_DETAIL_LIMIT + assert (response.status_code, response.json()) == (400, _openai_body(text)), ( + f"OpenAI body: {KEEP}" + ) + assert _counter_value(counter) == rejected_before + 1 + warnings = _validation_warnings(warning_records) + assert len(warnings) == 1, warnings + assert "[validation] POST /v1/chat/completions 400 (n=1): [{'loc':" in warnings[0], warnings + assert "m" * 100 not in warnings[0], "the log carries the error list, not the echoed input" + + # Throttled: the 1000th rejection is logged, the 1001st is not. + server._val_err_n = 999 + client.post("/v1/chat/completions", json=INVALID_BODY) + client.post("/v1/chat/completions", json=INVALID_BODY) + warnings = _validation_warnings(warning_records) + assert len(warnings) == 2, warnings + assert "(n=1000):" in warnings[1], warnings + assert _counter_value(counter) == rejected_before + 3 + + # An Anthropic route: the envelope, counted, never logged. The log counter is reset first, + # so an OpenAI-shaped rejection in its place would be the first one and would log: the + # Anthropic rejection neither logs nor advances that counter. + server._val_err_n = 0 + response = client.post("/v1/messages", json=INVALID_BODY) + text = str(raised_validation_errors[-1]) + assert (response.status_code, response.json()) == (400, _anthropic_body(text)), ( + f"Anthropic body: {KEEP}" + ) + assert _counter_value(counter) == rejected_before + 4 + assert server._val_err_n == 0, "an Anthropic rejection must not advance the OpenAI log counter" + assert len(_validation_warnings(warning_records)) == 2, "Anthropic rejections are not logged" + # Positive control: the next OpenAI-shaped rejection is the first one again and is logged. + client.post("/v1/chat/completions", json=INVALID_BODY) + assert server._val_err_n == 1 + warnings = _validation_warnings(warning_records) + assert len(warnings) == 3 and "(n=1):" in warnings[-1], warnings + + # Control characters in the method and the path are escaped in the log line. Such a path + # matches no route over HTTP, so the app's own handler is called directly. + server._val_err_n = 0 + handler = server.app.exception_handlers[RequestValidationError] + scope = { + "type": "http", + "method": "P\x01ST", + "path": "/v1/chat/completions\x7f", + "headers": [], + "query_string": b"", + "app": server.app, + } + error = RequestValidationError( + [{"type": "missing", "loc": ("body", "messages"), "msg": "Field required", "input": {}}] + ) + response = asyncio.run(handler(Request(scope), error)) + assert response.status_code == 400 + assert ( + "[validation] P\\x01ST /v1/chat/completions\\x7f 400 (n=1): " + in _validation_warnings(warning_records)[-1] + ), warning_records From fba31958519ce768748e592ecd55fc5c2622c17c Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Tue, 3 Nov 2026 17:00:30 +0000 Subject: [PATCH 24/26] [TRTLLM-17269][refactor] Add the dialect-neutral serving errors, the error boundary and the per-dialect error registry tensorrt_llm/serve/_base/errors.py gains the dialect-neutral error types a chat backend raises for any wire dialect to render (DialectNeutralError with InvalidRequestError, ContextLengthExceededError, UpstreamError, UnavailableError, InternalServingError; neutral_error picks the type for an HTTP status), ErrorRule and serving_error_boundary, the error ladder of an HTTP handler as a declaration, with CHAT_LADDER for the chat completion wrapper. The OpenAI body keeps its single builder, error_body. tensorrt_llm/serve/exception_handling/register.py is the per-app registry of error renderers keyed by (dialect, exception type): route_dialect marks an endpoint's dialect, dialect_of reads it from the matched endpoint, default_error_renderers gives the OpenAI and Anthropic renderers for validation errors and the typed errors, and install_exception_handlers puts the registry on app.state and installs the app-level validation handler with a per-server observer for side effects. The module imports no server or VisualGen code; the layering test holds that. Typed (mypy strict), with unit tests for the registry, the boundary and the neutral errors, including that the OpenAI rendering of a neutral error equals OpenAIServer.create_error_response and that the chat completion wrapper answers as before through the boundary. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- pyproject.toml | 1 + tensorrt_llm/serve/_base/errors.py | 278 +++++++++- .../serve/exception_handling/__init__.py | 14 + .../serve/exception_handling/register.py | 286 ++++++++++ .../serve/_base/test_serving_errors.py | 523 ++++++++++++++++++ .../test_dialect_error_renderers.py | 335 +++++++++++ tests/unittest/serve/test_layering.py | 19 + 7 files changed, 1452 insertions(+), 4 deletions(-) create mode 100644 tensorrt_llm/serve/exception_handling/__init__.py create mode 100644 tensorrt_llm/serve/exception_handling/register.py create mode 100644 tests/unittest/serve/_base/test_serving_errors.py create mode 100644 tests/unittest/serve/exception_handling/test_dialect_error_renderers.py diff --git a/pyproject.toml b/pyproject.toml index bacd8c111201..deb7b0ed383f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -871,6 +871,7 @@ files = [ "tensorrt_llm/serve/_base/routing.py", "tensorrt_llm/serve/api/ops/system/api_router.py", "tensorrt_llm/serve/api/ops/system/serving.py", + "tensorrt_llm/serve/exception_handling/register.py", "tensorrt_llm/serve/middleware/register.py", "tests/unittest/_torch/sampler/test_torch_sampler.py", "tests/unittest/_torch/sampler/test_token_ban.py", diff --git a/tensorrt_llm/serve/_base/errors.py b/tensorrt_llm/serve/_base/errors.py index e3fbd54b05a3..780ab4122e8c 100644 --- a/tensorrt_llm/serve/_base/errors.py +++ b/tensorrt_llm/serve/_base/errors.py @@ -17,30 +17,64 @@ 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"). + +The dialect-neutral errors (DialectNeutralError and its subclasses) are what a chat backend +raises for any wire dialect to render; ``neutral_error`` picks the type for an HTTP status. +``serving_error_boundary`` is the error ladder of an HTTP handler written as a declaration: it +renders the typed errors through the dialect's registered renderers +(``tensorrt_llm/serve/exception_handling/register.py``) and maps the handler's other exceptions +with its ``ErrorRule`` ladder. """ +import functools +import inspect +import traceback +from dataclasses import dataclass from http import HTTPStatus -from typing import Union +from typing import ( + TYPE_CHECKING, + Any, + Awaitable, + Callable, + Literal, + Optional, + Sequence, + Tuple, + Type, + TypeVar, + Union, + cast, +) + +from starlette.requests import Request +from starlette.responses import Response from tensorrt_llm.executor.utils import ( CONTEXT_LENGTH_EXCEEDED_CODE, is_context_length_exceeded_message, ) +from tensorrt_llm.logger import logger from tensorrt_llm.serve.openai_protocol import ErrorResponse +if TYPE_CHECKING: + from tensorrt_llm.serve.exception_handling.register import ErrorRenderers + +Handler = TypeVar("Handler", bound=Callable[..., Awaitable[Any]]) + def error_body( message: str, err_type: str = "BadRequestError", - status_code: HTTPStatus = HTTPStatus.BAD_REQUEST, + status_code: Union[HTTPStatus, int] = 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. + exception type survives to here. ``status_code`` may be a plain int for a status with no + HTTPStatus member (a worker's 499, for example). """ - code: Union[int, str] = status_code.value + code: Union[int, str] = int(status_code) if is_context_length_exceeded_message(message): code = CONTEXT_LENGTH_EXCEEDED_CODE return ErrorResponse(message=message, type=err_type, code=code) @@ -74,3 +108,239 @@ class EngineFatalError(Exception): It carries no response; each caller decides what to answer. """ + + +class DialectNeutralError(Exception): + """An error a chat backend raises for any wire dialect to render. + + Attributes: + message: The client-facing message. + status: The HTTP status to answer with. + err_type: The OpenAI error ``type``. + code: Overrides the OpenAI body's ``code`` when set; otherwise error_body decides. + """ + + status: int = 500 + err_type: str = "InternalServerError" + code: Optional[str] = None + + def __init__(self, message: str) -> None: + super().__init__(message) + self.message = message + + +class InvalidRequestError(DialectNeutralError): + """The request cannot be served as written (HTTP 400).""" + + status = 400 + err_type = "BadRequestError" + + +class ContextLengthExceededError(InvalidRequestError): + """The prompt does not fit the context window (HTTP 400, code ``context_length_exceeded``).""" + + code = CONTEXT_LENGTH_EXCEEDED_CODE + + +class UpstreamError(DialectNeutralError): + """A worker the request was forwarded to failed it. + + ``status`` is the worker's status, or 502 when it had none; a 4xx keeps the OpenAI ``type`` of + a rejected request, any other status that of a server error. + """ + + def __init__(self, message: str, status: Optional[int] = None) -> None: + super().__init__(message) + self.status = status or 502 + self.err_type = "BadRequestError" if 400 <= self.status < 500 else "InternalServerError" + + +class UnavailableError(DialectNeutralError): + """The server cannot take the request now (HTTP 503).""" + + status = 503 + + def __init__(self, message: str, retryable: bool = True) -> None: + super().__init__(message) + self.retryable = retryable + + +class InternalServingError(DialectNeutralError): + """The server failed while serving the request (HTTP 500).""" + + status = 500 + + +def neutral_error(status: int, message: str) -> DialectNeutralError: + """The dialect-neutral error for an HTTP status and a final message. + + 400 with the context-length message: ContextLengthExceededError; 400: InvalidRequestError; + 500: InternalServingError; 503: UnavailableError; any other status: UpstreamError(message, + status). Status and message are kept as given. + """ + if status == 400: + if is_context_length_exceeded_message(message): + return ContextLengthExceededError(message) + return InvalidRequestError(message) + if status == 500: + return InternalServingError(message) + if status == 503: + return UnavailableError(message) + return UpstreamError(message, status) + + +@dataclass(frozen=True) +class ErrorRule: + """One step of an error ladder. + + Attributes: + exc_type: The exception type (or types) the step matches, with their subclasses. + status: The HTTP status the answer carries. + err_type: The OpenAI error ``type`` of the answer. + log: What to log at ERROR: nothing, ``str(error)``, or the traceback. + """ + + exc_type: Union[Type[Exception], Tuple[Type[Exception], ...]] + status: HTTPStatus + err_type: str = "BadRequestError" + log: Literal["none", "message", "traceback"] = "none" + + +# The chat completion handler's ladder: an unencodable response (ValueError) is answered 400 +# without a log; any other failure is answered 400 with its traceback logged. +CHAT_LADDER: Tuple[ErrorRule, ...] = ( + ErrorRule(ValueError, HTTPStatus.BAD_REQUEST), + ErrorRule(Exception, HTTPStatus.BAD_REQUEST, log="traceback"), +) + + +def serving_error_boundary( + dialect: str, + *, + ladder: Sequence[ErrorRule] = (), + on_engine_fatal: Optional[Callable[[], Optional[Response]]] = None, + renderers: Optional["ErrorRenderers"] = None, +) -> Callable[[Handler], Handler]: + """The error ladder as a declaration on an async handler. + + On an exception from the handler, in this order: + + 1. ServingError and DialectNeutralError: rendered for ``dialect`` from their own fields + (the status travels with the error, never through ErrorResponse.code); + 2. EngineFatalError: ``on_engine_fatal()`` when given (chat: ``lambda: None``, which FastAPI + sends as HTTP 200 with body null), else re-raised; + 3. the first rule whose exc_type matches: the dialect's body with the rule's status and + err_type and str(error) as the message; logged as the rule says (message: + logger.error(str); traceback: logger.error(traceback.format_exc())); + 4. anything else is re-raised. + + The request handed to the renderer is the value of the handler's first parameter whose + annotation is starlette's Request (fastapi.Request is the same class) or, under + ``from __future__ import annotations``, the string ``"Request"``; else None. It is found by + annotation only: route functions also take a body model named ``request``. + + The registry is ``renderers`` when given; else the app of the request's ASGI scope + (``request.app.state.error_renderers``, which install_exception_handlers sets) when the + scope has one; else ``default_error_renderers()``. + + The wrapper keeps the handler's signature, name and docstring (``functools.wraps``), so a + route built from it has the same parameters, operationId and description. + """ + + def decorate(handler: Handler) -> Handler: + signature = inspect.signature(handler) + request_parameter = _request_parameter(signature) + + @functools.wraps(handler) + async def boundary(*args: Any, **kwargs: Any) -> Any: + try: + return await handler(*args, **kwargs) + except (ServingError, DialectNeutralError) as error: + request = _bound_request(signature, request_parameter, args, kwargs) + return _renderers_for(request, renderers).render(dialect, request, error) + except EngineFatalError: + if on_engine_fatal is None: + raise + return on_engine_fatal() + except Exception as error: + rule = _matching_rule(ladder, error) + if rule is None: + raise + _log(rule, error) + request = _bound_request(signature, request_parameter, args, kwargs) + rendered = _rule_error(rule, error) + return _renderers_for(request, renderers).render(dialect, request, rendered) + + return cast(Handler, boundary) + + return decorate + + +def _request_parameter(signature: inspect.Signature) -> Optional[str]: + """The name of the first parameter annotated as a Request, else None.""" + for parameter in signature.parameters.values(): + if parameter.annotation is Request or parameter.annotation == "Request": + return parameter.name + return None + + +def _bound_request( + signature: inspect.Signature, + name: Optional[str], + args: Tuple[Any, ...], + kwargs: dict[str, Any], +) -> Optional[Request]: + if name is None: + return None + try: + bound = signature.bind_partial(*args, **kwargs) + except TypeError: + return None + request: Optional[Request] = bound.arguments.get(name) + return request + + +def _renderers_for( + request: Optional[Request], renderers: Optional["ErrorRenderers"] +) -> "ErrorRenderers": + if renderers is not None: + return renderers + scope = getattr(request, "scope", None) + app = scope.get("app") if scope is not None else None + if app is None: + # Imported here: the registry module imports this one. + from tensorrt_llm.serve.exception_handling.register import default_error_renderers + + return default_error_renderers() + registry: Optional[ErrorRenderers] = getattr(app.state, "error_renderers", None) + if registry is None: + raise RuntimeError( + "the request's app has no error renderers: build the app with build_http_app, " + "or call install_exception_handlers on it before mounting handlers that use " + "serving_error_boundary" + ) + return registry + + +def _matching_rule(ladder: Sequence[ErrorRule], error: Exception) -> Optional[ErrorRule]: + for rule in ladder: + if isinstance(error, rule.exc_type): + return rule + return None + + +def _log(rule: ErrorRule, error: Exception) -> None: + # The TRT-LLM logger's methods carry no annotations. + if rule.log == "message": + logger.error(str(error)) # type: ignore[no-untyped-call] + elif rule.log == "traceback": + logger.error(traceback.format_exc()) # type: ignore[no-untyped-call] + + +def _rule_error(rule: ErrorRule, error: Exception) -> DialectNeutralError: + """The dialect-neutral error that renders ``error`` as ``rule`` says.""" + rendered = DialectNeutralError(str(error)) + rendered.status = rule.status.value + rendered.err_type = rule.err_type + rendered.__cause__ = error + return rendered diff --git a/tensorrt_llm/serve/exception_handling/__init__.py b/tensorrt_llm/serve/exception_handling/__init__.py new file mode 100644 index 000000000000..467079831e16 --- /dev/null +++ b/tensorrt_llm/serve/exception_handling/__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/tensorrt_llm/serve/exception_handling/register.py b/tensorrt_llm/serve/exception_handling/register.py new file mode 100644 index 000000000000..841b5e72a71d --- /dev/null +++ b/tensorrt_llm/serve/exception_handling/register.py @@ -0,0 +1,286 @@ +# 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. +"""Error renderers registered per wire dialect. + +A trtllm-serve app speaks up to three wire dialects: OpenAI, Anthropic and VisualGen. +Each renders an error in its own shape, so the renderer is chosen by the dialect of the +endpoint that matched the request, never by the request path. An endpoint declares its +dialect with ``@route_dialect(...)``; endpoints without a marker take the app's default, +which the composition root sets per server role. + +Every app keeps one :class:`ErrorRenderers` registry on ``app.state.error_renderers`` +(``install_exception_handlers`` puts it there). Code that registers later, for a dialect the +app only speaks in one role, reads it from ``app.state`` without a default: the VisualGen +role registers its validation renderer from its own route registration. + +Only request validation errors are rendered at app level, because they are raised before +the endpoint runs. Errors raised inside an endpoint are rendered where the dialect is known +statically: ``serving_error_boundary`` in ``tensorrt_llm/serve/_base/errors.py``. + +This module imports FastAPI, ``tensorrt_llm/serve/_base/errors.py`` and, inside the +Anthropic renderers only, the Anthropic adapter. It must not import either server module or +any VisualGen module, at any scope: VisualGen code is loaded only on the VisualGen role. +""" + +from typing import Any, Callable, Dict, Literal, Optional, Tuple, Type, TypeVar, cast + +from fastapi import FastAPI, Request +from fastapi.exceptions import RequestValidationError +from fastapi.responses import JSONResponse, Response + +from tensorrt_llm.serve._base.errors import DialectNeutralError, ServingError, error_body + +__all__ = [ + "ANTHROPIC", + "OPENAI", + "VISUAL_GEN", + "ErrorRenderer", + "ErrorRenderers", + "ValidationObserver", + "default_error_renderers", + "dialect_of", + "install_exception_handlers", + "openai_validation_error_renderer", + "route_dialect", +] + +OPENAI = "openai" +ANTHROPIC = "anthropic" +VISUAL_GEN = "visual_gen" + +# A renderer receives the request (None when the error is rendered outside a request) and +# the error; it may return any Response, a StreamingResponse included, so that an error found +# after a stream has opened can be answered inside the stream. +ErrorRenderer = Callable[[Optional[Request], BaseException], Response] +# Called before a validation error is rendered, with the dialect it is rendered for; side +# effects only (counters, logs). +ValidationObserver = Callable[[Request, RequestValidationError, str], None] + +F = TypeVar("F", bound=Callable[..., Any]) +E = TypeVar("E", bound=BaseException) + +_DIALECT_ATTRIBUTE = "__trtllm_dialect__" + + +def route_dialect(dialect: str) -> Callable[[F], F]: + """Mark an endpoint (function or method) as speaking ``dialect``. + + A bound method forwards attribute reads to its function, so a marked method stays + marked once it is bound to a server. + """ + + def mark(endpoint: F) -> F: + setattr(endpoint, _DIALECT_ATTRIBUTE, dialect) + return endpoint + + return mark + + +def dialect_of(request: Request, default: str) -> str: + """The dialect of the endpoint that matched ``request``, else ``default``. + + Starlette records the matched endpoint in ``scope["endpoint"]`` (FastAPI also records the + route in ``scope["route"]``) before the route runs, so both are set when an app-level + exception handler is called for a request-validation error. + """ + endpoint = request.scope.get("endpoint") + if endpoint is None: + endpoint = getattr(request.scope.get("route"), "endpoint", None) + dialect: str = getattr(endpoint, _DIALECT_ATTRIBUTE, default) + return dialect + + +class ErrorRenderers: + """Error renderers keyed by (dialect, exception type).""" + + def __init__(self) -> None: + self._renderers: Dict[Tuple[str, type], ErrorRenderer] = {} + + def register( + self, dialect: str, exc_type: Type[BaseException], renderer: ErrorRenderer + ) -> None: + """Register ``renderer`` for ``exc_type`` and its subclasses under ``dialect``.""" + self._renderers[(dialect, exc_type)] = renderer + + def has(self, dialect: str, exc_type: Type[BaseException]) -> bool: + """Whether a renderer is registered for exactly ``(dialect, exc_type)``.""" + return (dialect, exc_type) in self._renderers + + def render(self, dialect: str, request: Optional[Request], error: BaseException) -> Response: + """Render ``error`` for ``dialect``. + + Uses the renderer registered for the dialect and the closest class in + ``type(error).__mro__``; when the dialect has none for it, the OpenAI one. + + Raises: + KeyError: Neither the dialect nor the OpenAI dialect has a renderer for the error. + """ + renderer = self._lookup(dialect, type(error)) + if renderer is None and dialect != OPENAI: + renderer = self._lookup(OPENAI, type(error)) + if renderer is None: + raise KeyError( + f"no error renderer for dialect {dialect!r} and {type(error).__name__}; " + "register one with ErrorRenderers.register" + ) + return renderer(request, error) + + def _lookup(self, dialect: str, exc_type: type) -> Optional[ErrorRenderer]: + for cls in exc_type.__mro__: + renderer = self._renderers.get((dialect, cls)) + if renderer is not None: + return renderer + return None + + +def _renderer_for( + exc_type: Type[E], render: Callable[[Optional[Request], E], Response] +) -> ErrorRenderer: + """Adapt a renderer written for one exception type to the registry's signature.""" + + def renderer(request: Optional[Request], error: BaseException) -> Response: + if not isinstance(error, exc_type): + raise TypeError( + f"renderer for {exc_type.__name__} cannot render {type(error).__name__}" + ) + return render(request, error) + + return renderer + + +def openai_validation_error_renderer(detail_limit: Optional[int] = None) -> ErrorRenderer: + """The OpenAI rendering of a request-validation error: 400 with ``{"error": }``. + + ```` is ``str(exc)``. With ``detail_limit``, a longer text is cut at that many + characters and the cut is noted (pydantic echoes each failing input verbatim, which can be + megabytes of base64; the error locations are what matters). The aggregated server sets a + limit; the disaggregated frontend sets none. + """ + + def render(request: Optional[Request], exc: RequestValidationError) -> Response: + detail = str(exc) + if detail_limit is not None and len(detail) > detail_limit: + detail = detail[:detail_limit] + " ... [truncated %d bytes]" % ( + len(detail) - detail_limit + ) + return JSONResponse(status_code=400, content={"error": detail}) + + return _renderer_for(RequestValidationError, render) + + +def _render_openai_serving_error(request: Optional[Request], error: ServingError) -> Response: + return JSONResponse(content=error.body.model_dump(), status_code=error.status) + + +def _render_openai_neutral_error( + request: Optional[Request], error: DialectNeutralError +) -> Response: + body = error_body(error.message, error.err_type, error.status) + if error.code is not None: + body = body.model_copy(update={"code": error.code}) + return JSONResponse(content=body.model_dump(), status_code=error.status) + + +def _anthropic_error(message: str, status: int) -> Response: + # Imported here so that importing this module does not load the Anthropic adapter. + from tensorrt_llm.serve.anthropic_adapter import anthropic_error_response + + error_type: Literal["invalid_request_error", "api_error"] = ( + "invalid_request_error" if 400 <= status < 500 else "api_error" + ) + return anthropic_error_response(message, error_type, status) + + +def _render_anthropic_validation_error( + request: Optional[Request], exc: RequestValidationError +) -> Response: + return _anthropic_error(str(exc), 400) + + +def _render_anthropic_serving_error(request: Optional[Request], error: ServingError) -> Response: + return _anthropic_error(error.body.message, error.status) + + +def _render_anthropic_neutral_error( + request: Optional[Request], error: DialectNeutralError +) -> Response: + return _anthropic_error(error.message, error.status) + + +def default_error_renderers( + *, openai_validation_detail_limit: Optional[int] = None +) -> ErrorRenderers: + """The OpenAI and Anthropic renderers for validation errors and the typed serving errors. + + VisualGen's renderers are not here: VisualGen code registers them, for the VisualGen role + only. ``openai_validation_detail_limit`` is passed to + :func:`openai_validation_error_renderer`. + """ + renderers = ErrorRenderers() + renderers.register( + OPENAI, + RequestValidationError, + openai_validation_error_renderer(openai_validation_detail_limit), + ) + renderers.register( + OPENAI, ServingError, _renderer_for(ServingError, _render_openai_serving_error) + ) + renderers.register( + OPENAI, + DialectNeutralError, + _renderer_for(DialectNeutralError, _render_openai_neutral_error), + ) + renderers.register( + ANTHROPIC, + RequestValidationError, + _renderer_for(RequestValidationError, _render_anthropic_validation_error), + ) + renderers.register( + ANTHROPIC, ServingError, _renderer_for(ServingError, _render_anthropic_serving_error) + ) + renderers.register( + ANTHROPIC, + DialectNeutralError, + _renderer_for(DialectNeutralError, _render_anthropic_neutral_error), + ) + return renderers + + +def install_exception_handlers( + app: FastAPI, + *, + renderers: ErrorRenderers, + default_dialect: str, + on_validation_error: Optional[ValidationObserver] = None, +) -> None: + """Put ``renderers`` on ``app.state.error_renderers`` and install the validation handler. + + The handler resolves the request's dialect (``default_dialect`` for an unmarked + endpoint), calls ``on_validation_error`` with it, and renders the error with the registry + read from ``request.app.state`` at request time, so renderers registered after this call + are used too. + """ + app.state.error_renderers = renderers + + async def render_request_validation_error(request: Request, exc: Exception) -> Response: + # Registered for RequestValidationError only; Starlette types every handler on Exception. + validation_error = cast(RequestValidationError, exc) + dialect = dialect_of(request, default_dialect) + if on_validation_error is not None: + on_validation_error(request, validation_error, dialect) + registry: ErrorRenderers = request.app.state.error_renderers + return registry.render(dialect, request, validation_error) + + app.add_exception_handler(RequestValidationError, render_request_validation_error) diff --git a/tests/unittest/serve/_base/test_serving_errors.py b/tests/unittest/serve/_base/test_serving_errors.py new file mode 100644 index 000000000000..0bf5054d76a2 --- /dev/null +++ b/tests/unittest/serve/_base/test_serving_errors.py @@ -0,0 +1,523 @@ +# 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 dialect-neutral errors and the serving error boundary (``tensorrt_llm/serve/_base/errors.py``). + +The boundary renders the typed errors through the registry of +``tensorrt_llm/serve/exception_handling/register.py`` and maps a handler's other exceptions +with its ladder. ``test_chat_wrapper_answers_as_before`` runs the chat completion handler on a +server built through the real constructor and checks that each of its former ``except`` +branches still gives the same answer. +""" + +import asyncio +import inspect +import json +import logging +import signal +from http import HTTPStatus +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest +from fastapi import FastAPI, Request +from fastapi.responses import JSONResponse, Response +from pydantic import BaseModel +from test_common.serve.app_factory import AggConfig +from test_common.serve.fakes import ScriptedOutput + +from tensorrt_llm.executor import CppExecutorError +from tensorrt_llm.executor.utils import context_length_exceeded_message +from tensorrt_llm.logger import logger +from tensorrt_llm.serve import openai_server +from tensorrt_llm.serve._base.errors import ( + CHAT_LADDER, + ContextLengthExceededError, + DialectNeutralError, + EngineFatalError, + ErrorRule, + InternalServingError, + InvalidRequestError, + ServingError, + UnavailableError, + UpstreamError, + error_body, + neutral_error, + serving_error, + serving_error_boundary, +) +from tensorrt_llm.serve.exception_handling import register as register_module +from tensorrt_llm.serve.exception_handling.register import ( + ANTHROPIC, + OPENAI, + ErrorRenderers, + default_error_renderers, + install_exception_handlers, +) +from tensorrt_llm.serve.openai_protocol import ( + ChatCompletionResponse, + ChatCompletionResponseChoice, + ChatMessage, + UsageInfo, +) +from tensorrt_llm.serve.openai_server import OpenAIServer + +# threadleak is off: TestClients are created in the test body and closed by the serve_env +# fixture's teardown, which runs after the leak check. +pytestmark = [pytest.mark.cpu_only, pytest.mark.threadleak(enabled=False)] + +CONTEXT_LENGTH_MESSAGE = context_length_exceeded_message(2048, 4096) +MODEL = "fake-model" +REQUEST_RULE = ( + "serving_error_boundary must find the request by its Request annotation and read " + "app.state.error_renderers only when the scope has an app" +) + + +class _Body(BaseModel): + value: int + + +def _openai_body(message, err_type, code): + return {"object": "error", "message": message, "type": err_type, "param": None, "code": code} + + +def _anthropic_body(message, status): + error_type = "invalid_request_error" if 400 <= status < 500 else "api_error" + return {"type": "error", "error": {"type": error_type, "message": message}} + + +def _json(response): + return json.loads(response.body) + + +@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 instead. ERROR passes the logger's default level. + """ + 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) + + +class _Recording(ErrorRenderers): + """A registry that records what it is asked to render.""" + + def __init__(self): + super().__init__() + self.calls = [] + + def render(self, dialect, request, error): + self.calls.append((dialect, request, error)) + return JSONResponse(status_code=499, content={"recorded": True}) + + +# --------------------------------------------------------------------------- +# The dialect-neutral error types +# --------------------------------------------------------------------------- + + +def test_neutral_error_defaults(): + invalid = InvalidRequestError("bad") + assert (invalid.status, invalid.err_type, invalid.code, invalid.message, str(invalid)) == ( + 400, + "BadRequestError", + None, + "bad", + "bad", + ) + + too_long = ContextLengthExceededError("too long") + assert isinstance(too_long, InvalidRequestError) + assert (too_long.status, too_long.err_type, too_long.code) == ( + 400, + "BadRequestError", + "context_length_exceeded", + ) + + assert (UpstreamError("no", 404).status, UpstreamError("no", 404).err_type) == ( + 404, + "BadRequestError", + ) + assert (UpstreamError("no", 500).status, UpstreamError("no", 500).err_type) == ( + 500, + "InternalServerError", + ) + assert (UpstreamError("no", None).status, UpstreamError("no", None).err_type) == ( + 502, + "InternalServerError", + ) + + busy = UnavailableError("busy") + assert (busy.status, busy.err_type, busy.retryable) == (503, "InternalServerError", True) + assert UnavailableError("busy", retryable=False).retryable is False + + assert (InternalServingError("oops").status, InternalServingError("oops").err_type) == ( + 500, + "InternalServerError", + ) + + for error in (invalid, too_long, busy, UpstreamError("no"), InternalServingError("oops")): + assert isinstance(error, DialectNeutralError) + assert not issubclass(EngineFatalError, DialectNeutralError) + assert not issubclass(ServingError, DialectNeutralError) + + expected_types = { + (400, "bad"): InvalidRequestError, + (400, CONTEXT_LENGTH_MESSAGE): ContextLengthExceededError, + (499, "gone"): UpstreamError, + (500, "oops"): InternalServingError, + (502, "upstream"): UpstreamError, + (503, "busy"): UnavailableError, + } + for (status, message), expected in expected_types.items(): + error = neutral_error(status, message) + assert type(error) is expected, (status, message) + assert (error.status, error.message) == (status, message) + + +NEUTRAL_CASES = [ + pytest.param(serving_error("bad"), 400, "BadRequestError", 400, id="serving-error"), + pytest.param( + serving_error("gone", "NotFoundError", HTTPStatus.NOT_FOUND), + 404, + "NotFoundError", + 404, + id="serving-404", + ), + pytest.param(InvalidRequestError("bad"), 400, "BadRequestError", 400, id="invalid-request"), + pytest.param( + ContextLengthExceededError("too long"), + 400, + "BadRequestError", + "context_length_exceeded", + id="context-length", + ), + pytest.param( + UpstreamError("worker said no", 404), 404, "BadRequestError", 404, id="upstream-4xx" + ), + pytest.param( + UpstreamError("worker down", None), 502, "InternalServerError", 502, id="upstream-502" + ), + pytest.param( + UpstreamError("client went away", 499), 499, "BadRequestError", 499, id="upstream-499" + ), + pytest.param(UnavailableError("busy"), 503, "InternalServerError", 503, id="unavailable"), + pytest.param(InternalServingError("oops"), 500, "InternalServerError", 500, id="internal"), +] + + +@pytest.mark.parametrize("dialect", [OPENAI, ANTHROPIC]) +@pytest.mark.parametrize("error, status, err_type, code", NEUTRAL_CASES) +def test_boundary_renders_serving_and_neutral_errors(dialect, error, status, err_type, code): + @serving_error_boundary(dialect, renderers=default_error_renderers()) + async def handler(raw_request: Request) -> Response: + raise error + + response = asyncio.run(handler(None)) + message = error.body.message if isinstance(error, ServingError) else error.message + if dialect == OPENAI: + expected = _openai_body(message, err_type, code) + else: + expected = _anthropic_body(message, status) + assert (response.status_code, _json(response)) == (status, expected) + + +@pytest.mark.parametrize( + "message", ["bad", CONTEXT_LENGTH_MESSAGE], ids=["plain", "context-length"] +) +def test_openai_rendering_equals_create_error_response(message): + rendered = default_error_renderers().render(OPENAI, None, InvalidRequestError(message)) + reference = OpenAIServer.create_error_response(message) + assert (rendered.status_code, rendered.body) == (reference.status_code, reference.body), ( + "the OPENAI rendering of a dialect-neutral error must use error_body " + "(tensorrt_llm/serve/_base/errors.py) exactly as create_error_response does" + ) + if message is CONTEXT_LENGTH_MESSAGE: + assert _json(rendered)["code"] == "context_length_exceeded" + + +# --------------------------------------------------------------------------- +# The boundary +# --------------------------------------------------------------------------- + + +def test_boundary_finds_the_request_by_annotation(monkeypatch): + registry = _Recording() + raw = Request({"type": "http", "headers": []}) + + class Serving: + @serving_error_boundary(OPENAI, renderers=registry) + async def typed(self, request: _Body, raw_request: Request) -> Response: + raise InvalidRequestError("bad") + + @serving_error_boundary(OPENAI, renderers=registry) + async def stringly(self, request: _Body, raw_request: "Request") -> Response: + raise InvalidRequestError("bad") + + @serving_error_boundary(OPENAI) + async def unpinned(self, request: _Body, raw_request: Request) -> Response: + raise InvalidRequestError("bad") + + body = _Body(value=1) + asyncio.run(Serving().typed(body, raw)) + asyncio.run(Serving().stringly(body, raw_request=raw)) + asyncio.run(Serving().typed(body, None)) + + # An app that never installed the registry is a wiring error, named as such rather than + # silently answered with the default renderers. + bare = Request({"type": "http", "headers": [], "app": FastAPI()}) + with pytest.raises(RuntimeError, match="install_exception_handlers"): + asyncio.run(Serving().unpinned(body, bare)) + assert [call[1] for call in registry.calls] == [raw, raw, None], REQUEST_RULE + assert all(call[1] is not body for call in registry.calls), REQUEST_RULE + + # Without a registry argument: the request's app registry when the scope has an app, + # else the default renderers (patched here to tell them apart). + fallback = _Recording() + monkeypatch.setattr(register_module, "default_error_renderers", lambda **kwargs: fallback) + app = FastAPI() + app_registry = _Recording() + install_exception_handlers(app, renderers=app_registry, default_dialect=OPENAI) + + asyncio.run(Serving().unpinned(body, None)) + asyncio.run(Serving().unpinned(body, Request({"type": "http", "headers": []}))) + assert len(fallback.calls) == 2 and app_registry.calls == [], REQUEST_RULE + + with_app = Request({"type": "http", "headers": [], "app": app}) + asyncio.run(Serving().unpinned(body, with_app)) + assert len(fallback.calls) == 2, REQUEST_RULE + assert [call[1] for call in app_registry.calls] == [with_app], REQUEST_RULE + + +def test_boundary_keeps_signature_and_docstring(): + class Plain: + async def handle(self, request: _Body, raw_request: Request) -> Response: + """Handle a request.""" + return JSONResponse(content={"ok": True}) + + class Guarded: + @serving_error_boundary(OPENAI) + async def handle(self, request: _Body, raw_request: Request) -> Response: + """Handle a request.""" + return JSONResponse(content={"ok": True}) + + assert inspect.signature(Guarded.handle) == inspect.signature(Plain.handle) + assert inspect.signature(Guarded().handle) == inspect.signature(Plain().handle) + assert (Guarded.handle.__name__, Guarded.handle.__doc__) == ("handle", "Handle a request.") + + def operation(handler): + app = FastAPI() + app.add_api_route("/x", handler, methods=["POST"]) + return app.openapi()["paths"]["/x"]["post"] + + plain, guarded = operation(Plain().handle), operation(Guarded().handle) + for key in ("operationId", "summary", "description", "requestBody"): + assert guarded.get(key) == plain.get(key), key + assert guarded == plain + + +def test_boundary_engine_fatal(): + @serving_error_boundary(OPENAI, on_engine_fatal=lambda: None) + async def answered() -> Response: + raise EngineFatalError() + + @serving_error_boundary(OPENAI) + async def propagated() -> Response: + raise EngineFatalError() + + assert asyncio.run(answered()) is None + with pytest.raises(EngineFatalError): + asyncio.run(propagated()) + + +@pytest.mark.parametrize("ladder", [CHAT_LADDER], ids=["chat"]) +def test_boundary_ladders_match_today(ladder, error_records): + @serving_error_boundary(OPENAI, ladder=ladder, renderers=default_error_renderers()) + async def handler(error: Exception) -> Response: + raise error + + response = asyncio.run(handler(ValueError("bad"))) + assert (response.status_code, _json(response)) == ( + 400, + _openai_body("bad", "BadRequestError", 400), + ) + assert error_records == [], "a ValueError is answered without a log" + + response = asyncio.run(handler(RuntimeError("boom"))) + assert (response.status_code, _json(response)) == ( + 400, + _openai_body("boom", "BadRequestError", 400), + ) + assert len(error_records) == 1 and "RuntimeError: boom" in error_records[0], error_records + + +def test_boundary_rule_options(error_records): + ladder = ( + ErrorRule(KeyError, HTTPStatus.NOT_FOUND, "NotFoundError", log="message"), + ErrorRule( + (TypeError, LookupError), HTTPStatus.INTERNAL_SERVER_ERROR, "InternalServerError" + ), + ) + + @serving_error_boundary(ANTHROPIC, ladder=ladder, renderers=default_error_renderers()) + async def handler(error: Exception) -> Response: + raise error + + response = asyncio.run(handler(KeyError("k"))) + assert (response.status_code, _json(response)) == (404, _anthropic_body("'k'", 404)) + assert len(error_records) == 1 and error_records[0].endswith(" 'k'"), ( + "log='message' logs str(error) at ERROR", + error_records, + ) + + response = asyncio.run(handler(IndexError("i"))) + assert (response.status_code, _json(response)) == (500, _anthropic_body("i", 500)) + assert len(error_records) == 1, ("log='none' logs nothing", error_records) + + +def test_boundary_reraises_unmatched(): + @serving_error_boundary(OPENAI, ladder=(ErrorRule(ValueError, HTTPStatus.BAD_REQUEST),)) + async def handler() -> Response: + raise KeyError("unmatched") + + with pytest.raises(KeyError): + asyncio.run(handler()) + + +# --------------------------------------------------------------------------- +# The chat completion handler through the boundary +# --------------------------------------------------------------------------- + +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): + body = { + "model": MODEL, + "messages": [{"role": "user", "content": "hi"}], + "prompt_token_ids": [1, 2, 3], + } + body.update(overrides) + return body + + +def _chat_response(**choice_fields) -> ChatCompletionResponse: + return ChatCompletionResponse( + id="chatcmpl-boundary-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), + ) + + +class _UnserializableResponse(ChatCompletionResponse): + def model_dump(self, *args, **kwargs): + raise RuntimeError("boom") + + +@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 test_chat_wrapper_answers_as_before(serve_env, error_records, engine_failure_recorder): + server = serve_env.build_openai_server(AggConfig()) + server.await_disconnected = AsyncMock() + client = serve_env.client(server.app) + + # The core's rejection: 400 with the body create_error_response gives. + expected = OpenAIServer.create_error_response(NAMED_TOOL_ON_TOKEN_IDS_MESSAGE) + response = client.post( + "/v1/chat/completions", + json=_chat_body( + tools=[_TOOL], tool_choice={"type": "function", "function": {"name": "get_weather"}} + ), + ) + assert (response.status_code, response.content) == (400, expected.body) + assert error_records == [] + + # A response that cannot be encoded (ValueError): 400, no traceback. + server._create_chat_response = AsyncMock( + return_value=_chat_response(avg_decoded_tokens_per_iter=float("nan")) + ) + response = client.post("/v1/chat/completions", json=_chat_body()) + message = response.json()["message"] + assert message.startswith("Out of range float values are not JSON compliant") + assert (response.status_code, response.content) == ( + 400, + OpenAIServer.create_error_response(message).body, + ) + assert error_records == [] + + # Any other failure while answering (RuntimeError): 400, one traceback at ERROR. + server._create_chat_response = AsyncMock( + return_value=_UnserializableResponse(**_chat_response().model_dump()) + ) + response = client.post("/v1/chat/completions", json=_chat_body()) + assert (response.status_code, response.content) == ( + 400, + OpenAIServer.create_error_response("boom").body, + ) + assert len(error_records) == 1 and "RuntimeError: boom" in error_records[0], error_records + + # An engine failure: HTTP 200 with body null, after SIGINT. + server._create_chat_response = None + server.generator.script(ScriptedOutput(raise_on_call=CppExecutorError("dead"))) + response = client.post("/v1/chat/completions", json=_chat_body()) + assert (response.status_code, response.content) == (200, b"null") + assert engine_failure_recorder.signals == [signal.SIGINT] + assert len(engine_failure_recorder.terminations) == 1 + + +def test_error_body_is_the_only_openai_body_builder(): + """A neutral error's OpenAI body is error_body's for the same three values.""" + # 499 has no HTTPStatus member: the body is built from the status itself and renders. + error = UpstreamError("client went away", 499) + rendered = default_error_renderers().render(OPENAI, None, error) + assert rendered.status_code == 499 + assert _json(rendered) == error_body(error.message, error.err_type, error.status).model_dump() diff --git a/tests/unittest/serve/exception_handling/test_dialect_error_renderers.py b/tests/unittest/serve/exception_handling/test_dialect_error_renderers.py new file mode 100644 index 000000000000..772d16c9f02d --- /dev/null +++ b/tests/unittest/serve/exception_handling/test_dialect_error_renderers.py @@ -0,0 +1,335 @@ +# 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 per-dialect error registry (``tensorrt_llm/serve/exception_handling/register.py``).""" + +import ast +from pathlib import Path + +import pytest +from fastapi import FastAPI, Request +from fastapi.exceptions import RequestValidationError +from fastapi.responses import JSONResponse, StreamingResponse +from fastapi.testclient import TestClient +from pydantic import BaseModel + +from tensorrt_llm.serve._base.routing import make_router +from tensorrt_llm.serve.exception_handling.register import ( + ANTHROPIC, + OPENAI, + VISUAL_GEN, + ErrorRenderers, + default_error_renderers, + dialect_of, + install_exception_handlers, + openai_validation_error_renderer, + route_dialect, +) + +pytestmark = pytest.mark.cpu_only + +REPO_ROOT = Path(__file__).resolve().parents[4] + +MARKER_RULE = ( + "dialect_of must read the endpoint marker (scope['endpoint'] or scope['route'].endpoint)" +) + + +class _Body(BaseModel): + value: int + + +def _validation_error() -> RequestValidationError: + return RequestValidationError( + [ + { + "type": "int_parsing", + "loc": ("body", "value"), + "msg": "Input should be an integer", + "input": "x", + } + ] + ) + + +def _request(**scope) -> Request: + return Request({"type": "http", "headers": [], **scope}) + + +def _installed_app(registry=None, default_dialect=OPENAI, observer=None): + app = FastAPI() + install_exception_handlers( + app, + renderers=registry if registry is not None else default_error_renderers(), + default_dialect=default_dialect, + on_validation_error=observer, + ) + return app + + +# --------------------------------------------------------------------------- +# Dialect resolution +# --------------------------------------------------------------------------- + + +def test_marked_endpoint_resolves_its_dialect(): + @route_dialect(ANTHROPIC) + async def marked(body: _Body): + return {} + + async def unmarked(body: _Body): + return {} + + class Holder: + @route_dialect(ANTHROPIC) + async def method(self, body: _Body): + return {} + + bound = Holder().method + assert dialect_of(_request(endpoint=marked), OPENAI) == ANTHROPIC, MARKER_RULE + assert dialect_of(_request(endpoint=bound), OPENAI) == ANTHROPIC, MARKER_RULE + assert dialect_of(_request(endpoint=unmarked), OPENAI) == OPENAI, MARKER_RULE + assert dialect_of(_request(endpoint=unmarked), VISUAL_GEN) == VISUAL_GEN, MARKER_RULE + assert dialect_of(_request(), OPENAI) == OPENAI, MARKER_RULE + + # Through include_router(make_router()): the leaf endpoint is what the scope carries. + seen = [] + app = _installed_app( + observer=lambda request, exc, dialect: seen.append((request.url.path, dialect)) + ) + router = make_router() + router.add_api_route("/marked", marked, methods=["POST"]) + router.add_api_route("/unmarked", unmarked, methods=["POST"]) + router.add_api_route("/method", bound, methods=["POST"]) + app.include_router(router) + client = TestClient(app) + for path in ("/marked", "/unmarked", "/method"): + assert client.post(path, json={"value": "x"}).status_code == 400 + assert seen == [("/marked", ANTHROPIC), ("/unmarked", OPENAI), ("/method", ANTHROPIC)], ( + MARKER_RULE + ) + + # The route fallback, for a scope that carries the route but no endpoint. + route = next(r for _, r in _routes(app) if getattr(r, "path", "") == "/marked") + assert dialect_of(_request(route=route), OPENAI) == ANTHROPIC, MARKER_RULE + + +def _routes(app): + from test_common.serve.routes import effective_routes + + return effective_routes(app) + + +# --------------------------------------------------------------------------- +# Renderers and the registry +# --------------------------------------------------------------------------- + + +def test_validation_bodies_per_dialect(): + exc = _validation_error() + registry = default_error_renderers() + + openai = registry.render(OPENAI, None, exc) + assert (openai.status_code, openai.body) == ( + 400, + JSONResponse(content={"error": str(exc)}).body, + ), f"openai validation body changed: {openai.body!r}" + + anthropic = registry.render(ANTHROPIC, None, exc) + expected = {"type": "error", "error": {"type": "invalid_request_error", "message": str(exc)}} + assert (anthropic.status_code, anthropic.body) == (400, JSONResponse(content=expected).body), ( + f"anthropic validation body changed: {anthropic.body!r}" + ) + + envelope = { + "object": "error", + "message": "value: bad", + "type": "BadRequestError", + "param": None, + "code": 422, + } + registry.register( + VISUAL_GEN, + RequestValidationError, + lambda request, exc: JSONResponse(content=envelope, status_code=422), + ) + visual_gen = registry.render(VISUAL_GEN, None, exc) + assert (visual_gen.status_code, visual_gen.body) == ( + 422, + JSONResponse(content=envelope).body, + ), f"visual_gen validation body changed: {visual_gen.body!r}" + + +def test_openai_validation_text_cap(): + exc = _validation_error() + text = str(exc) + uncapped = openai_validation_error_renderer(None)(None, exc) + assert uncapped.body == JSONResponse(content={"error": text}).body + + limit = 10 + capped = openai_validation_error_renderer(limit)(None, exc) + cut = text[:limit] + " ... [truncated %d bytes]" % (len(text) - limit) + assert capped.body == JSONResponse(content={"error": cut}).body + assert capped.status_code == 400 + + at_limit = openai_validation_error_renderer(len(text))(None, exc) + assert at_limit.body == uncapped.body, "a text exactly at the limit is not cut" + + +def test_observer_runs_before_rendering_with_the_dialect(): + events = [] + registry = default_error_renderers() + + def render(request, exc): + events.append("render") + return JSONResponse(status_code=400, content={"rendered": True}) + + registry.register(ANTHROPIC, RequestValidationError, render) + app = _installed_app( + registry, observer=lambda request, exc, dialect: events.append(("observe", dialect)) + ) + + @route_dialect(ANTHROPIC) + async def endpoint(body: _Body): + return {} + + app.add_api_route("/a", endpoint, methods=["POST"]) + response = TestClient(app).post("/a", json={"value": "x"}) + assert response.json() == {"rendered": True} + assert events == [("observe", ANTHROPIC), "render"] + + +def test_render_falls_back_to_openai(): + registry = default_error_renderers() + exc = _validation_error() + assert not registry.has("some_other_dialect", RequestValidationError) + fallback = registry.render("some_other_dialect", None, exc) + assert fallback.body == registry.render(OPENAI, None, exc).body + + # The closest class in the error's MRO wins; a registration for a base class covers subclasses. + class Special(RequestValidationError): + pass + + registry.register( + OPENAI, Special, lambda request, exc: JSONResponse(status_code=418, content={}) + ) + assert registry.render(OPENAI, None, Special([])).status_code == 418 + assert registry.render(OPENAI, None, exc).status_code == 400 + assert registry.render(ANTHROPIC, None, Special([])).status_code == 400, ( + "a dialect with a renderer for a base class uses it before falling back to openai" + ) + + with pytest.raises(KeyError): + ErrorRenderers().render(OPENAI, None, exc) + + +def test_registry_is_on_app_state(): + registry = default_error_renderers() + app = _installed_app(registry) + assert app.state.error_renderers is registry + + async def endpoint(body: _Body): + return {} + + app.add_api_route("/a", endpoint, methods=["POST"]) + client = TestClient(app) + assert client.post("/a", json={"value": "x"}).json().keys() == {"error"} + + # Registered after the install: the handler must read the registry from app.state. + registry.register( + OPENAI, + RequestValidationError, + lambda request, exc: JSONResponse(status_code=400, content={"later": True}), + ) + assert client.post("/a", json={"value": "x"}).json() == {"later": True}, ( + "VisualGen and the Responses stream register after install; the handler must read " + "app.state.error_renderers" + ) + + # A registry swapped in after the install is the one used: the handler reads + # app.state.error_renderers on each request, not the registry it was installed with. + second = ErrorRenderers() + second.register( + OPENAI, + RequestValidationError, + lambda request, exc: JSONResponse(status_code=400, content={"second": True}), + ) + app.state.error_renderers = second + assert client.post("/a", json={"value": "x"}).json() == {"second": True}, ( + "the validation handler must read app.state.error_renderers on each request, not a " + "closure over the registry it was installed with" + ) + + +def test_renderer_may_return_a_stream(): + registry = default_error_renderers() + app = _installed_app(registry) + + async def frames(): + yield b"event: error\n" + yield b"data: {}\n\n" + + registry.register( + OPENAI, + RequestValidationError, + lambda request, exc: StreamingResponse(frames(), media_type="text/event-stream"), + ) + + async def endpoint(body: _Body): + return {} + + app.add_api_route("/a", endpoint, methods=["POST"]) + response = TestClient(app).post("/a", json={"value": "x"}) + assert response.headers["content-type"].startswith("text/event-stream") and response.text == ( + "event: error\ndata: {}\n\n" + ), ( + "error renderers must be allowed to return a stream (an error found after a Responses " + "stream has opened is answered inside the stream)" + ) + + +# --------------------------------------------------------------------------- +# Import isolation +# --------------------------------------------------------------------------- + +_SHARED_MODULES = ( + "tensorrt_llm/serve/exception_handling/register.py", + "tensorrt_llm/serve/_base/errors.py", +) +_VISUAL_GEN_MODULES = ( + "tensorrt_llm.serve.api.visual_gen", + "tensorrt_llm.serve._visual_gen_errors", + "tensorrt_llm.serve.openai_video_routes", +) + + +def _imported_modules(path: Path): + for node in ast.walk(ast.parse(path.read_text())): + if isinstance(node, ast.Import): + yield from (alias.name for alias in node.names) + elif isinstance(node, ast.ImportFrom) and node.module: + yield node.module + + +def test_registry_modules_import_no_visual_gen(): + for module in _SHARED_MODULES: + for imported in _imported_modules(REPO_ROOT / module): + assert not any( + imported == forbidden or imported.startswith(forbidden + ".") + for forbidden in _VISUAL_GEN_MODULES + ), ( + f"{module} imports {imported}: serve/exception_handling and serve/_base must not " + "import VisualGen code, which only the VisualGen role may load; VisualGen registers " + "its renderer from its own route registration" + ) diff --git a/tests/unittest/serve/test_layering.py b/tests/unittest/serve/test_layering.py index 2eae6a2d6e73..4fb58cd70528 100644 --- a/tests/unittest/serve/test_layering.py +++ b/tests/unittest/serve/test_layering.py @@ -117,6 +117,25 @@ "what they need. Move the shared helper to serve/_base/ instead." ), ), + Contract( + name="serve/exception_handling/ imports no server, VisualGen or Responses module", + importers=("tensorrt_llm.serve.exception_handling",), + forbidden=( + "tensorrt_llm.serve.openai_server", + "tensorrt_llm.serve.openai_disagg_server", + "tensorrt_llm.serve.app", + "tensorrt_llm.serve.api.visual_gen", + "tensorrt_llm.serve._visual_gen_errors", + "tensorrt_llm.serve.openai_video_routes", + "tensorrt_llm.serve.responses_utils", + ), + fix=( + "The error registry is shared by both apps and every role: the composition root " + "injects what it needs, and VisualGen code registers the VisualGen renderer on " + "app.state.error_renderers from its own route registration, which only the " + "VisualGen role runs." + ), + ), ) From 1c1172c64923d17e1228a20a31cba1bcb7172db5 Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Tue, 3 Nov 2026 17:02:35 +0000 Subject: [PATCH 25/26] [TRTLLM-17269][refactor] Register trtllm-serve error renderers per wire dialect Both servers replace their validation handler's path test (request.url.path.startswith("/v1/messages")) and the aggregated server's VisualGen role test with the per-dialect registry: install_exception_handlers resolves the dialect of the matched endpoint (the eight aggregated and two disaggregated Anthropic endpoints carry @route_dialect(ANTHROPIC)), falls back to the server's default dialect (VisualGen on the VisualGen role, OpenAI otherwise), calls the server's _observe_validation_error for its side effects and renders through the registry. Side effects stay as they were: the aggregated server counts a rejection on its metrics collector and caps the OpenAI-shaped text at 4096 characters; the disaggregated frontend counts on its Prometheus counter and logs the first and every 1000th OpenAI-shaped rejection. The VisualGen 422 renderer is registered with the VisualGen routes, for that role only, and a VisualGen server without it fails at startup. The chat completion wrapper's four except branches become @serving_error_boundary(OPENAI, ladder=CHAT_LADDER, on_engine_fatal=lambda: None) with the same answers. The Anthropic route tests install the servers' own handler instead of a copy of it. The multi-GPU trigger list gains tensorrt_llm/serve/exception_handling/. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- jenkins/L0_MergeRequest.groovy | 2 + tensorrt_llm/serve/openai_disagg_server.py | 69 ++++++---- tensorrt_llm/serve/openai_server.py | 118 ++++++++++-------- .../llmapi/apps/_test_openai_embeddings.py | 3 +- .../llmapi/apps/test_anthropic_routes.py | 30 ++--- .../llmapi/apps/test_serving_extensions.py | 15 ++- .../test_server_validation_errors.py | 56 +++++++++ 7 files changed, 191 insertions(+), 102 deletions(-) diff --git a/jenkins/L0_MergeRequest.groovy b/jenkins/L0_MergeRequest.groovy index fae12c998df7..dc1b5ffef085 100644 --- a/jenkins/L0_MergeRequest.groovy +++ b/jenkins/L0_MergeRequest.groovy @@ -1954,6 +1954,8 @@ def getMultiGpuFileChanged(pipeline, testFilter, globalVars) "tensorrt_llm/serve/api/ops/", "tensorrt_llm/serve/cluster_storage.py", "tensorrt_llm/serve/disagg_auto_scaling.py", + // The error renderers both apps install; the disaggregated frontend runs them. + "tensorrt_llm/serve/exception_handling/", "tensorrt_llm/serve/extensions/", "tensorrt_llm/serve/metadata_server.py", "tensorrt_llm/serve/middleware/", diff --git a/tensorrt_llm/serve/openai_disagg_server.py b/tensorrt_llm/serve/openai_disagg_server.py index 6808f38d579f..c7297a14b8a8 100644 --- a/tensorrt_llm/serve/openai_disagg_server.py +++ b/tensorrt_llm/serve/openai_disagg_server.py @@ -15,6 +15,7 @@ # yapf: disable import asyncio +import functools import json import os import signal @@ -56,6 +57,9 @@ resolve_request_conversation_id) from tensorrt_llm.serve.disagg_coordinator import (CoordinatorClient, DisaggCoordinatorService) +from tensorrt_llm.serve.exception_handling.register import ( + ANTHROPIC, OPENAI, default_error_renderers, install_exception_handlers, + route_dialect) from tensorrt_llm.serve.kv_hints import KvHintsRoute from tensorrt_llm.serve.middleware.register import register_middlewares from tensorrt_llm.serve.openai_client import OpenAIClient, OpenAIHttpClient @@ -289,32 +293,12 @@ async def lifespan(app) -> None: await self._service.teardown() await self._perf_metrics_writer.close() - # Log request-body validation failures so a client/server schema mismatch - # shows up server-side. Throttled (first, then every 1000th) to avoid - # flooding the event loop when every request fails identically. + # Validation errors are rendered per wire dialect (the marked Anthropic + # routes answer with the Anthropic envelope, every other route with the + # OpenAI shape) after _observe_validation_error has counted and logged + # them. self._val_err_n = 0 - async def validation_exception_handler(request: Request, exc): - self._perf_metrics_collector.validation_exceptions.inc() - # Anthropic clients parse the error envelope, so every /v1/messages - # route must fail in that shape rather than the generic one below. - if request.url.path.startswith("/v1/messages"): - return anthropic_error_response(str(exc), - "invalid_request_error", 400) - self._val_err_n += 1 - if self._val_err_n == 1 or self._val_err_n % 1000 == 0: - try: - errs = exc.errors() - # Compact: [{loc, type, msg}] -- drops the (large) echoed input. - brief = [{"loc": e.get("loc"), "type": e.get("type"), - "msg": e.get("msg")} for e in errs][:8] - except Exception: # noqa: BLE001 - brief = str(exc)[:500] - method = request.method.translate(_LOG_CONTROL_CHARACTERS) - path = request.url.path.translate(_LOG_CONTROL_CHARACTERS) - logger.warning( - f"[validation] {method} {path} 400 " - f"(n={self._val_err_n}): {brief}") - return JSONResponse(status_code=400, content={"error": str(exc)}) + self._error_renderers = default_error_renderers() # KvHintsRoute: the frontend's public port strips kv_hints from request # bodies unless this process trusts its ingress (TRTLLM_TRUST_KV_HINTS=1); @@ -322,8 +306,11 @@ async def validation_exception_handler(request: Request, exc): self.app = build_http_app( lifespan=lifespan, route_class=KvHintsRoute, - dialect_error_renderers=lambda app: app.add_exception_handler( - RequestValidationError, validation_exception_handler), + dialect_error_renderers=functools.partial( + install_exception_handlers, + renderers=self._error_renderers, + default_dialect=OPENAI, + on_validation_error=self._observe_validation_error), routers=(self._attach_routes,), middleware=lambda app: register_middlewares( app, @@ -338,6 +325,32 @@ def _attach_routes(self, app: FastAPI) -> None: self.app = app self.register_routes() + def _observe_validation_error(self, request: Request, + exc: RequestValidationError, + dialect: str) -> None: + self._perf_metrics_collector.validation_exceptions.inc() + # Anthropic clients parse the error envelope; only the OpenAI-shaped + # rejections are logged. + if dialect == ANTHROPIC: + return + # Log request-body validation failures so a client/server schema mismatch + # shows up server-side. Throttled (first, then every 1000th) to avoid + # flooding the event loop when every request fails identically. + self._val_err_n += 1 + if self._val_err_n == 1 or self._val_err_n % 1000 == 0: + try: + errs = exc.errors() + # Compact: [{loc, type, msg}] -- drops the (large) echoed input. + brief = [{"loc": e.get("loc"), "type": e.get("type"), + "msg": e.get("msg")} for e in errs][:8] + except Exception: # noqa: BLE001 + brief = str(exc)[:500] + method = request.method.translate(_LOG_CONTROL_CHARACTERS) + path = request.url.path.translate(_LOG_CONTROL_CHARACTERS) + logger.warning( + f"[validation] {method} {path} 400 " + f"(n={self._val_err_n}): {brief}") + def _create_client(self, router: Router, role: ServerRole, max_retries: int = 1) -> OpenAIClient: async def disagg_id_generator(): return await self._coordinator.get_disagg_request_id() @@ -470,6 +483,7 @@ async def wrapper(req: request_type, raw_req: Request) -> Response: return self._handle_exception(e) return wrapper + @route_dialect(ANTHROPIC) async def anthropic_messages(self, request: AnthropicMessagesRequest, raw_request: Request) -> Response: """Serve Anthropic Messages through the disaggregated chat pipeline.""" @@ -540,6 +554,7 @@ async def get_model(self) -> Response: code=status, ).model_dump()) + @route_dialect(ANTHROPIC) async def anthropic_count_tokens( self, request: AnthropicCountTokensRequest) -> Response: """Count Anthropic input tokens on a context worker.""" diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py index 1ad792a3a77f..477ec8c019a5 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -65,8 +65,10 @@ get_effective_kv_cache_event_hash_algo from tensorrt_llm.sampling_params import GuidedDecodingParams, SamplingParams from tensorrt_llm.serve._base.app_builder import build_http_app, serve_uvicorn -from tensorrt_llm.serve._base.errors import (EngineFatalError, ServingError, - error_body, serving_error) +from tensorrt_llm.serve._base.errors import (CHAT_LADDER, EngineFatalError, + ServingError, error_body, + serving_error, + serving_error_boundary) from tensorrt_llm.serve._base.msgspec_route import _MsgspecRoute from tensorrt_llm.serve._base.routing import ServerCapabilities, attach_routers from tensorrt_llm.serve._base.sse import chat_stream_generator @@ -96,6 +98,9 @@ from tensorrt_llm.serve.disagg_auto_scaling import DisaggClusterWorker from tensorrt_llm.serve.encode_batcher import (EncodeBatcher, InputTooLongError, QueueFullError) +from tensorrt_llm.serve.exception_handling.register import ( + ANTHROPIC, OPENAI, VISUAL_GEN, default_error_renderers, + install_exception_handlers, route_dialect) from tensorrt_llm.serve.metadata_server import create_metadata_server from tensorrt_llm.serve.middleware.register import register_middlewares from tensorrt_llm.serve.multi_frontend import (FORWARDED_ROUTES, @@ -544,6 +549,12 @@ def resolve_spec_decode_num_spec_tokens(args: Any) -> Optional[int]: return getattr(spec_config, "max_draft_len", None) +# Characters of a validation error's text kept in the OpenAI-shaped 400 body: +# pydantic echoes each failing input verbatim, which can be megabytes of +# base64; the error locations are what matters. +_VALIDATION_ERROR_DETAIL_LIMIT = 4096 + + class OpenAIServer(_VideoRoutesMixin): @staticmethod @@ -850,35 +861,22 @@ async def lifespan(app: FastAPI): self._metrics_stats_executor.shutdown(wait=False, cancel_futures=True) - async def validation_exception_handler(request, exc): - if request.url.path.startswith("/v1/messages"): - # Anthropic clients expect a 400 with the Anthropic error - # envelope, not FastAPI's 422 shape. - if self.metrics_collector: - self.metrics_collector.log_request_error(http_code=400) - return anthropic_error_response(str(exc), - "invalid_request_error", 400) - if self.server_role is ServerRole.VISUAL_GEN: - return self._create_visual_gen_validation_error_response(exc) - # Non-visual-gen roles keep the shared 400 + ``{"error": ...}`` - # response shape that integration tests (e.g. - # ``test_malformed_json_request``) and existing clients - # expect. - if self.metrics_collector: - self.metrics_collector.log_request_error(http_code=400) - # pydantic echoes each failing input verbatim, which can be - # megabytes of base64; the error locations are what matters. - detail = str(exc) - if len(detail) > 4096: - detail = detail[:4096] + " ... [truncated %d bytes]" % ( - len(detail) - 4096) - return JSONResponse(status_code=400, content={"error": detail}) - + # Validation errors are rendered per wire dialect: the dialect of the + # endpoint that matched, else the dialect of the server's role. The + # OpenAI-shaped 400 body keeps the shape integration tests (e.g. + # ``test_malformed_json_request``) and existing clients expect. + self._error_renderers = default_error_renderers( + openai_validation_detail_limit=_VALIDATION_ERROR_DETAIL_LIMIT) + default_dialect = (VISUAL_GEN if self.server_role + is ServerRole.VISUAL_GEN else OPENAI) self.app = build_http_app( lifespan=lifespan, route_class=_MsgspecRoute, - dialect_error_renderers=lambda app: app.add_exception_handler( - RequestValidationError, validation_exception_handler), + dialect_error_renderers=functools.partial( + install_exception_handlers, + renderers=self._error_renderers, + default_dialect=default_dialect, + on_validation_error=self._observe_validation_error), routers=(self._attach_role_routes, ), middleware=lambda app: register_middlewares( app, @@ -889,6 +887,21 @@ async def validation_exception_handler(request, exc): ) self.app.state.trust_kv_hints = os.getenv("TRTLLM_TRUST_KV_HINTS", "0") == "1" + if (self.server_role is ServerRole.VISUAL_GEN + and not self._error_renderers.has(VISUAL_GEN, + RequestValidationError)): + raise RuntimeError( + "VisualGen role started without its validation-error " + "renderer; the VisualGen route registration must register " + "it for the visual_gen dialect") + + def _observe_validation_error(self, request: Request, + exc: RequestValidationError, + dialect: str) -> None: + # A rejected request counts as a 400 whichever dialect answers it; + # the VisualGen role has no metrics collector and counts nothing. + if dialect != VISUAL_GEN and self.metrics_collector: + self.metrics_collector.log_request_error(http_code=400) def _attach_role_routes(self, app: FastAPI) -> None: # The register_* methods below add their routes through self.app. @@ -906,6 +919,12 @@ def _attach_role_routes(self, app: FastAPI) -> None: if self.server_role is ServerRole.VISUAL_GEN: assert self._is_visual_gen, \ "generator must be a VisualGen for VISUAL_GEN server" + # VisualGen answers a validation error with its own 422 envelope; + # the renderer is registered with the VisualGen routes, so only + # the VisualGen role carries it. + self._error_renderers.register( + VISUAL_GEN, RequestValidationError, lambda request, exc: self. + _create_visual_gen_validation_error_response(exc)) self.register_visual_gen_routes() elif self.server_role is ServerRole.MM_ENCODER: assert isinstance( @@ -2009,29 +2028,22 @@ def _record_iteration_stat(self, llm_stat: dict) -> None: if self._iteration_stats_buffer is not None: self._iteration_stats_buffer.append(llm_stat) + # The boundary renders the core's ServingError with its own status, answers + # an EngineFatalError as HTTP 200 with body null (FastAPI's rendering of + # None), which the in-process callers map to their own errors, and maps the + # rest with CHAT_LADDER: a response that cannot be encoded (JSONResponse + # encodes in its constructor, allow_nan=False) raises ValueError and is + # answered 400 without a log; any other failure is answered 400 with its + # traceback logged. + @serving_error_boundary(OPENAI, + ladder=CHAT_LADDER, + on_engine_fatal=lambda: None) async def openai_chat(self, request: ChatCompletionRequest, raw_request: Request) -> Response: - try: - result = await self.create_chat_completion(request, 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)) + result = await self.create_chat_completion(request, raw_request) + if isinstance(result, ChatCompletionResponse): + return JSONResponse(content=result.model_dump()) + return StreamingResponse(content=result, media_type="text/event-stream") async def create_chat_completion( self, request: ChatCompletionRequest, raw_request: Request @@ -2490,6 +2502,7 @@ def get_role() -> str: logger.error(traceback.format_exc()) raise serving_error(str(e)) from e + @route_dialect(ANTHROPIC) async def anthropic_messages(self, request: AnthropicMessagesRequest, raw_request: Request) -> Response: """Serve the Anthropic Messages API on top of the openai_chat path. @@ -2543,6 +2556,7 @@ async def anthropic_messages(self, request: AnthropicMessagesRequest, return JSONResponse(content=anthropic_response.model_dump( exclude_none=True)) + @route_dialect(ANTHROPIC) async def anthropic_count_tokens( self, request: AnthropicCountTokensRequest) -> Response: """Report how many input tokens a Messages request would consume. @@ -2695,6 +2709,7 @@ async def _run_batched_message( } return "succeeded", anthropic_response.model_dump(exclude_none=True) + @route_dialect(ANTHROPIC) async def anthropic_create_batch( self, request: AnthropicCreateBatchRequest) -> Response: try: @@ -2708,6 +2723,7 @@ async def anthropic_create_batch( 400) return JSONResponse(content=batch.model_dump()) + @route_dialect(ANTHROPIC) async def anthropic_list_batches( self, # Validated here rather than clamped in the store: clamping turned @@ -2740,18 +2756,21 @@ async def anthropic_list_batches( ) return JSONResponse(content=listing.model_dump()) + @route_dialect(ANTHROPIC) 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()) + @route_dialect(ANTHROPIC) 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()) + @route_dialect(ANTHROPIC) async def anthropic_delete_batch(self, batch_id: str) -> Response: try: deleted = self._batch_store().delete(batch_id) @@ -2763,6 +2782,7 @@ async def anthropic_delete_batch(self, batch_id: str) -> Response: return JSONResponse(content=AnthropicBatchDeleteResponse( id=batch_id).model_dump()) + @route_dialect(ANTHROPIC) async def anthropic_batch_results(self, batch_id: str) -> Response: store = self._batch_store() results = store.results(batch_id) diff --git a/tests/unittest/llmapi/apps/_test_openai_embeddings.py b/tests/unittest/llmapi/apps/_test_openai_embeddings.py index 6aabb8e4253a..c67388fbf253 100644 --- a/tests/unittest/llmapi/apps/_test_openai_embeddings.py +++ b/tests/unittest/llmapi/apps/_test_openai_embeddings.py @@ -132,7 +132,8 @@ def test_dimensions_must_be_positive(server: RemoteEmbeddingServer): # `dimensions` is a Pydantic PositiveInt, so a non-positive value fails # request-model validation. The server's RequestValidationError handler # renders that as a 400 (consistent `{"error": ...}` envelope) rather than - # FastAPI's default 422 — see openai_server.validation_exception_handler. + # FastAPI's default 422 — see the OpenAI validation renderer in + # tensorrt_llm/serve/exception_handling/register.py. resp = requests.post( server.url_for("v1", "embeddings"), json={"model": BERT_MODEL, "input": "x", "dimensions": 0}, diff --git a/tests/unittest/llmapi/apps/test_anthropic_routes.py b/tests/unittest/llmapi/apps/test_anthropic_routes.py index 82620eb57553..0429aa1d8a1a 100644 --- a/tests/unittest/llmapi/apps/test_anthropic_routes.py +++ b/tests/unittest/llmapi/apps/test_anthropic_routes.py @@ -28,6 +28,11 @@ from tensorrt_llm.inputs.chat_template_guard import UnusedChatTemplateKwargsError from tensorrt_llm.serve._base import sse as sse_module +from tensorrt_llm.serve.exception_handling.register import ( + OPENAI, + default_error_renderers, + install_exception_handlers, +) from tensorrt_llm.serve.openai_disagg_server import OpenAIDisaggServer from tensorrt_llm.serve.openai_protocol import ( ChatCompletionRequest, @@ -125,7 +130,7 @@ def _make_route_client(server_kind, openai_response): # 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) + install_exception_handlers(app, renderers=default_error_renderers(), default_dialect=OPENAI) if server_kind == "standard": server = object.__new__(OpenAIServer) server.model = MODEL @@ -530,7 +535,7 @@ async def ok(request): return "succeeded", {"type": "message", "role": "assistant", "content": []} app = FastAPI() - _anthropic_validation_handler(app) + install_exception_handlers(app, renderers=default_error_renderers(), default_dialect=OPENAI) server = object.__new__(OpenAIServer) server.model = MODEL server._anthropic_batch_store = AnthropicBatchStore(runner=runner or ok) @@ -841,28 +846,9 @@ 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. - - 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. - """ - 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 - - def _count_tokens_client(): app = FastAPI() - _anthropic_validation_handler(app) + install_exception_handlers(app, renderers=default_error_renderers(), default_dialect=OPENAI) server = object.__new__(OpenAIServer) server.model = MODEL server.tokenizer = _FakeTokenizer() diff --git a/tests/unittest/llmapi/apps/test_serving_extensions.py b/tests/unittest/llmapi/apps/test_serving_extensions.py index f8806b2f8c45..5d43d1b6b976 100644 --- a/tests/unittest/llmapi/apps/test_serving_extensions.py +++ b/tests/unittest/llmapi/apps/test_serving_extensions.py @@ -402,8 +402,17 @@ def generate_async(*, inputs, **kwargs): rendered["create_chat_response"] = server._create_chat_response + from tensorrt_llm.serve.exception_handling.register import ( + OPENAI, + default_error_renderers, + install_exception_handlers, + ) + app = FastAPI() app.state.server = server + # The handler renders its errors through the app's registry, which the + # servers install when they build their app; a bare app installs it here. + install_exception_handlers(app, renderers=default_error_renderers(), default_dialect=OPENAI) app.add_api_route("/v1/chat/completions", server.openai_chat, methods=["POST"]) return TestClient(app), events, rendered @@ -596,9 +605,9 @@ def test_required_attaches_mandatory_grammar(self, required_chat_client, thinkin def test_required_rejects_conflicting_requests(self, required_chat_client, change): client, server, rendered = required_chat_client response = client.post("/v1/chat/completions", json={**self.required_body(), **change}) - # The bare app uses FastAPI's 422 for request-schema errors. - expected_status = 422 if change.get("tools") == [] else 400 - assert response.status_code == expected_status, response.text + # Request-schema errors (`tools: []`) answer 400 too: the app carries the servers' + # validation handler, not FastAPI's default 422. + assert response.status_code == 400, response.text assert "generation" not in rendered @pytest.mark.parametrize( diff --git a/tests/unittest/serve/exception_handling/test_server_validation_errors.py b/tests/unittest/serve/exception_handling/test_server_validation_errors.py index e8141a83e31a..1af7cfe8eec2 100644 --- a/tests/unittest/serve/exception_handling/test_server_validation_errors.py +++ b/tests/unittest/serve/exception_handling/test_server_validation_errors.py @@ -34,9 +34,17 @@ from fastapi import Request from fastapi.exceptions import RequestValidationError from test_common.serve.app_factory import AggConfig, DisaggConfig +from test_common.serve.routes import effective_routes from tensorrt_llm.llmapi.disagg_utils import ServerRole from tensorrt_llm.logger import logger +from tensorrt_llm.serve.exception_handling.register import ( + ANTHROPIC, + OPENAI, + VISUAL_GEN, + ErrorRenderers, + dialect_of, +) # threadleak is off: TestClients are created in the test body and closed by the serve_env # fixture's teardown, which runs after the leak check. @@ -276,3 +284,51 @@ def test_disagg_bodies_counter_and_throttled_log( "[validation] P\\x01ST /v1/chat/completions\\x7f 400 (n=1): " in _validation_warnings(warning_records)[-1] ), warning_records + + +# --------------------------------------------------------------------------- +# Dialect resolution on the real apps +# --------------------------------------------------------------------------- + + +def test_visual_gen_role_requires_its_renderer(serve_env, monkeypatch): + """A VisualGen app whose renderer registration was lost fails at startup, not per request.""" + register = ErrorRenderers.register + + def dropping_visual_gen(self, dialect, exc_type, renderer): + if dialect != VISUAL_GEN: + register(self, dialect, exc_type, renderer) + + monkeypatch.setattr(ErrorRenderers, "register", dropping_visual_gen) + with pytest.raises( + RuntimeError, match="VisualGen role started without its validation-error renderer" + ): + serve_env.build_openai_server(AggConfig(role=ServerRole.VISUAL_GEN)) + + +@pytest.mark.parametrize("kind", ["agg_llm", "disagg"]) +def test_every_messages_route_resolves_to_anthropic(serve_env, kind): + """Every ``/v1/messages`` route is marked Anthropic; every other route takes the app default.""" + if kind == "agg_llm": + app = serve_env.build_openai_server(AggConfig()).app + else: + app = serve_env.build_disagg_server(DisaggConfig()).app + default = OPENAI + checked = 0 + for path, route in effective_routes(app): + endpoint = getattr(route, "endpoint", None) + if endpoint is None: # a Mount + continue + expected = ANTHROPIC if path.startswith("/v1/messages") else default + methods = ",".join(sorted(getattr(route, "methods", None) or ())) + for scope in ({"endpoint": endpoint}, {"route": route}): + dialect = dialect_of(Request({"type": "http", "headers": [], **scope}), default) + assert dialect == expected, ( + f"route {methods} {path} resolves to {dialect}; Anthropic endpoints must be marked " + "with @route_dialect(ANTHROPIC) (tensorrt_llm/serve/exception_handling/register.py)" + ) + checked += 1 + anthropic_routes = [ + path for path, _ in effective_routes(app) if path.startswith("/v1/messages") + ] + assert checked > len(anthropic_routes) > 0, (checked, anthropic_routes) From 0bbc9eed404700ca4c90379e111dfbf5c5bf32d8 Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Wed, 4 Nov 2026 07:05:06 +0000 Subject: [PATCH 26/26] [TRTLLM-17269][test] Check the error modules' VisualGen imports with a layering contract test_registry_modules_import_no_visual_gen scanned every import statement of serve/exception_handling/register.py and serve/_base/errors.py, at any scope including if TYPE_CHECKING:, for three VisualGen modules. The layering contract for serve/exception_handling/ already forbids those modules for register.py at run time, so the scan's own part was serve/_base/errors.py and the type-checking imports of both files. A layering contract over serve/exception_handling/ and serve/_base/errors.py, with the same three modules and type-checking imports included, replaces the test and its helpers. The contract also resolves relative imports and from-imports of a submodule through its parent package (from tensorrt_llm.serve import openai_video_routes), which the scan, matching module names only, did not see. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- .../test_dialect_error_renderers.py | 41 ------------------- tests/unittest/serve/test_layering.py | 15 +++++++ 2 files changed, 15 insertions(+), 41 deletions(-) diff --git a/tests/unittest/serve/exception_handling/test_dialect_error_renderers.py b/tests/unittest/serve/exception_handling/test_dialect_error_renderers.py index 772d16c9f02d..b3f13dbbceb9 100644 --- a/tests/unittest/serve/exception_handling/test_dialect_error_renderers.py +++ b/tests/unittest/serve/exception_handling/test_dialect_error_renderers.py @@ -14,9 +14,6 @@ # limitations under the License. """The per-dialect error registry (``tensorrt_llm/serve/exception_handling/register.py``).""" -import ast -from pathlib import Path - import pytest from fastapi import FastAPI, Request from fastapi.exceptions import RequestValidationError @@ -39,8 +36,6 @@ pytestmark = pytest.mark.cpu_only -REPO_ROOT = Path(__file__).resolve().parents[4] - MARKER_RULE = ( "dialect_of must read the endpoint marker (scope['endpoint'] or scope['route'].endpoint)" ) @@ -297,39 +292,3 @@ async def endpoint(body: _Body): "error renderers must be allowed to return a stream (an error found after a Responses " "stream has opened is answered inside the stream)" ) - - -# --------------------------------------------------------------------------- -# Import isolation -# --------------------------------------------------------------------------- - -_SHARED_MODULES = ( - "tensorrt_llm/serve/exception_handling/register.py", - "tensorrt_llm/serve/_base/errors.py", -) -_VISUAL_GEN_MODULES = ( - "tensorrt_llm.serve.api.visual_gen", - "tensorrt_llm.serve._visual_gen_errors", - "tensorrt_llm.serve.openai_video_routes", -) - - -def _imported_modules(path: Path): - for node in ast.walk(ast.parse(path.read_text())): - if isinstance(node, ast.Import): - yield from (alias.name for alias in node.names) - elif isinstance(node, ast.ImportFrom) and node.module: - yield node.module - - -def test_registry_modules_import_no_visual_gen(): - for module in _SHARED_MODULES: - for imported in _imported_modules(REPO_ROOT / module): - assert not any( - imported == forbidden or imported.startswith(forbidden + ".") - for forbidden in _VISUAL_GEN_MODULES - ), ( - f"{module} imports {imported}: serve/exception_handling and serve/_base must not " - "import VisualGen code, which only the VisualGen role may load; VisualGen registers " - "its renderer from its own route registration" - ) diff --git a/tests/unittest/serve/test_layering.py b/tests/unittest/serve/test_layering.py index 4fb58cd70528..3f113732454c 100644 --- a/tests/unittest/serve/test_layering.py +++ b/tests/unittest/serve/test_layering.py @@ -136,6 +136,21 @@ "VisualGen role runs." ), ), + Contract( + name="serve/exception_handling/ and serve/_base/errors.py import no VisualGen module", + importers=("tensorrt_llm.serve.exception_handling", "tensorrt_llm.serve._base.errors"), + forbidden=( + "tensorrt_llm.serve.api.visual_gen", + "tensorrt_llm.serve._visual_gen_errors", + "tensorrt_llm.serve.openai_video_routes", + ), + include_type_checking=True, + fix=( + "The error registry and serve/_base/errors.py serve every role, and only the " + "VisualGen role may load VisualGen code, not even for type checking: VisualGen " + "code registers its renderer from its own route registration." + ), + ), )