import argparse import csv import itertools import json import os import shlex import signal import subprocess import sys import time from copy import deepcopy from types import SimpleNamespace from typing import Any, Callable, Dict, Iterable, List, Optional, Sequence, Tuple import yaml from tqdm.auto import tqdm from sglang.benchmark.datasets import get_dataset from sglang.benchmark.datasets.autobench import ( sample_autobench_requests, serialize_dataset_row_to_autobench, ) from sglang.benchmark.utils import get_tokenizer SUPPORTED_DATASETS = { "sharegpt", "custom", "random", "generated-shared-prefix", } FLAG_ALIASES = { "tp": "tp_size", "pp": "pp_size", "dp": "dp_size", "ep": "ep_size", } OOM_HINT = "Candidate likely OOMed. Increase GPU count or use GPUs with larger memory." PROGRESS_FLAG_KEYS = ( "tp_size", "dp_size", "ep_size", "pp_size", "prefill_attention_backend", "decode_attention_backend", "attention_backend", "sampling_backend", "grammar_backend", "mem_fraction_static", "chunked_prefill_size", "prefill_max_requests", "max_prefill_tokens", "max_running_requests", "max_queued_requests", "schedule_policy", "schedule_conservativeness", "num_continuous_decode_steps", "stream_interval", "page_size", "cuda_graph_max_bs_decode", "speculative_num_steps", "speculative_eagle_topk", "speculative_num_draft_tokens", ) PROGRESS_FLAG_ALIASES = { "tp_size": "tp", "dp_size": "dp", "ep_size": "ep", "pp_size": "pp", "prefill_attention_backend": "prefill", "decode_attention_backend": "decode", "attention_backend": "attn", "sampling_backend": "sampling", "grammar_backend": "grammar", "mem_fraction_static": "mfs", "chunked_prefill_size": "chunk", "prefill_max_requests": "prefill_req", "max_prefill_tokens": "prefill_tok", "max_running_requests": "mrr", "max_queued_requests": "mqr", "schedule_policy": "sched", "schedule_conservativeness": "sched_cons", "num_continuous_decode_steps": "decode_steps", "stream_interval": "stream", "page_size": "page", "cuda_graph_max_bs_decode": "cg_bs", "speculative_num_steps": "spec_steps", "speculative_eagle_topk": "eagle_topk", "speculative_num_draft_tokens": "draft_tok", } SENSITIVE_ENV_MARKERS = ("TOKEN", "KEY", "SECRET", "PASSWORD") DEFAULT_MAX_CANDIDATES = 8 MAX_BINARY_SEARCH_ROUNDS = 5 DEFAULT_BINARY_SEARCH_ROUNDS = 5 MAX_SEARCH_DURATION_HOURS = 12.0 DEFAULT_SEARCH_DURATION_HOURS = 12.0 class SearchDeadlineExceeded(RuntimeError): """Raised when the auto benchmark exhausts its global search budget.""" def load_yaml(path: str) -> Dict[str, Any]: with open(path, "r", encoding="utf-8") as f: return yaml.safe_load(f) def as_list(value: Any) -> List[Any]: return value if isinstance(value, list) else [value] def slugify(text: str) -> str: return "".join(ch.lower() if ch.isalnum() else "-" for ch in text).strip("-") def canonical_flag_name(name: str) -> str: return FLAG_ALIASES.get(name, name) def canonicalize_flags(flags: Dict[str, Any]) -> Dict[str, Any]: return {canonical_flag_name(key): value for key, value in flags.items()} def flatten(data: Dict[str, Any], prefix: str = "") -> Dict[str, Any]: flat: Dict[str, Any] = {} for key, value in data.items(): name = f"{prefix}.{key}" if prefix else key if isinstance(value, dict): flat.update(flatten(value, name)) else: flat[name] = value return flat def log_line(message: str) -> None: tqdm.write(message) def detect_current_cuda_capability() -> Optional[Tuple[int, int]]: try: import torch except ModuleNotFoundError: return None if not torch.cuda.is_available(): return None major, minor = torch.cuda.get_device_capability() return int(major), int(minor) def is_attention_backend_supported( backend: Any, capability: Optional[Tuple[int, int]] ) -> bool: if capability is None or backend in (None, ""): return True major, _minor = capability if backend == "fa3": return major in (8, 9) return True def is_candidate_supported_on_current_device( candidate: Dict[str, Any], capability: Optional[Tuple[int, int]] ) -> bool: backend_keys = ( "attention_backend", "prefill_attention_backend", "decode_attention_backend", ) return all( is_attention_backend_supported(candidate.get(key), capability) for key in backend_keys ) def append_jsonl(path: str, records: Iterable[Dict[str, Any]]) -> None: with open(path, "a", encoding="utf-8") as f: for record in records: f.write(json.dumps(record, ensure_ascii=False) + "\n") def read_jsonl(path: str) -> List[Dict[str, Any]]: if not path or not os.path.isfile(path): return [] records: List[Dict[str, Any]] = [] with open(path, "r", encoding="utf-8") as f: for line in f: line = line.strip() if not line: continue records.append(json.loads(line)) return records def describe_search_tier(tier: int) -> str: descriptions = { 1: "tier 1: smallest and fastest sanity sweep", 2: "tier 2: balanced default sweep", 3: "tier 3: largest and slowest full search", } return descriptions.get(tier, f"tier {tier}") def install_interrupt_handlers() -> Dict[signal.Signals, Any]: previous = {} def handler(signum, _frame): # type: ignore[no-untyped-def] raise KeyboardInterrupt(f"Interrupted by signal {signum}") for sig in (signal.SIGINT, signal.SIGTERM): try: previous[sig] = signal.getsignal(sig) signal.signal(sig, handler) except Exception: continue return previous def restore_interrupt_handlers(previous: Dict[signal.Signals, Any]) -> None: for sig, handler in previous.items(): try: signal.signal(sig, handler) except Exception: continue def collect_stale_server_pids(port: int) -> List[int]: patterns = [ ["lsof", "-ti", f"tcp:{port}", "-sTCP:LISTEN"], ["pgrep", "-f", f"sglang.launch_server.*--port {port}"], ["pgrep", "-f", f"sglang.launch_server.*--port={port}"], ["pgrep", "-f", f"sglang serve .*--port {port}"], ["pgrep", "-f", f"sglang serve .*--port={port}"], ] pids = set() for command in patterns: try: result = subprocess.run( command, capture_output=True, text=True, check=False ) except FileNotFoundError: continue if result.returncode not in (0, 1): continue for line in result.stdout.splitlines(): line = line.strip() if line.isdigit(): pids.add(int(line)) return sorted(pids) def kill_pid_or_group(pid: int) -> None: try: pgid = os.getpgid(pid) except ProcessLookupError: return for sig, delay in ((signal.SIGTERM, 1.0), (signal.SIGKILL, 0.0)): try: os.killpg(pgid, sig) except ProcessLookupError: return except PermissionError: try: os.kill(pid, sig) except ProcessLookupError: return if delay: time.sleep(delay) def preclean_stale_server(port: int) -> None: stale_pids = collect_stale_server_pids(port) if not stale_pids: return log_line(f"preclean_port={port} stale_pids={stale_pids}") for pid in stale_pids: kill_pid_or_group(pid) def normalize_binary_search_rounds(value: Any) -> int: if value is None: return DEFAULT_BINARY_SEARCH_ROUNDS return max(1, min(int(value), MAX_BINARY_SEARCH_ROUNDS)) def resolve_max_candidates(search_cfg: Dict[str, Any]) -> Optional[int]: if "max_candidates" not in search_cfg: return DEFAULT_MAX_CANDIDATES configured = search_cfg.get("max_candidates") if configured is None: return None value = int(configured) if value < 1: raise ValueError("search.max_candidates must be >= 1 or null.") return value def estimate_binary_search_trials( lower: float, upper: float, tolerance: float, max_rounds: int ) -> int: if upper <= lower or tolerance <= 0: return 1 trials = 0 lo, hi = float(lower), float(upper) while hi - lo > tolerance and trials < max_rounds: qps = pick_qps_midpoint(lo, hi) if qps <= lo or qps >= hi: break hi = qps trials += 1 return max(trials, 1) def pick_qps_midpoint(lower: float, upper: float) -> float: midpoint = round((lower + upper) / 2, 4) if lower < midpoint < upper: return midpoint return (lower + upper) / 2 def estimate_trials_per_candidate(benchmark_cfg: Dict[str, Any]) -> int: mode, values, tolerance, max_rounds = build_qps_plan(benchmark_cfg) max_concurrency_values = as_list(benchmark_cfg.get("max_concurrency", [None])) if mode == "fixed": per_concurrency = len(values) else: per_concurrency = estimate_binary_search_trials( values[0], values[1], tolerance, max_rounds ) return max(1, per_concurrency) * len(max_concurrency_values) def describe_qps_plan(benchmark_cfg: Dict[str, Any]) -> str: mode, values, tolerance, max_rounds = build_qps_plan(benchmark_cfg) if mode == "fixed": return f"fixed qps values={values}" return ( f"binary search qps lower={values[0]} upper={values[1]} " f"tolerance={tolerance} max_rounds={max_rounds} " "estimated_trials_per_max_concurrency=" f"{estimate_binary_search_trials(values[0], values[1], tolerance, max_rounds)}" ) def scenario_plan_text(scenario: Dict[str, Any]) -> str: cfg = scenario["cfg"] parts = [f"kind={cfg['kind']}", f"num_prompts={cfg.get('num_prompts', '')}"] if cfg["kind"] == "random": parts.append(f"input_len={cfg['random_input_len']}") parts.append(f"output_len={cfg['random_output_len']}") elif cfg.get("path"): parts.append(f"path={cfg['path']}") return ", ".join(str(part) for part in parts if part != "") def print_run_plan( config_path: str, output_dir: str, tier: int, max_candidates: Optional[int], benchmark_cfg: Dict[str, Any], scenarios: Sequence[Dict[str, Any]], server_cfg: Dict[str, Any], base_candidates: Sequence[Dict[str, Any]], speculative_enabled: bool, search_budget_hours: float, search_deadline: float, ) -> None: estimated_base_trials = ( len(scenarios) * len(base_candidates) * estimate_trials_per_candidate(benchmark_cfg) ) log_line("=== Auto Benchmark Plan ===") log_line(f"config={config_path}") log_line(f"output_dir={output_dir}") log_line(f"search.tier={tier} ({describe_search_tier(tier)})") log_line( "search.max_candidates=" f"{max_candidates if max_candidates is not None else 'unbounded'}" ) log_line( f"search.max_duration_hours={search_budget_hours:.1f} " f"(deadline {time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(search_deadline))})" ) log_line(f"qps_plan={describe_qps_plan(benchmark_cfg)}") log_line( "max_concurrency=" f"{json.dumps(as_list(benchmark_cfg.get('max_concurrency', [None])), ensure_ascii=False)}" ) log_line(f"estimated_base_trials={estimated_base_trials}") log_line("Planned scenarios:") for index, scenario in enumerate(scenarios, start=1): log_line( f" [{index}/{len(scenarios)}] {scenario['display_name']}: " f"{scenario_plan_text(scenario)}" ) log_line("Planned base candidates:") for index, candidate in enumerate(base_candidates, start=1): rendered = merge_host_port(server_cfg, candidate) log_line( f" [{index}/{len(base_candidates)}] {json.dumps(rendered, ensure_ascii=False)}" ) if speculative_enabled: log_line( "Speculative stage is enabled. Its candidate list will be printed after " "the best base configuration is known." ) def estimated_finish_time( start_time: float, completed: int, total: Optional[int] ) -> str: if not total or completed <= 0: return "?" remaining_seconds = max( 0.0, (time.time() - start_time) * (total - completed) / completed ) return time.strftime( "%Y-%m-%d %H:%M:%S", time.localtime(time.time() + remaining_seconds) ) def current_time_text() -> str: return time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()) def resolve_search_budget_hours(search_cfg: Dict[str, Any]) -> float: configured = search_cfg.get("max_duration_hours", DEFAULT_SEARCH_DURATION_HOURS) return max(0.0, min(float(configured), MAX_SEARCH_DURATION_HOURS)) def format_timestamp(timestamp: float) -> str: return time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(timestamp)) def remaining_search_seconds(search_deadline: Optional[float]) -> Optional[float]: if search_deadline is None: return None return max(0.0, search_deadline - time.time()) def raise_if_search_deadline_reached( search_deadline: Optional[float], budget_hours: float ) -> None: remaining = remaining_search_seconds(search_deadline) if remaining is None or remaining > 0: return raise SearchDeadlineExceeded( "search budget of " f"{budget_hours:.1f}h reached before the full search completed " f"(deadline {format_timestamp(search_deadline)})" ) def summarize_progress_flags(server_flags: Dict[str, Any], limit: int = 6) -> str: parts = [] for key in PROGRESS_FLAG_KEYS: if key not in server_flags: continue value = server_flags[key] if value in (None, "", False): continue alias = PROGRESS_FLAG_ALIASES.get(key, key) parts.append(f"{alias}={value}") if len(parts) >= limit: break if not parts and server_flags.get("candidate_id") is not None: return f"candidate={server_flags['candidate_id']}" return ",".join(parts) def format_best_progress(record: Optional[Dict[str, Any]]) -> str: if not record or not record.get("metrics"): return "best pending" metrics = record["metrics"] flags = dict(record.get("server_flags", {})) flags["candidate_id"] = record.get("candidate_id") return ( "best " f"qps={record.get('requested_qps', 0.0):.4f} " f"tok/s={metrics.get('output_throughput', 0.0):.1f} " f"ttft={metrics.get('mean_ttft_ms', 0.0):.1f}ms " f"tpot={metrics.get('mean_tpot_ms', 0.0):.1f}ms " f"cfg[{summarize_progress_flags(flags)}]" ) def refresh_progress_eta( pbar: tqdm, start_time: float, best_record: Optional[Dict[str, Any]] = None ) -> None: pbar.set_postfix_str( f"now {current_time_text()} | " f"finish {estimated_finish_time(start_time, int(pbar.n), pbar.total)} | " f"{format_best_progress(best_record)}", refresh=False, ) def make_progress_bar( desc: str, total: int, position: int, leave: bool ) -> Tuple[tqdm, float]: start_time = time.time() pbar = tqdm( total=total, desc=desc, dynamic_ncols=True, mininterval=1.0, position=position, leave=leave, bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}] {postfix}", ) refresh_progress_eta(pbar, start_time) return pbar, start_time def advance_progress( pbar: tqdm, start_time: float, count: int = 1, best_record: Optional[Dict[str, Any]] = None, ) -> None: if pbar.total is not None and pbar.n + count > pbar.total: pbar.total = pbar.n + count pbar.update(count) refresh_progress_eta(pbar, start_time, best_record) def tail_text(path: str, limit: int = 4000) -> str: if not path or not os.path.isfile(path): return "" with open(path, "r", encoding="utf-8", errors="ignore") as f: text = f.read() return text[-limit:] def cli_args(flags: Dict[str, Any]) -> List[str]: args: List[str] = [] for key, value in flags.items(): if value is None or value is False: continue flag = f"--{key.replace('_', '-')}" if value is True: args.append(flag) elif isinstance(value, list): args.append(flag) args.extend(str(item) for item in value) else: args.extend([flag, str(value)]) return args def classify_failure(message: str) -> Tuple[Optional[str], Optional[str]]: lower = message.lower() oom_markers = ( "out of memory", "cuda out of memory", "hip out of memory", "cudnn_status_alloc_failed", "std::bad_alloc", "memoryerror", "memory allocation", "no available memory", ) if any(marker in lower for marker in oom_markers): return "oom", OOM_HINT return None, None def prompt_kind(prompt: Any) -> str: if isinstance(prompt, str): return "prompt" if isinstance(prompt, list) and prompt: if isinstance(prompt[0], dict): return "messages" if isinstance(prompt[0], str): return "multi_turn" if isinstance(prompt[0], int): return "token_ids" return "unknown" def summarize_rows(rows: Sequence[Any]) -> Dict[str, Any]: kinds: Dict[str, int] = {} output_lens = [row.output_len for row in rows] for row in rows: kind = prompt_kind(row.prompt) kinds[kind] = kinds.get(kind, 0) + 1 return { "num_requests": len(rows), "prompt_kinds": kinds, "output_len_min": min(output_lens) if output_lens else 0, "output_len_max": max(output_lens) if output_lens else 0, "output_len_avg": ( round(sum(output_lens) / len(output_lens), 2) if output_lens else 0.0 ), } def infer_backend(backend: str, rows: Sequence[Any]) -> str: if backend != "auto": return backend kinds = {prompt_kind(row.prompt) for row in rows} if kinds <= {"messages", "multi_turn"}: return "sglang-oai-chat" if kinds <= {"prompt"}: return "sglang-oai" if kinds <= {"token_ids"}: return "sglang" raise ValueError( f"Cannot infer backend for mixed prompt kinds: {sorted(kinds)}. " "Set benchmark.backend explicitly." ) def looks_like_autobench(path: str) -> bool: if not path or not os.path.isfile(path): return False with open(path, "r", encoding="utf-8") as f: for line in f: line = line.strip() if not line: continue try: row = json.loads(line) except json.JSONDecodeError: return False return isinstance(row, dict) and any( key in row for key in ("prompt", "messages", "prompt_origin", "system") ) return False def write_autobench_jsonl( path: str, rows: Sequence[Any], metadata: Optional[Dict[str, Any]] = None ) -> None: directory = os.path.dirname(path) if directory: os.makedirs(directory, exist_ok=True) with open(path, "w", encoding="utf-8") as f: for row in rows: record = serialize_dataset_row_to_autobench(row, metadata=metadata) f.write(json.dumps(record, ensure_ascii=False) + "\n") def normalize_dataset_cfg( dataset_cfg: Optional[Dict[str, Any]], benchmark_cfg: Dict[str, Any] ) -> Dict[str, Any]: raw = {} if dataset_cfg is None else dataset_cfg if isinstance(raw, str): raw = {"kind": raw} cfg = dict(raw) if "kind" not in cfg and cfg.get("path") in SUPPORTED_DATASETS: cfg["kind"] = cfg["path"] cfg["path"] = "" if "kind" not in cfg and benchmark_cfg.get("dataset_path"): cfg["kind"] = "custom" cfg["path"] = benchmark_cfg["dataset_path"] if "num_prompts" not in cfg and benchmark_cfg.get("num_prompts") is not None: cfg["num_prompts"] = benchmark_cfg["num_prompts"] cfg["kind"] = cfg.get("kind", "custom") if cfg["kind"] == "autobench": cfg["kind"] = "custom" if cfg["kind"] not in SUPPORTED_DATASETS: raise ValueError( f"Unsupported dataset kind: {cfg['kind']}. " f"Supported: {sorted(SUPPORTED_DATASETS)}" ) if cfg["kind"] == "custom" and not cfg.get("path"): raise ValueError("dataset.path is required when dataset.kind=custom.") return cfg def expand_dataset_scenarios(dataset_cfg: Dict[str, Any]) -> List[Dict[str, Any]]: if dataset_cfg["kind"] != "random": name = dataset_cfg.get("scenario_name", "default") return [ { "name": slugify(str(name)) or "default", "display_name": str(name), "cfg": dataset_cfg, } ] input_lens = as_list( dataset_cfg.get("input_len", dataset_cfg.get("random_input_len", 1024)) ) output_lens = as_list( dataset_cfg.get("output_len", dataset_cfg.get("random_output_len", 256)) ) if len(input_lens) != len(output_lens): raise ValueError( "random dataset input_len and output_len must have the same number of elements." ) scenario_names = dataset_cfg.get("scenario_names") if scenario_names is not None and len(as_list(scenario_names)) != len(input_lens): raise ValueError( "dataset.scenario_names must match the length of input_len/output_len." ) names = as_list(scenario_names) if scenario_names is not None else None scenarios = [] for index, (input_len, output_len) in enumerate(zip(input_lens, output_lens)): cfg = dict(dataset_cfg) cfg["random_input_len"] = int(input_len) cfg["random_output_len"] = int(output_len) cfg["input_len"] = int(input_len) cfg["output_len"] = int(output_len) display_name = ( str(names[index]) if names is not None else f"input{int(input_len)}-output{int(output_len)}" ) scenarios.append( { "name": slugify(display_name) or f"scenario-{index + 1}", "display_name": display_name, "cfg": cfg, } ) return scenarios def build_dataset_args( dataset_cfg: Dict[str, Any], tokenizer_path: str, model: Optional[str] ) -> SimpleNamespace: dataset_path = dataset_cfg.get("path", "") if dataset_cfg["kind"] == "sharegpt" and dataset_path in ("", None, "sharegpt"): dataset_path = "" is_random = dataset_cfg["kind"] == "random" return SimpleNamespace( dataset_name=dataset_cfg["kind"], dataset_path=dataset_path, tokenizer=tokenizer_path, model=model, num_prompts=int(dataset_cfg.get("num_prompts", 1000)), sharegpt_output_len=(dataset_cfg.get("output_len") if not is_random else None), sharegpt_context_len=dataset_cfg.get("context_len"), random_input_len=int( dataset_cfg.get("input_len", dataset_cfg.get("random_input_len", 1024)) ), random_output_len=int( dataset_cfg.get("output_len", dataset_cfg.get("random_output_len", 256)) ), random_range_ratio=float(dataset_cfg.get("random_range_ratio", 0.0)), prompt_suffix=dataset_cfg.get("prompt_suffix", ""), apply_chat_template=bool(dataset_cfg.get("apply_chat_template", False)), gsp_num_groups=int(dataset_cfg.get("gsp_num_groups", 64)), gsp_prompts_per_group=int(dataset_cfg.get("gsp_prompts_per_group", 16)), gsp_system_prompt_len=int(dataset_cfg.get("gsp_system_prompt_len", 2048)), gsp_question_len=int(dataset_cfg.get("gsp_question_len", 128)), gsp_output_len=int(dataset_cfg.get("gsp_output_len", 256)), gsp_range_ratio=float(dataset_cfg.get("gsp_range_ratio", 1.0)), gsp_fast_prepare=bool(dataset_cfg.get("gsp_fast_prepare", False)), gsp_send_routing_key=bool(dataset_cfg.get("gsp_send_routing_key", False)), gsp_num_turns=int(dataset_cfg.get("gsp_num_turns", 1)), gsp_ordered=bool(dataset_cfg.get("gsp_ordered", False)), seed=int(dataset_cfg.get("seed", 1)), ) def load_autobench_rows( dataset_path: str, tokenizer_path: str, num_prompts: int = 0, output_len: Optional[int] = None, ) -> List[Any]: return sample_autobench_requests( dataset_path=dataset_path, num_requests=num_prompts, tokenizer=get_tokenizer(tokenizer_path), fixed_output_len=output_len, ) def prepare_dataset( dataset_cfg: Dict[str, Any], tokenizer_path: str, model: Optional[str], output_path: str, ) -> Tuple[str, List[Any], Dict[str, Any]]: dataset_cfg = normalize_dataset_cfg(dataset_cfg, {}) if dataset_cfg["kind"] == "custom" and looks_like_autobench( dataset_cfg.get("path", "") ): rows = load_autobench_rows( dataset_path=dataset_cfg["path"], tokenizer_path=tokenizer_path, num_prompts=int(dataset_cfg.get("num_prompts", 0)), output_len=dataset_cfg.get("output_len"), ) else: tokenizer = get_tokenizer(tokenizer_path) dataset_args = build_dataset_args(dataset_cfg, tokenizer_path, model) rows = get_dataset(dataset_args, tokenizer=tokenizer, model_id=model) if not rows: raise ValueError("Prepared dataset is empty.") write_autobench_jsonl( output_path, rows, metadata={ "source_dataset_name": dataset_cfg["kind"], "source_dataset_path": dataset_cfg.get("path") or dataset_cfg["kind"], }, ) return output_path, rows, summarize_rows(rows) def infer_total_gpus(server_cfg: Dict[str, Any]) -> Optional[int]: parallel_cfg = server_cfg.get("parallel", {}) for key in ("gpu_count",): value = parallel_cfg.get(key, server_cfg.get(key)) if value is not None: return int(value) env = server_cfg.get("env", {}) for key in ( "CUDA_VISIBLE_DEVICES", "ROCR_VISIBLE_DEVICES", "HIP_VISIBLE_DEVICES", "NVIDIA_VISIBLE_DEVICES", ): value = env.get(key) if value is None: continue value = str(value).strip() if not value or value.lower() in {"all", "none", "void"}: continue return len([item for item in value.split(",") if item.strip()]) return None def resolve_parallelism( server_cfg: Dict[str, Any], flags: Dict[str, Any], parallel_requested: bool ) -> Dict[str, Any]: flags = canonicalize_flags(flags) if not parallel_requested: return flags tp_size = int(flags.get("tp_size", 1)) pp_size = int(flags.get("pp_size", 1)) if "dp_size" in flags: return flags total_gpus = infer_total_gpus(server_cfg) if total_gpus is None: raise ValueError( "Cannot infer total GPU count for parallel search. " "Set server.parallel.gpu_count or server.env.CUDA_VISIBLE_DEVICES." ) shard_size = tp_size * pp_size if shard_size <= 0 or total_gpus % shard_size != 0: raise ValueError( f"Cannot derive dp_size: total_gpus={total_gpus}, " f"tp_size={tp_size}, pp_size={pp_size}." ) flags["dp_size"] = total_gpus // shard_size return flags def build_server_candidates( server_cfg: Dict[str, Any], tier: int, max_candidates: Optional[int] ) -> List[Dict[str, Any]]: base_flags = canonicalize_flags(deepcopy(server_cfg.get("base_flags", {}))) search_space = canonicalize_flags(deepcopy(server_cfg.get("search_space", {}))) parallel_cfg = canonicalize_flags(deepcopy(server_cfg.get("parallel", {}))) parallel_requested = bool(parallel_cfg) for key, value in parallel_cfg.items(): if key == "gpu_count": continue values = as_list(value) if values: base_flags.setdefault(key, values[0]) search_space.update( {key: value for key, value in parallel_cfg.items() if key != "gpu_count"} ) candidates = build_candidates( base_flags=base_flags, search_space=search_space, tier=tier, max_candidates=max_candidates, ) return [ resolve_parallelism(server_cfg, candidate, parallel_requested) for candidate in candidates ] def build_candidates( base_flags: Dict[str, Any], search_space: Dict[str, Sequence[Any]], tier: int, max_candidates: Optional[int], ) -> List[Dict[str, Any]]: base_flags = canonicalize_flags(base_flags) search_space = canonicalize_flags(search_space) capability = detect_current_cuda_capability() items = [(key, as_list(values)) for key, values in search_space.items()] if tier == 1: items = [(k, v[:2]) for k, v in items[:6]] elif tier == 2: items = [(k, v[:3]) for k, v in items[:8]] candidates = [deepcopy(base_flags)] if tier == 1: for key, values in items: for value in values: candidates.append(deepcopy(base_flags) | {key: value}) elif tier == 2 and items: head, tail = items[:3], items[3:] for combo in itertools.product(*[values for _, values in head]): candidate = deepcopy(base_flags) for (key, _), value in zip(head, combo): candidate[key] = value candidates.append(candidate) for key, values in tail: for value in values: candidates.append(deepcopy(base_flags) | {key: value}) elif tier == 3 and items: for combo in itertools.product(*[values for _, values in items]): candidate = deepcopy(base_flags) for (key, _), value in zip(items, combo): candidate[key] = value candidates.append(candidate) deduped: List[Dict[str, Any]] = [] seen = set() for candidate in candidates: if not is_candidate_supported_on_current_device(candidate, capability): continue key = json.dumps(candidate, sort_keys=True, ensure_ascii=False) if key in seen: continue seen.add(key) deduped.append(candidate) if max_candidates is not None and len(deduped) >= max_candidates: break return deduped def build_qps_plan( benchmark_cfg: Dict[str, Any], ) -> Tuple[str, List[float], float, int]: qps_cfg = benchmark_cfg.get("qps", benchmark_cfg.get("request_rate")) if isinstance(qps_cfg, (int, float)): return "fixed", [float(qps_cfg)], 0.0, 0 if isinstance(qps_cfg, list): return "fixed", [float(value) for value in qps_cfg], 0.0, 0 if isinstance(qps_cfg, dict) and "values" in qps_cfg: return "fixed", [float(value) for value in qps_cfg["values"]], 0.0, 0 if isinstance(qps_cfg, dict) and {"lower", "upper"} <= set(qps_cfg): return ( "search", [float(qps_cfg["lower"]), float(qps_cfg["upper"])], float(qps_cfg.get("tolerance", 0.1)), normalize_binary_search_rounds(qps_cfg.get("max_rounds")), ) raise ValueError("benchmark.qps must be a list or a {lower, upper, tolerance} map.") def trial_key( stage_name: str, candidate_id: int, request_rate: float, max_concurrency: Optional[int], server_flags: Dict[str, Any], ) -> str: return json.dumps( { "stage": stage_name, "candidate_id": candidate_id, "requested_qps": request_rate, "max_concurrency": max_concurrency, "server_flags": canonicalize_flags(server_flags), }, sort_keys=True, ensure_ascii=False, ) def record_trial_key(record: Dict[str, Any]) -> str: return trial_key( stage_name=str(record.get("stage", "")), candidate_id=int(record.get("candidate_id", 0)), request_rate=float(record.get("requested_qps", 0.0)), max_concurrency=record.get("max_concurrency"), server_flags=record.get("server_flags", {}), ) def meets_sla(result: Dict[str, Any], benchmark_cfg: Dict[str, Any]) -> bool: sla = benchmark_cfg.get("sla", {}) max_ttft_ms = sla.get("max_ttft_ms") max_tpot_ms = sla.get("max_tpot_ms") if ( max_ttft_ms is not None and result.get("mean_ttft_ms", float("inf")) > max_ttft_ms ): return False if ( max_tpot_ms is not None and result.get("mean_tpot_ms", float("inf")) > max_tpot_ms ): return False return True def result_sort_key(record: Dict[str, Any]) -> Tuple[Any, ...]: return ( 1 if record.get("sla_passed") else 0, record.get("requested_qps", 0.0), record.get("metrics", {}).get("output_throughput", 0.0), -record.get("metrics", {}).get("mean_ttft_ms", float("inf")), -record.get("metrics", {}).get("mean_tpot_ms", float("inf")), ) def launch_server( server_cfg: Dict[str, Any], server_flags: Dict[str, Any], log_path: str ) -> subprocess.Popen: command_prefix = server_cfg.get("command_prefix") if command_prefix is None: command = [sys.executable, "-m", "sglang.launch_server"] elif isinstance(command_prefix, str): command = shlex.split(command_prefix) else: command = [str(item) for item in command_prefix] command.extend(cli_args(server_flags)) command.extend(str(item) for item in server_cfg.get("extra_args", [])) env = os.environ.copy() env.update({key: str(value) for key, value in server_cfg.get("env", {}).items()}) log_file = open(log_path, "w", encoding="utf-8") try: process = subprocess.Popen( command, stdout=log_file, stderr=subprocess.STDOUT, env=env, start_new_session=True, ) except Exception: log_file.close() raise process._autobench_log_file = log_file # type: ignore[attr-defined] return process def stop_server(process: Optional[subprocess.Popen]) -> None: if process is None: return try: os.killpg(process.pid, signal.SIGTERM) process.wait(timeout=20) except Exception: try: os.killpg(process.pid, signal.SIGKILL) except Exception: pass finally: log_file = getattr(process, "_autobench_log_file", None) if log_file is not None: log_file.close() def build_bench_command( benchmark_cfg: Dict[str, Any], dataset_summary: Dict[str, Any], backend: str, base_url: str, dataset_path: str, tokenizer_path: str, request_rate: float, max_concurrency: Optional[int], output_file: str, ) -> List[str]: command = [ sys.executable, "-m", "sglang.benchmark.serving", "--backend", backend, "--base-url", base_url, "--dataset-name", "autobench", "--dataset-path", dataset_path, "--tokenizer", tokenizer_path, "--num-prompts", str(dataset_summary["num_requests"]), "--request-rate", str(request_rate), "--output-file", output_file, "--seed", str(int(benchmark_cfg.get("seed", 1))), "--ready-check-timeout-sec", str(int(benchmark_cfg.get("ready_check_timeout_sec", 600))), ] if benchmark_cfg.get("model"): command.extend(["--model", str(benchmark_cfg["model"])]) if benchmark_cfg.get("served_model_name"): command.extend(["--served-model-name", str(benchmark_cfg["served_model_name"])]) if benchmark_cfg.get("disable_tqdm", True): command.append("--disable-tqdm") if benchmark_cfg.get("output_details"): command.append("--output-details") if benchmark_cfg.get("disable_stream"): command.append("--disable-stream") if benchmark_cfg.get("disable_ignore_eos"): command.append("--disable-ignore-eos") if benchmark_cfg.get("pd_separated"): command.append("--pd-separated") if benchmark_cfg.get("flush_cache"): command.append("--flush-cache") if benchmark_cfg.get("tag"): command.extend(["--tag", str(benchmark_cfg["tag"])]) if max_concurrency is not None: command.extend(["--max-concurrency", str(max_concurrency)]) if benchmark_cfg.get("warmup_requests") is not None: command.extend( ["--warmup-requests", str(int(benchmark_cfg["warmup_requests"]))] ) if benchmark_cfg.get("extra_request_body") is not None: command.extend( [ "--extra-request-body", json.dumps(benchmark_cfg["extra_request_body"]), ] ) return command def run_bench_command( command: List[str], timeout_sec: Optional[float] = None ) -> Dict[str, Any]: try: result = subprocess.run( command, capture_output=True, text=True, timeout=timeout_sec ) except subprocess.TimeoutExpired as exc: raise SearchDeadlineExceeded( f"search budget expired while waiting for bench_serving: {exc.cmd}" ) from exc if result.returncode != 0: message = (result.stderr or result.stdout).strip() if len(message) > 4000: head = message[:2000].rstrip() tail = message[-2000:].lstrip() message = f"{head}\n...\n{tail}" raise RuntimeError(message) output_file = command[command.index("--output-file") + 1] with open(output_file, "r", encoding="utf-8") as f: lines = [line.strip() for line in f if line.strip()] if not lines: raise RuntimeError("bench_serving produced no JSONL output") return json.loads(lines[-1]) def run_trial( stage_name: str, candidate_id: int, server_cfg: Dict[str, Any], benchmark_cfg: Dict[str, Any], dataset_summary: Dict[str, Any], backend: str, dataset_path: str, tokenizer_path: str, server_flags: Dict[str, Any], output_dir: str, request_rate: float, max_concurrency: Optional[int], search_deadline: Optional[float] = None, search_budget_hours: float = DEFAULT_SEARCH_DURATION_HOURS, ) -> Dict[str, Any]: process = None log_path = os.path.join( output_dir, f"server_{stage_name}_cand{candidate_id}_mc{max_concurrency}_q{request_rate}.log", ) bench_path = os.path.join( output_dir, f"bench_{stage_name}_cand{candidate_id}_mc{max_concurrency}_q{request_rate}.jsonl", ) host = server_cfg.get("host", "127.0.0.1") port = int(server_flags.get("port", server_cfg.get("port", 30000))) base_url = benchmark_cfg.get("base_url", f"http://{host}:{port}") record = { "stage": stage_name, "candidate_id": candidate_id, "requested_qps": request_rate, "max_concurrency": max_concurrency, "server_flags": deepcopy(server_flags), "sla_passed": False, } try: raise_if_search_deadline_reached(search_deadline, search_budget_hours) if server_cfg.get("launch", True): preclean_stale_server(port) process = launch_server(server_cfg, server_flags, log_path) metrics = run_bench_command( build_bench_command( benchmark_cfg=benchmark_cfg, dataset_summary=dataset_summary, backend=backend, base_url=base_url, dataset_path=dataset_path, tokenizer_path=tokenizer_path, request_rate=request_rate, max_concurrency=max_concurrency, output_file=bench_path, ), timeout_sec=remaining_search_seconds(search_deadline), ) record["sla_passed"] = meets_sla(metrics, benchmark_cfg) record["metrics"] = metrics except SearchDeadlineExceeded: raise except Exception as exc: # noqa: BLE001 record["error"] = repr(exc) diagnosis, hint = classify_failure( "\n".join(part for part in [repr(exc), tail_text(log_path)] if part) ) if diagnosis: record["diagnosis"] = diagnosis if hint: record["hint"] = hint finally: stop_server(process) return record def merge_host_port( server_cfg: Dict[str, Any], flags: Dict[str, Any] ) -> Dict[str, Any]: merged = canonicalize_flags(deepcopy(flags)) if server_cfg.get("host") is not None and "host" not in merged: merged["host"] = server_cfg["host"] if server_cfg.get("port") is not None and "port" not in merged: merged["port"] = server_cfg["port"] return merged def run_candidate( stage_name: str, candidate_id: int, server_cfg: Dict[str, Any], benchmark_cfg: Dict[str, Any], dataset_summary: Dict[str, Any], backend: str, dataset_path: str, tokenizer_path: str, server_flags: Dict[str, Any], output_dir: str, incumbent_record: Optional[Dict[str, Any]] = None, progress_callback: Optional[Callable[[Dict[str, Any]], None]] = None, record_callback: Optional[Callable[[Dict[str, Any]], None]] = None, existing_records: Optional[Sequence[Dict[str, Any]]] = None, search_deadline: Optional[float] = None, search_budget_hours: float = DEFAULT_SEARCH_DURATION_HOURS, ) -> List[Dict[str, Any]]: mode, values, tolerance, max_rounds = build_qps_plan(benchmark_cfg) max_concurrency_values = as_list(benchmark_cfg.get("max_concurrency", [None])) records: List[Dict[str, Any]] = [] existing_by_key = { record_trial_key(record): deepcopy(record) for record in (existing_records or []) } def one_trial( request_rate: float, max_concurrency: Optional[int] ) -> Tuple[Dict[str, Any], bool]: key = trial_key( stage_name=stage_name, candidate_id=candidate_id, request_rate=request_rate, max_concurrency=max_concurrency, server_flags=server_flags, ) if key in existing_by_key: return deepcopy(existing_by_key[key]), True return ( run_trial( stage_name=stage_name, candidate_id=candidate_id, server_cfg=server_cfg, benchmark_cfg=benchmark_cfg, dataset_summary=dataset_summary, backend=backend, dataset_path=dataset_path, tokenizer_path=tokenizer_path, server_flags=server_flags, output_dir=output_dir, request_rate=request_rate, max_concurrency=max_concurrency, search_deadline=search_deadline, search_budget_hours=search_budget_hours, ), False, ) for max_concurrency in max_concurrency_values: raise_if_search_deadline_reached(search_deadline, search_budget_hours) if mode == "fixed": incumbent_qps = None if ( incumbent_record and incumbent_record.get("metrics") and incumbent_record.get("sla_passed") ): incumbent_qps = float(incumbent_record.get("requested_qps", 0.0)) for qps in values: if incumbent_qps is not None and qps < incumbent_qps: continue record, reused = one_trial(qps, max_concurrency) records.append(record) if record_callback is not None and not reused: record_callback(record) if progress_callback is not None: progress_callback(record) continue lower, upper = values best: Optional[Dict[str, Any]] = None incumbent_qps = None if ( incumbent_record and incumbent_record.get("metrics") and incumbent_record.get("sla_passed") ): incumbent_qps = float(incumbent_record.get("requested_qps", 0.0)) if incumbent_qps is not None and lower < incumbent_qps <= upper: probe_record, reused = one_trial(incumbent_qps, max_concurrency) records.append(probe_record) if record_callback is not None and not reused: record_callback(probe_record) if progress_callback is not None: progress_callback(probe_record) if probe_record.get("metrics") and probe_record["sla_passed"]: lower = max(lower, incumbent_qps) best = probe_record else: probe_record["heuristic_pruned"] = True probe_record["heuristic_reason"] = ( "Failed incumbent probe; skipped lower-QPS search because " "it cannot beat the current best candidate." ) log_line( f"[{stage_name}] heuristic prune candidate={candidate_id} " f"mc={max_concurrency} incumbent_qps={incumbent_qps:.4f}" ) continue rounds_run = 0 while upper - lower > tolerance and rounds_run < max_rounds: qps = pick_qps_midpoint(lower, upper) if qps <= lower or qps >= upper: break record, reused = one_trial(qps, max_concurrency) records.append(record) if record_callback is not None and not reused: record_callback(record) if progress_callback is not None: progress_callback(record) if record.get("metrics") and record["sla_passed"]: lower = qps best = record else: upper = qps rounds_run += 1 if best is not None: best["best_for_candidate"] = True return records def write_jsonl(path: str, records: Iterable[Dict[str, Any]]) -> None: if os.path.exists(path): os.remove(path) append_jsonl(path, records) def write_csv(path: str, records: Sequence[Dict[str, Any]]) -> None: if not records: return rows = [flatten(record) for record in records] headers = sorted({header for row in rows for header in row}) with open(path, "w", newline="", encoding="utf-8") as f: writer = csv.DictWriter(f, fieldnames=headers) writer.writeheader() writer.writerows(rows) def best_record(records: Sequence[Dict[str, Any]]) -> Optional[Dict[str, Any]]: successful = [record for record in records if record.get("metrics")] return max(successful, key=result_sort_key) if successful else None def rendered_launch_command( server_cfg: Dict[str, Any], server_flags: Dict[str, Any] ) -> str: prefix = server_cfg.get("command_prefix") if prefix is None: command = ["python", "-m", "sglang.launch_server"] elif isinstance(prefix, str): command = shlex.split(prefix) else: command = [str(item) for item in prefix] command.extend(cli_args(server_flags)) command.extend(str(item) for item in server_cfg.get("extra_args", [])) env_parts = [] for key, value in sorted(server_cfg.get("env", {}).items()): if any(marker in key.upper() for marker in SENSITIVE_ENV_MARKERS): continue env_parts.append(f"{key}={shlex.quote(str(value))}") parts: List[str] = env_parts i = 0 while i < len(command): token = str(command[i]) if token.startswith("--") and i + 1 < len(command): nxt = str(command[i + 1]) if not nxt.startswith("--"): parts.append(f"{shlex.quote(token)} {shlex.quote(nxt)}") i += 2 continue parts.append(shlex.quote(token)) i += 1 return " \\\n ".join(parts) def write_markdown_summary( path: str, scenario: Dict[str, Any], dataset_cfg: Dict[str, Any], dataset_summary: Dict[str, Any], records: Sequence[Dict[str, Any]], best: Optional[Dict[str, Any]], server_cfg: Dict[str, Any], partial_reason: Optional[str] = None, ) -> None: lines = [f"# Auto Benchmark Summary: {scenario['display_name']}", ""] lines.append(f"- Dataset kind: `{dataset_cfg['kind']}`") lines.append(f"- Requests: `{dataset_summary['num_requests']}`") if partial_reason: lines.append(f"- Status: `partial` ({partial_reason})") if dataset_cfg["kind"] == "random": lines.append( f"- Random distribution: input `{dataset_cfg['random_input_len']}`, output `{dataset_cfg['random_output_len']}`" ) lines.append("") if best is not None: lines.extend(["## Best Launch Command", "", "```bash"]) lines.append(rendered_launch_command(server_cfg, best["server_flags"])) lines.extend(["```", ""]) lines.extend( [ "## Results", "", "| Candidate | Stage | QPS | Max Conc | Prefill | Decode | TP | EP | PP | Output tok/s | TTFT ms | TPOT ms | SLA | Note |", "|---|---:|---:|---:|---|---|---:|---:|---:|---:|---:|---:|---|---|", ] ) for record in sorted(records, key=result_sort_key, reverse=True): flags = record["server_flags"] metrics = record.get("metrics", {}) note = record.get("diagnosis") or record.get("hint") or record.get("error", "") note = note.splitlines()[0][:120] if note else "" lines.append( "| {candidate_id} | {stage} | {qps} | {mc} | {prefill} | {decode} | {tp} | {ep} | {pp} | {throughput} | {ttft} | {tpot} | {sla} | {note} |".format( candidate_id=record["candidate_id"], stage=record["stage"], qps=record["requested_qps"], mc=record["max_concurrency"], prefill=flags.get("prefill_attention_backend", ""), decode=flags.get("decode_attention_backend", ""), tp=flags.get("tp_size", 1), ep=flags.get("ep_size", ""), pp=flags.get("pp_size", 1), throughput=( round(metrics.get("output_throughput", 0.0), 2) if metrics else "" ), ttft=round(metrics.get("mean_ttft_ms", 0.0), 2) if metrics else "", tpot=round(metrics.get("mean_tpot_ms", 0.0), 2) if metrics else "", sla="pass" if record.get("sla_passed") else "fail", note=note.replace("|", "/"), ) ) with open(path, "w", encoding="utf-8") as f: f.write("\n".join(lines) + "\n") def render_scenario_summary_markdown( summary_rows: Sequence[Dict[str, Any]], run_partial_reason: Optional[str] = None, ) -> str: lines = ["# Scenario Summary", ""] if run_partial_reason: lines.extend([f"- Status: `partial` ({run_partial_reason})", ""]) lines.extend( [ "| Scenario | Status | QPS | Output tok/s | TTFT ms | TPOT ms | Summary |", "|---|---|---:|---:|---:|---:|---|", ] ) for row in summary_rows: summary_path = os.path.join(row["scenario_dir"], "summary.md") lines.append( "| {name} | {status} | {qps} | {throughput} | {ttft} | {tpot} | `{path}` |".format( name=row["scenario_name"], status=row["status"], qps=row.get("requested_qps") or "", throughput=( round(row.get("output_throughput", 0.0), 2) if row.get("output_throughput") is not None else "" ), ttft=( round(row.get("mean_ttft_ms", 0.0), 2) if row.get("mean_ttft_ms") is not None else "" ), tpot=( round(row.get("mean_tpot_ms", 0.0), 2) if row.get("mean_tpot_ms") is not None else "" ), path=summary_path, ) ) for row in summary_rows: if row.get("launch_command"): lines.extend( [ "", f"## {row['scenario_name']}", "", "```bash", row["launch_command"], "```", ] ) elif row["status"] == "no_successful_runs": lines.extend( [ "", f"## {row['scenario_name']}", "", "No successful run with metrics was produced for this scenario.", ] ) return "\n".join(lines) + "\n" def run_stage( scenario_name: str, stage_name: str, candidates: Sequence[Dict[str, Any]], server_cfg: Dict[str, Any], benchmark_cfg: Dict[str, Any], dataset_summary: Dict[str, Any], backend: str, dataset_path: str, tokenizer_path: str, output_dir: str, live_results_path: Optional[str] = None, existing_records: Optional[Sequence[Dict[str, Any]]] = None, search_deadline: Optional[float] = None, search_budget_hours: float = DEFAULT_SEARCH_DURATION_HOURS, ) -> Tuple[List[Dict[str, Any]], Optional[Dict[str, Any]]]: records: List[Dict[str, Any]] = [] existing_stage_records = [ deepcopy(record) for record in (existing_records or []) if record.get("stage") == stage_name ] current_best: Optional[Dict[str, Any]] = best_record(existing_stage_records) stage_label = f"{scenario_name} {stage_name}" candidate_pbar, candidate_started_at = make_progress_bar( desc=f"{stage_label} candidates", total=len(candidates), position=1, leave=True, ) trial_pbar, trial_started_at = make_progress_bar( desc=f"{stage_label} trials", total=len(candidates) * estimate_trials_per_candidate(benchmark_cfg), position=2, leave=False, ) try: for candidate_id, candidate_flags in enumerate(candidates): raise_if_search_deadline_reached(search_deadline, search_budget_hours) merged = merge_host_port(server_cfg, candidate_flags) log_line( f"[{stage_name}] scenario={scenario_name} " f"candidate {candidate_id + 1}/{len(candidates)}: " f"{json.dumps(merged, ensure_ascii=False)}" ) def on_trial(record: Dict[str, Any]) -> None: nonlocal current_best if record.get("metrics") and ( current_best is None or result_sort_key(record) > result_sort_key(current_best) ): current_best = record advance_progress(trial_pbar, trial_started_at, best_record=current_best) refresh_progress_eta( candidate_pbar, candidate_started_at, best_record=current_best ) def on_record(record: Dict[str, Any]) -> None: if live_results_path is not None: append_jsonl(live_results_path, [record]) candidate_records = run_candidate( stage_name=stage_name, candidate_id=candidate_id, server_cfg=server_cfg, benchmark_cfg=benchmark_cfg, dataset_summary=dataset_summary, backend=backend, dataset_path=dataset_path, tokenizer_path=tokenizer_path, server_flags=merged, output_dir=output_dir, incumbent_record=current_best, progress_callback=on_trial, record_callback=on_record, existing_records=existing_stage_records, search_deadline=search_deadline, search_budget_hours=search_budget_hours, ) records.extend(candidate_records) advance_progress( candidate_pbar, candidate_started_at, best_record=current_best, ) finally: if trial_pbar.total is not None and trial_pbar.n < trial_pbar.total: trial_pbar.total = trial_pbar.n refresh_progress_eta(trial_pbar, trial_started_at, current_best) candidate_pbar.close() trial_pbar.close() return records, current_best def persist_scenario_outputs( scenario_output_dir: str, scenario: Dict[str, Any], scenario_cfg: Dict[str, Any], dataset_summary: Dict[str, Any], records: Sequence[Dict[str, Any]], server_cfg: Dict[str, Any], partial_reason: Optional[str] = None, ) -> Optional[Dict[str, Any]]: if not records: return None results_jsonl = os.path.join(scenario_output_dir, "results.jsonl") results_csv = os.path.join(scenario_output_dir, "results.csv") best = best_record(records) write_jsonl(results_jsonl, records) write_csv(results_csv, records) write_markdown_summary( path=os.path.join(scenario_output_dir, "summary.md"), scenario=scenario, dataset_cfg=scenario_cfg, dataset_summary=dataset_summary, records=records, best=best, server_cfg=server_cfg, partial_reason=partial_reason, ) log_line(f"results_jsonl={results_jsonl}") log_line(f"results_csv={results_csv}") return best def run_auto_benchmark(config_path: str) -> str: config = load_yaml(config_path) server_cfg = config["server"] benchmark_cfg = config["benchmark"] search_cfg = config.get("search", {}) timestamp = time.strftime("%Y%m%d-%H%M%S") output_dir = benchmark_cfg.get("output_dir") or os.path.join( os.getcwd(), "auto_benchmark_results", timestamp ) os.makedirs(output_dir, exist_ok=True) tokenizer_path = benchmark_cfg.get("tokenizer") or server_cfg.get( "base_flags", {} ).get("model_path") model = benchmark_cfg.get("model") or server_cfg.get("base_flags", {}).get( "model_path" ) if tokenizer_path is None: raise ValueError( "benchmark.tokenizer or server.base_flags.model_path is required." ) dataset_cfg = normalize_dataset_cfg(config.get("dataset"), benchmark_cfg) scenarios = expand_dataset_scenarios(dataset_cfg) tier = int(search_cfg.get("tier", 2)) max_candidates = resolve_max_candidates(search_cfg) resume_enabled = bool(search_cfg.get("resume", True)) base_candidates = build_server_candidates(server_cfg, tier, max_candidates) search_budget_hours = resolve_search_budget_hours(search_cfg) search_deadline = time.time() + (search_budget_hours * 3600) scenario_records: List[Dict[str, Any]] = [] interrupted = False run_partial_reason: Optional[str] = None print_run_plan( config_path=config_path, output_dir=output_dir, tier=tier, max_candidates=max_candidates, benchmark_cfg=benchmark_cfg, scenarios=scenarios, server_cfg=server_cfg, base_candidates=base_candidates, speculative_enabled=bool(config.get("speculative", {}).get("enabled")), search_budget_hours=search_budget_hours, search_deadline=search_deadline, ) scenario_pbar, scenario_started_at = make_progress_bar( desc="scenarios", total=len(scenarios), position=0, leave=True, ) previous_handlers = install_interrupt_handlers() try: for scenario in scenarios: raise_if_search_deadline_reached(search_deadline, search_budget_hours) scenario_output_dir = ( output_dir if len(scenarios) == 1 else os.path.join(output_dir, scenario["name"]) ) os.makedirs(scenario_output_dir, exist_ok=True) live_results_path = os.path.join(scenario_output_dir, "live_results.jsonl") if os.path.exists(live_results_path) and not resume_enabled: os.remove(live_results_path) prepared_dataset_path = os.path.join( scenario_output_dir, "prepared_dataset.jsonl" ) existing_records = read_jsonl(live_results_path) if resume_enabled and os.path.exists(prepared_dataset_path): rows = load_autobench_rows( dataset_path=prepared_dataset_path, tokenizer_path=tokenizer_path, num_prompts=0, ) dataset_summary = summarize_rows(rows) else: prepared_dataset_path, rows, dataset_summary = prepare_dataset( dataset_cfg=scenario["cfg"], tokenizer_path=tokenizer_path, model=model, output_path=prepared_dataset_path, ) backend = infer_backend(benchmark_cfg.get("backend", "auto"), rows) log_line(f"scenario={scenario['display_name']}") log_line(f"prepared_dataset={prepared_dataset_path}") log_line( f"dataset_summary={json.dumps(dataset_summary, ensure_ascii=False)}" ) log_line(f"selected_backend={backend}") if resume_enabled and existing_records: log_line( f"resume=true loaded_records={len(existing_records)} " f"scenario={scenario['display_name']}" ) all_records: List[Dict[str, Any]] = [] scenario_partial_reason: Optional[str] = None try: all_records, best_base = run_stage( scenario_name=scenario["display_name"], stage_name="base", candidates=base_candidates, server_cfg=server_cfg, benchmark_cfg=benchmark_cfg, dataset_summary=dataset_summary, backend=backend, dataset_path=prepared_dataset_path, tokenizer_path=tokenizer_path, output_dir=scenario_output_dir, live_results_path=live_results_path, existing_records=existing_records, search_deadline=search_deadline, search_budget_hours=search_budget_hours, ) speculative_cfg = config.get("speculative", {}) if speculative_cfg.get("enabled"): if best_base is None: raise ValueError( "Speculative search requires at least one successful base run." ) if not speculative_cfg.get("draft_model_path"): raise ValueError("speculative.draft_model_path is required.") spec_base_flags = deepcopy(best_base["server_flags"]) spec_base_flags.update( deepcopy(speculative_cfg.get("base_flags", {})) ) spec_base_flags["speculative_algorithm"] = speculative_cfg.get( "algorithm", "EAGLE" ) spec_base_flags["speculative_draft_model_path"] = speculative_cfg[ "draft_model_path" ] spec_candidates = build_candidates( base_flags=canonicalize_flags(spec_base_flags), search_space=deepcopy(speculative_cfg.get("search_space", {})), tier=tier, max_candidates=max_candidates, ) log_line( f"Planned speculative candidates for scenario={scenario['display_name']}:" ) for index, candidate in enumerate(spec_candidates, start=1): log_line( f" [{index}/{len(spec_candidates)}] " f"{json.dumps(merge_host_port(server_cfg, candidate), ensure_ascii=False)}" ) spec_records, _ = run_stage( scenario_name=scenario["display_name"], stage_name="speculative", candidates=spec_candidates, server_cfg=server_cfg, benchmark_cfg=benchmark_cfg, dataset_summary=dataset_summary, backend=backend, dataset_path=prepared_dataset_path, tokenizer_path=tokenizer_path, output_dir=scenario_output_dir, live_results_path=live_results_path, existing_records=read_jsonl(live_results_path), search_deadline=search_deadline, search_budget_hours=search_budget_hours, ) all_records.extend(spec_records) except SearchDeadlineExceeded as exc: interrupted = True scenario_partial_reason = str(exc) run_partial_reason = scenario_partial_reason log_line( f"search_deadline_reached=true scenario={scenario['display_name']} " f"detail={scenario_partial_reason}" ) except KeyboardInterrupt: interrupted = True scenario_partial_reason = "interrupted before the full search completed" run_partial_reason = scenario_partial_reason log_line( f"interrupt_received=true scenario={scenario['display_name']} " "saving partial results before exit" ) finally: persisted_records = all_records live_records = read_jsonl(live_results_path) if len(live_records) > len(persisted_records): persisted_records = live_records best = persist_scenario_outputs( scenario_output_dir=scenario_output_dir, scenario=scenario, scenario_cfg=scenario["cfg"], dataset_summary=dataset_summary, records=persisted_records, server_cfg=server_cfg, partial_reason=scenario_partial_reason, ) if persisted_records: scenario_records.append( { "scenario_name": scenario["display_name"], "scenario_dir": scenario_output_dir, "best_record": best, "has_records": True, } ) if interrupted: break advance_progress(scenario_pbar, scenario_started_at) except SearchDeadlineExceeded as exc: interrupted = True run_partial_reason = str(exc) log_line(f"search_deadline_reached=true detail={run_partial_reason}") finally: scenario_pbar.close() restore_interrupt_handlers(previous_handlers) if scenario_records and len(scenarios) > 1: summary_rows = [] for item in scenario_records: record = item["best_record"] metrics = record.get("metrics", {}) if record else {} summary_rows.append( { "scenario_name": item["scenario_name"], "scenario_dir": item["scenario_dir"], "status": ( "ok" if record and record.get("metrics") else "no_successful_runs" ), "requested_qps": record.get("requested_qps") if record else None, "mean_ttft_ms": metrics.get("mean_ttft_ms"), "mean_tpot_ms": metrics.get("mean_tpot_ms"), "output_throughput": metrics.get("output_throughput"), "launch_command": ( rendered_launch_command(server_cfg, record["server_flags"]) if record else "" ), } ) write_jsonl(os.path.join(output_dir, "scenario_summary.jsonl"), summary_rows) write_csv(os.path.join(output_dir, "scenario_summary.csv"), summary_rows) with open(os.path.join(output_dir, "SUMMARY.md"), "w", encoding="utf-8") as f: f.write(render_scenario_summary_markdown(summary_rows, run_partial_reason)) if interrupted: log_line(f"interrupted=true partial_output_dir={output_dir}") return output_dir def convert_dataset(args: argparse.Namespace) -> None: dataset_cfg = normalize_dataset_cfg( { key: value for key, value in vars(args).items() if key not in {"command", "output", "tokenizer", "model"} }, {}, ) output_path, rows, summary = prepare_dataset( dataset_cfg=dataset_cfg, tokenizer_path=args.tokenizer, model=args.model, output_path=args.output, ) print(f"prepared_dataset={output_path}") print(f"rows={len(rows)}") print(json.dumps(summary, ensure_ascii=False, indent=2)) def validate_dataset(args: argparse.Namespace) -> None: rows = load_autobench_rows(args.dataset_path, args.tokenizer, num_prompts=0) print(json.dumps(summarize_rows(rows), ensure_ascii=False, indent=2))