diff --git a/docs/serving/online_serving/structured_decisions.md b/docs/serving/online_serving/structured_decisions.md index 2194b54e47..a1f58a3769 100644 --- a/docs/serving/online_serving/structured_decisions.md +++ b/docs/serving/online_serving/structured_decisions.md @@ -135,6 +135,11 @@ the request lists them, a `noul` labels its options `yes` and `no`, and a | `questions` | question id to `{type, instructions, criteria}`, asked in this order | | `instructions` | optional context placed ahead of the questions | | `chat_template_kwargs` | passed to the chat template, for example `{"enable_thinking": false}` | +| `cache_salt` | optional [prefix-cache salt](../../usage/security.md#cache-salting) applied to every question | A question with any field other than `type`, `instructions` and `criteria` is rejected with a 400. + +When [tracing](../../../examples/observability/opentelemetry/README.md) is enabled +with `--otlp-traces-endpoint`, the HTTP `traceparent` and `tracestate` headers are +forwarded to every question's generation request. diff --git a/docs/usage/security.md b/docs/usage/security.md index c0d2c79700..a74ed1d6eb 100644 --- a/docs/usage/security.md +++ b/docs/usage/security.md @@ -572,7 +572,7 @@ An attacker sharing the same backend can measure differences in Time to First To vLLM accepts an optional `cache_salt` parameter on requests. The salt is mixed into the hash of the first KV cache block, so only requests carrying the same salt can share cached prefix blocks. See [Automatic Prefix Caching](../design/prefix_caching.md) for the implementation details. -`cache_salt` is accepted by the OpenAI-compatible chat completions, completions, responses, and pooling (embeddings, classification, scoring) endpoints, and by the Anthropic `/v1/messages` endpoint. +`cache_salt` is accepted by the OpenAI-compatible chat completions, completions, responses, and pooling (embeddings, classification, scoring) endpoints, the Anthropic `/v1/messages` endpoint, and the structured decisions `/v1/systemone` endpoint. #### Usage with the OpenAI Python client diff --git a/tests/entrypoints/generate/structured_decisions/test_strategies.py b/tests/entrypoints/generate/structured_decisions/test_strategies.py index 3d9538e4dc..5dcb8edb30 100644 --- a/tests/entrypoints/generate/structured_decisions/test_strategies.py +++ b/tests/entrypoints/generate/structured_decisions/test_strategies.py @@ -2,24 +2,46 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from types import SimpleNamespace from typing import Any +from unittest.mock import MagicMock import pytest -from fastapi import FastAPI +from fastapi import FastAPI, Request from fastapi.testclient import TestClient +from vllm.config import ModelConfig +from vllm.engine.protocol import EngineClient from vllm.entrypoints.generate.structured_decisions.api_router import ( register_structured_decisions_api_router, ) +from vllm.entrypoints.generate.structured_decisions.protocol import ( + StructuredDecisionRequest, + StructuredDecisionResponse, +) from vllm.entrypoints.generate.structured_decisions.question_types import LABELS -from vllm.entrypoints.generate.structured_decisions.serving import state_text +from vllm.entrypoints.generate.structured_decisions.serving import ( + ServingStructuredDecisions, + state_text, +) from vllm.entrypoints.generate.structured_decisions.strategies import ( NextTokenStrategy, + ReadContext, reply_label_ids, select_read_strategy, ) +from vllm.entrypoints.openai.models.protocol import BaseModelPath +from vllm.entrypoints.openai.models.serving import OpenAIServingModels from vllm.entrypoints.serve.exception_handling.register import ( init_exception_handler, ) +from vllm.logprobs import Logprob +from vllm.outputs import CompletionOutput, RequestOutput +from vllm.renderers.hf import HfRenderer +from vllm.renderers.online_renderer import OnlineRenderer + +TRACE_HEADERS = { + "traceparent": "00-0123456789abcdef0123456789abcdef-0123456789abcdef-01", + "tracestate": "vendor=value", +} def model(architecture: str, logprobs_mode: str = "raw_logprobs") -> Any: @@ -69,9 +91,9 @@ def test_structured_state_keeps_unicode_in_prompt(): @pytest.fixture(scope="module") def qwen(): - from transformers import AutoTokenizer + from vllm.tokenizers import get_tokenizer - tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B") + tokenizer = get_tokenizer("Qwen/Qwen3-0.6B") prompt_ids = tokenizer.apply_chat_template( [{"role": "user", "content": "state"}], add_generation_prompt=True, @@ -93,3 +115,108 @@ def test_labels_start_the_reply(qwen): # After a colon, Qwen writes ":A" as one token, so "A" is not one token. with pytest.raises(ValueError, match="not one distinct token"): reply_label_ids(tokenizer, tokenizer.encode("team:")) + + +@pytest.fixture +def serving_and_engine(qwen): + tokenizer, _ = qwen + model_config = ModelConfig(model="Qwen/Qwen3-0.6B", dtype="float32") + renderer = HfRenderer( + SimpleNamespace( + model_config=model_config, + parallel_config=SimpleNamespace(_api_process_rank=0), + ), + tokenizer, + ) + engine = MagicMock(spec=EngineClient) + engine.errored = False + engine.model_config = model_config + engine.renderer = renderer + engine.input_processor = MagicMock() + online_renderer = OnlineRenderer( + model_config, + renderer, + request_logger=None, + chat_template=None, + chat_template_content_format="auto", + ) + + async def generate(prompt, params, request_id, **kwargs): + label_ids = params.logprob_token_ids + yield RequestOutput( + request_id=request_id, + prompt=None, + prompt_token_ids=prompt["prompt_token_ids"], + prompt_logprobs=None, + outputs=[ + CompletionOutput( + index=0, + text="A", + token_ids=[label_ids[0]], + cumulative_logprob=None, + logprobs=[{token_id: Logprob(-1.0) for token_id in label_ids}], + finish_reason="length", + ) + ], + finished=True, + ) + + engine.generate.side_effect = generate + strategy = NextTokenStrategy(ReadContext(engine, online_renderer, None, "auto", {})) + models = OpenAIServingModels( + engine, [BaseModelPath(name=model_config.model, model_path=model_config.model)] + ) + yield ServingStructuredDecisions(models, strategy), engine + renderer.shutdown() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "cache_salt,tracing_enabled,request_headers", + [ + ("tenant-a", True, {"authorization": "Bearer test", **TRACE_HEADERS}), + ("tenant-b", False, {"authorization": "Bearer test", **TRACE_HEADERS}), + (None, True, {"authorization": "Bearer test", **TRACE_HEADERS}), + (None, True, {}), + ("tenant-direct", True, None), + ], +) +async def test_request_metadata_reaches_every_question( + serving_and_engine, cache_salt, tracing_enabled, request_headers +): + serving, engine = serving_and_engine + engine.is_tracing_enabled.return_value = tracing_enabled + raw_request = None + if request_headers is not None: + raw_request = Request( + { + "type": "http", + "headers": [ + (k.encode(), v.encode()) for k, v in request_headers.items() + ], + } + ) + request = StructuredDecisionRequest( + state={"team": "billing", "language": "English"}, + questions={ + "team": { + "type": "choice", + "instructions": "Which team does the message name?", + "criteria": {"billing": "payments", "shipping": "deliveries"}, + }, + "lang": { + "type": "choice", + "instructions": "Which language does the message name?", + "criteria": {"English": "English text", "French": "French text"}, + }, + }, + cache_salt=cache_salt, + ) + response = await serving.create_decision(request, raw_request) + assert isinstance(response, StructuredDecisionResponse) + assert engine.generate.call_count == 2 + for call in engine.generate.call_args_list: + assert call.args[0].get("cache_salt") == cache_salt + assert call.kwargs["trace_headers"] == ( + TRACE_HEADERS if tracing_enabled and request_headers else None + ) diff --git a/tests/entrypoints/unit_tests/test_non_object_body_validation.py b/tests/entrypoints/unit_tests/test_non_object_body_validation.py index 5bdc1d5787..dd0e3eafd0 100644 --- a/tests/entrypoints/unit_tests/test_non_object_body_validation.py +++ b/tests/entrypoints/unit_tests/test_non_object_body_validation.py @@ -13,6 +13,9 @@ request models whose before-validators were missing the same guard. import pytest from pydantic import ValidationError +from vllm.entrypoints.generate.structured_decisions.protocol import ( + StructuredDecisionRequest, +) from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest from vllm.entrypoints.openai.completion.protocol import CompletionRequest from vllm.entrypoints.openai.responses.protocol import ResponsesRequest @@ -74,6 +77,18 @@ def test_tokenize_chat_request_still_validates_dict_bodies(): CACHE_SALT_MODELS = [ + ( + StructuredDecisionRequest, + { + "state": "hello", + "questions": { + "lang": { + "type": "choice", + "criteria": {"English": "Text written in English."}, + } + }, + }, + ), (CompletionRequest, {"prompt": "hello"}), (ChatCompletionRequest, {"messages": [{"role": "user", "content": "hello"}]}), (ResponsesRequest, {"input": "hello"}), diff --git a/vllm/entrypoints/generate/structured_decisions/protocol.py b/vllm/entrypoints/generate/structured_decisions/protocol.py index 3a193e203b..b15fee171f 100644 --- a/vllm/entrypoints/generate/structured_decisions/protocol.py +++ b/vllm/entrypoints/generate/structured_decisions/protocol.py @@ -5,10 +5,11 @@ import time from typing import Any, Literal -from pydantic import Field +from pydantic import Field, model_validator from vllm.config import ModelConfig from vllm.entrypoints.chat_utils import ChatTemplateContentFormatOption +from vllm.entrypoints.generate.base.protocol import validate_cache_salt from vllm.entrypoints.serve.engine.protocol import OpenAIBaseModel from vllm.renderers import ChatParams, TokenizeParams, merge_kwargs from vllm.utils import random_uuid @@ -30,9 +31,27 @@ class StructuredDecisionRequest(OpenAIBaseModel): default=None, description="Context placed ahead of the questions." ) chat_template_kwargs: dict[str, Any] | None = None + cache_salt: str | None = Field( + default=None, + description=( + "If specified, the prefix cache will be salted with the provided " + "string to prevent an attacker to guess prompts in multi-user " + "environments. The salt should be random, protected from " + "access by 3rd parties, and long enough to be " + "unpredictable (e.g., 43 characters base64-encoded, corresponding " + "to 256 bit)." + ), + ) priority: int = Field(default=0, ge=-(2**63), le=2**63 - 1) request_id: str = Field(default_factory=random_uuid) + @model_validator(mode="before") + @classmethod + def check_cache_salt_support(cls, data: Any) -> Any: + if isinstance(data, dict): + validate_cache_salt(data.get("cache_salt")) + return data + class DecisionUsage(OpenAIBaseModel): input_tokens: int @@ -65,6 +84,7 @@ class ReadPromptRequest(OpenAIBaseModel): would follow a thought rather than start the reply.""" chat_template_kwargs: dict[str, Any] | None = None + cache_salt: str | None = None def build_chat_params( self, diff --git a/vllm/entrypoints/generate/structured_decisions/serving.py b/vllm/entrypoints/generate/structured_decisions/serving.py index 4921bfcf35..e6e990bd46 100644 --- a/vllm/entrypoints/generate/structured_decisions/serving.py +++ b/vllm/entrypoints/generate/structured_decisions/serving.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import asyncio import json +from collections.abc import Mapping from typing import Any from fastapi import Request @@ -11,6 +12,11 @@ from vllm.entrypoints.serve.engine.protocol import ErrorResponse from vllm.entrypoints.serve.engine.serving import BaseServing from vllm.entrypoints.serve.utils.request_logger import RequestLogger from vllm.logger import init_logger +from vllm.tracing import ( + contains_trace_headers, + extract_trace_headers, + log_tracing_disabled_warning, +) from .protocol import ( DecisionUsage, @@ -71,6 +77,16 @@ class ServingStructuredDecisions(BaseServing): self.strategy = strategy self.limits = strategy.limits() + async def _get_trace_headers( + self, headers: Mapping[str, str] + ) -> Mapping[str, str] | None: + if not contains_trace_headers(headers): + return None + if not await self.strategy.context.engine_client.is_tracing_enabled(): + log_tracing_disabled_warning() + return None + return extract_trace_headers(headers) + async def create_decision( self, request: StructuredDecisionRequest, @@ -88,6 +104,11 @@ class ServingStructuredDecisions(BaseServing): questions = parse_questions(request, self.limits) lora_request = self._maybe_get_adapters(request) # type: ignore[arg-type] engine_client.check_admission(len(questions)) + trace_headers = ( + None + if raw_request is None + else await self._get_trace_headers(raw_request.headers) + ) reads = await self.strategy.read( questions, request.instructions, @@ -96,6 +117,8 @@ class ServingStructuredDecisions(BaseServing): chat_template_kwargs=request.chat_template_kwargs, lora_request=lora_request, priority=request.priority, + cache_salt=request.cache_salt, + trace_headers=trace_headers, ) except StructuredDecisionError as e: return self.create_error_response(e) diff --git a/vllm/entrypoints/generate/structured_decisions/strategies.py b/vllm/entrypoints/generate/structured_decisions/strategies.py index af588e13a2..65b5a9dd7b 100644 --- a/vllm/entrypoints/generate/structured_decisions/strategies.py +++ b/vllm/entrypoints/generate/structured_decisions/strategies.py @@ -7,7 +7,7 @@ Only models in NEXT_TOKEN_ARCHITECTURES are currently supported. import math from abc import ABC, abstractmethod -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from dataclasses import dataclass from typing import Any @@ -74,6 +74,8 @@ class ReadStrategy(ABC): chat_template_kwargs: dict[str, Any] | None, lora_request: LoRARequest | None, priority: int, + cache_salt: str | None = None, + trace_headers: Mapping[str, str] | None = None, ) -> list[QuestionRead]: """One read per question, in order.""" @@ -136,13 +138,15 @@ class NextTokenStrategy(ReadStrategy): return DecisionLimits(max_questions=64, max_options=len(LABELS)) def _read_request( - self, chat_template_kwargs: dict[str, Any] | None + self, chat_template_kwargs: dict[str, Any] | None, cache_salt: str | None = None ) -> ReadPromptRequest: if (chat_template_kwargs or {}).get("enable_thinking"): raise StructuredDecisionError( "a read is the reply's first token, so thinking must be off" ) - return ReadPromptRequest(chat_template_kwargs=chat_template_kwargs) + return ReadPromptRequest( + chat_template_kwargs=chat_template_kwargs, cache_salt=cache_salt + ) async def _render( self, read_request: ReadPromptRequest, messages: list[dict[str, Any]] @@ -170,9 +174,11 @@ class NextTokenStrategy(ReadStrategy): chat_template_kwargs: dict[str, Any] | None, lora_request: LoRARequest | None, priority: int, + cache_salt: str | None = None, + trace_headers: Mapping[str, str] | None = None, ) -> list[QuestionRead]: ctx = self.context - read_request = self._read_request(chat_template_kwargs) + read_request = self._read_request(chat_template_kwargs, cache_salt) slots, engine_inputs = [], [] for q in questions: @@ -197,6 +203,7 @@ class NextTokenStrategy(ReadStrategy): ], request_id, lora_request=lora_request, + trace_headers=trace_headers, priority=priority, )