[FEAT] Add optional prewarm before trace replay benchmarks

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)."""

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}",
+    ]

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.

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