mirror of https://github.com/razor-ai/soup.git
279 lines
9.5 KiB
Python
279 lines
9.5 KiB
Python
"""Tests for speculative decoding — CLI flags, model loading, serve integration."""
|
|
|
|
import re
|
|
from unittest.mock import MagicMock
|
|
from unittest.mock import patch as mock_patch
|
|
|
|
import pytest
|
|
|
|
|
|
def _strip_ansi(text: str) -> str:
|
|
"""Remove ANSI escape codes from Rich-formatted output."""
|
|
return re.sub(r'\x1b\[[0-9;]*m', '', text)
|
|
|
|
|
|
# ─── CLI Flag Tests ──────────────────────────────────────────────────────
|
|
|
|
|
|
class TestSpeculativeDecodingCLI:
|
|
"""Test that --speculative-decoding flag is accepted by soup serve."""
|
|
|
|
def test_serve_speculative_flag_exists(self):
|
|
"""The serve command should accept --speculative-decoding flag."""
|
|
from typer.testing import CliRunner
|
|
|
|
from soup_cli.cli import app
|
|
|
|
runner = CliRunner()
|
|
result = runner.invoke(app, ["serve", "--help"])
|
|
assert result.exit_code == 0
|
|
assert "--speculative-decoding" in _strip_ansi(result.output)
|
|
|
|
def test_serve_spec_tokens_flag_exists(self):
|
|
"""The serve command should accept --num-speculative-tokens flag."""
|
|
from typer.testing import CliRunner
|
|
|
|
from soup_cli.cli import app
|
|
|
|
runner = CliRunner()
|
|
result = runner.invoke(app, ["serve", "--help"])
|
|
assert result.exit_code == 0
|
|
assert "--num-speculative-tokens" in _strip_ansi(result.output)
|
|
|
|
|
|
# ─── Draft Model Loading Tests ────────────────────────────────────────────
|
|
|
|
|
|
class TestDraftModelLoading:
|
|
"""Test draft model loading for speculative decoding."""
|
|
|
|
def test_load_draft_model_calls_from_pretrained(self):
|
|
"""_load_draft_model should call AutoModelForCausalLM.from_pretrained."""
|
|
mock_model = MagicMock()
|
|
mock_model.eval = MagicMock()
|
|
|
|
with mock_patch(
|
|
"transformers.AutoModelForCausalLM.from_pretrained",
|
|
return_value=mock_model,
|
|
):
|
|
from soup_cli.commands.serve import _load_draft_model
|
|
|
|
result = _load_draft_model("small-model", "cuda")
|
|
mock_model.eval.assert_called_once()
|
|
assert result == mock_model
|
|
|
|
def test_load_draft_model_cpu_device(self):
|
|
"""CPU device should use device_map='cpu'."""
|
|
mock_model = MagicMock()
|
|
mock_model.eval = MagicMock()
|
|
|
|
with mock_patch(
|
|
"transformers.AutoModelForCausalLM.from_pretrained",
|
|
return_value=mock_model,
|
|
) as mock_load:
|
|
from soup_cli.commands.serve import _load_draft_model
|
|
|
|
_load_draft_model("small-model", "cpu")
|
|
call_kwargs = mock_load.call_args[1]
|
|
assert call_kwargs["device_map"] == "cpu"
|
|
|
|
def test_load_draft_model_blocks_urls(self):
|
|
"""_load_draft_model should reject URL-based model paths (SSRF)."""
|
|
import typer
|
|
|
|
from soup_cli.commands.serve import _load_draft_model
|
|
|
|
with pytest.raises(typer.Exit):
|
|
_load_draft_model("https://evil.com/model", "cuda")
|
|
|
|
def test_load_draft_model_blocks_http_urls(self):
|
|
"""_load_draft_model should reject http:// URLs."""
|
|
import typer
|
|
|
|
from soup_cli.commands.serve import _load_draft_model
|
|
|
|
with pytest.raises(typer.Exit):
|
|
_load_draft_model("http://evil.com/model", "cpu")
|
|
|
|
def test_load_draft_model_cuda_dtype(self):
|
|
"""CUDA device should use device_map='auto' and float16."""
|
|
import torch
|
|
|
|
mock_model = MagicMock()
|
|
mock_model.eval = MagicMock()
|
|
|
|
with mock_patch(
|
|
"transformers.AutoModelForCausalLM.from_pretrained",
|
|
return_value=mock_model,
|
|
) as mock_load:
|
|
from soup_cli.commands.serve import _load_draft_model
|
|
|
|
_load_draft_model("small-model", "cuda")
|
|
call_kwargs = mock_load.call_args[1]
|
|
assert call_kwargs["device_map"] == "auto"
|
|
assert call_kwargs["dtype"] == torch.float16
|
|
|
|
def test_load_draft_model_no_trust_remote_code(self):
|
|
"""_load_draft_model should NOT pass trust_remote_code."""
|
|
mock_model = MagicMock()
|
|
mock_model.eval = MagicMock()
|
|
|
|
with mock_patch(
|
|
"transformers.AutoModelForCausalLM.from_pretrained",
|
|
return_value=mock_model,
|
|
) as mock_load:
|
|
from soup_cli.commands.serve import _load_draft_model
|
|
|
|
_load_draft_model("small-model", "cuda")
|
|
call_kwargs = mock_load.call_args[1]
|
|
assert "trust_remote_code" not in call_kwargs
|
|
|
|
|
|
# ─── Generate Response with Draft Model Tests ────────────────────────────
|
|
|
|
|
|
class TestGenerateWithDraftModel:
|
|
"""Test that _generate_response passes assistant_model correctly."""
|
|
|
|
def test_generate_passes_assistant_model(self):
|
|
"""When assistant_model is provided, it should be in generate kwargs."""
|
|
import torch
|
|
|
|
mock_model = MagicMock()
|
|
mock_model.device = "cpu"
|
|
mock_outputs = torch.tensor([[1, 2, 3, 4, 5]])
|
|
mock_model.generate.return_value = mock_outputs
|
|
|
|
mock_tokenizer = MagicMock()
|
|
mock_tokenizer.chat_template = None
|
|
mock_tokenizer.pad_token_id = 0
|
|
mock_tokenizer.return_value = {
|
|
"input_ids": torch.tensor([[1, 2]]),
|
|
"attention_mask": torch.tensor([[1, 1]]),
|
|
}
|
|
mock_tokenizer.decode.return_value = "response"
|
|
|
|
mock_draft = MagicMock()
|
|
|
|
from soup_cli.commands.serve import _generate_response
|
|
|
|
_generate_response(
|
|
mock_model, mock_tokenizer,
|
|
[{"role": "user", "content": "hello"}],
|
|
max_tokens=10,
|
|
assistant_model=mock_draft,
|
|
num_assistant_tokens=3,
|
|
)
|
|
|
|
gen_call = mock_model.generate.call_args
|
|
assert gen_call[1]["assistant_model"] == mock_draft
|
|
assert gen_call[1]["num_assistant_tokens"] == 3
|
|
|
|
def test_generate_without_assistant_model(self):
|
|
"""When assistant_model is None, it should NOT be in generate kwargs."""
|
|
import torch
|
|
|
|
mock_model = MagicMock()
|
|
mock_model.device = "cpu"
|
|
mock_outputs = torch.tensor([[1, 2, 3, 4, 5]])
|
|
mock_model.generate.return_value = mock_outputs
|
|
|
|
mock_tokenizer = MagicMock()
|
|
mock_tokenizer.chat_template = None
|
|
mock_tokenizer.pad_token_id = 0
|
|
mock_tokenizer.return_value = {
|
|
"input_ids": torch.tensor([[1, 2]]),
|
|
"attention_mask": torch.tensor([[1, 1]]),
|
|
}
|
|
mock_tokenizer.decode.return_value = "response"
|
|
|
|
from soup_cli.commands.serve import _generate_response
|
|
|
|
_generate_response(
|
|
mock_model, mock_tokenizer,
|
|
[{"role": "user", "content": "hello"}],
|
|
max_tokens=10,
|
|
)
|
|
|
|
gen_call = mock_model.generate.call_args
|
|
assert "assistant_model" not in gen_call[1]
|
|
|
|
|
|
# ─── Create App with Draft Model Tests ──────────────────────────────────
|
|
|
|
|
|
class TestCreateAppWithDraftModel:
|
|
"""Test that _create_app accepts and uses draft_model parameter."""
|
|
|
|
def test_create_app_accepts_draft_model(self):
|
|
"""_create_app should accept draft_model parameter."""
|
|
try:
|
|
import fastapi # noqa: F401
|
|
except ImportError:
|
|
pytest.skip("FastAPI not installed")
|
|
|
|
mock_model = MagicMock()
|
|
mock_tokenizer = MagicMock()
|
|
mock_draft = MagicMock()
|
|
|
|
from soup_cli.commands.serve import _create_app
|
|
|
|
app = _create_app(
|
|
model_obj=mock_model,
|
|
tokenizer=mock_tokenizer,
|
|
device="cpu",
|
|
model_name="test",
|
|
max_tokens_default=512,
|
|
draft_model=mock_draft,
|
|
num_speculative_tokens=5,
|
|
)
|
|
assert app is not None
|
|
|
|
def test_create_app_without_draft_model(self):
|
|
"""_create_app should work without draft_model."""
|
|
try:
|
|
import fastapi # noqa: F401
|
|
except ImportError:
|
|
pytest.skip("FastAPI not installed")
|
|
|
|
mock_model = MagicMock()
|
|
mock_tokenizer = MagicMock()
|
|
|
|
from soup_cli.commands.serve import _create_app
|
|
|
|
app = _create_app(
|
|
model_obj=mock_model,
|
|
tokenizer=mock_tokenizer,
|
|
device="cpu",
|
|
model_name="test",
|
|
max_tokens_default=512,
|
|
)
|
|
assert app is not None
|
|
|
|
|
|
# ─── vLLM Speculative Decoding Tests ────────────────────────────────────
|
|
|
|
|
|
class TestVllmSpeculativeDecoding:
|
|
"""Test vLLM engine creation with speculative decoding params."""
|
|
|
|
def test_vllm_engine_spec_params_signature(self):
|
|
"""create_vllm_engine should accept speculative model params."""
|
|
import inspect
|
|
|
|
from soup_cli.utils.vllm import create_vllm_engine
|
|
|
|
sig = inspect.signature(create_vllm_engine)
|
|
assert "speculative_model" in sig.parameters
|
|
assert "num_speculative_tokens" in sig.parameters
|
|
|
|
def test_vllm_engine_spec_params_defaults(self):
|
|
"""Default speculative params should be None and 5."""
|
|
import inspect
|
|
|
|
from soup_cli.utils.vllm import create_vllm_engine
|
|
|
|
sig = inspect.signature(create_vllm_engine)
|
|
assert sig.parameters["speculative_model"].default is None
|
|
assert sig.parameters["num_speculative_tokens"].default == 5
|