wehub-resource-sync
81 lines
3.2 KiB
Python
81 lines
3.2 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Unit tests for the token-offsets request/response protocol wiring:
|
|
the request flag flowing into ``TokenizeParams`` and the ``GenerateRequest``
|
|
serialization boundary. End-to-end behavior is covered by
|
|
``tests/entrypoints/scale_out/render/test_render.py``; plain Pydantic field
|
|
storage is not retested here.
|
|
"""
|
|
|
|
from unittest.mock import Mock
|
|
|
|
from vllm.config import ModelConfig
|
|
from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
|
|
from vllm.entrypoints.openai.completion.protocol import CompletionRequest
|
|
from vllm.entrypoints.scale_out.token_in_token_out.protocol import GenerateRequest
|
|
from vllm.sampling_params import SamplingParams
|
|
|
|
|
|
def _model_config() -> Mock:
|
|
model_config = Mock(spec=ModelConfig)
|
|
model_config.max_model_len = 128
|
|
return model_config
|
|
|
|
|
|
def test_completion_flag_forwarded_to_tok_params():
|
|
"""build_tok_params must forward return_token_offsets, defaulting to
|
|
False (zero behavioral change for existing callers) and coercing JSON
|
|
null to False via the bool() guard."""
|
|
cfg = _model_config()
|
|
|
|
default = CompletionRequest(model="m", prompt="hi")
|
|
assert default.build_tok_params(cfg).return_token_offsets is False
|
|
|
|
on = CompletionRequest(model="m", prompt="hi", return_token_offsets=True)
|
|
assert on.build_tok_params(cfg).return_token_offsets is True
|
|
|
|
null = CompletionRequest(model="m", prompt="hi", return_token_offsets=None)
|
|
assert null.build_tok_params(cfg).return_token_offsets is False
|
|
|
|
|
|
def test_chat_flag_forwarded_to_tok_params():
|
|
"""Chat build_tok_params has its own (max_completion_tokens) branch, so
|
|
its return_token_offsets forwarding is verified independently."""
|
|
cfg = _model_config()
|
|
messages = [{"role": "user", "content": "hi"}]
|
|
|
|
default = ChatCompletionRequest(model="m", messages=messages)
|
|
assert default.build_tok_params(cfg).return_token_offsets is False
|
|
|
|
on = ChatCompletionRequest(model="m", messages=messages, return_token_offsets=True)
|
|
assert on.build_tok_params(cfg).return_token_offsets is True
|
|
|
|
null = ChatCompletionRequest(
|
|
model="m", messages=messages, return_token_offsets=None
|
|
)
|
|
assert null.build_tok_params(cfg).return_token_offsets is False
|
|
|
|
|
|
def test_generate_request_token_offsets_default_none():
|
|
"""Defaults to None so existing /v1/.../render responses are unchanged."""
|
|
req = GenerateRequest(token_ids=[1, 2, 3], sampling_params=SamplingParams())
|
|
assert req.token_offsets is None
|
|
|
|
|
|
def test_generate_request_token_offsets_survive_json_round_trip():
|
|
"""GenerateRequest crosses the disagg serialization boundary; the
|
|
tuple[int, int] offsets must survive model_dump and re-validate."""
|
|
req = GenerateRequest(
|
|
token_ids=[10, 20],
|
|
sampling_params=SamplingParams(),
|
|
token_offsets=[(0, 1), (1, 3)],
|
|
)
|
|
dumped = req.model_dump()
|
|
assert dumped["token_offsets"] == [(0, 1), (1, 3)]
|
|
# Re-validate from the dumped dict (sampling_params doesn't round-trip
|
|
# cleanly via dump, so re-inject a fresh instance).
|
|
again = GenerateRequest.model_validate(
|
|
{**dumped, "sampling_params": SamplingParams()}
|
|
)
|
|
assert again.token_offsets == [(0, 1), (1, 3)]
|