wehub-resource-sync
427 lines
14 KiB
Python
427 lines
14 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Data-driven replay harness for parser engine testing.
|
|
|
|
Replays token sequences through parsers at different chunk sizes to
|
|
verify chunk-size invariance: the same token sequence must produce
|
|
identical output regardless of how tokens are batched.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from collections.abc import Sequence
|
|
from dataclasses import dataclass, field
|
|
|
|
from vllm.entrypoints.openai.chat_completion.protocol import (
|
|
ChatCompletionRequest,
|
|
)
|
|
from vllm.entrypoints.openai.engine.protocol import DeltaMessage
|
|
|
|
|
|
@dataclass
|
|
class Sample:
|
|
"""One test sample loaded from a JSONL file."""
|
|
|
|
id: str
|
|
description: str
|
|
source: str
|
|
vocab: dict[str, int]
|
|
tokens: list[tuple[int, str]]
|
|
expected_reasoning: str | None
|
|
expected_content: str | None
|
|
expected_tool_calls: list[dict] | None
|
|
tools: list[dict] | None = None
|
|
chat_template_kwargs: dict | None = None
|
|
prompt_token_ids: list[int] | None = None
|
|
|
|
|
|
@dataclass
|
|
class ParseOutput:
|
|
"""Accumulated parse output from replaying a token stream."""
|
|
|
|
reasoning: str = ""
|
|
content: str = ""
|
|
tool_calls: list[dict] = field(default_factory=list)
|
|
|
|
|
|
class MockTokenizer:
|
|
"""Lightweight tokenizer mock that avoids unittest.mock overhead.
|
|
|
|
Used by ``benchmarks/benchmark_parsers.py`` in tight timing loops,
|
|
so hot-path methods (``decode``, ``get_vocab``) must be cheap.
|
|
MagicMock's call-recording machinery added ~40% overhead to small-
|
|
sample benchmarks, inflating the per-token cost of the parser engine.
|
|
"""
|
|
|
|
__slots__ = (
|
|
"_vocab",
|
|
"_token_ids",
|
|
"_token_decode_map",
|
|
"_special_ids",
|
|
"eos_token_id",
|
|
"bos_token_id",
|
|
"pad_token_id",
|
|
"_all_special_tokens",
|
|
)
|
|
|
|
def __init__(
|
|
self,
|
|
vocab: dict[str, int],
|
|
tokens: list[tuple[int, str]],
|
|
) -> None:
|
|
self._vocab = vocab
|
|
self._token_ids = [tid for tid, _ in tokens]
|
|
self._token_decode_map = {tid: text for tid, text in tokens}
|
|
self._special_ids = set(vocab.values())
|
|
self._all_special_tokens = list(vocab.keys())
|
|
self.eos_token_id = None
|
|
self.bos_token_id = None
|
|
self.pad_token_id = None
|
|
|
|
def set_vocab(self, vocab: dict[str, int]) -> None:
|
|
self._vocab = vocab
|
|
|
|
def get_vocab(self) -> dict[str, int]:
|
|
return self._vocab
|
|
|
|
@property
|
|
def all_special_tokens(self) -> list[str]:
|
|
return self._all_special_tokens
|
|
|
|
@property
|
|
def all_special_ids(self) -> list[int]:
|
|
return [self._vocab[t] for t in self._all_special_tokens if t in self._vocab]
|
|
|
|
def encode(self, text: str, **kwargs) -> list[int]:
|
|
return self._token_ids
|
|
|
|
def decode(self, ids: list[int], skip_special_tokens: bool = False) -> str:
|
|
parts: list[str] = []
|
|
for tid in ids:
|
|
if skip_special_tokens and tid in self._special_ids:
|
|
continue
|
|
text = self._token_decode_map.get(tid, f"?{tid}?")
|
|
parts.append(text)
|
|
return "".join(parts)
|
|
|
|
|
|
CHUNK_SIZES = [1, 2, 3, 5, 11, 23, None]
|
|
|
|
|
|
def make_mock_tokenizer(sample: Sample) -> MockTokenizer:
|
|
"""Build a mock tokenizer from a sample's vocab and token data."""
|
|
return MockTokenizer(
|
|
vocab=dict(sample.vocab),
|
|
tokens=sample.tokens,
|
|
)
|
|
|
|
|
|
def _test_request(
|
|
tools: list[dict] | None = None,
|
|
) -> ChatCompletionRequest:
|
|
return ChatCompletionRequest(
|
|
model="test-model",
|
|
messages=[{"role": "user", "content": "test"}],
|
|
tools=tools,
|
|
)
|
|
|
|
|
|
DUMMY_TOOLS = [
|
|
{
|
|
"type": "function",
|
|
"function": {"name": "stub", "parameters": {"type": "object"}},
|
|
},
|
|
]
|
|
|
|
|
|
def parse_non_streaming(
|
|
parser,
|
|
sample: Sample,
|
|
request: ChatCompletionRequest,
|
|
) -> ParseOutput:
|
|
"""Run ``parser.parse()`` and return a :class:`ParseOutput`."""
|
|
full_text = "".join(text for _, text in sample.tokens)
|
|
reasoning, content, tool_calls = parser.parse(
|
|
full_text,
|
|
request,
|
|
enable_auto_tools=True,
|
|
)
|
|
tc_list: list[dict] = []
|
|
if tool_calls:
|
|
for tc in tool_calls:
|
|
tc_list.append({"name": tc.name, "arguments": tc.arguments})
|
|
return ParseOutput(
|
|
reasoning=reasoning or "",
|
|
content=content or "",
|
|
tool_calls=tc_list,
|
|
)
|
|
|
|
|
|
def replay_streaming(
|
|
parser,
|
|
tokens: list[tuple[int, str]],
|
|
chunk_size: int | None = None,
|
|
holdback_chars: int = 0,
|
|
finished_on_last: bool = False,
|
|
tools: list[dict] | None = None,
|
|
prompt_token_ids: list[int] | None = None,
|
|
) -> list[DeltaMessage | None]:
|
|
"""Feed tokens through ``parser.parse_delta()`` at a given chunk size.
|
|
|
|
Args:
|
|
parser: A :class:`Parser` instance with ``parse_delta()`` method.
|
|
tokens: List of ``(token_id, decoded_text)`` pairs.
|
|
chunk_size: Number of tokens per batch. ``None`` means all at once.
|
|
holdback_chars: Simulate detokenizer holdback by holding back
|
|
this many characters of decoded text between batches.
|
|
finished_on_last: When True, pass ``finished=True`` on the last
|
|
``parse_delta()`` call, matching real server behavior.
|
|
tools: Optional tool definitions to include on the request,
|
|
matching the serving layer where tools set
|
|
``tool_choice`` to ``"auto"``.
|
|
|
|
Returns:
|
|
List of ``DeltaMessage`` results from each ``parse_delta()`` call.
|
|
"""
|
|
if chunk_size is None:
|
|
chunk_size = len(tokens)
|
|
|
|
results: list[DeltaMessage | None] = []
|
|
all_ids = [tid for tid, _ in tokens]
|
|
all_texts = [text for _, text in tokens]
|
|
|
|
request = _test_request(tools=tools)
|
|
first_prompt_ids = prompt_token_ids if prompt_token_ids is not None else []
|
|
|
|
if holdback_chars <= 0:
|
|
chunks = list(range(0, len(tokens), chunk_size))
|
|
for i, start in enumerate(chunks):
|
|
batch_end = min(start + chunk_size, len(tokens))
|
|
batch_ids = all_ids[start:batch_end]
|
|
delta_text = "".join(all_texts[start:batch_end])
|
|
is_last = i == len(chunks) - 1
|
|
|
|
result = parser.parse_delta(
|
|
delta_text,
|
|
batch_ids,
|
|
request,
|
|
prompt_token_ids=first_prompt_ids if start == 0 else None,
|
|
finished=finished_on_last and is_last,
|
|
)
|
|
results.append(result)
|
|
return results
|
|
|
|
emitted_up_to = 0
|
|
is_first = True
|
|
|
|
for start in range(0, len(tokens), chunk_size):
|
|
batch_end = min(start + chunk_size, len(tokens))
|
|
|
|
if batch_end < len(tokens):
|
|
held_chars = 0
|
|
safe_end = batch_end
|
|
while safe_end > emitted_up_to and held_chars < holdback_chars:
|
|
safe_end -= 1
|
|
held_chars += len(all_texts[safe_end])
|
|
else:
|
|
safe_end = batch_end
|
|
|
|
if safe_end <= emitted_up_to:
|
|
continue
|
|
|
|
batch_ids = all_ids[emitted_up_to:safe_end]
|
|
delta_text = "".join(all_texts[emitted_up_to:safe_end])
|
|
emitted_up_to = safe_end
|
|
|
|
is_last_chunk = batch_end >= len(tokens)
|
|
result = parser.parse_delta(
|
|
delta_text,
|
|
batch_ids,
|
|
request,
|
|
prompt_token_ids=first_prompt_ids if is_first else None,
|
|
finished=finished_on_last and is_last_chunk,
|
|
)
|
|
results.append(result)
|
|
is_first = False
|
|
|
|
if emitted_up_to < len(tokens):
|
|
batch_ids = all_ids[emitted_up_to:]
|
|
delta_text = "".join(all_texts[emitted_up_to:])
|
|
result = parser.parse_delta(
|
|
delta_text,
|
|
batch_ids,
|
|
request,
|
|
prompt_token_ids=first_prompt_ids if is_first else None,
|
|
finished=finished_on_last,
|
|
)
|
|
results.append(result)
|
|
|
|
return results
|
|
|
|
|
|
def replay_with_text_holdback(
|
|
parser,
|
|
tokens: list[tuple[int, str]],
|
|
text_delay: int = 1,
|
|
tools: list[dict] | None = None,
|
|
prompt_token_ids: list[int] | None = None,
|
|
) -> list[DeltaMessage | None]:
|
|
"""Replay token-by-token with text arriving *text_delay* steps late.
|
|
|
|
Simulates the production detokenizer holdback where token IDs arrive
|
|
immediately but decoded text is delayed. On the last token all
|
|
remaining held-back text is flushed, matching real server behavior::
|
|
|
|
step 0: ids=[tok0], text="" (held back)
|
|
step 1: ids=[tok1], text=tok0_text (tok0 released)
|
|
...
|
|
step N-1: ids=[tokN-1], text=remaining_texts (flush all)
|
|
|
|
This exercises the TokenIDScanner deferred-terminal path that
|
|
``replay_streaming`` (which keeps text and IDs aligned) does not.
|
|
"""
|
|
results: list[DeltaMessage | None] = []
|
|
request = _test_request(tools=tools)
|
|
first_prompt_ids = prompt_token_ids if prompt_token_ids is not None else []
|
|
|
|
n = len(tokens)
|
|
held_texts: list[str] = []
|
|
|
|
for i in range(n):
|
|
token_id = tokens[i][0]
|
|
held_texts.append(tokens[i][1])
|
|
|
|
is_last = i == n - 1
|
|
if is_last:
|
|
delta_text = "".join(held_texts)
|
|
held_texts.clear()
|
|
elif len(held_texts) > text_delay:
|
|
delta_text = held_texts.pop(0)
|
|
else:
|
|
delta_text = ""
|
|
|
|
result = parser.parse_delta(
|
|
delta_text,
|
|
[token_id],
|
|
request,
|
|
prompt_token_ids=first_prompt_ids if i == 0 else None,
|
|
finished=is_last,
|
|
)
|
|
results.append(result)
|
|
|
|
return results
|
|
|
|
|
|
def accumulate_deltas(
|
|
deltas: Sequence[DeltaMessage | None],
|
|
) -> dict:
|
|
reasoning_parts: list[str] = []
|
|
content_parts: list[str] = []
|
|
tool_calls_by_idx: dict[int, dict] = {}
|
|
|
|
for delta in deltas:
|
|
if delta is None:
|
|
continue
|
|
if delta.reasoning:
|
|
reasoning_parts.append(delta.reasoning)
|
|
if delta.content:
|
|
content_parts.append(delta.content)
|
|
if delta.tool_calls:
|
|
for tc in delta.tool_calls:
|
|
if tc.function and tc.function.name:
|
|
existing = tool_calls_by_idx.get(tc.index)
|
|
if existing is None:
|
|
tool_calls_by_idx[tc.index] = {
|
|
"name": tc.function.name,
|
|
"_args_parts": [tc.function.arguments or ""],
|
|
}
|
|
else:
|
|
existing["_args_parts"].append(tc.function.arguments or "")
|
|
elif tc.function and tc.function.arguments:
|
|
existing = tool_calls_by_idx.get(tc.index)
|
|
if existing is not None:
|
|
existing["_args_parts"].append(tc.function.arguments)
|
|
|
|
return {
|
|
"reasoning": "".join(reasoning_parts),
|
|
"content": "".join(content_parts),
|
|
"tool_calls": [
|
|
{"name": tc["name"], "arguments": "".join(tc["_args_parts"])}
|
|
for tc in tool_calls_by_idx.values()
|
|
],
|
|
}
|
|
|
|
|
|
def collect_output(results: list[DeltaMessage | None]) -> ParseOutput:
|
|
"""Accumulate ``DeltaMessage`` results into a :class:`ParseOutput`."""
|
|
result = accumulate_deltas(results)
|
|
return ParseOutput(
|
|
reasoning=result["reasoning"],
|
|
content=result["content"],
|
|
tool_calls=result["tool_calls"],
|
|
)
|
|
|
|
|
|
def assert_parse_output(actual: ParseOutput, sample: Sample) -> None:
|
|
"""Compare actual parse output against expected values from a sample."""
|
|
if sample.expected_reasoning is not None:
|
|
assert actual.reasoning == sample.expected_reasoning, (
|
|
f"Reasoning mismatch:\n"
|
|
f" expected: {sample.expected_reasoning!r}\n"
|
|
f" actual: {actual.reasoning!r}"
|
|
)
|
|
|
|
if sample.expected_content is not None:
|
|
assert actual.content == sample.expected_content, (
|
|
f"Content mismatch:\n"
|
|
f" expected: {sample.expected_content!r}\n"
|
|
f" actual: {actual.content!r}"
|
|
)
|
|
if sample.expected_tool_calls is not None:
|
|
assert len(actual.tool_calls) == len(sample.expected_tool_calls), (
|
|
f"Tool call count mismatch: "
|
|
f"expected {len(sample.expected_tool_calls)}, "
|
|
f"got {len(actual.tool_calls)}"
|
|
)
|
|
for i, (expected_tc, actual_tc) in enumerate(
|
|
zip(sample.expected_tool_calls, actual.tool_calls)
|
|
):
|
|
assert actual_tc["name"] == expected_tc["name"], (
|
|
f"Tool call {i} name mismatch: "
|
|
f"expected {expected_tc['name']!r}, "
|
|
f"got {actual_tc['name']!r}"
|
|
)
|
|
if "arguments" in expected_tc:
|
|
expected_args = expected_tc["arguments"]
|
|
actual_args_str = actual_tc.get("arguments", "{}")
|
|
if isinstance(expected_args, dict):
|
|
try:
|
|
actual_args = json.loads(actual_args_str)
|
|
except json.JSONDecodeError as e:
|
|
raise AssertionError(
|
|
f"Tool call {i} arguments not valid JSON: "
|
|
f"{actual_args_str!r}"
|
|
) from e
|
|
assert actual_args == expected_args, (
|
|
f"Tool call {i} arguments mismatch:\n"
|
|
f" expected: {expected_args}\n"
|
|
f" actual: {actual_args}"
|
|
)
|
|
|
|
|
|
def assert_no_terminal_leakage(
|
|
actual: ParseOutput,
|
|
terminals: list[str],
|
|
context: str = "",
|
|
) -> None:
|
|
"""Assert that none of *terminals* appear in reasoning or content."""
|
|
suffix = f" ({context})" if context else ""
|
|
for terminal in terminals:
|
|
assert terminal not in actual.reasoning, (
|
|
f"{terminal!r} leaked into reasoning{suffix}"
|
|
)
|
|
assert terminal not in actual.content, (
|
|
f"{terminal!r} leaked into content{suffix}"
|
|
)
|