From f838adb7d497f46ba8b04fdd846a501b9123454a Mon Sep 17 00:00:00 2001 From: Jimmy Shong <69131491+Jiminator@users.noreply.github.com> Date: Thu, 28 May 2026 14:39:46 -0700 Subject: [PATCH] bench_serving: add Zipfian shared-prefix sampling to generated-shared-prefix (#26378) --- .../docs/developer_guide/bench_serving.mdx | 22 + python/sglang/bench_serving.py | 71 +++ .../datasets/generated_shared_prefix.py | 167 ++++-- .../bench_fn/test_benchmark_datasets_api.py | 504 +++++++++++++++++- 4 files changed, 728 insertions(+), 36 deletions(-) diff --git a/docs_new/docs/developer_guide/bench_serving.mdx b/docs_new/docs/developer_guide/bench_serving.mdx index 4513572b4..e0dd504f7 100644 --- a/docs_new/docs/developer_guide/bench_serving.mdx +++ b/docs_new/docs/developer_guide/bench_serving.mdx @@ -83,6 +83,8 @@ Generated Shared Prefix flags (for `generated-shared-prefix`): - `--gsp-system-prompt-len` - `--gsp-question-len` - `--gsp-output-len` +- `--gsp-group-distribution {uniform,zipf}`: per-request prefix-group sampling distribution (default: `uniform`). With `zipf`, each request's group is sampled by rank with `p(rank) = (1/rank**alpha) / sum_k(1/k**alpha)`; rank starts at 1 and group index 0 is the hottest. The on-disk dataset cache uses a distinct key per `(group_distribution, zipf_alpha)`, so uniform-mode caches are never mixed with zipf-mode caches. +- `--gsp-zipf-alpha FLOAT`: Zipf exponent for `--gsp-group-distribution=zipf`. Must be a finite float strictly greater than 0; larger values concentrate requests on lower-ranked (hotter) groups. Required when the distribution is `zipf`; must be omitted otherwise. Image dataset flags (for `image`): @@ -318,6 +320,26 @@ python3 -m sglang.bench_serving \ --num-prompts 1024 ``` +Zipfian / power-law prefix popularity (opt-in via `--gsp-group-distribution=zipf`): + +```bash Command +python3 -m sglang.bench_serving \ + --backend sglang \ + --host 127.0.0.1 --port 30000 \ + --model meta-llama/Llama-3.1-8B-Instruct \ + --dataset-name generated-shared-prefix \ + --gsp-num-groups 64 --gsp-prompts-per-group 16 \ + --gsp-system-prompt-len 2048 --gsp-question-len 128 --gsp-output-len 256 \ + --gsp-group-distribution zipf --gsp-zipf-alpha 1.2 \ + --seed 42 +``` + +`zipf` mode samples each request's prefix group from the rank-based distribution `p(rank) = (1/rank**alpha) / sum_k(1/k**alpha)` with rank starting at 1, so group index 0 is the hottest. The total request count stays `num_groups * prompts_per_group` — identical to `uniform` mode — and only the per-request group assignment changes. `alpha` must be a finite float strictly greater than 0; larger values concentrate requests on lower-ranked (hotter) groups. + +The on-disk dataset cache at `~/.cache/sglang/benchmark/gen_shared_prefix_*.pkl` includes `group_distribution` and `zipf_alpha` in its key, so uniform-mode and zipf-mode runs (or two zipf runs with different alpha) never share a cache file. Uniform-mode filenames are unchanged from the legacy format, so existing caches remain valid. + +This flag controls prefix-popularity shape only. It does not by itself reproduce any production trace or guarantee an observed cache-hit rate for a given engine. + 6) Tokenized prompts (ids) for strict length control (sglang only): ```bash Command diff --git a/python/sglang/bench_serving.py b/python/sglang/bench_serving.py index 53e4ce18f..ea3534fa6 100644 --- a/python/sglang/bench_serving.py +++ b/python/sglang/bench_serving.py @@ -17,6 +17,7 @@ import asyncio import copy import importlib.util import json +import math import os import random import shutil @@ -1930,6 +1931,44 @@ def run_benchmark(args_: argparse.Namespace): ) +def _finite_positive_float(value) -> float: + """argparse type for a finite, strictly positive float.""" + try: + parsed = float(value) + except (TypeError, ValueError) as exc: + raise argparse.ArgumentTypeError( + f"expected a finite float > 0, got {value!r}" + ) from exc + if not math.isfinite(parsed) or parsed <= 0: + raise argparse.ArgumentTypeError(f"expected a finite float > 0, got {value!r}") + return parsed + + +def _validate_parsed_gsp_args( + parser: argparse.ArgumentParser, args: argparse.Namespace +) -> None: + """Reject malformed GSP distribution/alpha combinations at parse time. + + Invoked from the CLI entry point right after ``parser.parse_args()`` so + users see a clear argparse-style error before any server, model, or + tokenizer setup runs and masks the real cause with an unrelated network + failure. + """ + distribution = getattr(args, "gsp_group_distribution", None) + alpha = getattr(args, "gsp_zipf_alpha", None) + if distribution == "zipf" and alpha is None: + parser.error( + "--gsp-group-distribution=zipf requires --gsp-zipf-alpha " + "(a finite float > 0)" + ) + if distribution == "uniform" and alpha is not None: + parser.error( + "--gsp-zipf-alpha is only meaningful with " + "--gsp-group-distribution=zipf; remove --gsp-zipf-alpha " + "or set --gsp-group-distribution=zipf" + ) + + class LoRAPathAction(argparse.Action): def __call__(self, parser, namespace, values, option_string=None): setattr(namespace, self.dest, []) @@ -2366,6 +2405,37 @@ if __name__ == "__main__": action="store_true", help="Keep requests in order without shuffling. By default, requests are shuffled randomly.", ) + group.add_argument( + "--gsp-group-distribution", + type=str, + choices=["uniform", "zipf"], + default="uniform", + help=( + "Prefix-group sampling distribution for generated-shared-prefix. " + "'uniform' (default) assigns each group an equal number of requests. " + "'zipf' samples each request's group by rank with " + "p(rank) = (1/rank**alpha) / sum_k(1/k**alpha); rank starts at 1 " + "and group index 0 is the hottest. Requires --gsp-zipf-alpha " + "(a finite float > 0) when set to 'zipf'. Total request count is " + "still num_groups * prompts_per_group, identical to uniform mode; " + "only the per-request group assignment changes. The on-disk " + "dataset cache uses a distinct key per (group_distribution, " + "zipf_alpha), so uniform-mode caches are never mixed with " + "zipf-mode caches and zipf runs with different alpha use " + "separate files." + ), + ) + group.add_argument( + "--gsp-zipf-alpha", + type=_finite_positive_float, + default=None, + help=( + "Zipf exponent alpha for --gsp-group-distribution=zipf, with " + "p(rank) = (1/rank**alpha) / sum_k(1/k**alpha) and rank starting " + "at 1. Must be a finite float strictly greater than 0; larger " + "values concentrate requests on lower-ranked (hotter) groups." + ), + ) mooncake_group = parser.add_argument_group("mooncake dataset arguments") mooncake_group.add_argument( "--mooncake-slowdown-factor", @@ -2413,4 +2483,5 @@ if __name__ == "__main__": help="Custom HTTP headers in Key=Value format. Example: --header MyHeader=MY_VALUE MyAnotherHeader=myanothervalue", ) args = parser.parse_args() + _validate_parsed_gsp_args(parser, args) run_benchmark(args) diff --git a/python/sglang/benchmark/datasets/generated_shared_prefix.py b/python/sglang/benchmark/datasets/generated_shared_prefix.py index 51c4e18ae..c739b5533 100644 --- a/python/sglang/benchmark/datasets/generated_shared_prefix.py +++ b/python/sglang/benchmark/datasets/generated_shared_prefix.py @@ -1,3 +1,4 @@ +import math import pickle import random import uuid @@ -5,7 +6,7 @@ from argparse import Namespace from dataclasses import dataclass from datetime import datetime from pathlib import Path -from typing import List +from typing import List, Optional import numpy as np from tqdm.asyncio import tqdm @@ -19,6 +20,22 @@ from sglang.benchmark.datasets.common import ( ) +def _zipf_group_probs(num_groups: int, alpha: float) -> np.ndarray: + """Rank-based Zipf probability vector with rank starting at 1. + + weight(rank) = 1 / rank ** alpha (rank in 1..num_groups) + probability(rank) = weight(rank) / sum_over_all_ranks(weight) + + The returned array has length num_groups; element i corresponds to + group index i (rank i + 1), so group 0 is the hottest. + """ + if num_groups <= 0: + raise ValueError(f"num_groups must be > 0, got {num_groups}") + ranks = np.arange(1, num_groups + 1, dtype=np.float64) + weights = 1.0 / (ranks**alpha) + return weights / weights.sum() + + @dataclass class GeneratedSharedPrefixDataset(BaseDataset): num_groups: int @@ -32,10 +49,40 @@ class GeneratedSharedPrefixDataset(BaseDataset): send_routing_key: bool num_turns: int ordered: bool + group_distribution: str = "uniform" + zipf_alpha: Optional[float] = None @classmethod def from_args(cls, args: Namespace) -> "GeneratedSharedPrefixDataset": assert not getattr(args, "tokenize_prompt", False) + group_distribution = args.gsp_group_distribution + zipf_alpha = args.gsp_zipf_alpha + + # Defensive validation for in-process callers that construct a + # Namespace by hand and bypass the argparse boundary in + # bench_serving.py. The CLI hook enforces the same rules first. + if group_distribution not in ("uniform", "zipf"): + raise ValueError( + f"--gsp-group-distribution must be 'uniform' or 'zipf', " + f"got {group_distribution!r}" + ) + if group_distribution == "zipf": + if zipf_alpha is None: + raise ValueError( + "--gsp-group-distribution=zipf requires --gsp-zipf-alpha " + "(a finite float > 0)" + ) + if not math.isfinite(zipf_alpha) or zipf_alpha <= 0: + raise ValueError( + f"--gsp-zipf-alpha must be a finite float > 0, got {zipf_alpha!r}" + ) + elif zipf_alpha is not None: + raise ValueError( + "--gsp-zipf-alpha is only meaningful with " + "--gsp-group-distribution=zipf; remove --gsp-zipf-alpha " + "or set --gsp-group-distribution=zipf" + ) + return cls( num_groups=args.gsp_num_groups, prompts_per_group=args.gsp_prompts_per_group, @@ -48,6 +95,8 @@ class GeneratedSharedPrefixDataset(BaseDataset): send_routing_key=getattr(args, "gsp_send_routing_key", False), num_turns=getattr(args, "gsp_num_turns", 1), ordered=getattr(args, "gsp_ordered", False), + group_distribution=group_distribution, + zipf_alpha=zipf_alpha, ) def load( @@ -66,6 +115,8 @@ class GeneratedSharedPrefixDataset(BaseDataset): num_turns=self.num_turns, fast_prepare=self.fast_prepare, ordered=self.ordered, + group_distribution=self.group_distribution, + zipf_alpha=self.zipf_alpha, ) @@ -77,13 +128,24 @@ def get_gen_prefix_cache_path( question_len: int, output_len: int, tokenizer, + group_distribution: str = "uniform", + zipf_alpha: Optional[float] = None, ): - """Create cache directory under ~/.cache/sglang/benchmark""" + """Create cache directory under ~/.cache/sglang/benchmark. + + The uniform-mode filename is preserved exactly as before so existing + on-disk caches remain valid. Non-default sampling modes get an extra + suffix encoding the parameters that affect the cached payload. + """ cache_dir = Path.home() / ".cache" / "sglang" / "benchmark" + suffix = "" + if group_distribution != "uniform": + suffix = f"_{group_distribution}_{zipf_alpha}" + cache_key = ( f"gen_shared_prefix_{seed}_{num_groups}_{prompts_per_group}_" - f"{system_prompt_len}_{question_len}_{output_len}_" + f"{system_prompt_len}_{question_len}_{output_len}{suffix}_" f"{tokenizer.__class__.__name__}.pkl" ) return cache_dir / cache_key @@ -102,8 +164,22 @@ def sample_generated_shared_prefix_requests( num_turns: int = 1, fast_prepare: bool = False, ordered: bool = False, + group_distribution: str = "uniform", + zipf_alpha: Optional[float] = None, ) -> List[DatasetRow]: - """Generate benchmark requests with shared system prompts using random tokens and caching.""" + """Generate benchmark requests with shared system prompts using random tokens and caching. + + When group_distribution is "uniform" (default), each group receives exactly + prompts_per_group requests; behavior matches the legacy generator. + + When group_distribution is "zipf", each request's group is sampled by rank + with probability 1/rank**zipf_alpha / sum_k(1/k**zipf_alpha); rank starts at + 1 and group index 0 is the hottest. Sampling uses an isolated + numpy.random.default_rng(seed) so the shared question/system-prompt pool + stays byte-identical to uniform mode for the same seed and other args. + Zipf mode is cached on disk under a distinct key per (group_distribution, + zipf_alpha) value. + """ cache_path = get_gen_prefix_cache_path( seed, num_groups, @@ -112,18 +188,25 @@ def sample_generated_shared_prefix_requests( question_len, output_len, tokenizer, + group_distribution=group_distribution, + zipf_alpha=zipf_alpha, ) - should_cache = (range_ratio == 1) and not send_routing_key and num_turns == 1 + # range_ratio != 1 / num_turns > 1 perturb the payload but are not in the + # cache key; send_routing_key embeds a per-run uuid + timestamp that is + # meaningless to cache. Bypass for these pre-existing reasons only. + should_cache = range_ratio == 1 and not send_routing_key and num_turns == 1 - # Try to load from cache first - if cache_path.exists() and should_cache: + if should_cache and cache_path.exists(): print(f"\nLoading cached generated input data from {cache_path}") with open(cache_path, "rb") as f: return pickle.load(f) + if not should_cache: + print(f"\nCache bypassed ({range_ratio=}, {send_routing_key=}, {num_turns=})") + print( f"\nGenerating new input data... " - f"({num_groups=}, {prompts_per_group}, {system_prompt_len=}, {question_len=}, {output_len=}, {range_ratio=}, {num_turns=})" + f"({num_groups=}, {prompts_per_group}, {system_prompt_len=}, {question_len=}, {output_len=}, {range_ratio=}, {num_turns=}, {group_distribution=}, {zipf_alpha=})" ) run_random_str = uuid.uuid4().hex[:8] @@ -150,12 +233,11 @@ def sample_generated_shared_prefix_requests( ).reshape(num_groups, prompts_per_group) del system_prompt_len, question_len, output_len - # Generate system prompts for each group system_prompts = [ gen_prompt(tokenizer, system_prompt_lens[i]) for i in range(num_groups) ] - # Generate questions: shape (num_groups, prompts_per_group, num_turns) + # shape: (num_groups, prompts_per_group, num_turns) questions = [ [ [ @@ -167,48 +249,64 @@ def sample_generated_shared_prefix_requests( for g in range(num_groups) ] - # Combine system prompts with questions + # Per-slot group assignment. Uniform mode is the identity assignment + # [0,0,...,1,1,...,N-1,N-1]; zipf mode samples from the rank distribution + # using an isolated RNG so the module-level random / numpy.random state + # that compute_random_lens / gen_prompt rely on is never perturbed -- this + # keeps the system-prompt and question pool byte-identical to uniform mode + # for the same seed and other args. + total_slots = num_groups * prompts_per_group + if group_distribution == "uniform": + assignment = np.repeat(np.arange(num_groups), prompts_per_group) + else: # "zipf" + rng = np.random.default_rng(seed) + probs = _zipf_group_probs(num_groups, zipf_alpha) + assignment = rng.choice(num_groups, size=total_slots, replace=True, p=probs) + input_requests = [] total_input_tokens = 0 total_output_tokens = 0 + for slot_idx, sampled_g in enumerate( + tqdm(assignment, desc="Generating shared-prefix prompts") + ): + # src_(g,p) walks the question pool in uniform-enumeration order, so + # per-slot question text is reproducibly identical across modes. + src_g, src_p = divmod(slot_idx, prompts_per_group) + sampled_g = int(sampled_g) - for group_idx in tqdm(range(num_groups), desc="Generating system prompt"): - system_prompt = system_prompts[group_idx] + system_prompt = system_prompts[sampled_g] routing_key = ( - f"{run_random_str}_{run_start_timestamp}_{group_idx}" + f"{run_random_str}_{run_start_timestamp}_{sampled_g}" if send_routing_key else None ) - for prompt_idx in tqdm( - range(prompts_per_group), desc="Generating questions", leave=False - ): - turn_questions = questions[group_idx][prompt_idx] - turn_prompts = [f"{system_prompt}\n\n{turn_questions[0]}"] + turn_questions[ - 1: - ] - full_prompt = turn_prompts[0] if num_turns == 1 else turn_prompts - prompt_len = 1 if fast_prepare else len(tokenizer.encode(turn_prompts[0])) - output_len_val = int(output_lens[group_idx, prompt_idx]) + turn_questions = questions[src_g][src_p] + turn_prompts = [f"{system_prompt}\n\n{turn_questions[0]}"] + turn_questions[1:] + full_prompt = turn_prompts[0] if num_turns == 1 else turn_prompts + prompt_len = 1 if fast_prepare else len(tokenizer.encode(turn_prompts[0])) + output_len_val = int(output_lens[src_g, src_p]) - input_requests.append( - DatasetRow( - prompt=full_prompt, - prompt_len=prompt_len, - output_len=output_len_val, - routing_key=routing_key, - ) + input_requests.append( + DatasetRow( + prompt=full_prompt, + prompt_len=prompt_len, + output_len=output_len_val, + routing_key=routing_key, ) - total_input_tokens += prompt_len - total_output_tokens += output_len_val + ) + total_input_tokens += prompt_len + total_output_tokens += output_len_val if not ordered: random.shuffle(input_requests) - # Print statistics print(f"\nGenerated shared prefix dataset statistics:") print(f"Number of groups: {num_groups}") print(f"Prompts per group: {prompts_per_group}") print(f"Number of turns: {num_turns}") + print(f"Group distribution: {group_distribution}") + if group_distribution == "zipf": + print(f"Zipf alpha: {zipf_alpha}") print(f"Total prompts: {len(input_requests)}") if not fast_prepare: print(f"Total input tokens: {total_input_tokens}") @@ -221,7 +319,6 @@ def sample_generated_shared_prefix_requests( f"Average question length: {sum(len(tokenizer.encode(q)) for q in all_questions) / len(all_questions):.1f} tokens\n" ) - # Save to cache if should_cache: cache_path.parent.mkdir(parents=True, exist_ok=True) print(f"Caching generated input data to {cache_path}") diff --git a/test/registered/bench_fn/test_benchmark_datasets_api.py b/test/registered/bench_fn/test_benchmark_datasets_api.py index 777bb6741..d56babfee 100644 --- a/test/registered/bench_fn/test_benchmark_datasets_api.py +++ b/test/registered/bench_fn/test_benchmark_datasets_api.py @@ -1,11 +1,17 @@ import asyncio 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 patch +import numpy as np from PIL import Image from tokenizers import Tokenizer from tokenizers.models import WordLevel @@ -16,6 +22,9 @@ from sglang.benchmark.datasets import DATASET_MAPPING, get_dataset from sglang.benchmark.datasets.common import DatasetRow 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 sample_image_requests @@ -26,7 +35,7 @@ from sglang.benchmark.datasets.random import sample_random_requests from sglang.benchmark.datasets.sharegpt import sample_sharegpt_requests from sglang.test.ci.ci_register import register_cpu_ci -register_cpu_ci(est_time=6, suite="base-a-test-cpu") +register_cpu_ci(est_time=40, suite="base-a-test-cpu") register_cpu_ci(est_time=7, suite="base-b-test-cpu") @@ -133,6 +142,8 @@ def make_args(**overrides): "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, @@ -148,8 +159,19 @@ class TestBenchmarkDatasetsAPI(unittest.TestCase): 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): @@ -554,6 +576,486 @@ class TestBenchmarkDatasetsAPI(unittest.TestCase): 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___

____.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 "__". + 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.bench_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.bench_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.bench_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.bench_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()