feat: add aiperf_package config for controlling aiperf version

Add benchmark.aiperf_package field to specify the pip install spec for
aiperf (e.g., "aiperf>=0.7.0"). Passed as AIPERF_PACKAGE env var to
bench.sh which does pip install --upgrade. Defaults to "aiperf" if not
set. Always installs tiktoken alongside.

Needed because container-bundled aiperf may predate fixes like
trust-remote-code propagation to pool workers (aiperf PR #744).

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: alec-flowers

src/srtctl/benchmarks/base.py

@@ -61,7 +61,20 @@ def build_command(
 
 
 class AIPerfBenchmarkRunner(BenchmarkRunner):
-    """Marker base class for AIPerf-driven benchmarks."""
+    """Base class for AIPerf-driven benchmarks.
+
+    Provides shared aiperf_args handling for subclasses.
+    """
+
+    def append_aiperf_args(self, cmd: list[str], config: SrtConfig) -> list[str]:
+        """Append aiperf_args from config as CLI flags."""
+        for key, value in config.benchmark.aiperf_args.items():
+            if isinstance(value, bool):
+                if value:
+                    cmd.append(f"--{key}")
+            else:
+                cmd.extend([f"--{key}", str(value)])
+        return cmd
 
 
 # Registry of benchmark runners

src/srtctl/benchmarks/mooncake_router.py

@@ -101,7 +101,7 @@ def build_command(
         # For HF models, use the model ID directly so transformers downloads it
         tokenizer_path = str(runtime.model_path) if runtime.is_hf_model else "/model"
 
-        return [
+        cmd = [
             "bash",
             self.script_path,
             endpoint,
@@ -111,3 +111,7 @@ def build_command(
             str(itl_threshold),
             tokenizer_path,
         ]
+
+        self.append_aiperf_args(cmd, config)
+
+        return cmd

src/srtctl/benchmarks/scripts/trace-replay/bench.sh

@@ -67,14 +67,24 @@ if [ ! -f "${TRACE_FILE}" ]; then
     exit 1
 fi
 
-# Install aiperf if not present
-if ! command -v aiperf &> /dev/null; then
-    echo "Installing aiperf..."
-    pip install aiperf
+# Create isolated aiperf environment (avoids polluting container packages)
+# AIPERF_PACKAGE env var controls the version (e.g., "aiperf>=0.7.0")
+AIPERF_SPEC="${AIPERF_PACKAGE:-aiperf}"
+AIPERF_VENV="/tmp/aiperf-${SLURM_JOB_ID:-$$}"
+
+echo "Setting up aiperf environment: ${AIPERF_SPEC}"
+
+# Install uv if not in container
+if ! command -v uv &> /dev/null; then
+    echo "Installing uv..."
+    curl -LsSf https://astral.sh/uv/install.sh | sh
+    export PATH="$HOME/.local/bin:$PATH"
 fi
 
-# Install tiktoken if not present (needed for custom tokenizers like Kimi)
-python3 -c "import tiktoken" 2>/dev/null || pip install tiktoken
+uv venv "${AIPERF_VENV}"
+uv pip install -p "${AIPERF_VENV}" "${AIPERF_SPEC}" tiktoken
+export PATH="${AIPERF_VENV}/bin:${PATH}"
+echo "aiperf $(aiperf --version 2>/dev/null || echo 'installed') in ${AIPERF_VENV}"
 
 # Run small benchmark for warmup
 echo "Running warmup..."

src/srtctl/benchmarks/trace_replay.py

@@ -94,12 +94,6 @@ def build_command(
             tokenizer_path,
         ]
 
-        # Pass through extra aiperf CLI flags from config
-        for key, value in b.aiperf_args.items():
-            if isinstance(value, bool):
-                if value:
-                    cmd.append(f"--{key}")
-            else:
-                cmd.extend([f"--{key}", str(value)])
+        self.append_aiperf_args(cmd, config)
 
         return cmd

src/srtctl/cli/mixins/benchmark_stage.py

@@ -285,8 +285,10 @@ def _get_benchmark_env(self, runner: "BenchmarkRunner") -> dict[str, str]:
         env = self._get_benchmark_profiling_env(runner)
         env["SRTCTL_FRONTEND_TYPE"] = self.config.frontend.type
 
-        # Add AIPerf metrics URLs for AIPerf-driven benchmarks
+        # Add AIPerf-specific env vars for AIPerf-driven benchmarks only
         if isinstance(runner, AIPerfBenchmarkRunner):
             env.update(self._get_aiperf_server_metrics_env())
+            if self.config.benchmark.aiperf_package:
+                env["AIPERF_PACKAGE"] = self.config.benchmark.aiperf_package
 
         return env

src/srtctl/core/schema.py

@@ -546,6 +546,9 @@ class BenchmarkConfig:
     trace_file: str | None = None  # Path to trace JSONL file (container path, e.g., /traces/dataset.jsonl)
     custom_tokenizer: str | None = None  # Custom tokenizer class (e.g., "module.path.ClassName")
     use_chat_template: bool = True  # Pass --use-chat-template to benchmark (default: true)
+    # aiperf pip install spec (e.g., "aiperf>=0.7.0", "aiperf @ git+https://...@commit")
+    # If set, runs pip install <spec> before benchmarking. Upgrades if already installed.
+    aiperf_package: str | None = None
     # Extra aiperf CLI flags passed through to bench.sh (e.g., benchmark-duration: 600, workers-max: 200)
     aiperf_args: dict[str, Any] = field(default_factory=dict)