wehub-resource-sync
1279 lines
42 KiB
Python
1279 lines
42 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Tests for the admin benchmark module."""
|
|
|
|
import asyncio
|
|
from types import SimpleNamespace
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
from omlx.admin.benchmark import (
|
|
VALID_BATCH_SIZES,
|
|
VALID_PROMPT_LENGTHS,
|
|
BenchmarkRequest,
|
|
BenchmarkRun,
|
|
_clean_model_name,
|
|
_compute_single_metrics,
|
|
_detect_experimental_features,
|
|
_detect_quantization,
|
|
_generate_prompt,
|
|
_run_single_test,
|
|
cleanup_old_runs,
|
|
create_run,
|
|
get_run,
|
|
run_benchmark,
|
|
)
|
|
|
|
|
|
# =============================================================================
|
|
# BenchmarkRequest validation tests
|
|
# =============================================================================
|
|
|
|
|
|
class TestBenchmarkRequest:
|
|
def test_valid_request(self):
|
|
req = BenchmarkRequest(
|
|
model_id="test-model",
|
|
prompt_lengths=[1024, 4096],
|
|
generation_length=128,
|
|
batch_sizes=[2, 4],
|
|
)
|
|
assert req.model_id == "test-model"
|
|
assert req.prompt_lengths == [1024, 4096]
|
|
assert req.batch_sizes == [2, 4]
|
|
|
|
def test_prompt_lengths_sorted(self):
|
|
req = BenchmarkRequest(
|
|
model_id="test-model",
|
|
prompt_lengths=[8192, 1024, 4096],
|
|
)
|
|
assert req.prompt_lengths == [1024, 4096, 8192]
|
|
|
|
def test_empty_prompt_lengths_rejected(self):
|
|
with pytest.raises(ValueError, match="At least one prompt length"):
|
|
BenchmarkRequest(model_id="test-model", prompt_lengths=[])
|
|
|
|
def test_invalid_prompt_length_rejected(self):
|
|
with pytest.raises(ValueError, match="Invalid prompt length 512"):
|
|
BenchmarkRequest(model_id="test-model", prompt_lengths=[512])
|
|
|
|
def test_invalid_batch_size_rejected(self):
|
|
with pytest.raises(ValueError, match="Invalid batch size 3"):
|
|
BenchmarkRequest(
|
|
model_id="test-model",
|
|
prompt_lengths=[1024],
|
|
batch_sizes=[3],
|
|
)
|
|
|
|
def test_empty_batch_sizes_allowed(self):
|
|
req = BenchmarkRequest(
|
|
model_id="test-model",
|
|
prompt_lengths=[1024],
|
|
batch_sizes=[],
|
|
)
|
|
assert req.batch_sizes == []
|
|
|
|
def test_batch_sizes_sorted(self):
|
|
req = BenchmarkRequest(
|
|
model_id="test-model",
|
|
prompt_lengths=[1024],
|
|
batch_sizes=[8, 2, 4],
|
|
)
|
|
assert req.batch_sizes == [2, 4, 8]
|
|
|
|
def test_default_generation_length(self):
|
|
req = BenchmarkRequest(
|
|
model_id="test-model",
|
|
prompt_lengths=[1024],
|
|
)
|
|
assert req.generation_length == 128
|
|
|
|
def test_force_lm_engine_defaults_to_false(self):
|
|
req = BenchmarkRequest(
|
|
model_id="test-model",
|
|
prompt_lengths=[1024],
|
|
)
|
|
assert req.force_lm_engine is False
|
|
|
|
def test_force_lm_engine_can_be_enabled(self):
|
|
req = BenchmarkRequest(
|
|
model_id="test-model",
|
|
prompt_lengths=[1024],
|
|
force_lm_engine=True,
|
|
)
|
|
assert req.force_lm_engine is True
|
|
|
|
|
|
# =============================================================================
|
|
# Prompt generation tests
|
|
# =============================================================================
|
|
|
|
|
|
class TestGeneratePrompt:
|
|
def test_exact_token_count(self):
|
|
"""Verify prompt generates exact number of tokens."""
|
|
tokenizer = MagicMock()
|
|
|
|
# Simulate tokenizer behavior
|
|
def mock_encode(text):
|
|
# Return roughly 1 token per 4 chars
|
|
return list(range(len(text) // 4))
|
|
|
|
def mock_decode(tokens):
|
|
return "x" * len(tokens) * 4
|
|
|
|
tokenizer.encode = mock_encode
|
|
tokenizer.decode = mock_decode
|
|
|
|
prompt = _generate_prompt(tokenizer, 1024)
|
|
|
|
# Verify encode was called and result was truncated
|
|
encoded = tokenizer.encode(prompt)
|
|
assert len(encoded) == 1024
|
|
|
|
def test_uuid_prefix_uniqueness(self):
|
|
"""Verify each generated prompt has a unique UUID prefix."""
|
|
tokenizer = MagicMock()
|
|
tokenizer.encode = lambda text: list(range(2048))
|
|
tokenizer.decode = lambda tokens: f"decoded-{len(tokens)}"
|
|
|
|
prompts = set()
|
|
for _ in range(10):
|
|
# We can't easily verify uniqueness since decode is mocked,
|
|
# but we verify encode is called with text containing "BENCH-"
|
|
prompt = _generate_prompt(tokenizer, 100)
|
|
prompts.add(prompt)
|
|
|
|
# With mock decode they'll all be the same, but in real usage
|
|
# the UUID prefix ensures cache isolation
|
|
|
|
|
|
# =============================================================================
|
|
# Metrics computation tests
|
|
# =============================================================================
|
|
|
|
|
|
class TestComputeMetrics:
|
|
def test_basic_metrics(self):
|
|
"""Test metric computation with known values."""
|
|
metrics = _compute_single_metrics(
|
|
prompt_tokens=1024,
|
|
completion_tokens=128,
|
|
start_time=0.0,
|
|
first_token_time=0.1, # 100ms TTFT
|
|
end_time=1.38, # 1.28s generation
|
|
peak_memory=4 * 1024 * 1024 * 1024, # 4GB
|
|
cached_tokens=0,
|
|
)
|
|
|
|
assert metrics["ttft_ms"] == pytest.approx(100.0, abs=0.1)
|
|
assert metrics["prompt_tokens"] == 1024
|
|
assert metrics["completion_tokens"] == 128
|
|
assert metrics["cached_tokens"] == 0
|
|
assert metrics["peak_memory_bytes"] == 4 * 1024 * 1024 * 1024
|
|
|
|
# Gen TPS = 128 / 1.28 = 100 tok/s
|
|
assert metrics["gen_tps"] == pytest.approx(100.0, abs=0.1)
|
|
|
|
# Processing TPS = 1024 / 0.1 = 10240 tok/s
|
|
assert metrics["processing_tps"] == pytest.approx(10240.0, abs=1.0)
|
|
|
|
# TPOT = 1280ms / 127 = ~10.08 ms/tok
|
|
assert metrics["tpot_ms"] == pytest.approx(10.08, abs=0.1)
|
|
|
|
# E2E = 1.38s
|
|
assert metrics["e2e_latency_s"] == pytest.approx(1.38, abs=0.001)
|
|
|
|
# Total throughput = (1024 + 128) / 1.38 = ~834.8 tok/s
|
|
assert metrics["total_throughput"] == pytest.approx(834.8, abs=1.0)
|
|
|
|
def test_zero_duration_safety(self):
|
|
"""Single-token outputs do not report bogus decode throughput."""
|
|
metrics = _compute_single_metrics(
|
|
prompt_tokens=100,
|
|
completion_tokens=1,
|
|
start_time=0.0,
|
|
first_token_time=0.0,
|
|
end_time=0.0,
|
|
peak_memory=0,
|
|
cached_tokens=0,
|
|
)
|
|
# Should not raise, values should be finite
|
|
assert metrics["ttft_ms"] == 0.0
|
|
assert metrics["gen_tps"] == 0.0
|
|
assert metrics["tpot_ms"] == 0.0
|
|
|
|
def test_native_duration_overrides(self):
|
|
"""Native engine timings can override streaming timing artifacts."""
|
|
metrics = _compute_single_metrics(
|
|
prompt_tokens=1024,
|
|
completion_tokens=128,
|
|
start_time=0.0,
|
|
first_token_time=1.0,
|
|
end_time=1.0,
|
|
peak_memory=0,
|
|
cached_tokens=0,
|
|
prefill_duration_s=4.0,
|
|
generation_duration_s=8.0,
|
|
)
|
|
|
|
assert metrics["processing_tps"] == pytest.approx(256.0)
|
|
assert metrics["gen_tps"] == pytest.approx(16.0)
|
|
assert metrics["tpot_ms"] == pytest.approx(62.99, abs=0.01)
|
|
|
|
def test_single_token_completion_has_no_decode_rate(self):
|
|
"""Immediate stop has no inter-token decode interval to benchmark."""
|
|
metrics = _compute_single_metrics(
|
|
prompt_tokens=16384,
|
|
completion_tokens=1,
|
|
start_time=0.0,
|
|
first_token_time=8.629,
|
|
end_time=8.629001,
|
|
peak_memory=0,
|
|
cached_tokens=0,
|
|
)
|
|
|
|
assert metrics["gen_tps"] == 0.0
|
|
assert metrics["tpot_ms"] == 0.0
|
|
assert metrics["total_throughput"] == pytest.approx(1898.7, abs=0.1)
|
|
|
|
|
|
class TestRunSingleTest:
|
|
@pytest.mark.asyncio
|
|
async def test_uses_native_diffusion_metrics_for_chunked_stream(self):
|
|
"""Diffusion streams by canvas, so benchmark must not use chunk timing."""
|
|
|
|
class ChunkedDiffusionEngine:
|
|
async def stream_generate(self, **kwargs):
|
|
yield SimpleNamespace(
|
|
completion_tokens=128,
|
|
prompt_tokens=1024,
|
|
cached_tokens=0,
|
|
new_text="x" * 10,
|
|
finished=True,
|
|
finish_reason="length",
|
|
prompt_tps=256.0,
|
|
generation_tps=32.0,
|
|
)
|
|
|
|
metrics = await _run_single_test(
|
|
ChunkedDiffusionEngine(),
|
|
prompt="prompt",
|
|
max_tokens=128,
|
|
pp_len=1024,
|
|
)
|
|
|
|
assert metrics["processing_tps"] == pytest.approx(256.0)
|
|
assert metrics["gen_tps"] == pytest.approx(32.0)
|
|
assert metrics["gen_tps"] < 1000.0
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_uses_diffusion_canvas_metrics_when_eos_stops_early(self):
|
|
"""Diffusion benchmark TG should measure canvas work, not early EOS text."""
|
|
|
|
class EarlyStopDiffusionEngine:
|
|
async def stream_generate(self, **kwargs):
|
|
yield SimpleNamespace(
|
|
completion_tokens=16,
|
|
prompt_tokens=1024,
|
|
cached_tokens=0,
|
|
new_text="short answer",
|
|
finished=True,
|
|
finish_reason="stop",
|
|
prompt_tps=256.0,
|
|
generation_tps=2.0,
|
|
diffusion_canvas_tokens=128,
|
|
diffusion_canvas_tps=64.0,
|
|
)
|
|
|
|
metrics = await _run_single_test(
|
|
EarlyStopDiffusionEngine(),
|
|
prompt="prompt",
|
|
max_tokens=128,
|
|
pp_len=1024,
|
|
)
|
|
|
|
assert metrics["completion_tokens"] == 128
|
|
assert metrics["gen_tps"] == pytest.approx(64.0)
|
|
assert metrics["tpot_ms"] == pytest.approx(15.75, abs=0.01)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_uses_producer_timestamps_for_aggregated_output(self):
|
|
"""Aggregated chunks should use producer-side decode timing."""
|
|
|
|
class AggregatedEngine:
|
|
async def stream_generate(self, **kwargs):
|
|
yield SimpleNamespace(
|
|
completion_tokens=128,
|
|
prompt_tokens=1024,
|
|
cached_tokens=0,
|
|
new_text="x" * 128,
|
|
finished=True,
|
|
finish_reason="length",
|
|
generated_at=0.2,
|
|
generated_until=1.2,
|
|
)
|
|
|
|
with patch("omlx.admin.benchmark.time.perf_counter", side_effect=[0.0, 1.3]):
|
|
metrics = await _run_single_test(
|
|
AggregatedEngine(),
|
|
prompt="prompt",
|
|
max_tokens=128,
|
|
pp_len=1024,
|
|
)
|
|
|
|
assert metrics["ttft_ms"] == pytest.approx(200.0)
|
|
assert metrics["gen_tps"] == pytest.approx(128.0)
|
|
assert metrics["tpot_ms"] == pytest.approx(7.87, abs=0.01)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_aggregated_output_without_end_time_has_no_decode_rate(self):
|
|
"""Single aggregate chunks need a producer-side end time for tg TPS."""
|
|
|
|
class AggregatedEngine:
|
|
async def stream_generate(self, **kwargs):
|
|
yield SimpleNamespace(
|
|
completion_tokens=128,
|
|
prompt_tokens=1024,
|
|
cached_tokens=0,
|
|
new_text="x" * 128,
|
|
finished=True,
|
|
finish_reason="length",
|
|
generated_at=0.2,
|
|
)
|
|
|
|
with patch("omlx.admin.benchmark.time.perf_counter", side_effect=[0.0, 1.2]):
|
|
metrics = await _run_single_test(
|
|
AggregatedEngine(),
|
|
prompt="prompt",
|
|
max_tokens=128,
|
|
pp_len=1024,
|
|
)
|
|
|
|
assert metrics["ttft_ms"] == pytest.approx(200.0)
|
|
assert metrics["gen_tps"] == 0.0
|
|
assert metrics["tpot_ms"] == 0.0
|
|
|
|
|
|
# =============================================================================
|
|
# BenchmarkRun lifecycle tests
|
|
# =============================================================================
|
|
|
|
|
|
class TestBenchmarkRunLifecycle:
|
|
def test_create_run(self):
|
|
req = BenchmarkRequest(
|
|
model_id="test-model",
|
|
prompt_lengths=[1024],
|
|
)
|
|
run = create_run(req)
|
|
assert run.bench_id.startswith("bench-")
|
|
assert run.status == "running"
|
|
assert run.results == []
|
|
|
|
def test_get_run(self):
|
|
req = BenchmarkRequest(
|
|
model_id="test-model",
|
|
prompt_lengths=[1024],
|
|
)
|
|
run = create_run(req)
|
|
found = get_run(run.bench_id)
|
|
assert found is run
|
|
|
|
def test_get_nonexistent_run(self):
|
|
assert get_run("nonexistent") is None
|
|
|
|
def test_cleanup_old_runs(self):
|
|
# Create many completed runs
|
|
for _ in range(15):
|
|
req = BenchmarkRequest(
|
|
model_id="test-model",
|
|
prompt_lengths=[1024],
|
|
)
|
|
run = create_run(req)
|
|
run.status = "completed"
|
|
|
|
cleanup_old_runs(max_runs=5)
|
|
|
|
# Should have at most ~5 completed + any running ones
|
|
from omlx.admin.benchmark import _benchmark_runs
|
|
|
|
completed = [r for r in _benchmark_runs.values() if r.status == "completed"]
|
|
assert len(completed) <= 5
|
|
|
|
|
|
class _FakeBenchTokenizer:
|
|
def encode(self, text):
|
|
return list(range(2048))
|
|
|
|
def decode(self, tokens):
|
|
return "prompt"
|
|
|
|
|
|
class _FakeBenchEngine:
|
|
tokenizer = _FakeBenchTokenizer()
|
|
_engine = None
|
|
|
|
async def stream_generate(self, **kwargs):
|
|
yield SimpleNamespace(
|
|
completion_tokens=1,
|
|
prompt_tokens=1024,
|
|
cached_tokens=0,
|
|
new_text="x",
|
|
finished=True,
|
|
finish_reason="length",
|
|
)
|
|
|
|
|
|
class _FakeSettingsManager:
|
|
def __init__(self, settings):
|
|
self.settings = settings
|
|
|
|
def get_settings(self, model_id):
|
|
return self.settings
|
|
|
|
|
|
class _FakeBenchEnginePool:
|
|
def __init__(self, settings=None, engine=None):
|
|
self._settings_manager = (
|
|
_FakeSettingsManager(settings) if settings is not None else None
|
|
)
|
|
self._engine = engine or _FakeBenchEngine()
|
|
self.force_lm_values = []
|
|
|
|
def get_loaded_model_ids(self):
|
|
return []
|
|
|
|
async def get_engine(self, model_id, force_lm=False):
|
|
self.force_lm_values.append(force_lm)
|
|
return self._engine
|
|
|
|
async def _unload_engine(self, model_id):
|
|
pass
|
|
|
|
|
|
class TestBenchmarkEngineSelection:
|
|
async def _run(self, *, settings=None, force_lm_engine=False):
|
|
run = BenchmarkRun(
|
|
bench_id="bench-test",
|
|
request=BenchmarkRequest(
|
|
model_id="test-model",
|
|
prompt_lengths=[1024],
|
|
generation_length=1,
|
|
force_lm_engine=force_lm_engine,
|
|
),
|
|
)
|
|
pool = _FakeBenchEnginePool(settings)
|
|
with patch("omlx.admin.benchmark._upload_to_omlx_ai", AsyncMock()):
|
|
await run_benchmark(run, pool)
|
|
return run, pool
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_auto_uses_vlm_engine_for_vlm_mtp_with_drafter(self):
|
|
settings = SimpleNamespace(
|
|
vlm_mtp_enabled=True,
|
|
vlm_mtp_draft_model="draft-model",
|
|
)
|
|
|
|
run, pool = await self._run(settings=settings)
|
|
|
|
assert pool.force_lm_values == [False]
|
|
assert run.experimental_features == ["vlm_mtp"]
|
|
assert run.status == "completed"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_force_lm_engine_overrides_vlm_mtp_auto(self):
|
|
settings = SimpleNamespace(
|
|
vlm_mtp_enabled=True,
|
|
vlm_mtp_draft_model="draft-model",
|
|
)
|
|
|
|
run, pool = await self._run(settings=settings, force_lm_engine=True)
|
|
|
|
assert pool.force_lm_values == [True]
|
|
assert run.experimental_features == ["vlm_mtp"]
|
|
assert run.status == "completed"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_auto_keeps_lm_engine_without_vlm_mtp_drafter(self):
|
|
settings = SimpleNamespace(
|
|
vlm_mtp_enabled=True,
|
|
vlm_mtp_draft_model=None,
|
|
)
|
|
|
|
run, pool = await self._run(settings=settings)
|
|
|
|
assert pool.force_lm_values == [True]
|
|
assert run.experimental_features == ["vlm_mtp"]
|
|
assert run.status == "completed"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_batch_request_skips_engine_with_none_scheduler_core(self):
|
|
run = BenchmarkRun(
|
|
bench_id="bench-test",
|
|
request=BenchmarkRequest(
|
|
model_id="test-model",
|
|
prompt_lengths=[1024],
|
|
generation_length=1,
|
|
batch_sizes=[2],
|
|
),
|
|
)
|
|
pool = _FakeBenchEnginePool(engine=_FakeBenchEngine())
|
|
|
|
with patch("omlx.admin.benchmark._upload_to_omlx_ai", AsyncMock()):
|
|
await run_benchmark(run, pool)
|
|
|
|
assert run.status == "completed"
|
|
assert [r["test_type"] for r in run.results] == ["single"]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_diffusion_warmup_uses_benchmark_generation_length(self):
|
|
class FakeDiffusionEngine(_FakeBenchEngine):
|
|
is_diffusion_model = True
|
|
|
|
def __init__(self):
|
|
self.calls = []
|
|
|
|
async def stream_generate(self, **kwargs):
|
|
self.calls.append(kwargs)
|
|
yield SimpleNamespace(
|
|
completion_tokens=128,
|
|
prompt_tokens=1024,
|
|
cached_tokens=0,
|
|
new_text="x",
|
|
finished=True,
|
|
finish_reason="length",
|
|
prompt_tps=256.0,
|
|
generation_tps=32.0,
|
|
diffusion_canvas_tokens=128,
|
|
diffusion_canvas_tps=64.0,
|
|
)
|
|
|
|
engine = FakeDiffusionEngine()
|
|
run = BenchmarkRun(
|
|
bench_id="bench-test",
|
|
request=BenchmarkRequest(
|
|
model_id="test-model",
|
|
prompt_lengths=[1024],
|
|
generation_length=128,
|
|
),
|
|
)
|
|
pool = _FakeBenchEnginePool(engine=engine)
|
|
|
|
with patch("omlx.admin.benchmark._upload_to_omlx_ai", AsyncMock()):
|
|
await run_benchmark(run, pool)
|
|
|
|
assert run.status == "completed"
|
|
assert engine.calls[0]["max_tokens"] == 128
|
|
|
|
|
|
# =============================================================================
|
|
# Experimental feature detection tests
|
|
# =============================================================================
|
|
|
|
|
|
class TestExperimentalFeatureDetection:
|
|
def test_detects_all_upload_skipping_features(self):
|
|
settings = SimpleNamespace(
|
|
dflash_enabled=True,
|
|
specprefill_enabled=True,
|
|
turboquant_kv_enabled=True,
|
|
mtp_enabled=True,
|
|
vlm_mtp_enabled=True,
|
|
)
|
|
|
|
assert _detect_experimental_features(settings) == [
|
|
"dflash",
|
|
"specprefill",
|
|
"turboquant",
|
|
"mtp",
|
|
"vlm_mtp",
|
|
]
|
|
|
|
def test_missing_flags_are_treated_as_disabled(self):
|
|
assert _detect_experimental_features(SimpleNamespace()) == []
|
|
|
|
|
|
# =============================================================================
|
|
# SSE event format tests
|
|
# =============================================================================
|
|
|
|
|
|
class TestSSEEventFormat:
|
|
@pytest.mark.asyncio
|
|
async def test_send_event(self):
|
|
"""Test that events are properly queued."""
|
|
from omlx.admin.benchmark import _send_event
|
|
|
|
run = BenchmarkRun(
|
|
bench_id="test",
|
|
request=BenchmarkRequest(
|
|
model_id="test-model",
|
|
prompt_lengths=[1024],
|
|
),
|
|
)
|
|
|
|
await _send_event(run, {
|
|
"type": "progress",
|
|
"phase": "single",
|
|
"message": "Testing",
|
|
"current": 1,
|
|
"total": 3,
|
|
})
|
|
|
|
# SSE delivery: events are appended to `run.events` (replay log).
|
|
assert len(run.events) == 1
|
|
event = run.events[0]
|
|
assert event["type"] == "progress"
|
|
assert event["phase"] == "single"
|
|
assert event["current"] == 1
|
|
assert event["total"] == 3
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_result_event_format(self):
|
|
from omlx.admin.benchmark import _send_event
|
|
|
|
run = BenchmarkRun(
|
|
bench_id="test",
|
|
request=BenchmarkRequest(
|
|
model_id="test-model",
|
|
prompt_lengths=[1024],
|
|
),
|
|
)
|
|
|
|
result_data = {
|
|
"test_type": "single",
|
|
"pp": 1024,
|
|
"tg": 128,
|
|
"ttft_ms": 45.2,
|
|
"gen_tps": 81.3,
|
|
}
|
|
await _send_event(run, {"type": "result", "data": result_data})
|
|
|
|
assert len(run.events) == 1
|
|
event = run.events[0]
|
|
assert event["type"] == "result"
|
|
assert event["data"]["test_type"] == "single"
|
|
assert event["data"]["pp"] == 1024
|
|
|
|
|
|
# =============================================================================
|
|
# Quantization detection tests
|
|
# =============================================================================
|
|
|
|
|
|
class TestDetectQuantization:
|
|
def test_from_config_json(self, tmp_path):
|
|
config = {"quantization_config": {"quant_method": "awq", "bits": 4}}
|
|
(tmp_path / "config.json").write_text(
|
|
__import__("json").dumps(config)
|
|
)
|
|
assert _detect_quantization(str(tmp_path)) == "4bit"
|
|
|
|
def test_from_config_json_8bit(self, tmp_path):
|
|
config = {"quantization_config": {"bits": 8}}
|
|
(tmp_path / "config.json").write_text(
|
|
__import__("json").dumps(config)
|
|
)
|
|
assert _detect_quantization(str(tmp_path)) == "8bit"
|
|
|
|
def test_from_dirname_4bit(self, tmp_path):
|
|
model_dir = tmp_path / "Qwen3-30B-A3B-4bit"
|
|
model_dir.mkdir()
|
|
assert _detect_quantization(str(model_dir)) == "4bit"
|
|
|
|
def test_from_dirname_fp16(self, tmp_path):
|
|
model_dir = tmp_path / "Llama-3-8B-fp16"
|
|
model_dir.mkdir()
|
|
assert _detect_quantization(str(model_dir)) == "fp16"
|
|
|
|
def test_from_dirname_bf16(self, tmp_path):
|
|
model_dir = tmp_path / "Model-bf16"
|
|
model_dir.mkdir()
|
|
assert _detect_quantization(str(model_dir)) == "bf16"
|
|
|
|
def test_from_dirname_mxfp4(self, tmp_path):
|
|
model_dir = tmp_path / "gpt-oss-120b-MXFP4"
|
|
model_dir.mkdir()
|
|
assert _detect_quantization(str(model_dir)) == "mxfp4"
|
|
|
|
def test_from_dirname_nvfp4(self, tmp_path):
|
|
model_dir = tmp_path / "Model-NVFP4"
|
|
model_dir.mkdir()
|
|
assert _detect_quantization(str(model_dir)) == "nvfp4"
|
|
|
|
def test_unknown_fallback(self, tmp_path):
|
|
model_dir = tmp_path / "SomeModel"
|
|
model_dir.mkdir()
|
|
assert _detect_quantization(str(model_dir)) == "unknown"
|
|
|
|
def test_config_takes_priority_over_dirname(self, tmp_path):
|
|
model_dir = tmp_path / "Model-4bit"
|
|
model_dir.mkdir()
|
|
config = {"quantization_config": {"bits": 8}}
|
|
(model_dir / "config.json").write_text(
|
|
__import__("json").dumps(config)
|
|
)
|
|
assert _detect_quantization(str(model_dir)) == "8bit"
|
|
|
|
|
|
# =============================================================================
|
|
# Model name cleaning tests
|
|
# =============================================================================
|
|
|
|
|
|
class TestCleanModelName:
|
|
def test_strip_4bit(self):
|
|
assert _clean_model_name("Qwen3-30B-A3B-4bit", "4bit") == "Qwen3-30B-A3B"
|
|
|
|
def test_strip_8bit(self):
|
|
assert _clean_model_name("Llama-3-8B-8bit", "8bit") == "Llama-3-8B"
|
|
|
|
def test_strip_fp16(self):
|
|
assert _clean_model_name("Model-fp16", "fp16") == "Model"
|
|
|
|
def test_strip_mlx_marker(self):
|
|
assert _clean_model_name("Qwen3-30B-MLX-4bit", "4bit") == "Qwen3-30B"
|
|
|
|
def test_strip_mxfp4(self):
|
|
assert _clean_model_name("gpt-oss-120b-MXFP4", "mxfp4") == "gpt-oss-120b"
|
|
|
|
def test_strip_nvfp4(self):
|
|
assert _clean_model_name("Model-NVFP4", "nvfp4") == "Model"
|
|
|
|
def test_no_quant_suffix(self):
|
|
assert _clean_model_name("Qwen3-30B-A3B", "unknown") == "Qwen3-30B-A3B"
|
|
|
|
def test_preserves_model_size(self):
|
|
assert _clean_model_name("DeepSeek-R1-0528-Qwen3-8B-4bit", "4bit") == "DeepSeek-R1-0528-Qwen3-8B"
|
|
|
|
|
|
# =============================================================================
|
|
# Upload integration tests (mocked HTTP)
|
|
# =============================================================================
|
|
|
|
|
|
class TestUploadToOmlxAi:
|
|
@pytest.mark.asyncio
|
|
async def test_upload_success(self):
|
|
"""Test successful upload sends correct SSE events."""
|
|
from omlx.admin.benchmark import _upload_to_omlx_ai
|
|
|
|
run = BenchmarkRun(
|
|
bench_id="test-bench",
|
|
request=BenchmarkRequest(
|
|
model_id="Qwen3-30B-4bit",
|
|
prompt_lengths=[1024],
|
|
),
|
|
)
|
|
run.results = [
|
|
{
|
|
"test_type": "single",
|
|
"pp": 1024,
|
|
"tg": 128,
|
|
"processing_tps": 500.0,
|
|
"gen_tps": 50.0,
|
|
"ttft_ms": 100.0,
|
|
"peak_memory_bytes": 8 * 1024**3,
|
|
},
|
|
]
|
|
|
|
mock_entry = MagicMock()
|
|
mock_entry.model_path = "/models/Qwen3-30B-4bit"
|
|
mock_pool = MagicMock()
|
|
mock_pool.get_entry.return_value = mock_entry
|
|
mock_pool._settings_manager = None
|
|
|
|
mock_response = MagicMock()
|
|
mock_response.status_code = 201
|
|
mock_response.json.return_value = {
|
|
"id": "abc12345",
|
|
"url": "https://omlx.ai/benchmarks/abc12345",
|
|
}
|
|
|
|
mock_to_thread = AsyncMock(return_value=mock_response)
|
|
|
|
with patch("asyncio.to_thread", mock_to_thread):
|
|
await _upload_to_omlx_ai(run, mock_pool)
|
|
|
|
# Collect all events from the replay log.
|
|
events = list(run.events)
|
|
|
|
# Should have: progress, upload, upload_done
|
|
event_types = [e["type"] for e in events]
|
|
assert "progress" in event_types
|
|
assert "upload" in event_types
|
|
assert "upload_done" in event_types
|
|
|
|
upload_event = next(e for e in events if e["type"] == "upload")
|
|
assert upload_event["data"]["context_length"] == 1024
|
|
assert upload_event["data"]["id"] == "abc12345"
|
|
|
|
done_event = next(e for e in events if e["type"] == "upload_done")
|
|
assert done_event["data"]["success"] == 1
|
|
assert done_event["data"]["failed"] == 0
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_upload_duplicate(self):
|
|
"""Test 409 duplicate response is handled as success."""
|
|
from omlx.admin.benchmark import _upload_to_omlx_ai
|
|
|
|
run = BenchmarkRun(
|
|
bench_id="test-bench",
|
|
request=BenchmarkRequest(
|
|
model_id="Qwen3-30B-4bit",
|
|
prompt_lengths=[1024],
|
|
),
|
|
)
|
|
run.results = [
|
|
{
|
|
"test_type": "single",
|
|
"pp": 1024,
|
|
"tg": 128,
|
|
"processing_tps": 500.0,
|
|
"gen_tps": 50.0,
|
|
"ttft_ms": 100.0,
|
|
"peak_memory_bytes": 0,
|
|
},
|
|
]
|
|
|
|
mock_entry = MagicMock()
|
|
mock_entry.model_path = "/models/Qwen3-30B-4bit"
|
|
mock_pool = MagicMock()
|
|
mock_pool.get_entry.return_value = mock_entry
|
|
mock_pool._settings_manager = None
|
|
|
|
mock_response = MagicMock()
|
|
mock_response.status_code = 409
|
|
mock_response.json.return_value = {
|
|
"error": "Duplicate",
|
|
"existing_id": "xyz789",
|
|
"existing_url": "https://omlx.ai/benchmarks/xyz789",
|
|
}
|
|
|
|
mock_to_thread = AsyncMock(return_value=mock_response)
|
|
|
|
with patch("asyncio.to_thread", mock_to_thread):
|
|
await _upload_to_omlx_ai(run, mock_pool)
|
|
|
|
events = list(run.events)
|
|
|
|
upload_event = next(e for e in events if e["type"] == "upload")
|
|
assert upload_event["data"]["duplicate"] is True
|
|
assert upload_event["data"]["id"] == "xyz789"
|
|
|
|
done_event = next(e for e in events if e["type"] == "upload_done")
|
|
assert done_event["data"]["success"] == 1
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_upload_skips_unmeasurable_generation_results(self):
|
|
"""Rows without a measured decode rate are not uploaded."""
|
|
from omlx.admin.benchmark import _upload_to_omlx_ai
|
|
|
|
run = BenchmarkRun(
|
|
bench_id="test-bench",
|
|
request=BenchmarkRequest(
|
|
model_id="Qwen3-30B-4bit",
|
|
prompt_lengths=[1024, 8192],
|
|
),
|
|
)
|
|
run.results = [
|
|
{
|
|
"test_type": "single",
|
|
"pp": 1024,
|
|
"tg": 128,
|
|
"processing_tps": 500.0,
|
|
"gen_tps": 0.0,
|
|
"ttft_ms": 100.0,
|
|
"peak_memory_bytes": 0,
|
|
},
|
|
{
|
|
"test_type": "single",
|
|
"pp": 8192,
|
|
"tg": 128,
|
|
"processing_tps": 500.0,
|
|
"gen_tps": 50.0,
|
|
"ttft_ms": 500.0,
|
|
"peak_memory_bytes": 0,
|
|
},
|
|
]
|
|
|
|
mock_entry = MagicMock()
|
|
mock_entry.model_path = "/models/Qwen3-30B-4bit"
|
|
mock_pool = MagicMock()
|
|
mock_pool.get_entry.return_value = mock_entry
|
|
mock_pool._settings_manager = None
|
|
|
|
mock_response = MagicMock()
|
|
mock_response.status_code = 201
|
|
mock_response.json.return_value = {
|
|
"id": "abc12345",
|
|
"url": "https://omlx.ai/benchmarks/abc12345",
|
|
}
|
|
mock_to_thread = AsyncMock(return_value=mock_response)
|
|
|
|
with patch("asyncio.to_thread", mock_to_thread):
|
|
await _upload_to_omlx_ai(run, mock_pool)
|
|
|
|
assert mock_to_thread.await_count == 1
|
|
|
|
upload_events = [e for e in run.events if e["type"] == "upload"]
|
|
assert [e["data"]["context_length"] for e in upload_events] == [8192]
|
|
|
|
done_event = next(e for e in run.events if e["type"] == "upload_done")
|
|
assert done_event["data"]["total"] == 1
|
|
assert done_event["data"]["success"] == 1
|
|
assert done_event["data"]["failed"] == 0
|
|
assert done_event["data"]["skipped"] == 1
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_upload_skipped_when_experimental_features_enabled(self):
|
|
"""Upload is skipped (no HTTP call) when experimental features were active."""
|
|
from omlx.admin.benchmark import _upload_to_omlx_ai
|
|
|
|
run = BenchmarkRun(
|
|
bench_id="test-bench",
|
|
request=BenchmarkRequest(
|
|
model_id="Qwen3-30B-4bit",
|
|
prompt_lengths=[1024],
|
|
),
|
|
experimental_features=["dflash", "turboquant"],
|
|
)
|
|
run.results = [
|
|
{
|
|
"test_type": "single",
|
|
"pp": 1024,
|
|
"tg": 128,
|
|
"processing_tps": 500.0,
|
|
"gen_tps": 50.0,
|
|
"ttft_ms": 100.0,
|
|
"peak_memory_bytes": 0,
|
|
},
|
|
]
|
|
|
|
mock_pool = MagicMock()
|
|
mock_pool._settings_manager = None
|
|
mock_to_thread = AsyncMock()
|
|
|
|
with patch("asyncio.to_thread", mock_to_thread):
|
|
await _upload_to_omlx_ai(run, mock_pool)
|
|
|
|
events = list(run.events)
|
|
|
|
# Only an upload_skipped event is emitted, no progress / upload / upload_done
|
|
event_types = [e["type"] for e in events]
|
|
assert "upload_skipped" in event_types
|
|
assert "upload" not in event_types
|
|
assert "upload_done" not in event_types
|
|
assert "progress" not in event_types
|
|
|
|
skipped = next(e for e in events if e["type"] == "upload_skipped")
|
|
assert skipped["reason"] == "experimental_features"
|
|
assert skipped["features"] == ["dflash", "turboquant"]
|
|
|
|
# No HTTP call was made
|
|
mock_to_thread.assert_not_called()
|
|
|
|
|
|
_CF_INTERSTITIAL = (
|
|
'<!DOCTYPE html><html lang="en-US"><head><title>Just a moment...</title>'
|
|
'<meta http-equiv="refresh" content="360"></head><body><div class="main-wrapper">'
|
|
+ ("x" * 5000)
|
|
+ "</div></body></html>"
|
|
)
|
|
|
|
|
|
class TestSanitizeUploadError:
|
|
"""The Cloudflare interstitial pollutes the dashboard upload panel when
|
|
omlx.ai's API endpoint is gated behind a managed challenge. The
|
|
sanitizer must detect that case and surface an actionable message
|
|
without dumping the full 5KB HTML body."""
|
|
|
|
def _resp(self, status=403, headers=None, text="", json_raises=True, json_data=None):
|
|
from unittest.mock import MagicMock
|
|
resp = MagicMock()
|
|
resp.status_code = status
|
|
resp.headers = headers or {}
|
|
resp.text = text
|
|
if json_raises:
|
|
resp.json.side_effect = ValueError("not JSON")
|
|
else:
|
|
resp.json.return_value = json_data or {}
|
|
return resp
|
|
|
|
def test_cloudflare_challenge_via_header(self):
|
|
from omlx.admin.benchmark import _sanitize_upload_error
|
|
resp = self._resp(
|
|
status=403,
|
|
headers={"cf-mitigated": "challenge"},
|
|
text=_CF_INTERSTITIAL,
|
|
)
|
|
msg = _sanitize_upload_error(resp)
|
|
assert "Cloudflare" in msg
|
|
assert "403" in msg
|
|
# The raw HTML body must NOT appear in the error message.
|
|
assert "<!DOCTYPE" not in msg
|
|
assert "Just a moment" not in msg
|
|
assert len(msg) < 300
|
|
|
|
def test_cloudflare_challenge_via_body_sniff(self):
|
|
"""Header missing but body still contains the interstitial — covers
|
|
edge transports / proxies that strip cf-mitigated."""
|
|
from omlx.admin.benchmark import _sanitize_upload_error
|
|
resp = self._resp(status=403, headers={}, text=_CF_INTERSTITIAL)
|
|
msg = _sanitize_upload_error(resp)
|
|
assert "Cloudflare" in msg
|
|
assert "<!DOCTYPE" not in msg
|
|
|
|
def test_json_error_field_extracted(self):
|
|
from omlx.admin.benchmark import _sanitize_upload_error
|
|
resp = self._resp(
|
|
status=400,
|
|
json_raises=False,
|
|
json_data={"error": "Invalid model_name"},
|
|
text='{"error": "Invalid model_name"}',
|
|
)
|
|
assert _sanitize_upload_error(resp) == "Invalid model_name"
|
|
|
|
def test_json_detail_field_extracted(self):
|
|
from omlx.admin.benchmark import _sanitize_upload_error
|
|
resp = self._resp(
|
|
status=422,
|
|
json_raises=False,
|
|
json_data={"detail": "context_length out of range"},
|
|
text='{"detail": "context_length out of range"}',
|
|
)
|
|
assert _sanitize_upload_error(resp) == "context_length out of range"
|
|
|
|
def test_html_body_without_cf_signals_collapses_to_hint(self):
|
|
"""Non-CF HTML body (e.g. nginx 502 page) should not be dumped raw."""
|
|
from omlx.admin.benchmark import _sanitize_upload_error
|
|
resp = self._resp(
|
|
status=502,
|
|
text="<html><body>502 Bad Gateway</body></html>",
|
|
)
|
|
msg = _sanitize_upload_error(resp)
|
|
assert "<html>" not in msg
|
|
assert "502" in msg
|
|
|
|
def test_plain_text_short_body_passes_through(self):
|
|
from omlx.admin.benchmark import _sanitize_upload_error
|
|
resp = self._resp(status=500, text="upstream connection refused")
|
|
assert _sanitize_upload_error(resp) == "upstream connection refused"
|
|
|
|
def test_empty_body_falls_back_to_status(self):
|
|
from omlx.admin.benchmark import _sanitize_upload_error
|
|
resp = self._resp(status=503, text="")
|
|
assert _sanitize_upload_error(resp) == "HTTP 503"
|
|
|
|
|
|
# =============================================================================
|
|
# External endpoint benchmark tests
|
|
# =============================================================================
|
|
|
|
|
|
def _external_config():
|
|
from omlx.admin.external_api import ExternalEndpointConfig
|
|
|
|
return ExternalEndpointConfig(
|
|
base_url="http://localhost:8001/v1",
|
|
api_key="sk-test",
|
|
model="remote-model",
|
|
)
|
|
|
|
|
|
def _stream_stats(prompt=1000, completion=128, cached=0):
|
|
from omlx.admin.external_api import StreamStats
|
|
|
|
return StreamStats(
|
|
prompt_tokens=prompt,
|
|
completion_tokens=completion,
|
|
cached_tokens=cached,
|
|
start_time=0.0,
|
|
first_content_time=0.5,
|
|
last_content_time=1.5,
|
|
end_time=1.6,
|
|
text="x" * 16,
|
|
)
|
|
|
|
|
|
class TestExternalBenchmarkRequest:
|
|
def test_external_accepted(self):
|
|
req = BenchmarkRequest(
|
|
model_id="remote-model",
|
|
prompt_lengths=[1024],
|
|
external=_external_config(),
|
|
)
|
|
assert req.external is not None
|
|
assert req.external.model == "remote-model"
|
|
|
|
def test_external_defaults_to_none(self):
|
|
req = BenchmarkRequest(model_id="m", prompt_lengths=[1024])
|
|
assert req.external is None
|
|
|
|
def test_external_invalid_base_url_rejected(self):
|
|
with pytest.raises(ValueError, match="http:// or https://"):
|
|
BenchmarkRequest(
|
|
model_id="m",
|
|
prompt_lengths=[1024],
|
|
external={"base_url": "localhost:8001", "model": "m"},
|
|
)
|
|
|
|
|
|
class TestGenerateExternalPrompt:
|
|
def test_scales_with_target(self):
|
|
from omlx.admin.benchmark import _generate_external_prompt
|
|
|
|
short = _generate_external_prompt(1024)
|
|
long = _generate_external_prompt(4096)
|
|
assert len(long) > len(short) * 3
|
|
|
|
def test_unique_prefix(self):
|
|
from omlx.admin.benchmark import _generate_external_prompt
|
|
|
|
a = _generate_external_prompt(1024)
|
|
b = _generate_external_prompt(1024)
|
|
assert a.startswith("BENCH-")
|
|
assert a != b
|
|
|
|
|
|
class TestRunExternalBenchmark:
|
|
def _make_run(self, prompt_lengths=None, batch_sizes=None):
|
|
return BenchmarkRun(
|
|
bench_id="bench-ext",
|
|
request=BenchmarkRequest(
|
|
model_id="remote-model",
|
|
prompt_lengths=prompt_lengths or [1024],
|
|
batch_sizes=batch_sizes or [],
|
|
external=_external_config(),
|
|
),
|
|
)
|
|
|
|
def _mock_client(self, stats=None):
|
|
client = MagicMock()
|
|
client.stream_chat_completion = AsyncMock(
|
|
return_value=stats or _stream_stats()
|
|
)
|
|
client.aclose = AsyncMock()
|
|
return client
|
|
|
|
async def test_never_touches_engine_pool(self):
|
|
run = self._make_run(batch_sizes=[2])
|
|
pool = MagicMock()
|
|
client = self._mock_client()
|
|
|
|
with patch(
|
|
"omlx.admin.benchmark.ExternalAPIClient", return_value=client
|
|
):
|
|
await run_benchmark(run, pool)
|
|
|
|
assert run.status == "completed"
|
|
pool.get_engine.assert_not_called()
|
|
pool.get_loaded_model_ids.assert_not_called()
|
|
pool._unload_engine.assert_not_called()
|
|
|
|
async def test_event_sequence_and_upload_skipped(self):
|
|
run = self._make_run(prompt_lengths=[1024], batch_sizes=[2])
|
|
client = self._mock_client()
|
|
|
|
with patch(
|
|
"omlx.admin.benchmark.ExternalAPIClient", return_value=client
|
|
):
|
|
await run_benchmark(run, MagicMock())
|
|
|
|
event_types = [e["type"] for e in run.events]
|
|
assert event_types == [
|
|
"progress", # warmup
|
|
"progress", # single
|
|
"result",
|
|
"progress", # batch
|
|
"result",
|
|
"done",
|
|
"upload_skipped",
|
|
]
|
|
skipped = run.events[-1]
|
|
assert skipped["reason"] == "external_endpoint"
|
|
assert run.upload_state["phase"] == "skipped"
|
|
assert run.upload_state["skipped_reason"] == "external_endpoint"
|
|
client.aclose.assert_awaited()
|
|
|
|
async def test_single_result_uses_usage_token_counts(self):
|
|
run = self._make_run(prompt_lengths=[4096])
|
|
client = self._mock_client(_stream_stats(prompt=3900, completion=128))
|
|
|
|
with patch(
|
|
"omlx.admin.benchmark.ExternalAPIClient", return_value=client
|
|
):
|
|
await run_benchmark(run, MagicMock())
|
|
|
|
result = run.results[0]
|
|
assert result["test_type"] == "single"
|
|
assert result["pp"] == 4096
|
|
# Actual usage counts, not the nominal pp target
|
|
assert result["prompt_tokens"] == 3900
|
|
assert result["completion_tokens"] == 128
|
|
assert result["peak_memory_bytes"] is None
|
|
# gen duration 1.0s (first 0.5 → last 1.5) with 128 tokens
|
|
assert result["gen_tps"] == 128.0
|
|
assert result["ttft_ms"] == 500.0
|
|
|
|
async def test_batch_result_aggregates_usage_counts(self):
|
|
run = self._make_run(prompt_lengths=[1024], batch_sizes=[2])
|
|
client = self._mock_client(_stream_stats(prompt=1000, completion=100))
|
|
|
|
with patch(
|
|
"omlx.admin.benchmark.ExternalAPIClient", return_value=client
|
|
):
|
|
await run_benchmark(run, MagicMock())
|
|
|
|
batch = [r for r in run.results if r["test_type"] == "batch"][0]
|
|
assert batch["batch_size"] == 2
|
|
assert batch["total_gen_tokens"] == 200
|
|
|
|
async def test_missing_usage_fails_run(self):
|
|
from omlx.admin.external_api import ExternalEndpointError
|
|
|
|
run = self._make_run()
|
|
client = MagicMock()
|
|
client.stream_chat_completion = AsyncMock(
|
|
side_effect=ExternalEndpointError(
|
|
"External endpoint does not support stream usage"
|
|
)
|
|
)
|
|
client.aclose = AsyncMock()
|
|
|
|
with patch(
|
|
"omlx.admin.benchmark.ExternalAPIClient", return_value=client
|
|
):
|
|
await run_benchmark(run, MagicMock())
|
|
|
|
assert run.status == "error"
|
|
assert "stream usage" in run.error_message
|
|
error_events = [e for e in run.events if e["type"] == "error"]
|
|
assert error_events and "stream usage" in error_events[0]["message"]
|
|
|
|
async def test_cancellation_mid_run(self):
|
|
run = self._make_run()
|
|
started = asyncio.Event()
|
|
|
|
async def hang(**kwargs):
|
|
started.set()
|
|
await asyncio.Event().wait()
|
|
|
|
client = MagicMock()
|
|
client.stream_chat_completion = AsyncMock(side_effect=hang)
|
|
client.aclose = AsyncMock()
|
|
|
|
with patch(
|
|
"omlx.admin.benchmark.ExternalAPIClient", return_value=client
|
|
):
|
|
task = asyncio.create_task(run_benchmark(run, MagicMock()))
|
|
await started.wait()
|
|
task.cancel()
|
|
await task
|
|
|
|
assert run.status == "cancelled"
|
|
assert run.events[-1]["type"] == "error"
|
|
assert "cancelled" in run.events[-1]["message"].lower()
|
|
client.aclose.assert_awaited()
|