wehub-resource-sync
404 lines
14 KiB
Python
404 lines
14 KiB
Python
"""Tests for benchmark engine components: TextGenerator, Conversation, conversation_factory, BenchmarkState."""
|
|
|
|
import asyncio
|
|
import sys
|
|
|
|
import numpy as np
|
|
import pytest
|
|
from transformers import AutoTokenizer
|
|
|
|
from ray.llm._internal.serve.benchmark.models import TurnMetric, WorkloadSpec
|
|
from ray.llm._internal.serve.benchmark.runners import BenchmarkState
|
|
from ray.llm._internal.serve.benchmark.text_gen import (
|
|
Conversation,
|
|
TextGenerator,
|
|
conversation_factory,
|
|
)
|
|
|
|
# ============================================================================
|
|
# Fixtures
|
|
# ============================================================================
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def tokenizer():
|
|
return AutoTokenizer.from_pretrained("gpt2")
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def text_gen(tokenizer):
|
|
return TextGenerator(tokenizer)
|
|
|
|
|
|
def _make_spec(**overrides) -> WorkloadSpec:
|
|
"""Create a resolved WorkloadSpec with sensible defaults."""
|
|
defaults = dict(
|
|
isl=1000,
|
|
hit_rate=0.7,
|
|
num_turns=3,
|
|
osl=50,
|
|
shared_system_prompt_ratio=0.5,
|
|
concurrency=4,
|
|
num_sessions=10,
|
|
)
|
|
defaults.update(overrides)
|
|
return WorkloadSpec(**defaults).resolve()
|
|
|
|
|
|
# ============================================================================
|
|
# TextGenerator
|
|
# ============================================================================
|
|
|
|
|
|
class TestTextGenerator:
|
|
def test_generate_exact_token_count(self, text_gen, tokenizer):
|
|
"""Generated text should tokenize to exactly the requested count."""
|
|
np.random.seed(42)
|
|
for target in [1, 10, 50, 200]:
|
|
text = text_gen.generate(target)
|
|
actual = len(tokenizer.encode(text, add_special_tokens=False))
|
|
assert actual == target, f"Expected {target} tokens, got {actual}"
|
|
|
|
def test_generate_zero_tokens(self, text_gen):
|
|
assert text_gen.generate(0) == ""
|
|
|
|
def test_generate_negative_tokens(self, text_gen):
|
|
assert text_gen.generate(-5) == ""
|
|
|
|
def test_generate_token_ids_length(self, text_gen):
|
|
"""generate_token_ids should return exactly N IDs."""
|
|
np.random.seed(42)
|
|
ids = text_gen.generate_token_ids(100)
|
|
assert len(ids) == 100
|
|
|
|
def test_generate_token_ids_in_vocab_range(self, text_gen, tokenizer):
|
|
np.random.seed(42)
|
|
ids = text_gen.generate_token_ids(500)
|
|
assert all(0 <= i < tokenizer.vocab_size for i in ids)
|
|
|
|
def test_generate_token_ids_zero(self, text_gen):
|
|
assert text_gen.generate_token_ids(0) == []
|
|
|
|
def test_generate_token_ids_negative(self, text_gen):
|
|
assert text_gen.generate_token_ids(-1) == []
|
|
|
|
|
|
# ============================================================================
|
|
# Conversation
|
|
# ============================================================================
|
|
|
|
|
|
class TestConversation:
|
|
def test_turn_zero_messages(self):
|
|
conv = Conversation(
|
|
session_id="s0",
|
|
system_prompt="You are helpful.",
|
|
user_messages=["Hello", "Follow up"],
|
|
num_turns=2,
|
|
)
|
|
msgs = conv.get_turn_messages(0)
|
|
assert len(msgs) == 2
|
|
assert msgs[0] == {"role": "system", "content": "You are helpful."}
|
|
assert msgs[1] == {"role": "user", "content": "Hello"}
|
|
|
|
def test_turn_one_with_injected_response(self):
|
|
conv = Conversation(
|
|
session_id="s0",
|
|
system_prompt="sys",
|
|
user_messages=["u0", "u1"],
|
|
num_turns=2,
|
|
)
|
|
conv.inject_assistant_response(0, "a0")
|
|
msgs = conv.get_turn_messages(1)
|
|
assert len(msgs) == 4
|
|
assert msgs[0] == {"role": "system", "content": "sys"}
|
|
assert msgs[1] == {"role": "user", "content": "u0"}
|
|
assert msgs[2] == {"role": "assistant", "content": "a0"}
|
|
assert msgs[3] == {"role": "user", "content": "u1"}
|
|
|
|
def test_turn_one_without_response_uses_placeholder(self):
|
|
conv = Conversation(
|
|
session_id="s0",
|
|
system_prompt="sys",
|
|
user_messages=["u0", "u1"],
|
|
num_turns=2,
|
|
)
|
|
msgs = conv.get_turn_messages(1)
|
|
assert msgs[2] == {"role": "assistant", "content": "(placeholder)"}
|
|
|
|
def test_no_system_prompt(self):
|
|
conv = Conversation(
|
|
session_id="s0",
|
|
system_prompt="",
|
|
user_messages=["u0"],
|
|
num_turns=1,
|
|
)
|
|
msgs = conv.get_turn_messages(0)
|
|
assert len(msgs) == 1
|
|
assert msgs[0]["role"] == "user"
|
|
|
|
def test_inject_sequential(self):
|
|
conv = Conversation(
|
|
session_id="s0",
|
|
system_prompt="sys",
|
|
user_messages=["u0", "u1", "u2"],
|
|
num_turns=3,
|
|
)
|
|
conv.inject_assistant_response(0, "a0")
|
|
conv.inject_assistant_response(1, "a1")
|
|
msgs = conv.get_turn_messages(2)
|
|
roles = [m["role"] for m in msgs]
|
|
assert roles == ["system", "user", "assistant", "user", "assistant", "user"]
|
|
|
|
def test_inject_overwrite(self):
|
|
conv = Conversation(
|
|
session_id="s0",
|
|
system_prompt="sys",
|
|
user_messages=["u0"],
|
|
num_turns=1,
|
|
)
|
|
conv.inject_assistant_response(0, "first")
|
|
conv.inject_assistant_response(0, "second")
|
|
assert conv._assistant_responses[0] == "second"
|
|
|
|
def test_inject_out_of_order_raises(self):
|
|
conv = Conversation(
|
|
session_id="s0",
|
|
system_prompt="sys",
|
|
user_messages=["u0", "u1", "u2"],
|
|
num_turns=3,
|
|
)
|
|
with pytest.raises(ValueError, match="Cannot inject"):
|
|
conv.inject_assistant_response(2, "skipped turn 0 and 1")
|
|
|
|
|
|
# ============================================================================
|
|
# conversation_factory
|
|
# ============================================================================
|
|
|
|
|
|
class TestConversationFactory:
|
|
def test_creates_valid_conversation(self, text_gen):
|
|
np.random.seed(42)
|
|
spec = _make_spec()
|
|
shared_text = text_gen.generate(spec.shared_s) if spec.shared_s > 0 else ""
|
|
conv = conversation_factory(0, spec, shared_text, text_gen)
|
|
|
|
assert conv.session_id == "session-000000"
|
|
assert conv.num_turns == spec.num_turns
|
|
assert len(conv.user_messages) == spec.num_turns
|
|
|
|
def test_user_messages_have_correct_token_count(self, text_gen, tokenizer):
|
|
np.random.seed(42)
|
|
spec = _make_spec()
|
|
shared_text = text_gen.generate(spec.shared_s) if spec.shared_s > 0 else ""
|
|
conv = conversation_factory(0, spec, shared_text, text_gen)
|
|
|
|
for msg in conv.user_messages:
|
|
token_count = len(tokenizer.encode(msg, add_special_tokens=False))
|
|
assert (
|
|
token_count == spec.user_tokens
|
|
), f"Expected {spec.user_tokens} tokens, got {token_count}"
|
|
|
|
def test_cross_session_shared_prefix(self, text_gen):
|
|
"""Two sessions with cross_sharing > 0 should share the same prefix."""
|
|
np.random.seed(42)
|
|
spec = _make_spec(shared_system_prompt_ratio=0.8)
|
|
shared_text = text_gen.generate(spec.shared_s) if spec.shared_s > 0 else ""
|
|
|
|
conv0 = conversation_factory(0, spec, shared_text, text_gen)
|
|
conv1 = conversation_factory(1, spec, shared_text, text_gen)
|
|
|
|
assert conv0.system_prompt.startswith(shared_text)
|
|
assert conv1.system_prompt.startswith(shared_text)
|
|
|
|
def test_zero_cross_sharing_no_shared_prefix(self, text_gen):
|
|
np.random.seed(42)
|
|
spec = _make_spec(shared_system_prompt_ratio=0.0, isl=2000, hit_rate=0.6)
|
|
assert spec.shared_s == 0
|
|
|
|
def test_session_id_format(self, text_gen):
|
|
np.random.seed(42)
|
|
spec = _make_spec()
|
|
shared_text = text_gen.generate(spec.shared_s) if spec.shared_s > 0 else ""
|
|
conv = conversation_factory(42, spec, shared_text, text_gen)
|
|
assert conv.session_id == "session-000042"
|
|
|
|
|
|
# ============================================================================
|
|
# BenchmarkState
|
|
# ============================================================================
|
|
|
|
|
|
class TestBenchmarkState:
|
|
def _run(self, coro):
|
|
return asyncio.get_event_loop().run_until_complete(coro)
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _setup_loop(self):
|
|
loop = asyncio.new_event_loop()
|
|
asyncio.set_event_loop(loop)
|
|
yield
|
|
loop.close()
|
|
|
|
def test_get_next_session_increments(self):
|
|
spec = _make_spec(num_sessions=5)
|
|
state = BenchmarkState(spec)
|
|
ids = [self._run(state.get_next_session()) for _ in range(5)]
|
|
assert ids == [0, 1, 2, 3, 4]
|
|
|
|
def test_get_next_session_exhausted(self):
|
|
spec = _make_spec(num_sessions=2)
|
|
state = BenchmarkState(spec)
|
|
self._run(state.get_next_session())
|
|
self._run(state.get_next_session())
|
|
assert self._run(state.get_next_session()) is None
|
|
|
|
def test_get_next_session_unlimited(self):
|
|
spec = _make_spec(
|
|
num_sessions=None, duration_s=60.0, request_rate=1.0, concurrency=None
|
|
)
|
|
state = BenchmarkState(spec)
|
|
ids = [self._run(state.get_next_session()) for _ in range(100)]
|
|
assert ids == list(range(100))
|
|
|
|
def test_record_metric_during_warmup(self):
|
|
spec = _make_spec()
|
|
state = BenchmarkState(spec)
|
|
metric = TurnMetric(
|
|
session_id="s0",
|
|
turn=0,
|
|
ttft_ms=1.0,
|
|
fc_ms=1.0,
|
|
itl_ms=1.0,
|
|
e2e_latency_ms=10.0,
|
|
input_tokens=100,
|
|
output_tokens=50,
|
|
start_time_ms=0.0,
|
|
)
|
|
self._run(state.record_metric(metric))
|
|
assert len(state.warmup_metrics) == 1
|
|
assert len(state.measured_metrics) == 0
|
|
|
|
def test_record_metric_after_warmup(self):
|
|
spec = _make_spec()
|
|
state = BenchmarkState(spec)
|
|
self._run(state.mark_warmup_complete())
|
|
|
|
metric = TurnMetric(
|
|
session_id="s0",
|
|
turn=0,
|
|
ttft_ms=1.0,
|
|
fc_ms=1.0,
|
|
itl_ms=1.0,
|
|
e2e_latency_ms=10.0,
|
|
input_tokens=100,
|
|
output_tokens=50,
|
|
start_time_ms=0.0,
|
|
)
|
|
self._run(state.record_metric(metric))
|
|
assert len(state.warmup_metrics) == 0
|
|
assert len(state.measured_metrics) == 1
|
|
|
|
def test_inflight_tracking(self):
|
|
spec = _make_spec()
|
|
state = BenchmarkState(spec)
|
|
|
|
self._run(state.track_inflight_start())
|
|
self._run(state.track_inflight_start())
|
|
assert state.inflight == 2
|
|
assert state.max_inflight == 2
|
|
|
|
self._run(state.track_inflight_end())
|
|
assert state.inflight == 1
|
|
assert state.max_inflight == 2
|
|
|
|
def test_mark_session_complete(self):
|
|
spec = _make_spec()
|
|
state = BenchmarkState(spec)
|
|
state.active_turns["session-000000"] = 0
|
|
state.conversations[0] = Conversation(
|
|
session_id="session-000000",
|
|
system_prompt="sys",
|
|
user_messages=["u0"],
|
|
num_turns=1,
|
|
)
|
|
|
|
self._run(state.mark_session_complete(0, "session-000000"))
|
|
assert state.completed_sessions == 1
|
|
assert "session-000000" not in state.active_turns
|
|
assert 0 not in state.conversations
|
|
|
|
def test_warmup_entropy_check_single_turn_spec(self):
|
|
"""With num_turns=1, max_entropy=0 so warmup completes immediately."""
|
|
spec = _make_spec(num_turns=1, hit_rate=0.3)
|
|
state = BenchmarkState(spec)
|
|
state.active_turns["s0"] = 0
|
|
result = self._run(state.check_and_complete_warmup())
|
|
assert result is True
|
|
assert state.warmup_complete is True
|
|
|
|
def test_warmup_entropy_not_reached(self):
|
|
"""All sessions on same turn -> entropy=0 -> warmup not complete."""
|
|
spec = _make_spec(num_turns=5)
|
|
state = BenchmarkState(spec)
|
|
for i in range(10):
|
|
state.active_turns[f"s{i}"] = 0
|
|
result = self._run(state.check_and_complete_warmup())
|
|
assert result is False
|
|
assert state.warmup_complete is False
|
|
|
|
def test_warmup_entropy_reached(self):
|
|
"""Spread sessions across all turns -> high entropy -> warmup complete."""
|
|
spec = _make_spec(num_turns=4)
|
|
state = BenchmarkState(spec)
|
|
for i in range(20):
|
|
state.active_turns[f"s{i}"] = i % 4
|
|
result = self._run(state.check_and_complete_warmup())
|
|
assert result is True
|
|
assert state.warmup_complete is True
|
|
|
|
def test_has_remaining_sessions_with_limit(self):
|
|
spec = _make_spec(num_sessions=3)
|
|
state = BenchmarkState(spec)
|
|
assert self._run(state.has_remaining_sessions()) is True
|
|
self._run(state.get_next_session())
|
|
self._run(state.get_next_session())
|
|
assert self._run(state.has_remaining_sessions()) is True
|
|
self._run(state.get_next_session())
|
|
assert self._run(state.has_remaining_sessions()) is False
|
|
|
|
def test_has_remaining_sessions_unlimited(self):
|
|
spec = _make_spec(
|
|
num_sessions=None, duration_s=60.0, request_rate=1.0, concurrency=None
|
|
)
|
|
state = BenchmarkState(spec)
|
|
for _ in range(100):
|
|
self._run(state.get_next_session())
|
|
assert self._run(state.has_remaining_sessions()) is True
|
|
|
|
def test_get_inflight(self):
|
|
spec = _make_spec()
|
|
state = BenchmarkState(spec)
|
|
assert self._run(state.get_inflight()) == 0
|
|
self._run(state.track_inflight_start())
|
|
self._run(state.track_inflight_start())
|
|
assert self._run(state.get_inflight()) == 2
|
|
self._run(state.track_inflight_end())
|
|
assert self._run(state.get_inflight()) == 1
|
|
|
|
def test_get_stats(self):
|
|
spec = _make_spec()
|
|
state = BenchmarkState(spec)
|
|
stats = self._run(state.get_stats())
|
|
assert "warmup_metrics" in stats
|
|
assert "measured_metrics" in stats
|
|
assert "completed_sessions" in stats
|
|
assert "max_inflight" in stats
|
|
assert "warmup_complete" in stats
|
|
assert stats["warmup_complete"] is False
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(pytest.main(["-v", __file__]))
|