From 4a38ad70b0409159f0d1bf9b92a4c472795ff01b Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Tue, 13 Oct 2026 22:56:40 -0700 Subject: [PATCH] [#18151][feat] Admit trusted KV hints and carry immutable request metadata Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- docs/source/features/kv-hints.md | 246 +++++ docs/source/index.rst | 1 + .../disaggregation/orchestration/_kv_hints.py | 108 +++ tensorrt_llm/_torch/pyexecutor/llm_request.py | 4 + .../_torch/pyexecutor/request_utils.py | 1 + tensorrt_llm/executor/base_worker.py | 2 + tensorrt_llm/executor/executor.py | 5 +- tensorrt_llm/executor/request.py | 8 + tensorrt_llm/kv_hints/__init__.py | 8 + tensorrt_llm/kv_hints/metrics.py | 34 + tensorrt_llm/kv_hints/protocol.py | 157 ++++ tensorrt_llm/llmapi/__init__.py | 4 + tensorrt_llm/llmapi/llm.py | 13 + tensorrt_llm/serve/kv_hints.py | 61 ++ tensorrt_llm/serve/openai_disagg_server.py | 4 + tensorrt_llm/serve/openai_protocol.py | 34 + tensorrt_llm/serve/openai_server.py | 25 +- tests/integration/defs/perf/agentx_client.py | 14 +- .../integration/test_lists/test-db/l0_cpu.yml | 1 + .../visual_gen/test_attention_importable.py | 2 +- .../api_stability/api_stability_core.py | 2 +- .../api_stability/references/llm.yaml | 8 + .../references/trtllm_serve_api.yaml | 18 + .../disaggregated/test_kv_hint_selection.py | 224 +++++ tests/unittest/llmapi/test_kv_hints.py | 842 ++++++++++++++++++ .../test_trtllm_serve_config_schemas.py | 6 +- .../unittest/others/test_maintain_prims_ts.py | 1 + tests/unittest/others/test_vendor_bot.py | 3 +- tests/unittest/scripts/test_bolt.py | 5 +- .../scripts/test_perf_agentx_aiperf_probe.py | 40 + .../scripts/test_perf_agentx_client.py | 8 +- 31 files changed, 1861 insertions(+), 28 deletions(-) create mode 100644 docs/source/features/kv-hints.md create mode 100644 tensorrt_llm/_torch/disaggregation/orchestration/_kv_hints.py create mode 100644 tensorrt_llm/kv_hints/__init__.py create mode 100644 tensorrt_llm/kv_hints/metrics.py create mode 100644 tensorrt_llm/kv_hints/protocol.py create mode 100644 tensorrt_llm/serve/kv_hints.py create mode 100644 tests/unittest/disaggregated/test_kv_hint_selection.py create mode 100644 tests/unittest/llmapi/test_kv_hints.py diff --git a/docs/source/features/kv-hints.md b/docs/source/features/kv-hints.md new file mode 100644 index 000000000000..91ff1b64c160 --- /dev/null +++ b/docs/source/features/kv-hints.md @@ -0,0 +1,246 @@ +# KV hint request metadata + +The prototype `kv_hints` request field carries a versioned envelope from the +Python LLM API or the OpenAI-compatible completions, chat completions and Responses +endpoints to the PyTorch executor's request metadata. + +This is a transport foundation. Receiving an envelope does **not** execute its +actions, contact the named endpoint, activate a backend, choose a source or change +fallback behavior. Source permissions, required-source semantics and action +execution require a separately agreed consumer integration. Do not rely on this +field to enforce a routing requirement. + +The wire format matches Dynamo's +[KV hint envelope](https://github.com/ai-dynamo/dynamo/blob/3a4c497204c9ff309735a9a052707159cb99c346/lib/kv-router/src/kv_hints.rs). +The following example transports an action without applying it: + +```json +{ + "protocol_version": "0.1", + "message_id": "msg-123", + "actions": [ + { + "action_id": "a1", + "action_type": "kv.fetch", + "action_version": "1.0", + "payload": { + "source_control_endpoint": "tcp://127.0.0.1:23280", + "block_hashes": [11, 22] + } + } + ] +} +``` + +For the Python API, import `KvHintsEnvelope` and `KvHintAction` from either +`tensorrt_llm.llmapi` or `tensorrt_llm.kv_hints`. Pass a typed envelope or raw JSON +object to `llm.generate_async(..., kv_hints=envelope)` or +`llm.generate(..., kv_hints=envelope)`. Python callers are trusted application +code; public applications must establish their own ingress policy before calling +these APIs. Raw envelopes allow admission to ignore a future protocol version +before interpreting that version's fields. + +For batched `generate`, a scalar envelope repeats the **same envelope and +message_id** on every prompt. Use an aligned list for distinct messages or to +omit hints on selected prompts: + +```python +llm.generate([prompt_a, prompt_b], kv_hints=[envelope_a, None]) +``` + +For HTTP, put the JSON object under `kv_hints` in the body (or the client's +`extra_body`). The completions, chat completions (including Harmony) and Responses +paths carry this metadata, including through the disaggregated proxy. + +## Known limitations + +**Current gRPC request paths do not forward `kv_hints` to the executor.** This +affects clients and routers that submit requests to workers over gRPC, including +routers that accept HTTP requests but use gRPC for the worker connection. +TensorRT-LLM emits no hint-specific warning or drop metric for this missing +transport support; the admission counters below do not report this gap. +gRPC parity remains follow-up work. Use a supported HTTP or Python path when +hint transport is required. + +## Ingress and admission + +Worker and disaggregated-proxy HTTP servers strip client hints by default, +**before nested schema validation**, and count non-null stripped hints. Malformed +optional hints therefore cannot fail inference when ingress is disabled. To +accept hints, explicitly set `TRTLLM_TRUST_KV_HINTS=1` in each receiving server +process, including proxy and workers. This process-wide switch is +**prototype-only**. It trusts all traffic that reaches the process and does not +authenticate senders. + +Before enabling a hint consumer in production, require an authenticated, trusted +sender at ingress and enforce the +[endpoint registry check](#candidate-collection-and-consumer-boundary) before any +hint-directed lookup, connection or SDK submission. The switch does not provide +either control. Required P/D handoff keeps its existing authentication and +validation rules. + +At trusted HTTP and Python admission, unknown envelope versions, envelopes +larger than **65,536 UTF-8 JSON bytes (64 KiB)** and malformed hints become no +hint, so inference can continue. The limit applies to the compact serialized +envelope, not character count. Raw supported envelopes are bounded before +nested validation and the validated encoding is checked again. + +Structural errors in supported version `0.1`, such as a missing `message_id` +or non-object action payload, discard the entire envelope and count reason +`malformed`. Non-object hint values and non-JSON Python values follow the same +policy. Admission does not salvage individual actions from an invalid envelope. +Accepted envelopes must also fit the proxy's MessagePack transport. Admission +discards an envelope if any integer in its payload or extra metadata is outside +the range `-2**63` through `2**64 - 1`, inclusive, and counts it as `malformed`. +Admission also checks the serialized envelope with the snapshot's JSON decoder. +Envelopes that exceed the decoder's nesting limit are discarded as `malformed`, +even when they fit the byte limit. Snapshots retain the checked JSON bytes. +This policy applies only to optional hints; invalid request bodies, other +request fields and required P/D handoff parameters retain their validation. +Direct construction of a typed envelope still validates its fields and can +raise an error. Pass a raw JSON object when the future envelope shape is unknown. + +For Dynamo's root-aligned `block_hashes`, each array entry represents a block +starting from the request's first block. With typical decimal `u64` hashes, a +single-action envelope with short metadata holds about **3,100–3,200 hashes**, +or roughly **100,000 tokens of represented prefix at 32 tokens per block**. +This is an estimate, not a fixed token limit. Hash values, identifiers, endpoint +length, extra metadata and additional actions affect the available capacity. + +The size limit bounds the represented cache prefix, not the total prompt length. +A longer prompt can still carry a shorter reusable-prefix hint. If an envelope +exceeds the byte limit, admission discards the entire envelope without truncating +its hash array; inference continues without that hint. Larger-prefix support +requires agreement with the producer on the size budget and representation. + +`trtllm_kv_hints_ignored_total{reason="untrusted|unknown_version|oversized|malformed"}` +counts ignored envelopes at the admission edge where they are discarded. +Rejected metadata is removed before queueing, proxy forwarding or rank +broadcast; deserialization and rank consumers do not recount it. The metric is +available through the server's existing Prometheus metrics configuration. + +Admission freezes validated JSON before queueing a generation request. Caller +mutation cannot change queued metadata. Serialization and the Python-object +broadcast preserve the snapshot across tensor/context and pipeline parallel +ranks and on child requests. `LlmRequest.py_kv_hints` holds the snapshot; decoding +returns an independent copy. Unknown action types/versions and optional JSON +metadata remain intact. + +HTTP request text representations omit `kv_hints`, including when a request is +formatted in an error log. JSON serialization and proxy forwarding retain the +metadata. Explicitly logging a serialized request or the hints themselves still +requires separate redaction. + +## Candidate collection and consumer boundary + +The initial hinted-fetch consumer must ignore hints on requests that use a LoRA +adapter. Orchestration must check request eligibility before candidate +collection. The private helper remains independent of request state. This +restriction applies to optional hinted fetch; it does not change local inference +or required P/D handoff. + +Assembly supplies the supported `(action_type, action_version)` pairs for each +configured backend family. Keep this capability snapshot fixed for the request. +A family is a private integration key for action support. It is not a producer +wire field, a concrete backend binding or an addressing mode. Several bindings +can belong to one family. This does not add a public Backend capability API. + +The private helper collects every supported `kv.fetch` action in producer order. +Type and version strings must match exactly, independently of the envelope +protocol version. Each candidate contains the unordered set of supporting +families. An action appears once even if several families support it. Repeated +action identifiers do not remove producer entries. Empty capabilities or no +matching actions produce no candidates. + +Candidates share the immutable admitted envelope and retain their action +positions. When a consumer needs one candidate, it can obtain an immutable +one-action envelope. This preserves the envelope protocol, message identifier, +extra metadata and complete action, including its identifier, version and opaque +payload, without copying envelope metadata for every retained candidate. + +The helper reports separate counts for unknown action types and fetch actions +with no matching version. Each action is counted once, regardless of the number +of configured families. Orchestration accounts for these counts once per +request. The helper has no global counters and does not count retained +candidates as ignored. + +The initial orchestration policy selects the first candidate. A trusted mapping +must resolve exactly one configured, authorized binding within its supporting +families. Check that binding's actual capability; a family match does not prove +that every binding in the family can process the action. Family registration +order does not select a binding. Missing, ambiguous or unauthorized mappings, +payload validation failure and provider rejection discard optional reuse. They +must not select another candidate, family or binding. Retained candidates are +not an automatic retry list. + +The helper performs no I/O and does not interpret provider payloads. +Action-to-source integration authorizes the source and intent, then supplies +the final immutable backend hint. The holder forwards that value unchanged. +A provider can accept the selected one-action envelope itself as its hint. + +For generation-only requests, a hint opts into named fetch before the paired +source only when the request is eligible and the selected action resolves to +one configured, authorized backend binding. A `kv_hints` field alone does not +enable named fetch. Without a usable hint, keep the existing paired path and +add no named lookup or transfer work. This includes hints discarded by +admission, unsupported actions, excluded requests and missing, ambiguous or +unauthorized bindings. The hint does not replace required P/D handoff. + +Action-to-plan wiring and execution remain integration work. The holder and +shared transport do not interpret provider payloads or contact endpoints. + +Every `kv.fetch` action in `kv_hints` requests KV by content name. Request-paired +handoff uses separate P/D parameters and never comes from these actions. +Orchestration must reject a `kv.fetch` mapping to a request-paired binding before +provider contact. Discard that optional reuse and keep the permitted no-hint +path. A peer worker can serve named KV; the distinction is the addressing mode, +not the worker location or transport library. + +KVCM derives content names locally from the request's own tokens and state. +Producer `block_hashes` guide routing/range selection; they are not content +names or storage object keys. + +Before any hint-directed connection or SDK submission, orchestration's +action-to-source integration must check the selected `kv.fetch` endpoint +against one server-controlled trusted registry. The consumer adapter interprets +the endpoint field; generic envelope admission and the holder keep provider +payloads opaque. If the endpoint is not listed, discard the optional hint and +continue inference. This does not relax required P/D handoff rules. + +The provider must independently validate its action version, payload and +destination before connecting or submitting SDK work. It must not rely solely +on runtime approval. A supplied address never grants permission to contact it. + +Consumer activation requires tests that bypass each check separately. The +runtime check must block an unauthorized destination when provider validation is +bypassed. The provider check must block it when runtime validation is bypassed. +Verify that neither case starts a connection or SDK operation, and that rejection +of optional reuse lets a runnable request continue inference. + +Before consumer activation, map a supported `kv.fetch` action to a request-paired +binding in a test. Verify that orchestration rejects optional reuse without +starting a hint-directed lookup, connection or SDK operation. Verify that +required P/D handoff still uses its original parameters. + +Before consumer activation, test LoRA requests with a supported `kv.fetch` hint. +Verify that orchestration does not collect candidates, open a hint-directed +connection or submit hint-directed SDK work. Verify that eligible requests +without a LoRA adapter can still reach candidate collection. The LoRA gate must +preserve local inference and required P/D handoff behavior. + +Before consumer activation, test generation-only requests with and without a +usable hint. Verify that an eligible request with a selected, resolved and +authorized hint attempts named fetch before the paired source. Cover absent, +discarded and unsupported hints, excluded requests and failed binding resolution; +verify that these cases keep the paired path without added named-fetch work. +Verify that required P/D handoff rules still apply when named fetch succeeds or +fails. + +Before consumer activation, test overlapping family support and several +bindings in one family. Verify that trusted resolution selects one supported, +authorized binding without using registration order. Reject missing or +ambiguous mappings. Verify that authorization, payload or provider rejection +does not select another candidate, family or binding. + +This metadata-only PR defines these requirements; consumer integration must +implement and test them before activation. diff --git a/docs/source/index.rst b/docs/source/index.rst index 9920583eb754..d22ce9317121 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -71,6 +71,7 @@ Welcome to TensorRT LLM's Documentation! features/subagent-routing.md features/embeddings.md features/kvcache.md + features/kv-hints.md features/kv-cache-compression.md features/long-sequence.md features/lora.md diff --git a/tensorrt_llm/_torch/disaggregation/orchestration/_kv_hints.py b/tensorrt_llm/_torch/disaggregation/orchestration/_kv_hints.py new file mode 100644 index 000000000000..edcd71fd6860 --- /dev/null +++ b/tensorrt_llm/_torch/disaggregation/orchestration/_kv_hints.py @@ -0,0 +1,108 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Private capability matching for hint-to-source orchestration.""" + +from collections.abc import Mapping +from dataclasses import dataclass, field + +from tensorrt_llm.kv_hints.protocol import KvHintsSnapshot + + +@dataclass(frozen=True, slots=True) +class _KvFetchCandidate: + """One supported action with its possible backend families. + + Candidates share the immutable input envelope and retain action positions. + Materialize a one-action snapshot only when a consumer needs it, so retained + candidates do not each copy the envelope metadata. + + Attributes: + families: Families that support this action's type and version. This + unordered set does not select or authorize a concrete binding. + """ + + families: frozenset[str] + _envelope_snapshot: KvHintsSnapshot = field(repr=False) + _action_index: int + + def snapshot(self) -> KvHintsSnapshot: + """Materialize this action without exposing mutable retained state. + + Returns: + An immutable one-action envelope with all envelope and action + metadata preserved. Each call returns an independent snapshot. + """ + envelope = self._envelope_snapshot.envelope() + envelope.actions = [envelope.actions[self._action_index]] + return KvHintsSnapshot(envelope.model_dump_json().encode("utf-8")) + + +@dataclass(frozen=True, slots=True) +class _KvFetchCandidates: + """Capability matches and per-action counts for one admitted envelope. + + Attributes: + candidates: All supported fetch actions in producer order, one candidate + per action even when several families support it. + unknown_actions: Actions whose type is not recognized by this consumer. + unsupported_version_actions: Fetch actions whose exact type/version pair + is not supported by any configured family. + """ + + candidates: tuple[_KvFetchCandidate, ...] + unknown_actions: int = 0 + unsupported_version_actions: int = 0 + + +def _collect_kv_fetch_candidates( + snapshot: KvHintsSnapshot | None, + *, + family_capabilities: Mapping[str, frozenset[tuple[str, str]]], +) -> _KvFetchCandidates: + """Collect supported fetch actions without selecting a source or binding. + + Args: + snapshot: Immutable admitted metadata. Orchestration collects candidates + once per request and accounts for the returned counts. + family_capabilities: Assembly-supplied family IDs mapped to supported + (action_type, action_version) pairs. Keep this capability snapshot + fixed for the request. A family can have multiple bindings; matching + a family does not prove a specific binding's capability or authority. + + Returns: + All supported fetch actions in producer order, with unordered family + sets and disjoint unknown/unsupported counts. Count each action once, + regardless of the number of matching families. Empty capabilities match + nothing. Action identifiers do not deduplicate producer entries. + + Orchestration applies the first-candidate policy and resolves one trusted + binding before provider contact. Authorization, payload validation or + provider rejection must not trigger another candidate, family or binding. + This helper interprets no provider payloads and performs no I/O or global + accounting. It does not classify retained candidates as ignored. + """ + if snapshot is None: + return _KvFetchCandidates(()) + capabilities = tuple( + (family, frozenset(actions)) for family, actions in family_capabilities.items() + ) + envelope = snapshot.envelope() + candidates = [] + unknown_actions = 0 + unsupported_version_actions = 0 + for index, action in enumerate(envelope.actions): + if action.action_type != "kv.fetch": + unknown_actions += 1 + continue + action_key = (action.action_type, action.action_version) + families = frozenset(family for family, actions in capabilities if action_key in actions) + if not families: + unsupported_version_actions += 1 + continue + candidates.append(_KvFetchCandidate(families, snapshot, index)) + return _KvFetchCandidates( + candidates=tuple(candidates), + unknown_actions=unknown_actions, + unsupported_version_actions=unsupported_version_actions, + ) diff --git a/tensorrt_llm/_torch/pyexecutor/llm_request.py b/tensorrt_llm/_torch/pyexecutor/llm_request.py index a4e42209a2de..87c273f00a3a 100644 --- a/tensorrt_llm/_torch/pyexecutor/llm_request.py +++ b/tensorrt_llm/_torch/pyexecutor/llm_request.py @@ -17,6 +17,7 @@ from tensorrt_llm.inputs.multimodal import (MultimodalEncoderRequestState, MultimodalParams, get_multimodal_item_metadata) +from tensorrt_llm.kv_hints.protocol import KvHintsSnapshot from tensorrt_llm.sampling_params import LogprobMode SamplingConfig = tllm_executor.SamplingConfig @@ -865,6 +866,7 @@ def __init__( self.py_original_end_id = self.py_end_id self.py_bad_words = None self.py_no_repeat_ngram_size = None + self.py_kv_hints: Optional[KvHintsSnapshot] = None self.py_num_connector_matched_tokens = 0 @@ -1256,6 +1258,7 @@ def create_child_request(self, child_id: int) -> None: child.py_original_end_id = self.py_original_end_id child.py_disaggregated_params = deepcopy(self.py_disaggregated_params) child.py_conversation_params = deepcopy(self.py_conversation_params) + child.py_kv_hints = self.py_kv_hints child.py_uncached_prefill_tokens = self.py_uncached_prefill_tokens child.py_disable_speculative_decoding = self.py_disable_speculative_decoding child.is_attention_dp_dummy = self.is_attention_dp_dummy @@ -1753,6 +1756,7 @@ def executor_request_to_llm_request( # Published by the attention-DP router before this conversion, if at all. llm_request.py_uncached_prefill_tokens = getattr( executor_request, "py_uncached_prefill_tokens", None) + llm_request.py_kv_hints = getattr(executor_request, "py_kv_hints", None) if child_req_ids: for child_id in child_req_ids: llm_request.create_child_request(child_id) diff --git a/tensorrt_llm/_torch/pyexecutor/request_utils.py b/tensorrt_llm/_torch/pyexecutor/request_utils.py index aa8f5636bd6c..bb9f33850c42 100644 --- a/tensorrt_llm/_torch/pyexecutor/request_utils.py +++ b/tensorrt_llm/_torch/pyexecutor/request_utils.py @@ -528,6 +528,7 @@ def _collect_py_objects(self, new_requests: List) -> Tuple: "py_logprobs_simple_format", "py_return_routed_experts", "py_arrival_time", + "py_kv_hints", ) ] py_objects.extend( diff --git a/tensorrt_llm/executor/base_worker.py b/tensorrt_llm/executor/base_worker.py index 92c3cfe1bd54..8539e719a1c0 100644 --- a/tensorrt_llm/executor/base_worker.py +++ b/tensorrt_llm/executor/base_worker.py @@ -522,6 +522,8 @@ def _deduce_max_tokens(request: GenerationRequest) -> int: if request.conversation_params is not None: executor_request.py_conversation_params = request.conversation_params + executor_request.py_kv_hints = request.kv_hints + if request.arrival_time is not None: executor_request.py_arrival_time = request.arrival_time diff --git a/tensorrt_llm/executor/executor.py b/tensorrt_llm/executor/executor.py index 48ecac324a97..3e41a1b31df1 100644 --- a/tensorrt_llm/executor/executor.py +++ b/tensorrt_llm/executor/executor.py @@ -21,6 +21,7 @@ import torch from tensorrt_llm.inputs.multimodal import MultimodalParams +from tensorrt_llm.kv_hints import KvHintsInput from tensorrt_llm.logger import logger, set_level from .._utils import mpi_world_size @@ -132,6 +133,7 @@ def generate_async( encoder_input_token_ids: Optional[Union[torch.Tensor, np.ndarray, list]] = None, priority: float = DEFAULT_REQUEST_PRIORITY, + kv_hints: Optional[KvHintsInput] = None, ) -> GenerationResult: """Generate output for the given prompt token ids in the asynchronous mode. Asynchronous generation accepts single prompt only. @@ -161,7 +163,8 @@ def generate_async( cache_salt=cache_salt, arrival_time=arrival_time, encoder_input_token_ids=encoder_input_token_ids, - priority=priority) + priority=priority, + kv_hints=kv_hints) result = self.submit(request) # release memory in time if hasattr(request, "multimodal_params"): diff --git a/tensorrt_llm/executor/request.py b/tensorrt_llm/executor/request.py index 65aa5a153349..aa3c7f69a9ee 100644 --- a/tensorrt_llm/executor/request.py +++ b/tensorrt_llm/executor/request.py @@ -1,3 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + import os from collections.abc import Mapping from dataclasses import dataclass @@ -7,6 +10,8 @@ import torch from tensorrt_llm.inputs.multimodal import MultimodalParams +from tensorrt_llm.kv_hints import KvHintsInput +from tensorrt_llm.kv_hints.protocol import snapshot_kv_hints from ..conversation_params import ConversationParams from ..disaggregated_params import DisaggregatedParams @@ -113,6 +118,7 @@ def __init__( encoder_input_token_ids: Optional[Union[torch.Tensor, np.ndarray, list]] = None, priority: float = DEFAULT_REQUEST_PRIORITY, + kv_hints: Optional[KvHintsInput] = None, ): if isinstance(prompt_token_ids, list): self.prompt_token_ids = prompt_token_ids @@ -150,6 +156,7 @@ def __init__( self.trace_headers = trace_headers self.scheduling_params = scheduling_params self.conversation_params = conversation_params + self.kv_hints = snapshot_kv_hints(kv_hints) if cache_salt is not None: if not isinstance(cache_salt, str): raise TypeError( @@ -251,6 +258,7 @@ def __getstate__(self): return state def __setstate__(self, state): + state.setdefault("kv_hints", None) buf = None pt = state.get("_prompt_token_ids") if type(pt) is tuple and len( diff --git a/tensorrt_llm/kv_hints/__init__.py b/tensorrt_llm/kv_hints/__init__.py new file mode 100644 index 000000000000..575a827f29ae --- /dev/null +++ b/tensorrt_llm/kv_hints/__init__.py @@ -0,0 +1,8 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Versioned, opaque request metadata for KV-cache consumers.""" + +from .protocol import KvHintAction, KvHintsEnvelope, KvHintsInput + +__all__ = ["KvHintAction", "KvHintsEnvelope", "KvHintsInput"] diff --git a/tensorrt_llm/kv_hints/metrics.py b/tensorrt_llm/kv_hints/metrics.py new file mode 100644 index 000000000000..30a7384adec5 --- /dev/null +++ b/tensorrt_llm/kv_hints/metrics.py @@ -0,0 +1,34 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Admission metrics, initialized after serving configures multiprocessing.""" + +from threading import Lock +from typing import TYPE_CHECKING, Literal + +if TYPE_CHECKING: + from prometheus_client import Counter + +_counter: "Counter | None" = None +_counter_lock = Lock() + + +def count_ignored_kv_hints( + reason: Literal["untrusted", "unknown_version", "oversized", "malformed"], +) -> None: + """Count one ignored envelope at its admission edge, before rank broadcast. + + Args: + reason: Fixed rejection category; never contains producer metadata. + """ + global _counter + with _counter_lock: + if _counter is None: + from prometheus_client import Counter + + _counter = Counter( + "trtllm_kv_hints_ignored_total", + "Optional KV hint envelopes ignored at request ingress.", + ("reason",), + ) + _counter.labels(reason=reason).inc() diff --git a/tensorrt_llm/kv_hints/protocol.py b/tensorrt_llm/kv_hints/protocol.py new file mode 100644 index 000000000000..d45836843200 --- /dev/null +++ b/tensorrt_llm/kv_hints/protocol.py @@ -0,0 +1,157 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""KV hint wire envelopes and immutable request-admission snapshots. + +The envelope validates transport metadata. Action payloads stay opaque; accepting +an envelope does not activate a backend, perform a fetch, or apply source policy. +""" + +import json +from dataclasses import dataclass +from typing import TypeAlias + +import msgspec +from pydantic import BaseModel, ConfigDict, Field, JsonValue + +from .metrics import count_ignored_kv_hints + +__all__ = [ + "KvHintAction", + "KvHintsEnvelope", + "KvHintsInput", + "KvHintsSnapshot", + "admit_kv_hints", + "snapshot_kv_hints", +] + +MAX_KV_HINTS_BYTES = 64 * 1024 +_SUPPORTED_PROTOCOL_VERSION = "0.1" + + +class KvHintAction(BaseModel): + """One producer-versioned action carried without interpreting its payload.""" + + model_config = ConfigDict(extra="allow", allow_inf_nan=False, revalidate_instances="always") + __pydantic_extra__: dict[str, JsonValue] + + action_id: str = Field(min_length=1, description="Producer's action identifier.") + action_type: str = Field(min_length=1, description="Opaque action type, such as kv.fetch.") + action_version: str = Field(min_length=1, description="Producer's action schema version.") + payload: dict[str, JsonValue] = Field( + description="Opaque JSON object owned by the action consumer." + ) + + +class KvHintsEnvelope(BaseModel): + """Versioned KV hints transported as prototype per-request metadata. + + Unknown action types, action versions and optional JSON fields are preserved. Consumers must + negotiate and validate executable actions separately from this wire format. + """ + + model_config = ConfigDict(extra="allow", allow_inf_nan=False, revalidate_instances="always") + __pydantic_extra__: dict[str, JsonValue] + + protocol_version: str = Field( + min_length=1, + description="KV hint envelope protocol version; unsupported versions are ignored at admission.", + ) + message_id: str = Field(min_length=1, description="Producer's hint message identifier.") + actions: list[KvHintAction] = Field(description="Actions in their producer-specified order.") + + +KvHintsInput: TypeAlias = KvHintsEnvelope | dict[str, JsonValue] + + +@dataclass(frozen=True, slots=True) +class KvHintsSnapshot: + """Immutable admitted metadata, including every nested payload value. + + Attributes: + wire_json: Validated envelope encoded as UTF-8 JSON bytes. + """ + + wire_json: bytes + + def envelope(self) -> KvHintsEnvelope: + """Decode an independent copy for a consumer without exposing mutable state. + + Returns: + A validated envelope whose mutation cannot change this snapshot. + """ + return KvHintsEnvelope.model_validate_json(self.wire_json) + + +def admit_kv_hints(hints: KvHintsInput | None) -> KvHintsEnvelope | None: + """Apply version, size, schema and proxy transport checks to optional hints. + + Args: + hints: A typed envelope or raw JSON object from a trusted caller. + + Returns: + A validated envelope, or None after counting an unsupported, oversized + or malformed hint. Unsupported raw versions need not follow the supported + schema. Direct construction of a typed envelope still validates its fields. + """ + admitted = _admit_kv_hints(hints) + return None if admitted is None else admitted[0] + + +def _admit_kv_hints(hints: KvHintsInput | None) -> tuple[KvHintsEnvelope, bytes] | None: + """Return the admitted model and its checked, decodable JSON encoding.""" + if hints is None: + return None + if isinstance(hints, KvHintsEnvelope): + version = getattr(hints, "protocol_version", None) + elif isinstance(hints, dict): + version = hints.get("protocol_version") + else: + count_ignored_kv_hints("malformed") + return None + if isinstance(version, str) and version and version != _SUPPORTED_PROTOCOL_VERSION: + count_ignored_kv_hints("unknown_version") + return None + # Raw HTTP/Python input is bounded before nested validation. The validated + # encoding below enforces the same limit after any model normalization. + if isinstance(hints, dict): + try: + wire_json = json.dumps( + hints, ensure_ascii=False, separators=(",", ":"), allow_nan=False + ).encode("utf-8") + except (TypeError, ValueError, RecursionError): + count_ignored_kv_hints("malformed") + return None + if len(wire_json) > MAX_KV_HINTS_BYTES: + count_ignored_kv_hints("oversized") + return None + try: + validated = KvHintsEnvelope.model_validate(hints) + wire_json = validated.model_dump_json().encode("utf-8") + except (TypeError, ValueError, RecursionError): + count_ignored_kv_hints("malformed") + return None + if len(wire_json) > MAX_KV_HINTS_BYTES: + count_ignored_kv_hints("oversized") + return None + try: + validated = KvHintsEnvelope.model_validate_json(wire_json) + msgspec.msgpack.encode(validated.model_dump(mode="json")) + except (OverflowError, TypeError, ValueError, RecursionError): + count_ignored_kv_hints("malformed") + return None + return validated, wire_json + + +def snapshot_kv_hints(hints: KvHintsInput | None) -> KvHintsSnapshot | None: + """Validate and freeze optional hints before queueing a generation request. + + Args: + hints: The optional producer envelope or raw JSON object. + + Returns: + Immutable UTF-8 JSON, or None for absent or ignored hints. Malformed + metadata is counted and discarded before the request is queued. + """ + admitted = _admit_kv_hints(hints) + return None if admitted is None else KvHintsSnapshot(admitted[1]) diff --git a/tensorrt_llm/llmapi/__init__.py b/tensorrt_llm/llmapi/__init__.py index 8d1dfd245852..9bd914d608b8 100644 --- a/tensorrt_llm/llmapi/__init__.py +++ b/tensorrt_llm/llmapi/__init__.py @@ -5,6 +5,7 @@ from ..conversation_params import ConversationParams from ..disaggregated_params import DisaggregatedParams, DisaggScheduleStyle from ..executor import CompletionOutput, LoRARequest, RequestError +from ..kv_hints import KvHintAction, KvHintsEnvelope, KvHintsInput from ..sampling_params import GuidedDecodingParams, SamplingParams from ..scheduling_params import SchedulingParams from .llm import LLM, RequestOutput @@ -50,6 +51,9 @@ 'RequestOutput', 'GuidedDecodingParams', 'SamplingParams', + 'KvHintAction', + 'KvHintsEnvelope', + 'KvHintsInput', 'DisaggregatedParams', 'ConversationParams', 'DisaggScheduleStyle', diff --git a/tensorrt_llm/llmapi/llm.py b/tensorrt_llm/llmapi/llm.py index ca6bec788a23..df9f32cd75fd 100644 --- a/tensorrt_llm/llmapi/llm.py +++ b/tensorrt_llm/llmapi/llm.py @@ -33,6 +33,7 @@ from tensorrt_llm._utils import mpi_disabled from tensorrt_llm.inputs.multimodal import MultimodalParams from tensorrt_llm.inputs.registry import BaseMultimodalInputProcessor +from tensorrt_llm.kv_hints import KvHintsInput from tensorrt_llm.llmapi import tracing from tensorrt_llm.metrics.enums import MetricNames @@ -707,6 +708,8 @@ def generate( List[ConversationParams]]] = None, cache_salt: Optional[Union[str, Sequence[str]]] = None, priority: Union[float, List[float]] = DEFAULT_REQUEST_PRIORITY, + kv_hints: Optional[Union[KvHintsInput, + List[Optional[KvHintsInput]]]] = None, ) -> Union[RequestOutput, List[RequestOutput]]: """Generate output for the given prompts in the synchronous mode. Synchronous generation accepts either single prompt or batched prompts. @@ -731,6 +734,7 @@ def generate( Conversation parameters. Defaults to None. cache_salt (str, Sequence[str], optional): If specified, KV cache will be salted with the provided string to limit the kv cache reuse to the requests with the same string. Defaults to None. priority (float, List[float]): The scheduling priority for the request(s), in the range [0, 1]. Higher values indicate higher priority. Defaults to 0.5. + kv_hints (tensorrt_llm.kv_hints.protocol.KvHintsInput, List[Optional[tensorrt_llm.kv_hints.protocol.KvHintsInput]], optional): Prototype typed or raw JSON KV hint metadata. A scalar repeats the same envelope and message_id for every prompt; an aligned list supplies per-prompt hints. Unknown versions, oversized envelopes and malformed hints are ignored at admission. Transporting hints does not execute actions. Defaults to None. Returns: Union[tensorrt_llm.llmapi.llm.RequestOutput, List[tensorrt_llm.llmapi.llm.RequestOutput]]: The output data of the completion request to the LLM. @@ -748,6 +752,11 @@ def generate( else: request_inputs_list = [prompt_inputs(i) for i in inputs] + if isinstance(kv_hints, + list) and len(kv_hints) != len(request_inputs_list): + raise ValueError( + "kv_hints list length must match the number of prompts") + if isinstance(priority, list): if len(priority) != len(request_inputs_list): raise ValueError( @@ -776,6 +785,7 @@ def generate( conversation_params=self._item_at(conversation_params, i), cache_salt=self._item_at(cache_salt, i), priority=self._item_at(priority, i), + kv_hints=self._item_at(kv_hints, i), streaming=False, ) futures.append(future) @@ -807,6 +817,7 @@ def generate_async( conversation_params: Optional[ConversationParams] = None, cache_salt: Optional[str] = None, priority: float = DEFAULT_REQUEST_PRIORITY, + kv_hints: Optional[KvHintsInput] = None, ) -> RequestOutput: """Generate output for the given prompt in the asynchronous mode. Asynchronous generation accepts single prompt only. @@ -825,6 +836,7 @@ def generate_async( conversation_params (tensorrt_llm.conversation_params.ConversationParams, optional): Conversation parameters. Defaults to None. cache_salt (str, optional): If specified, KV cache will be salted with the provided string to limit the kv cache reuse to the requests with the same string. Defaults to None. priority (float): The scheduling priority for the request, in the range [0, 1]. Higher values indicate higher priority. Defaults to 0.5. + kv_hints (tensorrt_llm.kv_hints.protocol.KvHintsInput, optional): Prototype typed or raw JSON KV hint metadata from a trusted caller. Unknown versions, oversized envelopes and malformed hints are ignored at admission. Payloads are transported without executing actions or selecting a source. Defaults to None. Returns: tensorrt_llm.llmapi.llm.RequestOutput: The output data of the completion request to the LLM. @@ -897,6 +909,7 @@ def generate_async( arrival_time=arrival_time, encoder_input_token_ids=encoder_input_token_ids, priority=priority, + kv_hints=kv_hints, ) if sampling_params.return_perf_metrics: diff --git a/tensorrt_llm/serve/kv_hints.py b/tensorrt_llm/serve/kv_hints.py new file mode 100644 index 000000000000..0ad606e0a426 --- /dev/null +++ b/tensorrt_llm/serve/kv_hints.py @@ -0,0 +1,61 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Server-controlled hint trust policy applied before request model validation.""" + +from collections.abc import Awaitable, Callable + +from fastapi import Request, Response +from fastapi.routing import APIRoute + +from tensorrt_llm.kv_hints.metrics import count_ignored_kv_hints + + +class KvHintsRequest(Request): + """Request that strips untrusted optional hints before nested validation.""" + + def filter_kv_hints(self, body: object) -> object: + """Strip and count hints unless this server process trusts its ingress. + + Args: + body: Decoded request body; non-object bodies keep normal validation. + + Returns: + The decoded body with untrusted hints removed. Required P/D handoff + parameters and their independent authentication remain unchanged. + """ + if ( + isinstance(body, dict) + and "kv_hints" in body + and not getattr(self.app.state, "trust_kv_hints", False) + ): + if body.pop("kv_hints") is not None: + count_ignored_kv_hints("untrusted") + return body + + async def json(self) -> object: + """Decode JSON and apply the server's trust policy before validation. + + Returns: + The filtered JSON body, cached by the underlying request. + """ + return self.filter_kv_hints(await super().json()) + + +class KvHintsRoute(APIRoute): + """Apply hint ingress policy to typed FastAPI endpoints.""" + + request_type = KvHintsRequest + + def get_route_handler(self) -> Callable[[Request], Awaitable[Response]]: + """Wrap FastAPI validation with the configured request decoder. + + Returns: + An ASGI request handler that filters hints before model validation. + """ + original_route_handler = super().get_route_handler() + + async def route_handler(request: Request) -> Response: + return await original_route_handler(self.request_type(request.scope, request.receive)) + + return route_handler diff --git a/tensorrt_llm/serve/openai_disagg_server.py b/tensorrt_llm/serve/openai_disagg_server.py index 686612b2aa81..518839540d72 100644 --- a/tensorrt_llm/serve/openai_disagg_server.py +++ b/tensorrt_llm/serve/openai_disagg_server.py @@ -16,6 +16,7 @@ # yapf: disable import asyncio import json +import os import signal import socket import traceback @@ -56,6 +57,7 @@ resolve_request_conversation_id) from tensorrt_llm.serve.disagg_coordinator import (CoordinatorClient, DisaggCoordinatorService) +from tensorrt_llm.serve.kv_hints import KvHintsRoute from tensorrt_llm.serve.openai_client import OpenAIClient, OpenAIHttpClient from tensorrt_llm.serve.openai_disagg_service import ( OpenAIDisaggregatedService, ResponseHooks) @@ -290,6 +292,8 @@ async def lifespan(app) -> None: await self._perf_metrics_writer.close() self.app = FastAPI(lifespan=lifespan) + self.app.state.trust_kv_hints = os.getenv("TRTLLM_TRUST_KV_HINTS", "0") == "1" + self.app.router.route_class = KvHintsRoute if self._collect_perf_metrics: self.app.add_middleware( diff --git a/tensorrt_llm/serve/openai_protocol.py b/tensorrt_llm/serve/openai_protocol.py index 11c2ed3fc7dd..4c102f0da3b0 100644 --- a/tensorrt_llm/serve/openai_protocol.py +++ b/tensorrt_llm/serve/openai_protocol.py @@ -54,6 +54,8 @@ 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, @@ -151,6 +153,20 @@ 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. @@ -639,6 +655,12 @@ def _record_sampling_params_request_fields( 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 @@ -1066,6 +1088,12 @@ class ChatCompletionThinkingParam(OpenAIBaseModel): 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] @@ -1501,6 +1529,12 @@ def _materialize_validator_iterators(value, _depth=0): 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 # diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py index e3956e67c874..099002f6a114 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -31,7 +31,6 @@ from fastapi.exceptions import RequestValidationError from fastapi.responses import (FileResponse, JSONResponse, Response, StreamingResponse) -from fastapi.routing import APIRoute from pydantic import ValidationError from starlette.routing import Mount from transformers import AutoProcessor @@ -157,6 +156,7 @@ 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 @@ -181,7 +181,7 @@ def _is_visual_gen_instance(obj) -> bool: _msgpack_decoder = msgspec.msgpack.Decoder() -class _MsgspecRequest(Request): +class _MsgspecRequest(KvHintsRequest): """Request that decodes X-TRTLLM-Msgpack bodies with msgspec.""" async def json(self): @@ -196,20 +196,13 @@ async def json(self): raise json.JSONDecodeError(f"msgpack: {e}", "", 0) from e else: self._json_body = json.loads(body) - return self._json_body + return self.filter_kv_hints(self._json_body) -class _MsgspecRoute(APIRoute): - """APIRoute that parses request bodies via :class:`_MsgspecRequest`.""" +class _MsgspecRoute(KvHintsRoute): + """Parse msgpack/JSON and apply hint trust policy before model validation.""" - def get_route_handler(self): - original_route_handler = super().get_route_handler() - - async def route_handler(request: Request): - return await original_route_handler( - _MsgspecRequest(request.scope, request.receive)) - - return route_handler + request_type = _MsgspecRequest TIMEOUT_KEEP_ALIVE = 5 # seconds. @@ -883,6 +876,8 @@ async def lifespan(app: FastAPI): cancel_futures=True) self.app = FastAPI(lifespan=lifespan) + self.app.state.trust_kv_hints = os.getenv("TRTLLM_TRUST_KV_HINTS", + "0") == "1" self.app.router.route_class = _MsgspecRoute @self.app.exception_handler(RequestValidationError) @@ -2496,6 +2491,7 @@ async def chat_stream_generator( disaggregated_params=disaggregated_params, conversation_params=conversation_params, cache_salt=request.cache_salt, + kv_hints=request.kv_hints, trace_headers=trace_headers, scheduling_params=scheduling_params, priority=request.priority @@ -3135,6 +3131,7 @@ async def generator_wrapper(generator: AsyncIterator[Any]): disaggregated_params=disaggregated_params, conversation_params=conversation_params, trace_headers=trace_headers, + kv_hints=request.kv_hints, priority=request.priority if request.priority is not None else DEFAULT_REQUEST_PRIORITY) asyncio.create_task( @@ -3305,6 +3302,7 @@ async def create_streaming_generator(promise: RequestOutput, disaggregated_params=disaggregated_params, conversation_params=conversation_params, trace_headers=trace_headers, + kv_hints=request.kv_hints, priority=request.priority if request.priority is not None else DEFAULT_REQUEST_PRIORITY, ) @@ -3514,6 +3512,7 @@ async def create_streaming_generator(promise: RequestOutput, streaming=request.stream, disaggregated_params=disaggregated_params, conversation_params=conversation_params, + kv_hints=request.kv_hints, _postproc_params=postproc_params if self.postproc_worker_enabled else None, ) diff --git a/tests/integration/defs/perf/agentx_client.py b/tests/integration/defs/perf/agentx_client.py index 14244ab486cf..717cfc6d3766 100644 --- a/tests/integration/defs/perf/agentx_client.py +++ b/tests/integration/defs/perf/agentx_client.py @@ -215,9 +215,17 @@ def aiperf_is_usable() -> bool: f"if int(datasets.__version__.split('.')[0]) < {MIN_DATASETS_MAJOR}:\n" f" sys.exit(f'datasets {{datasets.__version__}} < {MIN_DATASETS_MAJOR}')\n" ) - result = subprocess.run( - [sys.executable, "-c", probe], stdout=subprocess.DEVNULL, stderr=subprocess.PIPE, text=True - ) + try: + result = subprocess.run( + [sys.executable, "-c", probe], + stdout=subprocess.DEVNULL, + stderr=subprocess.PIPE, + text=True, + timeout=120, + ) + except subprocess.TimeoutExpired as error: + _log(f"aiperf stack not usable: import probe timed out after {error.timeout}s") + return False if result.returncode != 0: lines = result.stderr.strip().splitlines() _log(f"aiperf stack not usable: {lines[-1] if lines else f'exit {result.returncode}'}") diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index a5655df8c510..e980badbf11f 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -276,6 +276,7 @@ l0_cpu: - unittest/llmapi/test_gc_utils.py - unittest/llmapi/test_whisper_suppress_tokens_processor.py - unittest/llmapi/test_kv_cache_dtype_override.py + - unittest/llmapi/test_kv_hints.py - unittest/llmapi/test_llm_args.py - unittest/llmapi/test_mx_args.py - unittest/llmapi/test_custom_tokenizer_aliases.py diff --git a/tests/unittest/_torch/visual_gen/test_attention_importable.py b/tests/unittest/_torch/visual_gen/test_attention_importable.py index 28b186bb22fe..356ea2731693 100644 --- a/tests/unittest/_torch/visual_gen/test_attention_importable.py +++ b/tests/unittest/_torch/visual_gen/test_attention_importable.py @@ -58,7 +58,7 @@ def test_import_order_is_free(first, second): f"import {MODULES} as m, {ATTENTION} as a; " "assert m.Attention is a.Attention and m.QKVMode is a.QKVMode" ) - subprocess.run([sys.executable, "-c", code], check=True) + subprocess.run([sys.executable, "-c", code], check=True, timeout=120) def test_modules_reexport_is_canonical(): diff --git a/tests/unittest/api_stability/api_stability_core.py b/tests/unittest/api_stability/api_stability_core.py index 66fd496635b1..8264f810b8b3 100644 --- a/tests/unittest/api_stability/api_stability_core.py +++ b/tests/unittest/api_stability/api_stability_core.py @@ -20,7 +20,7 @@ import transformers import yaml from annotated_types import Gt -from pydantic import BaseModel +from pydantic import BaseModel, JsonValue import tensorrt_llm from tensorrt_llm import LLM, DisaggregatedParams diff --git a/tests/unittest/api_stability/references/llm.yaml b/tests/unittest/api_stability/references/llm.yaml index e3e0e82ddcaa..98ec12dbb0a3 100644 --- a/tests/unittest/api_stability/references/llm.yaml +++ b/tests/unittest/api_stability/references/llm.yaml @@ -401,6 +401,10 @@ methods: return_annotation: None generate: parameters: + kv_hints: + annotation: Union[tensorrt_llm.kv_hints.protocol.KvHintsEnvelope, dict[str, JsonValue], List[Union[tensorrt_llm.kv_hints.protocol.KvHintsEnvelope, dict[str, JsonValue], NoneType]], NoneType] + default: null + status: prototype disaggregated_params: annotation: Union[tensorrt_llm.disaggregated_params.DisaggregatedParams, Sequence[tensorrt_llm.disaggregated_params.DisaggregatedParams], NoneType] default: null @@ -423,6 +427,10 @@ methods: return_annotation: Union[tensorrt_llm.llmapi.llm.RequestOutput, List[tensorrt_llm.llmapi.llm.RequestOutput]] generate_async: parameters: + kv_hints: + annotation: Union[tensorrt_llm.kv_hints.protocol.KvHintsEnvelope, dict[str, JsonValue], NoneType] + default: null + status: prototype disaggregated_params: annotation: Optional[tensorrt_llm.disaggregated_params.DisaggregatedParams] default: null diff --git a/tests/unittest/api_stability/references/trtllm_serve_api.yaml b/tests/unittest/api_stability/references/trtllm_serve_api.yaml index ed978e75c665..c1a46e6d7e64 100644 --- a/tests/unittest/api_stability/references/trtllm_serve_api.yaml +++ b/tests/unittest/api_stability/references/trtllm_serve_api.yaml @@ -55,6 +55,12 @@ models: CompletionRequest: fields: + kv_hints: + kind: extension + type: Optional[KvHintsEnvelope] + default: null + status: prototype + required: false model: kind: openai type: str @@ -508,6 +514,12 @@ models: ChatCompletionRequest: fields: + kv_hints: + kind: extension + type: Optional[KvHintsEnvelope] + default: null + status: prototype + required: false messages: kind: openai type: List[ChatCompletionMessageParam] @@ -1045,6 +1057,12 @@ models: ResponsesRequest: fields: + kv_hints: + kind: extension + type: Optional[KvHintsEnvelope] + default: null + status: prototype + required: false background: kind: openai type: Optional[bool] diff --git a/tests/unittest/disaggregated/test_kv_hint_selection.py b/tests/unittest/disaggregated/test_kv_hint_selection.py new file mode 100644 index 000000000000..5758bf0af66e --- /dev/null +++ b/tests/unittest/disaggregated/test_kv_hint_selection.py @@ -0,0 +1,224 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Private candidate collection preserves admitted metadata without side effects.""" + +import tracemalloc +from dataclasses import FrozenInstanceError +from unittest.mock import patch + +import pytest + +from tensorrt_llm._torch.disaggregation.orchestration._kv_hints import _collect_kv_fetch_candidates +from tensorrt_llm.kv_hints.protocol import snapshot_kv_hints + +pytestmark = pytest.mark.cpu_only + + +def _wire() -> dict: + """Return the producer's kv.fetch fixture plus an opaque future action.""" + return { + "protocol_version": "0.1", + "message_id": "msg-123", + "actions": [ + { + "action_id": "a1", + "action_type": "kv.fetch", + "action_version": "1.0", + "payload": { + "source_control_endpoint": "tcp://127.0.0.1:23280", + "block_hashes": [11, 22], + }, + }, + { + "action_id": "future", + "action_type": "kv.future_action", + "action_version": "7.3", + "payload": {"nested": {"enabled": True}, "items": [1, None, "x"]}, + }, + ], + } + + +def test_candidates_preserve_order_metadata_and_duplicate_action_ids() -> None: + """Each supported action keeps its position and metadata, even when IDs repeat.""" + wire = _wire() + first = wire["actions"][0] + first["optional_field"] = {"keep": [1, 2]} + later = dict(first, payload={"block_hashes": [33], "source_control_endpoint": "opaque"}) + wire["producer_metadata"] = {"trace": ["router"]} + wire["actions"] = [wire["actions"][1], first, later] + snapshot = snapshot_kv_hints(wire) + capabilities = {"store": frozenset({("kv.fetch", "1.0")})} + + with patch("tensorrt_llm.kv_hints.protocol.count_ignored_kv_hints") as count: + result = _collect_kv_fetch_candidates(snapshot, family_capabilities=capabilities) + materialized = [ + candidate.snapshot().envelope().model_dump() for candidate in result.candidates + ] + + count.assert_not_called() + assert isinstance(result.candidates, tuple) + assert materialized == [dict(wire, actions=[first]), dict(wire, actions=[later])] + assert all(candidate.families == frozenset({"store"}) for candidate in result.candidates) + assert result.unknown_actions == 1 + assert result.unsupported_version_actions == 0 + assert snapshot.envelope().model_dump() == wire + + +def test_family_capabilities_match_types_versions_and_count_each_action_once() -> None: + """Overlapping family support does not duplicate candidates or ignored counts.""" + wire = _wire() + fetch, future = wire["actions"] + versions = ["2.0", "1.0", "3.0", "1.00", "1.0"] + wire["actions"] = [ + dict(fetch, action_id=f"fetch-{index}", action_version=version) + for index, version in enumerate(versions) + ] + [future] + capabilities = { + "store": frozenset({("kv.fetch", "1.0"), ("kv.publish", "3.0")}), + "peer": frozenset({("kv.fetch", "1.0"), ("kv.fetch", "2.0")}), + "future": frozenset({("kv.future_action", "7.3")}), + } + snapshot = snapshot_kv_hints(wire) + + result = _collect_kv_fetch_candidates(snapshot, family_capabilities=capabilities) + reversed_result = _collect_kv_fetch_candidates( + snapshot, family_capabilities=dict(reversed(list(capabilities.items()))) + ) + + assert result == reversed_result + assert [candidate.families for candidate in result.candidates] == [ + frozenset({"peer"}), + frozenset({"store", "peer"}), + frozenset({"store", "peer"}), + ] + assert [ + candidate.snapshot().envelope().actions[0].action_id for candidate in result.candidates + ] == [ + "fetch-0", + "fetch-1", + "fetch-4", + ] + assert result.unknown_actions == 1 + assert result.unsupported_version_actions == 2 + + +@pytest.mark.parametrize( + "capabilities", + [ + {}, + {"store": frozenset()}, + {"store": frozenset({("kv.fetch", "2.0")})}, + {"store": frozenset({("kv.publish", "1.0")})}, + ], +) +def test_no_matching_family_produces_no_candidates( + capabilities: dict[str, frozenset[tuple[str, str]]], +) -> None: + """A missing exact capability never invents a candidate or a source.""" + result = _collect_kv_fetch_candidates( + snapshot_kv_hints(_wire()), family_capabilities=capabilities + ) + assert result.candidates == () + assert result.unknown_actions == 1 + assert result.unsupported_version_actions == 1 + + +def test_absent_or_unknown_only_envelope_produces_no_candidates() -> None: + """Absent input and unknown actions have different, disjoint counters.""" + capabilities = {"store": frozenset({("kv.fetch", "1.0")})} + absent = _collect_kv_fetch_candidates(None, family_capabilities=capabilities) + assert absent.candidates == () + assert absent.unknown_actions == 0 + assert absent.unsupported_version_actions == 0 + + wire = _wire() + wire["actions"] = wire["actions"][1:] + unknown = _collect_kv_fetch_candidates( + snapshot_kv_hints(wire), family_capabilities=capabilities + ) + assert unknown.candidates == () + assert unknown.unknown_actions == 1 + assert unknown.unsupported_version_actions == 0 + + wire["actions"] = [] + empty = _collect_kv_fetch_candidates(snapshot_kv_hints(wire), family_capabilities=capabilities) + assert empty == absent + + +def test_candidates_do_not_expose_mutable_input_or_result_state() -> None: + """Caller changes and decoded-envelope changes cannot alter retained candidates.""" + wire = _wire() + wire["producer_metadata"] = {"trace": ["router"]} + snapshot = snapshot_kv_hints(wire) + capabilities = {"store": frozenset({("kv.fetch", "1.0")})} + result = _collect_kv_fetch_candidates(snapshot, family_capabilities=capabilities) + candidate = result.candidates[0] + expected = candidate.snapshot().envelope().model_dump() + + capabilities.clear() + wire["producer_metadata"]["trace"].clear() + wire["actions"][0]["payload"].clear() + decoded = candidate.snapshot().envelope() + decoded.actions[0].payload.clear() + decoded.__pydantic_extra__["producer_metadata"]["trace"].clear() + + assert candidate.families == frozenset({"store"}) + assert candidate.snapshot().envelope().model_dump() == expected + assert snapshot.envelope().actions[0].payload == expected["actions"][0]["payload"] + with pytest.raises(FrozenInstanceError): + candidate.families = frozenset() + with pytest.raises(FrozenInstanceError): + result.candidates = () + + +def test_collection_does_not_validate_payloads_or_choose_an_action() -> None: + """An opaque first payload stays first; all supported candidates remain available.""" + wire = _wire() + fetch = wire["actions"][0] + first = dict(fetch, action_id="missing-provider-fields", payload={}) + wire["actions"] = [first, dict(fetch, action_id="later")] + snapshot = snapshot_kv_hints(wire) + capabilities = {"store": frozenset({("kv.fetch", "1.0")})} + + result = _collect_kv_fetch_candidates(snapshot, family_capabilities=capabilities) + + assert [ + candidate.snapshot().envelope().actions[0].action_id for candidate in result.candidates + ] == [ + "missing-provider-fields", + "later", + ] + assert result.candidates[0].snapshot().envelope().actions[0].payload == {} + assert result.unsupported_version_actions == 0 + assert _collect_kv_fetch_candidates(snapshot, family_capabilities=capabilities) == result + + +def test_collection_does_not_copy_large_envelope_metadata_per_candidate() -> None: + """Retaining many candidates must not multiply the envelope's shared metadata.""" + wire = _wire() + wire["producer_metadata"] = {"padding": "x" * 32_000} + fetch = dict(wire["actions"][0], payload={}) + wire["actions"] = [dict(fetch, action_id=str(index)) for index in range(200)] + snapshot = snapshot_kv_hints(wire) + assert snapshot is not None + capabilities = {"store": frozenset({("kv.fetch", "1.0")})} + + was_tracing = tracemalloc.is_tracing() + if not was_tracing: + tracemalloc.start() + try: + before, _ = tracemalloc.get_traced_memory() + result = _collect_kv_fetch_candidates(snapshot, family_capabilities=capabilities) + retained, _ = tracemalloc.get_traced_memory() + finally: + if not was_tracing: + tracemalloc.stop() + + assert len(result.candidates) == 200 + # Eagerly copying 32 KB of metadata for 200 candidates retains over 6 MB. + assert retained - before < 2 * 1024 * 1024 + selected = result.candidates[0].snapshot().envelope() + assert selected.__pydantic_extra__["producer_metadata"] == wire["producer_metadata"] + assert len(selected.actions) == 1 diff --git a/tests/unittest/llmapi/test_kv_hints.py b/tests/unittest/llmapi/test_kv_hints.py new file mode 100644 index 000000000000..4ada57b53836 --- /dev/null +++ b/tests/unittest/llmapi/test_kv_hints.py @@ -0,0 +1,842 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""KV hint transport through real request paths with inference boundaries faked.""" + +import asyncio +import json +import pickle +from contextlib import nullcontext +from dataclasses import FrozenInstanceError +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from pydantic import JsonValue, ValidationError + +from tensorrt_llm.executor.request import GenerationRequest +from tensorrt_llm.kv_hints import KvHintsEnvelope +from tensorrt_llm.kv_hints.protocol import MAX_KV_HINTS_BYTES, admit_kv_hints, snapshot_kv_hints +from tensorrt_llm.sampling_params import SamplingParams + +pytestmark = pytest.mark.cpu_only + + +class _StubModelConfig: + """Model identity needed by request parsing, without loading a model.""" + + model_type = "llama" + vocab_size = 1024 + + +def _wire() -> dict: + """Return the producer's kv.fetch fixture plus an opaque future action.""" + return { + "protocol_version": "0.1", + "message_id": "msg-123", + "actions": [ + { + "action_id": "a1", + "action_type": "kv.fetch", + "action_version": "1.0", + "payload": { + "source_control_endpoint": "tcp://127.0.0.1:23280", + "block_hashes": [11, 22], + }, + }, + { + "action_id": "future", + "action_type": "kv.future_action", + "action_version": "7.3", + "payload": {"nested": {"enabled": True}, "items": [1, None, "x"]}, + }, + ], + } + + +def _hints() -> KvHintsEnvelope: + """Return a valid envelope without interpreting its actions.""" + return KvHintsEnvelope.model_validate(_wire()) + + +def _request(hints: KvHintsEnvelope | None = None) -> GenerationRequest: + """Build a real generation request without an engine.""" + return GenerationRequest([1, 2, 3], SamplingParams(max_tokens=4), kv_hints=hints) + + +def test_wire_round_trip_preserves_unknown_actions_and_json() -> None: + """Action order, unknown versions and opaque nested values survive JSON.""" + envelope = KvHintsEnvelope.model_validate_json(json.dumps(_wire())) + assert json.loads(envelope.model_dump_json()) == _wire() + + +def test_optional_metadata_survives_admission_without_a_version_change() -> None: + """Unknown envelope and action metadata remain JSON-valid and immutable.""" + wire = _wire() + wire["producer_metadata"] = {"trace": ["router-a", 7]} + wire["actions"][0]["optional_metadata"] = {"future": True} + hints = KvHintsEnvelope.model_validate(wire) + request = _request(hints) + assert request.kv_hints.envelope().model_dump() == wire + hints.__pydantic_extra__["producer_metadata"]["trace"].append("mutated") + assert request.kv_hints.envelope().model_dump() == wire + + +@pytest.mark.parametrize("action_level", [False, True]) +@pytest.mark.parametrize("invalid", [float("nan"), object()]) +def test_optional_metadata_must_remain_json(action_level: bool, invalid: object) -> None: + """Forward-compatible metadata cannot carry non-JSON Python objects.""" + wire = _wire() + target = wire["actions"][0] if action_level else wire + target["optional_metadata"] = {"nested": invalid} + with pytest.raises(ValidationError): + KvHintsEnvelope.model_validate(wire) + with patch("tensorrt_llm.kv_hints.protocol.count_ignored_kv_hints") as count: + assert _request(wire).kv_hints is None + count.assert_called_once_with("malformed") + + +@pytest.mark.parametrize( + ("field", "value"), + [("protocol_version", ""), ("protocol_version", 7), ("message_id", ""), ("actions", None)], +) +def test_malformed_envelope_rejected(field: str, value: object) -> None: + """Construction rejects malformed envelopes; admission discards optional hints.""" + wire = _wire() + wire[field] = value + with pytest.raises(ValidationError): + KvHintsEnvelope.model_validate(wire) + with patch("tensorrt_llm.kv_hints.protocol.count_ignored_kv_hints") as count: + assert _request(wire).kv_hints is None + count.assert_called_once_with("malformed") + + +@pytest.mark.parametrize( + ("field", "value"), + [ + ("action_id", ""), + ("action_type", 7), + ("action_version", None), + ("payload", []), + ("payload", {"value": float("nan")}), + ("payload", {"value": b"not-json"}), + ], +) +def test_malformed_action_rejected(field: str, value: object) -> None: + """Opaque payloads still have to be valid JSON objects.""" + wire = _wire() + wire["actions"][0][field] = value + with pytest.raises(ValidationError): + KvHintsEnvelope.model_validate(wire) + with patch("tensorrt_llm.kv_hints.protocol.count_ignored_kv_hints") as count: + assert _request(wire).kv_hints is None + count.assert_called_once_with("malformed") + + +@pytest.mark.parametrize("hints", [b"x", "not-an-object", 7, ["hint"], object()]) +def test_non_object_hints_are_ignored_before_admission(hints: object) -> None: + """Invalid optional field types cannot reject an otherwise valid request.""" + with patch("tensorrt_llm.kv_hints.protocol.count_ignored_kv_hints") as count: + request = _request(hints) + assert request.kv_hints is None + assert request.prompt_token_ids == [1, 2, 3] + assert pickle.loads(pickle.dumps(request)).kv_hints is None + count.assert_called_once_with("malformed") + + +@pytest.mark.parametrize("invalid_json", ["circular", "unpaired_surrogate"]) +def test_non_json_raw_hints_are_ignored_before_admission(invalid_json: str) -> None: + """Invalid raw objects or UTF-8 strings cannot become admitted hint metadata.""" + wire = _wire() + wire["actions"][0]["payload"]["invalid"] = wire if invalid_json == "circular" else "\ud800" + with patch("tensorrt_llm.kv_hints.protocol.count_ignored_kv_hints") as count: + assert _request(wire).kv_hints is None + count.assert_called_once_with("malformed") + + +@pytest.mark.parametrize("typed", [False, True]) +@pytest.mark.parametrize("location", ["payload", "envelope_extra", "action_extra"]) +@pytest.mark.parametrize( + ("value", "accepted"), + [ + (-(2**63) - 1, False), + (-(2**63), True), + (2**63 - 1, True), + (2**63, True), + (2**64 - 1, True), + (2**64, False), + ], +) +def test_admission_enforces_msgpack_integer_bounds( + typed: bool, location: str, value: int, accepted: bool +) -> None: + """Admitted opaque JSON values must survive the proxy's MessagePack transport.""" + import msgspec + + wire = _wire() + targets = { + "payload": wire["actions"][0]["payload"], + "envelope_extra": wire, + "action_extra": wire["actions"][0], + } + targets[location]["nested"] = {"items": [value]} + hints = KvHintsEnvelope.model_validate(wire) if typed else wire + with patch("tensorrt_llm.kv_hints.protocol.count_ignored_kv_hints") as count: + request = _request(hints) + assert request.prompt_token_ids == [1, 2, 3] + if accepted: + assert request.kv_hints is not None + decoded = request.kv_hints.envelope().model_dump(mode="json") + assert msgspec.msgpack.decode(msgspec.msgpack.encode(decoded)) == wire + count.assert_not_called() + else: + assert request.kv_hints is None + count.assert_called_once_with("malformed") + + +@pytest.mark.parametrize("request_kind", ["completion", "chat", "responses"]) +@pytest.mark.parametrize("value", [-(2**63) - 1, 2**64]) +def test_proxy_forwarding_discards_non_msgpack_hints(request_kind: str, value: int) -> None: + """Invalid optional hints cannot stop a valid proxy request from being encoded.""" + import msgspec + + from tensorrt_llm.serve.openai_client import _encode_request_body + from tensorrt_llm.serve.openai_protocol import ( + ChatCompletionRequest, + CompletionRequest, + ResponsesRequest, + ) + + request_types = { + "completion": (CompletionRequest, {"prompt": [1, 2]}), + "chat": (ChatCompletionRequest, {"messages": [{"role": "user", "content": "hi"}]}), + "responses": (ResponsesRequest, {"input": "hi"}), + } + request_type, inputs = request_types[request_kind] + hints = _wire() + hints["actions"][0]["payload"]["nested"] = {"items": [value]} + handoff = {"request_type": "context_only", "disagg_request_id": 123} + with patch("tensorrt_llm.kv_hints.protocol.count_ignored_kv_hints") as count: + request = request_type.model_validate( + {"model": "test", **inputs, "kv_hints": hints, "disaggregated_params": handoff} + ) + forwarded = msgspec.msgpack.decode(_encode_request_body(request)) + assert request.kv_hints is None + assert forwarded["kv_hints"] is None + assert forwarded["disaggregated_params"]["disagg_request_id"] == 123 + count.assert_called_once_with("malformed") + + +@pytest.mark.parametrize("request_kind", ["completion", "chat", "responses"]) +@pytest.mark.parametrize("typed", [False, True]) +def test_http_request_repr_omits_hints_without_changing_transport( + request_kind: str, typed: bool +) -> None: + """Request logging omits hints while JSON and proxy forwarding retain them.""" + import msgspec + + from tensorrt_llm.serve.openai_client import _encode_request_body + from tensorrt_llm.serve.openai_protocol import ( + ChatCompletionRequest, + CompletionRequest, + ResponsesRequest, + ) + + request_types = { + "completion": (CompletionRequest, {"prompt": [1, 2]}), + "chat": (ChatCompletionRequest, {"messages": [{"role": "user", "content": "hi"}]}), + "responses": (ResponsesRequest, {"input": "hi"}), + } + request_type, inputs = request_types[request_kind] + wire = _wire() + sensitive_value = "private-hint-payload-marker" + wire["actions"][0]["payload"]["opaque_credential"] = sensitive_value + hints = KvHintsEnvelope.model_validate(wire) if typed else wire + request = request_type.model_validate({"model": "test", **inputs, "kv_hints": hints}) + + for representation in (str(request), repr(request)): + assert "kv_hints" not in representation + assert sensitive_value not in representation + assert request.model_dump(mode="json")["kv_hints"] == wire + assert json.loads(request.model_dump_json())["kv_hints"] == wire + assert msgspec.msgpack.decode(_encode_request_body(request))["kv_hints"] == wire + + +@pytest.mark.parametrize("typed", [False, True]) +@pytest.mark.parametrize("container_kind", ["dict", "list"]) +@pytest.mark.parametrize("location", ["payload", "envelope_extra", "action_extra"]) +@pytest.mark.parametrize(("depth", "accepted"), [(32, True), (200, False)]) +def test_admission_requires_decodable_nested_json( + typed: bool, container_kind: str, location: str, depth: int, accepted: bool +) -> None: + """Python-valid metadata must also fit the snapshot's JSON decoder.""" + nested: JsonValue = "leaf" + for _ in range(depth): + nested = {"child": nested} if container_kind == "dict" else [nested] + wire = _wire() + targets = { + "payload": wire["actions"][0]["payload"], + "envelope_extra": wire, + "action_extra": wire["actions"][0], + } + targets[location]["nested"] = nested + envelope = KvHintsEnvelope.model_validate(wire) + wire_json = envelope.model_dump_json().encode("utf-8") + assert len(wire_json) < MAX_KV_HINTS_BYTES + hints = envelope if typed else wire + + if not accepted: + with pytest.raises(ValidationError, match="recursion limit exceeded"): + KvHintsEnvelope.model_validate_json(wire_json) + + with patch("tensorrt_llm.kv_hints.protocol.count_ignored_kv_hints") as count: + admitted = admit_kv_hints(hints) + if accepted: + assert admitted is not None + assert admitted.model_dump() == wire + count.assert_not_called() + else: + assert admitted is None + count.assert_called_once_with("malformed") + + with patch("tensorrt_llm.kv_hints.protocol.count_ignored_kv_hints") as count: + request = _request(hints) + assert request.prompt_token_ids == [1, 2, 3] + if accepted: + assert request.kv_hints is not None + assert request.kv_hints.wire_json == wire_json + assert request.kv_hints.envelope().model_dump() == wire + count.assert_not_called() + else: + assert request.kv_hints is None + count.assert_called_once_with("malformed") + + +def test_invalid_request_input_is_not_hidden_by_ignored_hints() -> None: + """Hint fallback does not relax the generation request's own input contract.""" + with pytest.raises(TypeError, match="prompt_token_ids"): + GenerationRequest("not-token-ids", SamplingParams(max_tokens=4), kv_hints={}) + + +def test_admission_snapshot_is_deeply_immutable() -> None: + """Caller and consumer mutations cannot change a queued snapshot.""" + hints = _hints() + request = _request(hints) + hints.actions[0].payload["block_hashes"][0] = 99 + hints.actions.append(hints.actions[0]) + assert json.loads(request.kv_hints.wire_json) == _wire() + + decoded = request.kv_hints.envelope() + decoded.actions[1].payload["nested"]["enabled"] = False + assert request.kv_hints.envelope().model_dump() == _wire() + with pytest.raises(FrozenInstanceError): + request.kv_hints.wire_json = b"{}" + + +@pytest.mark.parametrize( + "mutation", + ["non_json_payload", "invalid_field_type", "missing_protocol_version", "unpaired_surrogate"], +) +def test_admission_revalidates_mutated_envelope(mutation: str) -> None: + """Admission discards a typed envelope that became invalid after construction.""" + hints = _hints() + if mutation == "non_json_payload": + hints.actions[0].payload["bad"] = object() + elif mutation == "invalid_field_type": + hints.actions = "invalid-actions" + elif mutation == "missing_protocol_version": + del hints.protocol_version + else: + hints.actions[0].payload["bad"] = "\ud800" + with patch("tensorrt_llm.kv_hints.protocol.count_ignored_kv_hints") as count: + assert _request(hints).kv_hints is None + count.assert_called_once_with("malformed") + + +def test_generation_request_pickle_preserves_snapshot() -> None: + """The existing request pickle path retains Python hint metadata.""" + request = _request(_hints()) + clone = pickle.loads(pickle.dumps(request)) + assert clone.kv_hints == request.kv_hints + assert clone.kv_hints.envelope().model_dump() == _wire() + + +def test_no_hint_and_legacy_request_state_default_to_none() -> None: + """Old serialized requests remain valid without a kv_hints field.""" + request = _request() + assert request.kv_hints is None + state = request.__getstate__() + del state["kv_hints"] + restored = GenerationRequest.__new__(GenerationRequest) + restored.__setstate__(state) + assert restored.kv_hints is None + + +def test_executor_admits_hint_snapshot() -> None: + """Executor admission passes hints into the real GenerationRequest.""" + from tensorrt_llm.executor.executor import GenerationExecutor + + executor = MagicMock() + GenerationExecutor.generate_async( + executor, [1, 2], SamplingParams(max_tokens=4), kv_hints=_hints() + ) + submitted = executor.submit.call_args.args[0] + assert submitted.kv_hints.envelope().model_dump() == _wire() + + +@pytest.mark.parametrize("batched", [False, True]) +def test_llm_generate_forwards_hint_metadata(batched: bool) -> None: + """Sync generation forwards a scalar envelope or one value per prompt.""" + from tensorrt_llm.llmapi.llm import BaseLLM + + llm = object.__new__(BaseLLM) + llm.generate_async = MagicMock() + hints = _hints() + values = [hints, None] if batched else hints + BaseLLM.generate(llm, [[1, 2], [3, 4]], kv_hints=values, use_tqdm=False) + calls = llm.generate_async.call_args_list + assert calls[0].kwargs["kv_hints"] is hints + assert calls[1].kwargs["kv_hints"] is (None if batched else hints) + + +def test_llm_rejects_hint_batch_length_before_submitting() -> None: + """A misaligned hint batch cannot partially submit requests.""" + from tensorrt_llm.llmapi.llm import BaseLLM + + llm = object.__new__(BaseLLM) + llm.generate_async = MagicMock() + with pytest.raises(ValueError, match="kv_hints list length"): + BaseLLM.generate(llm, [[1, 2], [3, 4]], kv_hints=[_hints()], use_tqdm=False) + llm.generate_async.assert_not_called() + + +def test_llm_async_forwards_hints_to_executor() -> None: + """The async public method retains metadata after prompt preprocessing.""" + from tensorrt_llm.llmapi.llm import BaseLLM, RequestOutput + + llm = MagicMock() + llm._encode_only = False + llm._executor.is_shutdown.return_value = False + llm._prepare_sampling_params.return_value = SamplingParams(max_tokens=4) + llm._preprocess.return_value = ([1, 2], None, None, None) + llm.args.return_perf_metrics = False + hints = _hints() + with patch.object(RequestOutput, "_from_generation_result"): + BaseLLM.generate_async(llm, [1, 2], kv_hints=hints) + assert llm._executor.generate_async.call_args.kwargs["kv_hints"] is hints + + +@pytest.mark.parametrize("with_hints", [False, True]) +def test_worker_and_llm_request_conversion(with_hints: bool) -> None: + """Real native request construction carries the Python side-channel data.""" + from tensorrt_llm._torch.pyexecutor.llm_request import executor_request_to_llm_request + from tensorrt_llm.executor.base_worker import BaseWorker + + request = _request(_hints() if with_hints else None) + request.set_id(42) + worker = MagicMock() + worker._is_pytorch_backend = True + worker.max_seq_len = None + worker.engine.enqueue_request.return_value = 42 + BaseWorker._enqueue_request(worker, request) + native_request = worker.engine.enqueue_request.call_args.args[0] + assert native_request.py_kv_hints == request.kv_hints + + llm_request = executor_request_to_llm_request( + 42, native_request, child_req_ids=[], exclude_last_generation_logits=False + ) + assert llm_request.py_kv_hints == request.kv_hints + + +def test_child_requests_inherit_hint_snapshot() -> None: + """Child cloning preserves the admitted snapshot independently of request IDs.""" + from tensorrt_llm._torch.pyexecutor.llm_request import executor_request_to_llm_request + from tensorrt_llm.bindings import executor as bindings + + request = bindings.Request( + input_token_ids=[1, 2, 3], + max_tokens=4, + sampling_config=bindings.SamplingConfig(num_return_sequences=2), + ) + request.py_kv_hints = snapshot_kv_hints(_hints()) + llm_request = executor_request_to_llm_request( + 42, request, child_req_ids=[43], exclude_last_generation_logits=False + ) + assert llm_request.child_requests[0].py_kv_hints == request.py_kv_hints + + +@pytest.mark.parametrize("broadcast_kind", ["single_round", "tensor_parallel", "pipeline_parallel"]) +def test_request_broadcast_transports_hints(broadcast_kind: str) -> None: + """Rank transport preserves immutable hints alongside other Python metadata.""" + from tensorrt_llm._torch.pyexecutor.executor_request_queue import RequestQueueItem + from tensorrt_llm._torch.pyexecutor.request_utils import ( + RequestBroadcaster, + attach_py_objects_to_requests, + ) + + hints = snapshot_kv_hints(_hints()) + items = [ + RequestQueueItem( + 1, + SimpleNamespace( + py_kv_hints=hints, + py_conversation_params=None, + py_lora_ckpt_source="hf", + ), + ), + RequestQueueItem(2, SimpleNamespace(py_kv_hints=None, py_conversation_params=None)), + ] + pipeline_parallel = broadcast_kind == "pipeline_parallel" + transport = MagicMock() + transport.world_size = 2 + transport.rank = 0 + transport.has_pp = pipeline_parallel + transport.supports_single_round_broadcast = broadcast_kind == "single_round" + transport.is_first_pp_rank = True + transport.is_last_pp_rank = False + transport.pp_size = 2 + transport.broadcast.side_effect = lambda value, root: pickle.loads(pickle.dumps(value)) + transport.tp_cp_broadcast.side_effect = transport.broadcast.side_effect + transport.broadcast_or_none.side_effect = transport.broadcast.side_effect + hang_detector = SimpleNamespace(pause=nullcontext) + broadcaster = RequestBroadcaster(transport, hang_detector) + collected = broadcaster._collect_py_objects(items) + with patch.dict("os.environ", {"TRTLLM_PP_REQ_SEND_ASYNC": "0"}): + if broadcast_kind == "single_round": + original_requests, _ = broadcaster.broadcast(items) + assert original_requests is items + sent = transport.broadcast_or_none.call_args.args[0] + receiver = MagicMock() + receiver.world_size = 2 + receiver.rank = 1 + receiver.has_pp = False + receiver.supports_single_round_broadcast = True + receiver.broadcast_or_none.return_value = pickle.loads(pickle.dumps(sent)) + _, received = RequestBroadcaster(receiver, hang_detector).broadcast([]) + transport.broadcast_or_none.assert_called_once() + receiver.broadcast_or_none.assert_called_once_with(None, root=0) + transport.broadcast_int64.assert_not_called() + receiver.broadcast_int64.assert_not_called() + else: + _, received = broadcaster._broadcast_requests(items, collected) + if pipeline_parallel: + transport.tp_cp_broadcast.assert_called_once() + transport.send_object.assert_called_once() + sent = transport.send_object.call_args.args[0] + receiver = MagicMock() + receiver.world_size = 2 + receiver.has_pp = True + receiver.is_first_pp_rank = False + receiver.is_last_pp_rank = True + receiver.pp_size = 2 + receiver.recv_object.return_value = pickle.loads(pickle.dumps(sent)) + _, received = RequestBroadcaster(receiver, hang_detector)._broadcast_requests([], None) + receiver.recv_object.assert_called_once() + elif broadcast_kind == "tensor_parallel": + transport.broadcast.assert_called_once() + + targets = [RequestQueueItem(1, SimpleNamespace()), RequestQueueItem(2, SimpleNamespace())] + attach_py_objects_to_requests(targets, received) + transported_hints = targets[0].request.py_kv_hints + assert transported_hints == hints + assert targets[0].request.py_lora_ckpt_source == "hf" + assert getattr(targets[1].request, "py_kv_hints", None) is None + transported_hints.envelope().actions[0].payload["block_hashes"].append(99) + assert transported_hints.envelope().model_dump() == _wire() + with pytest.raises(FrozenInstanceError): + transported_hints.wire_json = b"{}" + + +@pytest.mark.parametrize("request_kind", ["completion", "chat", "responses"]) +def test_http_request_models_validate_hints(request_kind: str) -> None: + """All supported HTTP request schemas share the same envelope validation.""" + from tensorrt_llm.serve.openai_protocol import ( + ChatCompletionRequest, + CompletionRequest, + ResponsesRequest, + ) + + constructors = { + "completion": (CompletionRequest, {"prompt": [1, 2]}), + "chat": (ChatCompletionRequest, {"messages": [{"role": "user", "content": "hi"}]}), + "responses": (ResponsesRequest, {"input": "hi"}), + } + request_type, fields = constructors[request_kind] + request = request_type(model="test", kv_hints=_wire(), **fields) + assert request.kv_hints.model_dump() == _wire() + assert request_type(model="test", **fields).kv_hints is None + invalid = _wire() + invalid["protocol_version"] = "unsupported" + assert request_type(model="test", kv_hints=invalid, **fields).kv_hints is None + + +@pytest.mark.parametrize("request_kind", ["completion", "chat", "harmony", "responses"]) +@pytest.mark.parametrize("with_hints", [False, True]) +def test_http_handlers_forward_hints(request_kind: str, with_hints: bool) -> None: + """Run each real HTTP handler until its model-generation boundary. + + The engine deliberately raises after recording the call. Tokenized inputs + avoid loading a tokenizer; Responses preprocessing is the only stubbed + adapter because it normally resolves model and conversation state. + """ + from tensorrt_llm.serve.openai_protocol import ( + ChatCompletionRequest, + CompletionRequest, + ResponsesRequest, + ) + from tensorrt_llm.serve.openai_server import OpenAIServer + + server = object.__new__(OpenAIServer) + server.model = "test" + server.allow_request_chat_template = False + server.model_config = _StubModelConfig() + server.processor = None + server.tokenizer = SimpleNamespace(tokenizer=SimpleNamespace(vocab_size=1024)) + server.chat_template = None + server.tool_parser = None + server.tool_call_id_type = "random" + server.multimodal_server_config = None + server.use_harmony = request_kind == "harmony" + server.harmony_adapter = SimpleNamespace(get_stop_tokens=lambda: [2]) + server.enable_store = False + server.conversation_store = None + server.generator = SimpleNamespace( + args=SimpleNamespace( + gather_generation_logits=False, + reasoning_parser=None, + backend="pytorch", + guided_decoding_backend=None, + num_postprocess_workers=0, + ), + generate_async=MagicMock(side_effect=ValueError("captured at engine boundary")), + ) + hints = _hints() if with_hints else None + if request_kind == "completion": + request = CompletionRequest(model="test", prompt=[1, 2], max_tokens=4, kv_hints=hints) + handler = server.openai_completion + elif request_kind == "responses": + request = ResponsesRequest(model="test", input="hi", max_output_tokens=4, kv_hints=hints) + handler = server.openai_responses + else: + request = ChatCompletionRequest( + model="test", + messages=[{"role": "user", "content": "hi"}], + prompt_token_ids=[1, 2], + max_tokens=4, + kv_hints=hints, + ) + handler = server.chat_harmony if request_kind == "harmony" else server.openai_chat + + with patch( + "tensorrt_llm.serve.openai_server.responses_api_request_preprocess", + new=AsyncMock(return_value=([1, 2], SamplingParams(max_tokens=4))), + ): + response = asyncio.run(handler(request, None)) + server.generator.generate_async.assert_called_once() + assert server.generator.generate_async.call_args.kwargs["kv_hints"] == hints + assert b"captured at engine boundary" in response.body + + +@pytest.mark.parametrize("typed", [False, True]) +def test_unknown_version_is_ignored_before_admission(typed: bool) -> None: + """Raw future envelopes do not have to conform to the current action schema.""" + wire = _wire() + wire["protocol_version"] = "future" + hints = ( + KvHintsEnvelope.model_validate(wire) + if typed + else {"protocol_version": "future", "future_payload": [7], "actions": "different format"} + ) + with patch("tensorrt_llm.kv_hints.protocol.count_ignored_kv_hints") as count: + request = _request(hints) + assert request.kv_hints is None + pickle.loads(pickle.dumps(request)) + count.assert_called_once_with("unknown_version") + + +@pytest.mark.parametrize("typed", [False, True]) +@pytest.mark.parametrize("extra_bytes", [0, 1]) +def test_snapshot_bounds_utf8_bytes(typed: bool, extra_bytes: int) -> None: + """The byte bound includes multibyte payloads and accepts the exact limit.""" + wire = _wire() + wire["actions"][0]["payload"] = {"padding": "雪"} + size = len(KvHintsEnvelope.model_validate(wire).model_dump_json().encode("utf-8")) + wire["actions"][0]["payload"]["padding"] += "x" * (MAX_KV_HINTS_BYTES - size + extra_bytes) + hints = KvHintsEnvelope.model_validate(wire) if typed else wire + with patch("tensorrt_llm.kv_hints.protocol.count_ignored_kv_hints") as count: + request = _request(hints) + if extra_bytes: + assert request.kv_hints is None + count.assert_called_once_with("oversized") + else: + assert len(request.kv_hints.wire_json) == MAX_KV_HINTS_BYTES + count.assert_not_called() + + +def test_hint_types_are_exported_with_request_parameters() -> None: + """Both public import paths identify the same request parameter types.""" + from tensorrt_llm import llmapi + from tensorrt_llm.kv_hints import KvHintAction, KvHintsInput + + for name, value in ( + ("KvHintAction", KvHintAction), + ("KvHintsEnvelope", KvHintsEnvelope), + ("KvHintsInput", KvHintsInput), + ): + assert name in llmapi.__all__ + assert getattr(llmapi, name) is value + + +@pytest.mark.parametrize( + "route_kind", ["worker_json", "worker_msgpack", "proxy", "proxy_responses"] +) +@pytest.mark.parametrize("trusted", [False, True]) +@pytest.mark.parametrize( + "hint_kind", ["valid", "unknown", "oversized", "malformed", "malformed_action", "non_object"] +) +def test_http_ingress_policy_before_nested_validation( + route_kind: str, + trusted: bool, + hint_kind: str, +) -> None: + """Real worker/proxy routes strip before parsing and never alter P/D fields.""" + import msgspec + from fastapi import FastAPI + from fastapi.testclient import TestClient + + 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": + hints = {"protocol_version": "future", "actions": "different schema"} + elif hint_kind == "oversized": + hints["actions"][0]["payload"] = {"large": "雪" * MAX_KV_HINTS_BYTES} + elif hint_kind == "malformed": + hints.pop("message_id") + elif hint_kind == "malformed_action": + hints["actions"][0]["action_type"] = 7 + elif hint_kind == "non_object": + hints = 7 + app = FastAPI() + app.state.trust_kv_hints = trusted + app.router.route_class = KvHintsRoute if route_kind.startswith("proxy") else _MsgspecRoute + + def capture(request: CompletionRequest | ResponsesRequest) -> dict: + """Return admitted hints and required handoff fields without inference.""" + admitted = _request(request.kv_hints) + return { + "kv_hints": None + if admitted.kv_hints is None + else admitted.kv_hints.envelope().model_dump(), + "disaggregated_params": request.disaggregated_params.model_dump(), + } + + async def capture_completions(request: CompletionRequest) -> dict: + """Validate completion inputs before recording their admitted metadata.""" + return capture(request) + + async def capture_responses(request: ResponsesRequest) -> dict: + """Validate Responses inputs before recording their admitted metadata.""" + return capture(request) + + is_responses = route_kind == "proxy_responses" + route_path = "/v1/responses" if is_responses else "/v1/completions" + endpoint = capture_responses if is_responses else capture_completions + app.add_api_route(route_path, endpoint, methods=["POST"]) + body = { + "model": "test", + **({"input": "hi"} if is_responses else {"prompt": [1, 2]}), + "kv_hints": hints, + "disaggregated_params": {"request_type": "context_only", "disagg_request_id": 123}, + } + with ( + TestClient(app) as client, + patch("tensorrt_llm.serve.kv_hints.count_ignored_kv_hints") as trust_count, + patch("tensorrt_llm.kv_hints.protocol.count_ignored_kv_hints") as admission_count, + ): + if route_kind == "worker_msgpack": + response = client.post( + route_path, + content=msgspec.msgpack.encode(body), + headers={"content-type": "application/json", "x-trtllm-msgpack": "1"}, + ) + else: + response = client.post(route_path, json=body) + assert response.status_code == 200 + output = response.json() + assert output["kv_hints"] == (_wire() if trusted and hint_kind == "valid" else None) + assert output["disaggregated_params"]["disagg_request_id"] == 123 + if not trusted: + trust_count.assert_called_once_with("untrusted") + admission_count.assert_not_called() + else: + trust_count.assert_not_called() + reasons = { + "unknown": "unknown_version", + "oversized": "oversized", + "malformed": "malformed", + "malformed_action": "malformed", + "non_object": "malformed", + } + if hint_kind == "valid": + admission_count.assert_not_called() + else: + admission_count.assert_called_once_with(reasons[hint_kind]) + + +@pytest.mark.parametrize("route_kind", ["worker_json", "worker_msgpack", "proxy"]) +@pytest.mark.parametrize("invalid_field", ["prompt", "handoff_request_type"]) +def test_ignored_hints_do_not_relax_http_request_validation( + route_kind: str, + invalid_field: str, +) -> None: + """Ordinary request and required handoff fields remain strict after hint fallback.""" + import msgspec + from fastapi import FastAPI + from fastapi.testclient import TestClient + + 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 + app.router.route_class = KvHintsRoute if route_kind == "proxy" else _MsgspecRoute + + def capture(request: CompletionRequest) -> dict: + """Fail if invalid request fields reach the endpoint.""" + raise AssertionError("Invalid request fields must fail before endpoint dispatch") + + app.add_api_route("/v1/completions", capture, methods=["POST"]) + body = { + "model": "test", + "prompt": [1, 2], + "kv_hints": {"protocol_version": "0.1", "actions": "malformed"}, + "disaggregated_params": {"request_type": "context_only", "disagg_request_id": 123}, + } + if invalid_field == "prompt": + body.pop("prompt") + error_location = ["body", "prompt"] + else: + body["disaggregated_params"].pop("request_type") + error_location = ["body", "disaggregated_params", "request_type"] + + with ( + TestClient(app) as client, + patch("tensorrt_llm.kv_hints.protocol.count_ignored_kv_hints") as count, + ): + if route_kind == "worker_msgpack": + response = client.post( + "/v1/completions", + content=msgspec.msgpack.encode(body), + headers={"content-type": "application/json", "x-trtllm-msgpack": "1"}, + ) + else: + response = client.post("/v1/completions", json=body) + count.assert_called_once_with("malformed") + assert response.status_code == 422 + errors = response.json()["detail"] + assert any(error["loc"] == error_location for error in errors) + assert all("kv_hints" not in error["loc"] for error in errors) diff --git a/tests/unittest/llmapi/test_trtllm_serve_config_schemas.py b/tests/unittest/llmapi/test_trtllm_serve_config_schemas.py index 5b9dbaa91d93..5f69fa8a46cd 100644 --- a/tests/unittest/llmapi/test_trtllm_serve_config_schemas.py +++ b/tests/unittest/llmapi/test_trtllm_serve_config_schemas.py @@ -682,12 +682,12 @@ def test_version_metadata_sync_needs_no_project_dependencies(tmp_path: Path) -> str(schema_dir), ] before_sync = {path.name: path.read_bytes() for path in schema_dir.iterdir()} - result = subprocess.run(command + ["--check"], capture_output=True, text=True) + result = subprocess.run(command + ["--check"], capture_output=True, text=True, timeout=60) assert result.returncode != 0 assert "Missing or stale configuration schemas" in result.stderr assert {path.name: path.read_bytes() for path in schema_dir.iterdir()} == before_sync - subprocess.run(command, check=True, capture_output=True, text=True) + subprocess.run(command, check=True, capture_output=True, text=True, timeout=60) version = generator._source_version() before_check = {} for filename, expected_command in generator._SCHEMA_COMMANDS.items(): @@ -699,7 +699,7 @@ def test_version_metadata_sync_needs_no_project_dependencies(tmp_path: Path) -> assert schema["sentinel"] == {"preserved": True} before_check[filename] = path.read_bytes() - subprocess.run(command + ["--check"], check=True, capture_output=True, text=True) + subprocess.run(command + ["--check"], check=True, capture_output=True, text=True, timeout=60) assert {path.name: path.read_bytes() for path in schema_dir.iterdir()} == before_check diff --git a/tests/unittest/others/test_maintain_prims_ts.py b/tests/unittest/others/test_maintain_prims_ts.py index 90e7fc2c2fee..b8a0533943fb 100644 --- a/tests/unittest/others/test_maintain_prims_ts.py +++ b/tests/unittest/others/test_maintain_prims_ts.py @@ -320,6 +320,7 @@ def api( ["git", "-C", str(remote), "rev-parse", "--verify", f"refs/heads/{branch}"], capture_output=True, text=True, + timeout=60, ) if result.returncode: raise RuntimeError("HTTP 404") diff --git a/tests/unittest/others/test_vendor_bot.py b/tests/unittest/others/test_vendor_bot.py index 4917be3a6243..68a456ec1a49 100644 --- a/tests/unittest/others/test_vendor_bot.py +++ b/tests/unittest/others/test_vendor_bot.py @@ -305,7 +305,7 @@ def test_unpaired_and_template() -> None: def _git(repo: Path, *arguments: str) -> str: result = subprocess.run( - ["git", "-C", str(repo), *arguments], capture_output=True, text=True, check=True + ["git", "-C", str(repo), *arguments], capture_output=True, text=True, check=True, timeout=60 ) return result.stdout.strip() @@ -2771,6 +2771,7 @@ def test_daemon_launch_status_stop_without_systemd(tmp_path: Path) -> None: capture_output=True, text=True, check=True, + timeout=60, ) state = json.loads(status.stdout) assert state["running"] diff --git a/tests/unittest/scripts/test_bolt.py b/tests/unittest/scripts/test_bolt.py index e16d16a24480..77a043ce4d82 100644 --- a/tests/unittest/scripts/test_bolt.py +++ b/tests/unittest/scripts/test_bolt.py @@ -30,8 +30,10 @@ import shutil import stat import subprocess +import sys import zipfile from pathlib import Path +from unittest.mock import patch import pytest @@ -42,7 +44,8 @@ def _load(name: str): spec = importlib.util.spec_from_file_location(name, BOLT_DIR / f"{name}.py") module = importlib.util.module_from_spec(spec) - spec.loader.exec_module(module) + with patch.object(sys, "path", sys.path.copy()): + spec.loader.exec_module(module) return module diff --git a/tests/unittest/scripts/test_perf_agentx_aiperf_probe.py b/tests/unittest/scripts/test_perf_agentx_aiperf_probe.py index bac577fb2609..70373618f057 100644 --- a/tests/unittest/scripts/test_perf_agentx_aiperf_probe.py +++ b/tests/unittest/scripts/test_perf_agentx_aiperf_probe.py @@ -17,8 +17,10 @@ import importlib.util import subprocess import sys +from contextlib import nullcontext from pathlib import Path from types import ModuleType +from unittest.mock import Mock import pytest @@ -112,6 +114,44 @@ def test_fork_with_old_datasets_is_not_usable(client: ModuleType, monkeypatch, t assert "datasets 3.1.0" in log +@pytest.mark.parametrize("recovered_after_install", [True, False]) +def test_probe_timeout_attempts_one_recovery( + client: ModuleType, + monkeypatch: pytest.MonkeyPatch, + capfd: pytest.CaptureFixture[str], + recovered_after_install: bool, +) -> None: + """A timed-out probe allows one install and still rejects an unusable stack.""" + timeout = subprocess.TimeoutExpired([sys.executable, "-c", "probe"], timeout=120) + after_install = subprocess.CompletedProcess([], 0) if recovered_after_install else timeout + run = Mock(side_effect=[timeout, after_install]) + install = Mock() + heartbeat = Mock(return_value=nullcontext()) + monkeypatch.delenv("AGENTX_FORCE_REINSTALL", raising=False) + monkeypatch.delenv("AGENTX_AIPERF_REF", raising=False) + monkeypatch.setattr(client.subprocess, "run", run) + monkeypatch.setattr(client, "_pip", install) + monkeypatch.setattr(client, "_Heartbeat", heartbeat) + + if recovered_after_install: + client.ensure_aiperf("http://h:1", "test-model") + else: + with pytest.raises(RuntimeError, match="still not importable after install"): + client.ensure_aiperf("http://h:1", "test-model") + + assert run.call_count == 2 + assert all(call.kwargs["timeout"] == 120 for call in run.call_args_list) + heartbeat.assert_called_once_with("http://h:1", "test-model") + aiperf_installs = [ + call + for call in install.call_args_list + if f"aiperf @ {client.DEFAULT_AIPERF_REF}" in call.args + ] + assert len(aiperf_installs) == 1 + expected_timeouts = 1 if recovered_after_install else 2 + assert capfd.readouterr().out.count("import probe timed out after 120s") == expected_timeouts + + def _default_cmd(client: ModuleType, monkeypatch) -> list[str]: monkeypatch.delenv("AGENTX_EXTRA_ARGS", raising=False) args = client.parse_args( diff --git a/tests/unittest/scripts/test_perf_agentx_client.py b/tests/unittest/scripts/test_perf_agentx_client.py index d2f7f9be9181..855bdb0d684c 100644 --- a/tests/unittest/scripts/test_perf_agentx_client.py +++ b/tests/unittest/scripts/test_perf_agentx_client.py @@ -69,10 +69,10 @@ def test_max_output_tokens_must_be_positive(client: ModuleType): client.parse_args(_BASE_ARGS + ["--max-output-tokens", "0"]) -def test_prefill_only_reports_absent_decode_intervals_as_zero(client: ModuleType, capsys): +def test_prefill_only_reports_absent_decode_intervals_as_zero(client: ModuleType, capfd): client.report_metrics(_export(client, with_decode_intervals=False), prefill_only=True) lines = dict( - line.rsplit(":", 1) for line in capsys.readouterr().out.splitlines() if line.strip() + line.rsplit(":", 1) for line in capfd.readouterr().out.splitlines() if line.strip() ) for stat in ("Mean", "Median", "P99"): assert float(lines[f"{stat} TPOT (ms)"]) == 0.0 @@ -82,9 +82,9 @@ def test_prefill_only_reports_absent_decode_intervals_as_zero(client: ModuleType assert float(lines["P99 E2EL (ms)"]) == 9.0 -def test_prefill_only_still_reports_measured_decode_intervals(client: ModuleType, capsys): +def test_prefill_only_still_reports_measured_decode_intervals(client: ModuleType, capfd): client.report_metrics(_export(client, with_decode_intervals=True), prefill_only=True) - assert "Mean TPOT (ms): 5.00" in capsys.readouterr().out + assert "Mean TPOT (ms): 5.00" in capfd.readouterr().out def test_missing_decode_intervals_fail_outside_prefill_only(client: ModuleType):