[Benchmark] Remove obsolete auto-benchmark remnants (#31941)
This commit is contained in:
Xiaoyu Zhang
@@ -1,82 +0,0 @@
|
||||
import argparse
|
||||
|
||||
from sglang.auto_benchmark_lib import (
|
||||
SUPPORTED_DATASETS,
|
||||
convert_dataset,
|
||||
run_auto_benchmark,
|
||||
validate_dataset,
|
||||
)
|
||||
|
||||
|
||||
def add_dataset_args(parser: argparse.ArgumentParser) -> None:
|
||||
parser.add_argument(
|
||||
"--kind",
|
||||
required=True,
|
||||
choices=sorted(SUPPORTED_DATASETS),
|
||||
help="Dataset kind: sharegpt, custom, random, or generated-shared-prefix.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--path",
|
||||
default="",
|
||||
help="Dataset file path. Leave empty for sharegpt auto-download.",
|
||||
)
|
||||
parser.add_argument("--tokenizer", required=True)
|
||||
parser.add_argument("--model", default=None)
|
||||
parser.add_argument("--num-prompts", type=int, default=1000)
|
||||
parser.add_argument("--output-len", type=int, default=None)
|
||||
parser.add_argument("--context-len", type=int, default=None)
|
||||
parser.add_argument("--prompt-suffix", type=str, default="")
|
||||
parser.add_argument("--apply-chat-template", action="store_true")
|
||||
parser.add_argument("--random-input-len", type=int, default=1024)
|
||||
parser.add_argument("--random-output-len", type=int, default=256)
|
||||
parser.add_argument("--random-range-ratio", type=float, default=0.0)
|
||||
parser.add_argument("--gsp-num-groups", type=int, default=64)
|
||||
parser.add_argument("--gsp-prompts-per-group", type=int, default=16)
|
||||
parser.add_argument("--gsp-system-prompt-len", type=int, default=2048)
|
||||
parser.add_argument("--gsp-question-len", type=int, default=128)
|
||||
parser.add_argument("--gsp-output-len", type=int, default=256)
|
||||
parser.add_argument("--gsp-range-ratio", type=float, default=1.0)
|
||||
parser.add_argument("--gsp-fast-prepare", action="store_true")
|
||||
parser.add_argument("--gsp-send-routing-key", action="store_true")
|
||||
parser.add_argument("--gsp-num-turns", type=int, default=1)
|
||||
parser.add_argument("--gsp-ordered", action="store_true")
|
||||
parser.add_argument("--seed", type=int, default=1)
|
||||
|
||||
|
||||
def build_parser() -> argparse.ArgumentParser:
|
||||
parser = argparse.ArgumentParser(description="SGLang auto benchmark utilities.")
|
||||
subparsers = parser.add_subparsers(dest="command", required=True)
|
||||
|
||||
run_parser = subparsers.add_parser(
|
||||
"run", help="Run auto benchmark from YAML config."
|
||||
)
|
||||
run_parser.add_argument("--config", required=True)
|
||||
|
||||
convert_parser = subparsers.add_parser(
|
||||
"convert",
|
||||
help="Prepare sharegpt/custom/random/generated-shared-prefix data into canonical autobench JSONL.",
|
||||
)
|
||||
add_dataset_args(convert_parser)
|
||||
convert_parser.add_argument("--output", required=True)
|
||||
|
||||
validate_parser = subparsers.add_parser(
|
||||
"validate", help="Validate a canonical autobench JSONL dataset."
|
||||
)
|
||||
validate_parser.add_argument("--dataset-path", required=True)
|
||||
validate_parser.add_argument("--tokenizer", required=True)
|
||||
|
||||
return parser
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = build_parser().parse_args()
|
||||
if args.command == "run":
|
||||
run_auto_benchmark(args.config)
|
||||
elif args.command == "convert":
|
||||
convert_dataset(args)
|
||||
elif args.command == "validate":
|
||||
validate_dataset(args)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,7 +1,6 @@
|
||||
from typing import Dict, Type
|
||||
|
||||
from sglang.benchmark.datasets.agentic_trace import AgenticTraceDataset
|
||||
from sglang.benchmark.datasets.autobench import AutoBenchmarkDataset
|
||||
from sglang.benchmark.datasets.common import BaseDataset, DatasetRow
|
||||
from sglang.benchmark.datasets.custom import CustomDataset
|
||||
from sglang.benchmark.datasets.generated_shared_prefix import (
|
||||
@@ -18,7 +17,6 @@ from sglang.benchmark.datasets.speed_bench import SpeedBenchDataset
|
||||
|
||||
DATASET_MAPPING: Dict[str, Type[BaseDataset]] = {
|
||||
"agentic-trace": AgenticTraceDataset,
|
||||
"autobench": AutoBenchmarkDataset,
|
||||
"sharegpt": ShareGPTDataset,
|
||||
"custom": CustomDataset,
|
||||
"openai": OpenAIDataset,
|
||||
|
||||
@@ -1,299 +0,0 @@
|
||||
import json
|
||||
from argparse import Namespace
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
|
||||
import numpy as np
|
||||
from transformers import PreTrainedTokenizerBase
|
||||
|
||||
from sglang.benchmark.datasets.common import BaseDataset, DatasetRow
|
||||
|
||||
AUTOBENCH_RESERVED_FIELDS = {
|
||||
"prompt",
|
||||
"messages",
|
||||
"prompt_origin",
|
||||
"output_len",
|
||||
"max_tokens",
|
||||
"max_completion_tokens",
|
||||
"completion_tokens",
|
||||
"prompt_len",
|
||||
"text_prompt_len",
|
||||
"vision_prompt_len",
|
||||
"image_data",
|
||||
"timestamp",
|
||||
"routing_key",
|
||||
"metadata",
|
||||
"extra_request_body",
|
||||
"param_send",
|
||||
}
|
||||
|
||||
|
||||
def _load_json_if_needed(value: Any) -> Any:
|
||||
if not isinstance(value, str):
|
||||
return value
|
||||
value = value.strip()
|
||||
if not value:
|
||||
return value
|
||||
if value[0] not in "[{":
|
||||
return value
|
||||
try:
|
||||
return json.loads(value)
|
||||
except json.JSONDecodeError:
|
||||
return value
|
||||
|
||||
|
||||
def _normalize_messages(messages: Any) -> Optional[List[Dict[str, Any]]]:
|
||||
messages = _load_json_if_needed(messages)
|
||||
if not isinstance(messages, list) or not messages:
|
||||
return None
|
||||
if not all(isinstance(message, dict) for message in messages):
|
||||
return None
|
||||
|
||||
normalized = []
|
||||
for message in messages:
|
||||
if "role" not in message:
|
||||
return None
|
||||
content = message.get("content")
|
||||
if content is None:
|
||||
return None
|
||||
normalized.append({"role": message["role"], "content": content})
|
||||
return normalized
|
||||
|
||||
|
||||
def _normalize_legacy_system_content(
|
||||
system_prompt: Any, content_list: Any
|
||||
) -> Optional[List[Dict[str, Any]]]:
|
||||
if not isinstance(content_list, list) or not content_list:
|
||||
return None
|
||||
|
||||
messages: List[Dict[str, Any]] = []
|
||||
if system_prompt:
|
||||
messages.append({"role": "system", "content": str(system_prompt)})
|
||||
|
||||
turns = [str(item) for item in content_list]
|
||||
# In the old auto_benchmark helpers, an even number of items usually means the
|
||||
# last assistant reply is present and should be removed before benchmarking.
|
||||
if len(turns) % 2 == 0:
|
||||
turns = turns[:-1]
|
||||
if not turns:
|
||||
return None
|
||||
|
||||
for index, turn in enumerate(turns):
|
||||
role = "user" if index % 2 == 0 else "assistant"
|
||||
messages.append({"role": role, "content": turn})
|
||||
return messages
|
||||
|
||||
|
||||
def _normalize_prompt(row: Dict[str, Any]) -> Tuple[Any, str]:
|
||||
prompt = row.get("prompt")
|
||||
messages = row.get("messages")
|
||||
prompt_origin = row.get("prompt_origin")
|
||||
|
||||
if messages is not None:
|
||||
normalized = _normalize_messages(messages)
|
||||
if normalized is not None:
|
||||
return normalized, "messages"
|
||||
|
||||
if prompt is not None:
|
||||
prompt = _load_json_if_needed(prompt)
|
||||
if isinstance(prompt, list) and prompt and isinstance(prompt[0], dict):
|
||||
normalized = _normalize_messages(prompt)
|
||||
if normalized is not None:
|
||||
return normalized, "messages"
|
||||
if (
|
||||
isinstance(prompt, list)
|
||||
and prompt
|
||||
and all(isinstance(item, str) for item in prompt)
|
||||
):
|
||||
return prompt, "multi_turn"
|
||||
if (
|
||||
isinstance(prompt, list)
|
||||
and prompt
|
||||
and all(
|
||||
isinstance(item, list)
|
||||
and item
|
||||
and all(
|
||||
isinstance(m, dict) and "role" in m and "content" in m for m in item
|
||||
)
|
||||
for item in prompt
|
||||
)
|
||||
):
|
||||
# Multi-turn with N messages per round (e.g. tool observations).
|
||||
return prompt, "multi_turn"
|
||||
if (
|
||||
isinstance(prompt, list)
|
||||
and prompt
|
||||
and all(isinstance(item, int) for item in prompt)
|
||||
):
|
||||
return prompt, "token_ids"
|
||||
if isinstance(prompt, str) and prompt:
|
||||
return prompt, "prompt"
|
||||
|
||||
if prompt_origin is not None:
|
||||
normalized = _normalize_messages(prompt_origin)
|
||||
if normalized is not None:
|
||||
return normalized, "messages"
|
||||
|
||||
if "system" in row and "content" in row:
|
||||
normalized = _normalize_legacy_system_content(
|
||||
row.get("system"), row.get("content")
|
||||
)
|
||||
if normalized is not None:
|
||||
return normalized, "messages"
|
||||
|
||||
raise ValueError("Unsupported auto benchmark row: missing prompt/messages")
|
||||
|
||||
|
||||
def _estimate_prompt_lens(
|
||||
prompt: Any,
|
||||
prompt_kind: str,
|
||||
tokenizer: PreTrainedTokenizerBase,
|
||||
row: Dict[str, Any],
|
||||
) -> Tuple[int, int, int]:
|
||||
if row.get("prompt_len") is not None:
|
||||
prompt_len = int(row["prompt_len"])
|
||||
text_prompt_len = int(row.get("text_prompt_len", prompt_len))
|
||||
vision_prompt_len = int(row.get("vision_prompt_len", 0))
|
||||
return prompt_len, text_prompt_len, vision_prompt_len
|
||||
|
||||
if prompt_kind == "messages":
|
||||
text_prompt_len = len(
|
||||
tokenizer.apply_chat_template(
|
||||
prompt, tokenize=True, add_generation_prompt=True
|
||||
)
|
||||
)
|
||||
vision_prompt_len = 0
|
||||
return text_prompt_len, text_prompt_len, vision_prompt_len
|
||||
|
||||
if prompt_kind == "prompt":
|
||||
prompt_len = len(tokenizer.encode(prompt, add_special_tokens=False))
|
||||
return prompt_len, prompt_len, 0
|
||||
|
||||
if prompt_kind == "token_ids":
|
||||
prompt_len = len(prompt)
|
||||
return prompt_len, prompt_len, 0
|
||||
|
||||
# Multi-turn prompt lists are handled specially by the serving benchmark and do not
|
||||
# contribute reliable static prompt lengths.
|
||||
return 0, 0, 0
|
||||
|
||||
|
||||
def _collect_extra_request_body(row: Dict[str, Any]) -> Dict[str, Any]:
|
||||
extra: Dict[str, Any] = {}
|
||||
|
||||
param_send = row.get("param_send")
|
||||
if param_send is not None:
|
||||
parsed = _load_json_if_needed(param_send)
|
||||
if isinstance(parsed, dict):
|
||||
extra.update(parsed)
|
||||
|
||||
for key, value in row.items():
|
||||
if key not in AUTOBENCH_RESERVED_FIELDS:
|
||||
extra[key] = value
|
||||
|
||||
explicit_extra = row.get("extra_request_body")
|
||||
explicit_extra = _load_json_if_needed(explicit_extra)
|
||||
if isinstance(explicit_extra, dict):
|
||||
extra.update(explicit_extra)
|
||||
|
||||
return extra
|
||||
|
||||
|
||||
def serialize_dataset_row_to_autobench(
|
||||
row: DatasetRow, metadata: Optional[Dict[str, Any]] = None
|
||||
) -> Dict[str, Any]:
|
||||
record: Dict[str, Any] = {
|
||||
"prompt": row.prompt,
|
||||
"output_len": row.output_len,
|
||||
}
|
||||
if row.prompt_len:
|
||||
record["prompt_len"] = row.prompt_len
|
||||
if row.text_prompt_len not in (None, row.prompt_len):
|
||||
record["text_prompt_len"] = row.text_prompt_len
|
||||
if row.vision_prompt_len:
|
||||
record["vision_prompt_len"] = row.vision_prompt_len
|
||||
if row.image_data:
|
||||
record["image_data"] = row.image_data
|
||||
if row.timestamp is not None:
|
||||
record["timestamp"] = row.timestamp
|
||||
if row.routing_key is not None:
|
||||
record["routing_key"] = row.routing_key
|
||||
if row.extra_request_body:
|
||||
record["extra_request_body"] = row.extra_request_body
|
||||
if metadata:
|
||||
record["metadata"] = metadata
|
||||
return record
|
||||
|
||||
|
||||
@dataclass
|
||||
class AutoBenchmarkDataset(BaseDataset):
|
||||
dataset_path: str
|
||||
num_requests: int
|
||||
fixed_output_len: Optional[int]
|
||||
|
||||
@classmethod
|
||||
def from_args(cls, args: Namespace) -> "AutoBenchmarkDataset":
|
||||
return cls(
|
||||
dataset_path=args.dataset_path,
|
||||
num_requests=args.num_prompts,
|
||||
fixed_output_len=args.sharegpt_output_len,
|
||||
)
|
||||
|
||||
def load(
|
||||
self, tokenizer: PreTrainedTokenizerBase, model_id=None
|
||||
) -> List[DatasetRow]:
|
||||
return sample_autobench_requests(
|
||||
dataset_path=self.dataset_path,
|
||||
num_requests=self.num_requests,
|
||||
tokenizer=tokenizer,
|
||||
fixed_output_len=self.fixed_output_len,
|
||||
)
|
||||
|
||||
|
||||
def sample_autobench_requests(
|
||||
dataset_path: str,
|
||||
num_requests: int,
|
||||
tokenizer: PreTrainedTokenizerBase,
|
||||
fixed_output_len: Optional[int] = None,
|
||||
) -> List[DatasetRow]:
|
||||
dataset: List[DatasetRow] = []
|
||||
|
||||
with open(dataset_path, "r", encoding="utf-8") as f:
|
||||
for line in f:
|
||||
if num_requests > 0 and len(dataset) >= num_requests:
|
||||
break
|
||||
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
|
||||
row = json.loads(line)
|
||||
prompt, prompt_kind = _normalize_prompt(row)
|
||||
prompt_len, text_prompt_len, vision_prompt_len = _estimate_prompt_lens(
|
||||
prompt, prompt_kind, tokenizer, row
|
||||
)
|
||||
|
||||
output_len = fixed_output_len or row.get("output_len")
|
||||
output_len = output_len or row.get("max_tokens")
|
||||
output_len = output_len or row.get("max_completion_tokens")
|
||||
output_len = output_len or row.get("completion_tokens")
|
||||
output_len = int(output_len or 256)
|
||||
|
||||
dataset.append(
|
||||
DatasetRow(
|
||||
prompt=prompt,
|
||||
prompt_len=prompt_len,
|
||||
output_len=output_len,
|
||||
text_prompt_len=text_prompt_len,
|
||||
vision_prompt_len=vision_prompt_len,
|
||||
image_data=row.get("image_data"),
|
||||
timestamp=row.get("timestamp"),
|
||||
routing_key=row.get("routing_key"),
|
||||
extra_request_body=_collect_extra_request_body(row),
|
||||
)
|
||||
)
|
||||
|
||||
print(f"Loaded {len(dataset)} auto benchmark requests")
|
||||
print(f"#Input tokens: {np.sum([x.prompt_len for x in dataset])}")
|
||||
print(f"#Output tokens: {np.sum([x.output_len for x in dataset])}")
|
||||
return dataset
|
||||
@@ -2187,7 +2187,6 @@ def cli_main():
|
||||
default="sharegpt",
|
||||
choices=[
|
||||
"agentic-trace",
|
||||
"autobench",
|
||||
"sharegpt",
|
||||
"custom",
|
||||
"openai",
|
||||
|
||||
@@ -1,188 +0,0 @@
|
||||
import json
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest import mock
|
||||
|
||||
from tokenizers import Tokenizer
|
||||
from tokenizers.models import WordLevel
|
||||
from tokenizers.pre_tokenizers import Whitespace
|
||||
from transformers import PreTrainedTokenizerFast
|
||||
|
||||
from sglang.auto_benchmark_lib import build_candidates, build_server_candidates
|
||||
|
||||
|
||||
def create_lightweight_tokenizer() -> PreTrainedTokenizerFast:
|
||||
vocab = {"[UNK]": 0, "[PAD]": 1, "[BOS]": 2, "[EOS]": 3}
|
||||
vocab.update({f"tok_{i}": i + 4 for i in range(4096)})
|
||||
|
||||
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'] }}: {{ message['content'] }}\n"
|
||||
"{% endfor %}"
|
||||
"{% if add_generation_prompt %}assistant:{% endif %}"
|
||||
)
|
||||
return hf_tokenizer
|
||||
|
||||
|
||||
class AutoBenchmarkTestCase(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.tmpdir = tempfile.TemporaryDirectory()
|
||||
self.tmpdir_path = Path(self.tmpdir.name)
|
||||
self.tokenizer = create_lightweight_tokenizer()
|
||||
self.tokenizer_dir = self.tmpdir_path / "tok"
|
||||
self.tokenizer.save_pretrained(self.tokenizer_dir)
|
||||
|
||||
def tearDown(self):
|
||||
self.tmpdir.cleanup()
|
||||
|
||||
def _write_autobench_jsonl(self) -> str:
|
||||
rows = [
|
||||
{"prompt": "tok_1 tok_2 tok_3", "output_len": 32},
|
||||
{
|
||||
"messages": [{"role": "user", "content": "tok_4 tok_5"}],
|
||||
"output_len": 24,
|
||||
"extra_request_body": {"temperature": 0.0},
|
||||
},
|
||||
{
|
||||
"system": "tok_6",
|
||||
"content": ["tok_7 tok_8", "tok_9", "tok_10 tok_11"],
|
||||
"output_len": 16,
|
||||
},
|
||||
]
|
||||
path = self.tmpdir_path / "sample.autobench.jsonl"
|
||||
with open(path, "w", encoding="utf-8") as f:
|
||||
for row in rows:
|
||||
f.write(json.dumps(row) + "\n")
|
||||
return str(path)
|
||||
|
||||
def _write_sharegpt_json(self) -> str:
|
||||
rows = [
|
||||
{
|
||||
"conversations": [
|
||||
{"value": "tok_1 tok_2 tok_3"},
|
||||
{"value": "tok_4 tok_5"},
|
||||
]
|
||||
},
|
||||
{
|
||||
"conversations": [
|
||||
{"value": "tok_6 tok_7"},
|
||||
{"value": "tok_8 tok_9 tok_10"},
|
||||
]
|
||||
},
|
||||
]
|
||||
path = self.tmpdir_path / "sharegpt.json"
|
||||
with open(path, "w", encoding="utf-8") as f:
|
||||
json.dump(rows, f)
|
||||
return str(path)
|
||||
|
||||
def _build_candidates_for_capability(
|
||||
self,
|
||||
base_flags,
|
||||
search_space,
|
||||
*,
|
||||
tier,
|
||||
max_candidates=None,
|
||||
capability=None,
|
||||
):
|
||||
with mock.patch(
|
||||
"sglang.auto_benchmark_lib.detect_current_cuda_capability",
|
||||
return_value=capability,
|
||||
):
|
||||
return build_candidates(
|
||||
base_flags,
|
||||
search_space,
|
||||
tier=tier,
|
||||
max_candidates=max_candidates,
|
||||
)
|
||||
|
||||
def _build_server_candidates_for_capability(
|
||||
self,
|
||||
server_cfg,
|
||||
*,
|
||||
tier=2,
|
||||
max_candidates=None,
|
||||
capability=None,
|
||||
):
|
||||
with mock.patch(
|
||||
"sglang.auto_benchmark_lib.detect_current_cuda_capability",
|
||||
return_value=capability,
|
||||
):
|
||||
return build_server_candidates(
|
||||
server_cfg,
|
||||
tier=tier,
|
||||
max_candidates=max_candidates,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _trial_record(
|
||||
request_rate,
|
||||
*,
|
||||
candidate_id=0,
|
||||
max_concurrency=None,
|
||||
server_flags=None,
|
||||
output_throughput=1.0,
|
||||
mean_ttft_ms=1.0,
|
||||
mean_tpot_ms=1.0,
|
||||
):
|
||||
return {
|
||||
"stage": "base",
|
||||
"candidate_id": candidate_id,
|
||||
"requested_qps": request_rate,
|
||||
"max_concurrency": max_concurrency,
|
||||
"server_flags": dict(server_flags or {"model_path": "/model"}),
|
||||
"sla_passed": True,
|
||||
"metrics": {
|
||||
"output_throughput": output_throughput,
|
||||
"mean_ttft_ms": mean_ttft_ms,
|
||||
"mean_tpot_ms": mean_tpot_ms,
|
||||
},
|
||||
}
|
||||
|
||||
def _make_run_trial_side_effect(
|
||||
self,
|
||||
calls,
|
||||
*,
|
||||
output_throughput=1.0,
|
||||
mean_ttft_ms=1.0,
|
||||
mean_tpot_ms=1.0,
|
||||
):
|
||||
def fake_run_trial(**kwargs):
|
||||
calls.append(kwargs["request_rate"])
|
||||
return self._trial_record(
|
||||
kwargs["request_rate"],
|
||||
candidate_id=kwargs["candidate_id"],
|
||||
max_concurrency=kwargs["max_concurrency"],
|
||||
server_flags=kwargs["server_flags"],
|
||||
output_throughput=output_throughput,
|
||||
mean_ttft_ms=mean_ttft_ms,
|
||||
mean_tpot_ms=mean_tpot_ms,
|
||||
)
|
||||
|
||||
return fake_run_trial
|
||||
|
||||
def _run_candidate_kwargs(self, benchmark_cfg, **overrides):
|
||||
kwargs = {
|
||||
"stage_name": "base",
|
||||
"candidate_id": 0,
|
||||
"server_cfg": {"host": "127.0.0.1", "port": 30000},
|
||||
"benchmark_cfg": benchmark_cfg,
|
||||
"dataset_summary": {"num_requests": 1},
|
||||
"backend": "sglang-oai",
|
||||
"dataset_path": str(self.tmpdir_path / "fake.jsonl"),
|
||||
"tokenizer_path": str(self.tokenizer_dir),
|
||||
"server_flags": {"model_path": "/model"},
|
||||
"output_dir": str(self.tmpdir_path),
|
||||
}
|
||||
kwargs.update(overrides)
|
||||
return kwargs
|
||||
@@ -1,104 +0,0 @@
|
||||
import json
|
||||
import sys
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
|
||||
CURRENT_DIR = Path(__file__).resolve().parent
|
||||
PARENT_DIR = CURRENT_DIR.parent
|
||||
if str(PARENT_DIR) not in sys.path:
|
||||
sys.path.insert(0, str(PARENT_DIR))
|
||||
|
||||
from auto_benchmark import AutoBenchmarkTestCase
|
||||
|
||||
from sglang.auto_benchmark_lib import infer_backend, prepare_dataset
|
||||
from sglang.benchmark.datasets.autobench import sample_autobench_requests
|
||||
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=6, stage="base-b", runner_config="1-gpu-small")
|
||||
register_amd_ci(est_time=6, suite="stage-b-test-1-gpu-small-amd")
|
||||
|
||||
|
||||
class TestAutoBenchmarkDatasetTools(AutoBenchmarkTestCase):
|
||||
def test_prepare_custom_autobench_dataset(self):
|
||||
dataset_path = self._write_autobench_jsonl()
|
||||
output_path = self.tmpdir_path / "prepared.autobench.jsonl"
|
||||
|
||||
prepared_path, rows, summary = prepare_dataset(
|
||||
dataset_cfg={
|
||||
"kind": "custom",
|
||||
"path": dataset_path,
|
||||
"num_prompts": 2,
|
||||
},
|
||||
tokenizer_path=str(self.tokenizer_dir),
|
||||
model=None,
|
||||
output_path=str(output_path),
|
||||
)
|
||||
|
||||
self.assertEqual(prepared_path, str(output_path))
|
||||
self.assertEqual(summary["num_requests"], 2)
|
||||
self.assertTrue(Path(prepared_path).exists())
|
||||
converted_rows = sample_autobench_requests(
|
||||
dataset_path=prepared_path,
|
||||
num_requests=0,
|
||||
tokenizer=self.tokenizer,
|
||||
)
|
||||
self.assertEqual(len(rows), 2)
|
||||
self.assertEqual(len(converted_rows), 2)
|
||||
|
||||
def test_invalid_json_like_prompt_falls_back_to_plain_text(self):
|
||||
path = self.tmpdir_path / "jsonlike.autobench.jsonl"
|
||||
path.write_text(
|
||||
json.dumps({"prompt": "[not actually json", "output_len": 8}) + "\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
rows = sample_autobench_requests(
|
||||
dataset_path=str(path),
|
||||
num_requests=0,
|
||||
tokenizer=self.tokenizer,
|
||||
)
|
||||
|
||||
self.assertEqual(len(rows), 1)
|
||||
self.assertEqual(rows[0].prompt, "[not actually json")
|
||||
|
||||
def test_prepare_sharegpt_dataset(self):
|
||||
sharegpt_path = self._write_sharegpt_json()
|
||||
output_path = self.tmpdir_path / "sharegpt.autobench.jsonl"
|
||||
|
||||
prepared_path, rows, summary = prepare_dataset(
|
||||
dataset_cfg={
|
||||
"kind": "sharegpt",
|
||||
"path": sharegpt_path,
|
||||
"num_prompts": 2,
|
||||
},
|
||||
tokenizer_path=str(self.tokenizer_dir),
|
||||
model=None,
|
||||
output_path=str(output_path),
|
||||
)
|
||||
|
||||
self.assertEqual(prepared_path, str(output_path))
|
||||
self.assertEqual(summary["num_requests"], 2)
|
||||
self.assertEqual(len(rows), 2)
|
||||
|
||||
def test_prepare_custom_dataset_requires_path(self):
|
||||
with self.assertRaisesRegex(ValueError, "dataset.path is required"):
|
||||
prepare_dataset(
|
||||
dataset_cfg={"kind": "custom"},
|
||||
tokenizer_path=str(self.tokenizer_dir),
|
||||
model=None,
|
||||
output_path=str(self.tmpdir_path / "missing.autobench.jsonl"),
|
||||
)
|
||||
|
||||
def test_infer_backend(self):
|
||||
prompt_rows = [SimpleNamespace(prompt="tok_1 tok_2")]
|
||||
chat_rows = [SimpleNamespace(prompt=[{"role": "user", "content": "tok_1"}])]
|
||||
token_id_rows = [SimpleNamespace(prompt=[1, 2, 3])]
|
||||
|
||||
self.assertEqual(infer_backend("auto", prompt_rows), "sglang-oai")
|
||||
self.assertEqual(infer_backend("auto", chat_rows), "sglang-oai-chat")
|
||||
self.assertEqual(infer_backend("auto", token_id_rows), "sglang")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,98 +0,0 @@
|
||||
import sys
|
||||
import time
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest import mock
|
||||
|
||||
CURRENT_DIR = Path(__file__).resolve().parent
|
||||
PARENT_DIR = CURRENT_DIR.parent
|
||||
if str(PARENT_DIR) not in sys.path:
|
||||
sys.path.insert(0, str(PARENT_DIR))
|
||||
|
||||
from auto_benchmark import AutoBenchmarkTestCase
|
||||
|
||||
from sglang.auto_benchmark_lib import SearchDeadlineExceeded, run_candidate
|
||||
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=6, stage="base-b", runner_config="1-gpu-small")
|
||||
register_amd_ci(est_time=6, suite="stage-b-test-1-gpu-small-amd")
|
||||
|
||||
|
||||
class TestAutoBenchmarkRunCandidate(AutoBenchmarkTestCase):
|
||||
def test_run_candidate_binary_search_avoids_rounding_loop(self):
|
||||
benchmark_cfg = {
|
||||
"qps": {"lower": 1.0, "upper": 1.00000001, "tolerance": 1e-12},
|
||||
"max_concurrency": [None],
|
||||
}
|
||||
calls = []
|
||||
|
||||
with mock.patch(
|
||||
"sglang.auto_benchmark_lib.run_trial",
|
||||
side_effect=self._make_run_trial_side_effect(calls),
|
||||
):
|
||||
records = run_candidate(**self._run_candidate_kwargs(benchmark_cfg))
|
||||
|
||||
self.assertLess(len(calls), 40)
|
||||
self.assertEqual(len(records), len(calls))
|
||||
|
||||
def test_run_candidate_binary_search_respects_max_rounds(self):
|
||||
benchmark_cfg = {
|
||||
"qps": {"lower": 1.0, "upper": 32.0, "tolerance": 1e-12, "max_rounds": 2},
|
||||
"max_concurrency": [None],
|
||||
}
|
||||
calls = []
|
||||
|
||||
with mock.patch(
|
||||
"sglang.auto_benchmark_lib.run_trial",
|
||||
side_effect=self._make_run_trial_side_effect(calls),
|
||||
):
|
||||
records = run_candidate(**self._run_candidate_kwargs(benchmark_cfg))
|
||||
|
||||
self.assertEqual(len(calls), 2)
|
||||
self.assertEqual(len(records), 2)
|
||||
|
||||
def test_run_candidate_stops_when_search_budget_is_exhausted(self):
|
||||
benchmark_cfg = {
|
||||
"qps": {"lower": 1.0, "upper": 2.0, "tolerance": 0.1},
|
||||
"max_concurrency": [None],
|
||||
}
|
||||
|
||||
with self.assertRaises(SearchDeadlineExceeded):
|
||||
run_candidate(
|
||||
**self._run_candidate_kwargs(
|
||||
benchmark_cfg,
|
||||
search_deadline=time.time() - 1.0,
|
||||
search_budget_hours=0.1,
|
||||
)
|
||||
)
|
||||
|
||||
def test_run_candidate_resume_skips_existing_fixed_trials(self):
|
||||
benchmark_cfg = {
|
||||
"qps": [1.0, 2.0],
|
||||
"max_concurrency": [None],
|
||||
}
|
||||
existing_records = [self._trial_record(1.0)]
|
||||
calls = []
|
||||
|
||||
with mock.patch(
|
||||
"sglang.auto_benchmark_lib.run_trial",
|
||||
side_effect=self._make_run_trial_side_effect(
|
||||
calls,
|
||||
output_throughput=2.0,
|
||||
mean_ttft_ms=2.0,
|
||||
mean_tpot_ms=2.0,
|
||||
),
|
||||
):
|
||||
records = run_candidate(
|
||||
**self._run_candidate_kwargs(
|
||||
benchmark_cfg,
|
||||
existing_records=existing_records,
|
||||
)
|
||||
)
|
||||
|
||||
self.assertEqual(calls, [2.0])
|
||||
self.assertEqual([record["requested_qps"] for record in records], [1.0, 2.0])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,306 +0,0 @@
|
||||
import json
|
||||
import sys
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
from unittest import mock
|
||||
|
||||
CURRENT_DIR = Path(__file__).resolve().parent
|
||||
PARENT_DIR = CURRENT_DIR.parent
|
||||
if str(PARENT_DIR) not in sys.path:
|
||||
sys.path.insert(0, str(PARENT_DIR))
|
||||
|
||||
from auto_benchmark import AutoBenchmarkTestCase
|
||||
|
||||
from sglang.auto_benchmark_lib import (
|
||||
append_jsonl,
|
||||
build_qps_plan,
|
||||
build_server_candidates,
|
||||
classify_failure,
|
||||
collect_stale_server_pids,
|
||||
describe_search_tier,
|
||||
estimate_trials_per_candidate,
|
||||
expand_dataset_scenarios,
|
||||
format_best_progress,
|
||||
render_scenario_summary_markdown,
|
||||
rendered_launch_command,
|
||||
resolve_max_candidates,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=6, stage="base-b", runner_config="1-gpu-small")
|
||||
register_amd_ci(est_time=6, suite="stage-b-test-1-gpu-small-amd")
|
||||
|
||||
|
||||
class TestAutoBenchmarkSearchTools(AutoBenchmarkTestCase):
|
||||
def test_build_candidates_by_tier(self):
|
||||
base_flags = {"model_path": "/model", "tp_size": 4}
|
||||
search_space = {
|
||||
"prefill_attention_backend": ["fa3", "flashinfer", "triton"],
|
||||
"decode_attention_backend": ["fa3", "flashinfer"],
|
||||
"chunked_prefill_size": [4096, 8192],
|
||||
"max_running_requests": [64, 128],
|
||||
"schedule_policy": ["lpm", "fcfs"],
|
||||
}
|
||||
|
||||
tier1 = self._build_candidates_for_capability(
|
||||
base_flags,
|
||||
search_space,
|
||||
tier=1,
|
||||
max_candidates=None,
|
||||
capability=None,
|
||||
)
|
||||
tier2 = self._build_candidates_for_capability(
|
||||
base_flags,
|
||||
search_space,
|
||||
tier=2,
|
||||
max_candidates=None,
|
||||
capability=None,
|
||||
)
|
||||
tier3 = self._build_candidates_for_capability(
|
||||
base_flags,
|
||||
search_space,
|
||||
tier=3,
|
||||
max_candidates=32,
|
||||
capability=None,
|
||||
)
|
||||
|
||||
self.assertGreater(len(tier1), 1)
|
||||
self.assertGreater(len(tier2), len(tier1))
|
||||
self.assertGreater(len(tier3), len(tier2))
|
||||
self.assertEqual(tier1[0]["model_path"], "/model")
|
||||
|
||||
def test_parallel_search_derives_dp_size(self):
|
||||
server_cfg = {
|
||||
"env": {"CUDA_VISIBLE_DEVICES": "0,1,2,3,4,5,6,7"},
|
||||
"base_flags": {"model_path": "/model"},
|
||||
"parallel": {
|
||||
"tp": [4, 2],
|
||||
"pp_size": [1],
|
||||
},
|
||||
"search_space": {},
|
||||
}
|
||||
|
||||
candidates = build_server_candidates(server_cfg, tier=2, max_candidates=None)
|
||||
tp_dp_pairs = {
|
||||
(candidate["tp_size"], candidate["dp_size"]) for candidate in candidates
|
||||
}
|
||||
self.assertIn((4, 2), tp_dp_pairs)
|
||||
self.assertIn((2, 4), tp_dp_pairs)
|
||||
|
||||
def test_build_server_candidates_filters_unsupported_fa3_on_sm100(self):
|
||||
server_cfg = {
|
||||
"base_flags": {"model_path": "/model", "tp_size": 1},
|
||||
"search_space": {
|
||||
"prefill_attention_backend": ["fa3", "flashinfer"],
|
||||
"decode_attention_backend": ["fa3", "flashinfer"],
|
||||
"chunked_prefill_size": [4096, 8192],
|
||||
},
|
||||
}
|
||||
|
||||
candidates = self._build_server_candidates_for_capability(
|
||||
server_cfg,
|
||||
tier=2,
|
||||
max_candidates=None,
|
||||
capability=(10, 0),
|
||||
)
|
||||
|
||||
self.assertGreater(len(candidates), 0)
|
||||
for candidate in candidates:
|
||||
self.assertNotEqual(candidate.get("attention_backend"), "fa3")
|
||||
self.assertNotEqual(candidate.get("prefill_attention_backend"), "fa3")
|
||||
self.assertNotEqual(candidate.get("decode_attention_backend"), "fa3")
|
||||
|
||||
def test_build_server_candidates_keeps_fa3_on_sm90(self):
|
||||
server_cfg = {
|
||||
"base_flags": {"model_path": "/model", "tp_size": 1},
|
||||
"search_space": {
|
||||
"prefill_attention_backend": ["fa3", "flashinfer"],
|
||||
"decode_attention_backend": ["fa3", "flashinfer"],
|
||||
},
|
||||
}
|
||||
|
||||
candidates = self._build_server_candidates_for_capability(
|
||||
server_cfg,
|
||||
tier=2,
|
||||
max_candidates=None,
|
||||
capability=(9, 0),
|
||||
)
|
||||
|
||||
self.assertTrue(
|
||||
any(
|
||||
candidate.get("prefill_attention_backend") == "fa3"
|
||||
or candidate.get("decode_attention_backend") == "fa3"
|
||||
for candidate in candidates
|
||||
)
|
||||
)
|
||||
|
||||
def test_ep_alias_and_oom_classification(self):
|
||||
server_cfg = {
|
||||
"base_flags": {"model_path": "/model", "tp_size": 8},
|
||||
"search_space": {"ep": [1, 4]},
|
||||
}
|
||||
|
||||
candidates = build_server_candidates(server_cfg, tier=2, max_candidates=None)
|
||||
ep_sizes = {candidate.get("ep_size", 1) for candidate in candidates}
|
||||
self.assertEqual(ep_sizes, {1, 4})
|
||||
|
||||
diagnosis, hint = classify_failure("RuntimeError: CUDA out of memory")
|
||||
self.assertEqual(diagnosis, "oom")
|
||||
self.assertIn("Increase GPU count", hint)
|
||||
|
||||
def test_expand_random_dataset_scenarios(self):
|
||||
scenarios = expand_dataset_scenarios(
|
||||
{
|
||||
"kind": "random",
|
||||
"scenario_names": ["chat", "summarization"],
|
||||
"input_len": [1000, 8000],
|
||||
"output_len": [1000, 1000],
|
||||
}
|
||||
)
|
||||
|
||||
self.assertEqual(len(scenarios), 2)
|
||||
self.assertEqual(scenarios[0]["name"], "chat")
|
||||
self.assertEqual(scenarios[0]["cfg"]["random_input_len"], 1000)
|
||||
self.assertEqual(scenarios[1]["cfg"]["random_input_len"], 8000)
|
||||
self.assertEqual(scenarios[1]["cfg"]["random_output_len"], 1000)
|
||||
|
||||
def test_estimate_trials_and_tier_descriptions(self):
|
||||
benchmark_cfg = {
|
||||
"qps": {"lower": 0.25, "upper": 4.0, "tolerance": 0.1},
|
||||
"max_concurrency": [None, 8, 16],
|
||||
}
|
||||
|
||||
self.assertEqual(estimate_trials_per_candidate(benchmark_cfg), 15)
|
||||
self.assertIn("default", describe_search_tier(2))
|
||||
self.assertIn("slowest", describe_search_tier(3))
|
||||
|
||||
def test_resolve_max_candidates_defaults_to_eight(self):
|
||||
self.assertEqual(resolve_max_candidates({}), 8)
|
||||
self.assertIsNone(resolve_max_candidates({"max_candidates": None}))
|
||||
|
||||
def test_resolve_max_candidates_rejects_non_positive_values(self):
|
||||
with self.assertRaisesRegex(ValueError, "search.max_candidates"):
|
||||
resolve_max_candidates({"max_candidates": 0})
|
||||
|
||||
def test_build_qps_plan_accepts_numeric_request_rate(self):
|
||||
mode, values, tolerance, max_rounds = build_qps_plan({"request_rate": 3.5})
|
||||
self.assertEqual(mode, "fixed")
|
||||
self.assertEqual(values, [3.5])
|
||||
self.assertEqual(tolerance, 0.0)
|
||||
self.assertEqual(max_rounds, 0)
|
||||
|
||||
def test_build_qps_plan_clamps_binary_rounds(self):
|
||||
mode, values, tolerance, max_rounds = build_qps_plan(
|
||||
{"qps": {"lower": 1.0, "upper": 16.0, "tolerance": 0.1, "max_rounds": 99}}
|
||||
)
|
||||
|
||||
self.assertEqual(mode, "search")
|
||||
self.assertEqual(values, [1.0, 16.0])
|
||||
self.assertEqual(tolerance, 0.1)
|
||||
self.assertEqual(max_rounds, 5)
|
||||
|
||||
def test_format_best_progress(self):
|
||||
text = format_best_progress(
|
||||
{
|
||||
"candidate_id": 3,
|
||||
"requested_qps": 3.5,
|
||||
"server_flags": {
|
||||
"tp_size": 4,
|
||||
"ep_size": 4,
|
||||
"mem_fraction_static": 0.84,
|
||||
"max_running_requests": 96,
|
||||
},
|
||||
"metrics": {
|
||||
"output_throughput": 1234.56,
|
||||
"mean_ttft_ms": 250.12,
|
||||
"mean_tpot_ms": 14.78,
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
self.assertIn("qps=3.5000", text)
|
||||
self.assertIn("tok/s=1234.6", text)
|
||||
self.assertIn("ttft=250.1ms", text)
|
||||
self.assertIn("tpot=14.8ms", text)
|
||||
self.assertIn("tp=4", text)
|
||||
self.assertIn("ep=4", text)
|
||||
|
||||
def test_append_jsonl(self):
|
||||
path = self.tmpdir_path / "live_results.jsonl"
|
||||
append_jsonl(
|
||||
str(path),
|
||||
[
|
||||
{"candidate_id": 1, "requested_qps": 2.0},
|
||||
{"candidate_id": 2, "requested_qps": 3.0},
|
||||
],
|
||||
)
|
||||
|
||||
lines = path.read_text(encoding="utf-8").strip().splitlines()
|
||||
self.assertEqual(len(lines), 2)
|
||||
self.assertEqual(json.loads(lines[0])["candidate_id"], 1)
|
||||
self.assertEqual(json.loads(lines[1])["requested_qps"], 3.0)
|
||||
|
||||
def test_collect_stale_server_pids_dedups(self):
|
||||
def fake_run(command, capture_output, text, check):
|
||||
stdout = "123\n" if command[0] == "lsof" else "123\n456\n"
|
||||
return SimpleNamespace(returncode=0, stdout=stdout)
|
||||
|
||||
with mock.patch(
|
||||
"sglang.auto_benchmark_lib.subprocess.run", side_effect=fake_run
|
||||
):
|
||||
self.assertEqual(collect_stale_server_pids(30000), [123, 456])
|
||||
|
||||
def test_rendered_launch_command_includes_env(self):
|
||||
text = rendered_launch_command(
|
||||
{
|
||||
"env": {
|
||||
"CUDA_VISIBLE_DEVICES": "0",
|
||||
"HF_TOKEN": "secret-value",
|
||||
},
|
||||
"extra_args": [],
|
||||
},
|
||||
{"model_path": "Qwen/Qwen3-32B", "tp_size": 1, "port": 30000},
|
||||
)
|
||||
|
||||
self.assertIn("CUDA_VISIBLE_DEVICES=0", text)
|
||||
self.assertIn("--model-path Qwen/Qwen3-32B", text)
|
||||
self.assertNotIn("HF_TOKEN", text)
|
||||
|
||||
def test_render_scenario_summary_markdown_keeps_rows_in_single_table(self):
|
||||
text = render_scenario_summary_markdown(
|
||||
[
|
||||
{
|
||||
"scenario_name": "chat",
|
||||
"scenario_dir": "/tmp/chat",
|
||||
"status": "ok",
|
||||
"requested_qps": 11.914,
|
||||
"output_throughput": 1867.28,
|
||||
"mean_ttft_ms": 99.58,
|
||||
"mean_tpot_ms": 21.09,
|
||||
"launch_command": "python -m sglang.launch_server --port 30000",
|
||||
},
|
||||
{
|
||||
"scenario_name": "summarization",
|
||||
"scenario_dir": "/tmp/summarization",
|
||||
"status": "ok",
|
||||
"requested_qps": 11.914,
|
||||
"output_throughput": 537.17,
|
||||
"mean_ttft_ms": 709.99,
|
||||
"mean_tpot_ms": 26.89,
|
||||
"launch_command": "python -m sglang.launch_server --port 30001",
|
||||
},
|
||||
]
|
||||
)
|
||||
|
||||
header = (
|
||||
"| Scenario | Status | QPS | Output tok/s | TTFT ms | TPOT ms | Summary |"
|
||||
)
|
||||
self.assertEqual(text.count(header), 1)
|
||||
self.assertLess(text.index("| chat |"), text.index("## chat"))
|
||||
self.assertLess(text.index("| summarization |"), text.index("## chat"))
|
||||
self.assertLess(text.index("| summarization |"), text.index("## summarization"))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user