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 diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index c2166e7bb04b..ffcb5862ea11 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -417,13 +417,30 @@ common-files: &common_files | tensorrt_llm/scaffolding/worker.py | tensorrt_llm/scheduling_params.py | tensorrt_llm/serve/__init__.py | + tensorrt_llm/serve/_base/msgspec_route.py | + tensorrt_llm/serve/_base/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 | 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 | @@ -1144,13 +1161,30 @@ legacy-files: &legacy_files | tensorrt_llm/scaffolding/worker.py | tensorrt_llm/scheduling_params.py | tensorrt_llm/serve/__init__.py | + tensorrt_llm/serve/_base/msgspec_route.py | + tensorrt_llm/serve/_base/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 | 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 | @@ -1620,6 +1654,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..e7a604c1851b 100644 --- a/jenkins/L0_MergeRequest.groovy +++ b/jenkins/L0_MergeRequest.groovy @@ -1948,6 +1948,14 @@ 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/", + // 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", "tensorrt_llm/serve/disagg_auto_scaling.py", "tensorrt_llm/serve/extensions/", diff --git a/legacy-files.txt b/legacy-files.txt index 9850e7c810b4..6c098106a80d 100644 --- a/legacy-files.txt +++ b/legacy-files.txt @@ -409,13 +409,30 @@ tensorrt_llm/scaffolding/task_collection.py tensorrt_llm/scaffolding/worker.py tensorrt_llm/scheduling_params.py tensorrt_llm/serve/__init__.py +tensorrt_llm/serve/_base/msgspec_route.py +tensorrt_llm/serve/_base/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 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 60e243d984b0..306dccf23275 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -474,13 +474,30 @@ exclude = [ "tensorrt_llm/scaffolding/worker.py", "tensorrt_llm/scheduling_params.py", "tensorrt_llm/serve/__init__.py", + "tensorrt_llm/serve/_base/msgspec_route.py", + "tensorrt_llm/serve/_base/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", "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", @@ -862,6 +879,13 @@ files = [ "tensorrt_llm/_torch/pyexecutor/sampler/token_ban.py", "tensorrt_llm/_torch/pyexecutor/sampler/ops/flashinfer.py", "tensorrt_llm/_torch/pyexecutor/sampler/ops/vanilla.py", + "tensorrt_llm/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", "tests/unittest/_torch/sampler/test_token_ban.py", "tests/unittest/_torch/sampler/test_beam_search.py", diff --git a/ruff-legacy.toml b/ruff-legacy.toml index dab55b4d0434..5752751aa372 100644 --- a/ruff-legacy.toml +++ b/ruff-legacy.toml @@ -426,13 +426,30 @@ include = [ "tensorrt_llm/scaffolding/worker.py", "tensorrt_llm/scheduling_params.py", "tensorrt_llm/serve/__init__.py", + "tensorrt_llm/serve/_base/msgspec_route.py", + "tensorrt_llm/serve/_base/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", "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/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index ce43e6718595..b110e290d467 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 @@ -2749,6 +2749,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/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/commands/serve.py b/tensorrt_llm/commands/serve.py index e8234ac09a2c..41934485bb22 100644 --- a/tensorrt_llm/commands/serve.py +++ b/tensorrt_llm/commands/serve.py @@ -553,8 +553,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 " @@ -2659,7 +2659,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 0519899987a3..60b1566e8579 100644 --- a/tensorrt_llm/executor/executor.py +++ b/tensorrt_llm/executor/executor.py @@ -340,6 +340,20 @@ 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 + + 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).""" + return self._fatal_error is not None + def check_health(self) -> bool: """Check whether the executor is healthy and able to process requests. @@ -408,6 +422,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 3e0f3f919ecc..172e75e3e655 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. @@ -671,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. @@ -1684,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 47278788dc33..9df67ef72abc 100644 --- a/tensorrt_llm/grpc/openengine/control.py +++ b/tensorrt_llm/grpc/openengine/control.py @@ -131,32 +131,26 @@ 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 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 --------------------------------------------------- @@ -222,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 @@ -660,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 6de981e35f41..7c3b0f0b9fc8 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 525ca2132802..fdba6d5209a0 100644 --- a/tensorrt_llm/grpc/openengine/server.py +++ b/tensorrt_llm/grpc/openengine/server.py @@ -197,7 +197,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; " @@ -225,11 +225,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 @@ -422,11 +418,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(): @@ -438,7 +430,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/openengine/streaming.py b/tensorrt_llm/grpc/openengine/streaming.py index 4c6944905b4c..f5a0b5de4ef9 100644 --- a/tensorrt_llm/grpc/openengine/streaming.py +++ b/tensorrt_llm/grpc/openengine/streaming.py @@ -36,7 +36,7 @@ class _OpenEngineFormatState: is_context_only: bool = False kv_transfer_backend: str = "" decode_preparation: str | None = None - # 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/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/inputs/registry.py b/tensorrt_llm/inputs/registry.py index 114b8b709cbf..ce89ed54426e 100644 --- a/tensorrt_llm/inputs/registry.py +++ b/tensorrt_llm/inputs/registry.py @@ -1153,6 +1153,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 efec9c945337..7d9b5fd7101a 100644 --- a/tensorrt_llm/llmapi/llm.py +++ b/tensorrt_llm/llmapi/llm.py @@ -42,8 +42,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) @@ -146,6 +147,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. @@ -340,6 +355,15 @@ 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). + 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. """ @@ -1242,10 +1266,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( @@ -1701,8 +1725,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]: @@ -1798,11 +1821,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") @@ -1812,6 +1844,164 @@ 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 + + @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/AGENTS.md b/tensorrt_llm/serve/AGENTS.md new file mode 100644 index 000000000000..0d5e18d877b8 --- /dev/null +++ b/tensorrt_llm/serve/AGENTS.md @@ -0,0 +1,105 @@ + + +# AGENTS.md - tensorrt_llm/serve + +Scope: `tensorrt_llm/serve/` and everything under it. Supplements the repo-root `AGENTS.md` +and `tests/AGENTS.md`; `tensorrt_llm/commands/AGENTS.md` points here for the serve CLI. +Where a check enforces a rule, the rule names it; a failing check says what to change. +Name files and symbols here, never line numbers. + +## Where code goes + +- New endpoint: a router and a serving class under `serve/api///` + (`api_router.py` + `serving.py`). Create the router with `make_router()` + (`serve/_base/routing.py`) and add a `RouterSpec` to `ROUTERS`. Never add a method or a + route to `OpenAIServer`. Check: `tests/unittest/api_stability/test_serve_routes.py`. +- Attach a router with `app.include_router(router)` and no `prefix=`; give the prefix to + `make_router(prefix=...)` instead. The composition root sets each serving object on + `app.state._serving_` before attaching; routers read it from + `request.app.state`, never from a closure. +- 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`. +- 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` 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 + regenerates it and needs the API review committee. + Check: `tests/unittest/api_stability/test_serve_routes.py`. +- Middleware order matters: `add_middleware` inserts at the front, so the last one added is + outermost. Do not reorder. Check: `tests/unittest/serve/test_middleware_order.py`. +- Import directions between layers are tests; a PR that creates a package adds its contract. + Check: `tests/unittest/serve/test_layering.py`. +- VisualGen serve tests live under a path that contains `visual_gen/` + (`tests/unittest/serve/api/visual_gen/`), or CI selection misses them. + +## Moving code + +- A move PR only moves code: no reformatting, no logic change. +- A move PR is redone on the latest `main` right before merge instead of being rebased, + and its PR body lists every moved definition with its old and new location, and every + definition the PR changes on purpose. +- In the same PR update the multi-GPU file list (`getMultiGpuFileChanged` in + `jenkins/L0_MergeRequest.groovy`), `.github/CODEOWNERS` and `legacy-files.txt`. A new file + takes the formatting group of the file its code comes from. +- Old import paths: a move deletes the old module (or the moved names) without a shim and + rewrites every import of it in the repository, tests included; module paths written as + strings (patch targets, subprocess command lines, `_LAZY_ATTRS` in `serve/__init__.py`) are + fixed by hand. The PR body lists "old path → new path" for the release notes. Only five old + paths keep a shim, because code outside the repository imports them: + `tensorrt_llm/llmapi/reasoning_parser.py`, `serve/tool_parser/`, + `serve/openai_disagg_server.py`, `serve/openai_server.py` and `commands/serve.py`. A shim + re-exports public names only, lists them in `__all__` (or `_LAZY_NAMES` in + `commands/serve.py`), warns once with `FutureWarning`, and has its rows in + `tests/unittest/serve/test_shim_identity.py`. +- `mock.patch("old.module.name")` does not follow moved code. Grep the tests for every moved + name, including `monkeypatch.setattr(module, "name", ...)` and dotted paths split across lines. + +## Tests + +- Register every new test file in a CI list in the same PR (`tests/AGENTS.md`). Serve tests + that need no GPU go in `tests/integration/test_lists/test-db/l0_cpu.yml` and must contain + `pytest.mark.cpu_only`, or the CPU stage collects nothing from them. +- New serve tests mirror the package: `tests/unittest/serve/` for `tensorrt_llm/serve/`. +- Build servers with the factory and fake generators in `tests/test_common/serve/` + (`isolated_serve_env`), never `object.__new__(OpenAIServer)` or + `object.__new__(OpenAIDisaggServer)`. Check: `tests/unittest/serve/test_no_object_new.py`. +- List an app's routes with `effective_routes()` from `tests/test_common/serve/routes.py`, not + `app.routes`: FastAPI 0.139 keeps included routers as objects without `.path`. +- The harmony tests guard whole files with try-imports. After renaming anything they import, + run them with `pytest -rs` and check that nothing is skipped as "not importable": a green run + can mean every test was skipped. + +## Checklists + +- Endpoint: package under `serve/api/`; `RouterSpec` in `ROUTERS`; regenerate the route + snapshot (API review); multi-GPU list entry if it runs on the disaggregated path; CPU test + registered in `tests/integration/test_lists/test-db/l0_cpu.yml`. + +## Commands + +- All serve CPU tests (the runner exits when `LLM_MODELS_ROOT` is unset): + `LLM_MODELS_ROOT=/path/to/models PYTHONPATH=tests python3 -m test_common.serve.cpu_tests` +- Route snapshot: `pytest -m cpu_only tests/unittest/api_stability/test_serve_routes.py tests/unittest/serve/api/visual_gen/test_routes_visual_gen.py` +- Regenerate the route snapshot: `PYTHONPATH=tests python3 -m test_common.serve.routes --write` +- API gates: `pytest tests/unittest/api_stability/test_serve_api.py tests/unittest/api_stability/test_serve_cli.py` +- Extensions and parsers: `pytest tests/unittest/llmapi/apps/test_serving_extensions.py tests/unittest/llmapi/apps/test_tool_parsers.py` +- Disaggregated router and coordinator (CPU): `pytest tests/unittest/disaggregated -m cpu_only -q` +- Which CI stage runs a test: `python scripts/test_to_stage_mapping.py --tests ""`. It + matches the text of list entries and does not expand directories: a file listed through its + directory maps only when you pass the directory entry, e.g. `--tests "unittest/serve"` for + anything under `tests/unittest/serve/`. +- Tests that load model files need `LLM_MODELS_ROOT` (the harmony tests' autouse fixture does). diff --git a/tensorrt_llm/serve/CLAUDE.md b/tensorrt_llm/serve/CLAUDE.md new file mode 100644 index 000000000000..40c6cf85d526 --- /dev/null +++ b/tensorrt_llm/serve/CLAUDE.md @@ -0,0 +1,2 @@ +# In ./CLAUDE.md +@AGENTS.md diff --git a/tensorrt_llm/serve/_base/__init__.py b/tensorrt_llm/serve/_base/__init__.py new file mode 100644 index 000000000000..0256eaf43af3 --- /dev/null +++ b/tensorrt_llm/serve/_base/__init__.py @@ -0,0 +1,15 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Shared serving infrastructure with no dependency on endpoint packages.""" diff --git a/tensorrt_llm/serve/_base/msgspec_route.py b/tensorrt_llm/serve/_base/msgspec_route.py new file mode 100644 index 000000000000..a7eaf33acc64 --- /dev/null +++ b/tensorrt_llm/serve/_base/msgspec_route.py @@ -0,0 +1,36 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Request and route classes that decode msgpack request bodies from the +disaggregated frontend.""" + +import json + +import msgspec + +from tensorrt_llm.serve.kv_hints import KvHintsRequest, KvHintsRoute + +_msgpack_decoder = msgspec.msgpack.Decoder() + + +class _MsgspecRequest(KvHintsRequest): + """Request that decodes X-TRTLLM-Msgpack bodies with msgspec.""" + + async def json(self): + if not hasattr(self, "_json_body"): + body = await self.body() + if not body: + self._json_body = {} + elif self.headers.get("x-trtllm-msgpack") == "1": + try: + self._json_body = _msgpack_decoder.decode(body) + except msgspec.DecodeError as e: + raise json.JSONDecodeError(f"msgpack: {e}", "", 0) from e + else: + self._json_body = json.loads(body) + return self.filter_kv_hints(self._json_body) + + +class _MsgspecRoute(KvHintsRoute): + """Parse msgpack/JSON and apply hint trust policy before model validation.""" + + request_type = _MsgspecRequest diff --git a/tensorrt_llm/serve/_base/routing.py b/tensorrt_llm/serve/_base/routing.py new file mode 100644 index 000000000000..70facf89f0ef --- /dev/null +++ b/tensorrt_llm/serve/_base/routing.py @@ -0,0 +1,114 @@ +# 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), + 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}) + ), +) + + +def make_router(**kwargs: Any) -> APIRouter: + """The only way to create a trtllm-serve router: every route decodes msgpack bodies. + + The disaggregated frontend sends worker requests as msgpack; a route of a plain + APIRouter keeps APIRoute and rejects them. + """ + if "route_class" in kwargs: + raise TypeError( + "make_router() always uses _MsgspecRoute; do not pass route_class. " + "A plain APIRoute breaks msgpack request bodies from the disaggregated frontend." + ) + return APIRouter(route_class=_MsgspecRoute, **kwargs) + + +def attach_routers( + app: FastAPI, caps: ServerCapabilities, routers: tuple[RouterSpec, ...] = ROUTERS +) -> None: + """Import and attach only the routers this role needs.""" + for spec in routers: + if caps.role_group in spec.roles and spec.when(caps): + importlib.import_module(spec.module).attach_router(app) diff --git a/tensorrt_llm/serve/_base/serving.py b/tensorrt_llm/serve/_base/serving.py new file mode 100644 index 000000000000..6feade1c0ed8 --- /dev/null +++ b/tensorrt_llm/serve/_base/serving.py @@ -0,0 +1,113 @@ +# 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.""" + +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: + """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. + """ + + @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.""" + + 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 + + @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/_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/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/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/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/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 new file mode 100644 index 000000000000..1c281a8ee935 --- /dev/null +++ b/tensorrt_llm/serve/api/openai/embeddings/serving.py @@ -0,0 +1,166 @@ +# 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``.""" + +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, + InputTooLongError, + QueueFullError) +from tensorrt_llm.serve.openai_protocol import (EmbeddingRequest, + EmbeddingResponse, + EmbeddingResponseData, + EmbeddingUsageInfo) + + +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() + + 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/__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/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/mm_encoder/serving.py b/tensorrt_llm/serve/api/openai/mm_encoder/serving.py new file mode 100644 index 000000000000..540cdd0adc74 --- /dev/null +++ b/tensorrt_llm/serve/api/openai/mm_encoder/serving.py @@ -0,0 +1,159 @@ +# 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.""" + +import asyncio +import signal +import traceback +from http import HTTPStatus +from typing import Any, List + +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): + """``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 + + 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/__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/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/api/openai/tokenize/serving.py b/tensorrt_llm/serve/api/openai/tokenize/serving.py new file mode 100644 index 000000000000..05ba379b1835 --- /dev/null +++ b/tensorrt_llm/serve/api/openai/tokenize/serving.py @@ -0,0 +1,30 @@ +# 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 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): + """``POST /_internal/tokenize``: token ids of a plain-text prompt.""" + + 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/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/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/chat_utils.py b/tensorrt_llm/serve/chat_utils.py index 19052ddd3f93..0ba7e8695bf4 100644 --- a/tensorrt_llm/serve/chat_utils.py +++ b/tensorrt_llm/serve/chat_utils.py @@ -165,7 +165,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=_load_request_media(image_io, str_content), @@ -184,7 +184,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=_load_request_media(video_io, str_content), @@ -193,7 +193,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=_load_request_media(audio_io, str_content), @@ -211,7 +211,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 d8ebec1a6ce7..c1830849f910 100644 --- a/tensorrt_llm/serve/openai_disagg_server.py +++ b/tensorrt_llm/serve/openai_disagg_server.py @@ -332,6 +332,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 deleted file mode 100644 index e2643c43d63c..000000000000 --- a/tensorrt_llm/serve/openai_protocol.py +++ /dev/null @@ -1,2657 +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 math -import time -import uuid -from typing import Any, Dict, List, Literal, Optional, Union - -import torch -import xgrammar -from fastapi import UploadFile -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 (AliasChoices, 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 -from tensorrt_llm.visual_gen.params import MediaRole - -_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) - - -# ============================================================================ -# 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] - -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/__init__.py b/tensorrt_llm/serve/openai_protocol/__init__.py new file mode 100644 index 000000000000..2411f8097c73 --- /dev/null +++ b/tensorrt_llm/serve/openai_protocol/__init__.py @@ -0,0 +1,331 @@ +# 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: 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", + "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..f15802ce4753 --- /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" diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py index eb79cdb18397..667a3142f70e 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 @@ -25,7 +24,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 @@ -43,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, @@ -67,6 +63,10 @@ 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._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, @@ -81,6 +81,12 @@ 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, + 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) @@ -90,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, @@ -101,16 +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, - 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, + 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, @@ -149,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, @@ -157,7 +155,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 +176,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 @@ -291,31 +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: - from tensorrt_llm.executor.proxy import GenerationExecutorProxy - - if isinstance(getattr(generator, "_executor", None), - GenerationExecutorProxy): - 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. @@ -779,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): @@ -907,6 +854,12 @@ async def validation_exception_handler(request, exc): len(detail) - 4096) return JSONResponse(status_code=400, content={"error": detail}) + # /health, /version and /v1/models for every role. The health + # probe stays here, with its private accesses, until the public + # health API exists. + self.app.state.ops_serving_system = SystemServing( + model=self.model, health_probe=self._health_probe) + if self.server_role is ServerRole.VISUAL_GEN: assert self._is_visual_gen, \ "generator must be a VisualGen for VISUAL_GEN server" @@ -915,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, @@ -961,11 +948,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) @@ -1093,41 +1076,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: @@ -1143,22 +1095,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( @@ -1211,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 @@ -1218,17 +1162,16 @@ 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"]) 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"]) @@ -1250,7 +1193,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( @@ -1326,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", @@ -1373,16 +1313,10 @@ 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, 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() @@ -1450,147 +1384,8 @@ 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) - - # 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 - self.embedding_batcher = EncodeBatcher( - encode_fn, - max_batch_size=engine.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, - ) - - 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"]) - - @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.""" - # 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"]) @@ -1699,17 +1494,17 @@ 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) - async def health(self) -> Response: + def _health_probe(self) -> bool: if self._check_health(): - return Response(status_code=200) + return True else: # If the engine has a fatal error, trigger server shutdown so the # pod doesn't linger as a zombie that accepts connections but never @@ -1718,20 +1513,15 @@ 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) - return Response( - status_code=503, - content= - "LLM is unavailable. Please check the server logs for more details." - ) + 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 False async def health_generate(self, raw_request: Request) -> Response: """Health check that performs a minimal generation.""" @@ -1782,14 +1572,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 @@ -2825,122 +2607,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: @@ -3630,17 +3296,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/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/postprocess_handlers.py b/tensorrt_llm/serve/postprocess_handlers.py index 221e4139e582..9ac7eb364e45 100644 --- a/tensorrt_llm/serve/postprocess_handlers.py +++ b/tensorrt_llm/serve/postprocess_handlers.py @@ -165,11 +165,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 @@ -605,8 +604,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: @@ -774,8 +772,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): @@ -905,14 +902,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], @@ -976,8 +972,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/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/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index 41e0cabd17ea..411768e89494 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 @@ -238,6 +239,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 @@ -301,6 +304,7 @@ l0_cpu: - unittest/llmapi/test_kv_hints.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. @@ -341,6 +345,15 @@ 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/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/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..a512c0d5dcd6 --- /dev/null +++ b/tests/test_common/serve/app_factory.py @@ -0,0 +1,370 @@ +# 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 + 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. + multi_frontend_attached: bool = False + + +@dataclass(frozen=True) +class DisaggConfig: + """One ``OpenAIDisaggServer`` configuration.""" + + fleet: bool = False # delegate to a remote coordinator (CoordinatorClient) + http_cluster_storage: bool = False + return_perf_metrics: bool = False + + +@dataclass(frozen=True) +class CoordinatorConfig: + """One ``CoordinatorServer`` configuration (the fleet-mode coordinator app).""" + + http_cluster_storage: bool = False + + +def _make_disagg_server_config( + *, return_perf_metrics: bool = False, http_cluster_storage: bool = False +) -> DisaggServerConfig: + # Loopback addresses on closed ports: construction is network-free, and + # nothing in these tests runs the disaggregated lifespan. + return DisaggServerConfig( + server_configs=[ + CtxGenServerConfig(type="ctx", hostname="127.0.0.1", port=9), + CtxGenServerConfig(type="gen", hostname="127.0.0.1", port=10), + ], + hostname="127.0.0.1", + port=0, + ctx_router_config=RouterConfig(server_role=ServerRole.CONTEXT), + gen_router_config=RouterConfig(server_role=ServerRole.GENERATION), + node_id=0, + return_perf_metrics=return_perf_metrics, + disagg_cluster_config=( + DisaggClusterConfig(cluster_uri="http://127.0.0.1:9", cluster_name="test") + if http_cluster_storage + else None + ), + ) + + +class _StubTelemetrySession: + """Collects termination observations for ``ServeTestEnv.record_terminations``.""" + + def __init__(self, sink: list) -> None: + self.owner_pid = os.getpid() + self._sink = sink + + def record_termination_observation(self, outcome, lifecycle_phase=None) -> None: + self._sink.append(outcome) + + +class ServeTestEnv: + """Handle returned by ``isolated_serve_env``; builds servers and clients.""" + + def __init__(self, tmp_dir: Path) -> None: + self.tmp_dir = Path(tmp_dir) + self._servers: List[OpenAIServer] = [] + self._entered_clients: List[TestClient] = [] + self._plain_clients: List[TestClient] = [] + + # -- builders --------------------------------------------------------- + + def _write_fake_checkpoint(self, model_type: str) -> Path: + """A directory with only a config.json, so load_pretrained_config works. + + No tokenizer files are written: they would make AutoProcessor succeed + and change ``server.processor`` from its fake-generator value (None). + """ + model_dir = self.tmp_dir / f"model-{model_type}" + model_dir.mkdir(parents=True, exist_ok=True) + architectures = { + "llama": ["LlamaForCausalLM"], + "gpt_oss": ["GptOssForCausalLM"], + }[model_type] + (model_dir / "config.json").write_text( + json.dumps({"model_type": model_type, "architectures": architectures}) + ) + return model_dir + + def build_openai_server(self, config: AggConfig) -> OpenAIServer: + from test_common.serve.fakes import FakeExecutor # local: keeps the import graph flat + + for name, value in config.env: + os.environ[name] = value + + args = FakeLlmArgs( + model=str(self._write_fake_checkpoint(config.model_type)), + return_perf_metrics=config.return_perf_metrics, + perf_metrics_output_dir=( + str(self.tmp_dir / "perf") if config.perf_metrics_output_dir else None + ), + sleep_config=SimpleNamespace() if config.runtime_control else None, + encode_only=config.role is ServerRole.EMBEDDING, + reasoning_parser=config.reasoning_parser, + num_postprocess_workers=config.num_postprocess_workers, + ) + executor = FakeExecutor(resource_governor_queue=object() if config.governor_queue else None) + if config.role is ServerRole.VISUAL_GEN: + # Lazy import: only VisualGen tests pay the visual_gen import cost, + # and importing it is what lets _is_visual_gen_instance say yes. + from test_common.serve.fakes_visual_gen import FakeVisualGen + + generator = FakeVisualGen(model=config.model_name) + elif config.role is ServerRole.MM_ENCODER: + generator = FakeMultimodalEncoder(args=args, executor=executor) + elif config.role is ServerRole.EMBEDDING: + generator = FakeEncodeOnlyLLM(args=args) + elif config.rl_control: + generator = FakeAsyncLLM(args=args, executor=executor) + else: + generator = FakeLLM(args=args, executor=executor) + + multi_frontend_serving = None + if config.multi_frontend_attached: + from tensorrt_llm.serve.multi_frontend import MultiFrontendServing + + ipc_dir = self.tmp_dir / "multi_frontend_ipc" + ipc_dir.mkdir(parents=True, exist_ok=True) + multi_frontend_serving = MultiFrontendServing( + launcher_uds=str(ipc_dir / "launcher.sock"), is_launcher=False + ) + + server = OpenAIServer( + generator=generator, + model=config.model_name, + tool_parser=config.tool_parser, + server_role=config.role, + metadata_server_cfg=None, + 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, + runtime_control_api_key="test-key" if config.runtime_control else None, + multi_frontend_serving=multi_frontend_serving, + ) + self._servers.append(server) + return server + + def build_disagg_server(self, config: DisaggConfig) -> OpenAIDisaggServer: + # At most one per environment: each registers the disaggregated server + # counters in the default Prometheus registry, and a second one fails + # with "Duplicated timeseries". + server_config = _make_disagg_server_config( + return_perf_metrics=config.return_perf_metrics, + http_cluster_storage=config.http_cluster_storage, + ) + return OpenAIDisaggServer( + server_config, + coordinator_url="http://127.0.0.1:9" if config.fleet else None, + ) + + def build_coordinator_server(self, config: CoordinatorConfig) -> CoordinatorServer: + server_config = _make_disagg_server_config(http_cluster_storage=config.http_cluster_storage) + coordinator = DisaggCoordinatorService(server_config, lambda *args, **kwargs: None) + return CoordinatorServer(coordinator) + + # -- request plumbing ------------------------------------------------- + + def client(self, app: FastAPI, *, lifespan: bool = True) -> TestClient: + """A TestClient closed at environment exit; runs the lifespan by default.""" + client = TestClient(app) + if lifespan: + client.__enter__() + self._entered_clients.append(client) + else: + self._plain_clients.append(client) + return client + + def record_terminations(self) -> list: + """Route termination observations of this process into the returned list. + + Installs a stub telemetry session; the environment's exit restores the + session that existed before entry. + """ + sink: list = [] + usage_lib._SESSION = _StubTelemetrySession(sink) + return sink + + # -- teardown helpers (called by isolated_serve_env) ------------------- + + def _close(self) -> None: + for client in self._entered_clients: + client.__exit__(None, None, None) + for client in self._plain_clients: + client.close() + self._entered_clients.clear() + self._plain_clients.clear() + for server in self._servers: + for pool_name in ( + "_input_proc_executor", + "_media_load_executor", + "_metrics_stats_executor", + ): + pool = getattr(server, pool_name, None) + if pool is not None: + pool.shutdown(wait=False, cancel_futures=True) + self._servers.clear() + + +@contextmanager +def isolated_serve_env(tmp_dir: Path) -> Iterator[ServeTestEnv]: + """Process-global state that building trtllm-serve apps touches, restored on exit. + + The only supported forms are ``with isolated_serve_env(tmp_path) as env:`` + and the ``serve_env`` fixture. The context-manager object has no builder + methods, so ``isolated_serve_env(tmp_path).build_openai_server(...)`` fails + with ``AttributeError`` instead of leaving global state behind. + + Build one configuration per environment when configurations differ in + environment variables or perf metrics. + """ + tmp_dir = Path(tmp_dir) + env_patcher = mock.patch.dict(os.environ) + env_patcher.start() + try: + os.environ["HF_HUB_OFFLINE"] = "1" + os.environ["TRANSFORMERS_OFFLINE"] = "1" + prometheus_dir = tmp_dir / "prometheus" + prometheus_dir.mkdir(parents=True, exist_ok=True) + os.environ["PROMETHEUS_MULTIPROC_DIR"] = str(prometheus_dir) + os.environ["TRTLLM_MEDIA_STORAGE_PATH"] = str(tmp_dir / "media") + for name in _CLEARED_ENV_VARS: + os.environ.pop(name, None) + + saved_media_executor = BaseMediaIO._executor + # The module global exists only after the first + # set_prometheus_multiproc_dir() call; save it with a sentinel. + saved_prometheus_dir = getattr(tensorrt_llm._utils, "prometheus_multiproc_dir", _MISSING) + # An unhealthy /health records a termination observation on the + # process-wide telemetry session, which keeps only its first + # observation; detach any session so tests cannot poison it. + saved_session = usage_lib._SESSION + usage_lib._SESSION = None + collectors_before = set(REGISTRY._collector_to_names.keys()) + + env = ServeTestEnv(tmp_dir) + try: + yield env + finally: + env._close() + for collector in list(REGISTRY._collector_to_names.keys()): + if collector not in collectors_before: + with contextlib.suppress(KeyError): + REGISTRY.unregister(collector) + BaseMediaIO._executor = saved_media_executor + if saved_prometheus_dir is _MISSING: + if hasattr(tensorrt_llm._utils, "prometheus_multiproc_dir"): + delattr(tensorrt_llm._utils, "prometheus_multiproc_dir") + else: + tensorrt_llm._utils.prometheus_multiproc_dir = saved_prometheus_dir + usage_lib._SESSION = saved_session + finally: + env_patcher.stop() diff --git a/tests/test_common/serve/cpu_tests.py b/tests/test_common/serve/cpu_tests.py new file mode 100644 index 000000000000..44575b7c24f1 --- /dev/null +++ b/tests/test_common/serve/cpu_tests.py @@ -0,0 +1,183 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""The one command that runs every serve CPU test, as the CPU CI stage runs them. + + LLM_MODELS_ROOT=... PYTHONPATH=tests python3 -m test_common.serve.cpu_tests [--list] [pytest args...] + +Reads the CPU test-db list, keeps the entries whose files are serve CPU tests, +and runs one pytest process per kept entry from ``tests/`` with +``tests/unittest`` on ``PYTHONPATH`` — one process per entry, as the CPU stage +does — so process-global state does not leak between entries. Extra arguments go +to every process. Exits with the first non-zero pytest exit code, or 0. +""" + +import argparse +import os +import subprocess +import sys +from dataclasses import dataclass +from pathlib import Path +from typing import List, Optional, Sequence, Tuple + +import yaml + +__all__ = ["Entry", "select_entries"] + +REPO_ROOT = Path(__file__).resolve().parents[3] +CPU_LIST = REPO_ROOT / "tests" / "integration" / "test_lists" / "test-db" / "l0_cpu.yml" + +_MODULE_MARKERS = ( + "tensorrt_llm.serve", + "tensorrt_llm.commands.serve", + "from tensorrt_llm.commands import serve", + "test_common.serve", +) + + +@dataclass(frozen=True) +class Entry: + """One kept test-db entry with the targets it runs.""" + + entry: str # the test-db line, e.g. "unittest/serve" + targets: Tuple[str, ...] # pytest targets relative to tests/ + + +def _iter_test_entries(data) -> List[str]: + """Every string test entry of a test-db document, in file order.""" + entries: List[str] = [] + + def _walk(node) -> None: + if isinstance(node, str): + entries.append(node) + elif isinstance(node, list): + for item in node: + _walk(item) + elif isinstance(node, dict): + for value in node.values(): + _walk(value) + + _walk(data) + return entries + + +def _is_serve_cpu_test(path: Path) -> bool: + """Whether one test file is a serve CPU test. + + The file must contain the literal ``pytest.mark.cpu_only`` (the CPU stage + collects nothing from it otherwise) and either live under + ``tests/unittest/serve/`` or mention a serve module or ``test_common.serve`` + in its text. The path condition catches files that name no module, such as + checks over the test tree itself. + """ + try: + text = path.read_text(errors="replace") + except OSError: + return False + if "pytest.mark.cpu_only" not in text: + return False + serve_dir = REPO_ROOT / "tests" / "unittest" / "serve" + if serve_dir in path.parents: + return True + return any(marker in text for marker in _MODULE_MARKERS) + + +def select_entries(cpu_list: Path = CPU_LIST) -> Tuple[List[Entry], List[str]]: + """(kept entries with targets, notes about skipped entries).""" + data = yaml.safe_load(cpu_list.read_text()) + tests_root = REPO_ROOT / "tests" + kept: List[Entry] = [] + notes: List[str] = [] + for raw_entry in _iter_test_entries(data): + entry = raw_entry.strip() + if not entry.startswith("unittest/"): + continue + target_part = entry.split("::")[0].split(" ")[0] + if entry != target_part and "::" not in entry: + # Entries with -k, -m, TIMEOUT or ISOLATION run with options this + # runner does not reproduce; run them through the CI stage instead. + notes.append(f"skipped (has options): {entry}") + continue + if "::" in entry and " " in entry: + notes.append(f"skipped (has options): {entry}") + continue + target = entry.split("::")[0].rstrip("/") + path = tests_root / target + if "::" in entry: + # A node-id entry runs only that node, never the whole file. + if path.is_file() and _is_serve_cpu_test(path): + kept.append(Entry(entry=raw_entry, targets=(entry,))) + continue + if path.is_file(): + if _is_serve_cpu_test(path): + kept.append(Entry(entry=raw_entry, targets=(target,))) + continue + if path.is_dir(): + selected = tuple( + str(candidate.relative_to(tests_root)) + for candidate in sorted(path.rglob("test_*.py")) + if _is_serve_cpu_test(candidate) + ) + if selected: + kept.append(Entry(entry=raw_entry, targets=selected)) + return kept, notes + + +def main(argv: Optional[Sequence[str]] = None) -> int: + parser = argparse.ArgumentParser( + prog="python3 -m test_common.serve.cpu_tests", + description="Run every serve CPU test, one pytest process per test-db entry.", + ) + parser.add_argument( + "--list", action="store_true", help="print the kept entries and their targets" + ) + args, pytest_args = parser.parse_known_args(argv) + + kept, notes = select_entries() + if args.list: + for note in notes: + print(note) + for entry in kept: + print(entry.entry) + for target in entry.targets: + print(f" {target}") + return 0 + + if not os.environ.get("LLM_MODELS_ROOT"): + print( + "LLM_MODELS_ROOT is not set. Some selected tests read model files through " + "autouse fixtures and would error in setup; set LLM_MODELS_ROOT and rerun.", + file=sys.stderr, + ) + return 2 + + env = dict(os.environ) + unittest_dir = str(REPO_ROOT / "tests" / "unittest") + env["PYTHONPATH"] = os.pathsep.join( + part for part in (unittest_dir, env.get("PYTHONPATH", "")) if part + ) + first_failure = 0 + for note in notes: + print(note) + for entry in kept: + command = [sys.executable, "-m", "pytest", "-m", "cpu_only", *entry.targets, *pytest_args] + print(f"[{entry.entry}] {' '.join(command)}") + result = subprocess.run(command, cwd=REPO_ROOT / "tests", env=env) + if result.returncode != 0 and first_failure == 0: + first_failure = result.returncode + return first_failure + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/test_common/serve/fakes.py b/tests/test_common/serve/fakes.py new file mode 100644 index 000000000000..c41bca32d041 --- /dev/null +++ b/tests/test_common/serve/fakes.py @@ -0,0 +1,507 @@ +# 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 + +import torch +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 + # What the result's ``disaggregated_params`` carries, for example the + # multimodal embedding handles the multimodal-encoder role answers with. + disaggregated_params: Any = None + + +@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 + self._disaggregated_params = script.disaggregated_params + 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. ``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: + 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: + """Pure ASGI no-op, used to check where ``--middleware`` lands in the stack.""" + + def __init__(self, app: Any) -> None: + self.app = app + + async def __call__(self, scope: Any, receive: Any, send: Any) -> None: + await self.app(scope, receive, send) diff --git a/tests/test_common/serve/fakes_visual_gen.py b/tests/test_common/serve/fakes_visual_gen.py new file mode 100644 index 000000000000..c1241fb6df77 --- /dev/null +++ b/tests/test_common/serve/fakes_visual_gen.py @@ -0,0 +1,44 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Fake VisualGen generator for CPU tests of the VISUAL_GEN serving role. + +Kept in its own module so that only tests that build a VisualGen app import the +``tensorrt_llm.visual_gen`` tree: ``OpenAIServer`` detects VisualGen generators +by probing ``sys.modules`` for ``tensorrt_llm.visual_gen.visual_gen``, so +importing this module is what makes ``_is_visual_gen_instance`` able to say yes. +""" + +import itertools +from typing import Any + +from tensorrt_llm.visual_gen.args import VisualGenArgs +from tensorrt_llm.visual_gen.visual_gen import VisualGen + +__all__ = ["FakeVisualGen"] + + +class FakeVisualGen(VisualGen): + """VisualGen generator fake; subclassing keeps ``isinstance`` checks true.""" + + def __init__(self, model: str = "fake-visual-gen") -> None: + # Never run VisualGen.__init__ (it loads a pipeline). VisualGen.shutdown, + # which __del__ calls, returns early when executor is None. + self.model = model + self.args = VisualGenArgs() + self.executor = None + self._req_counter = itertools.count() + + def generate_async(self, *args: Any, **kwargs: Any) -> Any: + raise NotImplementedError("FakeVisualGen does not generate.") diff --git a/tests/test_common/serve/layering.py b/tests/test_common/serve/layering.py new file mode 100644 index 000000000000..aba66f10df46 --- /dev/null +++ b/tests/test_common/serve/layering.py @@ -0,0 +1,222 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Import-direction contracts between serving layers, read from the source tree. + +The import graph is built with ``ast`` from the repository files, not the +installed package, and needs no new dependency. Each package a PR creates adds +its :class:`Contract` to ``tests/unittest/serve/test_layering.py``. +""" + +import ast +from dataclasses import dataclass, field +from pathlib import Path +from typing import FrozenSet, List, Optional, Sequence, Tuple + +__all__ = [ + "Contract", + "ImportEdge", + "build_import_graph", + "check_contracts", + "stale_allowances", +] + + +@dataclass(frozen=True) +class ImportEdge: + importer: str # dotted module of the importing file + imported: str # absolute dotted module (relative imports resolved) + names: Tuple[str, ...] # names of a from-import; () for "import x" + lineno: int + in_function: bool # inside a def or class body (runs at call time) + type_checking: bool # under "if TYPE_CHECKING:" + + +@dataclass(frozen=True) +class Contract: + name: str + importers: Tuple[str, ...] # module prefixes the rule applies to + forbidden: Tuple[str, ...] # imported-module prefixes that are not allowed + allowed: FrozenSet[Tuple[str, str, str]] = field(default_factory=frozenset) + include_type_checking: bool = False + fix: str = "" # printed on failure: which contract, where the code goes + + +def _matches_prefix(module: str, prefix: str) -> bool: + return module == prefix or module.startswith(prefix + ".") + + +class _ImportCollector(ast.NodeVisitor): + def __init__(self, module: str, package: str) -> None: + self._module = module + self._package = package + self._function_depth = 0 + self._type_checking_depth = 0 + self.edges: List[ImportEdge] = [] + + def _add(self, imported: Optional[str], names: Tuple[str, ...], lineno: int) -> None: + if imported is None: + return + self.edges.append( + ImportEdge( + importer=self._module, + imported=imported, + names=names, + lineno=lineno, + in_function=self._function_depth > 0, + type_checking=self._type_checking_depth > 0, + ) + ) + + def visit_Import(self, node: ast.Import) -> None: + for alias in node.names: + self._add(alias.name, (), node.lineno) + + def visit_ImportFrom(self, node: ast.ImportFrom) -> None: + imported = _resolve_relative(self._package, node.level, node.module) + self._add(imported, tuple(alias.name for alias in node.names), node.lineno) + + def visit_Call(self, node: ast.Call) -> None: + # importlib.import_module("") calls are edges with names=(). + target = node.func + is_import_module = ( + isinstance(target, ast.Attribute) + and target.attr == "import_module" + and isinstance(target.value, ast.Name) + and target.value.id == "importlib" + ) or (isinstance(target, ast.Name) and target.id == "import_module") + if is_import_module and node.args and isinstance(node.args[0], ast.Constant): + value = node.args[0].value + if isinstance(value, str): + self._add(value, (), node.lineno) + self.generic_visit(node) + + def visit_If(self, node: ast.If) -> None: + if _is_type_checking_test(node.test): + self._type_checking_depth += 1 + for child in node.body: + self.visit(child) + self._type_checking_depth -= 1 + for child in node.orelse: + self.visit(child) + return + self.generic_visit(node) + + def _visit_scope(self, node) -> None: + self._function_depth += 1 + self.generic_visit(node) + self._function_depth -= 1 + + def visit_FunctionDef(self, node: ast.FunctionDef) -> None: + self._visit_scope(node) + + def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> None: + self._visit_scope(node) + + def visit_ClassDef(self, node: ast.ClassDef) -> None: + self._visit_scope(node) + + +def _is_type_checking_test(test: ast.expr) -> bool: + if isinstance(test, ast.Name): + return test.id == "TYPE_CHECKING" + if isinstance(test, ast.Attribute): + return test.attr == "TYPE_CHECKING" + return False + + +def _resolve_relative(package: str, level: int, module: Optional[str]) -> Optional[str]: + if level == 0: + return module + parts = package.split(".") if package else [] + if level - 1 > len(parts): + return None + base = parts[: len(parts) - (level - 1)] + if module: + base.extend(module.split(".")) + return ".".join(base) if base else None + + +def build_import_graph(package_dir: Path, package: str) -> List[ImportEdge]: + """Every import edge of every ``*.py`` under ``package_dir``. + + Module-level and function-level imports, relative imports resolved, and + ``importlib.import_module("")`` calls recorded as edges with + ``names=()``. + """ + edges: List[ImportEdge] = [] + for path in sorted(package_dir.rglob("*.py")): + relative = path.relative_to(package_dir) + parts = [package, *relative.parts] + if parts[-1] == "__init__.py": + parts = parts[:-1] + else: + parts[-1] = parts[-1][: -len(".py")] + module = ".".join(parts) + file_package = module if path.name == "__init__.py" else module.rpartition(".")[0] + try: + tree = ast.parse(path.read_text(errors="replace")) + except SyntaxError: + continue + collector = _ImportCollector(module, file_package) + collector.visit(tree) + edges.extend(collector.edges) + return edges + + +def _edge_triples(edge: ImportEdge) -> List[Tuple[str, str, str]]: + names = edge.names or ("",) + return [(edge.importer, edge.imported, name) for name in names] + + +def check_contracts(edges: Sequence[ImportEdge], contracts: Sequence[Contract]) -> List[str]: + """Violations of each contract, as messages that name the contract and the fix. + + A from-import is matched as ``.``, so ``from pkg import sub`` + and ``from .. import sub`` count as imports of the submodule ``pkg.sub``. + """ + problems = [] + for contract in contracts: + for edge in edges: + if edge.type_checking and not contract.include_type_checking: + continue + if not any(_matches_prefix(edge.importer, prefix) for prefix in contract.importers): + continue + for triple in _edge_triples(edge): + importer, imported, name = triple + target = f"{imported}.{name}" if name else imported + if not any(_matches_prefix(target, prefix) for prefix in contract.forbidden): + continue + if triple in contract.allowed: + continue + problems.append( + f"contract '{contract.name}': {importer} (line {edge.lineno}) imports " + f"`{target}`. {contract.fix}" + ) + return problems + + +def stale_allowances(edges: Sequence[ImportEdge], contracts: Sequence[Contract]) -> List[str]: + """Baseline allowances that match no current import; the baseline only shrinks.""" + live = {triple for edge in edges for triple in _edge_triples(edge)} + problems = [] + for contract in contracts: + for triple in sorted(contract.allowed): + if triple not in live: + importer, imported, name = triple + problems.append( + f"contract '{contract.name}': the allowance ({importer}, {imported}, " + f"{name!r}) matches no import anymore; delete it, the baseline only shrinks." + ) + return problems diff --git a/tests/test_common/serve/middleware.py b/tests/test_common/serve/middleware.py new file mode 100644 index 000000000000..2cedec9a1741 --- /dev/null +++ b/tests/test_common/serve/middleware.py @@ -0,0 +1,42 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Middleware-order snapshot of a trtllm-serve app.""" + +from typing import Dict, List + +__all__ = ["middleware_snapshot"] + + +def middleware_snapshot(app) -> List[Dict]: + """One dict per middleware in ``app.user_middleware``, outermost first. + + Starlette's ``add_middleware`` inserts at position 0, so the last one added + is outermost and ``app.user_middleware`` already lists outermost first. For + a class whose module starts with ``tensorrt_llm.`` the keyword names are + recorded (plus ``expose_headers`` for ``PerfMetricsMiddleware``); for any + other class only the class name is, because third-party keyword names may + change between versions. + """ + snapshot = [] + for middleware in app.user_middleware: + cls = middleware.cls + if cls.__module__.startswith("tensorrt_llm."): + entry = {"class": cls.__name__, "kwargs": sorted(middleware.kwargs)} + if cls.__name__ == "PerfMetricsMiddleware": + entry["expose_headers"] = middleware.kwargs["expose_headers"] + else: + entry = {"class": cls.__name__} + snapshot.append(entry) + return snapshot diff --git a/tests/test_common/serve/routes.py b/tests/test_common/serve/routes.py new file mode 100644 index 000000000000..b4dc02823074 --- /dev/null +++ b/tests/test_common/serve/routes.py @@ -0,0 +1,494 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Route snapshot of every trtllm-serve app configuration. + +A route record is one string: ``" "``: + +- ````: the route's methods, sorted and comma-joined (``GET,HEAD`` for + FastAPI's documentation routes); ``MOUNT`` for a ``Mount``. +- ````: the effective path from :func:`effective_routes`. +- ````: the route's bare class name, on purpose: moving the route class + to another module must not change the snapshot. +- When two routes in one configuration share `` ``, both records + get `` #`` appended so their order can be compared. Names are not + recorded otherwise, so moving a handler into a router does not change the + snapshot. + +Regenerate the reference with:: + + PYTHONPATH=tests python3 -m test_common.serve.routes --write + +run from the repository root with the checkout's build importable. +""" + +import argparse +import sys +import tempfile +from collections import Counter +from dataclasses import dataclass +from pathlib import Path +from typing import Dict, List, Literal, Optional, Sequence, Tuple, Union + +import yaml +from fastapi import FastAPI + +from tensorrt_llm.llmapi.disagg_utils import ServerRole +from test_common.serve.app_factory import ( + AggConfig, + CoordinatorConfig, + DisaggConfig, + isolated_serve_env, +) + +__all__ = [ + "CONFIGS", + "VISUAL_GEN_CONFIGS", + "SnapshotConfig", + "compare", + "effective_routes", + "read_reference", + "reference_path", + "route_record", + "snapshot_config", + "snapshot_routes", + "write_reference", +] + +REPO_ROOT = Path(__file__).resolve().parents[3] + +REFERENCE_HEADER = """\ +# Routes of every trtllm-serve app configuration. Generated, do not edit by hand: +# PYTHONPATH=tests python3 -m test_common.serve.routes --write +# Each route is " ", in registration order. "#" is +# added only when two routes share methods and path. The checker compares the order +# only for routes that can match the same request. Changing this file changes the +# HTTP surface of trtllm-serve; the API review committee approves it. +""" + + +def reference_path(repo_root: Path = REPO_ROOT) -> Path: + return ( + repo_root + / "tests" + / "unittest" + / "api_stability" + / "references" + / "trtllm_serve_routes.yaml" + ) + + +# --------------------------------------------------------------------------- +# Walking routes +# --------------------------------------------------------------------------- + + +def effective_routes(app_or_routes, prefix: str = "") -> list: + """(effective path, leaf route) in match order, on FastAPI 0.136.x and 0.139.x. + + FastAPI 0.139 keeps one ``_IncludedRouter`` object per ``include_router`` + call in ``app.routes``; it has no ``path`` attribute, so reading ``.path`` + off ``app.routes`` raises once any router is included. This walks into such + objects and applies their include prefix. On 0.136.x included routes are + already flat, so the recursion never runs and the result equals + ``app.routes``. The duck-typed check below uses two attribute names of a + private FastAPI class; if a later FastAPI renames them, + ``test_effective_routes_flattens_included_routers`` fails first, naming this + function. + """ + routes = app_or_routes.routes if isinstance(app_or_routes, FastAPI) else app_or_routes + out = [] + for route in routes: + included = getattr(route, "original_router", None) + context = getattr(route, "include_context", None) + if included is not None and context is not None: # FastAPI 0.139 _IncludedRouter + out.extend(effective_routes(included.routes, prefix + context.prefix)) + continue + # Mounts keep their path and are not descended into. + out.append((prefix + getattr(route, "path", ""), route)) + return out + + +def route_record(path: str, route) -> str: + """``" "`` for one route, without the ``#`` suffix.""" + methods = getattr(route, "methods", None) + methods_text = ",".join(sorted(methods)) if methods else "MOUNT" + return f"{methods_text} {path} {type(route).__name__}" + + +def snapshot_routes(app: FastAPI) -> List[str]: + """Route records of ``app`` in registration order, disambiguated by name.""" + pairs = effective_routes(app) + records = [route_record(path, route) for path, route in pairs] + key_counts = Counter(record.rsplit(" ", 1)[0] for record in records) + out = [] + for record, (path, route) in zip(records, pairs): + methods_and_path = record.rsplit(" ", 1)[0] + if key_counts[methods_and_path] > 1: + record = f"{record} #{getattr(route, 'name', '')}" + out.append(record) + return out + + +# --------------------------------------------------------------------------- +# Configuration matrix +# --------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class SnapshotConfig: + name: str + kind: Literal["openai_server", "disagg_server", "coordinator_server"] + build: Union[AggConfig, DisaggConfig, CoordinatorConfig] + + +CONFIGS: Tuple[SnapshotConfig, ...] = ( + SnapshotConfig("llm", "openai_server", AggConfig()), + SnapshotConfig("llm_context", "openai_server", AggConfig(role=ServerRole.CONTEXT)), + SnapshotConfig("llm_generation", "openai_server", AggConfig(role=ServerRole.GENERATION)), + SnapshotConfig("llm_gpt_oss", "openai_server", AggConfig(model_type="gpt_oss")), + SnapshotConfig( + "llm_gpt_oss_harmony_disabled", + "openai_server", + AggConfig(model_type="gpt_oss", env=(("DISABLE_HARMONY_ADAPTER", "1"),)), + ), + SnapshotConfig("llm_governor", "openai_server", AggConfig(governor_queue=True)), + SnapshotConfig("llm_rl_control", "openai_server", AggConfig(rl_control=True)), + SnapshotConfig("llm_runtime_control", "openai_server", AggConfig(runtime_control=True)), + SnapshotConfig( + "llm_rl_and_runtime_control", + "openai_server", + AggConfig(rl_control=True, runtime_control=True), + ), + SnapshotConfig("llm_perf_metrics", "openai_server", AggConfig(return_perf_metrics=True)), + SnapshotConfig( + "llm_perf_metrics_enable_metrics_env", + "openai_server", + AggConfig(return_perf_metrics=True, env=(("ENABLE_METRICS", "true"),)), + ), + SnapshotConfig( + "llm_perf_metrics_output_dir_only", + "openai_server", + AggConfig(perf_metrics_output_dir=True), + ), + SnapshotConfig( + "llm_render_endpoints", + "openai_server", + AggConfig(env=(("TRTLLM_ENABLE_RENDER_ENDPOINTS", "1"),)), + ), + # An attached frontend of a multi-frontend server registers the launcher + # forwarding routes before the local table; same method and path, so + # registration order decides which handler serves them. + SnapshotConfig( + "llm_multi_frontend_attached", + "openai_server", + AggConfig(multi_frontend_attached=True), + ), + SnapshotConfig("mm_encoder", "openai_server", AggConfig(role=ServerRole.MM_ENCODER)), + SnapshotConfig( + "mm_encoder_rl_control", + "openai_server", + AggConfig(role=ServerRole.MM_ENCODER, rl_control=True), + ), + SnapshotConfig("embedding", "openai_server", AggConfig(role=ServerRole.EMBEDDING)), + SnapshotConfig( + "embedding_runtime_control", + "openai_server", + AggConfig(role=ServerRole.EMBEDDING, runtime_control=True), + ), + SnapshotConfig("disagg_local", "disagg_server", DisaggConfig()), + SnapshotConfig( + "disagg_local_perf_metrics", "disagg_server", DisaggConfig(return_perf_metrics=True) + ), + SnapshotConfig( + "disagg_local_http_cluster_storage", + "disagg_server", + DisaggConfig(http_cluster_storage=True), + ), + SnapshotConfig("disagg_fleet", "disagg_server", DisaggConfig(fleet=True)), + SnapshotConfig( + "disagg_fleet_perf_metrics", + "disagg_server", + DisaggConfig(fleet=True, return_perf_metrics=True), + ), + SnapshotConfig( + "disagg_fleet_http_cluster_storage", + "disagg_server", + DisaggConfig(fleet=True, http_cluster_storage=True), + ), + SnapshotConfig("coordinator", "coordinator_server", CoordinatorConfig()), + SnapshotConfig( + "coordinator_http_cluster_storage", + "coordinator_server", + CoordinatorConfig(http_cluster_storage=True), + ), +) + +VISUAL_GEN_CONFIGS: Tuple[SnapshotConfig, ...] = ( + SnapshotConfig( + "visual_gen", + "openai_server", + AggConfig(role=ServerRole.VISUAL_GEN, model_name="fake-visual-gen"), + ), +) + + +def snapshot_config(config: SnapshotConfig, tmp_dir: Path) -> Dict: + """``{"kind": ..., "routes": [...]}`` or ``{"kind": ..., "error": ": "}``.""" + with isolated_serve_env(tmp_dir) as env: + try: + if config.kind == "openai_server": + app = env.build_openai_server(config.build).app + elif config.kind == "disagg_server": + app = env.build_disagg_server(config.build).app + else: + app = env.build_coordinator_server(config.build).app + except ValueError as error: + return {"kind": config.kind, "error": f"{type(error).__name__}: {error}"} + return {"kind": config.kind, "routes": snapshot_routes(app)} + + +# --------------------------------------------------------------------------- +# Comparison +# --------------------------------------------------------------------------- + + +def _parse_record(record: str) -> Tuple[str, str, str, str]: + """(methods, path, class, name) of one record; name is "" when absent.""" + parts = record.split(" ") + methods, path, cls = parts[0], parts[1], parts[2] + name = parts[3][1:] if len(parts) > 3 and parts[3].startswith("#") else "" + return methods, path, cls, name + + +def _literal_prefix(path: str) -> str: + """The path up to its first ``{``.""" + brace = path.find("{") + return path if brace < 0 else path[:brace] + + +def _paths_can_match_same_request(path_a: str, path_b: str) -> bool: + segments_a = path_a.split("/") + segments_b = path_b.split("/") + for i in range(max(len(segments_a), len(segments_b))): + a = segments_a[i] if i < len(segments_a) else None + b = segments_b[i] if i < len(segments_b) else None + # A {name:path} parameter matches the rest of the path. + if a is not None and "{" in a and a.rstrip("}").endswith(":path"): + return True + if b is not None and "{" in b and b.rstrip("}").endswith(":path"): + return True + if a is None or b is None: + # Different segment counts with no path parameter: no overlap. + return False + if "{" in a or "{" in b: + continue # a parameter matches any literal + if a != b: + return False + return True + + +def _records_conflict(record_a: str, record_b: str) -> bool: + """Whether two records can match the same request (over-approximated). + + Two records conflict when (a) their methods intersect or either is a + mount, and (b) their path templates can match one request path. + Over-approximating conflicts only locks more pairs, which is safe. + """ + methods_a, path_a, _, _ = _parse_record(record_a) + methods_b, path_b, _, _ = _parse_record(record_b) + a_is_mount = methods_a == "MOUNT" + b_is_mount = methods_b == "MOUNT" + if not a_is_mount and not b_is_mount: + if not set(methods_a.split(",")) & set(methods_b.split(",")): + return False + return _paths_can_match_same_request(path_a, path_b) + # A mount at P conflicts with every route whose literal prefix starts + # with P, or that P starts with. + if a_is_mount and b_is_mount: + return path_a.startswith(path_b) or path_b.startswith(path_a) + mount_path = path_a if a_is_mount else path_b + other_prefix = _literal_prefix(path_b if a_is_mount else path_a) + return other_prefix.startswith(mount_path) or mount_path.startswith(other_prefix) + + +def compare(expected: Sequence[str], actual: Sequence[str]) -> List[str]: + """Human-readable differences between the reference and the live records.""" + differences: List[str] = [] + expected_counts = Counter(expected) + actual_counts = Counter(actual) + removed = list((expected_counts - actual_counts).elements()) + added = list((actual_counts - expected_counts).elements()) + + # A removed and an added record with the same methods and path but another + # class are one "route class changed" line. + def _methods_path(record: str) -> str: + methods, path, _, _ = _parse_record(record) + return f"{methods} {path}" + + remaining_added = list(added) + class_changes = [] + for removed_record in list(removed): + for added_record in remaining_added: + if _methods_path(added_record) == _methods_path(removed_record): + old_cls = _parse_record(removed_record)[2] + new_cls = _parse_record(added_record)[2] + why = ( + " (a plain APIRouter breaks msgpack bodies from the disaggregated frontend)" + if old_cls == "_MsgspecRoute" + else "" + ) + class_changes.append( + f"route class changed: {_methods_path(removed_record)}: " + f"{old_cls} -> {new_cls}{why}" + ) + removed.remove(removed_record) + remaining_added.remove(added_record) + break + for record in remaining_added: + differences.append(f"added: {record}") + for record in removed: + differences.append(f"removed: {record}") + differences.extend(class_changes) + + # Conflict-relevant order: for every conflicting pair a-before-b in + # expected, a must come before b in actual, when both are present. For + # repeated identical records, the k-th occurrence pairs with the k-th. + actual_positions: Dict[str, List[int]] = {} + for position, record in enumerate(actual): + actual_positions.setdefault(record, []).append(position) + expected_occurrence: List[Optional[int]] = [] + seen: Counter = Counter() + for record in expected: + positions = actual_positions.get(record, []) + occurrence = seen[record] + seen[record] += 1 + expected_occurrence.append(positions[occurrence] if occurrence < len(positions) else None) + for i, record_a in enumerate(expected): + position_a = expected_occurrence[i] + if position_a is None: + continue + for j in range(i + 1, len(expected)): + record_b = expected[j] + position_b = expected_occurrence[j] + if position_b is None or record_a == record_b: + continue + if _records_conflict(record_a, record_b) and position_a > position_b: + differences.append( + f"order changed: {record_a} must come before {record_b} " + "(FastAPI matches in registration order)" + ) + return differences + + +def format_failure(name: str, differences: Sequence[str]) -> str: + lines = "\n".join(f" {difference}" for difference in differences) + return ( + f"Route snapshot of configuration '{name}' differs from " + "tests/unittest/api_stability/references/trtllm_serve_routes.yaml:\n" + f"{lines}\n" + "If the change is intended, regenerate with " + "`PYTHONPATH=tests python3 -m test_common.serve.routes --write`, commit the YAML, " + "and request review from @NVIDIA/trt-llm-noncommitted-api-review-committee " + "(owner of tests/unittest/api_stability/). New endpoints are a router and a " + "serving class under tensorrt_llm/serve/api///: create the " + "router with make_router() and add a RouterSpec to ROUTERS " + "(tensorrt_llm/serve/_base/routing.py). Do not add routes or methods to " + "OpenAIServer; see tensorrt_llm/serve/AGENTS.md." + ) + + +# --------------------------------------------------------------------------- +# Reference file +# --------------------------------------------------------------------------- + + +def build_reference_data(tmp_root: Path) -> Dict: + configs: Dict[str, Dict] = {} + for index, config in enumerate(CONFIGS + VISUAL_GEN_CONFIGS): + configs[config.name] = snapshot_config(config, tmp_root / f"cfg{index}") + return {"version": 1, "configs": configs} + + +def write_reference(data: Dict, path: Optional[Path] = None) -> Path: + path = path or reference_path() + body = yaml.safe_dump(data, sort_keys=False, width=4096) + path.write_text(REFERENCE_HEADER + body) + return path + + +def read_reference(path: Optional[Path] = None) -> Dict: + path = path or reference_path() + return yaml.safe_load(path.read_text()) + + +def main(argv: Optional[Sequence[str]] = None) -> int: + parser = argparse.ArgumentParser( + prog="python3 -m test_common.serve.routes", + description="Check, rewrite or print the trtllm-serve route reference.", + ) + parser.add_argument( + "--write", action="store_true", help="rebuild every configuration and rewrite the reference" + ) + parser.add_argument( + "--print", + dest="print_name", + metavar="CONFIG", + help="print one configuration's live records", + ) + args = parser.parse_args(argv) + + by_name = {config.name: config for config in CONFIGS + VISUAL_GEN_CONFIGS} + with tempfile.TemporaryDirectory() as tmp: + tmp_root = Path(tmp) + if args.print_name: + if args.print_name not in by_name: + print(f"unknown configuration '{args.print_name}'; one of: {sorted(by_name)}") + return 2 + snapshot = snapshot_config(by_name[args.print_name], tmp_root / "print") + for line in snapshot.get("routes", [snapshot.get("error", "")]): + print(line) + return 0 + data = build_reference_data(tmp_root) + if args.write: + path = write_reference(data) + print(f"wrote {path}") + return 0 + reference = read_reference() + exit_code = 0 + for name, live in data["configs"].items(): + recorded = reference["configs"].get(name) + if recorded is None: + print(f"configuration '{name}' missing from the reference") + exit_code = 1 + continue + if recorded.get("error") != live.get("error"): + print( + f"configuration '{name}': error changed: " + f"{recorded.get('error')!r} -> {live.get('error')!r}" + ) + exit_code = 1 + differences = compare(recorded.get("routes", []), live.get("routes", [])) + if differences: + print(format_failure(name, differences)) + exit_code = 1 + stale = set(reference["configs"]) - set(data["configs"]) + if stale: + print(f"configurations only in the reference: {sorted(stale)}") + exit_code = 1 + return exit_code + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/test_common/serve/shims.py b/tests/test_common/serve/shims.py new file mode 100644 index 000000000000..98ae1a1cfa04 --- /dev/null +++ b/tests/test_common/serve/shims.py @@ -0,0 +1,376 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Shim identity, FutureWarning and old-path-import checks for moved serve modules. + +A module move keeps the old path only when code outside this repository imports +it. Such a kept shim re-exports the public names (each the same object as at the +new path), warns once with ``FutureWarning``, and code inside the repository +must import the new path. These checks hold all three properties; the PR that +creates a shim adds its :class:`ShimRow` rows to the tables in +``tests/unittest/serve/test_shim_identity.py``. +""" + +import ast +import importlib +import re +import sys +import warnings +from dataclasses import dataclass +from pathlib import Path +from typing import Dict, List, Optional, Sequence, Tuple + +__all__ = [ + "SCAN_ROOTS", + "ShimRow", + "check_identity", + "check_warnings", + "find_old_path_imports", +] + +SCAN_ROOTS = ("tensorrt_llm", "examples", "scripts", "triton_backend", "docs", "setup.py") + + +@dataclass(frozen=True) +class ShimRow: + """One re-exported name of one kept old path.""" + + old_module: str # "tensorrt_llm.serve.openai_disagg_server" + name: str # "OpenAIDisaggServer"; public names only, except in interface rows + new_module: str # "tensorrt_llm.serve.disagg.openai_disagg_server" + retire_in: str # release that removes the shim; "not retired" for an interface row + ban_old_path: bool = True # False: repo code may keep importing name from old_module + warns: bool = True # False: the re-export is the module's own interface, not a shim + + +# --------------------------------------------------------------------------- +# Identity +# --------------------------------------------------------------------------- + + +def check_identity(rows: Sequence[ShimRow]) -> List[str]: + """Each old-path name must be the same object as the new-path name. + + Identity only, never ``__module__``: some re-exported names are dicts or + constants. + """ + problems = [] + for row in rows: + old_module = importlib.import_module(row.old_module) + new_module = importlib.import_module(row.new_module) + try: + old_object = getattr(old_module, row.name) + except AttributeError: + problems.append( + f"`{row.old_module}.{row.name}` no longer exists: the re-export was removed, " + "probably by autoflake or ruff during pre-commit, which drop an import that " + f"`__all__` does not list. Restore `from {row.new_module} import {row.name}` in " + f"{row.old_module.replace('.', '/')}.py and list the name in `__all__` (a pure " + "shim), or restore its row in `_LAZY_NAMES` (`commands/serve.py`)." + ) + continue + new_object = getattr(new_module, row.name) + if old_object is not new_object: + problems.append( + f"`{row.old_module}.{row.name}` is not the same object as " + f"`{row.new_module}.{row.name}`. The old module must re-export the object, " + "not define a copy." + ) + return problems + + +# --------------------------------------------------------------------------- +# FutureWarning +# --------------------------------------------------------------------------- + +_IMPORT_WARNING = re.compile( + r"^(?P[A-Za-z0-9_.]+) has moved to (?P[A-Za-z0-9_.]+) " + r"and will be removed in (?P\S+?)\.$" +) +_READ_WARNING = re.compile( + r"^(?P[A-Za-z0-9_.]+)\.(?P[A-Za-z0-9_]+) has moved to " + r"(?P[A-Za-z0-9_.]+) and will be removed in (?P\S+?)\.$" +) + + +def _matching_warnings(caught, prefix: str) -> List[str]: + """Messages of the FutureWarnings that start with ``prefix``. + + Filtering by message, as ``tests/unittest/_torch/peft/test_shim_peft.py`` + does, also drops the warning of a parent shim package that a fresh import of + one of its submodules emits. + """ + return [ + str(record.message) + for record in caught + if issubclass(record.category, FutureWarning) and str(record.message).startswith(prefix) + ] + + +def _no_warning_message(old_module: str, where: str, name: str = "") -> str: + what = f"reading `{old_module}.{name}`" if where == "read" else f"importing `{old_module}`" + return ( + f"{what} emitted no FutureWarning. A kept shim warns once with " + "warnings.warn(..., FutureWarning, stacklevel=2): the TRT-LLM logger defaults to level " + "error, so a logged warning would not reach callers outside the repository. Restore " + "that call." + ) + + +def check_warnings(rows: Sequence[ShimRow]) -> List[str]: + """Each kept shim warns once with the expected FutureWarning text. + + A pure shim (no ``_LAZY_NAMES``) warns exactly once at import with + `` has moved to and will be removed in .``, + where ```` is the row's new module or a package that contains it, + and never on reads. A ``_LAZY_NAMES`` module warns never at import and + exactly once on the first read of each lazy name; a name it imports eagerly + never warns. Rows with ``warns=False`` (interface rows) are skipped. + + Imports each old module fresh (a module body runs once per process) and + puts any previously imported module object back in ``sys.modules`` and on + its parent package, so later tests keep the objects they imported. + """ + problems: List[str] = [] + by_module: Dict[str, List[ShimRow]] = {} + for row in rows: + if row.warns: + by_module.setdefault(row.old_module, []).append(row) + + for module_name, module_rows in by_module.items(): + saved = sys.modules.get(module_name) + if saved is not None: + del sys.modules[module_name] + try: + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + module = importlib.import_module(module_name) + import_messages = _matching_warnings(caught, f"{module_name} has moved to ") + lazy_names = getattr(module, "_LAZY_NAMES", None) + + if lazy_names is None: + problems.extend( + _check_pure_shim_import(module_name, module_rows, import_messages) + ) + elif import_messages: + problems.append( + f"importing `{module_name}` emitted a FutureWarning; trtllm-serve " + "imports commands/serve.py on every start, so it may warn only when a " + "moved name is read." + ) + + for row in module_rows: + is_lazy_name = lazy_names is not None and row.name in lazy_names + for read_index in (1, 2): + before = len(caught) + getattr(module, row.name) + read_messages = _matching_warnings( + caught[before:], f"{module_name}.{row.name} has moved to " + ) + if is_lazy_name and read_index == 1: + problems.extend(_check_lazy_first_read(row, read_messages)) + elif read_messages: + problems.append( + f"the FutureWarning of `{module_name}` says " + f"`{read_messages[0]}`; expected no warning on this read of " + f"`{row.name}`." + ) + finally: + if saved is not None: + sys.modules[module_name] = saved + parent_name, _, tail = module_name.rpartition(".") + if parent_name and parent_name in sys.modules: + setattr(sys.modules[parent_name], tail, saved) + return problems + + +def _check_pure_shim_import( + module_name: str, module_rows: Sequence[ShimRow], import_messages: List[str] +) -> List[str]: + if not import_messages: + return [_no_warning_message(module_name, "import")] + problems = [] + if len(import_messages) > 1: + problems.append( + f"importing `{module_name}` emitted {len(import_messages)} FutureWarnings; a pure " + "shim warns exactly once, at import." + ) + message = import_messages[0] + match = _IMPORT_WARNING.match(message) + retire_in = module_rows[0].retire_in + expected = f"{module_name} has moved to and will be removed in {retire_in}." + if match is None or match.group("release") != retire_in: + problems.append( + f"the FutureWarning of `{module_name}` says `{message}`; expected `{expected}`: it " + "names the new location and the removal release." + ) + return problems + location = match.group("location") + for row in module_rows: + if not (row.new_module == location or row.new_module.startswith(location + ".")): + problems.append( + f"the FutureWarning of `{module_name}` says `{message}`; expected a location " + f"that contains `{row.new_module}`: it names the new location and the removal " + "release." + ) + return problems + + +def _check_lazy_first_read(row: ShimRow, read_messages: List[str]) -> List[str]: + if not read_messages: + return [_no_warning_message(row.old_module, "read", row.name)] + expected = ( + f"{row.old_module}.{row.name} has moved to {row.new_module} " + f"and will be removed in {row.retire_in}." + ) + if read_messages != [expected]: + return [ + f"the FutureWarning of `{row.old_module}` says `{read_messages[0]}`; expected " + f"`{expected}`: it names the new location and the removal release." + ] + return [] + + +# --------------------------------------------------------------------------- +# Old-path import ban +# --------------------------------------------------------------------------- + + +def _module_name_of(path: Path, repo_root: Path) -> str: + relative = path.relative_to(repo_root) + parts = list(relative.parts) + if parts[-1] == "__init__.py": + parts = parts[:-1] + else: + parts[-1] = parts[-1][: -len(".py")] + return ".".join(parts) + + +def _package_of(path: Path, repo_root: Path) -> str: + module = _module_name_of(path, repo_root) + if path.name == "__init__.py": + return module + return module.rpartition(".")[0] + + +def _resolve_relative(package: str, level: int, module: Optional[str]) -> Optional[str]: + if level == 0: + return module + parts = package.split(".") if package else [] + if level - 1 > len(parts): + return None + base = parts[: len(parts) - (level - 1)] + if module: + base.append(module) + return ".".join(base) if base else None + + +def find_old_path_imports( + rows: Sequence[ShimRow], repo_root: Path, roots: Sequence[str] = SCAN_ROOTS +) -> List[str]: + """Imports of banned names from kept old paths, anywhere under ``roots``. + + Flags ``from import `` (absolute, or relative resolved + against the file's package) and ``.`` attribute access where + ```` is bound to the old module in that file. Importing an old module + that still has code is allowed; only the banned names are flagged. The old + module's own file is exempt, and ``tests/`` is not scanned. + """ + banned: Dict[str, set] = {} + for row in rows: + if row.ban_old_path: + banned.setdefault(row.old_module, set()).add(row.name) + if not banned: + return [] + + retire_by_module = {row.old_module: row.retire_in for row in rows} + new_by_module_name: Dict[Tuple[str, str], str] = { + (row.old_module, row.name): row.new_module for row in rows + } + problems: List[str] = [] + files: List[Path] = [] + for root in roots: + root_path = repo_root / root + if root_path.is_file(): + files.append(root_path) + elif root_path.is_dir(): + files.extend(sorted(root_path.rglob("*.py"))) + + for path in files: + text = path.read_text(errors="replace") + # Cheap prefilter: parse only files that mention an old module. + if not any(module.rsplit(".", 1)[-1] in text for module in banned): + continue + try: + tree = ast.parse(text) + except SyntaxError: + continue + file_module = _module_name_of(path, repo_root) + exempt_modules = {module for module in banned if file_module == module} + package = _package_of(path, repo_root) + + aliases: Dict[str, str] = {} # local name -> old module it is bound to + hits: List[Tuple[int, str, str]] = [] # (line, old module, name) + + for node in ast.walk(tree): + if isinstance(node, ast.Import): + for alias in node.names: + if alias.name in banned and alias.name not in exempt_modules and alias.asname: + aliases[alias.asname] = alias.name + elif isinstance(node, ast.ImportFrom): + resolved = _resolve_relative(package, node.level, node.module) + if resolved is None: + continue + for alias in node.names: + if resolved in banned and resolved not in exempt_modules: + if alias.name in banned[resolved]: + hits.append((node.lineno, resolved, alias.name)) + # `from import ` binds the module. + candidate = f"{resolved}.{alias.name}" + if candidate in banned and candidate not in exempt_modules: + aliases[alias.asname or alias.name] = candidate + + for node in ast.walk(tree): + if not isinstance(node, ast.Attribute): + continue + dotted = _dotted_name(node) + if dotted is None: + continue + prefix, _, attribute = dotted.rpartition(".") + bound = aliases.get(prefix) + if bound is None and prefix in banned and prefix not in exempt_modules: + bound = prefix # the dotted `import ` used as `.` + if bound is not None and attribute in banned.get(bound, ()): + hits.append((node.lineno, bound, attribute)) + + for lineno, module, name in sorted(set(hits)): + problems.append( + f"{path.relative_to(repo_root)}:{lineno} imports `{name}` from the old path " + f"`{module}`; import it from `{new_by_module_name[(module, name)]}`. The old " + "path is kept only for callers outside the repository until " + f"{retire_by_module[module]}." + ) + return problems + + +def _dotted_name(node: ast.Attribute) -> Optional[str]: + parts = [node.attr] + value = node.value + while isinstance(value, ast.Attribute): + parts.append(value.attr) + value = value.value + if isinstance(value, ast.Name): + parts.append(value.id) + return ".".join(reversed(parts)) + return None diff --git a/tests/unittest/_torch/pyexecutor/test_py_executor.py b/tests/unittest/_torch/pyexecutor/test_py_executor.py index bbe8b941859c..c56a76dbe756 100644 --- a/tests/unittest/_torch/pyexecutor/test_py_executor.py +++ b/tests/unittest/_torch/pyexecutor/test_py_executor.py @@ -1672,6 +1672,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/_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) diff --git a/tests/unittest/api_stability/api_stability_core.py b/tests/unittest/api_stability/api_stability_core.py index 8264f810b8b3..2541d02854f2 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 7c1adb05c659..001e55bfd51a 100644 --- a/tests/unittest/api_stability/references/llm.yaml +++ b/tests/unittest/api_stability/references/llm.yaml @@ -515,6 +515,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: @@ -536,6 +544,21 @@ methods: default: 2 return_annotation: tensorrt_llm.executor.result.IterationResult status: beta + check_health: + 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 @@ -592,3 +615,30 @@ 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 + 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/api_stability/references/trtllm_serve_routes.yaml b/tests/unittest/api_stability/references/trtllm_serve_routes.yaml new file mode 100644 index 000000000000..523342d4afae --- /dev/null +++ b/tests/unittest/api_stability/references/trtllm_serve_routes.yaml @@ -0,0 +1,663 @@ +# Routes of every trtllm-serve app configuration. Generated, do not edit by hand: +# PYTHONPATH=tests python3 -m test_common.serve.routes --write +# Each route is " ", in registration order. "#" is +# added only when two routes share methods and path. The checker compares the order +# only for routes that can match the same request. Changing this file changes the +# HTTP surface of trtllm-serve; the API review committee approves it. +version: 1 +configs: + llm: + kind: openai_server + routes: + - GET,HEAD /openapi.json Route + - GET,HEAD /docs Route + - GET,HEAD /docs/oauth2-redirect Route + - GET,HEAD /redoc Route + - GET /health _MsgspecRoute + - GET /health_generate _MsgspecRoute + - GET /version _MsgspecRoute + - GET /v1/models _MsgspecRoute + - GET /v1/data_transceiver_state _MsgspecRoute + - GET /metrics _MsgspecRoute + - GET /energy_metrics _MsgspecRoute + - GET /steady_clock_offset _MsgspecRoute + - POST /steady_clock_offset _MsgspecRoute + - POST /kv_cache_events _MsgspecRoute + - POST /v1/completions _MsgspecRoute + - POST /v1/chat/completions _MsgspecRoute + - POST /v1/messages _MsgspecRoute + - POST /v1/messages/count_tokens _MsgspecRoute + - POST /v1/messages/batches _MsgspecRoute + - GET /v1/messages/batches _MsgspecRoute + - GET /v1/messages/batches/{batch_id} _MsgspecRoute + - DELETE /v1/messages/batches/{batch_id} _MsgspecRoute + - POST /v1/messages/batches/{batch_id}/cancel _MsgspecRoute + - GET /v1/messages/batches/{batch_id}/results _MsgspecRoute + - POST /v1/responses _MsgspecRoute + - GET /v1/responses/{response_id} _MsgspecRoute + - DELETE /v1/responses/{response_id} _MsgspecRoute + - POST /_internal/tokenize _MsgspecRoute + - POST /start_profile _MsgspecRoute + - POST /stop_profile _MsgspecRoute + - GET /server_info _MsgspecRoute + llm_context: + kind: openai_server + routes: + - GET,HEAD /openapi.json Route + - GET,HEAD /docs Route + - GET,HEAD /docs/oauth2-redirect Route + - GET,HEAD /redoc Route + - GET /health _MsgspecRoute + - GET /health_generate _MsgspecRoute + - GET /version _MsgspecRoute + - GET /v1/models _MsgspecRoute + - GET /v1/data_transceiver_state _MsgspecRoute + - GET /metrics _MsgspecRoute + - GET /energy_metrics _MsgspecRoute + - GET /steady_clock_offset _MsgspecRoute + - POST /steady_clock_offset _MsgspecRoute + - POST /kv_cache_events _MsgspecRoute + - POST /v1/completions _MsgspecRoute + - POST /v1/chat/completions _MsgspecRoute + - POST /v1/messages _MsgspecRoute + - POST /v1/messages/count_tokens _MsgspecRoute + - POST /v1/messages/batches _MsgspecRoute + - GET /v1/messages/batches _MsgspecRoute + - GET /v1/messages/batches/{batch_id} _MsgspecRoute + - DELETE /v1/messages/batches/{batch_id} _MsgspecRoute + - POST /v1/messages/batches/{batch_id}/cancel _MsgspecRoute + - GET /v1/messages/batches/{batch_id}/results _MsgspecRoute + - POST /v1/responses _MsgspecRoute + - GET /v1/responses/{response_id} _MsgspecRoute + - DELETE /v1/responses/{response_id} _MsgspecRoute + - POST /_internal/tokenize _MsgspecRoute + - POST /start_profile _MsgspecRoute + - POST /stop_profile _MsgspecRoute + - GET /server_info _MsgspecRoute + llm_generation: + kind: openai_server + routes: + - GET,HEAD /openapi.json Route + - GET,HEAD /docs Route + - GET,HEAD /docs/oauth2-redirect Route + - GET,HEAD /redoc Route + - GET /health _MsgspecRoute + - GET /health_generate _MsgspecRoute + - GET /version _MsgspecRoute + - GET /v1/models _MsgspecRoute + - GET /v1/data_transceiver_state _MsgspecRoute + - GET /metrics _MsgspecRoute + - GET /energy_metrics _MsgspecRoute + - GET /steady_clock_offset _MsgspecRoute + - POST /steady_clock_offset _MsgspecRoute + - POST /kv_cache_events _MsgspecRoute + - POST /v1/completions _MsgspecRoute + - POST /v1/chat/completions _MsgspecRoute + - POST /v1/messages _MsgspecRoute + - POST /v1/messages/count_tokens _MsgspecRoute + - POST /v1/messages/batches _MsgspecRoute + - GET /v1/messages/batches _MsgspecRoute + - GET /v1/messages/batches/{batch_id} _MsgspecRoute + - DELETE /v1/messages/batches/{batch_id} _MsgspecRoute + - POST /v1/messages/batches/{batch_id}/cancel _MsgspecRoute + - GET /v1/messages/batches/{batch_id}/results _MsgspecRoute + - POST /v1/responses _MsgspecRoute + - GET /v1/responses/{response_id} _MsgspecRoute + - DELETE /v1/responses/{response_id} _MsgspecRoute + - POST /_internal/tokenize _MsgspecRoute + - POST /start_profile _MsgspecRoute + - POST /stop_profile _MsgspecRoute + - GET /server_info _MsgspecRoute + llm_gpt_oss: + kind: openai_server + routes: + - GET,HEAD /openapi.json Route + - GET,HEAD /docs Route + - GET,HEAD /docs/oauth2-redirect Route + - GET,HEAD /redoc Route + - GET /health _MsgspecRoute + - GET /health_generate _MsgspecRoute + - GET /version _MsgspecRoute + - GET /v1/models _MsgspecRoute + - GET /v1/data_transceiver_state _MsgspecRoute + - GET /metrics _MsgspecRoute + - GET /energy_metrics _MsgspecRoute + - GET /steady_clock_offset _MsgspecRoute + - POST /steady_clock_offset _MsgspecRoute + - POST /kv_cache_events _MsgspecRoute + - POST /v1/completions _MsgspecRoute + - POST /v1/chat/completions _MsgspecRoute + - POST /v1/responses _MsgspecRoute + - GET /v1/responses/{response_id} _MsgspecRoute + - DELETE /v1/responses/{response_id} _MsgspecRoute + - POST /_internal/tokenize _MsgspecRoute + - POST /start_profile _MsgspecRoute + - POST /stop_profile _MsgspecRoute + - GET /server_info _MsgspecRoute + llm_gpt_oss_harmony_disabled: + kind: openai_server + routes: + - GET,HEAD /openapi.json Route + - GET,HEAD /docs Route + - GET,HEAD /docs/oauth2-redirect Route + - GET,HEAD /redoc Route + - GET /health _MsgspecRoute + - GET /health_generate _MsgspecRoute + - GET /version _MsgspecRoute + - GET /v1/models _MsgspecRoute + - GET /v1/data_transceiver_state _MsgspecRoute + - GET /metrics _MsgspecRoute + - GET /energy_metrics _MsgspecRoute + - GET /steady_clock_offset _MsgspecRoute + - POST /steady_clock_offset _MsgspecRoute + - POST /kv_cache_events _MsgspecRoute + - POST /v1/completions _MsgspecRoute + - POST /v1/chat/completions _MsgspecRoute + - POST /v1/messages _MsgspecRoute + - POST /v1/messages/count_tokens _MsgspecRoute + - POST /v1/messages/batches _MsgspecRoute + - GET /v1/messages/batches _MsgspecRoute + - GET /v1/messages/batches/{batch_id} _MsgspecRoute + - DELETE /v1/messages/batches/{batch_id} _MsgspecRoute + - POST /v1/messages/batches/{batch_id}/cancel _MsgspecRoute + - GET /v1/messages/batches/{batch_id}/results _MsgspecRoute + - POST /v1/responses _MsgspecRoute + - GET /v1/responses/{response_id} _MsgspecRoute + - DELETE /v1/responses/{response_id} _MsgspecRoute + - POST /_internal/tokenize _MsgspecRoute + - POST /start_profile _MsgspecRoute + - POST /stop_profile _MsgspecRoute + - GET /server_info _MsgspecRoute + llm_governor: + kind: openai_server + routes: + - GET,HEAD /openapi.json Route + - GET,HEAD /docs Route + - GET,HEAD /docs/oauth2-redirect Route + - GET,HEAD /redoc Route + - GET /health _MsgspecRoute + - GET /health_generate _MsgspecRoute + - GET /version _MsgspecRoute + - GET /v1/models _MsgspecRoute + - GET /v1/data_transceiver_state _MsgspecRoute + - GET /metrics _MsgspecRoute + - GET /energy_metrics _MsgspecRoute + - GET /steady_clock_offset _MsgspecRoute + - POST /steady_clock_offset _MsgspecRoute + - POST /kv_cache_events _MsgspecRoute + - POST /_resource_governor/truncate _MsgspecRoute + - POST /_control/kv_cache/truncate _MsgspecRoute + - POST /_control/kv_cache/truncate_tokens _MsgspecRoute + - POST /v1/completions _MsgspecRoute + - POST /v1/chat/completions _MsgspecRoute + - POST /v1/messages _MsgspecRoute + - POST /v1/messages/count_tokens _MsgspecRoute + - POST /v1/messages/batches _MsgspecRoute + - GET /v1/messages/batches _MsgspecRoute + - GET /v1/messages/batches/{batch_id} _MsgspecRoute + - DELETE /v1/messages/batches/{batch_id} _MsgspecRoute + - POST /v1/messages/batches/{batch_id}/cancel _MsgspecRoute + - GET /v1/messages/batches/{batch_id}/results _MsgspecRoute + - POST /v1/responses _MsgspecRoute + - GET /v1/responses/{response_id} _MsgspecRoute + - DELETE /v1/responses/{response_id} _MsgspecRoute + - POST /_internal/tokenize _MsgspecRoute + - POST /start_profile _MsgspecRoute + - POST /stop_profile _MsgspecRoute + - GET /server_info _MsgspecRoute + llm_rl_control: + kind: openai_server + routes: + - GET,HEAD /openapi.json Route + - GET,HEAD /docs Route + - GET,HEAD /docs/oauth2-redirect Route + - GET,HEAD /redoc Route + - GET /health _MsgspecRoute + - GET /health_generate _MsgspecRoute + - GET /version _MsgspecRoute + - GET /v1/models _MsgspecRoute + - GET /v1/data_transceiver_state _MsgspecRoute + - GET /metrics _MsgspecRoute + - GET /energy_metrics _MsgspecRoute + - GET /steady_clock_offset _MsgspecRoute + - POST /steady_clock_offset _MsgspecRoute + - POST /kv_cache_events _MsgspecRoute + - POST /v1/completions _MsgspecRoute + - POST /v1/chat/completions _MsgspecRoute + - POST /v1/messages _MsgspecRoute + - POST /v1/messages/count_tokens _MsgspecRoute + - POST /v1/messages/batches _MsgspecRoute + - GET /v1/messages/batches _MsgspecRoute + - GET /v1/messages/batches/{batch_id} _MsgspecRoute + - DELETE /v1/messages/batches/{batch_id} _MsgspecRoute + - POST /v1/messages/batches/{batch_id}/cancel _MsgspecRoute + - GET /v1/messages/batches/{batch_id}/results _MsgspecRoute + - POST /v1/responses _MsgspecRoute + - GET /v1/responses/{response_id} _MsgspecRoute + - DELETE /v1/responses/{response_id} _MsgspecRoute + - POST /_internal/tokenize _MsgspecRoute + - POST /start_profile _MsgspecRoute + - POST /stop_profile _MsgspecRoute + - POST /release_memory _MsgspecRoute + - POST /resume_memory _MsgspecRoute + - POST /update_weights _MsgspecRoute + - GET /server_info _MsgspecRoute + llm_runtime_control: + kind: openai_server + routes: + - GET,HEAD /openapi.json Route + - GET,HEAD /docs Route + - GET,HEAD /docs/oauth2-redirect Route + - GET,HEAD /redoc Route + - GET /health _MsgspecRoute + - GET /health_generate _MsgspecRoute + - GET /version _MsgspecRoute + - GET /v1/models _MsgspecRoute + - GET /v1/data_transceiver_state _MsgspecRoute + - GET /metrics _MsgspecRoute + - GET /energy_metrics _MsgspecRoute + - GET /steady_clock_offset _MsgspecRoute + - POST /steady_clock_offset _MsgspecRoute + - POST /kv_cache_events _MsgspecRoute + - POST /v1/completions _MsgspecRoute + - POST /v1/chat/completions _MsgspecRoute + - POST /v1/messages _MsgspecRoute + - POST /v1/messages/count_tokens _MsgspecRoute + - POST /v1/messages/batches _MsgspecRoute + - GET /v1/messages/batches _MsgspecRoute + - GET /v1/messages/batches/{batch_id} _MsgspecRoute + - DELETE /v1/messages/batches/{batch_id} _MsgspecRoute + - POST /v1/messages/batches/{batch_id}/cancel _MsgspecRoute + - GET /v1/messages/batches/{batch_id}/results _MsgspecRoute + - POST /v1/responses _MsgspecRoute + - GET /v1/responses/{response_id} _MsgspecRoute + - DELETE /v1/responses/{response_id} _MsgspecRoute + - POST /_internal/tokenize _MsgspecRoute + - POST /start_profile _MsgspecRoute + - POST /stop_profile _MsgspecRoute + - POST /release_memory _MsgspecRoute + - POST /resume_memory _MsgspecRoute + - GET /memory_status _MsgspecRoute + - GET /server_info _MsgspecRoute + llm_rl_and_runtime_control: + kind: openai_server + error: 'ValueError: Legacy RL control and runtime control endpoints cannot both be enabled.' + llm_perf_metrics: + kind: openai_server + routes: + - GET,HEAD /openapi.json Route + - GET,HEAD /docs Route + - GET,HEAD /docs/oauth2-redirect Route + - GET,HEAD /redoc Route + - GET /health _MsgspecRoute + - GET /health_generate _MsgspecRoute + - GET /version _MsgspecRoute + - GET /v1/models _MsgspecRoute + - GET /v1/data_transceiver_state _MsgspecRoute + - GET /metrics _MsgspecRoute + - GET /energy_metrics _MsgspecRoute + - GET /steady_clock_offset _MsgspecRoute + - POST /steady_clock_offset _MsgspecRoute + - POST /kv_cache_events _MsgspecRoute + - POST /v1/completions _MsgspecRoute + - POST /v1/chat/completions _MsgspecRoute + - POST /v1/messages _MsgspecRoute + - POST /v1/messages/count_tokens _MsgspecRoute + - POST /v1/messages/batches _MsgspecRoute + - GET /v1/messages/batches _MsgspecRoute + - GET /v1/messages/batches/{batch_id} _MsgspecRoute + - DELETE /v1/messages/batches/{batch_id} _MsgspecRoute + - POST /v1/messages/batches/{batch_id}/cancel _MsgspecRoute + - GET /v1/messages/batches/{batch_id}/results _MsgspecRoute + - POST /v1/responses _MsgspecRoute + - GET /v1/responses/{response_id} _MsgspecRoute + - DELETE /v1/responses/{response_id} _MsgspecRoute + - POST /_internal/tokenize _MsgspecRoute + - POST /start_profile _MsgspecRoute + - POST /stop_profile _MsgspecRoute + - GET /server_info _MsgspecRoute + - MOUNT /prometheus/metrics Mount + llm_perf_metrics_enable_metrics_env: + kind: openai_server + routes: + - GET,HEAD /openapi.json Route + - GET,HEAD /docs Route + - GET,HEAD /docs/oauth2-redirect Route + - GET,HEAD /redoc Route + - GET /health _MsgspecRoute + - GET /health_generate _MsgspecRoute + - GET /version _MsgspecRoute + - GET /v1/models _MsgspecRoute + - GET /v1/data_transceiver_state _MsgspecRoute + - 'GET /metrics _MsgspecRoute #get_iteration_stats' + - GET /energy_metrics _MsgspecRoute + - GET /steady_clock_offset _MsgspecRoute + - POST /steady_clock_offset _MsgspecRoute + - POST /kv_cache_events _MsgspecRoute + - POST /v1/completions _MsgspecRoute + - POST /v1/chat/completions _MsgspecRoute + - POST /v1/messages _MsgspecRoute + - POST /v1/messages/count_tokens _MsgspecRoute + - POST /v1/messages/batches _MsgspecRoute + - GET /v1/messages/batches _MsgspecRoute + - GET /v1/messages/batches/{batch_id} _MsgspecRoute + - DELETE /v1/messages/batches/{batch_id} _MsgspecRoute + - POST /v1/messages/batches/{batch_id}/cancel _MsgspecRoute + - GET /v1/messages/batches/{batch_id}/results _MsgspecRoute + - POST /v1/responses _MsgspecRoute + - GET /v1/responses/{response_id} _MsgspecRoute + - DELETE /v1/responses/{response_id} _MsgspecRoute + - POST /_internal/tokenize _MsgspecRoute + - POST /start_profile _MsgspecRoute + - POST /stop_profile _MsgspecRoute + - GET /server_info _MsgspecRoute + - 'GET /metrics _MsgspecRoute #metrics' + - MOUNT /prometheus/metrics Mount + llm_perf_metrics_output_dir_only: + kind: openai_server + routes: + - GET,HEAD /openapi.json Route + - GET,HEAD /docs Route + - GET,HEAD /docs/oauth2-redirect Route + - GET,HEAD /redoc Route + - GET /health _MsgspecRoute + - GET /health_generate _MsgspecRoute + - GET /version _MsgspecRoute + - GET /v1/models _MsgspecRoute + - GET /v1/data_transceiver_state _MsgspecRoute + - GET /metrics _MsgspecRoute + - GET /energy_metrics _MsgspecRoute + - GET /steady_clock_offset _MsgspecRoute + - POST /steady_clock_offset _MsgspecRoute + - POST /kv_cache_events _MsgspecRoute + - POST /v1/completions _MsgspecRoute + - POST /v1/chat/completions _MsgspecRoute + - POST /v1/messages _MsgspecRoute + - POST /v1/messages/count_tokens _MsgspecRoute + - POST /v1/messages/batches _MsgspecRoute + - GET /v1/messages/batches _MsgspecRoute + - GET /v1/messages/batches/{batch_id} _MsgspecRoute + - DELETE /v1/messages/batches/{batch_id} _MsgspecRoute + - POST /v1/messages/batches/{batch_id}/cancel _MsgspecRoute + - GET /v1/messages/batches/{batch_id}/results _MsgspecRoute + - POST /v1/responses _MsgspecRoute + - GET /v1/responses/{response_id} _MsgspecRoute + - DELETE /v1/responses/{response_id} _MsgspecRoute + - POST /_internal/tokenize _MsgspecRoute + - POST /start_profile _MsgspecRoute + - POST /stop_profile _MsgspecRoute + - GET /server_info _MsgspecRoute + - MOUNT /prometheus/metrics Mount + llm_render_endpoints: + kind: openai_server + routes: + - GET,HEAD /openapi.json Route + - GET,HEAD /docs Route + - GET,HEAD /docs/oauth2-redirect Route + - GET,HEAD /redoc Route + - GET /health _MsgspecRoute + - GET /health_generate _MsgspecRoute + - GET /version _MsgspecRoute + - GET /v1/models _MsgspecRoute + - GET /v1/data_transceiver_state _MsgspecRoute + - GET /metrics _MsgspecRoute + - GET /energy_metrics _MsgspecRoute + - GET /steady_clock_offset _MsgspecRoute + - POST /steady_clock_offset _MsgspecRoute + - POST /kv_cache_events _MsgspecRoute + - POST /v1/completions _MsgspecRoute + - POST /v1/chat/completions _MsgspecRoute + - POST /v1/messages _MsgspecRoute + - POST /v1/messages/count_tokens _MsgspecRoute + - POST /v1/messages/batches _MsgspecRoute + - GET /v1/messages/batches _MsgspecRoute + - GET /v1/messages/batches/{batch_id} _MsgspecRoute + - DELETE /v1/messages/batches/{batch_id} _MsgspecRoute + - POST /v1/messages/batches/{batch_id}/cancel _MsgspecRoute + - GET /v1/messages/batches/{batch_id}/results _MsgspecRoute + - POST /v1/responses _MsgspecRoute + - GET /v1/responses/{response_id} _MsgspecRoute + - DELETE /v1/responses/{response_id} _MsgspecRoute + - POST /_internal/tokenize _MsgspecRoute + - POST /start_profile _MsgspecRoute + - POST /stop_profile _MsgspecRoute + - GET /server_info _MsgspecRoute + - POST /v1/chat/completions/render _MsgspecRoute + - POST /v1/completions/render _MsgspecRoute + llm_multi_frontend_attached: + kind: openai_server + routes: + - GET,HEAD /openapi.json Route + - GET,HEAD /docs Route + - GET,HEAD /docs/oauth2-redirect Route + - GET,HEAD /redoc Route + - 'GET /metrics _MsgspecRoute #forward' + - 'POST /kv_cache_events _MsgspecRoute #forward' + - 'POST /start_profile _MsgspecRoute #forward' + - 'POST /stop_profile _MsgspecRoute #forward' + - 'POST /steady_clock_offset _MsgspecRoute #forward' + - 'GET /v1/messages/batches _MsgspecRoute #forward' + - 'POST /v1/messages/batches _MsgspecRoute #forward' + - GET /v1/messages/batches/{rest:path} _MsgspecRoute + - POST /v1/messages/batches/{rest:path} _MsgspecRoute + - DELETE /v1/messages/batches/{rest:path} _MsgspecRoute + - GET /health _MsgspecRoute + - GET /health_generate _MsgspecRoute + - GET /version _MsgspecRoute + - GET /v1/models _MsgspecRoute + - GET /v1/data_transceiver_state _MsgspecRoute + - 'GET /metrics _MsgspecRoute #get_iteration_stats' + - GET /energy_metrics _MsgspecRoute + - GET /steady_clock_offset _MsgspecRoute + - 'POST /steady_clock_offset _MsgspecRoute #set_steady_clock_offset' + - 'POST /kv_cache_events _MsgspecRoute #get_kv_cache_events' + - POST /v1/completions _MsgspecRoute + - POST /v1/chat/completions _MsgspecRoute + - POST /v1/messages _MsgspecRoute + - POST /v1/messages/count_tokens _MsgspecRoute + - 'POST /v1/messages/batches _MsgspecRoute #anthropic_create_batch' + - 'GET /v1/messages/batches _MsgspecRoute #anthropic_list_batches' + - GET /v1/messages/batches/{batch_id} _MsgspecRoute + - DELETE /v1/messages/batches/{batch_id} _MsgspecRoute + - POST /v1/messages/batches/{batch_id}/cancel _MsgspecRoute + - GET /v1/messages/batches/{batch_id}/results _MsgspecRoute + - POST /v1/responses _MsgspecRoute + - GET /v1/responses/{response_id} _MsgspecRoute + - DELETE /v1/responses/{response_id} _MsgspecRoute + - POST /_internal/tokenize _MsgspecRoute + - 'POST /start_profile _MsgspecRoute #start_profile' + - 'POST /stop_profile _MsgspecRoute #stop_profile' + - GET /server_info _MsgspecRoute + mm_encoder: + kind: openai_server + routes: + - GET,HEAD /openapi.json Route + - GET,HEAD /docs Route + - GET,HEAD /docs/oauth2-redirect Route + - GET,HEAD /redoc Route + - GET /health _MsgspecRoute + - GET /version _MsgspecRoute + - GET /v1/models _MsgspecRoute + - GET /metrics _MsgspecRoute + - POST /v1/chat/completions _MsgspecRoute + mm_encoder_rl_control: + kind: openai_server + error: 'ValueError: RL control endpoints require AsyncLLM' + embedding: + kind: openai_server + routes: + - GET,HEAD /openapi.json Route + - GET,HEAD /docs Route + - GET,HEAD /docs/oauth2-redirect Route + - GET,HEAD /redoc Route + - GET /health _MsgspecRoute + - GET /version _MsgspecRoute + - GET /v1/models _MsgspecRoute + - POST /v1/embeddings _MsgspecRoute + embedding_runtime_control: + kind: openai_server + error: 'ValueError: Runtime control endpoints are not supported for embedding servers' + disagg_local: + kind: disagg_server + routes: + - GET,HEAD /openapi.json Route + - GET,HEAD /docs Route + - GET,HEAD /docs/oauth2-redirect Route + - GET,HEAD /redoc Route + - POST /v1/completions KvHintsRoute + - POST /v1/chat/completions KvHintsRoute + - POST /v1/messages KvHintsRoute + - POST /v1/messages/count_tokens KvHintsRoute + - POST /v1/responses KvHintsRoute + - GET /v1/models KvHintsRoute + - GET /health KvHintsRoute + - GET /cluster_info KvHintsRoute + - GET /version KvHintsRoute + - MOUNT /prometheus/metrics Mount + disagg_local_perf_metrics: + kind: disagg_server + routes: + - GET,HEAD /openapi.json Route + - GET,HEAD /docs Route + - GET,HEAD /docs/oauth2-redirect Route + - GET,HEAD /redoc Route + - POST /v1/completions KvHintsRoute + - POST /v1/chat/completions KvHintsRoute + - POST /v1/messages KvHintsRoute + - POST /v1/messages/count_tokens KvHintsRoute + - POST /v1/responses KvHintsRoute + - GET /v1/models KvHintsRoute + - GET /health KvHintsRoute + - GET /cluster_info KvHintsRoute + - GET /version KvHintsRoute + - MOUNT /prometheus/metrics Mount + disagg_local_http_cluster_storage: + kind: disagg_server + routes: + - GET,HEAD /openapi.json Route + - GET,HEAD /docs Route + - GET,HEAD /docs/oauth2-redirect Route + - GET,HEAD /redoc Route + - POST /v1/completions KvHintsRoute + - POST /v1/chat/completions KvHintsRoute + - POST /v1/messages KvHintsRoute + - POST /v1/messages/count_tokens KvHintsRoute + - POST /v1/responses KvHintsRoute + - GET /v1/models KvHintsRoute + - GET /health KvHintsRoute + - GET /cluster_info KvHintsRoute + - GET /version KvHintsRoute + - MOUNT /prometheus/metrics Mount + - POST /set KvHintsRoute + - GET /get KvHintsRoute + - DELETE /delete KvHintsRoute + - GET /expire KvHintsRoute + - GET /get_prefix KvHintsRoute + disagg_fleet: + kind: disagg_server + routes: + - GET,HEAD /openapi.json Route + - GET,HEAD /docs Route + - GET,HEAD /docs/oauth2-redirect Route + - GET,HEAD /redoc Route + - POST /v1/completions KvHintsRoute + - POST /v1/chat/completions KvHintsRoute + - POST /v1/messages KvHintsRoute + - POST /v1/messages/count_tokens KvHintsRoute + - POST /v1/responses KvHintsRoute + - GET /v1/models KvHintsRoute + - GET /health KvHintsRoute + - GET /cluster_info KvHintsRoute + - GET /version KvHintsRoute + - MOUNT /prometheus/metrics Mount + disagg_fleet_perf_metrics: + kind: disagg_server + routes: + - GET,HEAD /openapi.json Route + - GET,HEAD /docs Route + - GET,HEAD /docs/oauth2-redirect Route + - GET,HEAD /redoc Route + - POST /v1/completions KvHintsRoute + - POST /v1/chat/completions KvHintsRoute + - POST /v1/messages KvHintsRoute + - POST /v1/messages/count_tokens KvHintsRoute + - POST /v1/responses KvHintsRoute + - GET /v1/models KvHintsRoute + - GET /health KvHintsRoute + - GET /cluster_info KvHintsRoute + - GET /version KvHintsRoute + - MOUNT /prometheus/metrics Mount + disagg_fleet_http_cluster_storage: + kind: disagg_server + routes: + - GET,HEAD /openapi.json Route + - GET,HEAD /docs Route + - GET,HEAD /docs/oauth2-redirect Route + - GET,HEAD /redoc Route + - POST /v1/completions KvHintsRoute + - POST /v1/chat/completions KvHintsRoute + - POST /v1/messages KvHintsRoute + - POST /v1/messages/count_tokens KvHintsRoute + - POST /v1/responses KvHintsRoute + - GET /v1/models KvHintsRoute + - GET /health KvHintsRoute + - GET /cluster_info KvHintsRoute + - GET /version KvHintsRoute + - MOUNT /prometheus/metrics Mount + - POST /set KvHintsRoute + - GET /get KvHintsRoute + - DELETE /delete KvHintsRoute + - GET /expire KvHintsRoute + - GET /get_prefix KvHintsRoute + coordinator: + kind: coordinator_server + routes: + - GET,HEAD /openapi.json Route + - GET,HEAD /docs Route + - GET,HEAD /docs/oauth2-redirect Route + - GET,HEAD /redoc Route + - POST /select APIRoute + - POST /finish APIRoute + - POST /renew APIRoute + - GET /cluster_info APIRoute + - GET /health APIRoute + - GET /version APIRoute + coordinator_http_cluster_storage: + kind: coordinator_server + routes: + - GET,HEAD /openapi.json Route + - GET,HEAD /docs Route + - GET,HEAD /docs/oauth2-redirect Route + - GET,HEAD /redoc Route + - POST /select APIRoute + - POST /finish APIRoute + - POST /renew APIRoute + - GET /cluster_info APIRoute + - GET /health APIRoute + - GET /version APIRoute + - POST /set APIRoute + - GET /get APIRoute + - DELETE /delete APIRoute + - GET /expire APIRoute + - GET /get_prefix APIRoute + visual_gen: + kind: openai_server + routes: + - GET,HEAD /openapi.json Route + - GET,HEAD /docs Route + - GET,HEAD /docs/oauth2-redirect Route + - GET,HEAD /redoc Route + - GET /health _MsgspecRoute + - GET /version _MsgspecRoute + - GET /v1/models _MsgspecRoute + - GET /metrics _MsgspecRoute + - POST /v1/images/generations _MsgspecRoute + - POST /v1/images/edits _MsgspecRoute + - GET /v1/images/{image_id}/content _MsgspecRoute + - POST /v1/videos _MsgspecRoute + - POST /v1/videos/sync _MsgspecRoute + - POST /v1/videos/generations _MsgspecRoute + - GET /v1/videos _MsgspecRoute + - GET /v1/videos/{video_id} _MsgspecRoute + - GET /v1/videos/{video_id}/content _MsgspecRoute + - DELETE /v1/videos/{video_id} _MsgspecRoute diff --git a/tests/unittest/api_stability/test_serve_routes.py b/tests/unittest/api_stability/test_serve_routes.py new file mode 100644 index 000000000000..1392c1428ec3 --- /dev/null +++ b/tests/unittest/api_stability/test_serve_routes.py @@ -0,0 +1,251 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Route snapshot of every trtllm-serve configuration: a new api-stability gate. + +Each configuration's live routes must equal the generated reference +``references/trtllm_serve_routes.yaml``. Regenerating the reference cannot hide +the invariants at the bottom of this file: they are checked on the live apps. +The VisualGen configuration is checked by +``tests/unittest/serve/api/visual_gen/test_routes_visual_gen.py`` so that it +runs on VisualGen-only pull requests. +""" + +import pytest +from test_common.serve.routes import ( + CONFIGS, + VISUAL_GEN_CONFIGS, + compare, + format_failure, + read_reference, + snapshot_config, +) + +pytestmark = pytest.mark.cpu_only + + +@pytest.fixture(scope="module") +def reference(): + return read_reference() + + +@pytest.fixture(scope="module") +def live_snapshots(tmp_path_factory): + root = tmp_path_factory.mktemp("serve_routes") + return {config.name: snapshot_config(config, root / config.name) for config in CONFIGS} + + +def test_reference_configuration_names(reference): + expected = [config.name for config in CONFIGS + VISUAL_GEN_CONFIGS] + assert list(reference["configs"].keys()) == expected, ( + "The configuration names in references/trtllm_serve_routes.yaml must equal the " + "names in CONFIGS plus VISUAL_GEN_CONFIGS of tests/test_common/serve/routes.py, in " + "that order. Regenerate with `PYTHONPATH=tests python3 -m test_common.serve.routes " + "--write`." + ) + + +@pytest.mark.parametrize("name", [config.name for config in CONFIGS]) +def test_routes_match_reference(name, reference, live_snapshots): + recorded = reference["configs"][name] + live = live_snapshots[name] + assert recorded.get("error") == live.get("error"), ( + f"Configuration '{name}' now fails with {live.get('error')!r} instead of " + f"{recorded.get('error')!r}. If the change is intended, regenerate with " + "`PYTHONPATH=tests python3 -m test_common.serve.routes --write` and request review " + "from @NVIDIA/trt-llm-noncommitted-api-review-committee." + ) + differences = compare(recorded.get("routes", []), live.get("routes", [])) + assert not differences, format_failure(name, differences) + + +# --------------------------------------------------------------------------- +# Invariants checked on the live apps, independent of the reference. +# Regenerating the reference cannot hide these. +# --------------------------------------------------------------------------- + +_DOC_ROUTES = {"/openapi.json", "/docs", "/docs/oauth2-redirect", "/redoc"} + + +def _routes_of(live_snapshots, name): + return live_snapshots[name]["routes"] + + +def _parse(record): + parts = record.split(" ") + return parts[0], parts[1], parts[2] + + +def test_openai_server_routes_are_msgspec(live_snapshots): + for name, snapshot in live_snapshots.items(): + if snapshot.get("error") or snapshot["kind"] != "openai_server": + continue + for record in snapshot["routes"]: + methods, path, cls = _parse(record) + if path in _DOC_ROUTES: + assert cls == "Route", record + elif methods == "MOUNT": + assert cls == "Mount", record + else: + assert cls == "_MsgspecRoute", ( + f"configuration '{name}': {record}: every API route of the aggregated " + "server must keep the msgpack route class, or workers reject the " + "disaggregated frontend's msgpack bodies with 400/422 while plain JSON " + "still works." + ) + + +_API_ROUTE_CLASSES = { + "disagg_server": ( + "KvHintsRoute", + "every API route of the disaggregated frontend must keep KvHintsRoute: it strips " + "kv_hints from request bodies unless the server trusts its ingress " + "(TRTLLM_TRUST_KV_HINTS=1), and a plain APIRoute skips that policy.", + ), + "coordinator_server": ( + "APIRoute", + "the coordinator app registers FastAPI's plain routes; it takes neither the " + "disaggregated frontend's KvHintsRoute nor the aggregated server's msgpack route " + "class.", + ), +} + + +def test_disagg_and_coordinator_route_classes(live_snapshots): + for name, snapshot in live_snapshots.items(): + if snapshot.get("error") or snapshot["kind"] == "openai_server": + continue + expected, reason = _API_ROUTE_CLASSES[snapshot["kind"]] + for record in snapshot["routes"]: + methods, path, cls = _parse(record) + if path in _DOC_ROUTES: + assert cls == "Route", record + elif methods == "MOUNT": + assert cls == "Mount", record + else: + assert cls == expected, f"configuration '{name}': {record}: {reason}" + + +def test_control_route_groups(live_snapshots): + def paths(name): + return {_parse(record)[1] for record in _routes_of(live_snapshots, name)} + + rl = paths("llm_rl_control") + assert {"/release_memory", "/resume_memory", "/update_weights"} <= rl + assert "/memory_status" not in rl, ( + "the RL control group has no /memory_status; the runtime control group owns it" + ) + runtime = paths("llm_runtime_control") + assert {"/release_memory", "/resume_memory", "/memory_status"} <= runtime + assert "/update_weights" not in runtime, ( + "the runtime control group has no /update_weights; the RL control group owns it" + ) + control_paths = {"/release_memory", "/resume_memory", "/update_weights", "/memory_status"} + for name in ("llm", "llm_context", "llm_generation", "mm_encoder", "embedding"): + assert not control_paths & paths(name), ( + f"configuration '{name}' must register no memory-control route: those routes " + "exist only behind their construction-time switches." + ) + + +def test_metrics_presence_per_role(live_snapshots): + embedding_paths = {_parse(r)[1] for r in _routes_of(live_snapshots, "embedding")} + assert "/metrics" not in embedding_paths, ( + "the embedding server registers no /metrics today; keep the asymmetry until a " + "change deliberately adds it." + ) + mm_paths = {_parse(r)[1] for r in _routes_of(live_snapshots, "mm_encoder")} + assert "/metrics" in mm_paths + + +def test_launcher_forwarding_shadows_local_routes(live_snapshots): + """An attached frontend's forwarding routes precede the local handlers. + + Launcher-owned routes are registered before the local table with the same + method and path (or a path parameter that matches the same requests), and + FastAPI matches in registration order: flipping the order silently serves + launcher-owned state from the local process instead. + """ + records = _routes_of(live_snapshots, "llm_multi_frontend_attached") + base = _routes_of(live_snapshots, "llm") + assert len(records) == len(base) + 10, ( + "the attached-frontend configuration registers the ten launcher " + "forwarding routes on top of the default table" + ) + forward_positions = [i for i, r in enumerate(records) if r.endswith("#forward")] + assert len(forward_positions) == 7, records + for position in forward_positions: + methods_path = " ".join(records[position].split(" ")[:2]) + shadowed = [ + j + for j, record in enumerate(records) + if j != position and " ".join(record.split(" ")[:2]) == methods_path + ] + assert shadowed and all(position < j for j in shadowed), ( + f"{records[position]} must stay registered before the local handler of {methods_path}" + ) + rest_positions = [i for i, r in enumerate(records) if "{rest:path}" in _parse(r)[1]] + batch_id_positions = [i for i, r in enumerate(records) if "{batch_id}" in _parse(r)[1]] + assert rest_positions and batch_id_positions + assert max(rest_positions) < min(batch_id_positions), ( + "the batches forwarding routes (a path parameter that matches every " + "batch id) must stay registered before the local batch-id handlers" + ) + + +def test_metrics_shadowing_order(live_snapshots): + records = _routes_of(live_snapshots, "llm_perf_metrics_enable_metrics_env") + metrics = [r for r in records if _parse(r)[1] == "/metrics" and _parse(r)[0] == "GET"] + assert len(metrics) == 2, metrics + assert metrics[0].endswith("#get_iteration_stats") and metrics[1].endswith("#metrics"), ( + "the JSON /metrics route must stay registered before the Prometheus-text /metrics " + "route: FastAPI matches in registration order, so swapping them silently changes " + "the /metrics payload when ENABLE_METRICS is set." + ) + + +# The only configurations that fail at construction, with their exact errors. +_ERROR_CONFIGURATIONS = { + "llm_rl_and_runtime_control": ( + "ValueError: Legacy RL control and runtime control endpoints cannot both be enabled." + ), + "mm_encoder_rl_control": "ValueError: RL control endpoints require AsyncLLM", + "embedding_runtime_control": ( + "ValueError: Runtime control endpoints are not supported for embedding servers" + ), +} + + +def test_valid_configurations_build(live_snapshots): + failing = { + name: snapshot.get("error") + for name, snapshot in live_snapshots.items() + if name not in _ERROR_CONFIGURATIONS and "routes" not in snapshot + } + assert not failing, ( + f"these configurations must build but fail at construction: {failing}. Only the " + "configurations in _ERROR_CONFIGURATIONS may fail, and a regenerated reference that " + "records the error does not change that." + ) + + +def test_error_configurations(tmp_path): + by_name = {config.name: config for config in CONFIGS} + for name, message in _ERROR_CONFIGURATIONS.items(): + snapshot = snapshot_config(by_name[name], tmp_path / name) + assert snapshot.get("error") == message, ( + f"configuration '{name}' must fail at construction with exactly " + f"{message!r}, got {snapshot.get('error')!r}: the conflicting route groups are " + "rejected before any route is registered." + ) diff --git a/tests/unittest/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..f691dd23c5a8 --- /dev/null +++ b/tests/unittest/api_stability/visual_gen/test_visual_gen_serve_api.py @@ -0,0 +1,219 @@ +# 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, 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 +``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 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", + "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`` + 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() + models: dict[str, Any] = {} + for model_name, model_spec in (full.get("models", {}) or {}).items(): + 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, ( + 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, ( + 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]) + + +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, warnings + import tensorrt_llm.serve.openai_protocol as protocol + + target = {_VISUAL_GEN_PROTOCOL_MODULE!r} + 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: + 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']}{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 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: + 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/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/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/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..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", {}), ) @@ -610,7 +606,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..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]) @@ -176,9 +174,10 @@ def on_term(signum, frame): self.llm_id = "shared-engine" self.args = SimpleNamespace(guided_decoding_backend=None) self.tokenizer = None - self._executor = Mock(spec=GenerationExecutorProxy) - self._executor.multi_frontend_attach_info.return_value = {"mode": "classic"} - def _check_health(self): + self.engine_dead = False + 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" and (root / "initializing").exists() diff --git a/tests/unittest/grpc/openengine/test_metrics.py b/tests/unittest/grpc/openengine/test_metrics.py index 3e97cc47d4b5..d7faf31cca05 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)] @@ -151,7 +152,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): @@ -607,23 +608,31 @@ 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 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 - assert ( - any(route.path == "/prometheus/metrics" for route in server.app.routes) - == enabled - ) + # Match a request instead of reading .path off app.routes: FastAPI + # 0.139 keeps included routers there as objects without a path. + scope = { + "type": "http", + "method": "GET", + "path": "/prometheus/metrics", + "path_params": {}, + } + exposed = any(route.matches(scope)[0] == Match.FULL for route in server.app.routes) + assert exposed == enabled assert server._expose_perf_metrics is requested assert server._collect_perf_metrics is enabled assert args.return_perf_metrics is enabled diff --git a/tests/unittest/grpc/openengine/test_multimodal.py b/tests/unittest/grpc/openengine/test_multimodal.py index 864532a3a100..52ede9ed4767 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), ) @@ -423,7 +423,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), ) @@ -471,7 +471,7 @@ async def async_load_bytes(self, data): @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 @@ -768,7 +768,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() @@ -1071,7 +1071,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, @@ -1154,7 +1154,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, @@ -1162,11 +1162,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, @@ -1474,7 +1474,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: @@ -1512,7 +1512,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", @@ -1826,7 +1826,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: @@ -1870,7 +1870,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])) @@ -1959,7 +1959,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) @@ -2132,7 +2132,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 @@ -2178,7 +2178,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() @@ -2251,7 +2251,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() @@ -2326,7 +2326,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 @@ -2381,7 +2381,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") @@ -2431,7 +2431,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 3db03fdad5d9..4ad08de3ed6e 100644 --- a/tests/unittest/grpc/openengine/test_openengine.py +++ b/tests/unittest/grpc/openengine/test_openengine.py @@ -147,6 +147,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, @@ -156,11 +157,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, + check_health=lambda: True, + get_kv_cache_capacity=lambda: {}, + get_openengine_discovery=lambda: {}, ) async def exercise_server() -> None: @@ -207,10 +211,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: @@ -241,10 +246,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: @@ -311,10 +317,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: @@ -366,10 +373,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: @@ -420,13 +428,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") @@ -497,10 +514,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: @@ -551,16 +569,14 @@ 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( 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/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/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 5485db3d4342..02ec6c690edc 100644 --- a/tests/unittest/llmapi/apps/test_chat_utils.py +++ b/tests/unittest/llmapi/apps/test_chat_utils.py @@ -26,6 +26,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_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 2d6fc174a199..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. @@ -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/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_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/llmapi/apps/test_tool_parsers.py b/tests/unittest/llmapi/apps/test_tool_parsers.py index 64dee6ac0d5a..b1a3519089a7 100644 --- a/tests/unittest/llmapi/apps/test_tool_parsers.py +++ b/tests/unittest/llmapi/apps/test_tool_parsers.py @@ -6392,6 +6392,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) @@ -6417,6 +6418,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_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 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..6175d4a589ef --- /dev/null +++ b/tests/unittest/llmapi/test_llm_public_runtime_api.py @@ -0,0 +1,410 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""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 GenerationExecutorFrontendProxy, GenerationExecutorProxy +from tensorrt_llm.llmapi.engine_client import EngineClient +from tensorrt_llm.llmapi.llm import LLM, EncodeLimits + +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.workers_started = False + 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=_NoOpExecutor()).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=_NoOpExecutor()).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 + + +# --------------------------------------------------------------------------- +# 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..94d9b55cadcd 100644 --- a/tests/unittest/llmapi/test_serve_multi_frontend_helpers.py +++ b/tests/unittest/llmapi/test_serve_multi_frontend_helpers.py @@ -35,7 +35,11 @@ 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.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, AttachedFrontendWatchdog, @@ -103,7 +107,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 @@ -111,9 +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 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) + 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 @@ -241,12 +255,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 +281,27 @@ 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.""" + # 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) + + OpenAIServer._on_launcher_lost(server, 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..d3df07c17277 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,12 @@ 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.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 @@ -65,7 +70,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 @@ -73,10 +78,23 @@ def _llm_server_routes() -> list[tuple[str, str]]: server._enable_rl_control_endpoints = True server._enable_runtime_control_endpoints = True server.resource_governor = None + # As OpenAIServer.__init__ does: the 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) + 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( - (method, route.path) - for route in server.app.routes + (method, path) + for path, route in effective_routes(server.app) if getattr(route, "methods", None) for method in route.methods if method != "HEAD" diff --git a/tests/unittest/serve/__init__.py b/tests/unittest/serve/__init__.py new file mode 100644 index 000000000000..467079831e16 --- /dev/null +++ b/tests/unittest/serve/__init__.py @@ -0,0 +1,14 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/tests/unittest/serve/_base/__init__.py b/tests/unittest/serve/_base/__init__.py new file mode 100644 index 000000000000..467079831e16 --- /dev/null +++ b/tests/unittest/serve/_base/__init__.py @@ -0,0 +1,14 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/tests/unittest/serve/_base/test_routing.py b/tests/unittest/serve/_base/test_routing.py new file mode 100644 index 000000000000..37fbd7fec2c8 --- /dev/null +++ b/tests/unittest/serve/_base/test_routing.py @@ -0,0 +1,264 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""``make_router()``, ``attach_routers()`` and the role mapping of ``serve/_base/routing.py``.""" + +import ast +import importlib +import sys +import textwrap +from pathlib import Path + +import pytest +from fastapi import Depends, FastAPI, Request +from fastapi.routing import APIRoute +from fastapi.testclient import TestClient +from test_common.serve.app_factory import AggConfig +from test_common.serve.routes import effective_routes + +from tensorrt_llm.llmapi.disagg_utils import ServerRole +from tensorrt_llm.serve._base.msgspec_route import _MsgspecRoute +from tensorrt_llm.serve._base.routing import ( + ROUTERS, + RoleGroup, + RouterSpec, + ServerCapabilities, + attach_routers, + make_router, + role_group_of, +) +from tensorrt_llm.serve.openai_client import MSGPACK_HEADERS, _msgpack_encoder + +pytestmark = pytest.mark.cpu_only + + +def test_make_router_uses_msgspec_route(): + router = make_router() + assert router.route_class is _MsgspecRoute + + @router.get("/x") + async def x(): + return {} + + assert len(router.routes) == 1 + assert isinstance(router.routes[0], _MsgspecRoute) + + +def test_make_router_rejects_route_class(): + with pytest.raises(TypeError, match="_MsgspecRoute"): + make_router(route_class=APIRoute) + + +def test_role_group_of_every_role(): + expected = { + None: RoleGroup.LLM, + ServerRole.CONTEXT: RoleGroup.LLM, + ServerRole.GENERATION: RoleGroup.LLM, + ServerRole.MM_ENCODER: RoleGroup.MM_ENCODER, + ServerRole.EMBEDDING: RoleGroup.EMBEDDING, + ServerRole.VISUAL_GEN: RoleGroup.VISUAL_GEN, + } + # Every ServerRole member must have a row here: a new role must decide its + # role group (and its routes) explicitly. + assert set(expected) - {None} == set(ServerRole) + for role, group in expected.items(): + assert role_group_of(role) is group + caps = ServerCapabilities.for_role(role) + assert caps.role is role + assert caps.role_group is group + + +_ROUTER_MODULE_TEMPLATE = """ +from fastapi import FastAPI + +from tensorrt_llm.serve._base.routing import make_router + +router = make_router() + + +@router.get("{path}") +async def handler(): + return {{}} + + +def attach_router(app: FastAPI) -> None: + app.include_router(router) +""" + + +def _write_router_package(tmp_path, package: str, modules: dict) -> None: + pkg = tmp_path / package + pkg.mkdir() + (pkg / "__init__.py").write_text("") + for name, path in modules.items(): + (pkg / f"{name}.py").write_text(_ROUTER_MODULE_TEMPLATE.format(path=path)) + + +def test_attach_routers_imports_only_matching_specs(tmp_path, monkeypatch): + _write_router_package( + tmp_path, + "attachpkg", + {"llm_router": "/llm-only", "mm_router": "/mm-only", "gated_router": "/gated"}, + ) + monkeypatch.syspath_prepend(str(tmp_path)) + module_names = ("attachpkg.llm_router", "attachpkg.mm_router", "attachpkg.gated_router") + specs = ( + RouterSpec("attachpkg.llm_router", roles=frozenset({RoleGroup.LLM})), + RouterSpec("attachpkg.mm_router", roles=frozenset({RoleGroup.MM_ENCODER})), + RouterSpec( + "attachpkg.gated_router", + roles=frozenset({RoleGroup.LLM}), + when=lambda caps: False, + ), + ) + try: + app = FastAPI() + attach_routers(app, ServerCapabilities.for_role(None), routers=specs) + assert "attachpkg.llm_router" in sys.modules + assert "attachpkg.mm_router" not in sys.modules, "role mismatch must not be imported" + assert "attachpkg.gated_router" not in sys.modules, "a false `when` must not be imported" + paths = [path for path, _ in effective_routes(app)] + assert "/llm-only" in paths + assert "/mm-only" not in paths + assert "/gated" not in paths + finally: + for name in ("attachpkg", *module_names): + sys.modules.pop(name, None) + + +def test_router_routes_decode_msgpack(): + router = make_router() + + @router.post("/echo") + async def echo(raw_request: Request): + return {"body": await raw_request.json()} + + app = FastAPI() + app.include_router(router) + body = _msgpack_encoder.encode({"prompt": [1, 2, 3]}) + with TestClient(app) as client: + response = client.post("/echo", content=body, headers=MSGPACK_HEADERS) + assert response.status_code == 200 + assert response.json() == {"body": {"prompt": [1, 2, 3]}} + + +def _route_path(raw_request: Request) -> str: + return raw_request.scope["route"].path + + +def test_include_keeps_route_path(): + # A path prefix goes to make_router(prefix=...), never include_router(prefix=...): + # FastAPI 0.139 drops include prefixes from scope["route"].path. + router = make_router(prefix="/p") + + @router.get("/x") + async def x(route_path: str = Depends(_route_path)): + return {"route_path": route_path} + + app = FastAPI() + app.include_router(router) + with TestClient(app) as client: + response = client.get("/p/x") + assert response.status_code == 200 + assert response.json() == {"route_path": "/p/x"} + + +def test_every_router_spec_module_exposes_attach_router(): + # ROUTERS holds module paths as strings, so a renamed router module would + # break attachment only at server startup; import each one here instead. + assert ROUTERS, "ROUTERS must list at least the system router" + for spec in ROUTERS: + module = importlib.import_module(spec.module) + assert callable(getattr(module, "attach_router", None)), ( + f"{spec.module} must define attach_router(app)" + ) + tree = ast.parse(Path(module.__file__).read_text()) + for node in ast.walk(tree): + is_include_call = ( + isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and node.func.attr == "include_router" + ) + if is_include_call: + assert all(keyword.arg != "prefix" for keyword in node.keywords), ( + f"{spec.module} passes prefix= to include_router; give the prefix to " + "make_router(prefix=...) instead: FastAPI 0.139 drops include prefixes " + "from scope['route'].path." + ) + + +_ECHO_ROUTER_MODULE = textwrap.dedent( + ''' + """A new endpoint is a router plus a serving class: the router reads its + serving object from app.state; the composition root sets it and attaches.""" + + from fastapi import FastAPI, Request + + from tensorrt_llm.serve._base.routing import make_router + + router = make_router() + + + @router.post("/v1/test-echo") + async def echo(raw_request: Request): + serving = raw_request.app.state.test_serving_echo + return serving.render(await raw_request.json()) + + + def attach_router(app: FastAPI) -> None: + if not hasattr(app.state, "test_serving_echo"): + raise RuntimeError("app.state.test_serving_echo must be set before attaching") + app.include_router(router) + ''' +) + + +class _EchoServing: + def render(self, body): + return {"echo": body} + + +# The TestClient portal thread and the fake server's lazy worker-pool threads +# are alive when the leak check runs; the serve_env fixture closes them in its +# teardown (precedent: tests/unittest/serve/api/ops/system/test_system_endpoints.py). +@pytest.mark.threadleak(enabled=False) +def test_new_endpoint_lands_as_router_and_serving_class(tmp_path, monkeypatch, serve_env): + pkg = tmp_path / "newendpointpkg" + pkg.mkdir() + (pkg / "__init__.py").write_text("") + (pkg / "api_router.py").write_text(_ECHO_ROUTER_MODULE) + monkeypatch.syspath_prepend(str(tmp_path)) + spec = RouterSpec("newendpointpkg.api_router", roles=frozenset({RoleGroup.LLM})) + body = _msgpack_encoder.encode({"prompt": "hi"}) + try: + server = serve_env.build_openai_server(AggConfig()) + server.app.state.test_serving_echo = _EchoServing() + attach_routers(server.app, ServerCapabilities.for_role(None), routers=(spec,)) + client = serve_env.client(server.app) + response = client.post("/v1/test-echo", content=body, headers=MSGPACK_HEADERS) + assert response.status_code == 200 + assert response.json() == {"echo": {"prompt": "hi"}} + + # The same spec on a role outside its role set: not imported again, not + # attached, and the route is absent. + mm_server = serve_env.build_openai_server(AggConfig(role=ServerRole.MM_ENCODER)) + attach_routers( + mm_server.app, ServerCapabilities.for_role(ServerRole.MM_ENCODER), routers=(spec,) + ) + mm_client = serve_env.client(mm_server.app) + mm_response = mm_client.post("/v1/test-echo", content=body, headers=MSGPACK_HEADERS) + assert mm_response.status_code == 404 + finally: + for name in ("newendpointpkg", "newendpointpkg.api_router"): + sys.modules.pop(name, None) diff --git a/tests/unittest/serve/_base/test_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/__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/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/api/ops/__init__.py b/tests/unittest/serve/api/ops/__init__.py new file mode 100644 index 000000000000..467079831e16 --- /dev/null +++ b/tests/unittest/serve/api/ops/__init__.py @@ -0,0 +1,14 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/tests/unittest/serve/api/ops/system/__init__.py b/tests/unittest/serve/api/ops/system/__init__.py new file mode 100644 index 000000000000..467079831e16 --- /dev/null +++ b/tests/unittest/serve/api/ops/system/__init__.py @@ -0,0 +1,14 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/tests/unittest/serve/api/ops/system/test_system_endpoints.py b/tests/unittest/serve/api/ops/system/test_system_endpoints.py new file mode 100644 index 000000000000..12a88b166e90 --- /dev/null +++ b/tests/unittest/serve/api/ops/system/test_system_endpoints.py @@ -0,0 +1,132 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""/health, /version and /v1/models behavior, through the production handlers. + +These pin what the system endpoints do today, per serving role, so a later +change can show it changed nothing. An unhealthy engine makes ``/health`` raise +``SIGINT`` (after recording a termination observation) so the process does not +linger as a zombie; the signal is patched out here before any unhealthy request. +""" + +import signal + +import pytest +from test_common.serve.app_factory import AggConfig + +import tensorrt_llm.version +from tensorrt_llm.llmapi.disagg_utils import ServerRole + +# threadleak is off: these tests create TestClients and run requests inside the +# test body, so the client's portal thread and the fake server's worker-pool +# threads are still alive when the leak check runs; the serve_env fixture closes +# them in its teardown, which runs after that check (precedent: +# tests/unittest/llmapi/apps/test_serve_lazy_import.py). +pytestmark = [pytest.mark.cpu_only, pytest.mark.threadleak(enabled=False)] + +UNAVAILABLE_BODY = "LLM is unavailable. Please check the server logs for more details." + + +@pytest.fixture +def sigint_recorder(monkeypatch): + calls = [] + monkeypatch.setattr(signal, "raise_signal", calls.append) + return calls + + +@pytest.mark.parametrize( + "role", [None, ServerRole.MM_ENCODER, ServerRole.EMBEDDING], ids=["llm", "mm", "embedding"] +) +def test_health_returns_200_with_empty_body(role, serve_env): + server = serve_env.build_openai_server(AggConfig(role=role)) + response = serve_env.client(server.app).get("/health") + assert response.status_code == 200 + assert response.content == b"" + + +@pytest.mark.parametrize("role", [None, ServerRole.MM_ENCODER], ids=["llm", "mm"]) +def test_unhealthy_executor_returns_503_without_shutdown(role, serve_env, sigint_recorder): + server = serve_env.build_openai_server(AggConfig(role=role)) + server.generator._executor.healthy = False + response = serve_env.client(server.app).get("/health") + assert response.status_code == 503 + assert response.text == UNAVAILABLE_BODY + assert sigint_recorder == [], "an unhealthy executor without a fatal error must not SIGINT" + + +@pytest.mark.parametrize("role", [None, ServerRole.MM_ENCODER], ids=["llm", "mm"]) +def test_fatal_error_triggers_sigint_and_termination_record(role, serve_env, sigint_recorder): + server = serve_env.build_openai_server(AggConfig(role=role)) + client = serve_env.client(server.app) + seen = serve_env.record_terminations() + server.generator._executor._fatal_error = RuntimeError("boom") + server.generator._executor.doing_shutdown = False + response = client.get("/health") + assert response.status_code == 503 + assert sigint_recorder == [signal.SIGINT] + assert len(seen) == 1 + assert seen[0].termination_kind == "exception" + assert seen[0].component == "llm" + assert seen[0].reporting_source == "self" + + +@pytest.mark.parametrize("role", [None, ServerRole.MM_ENCODER], ids=["llm", "mm"]) +def test_fatal_error_during_shutdown_signals_once(role, serve_env, sigint_recorder): + server = serve_env.build_openai_server(AggConfig(role=role)) + client = serve_env.client(server.app) + seen = serve_env.record_terminations() + server.generator._executor._fatal_error = RuntimeError("boom") + server.generator._executor.doing_shutdown = True + response = client.get("/health") + assert response.status_code == 503 + assert sigint_recorder == [], "a teardown already in progress must not stack signals" + assert seen == [] + + +def test_embedding_health_depends_on_batcher_and_encoder(serve_env, sigint_recorder): + server = serve_env.build_openai_server(AggConfig(role=ServerRole.EMBEDDING)) + # Without the lifespan the encode batcher never started, so every + # /v1/embeddings request would hang; /health reports unhealthy. + without_lifespan = serve_env.client(server.app, lifespan=False) + assert without_lifespan.get("/health").status_code == 503 + with_lifespan = serve_env.client(server.app) + assert with_lifespan.get("/health").status_code == 200 + # Encode-only generators have no generation executor; health reads the + # encoder executor instead. + server.generator._encoder_executor = None + assert with_lifespan.get("/health").status_code == 503 + assert sigint_recorder == [], ( + "the embedding health path finds no generation executor, so it must not SIGINT" + ) + + +def test_version_reports_package_version(serve_env): + server = serve_env.build_openai_server(AggConfig()) + response = serve_env.client(server.app).get("/version") + assert response.status_code == 200 + assert response.json() == {"version": tensorrt_llm.version.__version__} + + +def test_models_lists_the_served_model(serve_env): + server = serve_env.build_openai_server(AggConfig()) + response = serve_env.client(server.app).get("/v1/models") + assert response.status_code == 200 + body = response.json() + assert body["object"] == "list" + assert len(body["data"]) == 1 + entry = body["data"][0] + assert entry["id"] == "fake-model" + assert entry["object"] == "model" + assert entry["owned_by"] == "tensorrt_llm" + assert isinstance(entry["created"], int) # the request time; never compare the value diff --git a/tests/unittest/serve/api/ops/system/test_system_router.py b/tests/unittest/serve/api/ops/system/test_system_router.py new file mode 100644 index 000000000000..36b680d1e7c2 --- /dev/null +++ b/tests/unittest/serve/api/ops/system/test_system_router.py @@ -0,0 +1,65 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""The ops.system router: attachment contract and route provenance per role.""" + +import pytest +from fastapi import FastAPI +from test_common.serve.app_factory import AggConfig +from test_common.serve.routes import effective_routes + +from tensorrt_llm.llmapi.disagg_utils import ServerRole +from tensorrt_llm.serve._base.msgspec_route import _MsgspecRoute +from tensorrt_llm.serve.api.ops.system.api_router import attach_router +from tensorrt_llm.serve.api.ops.system.serving import SystemServing + +pytestmark = pytest.mark.cpu_only + +SYSTEM_PATHS = ("/health", "/version", "/v1/models") +SYSTEM_ROUTER_MODULE = "tensorrt_llm.serve.api.ops.system.api_router" + + +def test_attach_requires_serving_state(): + # The composition root must set the serving object before attaching, so a + # wiring bug fails at startup, not with an AttributeError per request. + with pytest.raises(RuntimeError, match="ops_serving_system"): + attach_router(FastAPI()) + + +def test_system_serving_requires_arguments(): + # No defaults: a serving class built without its collaborators must fail at + # construction, not fall back silently. + with pytest.raises(TypeError): + SystemServing() + + +@pytest.mark.parametrize( + "role", + [None, ServerRole.MM_ENCODER, ServerRole.EMBEDDING], + ids=["llm", "mm_encoder", "embedding"], +) +def test_routes_come_from_ops_system(role, serve_env): + server = serve_env.build_openai_server(AggConfig(role=role)) + by_path = {path: [] for path in SYSTEM_PATHS} + for path, route in effective_routes(server.app): + if path in by_path: + by_path[path].append(route) + for path, routes in by_path.items(): + assert len(routes) == 1, f"{path} must be registered exactly once, found {len(routes)}" + route = routes[0] + assert isinstance(route, _MsgspecRoute), ( + f"{path} must be a _MsgspecRoute; a plain APIRoute breaks msgpack request " + "bodies from the disaggregated frontend" + ) + assert route.endpoint.__module__ == SYSTEM_ROUTER_MODULE diff --git a/tests/unittest/serve/api/visual_gen/__init__.py b/tests/unittest/serve/api/visual_gen/__init__.py new file mode 100644 index 000000000000..467079831e16 --- /dev/null +++ b/tests/unittest/serve/api/visual_gen/__init__.py @@ -0,0 +1,14 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/tests/unittest/serve/api/visual_gen/test_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/api/visual_gen/test_ops_system_router_visual_gen.py b/tests/unittest/serve/api/visual_gen/test_ops_system_router_visual_gen.py new file mode 100644 index 000000000000..488a5e621d05 --- /dev/null +++ b/tests/unittest/serve/api/visual_gen/test_ops_system_router_visual_gen.py @@ -0,0 +1,48 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Route provenance of the system endpoints for the VisualGen serving role. + +Lives under a ``visual_gen/`` path so that VisualGen-only pull requests, whose +CI keeps only test entries containing ``visual_gen/``, still run it. +""" + +import pytest +from test_common.serve.app_factory import isolated_serve_env +from test_common.serve.routes import VISUAL_GEN_CONFIGS, effective_routes + +from tensorrt_llm.serve._base.msgspec_route import _MsgspecRoute + +pytestmark = pytest.mark.cpu_only + +SYSTEM_PATHS = ("/health", "/version", "/v1/models") +SYSTEM_ROUTER_MODULE = "tensorrt_llm.serve.api.ops.system.api_router" + + +@pytest.mark.parametrize("config", VISUAL_GEN_CONFIGS, ids=lambda config: config.name) +def test_routes_come_from_ops_system_visual_gen(config, tmp_path): + with isolated_serve_env(tmp_path) as env: + server = env.build_openai_server(config.build) + by_path = {path: [] for path in SYSTEM_PATHS} + for path, route in effective_routes(server.app): + if path in by_path: + by_path[path].append(route) + for path, routes in by_path.items(): + assert len(routes) == 1, f"{path} must be registered exactly once, found {len(routes)}" + route = routes[0] + assert isinstance(route, _MsgspecRoute), ( + f"{path} must be a _MsgspecRoute; a plain APIRoute breaks msgpack request " + "bodies from the disaggregated frontend" + ) + assert route.endpoint.__module__ == SYSTEM_ROUTER_MODULE diff --git a/tests/unittest/serve/api/visual_gen/test_ops_system_visual_gen.py b/tests/unittest/serve/api/visual_gen/test_ops_system_visual_gen.py new file mode 100644 index 000000000000..205391381c13 --- /dev/null +++ b/tests/unittest/serve/api/visual_gen/test_ops_system_visual_gen.py @@ -0,0 +1,52 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""/health, /version and /v1/models of the VisualGen serving role.""" + +import pytest +from test_common.serve.app_factory import AggConfig, isolated_serve_env + +import tensorrt_llm.version +from tensorrt_llm.llmapi.disagg_utils import ServerRole + +pytestmark = pytest.mark.cpu_only + + +@pytest.fixture +def visual_gen_client(tmp_path): + with isolated_serve_env(tmp_path) as env: + server = env.build_openai_server( + AggConfig(role=ServerRole.VISUAL_GEN, model_name="fake-visual-gen") + ) + yield env.client(server.app) + + +def test_health_always_succeeds_today(visual_gen_client): + # The VisualGen generator exposes no health probe, so /health answers 200 + # unconditionally today; a change to that is a behavior change. + assert visual_gen_client.get("/health").status_code == 200 + + +def test_version(visual_gen_client): + response = visual_gen_client.get("/version") + assert response.status_code == 200 + assert response.json() == {"version": tensorrt_llm.version.__version__} + + +def test_models_lists_the_served_model(visual_gen_client): + response = visual_gen_client.get("/v1/models") + assert response.status_code == 200 + body = response.json() + assert body["object"] == "list" + assert [entry["id"] for entry in body["data"]] == ["fake-visual-gen"] diff --git a/tests/unittest/serve/api/visual_gen/test_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..40f7fc1cb9d4 --- /dev/null +++ b/tests/unittest/serve/api/visual_gen/test_private_access_lint_visual_gen.py @@ -0,0 +1,46 @@ +# 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 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 +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(): + 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/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..d767e5fc781a --- /dev/null +++ b/tests/unittest/serve/api/visual_gen/test_shim_identity_visual_gen.py @@ -0,0 +1,115 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""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 ShimRow, check_identity, check_warnings, find_old_path_imports + +pytestmark = pytest.mark.cpu_only + +REPO_ROOT = Path(__file__).resolve().parents[5] + +# 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(): + 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/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..c56e2000ee2b --- /dev/null +++ b/tests/unittest/serve/openai_protocol/test_openai_protocol_facade.py @@ -0,0 +1,466 @@ +# 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)), + }})) + """ + ) + culprits: list = [] + 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/private_access_baseline.yaml b/tests/unittest/serve/private_access_baseline.yaml new file mode 100644 index 000000000000..73fc557b69c0 --- /dev/null +++ b/tests/unittest/serve/private_access_baseline.yaml @@ -0,0 +1,27 @@ +# 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/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/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.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/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_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: _done, count: 1, owner: executor/result.py GenerationResultBase, replacement: GenerationResultBase.finished} +- {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_agents_md.py b/tests/unittest/serve/test_agents_md.py new file mode 100644 index 000000000000..b95dac68ff16 --- /dev/null +++ b/tests/unittest/serve/test_agents_md.py @@ -0,0 +1,132 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""The serve agent docs name only paths and commands that exist. + +Rules written for agents rot silently when the file or command they name moves; +this test keeps every inline-code path and command in the serve ``AGENTS.md`` +files real, and bans line-number anchors, which drift without failing anything. +""" + +import re +from pathlib import Path + +import pytest + +pytestmark = pytest.mark.cpu_only + +REPO_ROOT = Path(__file__).resolve().parents[3] + +DOCS = ("tensorrt_llm/serve/AGENTS.md", "tensorrt_llm/commands/AGENTS.md") + +PATH_ROOTS = ( + "tensorrt_llm/", + "serve/", + "commands/", + "tests/", + "scripts/", + "jenkins/", + ".github/", + "docs/", + "examples/", +) +ROOT_FILES = ("AGENTS.md", "CLAUDE.md", "legacy-files.txt", "CODING_GUIDELINES.md") +COMMAND_STARTS = ("pytest", "python", "python3", "PYTHONPATH=", "LLM_MODELS_ROOT=") + +_INLINE_CODE = re.compile(r"`([^`]+)`") +_LINE_ANCHOR = re.compile(r"\.(py|md|ya?ml|groovy|txt)\s*:\s*\d") +_BARE_ANCHOR_SPAN = re.compile(r"`:\d") + + +def _spans(doc: Path): + return _INLINE_CODE.findall(doc.read_text()) + + +def _resolve(path_text: str) -> Path: + if path_text.startswith(("serve/", "commands/")): + return REPO_ROOT / "tensorrt_llm" / path_text + return REPO_ROOT / path_text + + +def _is_template(token: str) -> bool: + return any(symbol in token for symbol in "<>*") + + +def _missing(doc: str, token: str) -> str: + return ( + f"{doc} names `{token}`, which does not exist. Update the rule in the PR that " + "moved or deleted it." + ) + + +def _check_command(doc: str, span: str, problems: list) -> None: + tokens = span.split(" ") + pythonpath_roots = [] + for token in tokens: + if token.startswith("PYTHONPATH="): + pythonpath_roots = token.split("=", 1)[1].split(":") + for index, token in enumerate(tokens): + if _is_template(token): + continue + if index > 0 and tokens[index - 1] == "-m" and "/" not in token and tokens[0] != "pytest": + # `python3 -m `: resolve to a file under the PYTHONPATH= + # roots or the repository root. + candidates = [] + for root in [*pythonpath_roots, "."]: + base = REPO_ROOT / root / Path(*token.split(".")) + candidates.extend((base.with_suffix(".py"), base / "__init__.py")) + if not any(candidate.exists() for candidate in candidates): + problems.append(_missing(doc, token)) + continue + if token.startswith(PATH_ROOTS): + target = token.split("::")[0] + if not _resolve(target).exists(): + problems.append(_missing(doc, token)) + + +@pytest.mark.parametrize("doc", DOCS) +def test_paths_and_commands_are_real(doc): + doc_path = REPO_ROOT / doc + problems = [] + for span in _spans(doc_path): + if " " not in span and (span.startswith(PATH_ROOTS) or span in ROOT_FILES): + if _is_template(span): + continue + target = _resolve(span.split("::")[0]) + if not target.exists(): + problems.append(_missing(doc, span)) + elif span.startswith(COMMAND_STARTS): + _check_command(doc, span, problems) + assert not problems, "\n".join(problems) + + +@pytest.mark.parametrize("doc", DOCS) +def test_no_line_number_anchors(doc): + text = (REPO_ROOT / doc).read_text() + assert not _LINE_ANCHOR.search(text) and not _BARE_ANCHOR_SPAN.search(text), ( + f"{doc} contains a line-number anchor. Name files and symbols, never line numbers: " + "anchors drift silently when the file changes." + ) + + +@pytest.mark.parametrize("doc", DOCS) +def test_sibling_claude_md(doc): + doc_path = REPO_ROOT / doc + sibling = doc_path.with_name("CLAUDE.md") + assert sibling.exists(), f"{doc} needs a sibling CLAUDE.md pointer file" + lines = [line.strip() for line in sibling.read_text().splitlines() if line.strip()] + assert lines == ["# In ./CLAUDE.md", "@AGENTS.md"], ( + f"{sibling.relative_to(REPO_ROOT)} must hold exactly the repo's two-line pointer " + "form: `# In ./CLAUDE.md` and `@AGENTS.md`." + ) diff --git a/tests/unittest/serve/test_layering.py b/tests/unittest/serve/test_layering.py new file mode 100644 index 000000000000..f70d9b335aa3 --- /dev/null +++ b/tests/unittest/serve/test_layering.py @@ -0,0 +1,378 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Import directions between serving layers, written as tests. + +Each pull request that creates a package under ``tensorrt_llm/serve`` adds its +contract here. A contract forbids imports by module prefix; a baseline allowance +records the imports that existed when the contract was written and only shrinks. +""" + +import ast +import textwrap +from pathlib import Path + +import pytest +from test_common.serve.layering import ( + Contract, + ImportEdge, + build_import_graph, + check_contracts, + stale_allowances, +) + +pytestmark = pytest.mark.cpu_only + +REPO_ROOT = Path(__file__).resolve().parents[3] + +CONTRACTS = ( + Contract( + name="render core has no HTTP or executor dependency", + importers=( + "tensorrt_llm.serve.render.chat", + "tensorrt_llm.serve.render.resources", + "tensorrt_llm.serve.render.fingerprint", + "tensorrt_llm.serve.render.prompt_types", + ), + forbidden=("fastapi", "starlette", "tensorrt_llm.executor"), + fix=( + "Keep HTTP code in serve/render/api_router.py or _http.py and pass engine data " + "in through RenderResources." + ), + ), + Contract( + name="gRPC frontends do not reach into serve", + importers=("tensorrt_llm.grpc",), + forbidden=("tensorrt_llm.serve",), + allowed=frozenset( + { + ( + "tensorrt_llm.grpc.openengine.server", + "tensorrt_llm.serve._frontend_processes", + name, + ) + for name in ( + "FrontendStartupCancelled", + "_init_multi_frontend_mode", + "_signal_frontend_ready", + "_spawn_attached_frontends", + "_terminate_attached_frontends", + ) + } + | { + ( + "tensorrt_llm.grpc.openengine.server", + "tensorrt_llm.serve._perf_metrics_config", + "resolve_perf_metrics_config", + ) + } + ), + fix=( + "gRPC frontends use serving code only through public entry points; add a public " + "API in tensorrt_llm/llmapi instead. The baseline only shrinks: delete an " + "allowance when its import goes away." + ), + ), + Contract( + name="serve/_base/ does not import serve/api/", + importers=("tensorrt_llm.serve._base",), + forbidden=("tensorrt_llm.serve.api",), + fix=( + "serve/_base/ is shared infrastructure and must not depend on endpoint " + "packages; pass what it needs in from the composition root." + ), + ), + Contract( + name="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." + ), + ), + 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." + ), + ), +) + + +@pytest.fixture(scope="module") +def edges(): + return build_import_graph(REPO_ROOT / "tensorrt_llm", "tensorrt_llm") + + +def test_contracts_hold(edges): + assert check_contracts(edges, CONTRACTS) == [] + + +def test_no_stale_allowances(edges): + assert stale_allowances(edges, CONTRACTS) == [] + + +_API_PREFIX = "tensorrt_llm.serve.api" + + +def _api_package_of(module: str): + """The endpoint package of a module under ``tensorrt_llm.serve.api``. + + The package is the first two components below ``api`` for ``openai`` and + ``ops`` (``openai.chat``, ``ops.system``) and the first one otherwise + (``anthropic``, ``visual_gen``). + """ + below = module.split(".")[3:] + if below and below[0] in ("openai", "ops"): + return tuple(below[:2]) + return tuple(below[:1]) + + +def _under_api(module: str) -> bool: + return module == _API_PREFIX or module.startswith(_API_PREFIX + ".") + + +def _cross_package_imports(edges) -> list: + """Imports between endpoint packages, matched per imported name as check_contracts does. + + ``from tensorrt_llm.serve import api`` names the api root package and + ``from tensorrt_llm.serve.api.ops import system`` names ``ops.system``: the + package is read off ``.``, not off the ``from`` module alone. + """ + problems = [] + for edge in edges: + if edge.type_checking or not _under_api(edge.importer): + continue + for name in edge.names or ("",): + target = f"{edge.imported}.{name}" if name else edge.imported + if not _under_api(target): + continue + if _api_package_of(target) != _api_package_of(edge.importer): + problems.append( + f"serve/api/ packages do not import each other: {edge.importer} " + f"(line {edge.lineno}) imports `{target}`. Share code through serve/_base/ " + "or inject it from the composition root." + ) + return problems + + +def test_api_packages_do_not_import_each_other(edges): + assert _cross_package_imports(edges) == [] + + +def test_cross_package_classification_reads_the_imported_name(): + def edge(importer, imported, names): + return ImportEdge( + importer=importer, + imported=imported, + names=names, + lineno=1, + in_function=False, + type_checking=False, + ) + + system = "tensorrt_llm.serve.api.ops.system.api_router" + # Parent-package forms resolve to the submodule they name: the same package. + assert _cross_package_imports([edge(system, "tensorrt_llm.serve.api.ops", ("system",))]) == [] + assert ( + _cross_package_imports([edge(system, "tensorrt_llm.serve.api.ops.system", ("serving",))]) + == [] + ) + # The api root, a sibling package and another dialect are cross-package, whatever the form. + assert len(_cross_package_imports([edge(system, "tensorrt_llm.serve", ("api",))])) == 1 + assert ( + len(_cross_package_imports([edge(system, "tensorrt_llm.serve.api.ops", ("metrics",))])) == 1 + ) + anthropic = "tensorrt_llm.serve.api.anthropic.serving" + assert len(_cross_package_imports([edge(system, anthropic, ("AnthropicServing",))])) == 1 + # Importers outside serve/api/ are not classified here. + root = "tensorrt_llm.serve.openai_server" + assert ( + _cross_package_imports( + [edge(root, "tensorrt_llm.serve.api.ops.system.serving", ("SystemServing",))] + ) + == [] + ) + + +def test_api_init_imports_nothing(): + tree = ast.parse((REPO_ROOT / "tensorrt_llm" / "serve" / "api" / "__init__.py").read_text()) + imports = [node for node in ast.walk(tree) if isinstance(node, (ast.Import, ast.ImportFrom))] + assert imports == [], ( + "serve/api/__init__.py imports nothing: attach_routers imports routers by role; an " + "import here would load them for every role and break the VisualGen CI selection." + ) + + +# --------------------------------------------------------------------------- +# Self-test on a synthetic package: the graph must see every import form. +# --------------------------------------------------------------------------- + + +@pytest.fixture +def synthetic_package(tmp_path): + pkg = tmp_path / "layerpkg" + pkg.mkdir() + (pkg / "__init__.py").write_text("") + (pkg / "inner.py").write_text("VALUE = 1\n") + (pkg / "top.py").write_text( + textwrap.dedent( + """ + from typing import TYPE_CHECKING + + import os + from json import dumps, loads + + from .inner import VALUE + + if TYPE_CHECKING: + import socket + + def late(): + import importlib + + importlib.import_module("collections.abc") + from ..sibling import thing # noqa: F401 + return VALUE + """ + ) + ) + return pkg + + +def test_import_graph_records_every_form(synthetic_package): + edges = build_import_graph(synthetic_package, "layerpkg") + by_imported = {edge.imported: edge for edge in edges} + + assert by_imported["os"].names == () + assert not by_imported["os"].in_function and not by_imported["os"].type_checking + assert by_imported["json"].names == ("dumps", "loads") + assert by_imported["layerpkg.inner"].names == ("VALUE",) # relative import resolved + assert by_imported["socket"].type_checking + assert by_imported["importlib"].in_function + assert by_imported["collections.abc"].names == () # importlib.import_module literal + assert by_imported["collections.abc"].in_function + assert by_imported["sibling"].names == ("thing",) # level-2 relative from layerpkg.top + + +def test_contract_violation_and_allowance(synthetic_package): + edges = build_import_graph(synthetic_package, "layerpkg") + contract = Contract( + name="top does not import json", + importers=("layerpkg.top",), + forbidden=("json",), + fix="Use the serializer helper instead.", + ) + problems = check_contracts(edges, [contract]) + assert len(problems) == 2 # one per from-import name + assert "contract 'top does not import json'" in problems[0] + assert "Use the serializer helper instead." in problems[0] + + allowed = Contract( + name="top does not import json", + importers=("layerpkg.top",), + forbidden=("json",), + allowed=frozenset({("layerpkg.top", "json", "dumps"), ("layerpkg.top", "json", "loads")}), + fix="Use the serializer helper instead.", + ) + assert check_contracts(edges, [allowed]) == [] + assert stale_allowances(edges, [allowed]) == [] + + stale = Contract( + name="top does not import json", + importers=("layerpkg.top",), + forbidden=("json",), + allowed=frozenset({("layerpkg.top", "json", "gone")}), + fix="Use the serializer helper instead.", + ) + messages = stale_allowances(edges, [stale]) + assert len(messages) == 1 and "the baseline only shrinks" in messages[0] + + +def test_parent_package_import_of_a_forbidden_module(synthetic_package): + # Both forms record the edge against the parent package, with the forbidden + # submodule only as an imported name. + (synthetic_package / "via_parent.py").write_text( + textwrap.dedent( + """ + from layerpkg import inner, top + from . import inner as relative_inner + """ + ) + ) + edges = build_import_graph(synthetic_package, "layerpkg") + contract = Contract( + name="no inner", + importers=("layerpkg.via_parent",), + forbidden=("layerpkg.inner",), + fix="Import VALUE through layerpkg.top.", + ) + problems = check_contracts(edges, [contract]) + assert len(problems) == 2 # `top` from the same statement is not flagged + assert "(line 2) imports `layerpkg.inner`." in problems[0] # from layerpkg import inner + assert "(line 3) imports `layerpkg.inner`." in problems[1] # from . import inner + + allowed = Contract( + name="no inner", + importers=("layerpkg.via_parent",), + forbidden=("layerpkg.inner",), + allowed=frozenset({("layerpkg.via_parent", "layerpkg", "inner")}), + fix="Import VALUE through layerpkg.top.", + ) + assert check_contracts(edges, [allowed]) == [] + assert stale_allowances(edges, [allowed]) == [] + + +def test_type_checking_imports_are_skipped_by_default(synthetic_package): + edges = build_import_graph(synthetic_package, "layerpkg") + contract = Contract( + name="no socket", + importers=("layerpkg",), + forbidden=("socket",), + fix="unused", + ) + assert check_contracts(edges, [contract]) == [] + including = Contract( + name="no socket", + importers=("layerpkg",), + forbidden=("socket",), + include_type_checking=True, + fix="unused", + ) + assert len(check_contracts(edges, [including])) == 1 diff --git a/tests/unittest/serve/test_middleware_order.py b/tests/unittest/serve/test_middleware_order.py new file mode 100644 index 000000000000..aec17dd4f79d --- /dev/null +++ b/tests/unittest/serve/test_middleware_order.py @@ -0,0 +1,133 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Middleware-order snapshot of the aggregated and disaggregated apps. + +Starlette's ``add_middleware`` inserts at position 0, so the last one added is +outermost. ``ServerArrivalTimeMiddleware`` assigns ``scope['state'] = {}`` and +``PerfMetricsMiddleware`` only ``setdefault``s into it, so their relative order +is load-bearing: swapping them wipes every perf-metrics record. VisualGen is not +listed: it takes the default code path without perf metrics, and no VisualGen +path can change middleware registration. +""" + +import pytest +from test_common.serve.app_factory import AggConfig, CoordinatorConfig, DisaggConfig +from test_common.serve.middleware import middleware_snapshot + +from tensorrt_llm.llmapi.disagg_utils import ServerRole + +pytestmark = pytest.mark.cpu_only + +_ARRIVAL_WITH_CLOCK = {"class": "ServerArrivalTimeMiddleware", "kwargs": ["adjusted_clock"]} +_ARRIVAL_PLAIN = {"class": "ServerArrivalTimeMiddleware", "kwargs": []} + + +def _perf(kwargs, expose_headers): + return { + "class": "PerfMetricsMiddleware", + "kwargs": sorted(kwargs), + "expose_headers": expose_headers, + } + + +# Outermost first, matching middleware_snapshot. +EXPECTED = { + "llm": (AggConfig(), [_ARRIVAL_WITH_CLOCK]), + "mm_encoder": (AggConfig(role=ServerRole.MM_ENCODER), [_ARRIVAL_WITH_CLOCK]), + "embedding": (AggConfig(role=ServerRole.EMBEDDING), [_ARRIVAL_WITH_CLOCK]), + "llm_perf_metrics": ( + AggConfig(return_perf_metrics=True), + [ + _ARRIVAL_WITH_CLOCK, + _perf(["adjusted_clock", "expose_headers", "writer"], True), + ], + ), + "llm_perf_metrics_output_dir_only": ( + AggConfig(perf_metrics_output_dir=True), + [ + _ARRIVAL_WITH_CLOCK, + _perf(["adjusted_clock", "expose_headers", "writer"], False), + ], + ), + "llm_perf_metrics_enable_metrics_env": ( + AggConfig(return_perf_metrics=True, env=(("ENABLE_METRICS", "true"),)), + [ + _ARRIVAL_WITH_CLOCK, + _perf(["adjusted_clock", "expose_headers", "writer"], True), + {"class": "PrometheusInstrumentatorMiddleware"}, + ], + ), + "disagg_local": (DisaggConfig(), [_ARRIVAL_PLAIN]), + "disagg_fleet": (DisaggConfig(fleet=True), [_ARRIVAL_PLAIN]), + "disagg_local_perf_metrics": ( + DisaggConfig(return_perf_metrics=True), + [_ARRIVAL_PLAIN, _perf(["expose_headers", "writer"], True)], + ), + "disagg_fleet_perf_metrics": ( + DisaggConfig(fleet=True, return_perf_metrics=True), + [_ARRIVAL_PLAIN, _perf(["expose_headers", "writer"], True)], + ), + "coordinator": (CoordinatorConfig(), []), +} + + +def _build(serve_env, config): + if isinstance(config, AggConfig): + return serve_env.build_openai_server(config).app + if isinstance(config, DisaggConfig): + return serve_env.build_disagg_server(config).app + return serve_env.build_coordinator_server(config).app + + +def _failure(name, expected, got): + return ( + f"Middleware stack of '{name}' changed: expected {expected}, got {got}. " + "Starlette's add_middleware inserts at position 0, so the last one added is " + "outermost. ServerArrivalTimeMiddleware assigns scope['state'] = {} and must stay " + "outside PerfMetricsMiddleware, or every perf-metrics record is wiped. Agg passes " + "adjusted_clock to both, disagg to neither. If the change is intended, update " + "EXPECTED in tests/unittest/serve/test_middleware_order.py and say why in the PR " + "body." + ) + + +@pytest.mark.parametrize("name", sorted(EXPECTED)) +def test_middleware_order(name, serve_env): + config, expected = EXPECTED[name] + snapshot = middleware_snapshot(_build(serve_env, config)) + assert snapshot == expected, _failure(name, expected, snapshot) + + +@pytest.mark.parametrize("name", sorted(EXPECTED)) +def test_arrival_time_stays_outside_perf_metrics(name, serve_env): + """When both are present, ServerArrivalTimeMiddleware is outermost.""" + config, _ = EXPECTED[name] + classes = [entry["class"] for entry in middleware_snapshot(_build(serve_env, config))] + if "PerfMetricsMiddleware" not in classes or "ServerArrivalTimeMiddleware" not in classes: + return + assert classes.index("ServerArrivalTimeMiddleware") < classes.index("PerfMetricsMiddleware"), ( + _failure(name, "ServerArrivalTimeMiddleware before PerfMetricsMiddleware", classes) + ) + + +def test_user_middleware_is_outermost(serve_env): + """``--middleware`` classes are applied after construction, so they end up outermost.""" + from tensorrt_llm.commands.serve import _apply_fastapi_middlewares + + server = serve_env.build_openai_server(AggConfig()) + _apply_fastapi_middlewares(server.app, ["test_common.serve.fakes.UserProbeMiddleware"]) + snapshot = middleware_snapshot(server.app) + expected = [{"class": "UserProbeMiddleware"}, _ARRIVAL_WITH_CLOCK] + assert snapshot == expected, _failure("llm_user_middleware", expected, snapshot) diff --git a/tests/unittest/serve/test_no_object_new.py b/tests/unittest/serve/test_no_object_new.py new file mode 100644 index 000000000000..0aa59460caf8 --- /dev/null +++ b/tests/unittest/serve/test_no_object_new.py @@ -0,0 +1,302 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Ratchet against new ``object.__new__()`` test sites. + +Tests that build a server with ``object.__new__`` bypass the constructor and pin +whatever attribute set the handler happens to read; they break or silently +mistest when handlers move. New tests build servers with ``isolated_serve_env`` +from ``tests/test_common/serve/app_factory.py`` instead. The existing sites stay +until the change that moves each handler rewrites them, so this test holds a +per-file baseline of sites that may only shrink. A site is named by its +enclosing function or class and the server class it builds, so a site added in +another function fails even when one is removed from the same file. + +The scan counts AST call nodes, not text, so a docstring or comment that +mentions ``object.__new__(OpenAIServer)`` does not count. A name bound by +``from ... import OpenAIServer as Server`` counts as ``OpenAIServer``. +""" + +import ast +from collections import Counter +from pathlib import Path + +import pytest + +pytestmark = pytest.mark.cpu_only + +REPO_ROOT = Path(__file__).resolve().parents[3] +TESTS_ROOT = REPO_ROOT / "tests" + +SERVER_CLASSES = ("OpenAIServer", "OpenAIDisaggServer", "CoordinatorServer") + +# Per-file sites that existed when this ratchet landed, as +# " " (see find_new_sites), repeated +# when one scope builds the class twice. The baseline only shrinks: rewrite a +# site on the factory, delete its entry. +BASELINE = { + "tests/unittest/disaggregated/test_coordinator_worker.py": [ + "test_coordinator_keep_alive_timeout_is_passed_to_uvicorn CoordinatorServer", + "test_coordinator_tcp_listener_uses_the_prebound_socket CoordinatorServer", + ], + "tests/unittest/disaggregated/test_disagg_internal_auth.py": [ + "test_aggregated_worker_ignores_affinity_header OpenAIServer", + "test_worker_affinity_survives_wire_roundtrip OpenAIServer", + "test_worker_rejects_protected_fields_without_cache_transceiver_config OpenAIServer", + ], + "tests/unittest/disaggregated/test_disagg_openai_client.py": [ + "test_worker_clock_calibration_uses_global_clock OpenAIServer", + ], + "tests/unittest/disaggregated/test_disagg_responses_api.py": [ + "test_a_context_worker_with_postprocessing_workers_returns_the_prompt OpenAIServer", + "test_a_models_failure_is_an_error_response OpenAIDisaggServer", + ], + "tests/unittest/disaggregated/test_openai_disagg_server.py": [ + "test_create_client_does_not_register_with_server_metrics_collector OpenAIDisaggServer", + "test_http_cluster_storage_request_is_proxied_to_coordinator OpenAIDisaggServer", + "test_server_keep_alive_timeout_is_passed_to_uvicorn OpenAIDisaggServer", + ], + "tests/unittest/disaggregated/test_openai_disagg_service.py": [ + "test_create_chat_response_sets_prompt_token_ids_for_context_only OpenAIServer", + "test_subagent_affinity_http_worker_adp_placement.make_worker OpenAIServer", + ], + "tests/unittest/disaggregated/test_openai_server_info.py": [ + "_make_server OpenAIServer", + ], + "tests/unittest/disaggregated/test_openai_server_stream_errors.py": [ + "_bare_server OpenAIServer", + ], + "tests/unittest/disaggregated/test_paired_request_cancellation.py": [ + "test_http_failure_reaches_paired_worker_abort OpenAIDisaggServer", + "test_http_failure_reaches_paired_worker_abort OpenAIServer", + ], + "tests/unittest/disaggregated/test_router.py": [ + "test_chat_harmony_preserves_original_tool_conversion_error OpenAIServer", + "test_gpt_oss_router_tokens_match_chat_harmony_server_input OpenAIServer", + ], + "tests/unittest/executor/test_spec_decode_stats_payload.py": [ + "TestServerOptIn._server OpenAIServer", + ], + "tests/unittest/llmapi/apps/_test_trtllm_serve_profile.py": [ + "server_with_mock_generator OpenAIServer", + ], + "tests/unittest/llmapi/apps/test_anthropic_routes.py": [ + "_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", + ], +} + + +def _class_name(node) -> str: + if isinstance(node, ast.Name): + return node.id + if isinstance(node, ast.Attribute): + return node.attr + return "" + + +def _server_class_names(tree: ast.AST) -> dict[str, str]: + """Local name -> server class, with the names bound by ``import ... as``.""" + names = {server_class: server_class for server_class in SERVER_CLASSES} + for node in ast.walk(tree): + if isinstance(node, ast.ImportFrom): + for alias in node.names: + if alias.name in SERVER_CLASSES and alias.asname: + names[alias.asname] = alias.name + return names + + +def _built_server_class(node: ast.AST, names: dict[str, str]) -> str: + """The server class that a ``__new__`` call node builds, or "" for any other node.""" + if not isinstance(node, ast.Call) or not isinstance(node.func, ast.Attribute): + return "" + if node.func.attr != "__new__": + return "" + owner = node.func.value + if _class_name(owner) == "object": + if not node.args: + return "" + owner = node.args[0] + if isinstance(owner, ast.Name): + return names.get(owner.id, "") + if isinstance(owner, ast.Attribute) and owner.attr in SERVER_CLASSES: + return owner.attr + return "" + + +def find_new_sites(tree: ast.AST) -> list[str]: + """``" "`` per ``__new__`` call that builds one. + + The scope is the dotted chain of enclosing ``def``/``class`` names, or + ```` at module level; the list is sorted. + """ + names = _server_class_names(tree) + sites = [] + stack = [(tree, "")] + while stack: + node, scope = stack.pop() + server_class = _built_server_class(node, names) + if server_class: + sites.append(f"{scope or ''} {server_class}") + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)): + scope = f"{scope}.{node.name}" if scope else node.name + stack.extend((child, scope) for child in ast.iter_child_nodes(node)) + return sorted(sites) + + +def scan_tests(tests_root: Path = TESTS_ROOT) -> dict[str, list[str]]: + sites_by_file = {} + for path in sorted(tests_root.rglob("*.py")): + try: + tree = ast.parse(path.read_text(errors="replace")) + except SyntaxError: + continue + sites = find_new_sites(tree) + if sites: + sites_by_file[str(path.relative_to(REPO_ROOT))] = sites + return sites_by_file + + +def ratchet_problems( + sites_by_file: dict[str, list[str]], baseline: dict[str, list[str]] +) -> list[str]: + """Sites missing from ``baseline``, then baseline entries with no site left, per file.""" + problems = [] + for file in sorted(set(sites_by_file) | set(baseline)): + live = Counter(sites_by_file.get(file, [])) + recorded = Counter(baseline.get(file, [])) + for site in sorted((live - recorded).elements()): + problems.append( + f"{file}: '{site}' builds a server with object.__new__ and is not in the " + "baseline; use `isolated_serve_env` from " + "tests/test_common/serve/app_factory.py instead; see " + "tensorrt_llm/serve/AGENTS.md." + ) + for site in sorted((recorded - live).elements()): + problems.append( + f"{file}: '{site}' no longer builds a server with object.__new__; delete it " + f"from BASELINE[{file!r}] in tests/unittest/serve/test_no_object_new.py; the " + "baseline only shrinks." + ) + return problems + + +def test_no_new_object_new_sites(): + problems = ratchet_problems(scan_tests(), BASELINE) + assert not problems, "\n".join(problems) + + +# --------------------------------------------------------------------------- +# Self-test on synthetic sources: the scan counts calls, not text. +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "source,expected", + [ + ('"""object.__new__(OpenAIServer) in a docstring"""', 0), + ("# object.__new__(OpenAIServer) in a comment\nx = 1", 0), + ("server = object.__new__(OpenAIServer)", 1), + ("server = object.__new__(mod.OpenAIDisaggServer)", 1), + ("server = CoordinatorServer.__new__(CoordinatorServer)", 1), + ("thing = object.__new__(SomethingElse)", 0), + ("from m import OpenAIServer as Server\nserver = object.__new__(Server)", 1), + ("from m import OpenAIDisaggServer as S\nserver = S.__new__(S)", 1), + ("from m import SomethingElse as Server\nthing = object.__new__(Server)", 0), + ], +) +def test_scan_counts_call_nodes(source, expected): + assert len(find_new_sites(ast.parse(source))) == expected + + +def test_scan_names_sites_by_scope_and_class(): + source = ( + "from tensorrt_llm.serve.openai_server import OpenAIServer as Server\n" + "SERVER = object.__new__(OpenAIServer)\n" + "class TestServer:\n" + " def test_build(self):\n" + " return object.__new__(Server)\n" + ) + assert find_new_sites(ast.parse(source)) == [ + " OpenAIServer", + "TestServer.test_build OpenAIServer", + ] + + +def test_ratchet_flags_a_replaced_site(): + # Removing one site and adding one in another function keeps the file's + # total; the ratchet still reports both. + file = "tests/unittest/test_example.py" + before = find_new_sites(ast.parse("def test_old():\n object.__new__(OpenAIServer)\n")) + after = find_new_sites(ast.parse("def test_new():\n object.__new__(OpenAIServer)\n")) + assert ratchet_problems({file: after}, {file: after}) == [] + problems = ratchet_problems({file: after}, {file: before}) + assert len(problems) == 2, problems + assert "'test_new OpenAIServer' builds a server" in problems[0] + assert "'test_old OpenAIServer' no longer builds a server" in problems[1] diff --git a/tests/unittest/serve/test_openai_server_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_private_access_lint.py b/tests/unittest/serve/test_private_access_lint.py new file mode 100644 index 000000000000..26603599f476 --- /dev/null +++ b/tests/unittest/serve/test_private_access_lint.py @@ -0,0 +1,410 @@ +# 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/") +OPENENGINE_SCOPE = "tensorrt_llm/grpc/openengine/" + +# Lower this when you remove accesses; never raise it. +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 +# 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()) + 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=""): + 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 but MAX_TOTAL={MAX_TOTAL}. " + "Lower MAX_TOTAL when the total shrinks. " + 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) + + +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(): + 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() 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..783d681f0c7a --- /dev/null +++ b/tests/unittest/serve/test_serve_runtime_api_calls.py @@ -0,0 +1,182 @@ +# 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`` 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 + +import pytest +from fastapi import FastAPI + +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 + + +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) + + +# --------------------------------------------------------------------------- +# EmbeddingServing builds its batcher from generator.encode_limits +# --------------------------------------------------------------------------- +def test_embedding_batcher_takes_the_limits_encode_checks(): + 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, + ) + + batcher = serving.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 = [] + # 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] + + +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" 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..9a5782a6d3b9 --- /dev/null +++ b/tests/unittest/serve/test_serve_test_helpers.py @@ -0,0 +1,433 @@ +# 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 ( + FakeEncodeOnlyLLM, + FakeLLM, + FakeLlmArgs, + ScriptedOutput, + ScriptedStep, + UserProbeMiddleware, +) +from test_common.serve.middleware import middleware_snapshot +from test_common.serve.routes import _records_conflict, compare, effective_routes + +import tensorrt_llm._utils +from tensorrt_llm.inputs.media_io import BaseMediaIO +from tensorrt_llm.sampling_params import SamplingParams +from tensorrt_llm.usage import usage_lib + +# threadleak is off: the scripted-generation tests create TestClients and run +# requests inside the test body, so the client's portal thread and the fake +# server's worker-pool threads are still alive when the leak check runs; the +# serve_env fixture closes them in its teardown, which runs after that check +# (precedent: tests/unittest/llmapi/apps/test_serve_lazy_import.py). +pytestmark = [pytest.mark.cpu_only, pytest.mark.threadleak(enabled=False)] + +REPO_ROOT = Path(__file__).resolve().parents[3] +_MISSING = object() + + +# --------------------------------------------------------------------------- +# effective_routes +# --------------------------------------------------------------------------- + + +def test_effective_routes_flattens_included_routers(): + app = FastAPI() + app.add_api_route("/direct", lambda: None, methods=["GET"]) + router = APIRouter() + router.add_api_route("/from_router", lambda: None, methods=["POST"]) + app.include_router(router) + prefixed = APIRouter() + prefixed.add_api_route("/suffix", lambda: None, methods=["GET"]) + app.include_router(prefixed, prefix="/pre") + paths = [path for path, _ in effective_routes(app)] + assert "/direct" in paths + assert "/from_router" in paths, ( + "effective_routes must walk into included routers; FastAPI 0.139 keeps them as " + "objects without a path attribute (duck-typed on original_router and " + "include_context in test_common/serve/routes.py)." + ) + assert "/pre/suffix" in paths, "an include_router prefix must be applied to the leaf paths" + + +def test_effective_routes_matches_app_routes_without_includes(): + app = FastAPI() + app.add_api_route("/only", lambda: None, methods=["GET"]) + assert [route for _, route in effective_routes(app)] == list(app.routes) + + +# --------------------------------------------------------------------------- +# compare and the conflict rule +# --------------------------------------------------------------------------- + + +def test_compare_reports_each_difference_kind(): + expected = [ + "GET /a _MsgspecRoute", + "GET /metrics _MsgspecRoute #get_iteration_stats", + "GET /metrics _MsgspecRoute #metrics", + "POST /v1/chat/completions _MsgspecRoute", + ] + actual = [ + "GET /metrics _MsgspecRoute #metrics", + "GET /metrics _MsgspecRoute #get_iteration_stats", + "POST /v1/chat/completions APIRoute", + "POST /new _MsgspecRoute", + ] + differences = compare(expected, actual) + assert any(d.startswith("added:") and "/new" in d for d in differences) + assert any(d.startswith("removed:") and "GET /a" in d for d in differences) + assert any( + "route class changed" in d + and "_MsgspecRoute -> APIRoute" in d + and "(a plain APIRouter breaks msgpack bodies from the disaggregated frontend)" in d + for d in differences + ) + assert any("order changed" in d and "#get_iteration_stats" in d for d in differences) + + +def test_compare_ignores_order_of_non_conflicting_routes(): + expected = ["GET /a APIRoute", "GET /b APIRoute"] + assert compare(expected, ["GET /b APIRoute", "GET /a APIRoute"]) == [] + assert compare(expected, list(expected)) == [] + + +# The first six rows use records of the real apps, the rest are synthetic. +@pytest.mark.parametrize( + "record_a,record_b,conflict", + [ + ( + "GET /metrics _MsgspecRoute #get_iteration_stats", + "GET /metrics _MsgspecRoute #metrics", + True, + ), + ( + "GET /v1/responses/{response_id} _MsgspecRoute", + "DELETE /v1/responses/{response_id} _MsgspecRoute", + False, + ), + ( + "GET /v1/messages/batches _MsgspecRoute", + "GET /v1/messages/batches/{batch_id} _MsgspecRoute", + False, + ), + ( + "GET /v1/images/{image_id}/content _MsgspecRoute", + "GET /v1/videos/{video_id}/content _MsgspecRoute", + False, + ), + ("MOUNT /prometheus/metrics Mount", "GET /metrics _MsgspecRoute", False), + ("GET,HEAD /docs Route", "GET,HEAD /docs/oauth2-redirect Route", False), + ("GET /v1/models _MsgspecRoute", "GET /v1/{name} APIRoute", True), + ("GET,HEAD /openapi.json Route", "GET /openapi.json APIRoute", True), + ("GET /files/{rest:path} APIRoute", "GET /files/a/b APIRoute", True), + ("MOUNT /prometheus/metrics Mount", "POST /prometheus/metrics/x APIRoute", True), + ], +) +def test_conflict_rule(record_a, record_b, conflict): + assert _records_conflict(record_a, record_b) is conflict + + +# --------------------------------------------------------------------------- +# middleware_snapshot format +# --------------------------------------------------------------------------- + + +def test_middleware_snapshot_format(): + from tensorrt_llm.serve.responses_utils import ServerArrivalTimeMiddleware + + app = FastAPI() + app.add_middleware(ServerArrivalTimeMiddleware, adjusted_clock=None) + app.add_middleware(UserProbeMiddleware) # added last, so outermost + snapshot = middleware_snapshot(app) + assert snapshot == [ + {"class": "UserProbeMiddleware"}, # not a tensorrt_llm class: name only + {"class": "ServerArrivalTimeMiddleware", "kwargs": ["adjusted_clock"]}, + ] + + +# --------------------------------------------------------------------------- +# isolated_serve_env +# --------------------------------------------------------------------------- + + +def test_isolated_serve_env_restores_global_state(tmp_path): + environ_before = dict(os.environ) + media_executor_before = BaseMediaIO._executor + prometheus_dir_before = getattr(tensorrt_llm._utils, "prometheus_multiproc_dir", _MISSING) + session_before = usage_lib._SESSION + collectors_before = set(REGISTRY._collector_to_names.keys()) + + with isolated_serve_env(tmp_path) as env: + env.build_openai_server(AggConfig(return_perf_metrics=True)) + os.environ["A2_RESTORE_SENTINEL"] = "1" + sink = env.record_terminations() + assert usage_lib._SESSION is not None and sink == [] + + assert dict(os.environ) == environ_before + assert BaseMediaIO._executor is media_executor_before + prometheus_dir_after = getattr(tensorrt_llm._utils, "prometheus_multiproc_dir", _MISSING) + if prometheus_dir_before is _MISSING: + assert prometheus_dir_after is _MISSING, ( + "isolated_serve_env must delete tensorrt_llm._utils.prometheus_multiproc_dir " + "when it did not exist before entry" + ) + else: + assert prometheus_dir_after is prometheus_dir_before + assert usage_lib._SESSION is session_before + assert set(REGISTRY._collector_to_names.keys()) == collectors_before, ( + "isolated_serve_env must unregister every Prometheus collector the build added" + ) + + +def test_isolated_serve_env_rejects_use_without_with(tmp_path): + with pytest.raises(AttributeError): + isolated_serve_env(tmp_path).build_openai_server(AggConfig()) + + +# --------------------------------------------------------------------------- +# cpu_tests selection +# --------------------------------------------------------------------------- + + +def test_cpu_tests_selection(): + kept, _ = select_entries() + targets = [target for entry in kept for target in entry.targets] + expected_new_files = [ + "unittest/api_stability/test_serve_routes.py", + "unittest/serve/test_agents_md.py", + "unittest/serve/test_layering.py", + "unittest/serve/test_middleware_order.py", + "unittest/serve/test_no_object_new.py", + "unittest/serve/test_serve_test_helpers.py", + "unittest/serve/test_shim_identity.py", + "unittest/serve/api/ops/system/test_system_endpoints.py", + "unittest/serve/api/visual_gen/test_ops_system_visual_gen.py", + "unittest/serve/api/visual_gen/test_routes_visual_gen.py", + "unittest/serve/api/visual_gen/test_shim_identity_visual_gen.py", + ] + for target in expected_new_files: + assert target in targets, f"the serve CPU runner must select {target}" + assert "unittest/llmapi/apps/test_anthropic_routes.py" in targets + for target in targets: + file_part = target.split("::")[0] + text = (REPO_ROOT / "tests" / file_part).read_text(errors="replace") + assert "pytest.mark.cpu_only" in text, ( + f"{file_part} is selected but carries no cpu_only marker; the CPU stage would " + "collect nothing from it" + ) + + +def test_cpu_tests_selection_keeps_node_ids(tmp_path): + # A node-id entry stays a node id: only that node runs, never the whole file. + node = "unittest/serve/test_serve_test_helpers.py::test_cpu_tests_selection" + cpu_list = tmp_path / "l0_cpu.yml" + cpu_list.write_text(f"l0_cpu:\n- tests:\n - {node}\n") + kept, notes = select_entries(cpu_list) + assert [entry.targets for entry in kept] == [(node,)], kept + assert not notes, notes + + +# --------------------------------------------------------------------------- +# Scripted generation through the production handlers +# --------------------------------------------------------------------------- + +_CHAT_BODY = { + "model": "fake-model", + "messages": [{"role": "user", "content": "hi"}], + "cache_salt": "salted", +} + + +def test_scripted_chat_non_streaming(serve_env): + server = serve_env.build_openai_server(AggConfig()) + client = serve_env.client(server.app) + response = client.post("/v1/chat/completions", json=_CHAT_BODY) + assert response.status_code == 200, response.text + body = response.json() + assert body["choices"][0]["message"]["content"] == "ok" # the default scripted text + generator = server.generator + assert len(generator.preprocess_calls) == 1 + call = generator.generate_calls[0] + assert isinstance(call.sampling_params, SamplingParams) + assert call.kwargs["streaming"] is False + assert call.kwargs["cache_salt"] == "salted" + + +def test_scripted_chat_streaming(serve_env): + server = serve_env.build_openai_server(AggConfig()) + server.generator.script( + ScriptedOutput( + steps=( + ScriptedStep(text="he", token_ids=(6,)), + ScriptedStep(text="llo", token_ids=(6,), finish_reason="stop"), + ) + ) + ) + client = serve_env.client(server.app) + with client.stream( + "POST", "/v1/chat/completions", json={**_CHAT_BODY, "stream": True} + ) as response: + assert response.status_code == 200 + lines = [line for line in response.iter_lines() if line.startswith("data: ")] + assert lines[-1] == "data: [DONE]" + deltas = [json.loads(line[len("data: ") :]) for line in lines[:-1]] + text = "".join( + choice["delta"].get("content") or "" for delta in deltas for choice in delta["choices"] + ) + assert text == "hello" + assert server.generator.generate_calls[0].kwargs["streaming"] is True + + +def test_scripted_completion_through_real_input_processor(serve_env): + server = serve_env.build_openai_server(AggConfig()) + server.generator.script( + ScriptedOutput(steps=(ScriptedStep(text="done", token_ids=(7,), finish_reason="stop"),)) + ) + client = serve_env.client(server.app) + response = client.post("/v1/completions", json={"model": "fake-model", "prompt": "hello hi"}) + assert response.status_code == 200, response.text + assert response.json()["choices"][0]["text"] == "done" + call = server.generator.generate_calls[0] + # A text prompt goes through the real DefaultInputProcessor, which + # tokenizes with the fake tokenizer: "hello hi" -> ids 6 and 5. + assert call.inputs == {"prompt_token_ids": [6, 5]} + + +def test_postprocess_worker_path_answers_through_postprocess_result(serve_env): + server = serve_env.build_openai_server(AggConfig(num_postprocess_workers=1)) + client = serve_env.client(server.app) + response = client.post("/v1/chat/completions", json=_CHAT_BODY) + assert response.status_code == 200, response.text + assert response.json()["choices"][0]["message"]["content"] == "ok" + promise = server.generator.generate_calls[0].result + assert promise.outputs[0]._postprocess_result is not None + with client.stream( + "POST", "/v1/chat/completions", json={**_CHAT_BODY, "stream": True} + ) as streaming_response: + lines = [line for line in streaming_response.iter_lines() if line.startswith("data: ")] + assert lines[-1] == "data: [DONE]" + assert any('"content"' in line for line in lines[:-1]) + + +def test_raise_on_call_returns_400_and_records_the_call(serve_env): + server = serve_env.build_openai_server(AggConfig()) + server.generator.script(ScriptedOutput(raise_on_call=RuntimeError("scripted failure"))) + client = serve_env.client(server.app) + response = client.post("/v1/chat/completions", json=_CHAT_BODY) + assert response.status_code == 400 + assert response.json()["message"] == "scripted failure" + call = server.generator.generate_calls[0] + assert call.result is None # recorded before the scripted exception fired + assert isinstance(call.sampling_params, SamplingParams) + + +def test_raise_after_steps_ends_a_stream_with_the_scripted_error(serve_env): + server = serve_env.build_openai_server(AggConfig()) + server.generator.script( + ScriptedOutput( + steps=(ScriptedStep(text="partial", token_ids=(5,)),), + raise_after_steps=RuntimeError("stream exploded"), + ) + ) + client = serve_env.client(server.app) + received = [] + with pytest.raises(RuntimeError, match="stream exploded"): + with client.stream( + "POST", "/v1/chat/completions", json={**_CHAT_BODY, "stream": True} + ) as response: + for line in response.iter_lines(): + received.append(line) + assert any("partial" in line for line in received), ( + "the first scripted chunk must reach the client before the scripted error" + ) + + +def test_hold_after_steps_release_and_cancellation(): + async def scenario(): + llm = FakeLLM(args=FakeLlmArgs(model="unused")) + llm.script( + ScriptedOutput( + steps=(ScriptedStep(text="first", token_ids=(5,)),), hold_after_steps=True + ) + ) + promise = llm.generate_async({"prompt_token_ids": [1, 2]}, SamplingParams()) + first = await promise.__anext__() + assert first.outputs[0].text == "first" + assert not promise.finished + + async def consume_rest(): + async for _ in promise: + pass + + consumer = asyncio.create_task(consume_rest()) + await asyncio.sleep(0) # the consumer reaches the hold + promise.release() + await consumer # release ends the wait with a normal end of stream + assert promise.finished + + llm.script( + ScriptedOutput(steps=(ScriptedStep(text="x", token_ids=(5,)),), hold_after_steps=True) + ) + held = llm.generate_async({"prompt_token_ids": [1]}, SamplingParams()) + await held.__anext__() + hanging = asyncio.create_task(held.__anext__()) + await asyncio.sleep(0) + hanging.cancel() # a client disconnect cancels the consuming task + with pytest.raises(asyncio.CancelledError): + await hanging + + asyncio.run(scenario()) + + +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 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