diff --git a/components/src/dynamo/frontend/prepost.py b/components/src/dynamo/frontend/prepost.py index 512450a987be..ef33adc168c6 100644 --- a/components/src/dynamo/frontend/prepost.py +++ b/components/src/dynamo/frontend/prepost.py @@ -5,10 +5,12 @@ import json import os +import re from collections.abc import Awaitable, Callable, Sequence from concurrent.futures import ThreadPoolExecutor from copy import deepcopy from dataclasses import dataclass +from functools import lru_cache from typing import Any, Protocol, cast from vllm.entrypoints.chat_utils import make_tool_call_id @@ -650,6 +652,22 @@ async def preprocess_chat_request( ) +@lru_cache(maxsize=64) +def _compile_marker_strip_re(strippable: tuple[str, ...]) -> "re.Pattern[str] | None": + """Alternation over the literals to remove, or None if there are none. + + Memoised: the set is a property of (tokenizer, active parser), so a server + sees a handful of distinct values, but a StreamingPostProcessor is built per + request and re.compile over ~15 escaped literals is ~12 us -- worth paying + once per process rather than once per request. + """ + if not strippable: + return None + # Longest-first so a marker that prefixes another cannot win. + ordered = sorted(strippable, key=len, reverse=True) + return re.compile("|".join(re.escape(m) for m in ordered)) + + class StreamingPostProcessor: def __init__( self, @@ -718,6 +736,29 @@ def __init__( t for t in getattr(tokenizer, "all_special_tokens", ()) if t ) + # Markers that may be removed from user-visible content: the tokenizer's + # actual special tokens, minus any the active parser owns. Membership in + # all_special_tokens is the test -- a literal that merely looks like a + # control token (`<|not_a_real_special_token|>`) is ordinary generated + # text and must survive. Built once: it cannot change for the lifetime of + # this postprocessor, and compiled once per process (see the memo above). + # + # Gate on the CONFIGURED parser, not on `self.reasoning_parser`. What forces + # skip_special_tokens=false is the deployment's parser NAME -- + # preprocessor.rs::parser_requires_special_tokens (PRE.1) matches static + # names only, and vLLM's ParserEngine.adjust_request likewise -- so neither + # sees this request's enable_thinking or response_reasoning_ended. Those are + # exactly what null self.reasoning_parser just above, so an instance-keyed + # gate is a no-op precisely on the requests that still receive raw markers. + self._marker_strip_re: re.Pattern[str] | None = None + if reasoning_parser_class is not None or self.tool_parser is not None: + # A tool parser's own markers are its wire format, consumed + # downstream; removing them here would blind the tool-start scan. + owned = set(self._tool_start_markers()) | set(self._tool_end_markers()) + self._marker_strip_re = _compile_marker_strip_re( + tuple(m for m in self._control_markers if m not in owned) + ) + self.previous_text = "" self.previous_token_ids: list[int] = [] self.in_progress_tool_calls: dict[int, DeltaToolCall] = {} @@ -915,6 +956,27 @@ def _tool_end_markers(self) -> tuple[str, ...]: ) ) + # A text-grammar reasoning parser can only split the stream if the model's + # content-kind control markers survive detokenisation, which is why + # ParserEngine.adjust_request forces `skip_special_tokens = False`. But nothing + # then removes the markers the parser did NOT consume, so terminal ones leak + # into user-visible content, e.g.: + # content = "<|end_message|>408<|end_message|>" + # As shipped the two settings are mutually exclusive: the default strips the + # markers before the parser can split (reasoning is silently dropped), and + # False leaks them into the answer. Keep the markers for the parser, then strip + # the residue afterwards. + # + # Scope is deliberately narrow: only when a parser is actually active (the only + # case where skip_special_tokens was forced off), and only over literals that + # are genuinely in `tokenizer.all_special_tokens`. `reasoning` is left untouched + # -- it is already clean, and rewriting it would risk mangling legitimate text + # the model quoted. + def _strip_control_markers(self, text: str | None) -> str | None: + if not text or self._marker_strip_re is None: + return text + return self._marker_strip_re.sub("", text) + @staticmethod def _compose_delta_message( reasoning: str | None, content: str | None @@ -1045,6 +1107,16 @@ def _emit_tool_calls_choice(self, output: Any) -> dict[str, Any]: def _build_choice(self, output: Any, delta: dict[str, Any]) -> dict[str, Any]: if delta.get("tool_calls"): self._tool_call_choices_emitted.add(output.index) + # The single strip point. Every choice this class emits passes through here, + # so the invariant "no unconsumed control marker reaches the client" is + # enforced once at the funnel rather than at each producer -- stripping at + # the producers as well only re-scans text that is about to be scanned. + if isinstance(delta.get("content"), str): + stripped = self._strip_control_markers(delta["content"]) + if stripped: + delta["content"] = stripped + else: + delta.pop("content", None) return { "index": output.index, "delta": delta, diff --git a/components/src/dynamo/frontend/tests/test_vllm_processor_unit.py b/components/src/dynamo/frontend/tests/test_vllm_processor_unit.py index b58888b4e878..3208d3ec89bf 100644 --- a/components/src/dynamo/frontend/tests/test_vllm_processor_unit.py +++ b/components/src/dynamo/frontend/tests/test_vllm_processor_unit.py @@ -24,6 +24,7 @@ ) from transformers import AutoTokenizer from vllm.sampling_params import SamplingParams, StructuredOutputsParams +from vllm.tool_parsers import ToolParser from dynamo.frontend import prepost as prepost_module from dynamo.frontend.prepost import ( @@ -1708,6 +1709,155 @@ async def test_root_thinking_suppresses_deployment_default(self, tokenizer): assert "thinking_mode" not in result.chat_template_kwargs +class _PassthroughStreamingToolParser(ToolParser): + """A tool parser that finds no tool call and passes the raw delta through. + + Subclasses the real ToolParser so every attribute the postprocessor touches + exists, and returns the unstripped `delta_text` the way a real parser does. + """ + + def extract_tool_calls_streaming(self, **kwargs): + return prepost_module.DeltaMessage(content=kwargs["delta_text"]) + + def get_structural_tag(self, request, *, reasoning=False): + return None + + +class _InactiveReasoningParser: + """A configured reasoning parser that never becomes active on a request. + + `enable_thinking=False` and `response_reasoning_ended=True` both leave + StreamingPostProcessor.reasoning_parser None, so this is only ever passed as a + class -- constructing it is the bug this fake exists to catch. + """ + + def __init__(self, tokenizer, *, chat_template_kwargs): + raise AssertionError("must not be constructed for an inactive request") + + +class _MarkerOwningToolParser(_PassthroughStreamingToolParser): + """A tool parser that owns two of the tokenizer's special tokens.""" + + tool_call_start_token = "<|im_start|>" + tool_call_end_token = "<|im_end|>" + + +class TestControlMarkerStrip: + """Unconsumed special tokens must not reach the client. + + A text-grammar reasoning parser can only split the stream if the model's + content-kind markers survive detokenisation, which is why + ParserEngine.adjust_request forces skip_special_tokens=False. Nothing then + removes the markers the parser did not consume, so terminal ones leak into + visible content -- the report was content = "<|end_message|>408<|end_message|>", + a special token on that model's tokenizer. These use `<|im_end|>`, which is + the equivalent on Qwen3-0.6B. + """ + + @staticmethod + def _post( + tokenizer, + *, + tool_parser=None, + reasoning_parser_class=None, + response_reasoning_ended=None, + chat_template_kwargs=None, + ): + request = json.loads(json.dumps(TOOL_REQUEST)) + request.pop("tools", None) + request, _, _, _, _ = _prepare_request( + request, tokenizer=tokenizer, tool_parser_class=None + ) + return StreamingPostProcessor( + tokenizer=tokenizer, + request_for_sampling=request, + sampling_params=SamplingParams(), + prompt_token_ids=[], + tool_parser=tool_parser, + reasoning_parser_class=reasoning_parser_class, + chat_template_kwargs=chat_template_kwargs or {}, + response_reasoning_ended=response_reasoning_ended, + stream_response=True, + ) + + @staticmethod + def _content(post, text): + choice = post.process_output( + SimpleNamespace( + index=0, + text=text, + token_ids=[], + finish_reason="stop", + logprobs=None, + ) + ) + assert choice is not None + return choice["delta"].get("content") + + def test_reported_leak_is_removed(self, tokenizer): + post = self._post( + tokenizer, tool_parser=_PassthroughStreamingToolParser(tokenizer) + ) + assert self._content(post, "<|im_end|>408<|im_end|>") == "408" + + def test_non_special_literal_of_the_same_shape_survives(self, tokenizer): + """Looking like a control token is not proof of being one. + + `<|not_a_real_special_token|>` is absent from all_special_tokens and + encodes to ordinary token IDs, so it is generated text and must reach the + client intact. + """ + post = self._post( + tokenizer, tool_parser=_PassthroughStreamingToolParser(tokenizer) + ) + literal = "The literal is <|not_a_real_special_token|>." + assert self._content(post, literal) == literal + + def test_parser_owned_markers_are_preserved(self, tokenizer): + """A marker the parser owns is wire format, consumed downstream.""" + post = self._post(tokenizer, tool_parser=_MarkerOwningToolParser(tokenizer)) + assert ( + post._strip_control_markers("<|im_start|>{}<|im_end|>") + == "<|im_start|>{}<|im_end|>" + ) + # A special token it does not own is still removed. + assert post._strip_control_markers("<|endoftext|>x") == "x" + + def test_no_parser_configured_is_passthrough(self, tokenizer): + """Only a *configured* parser forces skip_special_tokens off.""" + post = self._post(tokenizer) + assert post.tool_parser is None and post.reasoning_parser is None + assert self._content(post, "<|im_end|>408") == "<|im_end|>408" + + @pytest.mark.parametrize( + "inactive_kwargs", + [ + {"response_reasoning_ended": True}, + {"chat_template_kwargs": {"enable_thinking": False}}, + ], + ids=["reasoning-already-ended", "thinking-disabled"], + ) + def test_configured_but_inactive_reasoning_parser_still_strips( + self, tokenizer, inactive_kwargs + ): + """The CONFIGURED parser forces skip_special_tokens off, not the instance. + + preprocessor.rs::parser_requires_special_tokens keys the + skip_special_tokens=false default off the deployment's configured + reasoning_parser NAME; it never sees this request's enable_thinking or + response_reasoning_ended. But this class leaves self.reasoning_parser None on + exactly those requests, so gating the strip on the instance made it a no-op + precisely where markers still arrive undecoded. + """ + post = self._post( + tokenizer, + reasoning_parser_class=_InactiveReasoningParser, + **inactive_kwargs, + ) + assert post.reasoning_parser is None + assert self._content(post, "<|im_end|>408<|im_end|>") == "408" + + class TestToolCallGuidedDecoding: def _request(self, tokenizer, **overrides): raw_request = {**TOOL_REQUEST, "tool_choice": "auto", **overrides}