diff --git a/auto_tune_vllm/benchmarks/config.py b/auto_tune_vllm/benchmarks/config.py index 4068136..fdf7ea8 100644 --- a/auto_tune_vllm/benchmarks/config.py +++ b/auto_tune_vllm/benchmarks/config.py @@ -54,6 +54,9 @@ class BenchmarkConfig: # GuideLLM benchmark profile (concurrent or replay). Omitted => concurrent. profile: Optional[dict] = None + # Optional kernel prewarm before trace replay (guidellm_trace_replay only). + prewarm: Optional[dict] = None + # Set in benchmark section of study config # Logging level for GuideLLM logging_level: Literal["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"] = "INFO" @@ -91,10 +94,44 @@ def __post_init__(self) -> None: if self.rate <= 0: raise ValueError(f"benchmark.rate must be greater than 0; got {self.rate}") - from .profiles import profile_from_dict + self._validate_prewarm() + if self.profile is not None: + from .profiles import profile_from_dict + + profile = profile_from_dict(self.profile) + profile.validate(self) + + def _validate_prewarm(self) -> None: + """Validate optional prewarm settings for trace replay benchmarks.""" + if self.prewarm is None: + return + + if self.benchmark_type != "guidellm_trace_replay": + raise ValueError( + "benchmark.prewarm is only supported with " + "benchmark_type='guidellm_trace_replay'" + ) + + allowed_keys = frozenset({"duration", "concurrency"}) + unknown_keys = set(self.prewarm.keys()) - allowed_keys + if unknown_keys: + raise ValueError( + "benchmark.prewarm supports only duration and concurrency; " + f"got unexpected keys: {sorted(unknown_keys)}" + ) - profile = profile_from_dict(self.profile) - profile.validate(self) + for key in ("duration", "concurrency"): + if key not in self.prewarm: + raise ValueError(f"benchmark.prewarm.{key} is required") + value = self.prewarm[key] + if not isinstance(value, (int, float)): + raise ValueError( + f"benchmark.prewarm.{key} must be a number; got {value!r}" + ) + if value <= 0: + raise ValueError( + f"benchmark.prewarm.{key} must be greater than 0; got {value}" + ) def _normalize_benchmark_profile(self) -> None: """Align ``benchmark.profile`` with ``benchmark_type`` (single source of truth).""" diff --git a/auto_tune_vllm/benchmarks/profiles.py b/auto_tune_vllm/benchmarks/profiles.py index 5992ce9..0066e34 100644 --- a/auto_tune_vllm/benchmarks/profiles.py +++ b/auto_tune_vllm/benchmarks/profiles.py @@ -3,11 +3,15 @@ from __future__ import annotations import json +import logging import os from abc import ABC, abstractmethod from dataclasses import dataclass +from pathlib import Path from typing import Any, Literal +import pandas as pd + from .config import BenchmarkConfig TRACE_FORMATS = frozenset({"trace_synthetic", "mooncake"}) @@ -153,3 +157,147 @@ def render_replay_data_cli( ] ) return cmd + + +def _read_trace_dataframe(dataset_path: str) -> pd.DataFrame: + """Load a local trace file into a pandas DataFrame.""" + path = Path(dataset_path) + suffix = path.suffix.lower() + if suffix == ".jsonl": + return pd.read_json(dataset_path, lines=True) + if suffix == ".json": + return pd.read_json(dataset_path) + if suffix == ".csv": + return pd.read_csv(dataset_path) + if suffix == ".parquet": + return pd.read_parquet(dataset_path) + raise ValueError( + f"Unsupported trace dataset format {suffix!r}; " + "expected .jsonl, .json, .csv, or .parquet" + ) + + +def default_prewarm_token_stats(config: BenchmarkConfig) -> dict[str, int]: + """Fallback token statistics for prewarm synthetic data from benchmark defaults.""" + prompt_stdev = ( + config.prompt_tokens_stdev if config.prompt_tokens_stdev is not None else 1 + ) + output_stdev = ( + config.output_tokens_stdev if config.output_tokens_stdev is not None else 1 + ) + return { + "prompt_mean": max(1, config.prompt_tokens), + "prompt_stdev": max(1, prompt_stdev), + "output_mean": max(1, config.output_tokens), + "output_stdev": max(1, output_stdev), + } + + +def resolve_prewarm_token_stats( + config: BenchmarkConfig, + profile: ReplayProfile, + logger: logging.Logger | None = None, +) -> dict[str, int]: + """Derive prewarm token stats from a local trace file, or fall back to defaults.""" + dataset = config.dataset + if dataset is None: + raise ValueError("Trace replay requires benchmark.dataset to be set") + + if dataset.startswith("hf://"): + if logger is not None: + logger.warning( + "Cannot derive prewarm token stats from HuggingFace dataset %r; " + "using benchmark.prompt_tokens / benchmark.output_tokens defaults", + dataset, + ) + return default_prewarm_token_stats(config) + + if not os.path.exists(dataset): + if logger is not None: + logger.warning( + "Trace dataset file not found at %r; using benchmark.prompt_tokens / " + "benchmark.output_tokens defaults for prewarm", + dataset, + ) + return default_prewarm_token_stats(config) + + try: + return compute_trace_token_stats( + dataset, + profile.prompt_tokens_column, + profile.output_tokens_column, + ) + except (ValueError, FileNotFoundError) as exc: + if logger is not None: + logger.warning( + "Failed to derive prewarm token stats from trace file: %s; " + "using benchmark.prompt_tokens / benchmark.output_tokens defaults", + exc, + ) + return default_prewarm_token_stats(config) + + +def compute_trace_token_stats( + dataset_path: str, + prompt_col: str, + output_col: str, +) -> dict[str, int]: + """Compute mean and stdev of prompt/output token lengths from a trace file.""" + if not os.path.exists(dataset_path): + raise FileNotFoundError(f"Trace dataset file not found: {dataset_path}") + + df = _read_trace_dataframe(dataset_path) + for col in (prompt_col, output_col): + if col not in df.columns: + raise ValueError( + f"Trace dataset missing required column {col!r}; " + f"available columns: {list(df.columns)}" + ) + + prompt_series = df[prompt_col] + output_series = df[output_col] + + prompt_mean = max(1, int(round(prompt_series.mean()))) + output_mean = max(1, int(round(output_series.mean()))) + prompt_stdev = max(1, int(round(prompt_series.std(ddof=0) or 0.0))) + output_stdev = max(1, int(round(output_series.std(ddof=0) or 0.0))) + + return { + "prompt_mean": prompt_mean, + "prompt_stdev": prompt_stdev, + "output_mean": output_mean, + "output_stdev": output_stdev, + } + + +def render_prewarm_args( + config: BenchmarkConfig, + prewarm: dict[str, Any], + stats: dict[str, int], + results_file: str, +) -> list[str]: + """Build GuideLLM CLI args for a concurrent prewarm run before trace replay.""" + duration = prewarm["duration"] + concurrency = prewarm["concurrency"] + data_config: dict[str, Any] = { + "kind": "synthetic_text", + "prompt_tokens": stats["prompt_mean"], + "output_tokens": stats["output_mean"], + "prompt_tokens_stdev": stats["prompt_stdev"], + "output_tokens_stdev": stats["output_stdev"], + } + + return [ + "--profile", + f"kind=concurrent,streams={concurrency}", + "--constraint", + f"kind=max_duration,seconds={duration}", + "--metrics", + f"kind=generative,sample_size={config.sample_requests}", + "--output", + f"kind=json,path={results_file}", + "--data", + json.dumps(data_config), + "--data-loader", + f"kind=pytorch,samples={config.samples}", + ] diff --git a/auto_tune_vllm/benchmarks/providers.py b/auto_tune_vllm/benchmarks/providers.py index de28eeb..b456db3 100644 --- a/auto_tune_vllm/benchmarks/providers.py +++ b/auto_tune_vllm/benchmarks/providers.py @@ -233,7 +233,7 @@ def _validate_runtime(self, model_url: str, config: BenchmarkConfig) -> None: raise ValueError(f"Invalid model_url: {model_url!r} (expected http/https)") def _launch_subprocess( - self, cmd: list[str], config: BenchmarkConfig + self, cmd: list[str], config: BenchmarkConfig, log_path: str | None = None ) -> subprocess.Popen: """Launch a benchmark subprocess with the standard GuideLLM settings. @@ -244,7 +244,10 @@ def _launch_subprocess( env = os.environ.copy() env["GUIDELLM__LOGGING__CONSOLE_LOG_LEVEL"] = config.logging_level - self._benchmark_log_path = str(Path(self._results_file).with_suffix(".log")) + if log_path is not None: + self._benchmark_log_path = log_path + else: + self._benchmark_log_path = str(Path(self._results_file).with_suffix(".log")) # Parent handle can be closed right after Popen; the child keeps its # own duplicated fd. diff --git a/auto_tune_vllm/benchmarks/trace_replay.py b/auto_tune_vllm/benchmarks/trace_replay.py index 83d0d70..0e82014 100644 --- a/auto_tune_vllm/benchmarks/trace_replay.py +++ b/auto_tune_vllm/benchmarks/trace_replay.py @@ -2,14 +2,47 @@ from __future__ import annotations +import json +import signal +import subprocess +from pathlib import Path + from .config import BenchmarkConfig -from .profiles import ReplayProfile, profile_from_dict, render_replay_data_cli +from .profiles import ( + ReplayProfile, + profile_from_dict, + render_prewarm_args, + render_replay_data_cli, + resolve_prewarm_token_stats, +) from .providers import GuideLLMBenchmark class GuideLLMTraceReplayBenchmark(GuideLLMBenchmark): """Benchmark provider that replays trace files via GuideLLM ``kind=replay``.""" + def start_benchmark( + self, model_url: str, config: BenchmarkConfig + ) -> subprocess.Popen: + """Start trace replay benchmark, optionally after a concurrent prewarm run.""" + self._logger.info( + f"Starting GuideLLM trace replay benchmark for {config.model}" + ) + + self._validate_runtime(model_url, config) + + self._results_file = self._get_results_file_path() + + if config.prewarm is not None: + self._run_prewarm(model_url, config) + + cmd = self._build_guidellm_command(model_url, config, self._results_file) + + self._logger.info(f"Running: {' '.join(cmd)}") + self._logger.info(f"Results will be saved to: {self._results_file}") + + return self._launch_subprocess(cmd, config) + def _build_guidellm_command( self, model_url: str, config: BenchmarkConfig, results_file: str ) -> list[str]: @@ -22,3 +55,101 @@ def _build_guidellm_command( cmd = self._build_guidellm_base_command(model_url, config, results_file) cmd.extend(render_replay_data_cli(config, profile)) return cmd + + def _build_prewarm_command( + self, + model_url: str, + config: BenchmarkConfig, + profile: ReplayProfile, + prewarm_results_file: str, + ) -> list[str]: + """Build a short concurrent GuideLLM run to warm up vLLM kernels.""" + stats = resolve_prewarm_token_stats(config, profile, self._logger) + processor = config.processor if config.processor is not None else config.model + + cmd = [ + "guidellm", + "run", + "--backend", + f"kind=openai_http,target={model_url},model={config.model}", + "--tokenizer", + json.dumps( + { + "kind": "huggingface_auto", + "model": processor, + "load_kwargs": {"trust_remote_code": True}, + } + ), + ] + cmd.extend( + render_prewarm_args(config, config.prewarm, stats, prewarm_results_file) + ) + return cmd + + def _prewarm_results_path(self) -> str: + results_path = Path(self._results_file) + if results_path.name.endswith("_benchmark_results.json"): + return str( + results_path.with_name( + results_path.name.replace( + "_benchmark_results.json", "_prewarm_results.json" + ) + ) + ) + return str( + results_path.with_name(f"{results_path.stem}_prewarm{results_path.suffix}") + ) + + def _run_prewarm(self, model_url: str, config: BenchmarkConfig) -> None: + profile = profile_from_dict(config.profile) + if not isinstance(profile, ReplayProfile): + raise ValueError( + "benchmark_type 'guidellm_trace_replay' requires profile.kind='replay'" + ) + + prewarm_results_file = self._prewarm_results_path() + prewarm_log_path = str(Path(prewarm_results_file).with_suffix(".prewarm.log")) + cmd = self._build_prewarm_command( + model_url, config, profile, prewarm_results_file + ) + + duration = float(config.prewarm["duration"]) + self._logger.info( + f"Running prewarm ({duration:.0f}s, " + f"concurrency={config.prewarm['concurrency']}): {' '.join(cmd)}" + ) + + process = self._launch_subprocess(cmd, config, log_path=prewarm_log_path) + timeout = max(duration * 2, duration + 30) + + try: + returncode = process.wait(timeout=timeout) + except subprocess.TimeoutExpired: + self._logger.error( + f"Prewarm did not finish within {timeout:.0f}s; terminating prewarm" + ) + self.terminate_benchmark() + self._clear_process_handles() + raise RuntimeError( + f"Prewarm did not finish within {timeout:.0f}s; trial aborted" + ) + + self._clear_process_handles() + + if returncode == 0: + self._logger.info("Prewarm completed successfully") + return + + if returncode < 0 and -returncode in (signal.SIGTERM, signal.SIGKILL): + raise KeyboardInterrupt("Prewarm cancelled") + + log_tail = self.get_last_log_lines() + raise RuntimeError( + f"Prewarm failed with exit code {returncode}; trial aborted. " + f"Log tail:\n{log_tail}" + ) + + def _clear_process_handles(self) -> None: + self._process = None + self._process_pid = None + self._process_pgid = None diff --git a/docs/configuration.md b/docs/configuration.md index 8484c73..9328bae 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -399,7 +399,36 @@ Prompts are generated synthetically to match `input_length`; the replay profile | `profile.hash_ids_column` | Mooncake-only hash ID list column (default: `"hash_ids"`) | | `profile.hash_id_block_size` | Mooncake-only tokens per hash ID block (default: `512`) | -`warmup`, `cooldown`, and `rampup` are not supported with the replay profile. Baseline trials use the same `rate` / `time_scale` semantics as optimization trials. +`warmup`, `cooldown`, and `rampup` are not supported with the replay profile. Replay speed is controlled by `benchmark.rate` (or `profile.time_scale`), not by `baseline.concurrency_levels` — see [Baseline Configuration](#baseline-configuration). + +#### Prewarm (kernel warmup) + +Optional pre-run phase for `guidellm_trace_replay` only. Distinct from `benchmark.warmup`: prewarm runs **before** the replay benchmark as a separate short GuideLLM subprocess, so vLLM kernel autotuners (CUDA graphs, torch.compile, etc.) can settle before timed replay requests. + +When `benchmark.prewarm` is set, auto-tune-vllm: + +1. Derives mean/stdev of prompt and output token lengths from the local trace file when possible (using `profile.prompt_tokens_column` and `profile.output_tokens_column`). If the dataset is remote (`hf://`) or cannot be read locally, a warning is logged and `benchmark.prompt_tokens` / `benchmark.output_tokens` (plus optional stdev fields) are used instead. +2. Launches a concurrent GuideLLM run (`kind=concurrent`) for `duration` seconds with `concurrency` streams. +3. Uses `kind=synthetic_text` data whose token mean/stdev match those statistics. +4. Discards prewarm results and continues with the normal trace replay run only when prewarm succeeds. + +If prewarm fails (non-zero exit, timeout, or cancellation), the trial is marked as failed and trace replay does **not** run. + +```yaml +benchmark: + benchmark_type: "guidellm_trace_replay" + dataset: "examples/trace_replay/sample.jsonl" + prewarm: + duration: 30 # seconds for the prewarm run + concurrency: 4 # concurrent streams during prewarm +``` + +| Field | Description | +|-------|-------------| +| `prewarm.duration` | Wall-clock seconds for the prewarm subprocess (`> 0`) | +| `prewarm.concurrency` | Concurrent streams during prewarm (`> 0`) | + +Token statistics are derived from the trace file when it is available locally; otherwise defaults from `benchmark.prompt_tokens` / `benchmark.output_tokens` apply. See [examples/study_config_trace_replay.yaml](../examples/study_config_trace_replay.yaml) and [examples/trace_replay/sample.jsonl](../examples/trace_replay/sample.jsonl) for a full example. @@ -613,7 +642,7 @@ baseline: #### Configuration Fields: - **`enabled`** (boolean, default: `true`): Enable baseline trials. Set to `false` to disable. -- **`concurrency_levels`** (array, optional): List of concurrency levels to test baseline performance. If not specified, defaults to `[benchmark.rate]`. +- **`concurrency_levels`** (array, optional): For concurrent benchmarks (`guidellm`, `guidellm_multimodal`), list of concurrency levels to run as separate baseline trials. If not specified, defaults to `[benchmark.rate]`. For trace replay (`guidellm_trace_replay`), defaults to `[1]` (one baseline trial). **This field does not set replay speed** — replay `time_scale` comes from `benchmark.rate` / `profile.time_scale` instead. Baseline trials always reuse the benchmark section as-is; `concurrency_levels` only controls how many baseline trials are enqueued and affects `--max-num-seqs` when a level exceeds 256. - **`parameters`** (dict, optional): Custom vLLM parameters to use for all baseline trials - Will override static_parameters if defined in both ### Baseline Trial Behavior diff --git a/examples/study_config_trace_replay.yaml b/examples/study_config_trace_replay.yaml index f56e115..9a979bb 100644 --- a/examples/study_config_trace_replay.yaml +++ b/examples/study_config_trace_replay.yaml @@ -22,6 +22,12 @@ benchmark: dataset: "examples/trace_replay/sample.jsonl" profile: trace_format: trace_synthetic + # Optional prewarm: heats up vLLM kernels before the real replay run. + # Uses a short concurrent GuideLLM benchmark on synthetic data whose + # prompt/output token mean+stdev are auto-derived from the trace file. + # prewarm: + # duration: 30 + # concurrency: 4 logging: file_path: "./auto-tune-vllm-logs" diff --git a/tests/benchmarks/test_guidellm_trace_replay.py b/tests/benchmarks/test_guidellm_trace_replay.py index 9669857..49f093c 100644 --- a/tests/benchmarks/test_guidellm_trace_replay.py +++ b/tests/benchmarks/test_guidellm_trace_replay.py @@ -3,14 +3,21 @@ from __future__ import annotations import json +import logging +import subprocess +from unittest.mock import MagicMock, patch import pytest from auto_tune_vllm.benchmarks.config import BenchmarkConfig from auto_tune_vllm.benchmarks.profiles import ( ReplayProfile, + compute_trace_token_stats, + default_prewarm_token_stats, profile_from_dict, + render_prewarm_args, render_replay_data_cli, + resolve_prewarm_token_stats, ) from auto_tune_vllm.benchmarks.trace_replay import GuideLLMTraceReplayBenchmark @@ -173,3 +180,163 @@ def test_trace_replay_data_loader_when_data_samples_set(): assert "--data-loader" in cmd assert cmd[cmd.index("--data-loader") + 1] == "kind=pytorch,samples=100" + + +def test_prewarm_rejected_on_non_trace_replay_benchmark(): + with pytest.raises(ValueError, match="benchmark.prewarm is only supported"): + BenchmarkConfig( + model="test-model", + prewarm={"duration": 30, "concurrency": 4}, + ) + + +def test_prewarm_requires_duration_and_concurrency(): + with pytest.raises(ValueError, match="benchmark.prewarm.duration is required"): + BenchmarkConfig( + benchmark_type="guidellm_trace_replay", + model="test-model", + dataset=_TRACE_DATASET, + prewarm={"concurrency": 4}, + ) + + +def test_compute_trace_token_stats_from_sample_trace(): + stats = compute_trace_token_stats(_TRACE_DATASET, "input_length", "output_length") + + assert stats["prompt_mean"] >= 1 + assert stats["output_mean"] >= 1 + assert stats["prompt_stdev"] >= 1 + assert stats["output_stdev"] >= 1 + assert isinstance(stats["prompt_stdev"], int) + assert isinstance(stats["output_stdev"], int) + + +def test_resolve_prewarm_token_stats_uses_trace_file(): + config = BenchmarkConfig( + benchmark_type="guidellm_trace_replay", + model="test-model", + dataset=_TRACE_DATASET, + ) + profile = ReplayProfile() + + stats = resolve_prewarm_token_stats(config, profile) + + assert stats == compute_trace_token_stats( + _TRACE_DATASET, profile.prompt_tokens_column, profile.output_tokens_column + ) + + +def test_resolve_prewarm_token_stats_hf_dataset_uses_defaults(caplog): + config = BenchmarkConfig( + benchmark_type="guidellm_trace_replay", + model="test-model", + dataset="hf://demo/trace", + prompt_tokens=512, + output_tokens=128, + ) + profile = ReplayProfile() + logger = logging.getLogger("test.prewarm.hf") + + with caplog.at_level(logging.WARNING, logger="test.prewarm.hf"): + stats = resolve_prewarm_token_stats(config, profile, logger) + + assert stats == default_prewarm_token_stats(config) + assert "HuggingFace dataset" in caplog.text + + +def test_resolve_prewarm_token_stats_missing_file_uses_defaults(caplog): + config = BenchmarkConfig( + benchmark_type="guidellm_trace_replay", + model="test-model", + dataset="/tmp/nonexistent_trace.jsonl", + prompt_tokens=256, + output_tokens=64, + ) + profile = ReplayProfile() + logger = logging.getLogger("test.prewarm.missing") + + with caplog.at_level(logging.WARNING, logger="test.prewarm.missing"): + stats = resolve_prewarm_token_stats(config, profile, logger) + + assert stats == default_prewarm_token_stats(config) + assert "not found" in caplog.text + + +def test_render_prewarm_args_uses_token_stats(): + config = BenchmarkConfig( + benchmark_type="guidellm_trace_replay", + model="test-model", + dataset=_TRACE_DATASET, + samples=50, + sample_requests=0, + ) + stats = { + "prompt_mean": 400, + "prompt_stdev": 120, + "output_mean": 80, + "output_stdev": 30, + } + args = render_prewarm_args( + config, {"duration": 15, "concurrency": 2}, stats, "/tmp/prewarm.json" + ) + + assert args[args.index("--profile") + 1] == "kind=concurrent,streams=2" + assert args[args.index("--constraint") + 1] == "kind=max_duration,seconds=15" + data = json.loads(args[args.index("--data") + 1]) + assert data["prompt_tokens"] == 400 + assert data["output_tokens"] == 80 + assert data["prompt_tokens_stdev"] == 120 + assert data["output_tokens_stdev"] == 30 + assert isinstance(data["prompt_tokens_stdev"], int) + assert isinstance(data["output_tokens_stdev"], int) + + +def test_run_prewarm_failure_aborts_trial(): + benchmark = GuideLLMTraceReplayBenchmark() + benchmark._logger = logging.getLogger("test.prewarm.fail") + benchmark._results_file = "/tmp/trial_benchmark_results.json" + + config = BenchmarkConfig( + benchmark_type="guidellm_trace_replay", + model="test-model", + dataset=_TRACE_DATASET, + prewarm={"duration": 10, "concurrency": 2}, + ) + + mock_process = MagicMock() + mock_process.wait.return_value = 1 + + with patch.object( + GuideLLMTraceReplayBenchmark, + "_launch_subprocess", + return_value=mock_process, + ): + with pytest.raises(RuntimeError, match="Prewarm failed with exit code 1"): + benchmark._run_prewarm("http://localhost:8000/v1", config) + + +def test_run_prewarm_timeout_aborts_trial(): + benchmark = GuideLLMTraceReplayBenchmark() + benchmark._logger = logging.getLogger("test.prewarm.timeout") + benchmark._results_file = "/tmp/trial_benchmark_results.json" + + config = BenchmarkConfig( + benchmark_type="guidellm_trace_replay", + model="test-model", + dataset=_TRACE_DATASET, + prewarm={"duration": 10, "concurrency": 2}, + ) + + mock_process = MagicMock() + mock_process.wait.side_effect = subprocess.TimeoutExpired( + cmd=["guidellm"], timeout=50 + ) + + with patch.object( + GuideLLMTraceReplayBenchmark, + "_launch_subprocess", + return_value=mock_process, + ): + with patch.object(GuideLLMTraceReplayBenchmark, "terminate_benchmark"): + with pytest.raises(RuntimeError, match="Prewarm did not finish"): + benchmark._run_prewarm("http://localhost:8000/v1", config)