1399 lines
52 KiB
Python
1399 lines
52 KiB
Python
import asyncio
|
|
import base64
|
|
import io
|
|
import json
|
|
import pickle
|
|
import random
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
import unittest
|
|
from collections import Counter
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import numpy as np
|
|
from PIL import Image
|
|
from tokenizers import Tokenizer
|
|
from tokenizers.models import WordLevel
|
|
from tokenizers.pre_tokenizers import Whitespace
|
|
from transformers import PreTrainedTokenizerFast
|
|
|
|
from sglang.benchmark.datasets import DATASET_MAPPING, get_dataset
|
|
from sglang.benchmark.datasets.agentic_trace import (
|
|
DEFAULT_AGENTIC_OUTPUT_LEN,
|
|
AgenticTraceDataset,
|
|
)
|
|
from sglang.benchmark.datasets.common import DatasetRow, gen_mm_prompt
|
|
from sglang.benchmark.datasets.custom import sample_custom_requests
|
|
from sglang.benchmark.datasets.generated_shared_prefix import (
|
|
GeneratedSharedPrefixDataset,
|
|
_zipf_group_probs,
|
|
get_gen_prefix_cache_path,
|
|
sample_generated_shared_prefix_requests,
|
|
)
|
|
from sglang.benchmark.datasets.image import (
|
|
ImageDataset,
|
|
parse_random_image_resolution,
|
|
sample_image_requests,
|
|
)
|
|
from sglang.benchmark.datasets.mmmu import sample_mmmu_requests
|
|
from sglang.benchmark.datasets.mooncake import get_mooncake_request_over_time
|
|
from sglang.benchmark.datasets.openai_dataset import sample_openai_requests
|
|
from sglang.benchmark.datasets.random import sample_random_requests
|
|
from sglang.benchmark.datasets.sharegpt import sample_sharegpt_requests
|
|
from sglang.benchmark.serving import (
|
|
_BACKEND_API_PATHS,
|
|
_EMBEDDING_BACKENDS,
|
|
ASYNC_REQUEST_FUNCS,
|
|
async_request_openai_embeddings,
|
|
flush_server_cache,
|
|
)
|
|
from sglang.test.ci.ci_register import register_cpu_ci
|
|
|
|
register_cpu_ci(est_time=40, suite="base-a-test-cpu")
|
|
register_cpu_ci(est_time=7, suite="base-c-test-cpu")
|
|
|
|
|
|
class _DummyTokenTensor:
|
|
def __init__(self, value: int):
|
|
self.value = value
|
|
|
|
def numel(self) -> int:
|
|
return self.value
|
|
|
|
|
|
def create_lightweight_tokenizer() -> PreTrainedTokenizerFast:
|
|
"""Create a local lightweight tokenizer for CPU-only dataset tests."""
|
|
vocab = {"[UNK]": 0, "[PAD]": 1, "[BOS]": 2, "[EOS]": 3}
|
|
vocab.update({f"tok_{i}": i + 4 for i in range(2048)})
|
|
|
|
tokenizer = Tokenizer(WordLevel(vocab=vocab, unk_token="[UNK]"))
|
|
tokenizer.pre_tokenizer = Whitespace()
|
|
|
|
hf_tokenizer = PreTrainedTokenizerFast(
|
|
tokenizer_object=tokenizer,
|
|
unk_token="[UNK]",
|
|
pad_token="[PAD]",
|
|
bos_token="[BOS]",
|
|
eos_token="[EOS]",
|
|
)
|
|
hf_tokenizer.chat_template = (
|
|
"{% for message in messages %}"
|
|
"{{ message['role'] }}:"
|
|
"{% if message['content'] is string %}"
|
|
"{{ message['content'] }}"
|
|
"{% else %}"
|
|
"{% for item in message['content'] %}"
|
|
"{% if item['type'] == 'text' %}{{ item['text'] }}{% else %}[IMAGE]{% endif %}"
|
|
"{% endfor %}"
|
|
"{% endif %}\n"
|
|
"{% endfor %}"
|
|
"{% if add_generation_prompt %}assistant:{% endif %}"
|
|
)
|
|
return hf_tokenizer
|
|
|
|
|
|
class TestEmbeddingBenchmarkBackends(unittest.TestCase):
|
|
def test_vllm_embedding_reuses_the_openai_embedding_request_path(self):
|
|
self.assertIn("vllm-embedding", _EMBEDDING_BACKENDS)
|
|
self.assertIs(
|
|
ASYNC_REQUEST_FUNCS["vllm-embedding"], async_request_openai_embeddings
|
|
)
|
|
self.assertEqual(_BACKEND_API_PATHS["vllm-embedding"], "/v1/embeddings")
|
|
|
|
def test_embedding_cache_flush_uses_the_engine_specific_endpoint(self):
|
|
with (
|
|
patch("sglang.benchmark.serving.get_auth_headers", return_value={}),
|
|
patch("sglang.benchmark.serving.requests.post") as post,
|
|
):
|
|
post.return_value = MagicMock()
|
|
|
|
flush_server_cache("http://127.0.0.1:8000", "vllm-embedding")
|
|
post.assert_called_once_with(
|
|
"http://127.0.0.1:8000/reset_prefix_cache", headers={}
|
|
)
|
|
post.reset_mock()
|
|
|
|
flush_server_cache("http://127.0.0.1:30000", "sglang-embedding")
|
|
post.assert_called_once_with(
|
|
"http://127.0.0.1:30000/flush_cache", headers={}
|
|
)
|
|
|
|
|
|
class DummyProcessor:
|
|
def __init__(self, tokenizer: PreTrainedTokenizerFast):
|
|
self.tokenizer = tokenizer
|
|
self.image_token_id = None
|
|
|
|
def apply_chat_template(self, messages, add_generation_prompt=True, tokenize=False):
|
|
return self.tokenizer.apply_chat_template(
|
|
messages,
|
|
add_generation_prompt=add_generation_prompt,
|
|
tokenize=tokenize,
|
|
return_dict=False,
|
|
)
|
|
|
|
def __call__(self, text, images=None, padding=False, return_tensors="pt"):
|
|
text_len = len(self.tokenizer.encode(text[0]))
|
|
image_tokens = 4 * len(images) if images else 0
|
|
return {"input_ids": _DummyTokenTensor(text_len + image_tokens)}
|
|
|
|
|
|
class KimiK3Processor(DummyProcessor):
|
|
"""Mimics the Kimi K3 HF processor's media-kwargs interface (#32541)."""
|
|
|
|
def __init__(self, tokenizer: PreTrainedTokenizerFast):
|
|
super().__init__(tokenizer)
|
|
self.media_call_count = 0
|
|
|
|
def __call__(self, text, medias=None, **kwargs):
|
|
if medias is None:
|
|
raise ValueError("Kimi K3 requires medias with text")
|
|
self.media_call_count += 1
|
|
text_len = len(self.tokenizer.encode(text))
|
|
return {"input_ids": _DummyTokenTensor(text_len + 4 * len(medias))}
|
|
|
|
|
|
class _FakeMMMUDataset:
|
|
def __init__(self, records):
|
|
self.records = records
|
|
|
|
def __len__(self):
|
|
return len(self.records)
|
|
|
|
def select(self, indices):
|
|
if isinstance(indices, range):
|
|
indices = list(indices)
|
|
return _FakeMMMUDataset([self.records[i] for i in indices])
|
|
|
|
def __iter__(self):
|
|
return iter(self.records)
|
|
|
|
|
|
def make_args(**overrides):
|
|
args = {
|
|
"dataset_name": "sharegpt",
|
|
"dataset_path": "",
|
|
"num_prompts": 2,
|
|
"sharegpt_output_len": None,
|
|
"sharegpt_context_len": None,
|
|
"prompt_suffix": "",
|
|
"apply_chat_template": False,
|
|
"tokenize_prompt": False,
|
|
"random_input_len": 8,
|
|
"random_output_len": 4,
|
|
"random_range_ratio": 0.0,
|
|
"image_count": 1,
|
|
"random_image_count": False,
|
|
"image_format": "png",
|
|
"image_content": "blank",
|
|
"image_resolution": "8x8",
|
|
"backend": "sglang",
|
|
"gsp_num_groups": 2,
|
|
"gsp_prompts_per_group": 2,
|
|
"gsp_system_prompt_len": 8,
|
|
"gsp_question_len": 4,
|
|
"gsp_output_len": 4,
|
|
"gsp_range_ratio": 0.0,
|
|
"gsp_fast_prepare": False,
|
|
"gsp_send_routing_key": False,
|
|
"gsp_num_turns": 1,
|
|
"gsp_ordered": False,
|
|
"gsp_group_distribution": "uniform",
|
|
"gsp_zipf_alpha": None,
|
|
"seed": 1,
|
|
"mooncake_workload": "conversation",
|
|
"speed_bench_category": None,
|
|
"speed_bench_output_len": 512,
|
|
"dataset_offset": 0,
|
|
"agentic_max_turns": None,
|
|
}
|
|
args.update(overrides)
|
|
return SimpleNamespace(**args)
|
|
|
|
|
|
class TestBenchmarkDatasetsAPI(unittest.TestCase):
|
|
def setUp(self):
|
|
self.tokenizer = create_lightweight_tokenizer()
|
|
self.processor = DummyProcessor(self.tokenizer)
|
|
self.tmpdir = tempfile.TemporaryDirectory()
|
|
self.tmpdir_path = Path(self.tmpdir.name)
|
|
# Redirect ~ for the GSP on-disk cache to the per-test tempdir, so
|
|
# tests never read/write the real ~/.cache/sglang/benchmark. The Zipf
|
|
# tests in particular compare freshly generated rows against the
|
|
# uniform path, and a stale cache file from prior runs would silently
|
|
# short-circuit the uniform path and break that comparison.
|
|
self._home_patch = patch(
|
|
"sglang.benchmark.datasets.generated_shared_prefix.Path.home",
|
|
return_value=self.tmpdir_path,
|
|
)
|
|
self._home_patch.start()
|
|
|
|
def tearDown(self):
|
|
self._home_patch.stop()
|
|
self.tmpdir.cleanup()
|
|
|
|
def _write_sharegpt_json(self):
|
|
data = [
|
|
{
|
|
"conversations": [
|
|
{"value": "hello world"},
|
|
{"value": "answer one"},
|
|
]
|
|
},
|
|
{
|
|
"conversations": [
|
|
{"value": "how are you"},
|
|
{"value": "answer two"},
|
|
]
|
|
},
|
|
{
|
|
"conversations": [
|
|
{"value": "third prompt"},
|
|
{"value": "answer three"},
|
|
]
|
|
},
|
|
]
|
|
path = self.tmpdir_path / "sharegpt.json"
|
|
with open(path, "w") as f:
|
|
json.dump(data, f)
|
|
return str(path)
|
|
|
|
def _write_custom_jsonl(self):
|
|
rows = [
|
|
{
|
|
"conversations": [
|
|
{"content": "custom prompt 1"},
|
|
{"content": "custom answer 1"},
|
|
]
|
|
},
|
|
{
|
|
"conversations": [
|
|
{"value": "custom prompt 2"},
|
|
{"value": "custom answer 2"},
|
|
]
|
|
},
|
|
]
|
|
path = self.tmpdir_path / "custom.jsonl"
|
|
with open(path, "w") as f:
|
|
for row in rows:
|
|
f.write(json.dumps(row) + "\n")
|
|
return str(path)
|
|
|
|
def _write_openai_jsonl(self):
|
|
rows = [
|
|
{
|
|
"messages": [{"role": "user", "content": "What is 1+1?"}],
|
|
"max_tokens": 7,
|
|
"temperature": 0.3,
|
|
},
|
|
{
|
|
"messages": [{"role": "user", "content": "What is 2+2?"}],
|
|
"max_tokens": 8,
|
|
"tools": [{"type": "function", "function": {"name": "tool_a"}}],
|
|
},
|
|
]
|
|
path = self.tmpdir_path / "openai.jsonl"
|
|
with open(path, "w") as f:
|
|
for row in rows:
|
|
f.write(json.dumps(row) + "\n")
|
|
return str(path)
|
|
|
|
def _write_speed_bench_jsonl(self):
|
|
rows = [
|
|
{
|
|
"question_id": "sb_001",
|
|
"category": "low_entropy",
|
|
"turns": ["Complete this Python function: def add(a, b):"],
|
|
},
|
|
{
|
|
"question_id": "sb_002",
|
|
"category": "mixed",
|
|
"turns": [
|
|
"Explain the concept of attention mechanisms in transformers."
|
|
],
|
|
},
|
|
{
|
|
"question_id": "sb_003",
|
|
"category": "high_entropy",
|
|
"turns": ["Write a short story about a robot discovering music."],
|
|
},
|
|
{
|
|
"question_id": "sb_004",
|
|
"category": "low_entropy",
|
|
"turns": [
|
|
"Sort the following list in ascending order: [5, 2, 8, 1, 9]"
|
|
],
|
|
},
|
|
]
|
|
path = self.tmpdir_path / "speed_bench.jsonl"
|
|
with open(path, "w") as f:
|
|
for row in rows:
|
|
f.write(json.dumps(row) + "\n")
|
|
return str(path)
|
|
|
|
def _write_mooncake_jsonl(self):
|
|
rows = [
|
|
{"timestamp": 1000, "hash_ids": [1, 2], "output_length": 5},
|
|
{"timestamp": 2000, "hash_ids": [3, 4], "output_length": 6},
|
|
]
|
|
path = self.tmpdir_path / "mooncake.jsonl"
|
|
with open(path, "w") as f:
|
|
for row in rows:
|
|
f.write(json.dumps(row) + "\n")
|
|
return str(path)
|
|
|
|
def _write_agentic_trace_json(self):
|
|
trace = {
|
|
"metadata": {"source": "test"},
|
|
"conversations": [
|
|
[
|
|
{
|
|
"messages": [
|
|
{"role": "system", "content": "You are an agent."},
|
|
{"role": "user", "content": "Fix the bug."},
|
|
],
|
|
"prompt_tokens": 100,
|
|
},
|
|
{
|
|
"messages": [{"role": "user", "content": "Tool output: ok."}],
|
|
"prompt_tokens": 200,
|
|
},
|
|
{"messages": []},
|
|
],
|
|
[
|
|
{
|
|
"messages": [{"role": "user", "content": "Run the tests."}],
|
|
"prompt_tokens": 50,
|
|
},
|
|
],
|
|
],
|
|
}
|
|
path = self.tmpdir_path / "agentic_trace.json"
|
|
with open(path, "w") as f:
|
|
json.dump(trace, f)
|
|
return str(path)
|
|
|
|
async def _collect_mooncake_rows(self, records):
|
|
out = []
|
|
async for row in get_mooncake_request_over_time(
|
|
input_requests=records,
|
|
tokenizer=self.tokenizer,
|
|
slowdown_factor=0.0,
|
|
num_rounds=1,
|
|
):
|
|
out.append(row)
|
|
return out
|
|
|
|
def test_sharegpt_sampler(self):
|
|
dataset_path = self._write_sharegpt_json()
|
|
rows = sample_sharegpt_requests(
|
|
dataset_path=dataset_path,
|
|
num_requests=2,
|
|
tokenizer=self.tokenizer,
|
|
)
|
|
self.assertEqual(len(rows), 2)
|
|
self.assertTrue(all(isinstance(row, DatasetRow) for row in rows))
|
|
|
|
def test_random_sampler(self):
|
|
dataset_path = self._write_sharegpt_json()
|
|
rows_text = sample_random_requests(
|
|
input_len=8,
|
|
output_len=4,
|
|
num_prompts=2,
|
|
range_ratio=0.0,
|
|
tokenizer=self.tokenizer,
|
|
dataset_path=dataset_path,
|
|
random_sample=False,
|
|
return_text=True,
|
|
)
|
|
rows_ids = sample_random_requests(
|
|
input_len=8,
|
|
output_len=4,
|
|
num_prompts=2,
|
|
range_ratio=0.0,
|
|
tokenizer=self.tokenizer,
|
|
dataset_path=dataset_path,
|
|
random_sample=False,
|
|
return_text=False,
|
|
)
|
|
self.assertTrue(all(isinstance(row, DatasetRow) for row in rows_text))
|
|
self.assertTrue(all(isinstance(row.prompt, list) for row in rows_ids))
|
|
|
|
def test_custom_sampler(self):
|
|
dataset_path = self._write_custom_jsonl()
|
|
rows = sample_custom_requests(
|
|
dataset_path=dataset_path,
|
|
num_requests=2,
|
|
tokenizer=self.tokenizer,
|
|
)
|
|
self.assertEqual(len(rows), 2)
|
|
self.assertTrue(all(isinstance(row, DatasetRow) for row in rows))
|
|
|
|
def test_openai_sampler(self):
|
|
dataset_path = self._write_openai_jsonl()
|
|
rows = sample_openai_requests(
|
|
dataset_path=dataset_path,
|
|
num_requests=2,
|
|
tokenizer=self.tokenizer,
|
|
)
|
|
self.assertEqual(len(rows), 2)
|
|
self.assertIn("temperature", rows[0].extra_request_body)
|
|
self.assertIn("tools", rows[1].extra_request_body)
|
|
|
|
def test_generated_shared_prefix_sampler(self):
|
|
args = make_args(gsp_num_groups=2, gsp_prompts_per_group=2)
|
|
rows = sample_generated_shared_prefix_requests(
|
|
num_groups=args.gsp_num_groups,
|
|
prompts_per_group=args.gsp_prompts_per_group,
|
|
system_prompt_len=args.gsp_system_prompt_len,
|
|
question_len=args.gsp_question_len,
|
|
output_len=args.gsp_output_len,
|
|
range_ratio=args.gsp_range_ratio,
|
|
tokenizer=self.tokenizer,
|
|
seed=args.seed,
|
|
)
|
|
self.assertEqual(len(rows), 4)
|
|
self.assertTrue(all(isinstance(row, DatasetRow) for row in rows))
|
|
|
|
def test_image_sampler(self):
|
|
rows = sample_image_requests(
|
|
num_requests=2,
|
|
image_count=1,
|
|
input_len=8,
|
|
output_len=4,
|
|
range_ratio=0.0,
|
|
processor=self.processor,
|
|
image_content="blank",
|
|
image_format="png",
|
|
image_resolution="8x8",
|
|
backend="sglang",
|
|
random_image_count=False,
|
|
)
|
|
self.assertEqual(len(rows), 2)
|
|
self.assertTrue(all(isinstance(row, DatasetRow) for row in rows))
|
|
self.assertTrue(all(row.image_data for row in rows))
|
|
|
|
def test_image_sampler_chat_backends_use_raw_prompt(self):
|
|
for backend in ("sglang-oai-chat", "vllm-chat", "lmdeploy-chat"):
|
|
with self.subTest(backend=backend):
|
|
rows = sample_image_requests(
|
|
num_requests=1,
|
|
image_count=1,
|
|
input_len=8,
|
|
output_len=4,
|
|
range_ratio=0.0,
|
|
processor=self.processor,
|
|
image_content="blank",
|
|
image_format="png",
|
|
image_resolution="8x8",
|
|
backend=backend,
|
|
random_image_count=False,
|
|
)
|
|
self.assertEqual(len(rows), 1)
|
|
self.assertTrue(rows[0].image_data)
|
|
for marker in ("user:", "assistant:", "[IMAGE]"):
|
|
self.assertNotIn(marker, rows[0].prompt)
|
|
|
|
def test_image_sampler_uses_kimi_k3_media_contract(self):
|
|
processor = KimiK3Processor(self.tokenizer)
|
|
rows = sample_image_requests(
|
|
num_requests=1,
|
|
image_count=1,
|
|
input_len=8,
|
|
output_len=4,
|
|
range_ratio=0.0,
|
|
processor=processor,
|
|
image_content="blank",
|
|
image_format="png",
|
|
image_resolution="8x8",
|
|
backend="sglang-oai-chat",
|
|
random_image_count=False,
|
|
)
|
|
|
|
self.assertEqual(len(rows), 1)
|
|
self.assertEqual(processor.media_call_count, 1)
|
|
self.assertTrue(rows[0].image_data)
|
|
|
|
def test_image_sampler_random_resolution(self):
|
|
state = np.random.get_state()
|
|
np.random.seed(20260711)
|
|
try:
|
|
rows = sample_image_requests(
|
|
num_requests=4,
|
|
image_count=1,
|
|
input_len=8,
|
|
output_len=4,
|
|
range_ratio=0.0,
|
|
processor=self.processor,
|
|
image_content="blank",
|
|
image_format="png",
|
|
image_resolution="random:8x16-16x32",
|
|
backend="sglang",
|
|
)
|
|
finally:
|
|
np.random.set_state(state)
|
|
|
|
image_sizes = []
|
|
for row in rows:
|
|
encoded = row.image_data[0].split(",", maxsplit=1)[1]
|
|
with Image.open(io.BytesIO(base64.b64decode(encoded))) as image:
|
|
image_sizes.append(image.size)
|
|
|
|
self.assertGreater(len(set(image_sizes)), 1)
|
|
for width, height in image_sizes:
|
|
self.assertGreaterEqual(width, 16)
|
|
self.assertLessEqual(width, 32)
|
|
self.assertGreaterEqual(height, 8)
|
|
self.assertLessEqual(height, 16)
|
|
|
|
def test_image_dataset_seed_is_independent_of_processor_initialization(self):
|
|
dataset = ImageDataset.from_args(
|
|
make_args(
|
|
num_prompts=3,
|
|
image_resolution="random:8x16-16x32",
|
|
seed=20260717,
|
|
)
|
|
)
|
|
processor_init_count = 0
|
|
|
|
def get_processor_with_rng_side_effects(_model_id):
|
|
nonlocal processor_init_count
|
|
processor_init_count += 1
|
|
random.random()
|
|
np.random.random(processor_init_count)
|
|
return self.processor
|
|
|
|
with patch(
|
|
"sglang.benchmark.datasets.image.get_processor",
|
|
side_effect=get_processor_with_rng_side_effects,
|
|
):
|
|
first = dataset.load(model_id="test-model")
|
|
random.seed(999)
|
|
np.random.seed(999)
|
|
second = dataset.load(model_id="test-model")
|
|
|
|
self.assertEqual(
|
|
[
|
|
(row.prompt, row.prompt_len, row.output_len, row.image_data)
|
|
for row in first
|
|
],
|
|
[
|
|
(row.prompt, row.prompt_len, row.output_len, row.image_data)
|
|
for row in second
|
|
],
|
|
)
|
|
|
|
def test_parse_random_image_resolution(self):
|
|
self.assertEqual(
|
|
parse_random_image_resolution("random:256x384-1024x1536"),
|
|
((384, 256), (1536, 1024)),
|
|
)
|
|
self.assertIsNone(parse_random_image_resolution("256x384"))
|
|
with self.assertRaisesRegex(ValueError, "minimum cannot exceed"):
|
|
parse_random_image_resolution("random:1024x1024-256x256")
|
|
|
|
def test_gen_mm_prompt_excludes_special_tokens(self):
|
|
tokenizer = create_lightweight_tokenizer()
|
|
multimodal_special_tokens = [
|
|
"<|image_pad|>",
|
|
"<|video_pad|>",
|
|
"<|vision_start|>",
|
|
"<|vision_end|>",
|
|
"<|vision_pad|>",
|
|
]
|
|
tokenizer.add_special_tokens(
|
|
{"additional_special_tokens": multimodal_special_tokens}
|
|
)
|
|
special_token_ids = set(
|
|
tokenizer.convert_tokens_to_ids(multimodal_special_tokens)
|
|
)
|
|
image_pad_id = tokenizer.convert_tokens_to_ids("<|image_pad|>")
|
|
captured_population = {}
|
|
|
|
def fake_choices(population, k):
|
|
captured_population["tokens"] = population
|
|
return population[:k]
|
|
|
|
with patch(
|
|
"sglang.benchmark.datasets.common.random.choices",
|
|
side_effect=fake_choices,
|
|
):
|
|
gen_mm_prompt(tokenizer, image_pad_id, token_num=8)
|
|
|
|
sampled_pool = set(captured_population["tokens"])
|
|
self.assertFalse(special_token_ids & sampled_pool)
|
|
self.assertTrue(sampled_pool)
|
|
|
|
def test_gen_mm_prompt_is_independent_of_vocab_order(self):
|
|
class OrderedVocabTokenizer:
|
|
all_special_ids = []
|
|
|
|
def __init__(self, items):
|
|
self.vocab = dict(items)
|
|
|
|
def get_vocab(self):
|
|
return self.vocab
|
|
|
|
def decode(self, token_ids):
|
|
return " ".join(map(str, token_ids))
|
|
|
|
items = [(f"token_{token_id}", token_id) for token_id in range(32)]
|
|
first = OrderedVocabTokenizer(items)
|
|
second = OrderedVocabTokenizer(reversed(items))
|
|
|
|
random.seed(20260717)
|
|
first_prompt = gen_mm_prompt(first, image_pad_id=None, token_num=16)
|
|
random.seed(20260717)
|
|
second_prompt = gen_mm_prompt(second, image_pad_id=None, token_num=16)
|
|
|
|
self.assertEqual(first_prompt, second_prompt)
|
|
|
|
def test_mmmu_sampler(self):
|
|
fake_records = [
|
|
{"image_1": Image.new("RGB", (4, 4), color="white"), "question": "q1"},
|
|
{"image_1": Image.new("RGB", (4, 4), color="white"), "question": "q2"},
|
|
{"image_1": Image.new("RGB", (4, 4), color="white"), "question": "q3"},
|
|
]
|
|
fake_dataset = _FakeMMMUDataset(fake_records)
|
|
with patch(
|
|
"sglang.benchmark.datasets.mmmu.load_dataset", return_value=fake_dataset
|
|
):
|
|
rows = sample_mmmu_requests(
|
|
num_requests=2,
|
|
processor=self.processor,
|
|
backend="sglang",
|
|
fixed_output_len=6,
|
|
random_sample=False,
|
|
)
|
|
self.assertEqual(len(rows), 2)
|
|
self.assertTrue(all(isinstance(row, DatasetRow) for row in rows))
|
|
|
|
def test_mooncake_scheduler(self):
|
|
records = [
|
|
{"timestamp": 1000, "hash_ids": [1], "output_length": 5},
|
|
{"timestamp": 2000, "hash_ids": [2], "output_length": 6},
|
|
]
|
|
rows = asyncio.run(self._collect_mooncake_rows(records))
|
|
self.assertEqual(len(rows), 2)
|
|
self.assertTrue(all(isinstance(row, DatasetRow) for row in rows))
|
|
|
|
def test_speed_bench_sampler(self):
|
|
dataset_path = self._write_speed_bench_jsonl()
|
|
args = make_args(
|
|
dataset_name="speed-bench",
|
|
dataset_path=dataset_path,
|
|
num_prompts=3,
|
|
)
|
|
from sglang.benchmark.datasets.speed_bench import SpeedBenchDataset
|
|
|
|
dataset = SpeedBenchDataset.from_args(args)
|
|
rows = dataset.load(self.tokenizer)
|
|
self.assertEqual(len(rows), 3)
|
|
self.assertTrue(all(isinstance(row, DatasetRow) for row in rows))
|
|
self.assertTrue(all(row.output_len == 512 for row in rows))
|
|
self.assertTrue(all(row.prompt_len > 0 for row in rows))
|
|
|
|
def test_speed_bench_category_filter(self):
|
|
dataset_path = self._write_speed_bench_jsonl()
|
|
args = make_args(
|
|
dataset_name="speed-bench",
|
|
dataset_path=dataset_path,
|
|
num_prompts=2,
|
|
speed_bench_category="low_entropy",
|
|
)
|
|
from sglang.benchmark.datasets.speed_bench import SpeedBenchDataset
|
|
|
|
dataset = SpeedBenchDataset.from_args(args)
|
|
rows = dataset.load(self.tokenizer)
|
|
# Only 2 low_entropy rows in the fixture, num_prompts=2
|
|
self.assertEqual(len(rows), 2)
|
|
self.assertTrue(all(isinstance(row, DatasetRow) for row in rows))
|
|
|
|
def test_speed_bench_output_len_override(self):
|
|
dataset_path = self._write_speed_bench_jsonl()
|
|
args = make_args(
|
|
dataset_name="speed-bench",
|
|
dataset_path=dataset_path,
|
|
num_prompts=2,
|
|
speed_bench_output_len=128,
|
|
)
|
|
from sglang.benchmark.datasets.speed_bench import SpeedBenchDataset
|
|
|
|
dataset = SpeedBenchDataset.from_args(args)
|
|
rows = dataset.load(self.tokenizer)
|
|
self.assertEqual(len(rows), 2)
|
|
self.assertTrue(all(row.output_len == 128 for row in rows))
|
|
|
|
def test_speed_bench_empty_category_raises(self):
|
|
dataset_path = self._write_speed_bench_jsonl()
|
|
args = make_args(
|
|
dataset_name="speed-bench",
|
|
dataset_path=dataset_path,
|
|
num_prompts=1,
|
|
speed_bench_category="nonexistent_category",
|
|
)
|
|
from sglang.benchmark.datasets.speed_bench import SpeedBenchDataset
|
|
|
|
dataset = SpeedBenchDataset.from_args(args)
|
|
with self.assertRaises(ValueError):
|
|
dataset.load(self.tokenizer)
|
|
|
|
def test_speed_bench_no_path_raises(self):
|
|
args = make_args(
|
|
dataset_name="speed-bench",
|
|
dataset_path="",
|
|
num_prompts=1,
|
|
)
|
|
from sglang.benchmark.datasets.speed_bench import SpeedBenchDataset
|
|
|
|
with self.assertRaises(ValueError):
|
|
SpeedBenchDataset.from_args(args)
|
|
|
|
def test_agentic_trace_sampler(self):
|
|
dataset_path = self._write_agentic_trace_json()
|
|
args = make_args(
|
|
dataset_name="agentic-trace",
|
|
dataset_path=dataset_path,
|
|
num_prompts=10,
|
|
)
|
|
dataset = AgenticTraceDataset.from_args(args)
|
|
rows = dataset.load(self.tokenizer)
|
|
self.assertEqual(len(rows), 2)
|
|
self.assertTrue(all(isinstance(row, DatasetRow) for row in rows))
|
|
self.assertTrue(
|
|
all(row.output_len == DEFAULT_AGENTIC_OUTPUT_LEN for row in rows)
|
|
)
|
|
# Multi-turn shape: prompt is a list of per-turn message lists, with
|
|
# the empty third turn of the first conversation dropped.
|
|
self.assertEqual(len(rows[0].prompt), 2)
|
|
self.assertEqual(len(rows[1].prompt), 1)
|
|
self.assertEqual(rows[0].prompt[0][0]["role"], "system")
|
|
self.assertEqual(rows[0].prompt_len, 100)
|
|
self.assertEqual(rows[1].prompt_len, 50)
|
|
|
|
def test_agentic_trace_offset_and_max_turns(self):
|
|
dataset_path = self._write_agentic_trace_json()
|
|
args = make_args(
|
|
dataset_name="agentic-trace",
|
|
dataset_path=dataset_path,
|
|
num_prompts=10,
|
|
sharegpt_output_len=64,
|
|
dataset_offset=1,
|
|
agentic_max_turns=1,
|
|
)
|
|
dataset = AgenticTraceDataset.from_args(args)
|
|
rows = dataset.load(self.tokenizer)
|
|
self.assertEqual(len(rows), 2)
|
|
# offset=1 rotates the second (single-turn) conversation to the front.
|
|
self.assertEqual(rows[0].prompt_len, 50)
|
|
self.assertTrue(all(len(row.prompt) == 1 for row in rows))
|
|
self.assertTrue(all(row.output_len == 64 for row in rows))
|
|
|
|
def test_agentic_trace_invalid_input_raises(self):
|
|
args = make_args(
|
|
dataset_name="agentic-trace",
|
|
dataset_path=str(self.tmpdir_path / "missing.json"),
|
|
num_prompts=1,
|
|
)
|
|
with self.assertRaises(FileNotFoundError):
|
|
AgenticTraceDataset.from_args(args).load(self.tokenizer)
|
|
|
|
empty_path = self.tmpdir_path / "empty_trace.json"
|
|
with open(empty_path, "w") as f:
|
|
json.dump({"metadata": {}, "conversations": []}, f)
|
|
args = make_args(
|
|
dataset_name="agentic-trace",
|
|
dataset_path=str(empty_path),
|
|
num_prompts=1,
|
|
)
|
|
with self.assertRaises(ValueError):
|
|
AgenticTraceDataset.from_args(args).load(self.tokenizer)
|
|
|
|
def test_dataset_mapping_and_dispatch(self):
|
|
expected = {
|
|
"agentic-trace",
|
|
"sharegpt",
|
|
"custom",
|
|
"openai",
|
|
"random",
|
|
"random-ids",
|
|
"generated-shared-prefix",
|
|
"mmmu",
|
|
"image",
|
|
"mooncake",
|
|
"speed-bench",
|
|
}
|
|
self.assertTrue(expected.issubset(set(DATASET_MAPPING.keys())))
|
|
|
|
sharegpt_path = self._write_sharegpt_json()
|
|
mooncake_path = self._write_mooncake_jsonl()
|
|
|
|
random_args = make_args(dataset_name="random-ids", tokenize_prompt=True)
|
|
random_rows = get_dataset(random_args, self.tokenizer, model_id="dummy-model")
|
|
self.assertEqual(len(random_rows), random_args.num_prompts)
|
|
self.assertTrue(all(isinstance(row.prompt, list) for row in random_rows))
|
|
|
|
sharegpt_args = make_args(dataset_name="sharegpt", dataset_path=sharegpt_path)
|
|
sharegpt_rows = get_dataset(
|
|
sharegpt_args, self.tokenizer, model_id="dummy-model"
|
|
)
|
|
self.assertEqual(len(sharegpt_rows), sharegpt_args.num_prompts)
|
|
|
|
mooncake_args = make_args(
|
|
dataset_name="mooncake",
|
|
dataset_path=mooncake_path,
|
|
num_prompts=1,
|
|
)
|
|
mooncake_rows = get_dataset(
|
|
mooncake_args, self.tokenizer, model_id="dummy-model"
|
|
)
|
|
self.assertEqual(len(mooncake_rows), 1)
|
|
self.assertIsInstance(mooncake_rows[0], dict)
|
|
|
|
with patch(
|
|
"sglang.benchmark.datasets.image.get_processor",
|
|
return_value=self.processor,
|
|
):
|
|
image_args = make_args(dataset_name="image")
|
|
image_rows = get_dataset(image_args, self.tokenizer, model_id="dummy-model")
|
|
self.assertEqual(len(image_rows), image_args.num_prompts)
|
|
|
|
fake_mmmu_dataset = _FakeMMMUDataset(
|
|
[{"image_1": Image.new("RGB", (4, 4), color="white"), "question": "q"}]
|
|
)
|
|
with (
|
|
patch(
|
|
"sglang.benchmark.datasets.mmmu.get_processor",
|
|
return_value=self.processor,
|
|
),
|
|
patch(
|
|
"sglang.benchmark.datasets.mmmu.load_dataset",
|
|
return_value=fake_mmmu_dataset,
|
|
),
|
|
):
|
|
mmmu_args = make_args(dataset_name="mmmu", num_prompts=1)
|
|
mmmu_rows = get_dataset(mmmu_args, self.tokenizer, model_id="dummy-model")
|
|
self.assertEqual(len(mmmu_rows), 1)
|
|
|
|
gsp_args = make_args(
|
|
dataset_name="generated-shared-prefix",
|
|
gsp_num_groups=2,
|
|
gsp_prompts_per_group=2,
|
|
)
|
|
gsp_rows = get_dataset(gsp_args, self.tokenizer, model_id="dummy-model")
|
|
self.assertEqual(len(gsp_rows), 4)
|
|
self.assertTrue(all(isinstance(row, DatasetRow) for row in gsp_rows))
|
|
|
|
speed_bench_path = self._write_speed_bench_jsonl()
|
|
speed_bench_args = make_args(
|
|
dataset_name="speed-bench",
|
|
dataset_path=speed_bench_path,
|
|
num_prompts=2,
|
|
)
|
|
speed_bench_rows = get_dataset(
|
|
speed_bench_args, self.tokenizer, model_id="dummy-model"
|
|
)
|
|
self.assertEqual(len(speed_bench_rows), 2)
|
|
self.assertTrue(all(isinstance(row, DatasetRow) for row in speed_bench_rows))
|
|
|
|
agentic_args = make_args(
|
|
dataset_name="agentic-trace",
|
|
dataset_path=self._write_agentic_trace_json(),
|
|
num_prompts=2,
|
|
)
|
|
agentic_rows = get_dataset(agentic_args, self.tokenizer, model_id="dummy-model")
|
|
self.assertEqual(len(agentic_rows), 2)
|
|
self.assertTrue(all(isinstance(row, DatasetRow) for row in agentic_rows))
|
|
|
|
def test_get_dataset_unknown_dataset(self):
|
|
args = make_args(dataset_name="not-a-dataset")
|
|
with self.assertRaises(ValueError):
|
|
get_dataset(args, self.tokenizer, model_id="dummy-model")
|
|
|
|
# ------------------------------------------------------------------
|
|
# Generated-shared-prefix Zipf sampling
|
|
# ------------------------------------------------------------------
|
|
|
|
def _run_gsp(
|
|
self,
|
|
*,
|
|
mode="uniform",
|
|
alpha=None,
|
|
seed=42,
|
|
num_groups=4,
|
|
prompts_per_group=5,
|
|
num_turns=1,
|
|
send_routing_key=False,
|
|
ordered=True,
|
|
range_ratio=1.0,
|
|
system_prompt_len=4,
|
|
question_len=3,
|
|
output_len=2,
|
|
fast_prepare=True,
|
|
global_seed=None,
|
|
):
|
|
# GSP's own `seed` kwarg only feeds the cache filename; reproducibility
|
|
# of compute_random_lens / gen_prompt comes from seeding the module
|
|
# globals before calling. Tests must seed both random and numpy here.
|
|
seed_for_globals = global_seed if global_seed is not None else seed
|
|
random.seed(seed_for_globals)
|
|
np.random.seed(seed_for_globals)
|
|
return sample_generated_shared_prefix_requests(
|
|
num_groups=num_groups,
|
|
prompts_per_group=prompts_per_group,
|
|
system_prompt_len=system_prompt_len,
|
|
question_len=question_len,
|
|
output_len=output_len,
|
|
range_ratio=range_ratio,
|
|
tokenizer=self.tokenizer,
|
|
seed=seed,
|
|
send_routing_key=send_routing_key,
|
|
num_turns=num_turns,
|
|
fast_prepare=fast_prepare,
|
|
ordered=ordered,
|
|
group_distribution=mode,
|
|
zipf_alpha=alpha,
|
|
)
|
|
|
|
@staticmethod
|
|
def _row_fields(rows):
|
|
return [(r.prompt, r.prompt_len, r.output_len, r.routing_key) for r in rows]
|
|
|
|
def test_gsp_uniform_default_unchanged(self):
|
|
# Uniform mode returns the documented number of rows and is
|
|
# bit-reproducible under fixed seeding of the global RNGs.
|
|
rows_a = self._run_gsp(
|
|
mode="uniform", num_groups=3, prompts_per_group=4, seed=7
|
|
)
|
|
rows_b = self._run_gsp(
|
|
mode="uniform", num_groups=3, prompts_per_group=4, seed=7
|
|
)
|
|
self.assertEqual(len(rows_a), 3 * 4)
|
|
self.assertEqual(self._row_fields(rows_a), self._row_fields(rows_b))
|
|
|
|
def test_gsp_uniform_cache_path_format_unchanged(self):
|
|
# The uniform-mode cache filename keeps its existing
|
|
# gen_shared_prefix_<seed>_<N>_<P>_<sysL>_<qL>_<outL>_<TokenizerCls>.pkl
|
|
# shape. The trailing class name is a transformers/tokenizers internal
|
|
# detail (TokenizersBackend / PreTrainedTokenizerFast depending on
|
|
# version), so we only pin the deterministic numeric portion.
|
|
path = get_gen_prefix_cache_path(
|
|
seed=7,
|
|
num_groups=3,
|
|
prompts_per_group=4,
|
|
system_prompt_len=16,
|
|
question_len=8,
|
|
output_len=4,
|
|
tokenizer=self.tokenizer,
|
|
)
|
|
self.assertTrue(path.name.startswith("gen_shared_prefix_7_3_4_16_8_4_"))
|
|
self.assertTrue(path.name.endswith(".pkl"))
|
|
self.assertEqual(path.parent, Path.home() / ".cache" / "sglang" / "benchmark")
|
|
|
|
def test_zipf_group_probs_helper(self):
|
|
# Rank-based probability vector: weight(rank) = 1 / rank ** alpha,
|
|
# normalized to sum to 1, with rank starting at 1.
|
|
probs_n3_a1 = _zipf_group_probs(3, 1.0)
|
|
expected_n3_a1 = np.array([6.0, 3.0, 2.0]) / 11.0
|
|
np.testing.assert_allclose(probs_n3_a1, expected_n3_a1, atol=1e-12)
|
|
self.assertAlmostEqual(float(probs_n3_a1.sum()), 1.0, places=12)
|
|
|
|
probs_n4_a15 = _zipf_group_probs(4, 1.5)
|
|
ranks = np.arange(1, 5, dtype=np.float64)
|
|
ref = 1.0 / ranks**1.5
|
|
ref = ref / ref.sum()
|
|
np.testing.assert_allclose(probs_n4_a15, ref, atol=1e-12)
|
|
# Three-decimal pin against a hand-computable reference.
|
|
np.testing.assert_allclose(
|
|
np.round(probs_n4_a15, 3),
|
|
np.array([0.598, 0.212, 0.115, 0.075]),
|
|
atol=1e-3,
|
|
)
|
|
|
|
def test_zipf_group_probs_not_lora_skewed_formula(self):
|
|
# The helper must NOT use the LoRA `skewed` alpha**-i exponential
|
|
# formula; for alpha=1.5, N=4 the two formulas differ noticeably.
|
|
actual = _zipf_group_probs(4, 1.5)
|
|
lora_weights = np.array([1.5**-i for i in range(4)], dtype=np.float64)
|
|
lora_probs = lora_weights / lora_weights.sum()
|
|
self.assertFalse(
|
|
np.allclose(actual, lora_probs, atol=1e-3),
|
|
"Zipf helper must use rank-based 1/rank**alpha, not LoRA alpha**-i",
|
|
)
|
|
|
|
def test_zipf_reproducible_with_seed(self):
|
|
# Same seed + same args -> identical rows, including order, under
|
|
# both the in-order and shuffled paths.
|
|
kwargs = dict(
|
|
mode="zipf", alpha=1.7, seed=11, num_groups=4, prompts_per_group=10
|
|
)
|
|
rows_a = self._run_gsp(**kwargs)
|
|
rows_b = self._run_gsp(**kwargs)
|
|
self.assertEqual(len(rows_a), 4 * 10)
|
|
self.assertEqual(self._row_fields(rows_a), self._row_fields(rows_b))
|
|
|
|
# Also under the shuffled path.
|
|
rows_c = self._run_gsp(ordered=False, **kwargs)
|
|
rows_d = self._run_gsp(ordered=False, **kwargs)
|
|
self.assertEqual(self._row_fields(rows_c), self._row_fields(rows_d))
|
|
|
|
def test_zipf_different_seeds_differ(self):
|
|
# Different seeds -> at least one differing slot under Zipf sampling.
|
|
base = dict(mode="zipf", alpha=1.7, num_groups=4, prompts_per_group=10)
|
|
rows_a = self._run_gsp(seed=11, **base)
|
|
rows_b = self._run_gsp(seed=12, **base)
|
|
self.assertEqual(len(rows_a), len(rows_b))
|
|
self.assertNotEqual(self._row_fields(rows_a), self._row_fields(rows_b))
|
|
|
|
def test_zipf_does_not_perturb_global_random_state(self):
|
|
# The Zipf branch must consume zero draws from the global random /
|
|
# numpy.random state. Therefore the per-slot generated questions and
|
|
# system prompts under uniform and Zipf modes for the same args and
|
|
# the same global seed are byte-equal.
|
|
common = dict(
|
|
num_groups=4,
|
|
prompts_per_group=6,
|
|
system_prompt_len=4,
|
|
question_len=3,
|
|
output_len=2,
|
|
range_ratio=1.0,
|
|
seed=99,
|
|
ordered=True,
|
|
send_routing_key=False,
|
|
fast_prepare=True,
|
|
global_seed=99,
|
|
)
|
|
uniform_rows = self._run_gsp(mode="uniform", **common)
|
|
zipf_rows = self._run_gsp(mode="zipf", alpha=1.3, **common)
|
|
|
|
# Slot i in uniform mode pairs system_prompts[i // P] with
|
|
# questions[i // P][i % P], so the question substring after the
|
|
# delimiter is exactly the i-th question. Same construction is used by
|
|
# the Zipf branch (only the system prompt changes per slot), so the
|
|
# question substrings must match slot-by-slot under the same global
|
|
# seed.
|
|
delim = "\n\n"
|
|
|
|
def question_of(prompt):
|
|
return prompt.split(delim, 1)[1]
|
|
|
|
uniform_questions = [question_of(r.prompt) for r in uniform_rows]
|
|
zipf_questions = [question_of(r.prompt) for r in zipf_rows]
|
|
self.assertEqual(uniform_questions, zipf_questions)
|
|
|
|
# The set of system prompts (which the gen_prompt path generates) must
|
|
# also match between modes (set equality, since Zipf reuses prefixes).
|
|
def system_of(prompt):
|
|
return prompt.split(delim, 1)[0]
|
|
|
|
self.assertEqual(
|
|
set(system_of(r.prompt) for r in uniform_rows),
|
|
set(system_of(r.prompt) for r in zipf_rows),
|
|
)
|
|
|
|
def test_zipf_deterministic_per_group_counts(self):
|
|
# The per-group counts are deterministic and pinned for a known
|
|
# (num_groups, prompts_per_group, alpha, seed) tuple. Any drift in
|
|
# the Zipf sampling implementation will trip this assertion.
|
|
rows = self._run_gsp(
|
|
mode="zipf",
|
|
alpha=2.0,
|
|
seed=0,
|
|
num_groups=4,
|
|
prompts_per_group=25,
|
|
send_routing_key=True,
|
|
ordered=True,
|
|
)
|
|
self.assertEqual(len(rows), 4 * 25)
|
|
# routing_key format is "<uuid8>_<timestamp>_<group_idx>".
|
|
per_group = Counter(int(r.routing_key.rsplit("_", 1)[-1]) for r in rows)
|
|
# Pinned counts derived from the implementation for
|
|
# (N=4, P=25, alpha=2.0, seed=0) using numpy.random.default_rng(seed)
|
|
# and rng.choice over _zipf_group_probs(N, alpha).
|
|
self.assertEqual(
|
|
dict(per_group),
|
|
{0: 63, 1: 18, 2: 12, 3: 7},
|
|
)
|
|
# Independent skew sanity check: rank-1 (hottest) > rank-N (coldest).
|
|
self.assertGreater(per_group[0], per_group[3])
|
|
|
|
def test_zipf_uses_distinct_cache_from_uniform(self):
|
|
# The on-disk cache key includes group_distribution and zipf_alpha,
|
|
# so uniform mode, zipf alpha=1.0, and zipf alpha=2.0 each get their
|
|
# own file. Uniform mode never reads a zipf cache and vice versa.
|
|
from sglang.benchmark.datasets import generated_shared_prefix as gsp_mod
|
|
|
|
fake_home = self.tmpdir_path / "fakehome"
|
|
fake_home.mkdir()
|
|
|
|
common = dict(
|
|
num_groups=2,
|
|
prompts_per_group=3,
|
|
system_prompt_len=4,
|
|
question_len=3,
|
|
output_len=2,
|
|
range_ratio=1.0,
|
|
seed=5,
|
|
send_routing_key=False,
|
|
num_turns=1,
|
|
fast_prepare=True,
|
|
ordered=True,
|
|
)
|
|
|
|
with patch.object(gsp_mod.Path, "home", return_value=fake_home):
|
|
uniform_path = get_gen_prefix_cache_path(
|
|
seed=common["seed"],
|
|
num_groups=common["num_groups"],
|
|
prompts_per_group=common["prompts_per_group"],
|
|
system_prompt_len=common["system_prompt_len"],
|
|
question_len=common["question_len"],
|
|
output_len=common["output_len"],
|
|
tokenizer=self.tokenizer,
|
|
)
|
|
zipf_path_a = get_gen_prefix_cache_path(
|
|
seed=common["seed"],
|
|
num_groups=common["num_groups"],
|
|
prompts_per_group=common["prompts_per_group"],
|
|
system_prompt_len=common["system_prompt_len"],
|
|
question_len=common["question_len"],
|
|
output_len=common["output_len"],
|
|
tokenizer=self.tokenizer,
|
|
group_distribution="zipf",
|
|
zipf_alpha=1.5,
|
|
)
|
|
zipf_path_b = get_gen_prefix_cache_path(
|
|
seed=common["seed"],
|
|
num_groups=common["num_groups"],
|
|
prompts_per_group=common["prompts_per_group"],
|
|
system_prompt_len=common["system_prompt_len"],
|
|
question_len=common["question_len"],
|
|
output_len=common["output_len"],
|
|
tokenizer=self.tokenizer,
|
|
group_distribution="zipf",
|
|
zipf_alpha=2.0,
|
|
)
|
|
self.assertNotEqual(uniform_path, zipf_path_a)
|
|
self.assertNotEqual(zipf_path_a, zipf_path_b)
|
|
|
|
# Run each mode; each writes its own cache file.
|
|
self._run_gsp(mode="uniform", **common)
|
|
self._run_gsp(mode="zipf", alpha=1.5, **common)
|
|
self._run_gsp(mode="zipf", alpha=2.0, **common)
|
|
self.assertTrue(uniform_path.exists())
|
|
self.assertTrue(zipf_path_a.exists())
|
|
self.assertTrue(zipf_path_b.exists())
|
|
|
|
# Sentinel into the uniform cache: zipf must not read it.
|
|
sentinel = [DatasetRow(prompt="SENTINEL", prompt_len=1, output_len=1)]
|
|
with open(uniform_path, "wb") as f:
|
|
pickle.dump(sentinel, f)
|
|
zipf_rows = self._run_gsp(mode="zipf", alpha=1.5, **common)
|
|
self.assertNotEqual(zipf_rows, sentinel)
|
|
|
|
# Second zipf call with same args must load from cache (no
|
|
# regeneration). Mutate the zipf cache to a sentinel and confirm.
|
|
zipf_sentinel = [
|
|
DatasetRow(prompt="ZIPF_SENTINEL", prompt_len=1, output_len=1)
|
|
]
|
|
with open(zipf_path_a, "wb") as f:
|
|
pickle.dump(zipf_sentinel, f)
|
|
reloaded = self._run_gsp(mode="zipf", alpha=1.5, **common)
|
|
self.assertEqual(reloaded, zipf_sentinel)
|
|
|
|
def test_zipf_total_rows_and_unique_prompts(self):
|
|
# Total returned row count under Zipf equals num_groups *
|
|
# prompts_per_group (identical to uniform mode) and every prompt
|
|
# string is unique even when groups repeat.
|
|
rows = self._run_gsp(
|
|
mode="zipf",
|
|
alpha=2.5,
|
|
seed=3,
|
|
num_groups=4,
|
|
prompts_per_group=10,
|
|
send_routing_key=False,
|
|
)
|
|
self.assertEqual(len(rows), 4 * 10)
|
|
self.assertEqual(len({r.prompt for r in rows}), len(rows))
|
|
|
|
def test_zipf_ordered_preserves_generation_order(self):
|
|
# With ordered=True, output preserves the sampled order and matches
|
|
# an independently re-derived group sequence from default_rng(seed).
|
|
rows = self._run_gsp(
|
|
mode="zipf",
|
|
alpha=1.5,
|
|
seed=21,
|
|
num_groups=3,
|
|
prompts_per_group=8,
|
|
send_routing_key=True,
|
|
ordered=True,
|
|
)
|
|
observed_groups = [int(r.routing_key.rsplit("_", 1)[-1]) for r in rows]
|
|
|
|
# Independently reproduce the expected group sequence: an isolated
|
|
# default_rng(seed) over _zipf_group_probs(N, alpha) sampling
|
|
# N * P slots.
|
|
expected_rng = np.random.default_rng(21)
|
|
expected_probs = _zipf_group_probs(3, 1.5)
|
|
expected_groups = expected_rng.choice(
|
|
3, size=3 * 8, replace=True, p=expected_probs
|
|
).tolist()
|
|
self.assertEqual(observed_groups, expected_groups)
|
|
|
|
def test_zipf_shuffle_path_matches_uniform_shuffle(self):
|
|
# When ordered=False, both modes go through random.shuffle on a list
|
|
# of equal length, so the same global RNG seed yields the same
|
|
# permutation. Verified indirectly: two Zipf calls with the same
|
|
# global seed produce identical orderings.
|
|
kwargs = dict(
|
|
mode="zipf",
|
|
alpha=1.2,
|
|
seed=8,
|
|
num_groups=4,
|
|
prompts_per_group=6,
|
|
send_routing_key=False,
|
|
ordered=False,
|
|
global_seed=8,
|
|
)
|
|
rows_a = self._run_gsp(**kwargs)
|
|
rows_b = self._run_gsp(**kwargs)
|
|
self.assertEqual(self._row_fields(rows_a), self._row_fields(rows_b))
|
|
|
|
# ------------------------------------------------------------------
|
|
# CLI / from_args validation
|
|
# ------------------------------------------------------------------
|
|
|
|
def test_from_args_rejects_invalid_distribution_and_alpha(self):
|
|
# Defensive validation in from_args protects in-process callers
|
|
# that build a Namespace by hand and bypass the argparse boundary
|
|
# in bench_serving.py. Covers: unknown distribution, zipf without
|
|
# alpha, uniform with alpha, and non-finite/non-positive alpha.
|
|
cases = [
|
|
{"gsp_group_distribution": "not-a-distribution", "gsp_zipf_alpha": None},
|
|
{"gsp_group_distribution": "zipf", "gsp_zipf_alpha": None},
|
|
{"gsp_group_distribution": "uniform", "gsp_zipf_alpha": 1.0},
|
|
{"gsp_group_distribution": "zipf", "gsp_zipf_alpha": 0.0},
|
|
{"gsp_group_distribution": "zipf", "gsp_zipf_alpha": -0.5},
|
|
{"gsp_group_distribution": "zipf", "gsp_zipf_alpha": float("nan")},
|
|
{"gsp_group_distribution": "zipf", "gsp_zipf_alpha": float("inf")},
|
|
{"gsp_group_distribution": "zipf", "gsp_zipf_alpha": float("-inf")},
|
|
]
|
|
for case in cases:
|
|
args = make_args(dataset_name="generated-shared-prefix", **case)
|
|
with self.assertRaises(ValueError, msg=f"case={case}"):
|
|
GeneratedSharedPrefixDataset.from_args(args)
|
|
|
|
def test_bench_serving_help_and_invalid_choice_argparse(self):
|
|
# Subprocess-driven coverage of the live CLI: --help advertises both
|
|
# flags with the rank-based Zipf formula and the alpha constraint,
|
|
# and argparse rejects an unknown distribution choice.
|
|
help_res = subprocess.run(
|
|
[sys.executable, "-m", "sglang.benchmark.serving", "--help"],
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=90,
|
|
)
|
|
self.assertEqual(help_res.returncode, 0, help_res.stderr)
|
|
out = help_res.stdout
|
|
# Both new flags appear.
|
|
self.assertIn("--gsp-group-distribution", out)
|
|
self.assertIn("--gsp-zipf-alpha", out)
|
|
# Rank-based Zipf formula and alpha constraint are documented.
|
|
self.assertIn("1/rank**alpha", out)
|
|
self.assertIn("rank starts at 1", out)
|
|
self.assertIn("finite float", out)
|
|
|
|
# Argparse rejects unknown distribution choice.
|
|
bad_choice_res = subprocess.run(
|
|
[
|
|
sys.executable,
|
|
"-m",
|
|
"sglang.benchmark.serving",
|
|
"--dataset-name",
|
|
"generated-shared-prefix",
|
|
"--gsp-group-distribution",
|
|
"invalid_name",
|
|
],
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=90,
|
|
)
|
|
self.assertNotEqual(bad_choice_res.returncode, 0)
|
|
self.assertIn("invalid choice", (bad_choice_res.stderr + bad_choice_res.stdout))
|
|
|
|
def test_bench_serving_cli_rejects_zipf_without_alpha_before_server(self):
|
|
# Malformed CLI combinations (zipf with no alpha) must fail at
|
|
# argparse time so users see the GSP-flag error directly, not a
|
|
# downstream connection or model-fetch failure.
|
|
res = subprocess.run(
|
|
[
|
|
sys.executable,
|
|
"-m",
|
|
"sglang.benchmark.serving",
|
|
"--dataset-name",
|
|
"generated-shared-prefix",
|
|
"--gsp-group-distribution",
|
|
"zipf",
|
|
"--ready-check-timeout-sec",
|
|
"0",
|
|
],
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=90,
|
|
)
|
|
# parser.error() exits with code 2 (argparse convention).
|
|
self.assertEqual(res.returncode, 2, res.stderr)
|
|
stderr = res.stderr + res.stdout
|
|
self.assertIn("--gsp-group-distribution", stderr)
|
|
self.assertIn("--gsp-zipf-alpha", stderr)
|
|
# The error must mention the GSP flags directly, not a network or
|
|
# model-discovery problem masquerading as the failure.
|
|
for forbidden in [
|
|
"HTTPConnectionPool",
|
|
"HTTPSConnectionPool",
|
|
"Connection refused",
|
|
"Failed to fetch model",
|
|
"Traceback",
|
|
]:
|
|
self.assertNotIn(forbidden, stderr)
|
|
|
|
def test_bench_serving_cli_rejects_uniform_with_alpha_before_server(self):
|
|
# The complementary malformation: uniform distribution with an
|
|
# explicit alpha value. Must also fail at argparse time.
|
|
res = subprocess.run(
|
|
[
|
|
sys.executable,
|
|
"-m",
|
|
"sglang.benchmark.serving",
|
|
"--dataset-name",
|
|
"generated-shared-prefix",
|
|
"--gsp-group-distribution",
|
|
"uniform",
|
|
"--gsp-zipf-alpha",
|
|
"1.0",
|
|
"--ready-check-timeout-sec",
|
|
"0",
|
|
],
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=90,
|
|
)
|
|
self.assertEqual(res.returncode, 2, res.stderr)
|
|
stderr = res.stderr + res.stdout
|
|
self.assertIn("--gsp-group-distribution", stderr)
|
|
self.assertIn("--gsp-zipf-alpha", stderr)
|
|
for forbidden in [
|
|
"HTTPConnectionPool",
|
|
"HTTPSConnectionPool",
|
|
"Connection refused",
|
|
"Failed to fetch model",
|
|
"Traceback",
|
|
]:
|
|
self.assertNotIn(forbidden, stderr)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|