From 259fdf6ecbabc214f975088ed8b4735be06b9c53 Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Wed, 28 Oct 2026 03:51:22 +0000 Subject: [PATCH 01/30] [TRTLLM-17272][feat] Add public LLM health queries and a private-access lint for the serving frontends The HTTP server, the SMG gRPC frontend and the OpenEngine gRPC frontend learned engine health by reading private members (_check_health, _executor, _fatal_error, _engine_dead). Add the public health API on the classes that own those members and switch all three frontends to it, with no behavior change: - BaseLLM.check_health() (public rename of _check_health; the private name stays for one release as an alias for out-of-tree callers that look it up), BaseLLM.fatal_error, BaseLLM.shutdown_started, BaseLLM.engine_dead. - GenerationExecutor.fatal_error and engine_dead; GenerationExecutorProxy.engine_dead also reads the sticky worker-death flag set before the fatal error is recorded. Add a private-access lint (tests/unittest/serve/test_private_access_lint.py) that counts cross-package private accesses in tensorrt_llm/serve, tensorrt_llm/commands and tensorrt_llm/grpc against a checked-in baseline (tests/unittest/serve/private_access_baseline.yaml) that only shrinks, with companion test files under tests/unittest/serve/api/visual_gen/ and tests/unittest/grpc/openengine/ so VisualGen-only and OpenEngine-only pull requests keep running the check. Replace the three frontend health mirror tests with tests of the production code, rename the _check_health test doubles to the public name, and register the new CPU test files in l0_cpu.yml. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- tensorrt_llm/executor/executor.py | 10 + tensorrt_llm/executor/proxy.py | 9 + tensorrt_llm/grpc/openengine/control.py | 24 +- tensorrt_llm/grpc/openengine/server.py | 14 +- tensorrt_llm/grpc/smg/request_manager.py | 29 +- tensorrt_llm/llmapi/llm.py | 58 ++- tensorrt_llm/serve/openai_server.py | 25 +- .../integration/test_lists/test-db/l0_cpu.yml | 6 + .../api_stability/references/llm.yaml | 13 + .../executor/test_fatal_error_health_check.py | 162 +-------- .../unittest/grpc/openengine/test_control.py | 4 +- .../grpc/openengine/test_coordinator.py | 8 +- .../openengine/test_frontend_lifecycle.py | 3 +- .../unittest/grpc/openengine/test_metrics.py | 2 +- .../grpc/openengine/test_openengine.py | 42 ++- .../test_private_access_lint_openengine.py | 33 ++ .../unittest/grpc/openengine/test_servicer.py | 1 + tests/unittest/grpc/smg/test_smg_health.py | 71 ++++ tests/unittest/llmapi/test_llm_encode.py | 4 +- .../llmapi/test_llm_public_runtime_api.py | 180 +++++++++ tests/unittest/serve/__init__.py | 2 + tests/unittest/serve/api/__init__.py | 2 + .../unittest/serve/api/visual_gen/__init__.py | 2 + .../test_private_access_lint_visual_gen.py | 27 ++ .../serve/private_access_baseline.yaml | 53 +++ .../serve/test_private_access_lint.py | 344 ++++++++++++++++++ tests/unittest/serve/test_serve_health.py | 120 ++++++ 27 files changed, 1017 insertions(+), 231 deletions(-) create mode 100644 tests/unittest/grpc/openengine/test_private_access_lint_openengine.py create mode 100644 tests/unittest/grpc/smg/test_smg_health.py create mode 100644 tests/unittest/llmapi/test_llm_public_runtime_api.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/visual_gen/__init__.py create mode 100644 tests/unittest/serve/api/visual_gen/test_private_access_lint_visual_gen.py create mode 100644 tests/unittest/serve/private_access_baseline.yaml create mode 100644 tests/unittest/serve/test_private_access_lint.py create mode 100644 tests/unittest/serve/test_serve_health.py diff --git a/tensorrt_llm/executor/executor.py b/tensorrt_llm/executor/executor.py index 7e8b565683a8..6738aa4986a3 100644 --- a/tensorrt_llm/executor/executor.py +++ b/tensorrt_llm/executor/executor.py @@ -337,6 +337,16 @@ def is_shutdown(self) -> bool: """Return True if the executor is shutting down or fatally errored.""" return self.doing_shutdown or self._fatal_error is not None + @property + def fatal_error(self) -> Optional[BaseException]: + """The unrecoverable engine error recorded by ``_set_fatal_error``, or None.""" + return self._fatal_error + + @property + def engine_dead(self) -> bool: + """Whether the engine is known dead (an unrecoverable error was recorded).""" + return self._fatal_error is not None + def check_health(self) -> bool: """Check whether the executor is healthy and able to process requests. diff --git a/tensorrt_llm/executor/proxy.py b/tensorrt_llm/executor/proxy.py index 3e0f3f919ecc..765e9a001bef 100644 --- a/tensorrt_llm/executor/proxy.py +++ b/tensorrt_llm/executor/proxy.py @@ -480,6 +480,15 @@ def check_health(self) -> bool: return True + @property + def engine_dead(self) -> bool: + """Whether the engine is known dead (a worker died or a fatal error was recorded). + + ``_mark_engine_dead`` can set the sticky flag before the fatal error + is recorded, so check both. + """ + return self._engine_dead or self._fatal_error is not None + def _set_fatal_error(self, error: BaseException) -> None: """Record the fatal error, then unblock every pending request. diff --git a/tensorrt_llm/grpc/openengine/control.py b/tensorrt_llm/grpc/openengine/control.py index 47278788dc33..3275880862e8 100644 --- a/tensorrt_llm/grpc/openengine/control.py +++ b/tensorrt_llm/grpc/openengine/control.py @@ -140,23 +140,19 @@ async def _get_kv_cache_load(self) -> dict: def _engine_is_healthy(self) -> bool: """Readiness via the predicate the engine's own HTTP /health uses. - `_check_health` is stronger than the `is_shutdown()` guard generate_async + `check_health()` is stronger than the `is_shutdown()` guard generate_async raises on: it drains the executor's error queue, so a fatal error queued but not yet promoted is caught, and the proxy polls MPI worker liveness. - It also covers encode-only engines, where `_executor` is None. + It also covers encode-only engines, where there is no generation executor. """ - check = getattr(self._llm, "_check_health", None) - if callable(check): - try: - return bool(check()) - except (RuntimeError, ValueError, OSError) as error: - logger.warning(f"OpenEngine health check reported a failure: {error}") - return False - except (TypeError, AttributeError) as error: - logger.error(f"OpenEngine health check contract failed: {error}") - return False - executor = getattr(self._llm, "_executor", None) - return executor is not None and not executor.is_shutdown() + try: + return bool(self._llm.check_health()) + except (RuntimeError, ValueError, OSError) as error: + logger.warning(f"OpenEngine health check reported a failure: {error}") + return False + except (TypeError, AttributeError) as error: + logger.error(f"OpenEngine health check contract failed: {error}") + return False # -- Runtime metadata --------------------------------------------------- diff --git a/tensorrt_llm/grpc/openengine/server.py b/tensorrt_llm/grpc/openengine/server.py index 97ccb3210637..b495db5f6c1d 100644 --- a/tensorrt_llm/grpc/openengine/server.py +++ b/tensorrt_llm/grpc/openengine/server.py @@ -215,11 +215,7 @@ async def _refresh_inference_health(self) -> None: except (RuntimeError, ValueError, OSError, CoordinationError) as error: logger.warning(f"OpenEngine inference health check failed: {error}") self._engine_healthy = False - executor = getattr(self._llm, "_executor", None) - if executor is not None and ( - getattr(executor, "_fatal_error", None) is not None - or getattr(executor, "_engine_dead", False) - ): + if self._llm.engine_dead: self._fatal_engine_failure = True if self._stopping: return @@ -412,11 +408,7 @@ def healthy() -> bool: ) def terminal_engine_failure() -> bool: - executor = getattr(llm, "_executor", None) - return executor is not None and ( - getattr(executor, "_fatal_error", None) is not None - or getattr(executor, "_engine_dead", False) - ) + return llm.engine_dead async def supervise() -> None: while not stop_event.is_set(): @@ -428,7 +420,7 @@ async def supervise() -> None: return if server is None: engine_ready = await asyncio.wait_for( - asyncio.to_thread(llm._check_health), + asyncio.to_thread(llm.check_health), timeout=_HEALTH_CHECK_TIMEOUT_SECONDS, ) if server is None and not engine_ready: diff --git a/tensorrt_llm/grpc/smg/request_manager.py b/tensorrt_llm/grpc/smg/request_manager.py index 460d19804b0b..01a527691302 100644 --- a/tensorrt_llm/grpc/smg/request_manager.py +++ b/tensorrt_llm/grpc/smg/request_manager.py @@ -168,20 +168,21 @@ async def health_check(self) -> Tuple[bool, str]: if self.llm is None: return False, "LLM not initialized" - # Check executor health (includes error queue, MPI worker - # liveness, and fatal error state — not just doing_shutdown). - if hasattr(self.llm, "_executor"): - if self.llm._executor is None: - return False, "Executor is not available" - if not self.llm._executor.check_health(): - error_msg = "Executor is unhealthy" - if self.llm._executor._fatal_error is not None: - exc = self.llm._executor._fatal_error - lines = str(exc).splitlines() - short = (lines[0] if lines else type(exc).__name__)[:200] - error_msg = f"{type(exc).__name__}: {short}" - logger.error(f"Health check fatal error: {repr(exc)}") - return False, error_msg + # An encode-only LLM has no generation executor; SMG serves + # generate requests only. + if self.llm.args.encode_only is True: + return False, "Executor is not available" + # check_health() drains queued errors and checks MPI worker + # liveness and the fatal error state, not just shutdown. + if not self.llm.check_health(): + error_msg = "Executor is unhealthy" + exc = self.llm.fatal_error + if exc is not None: + lines = str(exc).splitlines() + short = (lines[0] if lines else type(exc).__name__)[:200] + error_msg = f"{type(exc).__name__}: {short}" + logger.error(f"Health check fatal error: {repr(exc)}") + return False, error_msg return True, "OK" except Exception as e: diff --git a/tensorrt_llm/llmapi/llm.py b/tensorrt_llm/llmapi/llm.py index 2d9dabd427d9..985e84978550 100644 --- a/tensorrt_llm/llmapi/llm.py +++ b/tensorrt_llm/llmapi/llm.py @@ -339,6 +339,9 @@ def _multimodal_params_have_encoder_features( llm_id (str): The unique ID of the LLM instance. disaggregated_params (dict): The disaggregated parameters of the LLM instance. startup_metrics (dict): The startup metrics reported by worker rank 0. + fatal_error (BaseException, optional): The unrecoverable engine error the executor recorded, or None. + shutdown_started (bool): Whether shutdown of the generation executor has begun. + engine_dead (bool): Whether the engine is known dead (a worker died or a fatal error was recorded). """ @@ -1785,11 +1788,20 @@ def _record_usage_shutdown(self) -> None: except Exception as exc: logger.debug(f"Usage telemetry shutdown tracking failed: {exc}") - def _check_health(self) -> bool: - """Check if the LLM is healthy. + @set_api_status("prototype") + def check_health(self) -> bool: + """Check whether the LLM can serve requests. + + For ``encode_only=True`` the LLM is healthy while its encoder executor + exists. Otherwise this delegates to the generation executor's + ``check_health()``, which drains queued background errors (a fatal one + is recorded, see ``fatal_error``, and starts executor shutdown) and, + for the multi-process proxy, polls worker liveness. Without a + generation executor (never built, or after ``shutdown()``) the LLM is + unhealthy. Returns: - bool: True if the executor is running and healthy, False otherwise. + bool: True if the LLM is healthy, False otherwise. """ if self._encode_only: return (hasattr(self, "_encoder_executor") @@ -1799,6 +1811,46 @@ def _check_health(self) -> bool: return False + # Kept for one release for out-of-tree callers that look this name up + # (ai-dynamo components/src/dynamo/trtllm/engine.py); use check_health(). + _check_health = check_health + + @property + @set_api_status("prototype") + def fatal_error(self) -> Optional[BaseException]: + """The unrecoverable engine error the executor recorded, if any. + + None when no fatal error was recorded or the LLM has no generation + executor (``encode_only=True``, or after ``shutdown()``). + """ + executor = getattr(self, "_executor", None) + return None if executor is None else executor.fatal_error + + @property + @set_api_status("prototype") + def shutdown_started(self) -> bool: + """Whether shutdown of the generation executor has begun. + + True once ``shutdown()`` was called or the executor started shutting + down after a fatal error. False for ``encode_only=True``, which has no + generation executor. + """ + executor = getattr(self, "_executor", None) + if executor is None: + return not getattr(self, "_encode_only", False) + return executor.doing_shutdown + + @property + @set_api_status("prototype") + def engine_dead(self) -> bool: + """Whether the engine is known dead (a worker died or a fatal error was recorded). + + False when the LLM has no generation executor (``encode_only=True``, + or after ``shutdown()``). + """ + executor = getattr(self, "_executor", None) + return False if executor is None else executor.engine_dead + @staticmethod def _shutdown_wrapper(self_ref): # Retrieve the instance if it still exists diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py index 5079b9b073e4..d8adde6eafd7 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -1227,9 +1227,11 @@ def _check_health(self) -> bool: batcher = self.embedding_batcher if batcher is not None and not batcher.is_alive(): return False - if hasattr(self.generator, '_check_health'): - return self.generator._check_health() - return True + if self._is_visual_gen: + # VisualGen has no check_health(), so its /health always + # answers 200. + return True + return self.generator.check_health() def register_routes(self): self.app.add_api_route("/health", self.health, methods=["GET"]) @@ -1727,15 +1729,14 @@ async def health(self) -> Response: # 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) + fatal_error = (None if self._is_visual_gen else + self.generator.fatal_error) + if fatal_error is not None and not self.generator.shutdown_started: + logger.error( + "Health check detected fatal engine error, initiating " + f"server shutdown: {fatal_error}") + _record_generator_termination(self.generator) + signal.raise_signal(signal.SIGINT) return Response( status_code=503, content= diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index 550796e3fa0c..69af33f55635 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -236,6 +236,8 @@ l0_cpu: - unittest/grpc/openengine/test_openengine.py - unittest/grpc/openengine/test_metrics.py - unittest/grpc/openengine/test_servicer.py + - unittest/grpc/openengine/test_private_access_lint_openengine.py + - unittest/grpc/smg/test_smg_health.py - unittest/grpc/smg/test_smg.py::test_streaming_deltas_preserve_sequences_logprobs_and_final_tokens - unittest/grpc/smg/test_smg.py::test_streaming_empty_token_history - unittest/grpc/test_grpc_optional.py::test_protobuf_warning_in_fresh_process @@ -295,6 +297,7 @@ l0_cpu: - unittest/llmapi/test_kv_cache_dtype_override.py - unittest/llmapi/test_llm_args.py - unittest/llmapi/test_gms_args.py + - unittest/llmapi/test_llm_public_runtime_api.py - unittest/llmapi/test_mx_args.py - unittest/llmapi/test_custom_tokenizer_aliases.py # Reads the Muse Spark Open tokenizer files from LLM_MODELS_ROOT. @@ -334,6 +337,9 @@ l0_cpu: - unittest/quantization/test_mode.py - unittest/scaffolding - unittest/scripts + - unittest/serve/api/visual_gen/test_private_access_lint_visual_gen.py + - unittest/serve/test_private_access_lint.py + - unittest/serve/test_serve_health.py - unittest/tools - unittest/usage - unittest/utils/test_logger.py diff --git a/tests/unittest/api_stability/references/llm.yaml b/tests/unittest/api_stability/references/llm.yaml index c75fe8fe084b..b1be771de618 100644 --- a/tests/unittest/api_stability/references/llm.yaml +++ b/tests/unittest/api_stability/references/llm.yaml @@ -528,6 +528,10 @@ methods: default: 2 return_annotation: tensorrt_llm.executor.result.IterationResult status: beta + check_health: + parameters: {} + return_annotation: bool + status: prototype get_data_transceiver_state: parameters: {} return_annotation: bytes @@ -584,3 +588,12 @@ properties: startup_metrics: annotation: dict default: None + fatal_error: + annotation: Optional[BaseException] + default: inspect._empty + shutdown_started: + annotation: bool + default: inspect._empty + engine_dead: + annotation: bool + default: inspect._empty diff --git a/tests/unittest/executor/test_fatal_error_health_check.py b/tests/unittest/executor/test_fatal_error_health_check.py index 3f8f7396a6ff..77a4b2698275 100644 --- a/tests/unittest/executor/test_fatal_error_health_check.py +++ b/tests/unittest/executor/test_fatal_error_health_check.py @@ -20,9 +20,11 @@ - GenerationExecutor.check_health(), _set_fatal_error(), is_shutdown() - GenerationExecutorProxy.check_health() with MPI worker liveness - GenerationExecutorProxy._error_monitor_loop() background detection -- GrpcRequestManager.health_check() integration -- OpenAIServer /health endpoint with fatal error shutdown -- BaseLLM._check_health() delegation + +Frontend health is tested against production code in +tests/unittest/llmapi/test_llm_public_runtime_api.py, +tests/unittest/serve/test_serve_health.py and +tests/unittest/grpc/smg/test_smg_health.py. """ # Import classify_error directly from the source file to avoid triggering @@ -30,14 +32,13 @@ import importlib.util import logging import pathlib -import signal import subprocess import sys import threading import time from concurrent.futures import Future from queue import Empty, Queue -from unittest.mock import MagicMock, Mock, patch +from unittest.mock import Mock, patch import pytest @@ -827,157 +828,6 @@ def test_stops_on_shutdown_flag(self): assert not t.is_alive() -# --------------------------------------------------------------------------- -# gRPC health check -# --------------------------------------------------------------------------- -class TestGrpcHealthCheck: - """Tests for gRPC health check integration with fatal error state.""" - - @staticmethod - async def _health_check(self): - """Mirrors GrpcRequestManager.health_check logic.""" - try: - if self.llm is None: - return False, "LLM not initialized" - if hasattr(self.llm, "_executor"): - if self.llm._executor is None: - return False, "Executor is not available" - if not self.llm._executor.check_health(): - error_msg = "Executor is unhealthy" - if self.llm._executor._fatal_error is not None: - exc = self.llm._executor._fatal_error - short = str(exc).splitlines()[0][:200] - error_msg = f"{type(exc).__name__}: {short}" - return False, error_msg - return True, "OK" - except Exception as e: - return False, f"Error: {e}" - - def _make_manager(self, executor=None): - """Create a mock gRPC manager with the given executor.""" - m = MagicMock() - m.llm = MagicMock() - m.llm._executor = executor - m.health_check = self._health_check.__get__(m) - return m - - @pytest.mark.asyncio - @pytest.mark.parametrize( - "setup,expected_healthy,msg_substr", - [ - ("healthy", True, "OK"), - ("fatal", False, "CUDA OOM"), - ("no_executor", False, "not available"), - ("no_llm", False, "not initialized"), - ("shutdown", False, "unhealthy"), - ], - ) - async def test_health_check(self, setup, expected_healthy, msg_substr): - """Verify gRPC health check returns correct status and message.""" - if setup == "no_llm": - m = MagicMock() - m.llm = None - m.health_check = self._health_check.__get__(m) - else: - ex = ConcreteExecutor() - if setup == "fatal": - ex._set_fatal_error(RuntimeError("CUDA OOM")) - elif setup == "no_executor": - ex = None - elif setup == "shutdown": - ex.doing_shutdown = True - m = self._make_manager(ex) - - healthy, msg = await m.health_check() - assert healthy is expected_healthy - assert msg_substr in msg - - -# --------------------------------------------------------------------------- -# OpenAI /health endpoint -# --------------------------------------------------------------------------- -class TestOpenAIHealthEndpoint: - """Tests for OpenAI /health endpoint with fatal error shutdown.""" - - @pytest.mark.asyncio - @pytest.mark.parametrize( - "check_healthy,fatal_error,expect_code,sigint", - [ - (True, None, 200, False), - (False, None, 503, False), - (False, RuntimeError("CUDA OOM"), 503, True), - ], - ) - async def test_health(self, check_healthy, fatal_error, expect_code, sigint): - """Verify status code and SIGINT behavior for different health states.""" - from starlette.responses import Response - - server = MagicMock() - server._check_health = Mock(return_value=check_healthy) - executor = Mock() - executor._fatal_error = fatal_error - executor.doing_shutdown = False - server.generator = Mock() - server.generator._executor = executor - - with patch.object(signal, "raise_signal") as mock_sig: - if server._check_health(): - response = Response(status_code=200) - else: - ex = getattr(server.generator, "_executor", None) - if ( - ex is not None - and getattr(ex, "_fatal_error", None) is not None - and not getattr(ex, "doing_shutdown", True) - ): - signal.raise_signal(signal.SIGINT) - response = Response(status_code=503) - - assert response.status_code == expect_code - if sigint: - mock_sig.assert_called_once_with(signal.SIGINT) - else: - mock_sig.assert_not_called() - - -# --------------------------------------------------------------------------- -# BaseLLM._check_health delegation -# --------------------------------------------------------------------------- -class TestBaseLLMCheckHealth: - """Tests for BaseLLM._check_health delegation to executor.""" - - @staticmethod - def _check_health(llm) -> bool: - """Mirrors BaseLLM._check_health logic.""" - if hasattr(llm, "_executor") and llm._executor is not None: - return llm._executor.check_health() - return False - - @pytest.mark.parametrize( - "has_executor,fatal,expected", - [ - (True, False, True), - (True, True, False), - (False, False, False), - ], - ) - def test_delegation(self, has_executor, fatal, expected): - """_check_health delegates to executor.check_health correctly.""" - llm = Mock() - if has_executor: - ex = ConcreteExecutor() - if fatal: - ex._set_fatal_error(RuntimeError("x")) - llm._executor = ex - else: - llm._executor = None - assert self._check_health(llm) is expected - - def test_no_executor_attr(self): - """Returns False when the object has no _executor attribute.""" - assert self._check_health(object()) is False - - # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- diff --git a/tests/unittest/grpc/openengine/test_control.py b/tests/unittest/grpc/openengine/test_control.py index 8723dde4d6eb..38149c34e7b1 100644 --- a/tests/unittest/grpc/openengine/test_control.py +++ b/tests/unittest/grpc/openengine/test_control.py @@ -128,7 +128,7 @@ def _llm(**overrides): supports_token_ids_input=overrides.get("_supports_token_ids_input", True) ), _executor=executor, - _check_health=lambda: executor.check_health(), + check_health=lambda: executor.check_health(), ) @@ -610,7 +610,7 @@ async def test_health_reports_not_ready_for_a_crashed_engine_that_is_not_shut_do async def test_health_survives_a_raising_health_check(error_type): """A probe reports state; it must not turn an engine fault into an RPC error.""" llm = _llm() - llm._check_health = lambda: (_ for _ in ()).throw(error_type("engine gone")) + llm.check_health = lambda: (_ for _ in ()).throw(error_type("engine gone")) servicer = OpenEngineControlServicer(llm, MODEL, _inference(), kv_transfer_backend="") response = await servicer.Health(lifecycle_pb2.HealthRequest(), FakeServicerContext()) diff --git a/tests/unittest/grpc/openengine/test_coordinator.py b/tests/unittest/grpc/openengine/test_coordinator.py index 22cade427513..01a83bfb8a83 100644 --- a/tests/unittest/grpc/openengine/test_coordinator.py +++ b/tests/unittest/grpc/openengine/test_coordinator.py @@ -202,6 +202,7 @@ async def __aiter__(self): handle = ResultHandle() llm = SimpleNamespace( + engine_dead=False, llm_id="shared-engine", args=SimpleNamespace(guided_decoding_backend=None), tokenizer=None, @@ -252,10 +253,11 @@ async def cleaned_up(): async def test_standard_health_tracks_frontend_group_readiness(): async with group() as (coordinator, (frontend, _), _): llm = SimpleNamespace( + engine_dead=False, llm_id="shared-engine", args=SimpleNamespace(guided_decoding_backend=None), tokenizer=None, - _check_health=lambda: True, + check_health=lambda: True, ) server = OpenEngineServer("127.0.0.1", 0, llm, "test-model", frontend=frontend) await server.start() @@ -314,10 +316,11 @@ async def aresult(self): return None llm = SimpleNamespace( + engine_dead=False, llm_id="shared-engine", args=SimpleNamespace(guided_decoding_backend=None), tokenizer=None, - _check_health=lambda: True, + check_health=lambda: True, generate_async=lambda *args, **kwargs: ProbeHandle(), ) original_request = frontend.request @@ -354,6 +357,7 @@ async def test_oversized_request_id_does_not_stop_group(request_id): """A rejected public ID must not kill unrelated requests or strand ownership.""" async with group() as (_, (first, second), failures): llm = SimpleNamespace( + engine_dead=False, llm_id="shared-engine", args=SimpleNamespace(guided_decoding_backend=None), tokenizer=None, diff --git a/tests/unittest/grpc/openengine/test_frontend_lifecycle.py b/tests/unittest/grpc/openengine/test_frontend_lifecycle.py index 6732fd5b76f7..ca4ef9a222af 100644 --- a/tests/unittest/grpc/openengine/test_frontend_lifecycle.py +++ b/tests/unittest/grpc/openengine/test_frontend_lifecycle.py @@ -176,9 +176,10 @@ def on_term(signum, frame): self.llm_id = "shared-engine" self.args = SimpleNamespace(guided_decoding_backend=None) self.tokenizer = None + self.engine_dead = False self._executor = Mock(spec=GenerationExecutorProxy) self._executor.multi_frontend_attach_info.return_value = {"mode": "classic"} - def _check_health(self): + def check_health(self): return not ( frontend == 0 and sys.argv[3] == "child-startup-engine-failure" and (root / "initializing").exists() diff --git a/tests/unittest/grpc/openengine/test_metrics.py b/tests/unittest/grpc/openengine/test_metrics.py index 3e97cc47d4b5..32dd0edb402e 100644 --- a/tests/unittest/grpc/openengine/test_metrics.py +++ b/tests/unittest/grpc/openengine/test_metrics.py @@ -151,7 +151,7 @@ async def __aiter__(self): llm = _FakeLlm([]) llm.args = _args() llm.llm_id = "test-engine" - llm._check_health = lambda: True + llm.check_health = lambda: True reads = [] def get_stats(timeout): diff --git a/tests/unittest/grpc/openengine/test_openengine.py b/tests/unittest/grpc/openengine/test_openengine.py index faf6e92ace17..4e6456a8cfe4 100644 --- a/tests/unittest/grpc/openengine/test_openengine.py +++ b/tests/unittest/grpc/openengine/test_openengine.py @@ -100,6 +100,7 @@ def test_openengine_server_serves_the_control_contract() -> None: answering rather than falling through to the generated base. """ llm = SimpleNamespace( + engine_dead=False, args=SimpleNamespace( max_seq_len=2048, max_batch_size=8, @@ -113,7 +114,7 @@ def test_openengine_server_serves_the_control_contract() -> None: ), llm_id="test-instance", tokenizer=object(), - _check_health=lambda: True, + check_health=lambda: True, ) async def exercise_server() -> None: @@ -160,10 +161,11 @@ def test_health_and_control_are_reachable_on_pod_address() -> None: pytest.skip("No non-loopback IPv4 address is available") llm = SimpleNamespace( + engine_dead=False, args=SimpleNamespace(guided_decoding_backend=None), llm_id="test-instance", tokenizer=None, - _check_health=lambda: True, + check_health=lambda: True, ) async def exercise_server() -> None: @@ -194,10 +196,11 @@ def test_standard_grpc_health_check_and_watch(monkeypatch: pytest.MonkeyPatch) - monkeypatch.setattr(oe_server, "_HEALTH_CHECK_INTERVAL_SECONDS", 0.01) engine = {"healthy": False} llm = SimpleNamespace( + engine_dead=False, args=SimpleNamespace(guided_decoding_backend=None), llm_id="test-instance", tokenizer=None, - _check_health=lambda: engine["healthy"], + check_health=lambda: engine["healthy"], ) async def exercise_server() -> None: @@ -264,10 +267,11 @@ async def check(service: str) -> int: def test_standard_grpc_health_watch_reports_shutdown() -> None: """Shutdown withdraws inference and finishes watchers without failing liveness.""" llm = SimpleNamespace( + engine_dead=False, args=SimpleNamespace(guided_decoding_backend=None), llm_id="test-instance", tokenizer=None, - _check_health=lambda: True, + check_health=lambda: True, ) async def exercise_server() -> None: @@ -319,10 +323,11 @@ def check_health() -> bool: return True llm = SimpleNamespace( + engine_dead=False, args=SimpleNamespace(guided_decoding_backend=None), llm_id="test-instance", tokenizer=None, - _check_health=check_health, + check_health=check_health, ) async def exercise_server() -> None: @@ -373,13 +378,22 @@ def test_fatal_engine_failure_fails_liveness_after_temporary_unavailability( monkeypatch.setattr(oe_server, "_HEALTH_CHECK_INTERVAL_SECONDS", 0.01) engine = {"healthy": True} executor = SimpleNamespace(_fatal_error=None, _engine_dead=False) - llm = SimpleNamespace( - args=SimpleNamespace(guided_decoding_backend=None), - llm_id="test-instance", - tokenizer=None, - _executor=executor, - _check_health=lambda: engine["healthy"] and executor._fatal_error is None, - ) + + class EngineFailureLLM: + """LLM double whose engine_dead derives from the executor's two fatal fields.""" + + args = SimpleNamespace(guided_decoding_backend=None) + llm_id = "test-instance" + tokenizer = None + + def check_health(self) -> bool: + return engine["healthy"] and executor._fatal_error is None + + @property + def engine_dead(self) -> bool: + return executor._fatal_error is not None or bool(executor._engine_dead) + + llm = EngineFailureLLM() async def exercise_server() -> None: server = OpenEngineServer("127.0.0.1", 0, llm, "test-model") @@ -450,10 +464,11 @@ def check_health() -> bool: return True llm = SimpleNamespace( + engine_dead=False, args=SimpleNamespace(guided_decoding_backend=None), llm_id="test-instance", tokenizer=None, - _check_health=check_health, + check_health=check_health, ) async def exercise_server() -> None: @@ -504,6 +519,7 @@ def test_event_startup_requires_loaded_publisher(publisher_enabled: bool) -> Non from tensorrt_llm.grpc.openengine.kv_events import KvEventsUnavailableError llm = SimpleNamespace( + engine_dead=False, args=SimpleNamespace( guided_decoding_backend=None, kv_cache_config=SimpleNamespace( diff --git a/tests/unittest/grpc/openengine/test_private_access_lint_openengine.py b/tests/unittest/grpc/openengine/test_private_access_lint_openengine.py new file mode 100644 index 000000000000..cee2100effc5 --- /dev/null +++ b/tests/unittest/grpc/openengine/test_private_access_lint_openengine.py @@ -0,0 +1,33 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""OpenEngine part of the private-access lint. + +Pull requests that change only tensorrt_llm/grpc/openengine/ are claimed by CBTS +Tier 1 and run only the test-db entries under unittest/grpc/openengine/. This file +applies the rule and baseline of tests/unittest/serve/test_private_access_lint.py +to tensorrt_llm/grpc/openengine/, so those pull requests are still checked. +""" + +import importlib.util +from pathlib import Path + +import pytest + +pytestmark = pytest.mark.cpu_only + +_MAIN = Path(__file__).resolve().parents[2] / "serve" / "test_private_access_lint.py" +_spec = importlib.util.spec_from_file_location("_serve_private_access_lint", _MAIN) +lint = importlib.util.module_from_spec(_spec) +_spec.loader.exec_module(lint) + + +def test_no_new_cross_package_private_access_in_openengine(): + assert (lint.REPO_ROOT / lint.OPENENGINE_SCOPE).is_dir(), ( + f"{lint.OPENENGINE_SCOPE} does not exist. If OpenEngine moved, point OPENENGINE_SCOPE " + "in tests/unittest/serve/test_private_access_lint.py at the new directory, and move " + "this file and its l0_cpu.yml entry to the test path the CBTS OpenEngine rule keeps " + "(jenkins/scripts/cbts/cbts/rules/openengine.py)." + ) + rows = lint.scan(prefixes=(lint.OPENENGINE_SCOPE,)) + errors = lint._compare(rows, lint.load_baseline(), prefix=lint.OPENENGINE_SCOPE) + assert not errors, "\n".join(errors) diff --git a/tests/unittest/grpc/openengine/test_servicer.py b/tests/unittest/grpc/openengine/test_servicer.py index a953bceba01f..178fdc51fbaa 100644 --- a/tests/unittest/grpc/openengine/test_servicer.py +++ b/tests/unittest/grpc/openengine/test_servicer.py @@ -95,6 +95,7 @@ def __init__(self, results: list[Any]) -> None: self.tokenizer = _FakeTokenizer() self.result_handle = _FakeResultHandle(results) self.generate_kwargs = None + self.engine_dead = False def generate_async(self, **kwargs: Any) -> _FakeResultHandle: self.generate_kwargs = kwargs diff --git a/tests/unittest/grpc/smg/test_smg_health.py b/tests/unittest/grpc/smg/test_smg_health.py new file mode 100644 index 000000000000..33467eadca8c --- /dev/null +++ b/tests/unittest/grpc/smg/test_smg_health.py @@ -0,0 +1,71 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Production tests of GrpcRequestManager.health_check. + +The (healthy, message) pair feeds HealthCheckResponse.status, so every message +string is wire-visible and pinned byte for byte, including the first-line rule +and the 200-character cut for fatal errors. +""" + +import asyncio +from types import SimpleNamespace + +import pytest + +pytestmark = pytest.mark.cpu_only + +pytest.importorskip("smg_grpc_proto") + +from tensorrt_llm.grpc.smg.request_manager import GrpcRequestManager # noqa: E402 + + +class _FakeLLM: + """LLM double with exactly the public health surface health_check reads.""" + + def __init__(self, healthy=True, fatal_error=None, encode_only=False, check_raises=None): + self.args = SimpleNamespace(encode_only=encode_only) + self.fatal_error = fatal_error + self._healthy = healthy + self._check_raises = check_raises + + def check_health(self): + if self._check_raises is not None: + raise self._check_raises + return self._healthy + + +def _health_check(llm): + return asyncio.run(GrpcRequestManager(llm).health_check()) + + +def test_healthy_llm_answers_ok(): + assert _health_check(_FakeLLM(healthy=True)) == (True, "OK") + + +def test_missing_llm_answers_not_initialized(): + assert _health_check(None) == (False, "LLM not initialized") + + +def test_encode_only_llm_answers_executor_not_available(): + assert _health_check(_FakeLLM(encode_only=True)) == (False, "Executor is not available") + + +def test_unhealthy_llm_without_fatal_error_answers_executor_unhealthy(): + assert _health_check(_FakeLLM(healthy=False)) == (False, "Executor is unhealthy") + + +def test_fatal_error_message_names_the_exception_type(): + result = _health_check(_FakeLLM(healthy=False, fatal_error=RuntimeError("CUDA OOM"))) + assert result == (False, "RuntimeError: CUDA OOM") + + +def test_fatal_error_message_keeps_the_first_line_cut_to_200_characters(): + long_first_line = "x" * 300 + result = _health_check( + _FakeLLM(healthy=False, fatal_error=RuntimeError(f"{long_first_line}\nsecond line")) + ) + assert result == (False, f"RuntimeError: {'x' * 200}") + + +def test_raising_health_check_answers_error(): + assert _health_check(_FakeLLM(check_raises=ValueError("boom"))) == (False, "Error: boom") diff --git a/tests/unittest/llmapi/test_llm_encode.py b/tests/unittest/llmapi/test_llm_encode.py index 44d9db731e1c..5714f41a2979 100644 --- a/tests/unittest/llmapi/test_llm_encode.py +++ b/tests/unittest/llmapi/test_llm_encode.py @@ -237,8 +237,8 @@ def test_encode_add_special_tokens_false(bert_encode_llm): def test_check_health_encoder_only(bert_encode_llm): - """_check_health() returns True for a live encoder-only LLM.""" - assert bert_encode_llm._check_health() is True + """check_health() returns True for a live encoder-only LLM.""" + assert bert_encode_llm.check_health() is True # --------------------------------------------------------------------------- # diff --git a/tests/unittest/llmapi/test_llm_public_runtime_api.py b/tests/unittest/llmapi/test_llm_public_runtime_api.py new file mode 100644 index 000000000000..ee2c036c50be --- /dev/null +++ b/tests/unittest/llmapi/test_llm_public_runtime_api.py @@ -0,0 +1,180 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Production tests of the public runtime health API the serving frontends use. + +Covers BaseLLM.check_health(), BaseLLM.fatal_error, BaseLLM.shutdown_started, +BaseLLM.engine_dead and the executor-side GenerationExecutor.fatal_error and +engine_dead, against the real classes (no mirrors; the LLM objects are built +with object.__new__ plus exactly the attributes BaseLLM.__init__ sets). +""" + +import pytest + +from tensorrt_llm.executor.executor import GenerationExecutor +from tensorrt_llm.executor.proxy import GenerationExecutorProxy +from tensorrt_llm.llmapi.llm import LLM + +pytestmark = pytest.mark.cpu_only + + +class _NoOpExecutor(GenerationExecutor): + """Minimal concrete GenerationExecutor; health logic comes from the base.""" + + def submit(self, request): + raise NotImplementedError("health tests never submit requests") + + def abort_request(self, request_id): + pass + + def shutdown(self): + self.doing_shutdown = True + + +def _llm(executor, encode_only=False, encoder_executor=None): + """An LLM with exactly the attributes BaseLLM.__init__ sets for health.""" + llm = object.__new__(LLM) + llm._executor = executor + llm._encode_only = encode_only + llm._encoder_executor = encoder_executor + return llm + + +def _proxy(engine_dead=False, fatal_error=None): + """A proxy with exactly the two fields its engine_dead property reads.""" + proxy = object.__new__(GenerationExecutorProxy) + proxy._engine_dead = engine_dead + proxy._fatal_error = fatal_error + return proxy + + +# --------------------------------------------------------------------------- +# BaseLLM.check_health() +# --------------------------------------------------------------------------- +def test_check_health_encode_only_with_encoder_executor(): + assert _llm(None, encode_only=True, encoder_executor=object()).check_health() is True + + +def test_check_health_encode_only_without_encoder_executor(): + assert _llm(None, encode_only=True, encoder_executor=None).check_health() is False + + +def test_check_health_healthy_executor(): + assert _llm(_NoOpExecutor()).check_health() is True + + +def test_check_health_fatal_error(): + executor = _NoOpExecutor() + executor._set_fatal_error(RuntimeError("CUDA OOM")) + assert _llm(executor).check_health() is False + + +def test_check_health_doing_shutdown(): + executor = _NoOpExecutor() + executor.doing_shutdown = True + assert _llm(executor).check_health() is False + + +def test_check_health_without_executor(): + assert _llm(None).check_health() is False + + +def test_check_health_drains_queued_fatal_error(): + """A fatal error still in the error queue is promoted by check_health().""" + executor = _NoOpExecutor() + queued = RuntimeError("engine died in a background thread") + executor._error_queue.put(queued) + + llm = _llm(executor) + assert llm.check_health() is False + assert llm.fatal_error is queued + assert llm.shutdown_started is True + + +def test_check_health_private_alias_is_the_same_function(): + """Kept for one release for out-of-tree callers that look the name up.""" + assert LLM._check_health is LLM.check_health + + +# --------------------------------------------------------------------------- +# BaseLLM.fatal_error +# --------------------------------------------------------------------------- +def test_fatal_error_none_without_executor(): + assert _llm(None).fatal_error is None + + +def test_fatal_error_none_on_healthy_executor(): + assert _llm(_NoOpExecutor()).fatal_error is None + + +def test_fatal_error_returns_recorded_exception_first_one_wins(): + executor = _NoOpExecutor() + first = RuntimeError("first fatal error") + executor._set_fatal_error(first) + executor._set_fatal_error(RuntimeError("second fatal error")) + assert _llm(executor).fatal_error is first + + +# --------------------------------------------------------------------------- +# BaseLLM.shutdown_started +# --------------------------------------------------------------------------- +def test_shutdown_started_follows_doing_shutdown(): + executor = _NoOpExecutor() + llm = _llm(executor) + assert llm.shutdown_started is False + executor.doing_shutdown = True + assert llm.shutdown_started is True + + +def test_shutdown_started_true_without_executor(): + assert _llm(None).shutdown_started is True + + +def test_shutdown_started_false_for_encode_only(): + assert _llm(None, encode_only=True, encoder_executor=object()).shutdown_started is False + + +# --------------------------------------------------------------------------- +# BaseLLM.engine_dead +# --------------------------------------------------------------------------- +def test_engine_dead_false_without_executor(): + assert _llm(None).engine_dead is False + + +def test_engine_dead_follows_the_executor(): + executor = _NoOpExecutor() + llm = _llm(executor) + assert llm.engine_dead is False + executor._set_fatal_error(RuntimeError("worker died")) + assert llm.engine_dead is True + + +# --------------------------------------------------------------------------- +# GenerationExecutor.fatal_error and engine_dead +# --------------------------------------------------------------------------- +def test_executor_fatal_error_property(): + executor = _NoOpExecutor() + assert executor.fatal_error is None + error = RuntimeError("fatal") + executor._set_fatal_error(error) + assert executor.fatal_error is error + + +def test_executor_engine_dead_property(): + executor = _NoOpExecutor() + assert executor.engine_dead is False + executor._set_fatal_error(RuntimeError("fatal")) + assert executor.engine_dead is True + + +@pytest.mark.parametrize( + "engine_dead,fatal_error,expected", + [ + (False, None, False), + (True, None, True), + (False, RuntimeError("fatal"), True), + (True, RuntimeError("fatal"), True), + ], +) +def test_proxy_engine_dead_reads_both_fields(engine_dead, fatal_error, expected): + """_mark_engine_dead can set the sticky flag before the error is recorded.""" + assert _proxy(engine_dead=engine_dead, fatal_error=fatal_error).engine_dead is expected diff --git a/tests/unittest/serve/__init__.py b/tests/unittest/serve/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tests/unittest/serve/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tests/unittest/serve/api/__init__.py b/tests/unittest/serve/api/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tests/unittest/serve/api/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 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..52a7a9daf028 --- /dev/null +++ b/tests/unittest/serve/api/visual_gen/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tests/unittest/serve/api/visual_gen/test_private_access_lint_visual_gen.py b/tests/unittest/serve/api/visual_gen/test_private_access_lint_visual_gen.py new file mode 100644 index 000000000000..b121bc555629 --- /dev/null +++ b/tests/unittest/serve/api/visual_gen/test_private_access_lint_visual_gen.py @@ -0,0 +1,27 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""VisualGen part of the private-access lint. + +Pull requests that change only VisualGen files are claimed by CBTS Tier 1 and run +only test-db entries whose path contains visual_gen/. This file applies the rule +and baseline of tests/unittest/serve/test_private_access_lint.py to +tensorrt_llm/serve/api/visual_gen/, so those pull requests are still checked. +""" + +import importlib.util +from pathlib import Path + +import pytest + +pytestmark = pytest.mark.cpu_only + +_MAIN = Path(__file__).resolve().parents[2] / "test_private_access_lint.py" +_spec = importlib.util.spec_from_file_location("_serve_private_access_lint", _MAIN) +lint = importlib.util.module_from_spec(_spec) +_spec.loader.exec_module(lint) + + +def test_no_new_cross_package_private_access_in_visual_gen(): + rows = lint.scan(prefixes=(lint.VISUAL_GEN_SCOPE,)) + errors = lint._compare(rows, lint.load_baseline(), prefix=lint.VISUAL_GEN_SCOPE) + assert not errors, "\n".join(errors) diff --git a/tests/unittest/serve/private_access_baseline.yaml b/tests/unittest/serve/private_access_baseline.yaml new file mode 100644 index 000000000000..123c9bdbd182 --- /dev/null +++ b/tests/unittest/serve/private_access_baseline.yaml @@ -0,0 +1,53 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# Cross-package private accesses from tensorrt_llm/serve (without serve/scripts/), +# tensorrt_llm/commands and tensorrt_llm/grpc, checked by test_private_access_lint.py. +# This list only shrinks. When you remove an access, lower `count` (delete the entry +# at 0) and lower MAX_TOTAL in the test. Never add an entry: add a public accessor on +# the class that owns the member. A pure move changes `path` and `scope` only. +# `owner` and `replacement` are for readers; the test uses path, scope, member, count. +version: 1 +entries: +- {path: tensorrt_llm/commands/serve.py, scope: _bind_launcher_uds, member: _executor, count: 1, owner: llmapi/llm.py BaseLLM, replacement: a public accessor for the multi-frontend IPC directory} +- {path: tensorrt_llm/commands/serve.py, scope: _bind_launcher_uds, member: _multi_frontend_ipc_dir, count: 1, owner: executor/proxy.py GenerationExecutorProxy, replacement: a public accessor for the multi-frontend IPC directory} +- {path: tensorrt_llm/commands/serve.py, scope: _run_fleet_worker_impl, member: _config, count: 2, owner: serve/openai_disagg_server.py OpenAIDisaggServer, replacement: OpenAIDisaggServer.address} +- {path: tensorrt_llm/grpc/openengine/capabilities.py, scope: _loaded_model_type, member: _hf_model_config, count: 1, owner: llmapi/llm.py BaseLLM, replacement: BaseLLM.hf_model_config} +- {path: tensorrt_llm/grpc/openengine/capabilities.py, scope: resolve_media_capabilities, member: _registered_model_type, count: 1, owner: inputs input-processor registry, replacement: a public registry accessor for the registered model type} +- {path: tensorrt_llm/grpc/openengine/control.py, scope: OpenEngineControlServicer.GetServerInfo, member: _executor, count: 1, owner: llmapi/llm.py BaseLLM, replacement: "BaseLLM.get_kv_cache_capacity(), BaseLLM.get_openengine_discovery()"} +- {path: tensorrt_llm/grpc/openengine/control.py, scope: OpenEngineControlServicer.GetServerInfo, member: _openengine_discovery, count: 1, owner: llmapi/llm_args.py BaseLlmArgs, replacement: BaseLlmArgs.openengine_discovery} +- {path: tensorrt_llm/grpc/openengine/control.py, scope: OpenEngineControlServicer.__init__, member: _enable_routing_load, count: 1, owner: llmapi/llm_args.py BaseLlmArgs, replacement: BaseLlmArgs.enable_routing_load} +- {path: tensorrt_llm/grpc/openengine/control.py, scope: OpenEngineControlServicer._get_kv_cache_load, member: _executor, count: 1, owner: llmapi/llm.py BaseLLM, replacement: BaseLLM.get_kv_cache_load()} +- {path: tensorrt_llm/grpc/openengine/control.py, scope: OpenEngineControlServicer._requested_sources, member: _executor, count: 1, owner: llmapi/llm.py BaseLLM, replacement: BaseLLM.get_openengine_discovery()} +- {path: tensorrt_llm/grpc/openengine/formatting.py, scope: _token_holdback, member: _stop_word_ids, count: 1, owner: sampling_params.py SamplingParams, replacement: SamplingParams.stop_word_ids} +- {path: tensorrt_llm/grpc/openengine/node.py, scope: initialize_node_discovery, member: _openengine_discovery, count: 2, owner: "llmapi/llm_args.py BaseLlmArgs (read); _torch/pyexecutor/py_executor.py PyExecutor (write)", replacement: "BaseLlmArgs.openengine_discovery; PyExecutor.set_openengine_discovery()"} +- {path: tensorrt_llm/grpc/openengine/node.py, scope: initialize_node_discovery, member: _openengine_node_server, count: 1, owner: _torch/pyexecutor/py_executor.py PyExecutor, replacement: PyExecutor.set_openengine_node_server()} +- {path: tensorrt_llm/grpc/openengine/request_mapping.py, scope: sampling_params_from_request, member: _set_request_provided_fields, count: 1, owner: sampling_params.py SamplingParams, replacement: SamplingParams.set_request_provided_fields()} +- {path: tensorrt_llm/grpc/openengine/server.py, scope: OpenEngineServer.start, member: _executor, count: 1, owner: llmapi/llm.py BaseLLM, replacement: BaseLLM.get_kv_cache_capacity()} +- {path: tensorrt_llm/serve/_frontend_processes.py, scope: _spawn_attached_frontends, member: _executor, count: 1, owner: llmapi/llm.py BaseLLM, replacement: BaseLLM.multi_frontend_attach_info()} +- {path: tensorrt_llm/serve/chat_utils.py, scope: parse_chat_message_content_part, member: _multimodal_server_config, count: 4, owner: inputs/utils.py MultimodalDataTracker, replacement: MultimodalDataTracker.multimodal_server_config} +- {path: tensorrt_llm/serve/openai_protocol.py, scope: _record_sampling_params_request_fields, member: _set_request_provided_fields, count: 1, owner: sampling_params.py SamplingParams, replacement: SamplingParams.set_request_provided_fields()} +- {path: tensorrt_llm/serve/openai_server.py, scope: OpenAIServer._create_chat_response, member: _postprocess_result, count: 1, owner: executor/result.py CompletionOutput, replacement: CompletionOutput.postprocess_result} +- {path: tensorrt_llm/serve/openai_server.py, scope: OpenAIServer._init_embedding_batcher, member: _encoder_executor, count: 1, owner: llmapi/llm.py BaseLLM, replacement: BaseLLM.encode_limits} +- {path: tensorrt_llm/serve/openai_server.py, scope: OpenAIServer._init_llm, member: _hf_model_dir, count: 1, owner: llmapi/llm.py BaseLLM, replacement: BaseLLM.hf_model_dir} +- {path: tensorrt_llm/serve/openai_server.py, scope: OpenAIServer._on_launcher_lost, member: _executor, count: 1, owner: llmapi/llm.py BaseLLM, replacement: public fatal-error recording on the LLM} +- {path: tensorrt_llm/serve/openai_server.py, scope: OpenAIServer._on_launcher_lost, member: _set_fatal_error, count: 1, owner: executor/executor.py GenerationExecutor, replacement: public fatal-error recording on the LLM} +- {path: tensorrt_llm/serve/openai_server.py, scope: OpenAIServer.chat_harmony.create_streaming_generator, member: _postprocess_result, count: 2, owner: executor/result.py CompletionOutput, replacement: CompletionOutput.postprocess_result} +- {path: tensorrt_llm/serve/openai_server.py, scope: OpenAIServer.openai_chat.chat_stream_generator, member: _postprocess_result, count: 2, owner: executor/result.py CompletionOutput, replacement: CompletionOutput.postprocess_result} +- {path: tensorrt_llm/serve/openai_server.py, scope: OpenAIServer.openai_completion.completion_generator, member: _postprocess_result, count: 1, owner: executor/result.py CompletionOutput, replacement: CompletionOutput.postprocess_result} +- {path: tensorrt_llm/serve/openai_server.py, scope: OpenAIServer.openai_completion.completion_response, member: _postprocess_result, count: 1, owner: executor/result.py CompletionOutput, replacement: CompletionOutput.postprocess_result} +- {path: tensorrt_llm/serve/openai_server.py, scope: OpenAIServer.openai_image_generation, member: _save_bytes, count: 1, owner: visual_gen/output.py, replacement: public VisualGen output encoder} +- {path: tensorrt_llm/serve/openai_server.py, scope: OpenAIServer.openai_responses.create_response, member: _postprocess_result, count: 1, owner: executor/result.py CompletionOutput, replacement: CompletionOutput.postprocess_result} +- {path: tensorrt_llm/serve/openai_server.py, scope: OpenAIServer.openai_responses.create_streaming_generator, member: _postprocess_result, count: 1, owner: executor/result.py CompletionOutput, replacement: CompletionOutput.postprocess_result} +- {path: tensorrt_llm/serve/openai_server.py, scope: OpenAIServer.register_routes, member: _executor, count: 1, owner: llmapi/llm.py BaseLLM, replacement: BaseLLM.resource_governor_queue} +- {path: tensorrt_llm/serve/openai_server.py, scope: _record_generator_termination, member: _executor, count: 1, owner: llmapi/llm.py BaseLLM, replacement: BaseLLM.engine_in_worker_process} +- {path: tensorrt_llm/serve/postprocess_handlers.py, scope: chat_harmony_streaming_post_processor, member: _done, count: 1, owner: executor/result.py GenerationResultBase, replacement: GenerationResultBase.finished} +- {path: tensorrt_llm/serve/postprocess_handlers.py, scope: chat_response_post_processor, member: _spec_dec_counters, count: 1, owner: executor/result.py CompletionOutput, replacement: CompletionOutput.spec_dec_counters} +- {path: tensorrt_llm/serve/postprocess_handlers.py, scope: chat_stream_post_processor, member: _done, count: 1, owner: executor/result.py GenerationResultBase, replacement: GenerationResultBase.finished} +- {path: tensorrt_llm/serve/postprocess_handlers.py, scope: chat_stream_post_processor, member: _spec_dec_counters, count: 1, owner: executor/result.py CompletionOutput, replacement: CompletionOutput.spec_dec_counters} +- {path: tensorrt_llm/serve/postprocess_handlers.py, scope: completion_response_post_processor, member: _spec_dec_counters, count: 1, owner: executor/result.py CompletionOutput, replacement: CompletionOutput.spec_dec_counters} +- {path: tensorrt_llm/serve/postprocess_handlers.py, scope: completion_stream_post_processor, member: _done, count: 1, owner: executor/result.py GenerationResultBase, replacement: GenerationResultBase.finished} +- {path: tensorrt_llm/serve/postprocess_handlers.py, scope: completion_stream_post_processor, member: _last_text_len, count: 1, owner: executor/result.py CompletionOutput, replacement: CompletionOutput.text_diff_offset} +- {path: tensorrt_llm/serve/postprocess_handlers.py, scope: completion_stream_post_processor, member: _spec_dec_counters, count: 1, owner: executor/result.py CompletionOutput, replacement: CompletionOutput.spec_dec_counters} +- {path: tensorrt_llm/serve/postprocess_handlers.py, scope: responses_api_streaming_post_processor, member: _done, count: 1, owner: executor/result.py GenerationResultBase, replacement: GenerationResultBase.finished} +- {path: tensorrt_llm/serve/responses_utils.py, scope: ResponsesStreamingProcessor.process_single_output, member: _done, count: 1, owner: executor/result.py GenerationResultBase, replacement: GenerationResultBase.finished} +- {path: tensorrt_llm/serve/router_utils.py, scope: BlockHashMixin._compute_block_hashes, member: _hash_cache_salt, count: 1, owner: runtime/kv_cache_manager_v2 KVCacheEventManager (nanobind), replacement: public hash_cache_salt binding} diff --git a/tests/unittest/serve/test_private_access_lint.py b/tests/unittest/serve/test_private_access_lint.py new file mode 100644 index 000000000000..ac48c7e81723 --- /dev/null +++ b/tests/unittest/serve/test_private_access_lint.py @@ -0,0 +1,344 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Private-access lint for the serving frontends. + +The HTTP server (tensorrt_llm/serve, without serve/scripts), the CLI +(tensorrt_llm/commands) and the gRPC frontends (tensorrt_llm/grpc) may use only +the public API of other TensorRT-LLM packages. This test counts every access to +a private member that belongs to another package and compares the result with +private_access_baseline.yaml, which lists the accesses that existed when the +lint was introduced. The baseline only shrinks. + +Counting rule: one AST node per access. An access is an attribute +`._name` or a getattr/hasattr/setattr/delattr call with the literal +"_name", where the receiver is not `self`, `cls` or `super()` and `_name` has one +leading underscore and is not a dunder. `a._b._c` is two accesses. Names +imported with `from m import _x` are not accesses. Ownership: OVERRIDES first, +then a name index over tensorrt_llm/**/*.py and *.pyi. + +Run `python tests/unittest/serve/test_private_access_lint.py --report` to print +the current inventory. +""" + +import ast +import collections +import sys +from pathlib import Path + +import pytest +import yaml + +pytestmark = pytest.mark.cpu_only + +REPO_ROOT = Path(__file__).resolve().parents[3] +BASELINE_FILE = Path(__file__).resolve().with_name("private_access_baseline.yaml") +SCOPES = ("tensorrt_llm/serve/", "tensorrt_llm/commands/", "tensorrt_llm/grpc/") +EXCLUDED = ("tensorrt_llm/serve/scripts/", "tensorrt_llm/grpc/openengine/_generated/") +VISUAL_GEN_SCOPE = "tensorrt_llm/serve/api/visual_gen/" +OPENENGINE_SCOPE = "tensorrt_llm/grpc/openengine/" + +# Lower this when you remove accesses; never raise it. +MAX_TOTAL = 50 + +# (path, member) -> (verdict, reason). Verdicts: "cross", "same-package", "third-party". +# Needed only where the name index cannot tell: the name is defined both inside and +# outside the scanning package, or a third-party name collides with a TensorRT-LLM name. +OVERRIDES = { + ("tensorrt_llm/serve/_telemetry.py", "_serve"): ("third-party", "uvicorn.Server._serve"), + ("tensorrt_llm/serve/_telemetry.py", "_captured_signals"): ( + "third-party", + "uvicorn.Server._captured_signals", + ), + ("tensorrt_llm/serve/cluster_storage.py", "_loop"): ("third-party", "asyncio.Queue._loop"), + ("tensorrt_llm/serve/render/fingerprint.py", "_mergeable_ranks"): ( + "third-party", + "tiktoken Encoding", + ), + ("tensorrt_llm/serve/render/fingerprint.py", "_special_tokens"): ( + "third-party", + "tiktoken Encoding", + ), + ("tensorrt_llm/serve/render/fingerprint.py", "_pat_str"): ("third-party", "tiktoken Encoding"), + ("tensorrt_llm/serve/router.py", "_lock"): ("same-package", "serve/router.py Router._lock"), + # Generic names of serve's own router objects: the name index classifies them + # correctly today, but a definition of the same name in another package would + # turn them into "unresolved" and fail unrelated pull requests. + ("tensorrt_llm/serve/router.py", "_servers"): ( + "same-package", + "serve/router.py Router._servers", + ), + ("tensorrt_llm/serve/router.py", "_server_info"): ( + "same-package", + "serve/router.py Router._server_info", + ), + ("tensorrt_llm/serve/disagg_coordinator.py", "_server_info"): ( + "same-package", + "serve/router.py Router._server_info", + ), + ("tensorrt_llm/serve/disagg_coordinator.py", "_local"): ( + "same-package", + "serve/router.py CoordinatorDelegatingRouter._local", + ), + ("tensorrt_llm/serve/openai_disagg_service.py", "_local"): ( + "same-package", + "serve/router.py CoordinatorDelegatingRouter._local", + ), + # _DecodedMediaReservation reads its own module's _DecodedMediaBudget; the + # generic names are also defined elsewhere, so the index cannot tell. + ("tensorrt_llm/grpc/openengine/request_mapping.py", "_condition"): ( + "same-package", + "grpc/openengine/request_mapping.py _DecodedMediaBudget._condition", + ), + ("tensorrt_llm/grpc/openengine/request_mapping.py", "_cancelled"): ( + "same-package", + "grpc/openengine/request_mapping.py _DecodedMediaBudget._cancelled", + ), +} + +FIX_HINT = ( + "Frontends use only public LLM, executor, result and SamplingParams APIs. Add a public " + "accessor on the class that owns the member (in the owning package) and call it here. " + "Do not add the access to tests/unittest/serve/private_access_baseline.yaml: the baseline " + "only shrinks. If the member belongs to this package or to a third-party library, add an " + "OVERRIDES entry in tests/unittest/serve/test_private_access_lint.py with the reason. " + "Runtime owners (@NVIDIA/trt-llm-runtime-devs) review both files." +) + + +def _is_private(name): + return name.startswith("_") and not (name.startswith("__") and name.endswith("__")) + + +def _package_of(rel): + parts = rel.split("/") + return parts[1][:-3] if len(parts) == 2 else parts[1] + + +def _is_self(node): + if isinstance(node, ast.Name) and node.id in ("self", "cls"): + return True + return ( + isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == "super" + ) + + +class _Candidates(ast.NodeVisitor): + def __init__(self): + self.scope, self.rows = [], [] + + def _qualname(self): + return ".".join(self.scope) or "" + + def visit_FunctionDef(self, node): + self.scope.append(node.name) + self.generic_visit(node) + self.scope.pop() + + visit_AsyncFunctionDef = visit_FunctionDef + visit_ClassDef = visit_FunctionDef + + def visit_Attribute(self, node): + if _is_private(node.attr) and not _is_self(node.value): + self.rows.append((node.lineno, node.attr, ast.unparse(node.value), self._qualname())) + self.generic_visit(node) + + def visit_Call(self, node): + func = node.func + if ( + isinstance(func, ast.Name) + and func.id in ("getattr", "hasattr", "setattr", "delattr") + and len(node.args) >= 2 + and isinstance(node.args[1], ast.Constant) + and isinstance(node.args[1].value, str) + and _is_private(node.args[1].value) + and not _is_self(node.args[0]) + ): + self.rows.append( + (node.lineno, node.args[1].value, ast.unparse(node.args[0]), self._qualname()) + ) + self.generic_visit(node) + + +class _Definitions(ast.NodeVisitor): + """Record which package defines each private member name.""" + + def __init__(self, rel, index): + self.rel, self.index = rel, index + + def _add(self, name): + if _is_private(name): + self.index[name].add(_package_of(self.rel)) + + def visit_ClassDef(self, node): + for stmt in node.body: + if isinstance(stmt, (ast.FunctionDef, ast.AsyncFunctionDef)): + self._add(stmt.name) + elif isinstance(stmt, ast.Assign): + for target in stmt.targets: + if isinstance(target, ast.Name): + self._add(target.id) + elif isinstance(stmt, ast.AnnAssign) and isinstance(stmt.target, ast.Name): + self._add(stmt.target.id) + self.generic_visit(node) + + def _targets(self, targets): + for target in targets: + if isinstance(target, ast.Attribute): + if isinstance(target.value, ast.Name) and target.value.id in ("self", "cls"): + self._add(target.attr) + elif not self.rel.startswith(SCOPES): + # `obj_or_class._x = ...` outside the frontends also defines `_x`. + self._add(target.attr) + elif isinstance(target, (ast.Tuple, ast.List)): + self._targets(target.elts) + + def visit_Assign(self, node): + self._targets(node.targets) + self.generic_visit(node) + + def visit_AnnAssign(self, node): + self._targets([node.target]) + self.generic_visit(node) + + def visit_AugAssign(self, node): + self._targets([node.target]) + self.generic_visit(node) + + +def _scanned_files(root, prefixes): + for prefix in prefixes: + base = root / prefix + if not base.is_dir(): + continue + for path in sorted(base.rglob("*.py")): + rel = path.relative_to(root).as_posix() + if not rel.startswith(EXCLUDED): + yield rel + + +def _name_index(root, names): + index = collections.defaultdict(set) + for path in (root / "tensorrt_llm").rglob("*.py*"): + if path.suffix not in (".py", ".pyi"): + continue + text = path.read_text(encoding="utf-8") + if not any(name in text for name in names): + continue + try: + tree = ast.parse(text) + except SyntaxError: + continue + _Definitions(path.relative_to(root).as_posix(), index).visit(tree) + return index + + +def scan(root=REPO_ROOT, prefixes=SCOPES): + """Return one row per access: (path, line, member, receiver, scope, verdict, reason).""" + candidates = [] + for rel in _scanned_files(root, prefixes): + visitor = _Candidates() + visitor.visit(ast.parse((root / rel).read_text(encoding="utf-8"))) + candidates.extend((rel,) + row for row in visitor.rows) + index = _name_index(root, {member for _, _, member, _, _ in candidates}) + rows = [] + for rel, line, member, receiver, scope in candidates: + package = _package_of(rel) + if (rel, member) in OVERRIDES: + verdict, reason = OVERRIDES[(rel, member)] + else: + owners = index.get(member, set()) + outside = sorted(owner for owner in owners if owner != package) + if package in owners and outside: + verdict, reason = "unresolved", ", ".join(sorted(owners)) + elif package in owners: + verdict, reason = "same-package", package + elif outside: + verdict, reason = "cross", ", ".join(outside) + else: + verdict, reason = "third-party", "not defined in tensorrt_llm" + rows.append((rel, line, member, receiver, scope, verdict, reason)) + return rows + + +def load_baseline(): + data = yaml.safe_load(BASELINE_FILE.read_text()) + return {(e["path"], e["scope"], e["member"]): e for e in data["entries"]} + + +def _compare(rows, baseline, prefix=""): + errors = [] + current = collections.defaultdict(list) + for row in rows: + if row[5] == "unresolved": + errors.append( + f"Cannot classify `{row[3]}.{row[2]}` at {row[0]}:{row[1]}: `{row[2]}` is defined " + f"both in this package and elsewhere ({row[6]}). If the receiver is an object from " + f"another package, this is a new cross-package private access. {FIX_HINT} If it " + "is this package's own member, add an OVERRIDES entry with verdict same-package." + ) + elif row[5] == "cross": + current[(row[0], row[4], row[2])].append(row) + for key, found in sorted(current.items()): + allowed = baseline.get(key, {}).get("count", 0) + if len(found) > allowed: + where = ", ".join(f"{r[0]}:{r[1]}" for r in found) + errors.append( + f"New private access to `{key[2]}` (owned by {found[0][6]}) in {key[1]} " + f"at {where}: {len(found)} found, baseline allows {allowed}. {FIX_HINT}" + ) + for key, entry in sorted(baseline.items()): + if not key[0].startswith(prefix): + continue + found = len(current.get(key, [])) + if found < entry["count"]: + errors.append( + f"Stale baseline entry: {key[0]} {key[1]} `{key[2]}` lists {entry['count']} " + f"access(es) but the code has {found}. Lower `count` (delete the entry at 0) in " + f"{BASELINE_FILE.name} and lower MAX_TOTAL in this PR. If the code moved, change " + "the entry's path and scope instead; the total must not grow." + ) + accessed = {(row[0], row[2]) for row in rows} + for path, member in sorted(OVERRIDES): + if path.startswith(prefix) and (path, member) not in accessed: + errors.append( + f"Stale OVERRIDES entry: {path} `{member}` matches no access. If the code moved, " + "change the key's path to the new location in this PR, as for a baseline entry; " + "if the access is gone, delete the entry." + ) + return errors + + +def test_no_new_cross_package_private_access(): + errors = _compare(scan(), load_baseline()) + assert not errors, "\n".join(errors) + + +def test_baseline_only_shrinks(): + total = sum(e["count"] for e in load_baseline().values()) + assert total <= MAX_TOTAL, ( + f"{BASELINE_FILE.name} lists {total} accesses, more than MAX_TOTAL={MAX_TOTAL}. " + f"The baseline only shrinks. {FIX_HINT}" + ) + + +@pytest.mark.parametrize( + "source,expected", + [ + ("self._x\ncls._x\nsuper()._x()", []), + ("llm._executor", ["_executor"]), + ("llm._executor._fatal_error", ["_executor", "_fatal_error"]), + ("getattr(self._llm, '_check_health', None)", ["_check_health"]), + ("hasattr(self, '_x')", []), + ("from m import _private\nobj.__dict__", []), + ], +) +def test_counting_rule(source, expected): + visitor = _Candidates() + visitor.visit(ast.parse(source)) + assert sorted(row[1] for row in visitor.rows) == sorted(expected) + + +if __name__ == "__main__": + if "--report" in sys.argv: + for row in scan(): + print("\t".join(str(x) for x in row)) diff --git a/tests/unittest/serve/test_serve_health.py b/tests/unittest/serve/test_serve_health.py new file mode 100644 index 000000000000..2717df0fa941 --- /dev/null +++ b/tests/unittest/serve/test_serve_health.py @@ -0,0 +1,120 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Production tests of the HTTP /health handler. + +Calls the production OpenAIServer._check_health and OpenAIServer.health on a +stand-in object that carries exactly the attributes they read (the shared +server fixture factory does not exist yet), with a fake generator class per +case instead of a MagicMock, which would answer any member with a truthy mock. +""" + +import asyncio +import types +from unittest import mock + +import pytest + +from tensorrt_llm.serve.openai_server import OpenAIServer + +pytestmark = pytest.mark.cpu_only + + +class _FakeGenerator: + """Generator double with exactly the public health surface /health reads.""" + + def __init__(self, healthy=True, fatal_error=None, shutdown_started=False): + self._healthy = healthy + self.fatal_error = fatal_error + self.shutdown_started = shutdown_started + + def check_health(self): + return self._healthy + + +class _NoHealthGenerator: + """A VisualGen-style generator with no health attributes at all.""" + + +class _DeadBatcher: + """An embedding batcher whose worker thread has exited.""" + + @staticmethod + def is_alive(): + return False + + +def _server(generator, is_visual_gen=False, embedding_batcher=None): + stub = types.SimpleNamespace( + embedding_batcher=embedding_batcher, + generator=generator, + _is_visual_gen=is_visual_gen, + ) + stub._check_health = types.MethodType(OpenAIServer._check_health, stub) + return stub + + +def _health(server): + with ( + mock.patch("tensorrt_llm.serve.openai_server.signal.raise_signal") as raise_signal, + mock.patch( + "tensorrt_llm.serve.openai_server._record_generator_termination" + ) as record_termination, + ): + response = asyncio.run(OpenAIServer.health(server)) + return response, raise_signal, record_termination + + +def test_health_returns_200_when_the_engine_is_healthy(): + response, raise_signal, _ = _health(_server(_FakeGenerator(healthy=True))) + assert response.status_code == 200 + raise_signal.assert_not_called() + + +def test_health_returns_503_without_sigint_when_unhealthy_with_no_fatal_error(): + response, raise_signal, record_termination = _health( + _server(_FakeGenerator(healthy=False, fatal_error=None)) + ) + assert response.status_code == 503 + assert response.body == (b"LLM is unavailable. Please check the server logs for more details.") + raise_signal.assert_not_called() + record_termination.assert_not_called() + + +def test_health_sends_sigint_exactly_once_on_a_fatal_error(): + import signal as signal_module + + generator = _FakeGenerator( + healthy=False, fatal_error=RuntimeError("CUDA OOM"), shutdown_started=False + ) + response, raise_signal, record_termination = _health(_server(generator)) + assert response.status_code == 503 + raise_signal.assert_called_once_with(signal_module.SIGINT) + record_termination.assert_called_once_with(generator) + + +def test_health_sends_no_sigint_while_shutdown_is_already_in_progress(): + response, raise_signal, record_termination = _health( + _server( + _FakeGenerator( + healthy=False, fatal_error=RuntimeError("CUDA OOM"), shutdown_started=True + ) + ) + ) + assert response.status_code == 503 + raise_signal.assert_not_called() + record_termination.assert_not_called() + + +def test_health_answers_200_for_a_generator_with_no_health_attributes(): + """A VisualGen generator has no check_health/fatal_error; /health stays 200.""" + response, raise_signal, _ = _health(_server(_NoHealthGenerator(), is_visual_gen=True)) + assert response.status_code == 200 + raise_signal.assert_not_called() + + +def test_health_returns_503_when_the_embedding_batcher_died(): + response, raise_signal, _ = _health( + _server(_FakeGenerator(healthy=True), embedding_batcher=_DeadBatcher()) + ) + assert response.status_code == 503 + raise_signal.assert_not_called() From ef64f31e12c084eb054f38248b13fafb10cb893f Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Sun, 1 Nov 2026 16:57:00 +0000 Subject: [PATCH 02/30] Fix frontend lint coverage and health test cleanup Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- .../llmapi/test_llm_public_runtime_api.py | 5 +- .../test_private_access_lint_visual_gen.py | 27 ++++++- .../serve/test_private_access_lint.py | 74 ++++++++++++++++++- 3 files changed, 96 insertions(+), 10 deletions(-) diff --git a/tests/unittest/llmapi/test_llm_public_runtime_api.py b/tests/unittest/llmapi/test_llm_public_runtime_api.py index ee2c036c50be..e1f9b321d1ee 100644 --- a/tests/unittest/llmapi/test_llm_public_runtime_api.py +++ b/tests/unittest/llmapi/test_llm_public_runtime_api.py @@ -42,6 +42,7 @@ def _llm(executor, encode_only=False, encoder_executor=None): def _proxy(engine_dead=False, fatal_error=None): """A proxy with exactly the two fields its engine_dead property reads.""" proxy = object.__new__(GenerationExecutorProxy) + proxy.workers_started = False proxy._engine_dead = engine_dead proxy._fatal_error = fatal_error return proxy @@ -51,7 +52,7 @@ def _proxy(engine_dead=False, fatal_error=None): # BaseLLM.check_health() # --------------------------------------------------------------------------- def test_check_health_encode_only_with_encoder_executor(): - assert _llm(None, encode_only=True, encoder_executor=object()).check_health() is True + assert _llm(None, encode_only=True, encoder_executor=_NoOpExecutor()).check_health() is True def test_check_health_encode_only_without_encoder_executor(): @@ -130,7 +131,7 @@ def test_shutdown_started_true_without_executor(): def test_shutdown_started_false_for_encode_only(): - assert _llm(None, encode_only=True, encoder_executor=object()).shutdown_started is False + assert _llm(None, encode_only=True, encoder_executor=_NoOpExecutor()).shutdown_started is False # --------------------------------------------------------------------------- diff --git a/tests/unittest/serve/api/visual_gen/test_private_access_lint_visual_gen.py b/tests/unittest/serve/api/visual_gen/test_private_access_lint_visual_gen.py index b121bc555629..40f7fc1cb9d4 100644 --- a/tests/unittest/serve/api/visual_gen/test_private_access_lint_visual_gen.py +++ b/tests/unittest/serve/api/visual_gen/test_private_access_lint_visual_gen.py @@ -4,8 +4,9 @@ Pull requests that change only VisualGen files are claimed by CBTS Tier 1 and run only test-db entries whose path contains visual_gen/. This file applies the rule -and baseline of tests/unittest/serve/test_private_access_lint.py to -tensorrt_llm/serve/api/visual_gen/, so those pull requests are still checked. +and baseline of tests/unittest/serve/test_private_access_lint.py to the existing +frontends, including the VisualGen handlers in serve/openai_server.py. Keep the +full frontend scope until those handlers move into their own directory. """ import importlib.util @@ -22,6 +23,24 @@ def test_no_new_cross_package_private_access_in_visual_gen(): - rows = lint.scan(prefixes=(lint.VISUAL_GEN_SCOPE,)) - errors = lint._compare(rows, lint.load_baseline(), prefix=lint.VISUAL_GEN_SCOPE) + for scope in lint.SCOPES: + assert (lint.REPO_ROOT / scope).is_dir(), f"Frontend scope does not exist: {scope}" + rows = lint.scan(root=lint.REPO_ROOT) + errors = lint._compare(rows, lint.load_baseline()) assert not errors, "\n".join(errors) + + +def test_visual_gen_check_detects_new_frontend_access(tmp_path, monkeypatch): + for scope in lint.SCOPES: + (tmp_path / scope).mkdir(parents=True) + (tmp_path / "tensorrt_llm/serve/openai_server.py").write_text("engine._private_state\n") + owner = tmp_path / "tensorrt_llm/llmapi/llm.py" + owner.parent.mkdir() + owner.write_text("class LLM:\n _private_state = None\n") + baseline = tmp_path / "baseline.yaml" + baseline.write_text("entries: []\n") + monkeypatch.setattr(lint, "REPO_ROOT", tmp_path) + monkeypatch.setattr(lint, "BASELINE_FILE", baseline) + monkeypatch.setattr(lint, "OVERRIDES", {}) + with pytest.raises(AssertionError, match="New private access to `_private_state`"): + test_no_new_cross_package_private_access_in_visual_gen() diff --git a/tests/unittest/serve/test_private_access_lint.py b/tests/unittest/serve/test_private_access_lint.py index ac48c7e81723..9bea0b46de81 100644 --- a/tests/unittest/serve/test_private_access_lint.py +++ b/tests/unittest/serve/test_private_access_lint.py @@ -34,7 +34,6 @@ BASELINE_FILE = Path(__file__).resolve().with_name("private_access_baseline.yaml") SCOPES = ("tensorrt_llm/serve/", "tensorrt_llm/commands/", "tensorrt_llm/grpc/") EXCLUDED = ("tensorrt_llm/serve/scripts/", "tensorrt_llm/grpc/openengine/_generated/") -VISUAL_GEN_SCOPE = "tensorrt_llm/serve/api/visual_gen/" OPENENGINE_SCOPE = "tensorrt_llm/grpc/openengine/" # Lower this when you remove accesses; never raise it. @@ -262,7 +261,13 @@ def scan(root=REPO_ROOT, prefixes=SCOPES): def load_baseline(): data = yaml.safe_load(BASELINE_FILE.read_text()) - return {(e["path"], e["scope"], e["member"]): e for e in data["entries"]} + baseline = {} + for entry in data["entries"]: + key = (entry["path"], entry["scope"], entry["member"]) + if key in baseline: + raise ValueError(f"Duplicate baseline entry: {key}") + baseline[key] = entry + return baseline def _compare(rows, baseline, prefix=""): @@ -315,8 +320,9 @@ def test_no_new_cross_package_private_access(): def test_baseline_only_shrinks(): total = sum(e["count"] for e in load_baseline().values()) - assert total <= MAX_TOTAL, ( - f"{BASELINE_FILE.name} lists {total} accesses, more than MAX_TOTAL={MAX_TOTAL}. " + assert total == MAX_TOTAL, ( + f"{BASELINE_FILE.name} lists {total} accesses but MAX_TOTAL={MAX_TOTAL}. " + "Lower MAX_TOTAL when the total shrinks. " f"The baseline only shrinks. {FIX_HINT}" ) @@ -338,6 +344,66 @@ def test_counting_rule(source, expected): assert sorted(row[1] for row in visitor.rows) == sorted(expected) +def test_duplicate_baseline_keys_are_rejected(tmp_path, monkeypatch): + baseline_file = tmp_path / "baseline.yaml" + entry = {"path": "tensorrt_llm/serve/server.py", "scope": "f", "member": "_x", "count": 1} + baseline_file.write_text(yaml.safe_dump({"entries": [entry, {**entry, "count": 2}]})) + monkeypatch.setattr(sys.modules[__name__], "BASELINE_FILE", baseline_file) + with pytest.raises(ValueError, match="Duplicate baseline entry"): + load_baseline() + + +def test_baseline_total_cannot_leave_unused_allowance(tmp_path, monkeypatch): + baseline_file = tmp_path / "baseline.yaml" + baseline_file.write_text("entries: []\n") + monkeypatch.setattr(sys.modules[__name__], "BASELINE_FILE", baseline_file) + monkeypatch.setattr(sys.modules[__name__], "MAX_TOTAL", 1) + with pytest.raises(AssertionError, match="Lower MAX_TOTAL"): + test_baseline_only_shrinks() + + +def test_scanner_ownership_and_comparison(tmp_path, monkeypatch): + frontend = "tensorrt_llm/serve/server.py" + sources = { + frontend: ( + "class Server:\n _ambiguous = None\n" + "def f(obj):\n obj._external\n obj._ambiguous\n" + " obj._third_party\n obj._override\n" + ), + "tensorrt_llm/executor/engine.py": ( + "class Engine:\n _external = None\n _ambiguous = None\n _override = None\n" + ), + } + for rel, source in sources.items(): + path = tmp_path / rel + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(source) + overrides = {(frontend, "_override"): ("same-package", "test receiver")} + monkeypatch.setattr(sys.modules[__name__], "OVERRIDES", overrides) + rows = scan(tmp_path, prefixes=("tensorrt_llm/serve/",)) + assert {row[2]: row[5] for row in rows} == { + "_external": "cross", + "_ambiguous": "unresolved", + "_third_party": "third-party", + "_override": "same-package", + } + errors = _compare(rows, {}) + assert len(errors) == 2 + assert any("New private access to `_external`" in error for error in errors) + assert any("Cannot classify `obj._ambiguous`" in error for error in errors) + + baseline = { + (frontend, "f", "_external"): {"count": 2}, + ("tensorrt_llm/grpc/server.py", "f", "_other"): {"count": 1}, + } + overrides[(frontend, "_removed")] = ("same-package", "removed test receiver") + errors = _compare(rows, baseline, prefix="tensorrt_llm/serve/") + assert len(errors) == 3 + assert any("Stale baseline entry" in error and "`_external`" in error for error in errors) + assert any("Stale OVERRIDES entry" in error and "`_removed`" in error for error in errors) + assert not any("_other" in error for error in errors) + + if __name__ == "__main__": if "--report" in sys.argv: for row in scan(): From 73c213bfb3b282fbc72584799becc7be619c6a7b Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Mon, 2 Nov 2026 10:25:35 +0000 Subject: [PATCH 03/30] [TRTLLM-17272][feat] Add public runtime accessors for the remaining serving-frontend private accesses The serving frontends (tensorrt_llm/serve, tensorrt_llm/commands, the OpenEngine gRPC frontend) read private members of LLM, the executors, CompletionOutput, SamplingParams, MultimodalDataTracker, OpenAIDisaggServer and PyExecutor. Give each remaining read a thin public entry point on the class that owns the member, switch the frontends to it with no behavior change, and shrink the private-access baseline from 50 recorded accesses to 19. The reads of the two private LlmArgs attributes stay (exposing them is an LLM-args change for its own PR); the remaining baseline entries belong to other in-flight work. New public members (status prototype where api-stability tracks them): LLM.record_fatal_error / multi_frontend_attach_info / multi_frontend_ipc_dir / engine_in_worker_process / hf_model_dir / hf_model_config / resource_governor_queue / encode_limits / get_kv_cache_load / get_openengine_discovery; GenerationExecutor.record_fatal_error and multi_frontend_ipc_dir; CompletionOutput.spec_dec_counters and text_diff_offset (and the SpecDecCounters rename); SamplingParams.set_request_provided_fields and stop_word_ids; MultimodalDataTracker.multimodal_server_config; OpenAIDisaggServer.address; PyExecutor.set_openengine_discovery and set_openengine_node_server; registered_model_type() in inputs/registry. Also adds tensorrt_llm/llmapi/engine_client.py: the EngineClient protocol naming the engine operations every frontend shares (generate, health, load), checked by tests rather than at runtime. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- pyproject.toml | 1 + tensorrt_llm/_torch/pyexecutor/py_executor.py | 17 +- tensorrt_llm/commands/serve.py | 6 +- tensorrt_llm/executor/executor.py | 9 + tensorrt_llm/executor/proxy.py | 17 +- tensorrt_llm/executor/result.py | 24 +- tensorrt_llm/grpc/openengine/capabilities.py | 6 +- tensorrt_llm/grpc/openengine/control.py | 28 +- tensorrt_llm/grpc/openengine/formatting.py | 6 +- tensorrt_llm/grpc/openengine/node.py | 4 +- .../grpc/openengine/request_mapping.py | 2 +- tensorrt_llm/grpc/openengine/server.py | 2 +- tensorrt_llm/grpc/openengine/streaming.py | 2 +- tensorrt_llm/inputs/registry.py | 9 + tensorrt_llm/inputs/utils.py | 5 + tensorrt_llm/llmapi/engine_client.py | 71 +++++ tensorrt_llm/llmapi/llm.py | 154 ++++++++++- tensorrt_llm/sampling_params.py | 23 +- tensorrt_llm/serve/_frontend_processes.py | 14 +- tensorrt_llm/serve/chat_utils.py | 8 +- tensorrt_llm/serve/openai_disagg_server.py | 5 + tensorrt_llm/serve/openai_protocol.py | 2 +- tensorrt_llm/serve/openai_server.py | 36 +-- tensorrt_llm/serve/postprocess_handlers.py | 19 +- .../integration/test_lists/test-db/l0_cpu.yml | 1 + .../_torch/pyexecutor/test_py_executor.py | 14 + .../api_stability/api_stability_core.py | 2 +- .../references/completion_output.yaml | 6 + .../api_stability/references/llm.yaml | 37 +++ .../references/sampling_params.yaml | 11 +- .../test_openai_disagg_service.py | 15 ++ .../unittest/grpc/openengine/test_control.py | 14 +- .../openengine/test_frontend_lifecycle.py | 6 +- .../unittest/grpc/openengine/test_metrics.py | 4 +- .../grpc/openengine/test_multimodal.py | 50 ++-- tests/unittest/grpc/openengine/test_node.py | 20 +- .../grpc/openengine/test_openengine.py | 8 +- .../llmapi/apps/openai_server_stub.py | 4 +- .../llmapi/apps/test_anthropic_routes.py | 4 +- tests/unittest/llmapi/apps/test_chat_utils.py | 10 + .../test_openai_server_iteration_stats.py | 2 +- .../apps/test_reasoning_prompt_resolution.py | 1 + .../unittest/llmapi/apps/test_tool_parsers.py | 2 + tests/unittest/llmapi/test_executor.py | 29 +++ .../llmapi/test_llm_public_runtime_api.py | 245 +++++++++++++++++- tests/unittest/llmapi/test_sampling_params.py | 30 +++ .../llmapi/test_serve_frontend_affinity.py | 4 +- .../test_serve_multi_frontend_helpers.py | 46 +++- .../test_serve_multi_frontend_route_gate.py | 2 +- .../serve/private_access_baseline.yaml | 28 +- .../serve/test_private_access_lint.py | 2 +- .../serve/test_serve_runtime_api_calls.py | 179 +++++++++++++ 52 files changed, 1041 insertions(+), 205 deletions(-) create mode 100644 tensorrt_llm/llmapi/engine_client.py create mode 100644 tests/unittest/serve/test_serve_runtime_api_calls.py diff --git a/pyproject.toml b/pyproject.toml index 60e243d984b0..533567d4f4d2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -862,6 +862,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/llmapi/engine_client.py", "tests/unittest/_torch/sampler/test_torch_sampler.py", "tests/unittest/_torch/sampler/test_token_ban.py", "tests/unittest/_torch/sampler/test_beam_search.py", diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index ae1c06313552..630bca381176 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -14,8 +14,8 @@ from contextlib import ExitStack, contextmanager, nullcontext from enum import IntEnum from queue import Queue -from typing import (TYPE_CHECKING, Callable, Dict, Iterable, Iterator, List, - Literal, NoReturn, Optional, Tuple, Union) +from typing import (TYPE_CHECKING, Any, Callable, Dict, Iterable, Iterator, + List, Literal, NoReturn, Optional, Tuple, Union) import torch from strenum import StrEnum @@ -2598,6 +2598,19 @@ def get_openengine_discovery(self) -> dict: """Return immutable-at-startup OpenEngine node and publisher discovery.""" return self._openengine_discovery + def set_openengine_discovery(self, discovery: dict) -> None: + """Record OpenEngine node and KV-event publisher discovery metadata. + + Called once at startup by the OpenEngine node-discovery hook + (tensorrt_llm/grpc/openengine/node.py); get_openengine_discovery() + returns it. + """ + self._openengine_discovery = discovery + + def set_openengine_node_server(self, server: Any) -> None: + """Adopt the OpenEngine node metadata server; shutdown closes it.""" + self._openengine_node_server = server + def get_kv_cache_load(self) -> dict: """Return the scheduler-owned KV load snapshot. diff --git a/tensorrt_llm/commands/serve.py b/tensorrt_llm/commands/serve.py index aba22131bb0d..f4fa10dd2cf4 100644 --- a/tensorrt_llm/commands/serve.py +++ b/tensorrt_llm/commands/serve.py @@ -555,8 +555,8 @@ def _bind_launcher_uds(llm) -> Tuple[socket.socket, str]: It lives in the multi-frontend ipc dir, which is private to the user and removed, socket file included, when the executor proxy shuts down. """ - # Only the classic IPC proxy has the attribute. - ipc_dir = getattr(llm._executor, "_multi_frontend_ipc_dir", None) + # Only the classic IPC proxy in multi-frontend mode has one. + ipc_dir = llm.multi_frontend_ipc_dir if not ipc_dir: raise ValueError( "num_serve_frontends > 1 requires the classic IPC executor " @@ -2661,7 +2661,7 @@ def _run_fleet_worker_impl(): """Build and run one fleet worker inside its telemetry boundary.""" _init_fleet_worker_process() server = _build_disagg_server_from_env() - host, port = server._config.hostname, server._config.port + host, port = server.address s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) diff --git a/tensorrt_llm/executor/executor.py b/tensorrt_llm/executor/executor.py index 6738aa4986a3..54cf6b91373e 100644 --- a/tensorrt_llm/executor/executor.py +++ b/tensorrt_llm/executor/executor.py @@ -342,6 +342,10 @@ def fatal_error(self) -> Optional[BaseException]: """The unrecoverable engine error recorded by ``_set_fatal_error``, or None.""" return self._fatal_error + def record_fatal_error(self, error: BaseException) -> None: + """Record an unrecoverable engine error; the first recorded error wins (see ``fatal_error``).""" + self._set_fatal_error(error) + @property def engine_dead(self) -> bool: """Whether the engine is known dead (an unrecoverable error was recorded).""" @@ -415,6 +419,11 @@ def resource_governor_queue(self): """Return the resource governor queue if this executor supports it.""" return None + @property + def multi_frontend_ipc_dir(self) -> Optional[str]: + """The multi-frontend IPC directory if this executor runs one, else None.""" + return None + @property def enable_postprocess_parallel(self) -> bool: return self.postproc_config.enabled diff --git a/tensorrt_llm/executor/proxy.py b/tensorrt_llm/executor/proxy.py index 765e9a001bef..172e75e3e655 100644 --- a/tensorrt_llm/executor/proxy.py +++ b/tensorrt_llm/executor/proxy.py @@ -680,6 +680,19 @@ def _setup_queues(self) -> WorkerCommIpcAddrs: frontend_result_queue_addrs=frontend_result_addrs, ) + @property + def multi_frontend_ipc_dir(self) -> Optional[str]: + """The launcher-owned directory of the multi-frontend IPC endpoints. + + None unless multi-frontend mode is active. The launcher places its + own Unix sockets (the attached frontends' forwarding endpoint) in it; + the directory is private to the user and removed, socket files + included, when this proxy shuts down. + """ + # GenerationExecutorFrontendProxy skips this class's __init__, so the + # attribute can be missing there; an attached frontend owns no ipc dir. + return getattr(self, "_multi_frontend_ipc_dir", None) + def multi_frontend_attach_info(self) -> Optional[dict]: """The attach payload consumed by attached serving frontends. @@ -1693,8 +1706,8 @@ def __init__( self.garbage_collection_gen0_threshold = None self.workers_started = False self.dispatch_result_thread: Optional[ManagedThread] = None - # Must be None: OpenAIServer reads the resource_governor_queue - # property at init; the governor lives with the launcher only. + # Must be None: LLM.resource_governor_queue reads this property at + # server init; the governor lives with the launcher only. self._resource_governor_queue = None hmac_key = bytes.fromhex(attach_info["hmac_key"]) diff --git a/tensorrt_llm/executor/result.py b/tensorrt_llm/executor/result.py index 43bcee12ff40..ef615b6bd6a1 100644 --- a/tensorrt_llm/executor/result.py +++ b/tensorrt_llm/executor/result.py @@ -67,7 +67,7 @@ class LogProbsResult(NamedTuple): prompt: Optional[TokenLogprobs | SimpleTokenLogprobs] = None -class _SpecDecCounters(NamedTuple): +class SpecDecCounters(NamedTuple): """Speculative-decoding counters reported by the request behind one sequence. Field names match the executor response so a formatter can read either. @@ -128,6 +128,8 @@ class CompletionOutput: logprobs_diff (TokenLogprobs | SimpleTokenLogprobs): Logprobs of newly generated tokens. text_diff (str): Newly generated tokens. routed_experts (Optional[torch.Tensor]): Per-token pre-EPLB logical top-k MoE expert ids (Router Replay / R3), or None when not requested. + spec_dec_counters (tensorrt_llm.executor.result.SpecDecCounters, optional): This sequence's speculative-decoding counters, or None. + text_diff_offset (int): Offset in text at which text_diff starts. """ index: int text: str = "" @@ -157,9 +159,9 @@ class CompletionOutput: # the result of result_handler passed to postprocess workers _postprocess_result: Any = None # this sequence's own spec-decode counters; see _handle_sequence - _spec_dec_counters: Optional[_SpecDecCounters] = field(default=None, - init=False, - repr=False) + _spec_dec_counters: Optional[SpecDecCounters] = field(default=None, + init=False, + repr=False) def __getstate__(self) -> dict: # _incremental_states holds a tokenizers.DecodeStream (a Rust object, @@ -208,6 +210,16 @@ def routed_experts(self) -> Optional[torch.Tensor]: def text_diff(self) -> str: return self.text[self._last_text_len:] + @property + def text_diff_offset(self) -> int: + """Offset in ``text`` at which ``text_diff`` starts.""" + return self._last_text_len + + @property + def spec_dec_counters(self) -> Optional[SpecDecCounters]: + """This sequence's own speculative-decoding counters, or None if the request did not draft.""" + return self._spec_dec_counters + @property def token_ids_diff(self) -> List[int]: return self.token_ids[self._last_token_ids_len:] @@ -382,7 +394,7 @@ def _handle_sequence(self, output.disaggregated_params = self.disaggregated_params output._last_token_ids_len = len(output.token_ids) output._last_logprobs_len = len(output.logprobs) - output._spec_dec_counters = _SpecDecCounters() + output._spec_dec_counters = SpecDecCounters() decoder_output_prefix = () if (self.sampling_params.exclude_input_from_output or getattr(self, "_streaming", False)): @@ -463,7 +475,7 @@ def _handle_sequence(self, # counters, while the request-level copies on self are overwritten by # every candidate's response. Capture this sequence's here so a choice # never reports whichever candidate happened to respond last. - output._spec_dec_counters = _SpecDecCounters( + output._spec_dec_counters = SpecDecCounters( per_pos_drafted=response_tensors.per_pos_drafted, per_pos_accepted=response_tensors.per_pos_accepted, spec_dec_totals=response_tensors.spec_dec_totals) diff --git a/tensorrt_llm/grpc/openengine/capabilities.py b/tensorrt_llm/grpc/openengine/capabilities.py index 303f7918f8e5..e8f5f5c94a21 100644 --- a/tensorrt_llm/grpc/openengine/capabilities.py +++ b/tensorrt_llm/grpc/openengine/capabilities.py @@ -11,6 +11,7 @@ from tensorrt_llm.inputs.registry import ( MULTIMODAL_PLACEHOLDER_REGISTRY, BaseMultimodalInputProcessor, + registered_model_type, ) from .kv_events import events_config @@ -67,7 +68,7 @@ def supported_guides(backend: str | None) -> frozenset[str]: def _loaded_model_type(llm: Any) -> str: """Return the top-level model type for the loaded input processor.""" input_processor = getattr(llm, "input_processor", None) - config = getattr(input_processor, "config", None) or getattr(llm, "_hf_model_config", None) + config = getattr(input_processor, "config", None) or getattr(llm, "hf_model_config", None) if config is None: return "" return resolve_top_level_model_type(config) @@ -102,8 +103,7 @@ def resolve_media_capabilities( if not model_type: return MediaCapabilities({}) if not isinstance(input_processor, BaseMultimodalInputProcessor): - registered_model_type = getattr(type(input_processor), "_registered_model_type", None) - if registered_model_type != model_type: + if registered_model_type(type(input_processor)) != model_type: return MediaCapabilities({}) if not callable(input_processor): return MediaCapabilities({}) diff --git a/tensorrt_llm/grpc/openengine/control.py b/tensorrt_llm/grpc/openengine/control.py index 3275880862e8..9df67ef72abc 100644 --- a/tensorrt_llm/grpc/openengine/control.py +++ b/tensorrt_llm/grpc/openengine/control.py @@ -131,9 +131,7 @@ async def _get_kv_cache_load(self) -> dict: async with self._load_snapshot_lock: now = time.monotonic() if now >= self._load_snapshot_deadline: - executor = getattr(self._llm, "_executor", None) - getter = getattr(executor, "get_kv_cache_load", None) - self._load_snapshot = await asyncio.to_thread(getter) if callable(getter) else {} + self._load_snapshot = await asyncio.to_thread(self._llm.get_kv_cache_load) self._load_snapshot_deadline = time.monotonic() + _LOAD_SNAPSHOT_CACHE_SECONDS return self._load_snapshot @@ -218,15 +216,17 @@ async def GetServerInfo( block_size = _positive_int(self._llm.args.kv_cache_config.tokens_per_block) if block_size is not None: capacity.kv_block_size = block_size - executor = getattr(self._llm, "_executor", None) - capacity_getter = getattr(executor, "get_kv_cache_capacity", None) - if callable(capacity_getter): - kv_capacity = await asyncio.to_thread(capacity_getter) - local_blocks = _positive_int(kv_capacity.get("maxNumBlocks")) - if local_blocks is not None and data_parallel_size(self._llm) == 1: - capacity.total_kv_blocks = local_blocks - discovery_getter = getattr(executor, "get_openengine_discovery", None) - discovery = await asyncio.to_thread(discovery_getter) if callable(discovery_getter) else {} + # An encode-only LLM has no generation executor and no KV capacity; + # LLM.get_kv_cache_capacity() raises for it, so skip the query. + kv_capacity = ( + {} + if self._llm.args.encode_only is True + else await asyncio.to_thread(self._llm.get_kv_cache_capacity) + ) + local_blocks = _positive_int(kv_capacity.get("maxNumBlocks")) + if local_blocks is not None and data_parallel_size(self._llm) == 1: + capacity.total_kv_blocks = local_blocks + discovery = await asyncio.to_thread(self._llm.get_openengine_discovery) events = events_config(self._llm) discovery_required = ( getattr(self._llm.args, "_openengine_discovery", None) is not None @@ -656,9 +656,7 @@ async def _requested_sources( """ if events_config(self._llm) is None: return [] - executor = getattr(self._llm, "_executor", None) - getter = getattr(executor, "get_openengine_discovery", None) - discovery = await asyncio.to_thread(getter) if callable(getter) else {} + discovery = await asyncio.to_thread(self._llm.get_openengine_discovery) descriptors = discovery.get("sources") if descriptors is None: await context.abort( diff --git a/tensorrt_llm/grpc/openengine/formatting.py b/tensorrt_llm/grpc/openengine/formatting.py index 9949b9ad2f4b..59ef02bf0cfd 100644 --- a/tensorrt_llm/grpc/openengine/formatting.py +++ b/tensorrt_llm/grpc/openengine/formatting.py @@ -116,15 +116,15 @@ def _stop_texts(sampling_params: SamplingParams) -> list[str]: def _token_holdback(sampling_params: SamplingParams) -> int: """Number of trailing tokens to withhold so a stop token id is not streamed. - NOTE: this depends on the private ``SamplingParams._stop_word_ids``, which the - LLM engine populates with the tokenized stop words. It is read after + NOTE: this depends on ``SamplingParams.stop_word_ids``, which the LLM + engine populates with the tokenized stop words. It is read after ``generate_async`` returns, when it is expected to be set. If the engine ever stops populating it (or populates it later), this degrades to 0 (no holdback) rather than failing — ``test_token_holdback_*`` guards the current behavior. """ if sampling_params.include_stop_str_in_output: return 0 - stop_word_ids = getattr(sampling_params, "_stop_word_ids", None) or [] + stop_word_ids = sampling_params.stop_word_ids or [] if not stop_word_ids: return 0 return max(1, max(len(token_ids) for token_ids in stop_word_ids) - 1) diff --git a/tensorrt_llm/grpc/openengine/node.py b/tensorrt_llm/grpc/openengine/node.py index e43bc1eb8ae5..abc1a1d2a93f 100644 --- a/tensorrt_llm/grpc/openengine/node.py +++ b/tensorrt_llm/grpc/openengine/node.py @@ -151,7 +151,7 @@ def initialize_node_discovery(executor) -> None: ), "source_owners": source_owners, } - executor._openengine_discovery = {"node": metadata, "sources": sources} + executor.set_openengine_discovery({"node": metadata, "sources": sources}) info = server_pb2.ServerInfo( engine_name="tensorrt_llm", engine_version=__version__, @@ -190,4 +190,4 @@ def initialize_node_discovery(executor) -> None: if server is not None: server.close() raise RuntimeError(f"OpenEngine node discovery startup failed: {errors}") - executor._openengine_node_server = server + executor.set_openengine_node_server(server) diff --git a/tensorrt_llm/grpc/openengine/request_mapping.py b/tensorrt_llm/grpc/openengine/request_mapping.py index 41334157914e..dc3fd854f8f1 100644 --- a/tensorrt_llm/grpc/openengine/request_mapping.py +++ b/tensorrt_llm/grpc/openengine/request_mapping.py @@ -347,7 +347,7 @@ def sampling_params_from_request( # `_record_sampling_params_request_fields`). Without this, the materialized # defaults above would look client-provided and suppress a model's # generation_config.json sampling defaults under generation_config="auto". - sampling_params._set_request_provided_fields(provided_fields) + sampling_params.set_request_provided_fields(provided_fields) return sampling_params diff --git a/tensorrt_llm/grpc/openengine/server.py b/tensorrt_llm/grpc/openengine/server.py index b495db5f6c1d..118a16ed1ad4 100644 --- a/tensorrt_llm/grpc/openengine/server.py +++ b/tensorrt_llm/grpc/openengine/server.py @@ -187,7 +187,7 @@ def __init__( async def start(self) -> None: """Start accepting OpenEngine requests.""" if events_config(self._llm) is not None: - capacity = await asyncio.to_thread(self._llm._executor.get_kv_cache_capacity) + capacity = await asyncio.to_thread(self._llm.get_kv_cache_capacity) if not capacity.get("kvEventsEnabled", False): raise KvEventsUnavailableError( "KV events are configured but the loaded engine has no streaming publisher; " diff --git a/tensorrt_llm/grpc/openengine/streaming.py b/tensorrt_llm/grpc/openengine/streaming.py index 926b4b64276a..3894d8df34b8 100644 --- a/tensorrt_llm/grpc/openengine/streaming.py +++ b/tensorrt_llm/grpc/openengine/streaming.py @@ -35,7 +35,7 @@ class _OpenEngineFormatState: sampling_params: Any = None is_context_only: bool = False kv_transfer_backend: str = "" - # None until first use: _token_holdback needs sampling_params._stop_word_ids, + # None until first use: _token_holdback needs sampling_params.stop_word_ids, # which generate_async only populates after the request is submitted, so it is # computed lazily on the first _format_result call (inline or in a worker). token_holdback: Optional[int] = None diff --git a/tensorrt_llm/inputs/registry.py b/tensorrt_llm/inputs/registry.py index b5c5539f4da7..6f9df23259b3 100644 --- a/tensorrt_llm/inputs/registry.py +++ b/tensorrt_llm/inputs/registry.py @@ -1111,6 +1111,15 @@ def wrapper(model_cls: N) -> N: return wrapper +def registered_model_type(processor_cls: type) -> Optional[str]: + """The model type an input-processor class was registered under, or None. + + ``register_input_processor`` stamps the type onto the processor class; + this is the public read of that stamp. + """ + return getattr(processor_cls, "_registered_model_type", None) + + def create_input_processor( model_path_or_dir: str, tokenizer, diff --git a/tensorrt_llm/inputs/utils.py b/tensorrt_llm/inputs/utils.py index 1360861ed743..93476892aaab 100644 --- a/tensorrt_llm/inputs/utils.py +++ b/tensorrt_llm/inputs/utils.py @@ -379,6 +379,11 @@ def __init__( # time; see `BaseMediaIO.merge_kwargs` in `inputs/media_io.py`. self._request_media_io_kwargs = request_media_io_kwargs + @property + def multimodal_server_config(self) -> MultimodalServerConfig: + """Server-wide media loading settings used for this request.""" + return self._multimodal_server_config + @property def request_media_io_kwargs(self) -> Optional[Dict[str, Dict[str, Any]]]: """Per-request `media_io_kwargs` override, or `None` if not set.""" diff --git a/tensorrt_llm/llmapi/engine_client.py b/tensorrt_llm/llmapi/engine_client.py new file mode 100644 index 000000000000..c4c794d22106 --- /dev/null +++ b/tensorrt_llm/llmapi/engine_client.py @@ -0,0 +1,71 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The engine operations the serving frontends share. + +The HTTP server (tensorrt_llm/serve), the SMG gRPC frontend and the OpenEngine +gRPC frontend each hold an in-process LLM and use the same core: generate, +health and load. EngineClient writes that core down; LLM satisfies it. +Operations only one role uses stay on LLM / AsyncLLM: encode() and +encode_limits, iteration stats, runtime memory control, RL collective_rpc, +profiling, multi_frontend_attach_info. Nothing checks the protocol at run time; +tests/unittest/serve/test_private_access_lint.py keeps the frontends off +private members. + +OpenEngine v1 alignment: + Inference.Generate -> generate_async + Control.Abort -> RequestOutput.abort() (no engine-level abort here) + Control.Health -> check_health, fatal_error, shutdown_started + Control.GetLoad -> get_kv_cache_load, get_kv_cache_capacity + GetServerInfo / GetModelInfo -> not here (LLM has no configuration snapshot method) +A transport-level client implementing this protocol adds failure modes the +in-process LLM does not have (timeouts, lost connections); its errors must be +dialect-neutral, so that each wire dialect (OpenAI, Anthropic) renders them in +its own format. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Mapping, Optional, Protocol, Union + +if TYPE_CHECKING: + from tensorrt_llm.conversation_params import ConversationParams + from tensorrt_llm.disaggregated_params import DisaggregatedParams + from tensorrt_llm.executor.request import LoRARequest, PromptAdapterRequest + from tensorrt_llm.inputs.data import PromptInputs + from tensorrt_llm.llmapi.llm import PreprocessedInputs, RequestOutput + from tensorrt_llm.llmapi.llm_utils import KvCacheRetentionConfig + from tensorrt_llm.sampling_params import SamplingParams + from tensorrt_llm.scheduling_params import SchedulingParams + + +class EngineClient(Protocol): + """The generate / health / load surface every serving frontend uses.""" + + def generate_async( + self, + inputs: Union[PromptInputs, PreprocessedInputs], + sampling_params: Optional[SamplingParams] = None, + *, + lora_request: Optional[LoRARequest] = None, + prompt_adapter_request: Optional[PromptAdapterRequest] = None, + streaming: bool = False, + kv_cache_retention_config: Optional[KvCacheRetentionConfig] = None, + disaggregated_params: Optional[DisaggregatedParams] = None, + trace_headers: Optional[Mapping[str, str]] = None, + scheduling_params: Optional[SchedulingParams] = None, + conversation_params: Optional[ConversationParams] = None, + cache_salt: Optional[str] = None, + priority: float = ..., + ) -> RequestOutput: ... + + def check_health(self) -> bool: ... + + @property + def fatal_error(self) -> Optional[BaseException]: ... + + @property + def shutdown_started(self) -> bool: ... + + def get_kv_cache_load(self) -> dict[str, Any]: ... + + def get_kv_cache_capacity(self) -> dict[str, Any]: ... diff --git a/tensorrt_llm/llmapi/llm.py b/tensorrt_llm/llmapi/llm.py index 985e84978550..5432c35412d9 100644 --- a/tensorrt_llm/llmapi/llm.py +++ b/tensorrt_llm/llmapi/llm.py @@ -41,8 +41,9 @@ from ..conversation_params import ConversationParams from ..disaggregated_params import DisaggregatedParams from ..executor import (DetokenizedGenerationResultBase, GenerationExecutor, - GenerationResult, IterationResult, LoRARequest, - PostprocWorkerConfig, PromptAdapterRequest) + GenerationExecutorProxy, GenerationResult, + IterationResult, LoRARequest, PostprocWorkerConfig, + PromptAdapterRequest) from ..executor.postproc_worker import PostprocParams from ..executor.postprocessor_hook import (PostProcessorHook, load_post_processor_hook) @@ -145,6 +146,20 @@ class EncoderOutput: prompt: Optional[str] = None +@dataclass(frozen=True) +class EncodeLimits: + """Limits that ``LLM.encode()`` checks every call against. + + Attributes: + max_batch_size (int): Maximum number of inputs per call. + max_num_tokens (int): Maximum number of tokens per call. + max_seq_len (int): Maximum number of tokens per input. + """ + max_batch_size: int + max_num_tokens: int + max_seq_len: int + + class _WhisperSuppressTokensLogitsProcessor(LogitsProcessor): """Apply Whisper suppress-token lists from the HF generation config. @@ -342,6 +357,12 @@ def _multimodal_params_have_encoder_features( fatal_error (BaseException, optional): The unrecoverable engine error the executor recorded, or None. shutdown_started (bool): Whether shutdown of the generation executor has begun. engine_dead (bool): Whether the engine is known dead (a worker died or a fatal error was recorded). + encode_limits (tensorrt_llm.llmapi.llm.EncodeLimits, optional): Limits encode() checks every call against; None unless encode_only=True. + engine_in_worker_process (bool): Whether generation runs through the IPC executor proxy, with the engine in separate MPI worker processes. + hf_model_config (transformers.PretrainedConfig, optional): The Hugging Face model config loaded from the checkpoint, or None. + hf_model_dir (Union[str, pathlib.Path]): The local Hugging Face checkpoint directory the model was loaded from. + multi_frontend_ipc_dir (str, optional): The directory of the multi-frontend IPC endpoints, or None outside multi-frontend mode. + resource_governor_queue (Any, optional): The generation executor's resource-governor queue, or None. """ @@ -1232,10 +1253,10 @@ def encode( if unbatched: inputs = [inputs] - engine = self._encoder_executor.model_engine - max_seq_len = engine.max_seq_len - max_num_tokens = engine.max_num_tokens - max_batch_size = engine.batch_size + limits = self.encode_limits + max_seq_len = limits.max_seq_len + max_num_tokens = limits.max_num_tokens + max_batch_size = limits.max_batch_size if len(inputs) > max_batch_size: raise ValueError( @@ -1691,8 +1712,7 @@ def _try_load_generation_config( def _try_load_generation_config_explicit_values(self) -> dict[str, Any]: if self.args.backend != "pytorch" or self.args.generation_config != "auto": return {} - model_dir = self._hf_model_dir or self.args.model - return ModelLoader.load_hf_generation_config_dict(model_dir) + return ModelLoader.load_hf_generation_config_dict(self.hf_model_dir) def _try_load_hf_model_config( self) -> Optional[transformers.PretrainedConfig]: @@ -1851,6 +1871,124 @@ def engine_dead(self) -> bool: executor = getattr(self, "_executor", None) return False if executor is None else executor.engine_dead + @set_api_status("prototype") + def record_fatal_error(self, error: BaseException) -> bool: + """Record an unrecoverable engine error on the generation executor. + + The first recorded error wins and is what ``fatal_error`` returns; + later calls keep the first error. + + Args: + error (BaseException): The exception to record as the fatal error. + + Returns: + bool: True if the executor took the call, False when the LLM has no generation executor (``encode_only=True``, or after ``shutdown()``). + """ + executor = getattr(self, "_executor", None) + if executor is None: + return False + executor.record_fatal_error(error) + return True + + @set_api_status("prototype") + def multi_frontend_attach_info(self) -> dict: + """Get the attach payload for serving frontends that share this LLM's engine. + + Only the classic IPC executor proxy in multi-frontend mode + (``num_serve_frontends > 1``) has one. + + Returns: + dict: Request and result lane addresses, HMAC keys and the rank-0 RPC endpoint. + + Raises: + ValueError: If the executor is not the classic IPC executor proxy in multi-frontend mode. + """ + executor = getattr(self, "_executor", None) + attach_info = (executor.multi_frontend_attach_info() if isinstance( + executor, GenerationExecutorProxy) else None) + if attach_info is None: + raise ValueError( + "num_serve_frontends > 1 requires the classic IPC executor " + f"proxy in multi-frontend mode, got {type(executor).__name__}") + return attach_info + + @property + @set_api_status("prototype") + def multi_frontend_ipc_dir(self) -> Optional[str]: + """The directory of the multi-frontend IPC endpoints, or None outside multi-frontend mode.""" + executor = getattr(self, "_executor", None) + return None if executor is None else executor.multi_frontend_ipc_dir + + @property + @set_api_status("prototype") + def engine_in_worker_process(self) -> bool: + """Whether generation runs through the IPC executor proxy, with the engine in separate MPI worker processes.""" + return isinstance(getattr(self, "_executor", None), + GenerationExecutorProxy) + + @property + @set_api_status("prototype") + def hf_model_dir(self) -> Union[str, Path]: + """The local Hugging Face checkpoint directory the model was loaded from (``args.model`` until it is resolved).""" + return self._hf_model_dir or self.args.model + + @property + @set_api_status("prototype") + def hf_model_config(self) -> Optional[transformers.PretrainedConfig]: + """The Hugging Face model config loaded from the checkpoint, or None when it could not be loaded.""" + return self._hf_model_config + + @property + @set_api_status("prototype") + def resource_governor_queue(self) -> Optional[Any]: + """The generation executor's resource-governor queue, or None when it has none. + + Raises: + RuntimeError: If the LLM has no generation executor (``encode_only=True``, or after ``shutdown()``). + """ + # Read the attribute directly: every built LLM has it, and on a bare + # instance the AttributeError is what inspect.getmembers (the + # api-stability snapshotter) tolerates, where a RuntimeError would not + # be. + executor = self._executor + if executor is None: + raise RuntimeError( + "resource_governor_queue is not available without a generation " + "executor (encode_only=True, or after shutdown()).") + return executor.resource_governor_queue + + @property + @set_api_status("prototype") + def encode_limits(self) -> Optional[EncodeLimits]: + """Limits ``encode()`` checks every call against; None unless ``encode_only=True`` and the encoder is up.""" + encoder = getattr(self, "_encoder_executor", None) + if not getattr(self, "_encode_only", False) or encoder is None: + return None + engine = encoder.model_engine + return EncodeLimits(max_batch_size=engine.batch_size, + max_num_tokens=engine.max_num_tokens, + max_seq_len=engine.max_seq_len) + + @set_api_status("prototype") + def get_kv_cache_load(self) -> dict: + """Get a non-destructive snapshot of the primary/GPU KV cache load. + + Returns: + dict: The executor's snapshot; empty without a generation executor or when the executor does not report load. + """ + executor = getattr(self, "_executor", None) + return {} if executor is None else executor.get_kv_cache_load() + + @set_api_status("prototype") + def get_openengine_discovery(self) -> dict: + """Get the OpenEngine node and KV-event publisher discovery metadata recorded at startup. + + Returns: + dict: The executor's discovery metadata; empty without a generation executor or when none was recorded. + """ + executor = getattr(self, "_executor", None) + return {} if executor is None else executor.get_openengine_discovery() + @staticmethod def _shutdown_wrapper(self_ref): # Retrieve the instance if it still exists diff --git a/tensorrt_llm/sampling_params.py b/tensorrt_llm/sampling_params.py index 49b28ebc679a..2369a4b4a8ad 100644 --- a/tensorrt_llm/sampling_params.py +++ b/tensorrt_llm/sampling_params.py @@ -269,6 +269,9 @@ class SamplingParams: truncate_prompt_tokens (int, optional): If set to an integer k, will use only the last k tokens from the prompt (i.e., left truncation). Defaults to None. skip_special_tokens (bool): Whether to skip special tokens in the output. Defaults to True. spaces_between_special_tokens (bool): Whether to add spaces between special tokens in the output. Defaults to True. + + Attributes: + stop_word_ids (List[List[int]], optional): Token ids of each stop string, filled in when the LLM sets up the request. """ # noqa: E501 # [TO DEVELOPER] This class provides an interface to LLMAPI users. @@ -467,10 +470,26 @@ def _validate(self): if self.logprobs_simple_format and self.use_beam_search: raise ValueError("logprobs_simple_format is not supported with beam search") - def _set_request_provided_fields(self, field_names: Iterable[str]) -> None: - """Record sampler fields explicitly supplied by a serving request.""" + def set_request_provided_fields(self, field_names: Iterable[str]) -> None: + """Record which sampler fields a serving request set explicitly. + + Serving adapters fill protocol defaults into every field; recording the + fields the request set keeps generation_config.json defaults for the rest. + + Args: + field_names (Iterable[str]): Names of the fields the request set; + names that are not sampler fields are ignored. + """ self._request_provided_fields = frozenset(field_names) & _GENERATION_CONFIG_SAMPLING_FIELDS + @property + def stop_word_ids(self) -> Optional[List[List[int]]]: + """Token ids of each ``stop`` string, filled in when the LLM sets up the request. + + None before that, or when the request has no ``stop`` strings. + """ + return self._stop_word_ids + def _apply_generation_config_defaults(self, generation_config: Mapping[str, Any]) -> None: """Apply compatible model defaults without overriding request values.""" for field_name in _GENERATION_CONFIG_SAMPLING_FIELDS: diff --git a/tensorrt_llm/serve/_frontend_processes.py b/tensorrt_llm/serve/_frontend_processes.py index 8f5ea4bf2b57..e514f8426fde 100644 --- a/tensorrt_llm/serve/_frontend_processes.py +++ b/tensorrt_llm/serve/_frontend_processes.py @@ -92,17 +92,9 @@ def _spawn_attached_frontends( bound for it (_frontend_affinity_sockets). Each child inherits its own socket; the launcher closes its copy and removes it from the dict. """ - from tensorrt_llm.executor.proxy import GenerationExecutorProxy - - executor = getattr(llm, "_executor", None) - if ( - not isinstance(executor, GenerationExecutorProxy) - or (attach_info := executor.multi_frontend_attach_info()) is None - ): - raise ValueError( - "num_serve_frontends > 1 requires the classic IPC executor " - f"proxy in multi-frontend mode, got {type(executor).__name__}" - ) + # Raises the startup ValueError when the executor is not the classic IPC + # proxy in multi-frontend mode. + attach_info = llm.multi_frontend_attach_info() # Carries the executor HMAC keys; the child deletes it from its env # once consumed (GenerationExecutor.create). attach_env = json.dumps(attach_info) diff --git a/tensorrt_llm/serve/chat_utils.py b/tensorrt_llm/serve/chat_utils.py index f2d26b3c442a..02cef1e53243 100644 --- a/tensorrt_llm/serve/chat_utils.py +++ b/tensorrt_llm/serve/chat_utils.py @@ -149,7 +149,7 @@ def parse_chat_message_content_part( if part_type == "image_url": str_content = cast(str, content) image_io = _make_media_io("image", - mm_data_tracker._multimodal_server_config, + mm_data_tracker.multimodal_server_config, mm_data_tracker.request_media_io_kwargs) return MultimodalData(modality="image", data=image_io.async_load(str_content), @@ -168,7 +168,7 @@ async def decode_image_embeds(): if part_type == "video_url": str_content = cast(str, content) video_io = _make_media_io("video", - mm_data_tracker._multimodal_server_config, + mm_data_tracker.multimodal_server_config, mm_data_tracker.request_media_io_kwargs) return MultimodalData(modality="video", data=video_io.async_load(str_content), @@ -177,7 +177,7 @@ async def decode_image_embeds(): if part_type == "audio_url": str_content = cast(str, content) audio_io = _make_media_io("audio", - mm_data_tracker._multimodal_server_config, + mm_data_tracker.multimodal_server_config, mm_data_tracker.request_media_io_kwargs) return MultimodalData(modality="audio", data=audio_io.async_load(str_content), @@ -195,7 +195,7 @@ async def decode_image_embeds(): audio_format = dict_content.get("format", "") audio_url = f"data:audio/{audio_format};base64,{audio_data}" audio_io = _make_media_io("audio", - mm_data_tracker._multimodal_server_config, + mm_data_tracker.multimodal_server_config, mm_data_tracker.request_media_io_kwargs) return MultimodalData(modality="audio", data=audio_io.async_load(audio_url), diff --git a/tensorrt_llm/serve/openai_disagg_server.py b/tensorrt_llm/serve/openai_disagg_server.py index 92a298d4260f..4acb9d61c4f5 100644 --- a/tensorrt_llm/serve/openai_disagg_server.py +++ b/tensorrt_llm/serve/openai_disagg_server.py @@ -328,6 +328,11 @@ async def validation_exception_handler(request: Request, exc): self.register_routes() + @property + def address(self) -> tuple[str, int]: + """The configured listen address ``(hostname, port)`` of this disaggregated server.""" + return self._config.hostname, self._config.port + def _create_client(self, router: Router, role: ServerRole, max_retries: int = 1) -> OpenAIClient: async def disagg_id_generator(): return await self._coordinator.get_disagg_request_id() diff --git a/tensorrt_llm/serve/openai_protocol.py b/tensorrt_llm/serve/openai_protocol.py index 3f65a9146628..75042c88262a 100644 --- a/tensorrt_llm/serve/openai_protocol.py +++ b/tensorrt_llm/serve/openai_protocol.py @@ -633,7 +633,7 @@ def _response_format_text_config_to_guided_decoding_params( def _record_sampling_params_request_fields( request: OpenAIBaseModel, sampling_params: SamplingParams) -> None: """Preserve explicitly supplied fields across protocol defaulting.""" - sampling_params._set_request_provided_fields( + sampling_params.set_request_provided_fields( field_name for field_name in request.model_fields_set if getattr(request, field_name, None) is not None) diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py index d8adde6eafd7..90269e40f2aa 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -303,10 +303,7 @@ def _record_generator_termination(generator) -> None: reporting_source = "self" termination_kind = "exception" try: - from tensorrt_llm.executor.proxy import GenerationExecutorProxy - - if isinstance(getattr(generator, "_executor", None), - GenerationExecutorProxy): + if generator.engine_in_worker_process: component = "engine_worker" reporting_source = "executor_proxy" termination_kind = "worker_failure" @@ -964,11 +961,7 @@ def _supports_image_edit(self) -> bool: def _init_llm(self, chat_template: Optional[str] = None): self.tokenizer = self.generator.tokenizer - hf_tokenizer_path = self.generator._hf_model_dir - if not hf_tokenizer_path: - hf_tokenizer_path = getattr( - self.tokenizer.tokenizer, "name_or_path", None) or getattr( - self.tokenizer, "name_or_path", None) + hf_tokenizer_path = self.generator.hf_model_dir trust_remote_code = self.generator.args.trust_remote_code checkpoint_format = getattr(self.generator.args, "checkpoint_format", None) @@ -1261,7 +1254,7 @@ def register_routes(self): self.app.add_api_route("/kv_cache_events", self.get_kv_cache_events, methods=["POST"]) - resource_governor_queue = self.generator._executor.resource_governor_queue + resource_governor_queue = self.generator.resource_governor_queue if resource_governor_queue is not None: from .resource_governor import ResourceGovernor self.resource_governor = ResourceGovernor( @@ -1474,20 +1467,15 @@ def encode_fn(token_ids_batch): # returns one EncoderOutput per item, in input order. return self.generator.encode(token_ids_batch) - # Source the batch limits from the resolved encoder engine — the same - # attributes encode() validates against (llm.py: engine.max_seq_len / - # max_num_tokens / batch_size). LlmArgs.max_seq_len/max_num_tokens default - # to None and are not written back to args, so reading them from args - # would leave the batcher's seq-len/token-budget guards inert (typed 400s - # would never fire) and could let it form a batch encode() then rejects. - engine = self.generator._encoder_executor.model_engine + # The limits encode() checks every call against; see LLM.encode_limits. + limits = self.generator.encode_limits self.embedding_batcher = EncodeBatcher( encode_fn, - max_batch_size=engine.batch_size, + max_batch_size=limits.max_batch_size, max_queue_delay=self._embedding_max_queue_delay, max_queue_size=self._embedding_max_queue_size, - max_num_tokens=engine.max_num_tokens, - max_seq_len=engine.max_seq_len, + max_num_tokens=limits.max_num_tokens, + max_seq_len=limits.max_seq_len, ) def register_embedding_routes(self): @@ -1710,11 +1698,11 @@ def _on_launcher_lost(self, error: BaseException) -> None: Mirrors the fatal-error branch of ``health``. """ - executor = self.generator._executor - if executor is None: # the LLM already shut down + # record_fatal_error keeps only the first error; False means the LLM + # already shut down. + if not self.generator.record_fatal_error(error): return - executor._set_fatal_error(error) # keeps only the first error - if not executor.doing_shutdown: + if not self.generator.shutdown_started: _record_generator_termination(self.generator) signal.raise_signal(signal.SIGINT) diff --git a/tensorrt_llm/serve/postprocess_handlers.py b/tensorrt_llm/serve/postprocess_handlers.py index b882a2b6e7f3..f0d38a27384d 100644 --- a/tensorrt_llm/serve/postprocess_handlers.py +++ b/tensorrt_llm/serve/postprocess_handlers.py @@ -153,11 +153,10 @@ def _build_spec_decode_stats( finish_reason: Optional[str]) -> Optional[SpeculativeDecodingStats]: """Derive per-request speculative-decoding acceptance for one sequence. - ``counters`` is the sequence's own ``CompletionOutput._spec_dec_counters``, + ``counters`` is the sequence's own ``CompletionOutput.spec_dec_counters``, never the request-level copies on the GenerationResult: with n > 1 every candidate reports its own counters and those copies hold whichever - candidate responded last. Callers read it with getattr, since outputs that - are not a CompletionOutput (test doubles, for one) lack the field. + candidate responded last. Returns None -- meaning the field is omitted entirely -- when the caller has not opted in, when the sequence has not finished (streaming carries this @@ -593,8 +592,7 @@ def yield_first_chat(num_tokens: int, 'avg_decoded_tokens_per_iter', None), speculative_decoding=_build_spec_decode_stats( - getattr(output, '_spec_dec_counters', None), args, - output.finish_reason), + output.spec_dec_counters, args, output.finish_reason), stop_reason=output.stop_reason, ) if args.return_logprobs: @@ -761,8 +759,7 @@ def chat_response_post_processor( 'avg_decoded_tokens_per_iter', None), speculative_decoding=_build_spec_decode_stats( - getattr(output, '_spec_dec_counters', None), args, - output.finish_reason), + output.spec_dec_counters, args, output.finish_reason), ) if output.finish_reason == "stop" and args.has_tool_call.get( output.index, False): @@ -892,14 +889,13 @@ def completion_stream_post_processor(rsp: DetokenizedGenerationResultBase, 'avg_decoded_tokens_per_iter', None), speculative_decoding=_build_spec_decode_stats( - getattr(output, '_spec_dec_counters', None), args, - output.finish_reason), + output.spec_dec_counters, args, output.finish_reason), ) if args.return_logprobs: logprobs = output.logprobs_diff token_ids = output.token_ids_diff choice.logprobs = create_completion_logprobs( - token_ids, args.tokenizer, logprobs, output._last_text_len) + token_ids, args.tokenizer, logprobs, output.text_diff_offset) chunk = CompletionStreamResponse(model=args.model, choices=[choice], @@ -963,8 +959,7 @@ def completion_response_post_processor( 'avg_decoded_tokens_per_iter', None), speculative_decoding=_build_spec_decode_stats( - getattr(output, '_spec_dec_counters', None), args, - output.finish_reason), + output.spec_dec_counters, args, output.finish_reason), ) if args.return_logprobs: logprobs = output.logprobs diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index 69af33f55635..53195bcb91fd 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -340,6 +340,7 @@ l0_cpu: - unittest/serve/api/visual_gen/test_private_access_lint_visual_gen.py - unittest/serve/test_private_access_lint.py - unittest/serve/test_serve_health.py + - unittest/serve/test_serve_runtime_api_calls.py - unittest/tools - unittest/usage - unittest/utils/test_logger.py diff --git a/tests/unittest/_torch/pyexecutor/test_py_executor.py b/tests/unittest/_torch/pyexecutor/test_py_executor.py index 7167cccd3bc3..9016b074469e 100644 --- a/tests/unittest/_torch/pyexecutor/test_py_executor.py +++ b/tests/unittest/_torch/pyexecutor/test_py_executor.py @@ -1648,6 +1648,20 @@ def get_primary_block_counts(self): assert executor.get_openengine_discovery() == {"node": {"engine_id": "test"}} +def test_openengine_setters_store_what_the_getter_and_shutdown_read(): + """The OpenEngine node-discovery hook records through these two setters.""" + executor = object.__new__(PyExecutor) + + discovery = {"node": {"engine_id": "incarnation"}, "sources": []} + executor.set_openengine_discovery(discovery) + assert executor.get_openengine_discovery() is discovery + + node_server = Mock() + executor.set_openengine_node_server(node_server) + # Shutdown closes the adopted server through this private field. + assert executor._openengine_node_server is node_server + + def test_adp_running_counts_update_without_a_new_block_sample(): """The final request count must update even inside the sampling interval.""" from tensorrt_llm._torch.pyexecutor.scheduler.rank_state import RankIterStatsPayload, RankState diff --git a/tests/unittest/api_stability/api_stability_core.py b/tests/unittest/api_stability/api_stability_core.py index 66fd496635b1..9549994057da 100644 --- a/tests/unittest/api_stability/api_stability_core.py +++ b/tests/unittest/api_stability/api_stability_core.py @@ -6,7 +6,7 @@ import inspect import os import pathlib -from collections.abc import Mapping +from collections.abc import Iterable, Mapping from dataclasses import _HAS_DEFAULT_FACTORY_CLASS, dataclass, fields from pprint import pprint from types import MethodType, NoneType diff --git a/tests/unittest/api_stability/references/completion_output.yaml b/tests/unittest/api_stability/references/completion_output.yaml index e183ebe20d0e..331e4a7f34c9 100644 --- a/tests/unittest/api_stability/references/completion_output.yaml +++ b/tests/unittest/api_stability/references/completion_output.yaml @@ -35,3 +35,9 @@ properties: routed_experts: annotation: Optional[torch.Tensor] default: inspect._empty + spec_dec_counters: + annotation: Optional[tensorrt_llm.executor.result.SpecDecCounters] + default: inspect._empty + text_diff_offset: + annotation: int + default: inspect._empty diff --git a/tests/unittest/api_stability/references/llm.yaml b/tests/unittest/api_stability/references/llm.yaml index b1be771de618..cf90ce180e25 100644 --- a/tests/unittest/api_stability/references/llm.yaml +++ b/tests/unittest/api_stability/references/llm.yaml @@ -507,6 +507,14 @@ methods: parameters: {} return_annotation: dict status: beta + get_kv_cache_load: + parameters: {} + return_annotation: dict + status: prototype + get_openengine_discovery: + parameters: {} + return_annotation: dict + status: prototype get_kv_cache_events_async: parameters: timeout: @@ -532,6 +540,17 @@ methods: parameters: {} return_annotation: bool status: prototype + record_fatal_error: + parameters: + error: + annotation: BaseException + default: inspect._empty + return_annotation: bool + status: prototype + multi_frontend_attach_info: + parameters: {} + return_annotation: dict + status: prototype get_data_transceiver_state: parameters: {} return_annotation: bytes @@ -597,3 +616,21 @@ properties: engine_dead: annotation: bool default: inspect._empty + encode_limits: + annotation: Optional[tensorrt_llm.llmapi.llm.EncodeLimits] + default: inspect._empty + engine_in_worker_process: + annotation: bool + default: inspect._empty + hf_model_config: + annotation: Optional[transformers.PretrainedConfig] + default: inspect._empty + hf_model_dir: + annotation: Union[str, pathlib.Path] + default: inspect._empty + multi_frontend_ipc_dir: + annotation: Optional[str] + default: inspect._empty + resource_governor_queue: + annotation: Optional[Any] + default: inspect._empty diff --git a/tests/unittest/api_stability/references/sampling_params.yaml b/tests/unittest/api_stability/references/sampling_params.yaml index ef870f77bde3..e99339f2fd60 100644 --- a/tests/unittest/api_stability/references/sampling_params.yaml +++ b/tests/unittest/api_stability/references/sampling_params.yaml @@ -32,4 +32,13 @@ methods: annotation: bool default: false return_annotation: None -properties: {} + set_request_provided_fields: + parameters: + field_names: + annotation: Iterable[str] + default: inspect._empty + return_annotation: None +properties: + stop_word_ids: + annotation: Optional[List[List[int]]] + default: inspect._empty diff --git a/tests/unittest/disaggregated/test_openai_disagg_service.py b/tests/unittest/disaggregated/test_openai_disagg_service.py index 46b063e71584..f605d96e31fa 100644 --- a/tests/unittest/disaggregated/test_openai_disagg_service.py +++ b/tests/unittest/disaggregated/test_openai_disagg_service.py @@ -104,6 +104,19 @@ def _make_service(schedule_style: str) -> OpenAIDisaggregatedService: return service +def test_disagg_server_address_is_the_configured_listen_address( + monkeypatch: pytest.MonkeyPatch, tmp_path +) -> None: + """The fleet launcher binds its SO_REUSEPORT socket to this address.""" + monkeypatch.setenv("PROMETHEUS_MULTIPROC_DIR", str(tmp_path)) + config = DisaggServerConfig(server_configs=[], hostname="127.0.0.1", port=18000) + + server = OpenAIDisaggServer(config=config) + + assert server.address == (config.hostname, config.port) + assert server.address == ("127.0.0.1", 18000) + + @pytest.mark.asyncio async def test_conditional_disagg_uses_selected_server_match_length(): service = _make_service("context_first") @@ -1091,6 +1104,7 @@ def test_generation_postprocessor_rewrites_usage_from_disaggregated_params(): stop_reason=None, finish_reason="stop", length=5, + spec_dec_counters=None, ) ], context_logits=None, @@ -1145,6 +1159,7 @@ def test_generation_chat_postprocessor_rewrites_usage_from_output_disaggregated_ ), stop_reason=None, finish_reason="stop", + spec_dec_counters=None, ) ], cached_tokens=3, diff --git a/tests/unittest/grpc/openengine/test_control.py b/tests/unittest/grpc/openengine/test_control.py index 38149c34e7b1..2aab5f815a30 100644 --- a/tests/unittest/grpc/openengine/test_control.py +++ b/tests/unittest/grpc/openengine/test_control.py @@ -104,6 +104,7 @@ def _llm(**overrides): guided_decoding_backend="xgrammar", max_beam_width=1, reasoning_parser=None, + encode_only=False, kv_cache_config=SimpleNamespace(tokens_per_block=32), _enable_routing_load=overrides.get("_enable_routing_load", "_kv_load" in overrides), ) @@ -113,13 +114,6 @@ def _llm(**overrides): setattr(args, key, value) shutdown = overrides.get("_executor_shutdown", False) unhealthy = overrides.get("_executor_unhealthy", False) - executor = SimpleNamespace( - is_shutdown=lambda: shutdown, - check_health=lambda: not (shutdown or unhealthy), - get_kv_cache_capacity=lambda: overrides.get("_kv_capacity", {}), - get_openengine_discovery=lambda: overrides.get("_openengine_discovery", {}), - get_kv_cache_load=lambda: overrides.get("_kv_load", {}), - ) return SimpleNamespace( args=args, llm_id="instance-1", @@ -127,8 +121,10 @@ def _llm(**overrides): input_processor=SimpleNamespace( supports_token_ids_input=overrides.get("_supports_token_ids_input", True) ), - _executor=executor, - check_health=lambda: executor.check_health(), + check_health=lambda: not (shutdown or unhealthy), + get_kv_cache_capacity=lambda: overrides.get("_kv_capacity", {}), + get_openengine_discovery=lambda: overrides.get("_openengine_discovery", {}), + get_kv_cache_load=lambda: overrides.get("_kv_load", {}), ) diff --git a/tests/unittest/grpc/openengine/test_frontend_lifecycle.py b/tests/unittest/grpc/openengine/test_frontend_lifecycle.py index ca4ef9a222af..404396ad7522 100644 --- a/tests/unittest/grpc/openengine/test_frontend_lifecycle.py +++ b/tests/unittest/grpc/openengine/test_frontend_lifecycle.py @@ -146,8 +146,6 @@ def select(self, timeout): import sys import time from types import SimpleNamespace -from unittest.mock import Mock -from tensorrt_llm.executor.proxy import GenerationExecutorProxy from tensorrt_llm.grpc.openengine import server root = Path(sys.argv[1]) @@ -177,8 +175,8 @@ def on_term(signum, frame): self.args = SimpleNamespace(guided_decoding_backend=None) self.tokenizer = None self.engine_dead = False - self._executor = Mock(spec=GenerationExecutorProxy) - self._executor.multi_frontend_attach_info.return_value = {"mode": "classic"} + def multi_frontend_attach_info(self): + return {"mode": "classic"} def check_health(self): return not ( frontend == 0 and sys.argv[3] == "child-startup-engine-failure" diff --git a/tests/unittest/grpc/openengine/test_metrics.py b/tests/unittest/grpc/openengine/test_metrics.py index 32dd0edb402e..36edd676daf8 100644 --- a/tests/unittest/grpc/openengine/test_metrics.py +++ b/tests/unittest/grpc/openengine/test_metrics.py @@ -614,9 +614,7 @@ def _http_enablement() -> None: args = _args() args.return_perf_metrics = requested args.perf_metrics_output_dir = output_dir - generator = SimpleNamespace( - args=args, _executor=SimpleNamespace(resource_governor_queue=None) - ) + generator = SimpleNamespace(args=args, resource_governor_queue=None) server = OpenAIServer(generator, "test-model", None, None, None) try: enabled = requested or output_dir is not None diff --git a/tests/unittest/grpc/openengine/test_multimodal.py b/tests/unittest/grpc/openengine/test_multimodal.py index 6e27fbc2dbd8..575923f46cc2 100644 --- a/tests/unittest/grpc/openengine/test_multimodal.py +++ b/tests/unittest/grpc/openengine/test_multimodal.py @@ -193,7 +193,7 @@ def abort(self) -> None: class _ImageLlm: def __init__(self, input_processor, config) -> None: self.input_processor = input_processor - self._hf_model_config = config + self.hf_model_config = config self.args = SimpleNamespace(guided_decoding_backend=None) self.tokenizer = _Tokenizer() output = SimpleNamespace( @@ -285,7 +285,7 @@ def is_valid(model_type, modality): ) return ( SimpleNamespace( - _hf_model_config=config, + hf_model_config=config, input_processor=input_processor, ), seen_model_types, @@ -310,7 +310,7 @@ def image_video_llm(monkeypatch): lambda model_type, modality: f"<{modality}_{{}}>", ) return SimpleNamespace( - _hf_model_config=config, + hf_model_config=config, input_processor=input_processor, ) @@ -333,7 +333,7 @@ def image_video_audio_llm(monkeypatch): lambda model_type, modality: f"<{modality}_{{}}>", ) return SimpleNamespace( - _hf_model_config=config, + hf_model_config=config, input_processor=input_processor, ) @@ -345,7 +345,7 @@ def test_media_capability_requires_the_loaded_multimodal_processor(image_llm): assert seen_model_types == [_ImageConfig.model_type] * 3 assert ( request_mapping.media_capabilities_for_model( - SimpleNamespace(_hf_model_config=llm._hf_model_config, input_processor=object()) + SimpleNamespace(hf_model_config=llm.hf_model_config, input_processor=object()) ).modalities == set() ) @@ -364,7 +364,7 @@ def test_media_capability_accepts_a_registered_non_base_processor(monkeypatch): lambda model_type, modality: "", ) llm = SimpleNamespace( - _hf_model_config=config, + hf_model_config=config, input_processor=_AudioInputProcessor(config), ) @@ -418,7 +418,7 @@ def test_media_capability_rejects_a_registered_non_callable_processor(monkeypatc config = _ImageConfig() monkeypatch.setattr(MULTIMODAL_PLACEHOLDER_REGISTRY, "is_valid", lambda *_args: True) llm = SimpleNamespace( - _hf_model_config=config, + hf_model_config=config, input_processor=_NonCallableRegisteredInputProcessor(config), ) @@ -429,7 +429,7 @@ def test_media_capability_rejects_a_registered_non_callable_processor(monkeypatc @pytest.mark.parametrize("disabled", [False, True]) async def test_disabled_mm_encoder_agrees_between_discovery_and_generate(image_llm, disabled): capability_llm, _ = image_llm - llm = _ImageLlm(capability_llm.input_processor, capability_llm._hf_model_config) + llm = _ImageLlm(capability_llm.input_processor, capability_llm.hf_model_config) llm.args.disable_mm_encoder = disabled llm.args.reasoning_parser = None llm.args.max_seq_len = 2048 @@ -726,7 +726,7 @@ def test_gemma4_missing_processor_config_fails_closed(config_state): processor = Gemma4InputProcessor.__new__(Gemma4InputProcessor) if config_state == "none": processor._config = None - llm = SimpleNamespace(input_processor=processor, _hf_model_config=config) + llm = SimpleNamespace(input_processor=processor, hf_model_config=config) assert processor.supported_modalities == frozenset() assert request_mapping.media_capabilities_for_model(llm).modalities == frozenset() @@ -1029,7 +1029,7 @@ def test_passthrough_placeholder_policy_follows_processor_input_path( from tensorrt_llm.inputs.content_format import ContentFormat capability_llm, _ = image_llm - llm = _ImageLlm(capability_llm.input_processor, capability_llm._hf_model_config) + llm = _ImageLlm(capability_llm.input_processor, capability_llm.hf_model_config) monkeypatch.setattr(llm.input_processor, "supports_token_id_mm_expansion", fast_path) monkeypatch.setattr( MULTIMODAL_PLACEHOLDER_REGISTRY, @@ -1112,7 +1112,7 @@ async def test_media_and_preprocessing_do_not_block_control_executor(image_llm, loop = asyncio.get_running_loop() capability_llm, _ = image_llm - llm = _ImageLlm(capability_llm.input_processor, capability_llm._hf_model_config) + llm = _ImageLlm(capability_llm.input_processor, capability_llm.hf_model_config) llm.args = SimpleNamespace( guided_decoding_backend=None, tensor_parallel_size=1, @@ -1120,11 +1120,11 @@ async def test_media_and_preprocessing_do_not_block_control_executor(image_llm, context_parallel_size=1, max_batch_size=8, max_num_tokens=2048, + encode_only=False, kv_cache_config=SimpleNamespace(tokens_per_block=32), ) - llm._executor = SimpleNamespace( - get_kv_cache_capacity=lambda: {}, get_openengine_discovery=lambda: {} - ) + llm.get_kv_cache_capacity = lambda: {} + llm.get_openengine_discovery = lambda: {} with ( ThreadPoolExecutor(max_workers=1, thread_name_prefix="test_input") as input_pool, ThreadPoolExecutor(max_workers=1, thread_name_prefix="test_media") as media_pool, @@ -1432,7 +1432,7 @@ async def test_request_validation_precedes_media_loading( capability_llm, _ = image_llm llm = _ImageLlm( capability_llm.input_processor, - capability_llm._hf_model_config, + capability_llm.hf_model_config, ) class _ImageIo: @@ -1470,7 +1470,7 @@ async def test_capability_error_precedes_other_request_validation(image_llm): capability_llm, _ = image_llm llm = _ImageLlm( capability_llm.input_processor, - capability_llm._hf_model_config, + capability_llm.hf_model_config, ) request = generation_pb2.GenerateRequest( request_id="unsupported-video", @@ -1784,7 +1784,7 @@ async def test_generate_submits_image_and_reports_expanded_prompt_usage(image_ll capability_llm, _ = image_llm llm = _ImageLlm( capability_llm.input_processor, - capability_llm._hf_model_config, + capability_llm.hf_model_config, ) class _ImageIo: @@ -1828,7 +1828,7 @@ async def test_placeholder_mismatch_does_not_poison_hashing( import torch capability_llm, _ = image_llm - llm = _ImageLlm(capability_llm.input_processor, capability_llm._hf_model_config) + llm = _ImageLlm(capability_llm.input_processor, capability_llm.hf_model_config) processor = llm.input_processor processor.multimodal_hashing_supported = None monkeypatch.setattr(processor, "get_mm_token_ids", lambda: torch.tensor([21, 23])) @@ -1917,7 +1917,7 @@ async def test_media_preprocessing_is_offloaded_and_abort_aware( image_video_llm, monkeypatch, outcome, modality ): capability_llm = image_video_llm - llm = _ImageLlm(capability_llm.input_processor, capability_llm._hf_model_config) + llm = _ImageLlm(capability_llm.input_processor, capability_llm.hf_model_config) loop = asyncio.get_running_loop() serving_thread = threading.get_ident() request_context = ContextVar("preprocess_request_context", default=None) @@ -2090,7 +2090,7 @@ async def collect_text(): ) async def test_preprocessing_errors_are_classified(image_llm, monkeypatch, error, expected_code): capability_llm, _ = image_llm - llm = _ImageLlm(capability_llm.input_processor, capability_llm._hf_model_config) + llm = _ImageLlm(capability_llm.input_processor, capability_llm.hf_model_config) class _ImageIo: @classmethod @@ -2136,7 +2136,7 @@ async def test_duplicate_request_id_is_reserved_while_media_loads(image_llm, mon capability_llm, _ = image_llm llm = _ImageLlm( capability_llm.input_processor, - capability_llm._hf_model_config, + capability_llm.hf_model_config, ) load_started = asyncio.Event() load_cancelled = asyncio.Event() @@ -2209,7 +2209,7 @@ async def test_media_loading_stops_for_coordinator_abort_or_rpc_cancellation( capability_llm, _ = image_llm llm = _ImageLlm( capability_llm.input_processor, - capability_llm._hf_model_config, + capability_llm.hf_model_config, ) load_started = asyncio.Event() load_cancelled = asyncio.Event() @@ -2284,7 +2284,7 @@ async def test_coordinator_abort_before_reservation_binding_prevents_media_loadi capability_llm, _ = image_llm llm = _ImageLlm( capability_llm.input_processor, - capability_llm._hf_model_config, + capability_llm.hf_model_config, ) decoder_creations = 0 @@ -2339,7 +2339,7 @@ async def test_generate_maps_media_errors_to_grpc_statuses(image_llm): capability_llm, _ = image_llm llm = _ImageLlm( capability_llm.input_processor, - capability_llm._hf_model_config, + capability_llm.hf_model_config, ) servicer = OpenEngineInferenceServicer(llm, model="image-model") @@ -2389,7 +2389,7 @@ async def test_generate_maps_transient_media_fetch_to_failed_precondition( capability_llm, _ = image_llm llm = _ImageLlm( capability_llm.input_processor, - capability_llm._hf_model_config, + capability_llm.hf_model_config, ) class _ImageIo: diff --git a/tests/unittest/grpc/openengine/test_node.py b/tests/unittest/grpc/openengine/test_node.py index 4aa518b71810..831099228754 100644 --- a/tests/unittest/grpc/openengine/test_node.py +++ b/tests/unittest/grpc/openengine/test_node.py @@ -21,6 +21,18 @@ pytestmark = pytest.mark.cpu_only +def _fake_node_executor(**kwargs): + """A PyExecutor double whose discovery setters store into the read attributes.""" + executor = SimpleNamespace(**kwargs) + executor.set_openengine_discovery = lambda discovery: setattr( + executor, "_openengine_discovery", discovery + ) + executor.set_openengine_node_server = lambda server: setattr( + executor, "_openengine_node_server", server + ) + return executor + + def _source(rank, host): return { "transport": "zmq", @@ -100,7 +112,7 @@ def close(self): [None, None, None], ] ) - executor = SimpleNamespace( + executor = _fake_node_executor( llm_args=SimpleNamespace( _openengine_discovery={ "engine_id": "incarnation", @@ -127,7 +139,7 @@ def close(self): def test_startup_rejects_missing_attention_dp_publisher(monkeypatch): monkeypatch.setattr(node.socket, "getfqdn", lambda: "follower") rows = [(0, "leader", _source(0, "leader"), None), (1, "follower", None, None)] - executor = SimpleNamespace( + executor = _fake_node_executor( llm_args=SimpleNamespace( _openengine_discovery={"engine_id": "incarnation", "host": "127.0.0.1", "port": 15051}, kv_cache_config=SimpleNamespace(kv_events_config=None), @@ -162,7 +174,7 @@ def close(self): [None, "another node could not bind"], ] ) - executor = SimpleNamespace( + executor = _fake_node_executor( llm_args=SimpleNamespace( _openengine_discovery={"engine_id": "incarnation", "host": "127.0.0.1", "port": 15051}, kv_cache_config=SimpleNamespace(kv_events_config=None), @@ -190,7 +202,7 @@ def allgather(value): gathered.append(value) return placements if len(gathered) == 1 else [None, value] - executor = SimpleNamespace( + executor = _fake_node_executor( llm_args=SimpleNamespace( _openengine_discovery={"engine_id": "incarnation", "port": 15051}, kv_cache_config=SimpleNamespace(kv_events_config=None), diff --git a/tests/unittest/grpc/openengine/test_openengine.py b/tests/unittest/grpc/openengine/test_openengine.py index 4e6456a8cfe4..1453302320d3 100644 --- a/tests/unittest/grpc/openengine/test_openengine.py +++ b/tests/unittest/grpc/openengine/test_openengine.py @@ -110,11 +110,14 @@ def test_openengine_server_serves_the_control_contract() -> None: context_parallel_size=1, guided_decoding_backend=None, reasoning_parser=None, + encode_only=False, kv_cache_config=SimpleNamespace(tokens_per_block=32), ), llm_id="test-instance", tokenizer=object(), check_health=lambda: True, + get_kv_cache_capacity=lambda: {}, + get_openengine_discovery=lambda: {}, ) async def exercise_server() -> None: @@ -526,10 +529,7 @@ def test_event_startup_requires_loaded_publisher(publisher_enabled: bool) -> Non kv_events_config=SimpleNamespace(enable_kv_cache_events=True, publisher="zmq") ), ), - _executor=SimpleNamespace( - get_kv_cache_capacity=lambda: {"kvEventsEnabled": True} if publisher_enabled else {}, - is_shutdown=lambda: False, - ), + get_kv_cache_capacity=lambda: {"kvEventsEnabled": True} if publisher_enabled else {}, ) async def exercise_server() -> None: diff --git a/tests/unittest/llmapi/apps/openai_server_stub.py b/tests/unittest/llmapi/apps/openai_server_stub.py index 12d00af5b73d..a633dd421650 100644 --- a/tests/unittest/llmapi/apps/openai_server_stub.py +++ b/tests/unittest/llmapi/apps/openai_server_stub.py @@ -55,8 +55,8 @@ def make_stub_openai_server( tokenizer=tokenizer, input_processor=None, llm_id="stub", - _hf_model_dir=str(tmp_path), - _executor=SimpleNamespace(resource_governor_queue=None), + hf_model_dir=str(tmp_path), + resource_governor_queue=None, generate_async=generate_async, ) previous_media_executor = BaseMediaIO._executor diff --git a/tests/unittest/llmapi/apps/test_anthropic_routes.py b/tests/unittest/llmapi/apps/test_anthropic_routes.py index b7e6cf62a3d6..bc18a3e52686 100644 --- a/tests/unittest/llmapi/apps/test_anthropic_routes.py +++ b/tests/unittest/llmapi/apps/test_anthropic_routes.py @@ -477,7 +477,7 @@ def test_standard_and_disagg_register_messages_route(prometheus_multiproc_dir): standard = object.__new__(OpenAIServer) standard.app = FastAPI() standard.generator = SimpleNamespace( - _executor=SimpleNamespace(resource_governor_queue=None), + resource_governor_queue=None, args=SimpleNamespace(return_perf_metrics=False), ) standard.use_harmony = False @@ -570,7 +570,7 @@ 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), + resource_governor_queue=None, args=SimpleNamespace(return_perf_metrics=False), ) standard.use_harmony = False diff --git a/tests/unittest/llmapi/apps/test_chat_utils.py b/tests/unittest/llmapi/apps/test_chat_utils.py index 12b51422784b..e44594f42509 100644 --- a/tests/unittest/llmapi/apps/test_chat_utils.py +++ b/tests/unittest/llmapi/apps/test_chat_utils.py @@ -25,6 +25,16 @@ def mock_mm_data_tracker(): return MagicMock() +def test_multimodal_server_config_property(): + """Chat parsing reads the tracker's media settings through this property.""" + config = MultimodalServerConfig() + assert MultimodalDataTracker("llava", config).multimodal_server_config is config + # Without an explicit config the tracker builds a default one. + assert isinstance( + MultimodalDataTracker("llava").multimodal_server_config, MultimodalServerConfig + ) + + class TestParseAssistantMessages: """Test suite for assistant role messages.""" diff --git a/tests/unittest/llmapi/apps/test_openai_server_iteration_stats.py b/tests/unittest/llmapi/apps/test_openai_server_iteration_stats.py index 9178486417c5..d31bcc935bc3 100644 --- a/tests/unittest/llmapi/apps/test_openai_server_iteration_stats.py +++ b/tests/unittest/llmapi/apps/test_openai_server_iteration_stats.py @@ -51,7 +51,7 @@ def __init__( enable_iter_perf_stats=enable_iter_perf_stats, return_perf_metrics=False, ) - self._executor = SimpleNamespace(resource_governor_queue=None) + self.resource_governor_queue = None self._stat_batches = deque(stat_batches) self._stats_gate = stats_gate self._stats_error = stats_error diff --git a/tests/unittest/llmapi/apps/test_reasoning_prompt_resolution.py b/tests/unittest/llmapi/apps/test_reasoning_prompt_resolution.py index 84149f07cc41..77a22a403d41 100644 --- a/tests/unittest/llmapi/apps/test_reasoning_prompt_resolution.py +++ b/tests/unittest/llmapi/apps/test_reasoning_prompt_resolution.py @@ -142,6 +142,7 @@ def __init__(self, text: str) -> None: self.finish_reason = "stop" self.stop_reason = None self.logprobs = None + self.spec_dec_counters = None self.disaggregated_params = LlmDisaggregatedParams(request_type="context_only") diff --git a/tests/unittest/llmapi/apps/test_tool_parsers.py b/tests/unittest/llmapi/apps/test_tool_parsers.py index 691820959447..729b80929ee8 100644 --- a/tests/unittest/llmapi/apps/test_tool_parsers.py +++ b/tests/unittest/llmapi/apps/test_tool_parsers.py @@ -6340,6 +6340,7 @@ def _fake_response(text, finish_reason="stop"): token_ids=[1, 2, 3], finish_reason=finish_reason, stop_reason=None, + spec_dec_counters=None, disaggregated_params=None) return SimpleNamespace(outputs=[output], cached_tokens=0) @@ -6365,6 +6366,7 @@ def _stream_deltas(args, chunks, finish_reason="stop"): finish_reason=finish_reason if last else None, stop_reason=None, length=chunk_index + 1, + spec_dec_counters=None, disaggregated_params=None, ) rsp = SimpleNamespace(outputs=[output], diff --git a/tests/unittest/llmapi/test_executor.py b/tests/unittest/llmapi/test_executor.py index 6dbe016150c0..9629f2905f07 100644 --- a/tests/unittest/llmapi/test_executor.py +++ b/tests/unittest/llmapi/test_executor.py @@ -26,6 +26,7 @@ GenerationResultBase, PostprocWorker) from tensorrt_llm.executor.ipc import FusedIpcQueue, ZeroMqQueue from tensorrt_llm.executor.proxy import GenerationExecutorProxy +from tensorrt_llm.executor.result import CompletionOutput, SpecDecCounters from tensorrt_llm.executor.utils import RequestError from tensorrt_llm.inputs.multimodal import MultimodalInput, MultimodalParams from tensorrt_llm.llmapi.tokenizer import TransformersTokenizer @@ -116,6 +117,34 @@ def test_GenerationResultBase(): assert result._done +def test_completion_output_spec_dec_counters_property(): + output = CompletionOutput(0) + assert output.spec_dec_counters is None + + counters = SpecDecCounters(per_pos_drafted=[2], + per_pos_accepted=[1], + spec_dec_totals=(2, 1)) + output._spec_dec_counters = counters + assert output.spec_dec_counters is counters + + # The result path fills the field per sequence; the property returns it. + sampling_params = SamplingParams(max_tokens=4) + result = GenerationResultBase(id=7, sampling_params=sampling_params) + result._handle_response(create_rsp(2, finished=True)) + sequence = result.outputs[0] + assert sequence.spec_dec_counters is sequence._spec_dec_counters + + +def test_completion_output_text_diff_offset(): + output = CompletionOutput(0) + assert output.text_diff_offset == 0 + + output.text = "hello world" + output._last_text_len = 6 + assert output.text_diff == "world" + assert output.text_diff_offset == len(output.text) - len(output.text_diff) + + @pytest.mark.parametrize("pytorch", [False, True]) @pytest.mark.parametrize("beam_search", [False, True]) @pytest.mark.parametrize("streaming", [False, True]) diff --git a/tests/unittest/llmapi/test_llm_public_runtime_api.py b/tests/unittest/llmapi/test_llm_public_runtime_api.py index e1f9b321d1ee..6175d4a589ef 100644 --- a/tests/unittest/llmapi/test_llm_public_runtime_api.py +++ b/tests/unittest/llmapi/test_llm_public_runtime_api.py @@ -1,18 +1,26 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Production tests of the public runtime health API the serving frontends use. - -Covers BaseLLM.check_health(), BaseLLM.fatal_error, BaseLLM.shutdown_started, -BaseLLM.engine_dead and the executor-side GenerationExecutor.fatal_error and -engine_dead, against the real classes (no mirrors; the LLM objects are built -with object.__new__ plus exactly the attributes BaseLLM.__init__ sets). +"""Production tests of the public runtime API the serving frontends use. + +Covers the health members (BaseLLM.check_health(), fatal_error, +shutdown_started, engine_dead and the executor-side equivalents), the runtime +accessors the frontends read (record_fatal_error, multi_frontend_attach_info, +multi_frontend_ipc_dir, engine_in_worker_process, hf_model_dir, +hf_model_config, resource_governor_queue, encode_limits, get_kv_cache_load, +get_openengine_discovery) and the EngineClient protocol conformance of LLM, +against the real classes (no mirrors; the LLM objects are built with +object.__new__ plus exactly the attributes BaseLLM.__init__ sets). """ +import inspect +from types import SimpleNamespace + import pytest from tensorrt_llm.executor.executor import GenerationExecutor -from tensorrt_llm.executor.proxy import GenerationExecutorProxy -from tensorrt_llm.llmapi.llm import LLM +from tensorrt_llm.executor.proxy import GenerationExecutorFrontendProxy, GenerationExecutorProxy +from tensorrt_llm.llmapi.engine_client import EngineClient +from tensorrt_llm.llmapi.llm import LLM, EncodeLimits pytestmark = pytest.mark.cpu_only @@ -179,3 +187,224 @@ def test_executor_engine_dead_property(): def test_proxy_engine_dead_reads_both_fields(engine_dead, fatal_error, expected): """_mark_engine_dead can set the sticky flag before the error is recorded.""" assert _proxy(engine_dead=engine_dead, fatal_error=fatal_error).engine_dead is expected + + +# --------------------------------------------------------------------------- +# BaseLLM.record_fatal_error +# --------------------------------------------------------------------------- +def test_record_fatal_error_without_executor_is_a_no_op(): + assert _llm(None).record_fatal_error(RuntimeError("lost")) is False + + +def test_record_fatal_error_records_on_the_executor_first_one_wins(): + llm = _llm(_NoOpExecutor()) + first = RuntimeError("first fatal error") + assert llm.record_fatal_error(first) is True + assert llm.record_fatal_error(RuntimeError("second fatal error")) is True + assert llm.fatal_error is first + + +# --------------------------------------------------------------------------- +# BaseLLM.multi_frontend_attach_info and multi_frontend_ipc_dir +# --------------------------------------------------------------------------- +def test_multi_frontend_attach_info_returns_the_proxy_payload(): + proxy = _proxy() + payload = {"mode": "classic"} + proxy.multi_frontend_attach_info = lambda: payload + assert _llm(proxy).multi_frontend_attach_info() is payload + + +def test_multi_frontend_attach_info_raises_for_a_single_frontend_proxy(): + """The real proxy method returns None outside multi-frontend mode.""" + proxy = _proxy() + proxy._num_frontends = 1 + with pytest.raises( + ValueError, + match=( + "num_serve_frontends > 1 requires the classic IPC executor " + "proxy in multi-frontend mode, got GenerationExecutorProxy" + ), + ): + _llm(proxy).multi_frontend_attach_info() + + +def test_multi_frontend_attach_info_raises_for_a_non_proxy_executor(): + with pytest.raises(ValueError, match="got _NoOpExecutor"): + _llm(_NoOpExecutor()).multi_frontend_attach_info() + + +def test_multi_frontend_attach_info_raises_without_an_executor(): + with pytest.raises(ValueError, match="got NoneType"): + _llm(None).multi_frontend_attach_info() + + +def test_multi_frontend_ipc_dir_reads_the_proxy_directory(): + proxy = _proxy() + proxy._multi_frontend_ipc_dir = "/tmp/trtllm_frontends_x" + assert _llm(proxy).multi_frontend_ipc_dir == "/tmp/trtllm_frontends_x" + + +def test_multi_frontend_ipc_dir_none_outside_multi_frontend_mode(): + assert _llm(None).multi_frontend_ipc_dir is None + assert _llm(_NoOpExecutor()).multi_frontend_ipc_dir is None + # An attached frontend skips GenerationExecutorProxy.__init__, so the + # private field can be missing entirely; the property stays tolerant. + frontend = object.__new__(GenerationExecutorFrontendProxy) + assert _llm(frontend).multi_frontend_ipc_dir is None + + +# --------------------------------------------------------------------------- +# BaseLLM.engine_in_worker_process +# --------------------------------------------------------------------------- +def test_engine_in_worker_process_true_for_the_ipc_proxies(): + assert _llm(_proxy()).engine_in_worker_process is True + frontend = object.__new__(GenerationExecutorFrontendProxy) + assert _llm(frontend).engine_in_worker_process is True + + +def test_engine_in_worker_process_false_for_in_process_executors(): + assert _llm(_NoOpExecutor()).engine_in_worker_process is False + assert _llm(None).engine_in_worker_process is False + + +# --------------------------------------------------------------------------- +# BaseLLM.hf_model_dir and hf_model_config +# --------------------------------------------------------------------------- +def test_hf_model_dir_returns_the_resolved_checkpoint_directory(): + llm = _llm(None) + llm._hf_model_dir = "/models/resolved" + llm.args = SimpleNamespace(model="org/model") + assert llm.hf_model_dir == "/models/resolved" + + +def test_hf_model_dir_falls_back_to_args_model_until_resolved(): + llm = _llm(None) + llm._hf_model_dir = None + llm.args = SimpleNamespace(model="org/model") + assert llm.hf_model_dir == "org/model" + + +def test_hf_model_config_returns_the_loaded_config(): + llm = _llm(None) + llm._hf_model_config = None + assert llm.hf_model_config is None + config = SimpleNamespace(model_type="llama") + llm._hf_model_config = config + assert llm.hf_model_config is config + + +# --------------------------------------------------------------------------- +# BaseLLM.resource_governor_queue +# --------------------------------------------------------------------------- +def test_resource_governor_queue_forwards_the_executor_value(): + assert _llm(_NoOpExecutor()).resource_governor_queue is None + proxy = _proxy() + queue = object() + proxy._resource_governor_queue = queue + assert _llm(proxy).resource_governor_queue is queue + + +def test_resource_governor_queue_raises_without_an_executor(): + with pytest.raises(RuntimeError, match="resource_governor_queue is not available"): + _llm(None).resource_governor_queue + + +def test_resource_governor_queue_attribute_error_on_a_bare_instance(): + """AttributeError (not RuntimeError) on an attribute-less instance. + + The api-stability snapshotter inspects such instances and tolerates only + AttributeError. + """ + with pytest.raises(AttributeError): + object.__new__(LLM).resource_governor_queue + + +# --------------------------------------------------------------------------- +# BaseLLM.encode_limits and encode() +# --------------------------------------------------------------------------- +def _encoder(batch_size=4, max_num_tokens=512, max_seq_len=128): + return SimpleNamespace( + model_engine=SimpleNamespace( + batch_size=batch_size, max_num_tokens=max_num_tokens, max_seq_len=max_seq_len + ) + ) + + +def test_encode_limits_reads_the_encoder_engine(): + llm = _llm(None, encode_only=True, encoder_executor=_encoder()) + assert llm.encode_limits == EncodeLimits(max_batch_size=4, max_num_tokens=512, max_seq_len=128) + + +def test_encode_limits_none_outside_encode_only_and_after_shutdown(): + assert _llm(_NoOpExecutor()).encode_limits is None + assert _llm(None, encode_only=True, encoder_executor=None).encode_limits is None + + +def test_encode_rejects_a_batch_above_encode_limits(): + llm = _llm(None, encode_only=True, encoder_executor=_encoder(batch_size=2)) + with pytest.raises(ValueError, match=r"Batch size \(3\) exceeds max_batch_size"): + llm.encode([[1], [2], [3]]) + + +# --------------------------------------------------------------------------- +# BaseLLM.get_kv_cache_load and get_openengine_discovery +# --------------------------------------------------------------------------- +def test_kv_load_and_discovery_forward_the_executor_values(): + class _ReportingExecutor(_NoOpExecutor): + def get_kv_cache_load(self): + return {"usedNumBlocks": 1} + + def get_openengine_discovery(self): + return {"node": {"engine_id": "x"}} + + llm = _llm(_ReportingExecutor()) + assert llm.get_kv_cache_load() == {"usedNumBlocks": 1} + assert llm.get_openengine_discovery() == {"node": {"engine_id": "x"}} + + +def test_kv_load_and_discovery_empty_without_an_executor(): + assert _llm(None).get_kv_cache_load() == {} + assert _llm(None).get_openengine_discovery() == {} + # The base executor reports nothing by default. + assert _llm(_NoOpExecutor()).get_kv_cache_load() == {} + assert _llm(_NoOpExecutor()).get_openengine_discovery() == {} + + +# --------------------------------------------------------------------------- +# EngineClient conformance +# --------------------------------------------------------------------------- +_CONFORMANCE_HINT = ( + "LLM no longer satisfies EngineClient: {member}; keep " + "tensorrt_llm/llmapi/engine_client.py and LLM in step " + "(API review committee approves both)." +) + + +def test_llm_satisfies_engine_client(): + """Every public EngineClient member exists on LLM with the same kind.""" + for name, member in vars(EngineClient).items(): + if name.startswith("_"): + continue + llm_member = inspect.getattr_static(LLM, name, None) + assert llm_member is not None, _CONFORMANCE_HINT.format(member=name) + if isinstance(member, property): + assert isinstance(llm_member, property), _CONFORMANCE_HINT.format(member=name) + else: + assert callable(llm_member), _CONFORMANCE_HINT.format(member=name) + + +def test_llm_generate_async_accepts_every_engine_client_parameter(): + """Each protocol parameter exists on LLM.generate_async with the same default. + + A protocol default of `...` is a placeholder for an implementation-defined + value and is not compared. + """ + protocol_params = inspect.signature(EngineClient.generate_async).parameters + llm_params = inspect.signature(LLM.generate_async).parameters + for name, parameter in protocol_params.items(): + assert name in llm_params, _CONFORMANCE_HINT.format(member=f"generate_async({name})") + if parameter.default in (inspect.Parameter.empty, Ellipsis): + continue + assert llm_params[name].default == parameter.default, _CONFORMANCE_HINT.format( + member=f"generate_async({name}) default" + ) diff --git a/tests/unittest/llmapi/test_sampling_params.py b/tests/unittest/llmapi/test_sampling_params.py index 4a10387338d2..43c428f8b859 100644 --- a/tests/unittest/llmapi/test_sampling_params.py +++ b/tests/unittest/llmapi/test_sampling_params.py @@ -171,6 +171,36 @@ def test_explicit_serve_request_overrides_generation_config(request_obj): assert prepared.temperature == 0.2 +def test_set_request_provided_fields_keeps_only_sampler_fields(): + params = SamplingParams(end_id=1) + params.temperature = 0.2 + params.set_request_provided_fields(["temperature", "not_a_field"]) + assert params._request_provided_fields == frozenset({"temperature"}) + + # A recorded field stays request-provided through model defaulting, exactly + # as the private recording did; unrecorded fields take the model default. + prepared = _apply_generation_config_sampling_defaults( + "auto", params, {"temperature": 0.7, "top_p": 0.9} + ) + assert prepared.temperature == 0.2 + assert prepared.top_p == 0.9 + + +def test_stop_word_ids_filled_in_by_request_setup(): + params = SamplingParams(end_id=1, stop="halt") + assert params.stop_word_ids is None + + class _Tokenizer: + eos_token_id = 1 + pad_token_id = 0 + + def encode(self, text, add_special_tokens=False): + return [10, 11] + + params._setup(_Tokenizer(), hf_model_config=None, generation_config=None) + assert params.stop_word_ids == [[10, 11]] + + def test_absent_generation_config_value_preserves_existing_defaults(): direct_params = _apply_generation_config_sampling_defaults("auto", SamplingParams(end_id=1), {}) serve_params = _apply_generation_config_sampling_defaults( diff --git a/tests/unittest/llmapi/test_serve_frontend_affinity.py b/tests/unittest/llmapi/test_serve_frontend_affinity.py index 90569b6b8750..9442ba553f94 100644 --- a/tests/unittest/llmapi/test_serve_frontend_affinity.py +++ b/tests/unittest/llmapi/test_serve_frontend_affinity.py @@ -93,8 +93,8 @@ def test_launch_server_serves_the_affinity_port(monkeypatch, tmp_path): monkeypatch.setenv("TRTLLM_RESPONSES_API_DISABLE_STORE", "1") monkeypatch.setenv("PROMETHEUS_MULTIPROC_DIR", str(tmp_path)) llm = Mock() - # The launcher binds its forwarding socket in the executor's ipc dir. - llm._executor._multi_frontend_ipc_dir = str(tmp_path) + # The launcher binds its forwarding socket in the multi-frontend ipc dir. + llm.multi_frontend_ipc_dir = str(tmp_path) monkeypatch.setattr(serve, "PyTorchLLM", Mock(return_value=llm)) handed_off = {} diff --git a/tests/unittest/llmapi/test_serve_multi_frontend_helpers.py b/tests/unittest/llmapi/test_serve_multi_frontend_helpers.py index 0a6f3b9055bb..feb83aaf8e04 100644 --- a/tests/unittest/llmapi/test_serve_multi_frontend_helpers.py +++ b/tests/unittest/llmapi/test_serve_multi_frontend_helpers.py @@ -103,7 +103,7 @@ async def lifespan(_: FastAPI): server = object.__new__(OpenAIServer) server.app = FastAPI(lifespan=lifespan) server.generator = SimpleNamespace( - _executor=SimpleNamespace(resource_governor_queue=None), + resource_governor_queue=None, args=SimpleNamespace(return_perf_metrics=False), ) server.use_harmony = False @@ -241,12 +241,25 @@ def test_launcher_logs_an_attached_frontend_that_dies(monkeypatch) -> None: assert "frontend 1 " in errors[0] and "signal 9" in errors[0] and "2 of 3" in errors[0] -@pytest.mark.parametrize("doing_shutdown", [False, True]) -def test_on_launcher_lost_fails_fast_then_shuts_down(monkeypatch, doing_shutdown: bool) -> None: - fatal_errors = [] - executor = SimpleNamespace(doing_shutdown=doing_shutdown, _set_fatal_error=fatal_errors.append) +class _FatalErrorRecorder: + """Generator double for the public fatal-error recording surface.""" + + def __init__(self, shutdown_started: bool, has_executor: bool = True) -> None: + self.shutdown_started = shutdown_started + self._has_executor = has_executor + self.fatal_errors = [] + + def record_fatal_error(self, error: BaseException) -> bool: + if not self._has_executor: + return False + self.fatal_errors.append(error) + return True + + +@pytest.mark.parametrize("shutdown_started", [False, True]) +def test_on_launcher_lost_fails_fast_then_shuts_down(monkeypatch, shutdown_started: bool) -> None: server = object.__new__(OpenAIServer) - server.generator = SimpleNamespace(_executor=executor) + server.generator = _FatalErrorRecorder(shutdown_started) signals, terminations = [], [] monkeypatch.setattr(openai_server.signal, "raise_signal", signals.append) monkeypatch.setattr(openai_server, "_record_generator_termination", terminations.append) @@ -254,10 +267,25 @@ def test_on_launcher_lost_fails_fast_then_shuts_down(monkeypatch, doing_shutdown error = RuntimeError("attached frontend lost its engine: test") server._on_launcher_lost(error) - assert fatal_errors == [error] + assert server.generator.fatal_errors == [error] # A shutdown already under way gets no second SIGINT (which would force-quit it). - assert signals == ([] if doing_shutdown else [signal.SIGINT]) - assert terminations == ([] if doing_shutdown else [server.generator]) + assert signals == ([] if shutdown_started else [signal.SIGINT]) + assert terminations == ([] if shutdown_started else [server.generator]) + + +def test_on_launcher_lost_is_a_no_op_after_llm_shutdown(monkeypatch) -> None: + """record_fatal_error returning False means the LLM already shut down.""" + server = object.__new__(OpenAIServer) + server.generator = _FatalErrorRecorder(shutdown_started=False, has_executor=False) + signals, terminations = [], [] + monkeypatch.setattr(openai_server.signal, "raise_signal", signals.append) + monkeypatch.setattr(openai_server, "_record_generator_termination", terminations.append) + + server._on_launcher_lost(RuntimeError("attached frontend lost its engine: test")) + + assert server.generator.fatal_errors == [] + assert signals == [] + assert terminations == [] def test_shared_clock_offset_publish_and_poll(tmp_path) -> None: 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..79aa3e774a06 100644 --- a/tests/unittest/llmapi/test_serve_multi_frontend_route_gate.py +++ b/tests/unittest/llmapi/test_serve_multi_frontend_route_gate.py @@ -65,7 +65,7 @@ def _llm_server_routes() -> list[tuple[str, str]]: server = object.__new__(OpenAIServer) server.app = FastAPI() server.generator = SimpleNamespace( - _executor=SimpleNamespace(resource_governor_queue=None), + resource_governor_queue=None, args=SimpleNamespace(return_perf_metrics=False), ) server.use_harmony = False diff --git a/tests/unittest/serve/private_access_baseline.yaml b/tests/unittest/serve/private_access_baseline.yaml index 123c9bdbd182..73fc557b69c0 100644 --- a/tests/unittest/serve/private_access_baseline.yaml +++ b/tests/unittest/serve/private_access_baseline.yaml @@ -8,29 +8,10 @@ # `owner` and `replacement` are for readers; the test uses path, scope, member, count. version: 1 entries: -- {path: tensorrt_llm/commands/serve.py, scope: _bind_launcher_uds, member: _executor, count: 1, owner: llmapi/llm.py BaseLLM, replacement: a public accessor for the multi-frontend IPC directory} -- {path: tensorrt_llm/commands/serve.py, scope: _bind_launcher_uds, member: _multi_frontend_ipc_dir, count: 1, owner: executor/proxy.py GenerationExecutorProxy, replacement: a public accessor for the multi-frontend IPC directory} -- {path: tensorrt_llm/commands/serve.py, scope: _run_fleet_worker_impl, member: _config, count: 2, owner: serve/openai_disagg_server.py OpenAIDisaggServer, replacement: OpenAIDisaggServer.address} -- {path: tensorrt_llm/grpc/openengine/capabilities.py, scope: _loaded_model_type, member: _hf_model_config, count: 1, owner: llmapi/llm.py BaseLLM, replacement: BaseLLM.hf_model_config} -- {path: tensorrt_llm/grpc/openengine/capabilities.py, scope: resolve_media_capabilities, member: _registered_model_type, count: 1, owner: inputs input-processor registry, replacement: a public registry accessor for the registered model type} -- {path: tensorrt_llm/grpc/openengine/control.py, scope: OpenEngineControlServicer.GetServerInfo, member: _executor, count: 1, owner: llmapi/llm.py BaseLLM, replacement: "BaseLLM.get_kv_cache_capacity(), BaseLLM.get_openengine_discovery()"} - {path: tensorrt_llm/grpc/openengine/control.py, scope: OpenEngineControlServicer.GetServerInfo, member: _openengine_discovery, count: 1, owner: llmapi/llm_args.py BaseLlmArgs, replacement: BaseLlmArgs.openengine_discovery} - {path: tensorrt_llm/grpc/openengine/control.py, scope: OpenEngineControlServicer.__init__, member: _enable_routing_load, count: 1, owner: llmapi/llm_args.py BaseLlmArgs, replacement: BaseLlmArgs.enable_routing_load} -- {path: tensorrt_llm/grpc/openengine/control.py, scope: OpenEngineControlServicer._get_kv_cache_load, member: _executor, count: 1, owner: llmapi/llm.py BaseLLM, replacement: BaseLLM.get_kv_cache_load()} -- {path: tensorrt_llm/grpc/openengine/control.py, scope: OpenEngineControlServicer._requested_sources, member: _executor, count: 1, owner: llmapi/llm.py BaseLLM, replacement: BaseLLM.get_openengine_discovery()} -- {path: tensorrt_llm/grpc/openengine/formatting.py, scope: _token_holdback, member: _stop_word_ids, count: 1, owner: sampling_params.py SamplingParams, replacement: SamplingParams.stop_word_ids} -- {path: tensorrt_llm/grpc/openengine/node.py, scope: initialize_node_discovery, member: _openengine_discovery, count: 2, owner: "llmapi/llm_args.py BaseLlmArgs (read); _torch/pyexecutor/py_executor.py PyExecutor (write)", replacement: "BaseLlmArgs.openengine_discovery; PyExecutor.set_openengine_discovery()"} -- {path: tensorrt_llm/grpc/openengine/node.py, scope: initialize_node_discovery, member: _openengine_node_server, count: 1, owner: _torch/pyexecutor/py_executor.py PyExecutor, replacement: PyExecutor.set_openengine_node_server()} -- {path: tensorrt_llm/grpc/openengine/request_mapping.py, scope: sampling_params_from_request, member: _set_request_provided_fields, count: 1, owner: sampling_params.py SamplingParams, replacement: SamplingParams.set_request_provided_fields()} -- {path: tensorrt_llm/grpc/openengine/server.py, scope: OpenEngineServer.start, member: _executor, count: 1, owner: llmapi/llm.py BaseLLM, replacement: BaseLLM.get_kv_cache_capacity()} -- {path: tensorrt_llm/serve/_frontend_processes.py, scope: _spawn_attached_frontends, member: _executor, count: 1, owner: llmapi/llm.py BaseLLM, replacement: BaseLLM.multi_frontend_attach_info()} -- {path: tensorrt_llm/serve/chat_utils.py, scope: parse_chat_message_content_part, member: _multimodal_server_config, count: 4, owner: inputs/utils.py MultimodalDataTracker, replacement: MultimodalDataTracker.multimodal_server_config} -- {path: tensorrt_llm/serve/openai_protocol.py, scope: _record_sampling_params_request_fields, member: _set_request_provided_fields, count: 1, owner: sampling_params.py SamplingParams, replacement: SamplingParams.set_request_provided_fields()} +- {path: tensorrt_llm/grpc/openengine/node.py, scope: initialize_node_discovery, member: _openengine_discovery, count: 1, owner: llmapi/llm_args.py BaseLlmArgs, replacement: BaseLlmArgs.openengine_discovery} - {path: tensorrt_llm/serve/openai_server.py, scope: OpenAIServer._create_chat_response, member: _postprocess_result, count: 1, owner: executor/result.py CompletionOutput, replacement: CompletionOutput.postprocess_result} -- {path: tensorrt_llm/serve/openai_server.py, scope: OpenAIServer._init_embedding_batcher, member: _encoder_executor, count: 1, owner: llmapi/llm.py BaseLLM, replacement: BaseLLM.encode_limits} -- {path: tensorrt_llm/serve/openai_server.py, scope: OpenAIServer._init_llm, member: _hf_model_dir, count: 1, owner: llmapi/llm.py BaseLLM, replacement: BaseLLM.hf_model_dir} -- {path: tensorrt_llm/serve/openai_server.py, scope: OpenAIServer._on_launcher_lost, member: _executor, count: 1, owner: llmapi/llm.py BaseLLM, replacement: public fatal-error recording on the LLM} -- {path: tensorrt_llm/serve/openai_server.py, scope: OpenAIServer._on_launcher_lost, member: _set_fatal_error, count: 1, owner: executor/executor.py GenerationExecutor, replacement: public fatal-error recording on the LLM} - {path: tensorrt_llm/serve/openai_server.py, scope: OpenAIServer.chat_harmony.create_streaming_generator, member: _postprocess_result, count: 2, owner: executor/result.py CompletionOutput, replacement: CompletionOutput.postprocess_result} - {path: tensorrt_llm/serve/openai_server.py, scope: OpenAIServer.openai_chat.chat_stream_generator, member: _postprocess_result, count: 2, owner: executor/result.py CompletionOutput, replacement: CompletionOutput.postprocess_result} - {path: tensorrt_llm/serve/openai_server.py, scope: OpenAIServer.openai_completion.completion_generator, member: _postprocess_result, count: 1, owner: executor/result.py CompletionOutput, replacement: CompletionOutput.postprocess_result} @@ -38,16 +19,9 @@ entries: - {path: tensorrt_llm/serve/openai_server.py, scope: OpenAIServer.openai_image_generation, member: _save_bytes, count: 1, owner: visual_gen/output.py, replacement: public VisualGen output encoder} - {path: tensorrt_llm/serve/openai_server.py, scope: OpenAIServer.openai_responses.create_response, member: _postprocess_result, count: 1, owner: executor/result.py CompletionOutput, replacement: CompletionOutput.postprocess_result} - {path: tensorrt_llm/serve/openai_server.py, scope: OpenAIServer.openai_responses.create_streaming_generator, member: _postprocess_result, count: 1, owner: executor/result.py CompletionOutput, replacement: CompletionOutput.postprocess_result} -- {path: tensorrt_llm/serve/openai_server.py, scope: OpenAIServer.register_routes, member: _executor, count: 1, owner: llmapi/llm.py BaseLLM, replacement: BaseLLM.resource_governor_queue} -- {path: tensorrt_llm/serve/openai_server.py, scope: _record_generator_termination, member: _executor, count: 1, owner: llmapi/llm.py BaseLLM, replacement: BaseLLM.engine_in_worker_process} - {path: tensorrt_llm/serve/postprocess_handlers.py, scope: chat_harmony_streaming_post_processor, member: _done, count: 1, owner: executor/result.py GenerationResultBase, replacement: GenerationResultBase.finished} -- {path: tensorrt_llm/serve/postprocess_handlers.py, scope: chat_response_post_processor, member: _spec_dec_counters, count: 1, owner: executor/result.py CompletionOutput, replacement: CompletionOutput.spec_dec_counters} - {path: tensorrt_llm/serve/postprocess_handlers.py, scope: chat_stream_post_processor, member: _done, count: 1, owner: executor/result.py GenerationResultBase, replacement: GenerationResultBase.finished} -- {path: tensorrt_llm/serve/postprocess_handlers.py, scope: chat_stream_post_processor, member: _spec_dec_counters, count: 1, owner: executor/result.py CompletionOutput, replacement: CompletionOutput.spec_dec_counters} -- {path: tensorrt_llm/serve/postprocess_handlers.py, scope: completion_response_post_processor, member: _spec_dec_counters, count: 1, owner: executor/result.py CompletionOutput, replacement: CompletionOutput.spec_dec_counters} - {path: tensorrt_llm/serve/postprocess_handlers.py, scope: completion_stream_post_processor, member: _done, count: 1, owner: executor/result.py GenerationResultBase, replacement: GenerationResultBase.finished} -- {path: tensorrt_llm/serve/postprocess_handlers.py, scope: completion_stream_post_processor, member: _last_text_len, count: 1, owner: executor/result.py CompletionOutput, replacement: CompletionOutput.text_diff_offset} -- {path: tensorrt_llm/serve/postprocess_handlers.py, scope: completion_stream_post_processor, member: _spec_dec_counters, count: 1, owner: executor/result.py CompletionOutput, replacement: CompletionOutput.spec_dec_counters} - {path: tensorrt_llm/serve/postprocess_handlers.py, scope: responses_api_streaming_post_processor, member: _done, count: 1, owner: executor/result.py GenerationResultBase, replacement: GenerationResultBase.finished} - {path: tensorrt_llm/serve/responses_utils.py, scope: ResponsesStreamingProcessor.process_single_output, member: _done, count: 1, owner: executor/result.py GenerationResultBase, replacement: GenerationResultBase.finished} - {path: tensorrt_llm/serve/router_utils.py, scope: BlockHashMixin._compute_block_hashes, member: _hash_cache_salt, count: 1, owner: runtime/kv_cache_manager_v2 KVCacheEventManager (nanobind), replacement: public hash_cache_salt binding} diff --git a/tests/unittest/serve/test_private_access_lint.py b/tests/unittest/serve/test_private_access_lint.py index 9bea0b46de81..26603599f476 100644 --- a/tests/unittest/serve/test_private_access_lint.py +++ b/tests/unittest/serve/test_private_access_lint.py @@ -37,7 +37,7 @@ OPENENGINE_SCOPE = "tensorrt_llm/grpc/openengine/" # Lower this when you remove accesses; never raise it. -MAX_TOTAL = 50 +MAX_TOTAL = 19 # (path, member) -> (verdict, reason). Verdicts: "cross", "same-package", "third-party". # Needed only where the name index cannot tell: the name is defined both inside and diff --git a/tests/unittest/serve/test_serve_runtime_api_calls.py b/tests/unittest/serve/test_serve_runtime_api_calls.py new file mode 100644 index 000000000000..a6cd1c05f2a1 --- /dev/null +++ b/tests/unittest/serve/test_serve_runtime_api_calls.py @@ -0,0 +1,179 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The HTTP server reads the engine through public LLM members only. + +Each test calls a production ``OpenAIServer`` method on a stand-in object (not +``object.__new__(OpenAIServer)``) whose generator double defines exactly the +public members the method reads, so a server method reaching for a private +engine member again fails loudly here. +""" + +from types import SimpleNamespace + +import pytest +from fastapi import FastAPI + +from tensorrt_llm.llmapi.llm import EncodeLimits +from tensorrt_llm.serve import openai_server +from tensorrt_llm.serve.openai_server import OpenAIServer, _record_generator_termination + +pytestmark = pytest.mark.cpu_only + + +class _ServerStub(SimpleNamespace): + """A SimpleNamespace whose unknown attributes are inert handlers. + + ``register_routes`` touches about thirty handler attributes; the routes + under test here are only the resource-governor ones, so every other + handler read resolves to a placeholder callable FastAPI can register. + """ + + def __getattr__(self, name): + def _handler(*args, **kwargs): + return None + + return _handler + + +def _generator(**kwargs): + kwargs.setdefault("args", SimpleNamespace(return_perf_metrics=False)) + return SimpleNamespace(**kwargs) + + +# --------------------------------------------------------------------------- +# _init_embedding_batcher reads generator.encode_limits +# --------------------------------------------------------------------------- +def test_embedding_batcher_takes_the_limits_encode_checks(): + stub = _ServerStub( + generator=_generator( + encode_limits=EncodeLimits(max_batch_size=4, max_num_tokens=512, max_seq_len=128), + encode=lambda batch: batch, + ), + _embedding_max_queue_delay=0.01, + _embedding_max_queue_size=8, + ) + + OpenAIServer._init_embedding_batcher(stub) + + batcher = stub.embedding_batcher + assert batcher._max_batch_size == 4 + assert batcher._max_num_tokens == 512 + assert batcher._max_seq_len == 128 + + +# --------------------------------------------------------------------------- +# register_routes reads generator.resource_governor_queue +# --------------------------------------------------------------------------- +class _RecordingGovernor: + def __init__(self, **kwargs): + self.kwargs = kwargs + self.registered_app = None + + def register_routes(self, app): + self.registered_app = app + + +def _register_routes_stub(resource_governor_queue): + return _ServerStub( + app=FastAPI(), + generator=_generator(resource_governor_queue=resource_governor_queue), + tokenizer=object(), + model_config=None, + processor=None, + allow_request_chat_template=False, + use_harmony=False, + chat_template=None, + ) + + +def test_register_routes_builds_the_governor_from_the_public_queue(monkeypatch): + queue = object() + stub = _register_routes_stub(resource_governor_queue=queue) + monkeypatch.setattr("tensorrt_llm.serve.resource_governor.ResourceGovernor", _RecordingGovernor) + + OpenAIServer.register_routes(stub) + + assert isinstance(stub.resource_governor, _RecordingGovernor) + assert stub.resource_governor.kwargs["resource_governor_queue"] is queue + assert stub.resource_governor.registered_app is stub.app + + +def test_register_routes_skips_the_governor_without_a_queue(monkeypatch): + stub = _register_routes_stub(resource_governor_queue=None) + monkeypatch.setattr("tensorrt_llm.serve.resource_governor.ResourceGovernor", _RecordingGovernor) + + OpenAIServer.register_routes(stub) + + assert stub.resource_governor is None + assert not any(route.path.startswith("/_resource_governor") for route in stub.app.routes) + + +# --------------------------------------------------------------------------- +# _init_llm reads generator.hf_model_dir +# --------------------------------------------------------------------------- +def test_init_llm_loads_processor_and_config_from_the_public_model_dir(monkeypatch): + class _Config: + model_type = "llama" + + processor_calls, config_calls = [], [] + monkeypatch.setattr( + openai_server.AutoProcessor, + "from_pretrained", + classmethod(lambda cls, path, **kwargs: processor_calls.append(path)), + ) + monkeypatch.setattr( + "tensorrt_llm._torch.pyexecutor.config_utils.load_pretrained_config", + lambda path, **kwargs: config_calls.append(path) or _Config(), + ) + stub = _ServerStub( + generator=_generator( + tokenizer=object(), + hf_model_dir="/models/resolved-dir", + args=SimpleNamespace( + trust_remote_code=False, + checkpoint_format=None, + post_processor_hook=None, + return_perf_metrics=False, + ), + ), + postproc_worker_enabled=False, + _launcher_forwarder=None, + ) + + OpenAIServer._init_llm(stub) + + assert processor_calls == ["/models/resolved-dir"] + assert config_calls == ["/models/resolved-dir"] + + +# --------------------------------------------------------------------------- +# _record_generator_termination reads generator.engine_in_worker_process +# --------------------------------------------------------------------------- +def _recorded_outcome(monkeypatch, generator): + outcomes = [] + monkeypatch.setattr(openai_server, "record_termination_observation", outcomes.append) + _record_generator_termination(generator) + assert len(outcomes) == 1 + return outcomes[0] + + +def test_termination_telemetry_blames_the_worker_behind_the_proxy(monkeypatch): + outcome = _recorded_outcome(monkeypatch, _generator(engine_in_worker_process=True)) + assert outcome.termination_kind == "worker_failure" + assert outcome.component == "engine_worker" + assert outcome.reporting_source == "executor_proxy" + + +def test_termination_telemetry_blames_the_llm_for_in_process_engines(monkeypatch): + outcome = _recorded_outcome(monkeypatch, _generator(engine_in_worker_process=False)) + assert outcome.termination_kind == "exception" + assert outcome.component == "llm" + assert outcome.reporting_source == "self" + + +def test_termination_telemetry_defaults_for_generators_without_the_member(monkeypatch): + """The telemetry guard must not raise for doubles lacking the property.""" + outcome = _recorded_outcome(monkeypatch, object()) + assert outcome.termination_kind == "exception" + assert outcome.component == "llm" + assert outcome.reporting_source == "self" From d4b6d1450ea9ecc19fb10a9e075c4cfa18e663bd Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Mon, 2 Nov 2026 14:50:39 +0000 Subject: [PATCH 04/30] [TRTLLM-17272][test] Call _on_launcher_lost on a stand-in instead of a bare OpenAIServer The shut-down-LLM case of the launcher-lost watchdog test now calls the production method unbound on a stand-in object whose generator double defines the public members the method reads, the form the file's other stand-in tests use. New tests must not build servers with object.__new__. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- .../unittest/llmapi/test_serve_multi_frontend_helpers.py | 8 +++++--- 1 file changed, 5 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 feb83aaf8e04..6d79502841cb 100644 --- a/tests/unittest/llmapi/test_serve_multi_frontend_helpers.py +++ b/tests/unittest/llmapi/test_serve_multi_frontend_helpers.py @@ -275,13 +275,15 @@ def test_on_launcher_lost_fails_fast_then_shuts_down(monkeypatch, shutdown_start def test_on_launcher_lost_is_a_no_op_after_llm_shutdown(monkeypatch) -> None: """record_fatal_error returning False means the LLM already shut down.""" - server = object.__new__(OpenAIServer) - server.generator = _FatalErrorRecorder(shutdown_started=False, has_executor=False) + # The production method on a stand-in: _on_launcher_lost reads only self.generator. + server = SimpleNamespace( + generator=_FatalErrorRecorder(shutdown_started=False, has_executor=False) + ) signals, terminations = [], [] monkeypatch.setattr(openai_server.signal, "raise_signal", signals.append) monkeypatch.setattr(openai_server, "_record_generator_termination", terminations.append) - server._on_launcher_lost(RuntimeError("attached frontend lost its engine: test")) + OpenAIServer._on_launcher_lost(server, RuntimeError("attached frontend lost its engine: test")) assert server.generator.fatal_errors == [] assert signals == [] 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 05/30] [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 06/30] [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 07/30] [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 08/30] [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 09/30] [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 10/30] [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 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 11/30] [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 12/30] [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 13/30] [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 14/30] [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 15/30] [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 16/30] [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 5bcd0a151d9e9bd864bc2bca08b75edce177e236 Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Tue, 3 Nov 2026 00:53:51 +0000 Subject: [PATCH 17/30] [TRTLLM-17264][refactor] Move the VisualGen protocol models into serve/api/visual_gen Move the eight VisualGen request and response models (ImageGenerationRequest, ImageEditRequest, ImageObject, ImageGenerationResponse, MediaReferenceItem, VideoGenerationRequest, VideoJob, VideoJobList) out of tensorrt_llm/serve/openai_protocol.py into the new package tensorrt_llm/serve/api/visual_gen/ (protocol.py), unchanged. openai_protocol keeps the eight names importable through a lazy module __getattr__: importing it no longer loads the VisualGen models or tensorrt_llm.visual_gen, and the re-export is part of its own interface, so it never warns and has no removal release. The imports only the moved block used (UploadFile, AliasChoices, MediaRole) leave openai_protocol with it. Every importer in the repository now imports the new module; the benchmarking guide links the new file. protocol.py joins legacy-files.txt (its code comes from a legacy-formatted file) and the derived lint configs are regenerated. No behavior change. OpenAPI component names may change with the move; they are not a contract. 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/visual_gen/__init__.py | 18 + tensorrt_llm/serve/api/visual_gen/protocol.py | 441 +++++++++++++++++ tensorrt_llm/serve/openai_protocol.py | 452 ++---------------- tensorrt_llm/serve/openai_server.py | 15 +- tensorrt_llm/serve/openai_video_routes.py | 6 +- .../serve/scripts/BENCHMARKING_VISUAL_GEN.md | 2 +- .../serve/scripts/benchmark_visual_gen.py | 2 +- tensorrt_llm/serve/visual_gen_utils.py | 2 +- .../visual_gen/test_trtllm_serve_endpoints.py | 2 +- .../visual_gen/test_visual_gen_params.py | 4 +- .../visual_gen/test_visual_gen_utils.py | 11 +- 15 files changed, 523 insertions(+), 437 deletions(-) create mode 100644 tensorrt_llm/serve/api/visual_gen/__init__.py create mode 100644 tensorrt_llm/serve/api/visual_gen/protocol.py diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 4aa5b0f82bcd..23df44e55e96 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/api/visual_gen/protocol.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/api/visual_gen/protocol.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..72f1c66e0acd 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/api/visual_gen/protocol.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..a65da3ab8465 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/api/visual_gen/protocol.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 1bd2bbfab0e5..6485e090c26b 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/api/visual_gen/protocol.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/visual_gen/__init__.py b/tensorrt_llm/serve/api/visual_gen/__init__.py new file mode 100644 index 000000000000..4d3a40bee9b0 --- /dev/null +++ b/tensorrt_llm/serve/api/visual_gen/__init__.py @@ -0,0 +1,18 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""VisualGen HTTP serving: request and response models, handlers and routers. + +Kept empty on purpose: importing one submodule must not load the others. +""" diff --git a/tensorrt_llm/serve/api/visual_gen/protocol.py b/tensorrt_llm/serve/api/visual_gen/protocol.py new file mode 100644 index 000000000000..bcdabe5d08ee --- /dev/null +++ b/tensorrt_llm/serve/api/visual_gen/protocol.py @@ -0,0 +1,441 @@ +# 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. + +# Adapted from +# https://github.com/vllm-project/vllm/blob/4db5176d9758b720b05460c50ace3c01026eb158/vllm/entrypoints/openai/protocol.py +"""Request and response models of the VisualGen endpoints. + +Image generation, image edit and video. ``tensorrt_llm.serve.openai_protocol`` +re-exports them lazily under their old names. +""" + +import time +from typing import Any, Dict, List, Literal, Optional, Union + +from fastapi import UploadFile +from pydantic import AliasChoices, Field, field_validator, model_validator + +from tensorrt_llm.serve.openai_protocol import OpenAIBaseModel +from tensorrt_llm.visual_gen.params import MediaRole + + +class ImageGenerationRequest(OpenAIBaseModel): + """OpenAI-compatible image generation request. + + Universal per-request fields map 1:1 to :class:`VisualGenParams`. + Model-specific knobs (``stg_scale``, ``guidance_rescale``, …) + travel through ``extra_params``; the executor validates each + key against the loaded pipeline's + ``extra_param_specs``. Unknown top-level fields are rejected + with HTTP 422 via the inherited ``extra="forbid"`` policy. + """ + + # Prompt + transport (OpenAI-standard, always honored) + prompt: str + response_format: Literal["url", "b64_json", "path"] = "url" + format: Literal["png", "webp", "jpeg", "safetensors", "pt"] = Field( + default="png", + description=( + "Generation content encoding format. Image encoders write " + "``png``/``webp``/``jpeg``; tensor encoders write " + "``safetensors``/``pt`` for programmatic post-processing."), + ) + seed: Optional[int] = Field(default=None, + ge=0, + description="Random seed for reproducibility.") + + # Resolution. ``size`` is OpenAI-shaped; ``width`` + ``height`` are an + # equivalent structured alternative. Exactly one of width/height is + # rejected by the paired validator below. Numeric fields use + # ``gt=0`` as a safety net so zero / negative inputs are rejected + # with HTTP 422 before reaching the pipeline. + size: Optional[str] = Field(default=None, pattern=r"^(\d+x\d+|auto)$") + width: Optional[int] = Field(default=None, gt=0) + height: Optional[int] = Field(default=None, gt=0) + + # TRT-LLM-supported per-request params (1:1 with VisualGenParams fields) + num_inference_steps: Optional[int] = Field(default=None, gt=0) + guidance_scale: Optional[float] = Field(default=None, gt=0) + max_sequence_length: Optional[int] = Field(default=None, gt=0) + negative_prompt: Optional[str] = None + n: Optional[int] = Field( + default=None, + gt=0, + le=10, + description=("Number of images to generate. Capped at 10 to match the " + "OpenAI images API and to bound GPU memory / disk usage."), + ) + + # Model-specific overflow + extra_params: Optional[Dict[str, Any]] = Field( + default=None, + description=( + "Model-specific parameters forwarded to the underlying pipeline. " + "See per-model docs for accepted keys."), + ) + + # Accepted-but-ignored OpenAI-shaped fields. The conversion no-ops; the + # server logs WARNING when a client sets ``quality`` or ``style``, and + # WARNING-on-mismatch for ``model``. Kept in the schema so OpenAI-SDK + # clients don't trip ``extra="forbid"``. + model: Optional[str] = None + quality: Optional[Literal["standard", "hd"]] = None + style: Optional[Literal["vivid", "natural"]] = None + user: Optional[str] = None + + @model_validator(mode="after") + def _check_paired_dimensions(self): + """Reject sending exactly one of ``width`` / ``height``. + + Either both are sent (structured resolution wins over ``size``) + or neither is sent (``size`` or pipeline default applies). + """ + if (self.width is None) != (self.height is None): + raise ValueError( + "width and height must be sent together; got width=" + f"{self.width!r}, height={self.height!r}") + return self + + +class ImageEditRequest(OpenAIBaseModel): + """OpenAI-compatible image editing request. + + The server accepts the OpenAI multipart shape and a JSON/base64 + shape for tests and non-SDK clients. Model-specific knobs travel + through ``extra_params`` and are validated by the loaded visual + generation pipeline. + """ + + prompt: str + image: Union[str, UploadFile, List[Union[str, UploadFile]]] = Field( + description="Input image or images to edit.") + mask: Optional[Union[str, UploadFile]] = Field( + default=None, + description= + "Optional edit mask. Currently accepted for compatibility but unsupported.", + ) + response_format: Literal["url", "b64_json", "path"] = "url" + output_format: Literal["png", "webp", "jpeg"] = Field( + default="png", + validation_alias=AliasChoices("output_format", "format"), + description="Edited image content encoding format.", + ) + seed: Optional[int] = Field(default=None, + ge=0, + description="Random seed for reproducibility.") + + size: Optional[str] = Field(default=None, pattern=r"^(\d+x\d+|auto)$") + width: Optional[int] = Field(default=None, gt=0) + height: Optional[int] = Field(default=None, gt=0) + + num_inference_steps: Optional[int] = Field(default=None, gt=0) + guidance_scale: Optional[float] = Field(default=None, gt=0) + max_sequence_length: Optional[int] = Field(default=None, gt=0) + negative_prompt: Optional[str] = None + n: Optional[int] = Field( + default=None, + gt=0, + le=10, + description=("Number of edited images to generate. Capped at 10 to " + "match the OpenAI images API."), + ) + + extra_params: Optional[Dict[str, Any]] = Field( + default=None, + description=( + "Model-specific parameters forwarded to the underlying pipeline. " + "See per-model docs for accepted keys."), + ) + + model: Optional[str] = None + quality: Optional[Literal["standard", "hd"]] = None + user: Optional[str] = None + + @model_validator(mode="after") + def _check_paired_dimensions(self): + if isinstance(self.image, list) and not self.image: + raise ValueError("image must contain at least one input image") + if (self.width is None) != (self.height is None): + raise ValueError( + "width and height must be sent together; got width=" + f"{self.width!r}, height={self.height!r}") + return self + + +class ImageObject(OpenAIBaseModel): + """Generated image object in the response.""" + b64_json: Optional[str] = None + url: Optional[str] = None + path: Optional[str] = None + revised_prompt: Optional[str] = None + + +class ImageGenerationResponse(OpenAIBaseModel): + """Response from image generation endpoint. + + ``output_format`` reports the encoding actually applied to the + returned bytes / files so clients can decode or label the payload + correctly. Image encoders are ``"png"``/``"webp"``/``"jpeg"``; + tensor formats are ``"safetensors"``/``"pt"``. + """ + + created: int = Field(default_factory=lambda: int(time.time())) + data: List[ImageObject] + output_format: Literal["png", "webp", "jpeg", "safetensors", "pt"] = "png" + quality: Literal["low", "medium", "high"] = "medium" + size: Optional[str] = None + + +class MediaReferenceItem(OpenAIBaseModel): + """One media reference (image / video / audio) for conditioning (mirrors ``MediaRef``). + + ``format`` declares how to read ``content``; it is required, so no wire form + is ever guessed. The request field it sits in (``image_reference`` / + ``video_reference`` / ``audio_reference``) fixes the modality. ``role`` is + required only when it is ambiguous — a model with more than one required role + for that modality; a single role, or a single required role (e.g. i2v + first_frame), is inferred. + """ + + content: str = Field( + description="The reference payload, in the form declared by ``format``." + ) + format: Literal["path", "url", "base64"] = Field(description=( + "Wire form of ``content``: ``path`` (a file readable by the server; a " + "``file://`` URI is also accepted), ``url`` (``http(s)``, fetched " + "through the SSRF-guarded loader), or ``base64`` (a ``data:`` URI is " + "also accepted). Raw bytes reach the server as a multipart upload, " + "which carries no ``format`` of its own. Distinct from the top-level " + "``format``, which selects the *output* encoding.")) + role: Optional[MediaRole] = Field( + default=None, + description="Which conditioning slot this reference fills. Required only " + "when the target model accepts this modality in more than one slot; omit " + "it when the model leaves no ambiguity.", + ) + + +class VideoGenerationRequest(OpenAIBaseModel): + """Video generation request (extended API). + + Universal per-request fields map 1:1 to :class:`VisualGenParams`. + Model-specific knobs travel through ``extra_params``. Unknown + top-level fields are rejected with HTTP 422 via the inherited + ``extra="forbid"`` policy. + """ + + # Prompt + transport + prompt: str + response_format: Literal["file", "path"] = "file" + format: Literal["mp4", "avi", "auto", "safetensors", "pt"] = Field( + default="auto", + description=( + "Generation content encoding format. Video encoders write " + "``mp4``/``avi``/``auto``; tensor encoders write " + "``safetensors``/``pt`` and carry video, audio, and scalar " + "metadata (frame rate, audio sample rate) in one payload."), + ) + seed: Optional[int] = Field(default=None, + ge=0, + description="Random seed for reproducibility.") + image_reference: Optional[Union[ + UploadFile, MediaReferenceItem, List[MediaReferenceItem]]] = Field( + default=None, + description= + ("Image reference(s) conditioning generation (e.g. image-to-video " + "first frame). Send a ``{content, format, role}`` object or a list " + "of them, where ``format`` is ``path`` / ``url`` / ``base64``; or " + "upload a single image file via multipart, whose form is implied. " + "PNG or JPEG only — HEIF/AVIF are not supported."), + ) + video_reference: Optional[Union[ + UploadFile, MediaReferenceItem, List[MediaReferenceItem]]] = Field( + default=None, + description= + ("Video reference(s) conditioning generation (video-to-video). Send " + "a ``{content, format}`` object or a list of them, where ``format`` " + "is ``path`` / ``url`` / ``base64``; or upload a single video file " + "via multipart, whose form is implied. MP4 or AVI, with H.264 the " + "tested codec and others best-effort."), + ) + audio_reference: Optional[Union[ + UploadFile, MediaReferenceItem, List[MediaReferenceItem]]] = Field( + default=None, + description= + ("Audio reference(s) conditioning generation. Send a " + "``{content, format}`` object or a list of them, where ``format`` " + "is ``path`` / ``url`` / ``base64``; or upload a single audio " + "file via multipart, whose form is implied. Accepted only by " + "models that declare an audio reference slot."), + ) + input_reference: Optional[Union[str, UploadFile]] = Field( + default=None, + description= + ("Deprecated. A single image or video reference, routed by content " + "signature to image-to-video or video-to-video. A JSON request carries " + "base64 bytes; a multipart request uploads the file. Kept for backward " + "compatibility; prefer the typed ``image_reference`` / " + "``video_reference`` fields, which take precedence — this field is " + "ignored whenever a typed ``image_reference`` or ``video_reference`` " + "is provided."), + ) + + # Resolution + size: Optional[str] = Field(default=None, pattern=r"^(\d+x\d+|auto)$") + width: Optional[int] = Field(default=None, gt=0) + height: Optional[int] = Field(default=None, gt=0) + + # Frame budget. ``num_frames`` is preferred; if absent the engine + # derives it from ``seconds * frame_rate``. ``frame_rate`` is the + # canonical name (matches the Python field); ``fps`` is an alias for + # OpenAI-shape clients via ``populate_by_name=True``. + # All three constrain to strictly positive values so a zero + # ``frame_rate`` (division-by-zero in the AVI fallback) or a + # negative ``num_frames`` are rejected with HTTP 422 before + # reaching the encoder. + # Upper bounds keep request-boundary protection against requests + # that can exhaust GPU memory or pin the server on unbounded work. + # The numbers are generous (a minute of video at 120 fps) so common + # workloads pass; clients that need larger budgets can lift the cap + # at deployment time. + num_frames: Optional[int] = Field(default=None, gt=0, le=7200) + seconds: Optional[float] = Field(default=None, gt=0, le=60.0) + frame_rate: Optional[float] = Field(default=None, + alias="fps", + gt=0, + le=120.0) + + # TRT-LLM-supported per-request params (1:1 with VisualGenParams) + num_inference_steps: Optional[int] = Field(default=None, gt=0) + guidance_scale: Optional[float] = Field(default=None, gt=0) + max_sequence_length: Optional[int] = Field(default=None, gt=0) + negative_prompt: Optional[str] = None + + # Model-specific overflow + extra_params: Optional[Dict[str, Any]] = Field( + default=None, + description=( + "Model-specific parameters forwarded to the underlying pipeline. " + "See per-model docs for accepted keys."), + ) + + # Accepted-but-ignored OpenAI-shaped field + model: Optional[str] = None + + @model_validator(mode="after") + def _check_paired_dimensions(self): + """Reject sending exactly one of ``width`` / ``height``.""" + if (self.width is None) != (self.height is None): + raise ValueError( + "width and height must be sent together; got width=" + f"{self.width!r}, height={self.height!r}") + return self + + @field_validator("response_format", mode="before") + @classmethod + def _reject_removed_response_format(cls, value): + """Give migrating callers an actionable error for removed values. + + ``url``/``b64_json`` were valid before the transport rewrite; run + before the ``Literal`` check so the error names the replacement + instead of the generic "Input should be 'file' or 'path'". + """ + removed = { + "url": + ("'url' was removed for video; use 'file' (raw bytes -- the " + "old 'url' behavior, renamed) or 'path' (server-side path)."), + "b64_json": ("'b64_json' was removed for video; use 'file' (raw " + "bytes) or 'path' (server-side path)."), + } + if isinstance(value, str) and value in removed: + raise ValueError(removed[value]) + return value + + +class VideoJob(OpenAIBaseModel): + """Metadata for an asynchronous video generation job. + + Follows the OpenAI Videos API specification: + https://platform.openai.com/docs/api-reference/videos + """ + completed_at: Optional[int] = Field( + default=None, description="Unix timestamp of completion") + created_at: int = Field(description="Unix timestamp of creation") + error: Optional[str] = Field(default=None, + description="Error message if failed") + expires_at: Optional[int] = Field( + default=None, description="Unix timestamp of expiration") + id: str = Field(description="Unique identifier for the video") + model: str = Field(description="The model used for generation") + object: str = Field(default="video", description="Object type") + progress: Optional[int] = Field( + default=None, + description="Progress of the video generation job (0-100)") + prompt: str = Field(description="The prompt used to generate the video") + status: Literal["queued", "generating", "postprocessing", "completed", + "failed"] = Field(description=( + "Current status of the video generation job. " + "``generating`` (model inference) becomes " + "``postprocessing`` (encode and/or write the output " + "file) when inference finishes, then ``completed`` " + "once downloadable via ``/content``.")) + + # Video properties + duration: Optional[float] = Field(default=None, + description="Video duration in seconds") + fps: Optional[float] = Field( + default=None, + description=( + "Frames per second. Float to preserve cinematic rates such " + "as 23.976 / 29.97 that some encoders / pipelines use."), + ) + size: Optional[str] = Field(default=None, + description="Video dimensions in 'WxH' format") + # exclude=True: internal file-location for /content resolution + delete; + # never on the wire (the path payload is the hand-built {id, output_path} + # envelope in /content), so status/list model_dump() stays status-only. + output_path: Optional[str] = Field( + default=None, + exclude=True, + description="Server-side saved path (internal; excluded from the wire)." + ) + output_paths: Optional[List[str]] = Field( + default=None, + exclude=True, + description= + "Server-side paths for n>1 (internal; excluded from the wire).") + # exclude=True internal timings, never on the wire (status/list + # model_dump() stays status-only). ``request_started`` carries the + # steady-clock ``server_arrival_time``; the background task computes + # ``total`` from it and stores the header timings + # (``generation``/``denoise``/``total``) in ``timing_metrics`` so + # ``/content`` emits the same Server-Timing header as the sync route. + request_started: Optional[float] = Field(default=None, exclude=True) + timing_metrics: Optional[Dict[str, float]] = Field(default=None, + exclude=True) + response_format: Optional[Literal["file", "path"]] = Field( + default=None, + description=( + "Transport the client requested. ``GET /v1/videos/{id}/content`` " + "honors this: ``path`` returns the server-side output path(s) in a " + "JSON envelope; ``file`` (or unset) returns the file as a " + "``FileResponse`` download."), + ) + + +class VideoJobList(OpenAIBaseModel): + """Response from listing video jobs endpoint.""" + data: List[VideoJob] = Field(description="List of video jobs") + object: str = Field(default="list", description="Object type") diff --git a/tensorrt_llm/serve/openai_protocol.py b/tensorrt_llm/serve/openai_protocol.py index e2643c43d63c..d0b4bcf5f4b2 100644 --- a/tensorrt_llm/serve/openai_protocol.py +++ b/tensorrt_llm/serve/openai_protocol.py @@ -16,6 +16,7 @@ # Adapted from # https://github.com/vllm-project/vllm/blob/4db5176d9758b720b05460c50ace3c01026eb158/vllm/entrypoints/openai/protocol.py import base64 +import importlib import math import time import uuid @@ -23,7 +24,6 @@ import torch import xgrammar -from fastapi import UploadFile from openai.types.chat import ChatCompletionAssistantMessageParam from openai.types.chat import \ ChatCompletionContentPartParam as OpenAIChatCompletionContentPartParam @@ -47,9 +47,8 @@ from openai.types.responses.tool import Tool from openai.types.shared import Metadata, Reasoning from openai_harmony import ReasoningEffort -from pydantic import (AliasChoices, BaseModel, ConfigDict, Field, - NonNegativeInt, PositiveInt, field_validator, - model_serializer, model_validator) +from pydantic import (BaseModel, ConfigDict, Field, NonNegativeInt, PositiveInt, + field_validator, model_serializer, model_validator) from typing_extensions import Annotated, Required, TypeAlias, TypedDict from tensorrt_llm.executor.request import LoRARequest @@ -65,7 +64,6 @@ validate_thinking_token_budget) from tensorrt_llm.scheduling_params import AgentHierarchy from tensorrt_llm.serve.serving_extensions import structured_output_format_for -from tensorrt_llm.visual_gen.params import MediaRole _LOGIT_BIAS_MIN = -100.0 _LOGIT_BIAS_MAX = 100.0 @@ -2210,421 +2208,6 @@ def to_llm_conversation_params( conversation_id=conversation_params.conversation_id) -# ============================================================================ -# Diffusion API Protocol Classes -# ============================================================================ - - -class ImageGenerationRequest(OpenAIBaseModel): - """OpenAI-compatible image generation request. - - Universal per-request fields map 1:1 to :class:`VisualGenParams`. - Model-specific knobs (``stg_scale``, ``guidance_rescale``, …) - travel through ``extra_params``; the executor validates each - key against the loaded pipeline's - ``extra_param_specs``. Unknown top-level fields are rejected - with HTTP 422 via the inherited ``extra="forbid"`` policy. - """ - - # Prompt + transport (OpenAI-standard, always honored) - prompt: str - response_format: Literal["url", "b64_json", "path"] = "url" - format: Literal["png", "webp", "jpeg", "safetensors", "pt"] = Field( - default="png", - description=( - "Generation content encoding format. Image encoders write " - "``png``/``webp``/``jpeg``; tensor encoders write " - "``safetensors``/``pt`` for programmatic post-processing."), - ) - seed: Optional[int] = Field(default=None, - ge=0, - description="Random seed for reproducibility.") - - # Resolution. ``size`` is OpenAI-shaped; ``width`` + ``height`` are an - # equivalent structured alternative. Exactly one of width/height is - # rejected by the paired validator below. Numeric fields use - # ``gt=0`` as a safety net so zero / negative inputs are rejected - # with HTTP 422 before reaching the pipeline. - size: Optional[str] = Field(default=None, pattern=r"^(\d+x\d+|auto)$") - width: Optional[int] = Field(default=None, gt=0) - height: Optional[int] = Field(default=None, gt=0) - - # TRT-LLM-supported per-request params (1:1 with VisualGenParams fields) - num_inference_steps: Optional[int] = Field(default=None, gt=0) - guidance_scale: Optional[float] = Field(default=None, gt=0) - max_sequence_length: Optional[int] = Field(default=None, gt=0) - negative_prompt: Optional[str] = None - n: Optional[int] = Field( - default=None, - gt=0, - le=10, - description=("Number of images to generate. Capped at 10 to match the " - "OpenAI images API and to bound GPU memory / disk usage."), - ) - - # Model-specific overflow - extra_params: Optional[Dict[str, Any]] = Field( - default=None, - description=( - "Model-specific parameters forwarded to the underlying pipeline. " - "See per-model docs for accepted keys."), - ) - - # Accepted-but-ignored OpenAI-shaped fields. The conversion no-ops; the - # server logs WARNING when a client sets ``quality`` or ``style``, and - # WARNING-on-mismatch for ``model``. Kept in the schema so OpenAI-SDK - # clients don't trip ``extra="forbid"``. - model: Optional[str] = None - quality: Optional[Literal["standard", "hd"]] = None - style: Optional[Literal["vivid", "natural"]] = None - user: Optional[str] = None - - @model_validator(mode="after") - def _check_paired_dimensions(self): - """Reject sending exactly one of ``width`` / ``height``. - - Either both are sent (structured resolution wins over ``size``) - or neither is sent (``size`` or pipeline default applies). - """ - if (self.width is None) != (self.height is None): - raise ValueError( - "width and height must be sent together; got width=" - f"{self.width!r}, height={self.height!r}") - return self - - -class ImageEditRequest(OpenAIBaseModel): - """OpenAI-compatible image editing request. - - The server accepts the OpenAI multipart shape and a JSON/base64 - shape for tests and non-SDK clients. Model-specific knobs travel - through ``extra_params`` and are validated by the loaded visual - generation pipeline. - """ - - prompt: str - image: Union[str, UploadFile, List[Union[str, UploadFile]]] = Field( - description="Input image or images to edit.") - mask: Optional[Union[str, UploadFile]] = Field( - default=None, - description= - "Optional edit mask. Currently accepted for compatibility but unsupported.", - ) - response_format: Literal["url", "b64_json", "path"] = "url" - output_format: Literal["png", "webp", "jpeg"] = Field( - default="png", - validation_alias=AliasChoices("output_format", "format"), - description="Edited image content encoding format.", - ) - seed: Optional[int] = Field(default=None, - ge=0, - description="Random seed for reproducibility.") - - size: Optional[str] = Field(default=None, pattern=r"^(\d+x\d+|auto)$") - width: Optional[int] = Field(default=None, gt=0) - height: Optional[int] = Field(default=None, gt=0) - - num_inference_steps: Optional[int] = Field(default=None, gt=0) - guidance_scale: Optional[float] = Field(default=None, gt=0) - max_sequence_length: Optional[int] = Field(default=None, gt=0) - negative_prompt: Optional[str] = None - n: Optional[int] = Field( - default=None, - gt=0, - le=10, - description=("Number of edited images to generate. Capped at 10 to " - "match the OpenAI images API."), - ) - - extra_params: Optional[Dict[str, Any]] = Field( - default=None, - description=( - "Model-specific parameters forwarded to the underlying pipeline. " - "See per-model docs for accepted keys."), - ) - - model: Optional[str] = None - quality: Optional[Literal["standard", "hd"]] = None - user: Optional[str] = None - - @model_validator(mode="after") - def _check_paired_dimensions(self): - if isinstance(self.image, list) and not self.image: - raise ValueError("image must contain at least one input image") - if (self.width is None) != (self.height is None): - raise ValueError( - "width and height must be sent together; got width=" - f"{self.width!r}, height={self.height!r}") - return self - - -class ImageObject(OpenAIBaseModel): - """Generated image object in the response.""" - b64_json: Optional[str] = None - url: Optional[str] = None - path: Optional[str] = None - revised_prompt: Optional[str] = None - - -class ImageGenerationResponse(OpenAIBaseModel): - """Response from image generation endpoint. - - ``output_format`` reports the encoding actually applied to the - returned bytes / files so clients can decode or label the payload - correctly. Image encoders are ``"png"``/``"webp"``/``"jpeg"``; - tensor formats are ``"safetensors"``/``"pt"``. - """ - - created: int = Field(default_factory=lambda: int(time.time())) - data: List[ImageObject] - output_format: Literal["png", "webp", "jpeg", "safetensors", "pt"] = "png" - quality: Literal["low", "medium", "high"] = "medium" - size: Optional[str] = None - - -class MediaReferenceItem(OpenAIBaseModel): - """One media reference (image / video / audio) for conditioning (mirrors ``MediaRef``). - - ``format`` declares how to read ``content``; it is required, so no wire form - is ever guessed. The request field it sits in (``image_reference`` / - ``video_reference`` / ``audio_reference``) fixes the modality. ``role`` is - required only when it is ambiguous — a model with more than one required role - for that modality; a single role, or a single required role (e.g. i2v - first_frame), is inferred. - """ - - content: str = Field( - description="The reference payload, in the form declared by ``format``." - ) - format: Literal["path", "url", "base64"] = Field(description=( - "Wire form of ``content``: ``path`` (a file readable by the server; a " - "``file://`` URI is also accepted), ``url`` (``http(s)``, fetched " - "through the SSRF-guarded loader), or ``base64`` (a ``data:`` URI is " - "also accepted). Raw bytes reach the server as a multipart upload, " - "which carries no ``format`` of its own. Distinct from the top-level " - "``format``, which selects the *output* encoding.")) - role: Optional[MediaRole] = Field( - default=None, - description="Which conditioning slot this reference fills. Required only " - "when the target model accepts this modality in more than one slot; omit " - "it when the model leaves no ambiguity.", - ) - - -class VideoGenerationRequest(OpenAIBaseModel): - """Video generation request (extended API). - - Universal per-request fields map 1:1 to :class:`VisualGenParams`. - Model-specific knobs travel through ``extra_params``. Unknown - top-level fields are rejected with HTTP 422 via the inherited - ``extra="forbid"`` policy. - """ - - # Prompt + transport - prompt: str - response_format: Literal["file", "path"] = "file" - format: Literal["mp4", "avi", "auto", "safetensors", "pt"] = Field( - default="auto", - description=( - "Generation content encoding format. Video encoders write " - "``mp4``/``avi``/``auto``; tensor encoders write " - "``safetensors``/``pt`` and carry video, audio, and scalar " - "metadata (frame rate, audio sample rate) in one payload."), - ) - seed: Optional[int] = Field(default=None, - ge=0, - description="Random seed for reproducibility.") - image_reference: Optional[Union[ - UploadFile, MediaReferenceItem, List[MediaReferenceItem]]] = Field( - default=None, - description= - ("Image reference(s) conditioning generation (e.g. image-to-video " - "first frame). Send a ``{content, format, role}`` object or a list " - "of them, where ``format`` is ``path`` / ``url`` / ``base64``; or " - "upload a single image file via multipart, whose form is implied. " - "PNG or JPEG only — HEIF/AVIF are not supported."), - ) - video_reference: Optional[Union[ - UploadFile, MediaReferenceItem, List[MediaReferenceItem]]] = Field( - default=None, - description= - ("Video reference(s) conditioning generation (video-to-video). Send " - "a ``{content, format}`` object or a list of them, where ``format`` " - "is ``path`` / ``url`` / ``base64``; or upload a single video file " - "via multipart, whose form is implied. MP4 or AVI, with H.264 the " - "tested codec and others best-effort."), - ) - audio_reference: Optional[Union[ - UploadFile, MediaReferenceItem, List[MediaReferenceItem]]] = Field( - default=None, - description= - ("Audio reference(s) conditioning generation. Send a " - "``{content, format}`` object or a list of them, where ``format`` " - "is ``path`` / ``url`` / ``base64``; or upload a single audio " - "file via multipart, whose form is implied. Accepted only by " - "models that declare an audio reference slot."), - ) - input_reference: Optional[Union[str, UploadFile]] = Field( - default=None, - description= - ("Deprecated. A single image or video reference, routed by content " - "signature to image-to-video or video-to-video. A JSON request carries " - "base64 bytes; a multipart request uploads the file. Kept for backward " - "compatibility; prefer the typed ``image_reference`` / " - "``video_reference`` fields, which take precedence — this field is " - "ignored whenever a typed ``image_reference`` or ``video_reference`` " - "is provided."), - ) - - # Resolution - size: Optional[str] = Field(default=None, pattern=r"^(\d+x\d+|auto)$") - width: Optional[int] = Field(default=None, gt=0) - height: Optional[int] = Field(default=None, gt=0) - - # Frame budget. ``num_frames`` is preferred; if absent the engine - # derives it from ``seconds * frame_rate``. ``frame_rate`` is the - # canonical name (matches the Python field); ``fps`` is an alias for - # OpenAI-shape clients via ``populate_by_name=True``. - # All three constrain to strictly positive values so a zero - # ``frame_rate`` (division-by-zero in the AVI fallback) or a - # negative ``num_frames`` are rejected with HTTP 422 before - # reaching the encoder. - # Upper bounds keep request-boundary protection against requests - # that can exhaust GPU memory or pin the server on unbounded work. - # The numbers are generous (a minute of video at 120 fps) so common - # workloads pass; clients that need larger budgets can lift the cap - # at deployment time. - num_frames: Optional[int] = Field(default=None, gt=0, le=7200) - seconds: Optional[float] = Field(default=None, gt=0, le=60.0) - frame_rate: Optional[float] = Field(default=None, - alias="fps", - gt=0, - le=120.0) - - # TRT-LLM-supported per-request params (1:1 with VisualGenParams) - num_inference_steps: Optional[int] = Field(default=None, gt=0) - guidance_scale: Optional[float] = Field(default=None, gt=0) - max_sequence_length: Optional[int] = Field(default=None, gt=0) - negative_prompt: Optional[str] = None - - # Model-specific overflow - extra_params: Optional[Dict[str, Any]] = Field( - default=None, - description=( - "Model-specific parameters forwarded to the underlying pipeline. " - "See per-model docs for accepted keys."), - ) - - # Accepted-but-ignored OpenAI-shaped field - model: Optional[str] = None - - @model_validator(mode="after") - def _check_paired_dimensions(self): - """Reject sending exactly one of ``width`` / ``height``.""" - if (self.width is None) != (self.height is None): - raise ValueError( - "width and height must be sent together; got width=" - f"{self.width!r}, height={self.height!r}") - return self - - @field_validator("response_format", mode="before") - @classmethod - def _reject_removed_response_format(cls, value): - """Give migrating callers an actionable error for removed values. - - ``url``/``b64_json`` were valid before the transport rewrite; run - before the ``Literal`` check so the error names the replacement - instead of the generic "Input should be 'file' or 'path'". - """ - removed = { - "url": - ("'url' was removed for video; use 'file' (raw bytes -- the " - "old 'url' behavior, renamed) or 'path' (server-side path)."), - "b64_json": ("'b64_json' was removed for video; use 'file' (raw " - "bytes) or 'path' (server-side path)."), - } - if isinstance(value, str) and value in removed: - raise ValueError(removed[value]) - return value - - -class VideoJob(OpenAIBaseModel): - """Metadata for an asynchronous video generation job. - - Follows the OpenAI Videos API specification: - https://platform.openai.com/docs/api-reference/videos - """ - completed_at: Optional[int] = Field( - default=None, description="Unix timestamp of completion") - created_at: int = Field(description="Unix timestamp of creation") - error: Optional[str] = Field(default=None, - description="Error message if failed") - expires_at: Optional[int] = Field( - default=None, description="Unix timestamp of expiration") - id: str = Field(description="Unique identifier for the video") - model: str = Field(description="The model used for generation") - object: str = Field(default="video", description="Object type") - progress: Optional[int] = Field( - default=None, - description="Progress of the video generation job (0-100)") - prompt: str = Field(description="The prompt used to generate the video") - status: Literal["queued", "generating", "postprocessing", "completed", - "failed"] = Field(description=( - "Current status of the video generation job. " - "``generating`` (model inference) becomes " - "``postprocessing`` (encode and/or write the output " - "file) when inference finishes, then ``completed`` " - "once downloadable via ``/content``.")) - - # Video properties - duration: Optional[float] = Field(default=None, - description="Video duration in seconds") - fps: Optional[float] = Field( - default=None, - description=( - "Frames per second. Float to preserve cinematic rates such " - "as 23.976 / 29.97 that some encoders / pipelines use."), - ) - size: Optional[str] = Field(default=None, - description="Video dimensions in 'WxH' format") - # exclude=True: internal file-location for /content resolution + delete; - # never on the wire (the path payload is the hand-built {id, output_path} - # envelope in /content), so status/list model_dump() stays status-only. - output_path: Optional[str] = Field( - default=None, - exclude=True, - description="Server-side saved path (internal; excluded from the wire)." - ) - output_paths: Optional[List[str]] = Field( - default=None, - exclude=True, - description= - "Server-side paths for n>1 (internal; excluded from the wire).") - # exclude=True internal timings, never on the wire (status/list - # model_dump() stays status-only). ``request_started`` carries the - # steady-clock ``server_arrival_time``; the background task computes - # ``total`` from it and stores the header timings - # (``generation``/``denoise``/``total``) in ``timing_metrics`` so - # ``/content`` emits the same Server-Timing header as the sync route. - request_started: Optional[float] = Field(default=None, exclude=True) - timing_metrics: Optional[Dict[str, float]] = Field(default=None, - exclude=True) - response_format: Optional[Literal["file", "path"]] = Field( - default=None, - description=( - "Transport the client requested. ``GET /v1/videos/{id}/content`` " - "honors this: ``path`` returns the server-side output path(s) in a " - "JSON envelope; ``file`` (or unset) returns the file as a " - "``FileResponse`` download."), - ) - - -class VideoJobList(OpenAIBaseModel): - """Response from listing video jobs endpoint.""" - data: List[VideoJob] = Field(description="List of video jobs") - object: str = Field(default="list", description="Object type") - - UCompletionRequest = Union[CompletionRequest, ChatCompletionRequest] UCompletionResponse = Union[CompletionResponse, ChatCompletionResponse] @@ -2655,3 +2238,32 @@ class StartProfileRequest(OpenAIBaseModel): description="Which activities to trace. Supported values: " "'CPU', 'GPU', 'CUDA_PROFILER'. Unknown values are rejected at " "schema validation time.") + + +# Old import path of names that moved to +# tensorrt_llm.serve.api.visual_gen.protocol. +# They are re-exported lazily (PEP 562): importing this module does not import +# their new module. +_LAZY_NAMES = { + "ImageEditRequest": "tensorrt_llm.serve.api.visual_gen.protocol", + "ImageGenerationRequest": "tensorrt_llm.serve.api.visual_gen.protocol", + "ImageGenerationResponse": "tensorrt_llm.serve.api.visual_gen.protocol", + "ImageObject": "tensorrt_llm.serve.api.visual_gen.protocol", + "MediaReferenceItem": "tensorrt_llm.serve.api.visual_gen.protocol", + "VideoGenerationRequest": "tensorrt_llm.serve.api.visual_gen.protocol", + "VideoJob": "tensorrt_llm.serve.api.visual_gen.protocol", + "VideoJobList": "tensorrt_llm.serve.api.visual_gen.protocol", +} + + +def __getattr__(name): + module_name = _LAZY_NAMES.get(name) + if module_name is None: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + value = getattr(importlib.import_module(module_name), name) + globals()[name] = value # cache: later access skips __getattr__ + return value + + +def __dir__(): + return sorted(set(globals()) | set(_LAZY_NAMES)) diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py index 92f2d9478baf..7d2038081dae 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -83,6 +83,10 @@ AnthropicCreateBatchRequest, AnthropicMessagesRequest) from tensorrt_llm.serve.api.ops.system.serving import SystemServing +from tensorrt_llm.serve.api.visual_gen.protocol import (ImageEditRequest, + ImageGenerationRequest, + ImageGenerationResponse, + ImageObject) 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) @@ -107,12 +111,11 @@ ChatCompletionToolsParam, ChatMessage, CompletionRequest, CompletionResponse, CompletionResponseChoice, EmbeddingRequest, EmbeddingResponse, EmbeddingResponseData, EmbeddingUsageInfo, ErrorResponse, - ImageEditRequest, ImageGenerationRequest, ImageGenerationResponse, - 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) + 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, diff --git a/tensorrt_llm/serve/openai_video_routes.py b/tensorrt_llm/serve/openai_video_routes.py index aa91f34fac4d..75c414c1aca5 100644 --- a/tensorrt_llm/serve/openai_video_routes.py +++ b/tensorrt_llm/serve/openai_video_routes.py @@ -30,7 +30,11 @@ from tensorrt_llm.logger import logger from tensorrt_llm.media.encoding import resolve_video_format from tensorrt_llm.media.tensor_payload import is_tensor_format -from tensorrt_llm.serve.openai_protocol import VideoGenerationRequest, VideoJob, VideoJobList +from tensorrt_llm.serve.api.visual_gen.protocol import ( + VideoGenerationRequest, + VideoJob, + VideoJobList, +) from tensorrt_llm.serve.visual_gen_metrics import ( build_visual_gen_server_timings, build_visual_gen_timing_headers, diff --git a/tensorrt_llm/serve/scripts/BENCHMARKING_VISUAL_GEN.md b/tensorrt_llm/serve/scripts/BENCHMARKING_VISUAL_GEN.md index 378deb04529d..951916103d2b 100644 --- a/tensorrt_llm/serve/scripts/BENCHMARKING_VISUAL_GEN.md +++ b/tensorrt_llm/serve/scripts/BENCHMARKING_VISUAL_GEN.md @@ -15,7 +15,7 @@ This page repeats no field description. Each lives with its definition: |---|---| | every flag, with its default | `--help` | | what a generation parameter means | `VisualGenParams` in [`visual_gen/params.py`](../../visual_gen/params.py) | -| which backend accepts which field | `ImageGenerationRequest` · `ImageEditRequest` · `VideoGenerationRequest` in [`openai_protocol.py`](../openai_protocol.py) | +| which backend accepts which field | `ImageGenerationRequest` · `ImageEditRequest` · `VideoGenerationRequest` in [`api/visual_gen/protocol.py`](../api/visual_gen/protocol.py) | | what a result-JSON key holds | `VisualGenBenchResult` and `VisualGenRequestRecord` in [`benchmark_visual_gen.py`](benchmark_visual_gen.py) | ## Benchmarking args diff --git a/tensorrt_llm/serve/scripts/benchmark_visual_gen.py b/tensorrt_llm/serve/scripts/benchmark_visual_gen.py index 5761ba4c093f..6f671a7116f1 100644 --- a/tensorrt_llm/serve/scripts/benchmark_visual_gen.py +++ b/tensorrt_llm/serve/scripts/benchmark_visual_gen.py @@ -52,7 +52,7 @@ from tqdm.asyncio import tqdm from tensorrt_llm.llmapi.utils import StrictBaseModel -from tensorrt_llm.serve.openai_protocol import ( +from tensorrt_llm.serve.api.visual_gen.protocol import ( ImageEditRequest, ImageGenerationRequest, MediaReferenceItem, diff --git a/tensorrt_llm/serve/visual_gen_utils.py b/tensorrt_llm/serve/visual_gen_utils.py index 0702470f6d4f..bd5114a1480f 100644 --- a/tensorrt_llm/serve/visual_gen_utils.py +++ b/tensorrt_llm/serve/visual_gen_utils.py @@ -13,7 +13,7 @@ from tensorrt_llm.inputs.media_io import sniff_media_kind from tensorrt_llm.logger import logger -from tensorrt_llm.serve.openai_protocol import ( +from tensorrt_llm.serve.api.visual_gen.protocol import ( ImageEditRequest, ImageGenerationRequest, VideoGenerationRequest, diff --git a/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py b/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py index f0ce2de17187..17b3662c30d2 100644 --- a/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py +++ b/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py @@ -37,7 +37,7 @@ from PIL import Image from tensorrt_llm._utils import AdjustedSteadyClock -from tensorrt_llm.serve.openai_protocol import VideoJob +from tensorrt_llm.serve.api.visual_gen.protocol import VideoJob from tensorrt_llm.serve.openai_server import _normalize_image_output from tensorrt_llm.serve.visual_gen_metrics import SERVER_TIMING_HEADER from tensorrt_llm.serve.visual_gen_utils import VIDEO_STORE diff --git a/tests/unittest/_torch/visual_gen/test_visual_gen_params.py b/tests/unittest/_torch/visual_gen/test_visual_gen_params.py index 555cd5679a6b..c27920b36d8a 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_params.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_params.py @@ -1568,14 +1568,14 @@ class TestDeprecatedInputReferenceStaysCompatible: """ def test_a_bare_string_is_still_accepted(self): - from tensorrt_llm.serve.openai_protocol import VideoGenerationRequest + from tensorrt_llm.serve.api.visual_gen.protocol import VideoGenerationRequest assert VideoGenerationRequest(prompt="x", input_reference="aGk=").input_reference == "aGk=" def test_the_typed_field_still_requires_a_format(self): from pydantic import ValidationError - from tensorrt_llm.serve.openai_protocol import VideoGenerationRequest + from tensorrt_llm.serve.api.visual_gen.protocol import VideoGenerationRequest with pytest.raises(ValidationError): VideoGenerationRequest(prompt="x", image_reference="aGk=") diff --git a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py index b8d9aa865d60..5ebcb31c3e25 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py @@ -24,7 +24,10 @@ from PIL import Image from tensorrt_llm.serve import visual_gen_utils -from tensorrt_llm.serve.openai_protocol import ImageGenerationRequest, VideoGenerationRequest +from tensorrt_llm.serve.api.visual_gen.protocol import ( + ImageGenerationRequest, + VideoGenerationRequest, +) from tensorrt_llm.serve.visual_gen_utils import ( _merge_extra_params, _warn_if_set_with_no_semantic, @@ -160,7 +163,7 @@ class TestSeedLowerBoundOnServeBoundary: """ def test_image_seed_accepts_zero_and_large_values(self): - from tensorrt_llm.serve.openai_protocol import ImageGenerationRequest + from tensorrt_llm.serve.api.visual_gen.protocol import ImageGenerationRequest assert ImageGenerationRequest(prompt="x", seed=0).seed == 0 large = 2**40 @@ -169,7 +172,7 @@ def test_image_seed_accepts_zero_and_large_values(self): def test_image_seed_rejects_negative(self): from pydantic import ValidationError - from tensorrt_llm.serve.openai_protocol import ImageGenerationRequest + from tensorrt_llm.serve.api.visual_gen.protocol import ImageGenerationRequest with pytest.raises(ValidationError): ImageGenerationRequest(prompt="x", seed=-1) @@ -177,7 +180,7 @@ def test_image_seed_rejects_negative(self): def test_video_seed_rejects_negative(self): from pydantic import ValidationError - from tensorrt_llm.serve.openai_protocol import VideoGenerationRequest + from tensorrt_llm.serve.api.visual_gen.protocol import VideoGenerationRequest with pytest.raises(ValidationError): VideoGenerationRequest(prompt="x", seed=-1) From 9b8661fe654bf2430214382999f506e6de10f5a5 Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Tue, 3 Nov 2026 00:55:21 +0000 Subject: [PATCH 18/30] [TRTLLM-17264][infra] Give tensorrt_llm/serve/api/visual_gen to the VisualGen team in CODEOWNERS VisualGen serving code now lives under tensorrt_llm/serve/api/visual_gen/. The new rule, placed after the /tensorrt_llm/serve rule so that it wins, makes @NVIDIA/trt-llm-torch-visual-gen-devs the owner of that directory: changes there need the VisualGen team's approval and no longer need runtime-devs. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- .github/CODEOWNERS | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index ee619f1a9a52..0452e36a3fea 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -371,6 +371,7 @@ /cpp/tensorrt_llm/thop/*lysses* @NVIDIA/trt-llm-torch-visual-gen-devs /cpp/tensorrt_llm/kernels/ulysses* @NVIDIA/trt-llm-torch-visual-gen-devs /tensorrt_llm/media @NVIDIA/trt-llm-torch-visual-gen-devs +/tensorrt_llm/serve/api/visual_gen @NVIDIA/trt-llm-torch-visual-gen-devs /tensorrt_llm/visual_gen @NVIDIA/trt-llm-torch-visual-gen-devs @NVIDIA/trt-llm-runtime-devs /tests/integration/defs/examples/visual_gen @NVIDIA/trt-llm-torch-visual-gen-devs /tests/integration/defs/visual_gen @NVIDIA/trt-llm-torch-visual-gen-devs From e122006cdeaa3221380026f05b6d62feb581f164 Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Tue, 3 Nov 2026 00:55:37 +0000 Subject: [PATCH 19/30] [TRTLLM-17264][test] Gate the VisualGen protocol models and their lazy re-export under a visual_gen path Add tests/unittest/api_stability/visual_gen/test_visual_gen_serve_api.py: it runs the checks of test_serve_api.py on the models of references/trtllm_serve_api.yaml that are defined in tensorrt_llm.serve.api.visual_gen.protocol, and checks that tensorrt_llm.serve.openai_protocol re-exports the eight VisualGen names lazily (importing it does not load the VisualGen module; unknown names still raise AttributeError). The file lives under a visual_gen/ path and has its own l0_cpu.yml entry so that CI selection for VisualGen-only pull requests keeps running it; the reference YAML is unchanged. Add the eight identity rows of the VisualGen names to the VisualGen shim identity table: identity only, no warning and no ban on the old path, because the re-export is openai_protocol's permanent interface. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- .../integration/test_lists/test-db/l0_cpu.yml | 1 + .../visual_gen/test_visual_gen_serve_api.py | 177 ++++++++++++++++++ .../test_shim_identity_visual_gen.py | 73 +++++++- 3 files changed, 249 insertions(+), 2 deletions(-) create mode 100644 tests/unittest/api_stability/visual_gen/test_visual_gen_serve_api.py diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index af02b6fa2746..12efab5ea2d4 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -205,6 +205,7 @@ l0_cpu: - unittest/_torch/visual_gen/test_wan_timestep_routing.py - unittest/_torch/weight_sharing - unittest/api_stability + - unittest/api_stability/visual_gen/test_visual_gen_serve_api.py - unittest/bindings - unittest/disaggregated - unittest/docs/test_llmapi_reference_docs.py diff --git a/tests/unittest/api_stability/visual_gen/test_visual_gen_serve_api.py b/tests/unittest/api_stability/visual_gen/test_visual_gen_serve_api.py new file mode 100644 index 000000000000..f5654ca84746 --- /dev/null +++ b/tests/unittest/api_stability/visual_gen/test_visual_gen_serve_api.py @@ -0,0 +1,177 @@ +# 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. +"""CI gate for the VisualGen request and response models of ``trtllm-serve``. + +Runs the checks of ``tests/unittest/api_stability/test_serve_api.py`` on the +models pinned in ``references/trtllm_serve_api.yaml`` that are defined in +``tensorrt_llm.serve.api.visual_gen.protocol``, and checks that +``tensorrt_llm.serve.openai_protocol`` keeps re-exporting them lazily under +their old names. + +The file lives under a ``visual_gen/`` path so that CI selection for +VisualGen-only pull requests still runs it: the directory entry +``unittest/api_stability`` is not kept for those. The reference YAML is +unchanged and stays with the API review committee, which also owns this file. +""" + +from __future__ import annotations + +import importlib +import importlib.util +import json +import subprocess +import sys +import textwrap +from pathlib import Path +from typing import Any + +import pytest + +# The L0 CPU stage invokes registered files with `-m cpu_only`; without this +# marker every test is deselected and the stage collects nothing from the file. +pytestmark = pytest.mark.cpu_only + +_VISUAL_GEN_PROTOCOL_MODULE = "tensorrt_llm.serve.api.visual_gen.protocol" + +# The VisualGen models that `tensorrt_llm.serve.openai_protocol` re-exports. +_VISUAL_GEN_MODEL_NAMES = ( + "ImageGenerationRequest", + "ImageEditRequest", + "ImageObject", + "ImageGenerationResponse", + "MediaReferenceItem", + "VideoGenerationRequest", + "VideoJob", + "VideoJobList", +) + +# The VisualGen models the reference YAML pins; a model pinned later is covered +# automatically because the fixture selects by defining module. +_PINNED_VISUAL_GEN_MODELS = frozenset( + { + "ImageGenerationRequest", + "ImageGenerationResponse", + "MediaReferenceItem", + "VideoGenerationRequest", + } +) + + +def _load_serve_api_checks(): + """Load ``test_serve_api.py`` by path, without pytest collecting it twice. + + The module is kept in a private name: a public ``TestServeAPIStability`` + in this module's namespace would be collected a second time, with the + whole reference instead of the VisualGen part of it. + """ + path = Path(__file__).resolve().parents[1] / "test_serve_api.py" + spec = importlib.util.spec_from_file_location("trtllm_serve_api_checks", path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +_serve_api = _load_serve_api_checks() + + +class TestVisualGenServeAPIStability(_serve_api.TestServeAPIStability): + """The serve API gate, restricted to the pinned VisualGen models.""" + + @pytest.fixture(scope="class") + def reference(self) -> dict[str, Any]: + """The reference YAML with only the models defined in the VisualGen package. + + Every pinned model is resolved through ``tensorrt_llm.serve.openai_protocol``, + as the full gate does; a model counts as VisualGen when the class it resolves + to is defined in ``tensorrt_llm.serve.api.visual_gen.protocol``. + """ + full = _serve_api._load_reference() + facade = importlib.import_module(_serve_api._PROTOCOL_MODULE) + models: dict[str, Any] = {} + for model_name, model_spec in (full.get("models", {}) or {}).items(): + cls = getattr(facade, model_name, None) + if cls is not None and cls.__module__ == _VISUAL_GEN_PROTOCOL_MODULE: + models[model_name] = model_spec + missing = sorted(_PINNED_VISUAL_GEN_MODELS - set(models)) + assert not missing, ( + f"trtllm_serve_api.yaml pins VisualGen models {missing} but they no longer " + f"resolve to {_VISUAL_GEN_PROTOCOL_MODULE} through {_serve_api._PROTOCOL_MODULE}. " + "Keep them in serve/api/visual_gen/protocol.py and in the lazy re-export of " + "openai_protocol; changes to the YAML need the API review committee." + ) + return {**full, "models": models} + + +def _run_in_fresh_interpreter(code: str) -> dict: + """Run ``code`` in a new interpreter and return the JSON it prints last.""" + result = subprocess.run( + [sys.executable, "-c", textwrap.dedent(code)], + capture_output=True, + text=True, + timeout=600, + ) + assert result.returncode == 0, result.stderr[-2000:] + return json.loads(result.stdout.strip().splitlines()[-1]) + + +class TestVisualGenProtocolLazyReexport: + """``openai_protocol`` keeps the VisualGen names without importing their module.""" + + def test_importing_openai_protocol_does_not_load_visual_gen(self) -> None: + out = _run_in_fresh_interpreter( + f""" + import json, sys + import tensorrt_llm.serve.openai_protocol as protocol + + target = {_VISUAL_GEN_PROTOCOL_MODULE!r} + before = target in sys.modules + importers = [] + if before: + loaded = sys.modules[target] + for name, module in list(sys.modules.items()): + if module is loaded or module is None: + continue + try: + values = list(vars(module).values()) + except TypeError: + continue + if any(v is loaded or getattr(v, "__module__", None) == target for v in values): + importers.append(name) + protocol.VideoJob + after = target in sys.modules + print(json.dumps({{"before": before, "after": after, "importers": sorted(importers)}})) + """ + ) + assert out["before"] is False, ( + f"importing tensorrt_llm.serve.openai_protocol loaded {_VISUAL_GEN_PROTOCOL_MODULE} " + f"eagerly; the modules holding its names are {out['importers']}. The VisualGen " + "models are re-exported lazily through the module __getattr__ of openai_protocol; " + "do not import them at module level there." + ) + assert out["after"] is True, ( + "reading openai_protocol.VideoJob did not import " + f"{_VISUAL_GEN_PROTOCOL_MODULE}; the lazy re-export no longer resolves the name there." + ) + + def test_reexported_names_are_listed_and_unknown_names_raise(self) -> None: + protocol = importlib.import_module(_serve_api._PROTOCOL_MODULE) + listed = set(dir(protocol)) + not_listed = [name for name in _VISUAL_GEN_MODEL_NAMES if name not in listed] + assert not not_listed, ( + f"dir(tensorrt_llm.serve.openai_protocol) lacks {not_listed}; add the names to the " + "lazy re-export table of openai_protocol." + ) + with pytest.raises(AttributeError): + getattr(protocol, "NotAVisualGenModel") 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 index 1e2d2d57bbe0..d767e5fc781a 100644 --- 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 @@ -25,13 +25,82 @@ from pathlib import Path import pytest -from test_common.serve.shims import check_identity, check_warnings, find_old_path_imports +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[5] -VISUAL_GEN_SHIM_ROWS: tuple = () +# The VisualGen request and response models live in +# ``tensorrt_llm.serve.api.visual_gen.protocol``; ``tensorrt_llm.serve.openai_protocol`` +# keeps re-exporting them lazily as part of its own interface, for good: no warning, no +# ban on the old path. These rows check identity only. +VISUAL_GEN_SHIM_ROWS: tuple = ( + ShimRow( + "tensorrt_llm.serve.openai_protocol", + "ImageGenerationRequest", + "tensorrt_llm.serve.api.visual_gen.protocol", + retire_in="not retired", + ban_old_path=False, + warns=False, + ), + ShimRow( + "tensorrt_llm.serve.openai_protocol", + "ImageEditRequest", + "tensorrt_llm.serve.api.visual_gen.protocol", + retire_in="not retired", + ban_old_path=False, + warns=False, + ), + ShimRow( + "tensorrt_llm.serve.openai_protocol", + "ImageObject", + "tensorrt_llm.serve.api.visual_gen.protocol", + retire_in="not retired", + ban_old_path=False, + warns=False, + ), + ShimRow( + "tensorrt_llm.serve.openai_protocol", + "ImageGenerationResponse", + "tensorrt_llm.serve.api.visual_gen.protocol", + retire_in="not retired", + ban_old_path=False, + warns=False, + ), + ShimRow( + "tensorrt_llm.serve.openai_protocol", + "MediaReferenceItem", + "tensorrt_llm.serve.api.visual_gen.protocol", + retire_in="not retired", + ban_old_path=False, + warns=False, + ), + ShimRow( + "tensorrt_llm.serve.openai_protocol", + "VideoGenerationRequest", + "tensorrt_llm.serve.api.visual_gen.protocol", + retire_in="not retired", + ban_old_path=False, + warns=False, + ), + ShimRow( + "tensorrt_llm.serve.openai_protocol", + "VideoJob", + "tensorrt_llm.serve.api.visual_gen.protocol", + retire_in="not retired", + ban_old_path=False, + warns=False, + ), + ShimRow( + "tensorrt_llm.serve.openai_protocol", + "VideoJobList", + "tensorrt_llm.serve.api.visual_gen.protocol", + retire_in="not retired", + ban_old_path=False, + warns=False, + ), +) def test_visual_gen_shim_identity(): From 18a8b2f9377ea752cd7e880c45d060bfd98dd632 Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Tue, 3 Nov 2026 04:01:33 +0000 Subject: [PATCH 20/30] [TRTLLM-17264][test] Harden the VisualGen api-stability gate and pin its no-warning contract Resolve every pinned model with the full gate's own loader, so a VisualGen model pinned in the reference YAML that is later renamed or removed fails this gate too instead of being dropped silently; the four known models stay a secondary check. Make the lazy-load probe's diagnostic best-effort and identity-based, so a failure prints which modules hold the VisualGen classes instead of crashing on module-level __getattr__ hooks while walking sys.modules. In the same fresh-interpreter probe, assert that reading each of the eight names from openai_protocol emits no FutureWarning or DeprecationWarning. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- .../visual_gen/test_visual_gen_serve_api.py | 104 ++++++++++++------ 1 file changed, 73 insertions(+), 31 deletions(-) diff --git a/tests/unittest/api_stability/visual_gen/test_visual_gen_serve_api.py b/tests/unittest/api_stability/visual_gen/test_visual_gen_serve_api.py index f5654ca84746..f691dd23c5a8 100644 --- a/tests/unittest/api_stability/visual_gen/test_visual_gen_serve_api.py +++ b/tests/unittest/api_stability/visual_gen/test_visual_gen_serve_api.py @@ -18,7 +18,7 @@ models pinned in ``references/trtllm_serve_api.yaml`` that are defined in ``tensorrt_llm.serve.api.visual_gen.protocol``, and checks that ``tensorrt_llm.serve.openai_protocol`` keeps re-exporting them lazily under -their old names. +their old names, without a warning. The file lives under a ``visual_gen/`` path so that CI selection for VisualGen-only pull requests still runs it: the directory entry @@ -57,8 +57,11 @@ "VideoJobList", ) -# The VisualGen models the reference YAML pins; a model pinned later is covered -# automatically because the fixture selects by defining module. +# The VisualGen models the reference YAML pins today. The fixture below resolves +# every pinned model with the full gate's own loader, which fails for a renamed or +# removed model whatever its name, and selects the VisualGen ones by their defining +# module, so a model pinned later is gated automatically. This set is a secondary +# check that the four known models did not quietly stop being VisualGen models. _PINNED_VISUAL_GEN_MODELS = frozenset( { "ImageGenerationRequest", @@ -93,16 +96,17 @@ class TestVisualGenServeAPIStability(_serve_api.TestServeAPIStability): def reference(self) -> dict[str, Any]: """The reference YAML with only the models defined in the VisualGen package. - Every pinned model is resolved through ``tensorrt_llm.serve.openai_protocol``, - as the full gate does; a model counts as VisualGen when the class it resolves - to is defined in ``tensorrt_llm.serve.api.visual_gen.protocol``. + Every pinned model is resolved through ``tensorrt_llm.serve.openai_protocol`` + with the full gate's own loader, so a pinned model that no longer resolves + there fails this gate too, whether or not it is a VisualGen model. A model + counts as VisualGen when the class it resolves to is defined in + ``tensorrt_llm.serve.api.visual_gen.protocol``. """ full = _serve_api._load_reference() - facade = importlib.import_module(_serve_api._PROTOCOL_MODULE) models: dict[str, Any] = {} for model_name, model_spec in (full.get("models", {}) or {}).items(): - cls = getattr(facade, model_name, None) - if cls is not None and cls.__module__ == _VISUAL_GEN_PROTOCOL_MODULE: + cls = _serve_api._load_model(model_name) + if cls.__module__ == _VISUAL_GEN_PROTOCOL_MODULE: models[model_name] = model_spec missing = sorted(_PINNED_VISUAL_GEN_MODELS - set(models)) assert not missing, ( @@ -122,7 +126,10 @@ def _run_in_fresh_interpreter(code: str) -> dict: text=True, timeout=600, ) - assert result.returncode == 0, result.stderr[-2000:] + assert result.returncode == 0, ( + f"the probe interpreter exited with {result.returncode}; stdout tail: " + f"{result.stdout[-1000:]!r}; stderr tail: {result.stderr[-2000:]!r}" + ) return json.loads(result.stdout.strip().splitlines()[-1]) @@ -132,37 +139,72 @@ class TestVisualGenProtocolLazyReexport: def test_importing_openai_protocol_does_not_load_visual_gen(self) -> None: out = _run_in_fresh_interpreter( f""" - import json, sys + import json, sys, warnings import tensorrt_llm.serve.openai_protocol as protocol target = {_VISUAL_GEN_PROTOCOL_MODULE!r} - before = target in sys.modules - importers = [] - if before: - loaded = sys.modules[target] - for name, module in list(sys.modules.items()): - if module is loaded or module is None: - continue + names = {list(_VISUAL_GEN_MODEL_NAMES)!r} + result = {{ + "before": target in sys.modules, + "importers": [], + "diagnostic_error": None, + "warned": [], + "errors": [], + }} + if result["before"]: + # Best effort and identity only: attribute access on arbitrary objects + # in sys.modules can raise (module-level __getattr__ hooks do), and the + # verdict below is printed whatever happens here. + try: + loaded = sys.modules[target] + held = [loaded] + [getattr(loaded, n) for n in names if hasattr(loaded, n)] + for module_name, module in list(sys.modules.items()): + if module is None or module is loaded: + continue + try: + values = list(vars(module).values()) + except TypeError: + continue + if any(value is obj for value in values for obj in held): + result["importers"].append(module_name) + result["importers"].sort() + except Exception as exc: + result["diagnostic_error"] = repr(exc) + with warnings.catch_warnings(): + warnings.simplefilter("error", FutureWarning) + warnings.simplefilter("error", DeprecationWarning) + for name in names: try: - values = list(vars(module).values()) - except TypeError: - continue - if any(v is loaded or getattr(v, "__module__", None) == target for v in values): - importers.append(name) - protocol.VideoJob - after = target in sys.modules - print(json.dumps({{"before": before, "after": after, "importers": sorted(importers)}})) + getattr(protocol, name) + except (FutureWarning, DeprecationWarning) as exc: + result["warned"].append(f"{{name}}: {{exc}}") + except Exception as exc: + result["errors"].append(f"{{name}}: {{exc!r}}") + result["after"] = target in sys.modules + print(json.dumps(result)) """ ) + diagnostic = ( + f" (diagnostic error: {out['diagnostic_error']})" if out["diagnostic_error"] else "" + ) assert out["before"] is False, ( f"importing tensorrt_llm.serve.openai_protocol loaded {_VISUAL_GEN_PROTOCOL_MODULE} " - f"eagerly; the modules holding its names are {out['importers']}. The VisualGen " - "models are re-exported lazily through the module __getattr__ of openai_protocol; " - "do not import them at module level there." + f"eagerly; the modules holding its names are {out['importers']}{diagnostic}. The " + "VisualGen models are re-exported lazily through the module __getattr__ of " + "openai_protocol; do not import them at module level there." + ) + assert out["errors"] == [], ( + "reading the lazily re-exported names from tensorrt_llm.serve.openai_protocol " + f"failed: {out['errors']}" + ) + assert out["warned"] == [], ( + f"reading {out['warned']} from tensorrt_llm.serve.openai_protocol emitted a " + "FutureWarning or DeprecationWarning; the lazy re-export is openai_protocol's own " + "permanent interface and must not warn." ) assert out["after"] is True, ( - "reading openai_protocol.VideoJob did not import " - f"{_VISUAL_GEN_PROTOCOL_MODULE}; the lazy re-export no longer resolves the name there." + "reading the VisualGen names from openai_protocol did not import " + f"{_VISUAL_GEN_PROTOCOL_MODULE}; the lazy re-export no longer resolves them there." ) def test_reexported_names_are_listed_and_unknown_names_raise(self) -> None: From 1fc2c43cdd797abe280405ec5282808a988c2c67 Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Wed, 4 Nov 2026 12:11:00 +0000 Subject: [PATCH 21/30] [TRTLLM-17265][refactor] Split openai_protocol.py into a package with one submodule per API tensorrt_llm/serve/openai_protocol.py held every OpenAI-dialect request and response model plus the helpers that turn them into engine parameters. It becomes the package tensorrt_llm/serve/openai_protocol/ with one submodule per API: common, sampling, tools, chat, completions, embeddings, responses, tokenize, models, kv_cache, control and disagg. Every definition moves unchanged (93 definitions; generated by an AST tool and checked definition by definition against the former module). The package __init__.py is a lazy facade: a name-to-module table and a module __getattr__ import a submodule only when one of its names is first read, so every import from tensorrt_llm.serve.openai_protocol keeps working for every name the former module had, the underscore helpers that tests import included, and the VisualGen models re-exported from tensorrt_llm.serve.api.visual_gen.protocol. The facade is the package's own interface, not a compatibility shim: no warning, no removal release, and no importer in the repository changes. Its TYPE_CHECKING block imports the public names from their submodules so type checkers resolve them; the VisualGen names are left out of it on purpose and appear only as strings in the table, so that importing the facade never loads the VisualGen package. DisaggScheduleStyle, which the former module imported and two importers took from it, stays reachable through the facade. Submodules import each other directly, never through the package, along an acyclic graph: common imports no sibling; sampling, tools, models, tokenize, embeddings and control import common; completions and responses import common and sampling; chat imports common, sampling and tools; kv_cache imports common, tools and chat; disagg imports common, completions and chat. torch is imported only by sampling and disagg, openai_harmony only by chat, and the serving-extension registry only by sampling, so code that needs only the tool types no longer loads the chat or Responses models. The 12 submodules are registered in legacy-files.txt (their code comes from a legacy-formatted file) and the deleted module's line is removed; the derived lint configs are regenerated with scripts/legacy_utils.py gen-configs. The facade is new code and stays under the default formatter. No behavior change. OpenAPI component names may change for classes that share a short name with an openai type, because pydantic qualifies them with the module path; they are not a contract. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- .pre-commit-config.yaml | 26 +- legacy-files.txt | 13 +- pyproject.toml | 13 +- ruff-legacy.toml | 13 +- tensorrt_llm/serve/openai_protocol.py | 2269 ----------------- .../serve/openai_protocol/__init__.py | 329 +++ tensorrt_llm/serve/openai_protocol/chat.py | 556 ++++ tensorrt_llm/serve/openai_protocol/common.py | 242 ++ .../serve/openai_protocol/completions.py | 279 ++ tensorrt_llm/serve/openai_protocol/control.py | 70 + tensorrt_llm/serve/openai_protocol/disagg.py | 221 ++ .../serve/openai_protocol/embeddings.py | 90 + .../serve/openai_protocol/kv_cache.py | 66 + tensorrt_llm/serve/openai_protocol/models.py | 36 + .../serve/openai_protocol/responses.py | 565 ++++ .../serve/openai_protocol/sampling.py | 221 ++ .../serve/openai_protocol/tokenize.py | 59 + tensorrt_llm/serve/openai_protocol/tools.py | 141 + 18 files changed, 2935 insertions(+), 2274 deletions(-) delete mode 100644 tensorrt_llm/serve/openai_protocol.py create mode 100644 tensorrt_llm/serve/openai_protocol/__init__.py create mode 100644 tensorrt_llm/serve/openai_protocol/chat.py create mode 100644 tensorrt_llm/serve/openai_protocol/common.py create mode 100644 tensorrt_llm/serve/openai_protocol/completions.py create mode 100644 tensorrt_llm/serve/openai_protocol/control.py create mode 100644 tensorrt_llm/serve/openai_protocol/disagg.py create mode 100644 tensorrt_llm/serve/openai_protocol/embeddings.py create mode 100644 tensorrt_llm/serve/openai_protocol/kv_cache.py create mode 100644 tensorrt_llm/serve/openai_protocol/models.py create mode 100644 tensorrt_llm/serve/openai_protocol/responses.py create mode 100644 tensorrt_llm/serve/openai_protocol/sampling.py create mode 100644 tensorrt_llm/serve/openai_protocol/tokenize.py create mode 100644 tensorrt_llm/serve/openai_protocol/tools.py diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 23df44e55e96..bf7b7b678c60 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -425,7 +425,18 @@ common-files: &common_files | tensorrt_llm/serve/harmony_adapter.py | tensorrt_llm/serve/metadata_server.py | tensorrt_llm/serve/openai_disagg_server.py | - tensorrt_llm/serve/openai_protocol.py | + tensorrt_llm/serve/openai_protocol/chat.py | + tensorrt_llm/serve/openai_protocol/common.py | + tensorrt_llm/serve/openai_protocol/completions.py | + tensorrt_llm/serve/openai_protocol/control.py | + tensorrt_llm/serve/openai_protocol/disagg.py | + tensorrt_llm/serve/openai_protocol/embeddings.py | + tensorrt_llm/serve/openai_protocol/kv_cache.py | + tensorrt_llm/serve/openai_protocol/models.py | + tensorrt_llm/serve/openai_protocol/responses.py | + tensorrt_llm/serve/openai_protocol/sampling.py | + tensorrt_llm/serve/openai_protocol/tokenize.py | + tensorrt_llm/serve/openai_protocol/tools.py | tensorrt_llm/serve/openai_server.py | tensorrt_llm/serve/postprocess_handlers.py | tensorrt_llm/serve/responses_utils.py | @@ -1154,7 +1165,18 @@ legacy-files: &legacy_files | tensorrt_llm/serve/harmony_adapter.py | tensorrt_llm/serve/metadata_server.py | tensorrt_llm/serve/openai_disagg_server.py | - tensorrt_llm/serve/openai_protocol.py | + tensorrt_llm/serve/openai_protocol/chat.py | + tensorrt_llm/serve/openai_protocol/common.py | + tensorrt_llm/serve/openai_protocol/completions.py | + tensorrt_llm/serve/openai_protocol/control.py | + tensorrt_llm/serve/openai_protocol/disagg.py | + tensorrt_llm/serve/openai_protocol/embeddings.py | + tensorrt_llm/serve/openai_protocol/kv_cache.py | + tensorrt_llm/serve/openai_protocol/models.py | + tensorrt_llm/serve/openai_protocol/responses.py | + tensorrt_llm/serve/openai_protocol/sampling.py | + tensorrt_llm/serve/openai_protocol/tokenize.py | + tensorrt_llm/serve/openai_protocol/tools.py | tensorrt_llm/serve/openai_server.py | tensorrt_llm/serve/postprocess_handlers.py | tensorrt_llm/serve/responses_utils.py | diff --git a/legacy-files.txt b/legacy-files.txt index 72f1c66e0acd..3c3a8846802b 100644 --- a/legacy-files.txt +++ b/legacy-files.txt @@ -417,7 +417,18 @@ tensorrt_llm/serve/disagg_auto_scaling.py tensorrt_llm/serve/harmony_adapter.py tensorrt_llm/serve/metadata_server.py tensorrt_llm/serve/openai_disagg_server.py -tensorrt_llm/serve/openai_protocol.py +tensorrt_llm/serve/openai_protocol/chat.py +tensorrt_llm/serve/openai_protocol/common.py +tensorrt_llm/serve/openai_protocol/completions.py +tensorrt_llm/serve/openai_protocol/control.py +tensorrt_llm/serve/openai_protocol/disagg.py +tensorrt_llm/serve/openai_protocol/embeddings.py +tensorrt_llm/serve/openai_protocol/kv_cache.py +tensorrt_llm/serve/openai_protocol/models.py +tensorrt_llm/serve/openai_protocol/responses.py +tensorrt_llm/serve/openai_protocol/sampling.py +tensorrt_llm/serve/openai_protocol/tokenize.py +tensorrt_llm/serve/openai_protocol/tools.py tensorrt_llm/serve/openai_server.py tensorrt_llm/serve/postprocess_handlers.py tensorrt_llm/serve/responses_utils.py diff --git a/pyproject.toml b/pyproject.toml index a65da3ab8465..8f59210c181c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -482,7 +482,18 @@ exclude = [ "tensorrt_llm/serve/harmony_adapter.py", "tensorrt_llm/serve/metadata_server.py", "tensorrt_llm/serve/openai_disagg_server.py", - "tensorrt_llm/serve/openai_protocol.py", + "tensorrt_llm/serve/openai_protocol/chat.py", + "tensorrt_llm/serve/openai_protocol/common.py", + "tensorrt_llm/serve/openai_protocol/completions.py", + "tensorrt_llm/serve/openai_protocol/control.py", + "tensorrt_llm/serve/openai_protocol/disagg.py", + "tensorrt_llm/serve/openai_protocol/embeddings.py", + "tensorrt_llm/serve/openai_protocol/kv_cache.py", + "tensorrt_llm/serve/openai_protocol/models.py", + "tensorrt_llm/serve/openai_protocol/responses.py", + "tensorrt_llm/serve/openai_protocol/sampling.py", + "tensorrt_llm/serve/openai_protocol/tokenize.py", + "tensorrt_llm/serve/openai_protocol/tools.py", "tensorrt_llm/serve/openai_server.py", "tensorrt_llm/serve/postprocess_handlers.py", "tensorrt_llm/serve/responses_utils.py", diff --git a/ruff-legacy.toml b/ruff-legacy.toml index 6485e090c26b..cca0b8dd4f9f 100644 --- a/ruff-legacy.toml +++ b/ruff-legacy.toml @@ -434,7 +434,18 @@ include = [ "tensorrt_llm/serve/harmony_adapter.py", "tensorrt_llm/serve/metadata_server.py", "tensorrt_llm/serve/openai_disagg_server.py", - "tensorrt_llm/serve/openai_protocol.py", + "tensorrt_llm/serve/openai_protocol/chat.py", + "tensorrt_llm/serve/openai_protocol/common.py", + "tensorrt_llm/serve/openai_protocol/completions.py", + "tensorrt_llm/serve/openai_protocol/control.py", + "tensorrt_llm/serve/openai_protocol/disagg.py", + "tensorrt_llm/serve/openai_protocol/embeddings.py", + "tensorrt_llm/serve/openai_protocol/kv_cache.py", + "tensorrt_llm/serve/openai_protocol/models.py", + "tensorrt_llm/serve/openai_protocol/responses.py", + "tensorrt_llm/serve/openai_protocol/sampling.py", + "tensorrt_llm/serve/openai_protocol/tokenize.py", + "tensorrt_llm/serve/openai_protocol/tools.py", "tensorrt_llm/serve/openai_server.py", "tensorrt_llm/serve/postprocess_handlers.py", "tensorrt_llm/serve/responses_utils.py", diff --git a/tensorrt_llm/serve/openai_protocol.py b/tensorrt_llm/serve/openai_protocol.py deleted file mode 100644 index d0b4bcf5f4b2..000000000000 --- a/tensorrt_llm/serve/openai_protocol.py +++ /dev/null @@ -1,2269 +0,0 @@ -# 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. - -# Adapted from -# https://github.com/vllm-project/vllm/blob/4db5176d9758b720b05460c50ace3c01026eb158/vllm/entrypoints/openai/protocol.py -import base64 -import importlib -import math -import time -import uuid -from typing import Any, Dict, List, Literal, Optional, Union - -import torch -import xgrammar -from openai.types.chat import ChatCompletionAssistantMessageParam -from openai.types.chat import \ - ChatCompletionContentPartParam as OpenAIChatCompletionContentPartParam -from openai.types.chat import \ - ChatCompletionMessageParam as OpenAIChatCompletionMessageParam -from openai.types.responses import ( - ResponseCodeInterpreterCallCodeDeltaEvent, - ResponseCodeInterpreterCallCodeDoneEvent, - ResponseCodeInterpreterCallCompletedEvent, - ResponseCodeInterpreterCallInProgressEvent, - ResponseCodeInterpreterCallInterpretingEvent, ResponseCompletedEvent, - ResponseContentPartAddedEvent, ResponseContentPartDoneEvent, - ResponseCreatedEvent, ResponseFormatTextConfig, ResponseFunctionToolCall, - ResponseInProgressEvent, ResponseInputItemParam, ResponseOutputItem, - ResponseOutputItemAddedEvent, ResponseOutputItemDoneEvent, ResponsePrompt, - ResponseReasoningItem, ResponseReasoningTextDeltaEvent, - ResponseReasoningTextDoneEvent, ResponseStatus, ResponseTextConfig, - ResponseWebSearchCallCompletedEvent, ResponseWebSearchCallInProgressEvent, - ResponseWebSearchCallSearchingEvent) -from openai.types.responses.response import IncompleteDetails, ToolChoice -from openai.types.responses.tool import Tool -from openai.types.shared import Metadata, Reasoning -from openai_harmony import ReasoningEffort -from pydantic import (BaseModel, ConfigDict, Field, NonNegativeInt, PositiveInt, - field_validator, model_serializer, model_validator) -from typing_extensions import Annotated, Required, TypeAlias, TypedDict - -from tensorrt_llm.executor.request import LoRARequest -from tensorrt_llm.inputs.media_io import MediaModality -from tensorrt_llm.kv_hints import KvHintsEnvelope -from tensorrt_llm.kv_hints.protocol import admit_kv_hints -from tensorrt_llm.llmapi import ConversationParams as LlmConversationParams -from tensorrt_llm.llmapi import DisaggregatedParams as LlmDisaggregatedParams -from tensorrt_llm.llmapi import (DisaggScheduleStyle, GuidedDecodingParams, - SamplingParams) -from tensorrt_llm.llmapi.reasoning_parser import ReasoningParserFactory -from tensorrt_llm.sampling_params import (check_logprobs_limit, - validate_thinking_token_budget) -from tensorrt_llm.scheduling_params import AgentHierarchy -from tensorrt_llm.serve.serving_extensions import structured_output_format_for - -_LOGIT_BIAS_MIN = -100.0 -_LOGIT_BIAS_MAX = 100.0 - -# Beam-search stopping mode as exposed over HTTP, mirroring the HuggingFace -# Transformers interface. The engine encodes it as an integer, but that encoding -# is an implementation detail and is kept out of the public schema. -EarlyStopping: TypeAlias = Union[bool, Literal["never"]] - - -def _early_stopping_to_int(value: Optional[EarlyStopping]) -> Optional[int]: - """Translate the HF-style tri-state into the engine's integer encoding. - - Mirrors ``BeamSearchEarlyStop``: ``False`` -> 0, ``True`` -> 1, - ``"never"`` -> 2. ``None`` stays unset so the engine picks its default. - """ - if value is None: - return None - if value == "never": - return 2 - # NB: bool is a subclass of int, so this yields 1/0 for True/False. - return int(value) - - -def ensure_request_chat_template_allowed(request: Any, - allow_request_chat_template: bool): - if (getattr(request, "chat_template", None) is not None - and not allow_request_chat_template): - raise ValueError( - "chat_template cannot be supplied per request unless request-level " - "chat templates are enabled at server startup.") - - -def _logit_bias_to_embedding_bias( - logit_bias: Optional[Dict[str, float]], - vocab_size: Optional[int]) -> Optional[torch.Tensor]: - """Convert OpenAI logit_bias dict to embedding_bias tensor for sampling.""" - if logit_bias is None: - return None - if vocab_size is None: - raise ValueError( - "logit_bias requires a tokenizer, but the server was started " - "without one (e.g. num_postprocess_workers > 0). " - "Remove logit_bias from your request or set num_postprocess_workers=0." - ) - elif vocab_size <= 0: - raise ValueError("vocab_size must be positive when logit_bias is used") - - # Create 1D zeros tensor as expected by executor API (will be unsqueezed to [1, vocab_size] internally) - embedding_bias = torch.zeros(vocab_size, dtype=torch.float32) - - # Apply biases for specified token IDs - for token_str, bias in logit_bias.items(): - try: - token_id = int(token_str) - except ValueError as e: - if "invalid literal" in str(e): - raise ValueError( - f"Invalid logit_bias key '{token_str}': must be a valid integer token ID" - ) - raise - if not 0 <= token_id < vocab_size: - raise ValueError( - f"Token ID {token_id} out of vocabulary range [0, {vocab_size})" - ) - try: - bias_value = float(bias) - except (TypeError, ValueError) as e: - raise ValueError( - f"logit_bias value for token ID {token_id} must be a number" - ) from e - if not math.isfinite(bias_value): - raise ValueError(f"logit_bias value for token ID {token_id} " - "must be finite") - if not _LOGIT_BIAS_MIN <= bias_value <= _LOGIT_BIAS_MAX: - raise ValueError( - f"logit_bias value for token ID {token_id} must be in " - f"[{_LOGIT_BIAS_MIN:g}, {_LOGIT_BIAS_MAX:g}]") - embedding_bias[token_id] = bias_value - - return embedding_bias - - -class OpenAIBaseModel(BaseModel): - # OpenAI API does not allow extra fields & allow to initialize by both alias and field name - model_config = ConfigDict(extra="forbid", populate_by_name=True) - - @field_validator("kv_hints", mode="before", check_fields=False) - @classmethod - def admit_hint_envelope(cls, value: Any) -> Optional[KvHintsEnvelope]: - """Apply optional-hint admission before nested schema validation. - - Args: - value: The decoded hint field from a trusted ingress. - - Returns: - Supported metadata or None for ignored hints, including malformed - envelopes. Validation of other request fields remains unchanged. - """ - return admit_kv_hints(value) - - -class SpeculativeDecodingStats(OpenAIBaseModel): - """Per-request speculative-decoding acceptance for one generated sequence. - - Opt-in: emitted only when the server sets per_request_spec_decode_stats. - Absent entirely when the request drafted nothing. PyTorch backend only. - - ``mean acceptance length`` is deliberately not a field here: it is already - reported per choice as ``avg_decoded_tokens_per_iter``, and duplicating it - would let the two drift. Consumers derive it as - ``1 + total_accepted_draft_tokens / num_spec_steps``. - - Three identities hold on every emitted record, and a consumer may rely on - them: - - * ``sum(acceptance_histogram) == num_spec_steps`` - * ``sum(j * acceptance_histogram[j]) == total_accepted_draft_tokens`` - * ``total_accepted_draft_tokens <= total_draft_tokens`` - """ - - acceptance_rate: float = Field( - description="Accepted draft tokens divided by proposed draft tokens.") - total_accepted_draft_tokens: int = Field( - description="Draft tokens accepted across the request, excluding the " - "always-accepted bonus token.") - total_draft_tokens: int = Field( - description="Draft tokens proposed across the request. For tree " - "drafting this counts paths, not tree nodes, mirroring the " - "getMaxDraftPathLen clamp in updateNumTokensPerIteration.") - num_spec_steps: int = Field( - description="Verify steps performed for the request. Equals the sum of " - "acceptance_histogram.") - acceptance_histogram: List[int] = Field( - description="Dense histogram indexed by accepted-draft count: entry j " - "is the number of verify steps that accepted exactly j draft tokens. " - "Tree-agnostic -- it records output lengths per step and encodes no " - "parent/child structure.") - num_spec_tokens: Optional[int] = Field( - default=None, - description="Maximum draft length per step, when the run has a fixed " - "bound. None under draft_len_schedule, where the bound varies by batch " - "size.") - - -class _OmitsAbsentSpecDecodeStats(OpenAIBaseModel): - """Drops ``speculative_decoding`` from serialized output when it is absent. - - Per-request spec-decode stats are off by default, so without this every - response on the paths that serialize with a plain ``model_dump()`` -- the - non-streaming chat and completions responses, and the completions stream - (``exclude_unset=False``) -- would gain ``"speculative_decoding": null`` for - every user, whether or not they enabled ``per_request_spec_decode_stats``. - - Scoped to this one field on purpose. Blanket ``exclude_none`` would also - strip unrelated optional fields that clients may rely on being present, and - the dump calls are spread across the serving layer rather than funnelled - through one place where an ``exclude=`` argument could be applied. - """ - - @model_serializer(mode="wrap") - def _omit_absent_spec_decode_stats(self, handler: Any) -> Any: - data = handler(self) - if isinstance(data, dict) and data.get("speculative_decoding", - ...) is None: - data.pop("speculative_decoding", None) - return data - - -class StreamOptions(OpenAIBaseModel): - include_usage: Optional[bool] = True - continuous_usage_stats: Optional[bool] = False - - -class PromptTokensDetails(OpenAIBaseModel): - cached_tokens: int = 0 - - -class UsageInfo(OpenAIBaseModel): - prompt_tokens: int = 0 - total_tokens: int = 0 - completion_tokens: Optional[int] = 0 - prompt_tokens_details: Optional[PromptTokensDetails] = None - - -class ModelCard(OpenAIBaseModel): - id: str - object: str = "model" - created: int = Field(default_factory=lambda: int(time.time())) - owned_by: str = "tensorrt_llm" - - -class ModelList(OpenAIBaseModel): - object: str = "list" - data: List[ModelCard] = Field(default_factory=list) - - -class TokenizeRequest(OpenAIBaseModel): - """Request body for the ``POST /_internal/tokenize`` endpoint. - - Asks the server to encode some text into the token-id sequence its - loaded tokenizer would produce, without running any generation. - - Only plain-text ``prompt`` tokenization is supported. A ``messages`` - path was intentionally left out: to be meaningful it must reuse the full - ``/v1/chat/completions`` preprocessing (server chat template, tools, - documents, template kwargs, model-type resolution, multimodal handling), - otherwise the returned token ids would silently diverge from the actual - inference prompt. It can be added once that path is shared. - - Fields: - model: Target model name. Accepted for compatibility with - multi-model routers but ignored -- the server always tokenizes - against its own loaded model. - prompt: Plain text to tokenize directly via ``tokenizer.encode``. - """ - - model: Optional[str] = None - prompt: str - - -class TokenizeResponse(OpenAIBaseModel): - """Response body for the ``POST /_internal/tokenize`` endpoint. - - Fields: - count: Number of tokens in the encoded sequence. Always populated; - clients needing only the length can ignore ``tokens``. - tokens: The encoded token ids. Kept ``Optional`` so future - count-only variants can omit it without breaking the schema. - """ - - count: int - tokens: Optional[List[int]] = None - - -class ResponseFormat(OpenAIBaseModel): - type: Literal["text", "json", "json_schema", "json_object", "regex", "ebnf", - "structural_tag"] - schema: Optional[dict] = None - json_schema: Optional[dict] = None - regex: Optional[str] = None - ebnf: Optional[str] = None - format: Optional[xgrammar.structural_tag.Format] = None - - -class DisaggregatedParams(OpenAIBaseModel): - request_type: str - first_gen_tokens: Optional[List[int]] = None - first_gen_log_probs: Optional[List] = None - first_gen_logits: Optional[List] = None - ctx_request_id: Optional[int] = None - encoded_opaque_state: Optional[str] = None - draft_tokens: Optional[List[int]] = None - disagg_request_id: Optional[int] = None - ctx_dp_rank: Optional[int] = None - ctx_info_endpoint: Optional[str] = None - schedule_style: Optional[DisaggScheduleStyle] = None - ctx_usage: Optional[UsageInfo] = None - remote_tail_start: Optional[NonNegativeInt] = Field( - default=None, - description= - "Prompt-token boundary transferred by a remote-tail context worker.") - multimodal_positions: Optional[List[NonNegativeInt]] = Field( - default=None, - description="Prompt-token offset of each multimodal item.") - multimodal_lengths: Optional[List[PositiveInt]] = Field( - default=None, description="Prompt-token count of each multimodal item.") - # TODO(TRTLLM-12407): Multimodal E/PD over trtllm-serve needs these protocol fields too: - # encoder embedding handles, multimodal hashes, and optional mRoPE handles. - # Add them here and in to_disaggregated_params()/to_llm_disaggregated_params() - # before routing MM encoder -> context -> generation through OpenAI protocol. - # Orchestrator -> context-worker instruction: return prompt_token_ids as a - # base64 int32 buffer (prompt_token_ids_b64) instead of a JSON int array. - return_prompt_token_ids_b64: bool = False - # Context worker -> generation worker: the reasoning mode the context - # worker read off the prompt it rendered. The generation worker only sees - # prompt_token_ids, so it cannot resolve this for itself. - resolved_thinking: Optional[bool] = None - # Orchestrator -> generation worker: seconds since arrival at the orchestrator. - orig_elapsed_s: Optional[float] = None - - @model_validator(mode="after") - def validate_multimodal_spans(self) -> "DisaggregatedParams": - if (self.multimodal_positions is None) != (self.multimodal_lengths - is None): - raise ValueError( - "multimodal_positions and multimodal_lengths must be provided together" - ) - if (self.multimodal_positions is not None and len( - self.multimodal_positions) != len(self.multimodal_lengths)): - raise ValueError( - "multimodal_positions and multimodal_lengths must have the same length" - ) - return self - - -class ConversationParams(OpenAIBaseModel): - model_config = ConfigDict(extra="forbid", - populate_by_name=True, - validate_assignment=True) - - conversation_id: str = Field( - description=("Stable multi-turn conversation id used for routing"), ) - - # Forwarded as a header; excluded from the request body for older workers. - subagent_affinity_id: Optional[str] = Field( - default=None, - exclude=True, - description=("Server-private parent-session id for sub-agent routing " - "affinity; forwarded to workers as a header, never in the " - "body.")) - - @field_validator("conversation_id", mode="before") - @classmethod - def validate_conversation_id(cls, value: Any) -> str: - if value is None: - raise ValueError("conversation_id must be non-empty") - conversation_id = str(value).strip() - if not conversation_id: - raise ValueError("conversation_id must be non-empty") - return conversation_id - - -class ErrorResponse(OpenAIBaseModel): - object: str = "error" - message: str - type: str - param: Optional[str] = None - # HTTP status by default; a machine-readable string for conditions that - # define one (e.g. "context_length_exceeded", matching OpenAI). - code: Union[int, str] - - -class CompletionLogProbs(OpenAIBaseModel): - text_offset: List[int] = Field(default_factory=list) - token_logprobs: List[Optional[float]] = Field(default_factory=list) - tokens: List[str] = Field(default_factory=list) - top_logprobs: List[Optional[Dict[str, float]]] = Field(default_factory=list) - - -class CompletionResponseChoice(_OmitsAbsentSpecDecodeStats): - index: int - text: str - token_ids: Optional[List[int]] = None - logprobs: Optional[CompletionLogProbs] = None - context_logits: Optional[Union[List[float], List[List[ - float]]]] = None # For reward models, the output is score logits instead of text. - finish_reason: Optional[str] = None - stop_reason: Optional[Union[int, str]] = Field( - default=None, - description=( - "The stop string or token id that caused the completion " - "to stop, None if the completion finished for some other reason " - "including encountering the EOS token"), - ) - disaggregated_params: Optional[DisaggregatedParams] = Field(default=None) - avg_decoded_tokens_per_iter: Optional[float] = Field(default=None) - speculative_decoding: Optional[SpeculativeDecodingStats] = Field( - default=None) - - -class CompletionResponse(OpenAIBaseModel): - id: str = Field(default_factory=lambda: f"cmpl-{str(uuid.uuid4().hex)}") - object: str = "text_completion" - created: int = Field(default_factory=lambda: int(time.time())) - model: str - choices: List[CompletionResponseChoice] - usage: UsageInfo - # Add prompt_tokens_ids to the response to remove the tokenization - # in the generation server in disaggreated serving - prompt_token_ids: Optional[Union[List[List[int]], List[int]]] = None - - -class CompletionResponseStreamChoice(_OmitsAbsentSpecDecodeStats): - index: int - text: str - token_ids: Optional[List[int]] = None - logprobs: Optional[CompletionLogProbs] = None - finish_reason: Optional[str] = None - stop_reason: Optional[Union[int, str]] = Field( - default=None, - description=( - "The stop string or token id that caused the completion " - "to stop, None if the completion finished for some other reason " - "including encountering the EOS token"), - ) - avg_decoded_tokens_per_iter: Optional[float] = Field(default=None) - speculative_decoding: Optional[SpeculativeDecodingStats] = Field( - default=None) - - -class CompletionStreamResponse(OpenAIBaseModel): - id: str = Field(default_factory=lambda: f"cmpl-{str(uuid.uuid4().hex)}") - object: str = "text_completion" - created: int = Field(default_factory=lambda: int(time.time())) - model: str - choices: List[CompletionResponseStreamChoice] - usage: Optional[UsageInfo] = Field(default=None) - - -class EmbeddingRequest(OpenAIBaseModel): - # OpenAI-compatible embeddings request. - # https://platform.openai.com/docs/api-reference/embeddings/create - model: str = Field(description="The model to use for the embedding.") - input: Union[str, List[str], List[int], List[List[int]]] = Field( - description="Text(s) or pre-tokenized token-id list(s) to embed. Accepts " - "the four OpenAI forms: str, list[str], list[int], list[list[int]].") - encoding_format: Literal["float", "base64"] = Field( - default="float", - description="Format of the returned embedding: a JSON array of floats " - "(\"float\") or a base64-encoded string of packed little-endian float32 " - "values (\"base64\").") - dimensions: Optional[PositiveInt] = Field( - default=None, - description="Number of dimensions the resulting embedding should have. " - "Only supported by Matryoshka-trained text-embedding models; rejected " - "otherwise.") - user: Optional[str] = Field( - default=None, - description="A stable identifier for the end-user (OpenAI-compatible; " - "accepted but unused).") - add_special_tokens: bool = Field( - default=True, - description= - "TRT-LLM extension: whether to add the model's special tokens " - "(e.g. [CLS]/[SEP]) when tokenizing string inputs. Default matches " - "llm.encode(); encoder models such as BERT generally need them.") - - -class EmbeddingResponseData(OpenAIBaseModel): - index: int = Field( - description="Index of this embedding in the request's input list.") - object: Literal["embedding"] = "embedding" - embedding: Union[List[float], str] = Field( - description= - "The embedding: a list of floats for encoding_format=\"float\", " - "or a base64 string of packed little-endian float32 values for " - "encoding_format=\"base64\".") - - -class EmbeddingUsageInfo(OpenAIBaseModel): - """Token usage for embeddings. - - OpenAI's embeddings `usage` object contains only `prompt_tokens` and - `total_tokens` (no `completion_tokens`), unlike chat/completions. - """ - prompt_tokens: int = Field(default=0, - description="Number of tokens in the input.") - total_tokens: int = Field( - default=0, - description="Total tokens consumed (equal to prompt_tokens for " - "embeddings).") - - -class EmbeddingResponse(OpenAIBaseModel): - id: str = Field(default_factory=lambda: f"embd-{str(uuid.uuid4().hex)}") - object: Literal["list"] = "list" - created: int = Field(default_factory=lambda: int(time.time())) - model: str = Field(description="The model used for the embedding.") - data: List[EmbeddingResponseData] = Field( - description="The list of embeddings, one per input item.") - usage: EmbeddingUsageInfo = Field( - description="Token usage for the request.") - - -def _response_format_to_guided_decoding_params( - response_format: Optional[ResponseFormat], - reasoning_parser: Optional[str] = None, - chat_template_kwargs: Optional[Dict[str, Any]] = None, -) -> Optional[GuidedDecodingParams]: - if response_format is None: - guided_decoding_params = None - elif response_format.type == "text": - guided_decoding_params = None - elif response_format.type == "json": - if response_format.schema is None: - raise ValueError( - f"response_format.schema is required for response_format.type == {response_format.type!r}, but got None." - ) - guided_decoding_params = GuidedDecodingParams( - json=response_format.schema) - elif response_format.type == "json_schema": - if response_format.json_schema is None: - raise ValueError( - f"response_format.json_schema is required for response_format.type == {response_format.type!r}, but got None." - ) - # OpenAI API spec wraps the actual JSON schema under a "schema" key. - # Extract the actual schema if the wrapper format is used. - json_schema = response_format.json_schema - if isinstance(json_schema, dict) and "schema" in json_schema: - json_schema = json_schema["schema"] - guided_decoding_params = GuidedDecodingParams(json=json_schema) - elif response_format.type == "json_object": - guided_decoding_params = GuidedDecodingParams(json_object=True) - elif response_format.type == "regex": - if response_format.regex is None: - raise ValueError( - f"response_format.regex is required for response_format.type == {response_format.type!r}, but got None." - ) - guided_decoding_params = GuidedDecodingParams( - regex=response_format.regex) - elif response_format.type == "ebnf": - if response_format.ebnf is None: - raise ValueError( - f"response_format.ebnf is required for response_format.type == {response_format.type!r}, but got None." - ) - guided_decoding_params = GuidedDecodingParams( - grammar=response_format.ebnf) - elif response_format.type == "structural_tag": - guided_decoding_params = GuidedDecodingParams( - structural_tag=response_format.model_dump_json(by_alias=True, - exclude_none=True)) - else: - raise ValueError(f"Unsupported response format: {response_format.type}") - - if guided_decoding_params is None or reasoning_parser is None: - return guided_decoding_params - - if guided_decoding_params.structural_tag is not None: - return guided_decoding_params - - # Adapt guided_decoding_params for reasoning parser - if guided_decoding_params.json is not None: - content = { - "type": "json_schema", - "json_schema": guided_decoding_params.json - } - elif guided_decoding_params.json_object: - content = {"type": "json_schema", "json_schema": {"type": "object"}} - elif guided_decoding_params.regex is not None: - content = {"type": "regex", "pattern": guided_decoding_params.regex} - elif guided_decoding_params.grammar is not None: - content = {"type": "grammar", "grammar": guided_decoding_params.grammar} - - extension_format = structured_output_format_for(reasoning_parser) - if extension_format is not None: - # A registered per-model serving extension owns the placement of the - # constraint relative to this model's reasoning markup. - stag_format = extension_format(content, chat_template_kwargs) - if stag_format is None: - return guided_decoding_params - else: - # Force thinking and then trigger user constraint - parser = ReasoningParserFactory.create_reasoning_parser( - reasoning_parser) - stag_format = { - "type": - "sequence", - "elements": [ - { - "type": "tag", - "begin": parser.reasoning_start, - "content": { - "type": "any_text" - }, - "end": parser.reasoning_end, - }, - content, - ], - } - - stag_format = ResponseFormat(type="structural_tag", format=stag_format) - return GuidedDecodingParams(structural_tag=stag_format.model_dump_json( - by_alias=True, exclude_none=True)) - - -def _response_format_text_config_to_guided_decoding_params( - text_format: Optional[ResponseFormatTextConfig], - reasoning_parser: Optional[str] = None, -) -> Optional[GuidedDecodingParams]: - if text_format is None: - return None - - resp_format = ResponseFormat(type=text_format.type, - json_schema=getattr(text_format, "schema_", - None)) - return _response_format_to_guided_decoding_params( - resp_format, reasoning_parser=reasoning_parser) - - -def _record_sampling_params_request_fields( - request: OpenAIBaseModel, sampling_params: SamplingParams) -> None: - """Preserve explicitly supplied fields across protocol defaulting.""" - sampling_params._set_request_provided_fields( - field_name for field_name in request.model_fields_set - if getattr(request, field_name, None) is not None) - - -class CompletionRequest(OpenAIBaseModel): - kv_hints: Optional[KvHintsEnvelope] = Field( - default=None, - repr=False, - description= - "Prototype KV hint metadata; transported without executing actions.", - ) - # Ordered by official OpenAI API documentation - # https://platform.openai.com/docs/api-reference/completions/create - model: str - prompt: Union[List[int], List[List[int]], str, List[str]] - best_of: Optional[int] = None - echo: Optional[bool] = False - frequency_penalty: Optional[float] = 0.0 - logit_bias: Optional[Dict[str, float]] = None - logprobs: Optional[int] = None - max_tokens: Optional[int] = None - n: int = 1 - presence_penalty: Optional[float] = 0.0 - seed: Optional[int] = Field(default=None) - stop: Optional[Union[str, List[str]]] = Field(default_factory=list) - stream: Optional[bool] = False - stream_options: Optional[StreamOptions] = None - suffix: Optional[str] = None - temperature: Optional[float] = None - top_p: Optional[float] = None - user: Optional[str] = None - lora_request: Optional[LoRARequest] = None - prompt_ignore_length: Optional[int] = 0 - - # doc: begin-completion-sampling-params - use_beam_search: bool = False - top_k: int = 0 - top_p_min: float = 0.0 - min_p: float = 0.0 - repetition_penalty: float = 1.0 - # Unset by default so the engine picks its own beam-search defaults, as it - # does for requests coming through the Python API. Sending concrete values - # here would put every served request on the length-normalized, exhaustive - # beam-search path. - length_penalty: Optional[float] = None - early_stopping: Optional[EarlyStopping] = None - stop_token_ids: Optional[List[int]] = Field(default_factory=list) - include_stop_str_in_output: bool = False - ignore_eos: bool = False - min_tokens: int = 0 - skip_special_tokens: bool = True - spaces_between_special_tokens: bool = True - truncate_prompt_tokens: Optional[Annotated[int, Field(ge=1)]] = None - return_context_logits: bool = False - detokenize: bool = True - thinking_token_budget: Optional[int] = None - # doc: end-completion-sampling-params - - # doc: begin-completion-extra-params - add_special_tokens: bool = Field( - default=True, - description=( - "If true (the default), special tokens (e.g. BOS) will be added to " - "the prompt."), - ) - response_format: Optional[ResponseFormat] = Field( - default=None, - description= - ("Similar to chat completion, this parameter specifies the format of output. " - "{'type': 'text'}, {'type': 'json'}, {'type': 'json_object'}, {'type': 'regex'}, " - "{'type': 'ebnf'}, {'type': 'structural_tag'} are supported."), - ) - - disaggregated_params: Optional[DisaggregatedParams] = Field( - default=None, - description=("Parameters for disaggregated serving"), - ) - priority: Optional[float] = Field( - default=None, - ge=0.0, - le=1.0, - description= - ("Scheduling priority in [0.0, 1.0]; higher is served first. Only honored " - "when the engine runs with scheduler_config.waiting_queue_policy=priority. " - "If unset, the engine default (0.5) is used."), - ) - conversation_params: Optional[ConversationParams] = Field( - default=None, - description=("Parameters for multi-turn conversation routing"), - ) - - # doc: end-completion-extra-params - - def to_sampling_params(self, - vocab_size: Optional[int] = None, - gather_generation_logits: bool = False, - backend: Optional[str] = None) -> SamplingParams: - sampling_logprobs = None - if self.logprobs: - if backend == "pytorch" or gather_generation_logits: - sampling_logprobs = self.logprobs - elif self.logprobs > 1: - raise ValueError( - "`logprobs` must be 1 or `gather_generation_logits` must be `True` to use `logprobs` > 1" - ) - - sampling_params = SamplingParams( - best_of=self.best_of, - frequency_penalty=self.frequency_penalty, - logprobs=sampling_logprobs, - max_tokens=self.max_tokens, - n=self.n, - presence_penalty=self.presence_penalty, - seed=self.seed, - stop=self.stop, - temperature=(self.temperature - if self.temperature is not None else 1.0), - top_p=(self.top_p if self.top_p is not None else 1.0), - prompt_ignore_length=self.prompt_ignore_length, - - # completion-sampling-params - use_beam_search=self.use_beam_search, - top_k=self.top_k, - top_p_min=self.top_p_min if self.top_p_min > 0 else None, - min_p=self.min_p, - repetition_penalty=self.repetition_penalty, - length_penalty=self.length_penalty, - early_stopping=_early_stopping_to_int(self.early_stopping), - stop_token_ids=self.stop_token_ids, - include_stop_str_in_output=self.include_stop_str_in_output, - ignore_eos=self.ignore_eos, - min_tokens=self.min_tokens, - skip_special_tokens=self.skip_special_tokens, - spaces_between_special_tokens=self.spaces_between_special_tokens, - truncate_prompt_tokens=self.truncate_prompt_tokens, - return_context_logits=self.return_context_logits, - guided_decoding=_response_format_to_guided_decoding_params( - self.response_format), - detokenize=self.detokenize, - thinking_token_budget=self.thinking_token_budget, - - # logits_bias - embedding_bias=_logit_bias_to_embedding_bias( - self.logit_bias, vocab_size), - - # completion-extra-params - add_special_tokens=self.add_special_tokens, - ) - _record_sampling_params_request_fields(self, sampling_params) - return sampling_params - - @model_validator(mode="before") - @classmethod - def check_logprobs(cls, data): - check_logprobs_limit("logprobs", data.get("logprobs")) - return data - - @field_validator("thinking_token_budget", mode="before") - @classmethod - def check_thinking_token_budget(cls, value): - return validate_thinking_token_budget(value) - - @model_validator(mode="before") - @classmethod - def validate_stream_options(cls, data): - if data.get("stream_options") and not data.get("stream"): - raise ValueError( - "Stream options can only be defined when stream is true.") - return data - - @model_validator(mode="before") - @classmethod - def check_suffix(cls, data): - if data.get("suffix"): - raise ValueError("suffix is not supported") - return data - - -class FunctionCall(OpenAIBaseModel): - name: str - arguments: str - - -class DeltaFunctionCall(OpenAIBaseModel): - name: Optional[str] = None - arguments: Optional[str] = None - - -class ToolCall(OpenAIBaseModel): - id: str = Field( - default_factory=lambda: f"chatcmpl-tool-{str(uuid.uuid4().hex)}") - type: Literal["function"] = "function" - function: FunctionCall - - -class DeltaToolCall(OpenAIBaseModel): - id: Optional[str] = None - type: Literal["function"] = "function" - index: int - function: Optional[DeltaFunctionCall] = None - - -class ChatMessage(OpenAIBaseModel): - role: str - content: Optional[str] = None - reasoning_content: Optional[str] = None - reasoning: Optional[str] = None - tool_calls: List[ToolCall] = Field(default_factory=list) - - -class ChatCompletionLogProb(OpenAIBaseModel): - token: str - logprob: float = -9999.0 - bytes: Optional[List[int]] = None - - -class ChatCompletionLogProbsContent(ChatCompletionLogProb): - top_logprobs: Optional[List[ChatCompletionLogProb]] = None - - -class CustomChatCompletionContentPartParam(TypedDict, total=False): - __pydantic_config__ = ConfigDict(extra="allow") # type: ignore - - type: Required[str] - """The type of the content part.""" - - -ChatCompletionContentPartParam = Union[OpenAIChatCompletionContentPartParam, - CustomChatCompletionContentPartParam] - - -class CustomChatCompletionMessageParam(TypedDict, total=False): - """Enables custom roles in the Chat Completion API.""" - - # This is so custom fields not in any of the `ChatCompletionMessageParam` defined by OpenAI - # are still allowed. - # Examples include: assistant messages with `reasoning` / `reasoning_content`. - __pydantic_config__ = ConfigDict(extra="allow") # type: ignore - - role: Required[str] - """The role of the message's author.""" - - content: Union[str, List[ChatCompletionContentPartParam], None] - """The contents of the message.""" - - name: str - """An optional name for the participant. - - Provides the model information to differentiate between participants of the - same role. - """ - - -class ReasoningAssistantMessage(ChatCompletionAssistantMessageParam): - """Assistant message that includes reasoning tokens.""" - reasoning: Optional[str] - # NOTE: some older benchmarks and chat templates assume the below, which has been deprecated - # in other inference frameworks in favor of the above `reasoning` field. - reasoning_content: Optional[str] - - -class DynamicToolsSystemMessageParam(TypedDict, total=False): - """System message carrying message-level (dynamic) tool declarations. - - Kimi-style templates render such messages as an in-conversation tool - declare block. Must come first in `ChatCompletionMessageParam`: the - stock OpenAI system-message TypedDict otherwise wins smart-union scoring - and silently drops the `tools` key. - """ - __pydantic_config__ = ConfigDict(extra="allow") # type: ignore - - role: Required[Literal["system"]] - tools: Required[List[dict]] - content: Union[str, List[ChatCompletionContentPartParam], None] - name: str - - -ChatCompletionMessageParam = Union[DynamicToolsSystemMessageParam, - OpenAIChatCompletionMessageParam, - CustomChatCompletionMessageParam, - ReasoningAssistantMessage] - - -class ChatCompletionLogProbs(OpenAIBaseModel): - content: Optional[List[ChatCompletionLogProbsContent]] = None - - -class ChatCompletionResponseChoice(_OmitsAbsentSpecDecodeStats): - index: int - message: ChatMessage - logprobs: Optional[ChatCompletionLogProbs] = None - finish_reason: Optional[str] = None - stop_reason: Optional[Union[int, str]] = None - # TODO: progressivly add more info like input_ids, specific_token_ids, mrope, mm_hashes, etc - # TODO: and use a JSON-safe handle to refer to the server-side output - mm_embedding_handle: Optional[Dict[str, Any]] = None - - disaggregated_params: Optional[DisaggregatedParams] = Field(default=None) - avg_decoded_tokens_per_iter: Optional[float] = Field(default=None) - speculative_decoding: Optional[SpeculativeDecodingStats] = Field( - default=None) - - -class ChatCompletionResponse(OpenAIBaseModel): - id: str = Field(default_factory=lambda: f"chatcmpl-{str(uuid.uuid4().hex)}") - object: Literal["chat.completion"] = "chat.completion" - created: int = Field(default_factory=lambda: int(time.time())) - model: str - choices: List[ChatCompletionResponseChoice] - usage: UsageInfo - # Add prompt_tokens_ids to the response to remove the tokenization - # in the generation server in disaggreated serving - prompt_token_ids: Optional[List[int]] = None - # base64 int32 buffer alternative to prompt_token_ids; set by the context - # worker so the orchestrator can relay a string instead of the int list. - prompt_token_ids_b64: Optional[str] = None - - -class DeltaMessage(OpenAIBaseModel): - role: Optional[str] = None - content: Optional[str] = None - reasoning_content: Optional[str] = None - # For GPT-OSS style reasoning - reasoning: Optional[str] = None - tool_calls: Optional[List[DeltaToolCall]] = None - - -class ChatCompletionResponseStreamChoice(_OmitsAbsentSpecDecodeStats): - index: int - delta: DeltaMessage - logprobs: Optional[ChatCompletionLogProbs] = None - finish_reason: Optional[str] = None - stop_reason: Optional[Union[int, str]] = None - avg_decoded_tokens_per_iter: Optional[float] = Field(default=None) - speculative_decoding: Optional[SpeculativeDecodingStats] = Field( - default=None) - - -class ChatCompletionStreamResponse(OpenAIBaseModel): - id: str = Field(default_factory=lambda: f"chatcmpl-{str(uuid.uuid4().hex)}") - object: Literal["chat.completion.chunk"] = "chat.completion.chunk" - created: int = Field(default_factory=lambda: int(time.time())) - model: str - choices: List[ChatCompletionResponseStreamChoice] - usage: Optional[UsageInfo] = Field(default=None) - - -# Maximum total number of `enum` values across all properties of one tool -# function's parameters schema. Mirrors the cap reference OpenAI-compatible -# platforms enforce; unbounded enums also inflate the guided-decoding -# grammar compiled for strict tool calls. -TOOL_PARAM_MAX_ENUM_VALUES = 1000 - -# JSON Schema keywords whose value is instance data rather than a subschema. -# An `enum` key nested inside these is a value (e.g. a default that happens to -# be `{"enum": [...]}`), not an enum constraint, so it must not count. -_SCHEMA_INSTANCE_KEYWORDS = frozenset({"default", "const", "examples"}) -# JSON Schema keywords whose value is a map of {name: subschema}. The names are -# user-controlled (a property can be literally named `default` or `enum`), so -# recurse into the values as subschemas without treating the names as keywords. -_SCHEMA_MAP_KEYWORDS = frozenset({ - "properties", "patternProperties", "$defs", "definitions", - "dependentSchemas" -}) - - -def _count_schema_enum_values(schema: Any) -> int: - """Recursively count `enum` *constraint* entries in a JSON-schema fragment. - - Only `enum` keywords in schema positions are counted. `enum` keys that are - instance data -- nested under `default`/`const`/`examples`, or the name of a - property -- are ignored, so a valid schema is not rejected for values that - are not enum constraints. - """ - count = 0 - if isinstance(schema, dict): - for key, value in schema.items(): - if key == "enum" and isinstance(value, list): - count += len(value) - elif key in _SCHEMA_INSTANCE_KEYWORDS: - continue - elif key in _SCHEMA_MAP_KEYWORDS and isinstance(value, dict): - for subschema in value.values(): - count += _count_schema_enum_values(subschema) - else: - count += _count_schema_enum_values(value) - elif isinstance(schema, list): - for item in schema: - count += _count_schema_enum_values(item) - return count - - -class FunctionDefinition(OpenAIBaseModel): - name: str - description: Optional[str] = None - parameters: Optional[Dict[str, Any]] = None - strict: Optional[bool] = None - - @model_validator(mode="after") - def check_enum_value_cap(self): - if self.parameters is not None: - num_enum_values = _count_schema_enum_values(self.parameters) - if num_enum_values > TOOL_PARAM_MAX_ENUM_VALUES: - raise ValueError( - f"tool function {self.name!r} declares {num_enum_values} " - f"enum values across its parameters schema; the maximum " - f"is {TOOL_PARAM_MAX_ENUM_VALUES}.") - return self - - -class ChatCompletionToolsParam(OpenAIBaseModel): - type: Literal["function"] = "function" - function: FunctionDefinition - - -class NamespacedFunctionToolParam(ChatCompletionToolsParam): - """A tool that a Responses request declared inside a namespace. - - ``function`` is what a chat template without namespaces is given: the - tool as a function named ``.``. That is all a dump - carries. The declaration itself is kept beside it, out of every dump, for - a model whose renderer reads namespaces (``ServingExtension.serialize_tool``). - """ - - namespace: str = Field(exclude=True) - namespace_description: Optional[str] = Field(default=None, exclude=True) - declared_name: str = Field(exclude=True) - declared_description: Optional[str] = Field(default=None, exclude=True) - - -class ChatCompletionNamedFunction(OpenAIBaseModel): - name: str - - -class ChatCompletionNamedToolChoiceParam(OpenAIBaseModel): - function: ChatCompletionNamedFunction - type: Literal["function"] = "function" - - -class ChatCompletionThinkingParam(OpenAIBaseModel): - """Kimi/Moonshot `thinking` extension controlling reasoning output. - - `keep` is fixed to `"all"` when thinking is enabled and ignored when - disabled; `effort` is only meaningful when enabled. An explicit - `effort` wins over a request-level `reasoning_effort`, which applies - only when `effort` is absent (and never for a disabled request). - """ - type: Literal["enabled", "disabled"] = "enabled" - keep: Optional[Literal["all"]] = None - effort: Optional[Literal["low", "high", "max"]] = None - - -class ChatCompletionRequest(OpenAIBaseModel): - kv_hints: Optional[KvHintsEnvelope] = Field( - default=None, - repr=False, - description= - "Prototype KV hint metadata; transported without executing actions.", - ) - # Ordered by official OpenAI API documentation - # https://platform.openai.com/docs/api-reference/chat/create - messages: List[ChatCompletionMessageParam] - # Add prompt_tokens_ids to the request to remove the tokenization - # in the generation server in disaggreated serving - prompt_token_ids: Optional[List[int]] = None - # base64 int32 buffer relayed by the orchestrator from the ctx response; - # decoded back to prompt_token_ids on the generation worker. Not for clients. - prompt_token_ids_b64: Optional[str] = None - model: str - frequency_penalty: Optional[float] = 0.0 - logit_bias: Optional[Dict[str, float]] = None - logprobs: Optional[bool] = False - top_logprobs: Optional[int] = 0 - max_completion_tokens: Optional[int] = Field(default=None, - validation_alias='max_tokens') - n: int = 1 - presence_penalty: Optional[float] = 0.0 - response_format: Optional[ResponseFormat] = None - seed: Optional[int] = Field(None) - service_tier: Optional[Literal["auto", "default", "flex", "scale", - "priority", "fast"]] = Field( - default=None, - description=( - "Accepted for OpenAI compatibility " - "and ignored: the server has no " - "service tiers. Scheduling priority " - "is the separate `priority` field.")) - stop: Optional[Union[str, List[str]]] = Field(default_factory=list) - # Standard OpenAI field, accepted for compatibility and otherwise ignored: - # the server keeps no copy of a chat completion, so there is nothing to - # store or to leave unstored. - store: Optional[bool] = None - stream: Optional[bool] = False - stream_options: Optional[StreamOptions] = None - temperature: Optional[float] = None - top_p: Optional[float] = None - tools: Optional[List[ChatCompletionToolsParam]] = None - tool_choice: Optional[Union[Literal["none", "auto", "required"], - ChatCompletionNamedToolChoiceParam]] = "none" - # Standard OpenAI field, accepted for compatibility. `false` is not - # enforced: the engine does not restrict how many tool calls the model - # emits per turn, so parallel emission remains model behavior either way. - parallel_tool_calls: Optional[bool] = None - user: Optional[str] = None - reasoning_effort: Optional[ReasoningEffort | Literal[ - "low", "medium", "high", "max", "none"]] = Field( - default=ReasoningEffort.LOW, - description=( - "The level of reasoning effort to use. Controls how much " - "reasoning is shown in the model's response. Options: " - "'low', 'medium', 'high' (harmony/gpt-oss), plus 'max' and " - "'none' for models with Kimi-style thinking control."), - ) - # Kimi/Moonshot extension: structured control of reasoning output. The - # serving layer maps it into the chat-template kwargs for models whose - # template understands it (e.g. kimi_k3); other models ignore it. - thinking: Optional[ChatCompletionThinkingParam] = None - thinking_token_budget: Optional[int] = None - prompt_ignore_length: Optional[int] = 0 - # Match Codex CLI's string-map contract; widen only for a client that needs - # other value types, even though the metadata is otherwise ignored. - client_metadata: Optional[Dict[str, str]] = Field( - default=None, - description=("Not an OpenAI field: opaque per-request metadata that " - "agent clients attach, a map of strings to strings. " - "Accepted and ignored.")) - - # doc: begin-chat-completion-sampling-params - best_of: Optional[int] = None - use_beam_search: bool = False - top_k: int = 0 - top_p_min: float = 0.0 - min_p: float = 0.0 - repetition_penalty: float = 1.0 - # Unset by default so the engine picks its own beam-search defaults, as it - # does for requests coming through the Python API. Sending concrete values - # here would put every served request on the length-normalized, exhaustive - # beam-search path. - length_penalty: Optional[float] = None - early_stopping: Optional[EarlyStopping] = None - stop_token_ids: Optional[List[int]] = Field(default_factory=list) - include_stop_str_in_output: bool = False - ignore_eos: bool = False - min_tokens: int = 0 - skip_special_tokens: bool = True - spaces_between_special_tokens: bool = True - truncate_prompt_tokens: Optional[Annotated[int, Field(ge=1)]] = None - lora_request: Optional[LoRARequest] = None - # doc: end-chat-completion-sampling-params - - # doc: begin-chat-completion-extra-params - echo: bool = Field( - default=False, - description=( - "If true, the new message will be prepended with the last message " - "if they belong to the same role."), - ) - add_generation_prompt: bool = Field( - default=True, - description= - ("If true, the generation prompt will be added to the chat template. " - "This is a parameter used by chat template in tokenizer config of the " - "model."), - ) - add_special_tokens: bool = Field( - default=False, - description=( - "If true, special tokens (e.g. BOS) will be added to the prompt " - "on top of what is added by the chat template. " - "For most models, the chat template takes care of adding the " - "special tokens so this should be set to false (as is the " - "default)."), - ) - documents: Optional[List[Dict[str, str]]] = Field( - default=None, - description= - ("A list of dicts representing documents that will be accessible to " - "the model if it is performing RAG (retrieval-augmented generation)." - " If the template does not support RAG, this argument will have no " - "effect. We recommend that each document should be a dict containing " - "\"title\" and \"text\" keys."), - ) - chat_template: Optional[str] = Field( - default=None, - description=( - "A Jinja template to use for this conversion. " - "If this is not passed, the model's default chat template will be " - "used instead."), - ) - chat_template_kwargs: Optional[Dict[str, Any]] = Field( - default=None, - description=("Additional kwargs to pass to the template renderer. " - "Will be accessible by the chat template."), - ) - injected_chat_template_kwargs: Optional[List[str]] = Field( - default=None, - description=( - "Server-populated. Keys of `chat_template_kwargs` that the server " - "derived from API-level fields (for example the Anthropic " - "`thinking` and `context_management` controls) rather than the " - "caller. The unused-kwargs guard exempts them, since the caller " - "cannot remove a control the server added. Relayed with the " - "request so a disaggregated worker rendering the prompt applies " - "the same exemption."), - ) - - media_io_kwargs: Optional[Dict[MediaModality, Dict[str, Any]]] = Field( - default=None, - description=( - "Per-request override for the server's `--media_io_kwargs`. " - "Shape: `{modality: {kwarg: value}}` with modality in " - "{\"image\", \"video\", \"audio\"}; unknown modality keys are " - "rejected. Per modality, request kwargs are shallow-merged " - "onto the server defaults (request wins per key). For " - "`video`, overriding only one of `fps`/`num_frames` drops " - "the other from the server default so the loader's built-in " - "is used. " - "Example: `{\"video\": {\"num_frames\": 32}}`."), - ) - - disaggregated_params: Optional[DisaggregatedParams] = Field( - default=None, - description=("Parameters for disaggregated serving"), - ) - conversation_params: Optional[ConversationParams] = Field( - default=None, - description=("Parameters for multi-turn conversation routing"), - ) - - cache_salt: Optional[str] = Field( - default=None, - description= - ("If specified, KV cache will be salted with the provided string " - "to limit the kv cache reuse on with the requests having the same string." - )) - priority: Optional[float] = Field( - default=None, - ge=0.0, - le=1.0, - description= - ("Scheduling priority in [0.0, 1.0]; higher is served first. Only honored " - "when the engine runs with scheduler_config.waiting_queue_policy=priority. " - "If unset, the engine default (0.5) is used."), - ) - - agent_hierarchy: Optional[AgentHierarchy] = Field( - default=None, description="Agent hierarchy ") - - mm_processor_kwargs: Optional[Dict[str, Any]] = Field( - default=None, - description= - "Per-request kwargs forwarded to the multimodal HF processor (e.g. num_frames for video models).", - ) - - # doc: end-chat-completion-extra-params - - def to_sampling_params(self, - vocab_size: Optional[int] = None, - gather_generation_logits: bool = False, - reasoning_parser: Optional[str] = None, - backend: Optional[str] = None) -> SamplingParams: - sampling_logprobs = None - if self.logprobs: - logprobs = 1 if not self.top_logprobs else self.top_logprobs - if backend == "pytorch" or gather_generation_logits: - sampling_logprobs = logprobs - elif self.top_logprobs: - raise ValueError( - "`gather_generation_logits` must be `True` to use `top_logprobs`" - ) - - sampling_params = SamplingParams( - frequency_penalty=self.frequency_penalty, - logprobs=sampling_logprobs, - max_tokens=self.max_completion_tokens, - n=self.n, - presence_penalty=self.presence_penalty, - seed=self.seed, - stop=self.stop, - temperature=(self.temperature - if self.temperature is not None else 1.0), - prompt_ignore_length=self.prompt_ignore_length, - - # chat-completion-sampling-params - best_of=self.best_of, - use_beam_search=self.use_beam_search, - top_k=self.top_k, - top_p=(self.top_p if self.top_p is not None else 1.0), - top_p_min=self.top_p_min if self.top_p_min > 0 else None, - min_p=self.min_p, - repetition_penalty=self.repetition_penalty, - length_penalty=self.length_penalty, - early_stopping=_early_stopping_to_int(self.early_stopping), - stop_token_ids=self.stop_token_ids, - include_stop_str_in_output=self.include_stop_str_in_output, - ignore_eos=self.ignore_eos, - min_tokens=self.min_tokens, - skip_special_tokens=self.skip_special_tokens, - spaces_between_special_tokens=self.spaces_between_special_tokens, - truncate_prompt_tokens=self.truncate_prompt_tokens, - guided_decoding=_response_format_to_guided_decoding_params( - self.response_format, - reasoning_parser=reasoning_parser, - chat_template_kwargs=self.chat_template_kwargs), - thinking_token_budget=self.thinking_token_budget, - - # logits_bias - embedding_bias=_logit_bias_to_embedding_bias( - self.logit_bias, vocab_size), - - # chat-completion-extra-params - add_special_tokens=self.add_special_tokens, - ) - _record_sampling_params_request_fields(self, sampling_params) - return sampling_params - - @model_validator(mode='before') - @classmethod - def validate_stream_options(cls, values): - if (values.get('stream_options') is not None - and not values.get('stream')): - raise ValueError("stream_options can only be set if stream is true") - return values - - @model_validator(mode="before") - @classmethod - def check_tool_choice(cls, data): - if not isinstance(data, dict): - return data - has_dynamic_tools = any( - isinstance(msg, dict) and msg.get("role") == "system" - and msg.get("tools") for msg in data.get("messages") or []) - if "tool_choice" not in data and (data.get("tools") - or has_dynamic_tools): - data["tool_choice"] = "auto" - # "none" and "auto" are meaningful without tools. "required" needs a - # non-empty tool set — request-level or message-level (dynamic) - # tools both count. A named function must be a request-level tool. - if "tool_choice" in data and data["tool_choice"] not in ("none", - "auto"): - satisfied = data.get("tools") or (data["tool_choice"] == "required" - and has_dynamic_tools) - if not satisfied: - raise ValueError( - "When using `tool_choice`, `tools` must be set.") - return data - - @model_validator(mode="before") - @classmethod - def check_message_tools_role(cls, data): - """Message-level tool declarations ride on system messages only. - - Union validation strips unknown keys from non-system messages, so - this raw-payload validator is the only layer that can reject the - misuse loudly instead of silently dropping the client's tools. The - rest of the dynamic-tools contract is model-specific and validated - in the serving layer (kimi_k3 only). - """ - if not isinstance(data, dict): - return data - for message in data.get("messages") or []: - if (isinstance(message, dict) and message.get("tools") is not None - and message.get("role") != "system"): - raise ValueError( - "Message-level `tools` are only allowed on system " - "messages.") - return data - - @model_validator(mode="before") - @classmethod - def check_logprobs(cls, data): - if (top_logprobs := data.get("top_logprobs")) is not None: - check_logprobs_limit("top_logprobs", top_logprobs) - if not data.get("logprobs"): - raise ValueError( - "logprobs must be true when using top_logprobs") - return data - - @field_validator("thinking_token_budget", mode="before") - @classmethod - def check_thinking_token_budget(cls, value): - return validate_thinking_token_budget(value) - - @model_validator(mode="before") - @classmethod - def check_suffix(cls, data): - if data.get("suffix"): - raise ValueError("suffix is not supported") - return data - - @field_validator("cache_salt") - @classmethod - def check_cache_salt_support(cls, v): - if v is not None: - if not isinstance(v, str) or not v.strip(): - raise ValueError( - "Parameter 'cache_salt' must be a non-empty string if provided." - ) - return v - - -class KVCacheTruncateRequest(OpenAIBaseModel): - model: str - messages: List[ChatCompletionMessageParam] = [] - messages_to_retain: List[ChatCompletionMessageParam] = [] - tools: Optional[List[ChatCompletionToolsParam]] = None - add_generation_prompt: Optional[bool] = True - documents: Optional[list] = None - chat_template: Optional[str] = None - chat_template_kwargs: Optional[dict] = None - reasoning_effort: Optional[str] = None - tool_choice: Optional[str] = None - - -class KVCacheTruncateTokensRequest(OpenAIBaseModel): - """Token-level analog of :class:`KVCacheTruncateRequest`. - - Posted to the ``/_control/kv_cache/truncate_tokens`` endpoint. Bypasses - the chat-template tokenization path entirely: callers (e.g. - trace replay) that already operate on raw token ids hand the server a - batch of full token-id sequences plus the prefix length to retain for - each, and the server trims the corresponding KV-cache blocks in the - radix tree. Each ``prefixes[i]`` becomes one - :class:`tensorrt_llm.executor.request.TruncateKVCacheRequest` on the - executor's KV-cache control queue, so one batched HTTP request fans out - into N atomic per-prefix radix-tree mutations while paying a single - network round-trip. - - Fields: - model: Target model name. Accepted for router compatibility but - ignored -- the server uses its own loaded model. - prefixes: Batch of full token-id sequences. Each entry is one - cached prefix whose KV-cache blocks should be truncated. - num_tokens_to_keep: Per-prefix number of leading tokens to retain, - parallel to ``prefixes`` (``num_tokens_to_keep[i]`` applies to - ``prefixes[i]``). Blocks beyond this length are dropped from the - radix tree. - """ - - model: Optional[str] = None - prefixes: List[List[int]] - num_tokens_to_keep: List[int] - - -# The trailing dict keeps a client's own item types from failing the request -# outright. Codex multi-agent sessions carry "agent_message" items, which no -# SDK model describes; without a permissive member the union rejects the whole -# input and every request from a spawned agent comes back 422, so agent -# collaboration cannot work at all. Known shapes still match their typed -# member first; see _response_output_item_to_chat_completion_message for how -# an unrecognised item is replayed. -ResponseInputOutputItem: TypeAlias = Union[ResponseInputItemParam, - ResponseReasoningItem, - ResponseFunctionToolCall, dict[str, - Any]] - -# Roles whose message items map to EasyInputMessageParam / Message, both of -# which forbid ``id``. An assistant message maps to ResponseOutputMessageParam -# instead, which *requires* ``id`` and ``status`` - stripping ``id`` there -# leaves the item matching no variant of the input union, so a conversation -# breaks as soon as it carries one assistant turn. -# -# Module level on purpose: pydantic turns a leading-underscore class attribute -# into a ModelPrivateAttr, so referring to it from a validator would raise -# "argument of type 'ModelPrivateAttr' is not iterable" at request time. -_ID_STRIPPED_ROLES = ("user", "system", "developer") - - -def _drop_explicit_nulls(value, _depth=0): - """Recursively drop dict keys whose value is an explicit ``null``. - - On the OpenAI wire an unset optional field is omitted, so an explicit - ``null`` carries no information — but the vendored Responses item types - validate ``null`` differently from omitted and reject it. Dropping - null-valued keys restores the omit-unset wire shape for clients that - serialize every unset optional as ``null`` (litellm among them). - """ - if _depth > 12: - return value - if isinstance(value, dict): - return { - k: _drop_explicit_nulls(v, _depth + 1) - for k, v in value.items() if v is not None - } - if isinstance(value, list): - return [_drop_explicit_nulls(v, _depth + 1) for v in value] - return value - - -def _materialize_validator_iterators(value, _depth=0): - """Recursively replace pydantic ValidatorIterator objects with lists. - - Matched by type name rather than by import: the class lives in the - pydantic_core extension module and is not part of its public API. - """ - if _depth > 12: - return value - if type(value).__name__ == "ValidatorIterator": - return [_materialize_validator_iterators(v, _depth + 1) for v in value] - if isinstance(value, dict): - for key, item in value.items(): - value[key] = _materialize_validator_iterators(item, _depth + 1) - return value - if isinstance(value, list): - return [_materialize_validator_iterators(v, _depth + 1) for v in value] - return value - - -class ResponsesRequest(OpenAIBaseModel): - kv_hints: Optional[KvHintsEnvelope] = Field( - default=None, - repr=False, - description= - "Prototype KV hint metadata; transported without executing actions.", - ) - # Ordered by official OpenAI API documentation - # https://platform.openai.com/docs/api-reference/responses/create - # - # Unlike the rest of the OpenAI models this one accepts unknown fields. - # Real Responses API clients attach evolving telemetry and routing keys - - # Codex CLI sends client_metadata and prompt_cache_key, for instance - and - # rejecting the whole request over a field the server would ignore anyway - # makes those clients unusable for no benefit. - model_config = ConfigDict(extra="allow", populate_by_name=True) - - background: Optional[bool] = False - include: Optional[list[ - Literal[ - "code_interpreter_call.outputs", - "computer_call_output.output.image_url", - "file_search_call.results", - "message.input_image.image_url", - "message.output_text.logprobs", - "reasoning.encrypted_content", - ], - ]] = None - input: Union[str, list[ResponseInputOutputItem]] - instructions: Optional[str] = None - - @field_validator("input", mode="before") - @classmethod - def _drop_unsupported_input_item_keys(cls, value): - """Strip per-item keys the vendored Responses item types reject. - - Clients echo back items exactly as the server emitted them, so input - message items arrive carrying the ``id`` the server assigned, which - ``EasyInputMessageParam`` / ``Message`` forbid. That fails the whole - request over a field with no prompt content. - - Assistant messages and tool-call items keep their ids: the former - needs it to validate, the latter uses it to pair calls with results. - """ - if not isinstance(value, list): - return value - - # Must run before the per-item shaping below: the item types reject - # explicit nulls that mean "unset" (see _drop_explicit_nulls). - value = [_drop_explicit_nulls(item) for item in value] - - def _with_annotations(part): - """output_text requires annotations; clients often omit it.""" - if (isinstance(part, dict) and part.get("type") == "output_text" - and "annotations" not in part): - return {**part, "annotations": []} - return part - - cleaned = [] - for item in value: - # A client may send a bare content part as a top-level item, not - # only nested inside a message. - item = _with_annotations(item) - if isinstance(item, dict) and item.get("type") in (None, "message"): - role = item.get("role") - if "id" in item and role in _ID_STRIPPED_ROLES: - item = {k: v for k, v in item.items() if k != "id"} - elif role == "assistant": - # ResponseOutputMessageParam requires status, and each - # output_text part requires annotations. Clients rebuild an - # assistant turn from the streamed deltas rather than from - # the server's content-part objects, so both routinely - # arrive absent and fail the whole request. Default them - # rather than reject a conversation over fields that carry - # no prompt content. - item = dict(item) - item.setdefault("status", "completed") - content = item.get("content") - if isinstance(content, list): - item["content"] = [ - { - **part, "annotations": part.get( - "annotations", []) - } if isinstance(part, dict) - and part.get("type") == "output_text" else part - for part in content - ] - cleaned.append(item) - return cleaned - - @field_validator("input", mode="after") - @classmethod - def _materialize_lazy_item_content(cls, value): - """Force lazily-validated sequences anywhere in ``input`` into lists. - - Several vendored Responses item types declare sequence fields as - ``Iterable[...]`` - ``ResponseOutputMessageParam.content`` and - ``ResponseOutputTextParam.annotations`` among them - and pydantic - validates an Iterable lazily into a ``ValidatorIterator``. That object - cannot be pickled, so a request carrying structured input dies with - - cannot pickle 'pydantic_core._pydantic_core.ValidatorIterator' - - the moment it is handed to a postprocess worker. It is also - single-consumption, so even with workers disabled anything that reads - the field twice sees an empty sequence the second time. - - The walk has to be recursive: the lazy fields are nested inside items, - e.g. input[1].content[0].annotations. - """ - return _materialize_validator_iterators(value) - - max_output_tokens: Optional[int] = None - max_tool_calls: Optional[int] = None - metadata: Optional[Metadata] = None - model: str - parallel_tool_calls: Optional[bool] = False - previous_response_id: Optional[str] = None - prompt: Optional[ResponsePrompt] = None - reasoning: Optional[Reasoning] = None - thinking_token_budget: Optional[int] = None - service_tier: Literal["auto", "default", "flex", "scale", - "priority"] = "auto" - store: Optional[bool] = True - stream: Optional[bool] = False - temperature: Optional[float] = None - text: Optional[ResponseTextConfig] = None - tool_choice: ToolChoice = "auto" - tools: list[Tool] = Field(default_factory=list) - top_logprobs: Optional[int] = 0 - top_p: Optional[float] = None - truncation: Optional[Literal["auto", "disabled"]] = "disabled" - user: Optional[str] = None - - request_id: str = Field( - default_factory=lambda: f"resp_{str(uuid.uuid4().hex)}", - description=( - "The request_id related to this request. If the caller does " - "not set it, a random_uuid will be generated. This id is used " - "through out the inference process and return in response."), - ) - - # doc: begin-responses-extra-params - # TensorRT-LLM extensions carrying the disaggregated-serving handoff between - # the orchestrator and the workers; clients do not set them. - disaggregated_params: Optional[DisaggregatedParams] = Field( - default=None, - description=("Parameters for disaggregated serving"), - ) - conversation_params: Optional[ConversationParams] = Field( - default=None, - description=("Parameters for multi-turn conversation routing"), - ) - # The context worker's tokenized prompt, so the generation worker does not - # render it again; the _b64 form is a base64 int32 buffer. - prompt_token_ids: Optional[List[int]] = None - prompt_token_ids_b64: Optional[str] = None - - # doc: end-responses-extra-params - - @field_validator("prompt_token_ids_b64") - @classmethod - def _check_prompt_token_ids_b64(cls, value: Optional[str]) -> Optional[str]: - # Decoded during routing and preprocessing, after validation; a bad - # value must be a 400 here rather than an internal error there. - if not value: - return value - try: - raw = base64.b64decode(value, validate=True) - except ValueError as error: - raise ValueError( - "prompt_token_ids_b64 is not valid base64") from error - if len(raw) % 4: - raise ValueError( - "prompt_token_ids_b64 is not a buffer of int32 token ids") - return value - - def relayed_prompt_token_ids(self) -> Optional[List[int]]: - """The prompt the orchestrator relayed from the context worker, if any. - - Only a generation-only handoff carries one. On any other request the - fields are ignored, so the prompt is rendered from the request and the - response echoes what was actually prefilled. - """ - params = self.disaggregated_params - if params is None or params.request_type != "generation_only": - return None - if self.prompt_token_ids is not None: - return self.prompt_token_ids - if not self.prompt_token_ids_b64: - return None - import numpy as np - return np.frombuffer(base64.b64decode(self.prompt_token_ids_b64), - dtype=np.int32).tolist() - - _DEFAULT_SAMPLING_PARAMS = { - "temperature": 1.0, - "top_p": 1.0, - } - - def to_sampling_params( - self, - default_sampling_params: Optional[dict] = None, - reasoning_parser: Optional[str] = None, - ) -> SamplingParams: - max_tokens = None - if self.max_output_tokens is not None: - max_tokens = self.max_output_tokens - - default_sampling_params = default_sampling_params or {} - if (temperature := self.temperature) is None: - temperature = default_sampling_params.get( - "temperature", self._DEFAULT_SAMPLING_PARAMS["temperature"]) - if (top_p := self.top_p) is None: - top_p = default_sampling_params.get( - "top_p", self._DEFAULT_SAMPLING_PARAMS["top_p"]) - stop_token_ids = default_sampling_params.get("stop_token_ids", None) - - # Structured output - guided_decoding = None - if self.text is not None and self.text.format is not None: - guided_decoding = _response_format_text_config_to_guided_decoding_params( - self.text.format, reasoning_parser=reasoning_parser) - - sampling_params = SamplingParams( - temperature=temperature, - top_p=top_p, - max_tokens=max_tokens, - logprobs=self.top_logprobs, - stop_token_ids=stop_token_ids, - guided_decoding=guided_decoding, - thinking_token_budget=self.thinking_token_budget, - ) - _record_sampling_params_request_fields(self, sampling_params) - return sampling_params - - @model_validator(mode="before") - @classmethod - def _null_top_level_fields_mean_unset(cls, data): - """Drop explicitly-null top-level fields so they take their defaults. - - Shallow on purpose: ``input`` items are scrubbed recursively by their - own validator, and the remaining structured fields are Optional - throughout. A null required field (``model``, ``input``) still - reports as missing. - """ - if not isinstance(data, dict): - return data - return {k: v for k, v in data.items() if v is not None} - - @model_validator(mode="before") - @classmethod - def hoist_additional_tools(cls, data: Any) -> Any: - """Move the tools of an ``additional_tools`` input item into ``tools``. - - Codex declares its tools that way; hoisted here, they are validated as - tools and resolve when the reply's calls are parsed. Malformed shapes - are left in place for validation to reject. - """ - if not isinstance(data, dict): - return data - items = data.get("input") - tools = data.get("tools") - if not isinstance(items, list) or not isinstance( - tools, (list, type(None))): - return data - hoisted = [] - kept = [] - for item in items: - if (isinstance(item, dict) - and item.get("type") == "additional_tools" - and isinstance(item.get("tools"), list)): - hoisted.extend(item["tools"]) - else: - kept.append(item) - if not hoisted: - return data - data = dict(data) - data["input"] = kept - data["tools"] = list(data.get("tools") or []) + hoisted - return data - - @model_validator(mode="before") - @classmethod - def validate_background(cls, data): - if not data.get("background"): - return data - if not data.get("store", True): - raise ValueError("background can only be used when `store` is true") - return data - - @model_validator(mode="before") - @classmethod - def validate_prompt(cls, data): - if data.get("prompt") is not None: - raise ValueError("prompt template is not supported") - return data - - @field_validator("thinking_token_budget", mode="before") - @classmethod - def check_thinking_token_budget(cls, value): - return validate_thinking_token_budget(value) - - -class InputTokensDetails(OpenAIBaseModel): - cached_tokens: int - # Required by the openai SDK models that re-validate this payload when it - # is embedded in a streaming event. Omitting it fails validation while the - # response is being streamed, which truncates the stream with no - # terminating event and leaves the client waiting forever. Prompt cache - # writes are not tracked separately, so this is reported as zero. - cache_write_tokens: int = 0 - - -class OutputTokensDetails(OpenAIBaseModel): - reasoning_tokens: int - - -class ResponseUsage(OpenAIBaseModel): - input_tokens: int - input_tokens_details: InputTokensDetails - output_tokens: int - output_tokens_details: OutputTokensDetails - total_tokens: int - - -StreamingResponsesResponse: TypeAlias = Union[ - ResponseCreatedEvent, - ResponseInProgressEvent, - ResponseCompletedEvent, - ResponseOutputItemAddedEvent, - ResponseOutputItemDoneEvent, - ResponseContentPartAddedEvent, - ResponseContentPartDoneEvent, - ResponseReasoningTextDeltaEvent, - ResponseReasoningTextDoneEvent, - ResponseCodeInterpreterCallInProgressEvent, - ResponseCodeInterpreterCallCodeDeltaEvent, - ResponseWebSearchCallInProgressEvent, - ResponseWebSearchCallSearchingEvent, - ResponseWebSearchCallCompletedEvent, - ResponseCodeInterpreterCallCodeDoneEvent, - ResponseCodeInterpreterCallInterpretingEvent, - ResponseCodeInterpreterCallCompletedEvent, -] - -_RESPONSES_HANDOFF_FIELDS = ( - "disaggregated_params", - "prompt_token_ids", - "prompt_token_ids_b64", - "finish_reason", -) - - -class ResponsesResponse(OpenAIBaseModel): - id: str = Field(default_factory=lambda: f"resp_{str(uuid.uuid4().hex)}") - created_at: int = Field(default_factory=lambda: int(time.time())) - # error: Optional[ResponseError] = None - incomplete_details: Optional[IncompleteDetails] = None - instructions: Optional[str] = None - metadata: Optional[Metadata] = None - model: str - object: Literal["response"] = "response" - output: list[ResponseOutputItem] - parallel_tool_calls: bool - temperature: float - tool_choice: ToolChoice - tools: list[Tool] - top_p: float - background: bool - max_output_tokens: Optional[int] = None - max_tool_calls: Optional[int] = None - previous_response_id: Optional[str] = None - prompt: Optional[ResponsePrompt] = None - reasoning: Optional[Reasoning] = None - service_tier: Literal["auto", "default", "flex", "scale", "priority"] - status: ResponseStatus - text: Optional[ResponseTextConfig] = None - top_logprobs: int - truncation: Literal["auto", "disabled"] - usage: Optional[ResponseUsage] = None - user: Optional[str] = None - - # TensorRT-LLM extensions, set only on a context-only response: the - # disaggregated handoff, and the engine's raw finish reason, which tells the - # orchestrator whether a generation phase is pending ("length" and - # "not_finished" both map to status "incomplete"). - disaggregated_params: Optional[DisaggregatedParams] = Field(default=None) - prompt_token_ids: Optional[List[int]] = None - prompt_token_ids_b64: Optional[str] = None - finish_reason: Optional[str] = None - - @model_serializer(mode="wrap") - def _omit_absent_handoff(self, handler: Any) -> Any: - # The handoff fields are unset on every response a client sees, so they - # are left out rather than serialized as null in each body and event. - data = handler(self) - if isinstance(data, dict): - for key in _RESPONSES_HANDOFF_FIELDS: - if data.get(key, ...) is None: - data.pop(key) - return data - - @classmethod - def from_request( - cls, - request: ResponsesRequest, - sampling_params: SamplingParams, - model_name: str, - created_time: int, - output: list[ResponseOutputItem], - status: ResponseStatus, - usage: Optional[ResponseUsage] = None, - ) -> "ResponsesResponse": - return cls( - id=request.request_id, - created_at=created_time, - instructions=request.instructions, - metadata=request.metadata, - model=model_name, - output=output, - parallel_tool_calls=request.parallel_tool_calls, - temperature=sampling_params.temperature, - tool_choice=request.tool_choice, - tools=request.tools, - top_p=sampling_params.top_p, - background=request.background, - max_output_tokens=sampling_params.max_tokens, - max_tool_calls=request.max_tool_calls, - previous_response_id=request.previous_response_id, - prompt=request.prompt, - reasoning=request.reasoning, - service_tier=request.service_tier, - status=status, - text=request.text, - top_logprobs=sampling_params.logprobs, - truncation=request.truncation, - user=request.user, - usage=usage, - ) - - -class ResponsesStreamResponse(OpenAIBaseModel): - response: ResponsesResponse - sequence_number: int - type: Literal["response.created", "response.in_progress", - "response.completed", "response.failed", - "response.incomplete"] - - -class MemoryUpdateRequest(OpenAIBaseModel): - tags: list[str] = Field( - min_length=1, - description="Memory tags to release/resume, e.g. ['model', 'kv_cache']") - - -class RuntimeMemoryUpdateRequest(OpenAIBaseModel): - tags: Optional[List[str]] = Field( - default=None, - description="Memory tags to release/resume; omit to select defaults") - - -class UpdateWeightsRequest(OpenAIBaseModel): - weights: Optional[Dict[str, str]] = Field( - default=None, - description="Weight handles dict, or None to finalize update") - - -def encode_opaque_state(opaque_state: Optional[bytes]) -> Optional[str]: - if opaque_state is None: - return None - return base64.b64encode(opaque_state).decode("utf-8") - - -def decode_opaque_state(encoded_opaque_state: Optional[str]) -> Optional[bytes]: - if encoded_opaque_state is None: - return None - return base64.b64decode(encoded_opaque_state) - - -def _serialize_first_gen_log_probs( - first_gen_log_probs: Optional[list]) -> Optional[List]: - """Serialize ``list[dict[int, Logprob]] | list[float]`` to a JSON-safe form. - - - Default (verbose) format: each position is a ``dict[int, Logprob]`` and is - serialized as a list of ``{token_id, logprob, rank}`` dicts. - - Simple format: each position is a ``float``; passed through verbatim. - """ - if first_gen_log_probs is None: - return None - if not isinstance(first_gen_log_probs, list): - raise ValueError("first_gen_log_probs must be a list") - result = [] - for i, pos in enumerate(first_gen_log_probs): - if isinstance(pos, dict): - result.append([{ - "token_id": tid, - "logprob": lp.logprob, - "rank": lp.rank - } for tid, lp in pos.items()]) - elif isinstance(pos, (float, int)): - # Simple format: per-token sampled logprob. - result.append(float(pos)) - else: - raise ValueError( - f"first_gen_log_probs[{i}] must be a dict or float, got {type(pos)}" - ) - return result - - -def _deserialize_first_gen_log_probs( - serialized: Optional[List], ) -> Optional[list]: - """Inverse of :func:`_serialize_first_gen_log_probs`. - - Returns either ``list[dict[int, Logprob]]`` (default format) or - ``list[float]`` (simple format) depending on the serialized payload. - """ - if serialized is None: - return None - from tensorrt_llm.executor.result import Logprob - result = [] - for i, pos in enumerate(serialized): - if isinstance(pos, (float, int)): - result.append(float(pos)) - continue - if not isinstance(pos, list): - raise ValueError( - f"first_gen_log_probs[{i}] must be a list or float, got {type(pos)}" - ) - token_map = {} - for j, item in enumerate(pos): - if not isinstance(item, dict): - raise ValueError( - f"first_gen_log_probs[{i}][{j}] must be a dict") - if "token_id" not in item or "logprob" not in item: - raise ValueError( - f"first_gen_log_probs[{i}][{j}] missing required keys " - "'token_id' and/or 'logprob'") - token_map[item["token_id"]] = Logprob(logprob=item["logprob"], - rank=item.get("rank")) - result.append(token_map) - return result - - -def _serialize_first_gen_logits( - first_gen_logits: Optional[list], ) -> Optional[List]: - """Serialize list[torch.Tensor] to JSON-safe list[dict] with base64 data.""" - if first_gen_logits is None: - return None - result = [] - for i, tensor in enumerate(first_gen_logits): - t = tensor.contiguous().cpu() - if t.dtype == torch.bfloat16: - t = t.to(torch.float16) - np_array = t.numpy() - result.append({ - "data": base64.b64encode(np_array.tobytes()).decode(), - "shape": list(np_array.shape), - "dtype": str(np_array.dtype), - }) - return result - - -def _deserialize_first_gen_logits( - serialized: Optional[List], ) -> Optional[list]: - """Deserialize JSON list[dict] back to list[torch.Tensor].""" - if serialized is None: - return None - import numpy as np - result = [] - for i, item in enumerate(serialized): - if not isinstance(item, dict): - raise ValueError( - f"first_gen_logits[{i}] must be a dict, got {type(item)}") - for key in ("data", "shape", "dtype"): - if key not in item: - raise ValueError( - f"first_gen_logits[{i}] missing required key '{key}'") - np_array = np.frombuffer( - base64.b64decode(item["data"]), - dtype=np.dtype(item["dtype"]), - ).reshape(item["shape"]) - result.append(torch.from_numpy(np_array.copy())) - return result - - -def to_disaggregated_params( - tllm_disagg_params: LlmDisaggregatedParams) -> DisaggregatedParams: - if tllm_disagg_params is None: - return None - ctx_usage = tllm_disagg_params.ctx_usage - if ctx_usage is not None and not isinstance(ctx_usage, UsageInfo): - ctx_usage = UsageInfo.model_validate(ctx_usage) - return DisaggregatedParams( - request_type=tllm_disagg_params.request_type, - first_gen_tokens=tllm_disagg_params.first_gen_tokens, - first_gen_log_probs=_serialize_first_gen_log_probs( - tllm_disagg_params.first_gen_log_probs), - first_gen_logits=_serialize_first_gen_logits( - tllm_disagg_params.first_gen_logits), - ctx_request_id=tllm_disagg_params.ctx_request_id, - encoded_opaque_state=encode_opaque_state( - tllm_disagg_params.opaque_state), - draft_tokens=tllm_disagg_params.draft_tokens, - disagg_request_id=tllm_disagg_params.disagg_request_id, - ctx_dp_rank=tllm_disagg_params.ctx_dp_rank, - ctx_info_endpoint=tllm_disagg_params.ctx_info_endpoint, - schedule_style=tllm_disagg_params.schedule_style, - ctx_usage=ctx_usage, - remote_tail_start=tllm_disagg_params.remote_tail_start, - multimodal_positions=tllm_disagg_params.multimodal_positions, - multimodal_lengths=tllm_disagg_params.multimodal_lengths, - ) - - -def to_llm_disaggregated_params( - disaggregated_params: DisaggregatedParams) -> LlmDisaggregatedParams: - if disaggregated_params is None: - return None - ctx_usage = disaggregated_params.ctx_usage - return LlmDisaggregatedParams( - request_type=disaggregated_params.request_type, - first_gen_tokens=disaggregated_params.first_gen_tokens, - first_gen_log_probs=_deserialize_first_gen_log_probs( - disaggregated_params.first_gen_log_probs), - first_gen_logits=_deserialize_first_gen_logits( - disaggregated_params.first_gen_logits), - ctx_request_id=disaggregated_params.ctx_request_id, - opaque_state=decode_opaque_state( - disaggregated_params.encoded_opaque_state), - draft_tokens=disaggregated_params.draft_tokens, - disagg_request_id=disaggregated_params.disagg_request_id, - ctx_dp_rank=disaggregated_params.ctx_dp_rank, - ctx_info_endpoint=disaggregated_params.ctx_info_endpoint, - schedule_style=disaggregated_params.schedule_style, - ctx_usage=None if ctx_usage is None else ctx_usage.model_dump(), - remote_tail_start=disaggregated_params.remote_tail_start, - multimodal_positions=disaggregated_params.multimodal_positions, - multimodal_lengths=disaggregated_params.multimodal_lengths, - orig_elapsed_s=disaggregated_params.orig_elapsed_s, - ) - - -def to_llm_conversation_params( - conversation_params: Optional[ConversationParams] -) -> Optional[LlmConversationParams]: - if conversation_params is None: - return None - return LlmConversationParams( - conversation_id=conversation_params.conversation_id) - - -UCompletionRequest = Union[CompletionRequest, ChatCompletionRequest] -UCompletionResponse = Union[CompletionResponse, ChatCompletionResponse] - -ProfileActivity = Literal["CPU", "GPU", "CUDA_PROFILER"] - - -class StartProfileRequest(OpenAIBaseModel): - """Request body for the POST /start_profile endpoint.""" - output_dir: Optional[str] = Field( - default=None, - description="Directory where chrome traces are written. Defaults " - "to the TLLM_TORCH_PROFILER_DIR environment variable, " - "then /tmp.") - num_steps: Optional[PositiveInt] = Field( - default=None, - description="Number of iterations to profile. Must be >= 1 if " - "provided; if omitted, profiling runs until /stop_profile is " - "called. ``num_steps == 0`` is rejected because the profile " - "window would never close — the stop iteration would equal the " - "start iteration, profile_step() would discard the stop marker " - "as stale, and the window would run forever.") - start_step: NonNegativeInt = Field( - default=0, - description="Skip this many iterations before profiling begins. " - "Must be >= 0.") - activities: List[ProfileActivity] = Field( - default_factory=lambda: ["CPU", "GPU"], - description="Which activities to trace. Supported values: " - "'CPU', 'GPU', 'CUDA_PROFILER'. Unknown values are rejected at " - "schema validation time.") - - -# Old import path of names that moved to -# tensorrt_llm.serve.api.visual_gen.protocol. -# They are re-exported lazily (PEP 562): importing this module does not import -# their new module. -_LAZY_NAMES = { - "ImageEditRequest": "tensorrt_llm.serve.api.visual_gen.protocol", - "ImageGenerationRequest": "tensorrt_llm.serve.api.visual_gen.protocol", - "ImageGenerationResponse": "tensorrt_llm.serve.api.visual_gen.protocol", - "ImageObject": "tensorrt_llm.serve.api.visual_gen.protocol", - "MediaReferenceItem": "tensorrt_llm.serve.api.visual_gen.protocol", - "VideoGenerationRequest": "tensorrt_llm.serve.api.visual_gen.protocol", - "VideoJob": "tensorrt_llm.serve.api.visual_gen.protocol", - "VideoJobList": "tensorrt_llm.serve.api.visual_gen.protocol", -} - - -def __getattr__(name): - module_name = _LAZY_NAMES.get(name) - if module_name is None: - raise AttributeError(f"module {__name__!r} has no attribute {name!r}") - value = getattr(importlib.import_module(module_name), name) - globals()[name] = value # cache: later access skips __getattr__ - return value - - -def __dir__(): - return sorted(set(globals()) | set(_LAZY_NAMES)) diff --git a/tensorrt_llm/serve/openai_protocol/__init__.py b/tensorrt_llm/serve/openai_protocol/__init__.py new file mode 100644 index 000000000000..a140aa29b70e --- /dev/null +++ b/tensorrt_llm/serve/openai_protocol/__init__.py @@ -0,0 +1,329 @@ +# 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. +"""OpenAI-dialect request and response models, one submodule per API. + +New code imports the submodule that defines a name, for example +``from tensorrt_llm.serve.openai_protocol.chat import ChatCompletionRequest``. +This package re-exports, on first access, every name the former single-file +module ``openai_protocol.py`` had, so existing imports keep working. Do not add +names here; tests/unittest/serve/openai_protocol/test_openai_protocol_facade.py +fails if the table grows. +""" + +import importlib +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + # Public names, so type checkers and IDEs resolve them; private names are + # reachable at run time through __getattr__. The VisualGen names are left + # out on purpose: this file names their module only as a string in + # _NAME_TO_MODULE (checked by the facade test under + # tests/unittest/serve/api/visual_gen/). + from .chat import ( + ChatCompletionContentPartParam, + ChatCompletionLogProb, + ChatCompletionLogProbs, + ChatCompletionLogProbsContent, + ChatCompletionMessageParam, + ChatCompletionRequest, + ChatCompletionResponse, + ChatCompletionResponseChoice, + ChatCompletionResponseStreamChoice, + ChatCompletionStreamResponse, + ChatCompletionThinkingParam, + ChatMessage, + CustomChatCompletionContentPartParam, + CustomChatCompletionMessageParam, + DeltaMessage, + DynamicToolsSystemMessageParam, + ReasoningAssistantMessage, + ) + from .common import ( + ConversationParams, + DisaggregatedParams, + DisaggScheduleStyle, + EarlyStopping, + ErrorResponse, + OpenAIBaseModel, + PromptTokensDetails, + ResponseFormat, + SpeculativeDecodingStats, + StreamOptions, + UsageInfo, + ensure_request_chat_template_allowed, + ) + from .completions import ( + CompletionLogProbs, + CompletionRequest, + CompletionResponse, + CompletionResponseChoice, + CompletionResponseStreamChoice, + CompletionStreamResponse, + ) + from .control import ( + MemoryUpdateRequest, + ProfileActivity, + RuntimeMemoryUpdateRequest, + StartProfileRequest, + UpdateWeightsRequest, + ) + from .disagg import ( + UCompletionRequest, + UCompletionResponse, + decode_opaque_state, + encode_opaque_state, + to_disaggregated_params, + to_llm_conversation_params, + to_llm_disaggregated_params, + ) + from .embeddings import ( + EmbeddingRequest, + EmbeddingResponse, + EmbeddingResponseData, + EmbeddingUsageInfo, + ) + from .kv_cache import KVCacheTruncateRequest, KVCacheTruncateTokensRequest + from .models import ModelCard, ModelList + from .responses import ( + InputTokensDetails, + OutputTokensDetails, + ResponseInputOutputItem, + ResponsesRequest, + ResponsesResponse, + ResponsesStreamResponse, + ResponseUsage, + StreamingResponsesResponse, + ) + from .tokenize import TokenizeRequest, TokenizeResponse + from .tools import ( + TOOL_PARAM_MAX_ENUM_VALUES, + ChatCompletionNamedFunction, + ChatCompletionNamedToolChoiceParam, + ChatCompletionToolsParam, + DeltaFunctionCall, + DeltaToolCall, + FunctionCall, + FunctionDefinition, + NamespacedFunctionToolParam, + ToolCall, + ) + +# Name -> module that provides it. +_NAME_TO_MODULE = { + "ChatCompletionContentPartParam": "tensorrt_llm.serve.openai_protocol.chat", + "ChatCompletionLogProb": "tensorrt_llm.serve.openai_protocol.chat", + "ChatCompletionLogProbs": "tensorrt_llm.serve.openai_protocol.chat", + "ChatCompletionLogProbsContent": "tensorrt_llm.serve.openai_protocol.chat", + "ChatCompletionMessageParam": "tensorrt_llm.serve.openai_protocol.chat", + "ChatCompletionNamedFunction": "tensorrt_llm.serve.openai_protocol.tools", + "ChatCompletionNamedToolChoiceParam": "tensorrt_llm.serve.openai_protocol.tools", + "ChatCompletionRequest": "tensorrt_llm.serve.openai_protocol.chat", + "ChatCompletionResponse": "tensorrt_llm.serve.openai_protocol.chat", + "ChatCompletionResponseChoice": "tensorrt_llm.serve.openai_protocol.chat", + "ChatCompletionResponseStreamChoice": "tensorrt_llm.serve.openai_protocol.chat", + "ChatCompletionStreamResponse": "tensorrt_llm.serve.openai_protocol.chat", + "ChatCompletionThinkingParam": "tensorrt_llm.serve.openai_protocol.chat", + "ChatCompletionToolsParam": "tensorrt_llm.serve.openai_protocol.tools", + "ChatMessage": "tensorrt_llm.serve.openai_protocol.chat", + "CompletionLogProbs": "tensorrt_llm.serve.openai_protocol.completions", + "CompletionRequest": "tensorrt_llm.serve.openai_protocol.completions", + "CompletionResponse": "tensorrt_llm.serve.openai_protocol.completions", + "CompletionResponseChoice": "tensorrt_llm.serve.openai_protocol.completions", + "CompletionResponseStreamChoice": "tensorrt_llm.serve.openai_protocol.completions", + "CompletionStreamResponse": "tensorrt_llm.serve.openai_protocol.completions", + "ConversationParams": "tensorrt_llm.serve.openai_protocol.common", + "CustomChatCompletionContentPartParam": "tensorrt_llm.serve.openai_protocol.chat", + "CustomChatCompletionMessageParam": "tensorrt_llm.serve.openai_protocol.chat", + "DeltaFunctionCall": "tensorrt_llm.serve.openai_protocol.tools", + "DeltaMessage": "tensorrt_llm.serve.openai_protocol.chat", + "DeltaToolCall": "tensorrt_llm.serve.openai_protocol.tools", + "DisaggregatedParams": "tensorrt_llm.serve.openai_protocol.common", + "DynamicToolsSystemMessageParam": "tensorrt_llm.serve.openai_protocol.chat", + "EarlyStopping": "tensorrt_llm.serve.openai_protocol.common", + "EmbeddingRequest": "tensorrt_llm.serve.openai_protocol.embeddings", + "EmbeddingResponse": "tensorrt_llm.serve.openai_protocol.embeddings", + "EmbeddingResponseData": "tensorrt_llm.serve.openai_protocol.embeddings", + "EmbeddingUsageInfo": "tensorrt_llm.serve.openai_protocol.embeddings", + "ErrorResponse": "tensorrt_llm.serve.openai_protocol.common", + "FunctionCall": "tensorrt_llm.serve.openai_protocol.tools", + "FunctionDefinition": "tensorrt_llm.serve.openai_protocol.tools", + "InputTokensDetails": "tensorrt_llm.serve.openai_protocol.responses", + "KVCacheTruncateRequest": "tensorrt_llm.serve.openai_protocol.kv_cache", + "KVCacheTruncateTokensRequest": "tensorrt_llm.serve.openai_protocol.kv_cache", + "MemoryUpdateRequest": "tensorrt_llm.serve.openai_protocol.control", + "ModelCard": "tensorrt_llm.serve.openai_protocol.models", + "ModelList": "tensorrt_llm.serve.openai_protocol.models", + "NamespacedFunctionToolParam": "tensorrt_llm.serve.openai_protocol.tools", + "OpenAIBaseModel": "tensorrt_llm.serve.openai_protocol.common", + "OutputTokensDetails": "tensorrt_llm.serve.openai_protocol.responses", + "ProfileActivity": "tensorrt_llm.serve.openai_protocol.control", + "PromptTokensDetails": "tensorrt_llm.serve.openai_protocol.common", + "ReasoningAssistantMessage": "tensorrt_llm.serve.openai_protocol.chat", + "ResponseFormat": "tensorrt_llm.serve.openai_protocol.common", + "ResponseInputOutputItem": "tensorrt_llm.serve.openai_protocol.responses", + "ResponseUsage": "tensorrt_llm.serve.openai_protocol.responses", + "ResponsesRequest": "tensorrt_llm.serve.openai_protocol.responses", + "ResponsesResponse": "tensorrt_llm.serve.openai_protocol.responses", + "ResponsesStreamResponse": "tensorrt_llm.serve.openai_protocol.responses", + "RuntimeMemoryUpdateRequest": "tensorrt_llm.serve.openai_protocol.control", + "SpeculativeDecodingStats": "tensorrt_llm.serve.openai_protocol.common", + "StartProfileRequest": "tensorrt_llm.serve.openai_protocol.control", + "StreamOptions": "tensorrt_llm.serve.openai_protocol.common", + "StreamingResponsesResponse": "tensorrt_llm.serve.openai_protocol.responses", + "TOOL_PARAM_MAX_ENUM_VALUES": "tensorrt_llm.serve.openai_protocol.tools", + "TokenizeRequest": "tensorrt_llm.serve.openai_protocol.tokenize", + "TokenizeResponse": "tensorrt_llm.serve.openai_protocol.tokenize", + "ToolCall": "tensorrt_llm.serve.openai_protocol.tools", + "UCompletionRequest": "tensorrt_llm.serve.openai_protocol.disagg", + "UCompletionResponse": "tensorrt_llm.serve.openai_protocol.disagg", + "UpdateWeightsRequest": "tensorrt_llm.serve.openai_protocol.control", + "UsageInfo": "tensorrt_llm.serve.openai_protocol.common", + "_ID_STRIPPED_ROLES": "tensorrt_llm.serve.openai_protocol.responses", + "_LOGIT_BIAS_MAX": "tensorrt_llm.serve.openai_protocol.sampling", + "_LOGIT_BIAS_MIN": "tensorrt_llm.serve.openai_protocol.sampling", + "_OmitsAbsentSpecDecodeStats": "tensorrt_llm.serve.openai_protocol.common", + "_RESPONSES_HANDOFF_FIELDS": "tensorrt_llm.serve.openai_protocol.responses", + "_SCHEMA_INSTANCE_KEYWORDS": "tensorrt_llm.serve.openai_protocol.tools", + "_SCHEMA_MAP_KEYWORDS": "tensorrt_llm.serve.openai_protocol.tools", + "_count_schema_enum_values": "tensorrt_llm.serve.openai_protocol.tools", + "_deserialize_first_gen_log_probs": "tensorrt_llm.serve.openai_protocol.disagg", + "_deserialize_first_gen_logits": "tensorrt_llm.serve.openai_protocol.disagg", + "_drop_explicit_nulls": "tensorrt_llm.serve.openai_protocol.responses", + "_early_stopping_to_int": "tensorrt_llm.serve.openai_protocol.sampling", + "_logit_bias_to_embedding_bias": "tensorrt_llm.serve.openai_protocol.sampling", + "_materialize_validator_iterators": "tensorrt_llm.serve.openai_protocol.responses", + "_record_sampling_params_request_fields": "tensorrt_llm.serve.openai_protocol.sampling", + "_response_format_text_config_to_guided_decoding_params": "tensorrt_llm.serve.openai_protocol.sampling", + "_response_format_to_guided_decoding_params": "tensorrt_llm.serve.openai_protocol.sampling", + "_serialize_first_gen_log_probs": "tensorrt_llm.serve.openai_protocol.disagg", + "_serialize_first_gen_logits": "tensorrt_llm.serve.openai_protocol.disagg", + "decode_opaque_state": "tensorrt_llm.serve.openai_protocol.disagg", + "encode_opaque_state": "tensorrt_llm.serve.openai_protocol.disagg", + "ensure_request_chat_template_allowed": "tensorrt_llm.serve.openai_protocol.common", + "to_disaggregated_params": "tensorrt_llm.serve.openai_protocol.disagg", + "to_llm_conversation_params": "tensorrt_llm.serve.openai_protocol.disagg", + "to_llm_disaggregated_params": "tensorrt_llm.serve.openai_protocol.disagg", + # Defined in the VisualGen serving package. Named here only as strings: an import, + # even under TYPE_CHECKING, would load that package into every LLM process. + "ImageEditRequest": "tensorrt_llm.serve.api.visual_gen.protocol", + "ImageGenerationRequest": "tensorrt_llm.serve.api.visual_gen.protocol", + "ImageGenerationResponse": "tensorrt_llm.serve.api.visual_gen.protocol", + "ImageObject": "tensorrt_llm.serve.api.visual_gen.protocol", + "MediaReferenceItem": "tensorrt_llm.serve.api.visual_gen.protocol", + "VideoGenerationRequest": "tensorrt_llm.serve.api.visual_gen.protocol", + "VideoJob": "tensorrt_llm.serve.api.visual_gen.protocol", + "VideoJobList": "tensorrt_llm.serve.api.visual_gen.protocol", + # Imported by the former module and imported from it by code in the repository. + "DisaggScheduleStyle": "tensorrt_llm.serve.openai_protocol.common", +} + +__all__ = [ + "ChatCompletionContentPartParam", + "ChatCompletionLogProb", + "ChatCompletionLogProbs", + "ChatCompletionLogProbsContent", + "ChatCompletionMessageParam", + "ChatCompletionNamedFunction", + "ChatCompletionNamedToolChoiceParam", + "ChatCompletionRequest", + "ChatCompletionResponse", + "ChatCompletionResponseChoice", + "ChatCompletionResponseStreamChoice", + "ChatCompletionStreamResponse", + "ChatCompletionThinkingParam", + "ChatCompletionToolsParam", + "ChatMessage", + "CompletionLogProbs", + "CompletionRequest", + "CompletionResponse", + "CompletionResponseChoice", + "CompletionResponseStreamChoice", + "CompletionStreamResponse", + "ConversationParams", + "CustomChatCompletionContentPartParam", + "CustomChatCompletionMessageParam", + "DeltaFunctionCall", + "DeltaMessage", + "DeltaToolCall", + "DisaggScheduleStyle", + "DisaggregatedParams", + "DynamicToolsSystemMessageParam", + "EarlyStopping", + "EmbeddingRequest", + "EmbeddingResponse", + "EmbeddingResponseData", + "EmbeddingUsageInfo", + "ErrorResponse", + "FunctionCall", + "FunctionDefinition", + "ImageEditRequest", + "ImageGenerationRequest", + "ImageGenerationResponse", + "ImageObject", + "InputTokensDetails", + "KVCacheTruncateRequest", + "KVCacheTruncateTokensRequest", + "MediaReferenceItem", + "MemoryUpdateRequest", + "ModelCard", + "ModelList", + "NamespacedFunctionToolParam", + "OpenAIBaseModel", + "OutputTokensDetails", + "ProfileActivity", + "PromptTokensDetails", + "ReasoningAssistantMessage", + "ResponseFormat", + "ResponseInputOutputItem", + "ResponseUsage", + "ResponsesRequest", + "ResponsesResponse", + "ResponsesStreamResponse", + "RuntimeMemoryUpdateRequest", + "SpeculativeDecodingStats", + "StartProfileRequest", + "StreamOptions", + "StreamingResponsesResponse", + "TOOL_PARAM_MAX_ENUM_VALUES", + "TokenizeRequest", + "TokenizeResponse", + "ToolCall", + "UCompletionRequest", + "UCompletionResponse", + "UpdateWeightsRequest", + "UsageInfo", + "VideoGenerationRequest", + "VideoJob", + "VideoJobList", + "decode_opaque_state", + "encode_opaque_state", + "ensure_request_chat_template_allowed", + "to_disaggregated_params", + "to_llm_conversation_params", + "to_llm_disaggregated_params", +] + + +def __getattr__(name): + module_name = _NAME_TO_MODULE.get(name) + if module_name is None: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + value = getattr(importlib.import_module(module_name), name) + globals()[name] = value # cache: later access skips __getattr__ + return value + + +def __dir__(): + return sorted(set(globals()) | set(_NAME_TO_MODULE)) diff --git a/tensorrt_llm/serve/openai_protocol/chat.py b/tensorrt_llm/serve/openai_protocol/chat.py new file mode 100644 index 000000000000..fae182bec979 --- /dev/null +++ b/tensorrt_llm/serve/openai_protocol/chat.py @@ -0,0 +1,556 @@ +# 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. + +# Adapted from +# https://github.com/vllm-project/vllm/blob/4db5176d9758b720b05460c50ace3c01026eb158/vllm/entrypoints/openai/protocol.py + +import time +import uuid +from typing import Any, Dict, List, Literal, Optional, Union + +from openai.types.chat import ChatCompletionAssistantMessageParam +from openai.types.chat import \ + ChatCompletionContentPartParam as OpenAIChatCompletionContentPartParam +from openai.types.chat import \ + ChatCompletionMessageParam as OpenAIChatCompletionMessageParam +from openai_harmony import ReasoningEffort +from pydantic import ConfigDict, Field, field_validator, model_validator +from typing_extensions import Annotated, Required, TypedDict + +from tensorrt_llm.executor.request import LoRARequest +from tensorrt_llm.inputs.media_io import MediaModality +from tensorrt_llm.kv_hints import KvHintsEnvelope +from tensorrt_llm.llmapi import SamplingParams +from tensorrt_llm.sampling_params import (check_logprobs_limit, + validate_thinking_token_budget) +from tensorrt_llm.scheduling_params import AgentHierarchy +from tensorrt_llm.serve.openai_protocol.common import ( + ConversationParams, DisaggregatedParams, EarlyStopping, OpenAIBaseModel, + ResponseFormat, SpeculativeDecodingStats, StreamOptions, UsageInfo, + _OmitsAbsentSpecDecodeStats) +from tensorrt_llm.serve.openai_protocol.sampling import ( + _early_stopping_to_int, _logit_bias_to_embedding_bias, + _record_sampling_params_request_fields, + _response_format_to_guided_decoding_params) +from tensorrt_llm.serve.openai_protocol.tools import ( + ChatCompletionNamedToolChoiceParam, ChatCompletionToolsParam, DeltaToolCall, + ToolCall) + + +class ChatMessage(OpenAIBaseModel): + role: str + content: Optional[str] = None + reasoning_content: Optional[str] = None + reasoning: Optional[str] = None + tool_calls: List[ToolCall] = Field(default_factory=list) + + +class ChatCompletionLogProb(OpenAIBaseModel): + token: str + logprob: float = -9999.0 + bytes: Optional[List[int]] = None + + +class ChatCompletionLogProbsContent(ChatCompletionLogProb): + top_logprobs: Optional[List[ChatCompletionLogProb]] = None + + +class CustomChatCompletionContentPartParam(TypedDict, total=False): + __pydantic_config__ = ConfigDict(extra="allow") # type: ignore + + type: Required[str] + """The type of the content part.""" + + +ChatCompletionContentPartParam = Union[OpenAIChatCompletionContentPartParam, + CustomChatCompletionContentPartParam] + + +class CustomChatCompletionMessageParam(TypedDict, total=False): + """Enables custom roles in the Chat Completion API.""" + + # This is so custom fields not in any of the `ChatCompletionMessageParam` defined by OpenAI + # are still allowed. + # Examples include: assistant messages with `reasoning` / `reasoning_content`. + __pydantic_config__ = ConfigDict(extra="allow") # type: ignore + + role: Required[str] + """The role of the message's author.""" + + content: Union[str, List[ChatCompletionContentPartParam], None] + """The contents of the message.""" + + name: str + """An optional name for the participant. + + Provides the model information to differentiate between participants of the + same role. + """ + + +class ReasoningAssistantMessage(ChatCompletionAssistantMessageParam): + """Assistant message that includes reasoning tokens.""" + reasoning: Optional[str] + # NOTE: some older benchmarks and chat templates assume the below, which has been deprecated + # in other inference frameworks in favor of the above `reasoning` field. + reasoning_content: Optional[str] + + +class DynamicToolsSystemMessageParam(TypedDict, total=False): + """System message carrying message-level (dynamic) tool declarations. + + Kimi-style templates render such messages as an in-conversation tool + declare block. Must come first in `ChatCompletionMessageParam`: the + stock OpenAI system-message TypedDict otherwise wins smart-union scoring + and silently drops the `tools` key. + """ + __pydantic_config__ = ConfigDict(extra="allow") # type: ignore + + role: Required[Literal["system"]] + tools: Required[List[dict]] + content: Union[str, List[ChatCompletionContentPartParam], None] + name: str + + +ChatCompletionMessageParam = Union[DynamicToolsSystemMessageParam, + OpenAIChatCompletionMessageParam, + CustomChatCompletionMessageParam, + ReasoningAssistantMessage] + + +class ChatCompletionLogProbs(OpenAIBaseModel): + content: Optional[List[ChatCompletionLogProbsContent]] = None + + +class ChatCompletionResponseChoice(_OmitsAbsentSpecDecodeStats): + index: int + message: ChatMessage + logprobs: Optional[ChatCompletionLogProbs] = None + finish_reason: Optional[str] = None + stop_reason: Optional[Union[int, str]] = None + # TODO: progressivly add more info like input_ids, specific_token_ids, mrope, mm_hashes, etc + # TODO: and use a JSON-safe handle to refer to the server-side output + mm_embedding_handle: Optional[Dict[str, Any]] = None + + disaggregated_params: Optional[DisaggregatedParams] = Field(default=None) + avg_decoded_tokens_per_iter: Optional[float] = Field(default=None) + speculative_decoding: Optional[SpeculativeDecodingStats] = Field( + default=None) + + +class ChatCompletionResponse(OpenAIBaseModel): + id: str = Field(default_factory=lambda: f"chatcmpl-{str(uuid.uuid4().hex)}") + object: Literal["chat.completion"] = "chat.completion" + created: int = Field(default_factory=lambda: int(time.time())) + model: str + choices: List[ChatCompletionResponseChoice] + usage: UsageInfo + # Add prompt_tokens_ids to the response to remove the tokenization + # in the generation server in disaggreated serving + prompt_token_ids: Optional[List[int]] = None + # base64 int32 buffer alternative to prompt_token_ids; set by the context + # worker so the orchestrator can relay a string instead of the int list. + prompt_token_ids_b64: Optional[str] = None + + +class DeltaMessage(OpenAIBaseModel): + role: Optional[str] = None + content: Optional[str] = None + reasoning_content: Optional[str] = None + # For GPT-OSS style reasoning + reasoning: Optional[str] = None + tool_calls: Optional[List[DeltaToolCall]] = None + + +class ChatCompletionResponseStreamChoice(_OmitsAbsentSpecDecodeStats): + index: int + delta: DeltaMessage + logprobs: Optional[ChatCompletionLogProbs] = None + finish_reason: Optional[str] = None + stop_reason: Optional[Union[int, str]] = None + avg_decoded_tokens_per_iter: Optional[float] = Field(default=None) + speculative_decoding: Optional[SpeculativeDecodingStats] = Field( + default=None) + + +class ChatCompletionStreamResponse(OpenAIBaseModel): + id: str = Field(default_factory=lambda: f"chatcmpl-{str(uuid.uuid4().hex)}") + object: Literal["chat.completion.chunk"] = "chat.completion.chunk" + created: int = Field(default_factory=lambda: int(time.time())) + model: str + choices: List[ChatCompletionResponseStreamChoice] + usage: Optional[UsageInfo] = Field(default=None) + + +class ChatCompletionThinkingParam(OpenAIBaseModel): + """Kimi/Moonshot `thinking` extension controlling reasoning output. + + `keep` is fixed to `"all"` when thinking is enabled and ignored when + disabled; `effort` is only meaningful when enabled. An explicit + `effort` wins over a request-level `reasoning_effort`, which applies + only when `effort` is absent (and never for a disabled request). + """ + type: Literal["enabled", "disabled"] = "enabled" + keep: Optional[Literal["all"]] = None + effort: Optional[Literal["low", "high", "max"]] = None + + +class ChatCompletionRequest(OpenAIBaseModel): + kv_hints: Optional[KvHintsEnvelope] = Field( + default=None, + repr=False, + description= + "Prototype KV hint metadata; transported without executing actions.", + ) + # Ordered by official OpenAI API documentation + # https://platform.openai.com/docs/api-reference/chat/create + messages: List[ChatCompletionMessageParam] + # Add prompt_tokens_ids to the request to remove the tokenization + # in the generation server in disaggreated serving + prompt_token_ids: Optional[List[int]] = None + # base64 int32 buffer relayed by the orchestrator from the ctx response; + # decoded back to prompt_token_ids on the generation worker. Not for clients. + prompt_token_ids_b64: Optional[str] = None + model: str + frequency_penalty: Optional[float] = 0.0 + logit_bias: Optional[Dict[str, float]] = None + logprobs: Optional[bool] = False + top_logprobs: Optional[int] = 0 + max_completion_tokens: Optional[int] = Field(default=None, + validation_alias='max_tokens') + n: int = 1 + presence_penalty: Optional[float] = 0.0 + response_format: Optional[ResponseFormat] = None + seed: Optional[int] = Field(None) + service_tier: Optional[Literal["auto", "default", "flex", "scale", + "priority", "fast"]] = Field( + default=None, + description=( + "Accepted for OpenAI compatibility " + "and ignored: the server has no " + "service tiers. Scheduling priority " + "is the separate `priority` field.")) + stop: Optional[Union[str, List[str]]] = Field(default_factory=list) + # Standard OpenAI field, accepted for compatibility and otherwise ignored: + # the server keeps no copy of a chat completion, so there is nothing to + # store or to leave unstored. + store: Optional[bool] = None + stream: Optional[bool] = False + stream_options: Optional[StreamOptions] = None + temperature: Optional[float] = None + top_p: Optional[float] = None + tools: Optional[List[ChatCompletionToolsParam]] = None + tool_choice: Optional[Union[Literal["none", "auto", "required"], + ChatCompletionNamedToolChoiceParam]] = "none" + # Standard OpenAI field, accepted for compatibility. `false` is not + # enforced: the engine does not restrict how many tool calls the model + # emits per turn, so parallel emission remains model behavior either way. + parallel_tool_calls: Optional[bool] = None + user: Optional[str] = None + reasoning_effort: Optional[ReasoningEffort | Literal[ + "low", "medium", "high", "max", "none"]] = Field( + default=ReasoningEffort.LOW, + description=( + "The level of reasoning effort to use. Controls how much " + "reasoning is shown in the model's response. Options: " + "'low', 'medium', 'high' (harmony/gpt-oss), plus 'max' and " + "'none' for models with Kimi-style thinking control."), + ) + # Kimi/Moonshot extension: structured control of reasoning output. The + # serving layer maps it into the chat-template kwargs for models whose + # template understands it (e.g. kimi_k3); other models ignore it. + thinking: Optional[ChatCompletionThinkingParam] = None + thinking_token_budget: Optional[int] = None + prompt_ignore_length: Optional[int] = 0 + # Match Codex CLI's string-map contract; widen only for a client that needs + # other value types, even though the metadata is otherwise ignored. + client_metadata: Optional[Dict[str, str]] = Field( + default=None, + description=("Not an OpenAI field: opaque per-request metadata that " + "agent clients attach, a map of strings to strings. " + "Accepted and ignored.")) + + # doc: begin-chat-completion-sampling-params + best_of: Optional[int] = None + use_beam_search: bool = False + top_k: int = 0 + top_p_min: float = 0.0 + min_p: float = 0.0 + repetition_penalty: float = 1.0 + # Unset by default so the engine picks its own beam-search defaults, as it + # does for requests coming through the Python API. Sending concrete values + # here would put every served request on the length-normalized, exhaustive + # beam-search path. + length_penalty: Optional[float] = None + early_stopping: Optional[EarlyStopping] = None + stop_token_ids: Optional[List[int]] = Field(default_factory=list) + include_stop_str_in_output: bool = False + ignore_eos: bool = False + min_tokens: int = 0 + skip_special_tokens: bool = True + spaces_between_special_tokens: bool = True + truncate_prompt_tokens: Optional[Annotated[int, Field(ge=1)]] = None + lora_request: Optional[LoRARequest] = None + # doc: end-chat-completion-sampling-params + + # doc: begin-chat-completion-extra-params + echo: bool = Field( + default=False, + description=( + "If true, the new message will be prepended with the last message " + "if they belong to the same role."), + ) + add_generation_prompt: bool = Field( + default=True, + description= + ("If true, the generation prompt will be added to the chat template. " + "This is a parameter used by chat template in tokenizer config of the " + "model."), + ) + add_special_tokens: bool = Field( + default=False, + description=( + "If true, special tokens (e.g. BOS) will be added to the prompt " + "on top of what is added by the chat template. " + "For most models, the chat template takes care of adding the " + "special tokens so this should be set to false (as is the " + "default)."), + ) + documents: Optional[List[Dict[str, str]]] = Field( + default=None, + description= + ("A list of dicts representing documents that will be accessible to " + "the model if it is performing RAG (retrieval-augmented generation)." + " If the template does not support RAG, this argument will have no " + "effect. We recommend that each document should be a dict containing " + "\"title\" and \"text\" keys."), + ) + chat_template: Optional[str] = Field( + default=None, + description=( + "A Jinja template to use for this conversion. " + "If this is not passed, the model's default chat template will be " + "used instead."), + ) + chat_template_kwargs: Optional[Dict[str, Any]] = Field( + default=None, + description=("Additional kwargs to pass to the template renderer. " + "Will be accessible by the chat template."), + ) + injected_chat_template_kwargs: Optional[List[str]] = Field( + default=None, + description=( + "Server-populated. Keys of `chat_template_kwargs` that the server " + "derived from API-level fields (for example the Anthropic " + "`thinking` and `context_management` controls) rather than the " + "caller. The unused-kwargs guard exempts them, since the caller " + "cannot remove a control the server added. Relayed with the " + "request so a disaggregated worker rendering the prompt applies " + "the same exemption."), + ) + + media_io_kwargs: Optional[Dict[MediaModality, Dict[str, Any]]] = Field( + default=None, + description=( + "Per-request override for the server's `--media_io_kwargs`. " + "Shape: `{modality: {kwarg: value}}` with modality in " + "{\"image\", \"video\", \"audio\"}; unknown modality keys are " + "rejected. Per modality, request kwargs are shallow-merged " + "onto the server defaults (request wins per key). For " + "`video`, overriding only one of `fps`/`num_frames` drops " + "the other from the server default so the loader's built-in " + "is used. " + "Example: `{\"video\": {\"num_frames\": 32}}`."), + ) + + disaggregated_params: Optional[DisaggregatedParams] = Field( + default=None, + description=("Parameters for disaggregated serving"), + ) + conversation_params: Optional[ConversationParams] = Field( + default=None, + description=("Parameters for multi-turn conversation routing"), + ) + + cache_salt: Optional[str] = Field( + default=None, + description= + ("If specified, KV cache will be salted with the provided string " + "to limit the kv cache reuse on with the requests having the same string." + )) + priority: Optional[float] = Field( + default=None, + ge=0.0, + le=1.0, + description= + ("Scheduling priority in [0.0, 1.0]; higher is served first. Only honored " + "when the engine runs with scheduler_config.waiting_queue_policy=priority. " + "If unset, the engine default (0.5) is used."), + ) + + agent_hierarchy: Optional[AgentHierarchy] = Field( + default=None, description="Agent hierarchy ") + + mm_processor_kwargs: Optional[Dict[str, Any]] = Field( + default=None, + description= + "Per-request kwargs forwarded to the multimodal HF processor (e.g. num_frames for video models).", + ) + + # doc: end-chat-completion-extra-params + + def to_sampling_params(self, + vocab_size: Optional[int] = None, + gather_generation_logits: bool = False, + reasoning_parser: Optional[str] = None, + backend: Optional[str] = None) -> SamplingParams: + sampling_logprobs = None + if self.logprobs: + logprobs = 1 if not self.top_logprobs else self.top_logprobs + if backend == "pytorch" or gather_generation_logits: + sampling_logprobs = logprobs + elif self.top_logprobs: + raise ValueError( + "`gather_generation_logits` must be `True` to use `top_logprobs`" + ) + + sampling_params = SamplingParams( + frequency_penalty=self.frequency_penalty, + logprobs=sampling_logprobs, + max_tokens=self.max_completion_tokens, + n=self.n, + presence_penalty=self.presence_penalty, + seed=self.seed, + stop=self.stop, + temperature=(self.temperature + if self.temperature is not None else 1.0), + prompt_ignore_length=self.prompt_ignore_length, + + # chat-completion-sampling-params + best_of=self.best_of, + use_beam_search=self.use_beam_search, + top_k=self.top_k, + top_p=(self.top_p if self.top_p is not None else 1.0), + top_p_min=self.top_p_min if self.top_p_min > 0 else None, + min_p=self.min_p, + repetition_penalty=self.repetition_penalty, + length_penalty=self.length_penalty, + early_stopping=_early_stopping_to_int(self.early_stopping), + stop_token_ids=self.stop_token_ids, + include_stop_str_in_output=self.include_stop_str_in_output, + ignore_eos=self.ignore_eos, + min_tokens=self.min_tokens, + skip_special_tokens=self.skip_special_tokens, + spaces_between_special_tokens=self.spaces_between_special_tokens, + truncate_prompt_tokens=self.truncate_prompt_tokens, + guided_decoding=_response_format_to_guided_decoding_params( + self.response_format, + reasoning_parser=reasoning_parser, + chat_template_kwargs=self.chat_template_kwargs), + thinking_token_budget=self.thinking_token_budget, + + # logits_bias + embedding_bias=_logit_bias_to_embedding_bias( + self.logit_bias, vocab_size), + + # chat-completion-extra-params + add_special_tokens=self.add_special_tokens, + ) + _record_sampling_params_request_fields(self, sampling_params) + return sampling_params + + @model_validator(mode='before') + @classmethod + def validate_stream_options(cls, values): + if (values.get('stream_options') is not None + and not values.get('stream')): + raise ValueError("stream_options can only be set if stream is true") + return values + + @model_validator(mode="before") + @classmethod + def check_tool_choice(cls, data): + if not isinstance(data, dict): + return data + has_dynamic_tools = any( + isinstance(msg, dict) and msg.get("role") == "system" + and msg.get("tools") for msg in data.get("messages") or []) + if "tool_choice" not in data and (data.get("tools") + or has_dynamic_tools): + data["tool_choice"] = "auto" + # "none" and "auto" are meaningful without tools. "required" needs a + # non-empty tool set — request-level or message-level (dynamic) + # tools both count. A named function must be a request-level tool. + if "tool_choice" in data and data["tool_choice"] not in ("none", + "auto"): + satisfied = data.get("tools") or (data["tool_choice"] == "required" + and has_dynamic_tools) + if not satisfied: + raise ValueError( + "When using `tool_choice`, `tools` must be set.") + return data + + @model_validator(mode="before") + @classmethod + def check_message_tools_role(cls, data): + """Message-level tool declarations ride on system messages only. + + Union validation strips unknown keys from non-system messages, so + this raw-payload validator is the only layer that can reject the + misuse loudly instead of silently dropping the client's tools. The + rest of the dynamic-tools contract is model-specific and validated + in the serving layer (kimi_k3 only). + """ + if not isinstance(data, dict): + return data + for message in data.get("messages") or []: + if (isinstance(message, dict) and message.get("tools") is not None + and message.get("role") != "system"): + raise ValueError( + "Message-level `tools` are only allowed on system " + "messages.") + return data + + @model_validator(mode="before") + @classmethod + def check_logprobs(cls, data): + if (top_logprobs := data.get("top_logprobs")) is not None: + check_logprobs_limit("top_logprobs", top_logprobs) + if not data.get("logprobs"): + raise ValueError( + "logprobs must be true when using top_logprobs") + return data + + @field_validator("thinking_token_budget", mode="before") + @classmethod + def check_thinking_token_budget(cls, value): + return validate_thinking_token_budget(value) + + @model_validator(mode="before") + @classmethod + def check_suffix(cls, data): + if data.get("suffix"): + raise ValueError("suffix is not supported") + return data + + @field_validator("cache_salt") + @classmethod + def check_cache_salt_support(cls, v): + if v is not None: + if not isinstance(v, str) or not v.strip(): + raise ValueError( + "Parameter 'cache_salt' must be a non-empty string if provided." + ) + return v diff --git a/tensorrt_llm/serve/openai_protocol/common.py b/tensorrt_llm/serve/openai_protocol/common.py new file mode 100644 index 000000000000..38340321e712 --- /dev/null +++ b/tensorrt_llm/serve/openai_protocol/common.py @@ -0,0 +1,242 @@ +# 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. + +# Adapted from +# https://github.com/vllm-project/vllm/blob/4db5176d9758b720b05460c50ace3c01026eb158/vllm/entrypoints/openai/protocol.py + +from typing import Any, List, Literal, Optional, Union + +import xgrammar +from pydantic import (BaseModel, ConfigDict, Field, NonNegativeInt, PositiveInt, + field_validator, model_serializer, model_validator) +from typing_extensions import TypeAlias + +from tensorrt_llm.kv_hints import KvHintsEnvelope +from tensorrt_llm.kv_hints.protocol import admit_kv_hints +from tensorrt_llm.llmapi import DisaggScheduleStyle + +# Beam-search stopping mode as exposed over HTTP, mirroring the HuggingFace +# Transformers interface. The engine encodes it as an integer, but that encoding +# is an implementation detail and is kept out of the public schema. +EarlyStopping: TypeAlias = Union[bool, Literal["never"]] + + +def ensure_request_chat_template_allowed(request: Any, + allow_request_chat_template: bool): + if (getattr(request, "chat_template", None) is not None + and not allow_request_chat_template): + raise ValueError( + "chat_template cannot be supplied per request unless request-level " + "chat templates are enabled at server startup.") + + +class OpenAIBaseModel(BaseModel): + # OpenAI API does not allow extra fields & allow to initialize by both alias and field name + model_config = ConfigDict(extra="forbid", populate_by_name=True) + + @field_validator("kv_hints", mode="before", check_fields=False) + @classmethod + def admit_hint_envelope(cls, value: Any) -> Optional[KvHintsEnvelope]: + """Apply optional-hint admission before nested schema validation. + + Args: + value: The decoded hint field from a trusted ingress. + + Returns: + Supported metadata or None for ignored hints, including malformed + envelopes. Validation of other request fields remains unchanged. + """ + return admit_kv_hints(value) + + +class SpeculativeDecodingStats(OpenAIBaseModel): + """Per-request speculative-decoding acceptance for one generated sequence. + + Opt-in: emitted only when the server sets per_request_spec_decode_stats. + Absent entirely when the request drafted nothing. PyTorch backend only. + + ``mean acceptance length`` is deliberately not a field here: it is already + reported per choice as ``avg_decoded_tokens_per_iter``, and duplicating it + would let the two drift. Consumers derive it as + ``1 + total_accepted_draft_tokens / num_spec_steps``. + + Three identities hold on every emitted record, and a consumer may rely on + them: + + * ``sum(acceptance_histogram) == num_spec_steps`` + * ``sum(j * acceptance_histogram[j]) == total_accepted_draft_tokens`` + * ``total_accepted_draft_tokens <= total_draft_tokens`` + """ + + acceptance_rate: float = Field( + description="Accepted draft tokens divided by proposed draft tokens.") + total_accepted_draft_tokens: int = Field( + description="Draft tokens accepted across the request, excluding the " + "always-accepted bonus token.") + total_draft_tokens: int = Field( + description="Draft tokens proposed across the request. For tree " + "drafting this counts paths, not tree nodes, mirroring the " + "getMaxDraftPathLen clamp in updateNumTokensPerIteration.") + num_spec_steps: int = Field( + description="Verify steps performed for the request. Equals the sum of " + "acceptance_histogram.") + acceptance_histogram: List[int] = Field( + description="Dense histogram indexed by accepted-draft count: entry j " + "is the number of verify steps that accepted exactly j draft tokens. " + "Tree-agnostic -- it records output lengths per step and encodes no " + "parent/child structure.") + num_spec_tokens: Optional[int] = Field( + default=None, + description="Maximum draft length per step, when the run has a fixed " + "bound. None under draft_len_schedule, where the bound varies by batch " + "size.") + + +class _OmitsAbsentSpecDecodeStats(OpenAIBaseModel): + """Drops ``speculative_decoding`` from serialized output when it is absent. + + Per-request spec-decode stats are off by default, so without this every + response on the paths that serialize with a plain ``model_dump()`` -- the + non-streaming chat and completions responses, and the completions stream + (``exclude_unset=False``) -- would gain ``"speculative_decoding": null`` for + every user, whether or not they enabled ``per_request_spec_decode_stats``. + + Scoped to this one field on purpose. Blanket ``exclude_none`` would also + strip unrelated optional fields that clients may rely on being present, and + the dump calls are spread across the serving layer rather than funnelled + through one place where an ``exclude=`` argument could be applied. + """ + + @model_serializer(mode="wrap") + def _omit_absent_spec_decode_stats(self, handler: Any) -> Any: + data = handler(self) + if isinstance(data, dict) and data.get("speculative_decoding", + ...) is None: + data.pop("speculative_decoding", None) + return data + + +class StreamOptions(OpenAIBaseModel): + include_usage: Optional[bool] = True + continuous_usage_stats: Optional[bool] = False + + +class PromptTokensDetails(OpenAIBaseModel): + cached_tokens: int = 0 + + +class UsageInfo(OpenAIBaseModel): + prompt_tokens: int = 0 + total_tokens: int = 0 + completion_tokens: Optional[int] = 0 + prompt_tokens_details: Optional[PromptTokensDetails] = None + + +class ResponseFormat(OpenAIBaseModel): + type: Literal["text", "json", "json_schema", "json_object", "regex", "ebnf", + "structural_tag"] + schema: Optional[dict] = None + json_schema: Optional[dict] = None + regex: Optional[str] = None + ebnf: Optional[str] = None + format: Optional[xgrammar.structural_tag.Format] = None + + +class DisaggregatedParams(OpenAIBaseModel): + request_type: str + first_gen_tokens: Optional[List[int]] = None + first_gen_log_probs: Optional[List] = None + first_gen_logits: Optional[List] = None + ctx_request_id: Optional[int] = None + encoded_opaque_state: Optional[str] = None + draft_tokens: Optional[List[int]] = None + disagg_request_id: Optional[int] = None + ctx_dp_rank: Optional[int] = None + ctx_info_endpoint: Optional[str] = None + schedule_style: Optional[DisaggScheduleStyle] = None + ctx_usage: Optional[UsageInfo] = None + remote_tail_start: Optional[NonNegativeInt] = Field( + default=None, + description= + "Prompt-token boundary transferred by a remote-tail context worker.") + multimodal_positions: Optional[List[NonNegativeInt]] = Field( + default=None, + description="Prompt-token offset of each multimodal item.") + multimodal_lengths: Optional[List[PositiveInt]] = Field( + default=None, description="Prompt-token count of each multimodal item.") + # TODO(TRTLLM-12407): Multimodal E/PD over trtllm-serve needs these protocol fields too: + # encoder embedding handles, multimodal hashes, and optional mRoPE handles. + # Add them here and in to_disaggregated_params()/to_llm_disaggregated_params() + # before routing MM encoder -> context -> generation through OpenAI protocol. + # Orchestrator -> context-worker instruction: return prompt_token_ids as a + # base64 int32 buffer (prompt_token_ids_b64) instead of a JSON int array. + return_prompt_token_ids_b64: bool = False + # Context worker -> generation worker: the reasoning mode the context + # worker read off the prompt it rendered. The generation worker only sees + # prompt_token_ids, so it cannot resolve this for itself. + resolved_thinking: Optional[bool] = None + # Orchestrator -> generation worker: seconds since arrival at the orchestrator. + orig_elapsed_s: Optional[float] = None + + @model_validator(mode="after") + def validate_multimodal_spans(self) -> "DisaggregatedParams": + if (self.multimodal_positions is None) != (self.multimodal_lengths + is None): + raise ValueError( + "multimodal_positions and multimodal_lengths must be provided together" + ) + if (self.multimodal_positions is not None and len( + self.multimodal_positions) != len(self.multimodal_lengths)): + raise ValueError( + "multimodal_positions and multimodal_lengths must have the same length" + ) + return self + + +class ConversationParams(OpenAIBaseModel): + model_config = ConfigDict(extra="forbid", + populate_by_name=True, + validate_assignment=True) + + conversation_id: str = Field( + description=("Stable multi-turn conversation id used for routing"), ) + + # Forwarded as a header; excluded from the request body for older workers. + subagent_affinity_id: Optional[str] = Field( + default=None, + exclude=True, + description=("Server-private parent-session id for sub-agent routing " + "affinity; forwarded to workers as a header, never in the " + "body.")) + + @field_validator("conversation_id", mode="before") + @classmethod + def validate_conversation_id(cls, value: Any) -> str: + if value is None: + raise ValueError("conversation_id must be non-empty") + conversation_id = str(value).strip() + if not conversation_id: + raise ValueError("conversation_id must be non-empty") + return conversation_id + + +class ErrorResponse(OpenAIBaseModel): + object: str = "error" + message: str + type: str + param: Optional[str] = None + # HTTP status by default; a machine-readable string for conditions that + # define one (e.g. "context_length_exceeded", matching OpenAI). + code: Union[int, str] diff --git a/tensorrt_llm/serve/openai_protocol/completions.py b/tensorrt_llm/serve/openai_protocol/completions.py new file mode 100644 index 000000000000..93fa39db2497 --- /dev/null +++ b/tensorrt_llm/serve/openai_protocol/completions.py @@ -0,0 +1,279 @@ +# 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. + +# Adapted from +# https://github.com/vllm-project/vllm/blob/4db5176d9758b720b05460c50ace3c01026eb158/vllm/entrypoints/openai/protocol.py + +import time +import uuid +from typing import Dict, List, Optional, Union + +from pydantic import Field, field_validator, model_validator +from typing_extensions import Annotated + +from tensorrt_llm.executor.request import LoRARequest +from tensorrt_llm.kv_hints import KvHintsEnvelope +from tensorrt_llm.llmapi import SamplingParams +from tensorrt_llm.sampling_params import (check_logprobs_limit, + validate_thinking_token_budget) +from tensorrt_llm.serve.openai_protocol.common import ( + ConversationParams, DisaggregatedParams, EarlyStopping, OpenAIBaseModel, + ResponseFormat, SpeculativeDecodingStats, StreamOptions, UsageInfo, + _OmitsAbsentSpecDecodeStats) +from tensorrt_llm.serve.openai_protocol.sampling import ( + _early_stopping_to_int, _logit_bias_to_embedding_bias, + _record_sampling_params_request_fields, + _response_format_to_guided_decoding_params) + + +class CompletionLogProbs(OpenAIBaseModel): + text_offset: List[int] = Field(default_factory=list) + token_logprobs: List[Optional[float]] = Field(default_factory=list) + tokens: List[str] = Field(default_factory=list) + top_logprobs: List[Optional[Dict[str, float]]] = Field(default_factory=list) + + +class CompletionResponseChoice(_OmitsAbsentSpecDecodeStats): + index: int + text: str + token_ids: Optional[List[int]] = None + logprobs: Optional[CompletionLogProbs] = None + context_logits: Optional[Union[List[float], List[List[ + float]]]] = None # For reward models, the output is score logits instead of text. + finish_reason: Optional[str] = None + stop_reason: Optional[Union[int, str]] = Field( + default=None, + description=( + "The stop string or token id that caused the completion " + "to stop, None if the completion finished for some other reason " + "including encountering the EOS token"), + ) + disaggregated_params: Optional[DisaggregatedParams] = Field(default=None) + avg_decoded_tokens_per_iter: Optional[float] = Field(default=None) + speculative_decoding: Optional[SpeculativeDecodingStats] = Field( + default=None) + + +class CompletionResponse(OpenAIBaseModel): + id: str = Field(default_factory=lambda: f"cmpl-{str(uuid.uuid4().hex)}") + object: str = "text_completion" + created: int = Field(default_factory=lambda: int(time.time())) + model: str + choices: List[CompletionResponseChoice] + usage: UsageInfo + # Add prompt_tokens_ids to the response to remove the tokenization + # in the generation server in disaggreated serving + prompt_token_ids: Optional[Union[List[List[int]], List[int]]] = None + + +class CompletionResponseStreamChoice(_OmitsAbsentSpecDecodeStats): + index: int + text: str + token_ids: Optional[List[int]] = None + logprobs: Optional[CompletionLogProbs] = None + finish_reason: Optional[str] = None + stop_reason: Optional[Union[int, str]] = Field( + default=None, + description=( + "The stop string or token id that caused the completion " + "to stop, None if the completion finished for some other reason " + "including encountering the EOS token"), + ) + avg_decoded_tokens_per_iter: Optional[float] = Field(default=None) + speculative_decoding: Optional[SpeculativeDecodingStats] = Field( + default=None) + + +class CompletionStreamResponse(OpenAIBaseModel): + id: str = Field(default_factory=lambda: f"cmpl-{str(uuid.uuid4().hex)}") + object: str = "text_completion" + created: int = Field(default_factory=lambda: int(time.time())) + model: str + choices: List[CompletionResponseStreamChoice] + usage: Optional[UsageInfo] = Field(default=None) + + +class CompletionRequest(OpenAIBaseModel): + kv_hints: Optional[KvHintsEnvelope] = Field( + default=None, + repr=False, + description= + "Prototype KV hint metadata; transported without executing actions.", + ) + # Ordered by official OpenAI API documentation + # https://platform.openai.com/docs/api-reference/completions/create + model: str + prompt: Union[List[int], List[List[int]], str, List[str]] + best_of: Optional[int] = None + echo: Optional[bool] = False + frequency_penalty: Optional[float] = 0.0 + logit_bias: Optional[Dict[str, float]] = None + logprobs: Optional[int] = None + max_tokens: Optional[int] = None + n: int = 1 + presence_penalty: Optional[float] = 0.0 + seed: Optional[int] = Field(default=None) + stop: Optional[Union[str, List[str]]] = Field(default_factory=list) + stream: Optional[bool] = False + stream_options: Optional[StreamOptions] = None + suffix: Optional[str] = None + temperature: Optional[float] = None + top_p: Optional[float] = None + user: Optional[str] = None + lora_request: Optional[LoRARequest] = None + prompt_ignore_length: Optional[int] = 0 + + # doc: begin-completion-sampling-params + use_beam_search: bool = False + top_k: int = 0 + top_p_min: float = 0.0 + min_p: float = 0.0 + repetition_penalty: float = 1.0 + # Unset by default so the engine picks its own beam-search defaults, as it + # does for requests coming through the Python API. Sending concrete values + # here would put every served request on the length-normalized, exhaustive + # beam-search path. + length_penalty: Optional[float] = None + early_stopping: Optional[EarlyStopping] = None + stop_token_ids: Optional[List[int]] = Field(default_factory=list) + include_stop_str_in_output: bool = False + ignore_eos: bool = False + min_tokens: int = 0 + skip_special_tokens: bool = True + spaces_between_special_tokens: bool = True + truncate_prompt_tokens: Optional[Annotated[int, Field(ge=1)]] = None + return_context_logits: bool = False + detokenize: bool = True + thinking_token_budget: Optional[int] = None + # doc: end-completion-sampling-params + + # doc: begin-completion-extra-params + add_special_tokens: bool = Field( + default=True, + description=( + "If true (the default), special tokens (e.g. BOS) will be added to " + "the prompt."), + ) + response_format: Optional[ResponseFormat] = Field( + default=None, + description= + ("Similar to chat completion, this parameter specifies the format of output. " + "{'type': 'text'}, {'type': 'json'}, {'type': 'json_object'}, {'type': 'regex'}, " + "{'type': 'ebnf'}, {'type': 'structural_tag'} are supported."), + ) + + disaggregated_params: Optional[DisaggregatedParams] = Field( + default=None, + description=("Parameters for disaggregated serving"), + ) + priority: Optional[float] = Field( + default=None, + ge=0.0, + le=1.0, + description= + ("Scheduling priority in [0.0, 1.0]; higher is served first. Only honored " + "when the engine runs with scheduler_config.waiting_queue_policy=priority. " + "If unset, the engine default (0.5) is used."), + ) + conversation_params: Optional[ConversationParams] = Field( + default=None, + description=("Parameters for multi-turn conversation routing"), + ) + + # doc: end-completion-extra-params + + def to_sampling_params(self, + vocab_size: Optional[int] = None, + gather_generation_logits: bool = False, + backend: Optional[str] = None) -> SamplingParams: + sampling_logprobs = None + if self.logprobs: + if backend == "pytorch" or gather_generation_logits: + sampling_logprobs = self.logprobs + elif self.logprobs > 1: + raise ValueError( + "`logprobs` must be 1 or `gather_generation_logits` must be `True` to use `logprobs` > 1" + ) + + sampling_params = SamplingParams( + best_of=self.best_of, + frequency_penalty=self.frequency_penalty, + logprobs=sampling_logprobs, + max_tokens=self.max_tokens, + n=self.n, + presence_penalty=self.presence_penalty, + seed=self.seed, + stop=self.stop, + temperature=(self.temperature + if self.temperature is not None else 1.0), + top_p=(self.top_p if self.top_p is not None else 1.0), + prompt_ignore_length=self.prompt_ignore_length, + + # completion-sampling-params + use_beam_search=self.use_beam_search, + top_k=self.top_k, + top_p_min=self.top_p_min if self.top_p_min > 0 else None, + min_p=self.min_p, + repetition_penalty=self.repetition_penalty, + length_penalty=self.length_penalty, + early_stopping=_early_stopping_to_int(self.early_stopping), + stop_token_ids=self.stop_token_ids, + include_stop_str_in_output=self.include_stop_str_in_output, + ignore_eos=self.ignore_eos, + min_tokens=self.min_tokens, + skip_special_tokens=self.skip_special_tokens, + spaces_between_special_tokens=self.spaces_between_special_tokens, + truncate_prompt_tokens=self.truncate_prompt_tokens, + return_context_logits=self.return_context_logits, + guided_decoding=_response_format_to_guided_decoding_params( + self.response_format), + detokenize=self.detokenize, + thinking_token_budget=self.thinking_token_budget, + + # logits_bias + embedding_bias=_logit_bias_to_embedding_bias( + self.logit_bias, vocab_size), + + # completion-extra-params + add_special_tokens=self.add_special_tokens, + ) + _record_sampling_params_request_fields(self, sampling_params) + return sampling_params + + @model_validator(mode="before") + @classmethod + def check_logprobs(cls, data): + check_logprobs_limit("logprobs", data.get("logprobs")) + return data + + @field_validator("thinking_token_budget", mode="before") + @classmethod + def check_thinking_token_budget(cls, value): + return validate_thinking_token_budget(value) + + @model_validator(mode="before") + @classmethod + def validate_stream_options(cls, data): + if data.get("stream_options") and not data.get("stream"): + raise ValueError( + "Stream options can only be defined when stream is true.") + return data + + @model_validator(mode="before") + @classmethod + def check_suffix(cls, data): + if data.get("suffix"): + raise ValueError("suffix is not supported") + return data diff --git a/tensorrt_llm/serve/openai_protocol/control.py b/tensorrt_llm/serve/openai_protocol/control.py new file mode 100644 index 000000000000..21d67dcb8afa --- /dev/null +++ b/tensorrt_llm/serve/openai_protocol/control.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. + +# Adapted from +# https://github.com/vllm-project/vllm/blob/4db5176d9758b720b05460c50ace3c01026eb158/vllm/entrypoints/openai/protocol.py + +from typing import Dict, List, Literal, Optional + +from pydantic import Field, NonNegativeInt, PositiveInt + +from tensorrt_llm.serve.openai_protocol.common import OpenAIBaseModel + + +class MemoryUpdateRequest(OpenAIBaseModel): + tags: list[str] = Field( + min_length=1, + description="Memory tags to release/resume, e.g. ['model', 'kv_cache']") + + +class RuntimeMemoryUpdateRequest(OpenAIBaseModel): + tags: Optional[List[str]] = Field( + default=None, + description="Memory tags to release/resume; omit to select defaults") + + +class UpdateWeightsRequest(OpenAIBaseModel): + weights: Optional[Dict[str, str]] = Field( + default=None, + description="Weight handles dict, or None to finalize update") + + +ProfileActivity = Literal["CPU", "GPU", "CUDA_PROFILER"] + + +class StartProfileRequest(OpenAIBaseModel): + """Request body for the POST /start_profile endpoint.""" + output_dir: Optional[str] = Field( + default=None, + description="Directory where chrome traces are written. Defaults " + "to the TLLM_TORCH_PROFILER_DIR environment variable, " + "then /tmp.") + num_steps: Optional[PositiveInt] = Field( + default=None, + description="Number of iterations to profile. Must be >= 1 if " + "provided; if omitted, profiling runs until /stop_profile is " + "called. ``num_steps == 0`` is rejected because the profile " + "window would never close — the stop iteration would equal the " + "start iteration, profile_step() would discard the stop marker " + "as stale, and the window would run forever.") + start_step: NonNegativeInt = Field( + default=0, + description="Skip this many iterations before profiling begins. " + "Must be >= 0.") + activities: List[ProfileActivity] = Field( + default_factory=lambda: ["CPU", "GPU"], + description="Which activities to trace. Supported values: " + "'CPU', 'GPU', 'CUDA_PROFILER'. Unknown values are rejected at " + "schema validation time.") diff --git a/tensorrt_llm/serve/openai_protocol/disagg.py b/tensorrt_llm/serve/openai_protocol/disagg.py new file mode 100644 index 000000000000..b721220b7d26 --- /dev/null +++ b/tensorrt_llm/serve/openai_protocol/disagg.py @@ -0,0 +1,221 @@ +# 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. + +# Adapted from +# https://github.com/vllm-project/vllm/blob/4db5176d9758b720b05460c50ace3c01026eb158/vllm/entrypoints/openai/protocol.py + +import base64 +from typing import List, Optional, Union + +import torch + +from tensorrt_llm.llmapi import ConversationParams as LlmConversationParams +from tensorrt_llm.llmapi import DisaggregatedParams as LlmDisaggregatedParams +from tensorrt_llm.serve.openai_protocol.chat import (ChatCompletionRequest, + ChatCompletionResponse) +from tensorrt_llm.serve.openai_protocol.common import (ConversationParams, + DisaggregatedParams, + UsageInfo) +from tensorrt_llm.serve.openai_protocol.completions import (CompletionRequest, + CompletionResponse) + + +def encode_opaque_state(opaque_state: Optional[bytes]) -> Optional[str]: + if opaque_state is None: + return None + return base64.b64encode(opaque_state).decode("utf-8") + + +def decode_opaque_state(encoded_opaque_state: Optional[str]) -> Optional[bytes]: + if encoded_opaque_state is None: + return None + return base64.b64decode(encoded_opaque_state) + + +def _serialize_first_gen_log_probs( + first_gen_log_probs: Optional[list]) -> Optional[List]: + """Serialize ``list[dict[int, Logprob]] | list[float]`` to a JSON-safe form. + + - Default (verbose) format: each position is a ``dict[int, Logprob]`` and is + serialized as a list of ``{token_id, logprob, rank}`` dicts. + - Simple format: each position is a ``float``; passed through verbatim. + """ + if first_gen_log_probs is None: + return None + if not isinstance(first_gen_log_probs, list): + raise ValueError("first_gen_log_probs must be a list") + result = [] + for i, pos in enumerate(first_gen_log_probs): + if isinstance(pos, dict): + result.append([{ + "token_id": tid, + "logprob": lp.logprob, + "rank": lp.rank + } for tid, lp in pos.items()]) + elif isinstance(pos, (float, int)): + # Simple format: per-token sampled logprob. + result.append(float(pos)) + else: + raise ValueError( + f"first_gen_log_probs[{i}] must be a dict or float, got {type(pos)}" + ) + return result + + +def _deserialize_first_gen_log_probs( + serialized: Optional[List], ) -> Optional[list]: + """Inverse of :func:`_serialize_first_gen_log_probs`. + + Returns either ``list[dict[int, Logprob]]`` (default format) or + ``list[float]`` (simple format) depending on the serialized payload. + """ + if serialized is None: + return None + from tensorrt_llm.executor.result import Logprob + result = [] + for i, pos in enumerate(serialized): + if isinstance(pos, (float, int)): + result.append(float(pos)) + continue + if not isinstance(pos, list): + raise ValueError( + f"first_gen_log_probs[{i}] must be a list or float, got {type(pos)}" + ) + token_map = {} + for j, item in enumerate(pos): + if not isinstance(item, dict): + raise ValueError( + f"first_gen_log_probs[{i}][{j}] must be a dict") + if "token_id" not in item or "logprob" not in item: + raise ValueError( + f"first_gen_log_probs[{i}][{j}] missing required keys " + "'token_id' and/or 'logprob'") + token_map[item["token_id"]] = Logprob(logprob=item["logprob"], + rank=item.get("rank")) + result.append(token_map) + return result + + +def _serialize_first_gen_logits( + first_gen_logits: Optional[list], ) -> Optional[List]: + """Serialize list[torch.Tensor] to JSON-safe list[dict] with base64 data.""" + if first_gen_logits is None: + return None + result = [] + for i, tensor in enumerate(first_gen_logits): + t = tensor.contiguous().cpu() + if t.dtype == torch.bfloat16: + t = t.to(torch.float16) + np_array = t.numpy() + result.append({ + "data": base64.b64encode(np_array.tobytes()).decode(), + "shape": list(np_array.shape), + "dtype": str(np_array.dtype), + }) + return result + + +def _deserialize_first_gen_logits( + serialized: Optional[List], ) -> Optional[list]: + """Deserialize JSON list[dict] back to list[torch.Tensor].""" + if serialized is None: + return None + import numpy as np + result = [] + for i, item in enumerate(serialized): + if not isinstance(item, dict): + raise ValueError( + f"first_gen_logits[{i}] must be a dict, got {type(item)}") + for key in ("data", "shape", "dtype"): + if key not in item: + raise ValueError( + f"first_gen_logits[{i}] missing required key '{key}'") + np_array = np.frombuffer( + base64.b64decode(item["data"]), + dtype=np.dtype(item["dtype"]), + ).reshape(item["shape"]) + result.append(torch.from_numpy(np_array.copy())) + return result + + +def to_disaggregated_params( + tllm_disagg_params: LlmDisaggregatedParams) -> DisaggregatedParams: + if tllm_disagg_params is None: + return None + ctx_usage = tllm_disagg_params.ctx_usage + if ctx_usage is not None and not isinstance(ctx_usage, UsageInfo): + ctx_usage = UsageInfo.model_validate(ctx_usage) + return DisaggregatedParams( + request_type=tllm_disagg_params.request_type, + first_gen_tokens=tllm_disagg_params.first_gen_tokens, + first_gen_log_probs=_serialize_first_gen_log_probs( + tllm_disagg_params.first_gen_log_probs), + first_gen_logits=_serialize_first_gen_logits( + tllm_disagg_params.first_gen_logits), + ctx_request_id=tllm_disagg_params.ctx_request_id, + encoded_opaque_state=encode_opaque_state( + tllm_disagg_params.opaque_state), + draft_tokens=tllm_disagg_params.draft_tokens, + disagg_request_id=tllm_disagg_params.disagg_request_id, + ctx_dp_rank=tllm_disagg_params.ctx_dp_rank, + ctx_info_endpoint=tllm_disagg_params.ctx_info_endpoint, + schedule_style=tllm_disagg_params.schedule_style, + ctx_usage=ctx_usage, + remote_tail_start=tllm_disagg_params.remote_tail_start, + multimodal_positions=tllm_disagg_params.multimodal_positions, + multimodal_lengths=tllm_disagg_params.multimodal_lengths, + ) + + +def to_llm_disaggregated_params( + disaggregated_params: DisaggregatedParams) -> LlmDisaggregatedParams: + if disaggregated_params is None: + return None + ctx_usage = disaggregated_params.ctx_usage + return LlmDisaggregatedParams( + request_type=disaggregated_params.request_type, + first_gen_tokens=disaggregated_params.first_gen_tokens, + first_gen_log_probs=_deserialize_first_gen_log_probs( + disaggregated_params.first_gen_log_probs), + first_gen_logits=_deserialize_first_gen_logits( + disaggregated_params.first_gen_logits), + ctx_request_id=disaggregated_params.ctx_request_id, + opaque_state=decode_opaque_state( + disaggregated_params.encoded_opaque_state), + draft_tokens=disaggregated_params.draft_tokens, + disagg_request_id=disaggregated_params.disagg_request_id, + ctx_dp_rank=disaggregated_params.ctx_dp_rank, + ctx_info_endpoint=disaggregated_params.ctx_info_endpoint, + schedule_style=disaggregated_params.schedule_style, + ctx_usage=None if ctx_usage is None else ctx_usage.model_dump(), + remote_tail_start=disaggregated_params.remote_tail_start, + multimodal_positions=disaggregated_params.multimodal_positions, + multimodal_lengths=disaggregated_params.multimodal_lengths, + orig_elapsed_s=disaggregated_params.orig_elapsed_s, + ) + + +def to_llm_conversation_params( + conversation_params: Optional[ConversationParams] +) -> Optional[LlmConversationParams]: + if conversation_params is None: + return None + return LlmConversationParams( + conversation_id=conversation_params.conversation_id) + + +UCompletionRequest = Union[CompletionRequest, ChatCompletionRequest] + +UCompletionResponse = Union[CompletionResponse, ChatCompletionResponse] diff --git a/tensorrt_llm/serve/openai_protocol/embeddings.py b/tensorrt_llm/serve/openai_protocol/embeddings.py new file mode 100644 index 000000000000..3f3c9af7dfe7 --- /dev/null +++ b/tensorrt_llm/serve/openai_protocol/embeddings.py @@ -0,0 +1,90 @@ +# 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. + +# Adapted from +# https://github.com/vllm-project/vllm/blob/4db5176d9758b720b05460c50ace3c01026eb158/vllm/entrypoints/openai/protocol.py + +import time +import uuid +from typing import List, Literal, Optional, Union + +from pydantic import Field, PositiveInt + +from tensorrt_llm.serve.openai_protocol.common import OpenAIBaseModel + + +class EmbeddingRequest(OpenAIBaseModel): + # OpenAI-compatible embeddings request. + # https://platform.openai.com/docs/api-reference/embeddings/create + model: str = Field(description="The model to use for the embedding.") + input: Union[str, List[str], List[int], List[List[int]]] = Field( + description="Text(s) or pre-tokenized token-id list(s) to embed. Accepts " + "the four OpenAI forms: str, list[str], list[int], list[list[int]].") + encoding_format: Literal["float", "base64"] = Field( + default="float", + description="Format of the returned embedding: a JSON array of floats " + "(\"float\") or a base64-encoded string of packed little-endian float32 " + "values (\"base64\").") + dimensions: Optional[PositiveInt] = Field( + default=None, + description="Number of dimensions the resulting embedding should have. " + "Only supported by Matryoshka-trained text-embedding models; rejected " + "otherwise.") + user: Optional[str] = Field( + default=None, + description="A stable identifier for the end-user (OpenAI-compatible; " + "accepted but unused).") + add_special_tokens: bool = Field( + default=True, + description= + "TRT-LLM extension: whether to add the model's special tokens " + "(e.g. [CLS]/[SEP]) when tokenizing string inputs. Default matches " + "llm.encode(); encoder models such as BERT generally need them.") + + +class EmbeddingResponseData(OpenAIBaseModel): + index: int = Field( + description="Index of this embedding in the request's input list.") + object: Literal["embedding"] = "embedding" + embedding: Union[List[float], str] = Field( + description= + "The embedding: a list of floats for encoding_format=\"float\", " + "or a base64 string of packed little-endian float32 values for " + "encoding_format=\"base64\".") + + +class EmbeddingUsageInfo(OpenAIBaseModel): + """Token usage for embeddings. + + OpenAI's embeddings `usage` object contains only `prompt_tokens` and + `total_tokens` (no `completion_tokens`), unlike chat/completions. + """ + prompt_tokens: int = Field(default=0, + description="Number of tokens in the input.") + total_tokens: int = Field( + default=0, + description="Total tokens consumed (equal to prompt_tokens for " + "embeddings).") + + +class EmbeddingResponse(OpenAIBaseModel): + id: str = Field(default_factory=lambda: f"embd-{str(uuid.uuid4().hex)}") + object: Literal["list"] = "list" + created: int = Field(default_factory=lambda: int(time.time())) + model: str = Field(description="The model used for the embedding.") + data: List[EmbeddingResponseData] = Field( + description="The list of embeddings, one per input item.") + usage: EmbeddingUsageInfo = Field( + description="Token usage for the request.") diff --git a/tensorrt_llm/serve/openai_protocol/kv_cache.py b/tensorrt_llm/serve/openai_protocol/kv_cache.py new file mode 100644 index 000000000000..19b28adf9149 --- /dev/null +++ b/tensorrt_llm/serve/openai_protocol/kv_cache.py @@ -0,0 +1,66 @@ +# 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. + +# Adapted from +# https://github.com/vllm-project/vllm/blob/4db5176d9758b720b05460c50ace3c01026eb158/vllm/entrypoints/openai/protocol.py + +from typing import List, Optional + +from tensorrt_llm.serve.openai_protocol.chat import ChatCompletionMessageParam +from tensorrt_llm.serve.openai_protocol.common import OpenAIBaseModel +from tensorrt_llm.serve.openai_protocol.tools import ChatCompletionToolsParam + + +class KVCacheTruncateRequest(OpenAIBaseModel): + model: str + messages: List[ChatCompletionMessageParam] = [] + messages_to_retain: List[ChatCompletionMessageParam] = [] + tools: Optional[List[ChatCompletionToolsParam]] = None + add_generation_prompt: Optional[bool] = True + documents: Optional[list] = None + chat_template: Optional[str] = None + chat_template_kwargs: Optional[dict] = None + reasoning_effort: Optional[str] = None + tool_choice: Optional[str] = None + + +class KVCacheTruncateTokensRequest(OpenAIBaseModel): + """Token-level analog of :class:`KVCacheTruncateRequest`. + + Posted to the ``/_control/kv_cache/truncate_tokens`` endpoint. Bypasses + the chat-template tokenization path entirely: callers (e.g. + trace replay) that already operate on raw token ids hand the server a + batch of full token-id sequences plus the prefix length to retain for + each, and the server trims the corresponding KV-cache blocks in the + radix tree. Each ``prefixes[i]`` becomes one + :class:`tensorrt_llm.executor.request.TruncateKVCacheRequest` on the + executor's KV-cache control queue, so one batched HTTP request fans out + into N atomic per-prefix radix-tree mutations while paying a single + network round-trip. + + Fields: + model: Target model name. Accepted for router compatibility but + ignored -- the server uses its own loaded model. + prefixes: Batch of full token-id sequences. Each entry is one + cached prefix whose KV-cache blocks should be truncated. + num_tokens_to_keep: Per-prefix number of leading tokens to retain, + parallel to ``prefixes`` (``num_tokens_to_keep[i]`` applies to + ``prefixes[i]``). Blocks beyond this length are dropped from the + radix tree. + """ + + model: Optional[str] = None + prefixes: List[List[int]] + num_tokens_to_keep: List[int] diff --git a/tensorrt_llm/serve/openai_protocol/models.py b/tensorrt_llm/serve/openai_protocol/models.py new file mode 100644 index 000000000000..b1410416e6fc --- /dev/null +++ b/tensorrt_llm/serve/openai_protocol/models.py @@ -0,0 +1,36 @@ +# 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. + +# Adapted from +# https://github.com/vllm-project/vllm/blob/4db5176d9758b720b05460c50ace3c01026eb158/vllm/entrypoints/openai/protocol.py + +import time +from typing import List + +from pydantic import Field + +from tensorrt_llm.serve.openai_protocol.common import OpenAIBaseModel + + +class ModelCard(OpenAIBaseModel): + id: str + object: str = "model" + created: int = Field(default_factory=lambda: int(time.time())) + owned_by: str = "tensorrt_llm" + + +class ModelList(OpenAIBaseModel): + object: str = "list" + data: List[ModelCard] = Field(default_factory=list) diff --git a/tensorrt_llm/serve/openai_protocol/responses.py b/tensorrt_llm/serve/openai_protocol/responses.py new file mode 100644 index 000000000000..22261787c42b --- /dev/null +++ b/tensorrt_llm/serve/openai_protocol/responses.py @@ -0,0 +1,565 @@ +# 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. + +# Adapted from +# https://github.com/vllm-project/vllm/blob/4db5176d9758b720b05460c50ace3c01026eb158/vllm/entrypoints/openai/protocol.py + +import base64 +import time +import uuid +from typing import Any, List, Literal, Optional, Union + +from openai.types.responses import ( + ResponseCodeInterpreterCallCodeDeltaEvent, + ResponseCodeInterpreterCallCodeDoneEvent, + ResponseCodeInterpreterCallCompletedEvent, + ResponseCodeInterpreterCallInProgressEvent, + ResponseCodeInterpreterCallInterpretingEvent, ResponseCompletedEvent, + ResponseContentPartAddedEvent, ResponseContentPartDoneEvent, + ResponseCreatedEvent, ResponseFunctionToolCall, ResponseInProgressEvent, + ResponseInputItemParam, ResponseOutputItem, ResponseOutputItemAddedEvent, + ResponseOutputItemDoneEvent, ResponsePrompt, ResponseReasoningItem, + ResponseReasoningTextDeltaEvent, ResponseReasoningTextDoneEvent, + ResponseStatus, ResponseTextConfig, ResponseWebSearchCallCompletedEvent, + ResponseWebSearchCallInProgressEvent, ResponseWebSearchCallSearchingEvent) +from openai.types.responses.response import IncompleteDetails, ToolChoice +from openai.types.responses.tool import Tool +from openai.types.shared import Metadata, Reasoning +from pydantic import (ConfigDict, Field, field_validator, model_serializer, + model_validator) +from typing_extensions import TypeAlias + +from tensorrt_llm.kv_hints import KvHintsEnvelope +from tensorrt_llm.llmapi import SamplingParams +from tensorrt_llm.sampling_params import validate_thinking_token_budget +from tensorrt_llm.serve.openai_protocol.common import (ConversationParams, + DisaggregatedParams, + OpenAIBaseModel) +from tensorrt_llm.serve.openai_protocol.sampling import ( + _record_sampling_params_request_fields, + _response_format_text_config_to_guided_decoding_params) + +# The trailing dict keeps a client's own item types from failing the request +# outright. Codex multi-agent sessions carry "agent_message" items, which no +# SDK model describes; without a permissive member the union rejects the whole +# input and every request from a spawned agent comes back 422, so agent +# collaboration cannot work at all. Known shapes still match their typed +# member first; see _response_output_item_to_chat_completion_message for how +# an unrecognised item is replayed. +ResponseInputOutputItem: TypeAlias = Union[ResponseInputItemParam, + ResponseReasoningItem, + ResponseFunctionToolCall, dict[str, + Any]] + +# Roles whose message items map to EasyInputMessageParam / Message, both of +# which forbid ``id``. An assistant message maps to ResponseOutputMessageParam +# instead, which *requires* ``id`` and ``status`` - stripping ``id`` there +# leaves the item matching no variant of the input union, so a conversation +# breaks as soon as it carries one assistant turn. +# +# Module level on purpose: pydantic turns a leading-underscore class attribute +# into a ModelPrivateAttr, so referring to it from a validator would raise +# "argument of type 'ModelPrivateAttr' is not iterable" at request time. +_ID_STRIPPED_ROLES = ("user", "system", "developer") + + +def _drop_explicit_nulls(value, _depth=0): + """Recursively drop dict keys whose value is an explicit ``null``. + + On the OpenAI wire an unset optional field is omitted, so an explicit + ``null`` carries no information — but the vendored Responses item types + validate ``null`` differently from omitted and reject it. Dropping + null-valued keys restores the omit-unset wire shape for clients that + serialize every unset optional as ``null`` (litellm among them). + """ + if _depth > 12: + return value + if isinstance(value, dict): + return { + k: _drop_explicit_nulls(v, _depth + 1) + for k, v in value.items() if v is not None + } + if isinstance(value, list): + return [_drop_explicit_nulls(v, _depth + 1) for v in value] + return value + + +def _materialize_validator_iterators(value, _depth=0): + """Recursively replace pydantic ValidatorIterator objects with lists. + + Matched by type name rather than by import: the class lives in the + pydantic_core extension module and is not part of its public API. + """ + if _depth > 12: + return value + if type(value).__name__ == "ValidatorIterator": + return [_materialize_validator_iterators(v, _depth + 1) for v in value] + if isinstance(value, dict): + for key, item in value.items(): + value[key] = _materialize_validator_iterators(item, _depth + 1) + return value + if isinstance(value, list): + return [_materialize_validator_iterators(v, _depth + 1) for v in value] + return value + + +class ResponsesRequest(OpenAIBaseModel): + kv_hints: Optional[KvHintsEnvelope] = Field( + default=None, + repr=False, + description= + "Prototype KV hint metadata; transported without executing actions.", + ) + # Ordered by official OpenAI API documentation + # https://platform.openai.com/docs/api-reference/responses/create + # + # Unlike the rest of the OpenAI models this one accepts unknown fields. + # Real Responses API clients attach evolving telemetry and routing keys - + # Codex CLI sends client_metadata and prompt_cache_key, for instance - and + # rejecting the whole request over a field the server would ignore anyway + # makes those clients unusable for no benefit. + model_config = ConfigDict(extra="allow", populate_by_name=True) + + background: Optional[bool] = False + include: Optional[list[ + Literal[ + "code_interpreter_call.outputs", + "computer_call_output.output.image_url", + "file_search_call.results", + "message.input_image.image_url", + "message.output_text.logprobs", + "reasoning.encrypted_content", + ], + ]] = None + input: Union[str, list[ResponseInputOutputItem]] + instructions: Optional[str] = None + + @field_validator("input", mode="before") + @classmethod + def _drop_unsupported_input_item_keys(cls, value): + """Strip per-item keys the vendored Responses item types reject. + + Clients echo back items exactly as the server emitted them, so input + message items arrive carrying the ``id`` the server assigned, which + ``EasyInputMessageParam`` / ``Message`` forbid. That fails the whole + request over a field with no prompt content. + + Assistant messages and tool-call items keep their ids: the former + needs it to validate, the latter uses it to pair calls with results. + """ + if not isinstance(value, list): + return value + + # Must run before the per-item shaping below: the item types reject + # explicit nulls that mean "unset" (see _drop_explicit_nulls). + value = [_drop_explicit_nulls(item) for item in value] + + def _with_annotations(part): + """output_text requires annotations; clients often omit it.""" + if (isinstance(part, dict) and part.get("type") == "output_text" + and "annotations" not in part): + return {**part, "annotations": []} + return part + + cleaned = [] + for item in value: + # A client may send a bare content part as a top-level item, not + # only nested inside a message. + item = _with_annotations(item) + if isinstance(item, dict) and item.get("type") in (None, "message"): + role = item.get("role") + if "id" in item and role in _ID_STRIPPED_ROLES: + item = {k: v for k, v in item.items() if k != "id"} + elif role == "assistant": + # ResponseOutputMessageParam requires status, and each + # output_text part requires annotations. Clients rebuild an + # assistant turn from the streamed deltas rather than from + # the server's content-part objects, so both routinely + # arrive absent and fail the whole request. Default them + # rather than reject a conversation over fields that carry + # no prompt content. + item = dict(item) + item.setdefault("status", "completed") + content = item.get("content") + if isinstance(content, list): + item["content"] = [ + { + **part, "annotations": part.get( + "annotations", []) + } if isinstance(part, dict) + and part.get("type") == "output_text" else part + for part in content + ] + cleaned.append(item) + return cleaned + + @field_validator("input", mode="after") + @classmethod + def _materialize_lazy_item_content(cls, value): + """Force lazily-validated sequences anywhere in ``input`` into lists. + + Several vendored Responses item types declare sequence fields as + ``Iterable[...]`` - ``ResponseOutputMessageParam.content`` and + ``ResponseOutputTextParam.annotations`` among them - and pydantic + validates an Iterable lazily into a ``ValidatorIterator``. That object + cannot be pickled, so a request carrying structured input dies with + + cannot pickle 'pydantic_core._pydantic_core.ValidatorIterator' + + the moment it is handed to a postprocess worker. It is also + single-consumption, so even with workers disabled anything that reads + the field twice sees an empty sequence the second time. + + The walk has to be recursive: the lazy fields are nested inside items, + e.g. input[1].content[0].annotations. + """ + return _materialize_validator_iterators(value) + + max_output_tokens: Optional[int] = None + max_tool_calls: Optional[int] = None + metadata: Optional[Metadata] = None + model: str + parallel_tool_calls: Optional[bool] = False + previous_response_id: Optional[str] = None + prompt: Optional[ResponsePrompt] = None + reasoning: Optional[Reasoning] = None + thinking_token_budget: Optional[int] = None + service_tier: Literal["auto", "default", "flex", "scale", + "priority"] = "auto" + store: Optional[bool] = True + stream: Optional[bool] = False + temperature: Optional[float] = None + text: Optional[ResponseTextConfig] = None + tool_choice: ToolChoice = "auto" + tools: list[Tool] = Field(default_factory=list) + top_logprobs: Optional[int] = 0 + top_p: Optional[float] = None + truncation: Optional[Literal["auto", "disabled"]] = "disabled" + user: Optional[str] = None + + request_id: str = Field( + default_factory=lambda: f"resp_{str(uuid.uuid4().hex)}", + description=( + "The request_id related to this request. If the caller does " + "not set it, a random_uuid will be generated. This id is used " + "through out the inference process and return in response."), + ) + + # doc: begin-responses-extra-params + # TensorRT-LLM extensions carrying the disaggregated-serving handoff between + # the orchestrator and the workers; clients do not set them. + disaggregated_params: Optional[DisaggregatedParams] = Field( + default=None, + description=("Parameters for disaggregated serving"), + ) + conversation_params: Optional[ConversationParams] = Field( + default=None, + description=("Parameters for multi-turn conversation routing"), + ) + # The context worker's tokenized prompt, so the generation worker does not + # render it again; the _b64 form is a base64 int32 buffer. + prompt_token_ids: Optional[List[int]] = None + prompt_token_ids_b64: Optional[str] = None + + # doc: end-responses-extra-params + + @field_validator("prompt_token_ids_b64") + @classmethod + def _check_prompt_token_ids_b64(cls, value: Optional[str]) -> Optional[str]: + # Decoded during routing and preprocessing, after validation; a bad + # value must be a 400 here rather than an internal error there. + if not value: + return value + try: + raw = base64.b64decode(value, validate=True) + except ValueError as error: + raise ValueError( + "prompt_token_ids_b64 is not valid base64") from error + if len(raw) % 4: + raise ValueError( + "prompt_token_ids_b64 is not a buffer of int32 token ids") + return value + + def relayed_prompt_token_ids(self) -> Optional[List[int]]: + """The prompt the orchestrator relayed from the context worker, if any. + + Only a generation-only handoff carries one. On any other request the + fields are ignored, so the prompt is rendered from the request and the + response echoes what was actually prefilled. + """ + params = self.disaggregated_params + if params is None or params.request_type != "generation_only": + return None + if self.prompt_token_ids is not None: + return self.prompt_token_ids + if not self.prompt_token_ids_b64: + return None + import numpy as np + return np.frombuffer(base64.b64decode(self.prompt_token_ids_b64), + dtype=np.int32).tolist() + + _DEFAULT_SAMPLING_PARAMS = { + "temperature": 1.0, + "top_p": 1.0, + } + + def to_sampling_params( + self, + default_sampling_params: Optional[dict] = None, + reasoning_parser: Optional[str] = None, + ) -> SamplingParams: + max_tokens = None + if self.max_output_tokens is not None: + max_tokens = self.max_output_tokens + + default_sampling_params = default_sampling_params or {} + if (temperature := self.temperature) is None: + temperature = default_sampling_params.get( + "temperature", self._DEFAULT_SAMPLING_PARAMS["temperature"]) + if (top_p := self.top_p) is None: + top_p = default_sampling_params.get( + "top_p", self._DEFAULT_SAMPLING_PARAMS["top_p"]) + stop_token_ids = default_sampling_params.get("stop_token_ids", None) + + # Structured output + guided_decoding = None + if self.text is not None and self.text.format is not None: + guided_decoding = _response_format_text_config_to_guided_decoding_params( + self.text.format, reasoning_parser=reasoning_parser) + + sampling_params = SamplingParams( + temperature=temperature, + top_p=top_p, + max_tokens=max_tokens, + logprobs=self.top_logprobs, + stop_token_ids=stop_token_ids, + guided_decoding=guided_decoding, + thinking_token_budget=self.thinking_token_budget, + ) + _record_sampling_params_request_fields(self, sampling_params) + return sampling_params + + @model_validator(mode="before") + @classmethod + def _null_top_level_fields_mean_unset(cls, data): + """Drop explicitly-null top-level fields so they take their defaults. + + Shallow on purpose: ``input`` items are scrubbed recursively by their + own validator, and the remaining structured fields are Optional + throughout. A null required field (``model``, ``input``) still + reports as missing. + """ + if not isinstance(data, dict): + return data + return {k: v for k, v in data.items() if v is not None} + + @model_validator(mode="before") + @classmethod + def hoist_additional_tools(cls, data: Any) -> Any: + """Move the tools of an ``additional_tools`` input item into ``tools``. + + Codex declares its tools that way; hoisted here, they are validated as + tools and resolve when the reply's calls are parsed. Malformed shapes + are left in place for validation to reject. + """ + if not isinstance(data, dict): + return data + items = data.get("input") + tools = data.get("tools") + if not isinstance(items, list) or not isinstance( + tools, (list, type(None))): + return data + hoisted = [] + kept = [] + for item in items: + if (isinstance(item, dict) + and item.get("type") == "additional_tools" + and isinstance(item.get("tools"), list)): + hoisted.extend(item["tools"]) + else: + kept.append(item) + if not hoisted: + return data + data = dict(data) + data["input"] = kept + data["tools"] = list(data.get("tools") or []) + hoisted + return data + + @model_validator(mode="before") + @classmethod + def validate_background(cls, data): + if not data.get("background"): + return data + if not data.get("store", True): + raise ValueError("background can only be used when `store` is true") + return data + + @model_validator(mode="before") + @classmethod + def validate_prompt(cls, data): + if data.get("prompt") is not None: + raise ValueError("prompt template is not supported") + return data + + @field_validator("thinking_token_budget", mode="before") + @classmethod + def check_thinking_token_budget(cls, value): + return validate_thinking_token_budget(value) + + +class InputTokensDetails(OpenAIBaseModel): + cached_tokens: int + # Required by the openai SDK models that re-validate this payload when it + # is embedded in a streaming event. Omitting it fails validation while the + # response is being streamed, which truncates the stream with no + # terminating event and leaves the client waiting forever. Prompt cache + # writes are not tracked separately, so this is reported as zero. + cache_write_tokens: int = 0 + + +class OutputTokensDetails(OpenAIBaseModel): + reasoning_tokens: int + + +class ResponseUsage(OpenAIBaseModel): + input_tokens: int + input_tokens_details: InputTokensDetails + output_tokens: int + output_tokens_details: OutputTokensDetails + total_tokens: int + + +StreamingResponsesResponse: TypeAlias = Union[ + ResponseCreatedEvent, + ResponseInProgressEvent, + ResponseCompletedEvent, + ResponseOutputItemAddedEvent, + ResponseOutputItemDoneEvent, + ResponseContentPartAddedEvent, + ResponseContentPartDoneEvent, + ResponseReasoningTextDeltaEvent, + ResponseReasoningTextDoneEvent, + ResponseCodeInterpreterCallInProgressEvent, + ResponseCodeInterpreterCallCodeDeltaEvent, + ResponseWebSearchCallInProgressEvent, + ResponseWebSearchCallSearchingEvent, + ResponseWebSearchCallCompletedEvent, + ResponseCodeInterpreterCallCodeDoneEvent, + ResponseCodeInterpreterCallInterpretingEvent, + ResponseCodeInterpreterCallCompletedEvent, +] + +_RESPONSES_HANDOFF_FIELDS = ( + "disaggregated_params", + "prompt_token_ids", + "prompt_token_ids_b64", + "finish_reason", +) + + +class ResponsesResponse(OpenAIBaseModel): + id: str = Field(default_factory=lambda: f"resp_{str(uuid.uuid4().hex)}") + created_at: int = Field(default_factory=lambda: int(time.time())) + # error: Optional[ResponseError] = None + incomplete_details: Optional[IncompleteDetails] = None + instructions: Optional[str] = None + metadata: Optional[Metadata] = None + model: str + object: Literal["response"] = "response" + output: list[ResponseOutputItem] + parallel_tool_calls: bool + temperature: float + tool_choice: ToolChoice + tools: list[Tool] + top_p: float + background: bool + max_output_tokens: Optional[int] = None + max_tool_calls: Optional[int] = None + previous_response_id: Optional[str] = None + prompt: Optional[ResponsePrompt] = None + reasoning: Optional[Reasoning] = None + service_tier: Literal["auto", "default", "flex", "scale", "priority"] + status: ResponseStatus + text: Optional[ResponseTextConfig] = None + top_logprobs: int + truncation: Literal["auto", "disabled"] + usage: Optional[ResponseUsage] = None + user: Optional[str] = None + + # TensorRT-LLM extensions, set only on a context-only response: the + # disaggregated handoff, and the engine's raw finish reason, which tells the + # orchestrator whether a generation phase is pending ("length" and + # "not_finished" both map to status "incomplete"). + disaggregated_params: Optional[DisaggregatedParams] = Field(default=None) + prompt_token_ids: Optional[List[int]] = None + prompt_token_ids_b64: Optional[str] = None + finish_reason: Optional[str] = None + + @model_serializer(mode="wrap") + def _omit_absent_handoff(self, handler: Any) -> Any: + # The handoff fields are unset on every response a client sees, so they + # are left out rather than serialized as null in each body and event. + data = handler(self) + if isinstance(data, dict): + for key in _RESPONSES_HANDOFF_FIELDS: + if data.get(key, ...) is None: + data.pop(key) + return data + + @classmethod + def from_request( + cls, + request: ResponsesRequest, + sampling_params: SamplingParams, + model_name: str, + created_time: int, + output: list[ResponseOutputItem], + status: ResponseStatus, + usage: Optional[ResponseUsage] = None, + ) -> "ResponsesResponse": + return cls( + id=request.request_id, + created_at=created_time, + instructions=request.instructions, + metadata=request.metadata, + model=model_name, + output=output, + parallel_tool_calls=request.parallel_tool_calls, + temperature=sampling_params.temperature, + tool_choice=request.tool_choice, + tools=request.tools, + top_p=sampling_params.top_p, + background=request.background, + max_output_tokens=sampling_params.max_tokens, + max_tool_calls=request.max_tool_calls, + previous_response_id=request.previous_response_id, + prompt=request.prompt, + reasoning=request.reasoning, + service_tier=request.service_tier, + status=status, + text=request.text, + top_logprobs=sampling_params.logprobs, + truncation=request.truncation, + user=request.user, + usage=usage, + ) + + +class ResponsesStreamResponse(OpenAIBaseModel): + response: ResponsesResponse + sequence_number: int + type: Literal["response.created", "response.in_progress", + "response.completed", "response.failed", + "response.incomplete"] diff --git a/tensorrt_llm/serve/openai_protocol/sampling.py b/tensorrt_llm/serve/openai_protocol/sampling.py new file mode 100644 index 000000000000..12fd044942cd --- /dev/null +++ b/tensorrt_llm/serve/openai_protocol/sampling.py @@ -0,0 +1,221 @@ +# 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. + +# Adapted from +# https://github.com/vllm-project/vllm/blob/4db5176d9758b720b05460c50ace3c01026eb158/vllm/entrypoints/openai/protocol.py + +import math +from typing import Any, Dict, Optional + +import torch +from openai.types.responses import ResponseFormatTextConfig + +from tensorrt_llm.llmapi import GuidedDecodingParams, SamplingParams +from tensorrt_llm.llmapi.reasoning_parser import ReasoningParserFactory +from tensorrt_llm.serve.openai_protocol.common import (EarlyStopping, + OpenAIBaseModel, + ResponseFormat) +from tensorrt_llm.serve.serving_extensions import structured_output_format_for + +_LOGIT_BIAS_MIN = -100.0 + +_LOGIT_BIAS_MAX = 100.0 + + +def _early_stopping_to_int(value: Optional[EarlyStopping]) -> Optional[int]: + """Translate the HF-style tri-state into the engine's integer encoding. + + Mirrors ``BeamSearchEarlyStop``: ``False`` -> 0, ``True`` -> 1, + ``"never"`` -> 2. ``None`` stays unset so the engine picks its default. + """ + if value is None: + return None + if value == "never": + return 2 + # NB: bool is a subclass of int, so this yields 1/0 for True/False. + return int(value) + + +def _logit_bias_to_embedding_bias( + logit_bias: Optional[Dict[str, float]], + vocab_size: Optional[int]) -> Optional[torch.Tensor]: + """Convert OpenAI logit_bias dict to embedding_bias tensor for sampling.""" + if logit_bias is None: + return None + if vocab_size is None: + raise ValueError( + "logit_bias requires a tokenizer, but the server was started " + "without one (e.g. num_postprocess_workers > 0). " + "Remove logit_bias from your request or set num_postprocess_workers=0." + ) + elif vocab_size <= 0: + raise ValueError("vocab_size must be positive when logit_bias is used") + + # Create 1D zeros tensor as expected by executor API (will be unsqueezed to [1, vocab_size] internally) + embedding_bias = torch.zeros(vocab_size, dtype=torch.float32) + + # Apply biases for specified token IDs + for token_str, bias in logit_bias.items(): + try: + token_id = int(token_str) + except ValueError as e: + if "invalid literal" in str(e): + raise ValueError( + f"Invalid logit_bias key '{token_str}': must be a valid integer token ID" + ) + raise + if not 0 <= token_id < vocab_size: + raise ValueError( + f"Token ID {token_id} out of vocabulary range [0, {vocab_size})" + ) + try: + bias_value = float(bias) + except (TypeError, ValueError) as e: + raise ValueError( + f"logit_bias value for token ID {token_id} must be a number" + ) from e + if not math.isfinite(bias_value): + raise ValueError(f"logit_bias value for token ID {token_id} " + "must be finite") + if not _LOGIT_BIAS_MIN <= bias_value <= _LOGIT_BIAS_MAX: + raise ValueError( + f"logit_bias value for token ID {token_id} must be in " + f"[{_LOGIT_BIAS_MIN:g}, {_LOGIT_BIAS_MAX:g}]") + embedding_bias[token_id] = bias_value + + return embedding_bias + + +def _response_format_to_guided_decoding_params( + response_format: Optional[ResponseFormat], + reasoning_parser: Optional[str] = None, + chat_template_kwargs: Optional[Dict[str, Any]] = None, +) -> Optional[GuidedDecodingParams]: + if response_format is None: + guided_decoding_params = None + elif response_format.type == "text": + guided_decoding_params = None + elif response_format.type == "json": + if response_format.schema is None: + raise ValueError( + f"response_format.schema is required for response_format.type == {response_format.type!r}, but got None." + ) + guided_decoding_params = GuidedDecodingParams( + json=response_format.schema) + elif response_format.type == "json_schema": + if response_format.json_schema is None: + raise ValueError( + f"response_format.json_schema is required for response_format.type == {response_format.type!r}, but got None." + ) + # OpenAI API spec wraps the actual JSON schema under a "schema" key. + # Extract the actual schema if the wrapper format is used. + json_schema = response_format.json_schema + if isinstance(json_schema, dict) and "schema" in json_schema: + json_schema = json_schema["schema"] + guided_decoding_params = GuidedDecodingParams(json=json_schema) + elif response_format.type == "json_object": + guided_decoding_params = GuidedDecodingParams(json_object=True) + elif response_format.type == "regex": + if response_format.regex is None: + raise ValueError( + f"response_format.regex is required for response_format.type == {response_format.type!r}, but got None." + ) + guided_decoding_params = GuidedDecodingParams( + regex=response_format.regex) + elif response_format.type == "ebnf": + if response_format.ebnf is None: + raise ValueError( + f"response_format.ebnf is required for response_format.type == {response_format.type!r}, but got None." + ) + guided_decoding_params = GuidedDecodingParams( + grammar=response_format.ebnf) + elif response_format.type == "structural_tag": + guided_decoding_params = GuidedDecodingParams( + structural_tag=response_format.model_dump_json(by_alias=True, + exclude_none=True)) + else: + raise ValueError(f"Unsupported response format: {response_format.type}") + + if guided_decoding_params is None or reasoning_parser is None: + return guided_decoding_params + + if guided_decoding_params.structural_tag is not None: + return guided_decoding_params + + # Adapt guided_decoding_params for reasoning parser + if guided_decoding_params.json is not None: + content = { + "type": "json_schema", + "json_schema": guided_decoding_params.json + } + elif guided_decoding_params.json_object: + content = {"type": "json_schema", "json_schema": {"type": "object"}} + elif guided_decoding_params.regex is not None: + content = {"type": "regex", "pattern": guided_decoding_params.regex} + elif guided_decoding_params.grammar is not None: + content = {"type": "grammar", "grammar": guided_decoding_params.grammar} + + extension_format = structured_output_format_for(reasoning_parser) + if extension_format is not None: + # A registered per-model serving extension owns the placement of the + # constraint relative to this model's reasoning markup. + stag_format = extension_format(content, chat_template_kwargs) + if stag_format is None: + return guided_decoding_params + else: + # Force thinking and then trigger user constraint + parser = ReasoningParserFactory.create_reasoning_parser( + reasoning_parser) + stag_format = { + "type": + "sequence", + "elements": [ + { + "type": "tag", + "begin": parser.reasoning_start, + "content": { + "type": "any_text" + }, + "end": parser.reasoning_end, + }, + content, + ], + } + + stag_format = ResponseFormat(type="structural_tag", format=stag_format) + return GuidedDecodingParams(structural_tag=stag_format.model_dump_json( + by_alias=True, exclude_none=True)) + + +def _response_format_text_config_to_guided_decoding_params( + text_format: Optional[ResponseFormatTextConfig], + reasoning_parser: Optional[str] = None, +) -> Optional[GuidedDecodingParams]: + if text_format is None: + return None + + resp_format = ResponseFormat(type=text_format.type, + json_schema=getattr(text_format, "schema_", + None)) + return _response_format_to_guided_decoding_params( + resp_format, reasoning_parser=reasoning_parser) + + +def _record_sampling_params_request_fields( + request: OpenAIBaseModel, sampling_params: SamplingParams) -> None: + """Preserve explicitly supplied fields across protocol defaulting.""" + sampling_params._set_request_provided_fields( + field_name for field_name in request.model_fields_set + if getattr(request, field_name, None) is not None) diff --git a/tensorrt_llm/serve/openai_protocol/tokenize.py b/tensorrt_llm/serve/openai_protocol/tokenize.py new file mode 100644 index 000000000000..19110ae6baad --- /dev/null +++ b/tensorrt_llm/serve/openai_protocol/tokenize.py @@ -0,0 +1,59 @@ +# 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. + +# Adapted from +# https://github.com/vllm-project/vllm/blob/4db5176d9758b720b05460c50ace3c01026eb158/vllm/entrypoints/openai/protocol.py + +from typing import List, Optional + +from tensorrt_llm.serve.openai_protocol.common import OpenAIBaseModel + + +class TokenizeRequest(OpenAIBaseModel): + """Request body for the ``POST /_internal/tokenize`` endpoint. + + Asks the server to encode some text into the token-id sequence its + loaded tokenizer would produce, without running any generation. + + Only plain-text ``prompt`` tokenization is supported. A ``messages`` + path was intentionally left out: to be meaningful it must reuse the full + ``/v1/chat/completions`` preprocessing (server chat template, tools, + documents, template kwargs, model-type resolution, multimodal handling), + otherwise the returned token ids would silently diverge from the actual + inference prompt. It can be added once that path is shared. + + Fields: + model: Target model name. Accepted for compatibility with + multi-model routers but ignored -- the server always tokenizes + against its own loaded model. + prompt: Plain text to tokenize directly via ``tokenizer.encode``. + """ + + model: Optional[str] = None + prompt: str + + +class TokenizeResponse(OpenAIBaseModel): + """Response body for the ``POST /_internal/tokenize`` endpoint. + + Fields: + count: Number of tokens in the encoded sequence. Always populated; + clients needing only the length can ignore ``tokens``. + tokens: The encoded token ids. Kept ``Optional`` so future + count-only variants can omit it without breaking the schema. + """ + + count: int + tokens: Optional[List[int]] = None diff --git a/tensorrt_llm/serve/openai_protocol/tools.py b/tensorrt_llm/serve/openai_protocol/tools.py new file mode 100644 index 000000000000..8c591aa1dbc1 --- /dev/null +++ b/tensorrt_llm/serve/openai_protocol/tools.py @@ -0,0 +1,141 @@ +# 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. + +# Adapted from +# https://github.com/vllm-project/vllm/blob/4db5176d9758b720b05460c50ace3c01026eb158/vllm/entrypoints/openai/protocol.py + +import uuid +from typing import Any, Dict, Literal, Optional + +from pydantic import Field, model_validator + +from tensorrt_llm.serve.openai_protocol.common import OpenAIBaseModel + + +class FunctionCall(OpenAIBaseModel): + name: str + arguments: str + + +class DeltaFunctionCall(OpenAIBaseModel): + name: Optional[str] = None + arguments: Optional[str] = None + + +class ToolCall(OpenAIBaseModel): + id: str = Field( + default_factory=lambda: f"chatcmpl-tool-{str(uuid.uuid4().hex)}") + type: Literal["function"] = "function" + function: FunctionCall + + +class DeltaToolCall(OpenAIBaseModel): + id: Optional[str] = None + type: Literal["function"] = "function" + index: int + function: Optional[DeltaFunctionCall] = None + + +# Maximum total number of `enum` values across all properties of one tool +# function's parameters schema. Mirrors the cap reference OpenAI-compatible +# platforms enforce; unbounded enums also inflate the guided-decoding +# grammar compiled for strict tool calls. +TOOL_PARAM_MAX_ENUM_VALUES = 1000 + +# JSON Schema keywords whose value is instance data rather than a subschema. +# An `enum` key nested inside these is a value (e.g. a default that happens to +# be `{"enum": [...]}`), not an enum constraint, so it must not count. +_SCHEMA_INSTANCE_KEYWORDS = frozenset({"default", "const", "examples"}) + +# JSON Schema keywords whose value is a map of {name: subschema}. The names are +# user-controlled (a property can be literally named `default` or `enum`), so +# recurse into the values as subschemas without treating the names as keywords. +_SCHEMA_MAP_KEYWORDS = frozenset({ + "properties", "patternProperties", "$defs", "definitions", + "dependentSchemas" +}) + + +def _count_schema_enum_values(schema: Any) -> int: + """Recursively count `enum` *constraint* entries in a JSON-schema fragment. + + Only `enum` keywords in schema positions are counted. `enum` keys that are + instance data -- nested under `default`/`const`/`examples`, or the name of a + property -- are ignored, so a valid schema is not rejected for values that + are not enum constraints. + """ + count = 0 + if isinstance(schema, dict): + for key, value in schema.items(): + if key == "enum" and isinstance(value, list): + count += len(value) + elif key in _SCHEMA_INSTANCE_KEYWORDS: + continue + elif key in _SCHEMA_MAP_KEYWORDS and isinstance(value, dict): + for subschema in value.values(): + count += _count_schema_enum_values(subschema) + else: + count += _count_schema_enum_values(value) + elif isinstance(schema, list): + for item in schema: + count += _count_schema_enum_values(item) + return count + + +class FunctionDefinition(OpenAIBaseModel): + name: str + description: Optional[str] = None + parameters: Optional[Dict[str, Any]] = None + strict: Optional[bool] = None + + @model_validator(mode="after") + def check_enum_value_cap(self): + if self.parameters is not None: + num_enum_values = _count_schema_enum_values(self.parameters) + if num_enum_values > TOOL_PARAM_MAX_ENUM_VALUES: + raise ValueError( + f"tool function {self.name!r} declares {num_enum_values} " + f"enum values across its parameters schema; the maximum " + f"is {TOOL_PARAM_MAX_ENUM_VALUES}.") + return self + + +class ChatCompletionToolsParam(OpenAIBaseModel): + type: Literal["function"] = "function" + function: FunctionDefinition + + +class NamespacedFunctionToolParam(ChatCompletionToolsParam): + """A tool that a Responses request declared inside a namespace. + + ``function`` is what a chat template without namespaces is given: the + tool as a function named ``.``. That is all a dump + carries. The declaration itself is kept beside it, out of every dump, for + a model whose renderer reads namespaces (``ServingExtension.serialize_tool``). + """ + + namespace: str = Field(exclude=True) + namespace_description: Optional[str] = Field(default=None, exclude=True) + declared_name: str = Field(exclude=True) + declared_description: Optional[str] = Field(default=None, exclude=True) + + +class ChatCompletionNamedFunction(OpenAIBaseModel): + name: str + + +class ChatCompletionNamedToolChoiceParam(OpenAIBaseModel): + function: ChatCompletionNamedFunction + type: Literal["function"] = "function" From df75a2bd03806d1f1895d66c7bef1c45b2212bac Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Tue, 3 Nov 2026 06:48:08 +0000 Subject: [PATCH 22/30] [TRTLLM-17265][test] Pin the openai_protocol facade: completeness, laziness and layering tests/unittest/serve/openai_protocol/test_openai_protocol_facade.py keeps the package tensorrt_llm/serve/openai_protocol/ honest: every name the former single-file module defined (plus DisaggScheduleStyle) is reachable through the facade and is the object its submodule defines; the facade table is frozen at those names; __all__, dir() and the TYPE_CHECKING block agree with the table; unknown names raise AttributeError; the submodule import graph stays inside the allowed acyclic graph and no submodule imports the package itself; importing the package loads no submodule, importing tools loads neither chat nor responses nor completions, and reading ChatCompletionRequest does not load the serving extensions; the private helpers that other tests import through the package stay importable. tests/unittest/serve/api/visual_gen/test_openai_protocol_facade_visual_gen.py holds the VisualGen assertions under a visual_gen/ path so that CI selection runs them for VisualGen-only changes: the eight VisualGen models map to tensorrt_llm.serve.api.visual_gen.protocol, nothing but a lookup of one of them loads that module, and the facade names it only as strings in its table, never in an import statement. tests/unittest/serve/test_layering.py gains the contract that the protocol package imports no handler module (servers, handlers and routers import the models, never the reverse), and tensorrt_llm/serve/AGENTS.md tells new code to import the submodule that defines a model rather than the package. Both test files are collected by the existing unittest/serve and unittest/serve/api/visual_gen/ entries of l0_cpu.yml and carry the cpu_only marker. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- tensorrt_llm/serve/AGENTS.md | 6 +- .../test_openai_protocol_facade_visual_gen.py | 181 +++++++ .../serve/openai_protocol/__init__.py | 14 + .../test_openai_protocol_facade.py | 465 ++++++++++++++++++ tests/unittest/serve/test_layering.py | 29 ++ 5 files changed, 693 insertions(+), 2 deletions(-) create mode 100644 tests/unittest/serve/api/visual_gen/test_openai_protocol_facade_visual_gen.py create mode 100644 tests/unittest/serve/openai_protocol/__init__.py create mode 100644 tests/unittest/serve/openai_protocol/test_openai_protocol_facade.py diff --git a/tensorrt_llm/serve/AGENTS.md b/tensorrt_llm/serve/AGENTS.md index 1e47ff2651b7..7272155ec6fa 100644 --- a/tensorrt_llm/serve/AGENTS.md +++ b/tensorrt_llm/serve/AGENTS.md @@ -23,9 +23,11 @@ Name files and symbols here, never line numbers. - 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 +- Request and response models: `serve/openai_protocol/.py`. Import the submodule, not the + package; the package's name table is frozen. 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`. + Check: `tests/unittest/api_stability/test_serve_api.py` and + `tests/unittest/serve/openai_protocol/test_openai_protocol_facade.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 diff --git a/tests/unittest/serve/api/visual_gen/test_openai_protocol_facade_visual_gen.py b/tests/unittest/serve/api/visual_gen/test_openai_protocol_facade_visual_gen.py new file mode 100644 index 000000000000..5a0659e88c7a --- /dev/null +++ b/tests/unittest/serve/api/visual_gen/test_openai_protocol_facade_visual_gen.py @@ -0,0 +1,181 @@ +# 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 VisualGen names of the ``tensorrt_llm.serve.openai_protocol`` facade. + +The facade maps the eight VisualGen request and response models to +``tensorrt_llm.serve.api.visual_gen.protocol`` and loads that module only when +one of them is read. It never imports it, not even under ``TYPE_CHECKING``, so +that CI selection can run only the VisualGen tests for VisualGen-only changes. +This file lives under a ``visual_gen/`` path so that those runs include it; the +other facade checks are in +``tests/unittest/serve/openai_protocol/test_openai_protocol_facade.py``. +""" + +import ast +import importlib +import json +import subprocess +import sys +import textwrap +from pathlib import Path + +import pytest + +pytestmark = [pytest.mark.cpu_only, pytest.mark.threadleak(enabled=False)] + +_FACADE_MODULE = "tensorrt_llm.serve.openai_protocol" +_FACADE_PATH = "tensorrt_llm/serve/openai_protocol/__init__.py" +_VISUAL_GEN_PACKAGE = "tensorrt_llm.serve.api.visual_gen" +_VISUAL_GEN_MODULE = "tensorrt_llm.serve.api.visual_gen.protocol" +_VISUAL_GEN_NAMES = ( + "ImageGenerationRequest", + "ImageEditRequest", + "ImageObject", + "ImageGenerationResponse", + "MediaReferenceItem", + "VideoGenerationRequest", + "VideoJob", + "VideoJobList", +) +_SUBMODULES = ( + "common", + "sampling", + "tools", + "models", + "tokenize", + "embeddings", + "completions", + "chat", + "kv_cache", + "responses", + "control", + "disagg", +) + + +def _run_in_fresh_interpreter(code: str) -> dict: + """Run ``code`` in a new interpreter and return the JSON it prints last.""" + result = subprocess.run( + [sys.executable, "-c", textwrap.dedent(code)], + capture_output=True, + text=True, + timeout=600, + ) + assert result.returncode == 0, ( + f"the probe interpreter exited with {result.returncode}; stdout tail: " + f"{result.stdout[-1000:]!r}; stderr tail: {result.stderr[-2000:]!r}" + ) + return json.loads(result.stdout.strip().splitlines()[-1]) + + +def test_visual_gen_names_map_to_their_module(): + facade = importlib.import_module(_FACADE_MODULE) + visual_gen = importlib.import_module(_VISUAL_GEN_MODULE) + problems = [] + for name in _VISUAL_GEN_NAMES: + mapped = facade._NAME_TO_MODULE.get(name) + same_object = getattr(facade, name, None) is getattr(visual_gen, name) + if mapped != _VISUAL_GEN_MODULE or not same_object: + problems.append( + f"{_FACADE_MODULE} must re-export {name} lazily from {_VISUAL_GEN_MODULE}; " + f"map '{name}' to that module in _NAME_TO_MODULE of {_FACADE_PATH}." + ) + assert problems == [], "\n".join(problems) + + +def test_protocol_submodules_do_not_load_visual_gen(): + out = _run_in_fresh_interpreter( + f""" + import importlib, json, sys + + target = {_VISUAL_GEN_MODULE!r} + loaded_by = [] + import tensorrt_llm.serve.openai_protocol + if target in sys.modules: + loaded_by.append("import tensorrt_llm.serve.openai_protocol") + for sub in {list(_SUBMODULES)!r}: + importlib.import_module(f"tensorrt_llm.serve.openai_protocol.{{sub}}") + if target in sys.modules: + loaded_by.append(f"import tensorrt_llm.serve.openai_protocol.{{sub}}") + from tensorrt_llm.serve.openai_protocol import ChatCompletionRequest + if target in sys.modules: + loaded_by.append("from tensorrt_llm.serve.openai_protocol import ChatCompletionRequest") + print(json.dumps({{"loaded_by": loaded_by}})) + """ + ) + assert out["loaded_by"] == [], ( + f"{out['loaded_by'][0]} loaded {_VISUAL_GEN_MODULE}; only a lookup of a VisualGen " + "name through the facade may load it." + ) + + +def _resolve(node: ast.ImportFrom, package: str) -> str: + """The absolute module an ``ImportFrom`` names, relative imports resolved against ``package``.""" + if node.level == 0: + return node.module or "" + parts = package.split(".") + base = parts[: len(parts) - (node.level - 1)] + if node.module: + base.append(node.module) + return ".".join(base) + + +def _under_visual_gen(module: str) -> bool: + return module == _VISUAL_GEN_PACKAGE or module.startswith(_VISUAL_GEN_PACKAGE + ".") + + +def test_facade_names_visual_gen_only_as_strings(): + facade = importlib.import_module(_FACADE_MODULE) + tree = ast.parse(Path(facade.__file__).read_text()) + problems = [] + for node in ast.walk(tree): + modules = [] + if isinstance(node, ast.Import): + modules = [alias.name for alias in node.names] + elif isinstance(node, ast.ImportFrom): + resolved = _resolve(node, _FACADE_MODULE) + modules = [resolved] + [f"{resolved}.{alias.name}" for alias in node.names] + for module in modules: + if _under_visual_gen(module): + problems.append( + f"{_FACADE_PATH} imports {module} at line {node.lineno}. The facade may name " + "VisualGen modules only as strings in _NAME_TO_MODULE; an import, even under " + "TYPE_CHECKING, breaks the VisualGen import isolation that lets CI run only " + "VisualGen tests for VisualGen-only changes." + ) + break + table_values = set() + for node in tree.body: + if isinstance(node, ast.Assign) and any( + isinstance(target, ast.Name) and target.id == "_NAME_TO_MODULE" + for target in node.targets + ): + table_values = { + id(constant) + for constant in ast.walk(node.value) + if isinstance(constant, ast.Constant) + } + for node in ast.walk(tree): + if ( + isinstance(node, ast.Constant) + and node.value == _VISUAL_GEN_MODULE + and id(node) not in table_values + ): + problems.append( + f"{_FACADE_PATH} names {_VISUAL_GEN_MODULE} at line {node.lineno} outside " + "_NAME_TO_MODULE; the facade refers to the VisualGen module only as the value of " + "a row of that table." + ) + assert problems == [], "\n".join(problems) diff --git a/tests/unittest/serve/openai_protocol/__init__.py b/tests/unittest/serve/openai_protocol/__init__.py new file mode 100644 index 000000000000..467079831e16 --- /dev/null +++ b/tests/unittest/serve/openai_protocol/__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/openai_protocol/test_openai_protocol_facade.py b/tests/unittest/serve/openai_protocol/test_openai_protocol_facade.py new file mode 100644 index 000000000000..f6b53a269a14 --- /dev/null +++ b/tests/unittest/serve/openai_protocol/test_openai_protocol_facade.py @@ -0,0 +1,465 @@ +# 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. +"""Completeness, laziness and layering of the ``tensorrt_llm.serve.openai_protocol`` facade. + +``tensorrt_llm/serve/openai_protocol/`` holds one submodule per API behind a +lazy facade (``__init__.py``) that re-exports, on first access, every name the +former single-file module had. These tests keep the facade's name table +complete and frozen, the facade lazy, and the submodule import graph acyclic. +The VisualGen names are checked by a sibling test under the VisualGen test +directory, so that CI selection for VisualGen-only pull requests runs it. +""" + +import ast +import importlib +import json +import re +import subprocess +import sys +import textwrap +from pathlib import Path + +import pytest +from test_common.serve.shims import ShimRow, check_identity + +pytestmark = [pytest.mark.cpu_only, pytest.mark.threadleak(enabled=False)] + +_PACKAGE = "tensorrt_llm.serve.openai_protocol" +_PACKAGE_DIR = "tensorrt_llm/serve/openai_protocol" +_FACADE_PATH = f"{_PACKAGE_DIR}/__init__.py" + +# Every name the former single-file module defined, with the module that defines it now, plus +# the one name it imported that code in the repository imports through it. Frozen: a new model +# is imported from its submodule (tensorrt_llm/serve/AGENTS.md), never added here. +_FORMER_NAMES: dict[str, str] = { + # common + "ConversationParams": "tensorrt_llm.serve.openai_protocol.common", + "DisaggregatedParams": "tensorrt_llm.serve.openai_protocol.common", + "EarlyStopping": "tensorrt_llm.serve.openai_protocol.common", + "ErrorResponse": "tensorrt_llm.serve.openai_protocol.common", + "OpenAIBaseModel": "tensorrt_llm.serve.openai_protocol.common", + "PromptTokensDetails": "tensorrt_llm.serve.openai_protocol.common", + "ResponseFormat": "tensorrt_llm.serve.openai_protocol.common", + "SpeculativeDecodingStats": "tensorrt_llm.serve.openai_protocol.common", + "StreamOptions": "tensorrt_llm.serve.openai_protocol.common", + "UsageInfo": "tensorrt_llm.serve.openai_protocol.common", + "_OmitsAbsentSpecDecodeStats": "tensorrt_llm.serve.openai_protocol.common", + "ensure_request_chat_template_allowed": "tensorrt_llm.serve.openai_protocol.common", + # sampling + "_LOGIT_BIAS_MAX": "tensorrt_llm.serve.openai_protocol.sampling", + "_LOGIT_BIAS_MIN": "tensorrt_llm.serve.openai_protocol.sampling", + "_early_stopping_to_int": "tensorrt_llm.serve.openai_protocol.sampling", + "_logit_bias_to_embedding_bias": "tensorrt_llm.serve.openai_protocol.sampling", + "_record_sampling_params_request_fields": "tensorrt_llm.serve.openai_protocol.sampling", + "_response_format_text_config_to_guided_decoding_params": ( + "tensorrt_llm.serve.openai_protocol.sampling" + ), + "_response_format_to_guided_decoding_params": "tensorrt_llm.serve.openai_protocol.sampling", + # tools + "ChatCompletionNamedFunction": "tensorrt_llm.serve.openai_protocol.tools", + "ChatCompletionNamedToolChoiceParam": "tensorrt_llm.serve.openai_protocol.tools", + "ChatCompletionToolsParam": "tensorrt_llm.serve.openai_protocol.tools", + "DeltaFunctionCall": "tensorrt_llm.serve.openai_protocol.tools", + "DeltaToolCall": "tensorrt_llm.serve.openai_protocol.tools", + "FunctionCall": "tensorrt_llm.serve.openai_protocol.tools", + "FunctionDefinition": "tensorrt_llm.serve.openai_protocol.tools", + "NamespacedFunctionToolParam": "tensorrt_llm.serve.openai_protocol.tools", + "TOOL_PARAM_MAX_ENUM_VALUES": "tensorrt_llm.serve.openai_protocol.tools", + "ToolCall": "tensorrt_llm.serve.openai_protocol.tools", + "_SCHEMA_INSTANCE_KEYWORDS": "tensorrt_llm.serve.openai_protocol.tools", + "_SCHEMA_MAP_KEYWORDS": "tensorrt_llm.serve.openai_protocol.tools", + "_count_schema_enum_values": "tensorrt_llm.serve.openai_protocol.tools", + # models + "ModelCard": "tensorrt_llm.serve.openai_protocol.models", + "ModelList": "tensorrt_llm.serve.openai_protocol.models", + # tokenize + "TokenizeRequest": "tensorrt_llm.serve.openai_protocol.tokenize", + "TokenizeResponse": "tensorrt_llm.serve.openai_protocol.tokenize", + # embeddings + "EmbeddingRequest": "tensorrt_llm.serve.openai_protocol.embeddings", + "EmbeddingResponse": "tensorrt_llm.serve.openai_protocol.embeddings", + "EmbeddingResponseData": "tensorrt_llm.serve.openai_protocol.embeddings", + "EmbeddingUsageInfo": "tensorrt_llm.serve.openai_protocol.embeddings", + # completions + "CompletionLogProbs": "tensorrt_llm.serve.openai_protocol.completions", + "CompletionRequest": "tensorrt_llm.serve.openai_protocol.completions", + "CompletionResponse": "tensorrt_llm.serve.openai_protocol.completions", + "CompletionResponseChoice": "tensorrt_llm.serve.openai_protocol.completions", + "CompletionResponseStreamChoice": "tensorrt_llm.serve.openai_protocol.completions", + "CompletionStreamResponse": "tensorrt_llm.serve.openai_protocol.completions", + # chat + "ChatCompletionContentPartParam": "tensorrt_llm.serve.openai_protocol.chat", + "ChatCompletionLogProb": "tensorrt_llm.serve.openai_protocol.chat", + "ChatCompletionLogProbs": "tensorrt_llm.serve.openai_protocol.chat", + "ChatCompletionLogProbsContent": "tensorrt_llm.serve.openai_protocol.chat", + "ChatCompletionMessageParam": "tensorrt_llm.serve.openai_protocol.chat", + "ChatCompletionRequest": "tensorrt_llm.serve.openai_protocol.chat", + "ChatCompletionResponse": "tensorrt_llm.serve.openai_protocol.chat", + "ChatCompletionResponseChoice": "tensorrt_llm.serve.openai_protocol.chat", + "ChatCompletionResponseStreamChoice": "tensorrt_llm.serve.openai_protocol.chat", + "ChatCompletionStreamResponse": "tensorrt_llm.serve.openai_protocol.chat", + "ChatCompletionThinkingParam": "tensorrt_llm.serve.openai_protocol.chat", + "ChatMessage": "tensorrt_llm.serve.openai_protocol.chat", + "CustomChatCompletionContentPartParam": "tensorrt_llm.serve.openai_protocol.chat", + "CustomChatCompletionMessageParam": "tensorrt_llm.serve.openai_protocol.chat", + "DeltaMessage": "tensorrt_llm.serve.openai_protocol.chat", + "DynamicToolsSystemMessageParam": "tensorrt_llm.serve.openai_protocol.chat", + "ReasoningAssistantMessage": "tensorrt_llm.serve.openai_protocol.chat", + # kv_cache + "KVCacheTruncateRequest": "tensorrt_llm.serve.openai_protocol.kv_cache", + "KVCacheTruncateTokensRequest": "tensorrt_llm.serve.openai_protocol.kv_cache", + # responses + "InputTokensDetails": "tensorrt_llm.serve.openai_protocol.responses", + "OutputTokensDetails": "tensorrt_llm.serve.openai_protocol.responses", + "ResponseInputOutputItem": "tensorrt_llm.serve.openai_protocol.responses", + "ResponseUsage": "tensorrt_llm.serve.openai_protocol.responses", + "ResponsesRequest": "tensorrt_llm.serve.openai_protocol.responses", + "ResponsesResponse": "tensorrt_llm.serve.openai_protocol.responses", + "ResponsesStreamResponse": "tensorrt_llm.serve.openai_protocol.responses", + "StreamingResponsesResponse": "tensorrt_llm.serve.openai_protocol.responses", + "_ID_STRIPPED_ROLES": "tensorrt_llm.serve.openai_protocol.responses", + "_RESPONSES_HANDOFF_FIELDS": "tensorrt_llm.serve.openai_protocol.responses", + "_drop_explicit_nulls": "tensorrt_llm.serve.openai_protocol.responses", + "_materialize_validator_iterators": "tensorrt_llm.serve.openai_protocol.responses", + # control + "MemoryUpdateRequest": "tensorrt_llm.serve.openai_protocol.control", + "ProfileActivity": "tensorrt_llm.serve.openai_protocol.control", + "RuntimeMemoryUpdateRequest": "tensorrt_llm.serve.openai_protocol.control", + "StartProfileRequest": "tensorrt_llm.serve.openai_protocol.control", + "UpdateWeightsRequest": "tensorrt_llm.serve.openai_protocol.control", + # disagg + "UCompletionRequest": "tensorrt_llm.serve.openai_protocol.disagg", + "UCompletionResponse": "tensorrt_llm.serve.openai_protocol.disagg", + "_deserialize_first_gen_log_probs": "tensorrt_llm.serve.openai_protocol.disagg", + "_deserialize_first_gen_logits": "tensorrt_llm.serve.openai_protocol.disagg", + "_serialize_first_gen_log_probs": "tensorrt_llm.serve.openai_protocol.disagg", + "_serialize_first_gen_logits": "tensorrt_llm.serve.openai_protocol.disagg", + "decode_opaque_state": "tensorrt_llm.serve.openai_protocol.disagg", + "encode_opaque_state": "tensorrt_llm.serve.openai_protocol.disagg", + "to_disaggregated_params": "tensorrt_llm.serve.openai_protocol.disagg", + "to_llm_conversation_params": "tensorrt_llm.serve.openai_protocol.disagg", + "to_llm_disaggregated_params": "tensorrt_llm.serve.openai_protocol.disagg", + # imported by the former module and imported from it by code in the repository + "DisaggScheduleStyle": "tensorrt_llm.llmapi", +} + +# The VisualGen models the former module also defined. Their identity rows live with the other +# VisualGen serve tests; here they are names only, so that this file names no VisualGen module. +_VISUAL_GEN_NAMES = frozenset( + { + "ImageGenerationRequest", + "ImageEditRequest", + "ImageObject", + "ImageGenerationResponse", + "MediaReferenceItem", + "VideoGenerationRequest", + "VideoJob", + "VideoJobList", + } +) + +# Private helpers that tests import through the package, with one importing test file each. +_TEST_IMPORTED_PRIVATE_NAMES = { + "_serialize_first_gen_log_probs": "tests/unittest/disaggregated/test_openai_disagg_service.py", + "_deserialize_first_gen_log_probs": ( + "tests/unittest/disaggregated/test_openai_disagg_service.py" + ), + "_serialize_first_gen_logits": "tests/unittest/disaggregated/test_openai_disagg_service.py", + "_deserialize_first_gen_logits": "tests/unittest/disaggregated/test_openai_disagg_service.py", + "_logit_bias_to_embedding_bias": "tests/unittest/llmapi/apps/test_harmony_parsing.py", + "_response_format_to_guided_decoding_params": ( + "tests/unittest/llmapi/apps/test_serving_extensions.py" + ), +} + +# Every name the facade re-exports: the former module's definitions and DisaggScheduleStyle. +_EXPECTED = frozenset(_FORMER_NAMES) | _VISUAL_GEN_NAMES + +# The facade is the package's own interface, not a compatibility shim: it never warns and the +# old path is not banned, so these rows check identity only. +_FACADE_ROWS = tuple( + ShimRow(_PACKAGE, name, module, retire_in="not retired", ban_old_path=False, warns=False) + for name, module in _FORMER_NAMES.items() +) + +# Which submodules each submodule may import ("a: {b}" means a imports from b). Acyclic. +_ALLOWED_SIBLING_IMPORTS = { + "common": frozenset(), + "sampling": frozenset({"common"}), + "tools": frozenset({"common"}), + "models": frozenset({"common"}), + "tokenize": frozenset({"common"}), + "embeddings": frozenset({"common"}), + "control": frozenset({"common"}), + "completions": frozenset({"common", "sampling"}), + "responses": frozenset({"common", "sampling"}), + "chat": frozenset({"common", "sampling", "tools"}), + "kv_cache": frozenset({"common", "tools", "chat"}), + "disagg": frozenset({"common", "completions", "chat"}), +} + + +@pytest.fixture(scope="module") +def facade(): + return importlib.import_module(_PACKAGE) + + +@pytest.fixture(scope="module") +def package_dir(facade) -> Path: + return Path(facade.__file__).resolve().parent + + +def _resolve(node: ast.ImportFrom, package: str) -> str: + """The absolute module an ``ImportFrom`` names, relative imports resolved against ``package``.""" + if node.level == 0: + return node.module or "" + parts = package.split(".") + base = parts[: len(parts) - (node.level - 1)] + if node.module: + base.append(node.module) + return ".".join(base) + + +def _run_in_fresh_interpreter(code: str) -> dict: + """Run ``code`` in a new interpreter and return the JSON it prints last.""" + result = subprocess.run( + [sys.executable, "-c", textwrap.dedent(code)], + capture_output=True, + text=True, + timeout=600, + ) + assert result.returncode == 0, ( + f"the probe interpreter exited with {result.returncode}; stdout tail: " + f"{result.stdout[-1000:]!r}; stderr tail: {result.stderr[-2000:]!r}" + ) + return json.loads(result.stdout.strip().splitlines()[-1]) + + +def _with_facade_fix(problem: str) -> str: + """Prefix a shim-identity message with what to change in the facade.""" + match = re.search(rf"`{re.escape(_PACKAGE)}\.(\w+)`", problem) + name = match.group(1) if match else "" + origin = _FORMER_NAMES.get(name, "") + if "no longer exists" in problem: + fix = ( + f"`{name}` was a top-level name of tensorrt_llm/serve/openai_protocol.py and must " + f"stay importable from {_PACKAGE}. Define it in {origin} and add '{name}': " + f"'{origin}' to _NAME_TO_MODULE in {_FACADE_PATH}." + ) + else: + fix = ( + f"{_PACKAGE}.{name} is not the object defined in {origin}; the facade must " + "re-export it, not copy or redefine it." + ) + return f"{fix} ({problem})" + + +def test_every_former_name_is_reachable_and_identical(): + problems = check_identity(_FACADE_ROWS) + assert problems == [], "\n".join(_with_facade_fix(problem) for problem in problems) + + +def test_facade_table_is_frozen(facade): + table = set(facade._NAME_TO_MODULE) + problems = [ + f"`{name}` is in the facade table but was not a name of the former module. New models " + f"are imported from their submodule (tensorrt_llm/serve/AGENTS.md); remove it from " + f"_NAME_TO_MODULE." + for name in sorted(table - _EXPECTED) + ] + problems += [ + f"`{name}` was a name of the former module but is not in _NAME_TO_MODULE. Add " + f"'{name}': '' to _NAME_TO_MODULE; underscore names belong in the table too." + for name in sorted(_EXPECTED - table) + ] + problems += [ + f"`{name}` is both a submodule of {_PACKAGE} and a row of _NAME_TO_MODULE; importing " + "the submodule would bind the package attribute to the module and shadow the " + "re-exported name." + for name in sorted(table & set(_ALLOWED_SIBLING_IMPORTS)) + ] + assert problems == [], "\n".join(problems) + + +def test_all_dir_and_type_checking_block(facade): + public = {name for name in facade._NAME_TO_MODULE if not name.startswith("_")} + problems = [] + if set(facade.__all__) != public: + problems.append( + f"__all__ differs from the public names of _NAME_TO_MODULE: " + f"{sorted(set(facade.__all__) ^ public)}" + ) + missing_from_dir = sorted(_EXPECTED - set(dir(facade))) + if missing_from_dir: + problems.append(f"dir({_PACKAGE}) lacks {missing_from_dir}") + + tree = ast.parse(Path(facade.__file__).read_text()) + type_checking_imports = set() + for node in tree.body: + if isinstance(node, ast.If) and ast.unparse(node.test) == "TYPE_CHECKING": + for statement in node.body: + if isinstance(statement, ast.ImportFrom): + module = _resolve(statement, _PACKAGE) + type_checking_imports |= {(module, alias.name) for alias in statement.names} + expected_imports = { + (facade._NAME_TO_MODULE[name], name) + for name in facade.__all__ + if name not in _VISUAL_GEN_NAMES + } + if type_checking_imports != expected_imports: + problems.append( + f"the TYPE_CHECKING imports and __all__ of {_FACADE_PATH} disagree with " + f"_NAME_TO_MODULE: {sorted(type_checking_imports ^ expected_imports)}. The block " + "imports every name of __all__ except the VisualGen names, which this file names " + "only as strings." + ) + for node in ast.walk(tree): + if isinstance(node, ast.Import): + modules = [alias.name for alias in node.names] + elif isinstance(node, ast.ImportFrom): + modules = [_resolve(node, _PACKAGE)] + else: + continue + for module in modules: + if module not in ("importlib", "typing") and not module.startswith(_PACKAGE + "."): + problems.append( + f"{_FACADE_PATH} imports {module} at line {node.lineno}; the facade imports " + "only importlib, typing and its own submodules (the latter under " + "TYPE_CHECKING)." + ) + assert problems == [], "\n".join(problems) + + +def test_unknown_name_raises_attribute_error(facade): + assert getattr(facade, "NoSuchModel", None) is None + with pytest.raises(ImportError): + exec(f"from {_PACKAGE} import NoSuchModel", {}) + + +def test_submodule_import_graph(package_dir): + present = {path.stem for path in package_dir.glob("*.py")} - {"__init__"} + problems = [ + f"{_PACKAGE_DIR}/{name}.py is not in the allowed import graph of this test; a new API's " + "submodule is added to _ALLOWED_SIBLING_IMPORTS here (importing common.py only, unless " + "it needs more) and the graph must stay acyclic." + for name in sorted(present - set(_ALLOWED_SIBLING_IMPORTS)) + ] + problems += [ + f"{_PACKAGE_DIR}/{name}.py is missing; the allowed import graph of this test lists it." + for name in sorted(set(_ALLOWED_SIBLING_IMPORTS) - present) + ] + for name in sorted(present & set(_ALLOWED_SIBLING_IMPORTS)): + tree = ast.parse((package_dir / f"{name}.py").read_text()) + for node in ast.walk(tree): + if isinstance(node, ast.Import): + modules = [alias.name for alias in node.names] + elif isinstance(node, ast.ImportFrom): + modules = [_resolve(node, _PACKAGE)] + else: + continue + for module in modules: + if module == _PACKAGE: + problems.append( + f"{_PACKAGE_DIR}/{name}.py imports the package {_PACKAGE} itself at " + f"line {node.lineno}; import the defining submodule, so that the facade " + "never sits in an import chain between submodules." + ) + elif module.startswith(_PACKAGE + "."): + other = module[len(_PACKAGE) + 1 :].split(".")[0] + if other not in _ALLOWED_SIBLING_IMPORTS[name]: + problems.append( + f"{_PACKAGE_DIR}/{name}.py imports {other}; allowed sibling imports " + f"are {sorted(_ALLOWED_SIBLING_IMPORTS[name])}. Move the shared " + "definition to common.py or sampling.py; the graph must stay acyclic." + ) + assert problems == [], "\n".join(problems) + + +def test_importing_the_package_loads_no_submodule(): + out = _run_in_fresh_interpreter( + f""" + import json, sys + import tensorrt_llm.serve.openai_protocol + prefix = {_PACKAGE!r} + "." + print(json.dumps(sorted(m for m in sys.modules if m.startswith(prefix)))) + """ + ) + assert out == [], ( + f"importing {_PACKAGE} loaded {out}; the facade imports a submodule only inside " + "__getattr__." + ) + + +def test_tools_does_not_load_chat_or_responses(): + out = _run_in_fresh_interpreter( + f""" + import json, sys + import tensorrt_llm.serve.openai_protocol.tools + prefix = {_PACKAGE!r} + "." + print(json.dumps(sorted(m[len(prefix):] for m in sys.modules if m.startswith(prefix)))) + """ + ) + heavy = sorted(set(out) & {"chat", "responses", "completions"}) + assert heavy == [], ( + f"importing {_PACKAGE}.tools loaded {heavy}; tools.py may import only common.py." + ) + + +def test_chat_models_do_not_load_extensions(package_dir): + out = _run_in_fresh_interpreter( + f""" + import json, sys + from tensorrt_llm.serve.openai_protocol import ChatCompletionRequest + prefix = {_PACKAGE!r} + "." + print(json.dumps({{ + "extensions": [m for m in sys.modules if m.startswith("tensorrt_llm.serve.extensions")], + "loaded": sorted(m[len(prefix):] for m in sys.modules if m.startswith(prefix)), + }})) + """ + ) + if out["extensions"]: + importers = sorted( + name + for name in out["loaded"] + if "tensorrt_llm.serve.extensions" in (package_dir / f"{name}.py").read_text() + ) + culprits = importers or out["loaded"] + assert out["extensions"] == [], ( + f"importing ChatCompletionRequest through {_PACKAGE} loaded {out['extensions']}; " + f"{[f'{_PACKAGE_DIR}/{name}.py' for name in culprits]} must not import " + "tensorrt_llm.serve.extensions: the built-in extensions import the protocol models, and " + "serving_extensions imports them only on the first lookup to keep that cycle open." + ) + + +def test_private_helpers_that_tests_import(): + problems = [] + for name, test_file in _TEST_IMPORTED_PRIVATE_NAMES.items(): + namespace: dict = {} + try: + exec(f"from {_PACKAGE} import {name}", namespace) + except ImportError as exc: + problems.append( + f"{test_file} imports {name} from {_PACKAGE}; keep it in _NAME_TO_MODULE (the " + f"table keeps underscore names too). Import failed with: {exc}" + ) + continue + expected = getattr(importlib.import_module(_FORMER_NAMES[name]), name) + if namespace[name] is not expected: + problems.append( + f"{test_file} imports {name} from {_PACKAGE}, which is not the object defined " + f"in {_FORMER_NAMES[name]}." + ) + assert problems == [], "\n".join(problems) diff --git a/tests/unittest/serve/test_layering.py b/tests/unittest/serve/test_layering.py index 357dcd4b02fc..f70d9b335aa3 100644 --- a/tests/unittest/serve/test_layering.py +++ b/tests/unittest/serve/test_layering.py @@ -103,6 +103,35 @@ "what they need. Move the shared helper to serve/_base/ instead." ), ), + Contract( + name="openai_protocol imports no handler module", + importers=("tensorrt_llm.serve.openai_protocol",), + forbidden=( + # Modules that define routes, build an app or run request handling today ... + "tensorrt_llm.serve.openai_server", + "tensorrt_llm.serve.openai_disagg_server", + "tensorrt_llm.serve.openai_disagg_service", + "tensorrt_llm.serve.openai_service", + "tensorrt_llm.serve.coordinator_server", + "tensorrt_llm.serve.openai_video_routes", + "tensorrt_llm.serve.responses_utils", + "tensorrt_llm.serve.postprocess_handlers", + "tensorrt_llm.serve.harmony_adapter", + "tensorrt_llm.serve.anthropic_adapter", + "tensorrt_llm.serve.anthropic_batches", + # ... and the serve modules for such code, whether or not they exist. + "tensorrt_llm.serve.api", + "tensorrt_llm.serve.app", + "tensorrt_llm.serve.disagg", + ), + include_type_checking=True, + fix=( + "serve/openai_protocol/ holds request and response models and the helpers that turn " + "them into engine parameters; servers, handlers and routers import it, never the " + "reverse. Move the code that needs the handler module out of the " + "protocol package, or have the caller pass the value in." + ), + ), ) From 8d62a8e256b806b5cb7a90699346e734ac24a110 Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Tue, 3 Nov 2026 09:20:08 +0000 Subject: [PATCH 23/30] [TRTLLM-17265][refactor] State why the facade keeps the VisualGen rows as strings and bind the test's culprit list unconditionally The table comment in tensorrt_llm/serve/openai_protocol/__init__.py said that a TYPE_CHECKING import of the VisualGen models would load their package at run time, which a TYPE_CHECKING block never does. The real reason is the static import-isolation check that lets CI run only the VisualGen tests for VisualGen-only changes: it forbids an import edge from this package to the VisualGen package even under TYPE_CHECKING, so the facade names that module only as a string. In tests/unittest/serve/openai_protocol/test_openai_protocol_facade.py the list of submodules that import the serving extensions was bound only on the failing path and read in the assertion message; it is now bound before the branch, so an unconditional message can never raise UnboundLocalError. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- tensorrt_llm/serve/openai_protocol/__init__.py | 6 ++++-- .../serve/openai_protocol/test_openai_protocol_facade.py | 1 + 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/serve/openai_protocol/__init__.py b/tensorrt_llm/serve/openai_protocol/__init__.py index a140aa29b70e..2411f8097c73 100644 --- a/tensorrt_llm/serve/openai_protocol/__init__.py +++ b/tensorrt_llm/serve/openai_protocol/__init__.py @@ -215,8 +215,10 @@ "to_disaggregated_params": "tensorrt_llm.serve.openai_protocol.disagg", "to_llm_conversation_params": "tensorrt_llm.serve.openai_protocol.disagg", "to_llm_disaggregated_params": "tensorrt_llm.serve.openai_protocol.disagg", - # Defined in the VisualGen serving package. Named here only as strings: an import, - # even under TYPE_CHECKING, would load that package into every LLM process. + # Defined in the VisualGen serving package. Named here only as strings: the static + # import-isolation check that lets CI run only the VisualGen tests for VisualGen-only + # changes forbids an import edge from this package to the VisualGen package, even under + # TYPE_CHECKING, so these names are not in the TYPE_CHECKING block above. "ImageEditRequest": "tensorrt_llm.serve.api.visual_gen.protocol", "ImageGenerationRequest": "tensorrt_llm.serve.api.visual_gen.protocol", "ImageGenerationResponse": "tensorrt_llm.serve.api.visual_gen.protocol", diff --git a/tests/unittest/serve/openai_protocol/test_openai_protocol_facade.py b/tests/unittest/serve/openai_protocol/test_openai_protocol_facade.py index f6b53a269a14..c56e2000ee2b 100644 --- a/tests/unittest/serve/openai_protocol/test_openai_protocol_facade.py +++ b/tests/unittest/serve/openai_protocol/test_openai_protocol_facade.py @@ -429,6 +429,7 @@ def test_chat_models_do_not_load_extensions(package_dir): }})) """ ) + culprits: list = [] if out["extensions"]: importers = sorted( name From 9c46c6fad9a6d70fb47b13a6fc5fe6c8e122e6a0 Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Tue, 3 Nov 2026 13:45:54 +0000 Subject: [PATCH 24/30] [TRTLLM-17267][refactor] Add the serving-class skeletons for embeddings, multimodal encoder and tokenize Create tensorrt_llm/serve/_base/serving.py with BaseServing and GenerateBaseServing, and the EmbeddingServing, MMEncoderServing and TokenizeServing skeletons under tensorrt_llm/serve/api/openai/. Every dependency is a keyword-only constructor argument. The new Group B files are registered in legacy-files.txt and the generated lint configs are regenerated. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- .pre-commit-config.yaml | 8 ++++++ legacy-files.txt | 4 +++ pyproject.toml | 4 +++ ruff-legacy.toml | 4 +++ tensorrt_llm/serve/_base/serving.py | 27 +++++++++++++++++++ tensorrt_llm/serve/api/openai/__init__.py | 2 ++ .../serve/api/openai/embeddings/__init__.py | 2 ++ .../serve/api/openai/embeddings/serving.py | 27 +++++++++++++++++++ .../serve/api/openai/mm_encoder/__init__.py | 2 ++ .../serve/api/openai/mm_encoder/serving.py | 25 +++++++++++++++++ .../serve/api/openai/tokenize/__init__.py | 2 ++ .../serve/api/openai/tokenize/serving.py | 15 +++++++++++ 12 files changed, 122 insertions(+) create mode 100644 tensorrt_llm/serve/_base/serving.py create mode 100644 tensorrt_llm/serve/api/openai/__init__.py create mode 100644 tensorrt_llm/serve/api/openai/embeddings/__init__.py create mode 100644 tensorrt_llm/serve/api/openai/embeddings/serving.py create mode 100644 tensorrt_llm/serve/api/openai/mm_encoder/__init__.py create mode 100644 tensorrt_llm/serve/api/openai/mm_encoder/serving.py create mode 100644 tensorrt_llm/serve/api/openai/tokenize/__init__.py create mode 100644 tensorrt_llm/serve/api/openai/tokenize/serving.py diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index bf7b7b678c60..ffcb5862ea11 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -418,6 +418,10 @@ 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/serving.py | + tensorrt_llm/serve/api/openai/embeddings/serving.py | + tensorrt_llm/serve/api/openai/mm_encoder/serving.py | + tensorrt_llm/serve/api/openai/tokenize/serving.py | tensorrt_llm/serve/api/visual_gen/protocol.py | tensorrt_llm/serve/chat_utils.py | tensorrt_llm/serve/cluster_storage.py | @@ -1158,6 +1162,10 @@ 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/serving.py | + tensorrt_llm/serve/api/openai/embeddings/serving.py | + tensorrt_llm/serve/api/openai/mm_encoder/serving.py | + tensorrt_llm/serve/api/openai/tokenize/serving.py | tensorrt_llm/serve/api/visual_gen/protocol.py | tensorrt_llm/serve/chat_utils.py | tensorrt_llm/serve/cluster_storage.py | diff --git a/legacy-files.txt b/legacy-files.txt index 3c3a8846802b..6c098106a80d 100644 --- a/legacy-files.txt +++ b/legacy-files.txt @@ -410,6 +410,10 @@ 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/serving.py +tensorrt_llm/serve/api/openai/embeddings/serving.py +tensorrt_llm/serve/api/openai/mm_encoder/serving.py +tensorrt_llm/serve/api/openai/tokenize/serving.py tensorrt_llm/serve/api/visual_gen/protocol.py tensorrt_llm/serve/chat_utils.py tensorrt_llm/serve/cluster_storage.py diff --git a/pyproject.toml b/pyproject.toml index 440136a97f0d..81d2e449c2b8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -475,6 +475,10 @@ exclude = [ "tensorrt_llm/scheduling_params.py", "tensorrt_llm/serve/__init__.py", "tensorrt_llm/serve/_base/msgspec_route.py", + "tensorrt_llm/serve/_base/serving.py", + "tensorrt_llm/serve/api/openai/embeddings/serving.py", + "tensorrt_llm/serve/api/openai/mm_encoder/serving.py", + "tensorrt_llm/serve/api/openai/tokenize/serving.py", "tensorrt_llm/serve/api/visual_gen/protocol.py", "tensorrt_llm/serve/chat_utils.py", "tensorrt_llm/serve/cluster_storage.py", diff --git a/ruff-legacy.toml b/ruff-legacy.toml index cca0b8dd4f9f..5752751aa372 100644 --- a/ruff-legacy.toml +++ b/ruff-legacy.toml @@ -427,6 +427,10 @@ include = [ "tensorrt_llm/scheduling_params.py", "tensorrt_llm/serve/__init__.py", "tensorrt_llm/serve/_base/msgspec_route.py", + "tensorrt_llm/serve/_base/serving.py", + "tensorrt_llm/serve/api/openai/embeddings/serving.py", + "tensorrt_llm/serve/api/openai/mm_encoder/serving.py", + "tensorrt_llm/serve/api/openai/tokenize/serving.py", "tensorrt_llm/serve/api/visual_gen/protocol.py", "tensorrt_llm/serve/chat_utils.py", "tensorrt_llm/serve/cluster_storage.py", diff --git a/tensorrt_llm/serve/_base/serving.py b/tensorrt_llm/serve/_base/serving.py new file mode 100644 index 000000000000..af505182c322 --- /dev/null +++ b/tensorrt_llm/serve/_base/serving.py @@ -0,0 +1,27 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Base classes of the serving layer: one serving class per endpoint.""" + +from typing import Any + + +class BaseServing: + """Base of every serving class. + + Holds only what all serving classes share; each subclass receives its + dependencies as explicit constructor arguments, never by reading the + server object. + """ + + +class GenerateBaseServing(BaseServing): + """Serving classes that submit requests to the generate pipeline.""" + + def __init__(self, *, model: str, generator: Any, tokenizer: Any, + model_config: Any) -> None: + # BaseServing has no constructor: set every attribute here, no + # super().__init__(...). + self.model = model + self.generator = generator + self.tokenizer = tokenizer + self.model_config = model_config diff --git a/tensorrt_llm/serve/api/openai/__init__.py b/tensorrt_llm/serve/api/openai/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/serve/api/openai/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/serve/api/openai/embeddings/__init__.py b/tensorrt_llm/serve/api/openai/embeddings/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/serve/api/openai/embeddings/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/serve/api/openai/embeddings/serving.py b/tensorrt_llm/serve/api/openai/embeddings/serving.py new file mode 100644 index 000000000000..2458829369ee --- /dev/null +++ b/tensorrt_llm/serve/api/openai/embeddings/serving.py @@ -0,0 +1,27 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""``POST /v1/embeddings`` on an encode-only ``LLM``.""" + +from typing import Any, Optional + +from tensorrt_llm.serve._base.serving import BaseServing +from tensorrt_llm.serve.api.openai.embeddings.batcher import EncodeBatcher + + +class EmbeddingServing(BaseServing): + """``POST /v1/embeddings`` on an encode-only ``LLM``. + + Concurrent requests are coalesced by ``EncodeBatcher``; the composition + root keeps a reference to the batcher, because the lifespan starts and + stops it and the health probe reports it. + """ + + def __init__(self, *, model: str, generator: Any, + embedding_max_queue_delay: float, + embedding_max_queue_size: int) -> None: + self.model = model + self.generator = generator + self._embedding_max_queue_delay = embedding_max_queue_delay + self._embedding_max_queue_size = embedding_max_queue_size + self.embedding_batcher: Optional[EncodeBatcher] = None + self._init_embedding_batcher() diff --git a/tensorrt_llm/serve/api/openai/mm_encoder/__init__.py b/tensorrt_llm/serve/api/openai/mm_encoder/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/serve/api/openai/mm_encoder/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/serve/api/openai/mm_encoder/serving.py b/tensorrt_llm/serve/api/openai/mm_encoder/serving.py new file mode 100644 index 000000000000..1623429878a8 --- /dev/null +++ b/tensorrt_llm/serve/api/openai/mm_encoder/serving.py @@ -0,0 +1,25 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""``POST /v1/chat/completions`` on the multimodal-encoder role.""" + +from typing import Any + +from tensorrt_llm.serve._base.serving import GenerateBaseServing +from tensorrt_llm.serve.render import RenderResources + + +class MMEncoderServing(GenerateBaseServing): + """``POST /v1/chat/completions`` on the multimodal-encoder role. + + Returns the multimodal embedding handles of the request instead of + generated text. Rendering uses the ``RenderResources`` the composition + root built once, never resources read back from this object. + """ + + def __init__(self, *, render_resources: RenderResources, + multimodal_server_config: Any, + allow_request_chat_template: bool, **base_kwargs: Any) -> None: + super().__init__(**base_kwargs) + self._render_resources = render_resources + self.multimodal_server_config = multimodal_server_config + self.allow_request_chat_template = allow_request_chat_template diff --git a/tensorrt_llm/serve/api/openai/tokenize/__init__.py b/tensorrt_llm/serve/api/openai/tokenize/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tensorrt_llm/serve/api/openai/tokenize/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tensorrt_llm/serve/api/openai/tokenize/serving.py b/tensorrt_llm/serve/api/openai/tokenize/serving.py new file mode 100644 index 000000000000..d3a04032e260 --- /dev/null +++ b/tensorrt_llm/serve/api/openai/tokenize/serving.py @@ -0,0 +1,15 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""``POST /_internal/tokenize``: token ids of a plain-text prompt.""" + +from typing import Any + +from tensorrt_llm.serve._base.serving import BaseServing + + +class TokenizeServing(BaseServing): + """``POST /_internal/tokenize``: token ids of a plain-text prompt.""" + + def __init__(self, *, model: str, tokenizer: Any) -> None: + self.model = model + self.tokenizer = tokenizer From 898dfd24d0a498a6bb0062ca2da5d00a8fff2e42 Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Tue, 3 Nov 2026 13:49:44 +0000 Subject: [PATCH 25/30] [TRTLLM-17267][refactor] Move the embeddings, multimodal-encoder and tokenize handlers into their serving classes Generated move. OpenAIServer loses openai_embedding, _normalize_embedding_input, _init_embedding_batcher (to EmbeddingServing), openai_mm_encoder (to MMEncoderServing) and tokenize (to TokenizeServing); register_embedding_routes is deleted, its route is served by the embeddings router. create_error_response, _attach_multimodal_data and await_disconnected move to the serving bases in tensorrt_llm/serve/_base/serving.py and stay reachable on OpenAIServer as aliases of the moved objects while chat and completions still call them. _record_generator_termination moves to the same module. tensorrt_llm/serve/encode_batcher.py becomes tensorrt_llm/serve/api/openai/embeddings/batcher.py; the old module is deleted and every import of it is rewritten. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- tensorrt_llm/serve/_base/serving.py | 88 ++++- .../openai/embeddings/batcher.py} | 0 .../serve/api/openai/embeddings/serving.py | 141 ++++++- .../serve/api/openai/mm_encoder/serving.py | 140 ++++++- .../serve/api/openai/tokenize/serving.py | 15 + tensorrt_llm/serve/openai_server.py | 354 +----------------- .../llmapi/apps/test_encode_batcher.py | 8 +- .../serve/test_serve_runtime_api_calls.py | 3 +- 8 files changed, 398 insertions(+), 351 deletions(-) rename tensorrt_llm/serve/{encode_batcher.py => api/openai/embeddings/batcher.py} (100%) diff --git a/tensorrt_llm/serve/_base/serving.py b/tensorrt_llm/serve/_base/serving.py index af505182c322..6feade1c0ed8 100644 --- a/tensorrt_llm/serve/_base/serving.py +++ b/tensorrt_llm/serve/_base/serving.py @@ -2,7 +2,18 @@ # SPDX-License-Identifier: Apache-2.0 """Base classes of the serving layer: one serving class per endpoint.""" -from typing import Any +import asyncio +from http import HTTPStatus +from typing import Any, Dict, List, Optional, Union + +from fastapi import Request +from fastapi.responses import JSONResponse, Response + +from tensorrt_llm.executor.utils import (CONTEXT_LENGTH_EXCEEDED_CODE, + is_context_length_exceeded_message) +from tensorrt_llm.logger import logger +from tensorrt_llm.serve.openai_protocol import ErrorResponse +from tensorrt_llm.usage import TerminalOutcome, record_termination_observation class BaseServing: @@ -13,6 +24,23 @@ class BaseServing: server object. """ + @staticmethod + 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) + return JSONResponse(content=error_response.model_dump(), + status_code=status_code.value) + class GenerateBaseServing(BaseServing): """Serving classes that submit requests to the generate pipeline.""" @@ -25,3 +53,61 @@ def __init__(self, *, model: str, generator: Any, tokenizer: Any, self.generator = generator self.tokenizer = tokenizer self.model_config = model_config + + @staticmethod + def _attach_multimodal_data( + prompt: Dict[str, Any], + request: Any, + mm_data: Any, + mm_item_order: Optional[List[Any]] = None) -> None: + """Attach decoded media to the prompt dict handed to the engine. + + When the request carried no ``prompt_token_ids``, ``prompt`` is the + chat template's output with one placeholder per media item already in + it, so ``chat_template_applied`` tells the input processor to use it + verbatim instead of rendering the template a second time. + """ + prompt["multi_modal_data"] = mm_data + if request.prompt_token_ids is None: + prompt["chat_template_applied"] = True + if mm_item_order: + prompt["mm_item_order"] = mm_item_order + + async def await_disconnected(self, raw_request: Request, promise): + if raw_request is None: + return + # A batched request is synthesised without a client socket, so + # is_disconnected() can never become true and this poll would run at + # 1Hz for the life of the process, holding the Request and its + # RequestOutput -- one leaked task per batched request. + if getattr(raw_request.state, "no_client_connection", False): + return + while not await raw_request.is_disconnected(): + await asyncio.sleep(1) + if not promise.finished: + promise.abort() + logger.info( + f"{raw_request.client} is disconnected, abort {promise.request_id}" + ) + + +def _record_generator_termination(generator) -> None: + """Classify a fatal generator error without exposing exception details.""" + component = "llm" + reporting_source = "self" + termination_kind = "exception" + try: + if generator.engine_in_worker_process: + component = "engine_worker" + reporting_source = "executor_proxy" + termination_kind = "worker_failure" + except Exception: + pass + record_termination_observation( + TerminalOutcome( + termination_kind=termination_kind, + component=component, + reporting_source=reporting_source, + exit_code_known=False + if termination_kind == "worker_failure" else None, + )) diff --git a/tensorrt_llm/serve/encode_batcher.py b/tensorrt_llm/serve/api/openai/embeddings/batcher.py similarity index 100% rename from tensorrt_llm/serve/encode_batcher.py rename to tensorrt_llm/serve/api/openai/embeddings/batcher.py diff --git a/tensorrt_llm/serve/api/openai/embeddings/serving.py b/tensorrt_llm/serve/api/openai/embeddings/serving.py index 2458829369ee..1c281a8ee935 100644 --- a/tensorrt_llm/serve/api/openai/embeddings/serving.py +++ b/tensorrt_llm/serve/api/openai/embeddings/serving.py @@ -2,10 +2,27 @@ # SPDX-License-Identifier: Apache-2.0 """``POST /v1/embeddings`` on an encode-only ``LLM``.""" +import array +import asyncio +import base64 +import traceback +from http import HTTPStatus from typing import Any, Optional +from fastapi import Request +from fastapi.responses import JSONResponse, Response + +from tensorrt_llm.inputs import prompt_inputs +from tensorrt_llm.logger import logger +from tensorrt_llm.sampling_params import SamplingParams from tensorrt_llm.serve._base.serving import BaseServing -from tensorrt_llm.serve.api.openai.embeddings.batcher import EncodeBatcher +from tensorrt_llm.serve.api.openai.embeddings.batcher import (EncodeBatcher, + InputTooLongError, + QueueFullError) +from tensorrt_llm.serve.openai_protocol import (EmbeddingRequest, + EmbeddingResponse, + EmbeddingResponseData, + EmbeddingUsageInfo) class EmbeddingServing(BaseServing): @@ -25,3 +42,125 @@ def __init__(self, *, model: str, generator: Any, self._embedding_max_queue_size = embedding_max_queue_size self.embedding_batcher: Optional[EncodeBatcher] = None self._init_embedding_batcher() + + def _init_embedding_batcher(self): + """Create the encode dynamic batcher for the embedding server. + + The batcher coalesces concurrent /v1/embeddings requests into a single + `llm.encode()` call. `encode()` is synchronous; the batcher runs it in + the default executor so the event loop stays responsive. + """ + + def encode_fn(token_ids_batch): + # token_ids_batch: list of pre-tokenized token-id lists. encode() + # returns one EncoderOutput per item, in input order. + return self.generator.encode(token_ids_batch) + + # The limits encode() checks every call against; see LLM.encode_limits. + limits = self.generator.encode_limits + self.embedding_batcher = EncodeBatcher( + encode_fn, + max_batch_size=limits.max_batch_size, + max_queue_delay=self._embedding_max_queue_delay, + max_queue_size=self._embedding_max_queue_size, + max_num_tokens=limits.max_num_tokens, + max_seq_len=limits.max_seq_len, + ) + + @staticmethod + def _normalize_embedding_input(inp): + """Normalize EmbeddingRequest.input into a list of items. + + Each item is either a string (to tokenize) or a pre-tokenized token-id + list. Mirrors the four OpenAI input forms: str | list[str] | list[int] | + list[list[int]]. + """ + if isinstance(inp, str): + return [inp] + if isinstance(inp, list): + if len(inp) == 0: + return [] + if isinstance(inp[0], int): + return [inp] # a single pre-tokenized input + return list(inp) # list of strings or list of token-id lists + raise TypeError(f"Unsupported embedding input type: {type(inp)}") + + async def openai_embedding(self, request: EmbeddingRequest, + raw_request: Request) -> Response: + try: + # `dimensions` is a Matryoshka-embedding knob (truncate-then-renormalize + # a sentence-embedding vector), matching vLLM, which rejects it for + # non-Matryoshka models. None of the models served here are + # Matryoshka-trained: BERT classifiers / reward models emit raw + # [num_labels] / [seq_len, num_labels] tensors, and Qwen3-Embedding emits + # a fixed-width pooled vector with no nested sub-dimensions. Truncating + # any of these is meaningless, so reject. (Honoring `dimensions` for a + # genuinely Matryoshka model would be a fast-follow.) + if request.dimensions is not None: + return self.create_error_response( + "`dimensions` is only supported by Matryoshka-trained " + "text-embedding models; none of the models served by this " + "endpoint are Matryoshka-trained.", + status_code=HTTPStatus.BAD_REQUEST) + + items = self._normalize_embedding_input(request.input) + if not items: + return self.create_error_response("`input` must not be empty.") + + # Tokenize strings via the same input processor used by /v1/completions, + # honoring add_special_tokens. Pre-tokenized inputs are used as-is. + sampling_params = SamplingParams( + add_special_tokens=request.add_special_tokens) + token_ids_list = [] + for item in items: + if isinstance(item, str): + token_ids, _ = await asyncio.to_thread( + self.generator.input_processor, prompt_inputs(item), + sampling_params) + else: + token_ids = item + token_ids_list.append(token_ids) + + try: + # Validate every input's length up front so an oversize item fails + # the whole request before any item is enqueued — otherwise gather() + # cancels the awaiters but their already-queued inputs still run + # encode() (wasted GPU work). + for token_ids in token_ids_list: + self.embedding_batcher.validate_input(token_ids) + results = await asyncio.gather(*[ + self.embedding_batcher.submit(token_ids) + for token_ids in token_ids_list + ]) + except InputTooLongError as e: + return self.create_error_response( + str(e), status_code=HTTPStatus.BAD_REQUEST) + except QueueFullError as e: + return self.create_error_response( + str(e), status_code=HTTPStatus.TOO_MANY_REQUESTS) + + data = [] + for idx, encoder_output in enumerate(results): + # Model-agnostic: serialize whatever per-request tensor encode() + # emits (classification logits, reward score, pooled vector, ...). + embedding = encoder_output.logits.flatten().tolist() + if request.encoding_format == "base64": + embedding = base64.b64encode( + array.array("f", embedding).tobytes()).decode("utf-8") + data.append( + EmbeddingResponseData(index=idx, embedding=embedding)) + + num_prompt_tokens = sum(len(t) for t in token_ids_list) + usage = EmbeddingUsageInfo(prompt_tokens=num_prompt_tokens, + total_tokens=num_prompt_tokens) + response = EmbeddingResponse(model=request.model, + data=data, + usage=usage) + return JSONResponse(content=response.model_dump()) + except Exception as e: + # Unexpected fault (not a client error): surface as 500, not 400. + logger.error(traceback.format_exc()) + return self.create_error_response( + str(e), + err_type="InternalServerError", + status_code=HTTPStatus.INTERNAL_SERVER_ERROR) diff --git a/tensorrt_llm/serve/api/openai/mm_encoder/serving.py b/tensorrt_llm/serve/api/openai/mm_encoder/serving.py index 1623429878a8..540cdd0adc74 100644 --- a/tensorrt_llm/serve/api/openai/mm_encoder/serving.py +++ b/tensorrt_llm/serve/api/openai/mm_encoder/serving.py @@ -2,10 +2,28 @@ # SPDX-License-Identifier: Apache-2.0 """``POST /v1/chat/completions`` on the multimodal-encoder role.""" -from typing import Any +import asyncio +import signal +import traceback +from http import HTTPStatus +from typing import Any, List -from tensorrt_llm.serve._base.serving import GenerateBaseServing -from tensorrt_llm.serve.render import RenderResources +from fastapi import Request +from fastapi.responses import JSONResponse, Response +from pydantic import ValidationError + +from tensorrt_llm.executor import CppExecutorError +from tensorrt_llm.inputs import prompt_inputs +from tensorrt_llm.inputs.utils import ConversationMessage +from tensorrt_llm.llmapi.llm import RequestOutput +from tensorrt_llm.logger import logger +from tensorrt_llm.serve._base.serving import (GenerateBaseServing, + _record_generator_termination) +from tensorrt_llm.serve.chat_utils import parse_chat_messages_coroutines +from tensorrt_llm.serve.openai_protocol import ( + ChatCompletionRequest, ChatCompletionResponse, ChatCompletionResponseChoice, + ChatMessage, UsageInfo, ensure_request_chat_template_allowed) +from tensorrt_llm.serve.render import RenderResources, render_conversation class MMEncoderServing(GenerateBaseServing): @@ -23,3 +41,119 @@ def __init__(self, *, render_resources: RenderResources, self._render_resources = render_resources self.multimodal_server_config = multimodal_server_config self.allow_request_chat_template = allow_request_chat_template + + async def openai_mm_encoder(self, request: ChatCompletionRequest, + raw_request: Request) -> Response: + + async def create_mm_embedding_response(promise: RequestOutput): + await promise.aresult() + # TODO: Replace mm_embedding_handles with a dedicated OpenAIBaseModel(JSON-safe), when enable multimodal disagg E2E + mm_embedding_handles = ( + promise.disaggregated_params.multimodal_embedding_handles + if promise.disaggregated_params else None) + if not mm_embedding_handles: + return self.create_error_response( + message="Multimodal embedding handle missing in response", + err_type="InternalServerError", + status_code=HTTPStatus.INTERNAL_SERVER_ERROR) + if any("tensor_size" not in h for h in mm_embedding_handles): + return self.create_error_response( + message="Multimodal embedding handle missing tensor_size", + err_type="InternalServerError", + status_code=HTTPStatus.INTERNAL_SERVER_ERROR) + mm_embedding_handle = (mm_embedding_handles[0] + if len(mm_embedding_handles) == 1 else + mm_embedding_handles) + num_tokens = sum( + int(h["tensor_size"][0]) for h in mm_embedding_handles) + return ChatCompletionResponse( + id=str(promise.request_id), + model=self.model, + choices=[ + ChatCompletionResponseChoice( + index=0, + message=ChatMessage(role="assistant", content="dummy"), + mm_embedding_handle=mm_embedding_handle, + finish_reason="length", + ) + ], + usage=UsageInfo( + prompt_tokens=num_tokens, + completion_tokens=1, + total_tokens=num_tokens + 1, + ), + ) + + try: + ensure_request_chat_template_allowed( + request, self.allow_request_chat_template) + conversation: List[ConversationMessage] = [] + # Same pipeline as the chat route: the server template and the + # model extension's tool serialization apply here too. + render_resources = self._render_resources + tool_dicts = None if request.tools is None else [ + render_resources.extension.serialize_tool(tool) + for tool in request.tools + ] + + try: + conversation, mm_coroutines, mm_placeholder_counts, mm_item_order = parse_chat_messages_coroutines( + request.messages, + self.model_config, + self.multimodal_server_config, + request_media_io_kwargs=request.media_io_kwargs, + ) + except ValidationError: + # ValidatorIterator rejects extra fields; fall back to raw JSON. + raw_body = await raw_request.json() + raw_messages = raw_body.get("messages", []) + conversation, mm_coroutines, mm_placeholder_counts, mm_item_order = parse_chat_messages_coroutines( + raw_messages, + self.model_config, + self.multimodal_server_config, + request_media_io_kwargs=request.media_io_kwargs, + ) + + if request.prompt_token_ids is not None: + prompt = request.prompt_token_ids + else: + render_task = render_conversation( + render_resources, + conversation=conversation, + mm_placeholder_counts=mm_placeholder_counts, + add_generation_prompt=request.add_generation_prompt, + tools=tool_dicts, + documents=request.documents, + chat_template=request.chat_template, + chat_template_kwargs=request.chat_template_kwargs, + injected_chat_template_kwargs=request. + injected_chat_template_kwargs, + ) + rendered, (mm_data, mm_embeddings) = await asyncio.gather( + render_task, mm_coroutines) + prompt = (rendered.text + if rendered.text is not None else rendered.token_ids) + prompt = prompt_inputs(prompt) + + if request.prompt_token_ids is not None: + mm_data, mm_embeddings = await mm_coroutines + if mm_embeddings: + raise ValueError("Cannot use multimodal embeddings as input") + if mm_data is not None: + self._attach_multimodal_data(prompt, request, mm_data, + mm_item_order) + + promise = self.generator.generate_async(inputs=prompt, ) + asyncio.create_task(self.await_disconnected(raw_request, promise)) + + response = await create_mm_embedding_response(promise) + return JSONResponse(content=response.model_dump()) + + except CppExecutorError: + logger.error(traceback.format_exc()) + # If internal executor error is raised, shutdown the server + _record_generator_termination(self.generator) + signal.raise_signal(signal.SIGINT) + except Exception as e: + logger.error(traceback.format_exc()) + return self.create_error_response(str(e)) diff --git a/tensorrt_llm/serve/api/openai/tokenize/serving.py b/tensorrt_llm/serve/api/openai/tokenize/serving.py index d3a04032e260..05ba379b1835 100644 --- a/tensorrt_llm/serve/api/openai/tokenize/serving.py +++ b/tensorrt_llm/serve/api/openai/tokenize/serving.py @@ -2,9 +2,13 @@ # SPDX-License-Identifier: Apache-2.0 """``POST /_internal/tokenize``: token ids of a plain-text prompt.""" +from http import HTTPStatus from typing import Any +from fastapi.responses import JSONResponse + from tensorrt_llm.serve._base.serving import BaseServing +from tensorrt_llm.serve.openai_protocol import TokenizeRequest, TokenizeResponse class TokenizeServing(BaseServing): @@ -13,3 +17,14 @@ class TokenizeServing(BaseServing): def __init__(self, *, model: str, tokenizer: Any) -> None: self.model = model self.tokenizer = tokenizer + + async def tokenize(self, request: TokenizeRequest) -> JSONResponse: + try: + token_ids = self.tokenizer.encode(request.prompt) + response = TokenizeResponse(count=len(token_ids), tokens=token_ids) + return JSONResponse(content=response.model_dump()) + except Exception as e: + return self.create_error_response( + message=str(e), + err_type="InvalidRequestError", + status_code=HTTPStatus.BAD_REQUEST) diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py index 994c2627eb5c..25cccf3b1e23 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -2,7 +2,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -import array import asyncio import base64 import functools @@ -42,8 +41,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, @@ -68,6 +65,8 @@ 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.serving import (BaseServing, GenerateBaseServing, + _record_generator_termination) from tensorrt_llm.serve._perf_metrics_config import resolve_perf_metrics_config from tensorrt_llm.serve.anthropic_adapter import ( AnthropicRequestError, AnthropicResponseError, anthropic_error_response, @@ -82,6 +81,7 @@ AnthropicCountTokensResponse, AnthropicCreateBatchRequest, AnthropicMessagesRequest) +from tensorrt_llm.serve.api.openai.embeddings.batcher import EncodeBatcher from tensorrt_llm.serve.api.ops.system.serving import SystemServing from tensorrt_llm.serve.api.visual_gen.protocol import (ImageEditRequest, ImageGenerationRequest, @@ -96,8 +96,6 @@ request_requires_internal_disagg_auth, validate_internal_disagg_request, validate_subagent_affinity) from tensorrt_llm.serve.disagg_auto_scaling import DisaggClusterWorker -from tensorrt_llm.serve.encode_batcher import (EncodeBatcher, InputTooLongError, - QueueFullError) from tensorrt_llm.serve.metadata_server import create_metadata_server from tensorrt_llm.serve.multi_frontend import (FORWARDED_ROUTES, AttachedFrontendWatchdog, @@ -107,14 +105,11 @@ follow_clock_offset) from tensorrt_llm.serve.openai_protocol import ( ChatCompletionNamedToolChoiceParam, ChatCompletionRequest, - ChatCompletionResponse, ChatCompletionResponseChoice, - ChatCompletionToolsParam, ChatMessage, CompletionRequest, - CompletionResponse, CompletionResponseChoice, EmbeddingRequest, - EmbeddingResponse, EmbeddingResponseData, EmbeddingUsageInfo, ErrorResponse, - MemoryUpdateRequest, PromptTokensDetails, ResponseFormat, ResponsesRequest, - ResponsesResponse, RuntimeMemoryUpdateRequest, StartProfileRequest, - TokenizeRequest, TokenizeResponse, UpdateWeightsRequest, UsageInfo, - ensure_request_chat_template_allowed, to_llm_conversation_params, + ChatCompletionResponse, ChatCompletionToolsParam, CompletionRequest, + CompletionResponse, CompletionResponseChoice, MemoryUpdateRequest, + PromptTokensDetails, ResponseFormat, ResponsesRequest, ResponsesResponse, + RuntimeMemoryUpdateRequest, StartProfileRequest, 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, @@ -153,7 +148,6 @@ build_visual_gen_server_timings, build_visual_gen_timing_headers) from tensorrt_llm.serve.visual_gen_utils import (local_media_path_is_disallowed, parse_visual_gen_params) -from tensorrt_llm.usage import TerminalOutcome, record_termination_observation from tensorrt_llm.version import __version__ as VERSION from .._utils import (AdjustedSteadyClock, @@ -267,28 +261,6 @@ def _configure_parser_special_token_decoding( sampling_params.spaces_between_special_tokens = False -def _record_generator_termination(generator) -> None: - """Classify a fatal generator error without exposing exception details.""" - component = "llm" - reporting_source = "self" - termination_kind = "exception" - try: - if generator.engine_in_worker_process: - component = "engine_worker" - reporting_source = "executor_proxy" - termination_kind = "worker_failure" - except Exception: - pass - record_termination_observation( - TerminalOutcome( - termination_kind=termination_kind, - component=component, - reporting_source=reporting_source, - exit_code_known=False - if termination_kind == "worker_failure" else None, - )) - - def _build_tool_strict_guided_decoding_params(tools, tool_parser_name): """Build GuidedDecodingParams with structural tags for tools with strict=True. @@ -1072,41 +1044,10 @@ def _has_cache_transceiver_config(self) -> bool: return (cache_transceiver_config is not None and getattr( cache_transceiver_config, "backend", None) is not None) - @staticmethod - def _attach_multimodal_data( - prompt: Dict[str, Any], - request: Any, - mm_data: Any, - mm_item_order: Optional[List[Any]] = None) -> None: - """Attach decoded media to the prompt dict handed to the engine. - - When the request carried no ``prompt_token_ids``, ``prompt`` is the - chat template's output with one placeholder per media item already in - it, so ``chat_template_applied`` tells the input processor to use it - verbatim instead of rendering the template a second time. - """ - prompt["multi_modal_data"] = mm_data - if request.prompt_token_ids is None: - prompt["chat_template_applied"] = True - if mm_item_order: - prompt["mm_item_order"] = mm_item_order - - async def await_disconnected(self, raw_request: Request, promise): - if raw_request is None: - return - # A batched request is synthesised without a client socket, so - # is_disconnected() can never become true and this poll would run at - # 1Hz for the life of the process, holding the Request and its - # RequestOutput -- one leaked task per batched request. - if getattr(raw_request.state, "no_client_connection", False): - return - while not await raw_request.is_disconnected(): - await asyncio.sleep(1) - if not promise.finished: - promise.abort() - logger.info( - f"{raw_request.client} is disconnected, abort {promise.request_id}" - ) + _attach_multimodal_data = staticmethod( + GenerateBaseServing._attach_multimodal_data) + + await_disconnected = GenerateBaseServing.await_disconnected @functools.cached_property def _is_visual_gen(self) -> bool: @@ -1122,22 +1063,7 @@ def postproc_worker_enabled(self) -> bool: return True if self.generator.args.num_postprocess_workers > 0 else False - @staticmethod - 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) - return JSONResponse(content=error_response.model_dump(), - status_code=status_code.value) + create_error_response = staticmethod(BaseServing.create_error_response) def _create_invalid_response_id_error(self, response_id: str) -> Response: return self.create_error_response( @@ -1425,133 +1351,6 @@ def _register_rl_control_routes(self) -> None: methods=["POST"], dependencies=dependencies) - def _init_embedding_batcher(self): - """Create the encode dynamic batcher for the embedding server. - - The batcher coalesces concurrent /v1/embeddings requests into a single - `llm.encode()` call. `encode()` is synchronous; the batcher runs it in - the default executor so the event loop stays responsive. - """ - - def encode_fn(token_ids_batch): - # token_ids_batch: list of pre-tokenized token-id lists. encode() - # returns one EncoderOutput per item, in input order. - return self.generator.encode(token_ids_batch) - - # The limits encode() checks every call against; see LLM.encode_limits. - limits = self.generator.encode_limits - self.embedding_batcher = EncodeBatcher( - encode_fn, - max_batch_size=limits.max_batch_size, - max_queue_delay=self._embedding_max_queue_delay, - max_queue_size=self._embedding_max_queue_size, - max_num_tokens=limits.max_num_tokens, - max_seq_len=limits.max_seq_len, - ) - - def register_embedding_routes(self): - self.app.add_api_route("/v1/embeddings", - self.openai_embedding, - methods=["POST"]) - - @staticmethod - def _normalize_embedding_input(inp): - """Normalize EmbeddingRequest.input into a list of items. - - Each item is either a string (to tokenize) or a pre-tokenized token-id - list. Mirrors the four OpenAI input forms: str | list[str] | list[int] | - list[list[int]]. - """ - if isinstance(inp, str): - return [inp] - if isinstance(inp, list): - if len(inp) == 0: - return [] - if isinstance(inp[0], int): - return [inp] # a single pre-tokenized input - return list(inp) # list of strings or list of token-id lists - raise TypeError(f"Unsupported embedding input type: {type(inp)}") - - async def openai_embedding(self, request: EmbeddingRequest, - raw_request: Request) -> Response: - try: - # `dimensions` is a Matryoshka-embedding knob (truncate-then-renormalize - # a sentence-embedding vector), matching vLLM, which rejects it for - # non-Matryoshka models. None of the models served here are - # Matryoshka-trained: BERT classifiers / reward models emit raw - # [num_labels] / [seq_len, num_labels] tensors, and Qwen3-Embedding emits - # a fixed-width pooled vector with no nested sub-dimensions. Truncating - # any of these is meaningless, so reject. (Honoring `dimensions` for a - # genuinely Matryoshka model would be a fast-follow.) - if request.dimensions is not None: - return self.create_error_response( - "`dimensions` is only supported by Matryoshka-trained " - "text-embedding models; none of the models served by this " - "endpoint are Matryoshka-trained.", - status_code=HTTPStatus.BAD_REQUEST) - - items = self._normalize_embedding_input(request.input) - if not items: - return self.create_error_response("`input` must not be empty.") - - # Tokenize strings via the same input processor used by /v1/completions, - # honoring add_special_tokens. Pre-tokenized inputs are used as-is. - sampling_params = SamplingParams( - add_special_tokens=request.add_special_tokens) - token_ids_list = [] - for item in items: - if isinstance(item, str): - token_ids, _ = await asyncio.to_thread( - self.generator.input_processor, prompt_inputs(item), - sampling_params) - else: - token_ids = item - token_ids_list.append(token_ids) - - try: - # Validate every input's length up front so an oversize item fails - # the whole request before any item is enqueued — otherwise gather() - # cancels the awaiters but their already-queued inputs still run - # encode() (wasted GPU work). - for token_ids in token_ids_list: - self.embedding_batcher.validate_input(token_ids) - results = await asyncio.gather(*[ - self.embedding_batcher.submit(token_ids) - for token_ids in token_ids_list - ]) - except InputTooLongError as e: - return self.create_error_response( - str(e), status_code=HTTPStatus.BAD_REQUEST) - except QueueFullError as e: - return self.create_error_response( - str(e), status_code=HTTPStatus.TOO_MANY_REQUESTS) - - data = [] - for idx, encoder_output in enumerate(results): - # Model-agnostic: serialize whatever per-request tensor encode() - # emits (classification logits, reward score, pooled vector, ...). - embedding = encoder_output.logits.flatten().tolist() - if request.encoding_format == "base64": - embedding = base64.b64encode( - array.array("f", embedding).tobytes()).decode("utf-8") - data.append( - EmbeddingResponseData(index=idx, embedding=embedding)) - - num_prompt_tokens = sum(len(t) for t in token_ids_list) - usage = EmbeddingUsageInfo(prompt_tokens=num_prompt_tokens, - total_tokens=num_prompt_tokens) - response = EmbeddingResponse(model=request.model, - data=data, - usage=usage) - return JSONResponse(content=response.model_dump()) - except Exception as e: - # Unexpected fault (not a client error): surface as 500, not 400. - logger.error(traceback.format_exc()) - return self.create_error_response( - str(e), - err_type="InternalServerError", - status_code=HTTPStatus.INTERNAL_SERVER_ERROR) - def register_visual_gen_routes(self): """Register routes for diffusion model serving.""" self.app.add_api_route("/metrics", @@ -2775,122 +2574,6 @@ def _no_such_batch(batch_id: str) -> 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: - - async def create_mm_embedding_response(promise: RequestOutput): - await promise.aresult() - # TODO: Replace mm_embedding_handles with a dedicated OpenAIBaseModel(JSON-safe), when enable multimodal disagg E2E - mm_embedding_handles = ( - promise.disaggregated_params.multimodal_embedding_handles - if promise.disaggregated_params else None) - if not mm_embedding_handles: - return self.create_error_response( - message="Multimodal embedding handle missing in response", - err_type="InternalServerError", - status_code=HTTPStatus.INTERNAL_SERVER_ERROR) - if any("tensor_size" not in h for h in mm_embedding_handles): - return self.create_error_response( - message="Multimodal embedding handle missing tensor_size", - err_type="InternalServerError", - status_code=HTTPStatus.INTERNAL_SERVER_ERROR) - mm_embedding_handle = (mm_embedding_handles[0] - if len(mm_embedding_handles) == 1 else - mm_embedding_handles) - num_tokens = sum( - int(h["tensor_size"][0]) for h in mm_embedding_handles) - return ChatCompletionResponse( - id=str(promise.request_id), - model=self.model, - choices=[ - ChatCompletionResponseChoice( - index=0, - message=ChatMessage(role="assistant", content="dummy"), - mm_embedding_handle=mm_embedding_handle, - finish_reason="length", - ) - ], - usage=UsageInfo( - prompt_tokens=num_tokens, - completion_tokens=1, - total_tokens=num_tokens + 1, - ), - ) - - try: - ensure_request_chat_template_allowed( - request, self.allow_request_chat_template) - conversation: List[ConversationMessage] = [] - # Same pipeline as the chat route: the server template and the - # model extension's tool serialization apply here too. - render_resources = RenderResources.from_server(self) - tool_dicts = None if request.tools is None else [ - render_resources.extension.serialize_tool(tool) - for tool in request.tools - ] - - try: - conversation, mm_coroutines, mm_placeholder_counts, mm_item_order = parse_chat_messages_coroutines( - request.messages, - self.model_config, - self.multimodal_server_config, - request_media_io_kwargs=request.media_io_kwargs, - ) - except ValidationError: - # ValidatorIterator rejects extra fields; fall back to raw JSON. - raw_body = await raw_request.json() - raw_messages = raw_body.get("messages", []) - conversation, mm_coroutines, mm_placeholder_counts, mm_item_order = parse_chat_messages_coroutines( - raw_messages, - self.model_config, - self.multimodal_server_config, - request_media_io_kwargs=request.media_io_kwargs, - ) - - if request.prompt_token_ids is not None: - prompt = request.prompt_token_ids - else: - render_task = render_conversation( - render_resources, - conversation=conversation, - mm_placeholder_counts=mm_placeholder_counts, - add_generation_prompt=request.add_generation_prompt, - tools=tool_dicts, - documents=request.documents, - chat_template=request.chat_template, - chat_template_kwargs=request.chat_template_kwargs, - injected_chat_template_kwargs=request. - injected_chat_template_kwargs, - ) - rendered, (mm_data, mm_embeddings) = await asyncio.gather( - render_task, mm_coroutines) - prompt = (rendered.text - if rendered.text is not None else rendered.token_ids) - prompt = prompt_inputs(prompt) - - if request.prompt_token_ids is not None: - mm_data, mm_embeddings = await mm_coroutines - if mm_embeddings: - raise ValueError("Cannot use multimodal embeddings as input") - if mm_data is not None: - self._attach_multimodal_data(prompt, request, mm_data, - mm_item_order) - - promise = self.generator.generate_async(inputs=prompt, ) - asyncio.create_task(self.await_disconnected(raw_request, promise)) - - response = await create_mm_embedding_response(promise) - return JSONResponse(content=response.model_dump()) - - except CppExecutorError: - logger.error(traceback.format_exc()) - # If internal executor error is raised, shutdown the server - _record_generator_termination(self.generator) - signal.raise_signal(signal.SIGINT) - except Exception as e: - logger.error(traceback.format_exc()) - return self.create_error_response(str(e)) - async def openai_completion(self, request: CompletionRequest, raw_request: Request) -> Response: @@ -3580,17 +3263,6 @@ async def openai_responses_delete_response( "deleted": True }) - async def tokenize(self, request: TokenizeRequest) -> JSONResponse: - try: - token_ids = self.tokenizer.encode(request.prompt) - response = TokenizeResponse(count=len(token_ids), tokens=token_ids) - return JSONResponse(content=response.model_dump()) - except Exception as e: - return self.create_error_response( - message=str(e), - err_type="InvalidRequestError", - status_code=HTTPStatus.BAD_REQUEST) - async def start_profile( self, request: Optional[StartProfileRequest] = None) -> JSONResponse: diff --git a/tests/unittest/llmapi/apps/test_encode_batcher.py b/tests/unittest/llmapi/apps/test_encode_batcher.py index 2d6fc174a199..5345360f4570 100644 --- a/tests/unittest/llmapi/apps/test_encode_batcher.py +++ b/tests/unittest/llmapi/apps/test_encode_batcher.py @@ -19,7 +19,7 @@ import pytest -from tensorrt_llm.serve.encode_batcher import EncodeBatcher +from tensorrt_llm.serve.api.openai.embeddings.batcher import EncodeBatcher pytestmark = pytest.mark.asyncio @@ -125,7 +125,7 @@ async def test_submit_raises_queue_full_when_queue_is_full(): """ import asyncio - from tensorrt_llm.serve.encode_batcher import QueueFullError + from tensorrt_llm.serve.api.openai.embeddings.batcher import QueueFullError def encode_fn(batch): return [0] * len(batch) @@ -279,7 +279,7 @@ async def test_submit_rejects_input_longer_than_max_seq_len(): submit raises InputTooLongError (mapped to HTTP 400) before the input ever enters the queue. """ - from tensorrt_llm.serve.encode_batcher import InputTooLongError + from tensorrt_llm.serve.api.openai.embeddings.batcher import InputTooLongError def encode_fn(batch): return [0] * len(batch) @@ -309,7 +309,7 @@ async def test_submit_rejects_single_input_over_token_budget(): than letting encode_fn reject the batch later. validate_input() exposes the same check for callers that want to pre-validate before enqueueing. """ - from tensorrt_llm.serve.encode_batcher import InputTooLongError + from tensorrt_llm.serve.api.openai.embeddings.batcher import InputTooLongError def encode_fn(batch): return [0] * len(batch) diff --git a/tests/unittest/serve/test_serve_runtime_api_calls.py b/tests/unittest/serve/test_serve_runtime_api_calls.py index a6cd1c05f2a1..35142dd775bf 100644 --- a/tests/unittest/serve/test_serve_runtime_api_calls.py +++ b/tests/unittest/serve/test_serve_runtime_api_calls.py @@ -15,7 +15,8 @@ from tensorrt_llm.llmapi.llm import EncodeLimits from tensorrt_llm.serve import openai_server -from tensorrt_llm.serve.openai_server import OpenAIServer, _record_generator_termination +from tensorrt_llm.serve._base.serving import _record_generator_termination +from tensorrt_llm.serve.openai_server import OpenAIServer pytestmark = pytest.mark.cpu_only From b00ab56dc47ebc7c4c7e016167ed640340cf7b7f Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Tue, 3 Nov 2026 13:57:15 +0000 Subject: [PATCH 26/30] [TRTLLM-17267][refactor] Wire the embeddings, multimodal-encoder and tokenize routers and serving objects Each new package gets an api_router.py built with make_router() and a RouterSpec in ROUTERS for exactly the roles the route had. The composition root builds RenderResources once after _init_llm and injects it, constructs the serving objects inside the role branches, puts them on app.state, and attaches the routers after the role switch, where every attach_router can check its app.state key. The two route registrations the routers now serve leave register_routes and register_mm_encoder_routes. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- tensorrt_llm/serve/_base/routing.py | 11 ++++ .../serve/api/openai/embeddings/api_router.py | 39 +++++++++++++ .../serve/api/openai/mm_encoder/api_router.py | 39 +++++++++++++ .../serve/api/openai/tokenize/api_router.py | 41 +++++++++++++ tensorrt_llm/serve/openai_server.py | 57 +++++++++++++++---- 5 files changed, 175 insertions(+), 12 deletions(-) create mode 100644 tensorrt_llm/serve/api/openai/embeddings/api_router.py create mode 100644 tensorrt_llm/serve/api/openai/mm_encoder/api_router.py create mode 100644 tensorrt_llm/serve/api/openai/tokenize/api_router.py diff --git a/tensorrt_llm/serve/_base/routing.py b/tensorrt_llm/serve/_base/routing.py index 3f461a029282..70facf89f0ef 100644 --- a/tensorrt_llm/serve/_base/routing.py +++ b/tensorrt_llm/serve/_base/routing.py @@ -77,6 +77,17 @@ 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.openai.embeddings.api_router", + roles=frozenset({RoleGroup.EMBEDDING}), + ), + RouterSpec( + "tensorrt_llm.serve.api.openai.mm_encoder.api_router", + roles=frozenset({RoleGroup.MM_ENCODER}), + ), + RouterSpec( + "tensorrt_llm.serve.api.openai.tokenize.api_router", roles=frozenset({RoleGroup.LLM}) + ), ) diff --git a/tensorrt_llm/serve/api/openai/embeddings/api_router.py b/tensorrt_llm/serve/api/openai/embeddings/api_router.py new file mode 100644 index 000000000000..426b6a28c528 --- /dev/null +++ b/tensorrt_llm/serve/api/openai/embeddings/api_router.py @@ -0,0 +1,39 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""HTTP routes of ``EmbeddingServing``. + +The composition root puts the serving object on ``app.state`` before the +router is attached; the route functions only map HTTP to it. +""" + +from fastapi import FastAPI, Request +from fastapi.responses import Response + +from tensorrt_llm.serve._base.routing import make_router +from tensorrt_llm.serve.api.openai.embeddings.serving import EmbeddingServing +from tensorrt_llm.serve.openai_protocol.embeddings import EmbeddingRequest + +APP_STATE_KEY = "openai_serving_embeddings" + +router = make_router() + + +def _serving(raw_request: Request) -> EmbeddingServing: + # No app.state default: wiring bugs fail loudly. + serving: EmbeddingServing = raw_request.app.state.openai_serving_embeddings + return serving + + +@router.post("/v1/embeddings") +async def openai_embedding(request: EmbeddingRequest, raw_request: Request) -> Response: + return await _serving(raw_request).openai_embedding(request, raw_request) + + +def attach_router(app: FastAPI) -> None: + if not hasattr(app.state, APP_STATE_KEY): + raise RuntimeError( + f"app.state.{APP_STATE_KEY} (an EmbeddingServing) must be set before attaching " + "tensorrt_llm.serve.api.openai.embeddings; the composition root sets it in the " + "EMBEDDING branch." + ) + app.include_router(router) diff --git a/tensorrt_llm/serve/api/openai/mm_encoder/api_router.py b/tensorrt_llm/serve/api/openai/mm_encoder/api_router.py new file mode 100644 index 000000000000..6b7e9573cd9c --- /dev/null +++ b/tensorrt_llm/serve/api/openai/mm_encoder/api_router.py @@ -0,0 +1,39 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""HTTP routes of ``MMEncoderServing``. + +The composition root puts the serving object on ``app.state`` before the +router is attached; the route functions only map HTTP to it. +""" + +from fastapi import FastAPI, Request +from fastapi.responses import Response + +from tensorrt_llm.serve._base.routing import make_router +from tensorrt_llm.serve.api.openai.mm_encoder.serving import MMEncoderServing +from tensorrt_llm.serve.openai_protocol.chat import ChatCompletionRequest + +APP_STATE_KEY = "openai_serving_mm_encoder" + +router = make_router() + + +def _serving(raw_request: Request) -> MMEncoderServing: + # No app.state default: wiring bugs fail loudly. + serving: MMEncoderServing = raw_request.app.state.openai_serving_mm_encoder + return serving + + +@router.post("/v1/chat/completions") +async def openai_mm_encoder(request: ChatCompletionRequest, raw_request: Request) -> Response: + return await _serving(raw_request).openai_mm_encoder(request, raw_request) + + +def attach_router(app: FastAPI) -> None: + if not hasattr(app.state, APP_STATE_KEY): + raise RuntimeError( + f"app.state.{APP_STATE_KEY} (an MMEncoderServing) must be set before attaching " + "tensorrt_llm.serve.api.openai.mm_encoder; the composition root sets it in the " + "MM_ENCODER branch." + ) + app.include_router(router) diff --git a/tensorrt_llm/serve/api/openai/tokenize/api_router.py b/tensorrt_llm/serve/api/openai/tokenize/api_router.py new file mode 100644 index 000000000000..c4f8607fcab5 --- /dev/null +++ b/tensorrt_llm/serve/api/openai/tokenize/api_router.py @@ -0,0 +1,41 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""HTTP routes of ``TokenizeServing``. + +The composition root puts the serving object on ``app.state`` before the +router is attached; the route functions only map HTTP to it. +""" + +from fastapi import FastAPI, Request +from fastapi.responses import JSONResponse + +from tensorrt_llm.serve._base.routing import make_router +from tensorrt_llm.serve.api.openai.tokenize.serving import TokenizeServing +from tensorrt_llm.serve.openai_protocol.tokenize import TokenizeRequest + +APP_STATE_KEY = "openai_serving_tokenize" + +router = make_router() + + +def _serving(raw_request: Request) -> TokenizeServing: + # No app.state default: wiring bugs fail loudly. + serving: TokenizeServing = raw_request.app.state.openai_serving_tokenize + return serving + + +@router.post("/_internal/tokenize") +async def tokenize(request: TokenizeRequest, raw_request: Request) -> JSONResponse: + # raw_request only locates the serving object; FastAPI leaves Request + # parameters out of the OpenAPI document, so the body schema is unchanged. + return await _serving(raw_request).tokenize(request) + + +def attach_router(app: FastAPI) -> None: + if not hasattr(app.state, APP_STATE_KEY): + raise RuntimeError( + f"app.state.{APP_STATE_KEY} (a TokenizeServing) must be set before attaching " + "tensorrt_llm.serve.api.openai.tokenize; the composition root sets it in the LLM " + "branch." + ) + app.include_router(router) diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py index 25cccf3b1e23..667a3142f70e 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -724,6 +724,8 @@ def __init__( self._init_visual_gen() else: self._init_llm(chat_template) + # Built once and injected into serving objects. + self._render_resources = RenderResources.from_server(self) @asynccontextmanager async def lifespan(app: FastAPI): @@ -857,10 +859,6 @@ 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)) - # 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, \ @@ -870,17 +868,51 @@ async def validation_exception_handler(request, exc): assert isinstance( self.generator, MultimodalEncoder ), "generator must be a MultimodalEncoder for multimodal encoder" + from tensorrt_llm.serve.api.openai.mm_encoder import \ + api_router as mm_encoder_router + from tensorrt_llm.serve.api.openai.mm_encoder.serving import \ + MMEncoderServing + setattr( + self.app.state, mm_encoder_router.APP_STATE_KEY, + MMEncoderServing( + render_resources=self._render_resources, + multimodal_server_config=self.multimodal_server_config, + allow_request_chat_template=self. + allow_request_chat_template, + **self._generate_serving_kwargs())) self.register_mm_encoder_routes() elif self.server_role is ServerRole.EMBEDDING: assert getattr(self.generator.args, "encode_only", False), ( "generator must be an encode_only=True LLM for the embedding " "server") - self._init_embedding_batcher() - self.register_embedding_routes() + from tensorrt_llm.serve.api.openai.embeddings import \ + api_router as embeddings_router + from tensorrt_llm.serve.api.openai.embeddings.serving import \ + EmbeddingServing + embedding_serving = EmbeddingServing( + model=self.model, + generator=self.generator, + embedding_max_queue_delay=self._embedding_max_queue_delay, + embedding_max_queue_size=self._embedding_max_queue_size) + # The lifespan starts and stops the batcher and the health probe + # reports it; both read it here. + self.embedding_batcher = embedding_serving.embedding_batcher + setattr(self.app.state, embeddings_router.APP_STATE_KEY, + embedding_serving) else: + from tensorrt_llm.serve.api.openai.tokenize import \ + api_router as tokenize_router + from tensorrt_llm.serve.api.openai.tokenize.serving import \ + TokenizeServing + setattr(self.app.state, tokenize_router.APP_STATE_KEY, + TokenizeServing(model=self.model, tokenizer=self.tokenizer)) self._register_launcher_forwarding_routes() self.register_routes() + # Attach routers after the role switch: every attach_router checks + # that its serving objects are on app.state. + attach_routers(self.app, ServerCapabilities.for_role(self.server_role)) + if self._collect_perf_metrics: self.app.add_middleware(PerfMetricsMiddleware, expose_headers=self._expose_perf_metrics, @@ -1116,6 +1148,13 @@ def _create_visual_gen_validation_error_response( status_code=HTTPStatus.UNPROCESSABLE_ENTITY, ) + def _generate_serving_kwargs(self) -> Dict[str, Any]: + """Constructor arguments of every GenerateBaseServing, read once.""" + return dict(model=self.model, + generator=self.generator, + tokenizer=self.tokenizer, + model_config=self.model_config) + def _check_health(self) -> bool: # An embedding server's requests flow through the batcher worker; if it # has exited the engine is up but every /v1/embeddings request hangs, so @@ -1230,9 +1269,6 @@ def register_routes(self): self.app.add_api_route('/v1/responses/{response_id}', self.openai_responses_delete_response, methods=["DELETE"]) - self.app.add_api_route("/_internal/tokenize", - self.tokenize, - methods=["POST"]) # Profiling endpoints (PyTorch backend only) self.app.add_api_route("/start_profile", @@ -1281,9 +1317,6 @@ def register_mm_encoder_routes(self): self.app.add_api_route("/metrics", self.get_iteration_stats, methods=["GET"]) - self.app.add_api_route("/v1/chat/completions", - self.openai_mm_encoder, - methods=["POST"]) async def _require_rl_control_auth(self, raw_request: Request) -> None: body = await raw_request.body() From 582820f1d28fd3052615ae08da34d342cadbe9a7 Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Tue, 3 Nov 2026 14:06:06 +0000 Subject: [PATCH 27/30] [TRTLLM-17267][test] Cover the embeddings, multimodal-encoder and tokenize serving classes on CPU Construction tests (every dependency is a required keyword-only argument), route tests through the real composition root and routers (msgpack bodies included), the batcher reference the health probe reads, the injected RenderResources carrying the server chat template, the serving bases and the borrowed-helper identity on OpenAIServer. The encode-only fake now encodes fixed embeddings and scripted outputs can carry disaggregated params; the app factory passes the chat-template knobs. The runtime-API tests build EmbeddingServing and patch the termination telemetry where the moved helper looks it up; the lazy-import guard also checks that importing tensorrt_llm.serve loads no endpoint package. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- tests/test_common/serve/app_factory.py | 4 + tests/test_common/serve/fakes.py | 22 ++- .../llmapi/apps/test_context_length_error.py | 16 ++ .../llmapi/apps/test_encode_batcher.py | 2 +- .../llmapi/apps/test_serve_lazy_import.py | 4 +- tests/unittest/serve/_base/test_serving.py | 133 ++++++++++++++ tests/unittest/serve/api/openai/__init__.py | 2 + .../serve/api/openai/embeddings/__init__.py | 2 + .../api/openai/embeddings/test_serving.py | 171 ++++++++++++++++++ .../serve/api/openai/mm_encoder/__init__.py | 2 + .../api/openai/mm_encoder/test_serving.py | 166 +++++++++++++++++ .../serve/api/openai/tokenize/__init__.py | 2 + .../serve/api/openai/tokenize/test_serving.py | 98 ++++++++++ .../test_openai_server_borrowed_helpers.py | 43 +++++ .../serve/test_serve_runtime_api_calls.py | 26 +-- .../unittest/serve/test_serve_test_helpers.py | 29 +++ 16 files changed, 704 insertions(+), 18 deletions(-) create mode 100644 tests/unittest/serve/_base/test_serving.py create mode 100644 tests/unittest/serve/api/openai/__init__.py create mode 100644 tests/unittest/serve/api/openai/embeddings/__init__.py create mode 100644 tests/unittest/serve/api/openai/embeddings/test_serving.py create mode 100644 tests/unittest/serve/api/openai/mm_encoder/__init__.py create mode 100644 tests/unittest/serve/api/openai/mm_encoder/test_serving.py create mode 100644 tests/unittest/serve/api/openai/tokenize/__init__.py create mode 100644 tests/unittest/serve/api/openai/tokenize/test_serving.py create mode 100644 tests/unittest/serve/test_openai_server_borrowed_helpers.py diff --git a/tests/test_common/serve/app_factory.py b/tests/test_common/serve/app_factory.py index 1e958e6b44f7..a512c0d5dcd6 100644 --- a/tests/test_common/serve/app_factory.py +++ b/tests/test_common/serve/app_factory.py @@ -106,6 +106,8 @@ class AggConfig: model_name: str = "fake-model" # OpenAIServer(model=...): the /v1/models id tool_parser: Optional[str] = None reasoning_parser: Optional[str] = None + chat_template: Optional[str] = None # server-side template, a literal Jinja string + allow_request_chat_template: bool = False # OpenAIServer(allow_request_chat_template=...) 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. @@ -240,6 +242,8 @@ def build_openai_server(self, config: AggConfig) -> OpenAIServer: tool_parser=config.tool_parser, server_role=config.role, metadata_server_cfg=None, + chat_template=config.chat_template, + allow_request_chat_template=config.allow_request_chat_template, 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, diff --git a/tests/test_common/serve/fakes.py b/tests/test_common/serve/fakes.py index 2318776cd4c8..c41bca32d041 100644 --- a/tests/test_common/serve/fakes.py +++ b/tests/test_common/serve/fakes.py @@ -37,6 +37,7 @@ from types import SimpleNamespace from typing import Any, AsyncIterator, Optional, Sequence, Union +import torch from tokenizers import Tokenizer from tokenizers.models import WordLevel from tokenizers.pre_tokenizers import Whitespace @@ -203,6 +204,9 @@ class ScriptedOutput: 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 + # What the result's ``disaggregated_params`` carries, for example the + # multimodal embedding handles the multimodal-encoder role answers with. + disaggregated_params: Any = None @dataclass @@ -264,6 +268,7 @@ def __init__( ) self._script = script self._fake_tokenizer = tokenizer + self._disaggregated_params = script.disaggregated_params if script.prompt_token_ids is not None: self.prompt_token_ids = list(script.prompt_token_ids) try: @@ -470,17 +475,26 @@ class FakeEncodeOnlyLLM(FakeLLM): 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. + ``_encoder_executor`` in encode-only mode. ``encode_limits`` is the real + ``LLM`` property, read over ``FakeEncoderExecutor.model_engine``. + + ``encode`` records every batch in ``encode_calls`` and answers each input + with the fixed two-value embedding ``[len(token_ids), first token id]``, + so a test can tell the inputs apart in the response. """ def __init__(self, *, args: FakeLlmArgs) -> None: self._setup_fake(args=args, executor=None, encode_only=True) self._encoder_executor = FakeEncoderExecutor() + self.encode_calls: list = [] def encode(self, token_ids_batch: Any) -> Any: - raise NotImplementedError( - "FakeEncodeOnlyLLM does not encode; script it when a test needs it." - ) + batch = [list(token_ids) for token_ids in token_ids_batch] + self.encode_calls.append(batch) + return [ + SimpleNamespace(logits=torch.tensor([float(len(ids)), float(ids[0]) if ids else 0.0])) + for ids in batch + ] class UserProbeMiddleware: diff --git a/tests/unittest/llmapi/apps/test_context_length_error.py b/tests/unittest/llmapi/apps/test_context_length_error.py index c1004c1f8fdc..5a2cdfa4452c 100644 --- a/tests/unittest/llmapi/apps/test_context_length_error.py +++ b/tests/unittest/llmapi/apps/test_context_length_error.py @@ -169,3 +169,19 @@ def test_disagg_route_other_upstream_errors_unchanged(): _disagg_server()._handle_exception(_upstream_error(worker_body)) assert excinfo.value.status_code == 400 assert excinfo.value.detail == "`max_tokens` (0) must be greater than 0" + + +def test_serving_classes_share_the_server_error_factory(): + """OpenAIServer and the serving classes render errors through one function. + + The one function on BaseServing keeps the envelope and the context-length + code rule from drifting between serve/api/ and the composition root. + """ + from tensorrt_llm.serve._base.serving import BaseServing + + assert OpenAIServer.create_error_response is BaseServing.create_error_response + msg = context_length_exceeded_message(2048, 4096) + response = BaseServing().create_error_response(msg) + assert response.status_code == HTTPStatus.BAD_REQUEST + assert _body(response) == _body(OpenAIServer.create_error_response(msg)) + assert _body(response)["code"] == CONTEXT_LENGTH_EXCEEDED_CODE diff --git a/tests/unittest/llmapi/apps/test_encode_batcher.py b/tests/unittest/llmapi/apps/test_encode_batcher.py index 5345360f4570..9f9abb67b1fd 100644 --- a/tests/unittest/llmapi/apps/test_encode_batcher.py +++ b/tests/unittest/llmapi/apps/test_encode_batcher.py @@ -11,7 +11,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. -"""Unit tests for the encode dynamic batcher (tensorrt_llm/serve/encode_batcher.py). +"""Unit tests for the encode dynamic batcher (tensorrt_llm/serve/api/openai/embeddings/batcher.py). These tests exercise the pure-async coalescing logic with an injected, synchronous `encode_fn`. They do not require a GPU or any TensorRT-LLM runtime. diff --git a/tests/unittest/llmapi/apps/test_serve_lazy_import.py b/tests/unittest/llmapi/apps/test_serve_lazy_import.py index 5fa30431da1b..53323c87f51e 100644 --- a/tests/unittest/llmapi/apps/test_serve_lazy_import.py +++ b/tests/unittest/llmapi/apps/test_serve_lazy_import.py @@ -46,10 +46,12 @@ def test_importing_the_package_does_not_load_the_servers() -> None: print(json.dumps({ "openai_server": "tensorrt_llm.serve.openai_server" in sys.modules, "openai_disagg_server": "tensorrt_llm.serve.openai_disagg_server" in sys.modules, + # The endpoint packages load per role, from the composition root. + "api": any(m.startswith("tensorrt_llm.serve.api.") for m in sys.modules), })) """ ) - assert out == {"openai_server": False, "openai_disagg_server": False} + assert out == {"openai_server": False, "openai_disagg_server": False, "api": False} def test_servers_load_on_first_access_and_are_the_real_classes() -> None: diff --git a/tests/unittest/serve/_base/test_serving.py b/tests/unittest/serve/_base/test_serving.py new file mode 100644 index 000000000000..397755a720c7 --- /dev/null +++ b/tests/unittest/serve/_base/test_serving.py @@ -0,0 +1,133 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The serving bases and the shared termination helper in ``serve/_base/serving.py``.""" + +import asyncio +from types import SimpleNamespace + +import pytest + +from tensorrt_llm.serve._base import serving as base_serving +from tensorrt_llm.serve._base.serving import ( + BaseServing, + GenerateBaseServing, + _record_generator_termination, +) + +pytestmark = pytest.mark.cpu_only + +GENERATE_REQUIRED = ("model", "generator", "tokenizer", "model_config") + + +def _generate_dependencies(): + return { + "model": "fake-model", + "generator": object(), + "tokenizer": object(), + "model_config": None, + } + + +def test_base_serving_has_no_constructor_and_no_state(): + assert "__init__" not in vars(BaseServing), ( + "BaseServing has no constructor: subclasses set every attribute they read themselves" + ) + assert vars(BaseServing()) == {} + + +def test_generate_base_serving_requires_each_dependency(): + dependencies = _generate_dependencies() + serving = GenerateBaseServing(**dependencies) # the control + assert vars(serving) == dependencies + for name in GENERATE_REQUIRED: + partial = {key: value for key, value in dependencies.items() if key != name} + try: + GenerateBaseServing(**partial) + except TypeError: + continue + pytest.fail( + f"GenerateBaseServing must require {name!r}: serving dependencies are keyword-only " + "constructor arguments without defaults (serve/AGENTS.md)" + ) + + +class _Request: + """What ``await_disconnected`` reads of a starlette request.""" + + def __init__(self, disconnected, no_client_connection=False): + self._disconnected = list(disconnected) # one answer per poll; a missing one is a test bug + self.state = SimpleNamespace() + if no_client_connection: + self.state.no_client_connection = True + self.client = ("127.0.0.1", 1234) + self.polls = 0 + + async def is_disconnected(self): + self.polls += 1 + return self._disconnected.pop(0) + + +class _Promise: + def __init__(self, finished): + self.finished = finished + self.request_id = 7 + self.aborted = False + + def abort(self): + self.aborted = True + + +def _await_disconnected(request, promise): + serving = GenerateBaseServing(**_generate_dependencies()) + asyncio.run(serving.await_disconnected(request, promise)) + + +def test_await_disconnected_returns_at_once_without_a_request(): + promise = _Promise(finished=False) + _await_disconnected(None, promise) + assert not promise.aborted + + +def test_await_disconnected_returns_at_once_for_a_synthetic_request(): + # A batched request has no client socket: it must not be polled at all. + request = _Request(disconnected=[], no_client_connection=True) + promise = _Promise(finished=False) + _await_disconnected(request, promise) + assert request.polls == 0 + assert not promise.aborted + + +def test_await_disconnected_aborts_an_unfinished_promise_once_the_client_leaves(monkeypatch): + sleeps = [] + real_sleep = asyncio.sleep + + async def fast_sleep(delay): + sleeps.append(delay) + await real_sleep(0) + + monkeypatch.setattr(base_serving.asyncio, "sleep", fast_sleep) + request = _Request(disconnected=[False, True]) + promise = _Promise(finished=False) + _await_disconnected(request, promise) + assert request.polls == 2 + assert sleeps == [1] # polled once a second + assert promise.aborted + + +def test_await_disconnected_leaves_a_finished_promise_alone(): + request = _Request(disconnected=[True]) + promise = _Promise(finished=True) + _await_disconnected(request, promise) + assert not promise.aborted + + +def test_record_generator_termination_blames_the_worker_behind_a_proxy(monkeypatch): + outcomes = [] + # Patched where the moved helper looks the name up, not on the composition root. + monkeypatch.setattr(base_serving, "record_termination_observation", outcomes.append) + _record_generator_termination(SimpleNamespace(engine_in_worker_process=True)) + (outcome,) = outcomes + assert outcome.termination_kind == "worker_failure" + assert outcome.component == "engine_worker" + assert outcome.reporting_source == "executor_proxy" + assert outcome.exit_code_known is False diff --git a/tests/unittest/serve/api/openai/__init__.py b/tests/unittest/serve/api/openai/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tests/unittest/serve/api/openai/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tests/unittest/serve/api/openai/embeddings/__init__.py b/tests/unittest/serve/api/openai/embeddings/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tests/unittest/serve/api/openai/embeddings/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tests/unittest/serve/api/openai/embeddings/test_serving.py b/tests/unittest/serve/api/openai/embeddings/test_serving.py new file mode 100644 index 000000000000..6a5e340a3414 --- /dev/null +++ b/tests/unittest/serve/api/openai/embeddings/test_serving.py @@ -0,0 +1,171 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""``POST /v1/embeddings`` through ``EmbeddingServing`` and its router. + +Pins what the embeddings route does today on an EMBEDDING server built by the +real composition root: the request path through the encode batcher, the error +statuses, the msgpack route class, and the batcher reference the health probe +reads. +""" + +import array +import base64 + +import pytest +from fastapi import FastAPI +from test_common.serve.app_factory import AggConfig +from test_common.serve.fakes import FakeEncodeOnlyLLM, FakeLlmArgs + +from tensorrt_llm.llmapi.disagg_utils import ServerRole +from tensorrt_llm.serve.api.openai.embeddings import api_router +from tensorrt_llm.serve.api.openai.embeddings.batcher import QueueFullError +from tensorrt_llm.serve.api.openai.embeddings.serving import EmbeddingServing +from tensorrt_llm.serve.openai_client import MSGPACK_HEADERS, _msgpack_encoder + +pytestmark = [pytest.mark.cpu_only, pytest.mark.threadleak(enabled=False)] + +REQUIRED = ("model", "generator", "embedding_max_queue_delay", "embedding_max_queue_size") +# "hi" and "hello" are single tokens (5 and 6) of the fake tokenizer. +BODY = {"model": "fake-model", "input": ["hi", "hello"]} + + +def _dependencies(tmp_path): + generator = FakeEncodeOnlyLLM(args=FakeLlmArgs(model=str(tmp_path), encode_only=True)) + return { + "model": "fake-model", + "generator": generator, + "embedding_max_queue_delay": 0.001, + "embedding_max_queue_size": 4, + } + + +def _serving(server) -> EmbeddingServing: + return getattr(server.app.state, api_router.APP_STATE_KEY) + + +def _embedding_server(serve_env): + server = serve_env.build_openai_server(AggConfig(role=ServerRole.EMBEDDING)) + return server, serve_env.client(server.app) # the lifespan starts the batcher + + +def test_constructor_requires_each_dependency(tmp_path): + dependencies = _dependencies(tmp_path) + # The control: with every dependency present, construction succeeds. + EmbeddingServing(**dependencies) + for name in REQUIRED: + partial = {key: value for key, value in dependencies.items() if key != name} + try: + EmbeddingServing(**partial) + except TypeError: + continue + pytest.fail( + f"EmbeddingServing must require {name!r}: serving dependencies are keyword-only " + "constructor arguments without defaults (serve/AGENTS.md)" + ) + with pytest.raises(TypeError): + EmbeddingServing(*dependencies.values()) # keyword-only: no positional construction + + +def test_batcher_built_from_encode_limits(tmp_path): + dependencies = _dependencies(tmp_path) + # The real LLM.encode_limits, read over the fake encoder executor. + limits = dependencies["generator"].encode_limits + assert limits is not None + batcher = EmbeddingServing(**dependencies).embedding_batcher + assert batcher._max_batch_size == limits.max_batch_size + assert batcher._max_num_tokens == limits.max_num_tokens + assert batcher._max_seq_len == limits.max_seq_len + assert batcher._max_queue_delay == 0.001 + assert batcher._queue.maxsize == 4 + + +def test_embeddings_route_end_to_end(serve_env): + server, client = _embedding_server(serve_env) + response = client.post("/v1/embeddings", json=BODY) + assert response.status_code == 200, response.text + body = response.json() + assert body["model"] == "fake-model" + assert [item["index"] for item in body["data"]] == [0, 1] + # The fake encodes each input as [len(token_ids), first token id]. + assert body["data"][0]["embedding"] == [1.0, 5.0] + assert body["data"][1]["embedding"] == [1.0, 6.0] + assert body["usage"]["prompt_tokens"] == 2 + assert body["usage"]["total_tokens"] == 2 + # Strings are tokenized through the generator's input processor before encode(). + assert [ids for batch in server.generator.encode_calls for ids in batch] == [[5], [6]] + + +def test_base64_encoding_format(serve_env): + _, client = _embedding_server(serve_env) + response = client.post( + "/v1/embeddings", json={"model": "fake-model", "input": "hi", "encoding_format": "base64"} + ) + assert response.status_code == 200, response.text + packed = base64.b64encode(array.array("f", [1.0, 5.0]).tobytes()).decode("utf-8") + assert response.json()["data"][0]["embedding"] == packed + + +def test_embeddings_route_decodes_msgpack(serve_env): + _, client = _embedding_server(serve_env) + response = client.post( + "/v1/embeddings", content=_msgpack_encoder.encode(BODY), headers=MSGPACK_HEADERS + ) + assert response.status_code == 200, response.text + assert len(response.json()["data"]) == 2 + + +def test_dimensions_and_empty_input_are_400(serve_env): + server, client = _embedding_server(serve_env) + response = client.post( + "/v1/embeddings", json={"model": "fake-model", "input": "hi", "dimensions": 8} + ) + assert response.status_code == 400 + assert "Matryoshka" in response.json()["message"] + response = client.post("/v1/embeddings", json={"model": "fake-model", "input": []}) + assert response.status_code == 400 + assert response.json()["message"] == "`input` must not be empty." + assert server.generator.encode_calls == [] + + +def test_input_too_long_is_400(serve_env): + server, client = _embedding_server(serve_env) + too_long = [5] * (server.generator.encode_limits.max_seq_len + 1) + response = client.post("/v1/embeddings", json={"model": "fake-model", "input": too_long}) + assert response.status_code == 400 + assert "exceeds max_seq_len" in response.json()["message"] + assert server.generator.encode_calls == [], "an oversize input is rejected before any encode" + + +def test_full_queue_is_429(serve_env, monkeypatch): + server, client = _embedding_server(serve_env) + + async def full(token_ids): + raise QueueFullError("Embedding request queue is full; retry later.") + + monkeypatch.setattr(_serving(server).embedding_batcher, "submit", full) + response = client.post("/v1/embeddings", json=BODY) + assert response.status_code == 429 + assert response.json()["message"] == "Embedding request queue is full; retry later." + + +def test_health_probe_sees_the_serving_batcher(serve_env, monkeypatch): + server, client = _embedding_server(serve_env) + serving = _serving(server) + message = ( + "the composition root must keep self.embedding_batcher bound to " + "EmbeddingServing.embedding_batcher, which the /health probe reads" + ) + assert server.embedding_batcher is serving.embedding_batcher, message + assert client.get("/health").status_code == 200 + monkeypatch.setattr(serving.embedding_batcher, "is_alive", lambda: False) + assert client.get("/health").status_code == 503, message + + +def test_embedding_role_has_no_metrics_route(serve_env): + _, client = _embedding_server(serve_env) + assert client.get("/metrics").status_code == 404 + + +def test_attach_router_requires_the_serving_object(): + with pytest.raises(RuntimeError, match=api_router.APP_STATE_KEY): + api_router.attach_router(FastAPI()) diff --git a/tests/unittest/serve/api/openai/mm_encoder/__init__.py b/tests/unittest/serve/api/openai/mm_encoder/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tests/unittest/serve/api/openai/mm_encoder/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tests/unittest/serve/api/openai/mm_encoder/test_serving.py b/tests/unittest/serve/api/openai/mm_encoder/test_serving.py new file mode 100644 index 000000000000..68c62f46d18a --- /dev/null +++ b/tests/unittest/serve/api/openai/mm_encoder/test_serving.py @@ -0,0 +1,166 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The multimodal-encoder ``POST /v1/chat/completions`` through ``MMEncoderServing``. + +Pins what the route does today on an MM_ENCODER server built by the real +composition root: the embedding handle in the response, the injected +``RenderResources`` (the server chat template reaches the prompt), the +request-template gate, the msgpack route class and the error shape. +""" + +import pytest +from fastapi import FastAPI +from test_common.serve.app_factory import AggConfig +from test_common.serve.fakes import FakeLlmArgs, FakeMultimodalEncoder, ScriptedOutput + +from tensorrt_llm.disaggregated_params import DisaggregatedParams +from tensorrt_llm.llmapi.disagg_utils import ServerRole +from tensorrt_llm.serve.api.openai.mm_encoder import api_router +from tensorrt_llm.serve.api.openai.mm_encoder.serving import MMEncoderServing +from tensorrt_llm.serve.openai_client import MSGPACK_HEADERS, _msgpack_encoder +from tensorrt_llm.serve.render import RenderResources + +pytestmark = [pytest.mark.cpu_only, pytest.mark.threadleak(enabled=False)] + +REQUIRED = ( + "render_resources", + "multimodal_server_config", + "allow_request_chat_template", + "model", + "generator", + "tokenizer", + "model_config", +) +CHAT_BODY = {"model": "fake-model", "messages": [{"role": "user", "content": "hi"}]} +# Distinct from the fake tokenizer's own template, so the rendered prompt tells +# which template produced it. +SERVER_TEMPLATE = "{% for m in messages %}[{{ m['role'] }}] {{ m['content'] }}{% endfor %}" +REQUEST_TEMPLATE = "{% for m in messages %}<{{ m['content'] }}>{% endfor %}" +HANDLE = {"tensor_size": [3, 16], "shm_name": "fake-handle"} + + +def _dependencies(tmp_path): + generator = FakeMultimodalEncoder(args=FakeLlmArgs(model=str(tmp_path))) + return { + "render_resources": RenderResources(tokenizer=generator.tokenizer), + "multimodal_server_config": None, + "allow_request_chat_template": False, + "model": "fake-model", + "generator": generator, + "tokenizer": generator.tokenizer, + "model_config": None, + } + + +def _serving(server) -> MMEncoderServing: + return getattr(server.app.state, api_router.APP_STATE_KEY) + + +def _mm_encoder_server(serve_env, **config): + server = serve_env.build_openai_server(AggConfig(role=ServerRole.MM_ENCODER, **config)) + return server, serve_env.client(server.app) + + +def _script_handles(server, *handles): + server.generator.script( + ScriptedOutput( + disaggregated_params=DisaggregatedParams(multimodal_embedding_handles=list(handles)) + ) + ) + + +def test_constructor_requires_each_dependency(tmp_path): + dependencies = _dependencies(tmp_path) + # The control: with every dependency present, construction succeeds. + MMEncoderServing(**dependencies) + for name in REQUIRED: + partial = {key: value for key, value in dependencies.items() if key != name} + try: + MMEncoderServing(**partial) + except TypeError: + continue + pytest.fail( + f"MMEncoderServing must require {name!r}: serving dependencies are keyword-only " + "constructor arguments without defaults (serve/AGENTS.md)" + ) + + +def test_render_resources_are_the_servers(serve_env): + server, _ = _mm_encoder_server(serve_env, chat_template=SERVER_TEMPLATE) + serving = _serving(server) + assert serving._render_resources is server._render_resources, ( + "MMEncoderServing must use the RenderResources the composition root built, never " + "RenderResources.from_server(serving), which silently drops the server chat template" + ) + assert server.chat_template == SERVER_TEMPLATE + assert server._render_resources.default_chat_template == server.chat_template + + +def test_mm_encoder_route_returns_embedding_handle(serve_env): + server, client = _mm_encoder_server(serve_env, chat_template=SERVER_TEMPLATE) + _script_handles(server, HANDLE) + response = client.post("/v1/chat/completions", json=CHAT_BODY) + assert response.status_code == 200, response.text + body = response.json() + assert body["model"] == "fake-model" + choice = body["choices"][0] + assert choice["mm_embedding_handle"] == HANDLE + assert choice["finish_reason"] == "length" + assert choice["message"]["role"] == "assistant" + assert choice["message"]["content"] == "dummy" + assert body["usage"]["prompt_tokens"] == 3 # tensor_size[0] of the one handle + assert body["usage"]["completion_tokens"] == 1 + assert body["usage"]["total_tokens"] == 4 + # The prompt was rendered with the server template the injected resources carry. + assert server.generator.generate_calls[0].inputs == {"prompt": "[user] hi"} + + +def test_mm_encoder_route_decodes_msgpack(serve_env): + server, client = _mm_encoder_server(serve_env) + _script_handles(server, HANDLE) + response = client.post( + "/v1/chat/completions", content=_msgpack_encoder.encode(CHAT_BODY), headers=MSGPACK_HEADERS + ) + assert response.status_code == 200, response.text + assert response.json()["choices"][0]["mm_embedding_handle"] == HANDLE + + +def test_missing_handle_is_an_error_response(serve_env): + server, client = _mm_encoder_server(serve_env) + # The default scripted output carries no disaggregated params, so no handle. + response = client.post("/v1/chat/completions", json=CHAT_BODY) + # Today's shape, kept as it is: the handler builds a 500 for the missing + # handle, then calls model_dump() on that response, and the resulting + # AttributeError is answered as a 400 BadRequestError. + assert response.status_code == 400 + body = response.json() + assert body["object"] == "error" + assert body["type"] == "BadRequestError" + assert len(server.generator.generate_calls) == 1 + + +def test_request_chat_template_rejected_unless_allowed(serve_env): + server, client = _mm_encoder_server(serve_env) + response = client.post( + "/v1/chat/completions", json={**CHAT_BODY, "chat_template": REQUEST_TEMPLATE} + ) + assert response.status_code == 400 + assert "chat_template cannot be supplied per request" in response.json()["message"] + assert server.generator.generate_calls == [] + + +def test_request_chat_template_renders_when_allowed(serve_env): + server, client = _mm_encoder_server( + serve_env, chat_template=SERVER_TEMPLATE, allow_request_chat_template=True + ) + _script_handles(server, HANDLE) + response = client.post( + "/v1/chat/completions", json={**CHAT_BODY, "chat_template": REQUEST_TEMPLATE} + ) + assert response.status_code == 200, response.text + assert server.generator.generate_calls[0].inputs == {"prompt": ""} + + +def test_attach_router_requires_the_serving_object(): + with pytest.raises(RuntimeError, match=api_router.APP_STATE_KEY): + api_router.attach_router(FastAPI()) diff --git a/tests/unittest/serve/api/openai/tokenize/__init__.py b/tests/unittest/serve/api/openai/tokenize/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tests/unittest/serve/api/openai/tokenize/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tests/unittest/serve/api/openai/tokenize/test_serving.py b/tests/unittest/serve/api/openai/tokenize/test_serving.py new file mode 100644 index 000000000000..5f5921da2822 --- /dev/null +++ b/tests/unittest/serve/api/openai/tokenize/test_serving.py @@ -0,0 +1,98 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""``POST /_internal/tokenize`` through ``TokenizeServing`` and its router. + +Pins what the route does today on the LLM roles (default, CONTEXT and +GENERATION) built by the real composition root, and that the other roles do +not serve it. +""" + +from types import SimpleNamespace + +import pytest +from fastapi import FastAPI +from test_common.serve.app_factory import AggConfig + +from tensorrt_llm.llmapi.disagg_utils import ServerRole +from tensorrt_llm.serve.api.openai.tokenize import api_router +from tensorrt_llm.serve.api.openai.tokenize.serving import TokenizeServing + +pytestmark = [pytest.mark.cpu_only, pytest.mark.threadleak(enabled=False)] + +REQUIRED = ("model", "tokenizer") + + +def _dependencies(): + return {"model": "fake-model", "tokenizer": SimpleNamespace(encode=lambda text: [1])} + + +def _llm_server(serve_env, role=None): + server = serve_env.build_openai_server(AggConfig(role=role)) + return server, serve_env.client(server.app) + + +def test_constructor_requires_each_dependency(): + dependencies = _dependencies() + # The control: with every dependency present, construction succeeds. + TokenizeServing(**dependencies) + for name in REQUIRED: + partial = {key: value for key, value in dependencies.items() if key != name} + try: + TokenizeServing(**partial) + except TypeError: + continue + pytest.fail( + f"TokenizeServing must require {name!r}: serving dependencies are keyword-only " + "constructor arguments without defaults (serve/AGENTS.md)" + ) + + +@pytest.mark.parametrize( + "role", + [None, ServerRole.CONTEXT, ServerRole.GENERATION], + ids=["llm", "context", "generation"], +) +def test_tokenize_route(role, serve_env): + server, client = _llm_server(serve_env, role) + response = client.post("/_internal/tokenize", json={"prompt": "hello"}) + assert response.status_code == 200, response.text + body = response.json() + assert body["tokens"] == list(server.tokenizer.encode("hello")) + assert body["count"] == len(body["tokens"]) == 1 + + +def test_tokenize_validation_error_has_the_servers_error_shape(serve_env): + _, client = _llm_server(serve_env) + response = client.post("/_internal/tokenize", json={}) + assert response.status_code == 400 + body = response.json() + assert set(body) == {"error"} # the validation handler of the composition root + assert "prompt" in body["error"] + + +def test_tokenizer_failure_is_400_invalid_request(serve_env, monkeypatch): + server, client = _llm_server(serve_env) + serving = getattr(server.app.state, api_router.APP_STATE_KEY) + + def boom(text): + raise ValueError("cannot tokenize") + + monkeypatch.setattr(serving, "tokenizer", SimpleNamespace(encode=boom)) + response = client.post("/_internal/tokenize", json={"prompt": "hello"}) + assert response.status_code == 400 + body = response.json() + assert body["type"] == "InvalidRequestError" + assert body["message"] == "cannot tokenize" + + +@pytest.mark.parametrize( + "role", [ServerRole.MM_ENCODER, ServerRole.EMBEDDING], ids=["mm_encoder", "embedding"] +) +def test_tokenize_is_not_served_on_other_roles(role, serve_env): + _, client = _llm_server(serve_env, role) + assert client.post("/_internal/tokenize", json={"prompt": "hello"}).status_code == 404 + + +def test_attach_router_requires_the_serving_object(): + with pytest.raises(RuntimeError, match=api_router.APP_STATE_KEY): + api_router.attach_router(FastAPI()) diff --git a/tests/unittest/serve/test_openai_server_borrowed_helpers.py b/tests/unittest/serve/test_openai_server_borrowed_helpers.py new file mode 100644 index 000000000000..917fa7c463c6 --- /dev/null +++ b/tests/unittest/serve/test_openai_server_borrowed_helpers.py @@ -0,0 +1,43 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""``OpenAIServer`` borrows the helpers that moved to the serving bases; it keeps no copies. + +A helper that moved into a serving base while handlers still on ``OpenAIServer`` +call it is reachable on the class as the moved object itself. A copy would +drift from the moved one, so identity is what this file pins. Each move that +takes the last ``OpenAIServer`` caller of a helper removes its row here. +""" + +from types import SimpleNamespace + +import pytest + +from tensorrt_llm.serve._base.serving import BaseServing, GenerateBaseServing +from tensorrt_llm.serve.openai_server import OpenAIServer + +pytestmark = pytest.mark.cpu_only + +# Helper name -> the base class it moved to. +BORROWED = { + "create_error_response": BaseServing, + "await_disconnected": GenerateBaseServing, + "_attach_multimodal_data": GenerateBaseServing, +} + + +@pytest.mark.parametrize("name", sorted(BORROWED)) +def test_openai_server_borrows_the_moved_helper(name): + base = BORROWED[name] + assert getattr(OpenAIServer, name) is getattr(base, name), ( + f"OpenAIServer.{name} must be the helper moved to {base.__name__} " + f"(staticmethod({base.__name__}.{name}) for a static method, {base.__name__}.{name} " + "otherwise); never copy a moved helper back: a copy drifts from the moved one" + ) + + +def test_static_helpers_survived_aliasing(): + # A bare `Base.helper` alias would bind the server as the first argument. + assert OpenAIServer.create_error_response("x").status_code == 400 + prompt = {} + OpenAIServer._attach_multimodal_data(prompt, SimpleNamespace(prompt_token_ids=None), "media") + assert prompt == {"multi_modal_data": "media", "chat_template_applied": True} diff --git a/tests/unittest/serve/test_serve_runtime_api_calls.py b/tests/unittest/serve/test_serve_runtime_api_calls.py index 35142dd775bf..783d681f0c7a 100644 --- a/tests/unittest/serve/test_serve_runtime_api_calls.py +++ b/tests/unittest/serve/test_serve_runtime_api_calls.py @@ -2,10 +2,10 @@ # SPDX-License-Identifier: Apache-2.0 """The HTTP server reads the engine through public LLM members only. -Each test calls a production ``OpenAIServer`` method on a stand-in object (not -``object.__new__(OpenAIServer)``) whose generator double defines exactly the -public members the method reads, so a server method reaching for a private -engine member again fails loudly here. +Each test calls a production ``OpenAIServer`` or serving-class method on a +stand-in object (not ``object.__new__(OpenAIServer)``) whose generator double +defines exactly the public members the method reads, so a server method +reaching for a private engine member again fails loudly here. """ from types import SimpleNamespace @@ -15,7 +15,9 @@ from tensorrt_llm.llmapi.llm import EncodeLimits from tensorrt_llm.serve import openai_server +from tensorrt_llm.serve._base import serving as base_serving from tensorrt_llm.serve._base.serving import _record_generator_termination +from tensorrt_llm.serve.api.openai.embeddings.serving import EmbeddingServing from tensorrt_llm.serve.openai_server import OpenAIServer pytestmark = pytest.mark.cpu_only @@ -42,21 +44,20 @@ def _generator(**kwargs): # --------------------------------------------------------------------------- -# _init_embedding_batcher reads generator.encode_limits +# EmbeddingServing builds its batcher from generator.encode_limits # --------------------------------------------------------------------------- def test_embedding_batcher_takes_the_limits_encode_checks(): - stub = _ServerStub( + serving = EmbeddingServing( + model="fake-model", generator=_generator( encode_limits=EncodeLimits(max_batch_size=4, max_num_tokens=512, max_seq_len=128), encode=lambda batch: batch, ), - _embedding_max_queue_delay=0.01, - _embedding_max_queue_size=8, + embedding_max_queue_delay=0.01, + embedding_max_queue_size=8, ) - OpenAIServer._init_embedding_batcher(stub) - - batcher = stub.embedding_batcher + batcher = serving.embedding_batcher assert batcher._max_batch_size == 4 assert batcher._max_num_tokens == 512 assert batcher._max_seq_len == 128 @@ -152,7 +153,8 @@ class _Config: # --------------------------------------------------------------------------- def _recorded_outcome(monkeypatch, generator): outcomes = [] - monkeypatch.setattr(openai_server, "record_termination_observation", outcomes.append) + # Patched where the helper looks the name up: it lives in serve/_base/serving.py. + monkeypatch.setattr(base_serving, "record_termination_observation", outcomes.append) _record_generator_termination(generator) assert len(outcomes) == 1 return outcomes[0] diff --git a/tests/unittest/serve/test_serve_test_helpers.py b/tests/unittest/serve/test_serve_test_helpers.py index aa70512d9986..9a5782a6d3b9 100644 --- a/tests/unittest/serve/test_serve_test_helpers.py +++ b/tests/unittest/serve/test_serve_test_helpers.py @@ -25,6 +25,7 @@ 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 ( + FakeEncodeOnlyLLM, FakeLLM, FakeLlmArgs, ScriptedOutput, @@ -402,3 +403,31 @@ async def consume_rest(): await hanging asyncio.run(scenario()) + + +def test_scripted_disaggregated_params_reach_the_result(): + llm = FakeLLM(args=FakeLlmArgs(model="unused")) + default = llm.generate_async({"prompt_token_ids": [1]}, SamplingParams()) + assert default.disaggregated_params is None + llm.script(ScriptedOutput(disaggregated_params="scripted params")) + scripted = llm.generate_async({"prompt_token_ids": [1]}, SamplingParams()) + assert scripted.disaggregated_params == "scripted params" + + +def test_encode_only_fake_answers_fixed_embeddings_and_records_batches(): + llm = FakeEncodeOnlyLLM(args=FakeLlmArgs(model="unused", encode_only=True)) + outputs = llm.encode([[5], [6, 7]]) + assert [output.logits.flatten().tolist() for output in outputs] == [[1.0, 5.0], [2.0, 6.0]] + assert llm.encode_calls == [[[5], [6, 7]]] + # The real LLM.encode_limits property, over the fake encoder executor. + limits = llm.encode_limits + assert (limits.max_batch_size, limits.max_num_tokens, limits.max_seq_len) == (8, 8192, 2048) + + +def test_agg_config_passes_the_chat_template_knobs(serve_env): + template = "{% for m in messages %}{{ m['content'] }}{% endfor %}" + server = serve_env.build_openai_server( + AggConfig(chat_template=template, allow_request_chat_template=True) + ) + assert server.chat_template == template + assert server.allow_request_chat_template is True From 63d51ce175474a257af1b34711305e2df0bb91d4 Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Tue, 3 Nov 2026 14:06:21 +0000 Subject: [PATCH 28/30] [TRTLLM-17267][infra] Register the new serving packages with CI, mypy and the serve guide The multi-GPU trigger table gains the three package prefixes, so code moved out of tensorrt_llm/serve/openai_server.py keeps its trigger; the three router modules join the mypy file list; the serve guide states the constructor rule of serving classes. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- jenkins/L0_MergeRequest.groovy | 4 ++++ pyproject.toml | 3 +++ tensorrt_llm/serve/AGENTS.md | 6 ++++++ 3 files changed, 13 insertions(+) diff --git a/jenkins/L0_MergeRequest.groovy b/jenkins/L0_MergeRequest.groovy index 541ec1cb7626..e7a604c1851b 100644 --- a/jenkins/L0_MergeRequest.groovy +++ b/jenkins/L0_MergeRequest.groovy @@ -1950,6 +1950,10 @@ def getMultiGpuFileChanged(pipeline, testFilter, globalVars) "tensorrt_llm/serve/__init__.py", // _MsgspecRoute decodes the disaggregated frontend's msgpack bodies. "tensorrt_llm/serve/_base/", + // Code moved out of tensorrt_llm/serve/openai_server.py keeps its trigger. + "tensorrt_llm/serve/api/openai/embeddings/", + "tensorrt_llm/serve/api/openai/mm_encoder/", + "tensorrt_llm/serve/api/openai/tokenize/", // /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 81d2e449c2b8..306dccf23275 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -881,6 +881,9 @@ files = [ "tensorrt_llm/_torch/pyexecutor/sampler/ops/vanilla.py", "tensorrt_llm/llmapi/engine_client.py", "tensorrt_llm/serve/_base/routing.py", + "tensorrt_llm/serve/api/openai/embeddings/api_router.py", + "tensorrt_llm/serve/api/openai/mm_encoder/api_router.py", + "tensorrt_llm/serve/api/openai/tokenize/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 7272155ec6fa..0d5e18d877b8 100644 --- a/tensorrt_llm/serve/AGENTS.md +++ b/tensorrt_llm/serve/AGENTS.md @@ -20,6 +20,12 @@ 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. +- Serving classes take every dependency as a keyword-only constructor argument without a + default; the composition root builds them and puts them on `app.state`. Never read an + optional attribute with `getattr(self, ...)` in new serving code: a missing dependency + must fail at construction, not fall back silently at request time. Check: the + constructor tests in `tests/unittest/serve/api/openai/embeddings/test_serving.py` and its + siblings. - 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`. From 64538b36adb020ff02743d201d644cfbea2b271a Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Tue, 3 Nov 2026 14:51:52 +0000 Subject: [PATCH 29/30] [TRTLLM-17267][test] Let the OpenEngine Prometheus test's constructor stub provide the tokenizer The test runs the real OpenAIServer constructor with _init_llm replaced by a stub. The constructor now builds the tokenize serving object for the LLM role from self.tokenizer, so the stub sets it too; a missing dependency fails at construction on purpose instead of at request time. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- tests/unittest/grpc/openengine/test_metrics.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/tests/unittest/grpc/openengine/test_metrics.py b/tests/unittest/grpc/openengine/test_metrics.py index 5e52a7d6be15..d7faf31cca05 100644 --- a/tests/unittest/grpc/openengine/test_metrics.py +++ b/tests/unittest/grpc/openengine/test_metrics.py @@ -608,9 +608,13 @@ def _http_enablement() -> None: with pytest.MonkeyPatch.context() as patch: # Exercise the real constructor and route registration without loading a model. - patch.setattr( - OpenAIServer, "_init_llm", lambda self, _: setattr(self, "use_harmony", False) - ) + # The constructor reads these after _init_llm: use_harmony for the chat routes + # and tokenizer for the tokenize serving object it builds for the LLM role. + def init_llm(self, _): + self.use_harmony = False + self.tokenizer = None + + patch.setattr(OpenAIServer, "_init_llm", init_llm) for requested, output_dir in [(False, None), (True, None), (False, "unused-jsonl")]: args = _args() args.return_perf_metrics = requested From 600bcd503bf3aba087cb194cf101bab697fc71c5 Mon Sep 17 00:00:00 2001 From: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> Date: Tue, 3 Nov 2026 16:00:11 +0000 Subject: [PATCH 30/30] [TRTLLM-17267][test] Let the multi-frontend tests re-enact the composition root with the tokenize serving object Both tests build a partial OpenAIServer and register its routes the way the constructor does. The constructor now puts the tokenize serving object on app.state and attaches the role's routers after the local table, so the re-enactments do the same; the tokenize router refuses to attach without its serving object. Signed-off-by: JunyiXu-nv <219237550+JunyiXu-nv@users.noreply.github.com> --- .../llmapi/test_serve_multi_frontend_helpers.py | 16 ++++++++++++---- .../test_serve_multi_frontend_route_gate.py | 15 ++++++++++++--- 2 files changed, 24 insertions(+), 7 deletions(-) diff --git a/tests/unittest/llmapi/test_serve_multi_frontend_helpers.py b/tests/unittest/llmapi/test_serve_multi_frontend_helpers.py index 5a0a88d372eb..94d9b55cadcd 100644 --- a/tests/unittest/llmapi/test_serve_multi_frontend_helpers.py +++ b/tests/unittest/llmapi/test_serve_multi_frontend_helpers.py @@ -37,6 +37,8 @@ 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.openai.tokenize import api_router as tokenize_router +from tensorrt_llm.serve.api.openai.tokenize.serving import TokenizeServing from tensorrt_llm.serve.api.ops.system.serving import SystemServing from tensorrt_llm.serve.multi_frontend import ( FORWARDED_ROUTES, @@ -113,13 +115,19 @@ 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, then the local table, then the routers of the role (the system + # router with /health, /version and /v1/models, and the tokenize router), + # which read their serving objects from app.state. server.app.state.ops_serving_system = SystemServing(model="m", health_probe=lambda: True) - attach_routers(server.app, ServerCapabilities.for_role(None)) + setattr( + server.app.state, + tokenize_router.APP_STATE_KEY, + TokenizeServing(model="m", tokenizer=None), + ) server._register_launcher_forwarding_routes() server.register_routes() + attach_routers(server.app, ServerCapabilities.for_role(None)) 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 b2cd8c1841a4..d3df07c17277 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,8 @@ from test_common.serve.routes import effective_routes from tensorrt_llm.serve._base.routing import ServerCapabilities, attach_routers +from tensorrt_llm.serve.api.openai.tokenize import api_router as tokenize_router +from tensorrt_llm.serve.api.openai.tokenize.serving import TokenizeServing 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,11 +78,18 @@ 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 serving objects the routers read go on + # app.state, the role's own table is registered, then the routers of the role + # (the system router with /health, /version and /v1/models, and the tokenize + # router) are attached. server.app.state.ops_serving_system = SystemServing(model="m", health_probe=lambda: True) - attach_routers(server.app, ServerCapabilities.for_role(None)) + setattr( + server.app.state, + tokenize_router.APP_STATE_KEY, + TokenizeServing(model="m", tokenizer=None), + ) server.register_routes() + attach_routers(server.app, ServerCapabilities.for_role(None)) # effective_routes walks included routers, which FastAPI 0.139 keeps as # objects without .path or .methods in app.routes. return sorted(