Add vLLM v0.18.1 source tree with KV transfer abort fix
third_party/vllm/ now tracked in git for direct patch management.
Based on vLLM v0.18.1 release with one patch applied:
vllm/v1/core/sched/scheduler.py:
Replace fatal assert with graceful skip when KV transfer callback
arrives for an already-aborted request during PD disaggregated serving.
Future vLLM modifications should be made directly in third_party/vllm/
and committed normally. The patches/ directory is kept as documentation
of what changed from upstream.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
0
third_party/vllm/tests/renderers/__init__.py
vendored
Normal file
0
third_party/vllm/tests/renderers/__init__.py
vendored
Normal file
0
third_party/vllm/tests/renderers/inputs/__init__.py
vendored
Normal file
0
third_party/vllm/tests/renderers/inputs/__init__.py
vendored
Normal file
41
third_party/vllm/tests/renderers/inputs/test_preprocess.py
vendored
Normal file
41
third_party/vllm/tests/renderers/inputs/test_preprocess.py
vendored
Normal file
@@ -0,0 +1,41 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
from vllm.renderers.inputs.preprocess import prompt_to_seq
|
||||
|
||||
|
||||
def test_empty_input():
|
||||
assert prompt_to_seq([]) == []
|
||||
assert prompt_to_seq([[]]) == [[]]
|
||||
assert prompt_to_seq([[], []]) == [[], []]
|
||||
|
||||
|
||||
def test_text_input():
|
||||
assert prompt_to_seq("foo") == ["foo"]
|
||||
assert prompt_to_seq(["foo"]) == ["foo"]
|
||||
assert prompt_to_seq(["foo", "bar"]) == ["foo", "bar"]
|
||||
|
||||
|
||||
def test_token_input():
|
||||
assert prompt_to_seq([1, 2]) == [[1, 2]]
|
||||
assert prompt_to_seq([[1, 2]]) == [[1, 2]]
|
||||
assert prompt_to_seq([[1, 2], [3, 4]]) == [[1, 2], [3, 4]]
|
||||
|
||||
|
||||
def test_text_token_input():
|
||||
assert prompt_to_seq([[1, 2], "foo"]) == [[1, 2], "foo"]
|
||||
assert prompt_to_seq(["foo", [1, 2]]) == ["foo", [1, 2]]
|
||||
|
||||
|
||||
def test_bytes_input():
|
||||
assert prompt_to_seq(b"foo") == [b"foo"]
|
||||
assert prompt_to_seq([b"foo"]) == [b"foo"]
|
||||
assert prompt_to_seq([b"foo", b"bar"]) == [b"foo", b"bar"]
|
||||
|
||||
|
||||
def test_dict_input():
|
||||
assert prompt_to_seq({"prompt": "foo"}) == [{"prompt": "foo"}]
|
||||
assert prompt_to_seq([{"prompt": "foo"}]) == [{"prompt": "foo"}]
|
||||
assert prompt_to_seq([{"prompt": "foo"}, {"prompt_token_ids": [1, 2]}]) == [
|
||||
{"prompt": "foo"},
|
||||
{"prompt_token_ids": [1, 2]},
|
||||
]
|
||||
500
third_party/vllm/tests/renderers/test_completions.py
vendored
Normal file
500
third_party/vllm/tests/renderers/test_completions.py
vendored
Normal file
@@ -0,0 +1,500 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import io
|
||||
from collections.abc import Sequence
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
import pybase64
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.config import ModelConfig
|
||||
from vllm.inputs import SingletonPrompt
|
||||
from vllm.renderers import TokenizeParams
|
||||
from vllm.renderers.hf import HfRenderer
|
||||
from vllm.renderers.inputs.preprocess import parse_model_prompt, prompt_to_seq
|
||||
from vllm.tokenizers.registry import tokenizer_args_from_config
|
||||
|
||||
MODEL_NAME = "openai-community/gpt2"
|
||||
|
||||
|
||||
@dataclass
|
||||
class MockHFConfig:
|
||||
model_type: str = "any"
|
||||
|
||||
|
||||
@dataclass
|
||||
class MockModelConfig:
|
||||
runner_type = "generate"
|
||||
model: str = MODEL_NAME
|
||||
tokenizer: str = MODEL_NAME
|
||||
trust_remote_code: bool = False
|
||||
tokenizer_revision = None
|
||||
tokenizer_mode = "auto"
|
||||
hf_config = MockHFConfig()
|
||||
encoder_config: dict[str, Any] | None = None
|
||||
enable_prompt_embeds: bool = True
|
||||
skip_tokenizer_init: bool = False
|
||||
is_encoder_decoder: bool = False
|
||||
is_multimodal_model: bool = False
|
||||
|
||||
|
||||
@dataclass
|
||||
class MockParallelConfig:
|
||||
_api_process_rank: int = 0
|
||||
|
||||
|
||||
@dataclass
|
||||
class MockVllmConfig:
|
||||
model_config: MockModelConfig
|
||||
parallel_config: MockParallelConfig
|
||||
|
||||
|
||||
@dataclass
|
||||
class DummyTokenizer:
|
||||
truncation_side: str = "left"
|
||||
max_chars_per_token: int = 1
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
self._captured_encode_kwargs: dict = {}
|
||||
|
||||
def decode(self, tokens: list[int]):
|
||||
return str(tokens)
|
||||
|
||||
def encode(self, text: str, **kwargs):
|
||||
self._captured_encode_kwargs = kwargs
|
||||
|
||||
in_length = len(text)
|
||||
truncation = kwargs.get("truncation")
|
||||
max_length = kwargs.get("max_length")
|
||||
if truncation and max_length is not None:
|
||||
return list(range(min(in_length, max_length)))
|
||||
|
||||
return list(range(in_length))
|
||||
|
||||
|
||||
def _build_renderer(
|
||||
model_config: MockModelConfig,
|
||||
*,
|
||||
truncation_side: str = "left",
|
||||
max_chars_per_token: int = 1,
|
||||
):
|
||||
_, tokenizer_name, _, kwargs = tokenizer_args_from_config(model_config)
|
||||
|
||||
renderer = HfRenderer(
|
||||
MockVllmConfig(model_config, parallel_config=MockParallelConfig()),
|
||||
tokenizer=(
|
||||
None
|
||||
if model_config.skip_tokenizer_init
|
||||
else DummyTokenizer(
|
||||
truncation_side=truncation_side,
|
||||
max_chars_per_token=max_chars_per_token,
|
||||
)
|
||||
),
|
||||
)
|
||||
|
||||
return renderer
|
||||
|
||||
|
||||
def _preprocess_prompt(
|
||||
model_config: ModelConfig,
|
||||
prompt_or_prompts: SingletonPrompt | bytes | Sequence[SingletonPrompt | bytes],
|
||||
):
|
||||
return [
|
||||
(
|
||||
prompt
|
||||
if isinstance(prompt, bytes)
|
||||
else parse_model_prompt(model_config, prompt)
|
||||
)
|
||||
for prompt in prompt_to_seq(prompt_or_prompts)
|
||||
]
|
||||
|
||||
|
||||
class TestValidatePrompt:
|
||||
def test_empty_input(self):
|
||||
renderer = _build_renderer(MockModelConfig())
|
||||
|
||||
with pytest.raises(ValueError, match="at least one prompt"):
|
||||
renderer.render_prompts(_preprocess_prompt(renderer.model_config, []))
|
||||
|
||||
def test_invalid_type(self):
|
||||
renderer = _build_renderer(MockModelConfig())
|
||||
|
||||
with pytest.raises(TypeError, match="should be a list of integers"):
|
||||
renderer.render_prompts(
|
||||
_preprocess_prompt(renderer.model_config, [[1, 2], ["foo", "bar"]]) # type: ignore[arg-type]
|
||||
)
|
||||
|
||||
|
||||
class TestRenderPrompt:
|
||||
def test_token_input(self):
|
||||
renderer = _build_renderer(MockModelConfig())
|
||||
|
||||
tokens = [101, 7592, 2088]
|
||||
prompts = renderer.render_prompts(
|
||||
_preprocess_prompt(renderer.model_config, tokens)
|
||||
)
|
||||
results = renderer.tokenize_prompts(
|
||||
prompts,
|
||||
TokenizeParams(max_total_tokens=100),
|
||||
)
|
||||
|
||||
assert len(results) == 1
|
||||
assert results[0]["prompt_token_ids"] == tokens
|
||||
|
||||
def test_token_list_input(self):
|
||||
renderer = _build_renderer(MockModelConfig())
|
||||
|
||||
token_lists = [[101, 7592, 2088], [102, 1234, 5678, 9012], [103, 4567]]
|
||||
prompts = renderer.render_prompts(
|
||||
_preprocess_prompt(renderer.model_config, token_lists)
|
||||
)
|
||||
results = renderer.tokenize_prompts(
|
||||
prompts,
|
||||
TokenizeParams(max_total_tokens=100),
|
||||
)
|
||||
|
||||
assert len(results) == 3
|
||||
assert results[0]["prompt_token_ids"] == [101, 7592, 2088]
|
||||
assert results[1]["prompt_token_ids"] == [102, 1234, 5678, 9012]
|
||||
assert results[2]["prompt_token_ids"] == [103, 4567]
|
||||
|
||||
def test_text_input(self):
|
||||
renderer = _build_renderer(MockModelConfig())
|
||||
|
||||
text_input = "x" * 10
|
||||
prompts = renderer.render_prompts(
|
||||
_preprocess_prompt(renderer.model_config, text_input)
|
||||
)
|
||||
results = renderer.tokenize_prompts(
|
||||
prompts,
|
||||
TokenizeParams(max_total_tokens=100),
|
||||
)
|
||||
|
||||
assert len(results) == 1
|
||||
assert len(results[0]["prompt_token_ids"]) == 10
|
||||
|
||||
def test_text_list_input(self):
|
||||
renderer = _build_renderer(MockModelConfig())
|
||||
|
||||
text_list_input = ["x" * 10, "x" * 12, "x" * 14]
|
||||
prompts = renderer.render_prompts(
|
||||
_preprocess_prompt(renderer.model_config, text_list_input)
|
||||
)
|
||||
results = renderer.tokenize_prompts(
|
||||
prompts,
|
||||
TokenizeParams(max_total_tokens=100),
|
||||
)
|
||||
|
||||
assert len(results) == 3
|
||||
for text_input, result in zip(text_list_input, results):
|
||||
assert len(result["prompt_token_ids"]) == len(text_input)
|
||||
|
||||
def test_zero_truncation(self):
|
||||
renderer = _build_renderer(MockModelConfig())
|
||||
|
||||
prompts = renderer.render_prompts(
|
||||
_preprocess_prompt(renderer.model_config, "x" * 200)
|
||||
)
|
||||
results = renderer.tokenize_prompts(
|
||||
prompts,
|
||||
TokenizeParams(max_total_tokens=100, truncate_prompt_tokens=0),
|
||||
)
|
||||
|
||||
assert len(results) == 1
|
||||
assert len(results[0]["prompt_token_ids"]) == 0
|
||||
|
||||
def test_pos_truncation(self):
|
||||
renderer = _build_renderer(MockModelConfig())
|
||||
|
||||
prompts = renderer.render_prompts(
|
||||
_preprocess_prompt(renderer.model_config, "x" * 200)
|
||||
)
|
||||
results = renderer.tokenize_prompts(
|
||||
prompts,
|
||||
TokenizeParams(max_total_tokens=100, truncate_prompt_tokens=50),
|
||||
)
|
||||
|
||||
assert len(results) == 1
|
||||
assert len(results[0]["prompt_token_ids"]) == 50
|
||||
|
||||
def test_neg_truncation(self):
|
||||
renderer = _build_renderer(MockModelConfig())
|
||||
|
||||
prompts = renderer.render_prompts(
|
||||
_preprocess_prompt(renderer.model_config, "x" * 200)
|
||||
)
|
||||
results = renderer.tokenize_prompts(
|
||||
prompts,
|
||||
TokenizeParams(max_total_tokens=100, truncate_prompt_tokens=-1),
|
||||
)
|
||||
|
||||
assert len(results) == 1
|
||||
assert len(results[0]["prompt_token_ids"]) == 100 # max_total_tokens
|
||||
|
||||
def test_truncation_left(self):
|
||||
renderer = _build_renderer(MockModelConfig(), truncation_side="left")
|
||||
|
||||
long_tokens = [100, 101, 102, 103, 104, 105, 106, 107, 108, 109] # 10 tokens
|
||||
prompts = renderer.render_prompts(
|
||||
_preprocess_prompt(renderer.model_config, long_tokens)
|
||||
)
|
||||
results = renderer.tokenize_prompts(
|
||||
prompts,
|
||||
TokenizeParams(max_total_tokens=100, truncate_prompt_tokens=5),
|
||||
)
|
||||
|
||||
assert len(results) == 1
|
||||
# Should keep the last 5 tokens: [105, 106, 107, 108, 109]
|
||||
assert results[0]["prompt_token_ids"] == [105, 106, 107, 108, 109]
|
||||
|
||||
def test_truncation_right(self):
|
||||
renderer = _build_renderer(MockModelConfig(), truncation_side="right")
|
||||
|
||||
long_tokens = [100, 101, 102, 103, 104, 105, 106, 107, 108, 109] # 10 tokens
|
||||
prompts = renderer.render_prompts(
|
||||
_preprocess_prompt(renderer.model_config, long_tokens)
|
||||
)
|
||||
results = renderer.tokenize_prompts(
|
||||
prompts,
|
||||
TokenizeParams(max_total_tokens=100, truncate_prompt_tokens=5),
|
||||
)
|
||||
|
||||
assert len(results) == 1
|
||||
# Should keep the first 5 tokens: [100, 101, 102, 103, 104]
|
||||
assert results[0]["prompt_token_ids"] == [100, 101, 102, 103, 104]
|
||||
|
||||
def test_text_max_length_exceeded_obvious(self):
|
||||
renderer = _build_renderer(MockModelConfig(), max_chars_per_token=1)
|
||||
|
||||
# Exceeds max_total_tokens and max_total_tokens * VLLM_MAX_CHARS_PER_TOKEN
|
||||
long_tokens = "x" * 150
|
||||
prompts = renderer.render_prompts(
|
||||
_preprocess_prompt(renderer.model_config, long_tokens)
|
||||
)
|
||||
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match="maximum context length is",
|
||||
):
|
||||
renderer.tokenize_prompts(
|
||||
prompts,
|
||||
TokenizeParams(max_total_tokens=100),
|
||||
)
|
||||
|
||||
# Should not even attempt tokenization
|
||||
assert renderer.tokenizer._captured_encode_kwargs == {}
|
||||
|
||||
def test_text_max_length_exceeded_nonobvious(self):
|
||||
renderer = _build_renderer(MockModelConfig(), max_chars_per_token=2)
|
||||
|
||||
# Exceeds max_total_tokens but not max_total_tokens * VLLM_MAX_CHARS_PER_TOKEN
|
||||
long_tokens = "x" * 150
|
||||
prompts = renderer.render_prompts(
|
||||
_preprocess_prompt(renderer.model_config, long_tokens)
|
||||
)
|
||||
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match="maximum context length is",
|
||||
):
|
||||
renderer.tokenize_prompts(
|
||||
prompts,
|
||||
TokenizeParams(max_total_tokens=100),
|
||||
)
|
||||
|
||||
# Should only tokenize the first max_total_tokens + 1 tokens
|
||||
assert renderer.tokenizer._captured_encode_kwargs["truncation"] is True
|
||||
assert renderer.tokenizer._captured_encode_kwargs["max_length"] == 101
|
||||
|
||||
def test_token_max_length_exceeded(self):
|
||||
renderer = _build_renderer(MockModelConfig())
|
||||
|
||||
long_tokens = list(range(150)) # Exceeds max_total_tokens=100
|
||||
prompts = renderer.render_prompts(
|
||||
_preprocess_prompt(renderer.model_config, long_tokens)
|
||||
)
|
||||
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match="maximum context length is",
|
||||
):
|
||||
renderer.tokenize_prompts(
|
||||
prompts,
|
||||
TokenizeParams(max_total_tokens=100, truncate_prompt_tokens=None),
|
||||
)
|
||||
|
||||
def test_no_tokenizer_for_text(self):
|
||||
renderer = _build_renderer(MockModelConfig(skip_tokenizer_init=True))
|
||||
|
||||
prompts = renderer.render_prompts(
|
||||
_preprocess_prompt(renderer.model_config, "Hello world")
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="`skip_tokenizer_init=True`"):
|
||||
renderer.tokenize_prompts(
|
||||
prompts,
|
||||
TokenizeParams(max_total_tokens=100),
|
||||
)
|
||||
|
||||
def test_token_input_with_needs_detokenization(self):
|
||||
renderer = _build_renderer(MockModelConfig())
|
||||
|
||||
tokens = [1, 2, 3, 4]
|
||||
prompts = renderer.render_prompts(
|
||||
_preprocess_prompt(renderer.model_config, tokens)
|
||||
)
|
||||
results = renderer.tokenize_prompts(
|
||||
prompts,
|
||||
TokenizeParams(
|
||||
max_total_tokens=100,
|
||||
needs_detokenization=True,
|
||||
),
|
||||
)
|
||||
|
||||
assert len(results) == 1
|
||||
assert results[0]["prompt_token_ids"] == tokens
|
||||
assert results[0]["prompt"] == "[1, 2, 3, 4]"
|
||||
|
||||
|
||||
class TestRenderEmbedPrompt:
|
||||
def _create_test_embed_bytes(self, tensor: torch.Tensor) -> bytes:
|
||||
"""Helper to create base64-encoded tensor bytes"""
|
||||
buffer = io.BytesIO()
|
||||
torch.save(tensor, buffer)
|
||||
buffer.seek(0)
|
||||
return pybase64.b64encode(buffer.read())
|
||||
|
||||
def test_single_prompt_embed(self):
|
||||
renderer = _build_renderer(MockModelConfig())
|
||||
|
||||
# Create a test tensor
|
||||
tensor_input = torch.randn(10, 768, dtype=torch.float32)
|
||||
embed_bytes = self._create_test_embed_bytes(tensor_input)
|
||||
|
||||
prompts = renderer.render_prompts(
|
||||
_preprocess_prompt(renderer.model_config, embed_bytes)
|
||||
)
|
||||
results = renderer.tokenize_prompts(
|
||||
prompts,
|
||||
TokenizeParams(max_total_tokens=100),
|
||||
)
|
||||
|
||||
assert len(results) == 1
|
||||
assert torch.equal(results[0]["prompt_embeds"], tensor_input)
|
||||
|
||||
def test_multiple_prompt_embeds(self):
|
||||
renderer = _build_renderer(MockModelConfig())
|
||||
|
||||
# Create multiple test tensors
|
||||
tensor_inputs = [
|
||||
torch.randn(8, 512, dtype=torch.float32),
|
||||
torch.randn(12, 512, dtype=torch.float32),
|
||||
]
|
||||
|
||||
prompts = renderer.render_prompts(
|
||||
_preprocess_prompt(
|
||||
renderer.model_config,
|
||||
[self._create_test_embed_bytes(t) for t in tensor_inputs],
|
||||
)
|
||||
)
|
||||
results = renderer.tokenize_prompts(
|
||||
prompts,
|
||||
TokenizeParams(max_total_tokens=100),
|
||||
)
|
||||
|
||||
assert len(results) == 2
|
||||
for i, result in enumerate(results):
|
||||
assert torch.allclose(result["prompt_embeds"], tensor_inputs[i])
|
||||
|
||||
def test_prompt_embed_truncation(self):
|
||||
renderer = _build_renderer(MockModelConfig())
|
||||
|
||||
# Create tensor with more tokens than truncation limit
|
||||
tensor_input = torch.randn(20, 768, dtype=torch.float32)
|
||||
|
||||
prompts = renderer.render_prompts(
|
||||
_preprocess_prompt(
|
||||
renderer.model_config, self._create_test_embed_bytes(tensor_input)
|
||||
)
|
||||
)
|
||||
results = renderer.tokenize_prompts(
|
||||
prompts,
|
||||
TokenizeParams(
|
||||
max_total_tokens=100,
|
||||
truncate_prompt_tokens=10,
|
||||
),
|
||||
)
|
||||
|
||||
assert len(results) == 1
|
||||
# Should keep last 10 tokens
|
||||
expected = tensor_input[-10:]
|
||||
assert torch.equal(results[0]["prompt_embeds"], expected)
|
||||
|
||||
def test_prompt_embed_different_dtypes(self):
|
||||
renderer = _build_renderer(MockModelConfig())
|
||||
|
||||
# Test different supported dtypes
|
||||
dtypes = [torch.float32, torch.float16, torch.bfloat16]
|
||||
|
||||
for dtype in dtypes:
|
||||
tensor_input = torch.randn(5, 256, dtype=dtype)
|
||||
|
||||
prompts = renderer.render_prompts(
|
||||
_preprocess_prompt(
|
||||
renderer.model_config, self._create_test_embed_bytes(tensor_input)
|
||||
)
|
||||
)
|
||||
results = renderer.tokenize_prompts(
|
||||
prompts,
|
||||
TokenizeParams(max_total_tokens=100),
|
||||
)
|
||||
|
||||
assert len(results) == 1
|
||||
assert results[0]["prompt_embeds"].dtype == dtype
|
||||
|
||||
def test_prompt_embed_squeeze_batch_dim(self):
|
||||
renderer = _build_renderer(MockModelConfig())
|
||||
|
||||
# Test tensor with batch dimension gets squeezed
|
||||
tensor_input = torch.randn(1, 10, 768, dtype=torch.float32)
|
||||
|
||||
prompts = renderer.render_prompts(
|
||||
_preprocess_prompt(
|
||||
renderer.model_config, self._create_test_embed_bytes(tensor_input)
|
||||
)
|
||||
)
|
||||
results = renderer.tokenize_prompts(
|
||||
prompts,
|
||||
TokenizeParams(max_total_tokens=100),
|
||||
)
|
||||
|
||||
assert len(results) == 1
|
||||
# Should be squeezed to 2D
|
||||
assert results[0]["prompt_embeds"].shape == (10, 768)
|
||||
|
||||
def test_both_prompts_and_embeds(self):
|
||||
renderer = _build_renderer(MockModelConfig())
|
||||
|
||||
text_input = "Hello world"
|
||||
tensor_input = torch.randn(5, 256, dtype=torch.float32)
|
||||
|
||||
prompts = renderer.render_prompts(
|
||||
_preprocess_prompt(
|
||||
renderer.model_config,
|
||||
[text_input, self._create_test_embed_bytes(tensor_input)],
|
||||
)
|
||||
)
|
||||
results = renderer.tokenize_prompts(
|
||||
prompts,
|
||||
TokenizeParams(max_total_tokens=100),
|
||||
)
|
||||
|
||||
assert len(results) == 2
|
||||
# First should be tokens prompt
|
||||
assert "prompt_token_ids" in results[0]
|
||||
assert len(results[0]["prompt_token_ids"]) == len(text_input)
|
||||
# Second should be embed prompt
|
||||
assert torch.equal(results[1]["prompt_embeds"], tensor_input)
|
||||
593
third_party/vllm/tests/renderers/test_hf.py
vendored
Normal file
593
third_party/vllm/tests/renderers/test_hf.py
vendored
Normal file
@@ -0,0 +1,593 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm.config import ModelConfig
|
||||
from vllm.entrypoints.chat_utils import load_chat_template
|
||||
from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
|
||||
from vllm.renderers.hf import (
|
||||
_get_hf_base_chat_template_params,
|
||||
_try_extract_ast,
|
||||
resolve_chat_template,
|
||||
resolve_chat_template_content_format,
|
||||
resolve_chat_template_kwargs,
|
||||
safe_apply_chat_template,
|
||||
)
|
||||
from vllm.tokenizers import get_tokenizer
|
||||
|
||||
from ..models.registry import HF_EXAMPLE_MODELS
|
||||
from ..utils import VLLM_PATH
|
||||
|
||||
EXAMPLES_DIR = VLLM_PATH / "examples"
|
||||
|
||||
chatml_jinja_path = VLLM_PATH / "examples/template_chatml.jinja"
|
||||
assert chatml_jinja_path.exists()
|
||||
|
||||
# Define models, templates, and their corresponding expected outputs
|
||||
MODEL_TEMPLATE_GENERATION_OUTPUT = [
|
||||
(
|
||||
"facebook/opt-125m",
|
||||
chatml_jinja_path,
|
||||
True,
|
||||
False,
|
||||
"""<|im_start|>user
|
||||
Hello<|im_end|>
|
||||
<|im_start|>assistant
|
||||
Hi there!<|im_end|>
|
||||
<|im_start|>user
|
||||
What is the capital of<|im_end|>
|
||||
<|im_start|>assistant
|
||||
""",
|
||||
),
|
||||
(
|
||||
"facebook/opt-125m",
|
||||
chatml_jinja_path,
|
||||
False,
|
||||
False,
|
||||
"""<|im_start|>user
|
||||
Hello<|im_end|>
|
||||
<|im_start|>assistant
|
||||
Hi there!<|im_end|>
|
||||
<|im_start|>user
|
||||
What is the capital of""",
|
||||
),
|
||||
(
|
||||
"facebook/opt-125m",
|
||||
chatml_jinja_path,
|
||||
False,
|
||||
True,
|
||||
"""<|im_start|>user
|
||||
Hello<|im_end|>
|
||||
<|im_start|>assistant
|
||||
Hi there!<|im_end|>
|
||||
<|im_start|>user
|
||||
What is the capital of<|im_end|>
|
||||
<|im_start|>assistant
|
||||
The capital of""",
|
||||
),
|
||||
]
|
||||
|
||||
TEST_MESSAGES = [
|
||||
{"role": "user", "content": "Hello"},
|
||||
{"role": "assistant", "content": "Hi there!"},
|
||||
{"role": "user", "content": "What is the capital of"},
|
||||
]
|
||||
ASSISTANT_MESSAGE_TO_CONTINUE = {"role": "assistant", "content": "The capital of"}
|
||||
|
||||
|
||||
def test_load_chat_template():
|
||||
# Testing chatml template
|
||||
template_content = load_chat_template(chat_template=chatml_jinja_path)
|
||||
|
||||
# Test assertions
|
||||
assert template_content is not None
|
||||
# Hard coded value for template_chatml.jinja
|
||||
assert (
|
||||
template_content
|
||||
== """{% for message in messages %}{{'<|im_start|>' + message['role'] + '\\n' + message['content']}}{% if (loop.last and add_generation_prompt) or not loop.last %}{{ '<|im_end|>' + '\\n'}}{% endif %}{% endfor %}
|
||||
{% if add_generation_prompt and messages[-1]['role'] != 'assistant' %}{{ '<|im_start|>assistant\\n' }}{% endif %}""" # noqa: E501
|
||||
)
|
||||
|
||||
|
||||
def test_no_load_chat_template_filelike():
|
||||
# Testing chatml template
|
||||
template = "../../examples/does_not_exist"
|
||||
|
||||
with pytest.raises(ValueError, match="looks like a file path"):
|
||||
load_chat_template(chat_template=template)
|
||||
|
||||
|
||||
def test_no_load_chat_template_literallike():
|
||||
# Testing chatml template
|
||||
template = "{{ messages }}"
|
||||
|
||||
template_content = load_chat_template(chat_template=template)
|
||||
|
||||
assert template_content == template
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model",
|
||||
[
|
||||
"Qwen/Qwen2-VL-2B-Instruct", # chat_template is of type str
|
||||
"NousResearch/Hermes-3-Llama-3.1-8B", # chat_template is of type dict
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("use_tools", [True, False])
|
||||
def test_resolve_chat_template(sample_json_schema, model, use_tools):
|
||||
"""checks that chat_template is a dict type for HF models."""
|
||||
model_info = HF_EXAMPLE_MODELS.find_hf_info(model)
|
||||
model_info.check_available_online(on_fail="skip")
|
||||
|
||||
model_config = ModelConfig(
|
||||
model,
|
||||
tokenizer=model_info.tokenizer or model,
|
||||
tokenizer_mode=model_info.tokenizer_mode,
|
||||
revision=model_info.revision,
|
||||
trust_remote_code=model_info.trust_remote_code,
|
||||
hf_overrides=model_info.hf_overrides,
|
||||
skip_tokenizer_init=model_info.require_embed_inputs,
|
||||
enable_prompt_embeds=model_info.require_embed_inputs,
|
||||
enable_mm_embeds=model_info.require_embed_inputs,
|
||||
enforce_eager=model_info.enforce_eager,
|
||||
dtype=model_info.dtype,
|
||||
)
|
||||
|
||||
# Build the tokenizer
|
||||
tokenizer = get_tokenizer(
|
||||
model,
|
||||
trust_remote_code=model_config.trust_remote_code,
|
||||
)
|
||||
|
||||
tools = (
|
||||
[
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "dummy_function_name",
|
||||
"description": "This is a dummy function",
|
||||
"parameters": sample_json_schema,
|
||||
},
|
||||
}
|
||||
]
|
||||
if use_tools
|
||||
else None
|
||||
)
|
||||
|
||||
# Test detecting the tokenizer's chat_template
|
||||
chat_template = resolve_chat_template(
|
||||
tokenizer,
|
||||
chat_template=None,
|
||||
tools=tools,
|
||||
model_config=model_config,
|
||||
)
|
||||
assert isinstance(chat_template, str)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model, expected_kwargs",
|
||||
[
|
||||
(
|
||||
"Qwen/Qwen2-VL-2B-Instruct",
|
||||
{
|
||||
"add_vision_id",
|
||||
"add_generation_prompt",
|
||||
"continue_final_message",
|
||||
"tools",
|
||||
},
|
||||
),
|
||||
(
|
||||
"Qwen/Qwen3-8B",
|
||||
{
|
||||
"enable_thinking",
|
||||
"add_generation_prompt",
|
||||
"continue_final_message",
|
||||
"tools",
|
||||
},
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_resolve_chat_template_kwargs(sample_json_schema, model, expected_kwargs):
|
||||
"""checks that chat_template is a dict type for HF models."""
|
||||
model_info = HF_EXAMPLE_MODELS.find_hf_info(model)
|
||||
model_info.check_available_online(on_fail="skip")
|
||||
|
||||
tools = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "dummy_function_name",
|
||||
"description": "This is a dummy function",
|
||||
"parameters": sample_json_schema,
|
||||
},
|
||||
}
|
||||
]
|
||||
|
||||
chat_template_kwargs = {
|
||||
# both unused
|
||||
"unused_kwargs_1": 123,
|
||||
"unused_kwargs_2": "abc",
|
||||
# should not appear
|
||||
"chat_template": "{% Hello world! %}",
|
||||
"tokenize": True,
|
||||
# used by tokenizer
|
||||
"continue_final_message": True,
|
||||
"tools": tools,
|
||||
# both used by Qwen2-VL and Qwen3
|
||||
"add_generation_prompt": True,
|
||||
# only used by Qwen2-VL
|
||||
"add_vision_id": True,
|
||||
# only used by Qwen3
|
||||
"enable_thinking": True,
|
||||
}
|
||||
|
||||
model_config = ModelConfig(
|
||||
model,
|
||||
tokenizer=model_info.tokenizer or model,
|
||||
tokenizer_mode=model_info.tokenizer_mode,
|
||||
revision=model_info.revision,
|
||||
trust_remote_code=model_info.trust_remote_code,
|
||||
hf_overrides=model_info.hf_overrides,
|
||||
skip_tokenizer_init=model_info.require_embed_inputs,
|
||||
enable_prompt_embeds=model_info.require_embed_inputs,
|
||||
enable_mm_embeds=model_info.require_embed_inputs,
|
||||
enforce_eager=model_info.enforce_eager,
|
||||
dtype=model_info.dtype,
|
||||
)
|
||||
|
||||
# Build the tokenizer
|
||||
tokenizer = get_tokenizer(
|
||||
model,
|
||||
trust_remote_code=model_config.trust_remote_code,
|
||||
)
|
||||
|
||||
# Test detecting the tokenizer's chat_template
|
||||
chat_template = resolve_chat_template(
|
||||
tokenizer,
|
||||
chat_template=None,
|
||||
tools=tools,
|
||||
model_config=model_config,
|
||||
)
|
||||
with pytest.raises(
|
||||
ValueError, match="Found unexpected chat template kwargs from request"
|
||||
):
|
||||
# should raise error if `chat_template_kwargs` contains
|
||||
# `chat_template` or `tokenize`
|
||||
resolve_chat_template_kwargs(
|
||||
tokenizer,
|
||||
chat_template=chat_template,
|
||||
chat_template_kwargs=chat_template_kwargs,
|
||||
)
|
||||
resolved_chat_template_kwargs = resolve_chat_template_kwargs(
|
||||
tokenizer,
|
||||
chat_template=chat_template,
|
||||
chat_template_kwargs=chat_template_kwargs,
|
||||
raise_on_unexpected=False,
|
||||
)
|
||||
assert set(resolved_chat_template_kwargs.keys()) == expected_kwargs
|
||||
|
||||
# Additional test: Verify HF base parameters work with **kwargs tokenizers
|
||||
# This validates the fix for tokenizers like Kimi K2 that use **kwargs
|
||||
# to receive standard HuggingFace parameters instead of declaring them explicitly
|
||||
hf_base_params = _get_hf_base_chat_template_params()
|
||||
# Verify common HF parameters are in the base class
|
||||
assert {"add_generation_prompt", "tools", "continue_final_message"}.issubset(
|
||||
hf_base_params
|
||||
), f"Expected HF base params not found in {hf_base_params}"
|
||||
|
||||
# Test with a mock tokenizer that uses **kwargs (like Kimi K2)
|
||||
class MockTokenizerWithKwargs:
|
||||
def apply_chat_template(self, conversation, **kwargs):
|
||||
return "mocked_output"
|
||||
|
||||
mock_tokenizer = MockTokenizerWithKwargs()
|
||||
mock_kwargs = {
|
||||
"add_generation_prompt": True,
|
||||
"tools": tools,
|
||||
"continue_final_message": False,
|
||||
"unknown_param": "should_be_filtered",
|
||||
}
|
||||
resolved_mock = resolve_chat_template_kwargs(
|
||||
mock_tokenizer, chat_template, mock_kwargs, raise_on_unexpected=False
|
||||
)
|
||||
# HF base params should pass through even with **kwargs tokenizer
|
||||
assert "add_generation_prompt" in resolved_mock
|
||||
assert "tools" in resolved_mock
|
||||
assert "continue_final_message" in resolved_mock
|
||||
# Unknown params should be filtered out
|
||||
assert "unknown_param" not in resolved_mock
|
||||
|
||||
|
||||
def test_resolve_chat_template_resolves_name():
|
||||
"""When chat_template is a name, resolve_chat_template should return
|
||||
the actual Jinja content so that kwargs detection works correctly."""
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
jinja_content = "{{ messages }}{% if tools %}{{ tools }}{% endif %}"
|
||||
tokenizer = MagicMock()
|
||||
tokenizer.get_chat_template.return_value = jinja_content
|
||||
|
||||
model_config = MagicMock()
|
||||
|
||||
result = resolve_chat_template(
|
||||
tokenizer,
|
||||
chat_template="tool_use",
|
||||
tools=None,
|
||||
model_config=model_config,
|
||||
)
|
||||
|
||||
assert result == jinja_content
|
||||
tokenizer.get_chat_template.assert_called_once_with("tool_use", tools=None)
|
||||
|
||||
|
||||
def test_resolve_chat_template_kwargs_with_template_name():
|
||||
"""Ensures template kwargs are not silently dropped when chat_template
|
||||
was originally a template name that has been resolved to Jinja content."""
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
jinja_content = (
|
||||
"{% for m in messages %}{{ m }}{% endfor %}"
|
||||
"{% if tools %}{{ tools }}{% endif %}"
|
||||
"{% if documents %}{{ documents }}{% endif %}"
|
||||
)
|
||||
|
||||
tokenizer = MagicMock()
|
||||
tokenizer.apply_chat_template = MagicMock()
|
||||
|
||||
kwargs = {
|
||||
"tools": [{"type": "function", "function": {"name": "f"}}],
|
||||
"documents": [{"title": "doc"}],
|
||||
"unknown_param": "should be dropped",
|
||||
}
|
||||
|
||||
resolved = resolve_chat_template_kwargs(
|
||||
tokenizer,
|
||||
chat_template=jinja_content,
|
||||
chat_template_kwargs=kwargs,
|
||||
raise_on_unexpected=False,
|
||||
)
|
||||
|
||||
# template vars "tools" and "documents" should be preserved
|
||||
assert "tools" in resolved
|
||||
assert "documents" in resolved
|
||||
# unknown param should be filtered
|
||||
assert "unknown_param" not in resolved
|
||||
|
||||
|
||||
# NOTE: Qwen2-Audio default chat template is specially defined inside
|
||||
# processor class instead of using `tokenizer_config.json`
|
||||
@pytest.mark.parametrize(
|
||||
("model", "expected_format"),
|
||||
[
|
||||
("microsoft/Phi-3.5-vision-instruct", "string"),
|
||||
("Qwen/Qwen2-VL-2B-Instruct", "openai"),
|
||||
("Qwen/Qwen2.5-VL-3B-Instruct", "openai"),
|
||||
("fixie-ai/ultravox-v0_5-llama-3_2-1b", "string"),
|
||||
("Qwen/Qwen2-Audio-7B-Instruct", "openai"),
|
||||
("meta-llama/Llama-Guard-3-1B", "openai"),
|
||||
],
|
||||
)
|
||||
def test_resolve_content_format_hf_defined(model, expected_format):
|
||||
model_info = HF_EXAMPLE_MODELS.find_hf_info(model)
|
||||
model_info.check_available_online(on_fail="skip")
|
||||
|
||||
model_config = ModelConfig(
|
||||
model,
|
||||
tokenizer=model_info.tokenizer or model,
|
||||
tokenizer_mode=model_info.tokenizer_mode,
|
||||
revision=model_info.revision,
|
||||
trust_remote_code=model_info.trust_remote_code,
|
||||
hf_overrides=model_info.hf_overrides,
|
||||
skip_tokenizer_init=model_info.require_embed_inputs,
|
||||
enable_prompt_embeds=model_info.require_embed_inputs,
|
||||
enable_mm_embeds=model_info.require_embed_inputs,
|
||||
enforce_eager=model_info.enforce_eager,
|
||||
dtype=model_info.dtype,
|
||||
)
|
||||
|
||||
tokenizer = get_tokenizer(
|
||||
model,
|
||||
trust_remote_code=model_config.trust_remote_code,
|
||||
)
|
||||
|
||||
# Test detecting the tokenizer's chat_template
|
||||
chat_template = resolve_chat_template(
|
||||
tokenizer,
|
||||
chat_template=None,
|
||||
tools=None,
|
||||
model_config=model_config,
|
||||
)
|
||||
assert isinstance(chat_template, str)
|
||||
|
||||
print("[TEXT]")
|
||||
print(chat_template)
|
||||
print("[AST]")
|
||||
print(_try_extract_ast(chat_template))
|
||||
|
||||
resolved_format = resolve_chat_template_content_format(
|
||||
None, # Test detecting the tokenizer's chat_template
|
||||
None,
|
||||
"auto",
|
||||
tokenizer,
|
||||
model_config=model_config,
|
||||
)
|
||||
|
||||
assert resolved_format == expected_format
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("model", "expected_format"),
|
||||
[
|
||||
("Salesforce/blip2-opt-2.7b", "string"),
|
||||
("facebook/chameleon-7b", "string"),
|
||||
("deepseek-ai/deepseek-vl2-tiny", "string"),
|
||||
("adept/fuyu-8b", "string"),
|
||||
("google/paligemma-3b-mix-224", "string"),
|
||||
("Qwen/Qwen-VL", "string"),
|
||||
("Qwen/Qwen-VL-Chat", "string"),
|
||||
],
|
||||
)
|
||||
def test_resolve_content_format_fallbacks(model, expected_format):
|
||||
model_info = HF_EXAMPLE_MODELS.find_hf_info(model)
|
||||
model_info.check_available_online(on_fail="skip")
|
||||
|
||||
model_config = ModelConfig(
|
||||
model,
|
||||
tokenizer=model_info.tokenizer or model,
|
||||
tokenizer_mode=model_info.tokenizer_mode,
|
||||
revision=model_info.revision,
|
||||
trust_remote_code=model_info.trust_remote_code,
|
||||
hf_overrides=model_info.hf_overrides,
|
||||
skip_tokenizer_init=model_info.require_embed_inputs,
|
||||
enable_prompt_embeds=model_info.require_embed_inputs,
|
||||
enable_mm_embeds=model_info.require_embed_inputs,
|
||||
enforce_eager=model_info.enforce_eager,
|
||||
dtype=model_info.dtype,
|
||||
)
|
||||
|
||||
tokenizer = get_tokenizer(
|
||||
model_config.tokenizer,
|
||||
trust_remote_code=model_config.trust_remote_code,
|
||||
)
|
||||
|
||||
# Test detecting the tokenizer's chat_template
|
||||
chat_template = resolve_chat_template(
|
||||
tokenizer,
|
||||
chat_template=None,
|
||||
tools=None,
|
||||
model_config=model_config,
|
||||
)
|
||||
assert isinstance(chat_template, str)
|
||||
|
||||
print("[TEXT]")
|
||||
print(chat_template)
|
||||
print("[AST]")
|
||||
print(_try_extract_ast(chat_template))
|
||||
|
||||
resolved_format = resolve_chat_template_content_format(
|
||||
None, # Test detecting the tokenizer's chat_template
|
||||
None,
|
||||
"auto",
|
||||
tokenizer,
|
||||
model_config=model_config,
|
||||
)
|
||||
|
||||
assert resolved_format == expected_format
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("template_path", "expected_format"),
|
||||
[
|
||||
("template_alpaca.jinja", "string"),
|
||||
("template_baichuan.jinja", "string"),
|
||||
("template_chatglm.jinja", "string"),
|
||||
("template_chatglm2.jinja", "string"),
|
||||
("template_chatml.jinja", "string"),
|
||||
("template_falcon_180b.jinja", "string"),
|
||||
("template_falcon.jinja", "string"),
|
||||
("template_inkbot.jinja", "string"),
|
||||
("template_teleflm.jinja", "string"),
|
||||
("pooling/embed/template/dse_qwen2_vl.jinja", "openai"),
|
||||
("pooling/embed/template/vlm2vec_phi3v.jinja", "openai"),
|
||||
("pooling/embed/template/vlm2vec_qwen2vl.jinja", "openai"),
|
||||
("tool_chat_template_granite_20b_fc.jinja", "string"),
|
||||
("tool_chat_template_hermes.jinja", "string"),
|
||||
("tool_chat_template_internlm2_tool.jinja", "string"),
|
||||
("tool_chat_template_llama3.1_json.jinja", "openai"),
|
||||
("tool_chat_template_llama3.2_json.jinja", "openai"),
|
||||
("tool_chat_template_mistral_parallel.jinja", "string"),
|
||||
("tool_chat_template_mistral.jinja", "string"),
|
||||
],
|
||||
)
|
||||
def test_resolve_content_format_examples(template_path, expected_format):
|
||||
model = "Qwen/Qwen2-VL-2B-Instruct" # Dummy
|
||||
model_config = ModelConfig(
|
||||
model,
|
||||
tokenizer=model,
|
||||
trust_remote_code=True,
|
||||
)
|
||||
|
||||
dummy_tokenizer = get_tokenizer(
|
||||
model,
|
||||
trust_remote_code=model_config.trust_remote_code,
|
||||
)
|
||||
dummy_tokenizer.chat_template = None
|
||||
|
||||
chat_template = load_chat_template(EXAMPLES_DIR / template_path)
|
||||
assert isinstance(chat_template, str)
|
||||
|
||||
print("[TEXT]")
|
||||
print(chat_template)
|
||||
print("[AST]")
|
||||
print(_try_extract_ast(chat_template))
|
||||
|
||||
resolved_format = resolve_chat_template_content_format(
|
||||
chat_template,
|
||||
None,
|
||||
"auto",
|
||||
dummy_tokenizer,
|
||||
model_config=model_config,
|
||||
)
|
||||
|
||||
assert resolved_format == expected_format
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model,template,add_generation_prompt,continue_final_message,expected_output",
|
||||
MODEL_TEMPLATE_GENERATION_OUTPUT,
|
||||
)
|
||||
def test_get_gen_prompt(
|
||||
model, template, add_generation_prompt, continue_final_message, expected_output
|
||||
):
|
||||
model_info = HF_EXAMPLE_MODELS.find_hf_info(model)
|
||||
model_info.check_available_online(on_fail="skip")
|
||||
|
||||
model_config = ModelConfig(
|
||||
model,
|
||||
tokenizer=model_info.tokenizer or model,
|
||||
tokenizer_mode=model_info.tokenizer_mode,
|
||||
trust_remote_code=model_info.trust_remote_code,
|
||||
revision=model_info.revision,
|
||||
hf_overrides=model_info.hf_overrides,
|
||||
skip_tokenizer_init=model_info.require_embed_inputs,
|
||||
enable_prompt_embeds=model_info.require_embed_inputs,
|
||||
enable_mm_embeds=model_info.require_embed_inputs,
|
||||
enforce_eager=model_info.enforce_eager,
|
||||
dtype=model_info.dtype,
|
||||
)
|
||||
|
||||
# Initialize the tokenizer
|
||||
tokenizer = get_tokenizer(
|
||||
tokenizer_name=model_config.tokenizer,
|
||||
trust_remote_code=model_config.trust_remote_code,
|
||||
)
|
||||
template_content = load_chat_template(chat_template=template)
|
||||
|
||||
# Create a mock request object using keyword arguments
|
||||
mock_request = ChatCompletionRequest(
|
||||
model=model,
|
||||
messages=TEST_MESSAGES + [ASSISTANT_MESSAGE_TO_CONTINUE]
|
||||
if continue_final_message
|
||||
else TEST_MESSAGES,
|
||||
add_generation_prompt=add_generation_prompt,
|
||||
continue_final_message=continue_final_message,
|
||||
)
|
||||
|
||||
# Call the function and get the result
|
||||
result = safe_apply_chat_template(
|
||||
model_config,
|
||||
tokenizer,
|
||||
mock_request.messages,
|
||||
tools=None,
|
||||
chat_template=mock_request.chat_template or template_content,
|
||||
add_generation_prompt=mock_request.add_generation_prompt,
|
||||
continue_final_message=mock_request.continue_final_message,
|
||||
tokenize=False,
|
||||
)
|
||||
|
||||
# Test assertion
|
||||
assert result == expected_output, (
|
||||
f"The generated prompt does not match the expected output for "
|
||||
f"model {model} and template {template}"
|
||||
)
|
||||
140
third_party/vllm/tests/renderers/test_mistral.py
vendored
Normal file
140
third_party/vllm/tests/renderers/test_mistral.py
vendored
Normal file
@@ -0,0 +1,140 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import asyncio
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
from mistral_common.tokens.tokenizers.base import SpecialTokenPolicy
|
||||
|
||||
from vllm.renderers import ChatParams
|
||||
from vllm.renderers.mistral import MistralRenderer, safe_apply_chat_template
|
||||
from vllm.tokenizers.mistral import MistralTokenizer
|
||||
|
||||
MODEL_NAME = "mistralai/Mistral-7B-Instruct-v0.3"
|
||||
|
||||
|
||||
@dataclass
|
||||
class MockHFConfig:
|
||||
model_type: str = "any"
|
||||
|
||||
|
||||
@dataclass
|
||||
class MockModelConfig:
|
||||
runner_type = "generate"
|
||||
model: str = MODEL_NAME
|
||||
tokenizer: str = MODEL_NAME
|
||||
trust_remote_code: bool = False
|
||||
max_model_len: int = 100
|
||||
tokenizer_revision = None
|
||||
tokenizer_mode = "mistral"
|
||||
hf_config = MockHFConfig()
|
||||
encoder_config: dict[str, Any] | None = None
|
||||
enable_prompt_embeds: bool = True
|
||||
skip_tokenizer_init: bool = False
|
||||
is_encoder_decoder: bool = False
|
||||
is_multimodal_model: bool = False
|
||||
|
||||
|
||||
@dataclass
|
||||
class MockParallelConfig:
|
||||
_api_process_rank: int = 0
|
||||
|
||||
|
||||
@dataclass
|
||||
class MockVllmConfig:
|
||||
model_config: MockModelConfig
|
||||
parallel_config: MockParallelConfig
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_mistral_tokenizer_does_not_block_event_loop():
|
||||
expected_tokens = [1, 2, 3]
|
||||
|
||||
# Mock the blocking version to sleep
|
||||
def mocked_apply_chat_template(*_args, **_kwargs):
|
||||
time.sleep(2)
|
||||
return expected_tokens
|
||||
|
||||
mock_model_config = MockModelConfig(skip_tokenizer_init=True)
|
||||
mock_tokenizer = Mock(spec=MistralTokenizer)
|
||||
mock_tokenizer.apply_chat_template = mocked_apply_chat_template
|
||||
mock_renderer = MistralRenderer(
|
||||
MockVllmConfig(mock_model_config, parallel_config=MockParallelConfig()),
|
||||
tokenizer=mock_tokenizer,
|
||||
)
|
||||
|
||||
task = mock_renderer.render_messages_async([], ChatParams())
|
||||
|
||||
# Ensure the event loop is not blocked
|
||||
blocked_count = 0
|
||||
for _i in range(20): # Check over ~2 seconds
|
||||
start = time.perf_counter()
|
||||
await asyncio.sleep(0)
|
||||
elapsed = time.perf_counter() - start
|
||||
|
||||
# an overly generous elapsed time for slow machines
|
||||
if elapsed >= 0.5:
|
||||
blocked_count += 1
|
||||
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
# Ensure task completes
|
||||
_, prompt = await task
|
||||
assert prompt["prompt_token_ids"] == expected_tokens, (
|
||||
"Mocked blocking tokenizer was not called"
|
||||
)
|
||||
assert blocked_count == 0, "Event loop blocked during tokenization"
|
||||
|
||||
|
||||
def test_apply_mistral_chat_template_thinking_chunk():
|
||||
messages = [
|
||||
{
|
||||
"role": "system",
|
||||
"content": [
|
||||
{"type": "text", "text": "You are a helpful assistant."},
|
||||
{
|
||||
"type": "thinking",
|
||||
"closed": True,
|
||||
"thinking": "Only return the answer when you are confident.",
|
||||
},
|
||||
],
|
||||
},
|
||||
{"role": "user", "content": "What is 2+2?"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{"type": "text", "text": "Let me think about it."},
|
||||
{"type": "thinking", "closed": True, "thinking": "2+2 = 4"},
|
||||
{
|
||||
"type": "text",
|
||||
"text": "The answer is 4.",
|
||||
},
|
||||
],
|
||||
},
|
||||
{"role": "user", "content": "Thanks, what is 3+3?"},
|
||||
]
|
||||
mistral_tokenizer = MistralTokenizer.from_pretrained(
|
||||
"mistralai/Magistral-Small-2509"
|
||||
)
|
||||
|
||||
tokens_ids = safe_apply_chat_template(
|
||||
mistral_tokenizer, messages, chat_template=None, tools=None
|
||||
)
|
||||
|
||||
string_tokens = mistral_tokenizer.mistral.decode(
|
||||
tokens_ids, special_token_policy=SpecialTokenPolicy.KEEP
|
||||
)
|
||||
|
||||
expected_tokens = (
|
||||
r"<s>[SYSTEM_PROMPT]You are a helpful assistant.[THINK]Only return the"
|
||||
r" answer when you are confident.[/THINK][/SYSTEM_PROMPT]"
|
||||
r"[INST]What is 2+2?[/INST]"
|
||||
r"Let me think about it.[THINK]2+2 = 4[/THINK]The answer is 4.</s>"
|
||||
r"[INST]Thanks, what is 3+3?[/INST]"
|
||||
)
|
||||
|
||||
assert string_tokens == expected_tokens
|
||||
184
third_party/vllm/tests/renderers/test_process_multi_modal_uuids.py
vendored
Normal file
184
third_party/vllm/tests/renderers/test_process_multi_modal_uuids.py
vendored
Normal file
@@ -0,0 +1,184 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm.assets.image import ImageAsset
|
||||
from vllm.assets.video import VideoAsset
|
||||
from vllm.config import CacheConfig, ModelConfig, VllmConfig
|
||||
from vllm.multimodal.parse import parse_mm_uuids
|
||||
from vllm.renderers.hf import HfRenderer
|
||||
from vllm.tokenizers.registry import tokenizer_args_from_config
|
||||
|
||||
cherry_pil_image = ImageAsset("cherry_blossom").pil_image
|
||||
stop_pil_image = ImageAsset("stop_sign").pil_image
|
||||
baby_reading_np_ndarrays = VideoAsset("baby_reading").np_ndarrays
|
||||
|
||||
|
||||
def _build_renderer(
|
||||
*, mm_cache_gb: float = 4.0, enable_prefix_caching: bool = True
|
||||
) -> HfRenderer:
|
||||
model_config = ModelConfig(
|
||||
model="Qwen/Qwen2.5-VL-3B-Instruct",
|
||||
max_model_len=128,
|
||||
mm_processor_cache_gb=mm_cache_gb,
|
||||
)
|
||||
|
||||
vllm_config = VllmConfig(
|
||||
model_config=model_config,
|
||||
cache_config=CacheConfig(enable_prefix_caching=enable_prefix_caching),
|
||||
)
|
||||
|
||||
_, tokenizer_name, _, kwargs = tokenizer_args_from_config(model_config)
|
||||
|
||||
return HfRenderer.from_config(
|
||||
vllm_config,
|
||||
tokenizer_kwargs={**kwargs, "tokenizer_name": tokenizer_name},
|
||||
)
|
||||
|
||||
|
||||
def test_multi_modal_uuids_length_mismatch_raises():
|
||||
renderer = _build_renderer()
|
||||
|
||||
mm_data = {"image": [cherry_pil_image, stop_pil_image]}
|
||||
|
||||
# Mismatch: 2 items but only 0 uuids provided
|
||||
mm_uuids = {"image": []} # type: ignore[var-annotated]
|
||||
|
||||
mm_processor = renderer.get_mm_processor()
|
||||
mm_data_items = mm_processor.info.parse_mm_data(mm_data)
|
||||
mm_uuid_items = parse_mm_uuids(mm_uuids)
|
||||
|
||||
with pytest.raises(ValueError, match="must have same length as"):
|
||||
renderer._process_mm_uuids(mm_data, mm_data_items, mm_uuid_items, "req-1a")
|
||||
|
||||
# Mismatch: 2 items but only 1 uuid provided
|
||||
mm_uuids = {"image": ["hash_cherry"]}
|
||||
|
||||
mm_processor = renderer.get_mm_processor()
|
||||
mm_data_items = mm_processor.info.parse_mm_data(mm_data)
|
||||
mm_uuid_items = parse_mm_uuids(mm_uuids)
|
||||
|
||||
with pytest.raises(ValueError, match="must have same length as"):
|
||||
renderer._process_mm_uuids(mm_data, mm_data_items, mm_uuid_items, "req-1b")
|
||||
|
||||
|
||||
def test_multi_modal_uuids_missing_modality_raises():
|
||||
renderer = _build_renderer()
|
||||
|
||||
mm_data = {
|
||||
"image": [cherry_pil_image],
|
||||
"video": None,
|
||||
}
|
||||
|
||||
# Only image uuids provided; video missing should raise
|
||||
mm_uuids = {"image": ["hash_cherry"]}
|
||||
|
||||
mm_processor = renderer.get_mm_processor()
|
||||
mm_data_items = mm_processor.info.parse_mm_data(mm_data)
|
||||
mm_uuid_items = parse_mm_uuids(mm_uuids)
|
||||
|
||||
with pytest.raises(ValueError, match="is empty but .* is missing"):
|
||||
renderer._process_mm_uuids(mm_data, mm_data_items, mm_uuid_items, "req-2")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"mm_cache_gb, enable_prefix_caching",
|
||||
[
|
||||
(4.0, True), # default behavior
|
||||
(4.0, False), # prefix caching disabled
|
||||
(0.0, True), # processor cache disabled
|
||||
],
|
||||
)
|
||||
def test_multi_modal_uuids_accepts_none_and_passes_through(
|
||||
mm_cache_gb: float, enable_prefix_caching: bool
|
||||
):
|
||||
renderer = _build_renderer(
|
||||
mm_cache_gb=mm_cache_gb,
|
||||
enable_prefix_caching=enable_prefix_caching,
|
||||
)
|
||||
|
||||
mm_data = {
|
||||
"image": [cherry_pil_image, stop_pil_image],
|
||||
"video": baby_reading_np_ndarrays,
|
||||
}
|
||||
|
||||
# Use a consistent two-image scenario across all configurations
|
||||
mm_uuids = {"image": [None, "hash_stop"], "video": None}
|
||||
|
||||
mm_processor = renderer.get_mm_processor()
|
||||
mm_data_items = mm_processor.info.parse_mm_data(mm_data)
|
||||
mm_uuid_items = parse_mm_uuids(mm_uuids)
|
||||
|
||||
processed_mm_uuids = renderer._process_mm_uuids(
|
||||
mm_data, mm_data_items, mm_uuid_items, "req-3"
|
||||
)
|
||||
|
||||
assert processed_mm_uuids == mm_uuids
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"mm_cache_gb, enable_prefix_caching",
|
||||
[
|
||||
(4.0, True), # default behavior
|
||||
(4.0, False), # prefix caching disabled
|
||||
(0.0, True), # processor cache disabled
|
||||
],
|
||||
)
|
||||
def test_multi_modal_uuids_accepts_empty(
|
||||
mm_cache_gb: float, enable_prefix_caching: bool
|
||||
):
|
||||
renderer = _build_renderer(
|
||||
mm_cache_gb=mm_cache_gb,
|
||||
enable_prefix_caching=enable_prefix_caching,
|
||||
)
|
||||
|
||||
# While None means cached multi-modal input requiring UUIDs
|
||||
# an empty list means no multi-modal input
|
||||
mm_data = {"image": [], "video": [], "audio": None} # type: ignore[var-annotated]
|
||||
mm_uuids = {"image": [], "video": None, "audio": []} # type: ignore[var-annotated]
|
||||
|
||||
mm_processor = renderer.get_mm_processor()
|
||||
mm_data_items = mm_processor.info.parse_mm_data(mm_data)
|
||||
mm_uuid_items = parse_mm_uuids(mm_uuids)
|
||||
|
||||
processed_mm_uuids = renderer._process_mm_uuids(
|
||||
mm_data, mm_data_items, mm_uuid_items, "req-4"
|
||||
)
|
||||
|
||||
assert processed_mm_uuids == mm_uuids
|
||||
|
||||
|
||||
def test_multi_modal_uuids_ignored_when_caching_disabled():
|
||||
# When both processor cache is 0 and prefix caching disabled, the
|
||||
# processor builds overrides from request id instead of using user UUIDs.
|
||||
renderer = _build_renderer(mm_cache_gb=0.0, enable_prefix_caching=False)
|
||||
|
||||
request_id = "req-42"
|
||||
mm_data = {
|
||||
"image": [cherry_pil_image, stop_pil_image],
|
||||
"video": baby_reading_np_ndarrays,
|
||||
}
|
||||
mm_uuids = {"image": ["hash_cherry", "hash_stop"], "video": ["hash_video"]}
|
||||
|
||||
mm_processor = renderer.get_mm_processor()
|
||||
mm_data_items = mm_processor.info.parse_mm_data(mm_data)
|
||||
mm_uuid_items = parse_mm_uuids(mm_uuids)
|
||||
|
||||
processed_mm_uuids = renderer._process_mm_uuids(
|
||||
mm_data, mm_data_items, mm_uuid_items, request_id
|
||||
)
|
||||
|
||||
# Expect request-id-based overrides are passed through
|
||||
assert set(mm_uuids.keys()) == {"image", "video"}
|
||||
assert len(mm_uuids["image"]) == 2
|
||||
assert len(mm_uuids["video"]) == 1
|
||||
assert processed_mm_uuids["image"][0].startswith(
|
||||
f"{request_id}-image-"
|
||||
) and processed_mm_uuids["image"][0].endswith("-0")
|
||||
assert processed_mm_uuids["image"][1].startswith(
|
||||
f"{request_id}-image-"
|
||||
) and processed_mm_uuids["image"][1].endswith("-1")
|
||||
assert processed_mm_uuids["video"][0].startswith(
|
||||
f"{request_id}-video-"
|
||||
) and processed_mm_uuids["video"][0].endswith("-0")
|
||||
290
third_party/vllm/tests/renderers/test_sparse_tensor_validation.py
vendored
Normal file
290
third_party/vllm/tests/renderers/test_sparse_tensor_validation.py
vendored
Normal file
@@ -0,0 +1,290 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Tests verify that malicious sparse tensors are rejected before they can trigger
|
||||
out-of-bounds memory writes during to_dense() operations.
|
||||
"""
|
||||
|
||||
import base64
|
||||
import io
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.multimodal.media import AudioEmbeddingMediaIO, ImageEmbeddingMediaIO
|
||||
from vllm.renderers.embed_utils import safe_load_prompt_embeds
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def model_config():
|
||||
"""Mock ModelConfig for testing."""
|
||||
from vllm.config import ModelConfig
|
||||
|
||||
return ModelConfig(
|
||||
model="facebook/opt-125m",
|
||||
tokenizer="facebook/opt-125m",
|
||||
tokenizer_mode="auto",
|
||||
trust_remote_code=False,
|
||||
dtype="float32",
|
||||
seed=0,
|
||||
enable_prompt_embeds=True, # Required for prompt embeds tests
|
||||
)
|
||||
|
||||
|
||||
def _encode_tensor(tensor: torch.Tensor) -> bytes:
|
||||
"""Helper to encode a tensor as base64 bytes."""
|
||||
buffer = io.BytesIO()
|
||||
torch.save(tensor, buffer)
|
||||
buffer.seek(0)
|
||||
return base64.b64encode(buffer.read())
|
||||
|
||||
|
||||
def _create_malicious_sparse_tensor() -> torch.Tensor:
|
||||
"""
|
||||
Create a malicious sparse COO tensor with out-of-bounds indices.
|
||||
|
||||
This tensor has indices that point beyond the declared shape, which would
|
||||
cause an out-of-bounds write when converted to dense format without
|
||||
validation.
|
||||
"""
|
||||
# Create a 3x3 sparse tensor but with indices pointing to (10, 10)
|
||||
indices = torch.tensor([[10], [10]]) # Out of bounds for 3x3 shape
|
||||
values = torch.tensor([1.0])
|
||||
shape = (3, 3)
|
||||
|
||||
# Create sparse tensor (this will be invalid)
|
||||
sparse_tensor = torch.sparse_coo_tensor(indices, values, shape, dtype=torch.float32)
|
||||
return sparse_tensor
|
||||
|
||||
|
||||
def _create_valid_sparse_tensor() -> torch.Tensor:
|
||||
"""Create a valid sparse COO tensor for baseline testing."""
|
||||
indices = torch.tensor([[0, 1, 2], [0, 1, 2]])
|
||||
values = torch.tensor([1.0, 2.0, 3.0])
|
||||
shape = (3, 3)
|
||||
|
||||
sparse_tensor = torch.sparse_coo_tensor(indices, values, shape, dtype=torch.float32)
|
||||
return sparse_tensor
|
||||
|
||||
|
||||
def _create_valid_dense_tensor() -> torch.Tensor:
|
||||
"""Create a valid dense tensor for baseline testing."""
|
||||
return torch.randn(10, 768, dtype=torch.float32) # (seq_len, hidden_size)
|
||||
|
||||
|
||||
class TestPromptEmbedsValidation:
|
||||
"""Test sparse tensor validation in prompt embeddings (Completions API)."""
|
||||
|
||||
def test_valid_dense_tensor_accepted(self, model_config):
|
||||
"""Baseline: Valid dense tensors should work normally."""
|
||||
valid_tensor = _create_valid_dense_tensor()
|
||||
encoded = _encode_tensor(valid_tensor)
|
||||
|
||||
# Should not raise any exception
|
||||
result = safe_load_prompt_embeds(model_config, encoded)
|
||||
assert result.shape == valid_tensor.shape
|
||||
|
||||
def test_valid_sparse_tensor_accepted(self):
|
||||
"""Baseline: Valid sparse tensors should load successfully."""
|
||||
io_handler = ImageEmbeddingMediaIO()
|
||||
|
||||
valid_sparse = _create_valid_sparse_tensor()
|
||||
encoded = _encode_tensor(valid_sparse)
|
||||
|
||||
# Should not raise any exception (sparse tensors remain sparse)
|
||||
result = io_handler.load_base64("", encoded.decode("utf-8"))
|
||||
assert result.shape == valid_sparse.shape
|
||||
|
||||
def test_malicious_sparse_tensor_rejected(self, model_config):
|
||||
"""Security: Malicious sparse tensors should be rejected."""
|
||||
malicious_tensor = _create_malicious_sparse_tensor()
|
||||
encoded = _encode_tensor(malicious_tensor)
|
||||
|
||||
# Should raise RuntimeError due to invalid sparse tensor
|
||||
with pytest.raises((RuntimeError, ValueError)) as exc_info:
|
||||
safe_load_prompt_embeds(model_config, encoded)
|
||||
|
||||
# Error should indicate sparse tensor validation failure
|
||||
error_msg = str(exc_info.value).lower()
|
||||
assert "sparse" in error_msg or "index" in error_msg or "bounds" in error_msg
|
||||
|
||||
def test_extremely_large_indices_rejected(self, model_config):
|
||||
"""Security: Sparse tensors with extremely large indices should be rejected."""
|
||||
# Create tensor with indices far beyond reasonable bounds
|
||||
indices = torch.tensor([[999999], [999999]])
|
||||
values = torch.tensor([1.0])
|
||||
shape = (10, 10)
|
||||
|
||||
malicious_tensor = torch.sparse_coo_tensor(
|
||||
indices, values, shape, dtype=torch.float32
|
||||
)
|
||||
encoded = _encode_tensor(malicious_tensor)
|
||||
|
||||
with pytest.raises((RuntimeError, ValueError)):
|
||||
safe_load_prompt_embeds(model_config, encoded)
|
||||
|
||||
def test_negative_indices_rejected(self, model_config):
|
||||
"""Security: Sparse tensors with negative indices should be rejected."""
|
||||
# Create tensor with negative indices
|
||||
indices = torch.tensor([[-1], [-1]])
|
||||
values = torch.tensor([1.0])
|
||||
shape = (10, 10)
|
||||
|
||||
malicious_tensor = torch.sparse_coo_tensor(
|
||||
indices, values, shape, dtype=torch.float32
|
||||
)
|
||||
encoded = _encode_tensor(malicious_tensor)
|
||||
|
||||
with pytest.raises((RuntimeError, ValueError)):
|
||||
safe_load_prompt_embeds(model_config, encoded)
|
||||
|
||||
|
||||
class TestImageEmbedsValidation:
|
||||
"""Test sparse tensor validation in image embeddings (Chat API)."""
|
||||
|
||||
def test_valid_dense_tensor_accepted(self):
|
||||
"""Baseline: Valid dense tensors should work normally."""
|
||||
io_handler = ImageEmbeddingMediaIO()
|
||||
|
||||
valid_tensor = _create_valid_dense_tensor()
|
||||
encoded = _encode_tensor(valid_tensor)
|
||||
|
||||
# Should not raise any exception
|
||||
result = io_handler.load_base64("", encoded.decode("utf-8"))
|
||||
assert result.shape == valid_tensor.shape
|
||||
|
||||
def test_valid_sparse_tensor_accepted(self):
|
||||
"""Baseline: Valid sparse tensors should load successfully."""
|
||||
io_handler = AudioEmbeddingMediaIO()
|
||||
|
||||
valid_sparse = _create_valid_sparse_tensor()
|
||||
encoded = _encode_tensor(valid_sparse)
|
||||
|
||||
# Should not raise any exception (sparse tensors remain sparse)
|
||||
result = io_handler.load_base64("", encoded.decode("utf-8"))
|
||||
assert result.shape == valid_sparse.shape
|
||||
|
||||
def test_malicious_sparse_tensor_rejected(self):
|
||||
"""Security: Malicious sparse tensors should be rejected."""
|
||||
io_handler = ImageEmbeddingMediaIO()
|
||||
|
||||
malicious_tensor = _create_malicious_sparse_tensor()
|
||||
encoded = _encode_tensor(malicious_tensor)
|
||||
|
||||
# Should raise RuntimeError due to invalid sparse tensor
|
||||
with pytest.raises((RuntimeError, ValueError)) as exc_info:
|
||||
io_handler.load_base64("", encoded.decode("utf-8"))
|
||||
|
||||
error_msg = str(exc_info.value).lower()
|
||||
assert "sparse" in error_msg or "index" in error_msg or "bounds" in error_msg
|
||||
|
||||
def test_load_bytes_validates(self):
|
||||
"""Security: Validation should also work for load_bytes method."""
|
||||
io_handler = ImageEmbeddingMediaIO()
|
||||
|
||||
malicious_tensor = _create_malicious_sparse_tensor()
|
||||
buffer = io.BytesIO()
|
||||
torch.save(malicious_tensor, buffer)
|
||||
buffer.seek(0)
|
||||
|
||||
with pytest.raises((RuntimeError, ValueError)):
|
||||
io_handler.load_bytes(buffer.read())
|
||||
|
||||
|
||||
class TestAudioEmbedsValidation:
|
||||
"""Test sparse tensor validation in audio embeddings (Chat API)."""
|
||||
|
||||
def test_valid_dense_tensor_accepted(self):
|
||||
"""Baseline: Valid dense tensors should work normally."""
|
||||
io_handler = AudioEmbeddingMediaIO()
|
||||
|
||||
valid_tensor = _create_valid_dense_tensor()
|
||||
encoded = _encode_tensor(valid_tensor)
|
||||
|
||||
# Should not raise any exception
|
||||
result = io_handler.load_base64("", encoded.decode("utf-8"))
|
||||
assert result.shape == valid_tensor.shape
|
||||
|
||||
def test_valid_sparse_tensor_accepted(self):
|
||||
"""Baseline: Valid sparse tensors should be converted successfully."""
|
||||
io_handler = AudioEmbeddingMediaIO()
|
||||
|
||||
valid_sparse = _create_valid_sparse_tensor()
|
||||
encoded = _encode_tensor(valid_sparse)
|
||||
|
||||
# Should not raise any exception
|
||||
result = io_handler.load_base64("", encoded.decode("utf-8"))
|
||||
assert result.is_sparse is False
|
||||
|
||||
def test_malicious_sparse_tensor_rejected(self):
|
||||
"""Security: Malicious sparse tensors should be rejected."""
|
||||
io_handler = AudioEmbeddingMediaIO()
|
||||
|
||||
malicious_tensor = _create_malicious_sparse_tensor()
|
||||
encoded = _encode_tensor(malicious_tensor)
|
||||
|
||||
# Should raise RuntimeError due to invalid sparse tensor
|
||||
with pytest.raises((RuntimeError, ValueError)) as exc_info:
|
||||
io_handler.load_base64("", encoded.decode("utf-8"))
|
||||
|
||||
error_msg = str(exc_info.value).lower()
|
||||
assert "sparse" in error_msg or "index" in error_msg or "bounds" in error_msg
|
||||
|
||||
def test_load_bytes_validates(self):
|
||||
"""Security: Validation should also work for load_bytes method."""
|
||||
io_handler = AudioEmbeddingMediaIO()
|
||||
|
||||
malicious_tensor = _create_malicious_sparse_tensor()
|
||||
buffer = io.BytesIO()
|
||||
torch.save(malicious_tensor, buffer)
|
||||
buffer.seek(0)
|
||||
|
||||
with pytest.raises((RuntimeError, ValueError)):
|
||||
io_handler.load_bytes(buffer.read())
|
||||
|
||||
|
||||
class TestSparseTensorValidationIntegration:
|
||||
"""
|
||||
These tests verify the complete attack chain is blocked at all entry points.
|
||||
"""
|
||||
|
||||
def test_attack_scenario_completions_api(self, model_config):
|
||||
"""
|
||||
Simulate a complete attack through the Completions API.
|
||||
|
||||
Attack scenario:
|
||||
1. Attacker crafts malicious sparse tensor
|
||||
2. Encodes it as base64
|
||||
3. Sends to /v1/completions with prompt_embeds parameter
|
||||
4. Server should reject before memory corruption occurs
|
||||
"""
|
||||
# Step 1-2: Attacker creates malicious payload
|
||||
attack_payload = _encode_tensor(_create_malicious_sparse_tensor())
|
||||
|
||||
# Step 3-4: Server processes and should reject
|
||||
with pytest.raises((RuntimeError, ValueError)):
|
||||
safe_load_prompt_embeds(model_config, attack_payload)
|
||||
|
||||
def test_attack_scenario_chat_api_image(self):
|
||||
"""
|
||||
Simulate attack through Chat API with image_embeds.
|
||||
|
||||
Verifies the image embeddings path is protected.
|
||||
"""
|
||||
io_handler = ImageEmbeddingMediaIO()
|
||||
attack_payload = _encode_tensor(_create_malicious_sparse_tensor())
|
||||
|
||||
with pytest.raises((RuntimeError, ValueError)):
|
||||
io_handler.load_base64("", attack_payload.decode("utf-8"))
|
||||
|
||||
def test_attack_scenario_chat_api_audio(self):
|
||||
"""
|
||||
Simulate attack through Chat API with audio_embeds.
|
||||
|
||||
Verifies the audio embeddings path is protected.
|
||||
"""
|
||||
io_handler = AudioEmbeddingMediaIO()
|
||||
attack_payload = _encode_tensor(_create_malicious_sparse_tensor())
|
||||
|
||||
with pytest.raises((RuntimeError, ValueError)):
|
||||
io_handler.load_base64("", attack_payload.decode("utf-8"))
|
||||
Reference in New Issue
Block a user