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/27] [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/27] [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/27] [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/27] [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/27] [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/27] [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 75fc8e7dac1c72ffb1bf2a536180c9850d0cb0a4 Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Mon, 2 Nov 2026 09:51:20 +0000 Subject: [PATCH 07/27] [TRTLLM-17262][test] Pin the disaggregated entry-point wrapper and its outside-subclass contract Every OpenAI request on the disaggregated frontend, and every Anthropic /v1/messages request, goes through OpenAIDisaggServer._wrap_entry_point. Before reorganizing that closure, pin what it does today: - tests/unittest/disaggregated/test_disagg_entry_point.py pins the per-request counters, the chat-template 400, the conversation identity from the headers, the serialization of streaming and non-streaming responses (including the guarded streamed Responses relay and model_dump(by_alias=True)), the error mapping through _handle_exception, and the single disaggregated_request trace span. - tests/unittest/disaggregated/test_disagg_server_contract.py pins the surface NeMo-RL's OpenAIDisaggServerAdaptor subclass relies on: the constructor (config, req_timeout_secs=..., coordinator_url=None), self.app with every route registered when __init__ returns, __call__(host, port, sockets=None), _wrap_entry_point, _handle_exception and _perf_metrics_collector.http_exceptions. - The OpenAIDisaggServer class docstring and tensorrt_llm/serve/AGENTS.md name that surface, so a change to it is agreed with the NeMo-RL owners first. Both files are CPU-only and collected by the unittest/disaggregated entry of tests/integration/test_lists/test-db/l0_cpu.yml. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- tensorrt_llm/serve/AGENTS.md | 6 + tensorrt_llm/serve/openai_disagg_server.py | 10 + .../disaggregated/test_disagg_entry_point.py | 457 ++++++++++++++++++ .../test_disagg_server_contract.py | 347 +++++++++++++ 4 files changed, 820 insertions(+) create mode 100644 tests/unittest/disaggregated/test_disagg_entry_point.py create mode 100644 tests/unittest/disaggregated/test_disagg_server_contract.py diff --git a/tensorrt_llm/serve/AGENTS.md b/tensorrt_llm/serve/AGENTS.md index bdf1fe8c23b6..cb2a0995dc08 100644 --- a/tensorrt_llm/serve/AGENTS.md +++ b/tensorrt_llm/serve/AGENTS.md @@ -46,6 +46,12 @@ Name files and symbols here, never line numbers. re-exports public names only, lists them in `__all__` (or `_LAZY_NAMES` in `commands/serve.py`), warns once with `FutureWarning`, and has its rows in `tests/unittest/serve/test_shim_identity.py`. +- `OpenAIDisaggServer` (`serve/openai_disagg_server.py`) is subclassed outside the repository + (NeMo-RL). Keep its constructor `(config, req_timeout_secs=..., coordinator_url=None)`, + `self.app` with every route registered when `__init__` returns, + `__call__(host, port, sockets=None)`, `_wrap_entry_point`, `_handle_exception` and + `_perf_metrics_collector.http_exceptions`, or agree the change with the NeMo-RL owners + first. Check: `tests/unittest/disaggregated/test_disagg_server_contract.py`. - `mock.patch("old.module.name")` does not follow moved code. Grep the tests for every moved name, including `monkeypatch.setattr(module, "name", ...)` and dotted paths split across lines. diff --git a/tensorrt_llm/serve/openai_disagg_server.py b/tensorrt_llm/serve/openai_disagg_server.py index d8ebec1a6ce7..29a7b0cb354c 100644 --- a/tensorrt_llm/serve/openai_disagg_server.py +++ b/tensorrt_llm/serve/openai_disagg_server.py @@ -214,6 +214,16 @@ def _upstream_error_message(error: aiohttp.ClientResponseError) -> str: class OpenAIDisaggServer: + """OpenAI- and Anthropic-compatible frontend of disaggregated serving. + + Extension surface used outside this repository (NeMo-RL subclasses this + class): the constructor (config, req_timeout_secs=..., + coordinator_url=None) with self.app built and every route registered when + it returns; __call__(host, port, sockets=None); _wrap_entry_point; + _handle_exception; and _perf_metrics_collector.http_exceptions. Agree a + change to any of them with the NeMo-RL owners first; the contract test is + tests/unittest/disaggregated/test_disagg_server_contract.py. + """ def __init__(self, config: DisaggServerConfig, req_timeout_secs: int = 180, diff --git a/tests/unittest/disaggregated/test_disagg_entry_point.py b/tests/unittest/disaggregated/test_disagg_entry_point.py new file mode 100644 index 000000000000..e239f2451ff0 --- /dev/null +++ b/tests/unittest/disaggregated/test_disagg_entry_point.py @@ -0,0 +1,457 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Behavior of the disaggregated server's entry-point wrapper, pinned. + +``OpenAIDisaggServer._wrap_entry_point`` turns a service entry point into the +FastAPI handler of the OpenAI routes: it counts the request, applies the +chat-template permission, resolves the conversation identity from the headers, +builds the per-request hooks, runs the entry point, counts the response, +serializes it, and sends every failure through ``_handle_exception`` while the +exception is being handled. These tests pin that behavior - the per-request +counters, the status codes, the serialization and the single trace span - so +the wrapper's internals can be reorganized without changing what a client or +a metrics scrape observes. +""" + +import asyncio +import contextlib +import json +import signal +from types import SimpleNamespace +from unittest.mock import AsyncMock, Mock + +import aiohttp +import pytest +from fastapi import HTTPException +from fastapi.responses import JSONResponse, StreamingResponse +from openai.types.responses import ResponseFormatTextJSONSchemaConfig, ResponseTextConfig +from opentelemetry.trace import SpanKind +from starlette.requests import Request +from test_common.serve.app_factory import DisaggConfig, isolated_serve_env + +import tensorrt_llm.serve.openai_disagg_server as disagg_server_module +from tensorrt_llm.executor import CppExecutorError +from tensorrt_llm.executor.utils import context_length_exceeded_message +from tensorrt_llm.llmapi import tracing +from tensorrt_llm.serve.openai_disagg_server import RawRequestResponseHooks +from tensorrt_llm.serve.openai_protocol import ( + ChatCompletionRequest, + ChatCompletionResponse, + ChatCompletionResponseChoice, + ChatMessage, + ResponsesRequest, + ResponsesResponse, + UsageInfo, +) +from tensorrt_llm.serve.openai_server import OpenAIServer + +pytestmark = pytest.mark.cpu_only + +COUNTER_NAMES = ( + "total_requests", + "stream_requests", + "nonstream_requests", + "validation_exceptions", + "http_exceptions", + "internal_errors", + "total_responses", +) + + +def _mock_collector() -> SimpleNamespace: + """The seven server counters and the queue-latency histogram, as mocks.""" + return SimpleNamespace( + **{name: SimpleNamespace(inc=Mock()) for name in COUNTER_NAMES}, + queue_latency_seconds=SimpleNamespace(observe=Mock()), + ) + + +def _server_with_counters(env): + server = env.build_disagg_server(DisaggConfig()) + counters = _mock_collector() + server._perf_metrics_collector = counters + return server, counters + + +def _assert_counts(counters: SimpleNamespace, case: str, **expected: int) -> None: + for name in COUNTER_NAMES: + actual = getattr(counters, name).inc.call_count + assert actual == expected.get(name, 0), ( + f"the {name} counter changed for {case}; _wrap_entry_point and its " + "steps must keep the request accounting exactly" + ) + + +def _raw_request(*headers: tuple) -> Request: + return Request( + { + "type": "http", + "headers": list(headers), + "state": {"server_arrival_time": 1.0, "perf_metrics_records": []}, + } + ) + + +def _chat_request(**overrides) -> ChatCompletionRequest: + payload = {"model": "m", "messages": [{"role": "user", "content": "hi"}]} + payload.update(overrides) + return ChatCompletionRequest(**payload) + + +def _chat_response() -> ChatCompletionResponse: + return ChatCompletionResponse( + id="chatcmpl-entry-point-test", + model="m", + choices=[ + ChatCompletionResponseChoice( + index=0, + message=ChatMessage(role="assistant", content="hello"), + finish_reason="stop", + ) + ], + usage=UsageInfo(prompt_tokens=3, completion_tokens=2, total_tokens=5), + ) + + +def test_non_streaming_request_is_counted_and_serialized(tmp_path): + with isolated_serve_env(tmp_path) as env: + server, counters = _server_with_counters(env) + chat_response = _chat_response() + entry_point = AsyncMock(return_value=chat_response) + wrapper = server._wrap_entry_point(entry_point, ChatCompletionRequest) + req = _chat_request() + raw_req = _raw_request() + + response = asyncio.run(wrapper(req, raw_req)) + + assert isinstance(response, JSONResponse) + assert json.loads(bytes(response.body)) == chat_response.model_dump(by_alias=True) + entry_point.assert_awaited_once() + called_req, called_hooks = entry_point.await_args.args + assert called_req is req + assert isinstance(called_hooks, RawRequestResponseHooks) + assert called_hooks.raw_req is raw_req + _assert_counts( + counters, + "a non-streaming chat request", + total_requests=1, + nonstream_requests=1, + total_responses=1, + ) + + +def test_streaming_request_returns_the_generator_unchanged(tmp_path): + with isolated_serve_env(tmp_path) as env: + server, counters = _server_with_counters(env) + + async def generated(): + yield b"data: [DONE]\n\n" + + generator = generated() + entry_point = AsyncMock(return_value=generator) + wrapper = server._wrap_entry_point(entry_point, ChatCompletionRequest) + + response = asyncio.run(wrapper(_chat_request(stream=True), _raw_request())) + + assert isinstance(response, StreamingResponse) + assert response.media_type == "text/event-stream" + assert response.body_iterator is generator + _assert_counts( + counters, + "a streaming chat request", + total_requests=1, + stream_requests=1, + total_responses=1, + ) + + +def test_forbidden_chat_template_is_a_400_before_any_worker(tmp_path): + with isolated_serve_env(tmp_path) as env: + server, counters = _server_with_counters(env) + entry_point = AsyncMock() + wrapper = server._wrap_entry_point(entry_point, ChatCompletionRequest) + + with pytest.raises(HTTPException) as exc_info: + asyncio.run(wrapper(_chat_request(chat_template="x"), _raw_request())) + + assert exc_info.value.status_code == 400 + entry_point.assert_not_awaited() + _assert_counts( + counters, + "a request with a forbidden chat template", + total_requests=1, + nonstream_requests=1, + http_exceptions=1, + ) + + +def test_conversation_id_comes_from_the_headers(tmp_path): + with isolated_serve_env(tmp_path) as env: + server, _ = _server_with_counters(env) + entry_point = AsyncMock(return_value=_chat_response()) + wrapper = server._wrap_entry_point(entry_point, ChatCompletionRequest) + + asyncio.run(wrapper(_chat_request(), _raw_request((b"x-session-id", b"s")))) + + (called_req, _), _ = entry_point.await_args + assert called_req.conversation_params is not None + assert called_req.conversation_params.conversation_id == "s" + + +def test_upstream_error_keeps_its_status(tmp_path): + with isolated_serve_env(tmp_path) as env: + server, counters = _server_with_counters(env) + entry_point = AsyncMock( + side_effect=aiohttp.ClientResponseError( + request_info=Mock(), history=(), status=404, message='{"message": "nope"}' + ) + ) + wrapper = server._wrap_entry_point(entry_point, ChatCompletionRequest) + + with pytest.raises(HTTPException) as exc_info: + asyncio.run(wrapper(_chat_request(), _raw_request())) + + assert exc_info.value.status_code == 404 + assert exc_info.value.detail == "nope" + _assert_counts( + counters, + "an upstream worker error", + total_requests=1, + nonstream_requests=1, + http_exceptions=1, + ) + + +def test_context_length_error_returns_the_worker_envelope(tmp_path): + with isolated_serve_env(tmp_path) as env: + server, counters = _server_with_counters(env) + worker_body = OpenAIServer.create_error_response( + context_length_exceeded_message(2048, 4096) + ).body.decode() + # The error the way openai_client raises one: post_json/_send_request + # use message=f"{reason}: {body[:2048]}" with the worker's response text. + entry_point = AsyncMock( + side_effect=aiohttp.ClientResponseError( + request_info=Mock(), + history=(), + status=400, + message=f"Bad Request: {worker_body}", + ) + ) + wrapper = server._wrap_entry_point(entry_point, ChatCompletionRequest) + + response = asyncio.run(wrapper(_chat_request(), _raw_request())) + + assert isinstance(response, JSONResponse) + assert response.status_code == 400 + assert json.loads(bytes(response.body))["code"] == "context_length_exceeded" + _assert_counts( + counters, + "a context-length rejection", + total_requests=1, + nonstream_requests=1, + http_exceptions=1, + ) + + +def test_engine_fatal_error_returns_none_after_sigint(tmp_path, monkeypatch): + with isolated_serve_env(tmp_path) as env: + server, counters = _server_with_counters(env) + recorder = Mock() + monkeypatch.setattr(signal, "raise_signal", recorder) + monkeypatch.setattr(disagg_server_module, "record_termination_observation", Mock()) + entry_point = AsyncMock(side_effect=CppExecutorError("x")) + wrapper = server._wrap_entry_point(entry_point, ChatCompletionRequest) + + response = asyncio.run(wrapper(_chat_request(), _raw_request())) + + assert response is None + recorder.assert_called_once_with(signal.SIGINT) + _assert_counts( + counters, + "an engine fatal error", + total_requests=1, + nonstream_requests=1, + ) + + +def test_unexpected_error_is_a_500(tmp_path): + with isolated_serve_env(tmp_path) as env: + server, counters = _server_with_counters(env) + entry_point = AsyncMock(side_effect=RuntimeError("boom")) + wrapper = server._wrap_entry_point(entry_point, ChatCompletionRequest) + + with pytest.raises(HTTPException) as exc_info: + asyncio.run(wrapper(_chat_request(), _raw_request())) + + assert exc_info.value.status_code == 500 + assert exc_info.value.detail == "Internal server error boom" + _assert_counts( + counters, + "an unexpected error", + total_requests=1, + nonstream_requests=1, + internal_errors=1, + ) + + +def test_serialization_failure_goes_to_the_error_path(tmp_path): + with isolated_serve_env(tmp_path) as env: + server, counters = _server_with_counters(env) + + class _ExplodingResponse: + def model_dump(self, **kwargs): + raise RuntimeError("cannot serialize") + + entry_point = AsyncMock(return_value=_ExplodingResponse()) + wrapper = server._wrap_entry_point(entry_point, ChatCompletionRequest) + + with pytest.raises(HTTPException) as exc_info: + asyncio.run(wrapper(_chat_request(), _raw_request())) + + assert exc_info.value.status_code == 500 + _assert_counts( + counters, + "a response that fails to serialize", + total_requests=1, + nonstream_requests=1, + total_responses=1, + internal_errors=1, + ) + + +class _FakeTracer: + def __init__(self): + self.spans = [] + + @contextlib.contextmanager + def start_as_current_span(self, name, kind=None, context=None): + self.spans.append((name, kind)) + yield SimpleNamespace(set_status=Mock(), record_exception=Mock()) + + +def test_one_trace_span_per_request(tmp_path, monkeypatch): + with isolated_serve_env(tmp_path) as env: + server, _ = _server_with_counters(env) + fake = _FakeTracer() + monkeypatch.setattr(tracing, "global_otlp_tracer", lambda: fake) + entry_point = AsyncMock(return_value=_chat_response()) + wrapper = server._wrap_entry_point(entry_point, ChatCompletionRequest) + + asyncio.run(wrapper(_chat_request(), _raw_request())) + assert fake.spans == [("disaggregated_request", SpanKind.SERVER)] + + with pytest.raises(HTTPException): + asyncio.run(wrapper(_chat_request(chat_template="x"), _raw_request())) + assert fake.spans == [("disaggregated_request", SpanKind.SERVER)] * 2 + + +def test_anthropic_messages_still_counts_like_chat(tmp_path): + with isolated_serve_env(tmp_path) as env: + server, counters = _server_with_counters(env) + server._service = SimpleNamespace( + openai_chat_completion=AsyncMock(return_value=_chat_response()) + ) + client = env.client(server.app, lifespan=False) + + response = client.post( + "/v1/messages", + json={ + "model": "m", + "max_tokens": 16, + "messages": [{"role": "user", "content": "hi"}], + }, + ) + + assert response.status_code == 200 + assert response.json()["content"][0]["text"] == "hello" + _assert_counts( + counters, + "an Anthropic Messages request", + total_requests=1, + nonstream_requests=1, + total_responses=1, + ) + + +def test_non_streaming_responses_request_serializes_by_alias(tmp_path): + with isolated_serve_env(tmp_path) as env: + server, counters = _server_with_counters(env) + responses_response = ResponsesResponse( + model="m", + output=[], + parallel_tool_calls=False, + temperature=1.0, + tool_choice="auto", + tools=[], + top_p=1.0, + background=False, + service_tier="auto", + status="completed", + text=ResponseTextConfig( + format=ResponseFormatTextJSONSchemaConfig( + name="shape", schema={"type": "object"}, type="json_schema" + ) + ), + top_logprobs=0, + truncation="disabled", + ) + entry_point = AsyncMock(return_value=responses_response) + wrapper = server._wrap_entry_point(entry_point, ResponsesRequest) + + response = asyncio.run(wrapper(ResponsesRequest(model="m", input="hi"), _raw_request())) + + assert isinstance(response, JSONResponse) + body = json.loads(bytes(response.body)) + # by_alias serialization: text.format.schema goes out under its wire + # name, not the Python field name schema_. + assert body["text"]["format"]["schema"] == {"type": "object"} + assert "schema_" not in body["text"]["format"] + assert body == responses_response.model_dump(by_alias=True) + _assert_counts( + counters, + "a non-streaming Responses request", + total_requests=1, + nonstream_requests=1, + total_responses=1, + ) + + +def test_streamed_responses_request_is_guarded(tmp_path): + with isolated_serve_env(tmp_path) as env: + server, counters = _server_with_counters(env) + + async def generated(): + yield b"data: {}\n\n" + + generator = generated() + entry_point = AsyncMock(return_value=generator) + wrapper = server._wrap_entry_point(entry_point, ResponsesRequest) + req = ResponsesRequest(model="m", input="hi", stream=True) + + response = asyncio.run(wrapper(req, _raw_request())) + + assert isinstance(response, StreamingResponse) + assert response.media_type == "text/event-stream" + # A relayed Responses stream is wrapped so that an early failure ends + # with an `error` event; the raw generator must not go out as-is. + assert response.body_iterator is not generator + _assert_counts( + counters, + "a streamed Responses request", + total_requests=1, + stream_requests=1, + total_responses=1, + ) diff --git a/tests/unittest/disaggregated/test_disagg_server_contract.py b/tests/unittest/disaggregated/test_disagg_server_contract.py new file mode 100644 index 000000000000..28a7121f2f78 --- /dev/null +++ b/tests/unittest/disaggregated/test_disagg_server_contract.py @@ -0,0 +1,347 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""What a subclass of ``OpenAIDisaggServer`` outside this repository relies on. + +NeMo-RL's ``OpenAIDisaggServerAdaptor`` subclasses ``OpenAIDisaggServer``. It +constructs the server as ``(config, req_timeout_secs=..., coordinator_url=None)`` +plus one keyword argument of its own, adds a pure-ASGI middleware to +``self.app`` right after ``super().__init__()`` (so the app must exist with +every route registered when ``__init__`` returns), replaces ``__call__(host, +port, sockets=None)`` to keep a handle on its own uvicorn server, overrides +``_wrap_entry_point`` (swapping the request type and re-wrapping the returned +handler) and ``_handle_exception`` (answering a 4xx upstream error itself, +counting it on ``_perf_metrics_collector.http_exceptions``, delegating the +rest to ``super()``). These tests build a subclass of the same shape and pin +that surface. A failing test here means the change needs the NeMo-RL owners' +agreement first; see the ``OpenAIDisaggServer`` class docstring. +""" + +import asyncio +import importlib +import inspect +import json +from types import SimpleNamespace +from unittest.mock import AsyncMock, Mock + +import aiohttp +import pytest +from fastapi import FastAPI, HTTPException, Request, Response +from fastapi.responses import JSONResponse +from test_common.serve.app_factory import isolated_serve_env +from test_common.serve.routes import effective_routes + +from tensorrt_llm.llmapi.disagg_utils import ( + CtxGenServerConfig, + DisaggServerConfig, + RouterConfig, + ServerRole, +) +from tensorrt_llm.serve.openai_protocol import ( + ChatCompletionRequest, + ChatCompletionResponse, + ChatCompletionResponseChoice, + ChatMessage, + CompletionRequest, + UCompletionRequest, + UsageInfo, +) + +pytestmark = pytest.mark.cpu_only + +# The import paths of the class. A move that keeps the old path as a shim adds +# the new path here and keeps the old one, so the contract holds through both. +CONTRACT_MODULES = ("tensorrt_llm.serve.openai_disagg_server",) + + +@pytest.fixture(params=CONTRACT_MODULES) +def server_cls(request): + return importlib.import_module(request.param).OpenAIDisaggServer + + +def _contract(item: str) -> str: + return ( + f"{item}: NeMo-RL subclasses OpenAIDisaggServer and overrides this; " + "agree the change with the NeMo-RL owners first " + "(see the OpenAIDisaggServer docstring)" + ) + + +class _HeaderMiddleware: + """Pure ASGI, the way NeMo-RL installs its request adaptor.""" + + def __init__(self, app): + self.app = app + + async def __call__(self, scope, receive, send): + if scope["type"] != "http": + await self.app(scope, receive, send) + return + + async def send_with_header(message): + if message["type"] == "http.response.start": + message.setdefault("headers", []).append((b"x-contract", b"1")) + await send(message) + + await self.app(scope, receive, send_with_header) + + +def _contract_subclass(server_cls): + """A subclass shaped like NeMo-RL's ``OpenAIDisaggServerAdaptor``.""" + + class _Adaptor(server_cls): + def __init__(self, *args, **kwargs): + self.marker = kwargs.pop("marker") + self.wrapped = [] + super().__init__(*args, **kwargs) + self.app.add_middleware(_HeaderMiddleware) + + async def __call__(self, host, port, sockets=None): + raise AssertionError("the contract tests never serve") + + def _handle_exception(self, exception): + if isinstance(exception, aiohttp.ClientResponseError) and 400 <= exception.status < 500: + self._perf_metrics_collector.http_exceptions.inc() + raise HTTPException(status_code=exception.status, detail=exception.message) + return super()._handle_exception(exception) + + def _wrap_entry_point(self, entry_point, request_type=UCompletionRequest): + self.wrapped.append((entry_point, request_type)) + inner = super()._wrap_entry_point(entry_point, request_type) + + async def wrapper(req: request_type, raw_req: Request) -> Response: + response = await inner(req, raw_req) + if isinstance(response, JSONResponse): + response.headers["x-wrapped"] = "1" + return response + + return wrapper + + return _Adaptor + + +def _disagg_server_config() -> DisaggServerConfig: + """Two workers on closed loopback ports: construction is network-free.""" + return DisaggServerConfig( + server_configs=[ + CtxGenServerConfig(type="ctx", hostname="127.0.0.1", port=9), + CtxGenServerConfig(type="gen", hostname="127.0.0.1", port=10), + ], + hostname="127.0.0.1", + port=0, + ctx_router_config=RouterConfig(server_role=ServerRole.CONTEXT), + gen_router_config=RouterConfig(server_role=ServerRole.GENERATION), + node_id=0, + ) + + +def _build_adaptor(server_cls): + """Construct the subclass with NeMo-RL's call shape.""" + config = _disagg_server_config() + return _contract_subclass(server_cls)( + config, req_timeout_secs=7, coordinator_url=None, marker=1 + ) + + +def _raw_request() -> Request: + return Request( + { + "type": "http", + "headers": [], + "state": {"server_arrival_time": 1.0, "perf_metrics_records": []}, + } + ) + + +def _run_through_subclass(coroutine): + """Await a wrapped request built by the subclass. + + A TypeError here means the base changed how it calls what the subclass + overrides: NeMo-RL's _handle_exception takes exactly one positional + argument (the exception) and its wrapper takes (req, raw_req). + """ + try: + return asyncio.run(coroutine) + except TypeError as error: + raise AssertionError( + _contract( + "_handle_exception called with exactly one positional argument (the " + "exception) and the wrapper called with (req, raw_req)" + ) + ) from error + + +def _chat_response() -> ChatCompletionResponse: + return ChatCompletionResponse( + id="chatcmpl-contract-test", + model="m", + choices=[ + ChatCompletionResponseChoice( + index=0, + message=ChatMessage(role="assistant", content="hello"), + finish_reason="stop", + ) + ], + usage=UsageInfo(prompt_tokens=3, completion_tokens=2, total_tokens=5), + ) + + +def test_constructor_takes_config_timeout_and_coordinator_url(tmp_path, server_cls): + config = _disagg_server_config() + try: + inspect.signature(server_cls).bind(config, req_timeout_secs=7, coordinator_url=None) + except TypeError as error: + raise AssertionError( + _contract( + "the constructor call shape (config, req_timeout_secs=..., coordinator_url=None)" + ) + ) from error + with isolated_serve_env(tmp_path): + server = _build_adaptor(server_cls) + assert server.marker == 1, _contract( + "a subclass keyword argument popped before super().__init__" + ) + + +def test_app_and_routes_exist_when_init_returns(tmp_path, server_cls): + with isolated_serve_env(tmp_path): + server = _build_adaptor(server_cls) + assert isinstance(server.app, FastAPI), _contract("self.app set when __init__ returns") + post_routes = { + path: route + for path, route in effective_routes(server.app) + if "POST" in (getattr(route, "methods", None) or ()) + } + for path in ("/v1/completions", "/v1/chat/completions"): + assert path in post_routes, _contract( + f"the {path} route registered when __init__ returns" + ) + assert post_routes[path].endpoint.__module__ == __name__, _contract( + f"the {path} endpoint is the wrapper the subclass returned" + ) + wrapped_types = [request_type for _, request_type in server.wrapped] + assert any(t is CompletionRequest for t in wrapped_types), _contract( + "_wrap_entry_point called with CompletionRequest, compared by identity" + ) + assert any(t is ChatCompletionRequest for t in wrapped_types), _contract( + "_wrap_entry_point called with ChatCompletionRequest, compared by identity" + ) + subclass_endpoints = [ + route + for _, route in effective_routes(server.app) + if getattr(route, "endpoint", None) is not None + and route.endpoint.__module__ == __name__ + ] + assert len(server.wrapped) == len(subclass_endpoints), _contract( + "every route built through _wrap_entry_point goes through the override" + ) + + +def test_middleware_added_after_init_runs(tmp_path, server_cls): + with isolated_serve_env(tmp_path) as env: + server = _build_adaptor(server_cls) + response = env.client(server.app, lifespan=False).get("/version") + assert response.status_code == 200, _contract( + "a request served by self.app after construction" + ) + assert response.headers.get("x-contract") == "1", _contract( + "middleware added to self.app right after super().__init__ runs" + ) + + +def test_wrapper_result_reaches_the_client(tmp_path, server_cls): + with isolated_serve_env(tmp_path): + server = _build_adaptor(server_cls) + chat_response = _chat_response() + wrapper = server._wrap_entry_point( + AsyncMock(return_value=chat_response), ChatCompletionRequest + ) + req = ChatCompletionRequest(model="m", messages=[{"role": "user", "content": "hi"}]) + + response = _run_through_subclass(wrapper(req, _raw_request())) + + assert isinstance(response, JSONResponse), _contract( + "the base wrapper's non-streaming result, re-wrappable by a subclass" + ) + assert response.headers.get("x-wrapped") == "1", _contract( + "the subclass wrapper around the base wrapper's result" + ) + assert json.loads(bytes(response.body)) == chat_response.model_dump(by_alias=True), ( + _contract("the response body the base wrapper serializes") + ) + + +def test_every_failure_reaches_the_overridden_handle_exception(tmp_path, server_cls): + with isolated_serve_env(tmp_path): + server = _build_adaptor(server_cls) + assert callable(server._perf_metrics_collector.http_exceptions.inc), _contract( + "_perf_metrics_collector.http_exceptions.inc" + ) + server._perf_metrics_collector = SimpleNamespace( + http_exceptions=SimpleNamespace(inc=Mock()), + internal_errors=SimpleNamespace(inc=Mock()), + queue_latency_seconds=SimpleNamespace(observe=Mock()), + total_requests=SimpleNamespace(inc=Mock()), + stream_requests=SimpleNamespace(inc=Mock()), + nonstream_requests=SimpleNamespace(inc=Mock()), + total_responses=SimpleNamespace(inc=Mock()), + ) + req = ChatCompletionRequest(model="m", messages=[{"role": "user", "content": "hi"}]) + + upstream = AsyncMock( + side_effect=aiohttp.ClientResponseError( + request_info=Mock(), history=(), status=413, message="too long" + ) + ) + with pytest.raises(HTTPException) as exc_info: + _run_through_subclass( + server._wrap_entry_point(upstream, ChatCompletionRequest)(req, _raw_request()) + ) + assert exc_info.value.status_code == 413, _contract( + "a 4xx upstream error answered by the overridden _handle_exception" + ) + assert exc_info.value.detail == "too long", _contract( + "the overridden _handle_exception's own error text, not the base's" + ) + assert server._perf_metrics_collector.http_exceptions.inc.call_count == 1, _contract( + "the override counting on _perf_metrics_collector.http_exceptions" + ) + + unexpected = AsyncMock(side_effect=RuntimeError("boom")) + with pytest.raises(HTTPException) as exc_info: + _run_through_subclass( + server._wrap_entry_point(unexpected, ChatCompletionRequest)(req, _raw_request()) + ) + assert exc_info.value.status_code == 500, _contract( + "errors the override delegates reach super()._handle_exception" + ) + assert exc_info.value.detail == "Internal server error boom", _contract( + "super()._handle_exception's 500 for a delegated error" + ) + assert server._perf_metrics_collector.internal_errors.inc.call_count == 1, _contract( + "super()._handle_exception counting internal errors" + ) + + +def test_call_serves_host_port_sockets(server_cls): + assert inspect.iscoroutinefunction(server_cls.__call__), _contract( + "__call__ as the coroutine that serves" + ) + parameters = inspect.signature(server_cls.__call__).parameters + assert list(parameters) == ["self", "host", "port", "sockets"], _contract( + "__call__(self, host, port, sockets)" + ) + assert parameters["sockets"].default is None, _contract( + "the sockets argument of __call__ defaulting to None" + ) From 533cb0b89db2bdade8d6dbefa0801b95e5fe1a71 Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Mon, 2 Nov 2026 09:57:07 +0000 Subject: [PATCH 08/27] [TRTLLM-17262][refactor] Split the disaggregated entry-point wrapper into prepare, run, finish and serialize steps No behavior change. The body of the closure OpenAIDisaggServer builds in _wrap_entry_point becomes a composition of five methods, so that later changes (disconnect watching, in-stream errors, request trace) land in the step they belong to instead of growing one closure: - _prepare_request: the request counters, the chat-template permission (HTTPException 400), the conversation identity, the per-request hooks - _run_entry_point: the call into the disaggregated service - _finish_request: success accounting (total_responses) - _finish_failed_request: every failure, through _handle_exception while the exception is being handled (its traceback logging needs that) - _serialize_response: StreamingResponse with the Responses stream guard, or JSONResponse(content=...model_dump(by_alias=True)), as before The trace span stays on the wrapper that FastAPI calls; the steps open no spans of their own. The statements keep their order inside the same try block: the counters run before the template check, total_responses before serialization, and a serialization failure still reaches _handle_exception. Some step parameters (raw_req of _run_entry_point, req and hooks of the finish steps) are not read yet; they are the inputs the known upcoming changes to each step need, so those changes will not have to alter the signatures again. The characterization tests and the subclass contract tests of the previous commit pass unchanged on both sides of the split. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- tensorrt_llm/serve/openai_disagg_server.py | 111 +++++++++++++++------ 1 file changed, 79 insertions(+), 32 deletions(-) diff --git a/tensorrt_llm/serve/openai_disagg_server.py b/tensorrt_llm/serve/openai_disagg_server.py index 29a7b0cb354c..db11cc39b26f 100644 --- a/tensorrt_llm/serve/openai_disagg_server.py +++ b/tensorrt_llm/serve/openai_disagg_server.py @@ -70,6 +70,7 @@ PerfMetricsMiddleware, combine_disagg_metrics) from tensorrt_llm.serve.responses_utils import (ServerArrivalTimeMiddleware, + UCompletionResponseOrGenerator, get_steady_clock_now_in_seconds, guard_responses_stream, stream_error_event) @@ -436,44 +437,90 @@ def _wrap_entry_point(self, entry_point: Callable, request_type: type = UComplet # annotation with request_type (as openai_server.py does). @tracing.trace_span("disaggregated_request") async def wrapper(req: request_type, raw_req: Request) -> Response: + hooks: Optional[RawRequestResponseHooks] = None try: - self._perf_metrics_collector.total_requests.inc() - if req.stream: - self._perf_metrics_collector.stream_requests.inc() - else: - self._perf_metrics_collector.nonstream_requests.inc() - try: - ensure_request_chat_template_allowed( - req, self._allow_request_chat_template) - except ValueError as e: - raise HTTPException(status_code=400, detail=str(e)) from e - self._extract_conversation_id( - req, raw_req, - self._config.conversation_affinity_header_for_subagents) - hooks = RawRequestResponseHooks( - raw_req, self._perf_metrics_collector.queue_latency_seconds, - self._collect_perf_metrics) - response_or_generator = await entry_point(req, hooks) - self._perf_metrics_collector.total_responses.inc() - if req.stream: - stream = response_or_generator - if isinstance(req, ResponsesRequest): - # A relayed Responses stream that fails early ends with - # an `error` event; other protocols end on a sentinel. - stream = guard_responses_stream(stream, - stream_error_event) - return StreamingResponse( - content=stream, - media_type="text/event-stream") - # by_alias: e.g. text.format.schema goes out under its wire name. - return JSONResponse( - content=response_or_generator.model_dump(by_alias=True)) + hooks = self._prepare_request(req, raw_req) + response_or_generator = await self._run_entry_point( + entry_point, req, hooks, raw_req) + self._finish_request(req, hooks) + return self._serialize_response(req, response_or_generator) except Exception as e: # Usually raises; returns a Response for worker errors that # carry a machine-readable code (context_length_exceeded). - return self._handle_exception(e) + return self._finish_failed_request(e, req, hooks) return wrapper + def _prepare_request(self, req: UCompletionRequest, + raw_req: Request) -> RawRequestResponseHooks: + """Account for a request and check it before any worker is involved. + + Counts the request, applies the chat-template permission (HTTPException + 400), resolves the conversation identity from the headers and builds + the per-request hooks. + """ + self._perf_metrics_collector.total_requests.inc() + if req.stream: + self._perf_metrics_collector.stream_requests.inc() + else: + self._perf_metrics_collector.nonstream_requests.inc() + try: + ensure_request_chat_template_allowed( + req, self._allow_request_chat_template) + except ValueError as e: + raise HTTPException(status_code=400, detail=str(e)) from e + self._extract_conversation_id( + req, raw_req, + self._config.conversation_affinity_header_for_subagents) + return RawRequestResponseHooks( + raw_req, self._perf_metrics_collector.queue_latency_seconds, + self._collect_perf_metrics) + + async def _run_entry_point(self, entry_point: Callable, + req: UCompletionRequest, + hooks: RawRequestResponseHooks, + raw_req: Request) -> UCompletionResponseOrGenerator: + """Run the service entry point for one request. + + A typed response for non-streaming requests, an async generator of the + worker's raw SSE bytes for streaming ones. + """ + return await entry_point(req, hooks) + + def _finish_request(self, req: UCompletionRequest, + hooks: RawRequestResponseHooks) -> None: + """Account for a request the service answered.""" + self._perf_metrics_collector.total_responses.inc() + + def _finish_failed_request( + self, error: Exception, req: UCompletionRequest, + hooks: Optional[RawRequestResponseHooks]) -> Optional[Response]: + """Answer a failed request through _handle_exception. + + Must be called while ``error`` is being handled: _handle_exception logs + with traceback.format_exc(). Raises HTTPException in most cases, returns + a JSONResponse for worker errors that carry a machine-readable code, and + returns None after an engine fatal error (SIGINT already sent). + """ + return self._handle_exception(error) + + def _serialize_response( + self, req: UCompletionRequest, + response_or_generator: UCompletionResponseOrGenerator) -> Response: + """The HTTP response of the OpenAI routes.""" + if req.stream: + stream = response_or_generator + if isinstance(req, ResponsesRequest): + # A relayed Responses stream that fails early ends with + # an `error` event; other protocols end on a sentinel. + stream = guard_responses_stream(stream, + stream_error_event) + return StreamingResponse( + content=stream, + media_type="text/event-stream") + # by_alias: e.g. text.format.schema goes out under its wire name. + return JSONResponse( + content=response_or_generator.model_dump(by_alias=True)) + async def anthropic_messages(self, request: AnthropicMessagesRequest, raw_request: Request) -> Response: """Serve Anthropic Messages through the disaggregated chat pipeline.""" 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/27] [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/27] [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/27] [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/27] [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/27] [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/27] [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/27] [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/27] [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/27] [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/27] [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/27] [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/27] [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 752da8bbdd275764ee196f56935dcedd73b6896c Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Tue, 3 Nov 2026 00:28:05 +0000 Subject: [PATCH 21/27] [TRTLLM-17263][refactor] Move the Anthropic Messages modules into serve/api/anthropic Move tensorrt_llm/serve/anthropic_adapter.py, anthropic_protocol.py and anthropic_batches.py whole into tensorrt_llm/serve/api/anthropic/ as adapter.py, protocol.py and batches.py, and rewrite every import of them in the repository. The old paths are deleted without a compatibility shim: no code outside the repository imports them. Generated change: the three files are copied byte for byte apart from their import lines; the importers change only in their import lines. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- .../anthropic/adapter.py} | 2 +- .../anthropic/batches.py} | 2 +- .../anthropic/protocol.py} | 0 tensorrt_llm/serve/openai_disagg_server.py | 16 ++++++------- tensorrt_llm/serve/openai_disagg_service.py | 2 +- tensorrt_llm/serve/openai_server.py | 18 +++++++------- .../llmapi/apps/test_anthropic_adapter.py | 6 ++--- .../llmapi/apps/test_anthropic_batches.py | 4 ++-- .../llmapi/apps/test_anthropic_routes.py | 24 +++++++++---------- .../apps/test_chat_template_kwargs_guard.py | 4 ++-- .../llmapi/apps/test_kimi_serve_extensions.py | 4 ++-- 11 files changed, 40 insertions(+), 42 deletions(-) rename tensorrt_llm/serve/{anthropic_adapter.py => api/anthropic/adapter.py} (99%) rename tensorrt_llm/serve/{anthropic_batches.py => api/anthropic/batches.py} (99%) rename tensorrt_llm/serve/{anthropic_protocol.py => api/anthropic/protocol.py} (100%) diff --git a/tensorrt_llm/serve/anthropic_adapter.py b/tensorrt_llm/serve/api/anthropic/adapter.py similarity index 99% rename from tensorrt_llm/serve/anthropic_adapter.py rename to tensorrt_llm/serve/api/anthropic/adapter.py index 267b77ce2d1d..74c5259f187f 100644 --- a/tensorrt_llm/serve/anthropic_adapter.py +++ b/tensorrt_llm/serve/api/anthropic/adapter.py @@ -37,7 +37,7 @@ from fastapi.responses import JSONResponse from tensorrt_llm.logger import logger -from tensorrt_llm.serve.anthropic_protocol import ( +from tensorrt_llm.serve.api.anthropic.protocol import ( AnthropicContentBlockDeltaEvent, AnthropicContentBlockStartEvent, AnthropicContentBlockStopEvent, diff --git a/tensorrt_llm/serve/anthropic_batches.py b/tensorrt_llm/serve/api/anthropic/batches.py similarity index 99% rename from tensorrt_llm/serve/anthropic_batches.py rename to tensorrt_llm/serve/api/anthropic/batches.py index 6d693df7c431..c6583367e91e 100644 --- a/tensorrt_llm/serve/anthropic_batches.py +++ b/tensorrt_llm/serve/api/anthropic/batches.py @@ -41,7 +41,7 @@ from typing import Any, Awaitable, Callable, Dict, List, Literal, Optional, Tuple from tensorrt_llm.logger import logger -from tensorrt_llm.serve.anthropic_protocol import ( +from tensorrt_llm.serve.api.anthropic.protocol import ( AnthropicBatchRequestCounts, AnthropicBatchRequestItem, AnthropicMessageBatch, diff --git a/tensorrt_llm/serve/anthropic_protocol.py b/tensorrt_llm/serve/api/anthropic/protocol.py similarity index 100% rename from tensorrt_llm/serve/anthropic_protocol.py rename to tensorrt_llm/serve/api/anthropic/protocol.py diff --git a/tensorrt_llm/serve/openai_disagg_server.py b/tensorrt_llm/serve/openai_disagg_server.py index d8ebec1a6ce7..a988886956bb 100644 --- a/tensorrt_llm/serve/openai_disagg_server.py +++ b/tensorrt_llm/serve/openai_disagg_server.py @@ -42,14 +42,14 @@ MetadataServerConfig, ServerRole) from tensorrt_llm.logger import logger from tensorrt_llm.serve._telemetry import create_uvicorn_server -from tensorrt_llm.serve.anthropic_adapter import (AnthropicRequestError, - AnthropicResponseError, - anthropic_error_response, - convert_anthropic_request, - convert_chat_response, - reframe_openai_stream) -from tensorrt_llm.serve.anthropic_protocol import (AnthropicCountTokensRequest, - AnthropicMessagesRequest) +from tensorrt_llm.serve.api.anthropic.adapter import (AnthropicRequestError, + AnthropicResponseError, + anthropic_error_response, + convert_anthropic_request, + convert_chat_response, + reframe_openai_stream) +from tensorrt_llm.serve.api.anthropic.protocol import ( + AnthropicCountTokensRequest, AnthropicMessagesRequest) from tensorrt_llm.serve.cluster_storage import ( HttpClusterStorageServer, create_cluster_storage, validate_http_cluster_storage_scope) diff --git a/tensorrt_llm/serve/openai_disagg_service.py b/tensorrt_llm/serve/openai_disagg_service.py index c62d62b31f8e..bc9e02945756 100644 --- a/tensorrt_llm/serve/openai_disagg_service.py +++ b/tensorrt_llm/serve/openai_disagg_service.py @@ -23,7 +23,7 @@ from tensorrt_llm._utils import get_global_steady_clock_now_in_seconds from tensorrt_llm.llmapi.disagg_utils import ConditionalDisaggConfig, DisaggServerConfig, ServerRole from tensorrt_llm.logger import logger -from tensorrt_llm.serve.anthropic_protocol import ( +from tensorrt_llm.serve.api.anthropic.protocol import ( AnthropicCountTokensRequest, AnthropicCountTokensResponse, ) diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py index 92f2d9478baf..284d55c08d15 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -69,19 +69,17 @@ 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 ( +from tensorrt_llm.serve.api.anthropic.adapter import ( AnthropicRequestError, AnthropicResponseError, anthropic_error_response, convert_anthropic_count_tokens_request, convert_anthropic_request, convert_chat_response, reframe_openai_stream) -from tensorrt_llm.serve.anthropic_batches import (AnthropicBatchStore, - BatchStoreFullError, - results_to_jsonl) -from tensorrt_llm.serve.anthropic_protocol import (AnthropicBatchDeleteResponse, - AnthropicBatchList, - AnthropicCountTokensRequest, - AnthropicCountTokensResponse, - AnthropicCreateBatchRequest, - AnthropicMessagesRequest) +from tensorrt_llm.serve.api.anthropic.batches import (AnthropicBatchStore, + BatchStoreFullError, + results_to_jsonl) +from tensorrt_llm.serve.api.anthropic.protocol import ( + AnthropicBatchDeleteResponse, AnthropicBatchList, + AnthropicCountTokensRequest, 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, diff --git a/tests/unittest/llmapi/apps/test_anthropic_adapter.py b/tests/unittest/llmapi/apps/test_anthropic_adapter.py index c06404d993dc..5cfe060fe8fd 100644 --- a/tests/unittest/llmapi/apps/test_anthropic_adapter.py +++ b/tests/unittest/llmapi/apps/test_anthropic_adapter.py @@ -27,7 +27,7 @@ from pydantic import ValidationError import tensorrt_llm.serve.anthropic_adapter as adapter -from tensorrt_llm.serve.anthropic_adapter import ( +from tensorrt_llm.serve.api.anthropic.adapter import ( AnthropicRequestError, AnthropicResponseError, AnthropicStreamReframer, @@ -37,7 +37,7 @@ map_stop_reason, reframe_openai_stream, ) -from tensorrt_llm.serve.anthropic_protocol import AnthropicMessagesRequest +from tensorrt_llm.serve.api.anthropic.protocol import AnthropicMessagesRequest from tensorrt_llm.serve.openai_protocol import ( ChatCompletionResponse, ChatCompletionResponseChoice, @@ -1758,7 +1758,7 @@ def test_every_known_block_type_selects_its_own_model(): catch-all -- where it would be reported as unsupported even though the server models it. Nothing else would notice. """ - from tensorrt_llm.serve.anthropic_protocol import ( + from tensorrt_llm.serve.api.anthropic.protocol import ( _CONTENT_BLOCK_MODELS, AnthropicMessage, AnthropicUnknownBlock, diff --git a/tests/unittest/llmapi/apps/test_anthropic_batches.py b/tests/unittest/llmapi/apps/test_anthropic_batches.py index 3c3a69ba69f9..a8a16b11327a 100644 --- a/tests/unittest/llmapi/apps/test_anthropic_batches.py +++ b/tests/unittest/llmapi/apps/test_anthropic_batches.py @@ -23,12 +23,12 @@ import pytest -from tensorrt_llm.serve.anthropic_batches import ( +from tensorrt_llm.serve.api.anthropic.batches import ( AnthropicBatchStore, BatchStoreFullError, results_to_jsonl, ) -from tensorrt_llm.serve.anthropic_protocol import ( +from tensorrt_llm.serve.api.anthropic.protocol import ( AnthropicBatchRequestItem, AnthropicCreateBatchRequest, ) diff --git a/tests/unittest/llmapi/apps/test_anthropic_routes.py b/tests/unittest/llmapi/apps/test_anthropic_routes.py index b7e6cf62a3d6..ba989f408193 100644 --- a/tests/unittest/llmapi/apps/test_anthropic_routes.py +++ b/tests/unittest/llmapi/apps/test_anthropic_routes.py @@ -523,7 +523,7 @@ def test_standard_and_disagg_register_messages_route(prometheus_multiproc_dir): def _batch_client(runner=None): """A TestClient over just the batch routes, with a stubbed item runner.""" - from tensorrt_llm.serve.anthropic_batches import AnthropicBatchStore + from tensorrt_llm.serve.api.anthropic.batches import AnthropicBatchStore async def ok(request): return "succeeded", {"type": "message", "role": "assistant", "content": []} @@ -630,8 +630,8 @@ def test_results_of_an_unfinished_batch_are_400_not_404(): still be in flight by the time the second request arrives, and the batch would have already ended. """ - from tensorrt_llm.serve.anthropic_batches import AnthropicBatchStore - from tensorrt_llm.serve.anthropic_protocol import AnthropicBatchRequestItem + from tensorrt_llm.serve.api.anthropic.batches import AnthropicBatchStore + from tensorrt_llm.serve.api.anthropic.protocol import AnthropicBatchRequestItem async def scenario(): blocked = asyncio.Event() @@ -673,8 +673,8 @@ def test_interrupted_batch_still_reports_counts_that_sum(): end the batch with fewer results than requests - and request_counts is the number clients use to decide whether every request was accounted for. """ - from tensorrt_llm.serve.anthropic_batches import AnthropicBatchStore - from tensorrt_llm.serve.anthropic_protocol import AnthropicBatchRequestItem + from tensorrt_llm.serve.api.anthropic.batches import AnthropicBatchStore + from tensorrt_llm.serve.api.anthropic.protocol import AnthropicBatchRequestItem async def scenario(): blocked = asyncio.Event() @@ -850,7 +850,7 @@ def _anthropic_validation_handler(app): """ from fastapi.exceptions import RequestValidationError - from tensorrt_llm.serve.anthropic_adapter import anthropic_error_response + from tensorrt_llm.serve.api.anthropic.adapter import anthropic_error_response @app.exception_handler(RequestValidationError) async def _handler(request, exc): # noqa: ANN001 @@ -983,7 +983,7 @@ def _disagg_service(servers, post_json=None): def test_count_tokens_forwards_to_a_context_worker(): """The disagg server has no tokenizer, so the count must come from a worker.""" - from tensorrt_llm.serve.anthropic_protocol import ( + from tensorrt_llm.serve.api.anthropic.protocol import ( AnthropicCountTokensRequest, AnthropicCountTokensResponse, ) @@ -1004,7 +1004,7 @@ def test_count_tokens_forwards_to_a_context_worker(): def test_count_tokens_spreads_over_context_workers(): """Round-robin: one worker must not absorb every count.""" - from tensorrt_llm.serve.anthropic_protocol import ( + from tensorrt_llm.serve.api.anthropic.protocol import ( AnthropicCountTokensRequest, AnthropicCountTokensResponse, ) @@ -1029,7 +1029,7 @@ def test_count_tokens_without_context_workers_is_an_error_not_a_zero(): Returning 0 would read as an empty prompt and silently skew a client's context budgeting. """ - from tensorrt_llm.serve.anthropic_protocol import AnthropicCountTokensRequest + from tensorrt_llm.serve.api.anthropic.protocol import AnthropicCountTokensRequest service = _disagg_service([]) request = AnthropicCountTokensRequest(model=MODEL, messages=[{"role": "user", "content": "hi"}]) @@ -1085,7 +1085,7 @@ def test_post_json_sends_msgpack_the_worker_can_decode(): """ import msgspec - from tensorrt_llm.serve.anthropic_protocol import ( + from tensorrt_llm.serve.api.anthropic.protocol import ( AnthropicCountTokensRequest, AnthropicCountTokensResponse, ) @@ -1117,7 +1117,7 @@ def test_post_json_raises_on_an_error_status(): """ import aiohttp - from tensorrt_llm.serve.anthropic_protocol import ( + from tensorrt_llm.serve.api.anthropic.protocol import ( AnthropicCountTokensRequest, AnthropicCountTokensResponse, ) @@ -1195,7 +1195,7 @@ def test_create_batch_at_capacity_is_429_not_400(): each request runs in its own event loop and a batch's task does not survive between them. """ - from tensorrt_llm.serve.anthropic_batches import BatchStoreFullError + from tensorrt_llm.serve.api.anthropic.batches import BatchStoreFullError app = FastAPI() server = object.__new__(OpenAIServer) diff --git a/tests/unittest/llmapi/apps/test_chat_template_kwargs_guard.py b/tests/unittest/llmapi/apps/test_chat_template_kwargs_guard.py index a3de0428a00a..34a55bb1990d 100644 --- a/tests/unittest/llmapi/apps/test_chat_template_kwargs_guard.py +++ b/tests/unittest/llmapi/apps/test_chat_template_kwargs_guard.py @@ -373,8 +373,8 @@ class TestServerInjectedControls: @staticmethod def _anthropic_chat_request(): - from tensorrt_llm.serve.anthropic_adapter import convert_anthropic_request - from tensorrt_llm.serve.anthropic_protocol import AnthropicMessagesRequest + from tensorrt_llm.serve.api.anthropic.adapter import convert_anthropic_request + from tensorrt_llm.serve.api.anthropic.protocol import AnthropicMessagesRequest return convert_anthropic_request( AnthropicMessagesRequest( diff --git a/tests/unittest/llmapi/apps/test_kimi_serve_extensions.py b/tests/unittest/llmapi/apps/test_kimi_serve_extensions.py index f2cc94d81a50..156d03085b13 100644 --- a/tests/unittest/llmapi/apps/test_kimi_serve_extensions.py +++ b/tests/unittest/llmapi/apps/test_kimi_serve_extensions.py @@ -506,8 +506,8 @@ class TestKimiAnthropicThinkingControls: """Anthropic ``thinking`` / ``output_config.effort`` reach the K3 template kwargs.""" def convert(self, **fields) -> ChatCompletionRequest: - from tensorrt_llm.serve.anthropic_adapter import convert_anthropic_request - from tensorrt_llm.serve.anthropic_protocol import AnthropicMessagesRequest + from tensorrt_llm.serve.api.anthropic.adapter import convert_anthropic_request + from tensorrt_llm.serve.api.anthropic.protocol import AnthropicMessagesRequest req = convert_anthropic_request( AnthropicMessagesRequest(model="m", max_tokens=4096, messages=[USER_MSG], **fields) From 147c1148977e427a74e0c833ee36ed812ed299aa Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Tue, 3 Nov 2026 00:40:05 +0000 Subject: [PATCH 22/27] [TRTLLM-17263][test] Pin both servers' Anthropic Messages answers and mark the new package Add tensorrt_llm/serve/api/anthropic/__init__.py, retarget the two module imports of the adapter in test_anthropic_adapter.py to the new path, and add tests/unittest/serve/api/anthropic/test_messages_behavior.py: characterization tests of what the aggregated and the disaggregated server answer on the Anthropic routes today (success, streaming, a stream without chunks, rejected requests, context length, upstream errors, unexpected errors, engine failures, an unconvertible chat response, and the aggregated Message Batches), with the served model and the requested model named differently so that each server's own answer is pinned. The follow-up PR that gives both servers one Anthropic Messages serving class must keep this file unchanged and green. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- tensorrt_llm/serve/api/anthropic/__init__.py | 15 + .../llmapi/apps/test_anthropic_adapter.py | 4 +- .../unittest/serve/api/anthropic/__init__.py | 14 + .../api/anthropic/test_messages_behavior.py | 630 ++++++++++++++++++ 4 files changed, 661 insertions(+), 2 deletions(-) create mode 100644 tensorrt_llm/serve/api/anthropic/__init__.py create mode 100644 tests/unittest/serve/api/anthropic/__init__.py create mode 100644 tests/unittest/serve/api/anthropic/test_messages_behavior.py diff --git a/tensorrt_llm/serve/api/anthropic/__init__.py b/tensorrt_llm/serve/api/anthropic/__init__.py new file mode 100644 index 000000000000..035527480c4f --- /dev/null +++ b/tensorrt_llm/serve/api/anthropic/__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. +"""Anthropic Messages wire dialect: protocol, adapter and batches. Import submodules directly.""" diff --git a/tests/unittest/llmapi/apps/test_anthropic_adapter.py b/tests/unittest/llmapi/apps/test_anthropic_adapter.py index 5cfe060fe8fd..7c53b736ca20 100644 --- a/tests/unittest/llmapi/apps/test_anthropic_adapter.py +++ b/tests/unittest/llmapi/apps/test_anthropic_adapter.py @@ -26,7 +26,7 @@ import pytest from pydantic import ValidationError -import tensorrt_llm.serve.anthropic_adapter as adapter +import tensorrt_llm.serve.api.anthropic.adapter as adapter from tensorrt_llm.serve.api.anthropic.adapter import ( AnthropicRequestError, AnthropicResponseError, @@ -1858,7 +1858,7 @@ def test_cached_tokens_exceeding_prompt_tokens_is_logged(monkeypatch): tensorrt_llm.logger does not propagate to the stdlib root, so caplog sees nothing and the test would pass without the warning existing. """ - import tensorrt_llm.serve.anthropic_adapter as adapter + import tensorrt_llm.serve.api.anthropic.adapter as adapter warnings = [] monkeypatch.setattr(adapter.logger, "warning", warnings.append) diff --git a/tests/unittest/serve/api/anthropic/__init__.py b/tests/unittest/serve/api/anthropic/__init__.py new file mode 100644 index 000000000000..467079831e16 --- /dev/null +++ b/tests/unittest/serve/api/anthropic/__init__.py @@ -0,0 +1,14 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/tests/unittest/serve/api/anthropic/test_messages_behavior.py b/tests/unittest/serve/api/anthropic/test_messages_behavior.py new file mode 100644 index 000000000000..ce384da56005 --- /dev/null +++ b/tests/unittest/serve/api/anthropic/test_messages_behavior.py @@ -0,0 +1,630 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""What the aggregated and the disaggregated server answer on the Anthropic routes, pinned. + +Both servers serve ``POST /v1/messages`` by translating the request into a chat +completion, running it through their own chat pipeline and translating the +answer back; the aggregated server also runs Message Batches over the same +path. The two pipelines answer some cases differently today (the model name of +a stream without chunks, the message of an unexpected error) and are meant to +keep doing so. These tests pin each server's own answer, case by case, so that +a change to how the Anthropic routes reach the chat pipelines can show it kept +every one of them. + +Only HTTP and the engine seams are used. The aggregated chat path is stubbed +where it meets the engine (``generate_async`` and the response builder +``_create_chat_response``); the disaggregated one where it meets the service +(``openai_chat_completion``). The served model is ``server-model`` and every +request names ``client-model``, so an answer that echoes the wrong one is +caught. +""" + +import json +import logging +import signal +import time +from types import SimpleNamespace +from unittest.mock import AsyncMock, Mock + +import aiohttp +import pytest +from test_common.serve.app_factory import AggConfig, DisaggConfig + +import tensorrt_llm.serve.openai_disagg_server as disagg_server_module +import tensorrt_llm.serve.openai_server as openai_server_module +from tensorrt_llm.executor import CppExecutorError +from tensorrt_llm.executor.utils import context_length_exceeded_message +from tensorrt_llm.inputs.chat_template_guard import UnusedChatTemplateKwargsError +from tensorrt_llm.logger import logger as trtllm_logger +from tensorrt_llm.serve.openai_protocol import ( + ChatCompletionResponse, + ChatCompletionResponseChoice, + ChatCompletionStreamResponse, + ChatMessage, + FunctionCall, + ToolCall, + UsageInfo, +) +from tensorrt_llm.serve.openai_server import OpenAIServer + +# threadleak is off: the TestClient's portal thread and the fake server's +# worker pools are still alive when the leak check runs; the serve_env fixture +# closes them in its teardown, which runs after that check. +pytestmark = [pytest.mark.cpu_only, pytest.mark.threadleak(enabled=False)] + +SIDES = ["aggregated", "disaggregated"] +SERVER_MODEL = "server-model" +CLIENT_MODEL = "client-model" +# Two chunks of streamed text. The accent gives the disaggregated byte stream a +# multi-byte character to be cut inside. +STREAMED_CHUNKS = ("hél", "lo") +STREAMED_TEXT = "".join(STREAMED_CHUNKS) +COUNTER_NAMES = ( + "total_requests", + "stream_requests", + "nonstream_requests", + "validation_exceptions", + "http_exceptions", + "internal_errors", + "total_responses", +) + + +def _changed(case: str, side: str, expected, got) -> str: + return ( + f"Anthropic {case} on the {side} server changed: expected {expected!r}, got {got!r}. " + "Each server must keep its own answer; the aggregated and disaggregated servers " + "differ on purpose." + ) + + +# -- requests and engine answers ----------------------------------------------- + + +def _request(text: str = "hello", **overrides) -> dict: + payload = { + "model": CLIENT_MODEL, + "max_tokens": 32, + "messages": [{"role": "user", "content": text}], + } + payload.update(overrides) + return payload + + +def _chat_response(*, tool_arguments=None) -> ChatCompletionResponse: + """The chat pipeline's non-streaming answer: the text "hello", optionally with a tool call.""" + tool_calls = [] + if tool_arguments is not None: + tool_calls = [ + ToolCall( + id="call_1", + function=FunctionCall(name="get_weather", arguments=tool_arguments), + ) + ] + return ChatCompletionResponse( + id="chatcmpl-messages-behavior", + model=SERVER_MODEL, + choices=[ + ChatCompletionResponseChoice( + index=0, + message=ChatMessage(role="assistant", content="hello", tool_calls=tool_calls), + finish_reason="tool_calls" if tool_calls else "stop", + ) + ], + usage=UsageInfo(prompt_tokens=3, completion_tokens=2, total_tokens=5), + ) + + +def _stream_frames(*texts: str) -> list: + """The chat pipeline's SSE frames for a stream of text chunks, without the final [DONE].""" + chunks = [ + ChatCompletionStreamResponse( + model=SERVER_MODEL, choices=[{"index": 0, "delta": {"content": text}}] + ) + for text in texts + ] + chunks.append( + ChatCompletionStreamResponse( + model=SERVER_MODEL, + choices=[{"index": 0, "delta": {}, "finish_reason": "stop"}], + usage=UsageInfo(prompt_tokens=3, completion_tokens=2, total_tokens=5), + ) + ) + return [f"data: {chunk.model_dump_json()}\n\n" for chunk in chunks] + + +def _promise() -> SimpleNamespace: + """The engine's handle for a non-streaming request; the response builder is stubbed.""" + return SimpleNamespace(prompt_token_ids=[1, 2, 3]) + + +class _StreamedGeneration: + """A streamed generation whose steps are already post-processed SSE frames.""" + + prompt_token_ids = [1, 2, 3] + + def __init__(self, steps): + self._steps = list(steps) + + def __aiter__(self): + return self + + async def __anext__(self): + if not self._steps: + raise StopAsyncIteration + frames = self._steps.pop(0) + return SimpleNamespace( + finished=not self._steps, + outputs=[SimpleNamespace(_postprocess_result=frames)], + ) + + +def _relayed_bytes(frames: list): + """The worker's SSE bytes as the disaggregated service relays them. + + Delivered in two reads, cut inside the multi-byte character of the streamed + text, as network reads need not align with characters or lines. + """ + payload = "".join(frames).encode("utf-8") + cut = payload.index("é".encode("utf-8")) + 1 + + async def relay(): + yield payload[:cut] + yield payload[cut:] + + return relay() + + +def _done_only(): + """A relayed stream finished by the context worker alone: only the [DONE] sentinel.""" + + async def relay(): + yield b"data: [DONE]\n\n" + + return relay() + + +def _upstream_error(worker_body: str, status: int) -> aiohttp.ClientResponseError: + """A worker rejection the way the disaggregated client raises it: reason, then the body.""" + return aiohttp.ClientResponseError( + request_info=Mock(), history=(), status=status, message=f"Bad Request: {worker_body}" + ) + + +# -- servers ------------------------------------------------------------------- + + +def _aggregated_server(env, *, generate_async, chat_response=None) -> OpenAIServer: + """The aggregated server with its chat path stubbed where it meets the engine. + + Post-processing runs in worker mode, so a streamed generation's steps are + the SSE frames themselves. The engine's preprocessing and the disconnect + watcher are not part of any answer and are left out. + """ + server = env.build_openai_server(AggConfig(model_name=SERVER_MODEL, num_postprocess_workers=1)) + server.generator.preprocess = None + server.generator.generate_async = generate_async + server.await_disconnected = AsyncMock() + if chat_response is not None: + server._create_chat_response = AsyncMock(return_value=chat_response) + return server + + +def _disaggregated_server(env, openai_chat_completion): + """The disaggregated server with the service stubbed and its request counters mocked.""" + server = env.build_disagg_server(DisaggConfig()) + counters = SimpleNamespace( + **{name: SimpleNamespace(inc=Mock()) for name in COUNTER_NAMES}, + queue_latency_seconds=SimpleNamespace(observe=Mock()), + ) + server._perf_metrics_collector = counters + server._service = SimpleNamespace(openai_chat_completion=openai_chat_completion) + return server, counters + + +def _post_messages(env, side: str, server, payload: dict): + # The disaggregated lifespan connects to workers; these tests never run it. + client = env.client(server.app, lifespan=side == "aggregated") + return client.post("/v1/messages", json=payload) + + +def _assert_error(response, *, case: str, side: str, status: int, error_type: str, message: str): + body = response.json() + got = (response.status_code, body.get("type"), body.get("error")) + expected = (status, "error", {"type": error_type, "message": message}) + assert got == expected, _changed(case, side, expected, got) + + +def _assert_counts(counters, case: str, **expected: int) -> None: + for name in COUNTER_NAMES: + got = getattr(counters, name).inc.call_count + assert got == expected.get(name, 0), _changed( + f"{case} ({name} counter)", "disaggregated", expected.get(name, 0), got + ) + + +def _events(body: str) -> list: + """The (event name, decoded data) pairs of an Anthropic SSE body.""" + events = [] + for frame in body.strip().split("\n\n"): + name, data = None, None + for line in frame.splitlines(): + if line.startswith("event: "): + name = line[len("event: ") :] + elif line.startswith("data: "): + data = json.loads(line[len("data: ") :]) + events.append((name, data)) + return events + + +def _streamed_text(events: list) -> str: + return "".join( + data["delta"]["text"] + for name, data in events + if name == "content_block_delta" and data["delta"]["type"] == "text_delta" + ) + + +@pytest.fixture +def error_records(): + """The ERROR records the TensorRT-LLM logger emits during the test.""" + records = [] + + class Capture(logging.Handler): + def emit(self, record): + records.append(record) + + handler = Capture(level=logging.ERROR) + trtllm_logger._logger.addHandler(handler) + try: + yield records + finally: + trtllm_logger._logger.removeHandler(handler) + + +@pytest.fixture +def engine_failure(monkeypatch): + """What an engine failure does to the process, recorded instead of done. + + Both servers answer an engine failure by recording the termination and + raising SIGINT; the recorder is patched in the module whose code raises it. + """ + recorded = SimpleNamespace(signals=[], aggregated=Mock(), disaggregated=Mock()) + monkeypatch.setattr(signal, "raise_signal", recorded.signals.append) + monkeypatch.setattr(openai_server_module, "_record_generator_termination", recorded.aggregated) + monkeypatch.setattr( + disagg_server_module, "record_termination_observation", recorded.disaggregated + ) + return recorded + + +# -- /v1/messages -------------------------------------------------------------- + + +@pytest.mark.parametrize("side", SIDES) +def test_non_streaming_success_is_the_translated_chat_response(side, serve_env): + if side == "aggregated": + server = _aggregated_server( + serve_env, generate_async=Mock(return_value=_promise()), chat_response=_chat_response() + ) + counters = None + else: + server, counters = _disaggregated_server( + serve_env, AsyncMock(return_value=_chat_response()) + ) + + response = _post_messages(serve_env, side, server, _request()) + + body = response.json() + got = (response.status_code, body.get("content"), body.get("model")) + expected = (200, [{"type": "text", "text": "hello"}], SERVER_MODEL) + assert got == expected, _changed("non-streaming success", side, expected, got) + if counters is not None: + _assert_counts( + counters, + "non-streaming success", + total_requests=1, + nonstream_requests=1, + total_responses=1, + ) + + +@pytest.mark.parametrize("side", SIDES) +def test_streaming_success_reports_the_chunks_model_and_text(side, serve_env): + frames = _stream_frames(*STREAMED_CHUNKS) + if side == "aggregated": + generation = _StreamedGeneration([[frame] for frame in frames]) + server = _aggregated_server(serve_env, generate_async=Mock(return_value=generation)) + else: + relayed = _relayed_bytes(frames + ["data: [DONE]\n\n"]) + server, _ = _disaggregated_server(serve_env, AsyncMock(return_value=relayed)) + + response = _post_messages(serve_env, side, server, _request(stream=True)) + + assert response.status_code == 200, _changed( + "streaming success", side, 200, (response.status_code, response.text) + ) + assert response.headers["content-type"].startswith("text/event-stream") + events = _events(response.text) + got = (events[0][0], events[0][1]["message"]["model"], _streamed_text(events), events[-1][0]) + expected = ("message_start", SERVER_MODEL, STREAMED_TEXT, "message_stop") + assert got == expected, _changed("streaming success", side, expected, got) + + +@pytest.mark.parametrize( + "side, expected_model", + [("aggregated", SERVER_MODEL), ("disaggregated", CLIENT_MODEL)], +) +def test_stream_without_chunks_names_each_servers_own_model(side, expected_model, serve_env): + """With no chunk to take the model from, each server falls back to its own choice.""" + if side == "aggregated": + generation = _StreamedGeneration([[]]) + server = _aggregated_server(serve_env, generate_async=Mock(return_value=generation)) + else: + server, _ = _disaggregated_server(serve_env, AsyncMock(return_value=_done_only())) + + response = _post_messages(serve_env, side, server, _request(stream=True)) + + assert response.status_code == 200, _changed( + "stream without chunks", side, 200, (response.status_code, response.text) + ) + events = _events(response.text) + got = (events[0][0], events[0][1]["message"]["model"], events[-1][0]) + expected = ("message_start", expected_model, "message_stop") + assert got == expected, _changed("stream without chunks", side, expected, got) + + +@pytest.mark.parametrize("side", SIDES) +def test_rejected_request_is_a_400_with_the_rejection(side, serve_env): + if side == "aggregated": + server = _aggregated_server(serve_env, generate_async=Mock(side_effect=ValueError("bad"))) + else: + server, _ = _disaggregated_server( + serve_env, AsyncMock(side_effect=UnusedChatTemplateKwargsError("bad")) + ) + + response = _post_messages(serve_env, side, server, _request()) + + _assert_error( + response, + case="rejected request", + side=side, + status=400, + error_type="invalid_request_error", + message="bad", + ) + + +@pytest.mark.parametrize("side", SIDES) +def test_context_length_error_keeps_its_message(side, serve_env): + message = context_length_exceeded_message(2048, 4096) + if side == "aggregated": + server = _aggregated_server(serve_env, generate_async=Mock(side_effect=ValueError(message))) + else: + worker_body = OpenAIServer.create_error_response(message).body.decode() + server, _ = _disaggregated_server( + serve_env, AsyncMock(side_effect=_upstream_error(worker_body, status=400)) + ) + + response = _post_messages(serve_env, side, server, _request()) + + _assert_error( + response, + case="context length", + side=side, + status=400, + error_type="invalid_request_error", + message=message, + ) + + +def test_upstream_error_keeps_the_worker_status_and_message(serve_env): + error = aiohttp.ClientResponseError( + request_info=Mock(), + history=(), + status=429, + message='Too Many Requests: {"error": {"message": "slow down"}}', + ) + server, counters = _disaggregated_server(serve_env, AsyncMock(side_effect=error)) + + response = _post_messages(serve_env, "disaggregated", server, _request()) + + _assert_error( + response, + case="upstream error", + side="disaggregated", + status=429, + error_type="invalid_request_error", + message="slow down", + ) + _assert_counts( + counters, "upstream error", total_requests=1, nonstream_requests=1, http_exceptions=1 + ) + + +@pytest.mark.parametrize( + "side, status, error_type, message", + [ + ("aggregated", 400, "invalid_request_error", "boom"), + ("disaggregated", 500, "api_error", "Internal server error boom"), + ], +) +def test_unexpected_error_is_answered_each_servers_way( + side, status, error_type, message, serve_env +): + if side == "aggregated": + server = _aggregated_server( + serve_env, generate_async=Mock(side_effect=RuntimeError("boom")) + ) + counters = None + else: + server, counters = _disaggregated_server( + serve_env, AsyncMock(side_effect=RuntimeError("boom")) + ) + + response = _post_messages(serve_env, side, server, _request()) + + _assert_error( + response, + case="unexpected error", + side=side, + status=status, + error_type=error_type, + message=message, + ) + if counters is not None: + _assert_counts( + counters, "unexpected error", total_requests=1, nonstream_requests=1, internal_errors=1 + ) + + +@pytest.mark.parametrize("side", SIDES) +def test_engine_failure_is_a_500_after_one_sigint(side, serve_env, engine_failure): + if side == "aggregated": + server = _aggregated_server( + serve_env, generate_async=Mock(side_effect=CppExecutorError("dead")) + ) + recorder = engine_failure.aggregated + else: + server, _ = _disaggregated_server( + serve_env, AsyncMock(side_effect=CppExecutorError("dead")) + ) + recorder = engine_failure.disaggregated + + response = _post_messages(serve_env, side, server, _request()) + + _assert_error( + response, + case="engine failure", + side=side, + status=500, + error_type="api_error", + message="Internal server error", + ) + got = (engine_failure.signals, recorder.call_count) + expected = ([signal.SIGINT], 1) + assert got == expected, _changed( + "engine failure (SIGINT, termination record)", side, expected, got + ) + + +@pytest.mark.parametrize("side", SIDES) +def test_unconvertible_chat_response_is_a_500_with_one_error_log(side, serve_env, error_records): + response_with_bad_tool_call = _chat_response(tool_arguments="{not-json") + if side == "aggregated": + server = _aggregated_server( + serve_env, + generate_async=Mock(return_value=_promise()), + chat_response=response_with_bad_tool_call, + ) + else: + server, _ = _disaggregated_server( + serve_env, AsyncMock(return_value=response_with_bad_tool_call) + ) + + response = _post_messages(serve_env, side, server, _request()) + + _assert_error( + response, + case="unconvertible response", + side=side, + status=500, + error_type="api_error", + message="Internal server error", + ) + messages = [record.getMessage() for record in error_records] + assert len(messages) == 1 and "Invalid response from OpenAI chat pipeline" in messages[0], ( + _changed( + "unconvertible response (error log)", + side, + "one record naming the invalid response", + messages, + ) + ) + + +# -- /v1/messages/batches (aggregated only) -------------------------------------- + + +def test_batch_items_report_each_outcome(serve_env, engine_failure): + def generate_async(inputs, **kwargs): + prompt = inputs.get("prompt") or "" + if "reject me" in prompt: + raise ValueError("bad") + if "kill the engine" in prompt: + raise CppExecutorError("dead") + return _promise() + + server = _aggregated_server( + serve_env, generate_async=generate_async, chat_response=_chat_response() + ) + client = serve_env.client(server.app) + items = [ + {"custom_id": custom_id, "params": _request(text)} + for custom_id, text in ( + ("ok", "hello"), + ("rejected", "reject me"), + ("engine", "kill the engine"), + ) + ] + + created = client.post("/v1/messages/batches", json={"requests": items}) + assert created.status_code == 200, created.text + batch_id = created.json()["id"] + deadline = time.monotonic() + 30 + while True: + status = client.get(f"/v1/messages/batches/{batch_id}").json()["processing_status"] + if status == "ended": + break + assert time.monotonic() < deadline, f"batch {batch_id} is still {status} after 30 s" + time.sleep(0.01) + results = client.get(f"/v1/messages/batches/{batch_id}/results") + + assert results.status_code == 200, results.text + by_id = { + line["custom_id"]: line["result"] for line in map(json.loads, results.text.splitlines()) + } + got = { + "ok": ( + by_id["ok"]["type"], + by_id["ok"]["message"]["content"], + by_id["ok"]["message"]["model"], + ), + "rejected": by_id["rejected"], + "engine": by_id["engine"], + } + expected = { + "ok": ("succeeded", [{"type": "text", "text": "hello"}], SERVER_MODEL), + "rejected": { + "type": "errored", + "error": { + "type": "error", + "error": {"type": "invalid_request_error", "message": "bad"}, + "request_id": None, + }, + }, + "engine": { + "type": "errored", + "error": { + "type": "error", + "error": {"type": "api_error", "message": "Internal server error"}, + "request_id": None, + }, + }, + } + assert got == expected, _changed("batch item outcomes", "aggregated", expected, got) + assert engine_failure.signals == [signal.SIGINT], _changed( + "batch engine failure (SIGINT)", "aggregated", [signal.SIGINT], engine_failure.signals + ) From 31fa80bfb0cc42f5e3d48b4977b07577b54645d0 Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Tue, 3 Nov 2026 06:30:14 +0000 Subject: [PATCH 23/27] [TRTLLM-17263][refactor] Add the Anthropic count-tokens and batches serving classes The two classes receive the handlers that leave OpenAIServer in the next commit. Their constructors take what those handlers read from the server (the render-resources factory, the batch item runner), so missing wiring fails at construction instead of falling back to a default. The module is registered as a legacy-formatted file before any moved code lands in it, so the formatters keep the moved methods as they are; ruff-legacy.toml and the managed blocks of pyproject.toml and .pre-commit-config.yaml are regenerated from legacy-files.txt. 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/api/anthropic/serving.py | 44 +++++++++++++++++++++ 5 files changed, 49 insertions(+) create mode 100644 tensorrt_llm/serve/api/anthropic/serving.py diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index aeabd348c165..a5a96a2ea8a3 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -419,6 +419,7 @@ common-files: &common_files | tensorrt_llm/serve/__init__.py | tensorrt_llm/serve/_base/msgspec_route.py | tensorrt_llm/serve/_base/sse.py | + tensorrt_llm/serve/api/anthropic/serving.py | tensorrt_llm/serve/chat_utils.py | tensorrt_llm/serve/cluster_storage.py | tensorrt_llm/serve/disagg_auto_scaling.py | @@ -1148,6 +1149,7 @@ legacy-files: &legacy_files | tensorrt_llm/serve/__init__.py | tensorrt_llm/serve/_base/msgspec_route.py | tensorrt_llm/serve/_base/sse.py | + tensorrt_llm/serve/api/anthropic/serving.py | tensorrt_llm/serve/chat_utils.py | tensorrt_llm/serve/cluster_storage.py | tensorrt_llm/serve/disagg_auto_scaling.py | diff --git a/legacy-files.txt b/legacy-files.txt index c1903c47e3f8..3dd21ca755bf 100644 --- a/legacy-files.txt +++ b/legacy-files.txt @@ -411,6 +411,7 @@ tensorrt_llm/scheduling_params.py tensorrt_llm/serve/__init__.py tensorrt_llm/serve/_base/msgspec_route.py tensorrt_llm/serve/_base/sse.py +tensorrt_llm/serve/api/anthropic/serving.py tensorrt_llm/serve/chat_utils.py tensorrt_llm/serve/cluster_storage.py tensorrt_llm/serve/disagg_auto_scaling.py diff --git a/pyproject.toml b/pyproject.toml index 86a343c97bc3..2b77d93e161b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -476,6 +476,7 @@ exclude = [ "tensorrt_llm/serve/__init__.py", "tensorrt_llm/serve/_base/msgspec_route.py", "tensorrt_llm/serve/_base/sse.py", + "tensorrt_llm/serve/api/anthropic/serving.py", "tensorrt_llm/serve/chat_utils.py", "tensorrt_llm/serve/cluster_storage.py", "tensorrt_llm/serve/disagg_auto_scaling.py", diff --git a/ruff-legacy.toml b/ruff-legacy.toml index 4ccd76cfbd7c..a091e2e35088 100644 --- a/ruff-legacy.toml +++ b/ruff-legacy.toml @@ -428,6 +428,7 @@ include = [ "tensorrt_llm/serve/__init__.py", "tensorrt_llm/serve/_base/msgspec_route.py", "tensorrt_llm/serve/_base/sse.py", + "tensorrt_llm/serve/api/anthropic/serving.py", "tensorrt_llm/serve/chat_utils.py", "tensorrt_llm/serve/cluster_storage.py", "tensorrt_llm/serve/disagg_auto_scaling.py", diff --git a/tensorrt_llm/serve/api/anthropic/serving.py b/tensorrt_llm/serve/api/anthropic/serving.py new file mode 100644 index 000000000000..8142a90efd1b --- /dev/null +++ b/tensorrt_llm/serve/api/anthropic/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. +"""Serving classes of the Anthropic Messages dialect. + +Messages, count_tokens and the Message Batches, served from values the +composition root injects; no class here reads a server attribute. +""" + +from typing import Callable, Optional + +from tensorrt_llm.serve.api.anthropic.batches import (AnthropicBatchStore, + BatchItemRunner) +from tensorrt_llm.serve.render import RenderResources + + +class AnthropicCountTokensServing: + """Counts the input tokens of a Messages request over the rendered prompt.""" + + def __init__(self, *, render_resources: Callable[[], + RenderResources]) -> None: + # Called per request: the resources are read at call time, as the + # chat route reads them. + self._render_resources = render_resources + + +class AnthropicBatchesServing: + """The Message Batches routes of the aggregated server.""" + + def __init__(self, *, runner: BatchItemRunner) -> None: + # The batch methods create the store on first use, with this runner. + self._run_batched_message = runner + self._anthropic_batch_store: Optional[AnthropicBatchStore] = None From 95bb1e1c80b3b74ca798e63ed2ef1f1e7a3ab0b3 Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Tue, 3 Nov 2026 06:37:26 +0000 Subject: [PATCH 24/27] [TRTLLM-17263][refactor] Move the Anthropic count-tokens and batch handlers out of OpenAIServer Generated move: anthropic_count_tokens goes to AnthropicCountTokensServing and _batch_store, the six Message Batches handlers and _no_such_batch go to AnthropicBatchesServing in tensorrt_llm/serve/api/anthropic/serving.py, verbatim except for the declared rewrite in count_tokens (RenderResources.from_server(self) becomes self._render_resources(), the injected factory). anthropic_messages, _run_batched_message and _synthetic_request are deleted here and replaced by the serving class and the request context added in the next commit. The import cleanup is the formatters' (autoflake, isort, yapf). Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- tensorrt_llm/serve/api/anthropic/serving.py | 170 ++++++++++- tensorrt_llm/serve/openai_server.py | 317 +------------------- 2 files changed, 173 insertions(+), 314 deletions(-) diff --git a/tensorrt_llm/serve/api/anthropic/serving.py b/tensorrt_llm/serve/api/anthropic/serving.py index 8142a90efd1b..def94f0fb542 100644 --- a/tensorrt_llm/serve/api/anthropic/serving.py +++ b/tensorrt_llm/serve/api/anthropic/serving.py @@ -18,11 +18,29 @@ composition root injects; no class here reads a server attribute. """ +import asyncio +import traceback from typing import Callable, Optional +from fastapi import Query +from fastapi.responses import JSONResponse, Response +from pydantic import ValidationError + +from tensorrt_llm.logger import logger +from tensorrt_llm.serve.api.anthropic.adapter import ( + AnthropicRequestError, anthropic_error_response, + convert_anthropic_count_tokens_request) from tensorrt_llm.serve.api.anthropic.batches import (AnthropicBatchStore, - BatchItemRunner) -from tensorrt_llm.serve.render import RenderResources + BatchItemRunner, + BatchStoreFullError, + results_to_jsonl) +from tensorrt_llm.serve.api.anthropic.protocol import ( + AnthropicBatchDeleteResponse, AnthropicBatchList, + AnthropicCountTokensRequest, AnthropicCountTokensResponse, + AnthropicCreateBatchRequest) +from tensorrt_llm.serve.render import (HarmonyConversionError, RenderResources, + UnsupportedRenderError, render_chat, + render_chat_for_length) class AnthropicCountTokensServing: @@ -34,6 +52,56 @@ def __init__(self, *, render_resources: Callable[[], # chat route reads them. self._render_resources = render_resources + async def anthropic_count_tokens( + self, request: AnthropicCountTokensRequest) -> Response: + """Report how many input tokens a Messages request would consume. + + The count is taken over the rendered prompt rather than over the raw + messages, so system blocks, tool definitions and the thinking prefix + are all included - those are exactly what make a client's own estimate + drift from the server's. + """ + try: + chat_request = convert_anthropic_count_tokens_request(request) + except AnthropicRequestError as e: + return anthropic_error_response(str(e), "invalid_request_error", + 400) + except ValidationError as e: + return anthropic_error_response(str(e), "invalid_request_error", + 400) + + try: + + # Both the render and the encode are synchronous and both scale + # with prompt size. Claude Code calls this before most turns with + # the whole conversation attached, so running them inline would + # block the event loop -- and therefore every other in-flight + # request on this server -- for a full template render plus a + # tokenizer pass. The chat path avoids this the same way. + def _count() -> int: + # The count is over the prompt tokens the chat route would + # execute: same pipeline, same template, same tokenization. + resources = self._render_resources() + try: + rendered = render_chat(chat_request, + resources, + tokenize=True) + except (UnsupportedRenderError, HarmonyConversionError): + # Media, an own-processor model or an unconvertible Harmony + # request: count the text, tools and template without the media. + rendered = render_chat_for_length(chat_request, resources) + return len(rendered.token_ids) + + input_tokens = await asyncio.to_thread(_count) + except Exception: + logger.error("count_tokens failed to render the prompt:\n" + f"{traceback.format_exc()}") + return anthropic_error_response("Internal server error", + "api_error", 500) + + return JSONResponse(content=AnthropicCountTokensResponse( + input_tokens=input_tokens).model_dump()) + class AnthropicBatchesServing: """The Message Batches routes of the aggregated server.""" @@ -42,3 +110,101 @@ def __init__(self, *, runner: BatchItemRunner) -> None: # The batch methods create the store on first use, with this runner. self._run_batched_message = runner self._anthropic_batch_store: Optional[AnthropicBatchStore] = None + + def _batch_store(self) -> AnthropicBatchStore: + """Created on first use so servers that never batch pay nothing.""" + if getattr(self, "_anthropic_batch_store", None) is None: + self._anthropic_batch_store = AnthropicBatchStore( + runner=self._run_batched_message) + return self._anthropic_batch_store + + async def anthropic_create_batch( + self, request: AnthropicCreateBatchRequest) -> Response: + try: + batch = self._batch_store().create(request.requests) + except BatchStoreFullError as e: + # 429, not 400: the request is valid and an identical one will + # succeed later, so the client should back off rather than edit it. + return anthropic_error_response(str(e), "rate_limit_error", 429) + except ValueError as e: + return anthropic_error_response(str(e), "invalid_request_error", + 400) + return JSONResponse(content=batch.model_dump()) + + async def anthropic_list_batches( + self, + # Validated here rather than clamped in the store: clamping turned + # limit=0 into 1 instead of a 400, and bounded the page against the + # retention limit, which is an unrelated quantity. + limit: int = Query(20, ge=1, le=100), + after_id: Optional[str] = None, + before_id: Optional[str] = None) -> Response: + try: + batches, has_more = self._batch_store().list(limit=limit, + after_id=after_id, + before_id=before_id) + except LookupError as e: + # Saying so beats an empty page: an empty page is indistinguishable + # from a finished walk, so the client stops early having silently + # skipped every batch past the cursor. + return anthropic_error_response( + f"Unknown pagination cursor {e.args[0]!r}. The batch may have " + "been deleted or expired; restart the listing without a cursor.", + "invalid_request_error", 400) + listing = AnthropicBatchList( + data=batches, + # Reported honestly rather than hardcoded false: a client that + # asked for N and received N cannot otherwise tell a full page from + # the end of the list, and silently stops early. It pages on with + # ?after_id=. + has_more=has_more, + first_id=batches[0].id if batches else None, + last_id=batches[-1].id if batches else None, + ) + return JSONResponse(content=listing.model_dump()) + + async def anthropic_get_batch(self, batch_id: str) -> Response: + batch = self._batch_store().get(batch_id) + if batch is None: + return self._no_such_batch(batch_id) + return JSONResponse(content=batch.model_dump()) + + async def anthropic_cancel_batch(self, batch_id: str) -> Response: + batch = self._batch_store().cancel(batch_id) + if batch is None: + return self._no_such_batch(batch_id) + return JSONResponse(content=batch.model_dump()) + + async def anthropic_delete_batch(self, batch_id: str) -> Response: + try: + deleted = self._batch_store().delete(batch_id) + except ValueError as e: + return anthropic_error_response(str(e), "invalid_request_error", + 400) + if not deleted: + return self._no_such_batch(batch_id) + return JSONResponse(content=AnthropicBatchDeleteResponse( + id=batch_id).model_dump()) + + async def anthropic_batch_results(self, batch_id: str) -> Response: + store = self._batch_store() + results = store.results(batch_id) + if results is None: + if not store.has(batch_id): + return self._no_such_batch(batch_id) + # Known batch, but not finished. 404 would say "no such batch", + # which is wrong and would stop a client from polling. + return anthropic_error_response( + f"Message Batch {batch_id} is still processing; results are " + "available once processing_status is 'ended'", + "invalid_request_error", 400) + return Response(content=results_to_jsonl(results), + media_type="application/x-ndjson") + + @staticmethod + def _no_such_batch(batch_id: str) -> Response: + # Batches are held in memory only, so an id from before a restart lands + # here. Say so, rather than leaving the caller to wonder. + return anthropic_error_response( + f"Message Batch {batch_id} not found. Batches are stored in memory " + "and do not survive a server restart.", "not_found_error", 404) diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py index 58bf07392f86..03ab10ef02fb 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -23,10 +23,10 @@ from http import HTTPStatus from pathlib import Path from typing import (TYPE_CHECKING, Annotated, Any, AsyncIterator, Dict, List, - Optional, Tuple, Union) + Optional, Union) import uvicorn -from fastapi import Body, Depends, FastAPI, HTTPException, Query, Request +from fastapi import Body, Depends, FastAPI, HTTPException, Request from fastapi.exceptions import RequestValidationError from fastapi.responses import (FileResponse, JSONResponse, Response, StreamingResponse) @@ -70,17 +70,7 @@ 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.api.anthropic.adapter import ( - AnthropicRequestError, AnthropicResponseError, anthropic_error_response, - convert_anthropic_count_tokens_request, convert_anthropic_request, - convert_chat_response, reframe_openai_stream) -from tensorrt_llm.serve.api.anthropic.batches import (AnthropicBatchStore, - BatchStoreFullError, - results_to_jsonl) -from tensorrt_llm.serve.api.anthropic.protocol import ( - AnthropicBatchDeleteResponse, AnthropicBatchList, - AnthropicCountTokensRequest, AnthropicCountTokensResponse, - AnthropicCreateBatchRequest, AnthropicMessagesRequest) +from tensorrt_llm.serve.api.anthropic.adapter import anthropic_error_response from tensorrt_llm.serve.api.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, @@ -123,11 +113,8 @@ chat_stream_post_processor, completion_response_post_processor, completion_stream_post_processor, responses_api_post_processor, responses_api_streaming_post_processor) -from tensorrt_llm.serve.render import (HarmonyConversionError, RenderResources, - UnsupportedRenderError, - mount_render_endpoints, - prepare_chat_request, render_chat, - render_chat_for_length, +from tensorrt_llm.serve.render import (RenderResources, mount_render_endpoints, + prepare_chat_request, render_conversation) from tensorrt_llm.serve.responses_utils import (ClosingStreamingResponse, ConversationHistoryStore, @@ -2484,302 +2471,8 @@ def get_role() -> str: logger.error(traceback.format_exc()) raise serving_error(str(e)) from e - async def anthropic_messages(self, request: AnthropicMessagesRequest, - raw_request: Request) -> Response: - """Serve the Anthropic Messages API on top of the openai_chat path. - - The Anthropic request is translated into a ChatCompletionRequest, - handed to ``openai_chat`` unchanged (chat template, tool parser and - post-processing are reused verbatim), and the OpenAI-shaped result is - translated back: JSON body for non-streaming, SSE reframing for - streaming. - """ - try: - chat_request = convert_anthropic_request(request) - except AnthropicRequestError as e: - return anthropic_error_response(str(e), "invalid_request_error", - 400) - except ValidationError as e: - return anthropic_error_response(str(e), "invalid_request_error", - 400) - - response = await self.openai_chat(chat_request, raw_request) - if response is None: - return anthropic_error_response("Internal server error", - "api_error", 500) - - if isinstance(response, StreamingResponse): - return StreamingResponse( - content=reframe_openai_stream(response.body_iterator, - model=self.model), - media_type="text/event-stream", - ) - - status = getattr(response, "status_code", 500) - if status != 200: - try: - payload = json.loads(response.body) - message = payload.get("message") or json.dumps(payload) - except (json.JSONDecodeError, AttributeError): - message = "Internal server error" - err_type = ("invalid_request_error" - if 400 <= status < 500 else "api_error") - return anthropic_error_response(message, err_type, status) - - try: - chat_response = ChatCompletionResponse(**json.loads(response.body)) - anthropic_response = convert_chat_response(chat_response) - except (AnthropicResponseError, ValidationError, json.JSONDecodeError): - logger.error("Invalid response from OpenAI chat pipeline:\n" - f"{traceback.format_exc()}") - return anthropic_error_response("Internal server error", - "api_error", 500) - return JSONResponse(content=anthropic_response.model_dump( - exclude_none=True)) - - async def anthropic_count_tokens( - self, request: AnthropicCountTokensRequest) -> Response: - """Report how many input tokens a Messages request would consume. - - The count is taken over the rendered prompt rather than over the raw - messages, so system blocks, tool definitions and the thinking prefix - are all included - those are exactly what make a client's own estimate - drift from the server's. - """ - try: - chat_request = convert_anthropic_count_tokens_request(request) - except AnthropicRequestError as e: - return anthropic_error_response(str(e), "invalid_request_error", - 400) - except ValidationError as e: - return anthropic_error_response(str(e), "invalid_request_error", - 400) - - try: - - # Both the render and the encode are synchronous and both scale - # with prompt size. Claude Code calls this before most turns with - # the whole conversation attached, so running them inline would - # block the event loop -- and therefore every other in-flight - # request on this server -- for a full template render plus a - # tokenizer pass. The chat path avoids this the same way. - def _count() -> int: - # The count is over the prompt tokens the chat route would - # execute: same pipeline, same template, same tokenization. - resources = RenderResources.from_server(self) - try: - rendered = render_chat(chat_request, - resources, - tokenize=True) - except (UnsupportedRenderError, HarmonyConversionError): - # Media, an own-processor model or an unconvertible Harmony - # request: count the text, tools and template without the media. - rendered = render_chat_for_length(chat_request, resources) - return len(rendered.token_ids) - - input_tokens = await asyncio.to_thread(_count) - except Exception: - logger.error("count_tokens failed to render the prompt:\n" - f"{traceback.format_exc()}") - return anthropic_error_response("Internal server error", - "api_error", 500) - - return JSONResponse(content=AnthropicCountTokensResponse( - input_tokens=input_tokens).model_dump()) - # -- Message Batches --------------------------------------------------- - @staticmethod - def _synthetic_request(body: Dict[str, Any]) -> Request: - """A Request for work with no client connection behind it. - - The chat path reads ``raw_request.state`` and may re-read the body via - ``.json()``. A batched request has no live connection, so rather than - duck-type a stub that drifts as that path evolves, build a real Request - over a synthetic ASGI scope - anything the chat path can legitimately - do to a Request keeps working. - """ - payload = json.dumps(body).encode() - - async def receive() -> Dict[str, Any]: - return {"type": "http.request", "body": payload, "more_body": False} - - return Request( - { - "type": "http", - "http_version": "1.1", - "method": "POST", - "scheme": "http", - "path": "/v1/messages", - "raw_path": b"/v1/messages", - "query_string": b"", - "root_path": "", - "headers": [(b"content-type", b"application/json")], - "client": None, - "server": None, - # PerfMetricsMiddleware seeds state["perf_metrics_records"] on a - # real connection (see perf_metrics.py). A batched request never - # passes through the middleware, so seed it here: with - # return_perf_metrics enabled, _extract_metrics appends to this - # list unconditionally and every batched request would otherwise - # die with AttributeError and come back "errored". - # No socket sits behind this request, so await_disconnected - # keys off no_client_connection to avoid polling forever on a - # connection that can never drop. - "state": { - "perf_metrics_records": [], - "no_client_connection": True - }, - }, - receive, - ) - - def _batch_store(self) -> AnthropicBatchStore: - """Created on first use so servers that never batch pay nothing.""" - if getattr(self, "_anthropic_batch_store", None) is None: - self._anthropic_batch_store = AnthropicBatchStore( - runner=self._run_batched_message) - return self._anthropic_batch_store - - async def _run_batched_message( - self, - request: AnthropicMessagesRequest) -> Tuple[str, Dict[str, Any]]: - """Run one batched Messages request through the ordinary chat path.""" - # Streaming makes no sense for a batch: the result is collected, not - # watched. Forcing it off here means a client that sets stream=true in - # params gets a usable message instead of an SSE body in its .jsonl. - request = request.model_copy(update={"stream": False}) - try: - chat_request = convert_anthropic_request(request) - except (AnthropicRequestError, ValidationError) as e: - return "errored", { - "type": "invalid_request_error", - "message": str(e) - } - - raw_request = self._synthetic_request( - chat_request.model_dump(exclude_none=True)) - response = await self.openai_chat(chat_request, raw_request) - if response is None: - return "errored", { - "type": "api_error", - "message": "Internal server error" - } - - status = getattr(response, "status_code", 500) - if status != 200: - try: - payload = json.loads(response.body) - message = payload.get("message") or json.dumps(payload) - except (json.JSONDecodeError, AttributeError): - message = "Internal server error" - err_type = "invalid_request_error" if 400 <= status < 500 else "api_error" - return "errored", {"type": err_type, "message": message} - - try: - chat_response = ChatCompletionResponse(**json.loads(response.body)) - anthropic_response = convert_chat_response(chat_response) - except (AnthropicResponseError, ValidationError, json.JSONDecodeError): - logger.error( - "Invalid response from OpenAI chat pipeline in batch:\n" - f"{traceback.format_exc()}") - return "errored", { - "type": "api_error", - "message": "Internal server error" - } - return "succeeded", anthropic_response.model_dump(exclude_none=True) - - async def anthropic_create_batch( - self, request: AnthropicCreateBatchRequest) -> Response: - try: - batch = self._batch_store().create(request.requests) - except BatchStoreFullError as e: - # 429, not 400: the request is valid and an identical one will - # succeed later, so the client should back off rather than edit it. - return anthropic_error_response(str(e), "rate_limit_error", 429) - except ValueError as e: - return anthropic_error_response(str(e), "invalid_request_error", - 400) - return JSONResponse(content=batch.model_dump()) - - async def anthropic_list_batches( - self, - # Validated here rather than clamped in the store: clamping turned - # limit=0 into 1 instead of a 400, and bounded the page against the - # retention limit, which is an unrelated quantity. - limit: int = Query(20, ge=1, le=100), - after_id: Optional[str] = None, - before_id: Optional[str] = None) -> Response: - try: - batches, has_more = self._batch_store().list(limit=limit, - after_id=after_id, - before_id=before_id) - except LookupError as e: - # Saying so beats an empty page: an empty page is indistinguishable - # from a finished walk, so the client stops early having silently - # skipped every batch past the cursor. - return anthropic_error_response( - f"Unknown pagination cursor {e.args[0]!r}. The batch may have " - "been deleted or expired; restart the listing without a cursor.", - "invalid_request_error", 400) - listing = AnthropicBatchList( - data=batches, - # Reported honestly rather than hardcoded false: a client that - # asked for N and received N cannot otherwise tell a full page from - # the end of the list, and silently stops early. It pages on with - # ?after_id=. - has_more=has_more, - first_id=batches[0].id if batches else None, - last_id=batches[-1].id if batches else None, - ) - return JSONResponse(content=listing.model_dump()) - - async def anthropic_get_batch(self, batch_id: str) -> Response: - batch = self._batch_store().get(batch_id) - if batch is None: - return self._no_such_batch(batch_id) - return JSONResponse(content=batch.model_dump()) - - async def anthropic_cancel_batch(self, batch_id: str) -> Response: - batch = self._batch_store().cancel(batch_id) - if batch is None: - return self._no_such_batch(batch_id) - return JSONResponse(content=batch.model_dump()) - - async def anthropic_delete_batch(self, batch_id: str) -> Response: - try: - deleted = self._batch_store().delete(batch_id) - except ValueError as e: - return anthropic_error_response(str(e), "invalid_request_error", - 400) - if not deleted: - return self._no_such_batch(batch_id) - return JSONResponse(content=AnthropicBatchDeleteResponse( - id=batch_id).model_dump()) - - async def anthropic_batch_results(self, batch_id: str) -> Response: - store = self._batch_store() - results = store.results(batch_id) - if results is None: - if not store.has(batch_id): - return self._no_such_batch(batch_id) - # Known batch, but not finished. 404 would say "no such batch", - # which is wrong and would stop a client from polling. - return anthropic_error_response( - f"Message Batch {batch_id} is still processing; results are " - "available once processing_status is 'ended'", - "invalid_request_error", 400) - return Response(content=results_to_jsonl(results), - media_type="application/x-ndjson") - - @staticmethod - def _no_such_batch(batch_id: str) -> Response: - # Batches are held in memory only, so an id from before a restart lands - # here. Say so, rather than leaving the caller to wonder. - return anthropic_error_response( - f"Message Batch {batch_id} not found. Batches are stored in memory " - "and do not survive a server restart.", "not_found_error", 404) - async def openai_mm_encoder(self, request: ChatCompletionRequest, raw_request: Request) -> Response: From d5e971e0d8b71f2ee8d651b20ebb8f493c44f557 Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Tue, 3 Nov 2026 08:37:11 +0000 Subject: [PATCH 25/27] [TRTLLM-17263][refactor] One Anthropic Messages serving class for the aggregated and disaggregated servers The Anthropic Messages orchestration existed three times (aggregated messages, the aggregated Message Batches runner, disaggregated messages), each serializing the chat result, parsing it back and sniffing the error body for a message; the batch runner entered the chat pipeline through a synthetic ASGI request. One AnthropicMessagesServing now runs it for any ChatBackend: the aggregated backend is the typed chat core create_chat_completion, the disaggregated backend is the entry-point wrapper's prepare, run and finish steps without the serialization step (the new _call_entry_point, in the same span). A RequestContext carries what the chat core reads from a request (headers, state, raw JSON, disconnect polling); the batch path builds one directly. The eight aggregated Anthropic routes come from the new router in tensorrt_llm/serve/api/anthropic/api_router.py, attached through the router table and gated by the new harmony_model capability as the route registration was gated. The launcher forwarding routes of an attached frontend register before the routers, so that they keep winning over the local Message Batches handlers. Each server keeps its own answers: the streamed model name and the error-message extraction come from its backend (stream_model, error_message); an engine failure stays a 500 on both. No behavior change. The new modules join the mypy file list; the two serving-class modules join the multi-GPU trigger table because they carry the disaggregated messages path. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- jenkins/L0_MergeRequest.groovy | 3 + pyproject.toml | 3 + tensorrt_llm/serve/AGENTS.md | 4 + tensorrt_llm/serve/_base/chat_backend.py | 64 ++++++++ tensorrt_llm/serve/_base/request_context.py | 149 ++++++++++++++++++ tensorrt_llm/serve/_base/routing.py | 12 ++ tensorrt_llm/serve/_base/sse.py | 7 +- tensorrt_llm/serve/api/anthropic/__init__.py | 5 +- .../serve/api/anthropic/api_router.py | 139 ++++++++++++++++ tensorrt_llm/serve/api/anthropic/protocol.py | 2 +- tensorrt_llm/serve/api/anthropic/serving.py | 119 +++++++++++++- tensorrt_llm/serve/openai_disagg_server.py | 147 ++++++++++------- tensorrt_llm/serve/openai_server.py | 112 +++++++------ 13 files changed, 649 insertions(+), 117 deletions(-) create mode 100644 tensorrt_llm/serve/_base/chat_backend.py create mode 100644 tensorrt_llm/serve/_base/request_context.py create mode 100644 tensorrt_llm/serve/api/anthropic/api_router.py diff --git a/jenkins/L0_MergeRequest.groovy b/jenkins/L0_MergeRequest.groovy index 541ec1cb7626..6d31bd8f7e27 100644 --- a/jenkins/L0_MergeRequest.groovy +++ b/jenkins/L0_MergeRequest.groovy @@ -1950,6 +1950,9 @@ def getMultiGpuFileChanged(pipeline, testFilter, globalVars) "tensorrt_llm/serve/__init__.py", // _MsgspecRoute decodes the disaggregated frontend's msgpack bodies. "tensorrt_llm/serve/_base/", + // The disaggregated /v1/messages runs through the Anthropic serving class. + "tensorrt_llm/serve/api/anthropic/api_router.py", + "tensorrt_llm/serve/api/anthropic/serving.py", // /health is the disaggregated readiness probe. "tensorrt_llm/serve/api/ops/", "tensorrt_llm/serve/cluster_storage.py", diff --git a/pyproject.toml b/pyproject.toml index 2b77d93e161b..20a951c8b770 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -865,8 +865,11 @@ files = [ "tensorrt_llm/_torch/pyexecutor/sampler/token_ban.py", "tensorrt_llm/_torch/pyexecutor/sampler/ops/flashinfer.py", "tensorrt_llm/_torch/pyexecutor/sampler/ops/vanilla.py", + "tensorrt_llm/serve/_base/chat_backend.py", "tensorrt_llm/serve/_base/errors.py", + "tensorrt_llm/serve/_base/request_context.py", "tensorrt_llm/serve/_base/routing.py", + "tensorrt_llm/serve/api/anthropic/api_router.py", "tensorrt_llm/serve/api/ops/system/api_router.py", "tensorrt_llm/serve/api/ops/system/serving.py", "tests/unittest/_torch/sampler/test_torch_sampler.py", diff --git a/tensorrt_llm/serve/AGENTS.md b/tensorrt_llm/serve/AGENTS.md index 927b35d9740e..6ff6ffeaf440 100644 --- a/tensorrt_llm/serve/AGENTS.md +++ b/tensorrt_llm/serve/AGENTS.md @@ -20,6 +20,10 @@ Name files and symbols here, never line numbers. `make_router(prefix=...)` instead. The composition root sets each serving object on `app.state._serving_` before attaching; routers read it from `request.app.state`, never from a closure. +- Anthropic Messages: `serve/api/anthropic/`. One `AnthropicMessagesServing` serves both + servers through a `ChatBackend`; the aggregated routes are in `api_router.py`, the + disaggregated frontend keeps its own `/v1/messages` route and calls the same class. + Check: `tests/unittest/serve/api/anthropic/test_messages_behavior.py`. - Model-specific behavior: a hook on `ServingExtension` (`serve/serving_extensions.py`, `serve/extensions/`). Do not branch on a model name in a handler. - Prompt rendering: `serve/render/` only. Handlers do not call `apply_chat_template`. diff --git a/tensorrt_llm/serve/_base/chat_backend.py b/tensorrt_llm/serve/_base/chat_backend.py new file mode 100644 index 000000000000..06e90c8856d2 --- /dev/null +++ b/tensorrt_llm/serve/_base/chat_backend.py @@ -0,0 +1,64 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""The chat implementation a wire-dialect serving class runs on.""" + +from __future__ import annotations + +from typing import AsyncIterator, Protocol, Union + +from tensorrt_llm.serve._base.errors import ServingError +from tensorrt_llm.serve._base.request_context import RequestContext +from tensorrt_llm.serve.openai_protocol import ChatCompletionRequest, ChatCompletionResponse + +# Non-streaming: the typed response. Streaming: the OpenAI SSE stream, as str +# frames (aggregated) or as the worker's raw bytes (disaggregated). +ChatResult = Union[ChatCompletionResponse, AsyncIterator[str], AsyncIterator[bytes]] + + +class ChatBackend(Protocol): + """The chat implementation a wire-dialect serving class runs on. + + The aggregated server implements it over its typed chat core, the + disaggregated frontend over its prepare, run and finish steps. Each side + keeps its own streamed model name and error-message extraction, which is + why both are methods of the backend and not of the serving class. + """ + + async def create_chat_completion( + self, request: ChatCompletionRequest, ctx: RequestContext + ) -> ChatResult: + """Run one chat completion. + + Non-streaming: the typed response. Streaming: the OpenAI SSE stream, as + str frames (aggregated) or the worker's raw bytes (disaggregated); one + stream never mixes the two. + + Raises: + ServingError: The request was rejected; the HTTP status travels + apart from the body. + EngineFatalError: The engine failed; SIGINT has been sent. + """ + ... + + def stream_model(self, request: ChatCompletionRequest) -> str: + """The model a translated stream announces when no chunk names one.""" + ... + + def error_message(self, error: ServingError) -> str: + """The client-facing message of an error. + + Each backend extracts it from the error body its own way. + """ + ... diff --git a/tensorrt_llm/serve/_base/request_context.py b/tensorrt_llm/serve/_base/request_context.py new file mode 100644 index 000000000000..2e2b5949267b --- /dev/null +++ b/tensorrt_llm/serve/_base/request_context.py @@ -0,0 +1,149 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""What the chat core reads from an HTTP request, as one explicit object.""" + +from __future__ import annotations + +import json +from typing import Any, Dict, List, Mapping, Optional + +from starlette.datastructures import Address, Headers, State +from starlette.requests import Request + +_UNSET: Any = object() + + +class RequestContext: + """What the chat core reads from an HTTP request, made explicit. + + ``from_http()`` wraps the request FastAPI passed to a route and delegates to + it lazily, so the middleware and the core share one state + (``server_arrival_time``, ``server_first_token_time``, + ``perf_metrics_records``) and msgpack bodies keep decoding through the + request's own ``json()``. ``internal()`` is a request with no client behind + it, for Anthropic message batches. + + The properties ``headers``, ``state``, ``client``, ``json()`` and + ``is_disconnected()`` carry the names ``Request`` uses, so helpers that + still take a ``Request`` accept a context unchanged. The object is always + truthy, as every ``Request`` is. + """ + + def __init__( + self, + *, + request: Optional[Request] = None, + state: Optional[State] = None, + raw_json: Optional[Dict[str, Any]] = None, + ) -> None: + """Use ``from_http()`` or ``internal()``.""" + if request is None and state is None: + raise ValueError( + "RequestContext needs a request to delegate to or an internal state; " + "build it with RequestContext.from_http(request) or " + "RequestContext.internal(raw_json)" + ) + self._request = request + self._state = state + self._raw_json = raw_json + self._json_cache: Any = _UNSET + + @classmethod + def from_http(cls, request: Request) -> RequestContext: + """Wrap ``request``; nothing is read from it until the core asks.""" + return cls(request=request) + + @classmethod + def internal(cls, raw_json: Dict[str, Any]) -> RequestContext: + """A request without headers or a client, whose raw JSON is ``raw_json``. + + The state carries what the middleware would have seeded on a real + connection: an empty ``perf_metrics_records`` list, and + ``no_client_connection`` so that the disconnect watcher does not poll a + connection that can never drop. + """ + return cls( + state=State({"perf_metrics_records": [], "no_client_connection": True}), + raw_json=raw_json, + ) + + @property + def http_request(self) -> Optional[Request]: + """The wrapped HTTP request; ``None`` for an internal context.""" + return self._request + + @property + def headers(self) -> Mapping[str, str]: + if self._request is None: + return Headers() + return self._request.headers + + @property + def state(self) -> State: + if self._request is not None: + return self._request.state + assert self._state is not None # the constructor requires one of the two + return self._state + + @property + def client(self) -> Optional[Address]: + if self._request is None: + return None + return self._request.client + + async def json(self) -> Any: + """The request body as JSON. + + On the HTTP path this is the request's own ``json()`` (msgpack-aware on + workers, cached by the request). An internal context answers its raw + JSON after one encode-decode round trip, computed once, which is what a + body read from the wire looks like. + """ + if self._request is not None: + return await self._request.json() + if self._json_cache is _UNSET: + self._json_cache = json.loads(json.dumps(self._raw_json)) + return self._json_cache + + async def is_disconnected(self) -> bool: + """Whether the client went away; always ``False`` without a client.""" + if self._request is None: + return False + return await self._request.is_disconnected() + + # Typed views of the state, for code written against the context. + + @property + def no_client_connection(self) -> bool: + return bool(getattr(self.state, "no_client_connection", False)) + + @property + def server_arrival_time(self) -> Optional[float]: + value: Optional[float] = getattr(self.state, "server_arrival_time", None) + return value + + @property + def server_first_token_time(self) -> Optional[float]: + value: Optional[float] = getattr(self.state, "server_first_token_time", None) + return value + + @server_first_token_time.setter + def server_first_token_time(self, value: Optional[float]) -> None: + self.state.server_first_token_time = value + + @property + def perf_metrics_records(self) -> List[Dict[str, Any]]: + records: List[Dict[str, Any]] = self.state.perf_metrics_records + return records diff --git a/tensorrt_llm/serve/_base/routing.py b/tensorrt_llm/serve/_base/routing.py index 3f461a029282..79e9b0b4a0d9 100644 --- a/tensorrt_llm/serve/_base/routing.py +++ b/tensorrt_llm/serve/_base/routing.py @@ -54,6 +54,9 @@ class ServerCapabilities: role: Optional[ServerRole] role_group: RoleGroup + # A harmony (gpt-oss) model serves chat through the harmony path, which + # the Anthropic Messages routes do not support; LLM role only. + harmony_model: bool = False @classmethod def for_role(cls, role: Optional[ServerRole]) -> ServerCapabilities: @@ -64,6 +67,10 @@ def _always(caps: ServerCapabilities) -> bool: return True +def _anthropic_routes_enabled(caps: ServerCapabilities) -> bool: + return not caps.harmony_model + + @dataclass(frozen=True) class RouterSpec: """One router module and the roles and capabilities that mount it.""" @@ -77,6 +84,11 @@ class RouterSpec: # order: a router that registers GET /metrics must come before the Prometheus mount. ROUTERS: tuple[RouterSpec, ...] = ( RouterSpec("tensorrt_llm.serve.api.ops.system.api_router", roles=ALL_ROLE_GROUPS), + RouterSpec( + "tensorrt_llm.serve.api.anthropic.api_router", + roles=frozenset({RoleGroup.LLM}), + when=_anthropic_routes_enabled, + ), ) diff --git a/tensorrt_llm/serve/_base/sse.py b/tensorrt_llm/serve/_base/sse.py index 6b91472ef45d..ef935a266d10 100644 --- a/tensorrt_llm/serve/_base/sse.py +++ b/tensorrt_llm/serve/_base/sse.py @@ -13,13 +13,12 @@ async def chat_stream_generator(promise: RequestOutput, postproc_params: PostprocParams, *, server, - raw_request) -> AsyncGenerator[str, None]: + ctx) -> AsyncGenerator[str, None]: try: if not server.postproc_worker_enabled: post_processor, args = postproc_params.post_processor, postproc_params.postproc_args first_response = await anext(promise) - raw_request.state.server_first_token_time = server._adjusted_steady_clock.now( - ) + ctx.state.server_first_token_time = server._adjusted_steady_clock.now() pp_results = first_response.outputs[ 0]._postprocess_result if server.postproc_worker_enabled else post_processor( first_response, args) @@ -34,7 +33,7 @@ async def chat_stream_generator(promise: RequestOutput, for pp_res in pp_results: yield pp_res yield "data: [DONE]\n\n" - await server._extract_metrics(res, raw_request) + await server._extract_metrics(res, ctx) nvtx_mark("generation ends") except Exception: # Only failures are logged: a consumer that stops reading closes diff --git a/tensorrt_llm/serve/api/anthropic/__init__.py b/tensorrt_llm/serve/api/anthropic/__init__.py index 035527480c4f..82f6904e8d05 100644 --- a/tensorrt_llm/serve/api/anthropic/__init__.py +++ b/tensorrt_llm/serve/api/anthropic/__init__.py @@ -12,4 +12,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""Anthropic Messages wire dialect: protocol, adapter and batches. Import submodules directly.""" +"""Anthropic Messages wire dialect: protocol, adapter, batches, serving classes and router. + +Import submodules directly. +""" diff --git a/tensorrt_llm/serve/api/anthropic/api_router.py b/tensorrt_llm/serve/api/anthropic/api_router.py new file mode 100644 index 000000000000..bf2b2432a812 --- /dev/null +++ b/tensorrt_llm/serve/api/anthropic/api_router.py @@ -0,0 +1,139 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""The Anthropic Messages router of the aggregated server. + +``/v1/messages``, ``/v1/messages/count_tokens`` and the six Message Batches +routes. The disaggregated frontend keeps its own ``/v1/messages`` route and +calls the same serving class. +""" + +from __future__ import annotations + +from typing import Optional + +from fastapi import FastAPI, Query, Request, Response + +from tensorrt_llm.serve._base.request_context import RequestContext +from tensorrt_llm.serve._base.routing import make_router +from tensorrt_llm.serve.api.anthropic.protocol import ( + AnthropicCountTokensRequest, + AnthropicCreateBatchRequest, + AnthropicMessagesRequest, +) +from tensorrt_llm.serve.api.anthropic.serving import ( + AnthropicBatchesServing, + AnthropicCountTokensServing, + AnthropicMessagesServing, +) + +# The serving objects the composition root sets on app.state before attaching. +APP_STATE_KEYS = ( + "anthropic_serving_messages", + "anthropic_serving_count_tokens", + "anthropic_serving_batches", +) + +router = make_router() + + +# No app.state defaults: wiring bugs fail loudly. +def _messages(raw_request: Request) -> AnthropicMessagesServing: + serving: AnthropicMessagesServing = raw_request.app.state.anthropic_serving_messages + return serving + + +def _count_tokens(raw_request: Request) -> AnthropicCountTokensServing: + serving: AnthropicCountTokensServing = raw_request.app.state.anthropic_serving_count_tokens + return serving + + +def _batches(raw_request: Request) -> AnthropicBatchesServing: + serving: AnthropicBatchesServing = raw_request.app.state.anthropic_serving_batches + return serving + + +async def anthropic_messages(request: AnthropicMessagesRequest, raw_request: Request) -> Response: + """Serve the Anthropic Messages API on top of the chat pipeline.""" + return await _messages(raw_request).create_message( + request, RequestContext.from_http(raw_request) + ) + + +async def anthropic_count_tokens( + request: AnthropicCountTokensRequest, raw_request: Request +) -> Response: + """Report how many input tokens a Messages request would consume.""" + return await _count_tokens(raw_request).anthropic_count_tokens(request) + + +async def anthropic_create_batch( + request: AnthropicCreateBatchRequest, raw_request: Request +) -> Response: + return await _batches(raw_request).anthropic_create_batch(request) + + +async def anthropic_list_batches( + raw_request: Request, + limit: int = Query(20, ge=1, le=100), + after_id: Optional[str] = None, + before_id: Optional[str] = None, +) -> Response: + return await _batches(raw_request).anthropic_list_batches( + limit=limit, after_id=after_id, before_id=before_id + ) + + +async def anthropic_get_batch(batch_id: str, raw_request: Request) -> Response: + return await _batches(raw_request).anthropic_get_batch(batch_id) + + +async def anthropic_delete_batch(batch_id: str, raw_request: Request) -> Response: + return await _batches(raw_request).anthropic_delete_batch(batch_id) + + +async def anthropic_cancel_batch(batch_id: str, raw_request: Request) -> Response: + return await _batches(raw_request).anthropic_cancel_batch(batch_id) + + +async def anthropic_batch_results(batch_id: str, raw_request: Request) -> Response: + return await _batches(raw_request).anthropic_batch_results(batch_id) + + +# Registration order is match order. The literal batch paths come before the +# parameterized ones, as they always did. +router.add_api_route("/v1/messages", anthropic_messages, methods=["POST"]) +# Claude Code calls count_tokens before most turns to decide whether to +# compact its context; an unregistered route 404s every call. +router.add_api_route("/v1/messages/count_tokens", anthropic_count_tokens, methods=["POST"]) +router.add_api_route("/v1/messages/batches", anthropic_create_batch, methods=["POST"]) +router.add_api_route("/v1/messages/batches", anthropic_list_batches, methods=["GET"]) +router.add_api_route("/v1/messages/batches/{batch_id}", anthropic_get_batch, methods=["GET"]) +router.add_api_route("/v1/messages/batches/{batch_id}", anthropic_delete_batch, methods=["DELETE"]) +router.add_api_route( + "/v1/messages/batches/{batch_id}/cancel", anthropic_cancel_batch, methods=["POST"] +) +router.add_api_route( + "/v1/messages/batches/{batch_id}/results", anthropic_batch_results, methods=["GET"] +) + + +def attach_router(app: FastAPI) -> None: + for key in APP_STATE_KEYS: + if not hasattr(app.state, key): + raise RuntimeError( + f"app.state.{key} must be set before attaching " + "tensorrt_llm.serve.api.anthropic.api_router; the composition root sets it." + ) + app.include_router(router) diff --git a/tensorrt_llm/serve/api/anthropic/protocol.py b/tensorrt_llm/serve/api/anthropic/protocol.py index 929f04d520e0..80eaacf94dac 100644 --- a/tensorrt_llm/serve/api/anthropic/protocol.py +++ b/tensorrt_llm/serve/api/anthropic/protocol.py @@ -344,7 +344,7 @@ class AnthropicCountTokensResponse(AnthropicBaseModel): AnthropicBatchResultType = Literal["succeeded", "errored", "canceled", "expired"] # Anthropic's documented ceiling. Enforced because every request in a batch is -# held in memory here (see anthropic_batches), so an unbounded batch is a way +# held in memory here (see batches), so an unbounded batch is a way # to exhaust the server's memory with a single call. ANTHROPIC_BATCH_MAX_REQUESTS = 100_000 diff --git a/tensorrt_llm/serve/api/anthropic/serving.py b/tensorrt_llm/serve/api/anthropic/serving.py index def94f0fb542..daeddfe2ed53 100644 --- a/tensorrt_llm/serve/api/anthropic/serving.py +++ b/tensorrt_llm/serve/api/anthropic/serving.py @@ -19,17 +19,22 @@ """ import asyncio +import json import traceback -from typing import Callable, Optional +from typing import Any, Callable, Dict, Optional, Tuple from fastapi import Query -from fastapi.responses import JSONResponse, Response +from fastapi.responses import JSONResponse, Response, StreamingResponse from pydantic import ValidationError from tensorrt_llm.logger import logger +from tensorrt_llm.serve._base.chat_backend import ChatBackend +from tensorrt_llm.serve._base.errors import EngineFatalError, ServingError +from tensorrt_llm.serve._base.request_context import RequestContext from tensorrt_llm.serve.api.anthropic.adapter import ( - AnthropicRequestError, anthropic_error_response, - convert_anthropic_count_tokens_request) + AnthropicRequestError, AnthropicResponseError, anthropic_error_response, + convert_anthropic_count_tokens_request, convert_anthropic_request, + convert_chat_response, reframe_openai_stream) from tensorrt_llm.serve.api.anthropic.batches import (AnthropicBatchStore, BatchItemRunner, BatchStoreFullError, @@ -37,12 +42,116 @@ from tensorrt_llm.serve.api.anthropic.protocol import ( AnthropicBatchDeleteResponse, AnthropicBatchList, AnthropicCountTokensRequest, AnthropicCountTokensResponse, - AnthropicCreateBatchRequest) + AnthropicCreateBatchRequest, AnthropicMessagesRequest) +from tensorrt_llm.serve.openai_protocol import ChatCompletionResponse from tensorrt_llm.serve.render import (HarmonyConversionError, RenderResources, UnsupportedRenderError, render_chat, render_chat_for_length) +def _error_type(status: int) -> str: + """The Anthropic error type both servers derive from an HTTP status.""" + return "invalid_request_error" if 400 <= status < 500 else "api_error" + + +class AnthropicMessagesServing: + """The Anthropic Messages orchestration, for any ChatBackend. + + One class serves the aggregated server (over its typed chat core), the + disaggregated frontend (over its prepare, run and finish steps) and the + aggregated Message Batches. What differs between the servers, the model a + stream without chunks announces and how an error's message is extracted, + comes from the backend. + """ + + def __init__(self, backend: ChatBackend) -> None: + self._backend = backend + + async def create_message(self, request: AnthropicMessagesRequest, + ctx: RequestContext) -> Response: + """Serve the Anthropic Messages API on top of the chat pipeline. + + The Anthropic request is translated into a ChatCompletionRequest, + handed to the chat backend unchanged (chat template, tool parser and + post-processing are reused verbatim), and the typed result is + translated back: JSON body for non-streaming, SSE reframing for + streaming. + """ + try: + chat_request = convert_anthropic_request(request) + except (AnthropicRequestError, ValidationError) as e: + return anthropic_error_response(str(e), "invalid_request_error", + 400) + try: + result = await self._backend.create_chat_completion( + chat_request, ctx) + except EngineFatalError: + return anthropic_error_response("Internal server error", + "api_error", 500) + except ServingError as e: + return anthropic_error_response(self._backend.error_message(e), + _error_type(e.status), e.status) + if not isinstance(result, ChatCompletionResponse): + return StreamingResponse( + content=reframe_openai_stream( + result, model=self._backend.stream_model(chat_request)), + media_type="text/event-stream", + ) + try: + anthropic_response = convert_chat_response(result) + except (AnthropicResponseError, ValidationError, json.JSONDecodeError): + logger.error("Invalid response from OpenAI chat pipeline:\n" + f"{traceback.format_exc()}") + return anthropic_error_response("Internal server error", + "api_error", 500) + return JSONResponse(content=anthropic_response.model_dump( + exclude_none=True)) + + async def run_batched_message( + self, + request: AnthropicMessagesRequest) -> Tuple[str, Dict[str, Any]]: + """Run one batched Messages request through the chat backend.""" + # Streaming makes no sense for a batch: the result is collected, not + # watched. Forcing it off here means a client that sets stream=true in + # params gets a usable message instead of an SSE body in its .jsonl. + request = request.model_copy(update={"stream": False}) + try: + chat_request = convert_anthropic_request(request) + except (AnthropicRequestError, ValidationError) as e: + return "errored", { + "type": "invalid_request_error", + "message": str(e) + } + # A batched request has no client connection behind it: the context + # carries the seeded state and the converted request as its raw JSON. + ctx = RequestContext.internal( + chat_request.model_dump(exclude_none=True)) + try: + result = await self._backend.create_chat_completion( + chat_request, ctx) + except EngineFatalError: + return "errored", { + "type": "api_error", + "message": "Internal server error" + } + except ServingError as e: + return "errored", { + "type": _error_type(e.status), + "message": self._backend.error_message(e) + } + try: + anthropic_response = convert_chat_response(result) + except (AnthropicResponseError, ValidationError, json.JSONDecodeError): + logger.error( + "Invalid response from OpenAI chat pipeline in batch:\n" + f"{traceback.format_exc()}") + return "errored", { + "type": "api_error", + "message": "Internal server error" + } + return "succeeded", anthropic_response.model_dump(exclude_none=True) + + class AnthropicCountTokensServing: """Counts the input tokens of a Messages request over the rendered prompt.""" diff --git a/tensorrt_llm/serve/openai_disagg_server.py b/tensorrt_llm/serve/openai_disagg_server.py index f4db651f268f..60f34f5f4ac7 100644 --- a/tensorrt_llm/serve/openai_disagg_server.py +++ b/tensorrt_llm/serve/openai_disagg_server.py @@ -22,14 +22,13 @@ import traceback import uuid from contextlib import asynccontextmanager -from typing import Callable, Optional +from typing import Callable, Optional, Union import aiohttp import uvicorn from fastapi import FastAPI, HTTPException, Request from fastapi.exceptions import RequestValidationError from fastapi.responses import JSONResponse, Response, StreamingResponse -from pydantic import ValidationError from tensorrt_llm.executor import CppExecutorError from tensorrt_llm.executor.executor import CppExecutorError @@ -41,15 +40,13 @@ from tensorrt_llm.llmapi.disagg_utils import (DisaggServerConfig, MetadataServerConfig, ServerRole) from tensorrt_llm.logger import logger +from tensorrt_llm.serve._base.errors import EngineFatalError, ServingError +from tensorrt_llm.serve._base.request_context import RequestContext from tensorrt_llm.serve._telemetry import create_uvicorn_server -from tensorrt_llm.serve.api.anthropic.adapter import (AnthropicRequestError, - AnthropicResponseError, - anthropic_error_response, - convert_anthropic_request, - convert_chat_response, - reframe_openai_stream) +from tensorrt_llm.serve.api.anthropic.adapter import anthropic_error_response from tensorrt_llm.serve.api.anthropic.protocol import ( AnthropicCountTokensRequest, AnthropicMessagesRequest) +from tensorrt_llm.serve.api.anthropic.serving import AnthropicMessagesServing from tensorrt_llm.serve.cluster_storage import ( HttpClusterStorageServer, create_cluster_storage, validate_http_cluster_storage_scope) @@ -62,9 +59,9 @@ from tensorrt_llm.serve.openai_disagg_service import ( OpenAIDisaggregatedService, ResponseHooks) from tensorrt_llm.serve.openai_protocol import ( - ChatCompletionRequest, ChatCompletionResponse, CompletionRequest, - ConversationParams, ErrorResponse, ResponsesRequest, UCompletionRequest, - UCompletionResponse, ensure_request_chat_template_allowed) + ChatCompletionRequest, CompletionRequest, ConversationParams, ErrorResponse, + ResponsesRequest, UCompletionRequest, UCompletionResponse, + ensure_request_chat_template_allowed) from tensorrt_llm.serve.perf_metrics import (DisaggPerfMetricsCollector, PerfMetricsJsonlWriter, PerfMetricsMiddleware, @@ -214,6 +211,61 @@ def _upstream_error_message(error: aiohttp.ClientResponseError) -> str: return f"context worker rejected the request with status {error.status}" +def _envelope_message(response: Response) -> str: + """The message of an error envelope, read the way the disaggregated /v1/messages did.""" + try: + payload = json.loads(response.body) + message = (payload.get("message") or payload.get("detail") + or payload.get("error") or json.dumps(payload)) + except (json.JSONDecodeError, AttributeError, TypeError): + message = "Internal server error" + return str(message) + + +def _disagg_serving_error(status: int, message: str) -> ServingError: + """A disaggregated pipeline failure as a ServingError; the message is final.""" + code: Union[int, str] = status + if is_context_length_exceeded_message(message): + code = CONTEXT_LENGTH_EXCEEDED_CODE + return ServingError(status, ErrorResponse( + message=message, + type="BadRequestError" if 400 <= status < 500 else "InternalServerError", + code=code)) + + +class _DisaggChatBackend: + """ChatBackend over the prepare, run and finish steps (no serialize).""" + + def __init__(self, server: "OpenAIDisaggServer") -> None: + self._server = server + + async def create_chat_completion(self, request: ChatCompletionRequest, + ctx: RequestContext): + try: + # The service is read per request, so tests may replace it. + outcome = await self._server._call_entry_point( + self._server._service.openai_chat_completion, request, + ctx.http_request) + except HTTPException as e: + raise _disagg_serving_error(e.status_code, str(e.detail)) from e + if outcome is None: + # Engine failure: _handle_exception has sent SIGINT. + raise EngineFatalError() + if isinstance(outcome, Response): + # The envelope _handle_exception returns for a worker error that + # carries a machine-readable code (context_length_exceeded). + raise _disagg_serving_error(outcome.status_code, + _envelope_message(outcome)) + return outcome + + def stream_model(self, request: ChatCompletionRequest) -> str: + # The client's name, as the disaggregated route always announced it. + return request.model + + def error_message(self, error: ServingError) -> str: + return error.body.message + + class OpenAIDisaggServer: """OpenAI- and Anthropic-compatible frontend of disaggregated serving. @@ -281,6 +333,11 @@ def __init__(self, self._service = OpenAIDisaggregatedService( self._config, self._coordinator, self._create_client, req_timeout_secs=self._req_timeout_secs) + # The Anthropic Messages orchestration is shared with the aggregated + # server; here it runs over the steps of _wrap_entry_point, without + # the serialization the OpenAI routes need. + self._anthropic_messages_serving = AnthropicMessagesServing( + _DisaggChatBackend(self)) try: otlp_cfg = config.otlp_config @@ -450,6 +507,26 @@ async def wrapper(req: request_type, raw_req: Request) -> Response: return self._finish_failed_request(e, req, hooks) return wrapper + @tracing.trace_span("disaggregated_request") + async def _call_entry_point(self, entry_point: Callable, + req: UCompletionRequest, raw_req: Request): + """Prepare, run and finish one request without serializing it. + + Returns the service's result (a typed response or the worker stream) + or, on failure, what _finish_failed_request returns (an error + envelope, or None after an engine failure); its HTTPException + propagates. + """ + hooks: Optional[RawRequestResponseHooks] = None + try: + hooks = self._prepare_request(req, raw_req) + result = await self._run_entry_point(entry_point, req, hooks, + raw_req) + self._finish_request(req, hooks) + return result + except Exception as e: + return self._finish_failed_request(e, req, hooks) + def _prepare_request(self, req: UCompletionRequest, raw_req: Request) -> RawRequestResponseHooks: """Account for a request and check it before any worker is involved. @@ -524,52 +601,8 @@ def _serialize_response( async def anthropic_messages(self, request: AnthropicMessagesRequest, raw_request: Request) -> Response: """Serve Anthropic Messages through the disaggregated chat pipeline.""" - try: - chat_request = convert_anthropic_request(request) - except (AnthropicRequestError, ValidationError) as e: - return anthropic_error_response(str(e), "invalid_request_error", - 400) - - try: - openai_response = await self._wrap_entry_point( - self._service.openai_chat_completion)(chat_request, raw_request) - except HTTPException as e: - error_type = ("invalid_request_error" - if 400 <= e.status_code < 500 else "api_error") - return anthropic_error_response(str(e.detail), error_type, - e.status_code) - - if isinstance(openai_response, StreamingResponse): - return StreamingResponse( - content=reframe_openai_stream(openai_response.body_iterator, - model=request.model), - media_type="text/event-stream", - ) - - status = getattr(openai_response, "status_code", 500) - if status != 200: - try: - payload = json.loads(openai_response.body) - message = (payload.get("message") or payload.get("detail") - or payload.get("error") or json.dumps(payload)) - except (json.JSONDecodeError, AttributeError, TypeError): - message = "Internal server error" - error_type = ("invalid_request_error" - if 400 <= status < 500 else "api_error") - return anthropic_error_response(str(message), error_type, status) - - try: - chat_response = ChatCompletionResponse( - **json.loads(openai_response.body)) - anthropic_response = convert_chat_response(chat_response) - except (AnthropicResponseError, ValidationError, json.JSONDecodeError): - logger.error( - "Invalid response from OpenAI chat pipeline:\n" - f"{traceback.format_exc()}") - return anthropic_error_response("Internal server error", - "api_error", 500) - return JSONResponse(content=anthropic_response.model_dump( - exclude_none=True)) + return await self._anthropic_messages_serving.create_message( + request, RequestContext.from_http(raw_request)) async def get_model(self) -> Response: """List the served model, as a context worker reports it.""" diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py index 03ab10ef02fb..242995bab904 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -5,6 +5,7 @@ import array import asyncio import base64 +import dataclasses import functools import json import os @@ -67,10 +68,15 @@ 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.request_context import RequestContext +from tensorrt_llm.serve._base.routing import (RoleGroup, ServerCapabilities, + attach_routers, role_group_of) from tensorrt_llm.serve._base.sse import chat_stream_generator from tensorrt_llm.serve._perf_metrics_config import resolve_perf_metrics_config from tensorrt_llm.serve.api.anthropic.adapter import anthropic_error_response +from tensorrt_llm.serve.api.anthropic.serving import ( + AnthropicBatchesServing, AnthropicCountTokensServing, + AnthropicMessagesServing) from tensorrt_llm.serve.api.ops.system.serving import SystemServing from tensorrt_llm.serve.chat_tokenization import tokenize_harmony_chat_request from tensorrt_llm.serve.chat_utils import (load_chat_template, @@ -530,6 +536,30 @@ def resolve_spec_decode_num_spec_tokens(args: Any) -> Optional[int]: return getattr(spec_config, "max_draft_len", None) +class _AggregatedChatBackend: + """ChatBackend over OpenAIServer's typed chat core. + + Transitional: deleted when the chat core moves out of OpenAIServer into a + serving class that implements stream_model and error_message itself. + """ + + def __init__(self, server: "OpenAIServer") -> None: + self._server = server + + async def create_chat_completion(self, request: ChatCompletionRequest, + ctx: RequestContext): + return await self._server.create_chat_completion(request, ctx) + + def stream_model(self, request: ChatCompletionRequest) -> str: + # The served name, as the aggregated route always announced it. + return self._server.model + + def error_message(self, error: ServingError) -> str: + # What the rejected chat route's body said: its message, else the + # whole body. + return error.body.message or json.dumps(error.body.model_dump()) + + class OpenAIServer(_VideoRoutesMixin): @staticmethod @@ -871,7 +901,26 @@ async def validation_exception_handler(request, exc): # 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)) + if role_group_of(self.server_role) is RoleGroup.LLM: + # An attached frontend of a multi-frontend server forwards the + # launcher-owned routes, the Message Batches routes among them. + # They register before the routers so that they keep winning + # over the local handlers: FastAPI matches in registration order. + self._register_launcher_forwarding_routes() + messages_serving = AnthropicMessagesServing( + _AggregatedChatBackend(self)) + self.app.state.anthropic_serving_messages = messages_serving + # Transitional: RenderResources per request until it is built once. + self.app.state.anthropic_serving_count_tokens = ( + AnthropicCountTokensServing( + render_resources=lambda: RenderResources.from_server(self))) + self.app.state.anthropic_serving_batches = AnthropicBatchesServing( + runner=messages_serving.run_batched_message) + attach_routers( + self.app, + dataclasses.replace(ServerCapabilities.for_role(self.server_role), + harmony_model=(not self._is_visual_gen) + and self.use_harmony)) # Runs before the role dispatch for now. It moves after the dispatch # once the serving objects attach_router checks for are built per # role. @@ -892,7 +941,6 @@ async def validation_exception_handler(request, exc): self._init_embedding_batcher() self.register_embedding_routes() else: - self._register_launcher_forwarding_routes() self.register_routes() if self._collect_perf_metrics: @@ -1243,38 +1291,6 @@ def register_routes(self): "/v1/chat/completions", self.openai_chat if not self.use_harmony else self.chat_harmony, methods=["POST"]) - # Anthropic Messages API adapter (e.g. Claude Code as a client). - # Not supported together with the harmony chat path (gpt-oss). - if not self.use_harmony: - self.app.add_api_route("/v1/messages", - self.anthropic_messages, - methods=["POST"]) - # Claude Code calls this before most turns to decide whether to - # compact its context. Leaving it unregistered 404s every call, - # which the client cannot act on. - self.app.add_api_route("/v1/messages/count_tokens", - self.anthropic_count_tokens, - methods=["POST"]) - # Message Batches. Registered before /v1/messages/{...} would be, - # so the literal paths win; FastAPI matches in registration order. - self.app.add_api_route("/v1/messages/batches", - self.anthropic_create_batch, - methods=["POST"]) - self.app.add_api_route("/v1/messages/batches", - self.anthropic_list_batches, - methods=["GET"]) - self.app.add_api_route("/v1/messages/batches/{batch_id}", - self.anthropic_get_batch, - methods=["GET"]) - self.app.add_api_route("/v1/messages/batches/{batch_id}", - self.anthropic_delete_batch, - methods=["DELETE"]) - self.app.add_api_route("/v1/messages/batches/{batch_id}/cancel", - self.anthropic_cancel_batch, - methods=["POST"]) - self.app.add_api_route("/v1/messages/batches/{batch_id}/results", - self.anthropic_batch_results, - methods=["GET"]) self.app.add_api_route("/v1/responses", self.openai_responses, methods=["POST"]) @@ -1993,7 +2009,8 @@ 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) + result = await self.create_chat_completion( + request, RequestContext.from_http(raw_request)) if isinstance(result, ChatCompletionResponse): return JSONResponse(content=result.model_dump()) return StreamingResponse(content=result, @@ -2015,7 +2032,7 @@ async def openai_chat(self, request: ChatCompletionRequest, return self.create_error_response(str(e)) async def create_chat_completion( - self, request: ChatCompletionRequest, raw_request: Request + self, request: ChatCompletionRequest, ctx: RequestContext ) -> Union[ChatCompletionResponse, AsyncIterator[str]]: """Run one chat completion and return a typed result. @@ -2243,7 +2260,7 @@ def get_role() -> str: # calls in the model output. postproc_args.tools = (postproc_args.tools or []) + dynamic_tool_params - self._validate_internal_disagg_request(request, raw_request) + self._validate_internal_disagg_request(request, ctx) disaggregated_params = to_llm_disaggregated_params( request.disaggregated_params) @@ -2275,7 +2292,7 @@ def get_role() -> str: ) except ValidationError: # ValidatorIterator rejects extra fields; fall back to raw JSON. - raw_body = await raw_request.json() + raw_body = await ctx.json() raw_messages = raw_body.get("messages", []) conversation, mm_coroutines, mm_placeholder_counts, mm_item_order = parse_chat_messages_coroutines( raw_messages, @@ -2396,15 +2413,14 @@ def get_role() -> str: postproc_args=postproc_args, ) - trace_headers = (None if raw_request is None else - tracing.extract_trace_headers(raw_request.headers)) + trace_headers = (None if ctx is None else + tracing.extract_trace_headers(ctx.headers)) resolve_request_conversation_id( - request, None if raw_request is None else raw_request.headers) + request, None if ctx is None else ctx.headers) conversation_params = to_llm_conversation_params( request.conversation_params) - scheduling_params = self._get_scheduling_params( - request, raw_request) + scheduling_params = self._get_scheduling_params(request, ctx) generate_inputs = prompt preprocess_fn = getattr(self.generator, "preprocess", None) @@ -2431,7 +2447,7 @@ def get_role() -> str: priority=request.priority if request.priority is not None else DEFAULT_REQUEST_PRIORITY, ) - asyncio.create_task(self.await_disconnected(raw_request, promise)) + asyncio.create_task(self.await_disconnected(ctx, promise)) if not self.postproc_worker_enabled: postproc_args.tokenizer = self.tokenizer postproc_args.num_prompt_tokens = len(promise.prompt_token_ids) @@ -2440,10 +2456,10 @@ def get_role() -> str: return chat_stream_generator(promise, postproc_params, server=self, - raw_request=raw_request) + ctx=ctx) else: response = await self._create_chat_response( - promise, postproc_params, raw_request, disaggregated_params) + promise, postproc_params, ctx, disaggregated_params) # Context-only: optionally return prompt_token_ids as a base64 # int32 buffer (one string) so the disagg orchestrator relays it # without materializing the int list on its event loop. The encode @@ -2471,8 +2487,6 @@ def get_role() -> str: logger.error(traceback.format_exc()) raise serving_error(str(e)) from e - # -- Message Batches --------------------------------------------------- - async def openai_mm_encoder(self, request: ChatCompletionRequest, raw_request: Request) -> Response: From 07016ffd78931cf6661299ce4d2092e33be674b4 Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Tue, 3 Nov 2026 08:38:58 +0000 Subject: [PATCH 26/27] [TRTLLM-17263][test] Test the Anthropic serving class, the request context and the router on the real apps New tests: tests/unittest/serve/api/anthropic/test_serving.py (the serving class over a fake chat backend: the stream is reframed from the backend's own iterator, the backend's model and message reach the answer, each error kind keeps its status, the batch path hands the backend a non-streaming request and an internal context whose JSON is the converted request); tests/unittest/serve/_base/test_request_context.py (delegation to the HTTP request including its msgpack decoding, the internal context, the disconnect watcher returning at once for it); tests/unittest/serve/api/anthropic/test_api_router.py (the eight routes come from the router with the msgpack route class on the LLM roles only, none on a harmony model or another role, count_tokens decodes a msgpack body, attaching needs the serving objects). All three carry the cpu_only marker and run through the unittest/serve entry of the CPU list. Rewritten on the fixture factory: test_anthropic_routes.py posts to the real apps and their own validation handler instead of a mirror handler on bare servers, drives the batch and count-tokens routes through the serving objects on app.state, and hands the synthetic-request test to test_request_context.py; TestCountTokens in test_render_callers.py builds the count-tokens serving over a namespace. The chat-core and SSE-generator tests follow the renamed parameter; the multi-frontend tests set the Anthropic serving objects before attaching the routers. The object.__new__ ratchet drops to 2 and 1 sites for the two rewritten files; a layering contract keeps the Anthropic package from importing a server. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- .../llmapi/apps/test_anthropic_routes.py | 586 ++++++++---------- .../llmapi/apps/test_render_callers.py | 25 +- .../test_serve_multi_frontend_helpers.py | 19 +- .../test_serve_multi_frontend_route_gate.py | 16 +- .../serve/_base/test_request_context.py | 171 +++++ tests/unittest/serve/_base/test_sse.py | 12 +- .../serve/api/anthropic/test_api_router.py | 140 +++++ .../serve/api/anthropic/test_serving.py | 355 +++++++++++ tests/unittest/serve/test_layering.py | 12 + tests/unittest/serve/test_no_object_new.py | 12 - .../serve/test_openai_server_chat_core.py | 9 +- 11 files changed, 978 insertions(+), 379 deletions(-) create mode 100644 tests/unittest/serve/_base/test_request_context.py create mode 100644 tests/unittest/serve/api/anthropic/test_api_router.py create mode 100644 tests/unittest/serve/api/anthropic/test_serving.py diff --git a/tests/unittest/llmapi/apps/test_anthropic_routes.py b/tests/unittest/llmapi/apps/test_anthropic_routes.py index 272097e43f12..f292c347e9ed 100644 --- a/tests/unittest/llmapi/apps/test_anthropic_routes.py +++ b/tests/unittest/llmapi/apps/test_anthropic_routes.py @@ -12,7 +12,13 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""Route-level tests for the Anthropic Messages compatibility handlers.""" +"""Route-level tests for the Anthropic Messages compatibility handlers. + +The servers are built through the real constructors with the fake generator +from ``tests/test_common/serve/``, so every request goes through the real app: +its validation handler, its routers and the serving classes. Only the chat core +(aggregated) or the service entry point (disaggregated) is stubbed. +""" import asyncio import json @@ -22,12 +28,17 @@ import aiohttp import pytest -from fastapi import FastAPI -from fastapi.responses import JSONResponse, StreamingResponse -from fastapi.testclient import TestClient +from fastapi.responses import StreamingResponse +from test_common.serve.app_factory import AggConfig, DisaggConfig, isolated_serve_env +from test_common.serve.routes import effective_routes from tensorrt_llm.inputs.chat_template_guard import UnusedChatTemplateKwargsError from tensorrt_llm.serve._base import sse as sse_module +from tensorrt_llm.serve.api.anthropic.protocol import AnthropicCreateBatchRequest +from tensorrt_llm.serve.api.anthropic.serving import ( + AnthropicBatchesServing, + AnthropicCountTokensServing, +) from tensorrt_llm.serve.openai_disagg_server import OpenAIDisaggServer from tensorrt_llm.serve.openai_protocol import ( ChatCompletionRequest, @@ -40,6 +51,7 @@ UsageInfo, ) from tensorrt_llm.serve.openai_server import OpenAIServer +from tensorrt_llm.serve.render import RenderResources from tensorrt_llm.serve.render import chat as render_chat_module # These tests are CPU-only (no GPU, engine or sockets) and run in the @@ -47,6 +59,15 @@ pytestmark = pytest.mark.cpu_only MODEL = "test-model" +COUNTER_NAMES = ( + "total_requests", + "stream_requests", + "nonstream_requests", + "validation_exceptions", + "http_exceptions", + "internal_errors", + "total_responses", +) def _request(**overrides): @@ -82,14 +103,8 @@ def _chat_response(*, tool_arguments=None): ) -def _json_chat_response(*, tool_arguments=None, status_code=200): - return JSONResponse( - content=_chat_response(tool_arguments=tool_arguments).model_dump(), - status_code=status_code, - ) - - -def _streaming_chat_response(): +def _streaming_chat_result(server_kind): + """The chat pipeline's stream: SSE strings on the aggregated server, the worker's bytes on the disaggregated one.""" chunks = [ ChatCompletionStreamResponse( model=MODEL, @@ -111,40 +126,46 @@ def _streaming_chat_response(): usage=UsageInfo(prompt_tokens=3, completion_tokens=2, total_tokens=5), ), ] + frames = [f"data: {chunk.model_dump_json()}\n\n" for chunk in chunks] + ["data: [DONE]\n\n"] async def source(): - for chunk in chunks: - yield f"data: {chunk.model_dump_json()}\n\n" - yield "data: [DONE]\n\n" + for frame in frames: + yield frame if server_kind == "standard" else frame.encode() - return StreamingResponse(source(), media_type="text/event-stream") + return source() -def _make_route_client(server_kind, openai_response): - app = FastAPI() - # Body validation fails before the route function runs, so without the - # server's own handler these tests would see FastAPI's 422 instead of the - # 400 + Anthropic envelope a client actually gets. - _anthropic_validation_handler(app) +def _counters(): + return SimpleNamespace( + **{name: SimpleNamespace(inc=Mock()) for name in COUNTER_NAMES}, + queue_latency_seconds=SimpleNamespace(observe=Mock()), + ) + + +def _route_client(env, server_kind, result): + """A client over the real app of one server, its chat pipeline stubbed to answer ``result``. + + Returns the client and the stub; the stub's first positional argument is + the chat request the Anthropic request was converted to. + """ + core = AsyncMock(return_value=result) if server_kind == "standard": - server = object.__new__(OpenAIServer) - server.model = MODEL - backend = AsyncMock(return_value=openai_response) - server.openai_chat = backend - else: - server = object.__new__(OpenAIDisaggServer) - server._service = SimpleNamespace(openai_chat_completion=object()) - backend = AsyncMock(return_value=openai_response) - server._wrap_entry_point = Mock(return_value=backend) - app.add_api_route("/v1/messages", server.anthropic_messages, methods=["POST"]) - return TestClient(app), backend + server = env.build_openai_server(AggConfig(model_name=MODEL)) + server.create_chat_completion = core + return env.client(server.app), core + server = env.build_disagg_server(DisaggConfig()) + server._perf_metrics_collector = _counters() + server._service = SimpleNamespace(openai_chat_completion=core) + # The disaggregated lifespan connects to workers; these tests never run it. + return env.client(server.app, lifespan=False), core @pytest.mark.parametrize("server_kind", ["standard", "disagg"]) -def test_messages_route_converts_nonstream_response(server_kind): - client, backend = _make_route_client(server_kind, _json_chat_response()) +def test_messages_route_converts_nonstream_response(server_kind, tmp_path): + with isolated_serve_env(tmp_path) as env: + client, core = _route_client(env, server_kind, _chat_response()) - response = client.post("/v1/messages", json=_request()) + response = client.post("/v1/messages", json=_request()) assert response.status_code == 200 assert response.json() | {"id": "ignored"} == { @@ -156,47 +177,49 @@ def test_messages_route_converts_nonstream_response(server_kind): "stop_reason": "end_turn", "usage": {"input_tokens": 3, "output_tokens": 2}, } - chat_request = backend.await_args.args[0] + chat_request = core.await_args.args[0] assert chat_request.model == MODEL assert chat_request.max_completion_tokens == 64 assert not chat_request.stream @pytest.mark.parametrize("server_kind", ["standard", "disagg"]) -def test_messages_route_accepts_thinking_display(server_kind): +def test_messages_route_accepts_thinking_display(server_kind, tmp_path): """The thinking config Claude Code sends is served, not answered with a 400.""" - client, backend = _make_route_client(server_kind, _json_chat_response()) + with isolated_serve_env(tmp_path) as env: + client, core = _route_client(env, server_kind, _chat_response()) - response = client.post( - "/v1/messages", - json=_request(max_tokens=4096, thinking={"type": "adaptive", "display": "updates"}), - ) + response = client.post( + "/v1/messages", + json=_request(max_tokens=4096, thinking={"type": "adaptive", "display": "updates"}), + ) assert response.status_code == 200 - assert backend.await_args.args[0].chat_template_kwargs["enable_thinking"] is True + assert core.await_args.args[0].chat_template_kwargs["enable_thinking"] is True @pytest.mark.parametrize("server_kind", ["standard", "disagg"]) -def test_messages_route_rejects_anthropic_server_tools(server_kind): - client, backend = _make_route_client(server_kind, _json_chat_response()) - - response = client.post( - "/v1/messages", - json=_request( - tools=[ - { - "name": "web_search", - "type": "web_search_20250305", - } - ] - ), - ) +def test_messages_route_rejects_anthropic_server_tools(server_kind, tmp_path): + with isolated_serve_env(tmp_path) as env: + client, core = _route_client(env, server_kind, _chat_response()) + + response = client.post( + "/v1/messages", + json=_request( + tools=[ + { + "name": "web_search", + "type": "web_search_20250305", + } + ] + ), + ) assert response.status_code == 400 assert response.json()["type"] == "error" assert response.json()["error"]["type"] == "invalid_request_error" assert "server tool" in response.json()["error"]["message"] - backend.assert_not_awaited() + core.assert_not_awaited() @pytest.mark.parametrize( @@ -214,7 +237,7 @@ def test_messages_route_rejects_anthropic_server_tools(server_kind): ), ], ) -def test_malformed_known_block_is_a_400_against_its_own_type(block, tag, field): +def test_malformed_known_block_is_a_400_against_its_own_type(block, tag, field, tmp_path): """A known type with bad fields must fail as that type, not as the catch-all. The union is tagged on `type`, so the block is checked against the model the @@ -223,14 +246,17 @@ def test_malformed_known_block_is_a_400_against_its_own_type(block, tag, field): the adapter's dispatch on `type`, and raise AttributeError on an attribute the block never had -- a 500 for a plain client error. - Errors carry `loc` as a tuple, so the tag and field appear adjacent. + Errors carry `loc` as a tuple, so the tag and field appear adjacent. The + 400 and the Anthropic envelope come from the server's own validation + handler, which the real app carries. """ - client, backend = _make_route_client("standard", _json_chat_response()) + with isolated_serve_env(tmp_path) as env: + client, core = _route_client(env, "standard", _chat_response()) - response = client.post( - "/v1/messages", - json=_request(messages=[{"role": "user", "content": [block]}]), - ) + response = client.post( + "/v1/messages", + json=_request(messages=[{"role": "user", "content": [block]}]), + ) assert response.status_code == 400 assert response.json()["error"]["type"] == "invalid_request_error" @@ -238,33 +264,37 @@ def test_malformed_known_block_is_a_400_against_its_own_type(block, tag, field): assert f"'{tag}', '{field}'" in message # The catch-all must not have absorbed it; that is the regression itself. assert "'unknown'," not in message - backend.assert_not_awaited() + core.assert_not_awaited() -def test_unknown_block_type_still_reaches_the_adapter(): +def test_unknown_block_type_still_reaches_the_adapter(tmp_path): """The catch-all exists for types this server does not model; keep it working. Tagging the union must not close it: an unmodelled type has to validate and be named by the adapter, rather than returning a wall of failed members. """ - client, backend = _make_route_client("standard", _json_chat_response()) + with isolated_serve_env(tmp_path) as env: + client, core = _route_client(env, "standard", _chat_response()) - response = client.post( - "/v1/messages", - json=_request(messages=[{"role": "user", "content": [{"type": "document", "source": {}}]}]), - ) + response = client.post( + "/v1/messages", + json=_request( + messages=[{"role": "user", "content": [{"type": "document", "source": {}}]}] + ), + ) assert response.status_code == 400 assert response.json()["error"]["type"] == "invalid_request_error" assert "document" in response.json()["error"]["message"] - backend.assert_not_awaited() + core.assert_not_awaited() @pytest.mark.parametrize("server_kind", ["standard", "disagg"]) -def test_messages_route_hides_invalid_generated_tool_arguments(server_kind): - client, _ = _make_route_client(server_kind, _json_chat_response(tool_arguments="{not-json")) +def test_messages_route_hides_invalid_generated_tool_arguments(server_kind, tmp_path): + with isolated_serve_env(tmp_path) as env: + client, _ = _route_client(env, server_kind, _chat_response(tool_arguments="{not-json")) - response = client.post("/v1/messages", json=_request()) + response = client.post("/v1/messages", json=_request()) assert response.status_code == 500 assert response.json() == { @@ -274,10 +304,11 @@ def test_messages_route_hides_invalid_generated_tool_arguments(server_kind): @pytest.mark.parametrize("server_kind", ["standard", "disagg"]) -def test_messages_route_reframes_streaming_response(server_kind): - client, backend = _make_route_client(server_kind, _streaming_chat_response()) +def test_messages_route_reframes_streaming_response(server_kind, tmp_path): + with isolated_serve_env(tmp_path) as env: + client, core = _route_client(env, server_kind, _streaming_chat_result(server_kind)) - response = client.post("/v1/messages", json=_request(stream=True)) + response = client.post("/v1/messages", json=_request(stream=True)) assert response.status_code == 200 assert response.headers["content-type"].startswith("text/event-stream") @@ -285,7 +316,7 @@ def test_messages_route_reframes_streaming_response(server_kind): assert "event: content_block_delta\n" in response.text assert '"text":"hello"' in response.text assert response.text.rstrip().endswith('event: message_stop\ndata: {"type":"message_stop"}') - assert backend.await_args.args[0].stream + assert core.await_args.args[0].stream class _StreamedGeneration: @@ -451,61 +482,19 @@ async def run(): assert "engine failed" in server_logger.error.call_args.args[0] -@pytest.fixture -def prometheus_multiproc_dir(tmp_path, monkeypatch): - """PROMETHEUS_MULTIPROC_DIR for one test. +def _paths(app) -> set: + return {path for path, _ in effective_routes(app)} - prometheus_client picks and caches its value class the first time a metric - is built with the variable set, so the class is restored as well as the - environment; otherwise later metrics in the process fail without a - directory. - """ - import prometheus_client.values - - monkeypatch.setenv("PROMETHEUS_MULTIPROC_DIR", str(tmp_path)) - saved_value_class = prometheus_client.values.ValueClass - try: - yield tmp_path - finally: - prometheus_client.values.ValueClass = saved_value_class - - -def test_standard_and_disagg_register_messages_route(prometheus_multiproc_dir): - # register_routes mounts a prometheus multiprocess collector, which needs a - # real PROMETHEUS_MULTIPROC_DIR; object.__new__ skips the startup that - # creates one. - - standard = object.__new__(OpenAIServer) - standard.app = FastAPI() - standard.generator = SimpleNamespace( - _executor=SimpleNamespace(resource_governor_queue=None), - args=SimpleNamespace(return_perf_metrics=False), - ) - standard.use_harmony = False - # register_routes() reads these to decide whether to mount the optional - # control endpoints. Building the server with object.__new__ skips - # __init__, so anything register_routes() touches has to be supplied. - standard._enable_runtime_control_endpoints = False - standard._enable_rl_control_endpoints = False - standard.register_routes() - - disagg = object.__new__(OpenAIDisaggServer) - disagg.app = FastAPI() - disagg._service = SimpleNamespace( - openai_completion=AsyncMock(), - openai_chat_completion=AsyncMock(), - anthropic_count_tokens=AsyncMock(), - openai_responses=AsyncMock(), - ) - disagg._perf_metrics_collector = SimpleNamespace(get_perf_metrics=AsyncMock()) - disagg._disagg_cluster_storage = None - # register_routes reads this; the server is built with object.__new__ here, - # so every attribute it touches has to be stubbed explicitly. - disagg._coordinator = None - disagg.register_routes() - - for server in (standard, disagg): - paths = {route.path for route in server.app.routes} + +def test_standard_and_disagg_register_messages_route(tmp_path): + # One server per environment: both register collectors in the default + # Prometheus registry, which the environment restores on exit. + with isolated_serve_env(tmp_path / "standard") as env: + standard_paths = _paths(env.build_openai_server(AggConfig()).app) + with isolated_serve_env(tmp_path / "disagg") as env: + disagg_paths = _paths(env.build_disagg_server(DisaggConfig()).app) + + for paths in (standard_paths, disagg_paths): assert "/v1/messages" in paths # Claude Code calls count_tokens before most turns to size its context. @@ -513,42 +502,22 @@ def test_standard_and_disagg_register_messages_route(prometheus_multiproc_dir): # so registration is asserted separately from /v1/messages -- and on both # servers, because a client pointed at a disaggregated deployment cannot # tell the difference and must not lose the endpoint. - for server in (standard, disagg): - paths = {route.path for route in server.app.routes} + for paths in (standard_paths, disagg_paths): assert "/v1/messages/count_tokens" in paths - disagg_paths = {route.path for route in disagg.app.routes} assert "/v1/responses" in disagg_paths assert "/v1/models" in disagg_paths -def _batch_client(runner=None): - """A TestClient over just the batch routes, with a stubbed item runner.""" - from tensorrt_llm.serve.api.anthropic.batches import AnthropicBatchStore +def _batch_client(env, runner=None): + """A client over the real app, with the batch serving's item runner stubbed.""" async def ok(request): return "succeeded", {"type": "message", "role": "assistant", "content": []} - app = FastAPI() - _anthropic_validation_handler(app) - server = object.__new__(OpenAIServer) - server.model = MODEL - server._anthropic_batch_store = AnthropicBatchStore(runner=runner or ok) - app.add_api_route("/v1/messages/batches", server.anthropic_create_batch, methods=["POST"]) - app.add_api_route("/v1/messages/batches", server.anthropic_list_batches, methods=["GET"]) - app.add_api_route( - "/v1/messages/batches/{batch_id}", server.anthropic_get_batch, methods=["GET"] - ) - app.add_api_route( - "/v1/messages/batches/{batch_id}", server.anthropic_delete_batch, methods=["DELETE"] - ) - app.add_api_route( - "/v1/messages/batches/{batch_id}/cancel", server.anthropic_cancel_batch, methods=["POST"] - ) - app.add_api_route( - "/v1/messages/batches/{batch_id}/results", server.anthropic_batch_results, methods=["GET"] - ) - return TestClient(app) + server = env.build_openai_server(AggConfig(model_name=MODEL)) + server.app.state.anthropic_serving_batches = AnthropicBatchesServing(runner=runner or ok) + return env.client(server.app) def _batch_body(*custom_ids): @@ -567,22 +536,10 @@ def _batch_body(*custom_ids): } -def test_batch_routes_are_registered_on_the_standard_server(): - standard = object.__new__(OpenAIServer) - standard.app = FastAPI() - standard.generator = SimpleNamespace( - _executor=SimpleNamespace(resource_governor_queue=None), - args=SimpleNamespace(return_perf_metrics=False), - ) - standard.use_harmony = False - # register_routes() reads these to decide whether to mount the optional - # control endpoints. Building the server with object.__new__ skips - # __init__, so anything register_routes() touches has to be supplied. - standard._enable_runtime_control_endpoints = False - standard._enable_rl_control_endpoints = False - standard.register_routes() - - paths = {route.path for route in standard.app.routes} +def test_batch_routes_are_registered_on_the_standard_server(tmp_path): + with isolated_serve_env(tmp_path) as env: + paths = _paths(env.build_openai_server(AggConfig()).app) + for path in ( "/v1/messages/batches", "/v1/messages/batches/{batch_id}", @@ -592,9 +549,10 @@ def test_batch_routes_are_registered_on_the_standard_server(): assert path in paths, f"{path} not registered" -def test_create_batch_returns_an_in_progress_batch(): - client = _batch_client() - response = client.post("/v1/messages/batches", json=_batch_body("a", "b")) +def test_create_batch_returns_an_in_progress_batch(tmp_path): + with isolated_serve_env(tmp_path) as env: + client = _batch_client(env) + response = client.post("/v1/messages/batches", json=_batch_body("a", "b")) assert response.status_code == 200 body = response.json() @@ -604,18 +562,20 @@ def test_create_batch_returns_an_in_progress_batch(): assert body["request_counts"]["processing"] + body["request_counts"]["succeeded"] == 2 -def test_create_batch_rejects_duplicate_custom_ids(): - client = _batch_client() - response = client.post("/v1/messages/batches", json=_batch_body("same", "same")) +def test_create_batch_rejects_duplicate_custom_ids(tmp_path): + with isolated_serve_env(tmp_path) as env: + client = _batch_client(env) + response = client.post("/v1/messages/batches", json=_batch_body("same", "same")) assert response.status_code == 400 assert response.json()["error"]["type"] == "invalid_request_error" -def test_unknown_batch_says_it_is_not_found_and_explains_why(): +def test_unknown_batch_says_it_is_not_found_and_explains_why(tmp_path): """The in-memory limitation is the likeliest cause, so name it.""" - client = _batch_client() - response = client.get("/v1/messages/batches/msgbatch_nope") + with isolated_serve_env(tmp_path) as env: + client = _batch_client(env) + response = client.get("/v1/messages/batches/msgbatch_nope") assert response.status_code == 404 error = response.json()["error"] @@ -626,12 +586,9 @@ def test_unknown_batch_says_it_is_not_found_and_explains_why(): def test_results_of_an_unfinished_batch_are_400_not_404(): """404 would tell a polling client the batch never existed. - Driven directly rather than through TestClient: each TestClient call runs - in its own portal, so the batch's background task is not guaranteed to - still be in flight by the time the second request arrives, and the batch - would have already ended. + Driven on the serving class directly rather than through a client, so the + batch's background task is still in flight when the results are asked for. """ - from tensorrt_llm.serve.api.anthropic.batches import AnthropicBatchStore from tensorrt_llm.serve.api.anthropic.protocol import AnthropicBatchRequestItem async def scenario(): @@ -641,8 +598,7 @@ async def never_finishes(request): await blocked.wait() return "succeeded", {"type": "message"} - server = object.__new__(OpenAIServer) - server._anthropic_batch_store = AnthropicBatchStore(runner=never_finishes) + serving = AnthropicBatchesServing(runner=never_finishes) item = AnthropicBatchRequestItem( custom_id="a", params={ @@ -651,15 +607,15 @@ async def never_finishes(request): "messages": [{"role": "user", "content": "hi"}], }, ) - batch = server._anthropic_batch_store.create([item]) + batch = serving._batch_store().create([item]) await asyncio.sleep(0) # let the worker start and block - response = await server.anthropic_batch_results(batch.id) + response = await serving.anthropic_batch_results(batch.id) assert response.status_code == 400 assert "still processing" in json.loads(response.body)["error"]["message"] # An id that was never issued is a genuine 404, and must stay distinct. - missing = await server.anthropic_batch_results("msgbatch_nope") + missing = await serving.anthropic_batch_results("msgbatch_nope") assert missing.status_code == 404 blocked.set() @@ -715,29 +671,32 @@ async def never_finishes(request): asyncio.run(scenario()) -def test_finished_batch_serves_ndjson_results(): - client = _batch_client() - created = client.post("/v1/messages/batches", json=_batch_body("a", "b")).json() +def test_finished_batch_serves_ndjson_results(tmp_path): + with isolated_serve_env(tmp_path) as env: + client = _batch_client(env) + created = client.post("/v1/messages/batches", json=_batch_body("a", "b")).json() - for _ in range(200): - if ( - client.get(f"/v1/messages/batches/{created['id']}").json()["processing_status"] - == "ended" - ): - break - time.sleep(0.01) + for _ in range(200): + if ( + client.get(f"/v1/messages/batches/{created['id']}").json()["processing_status"] + == "ended" + ): + break + time.sleep(0.01) + + response = client.get(f"/v1/messages/batches/{created['id']}/results") - response = client.get(f"/v1/messages/batches/{created['id']}/results") assert response.status_code == 200 assert response.headers["content-type"].startswith("application/x-ndjson") lines = [json.loads(line) for line in response.text.splitlines()] assert {line["custom_id"] for line in lines} == {"a", "b"} -def test_list_batches_returns_a_page(): - client = _batch_client() - client.post("/v1/messages/batches", json=_batch_body("a")) - body = client.get("/v1/messages/batches").json() +def test_list_batches_returns_a_page(tmp_path): + with isolated_serve_env(tmp_path) as env: + client = _batch_client(env) + client.post("/v1/messages/batches", json=_batch_body("a")) + body = client.get("/v1/messages/batches").json() assert isinstance(body["data"], list) and body["data"] assert body["first_id"] == body["data"][0]["id"] @@ -841,125 +800,104 @@ def encode(self, text, add_special_tokens=True, **kwargs): return text.split() -def _anthropic_validation_handler(app): - """Mirror the server's handler so boundary rejections keep the envelope. +def _count_tokens_client(env, resources): + """A client over the real app, counting over ``resources`` (what the renderer reads off a server). - Model-level validation fails before the route function runs, so without - this a bare test app answers with FastAPI's 422 shape while the real - server answers 400 with the Anthropic envelope -- the test would pin the - fixture's behaviour instead of the product's. + Transitional: the resources get built directly once the server assembly + stops reading them off the server per request. """ - from fastapi.exceptions import RequestValidationError - - from tensorrt_llm.serve.api.anthropic.adapter import anthropic_error_response - - @app.exception_handler(RequestValidationError) - async def _handler(request, exc): # noqa: ANN001 - if request.url.path.startswith("/v1/messages"): - return anthropic_error_response(str(exc), "invalid_request_error", 400) - raise exc + server = env.build_openai_server(AggConfig(model_name=MODEL)) + server.app.state.anthropic_serving_count_tokens = AnthropicCountTokensServing( + render_resources=lambda: RenderResources.from_server(resources) + ) + return env.client(server.app) -def _count_tokens_client(): - app = FastAPI() - _anthropic_validation_handler(app) - server = object.__new__(OpenAIServer) - server.model = MODEL - server.tokenizer = _FakeTokenizer() - app.add_api_route("/v1/messages/count_tokens", server.anthropic_count_tokens, methods=["POST"]) - return TestClient(app) +def _count_resources(chat_template=None): + return SimpleNamespace(tokenizer=_FakeTokenizer(), chat_template=chat_template) -def test_count_tokens_route_returns_a_count(): - client = _count_tokens_client() +def test_count_tokens_route_returns_a_count(tmp_path): + with isolated_serve_env(tmp_path) as env: + client = _count_tokens_client(env, _count_resources()) - response = client.post( - "/v1/messages/count_tokens", - json={ - "model": MODEL, - "messages": [{"role": "user", "content": "one two three"}], - }, - ) + response = client.post( + "/v1/messages/count_tokens", + json={ + "model": MODEL, + "messages": [{"role": "user", "content": "one two three"}], + }, + ) assert response.status_code == 200 assert response.json() == {"input_tokens": 3} -def test_count_tokens_counts_the_rendered_prompt_not_the_raw_messages(): +def test_count_tokens_counts_the_rendered_prompt_not_the_raw_messages(tmp_path): """Tools inflate the real prompt, so they have to inflate the count.""" - client = _count_tokens_client() body = { "model": MODEL, "messages": [{"role": "user", "content": "one two three"}], } - - bare = client.post("/v1/messages/count_tokens", json=body).json() - with_tools = client.post( - "/v1/messages/count_tokens", - json={ - **body, - "tools": [ - { - "name": "get_weather", - "input_schema": {"type": "object"}, - } - ], - }, - ).json() + with isolated_serve_env(tmp_path) as env: + client = _count_tokens_client(env, _count_resources()) + + bare = client.post("/v1/messages/count_tokens", json=body).json() + with_tools = client.post( + "/v1/messages/count_tokens", + json={ + **body, + "tools": [ + { + "name": "get_weather", + "input_schema": {"type": "object"}, + } + ], + }, + ).json() assert with_tools["input_tokens"] > bare["input_tokens"] -def test_count_tokens_honours_the_server_chat_template(): +def test_count_tokens_honours_the_server_chat_template(tmp_path): """A server started with --chat_template must count the prompt it renders. - The renderer only consults request.chat_template; openai_chat resolves + The renderer only consults request.chat_template; the chat route resolves `request.chat_template or self.chat_template`. Without the same resolution here, count_tokens reports a prompt the server never builds - the exact drift this endpoint exists to remove. """ - app = FastAPI() - server = object.__new__(OpenAIServer) - server.model = MODEL - server.tokenizer = _FakeTokenizer() # A template that measurably inflates the rendered prompt. - server.chat_template = "SERVER-TEMPLATE-MARKER {{ messages }}" - app.add_api_route("/v1/messages/count_tokens", server.anthropic_count_tokens, methods=["POST"]) - client = TestClient(app) - - with_template = client.post( - "/v1/messages/count_tokens", - json={ - "model": MODEL, - "messages": [{"role": "user", "content": "one two three"}], - }, - ).json() + resources = _count_resources(chat_template="SERVER-TEMPLATE-MARKER {{ messages }}") + body = { + "model": MODEL, + "messages": [{"role": "user", "content": "one two three"}], + } + with isolated_serve_env(tmp_path) as env: + client = _count_tokens_client(env, resources) - server.chat_template = None - without = client.post( - "/v1/messages/count_tokens", - json={ - "model": MODEL, - "messages": [{"role": "user", "content": "one two three"}], - }, - ).json() + with_template = client.post("/v1/messages/count_tokens", json=body).json() + + resources.chat_template = None + without = client.post("/v1/messages/count_tokens", json=body).json() assert with_template["input_tokens"] != without["input_tokens"], ( "server chat_template had no effect on the count, so it was ignored" ) -def test_count_tokens_rejects_a_malformed_request(): - client = _count_tokens_client() +def test_count_tokens_rejects_a_malformed_request(tmp_path): + with isolated_serve_env(tmp_path) as env: + client = _count_tokens_client(env, _count_resources()) - response = client.post( - "/v1/messages/count_tokens", - json={ - "model": MODEL, - "messages": [{"role": "user", "content": "hi"}], - "thinking": {"type": "bogus"}, - }, - ) + response = client.post( + "/v1/messages/count_tokens", + json={ + "model": MODEL, + "messages": [{"role": "user", "content": "hi"}], + "thinking": {"type": "bogus"}, + }, + ) assert response.status_code == 400 assert response.json()["error"]["type"] == "invalid_request_error" @@ -1191,29 +1129,28 @@ def test_create_batch_at_capacity_is_429_not_400(): instead of surfacing a hard failure to the user. This pins the status mapping only. Whether the store actually refuses at - capacity is store policy, covered in test_anthropic_batches.py -- driving a - real store to its limit through TestClient is not possible anyway, since - each request runs in its own event loop and a batch's task does not - survive between them. + capacity is store policy, covered in test_anthropic_batches.py. Driven on + the serving class directly: the mapping needs no server. """ from tensorrt_llm.serve.api.anthropic.batches import BatchStoreFullError - app = FastAPI() - server = object.__new__(OpenAIServer) - server.model = MODEL + async def ok(request): + return "succeeded", {"type": "message"} def full(_requests): raise BatchStoreFullError( "the server is holding 100 unfinished batches (limit 100); retry once one of them ends" ) - server._anthropic_batch_store = SimpleNamespace(create=full) - app.add_api_route("/v1/messages/batches", server.anthropic_create_batch, methods=["POST"]) + serving = AnthropicBatchesServing(runner=ok) + serving._anthropic_batch_store = SimpleNamespace(create=full) - response = TestClient(app).post("/v1/messages/batches", json=_batch_body("a")) + response = asyncio.run( + serving.anthropic_create_batch(AnthropicCreateBatchRequest(**_batch_body("a"))) + ) assert response.status_code == 429 - body = response.json() + body = json.loads(response.body) assert body["type"] == "error" assert body["error"]["type"] == "rate_limit_error" # The message has to say what to do about it; an "invalid request" shape @@ -1221,46 +1158,17 @@ def full(_requests): assert "retry" in body["error"]["message"].lower() -def test_await_disconnected_returns_immediately_for_a_synthetic_request(): - """A batched request has no socket, so the watchdog must not poll on it. - - is_disconnected() can never become true for a synthesised request, so - without an explicit opt-out this coroutine loops at 1Hz for the life of - the process -- one live task per batched request, each holding its Request - and RequestOutput. - """ - server = object.__new__(OpenAIServer) - raw_request = OpenAIServer._synthetic_request( - { - "model": MODEL, - "max_tokens": 8, - "messages": [{"role": "user", "content": "hi"}], - } - ) - promise = SimpleNamespace(finished=False, request_id="req", abort=Mock()) - - async def run(): - # A real socket-backed request would never let this return; a bounded - # wait keeps the failure a clear timeout rather than a hung suite. - await asyncio.wait_for( - OpenAIServer.await_disconnected(server, raw_request, promise), - timeout=5.0, - ) - - asyncio.run(run()) - # Returning early must not be mistaken for a disconnect. - promise.abort.assert_not_called() - - @pytest.mark.parametrize("limit", [0, -1, 101]) -def test_list_batches_rejects_out_of_range_limit(limit): +def test_list_batches_rejects_out_of_range_limit(limit, tmp_path): """Clamping hid both a client error and a surprising page size. limit=0 became 1 rather than an error, and the page was bounded by the retention limit -- how many batches are kept, which has nothing to do with - how many fit on a page. + how many fit on a page. The 400 and the envelope come from the server's own + validation handler. """ - response = _batch_client().get(f"/v1/messages/batches?limit={limit}") + with isolated_serve_env(tmp_path) as env: + response = _batch_client(env).get(f"/v1/messages/batches?limit={limit}") assert response.status_code == 400, response.text assert response.json()["error"]["type"] == "invalid_request_error" diff --git a/tests/unittest/llmapi/apps/test_render_callers.py b/tests/unittest/llmapi/apps/test_render_callers.py index 752b8794b5cf..fd2d92dd6939 100644 --- a/tests/unittest/llmapi/apps/test_render_callers.py +++ b/tests/unittest/llmapi/apps/test_render_callers.py @@ -664,22 +664,25 @@ def _client( from fastapi import FastAPI from fastapi.testclient import TestClient - from tensorrt_llm.serve.openai_server import OpenAIServer + from tensorrt_llm.serve.api.anthropic.serving import AnthropicCountTokensServing - app = FastAPI() - server = object.__new__(OpenAIServer) - server.model = "m" - server.tokenizer = tokenizer - server.chat_template = chat_template + # What RenderResources.from_server reads off a server, as a namespace. + # Transitional: the resources get built directly once the server + # assembly stops reading them off the server per request. + ns = SimpleNamespace(model="m", tokenizer=tokenizer, chat_template=chat_template) if use_harmony: - server.use_harmony = True - server.model_config = SimpleNamespace(model_type="gpt_oss") + ns.use_harmony = True + ns.model_config = SimpleNamespace(model_type="gpt_oss") elif model_type is not None: - server.model_config = SimpleNamespace(model_type=model_type) + ns.model_config = SimpleNamespace(model_type=model_type) if input_processor is not None: - server.generator = SimpleNamespace(input_processor=input_processor) + ns.generator = SimpleNamespace(input_processor=input_processor) + serving = AnthropicCountTokensServing( + render_resources=lambda: RenderResources.from_server(ns) + ) + app = FastAPI() app.add_api_route( - "/v1/messages/count_tokens", server.anthropic_count_tokens, methods=["POST"] + "/v1/messages/count_tokens", serving.anthropic_count_tokens, methods=["POST"] ) return TestClient(app) diff --git a/tests/unittest/llmapi/test_serve_multi_frontend_helpers.py b/tests/unittest/llmapi/test_serve_multi_frontend_helpers.py index add864bb829d..6bb5d3e5fbc1 100644 --- a/tests/unittest/llmapi/test_serve_multi_frontend_helpers.py +++ b/tests/unittest/llmapi/test_serve_multi_frontend_helpers.py @@ -37,6 +37,11 @@ from tensorrt_llm.serve import _frontend_processes, openai_server from tensorrt_llm.serve._base.routing import ServerCapabilities, attach_routers from tensorrt_llm.serve._frontend_processes import _watch_attached_frontends +from tensorrt_llm.serve.api.anthropic.serving import ( + AnthropicBatchesServing, + AnthropicCountTokensServing, + AnthropicMessagesServing, +) from tensorrt_llm.serve.api.ops.system.serving import SystemServing from tensorrt_llm.serve.multi_frontend import ( FORWARDED_ROUTES, @@ -113,12 +118,18 @@ async def lifespan(_: FastAPI): server._enable_runtime_control_endpoints = False server.resource_governor = None server._launcher_forwarder = forwarder - # The order OpenAIServer.__init__ registers them in: the system router - # (/health, /version, /v1/models) first, then the launcher forwarding - # routes, then the local table. + # The order OpenAIServer.__init__ registers them in: the launcher + # forwarding routes first, then the routers (system, Anthropic), then the + # local table. The routes are enumerated and forwarded, never served + # locally, so the Anthropic serving objects get unused collaborators. server.app.state.ops_serving_system = SystemServing(model="m", health_probe=lambda: True) - attach_routers(server.app, ServerCapabilities.for_role(None)) + server.app.state.anthropic_serving_messages = AnthropicMessagesServing(backend=None) + server.app.state.anthropic_serving_count_tokens = AnthropicCountTokensServing( + render_resources=lambda: None + ) + server.app.state.anthropic_serving_batches = AnthropicBatchesServing(runner=None) server._register_launcher_forwarding_routes() + attach_routers(server.app, ServerCapabilities.for_role(None)) server.register_routes() return server diff --git a/tests/unittest/llmapi/test_serve_multi_frontend_route_gate.py b/tests/unittest/llmapi/test_serve_multi_frontend_route_gate.py index f419d5458da0..b682e4f6b9e4 100644 --- a/tests/unittest/llmapi/test_serve_multi_frontend_route_gate.py +++ b/tests/unittest/llmapi/test_serve_multi_frontend_route_gate.py @@ -28,6 +28,11 @@ from test_common.serve.routes import effective_routes from tensorrt_llm.serve._base.routing import ServerCapabilities, attach_routers +from tensorrt_llm.serve.api.anthropic.serving import ( + AnthropicBatchesServing, + AnthropicCountTokensServing, + AnthropicMessagesServing, +) from tensorrt_llm.serve.api.ops.system.serving import SystemServing from tensorrt_llm.serve.multi_frontend import FORWARDED_ROUTES from tensorrt_llm.serve.openai_server import OpenAIServer @@ -76,9 +81,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. + # As OpenAIServer.__init__ does: the routers (the system router for every + # role, the Anthropic router for the LLM role) are attached before the + # role's own table. The routes are enumerated, never called, so the + # serving objects get collaborators that are never used. server.app.state.ops_serving_system = SystemServing(model="m", health_probe=lambda: True) + server.app.state.anthropic_serving_messages = AnthropicMessagesServing(backend=None) + server.app.state.anthropic_serving_count_tokens = AnthropicCountTokensServing( + render_resources=lambda: None + ) + server.app.state.anthropic_serving_batches = AnthropicBatchesServing(runner=None) attach_routers(server.app, ServerCapabilities.for_role(None)) server.register_routes() # effective_routes walks included routers, which FastAPI 0.139 keeps as diff --git a/tests/unittest/serve/_base/test_request_context.py b/tests/unittest/serve/_base/test_request_context.py new file mode 100644 index 000000000000..dfecf34ab828 --- /dev/null +++ b/tests/unittest/serve/_base/test_request_context.py @@ -0,0 +1,171 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""``RequestContext``: what the chat core reads from a request, on the HTTP and the internal path.""" + +import asyncio +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +from starlette.requests import Request +from test_common.serve.app_factory import AggConfig + +from tensorrt_llm.serve._base.msgspec_route import _MsgspecRequest +from tensorrt_llm.serve._base.request_context import RequestContext +from tensorrt_llm.serve.openai_client import MSGPACK_HEADERS, _msgpack_encoder + +pytestmark = pytest.mark.cpu_only + +BODY = {"model": "m", "messages": [{"role": "user", "content": "hi"}]} + + +def _http_request(body: bytes = b"", headers=(), state=None) -> Request: + async def receive(): + return {"type": "http.request", "body": body, "more_body": False} + + scope = { + "type": "http", + "method": "POST", + "path": "/v1/messages", + "headers": [(k.lower().encode(), v.encode()) for k, v in headers], + "state": {"perf_metrics_records": []} if state is None else state, + } + return Request(scope, receive) + + +# -- from_http: delegation to the request FastAPI passed ---------------------- + + +def test_from_http_shares_the_request_state(): + request = _http_request() + ctx = RequestContext.from_http(request) + + assert ctx.http_request is request + assert ctx.state is request.state + assert ctx.perf_metrics_records is request.state.perf_metrics_records + # A write through the context is a write on the request's state, which the + # perf-metrics middleware reads when the response goes out. + ctx.server_first_token_time = 2.0 + assert request.state.server_first_token_time == 2.0 + assert ctx.server_first_token_time == 2.0 + assert bool(ctx), "the chat core tests `if raw_request and ...`; a context is truthy" + + +def test_from_http_reads_headers_and_client_from_the_request(): + request = _http_request(headers=(("X-Conversation-Id", "c1"),)) + ctx = RequestContext.from_http(request) + + assert ctx.headers["x-conversation-id"] == "c1" + assert ctx.client is None # the scope above names no client + assert ctx.server_arrival_time is None + assert ctx.no_client_connection is False + + +def test_from_http_reads_nothing_at_construction(): + # Tests hand the chat core a bare namespace as the request; wrapping it + # must not touch any attribute. + ctx = RequestContext.from_http(SimpleNamespace()) + assert ctx.http_request is not None + + +def test_from_http_keeps_the_msgpack_decoding_of_the_route_request(): + """A worker's request decodes msgpack bodies; the context must reuse it, not rebuild one.""" + body = _msgpack_encoder.encode(BODY) + request = _MsgspecRequest( + _http_request(body=body, headers=tuple(MSGPACK_HEADERS.items())).scope, + _http_request(body=body).receive, + ) + ctx = RequestContext.from_http(request) + + assert asyncio.run(ctx.json()) == BODY + + +def test_from_http_is_disconnected_asks_the_request(): + request = _http_request() + disconnected = {"value": False} + + async def is_disconnected(): + return disconnected["value"] + + request.is_disconnected = is_disconnected # type: ignore[method-assign] + ctx = RequestContext.from_http(request) + + assert asyncio.run(ctx.is_disconnected()) is False + disconnected["value"] = True + assert asyncio.run(ctx.is_disconnected()) is True + + +# -- internal: a request with no client behind it ------------------------------ + + +def test_internal_context_has_no_client_and_seeds_the_state(): + ctx = RequestContext.internal(BODY) + + assert ctx.http_request is None + assert dict(ctx.headers) == {} + assert ctx.client is None + assert ctx.no_client_connection is True + assert ctx.perf_metrics_records == [] + assert ctx.server_arrival_time is None + assert ctx.server_first_token_time is None + assert asyncio.run(ctx.json()) == BODY + assert asyncio.run(ctx.is_disconnected()) is False + assert bool(ctx) + + +def test_internal_json_is_what_a_wire_body_decodes_to_and_is_computed_once(): + # A body read from the wire has been through JSON: tuples become lists, + # non-string keys become strings. The internal context answers the same. + ctx = RequestContext.internal({"a": (1, 2), 3: "x"}) + + first = asyncio.run(ctx.json()) + assert first == {"a": [1, 2], "3": "x"} + assert asyncio.run(ctx.json()) is first + + +def test_internal_state_writes_stay_in_the_context(): + ctx = RequestContext.internal(BODY) + ctx.server_first_token_time = 5.0 + ctx.perf_metrics_records.append({"k": "v"}) + + assert ctx.state.server_first_token_time == 5.0 + assert ctx.perf_metrics_records == [{"k": "v"}] + + +def test_a_context_needs_a_request_or_an_internal_state(): + with pytest.raises(ValueError, match="from_http"): + RequestContext() + + +def test_await_disconnected_returns_immediately_for_an_internal_context(serve_env): + """A batched request has no socket, so the watchdog must not poll on it. + + is_disconnected() can never become true without a client, so without the + no_client_connection opt-out this coroutine would loop at 1Hz for the life + of the process, one live task per batched request. + """ + server = serve_env.build_openai_server(AggConfig()) + promise = SimpleNamespace(finished=False, request_id="req", abort=Mock()) + + async def run(): + # A socket-backed request would never let this return; a bounded wait + # keeps the failure a clear timeout rather than a hung suite. + await asyncio.wait_for( + server.await_disconnected(RequestContext.internal(BODY), promise), timeout=5.0 + ) + + asyncio.run(run()) + # Returning early must not be mistaken for a disconnect. + promise.abort.assert_not_called() diff --git a/tests/unittest/serve/_base/test_sse.py b/tests/unittest/serve/_base/test_sse.py index b62aa6f53cb1..517d93d383a4 100644 --- a/tests/unittest/serve/_base/test_sse.py +++ b/tests/unittest/serve/_base/test_sse.py @@ -101,9 +101,7 @@ def test_frames_then_done_then_metrics(error_records): params = PostprocParams(post_processor=None, postproc_args=None) frames = asyncio.run( - _collect( - chat_stream_generator(_Promise(results), params, server=server, raw_request=raw_request) - ) + _collect(chat_stream_generator(_Promise(results), params, server=server, ctx=raw_request)) ) assert frames == ["data: a\n\n", "data: b\n\n", "data: [DONE]\n\n"] @@ -125,9 +123,7 @@ def post_processor(result, args): 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) - ) + _collect(chat_stream_generator(_Promise(results), params, server=server, ctx=raw_request)) ) assert frames == ["data: 1\n\n", "data: 2\n\n", "data: [DONE]\n\n"] @@ -143,9 +139,7 @@ def test_a_failure_is_logged_and_raised(error_records): with pytest.raises(RuntimeError, match="boom"): asyncio.run( - _collect( - chat_stream_generator(promise, params, server=server, raw_request=_raw_request()) - ) + _collect(chat_stream_generator(promise, params, server=server, ctx=_raw_request())) ) assert len(error_records) == 1, error_records diff --git a/tests/unittest/serve/api/anthropic/test_api_router.py b/tests/unittest/serve/api/anthropic/test_api_router.py new file mode 100644 index 000000000000..b7f654ce7e36 --- /dev/null +++ b/tests/unittest/serve/api/anthropic/test_api_router.py @@ -0,0 +1,140 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""The Anthropic router: attachment contract, route provenance per role, msgpack bodies.""" + +from types import SimpleNamespace + +import pytest +from fastapi import FastAPI +from test_common.serve.app_factory import AggConfig +from test_common.serve.routes import effective_routes + +from tensorrt_llm.llmapi.disagg_utils import ServerRole +from tensorrt_llm.serve._base.msgspec_route import _MsgspecRoute +from tensorrt_llm.serve.api.anthropic.api_router import attach_router +from tensorrt_llm.serve.api.anthropic.serving import ( + AnthropicBatchesServing, + AnthropicCountTokensServing, + AnthropicMessagesServing, +) +from tensorrt_llm.serve.openai_client import MSGPACK_HEADERS, _msgpack_encoder +from tensorrt_llm.serve.render import RenderResources + +pytestmark = pytest.mark.cpu_only + +ANTHROPIC_ROUTER_MODULE = "tensorrt_llm.serve.api.anthropic.api_router" +ANTHROPIC_ROUTES = ( + ("POST", "/v1/messages"), + ("POST", "/v1/messages/count_tokens"), + ("POST", "/v1/messages/batches"), + ("GET", "/v1/messages/batches"), + ("GET", "/v1/messages/batches/{batch_id}"), + ("DELETE", "/v1/messages/batches/{batch_id}"), + ("POST", "/v1/messages/batches/{batch_id}/cancel"), + ("GET", "/v1/messages/batches/{batch_id}/results"), +) + + +def _anthropic_routes(app: FastAPI) -> dict: + """(method, path) -> [routes] for every Anthropic path the app registers.""" + found = {key: [] for key in ANTHROPIC_ROUTES} + for path, route in effective_routes(app): + for method in getattr(route, "methods", None) or (): + if (method, path) in found: + found[(method, path)].append(route) + return found + + +def test_attach_requires_every_serving_object(): + # The composition root sets the three serving objects before attaching, so + # a wiring bug fails at startup, not with an AttributeError per request. + with pytest.raises(RuntimeError, match="anthropic_serving_messages"): + attach_router(FastAPI()) + + +def test_serving_classes_require_their_collaborators(): + # No defaults: a serving class built without its collaborators fails at + # construction instead of falling back silently. + with pytest.raises(TypeError): + AnthropicMessagesServing() + with pytest.raises(TypeError): + AnthropicCountTokensServing() + with pytest.raises(TypeError): + AnthropicBatchesServing() + + +@pytest.mark.parametrize( + "role", [None, ServerRole.CONTEXT, ServerRole.GENERATION], ids=["llm", "context", "generation"] +) +def test_the_eight_routes_come_from_the_anthropic_router(role, serve_env): + server = serve_env.build_openai_server(AggConfig(role=role)) + + for key, routes in _anthropic_routes(server.app).items(): + assert len(routes) == 1, f"{key} must be registered exactly once, found {len(routes)}" + route = routes[0] + assert isinstance(route, _MsgspecRoute), ( + f"{key} must be a _MsgspecRoute; a plain APIRoute breaks msgpack request bodies " + "from the disaggregated frontend" + ) + assert route.endpoint.__module__ == ANTHROPIC_ROUTER_MODULE + + +def test_a_harmony_model_has_no_anthropic_routes(serve_env): + server = serve_env.build_openai_server(AggConfig(model_type="gpt_oss")) + assert all(not routes for routes in _anthropic_routes(server.app).values()) + + +@pytest.mark.parametrize("role", [ServerRole.MM_ENCODER, ServerRole.EMBEDDING]) +def test_other_roles_have_no_anthropic_routes(role, serve_env): + server = serve_env.build_openai_server(AggConfig(role=role)) + assert all(not routes for routes in _anthropic_routes(server.app).values()) + + +class _WordTokenizer: + """Renders the conversation as the words of its messages and counts them.""" + + def get_chat_template(self, chat_template=None, tools=None): + return chat_template or "WORDS" + + def apply_chat_template(self, messages=None, **kwargs): + if messages is None: + messages = kwargs.pop("conversation") + return " ".join(str(m.get("content", "")) for m in messages) + + def encode(self, text, add_special_tokens=True, **kwargs): + return text.split() + + +# The TestClient portal thread and the fake server's lazy worker-pool threads +# are alive when the leak check runs; the serve_env fixture closes them in its +# teardown (precedent: tests/unittest/serve/api/ops/system/test_system_endpoints.py). +@pytest.mark.threadleak(enabled=False) +def test_count_tokens_decodes_the_disaggregated_frontends_msgpack_body(serve_env): + """The disaggregated frontend forwards count_tokens as msgpack to a context worker.""" + server = serve_env.build_openai_server(AggConfig()) + resources = SimpleNamespace(tokenizer=_WordTokenizer(), chat_template=None) + server.app.state.anthropic_serving_count_tokens = AnthropicCountTokensServing( + render_resources=lambda: RenderResources.from_server(resources) + ) + body = _msgpack_encoder.encode( + {"model": "m", "messages": [{"role": "user", "content": "one two three"}]} + ) + + response = serve_env.client(server.app).post( + "/v1/messages/count_tokens", content=body, headers=MSGPACK_HEADERS + ) + + assert response.status_code == 200, response.text + assert response.json() == {"input_tokens": 3} diff --git a/tests/unittest/serve/api/anthropic/test_serving.py b/tests/unittest/serve/api/anthropic/test_serving.py new file mode 100644 index 000000000000..bb50d5772344 --- /dev/null +++ b/tests/unittest/serve/api/anthropic/test_serving.py @@ -0,0 +1,355 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""``AnthropicMessagesServing`` over a fake ``ChatBackend``: what reaches the answer and how.""" + +import asyncio +import json +import logging + +import pytest +from fastapi.responses import JSONResponse, StreamingResponse + +from tensorrt_llm.logger import logger as trtllm_logger +from tensorrt_llm.serve._base.errors import EngineFatalError, ServingError +from tensorrt_llm.serve._base.request_context import RequestContext +from tensorrt_llm.serve.api.anthropic import serving as serving_module +from tensorrt_llm.serve.api.anthropic.protocol import AnthropicMessagesRequest +from tensorrt_llm.serve.api.anthropic.serving import AnthropicMessagesServing, _error_type +from tensorrt_llm.serve.openai_protocol import ( + ChatCompletionResponse, + ChatCompletionResponseChoice, + ChatCompletionStreamResponse, + ChatMessage, + ErrorResponse, + FunctionCall, + ToolCall, + UsageInfo, +) + +pytestmark = pytest.mark.cpu_only + +CLIENT_MODEL = "client-model" +RESPONSE_MODEL = "response-model" + + +class _Backend: + """A ChatBackend that answers with a fixed result or raises a fixed error.""" + + def __init__(self, *, result=None, error=None, model="x-model", message="x-message"): + self.result = result + self.error = error + self.model = model + self.message = message + self.calls = [] + + async def create_chat_completion(self, request, ctx): + self.calls.append((request, ctx)) + if self.error is not None: + raise self.error + return self.result + + def stream_model(self, request): + return self.model + + def error_message(self, error): + return self.message + + +def _request(**overrides) -> AnthropicMessagesRequest: + payload = { + "model": CLIENT_MODEL, + "max_tokens": 32, + "messages": [{"role": "user", "content": "hello"}], + } + payload.update(overrides) + return AnthropicMessagesRequest(**payload) + + +def _chat_response(*, tool_arguments=None) -> ChatCompletionResponse: + tool_calls = [] + if tool_arguments is not None: + tool_calls = [ + ToolCall( + id="call_1", function=FunctionCall(name="get_weather", arguments=tool_arguments) + ) + ] + return ChatCompletionResponse( + id="chatcmpl-serving-test", + model=RESPONSE_MODEL, + choices=[ + ChatCompletionResponseChoice( + index=0, + message=ChatMessage(role="assistant", content="hello", tool_calls=tool_calls), + finish_reason="tool_calls" if tool_calls else "stop", + ) + ], + usage=UsageInfo(prompt_tokens=3, completion_tokens=2, total_tokens=5), + ) + + +def _frames(*texts: str) -> list: + chunks = [ + ChatCompletionStreamResponse( + model=RESPONSE_MODEL, choices=[{"index": 0, "delta": {"content": text}}] + ) + for text in texts + ] + return [f"data: {chunk.model_dump_json()}\n\n" for chunk in chunks] + ["data: [DONE]\n\n"] + + +async def _stream(frames: list): + for frame in frames: + yield frame + + +def _serving_error(status: int) -> ServingError: + return ServingError(status, ErrorResponse(message="body message", type="t", code=status)) + + +def _create(backend: _Backend, request: AnthropicMessagesRequest): + serving = AnthropicMessagesServing(backend) + return asyncio.run(serving.create_message(request, RequestContext.internal({}))) + + +def _body(response: JSONResponse) -> dict: + return json.loads(bytes(response.body)) + + +def _events(response: StreamingResponse) -> list: + async def collect(): + return [frame async for frame in response.body_iterator] + + events = [] + for frame in "".join(asyncio.run(collect())).strip().split("\n\n"): + name, data = None, None + for line in frame.splitlines(): + if line.startswith("event: "): + name = line[len("event: ") :] + elif line.startswith("data: "): + data = json.loads(line[len("data: ") :]) + events.append((name, data)) + return events + + +@pytest.fixture +def error_records(): + """Messages of the ERROR records the TensorRT-LLM logger emits (it does not propagate to caplog).""" + records = [] + + class Capture(logging.Handler): + def emit(self, record): + records.append(record.getMessage()) + + handler = Capture(level=logging.ERROR) + trtllm_logger._logger.addHandler(handler) + try: + yield records + finally: + trtllm_logger._logger.removeHandler(handler) + + +# -- create_message -------------------------------------------------------------- + + +def test_non_streaming_result_is_translated_without_a_round_trip(): + backend = _Backend(result=_chat_response()) + + response = _create(backend, _request()) + + assert isinstance(response, JSONResponse) and response.status_code == 200 + body = _body(response) + assert body["content"] == [{"type": "text", "text": "hello"}] + assert body["model"] == RESPONSE_MODEL + assert body["usage"] == {"input_tokens": 3, "output_tokens": 2} + (chat_request, ctx), *_ = backend.calls + assert chat_request.model == CLIENT_MODEL and chat_request.stream is False + assert isinstance(ctx, RequestContext) + + +def test_streaming_result_is_reframed_as_event_stream_from_the_backends_iterator(monkeypatch): + stream = _stream(_frames("he", "llo")) + backend = _Backend(result=stream) + seen = {} + + async def fake_reframe(openai_sse, model): + seen["iterator"] = openai_sse + seen["model"] = model + yield "event: ping\ndata: {}\n\n" + + monkeypatch.setattr(serving_module, "reframe_openai_stream", fake_reframe) + + response = _create(backend, _request(stream=True)) + + assert isinstance(response, StreamingResponse) + assert response.media_type == "text/event-stream" + assert _events(response) == [("ping", {})] + # The backend's stream object goes to the reframer as it is: no copy, no + # wrapper, so the adapter reads the worker's bytes or the core's frames. + assert seen["iterator"] is stream + assert seen["model"] == "x-model" + + +def test_stream_without_chunks_names_the_backends_model(): + backend = _Backend(result=_stream(["data: [DONE]\n\n"]), model="x-model") + + response = _create(backend, _request(stream=True)) + + events = _events(response) + assert events[0][0] == "message_start" + assert events[0][1]["message"]["model"] == "x-model" + assert events[-1][0] == "message_stop" + + +def test_stream_with_chunks_adopts_the_chunks_model(): + backend = _Backend(result=_stream(_frames("he", "llo")), model="x-model") + + response = _create(backend, _request(stream=True)) + + events = _events(response) + assert events[0][1]["message"]["model"] == RESPONSE_MODEL + text = "".join( + data["delta"]["text"] + for name, data in events + if name == "content_block_delta" and data["delta"]["type"] == "text_delta" + ) + assert text == "hello" + + +@pytest.mark.parametrize( + "status, expected_type", + [(400, "invalid_request_error"), (429, "invalid_request_error"), (503, "api_error")], +) +def test_serving_error_is_rendered_with_the_backends_message_and_the_status(status, expected_type): + backend = _Backend(error=_serving_error(status), message="x-message") + + response = _create(backend, _request()) + + assert response.status_code == status + assert _body(response) == { + "type": "error", + "error": {"type": expected_type, "message": "x-message"}, + } + + +def test_engine_failure_is_a_500_api_error(): + backend = _Backend(error=EngineFatalError()) + + response = _create(backend, _request()) + + assert response.status_code == 500 + assert _body(response)["error"] == {"type": "api_error", "message": "Internal server error"} + + +def test_request_conversion_failure_is_a_400_before_the_backend_runs(): + backend = _Backend(result=_chat_response()) + + response = _create( + backend, _request(tools=[{"name": "web_search", "type": "web_search_20250305"}]) + ) + + assert response.status_code == 400 + assert _body(response)["error"]["type"] == "invalid_request_error" + assert "server tool" in _body(response)["error"]["message"] + assert backend.calls == [] + + +def test_unconvertible_chat_response_is_a_500_with_one_error_log(error_records): + backend = _Backend(result=_chat_response(tool_arguments="{not-json")) + + response = _create(backend, _request()) + + assert response.status_code == 500 + assert _body(response)["error"] == {"type": "api_error", "message": "Internal server error"} + assert len(error_records) == 1 + assert "Invalid response from OpenAI chat pipeline" in error_records[0] + + +# -- run_batched_message ---------------------------------------------------------- + + +def _run_batched(backend: _Backend, request: AnthropicMessagesRequest): + serving = AnthropicMessagesServing(backend) + return asyncio.run(serving.run_batched_message(request)) + + +def test_batched_message_runs_an_internal_context_without_streaming(): + backend = _Backend(result=_chat_response()) + + outcome, payload = _run_batched(backend, _request(stream=True)) + + assert outcome == "succeeded" + assert payload["content"] == [{"type": "text", "text": "hello"}] + (chat_request, ctx), *_ = backend.calls + # Streaming makes no sense for a batch: the result is collected, not watched. + assert chat_request.stream is False + assert isinstance(ctx, RequestContext) + assert ctx.http_request is None + assert dict(ctx.headers) == {} + assert ctx.no_client_connection is True + assert asyncio.run(ctx.json()) == chat_request.model_dump(exclude_none=True) + + +@pytest.mark.parametrize( + "error, expected", + [ + (_serving_error(400), {"type": "invalid_request_error", "message": "x-message"}), + (_serving_error(503), {"type": "api_error", "message": "x-message"}), + (EngineFatalError(), {"type": "api_error", "message": "Internal server error"}), + ], + ids=["rejected", "upstream_5xx", "engine_failure"], +) +def test_batched_message_reports_errors_as_errored_items(error, expected): + backend = _Backend(error=error, message="x-message") + + assert _run_batched(backend, _request()) == ("errored", expected) + + +def test_batched_message_conversion_failure_does_not_reach_the_backend(): + backend = _Backend(result=_chat_response()) + + outcome, payload = _run_batched( + backend, _request(tools=[{"name": "web_search", "type": "web_search_20250305"}]) + ) + + assert outcome == "errored" + assert payload["type"] == "invalid_request_error" + assert backend.calls == [] + + +def test_batched_unconvertible_response_is_an_errored_item_with_one_log(error_records): + backend = _Backend(result=_chat_response(tool_arguments="{not-json")) + + assert _run_batched(backend, _request()) == ( + "errored", + {"type": "api_error", "message": "Internal server error"}, + ) + assert len(error_records) == 1 + assert "Invalid response from OpenAI chat pipeline in batch" in error_records[0] + + +# -- _error_type ---------------------------------------------------------------- + + +@pytest.mark.parametrize( + "status, expected", + [ + (399, "api_error"), + (400, "invalid_request_error"), + (499, "invalid_request_error"), + (500, "api_error"), + ], +) +def test_error_type_boundaries(status, expected): + assert _error_type(status) == expected diff --git a/tests/unittest/serve/test_layering.py b/tests/unittest/serve/test_layering.py index 357dcd4b02fc..e09971963ba0 100644 --- a/tests/unittest/serve/test_layering.py +++ b/tests/unittest/serve/test_layering.py @@ -93,6 +93,18 @@ "packages; pass what it needs in from the composition root." ), ), + Contract( + name="the Anthropic package does not import a server", + importers=("tensorrt_llm.serve.api.anthropic",), + forbidden=( + "tensorrt_llm.serve.openai_server", + "tensorrt_llm.serve.openai_disagg_server", + ), + fix=( + "serve/api/anthropic/ serves both servers; take the chat backend and other " + "dependencies through constructors and never import a server module." + ), + ), Contract( name="serve/api/ and serve/_base/ do not import the composition root", importers=("tensorrt_llm.serve.api", "tensorrt_llm.serve._base"), diff --git a/tests/unittest/serve/test_no_object_new.py b/tests/unittest/serve/test_no_object_new.py index 0aa59460caf8..99fa1ed463aa 100644 --- a/tests/unittest/serve/test_no_object_new.py +++ b/tests/unittest/serve/test_no_object_new.py @@ -92,19 +92,8 @@ "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", @@ -113,7 +102,6 @@ "_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": [ diff --git a/tests/unittest/serve/test_openai_server_chat_core.py b/tests/unittest/serve/test_openai_server_chat_core.py index 3ed4257650a8..da96de12d1c0 100644 --- a/tests/unittest/serve/test_openai_server_chat_core.py +++ b/tests/unittest/serve/test_openai_server_chat_core.py @@ -42,6 +42,7 @@ from tensorrt_llm.logger import logger from tensorrt_llm.serve import openai_server from tensorrt_llm.serve._base.errors import EngineFatalError, ServingError +from tensorrt_llm.serve._base.request_context import RequestContext from tensorrt_llm.serve.openai_protocol import ( ChatCompletionRequest, ChatCompletionResponse, @@ -267,7 +268,9 @@ def _request(**overrides) -> ChatCompletionRequest: def _run_core(server, request): - return asyncio.run(server.create_chat_completion(request, _raw_request())) + return asyncio.run( + server.create_chat_completion(request, RequestContext.from_http(_raw_request())) + ) def test_typed_result_for_a_non_streaming_request(serve_env): @@ -288,7 +291,9 @@ def test_stream_is_returned_as_sse_strings(serve_env): ) async def run(): - result = await server.create_chat_completion(_request(stream=True), _raw_request()) + result = await server.create_chat_completion( + _request(stream=True), RequestContext.from_http(_raw_request()) + ) assert not isinstance(result, Response) return [frame async for frame in result] From 6998c2b9e6359af1b68460a600465e803a83b1f8 Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Wed, 4 Nov 2026 20:36:09 +0000 Subject: [PATCH 27/27] [TRTLLM-17263][test] Mount the Messages usage test on the Anthropic router and pass the kv-hints handlers a request tests/unittest/llmapi/apps/test_responses_usage.py mounted OpenAIServer.anthropic_messages, which no longer exists: the Anthropic routes live in tensorrt_llm/serve/api/anthropic/api_router.py and read their serving object from app.state. The test's bare app now mounts the router's route function and carries the AnthropicMessagesServing the server's constructor built, so the Messages route still runs over the same chat core as the chat route. tests/unittest/llmapi/test_kv_hints.py called the handlers with raw_request=None. The chat wrapper now wraps the request in a RequestContext, which needs an object to delegate to, so the handlers get a minimal request (empty headers, an empty state), as the chat core tests do. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- tests/unittest/llmapi/apps/test_responses_usage.py | 6 +++++- tests/unittest/llmapi/test_kv_hints.py | 6 +++++- 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/tests/unittest/llmapi/apps/test_responses_usage.py b/tests/unittest/llmapi/apps/test_responses_usage.py index 7b36f70a5d00..789b53eaa421 100644 --- a/tests/unittest/llmapi/apps/test_responses_usage.py +++ b/tests/unittest/llmapi/apps/test_responses_usage.py @@ -30,6 +30,7 @@ from tensorrt_llm.executor.result import CompletionOutput from tensorrt_llm.inputs.tokenization import tokenize_prompt +from tensorrt_llm.serve.api.anthropic.api_router import anthropic_messages from tensorrt_llm.serve.openai_disagg_service import _ctx_usage_info from tensorrt_llm.serve.openai_protocol import ( PromptTokensDetails, @@ -380,7 +381,10 @@ def generate_async(*, inputs, disaggregated_params=None, **kwargs): server._extract_metrics = AsyncMock() app = FastAPI() app.add_api_route("/v1/chat/completions", server.openai_chat, methods=["POST"]) - app.add_api_route("/v1/messages", server.anthropic_messages, methods=["POST"]) + # The Messages route function reads its serving object from app.state; the + # server's constructor built it over the same chat core the chat route runs. + app.state.anthropic_serving_messages = server.app.state.anthropic_serving_messages + app.add_api_route("/v1/messages", anthropic_messages, methods=["POST"]) app.add_api_route("/v1/responses", server.openai_responses, methods=["POST"]) return TestClient(app), prompts diff --git a/tests/unittest/llmapi/test_kv_hints.py b/tests/unittest/llmapi/test_kv_hints.py index e73a6097d2d8..1e2d4d51b827 100644 --- a/tests/unittest/llmapi/test_kv_hints.py +++ b/tests/unittest/llmapi/test_kv_hints.py @@ -662,11 +662,15 @@ def test_http_handlers_forward_hints(request_kind: str, with_hints: bool) -> Non ) handler = server.chat_harmony if request_kind == "harmony" else server.openai_chat + # The request FastAPI would pass, reduced to what the handlers read before + # the engine boundary. The chat handler wraps it in a RequestContext, which + # needs an object to delegate to. + raw_request = SimpleNamespace(state=SimpleNamespace(), headers={}) with patch( "tensorrt_llm.serve.openai_server.responses_api_request_preprocess", new=AsyncMock(return_value=([1, 2], SamplingParams(max_tokens=4))), ): - response = asyncio.run(handler(request, None)) + response = asyncio.run(handler(request, raw_request)) server.generator.generate_async.assert_called_once() assert server.generator.generate_async.call_args.kwargs["kv_hints"] == hints assert b"captured at engine boundary" in response.body