[smg][ci]: migrate benchmarks to e2e_test/benchmarks/, use parent conftest (#16597)

This commit is contained in:
Simo Lin
2026-01-06 20:15:20 -08:00
committed by GitHub
parent 913b688f21
commit d8b8198192
14 changed files with 1043 additions and 1963 deletions
@@ -35,6 +35,8 @@ from .gpu_allocator import (
nvml_context,
wait_for_gpu_memory_to_clear,
)
from .gpu_monitor import GPUMonitor
from .gpu_monitor import should_monitor as should_monitor_gpu
from .model_pool import ModelInstance, ModelPool
from .model_specs import ( # Default model paths; Model groups
CHAT_MODELS,
@@ -104,6 +106,9 @@ __all__ = [
"wait_for_health",
"wait_for_workers_ready",
"detect_ib_device",
# GPU monitoring
"GPUMonitor",
"should_monitor_gpu",
# Model management
"ModelInstance",
"ModelPool",
@@ -0,0 +1,329 @@
"""GPU utilization monitoring for benchmarks.
This module provides a low-impact GPU monitor that runs in a separate process
and collects utilization samples using NVML.
"""
from __future__ import annotations
import json
import logging
import os
import time
from multiprocessing import Process
from pathlib import Path
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from typing import Any
logger = logging.getLogger(__name__)
def _percentile(samples: list[float], p: float) -> float:
"""Calculate percentile from sorted samples."""
if not samples:
return 0.0
sorted_samples = sorted(samples)
idx = max(
0,
min(
len(sorted_samples) - 1, int(round((p / 100.0) * (len(sorted_samples) - 1)))
),
)
return float(sorted_samples[idx])
def _compute_stats(samples: list[float]) -> dict[str, float]:
"""Compute statistics for a list of samples."""
if not samples:
return {
"mean": 0.0,
"min": 0.0,
"max": 0.0,
"p5": 0.0,
"p10": 0.0,
"p25": 0.0,
"p50": 0.0,
"p75": 0.0,
"p90": 0.0,
"p95": 0.0,
"count": 0,
}
return {
"mean": sum(samples) / len(samples),
"min": min(samples),
"max": max(samples),
"p5": _percentile(samples, 5),
"p10": _percentile(samples, 10),
"p25": _percentile(samples, 25),
"p50": _percentile(samples, 50),
"p75": _percentile(samples, 75),
"p90": _percentile(samples, 90),
"p95": _percentile(samples, 95),
"count": len(samples),
}
def _monitor_loop(pid: int, output_path: str, interval: float) -> None:
"""Main monitoring loop - runs in separate process.
Monitors GPU utilization until the target process exits, then writes
results to output_path as JSON.
"""
# Lower process priority to minimize impact on benchmark
try:
os.nice(10)
except Exception:
pass
# Initialize NVML
try:
import pynvml
pynvml.nvmlInit()
except Exception as e:
logger.warning("Failed to initialize NVML: %s", e)
_write_empty_result(output_path)
return
# Get GPU handles
try:
device_count = pynvml.nvmlDeviceGetCount()
handles = [pynvml.nvmlDeviceGetHandleByIndex(i) for i in range(device_count)]
except Exception as e:
logger.warning("Failed to get GPU handles: %s", e)
_write_empty_result(output_path)
_shutdown_nvml()
return
# Collect samples
per_gpu_samples: dict[str, list[float]] = {str(i): [] for i in range(device_count)}
overall_samples: list[float] = []
try:
while _process_alive(pid):
try:
gpu_utils = []
for idx, handle in enumerate(handles):
try:
util = pynvml.nvmlDeviceGetUtilizationRates(handle).gpu
gpu_utils.append(float(util))
per_gpu_samples[str(idx)].append(float(util))
except Exception:
continue
if gpu_utils:
avg = sum(gpu_utils) / len(gpu_utils)
overall_samples.append(avg)
except Exception:
pass
time.sleep(interval)
finally:
# Write results
_write_result(output_path, pid, interval, overall_samples, per_gpu_samples)
_shutdown_nvml()
def _process_alive(pid: int) -> bool:
"""Check if process is still running."""
try:
os.kill(pid, 0)
return True
except (OSError, ProcessLookupError):
return False
def _write_empty_result(path: str) -> None:
"""Write empty result file."""
try:
os.makedirs(os.path.dirname(path), exist_ok=True)
with open(path, "w") as f:
json.dump(
{
"count": 0,
"overall": {"mean": 0.0},
"per_gpu": {},
"raw": {"overall": [], "per_gpu": {}},
},
f,
)
except Exception:
pass
def _write_result(
path: str,
pid: int,
interval: float,
overall_samples: list[float],
per_gpu_samples: dict[str, list[float]],
) -> None:
"""Write monitoring results to JSON file."""
try:
os.makedirs(os.path.dirname(path), exist_ok=True)
with open(path, "w") as f:
json.dump(
{
"bench_pid": pid,
"interval_sec": interval,
"count": len(overall_samples),
"overall": _compute_stats(overall_samples),
"per_gpu": {
k: _compute_stats(v) for k, v in per_gpu_samples.items()
},
"raw": {
"overall": overall_samples,
"per_gpu": per_gpu_samples,
},
},
f,
)
except Exception as e:
logger.warning("Failed to write GPU monitor results: %s", e)
def _shutdown_nvml() -> None:
"""Shutdown NVML."""
try:
import pynvml
pynvml.nvmlShutdown()
except Exception:
pass
class GPUMonitor:
"""GPU utilization monitor for benchmarks.
Usage:
monitor = GPUMonitor(output_dir="benchmark_results")
monitor.start(target_pid=12345)
# ... run benchmark ...
result = monitor.stop()
monitor.assert_thresholds({"gpu_util_p50_min": 99})
"""
def __init__(
self,
output_dir: str | Path = ".",
interval: float = 2.0,
):
self.output_dir = Path(output_dir)
self.interval = interval
self._process: Process | None = None
self._output_path: str | None = None
self._result: dict[str, Any] | None = None
@property
def output_path(self) -> str | None:
"""Path to the GPU utilization JSON file."""
return self._output_path
def start(self, target_pid: int) -> None:
"""Start monitoring GPU utilization for the target process."""
self._output_path = str(self.output_dir / "gpu_utilization.json")
self._result = None
self._process = Process(
target=_monitor_loop,
args=(target_pid, self._output_path, self.interval),
daemon=True,
)
self._process.start()
logger.debug("Started GPU monitor for PID %d", target_pid)
def stop(self, timeout: float = 5.0) -> dict[str, Any] | None:
"""Stop monitoring and return results."""
if self._process is None:
return None
try:
self._process.join(timeout=timeout)
except Exception:
pass
if self._process.is_alive():
try:
self._process.terminate()
except Exception:
pass
self._process = None
self._result = self._read_result()
return self._result
def _read_result(self) -> dict[str, Any] | None:
"""Read results from output file."""
if not self._output_path or not os.path.exists(self._output_path):
return None
try:
with open(self._output_path) as f:
return json.load(f)
except Exception as e:
logger.warning("Failed to read GPU monitor result: %s", e)
return None
def log_summary(self) -> None:
"""Log a summary of GPU utilization."""
result = self._result or self._read_result()
if not result or result.get("count", 0) <= 0:
logger.warning("GPU utilization monitor produced no samples")
return
overall = result.get("overall", {})
logger.info(
"GPU utilization: mean=%.2f%% p50=%.2f%% (samples=%d)",
overall.get("mean", 0.0),
overall.get("p50", 0.0),
result.get("count", 0),
)
def assert_thresholds(self, thresholds: dict[str, float] | None) -> None:
"""Assert GPU utilization meets thresholds.
Supported thresholds:
- gpu_util_mean_min: Minimum mean GPU utilization %
- gpu_util_p50_min: Minimum p50 GPU utilization %
"""
if not thresholds:
return
result = self._result or self._read_result()
if not result or result.get("count", 0) <= 0:
logger.warning("GPU utilization monitor produced no samples")
return
overall = result.get("overall", {})
mean_threshold = thresholds.get("gpu_util_mean_min")
if mean_threshold is not None:
mean_value = overall.get("mean", 0.0)
assert (
mean_value >= mean_threshold
), f"GPU utilization mean below threshold: {mean_value:.2f}% < {mean_threshold}%"
p50_threshold = thresholds.get("gpu_util_p50_min")
if p50_threshold is not None:
p50_value = overall.get("p50")
if p50_value is not None:
assert (
p50_value >= p50_threshold
), f"GPU utilization p50 below threshold: {p50_value:.2f}% < {p50_threshold}%"
def should_monitor(thresholds: dict[str, Any] | None) -> bool:
"""Check if GPU monitoring should be enabled.
Returns True if:
- thresholds contains gpu_util_mean_min or gpu_util_p50_min, OR
- GPU_UTIL_LOG environment variable is truthy
"""
if thresholds:
if thresholds.get("gpu_util_mean_min") is not None:
return True
if thresholds.get("gpu_util_p50_min") is not None:
return True
return os.environ.get("GPU_UTIL_LOG", "").lower() in ("1", "true", "yes")
+153 -17
View File
@@ -293,6 +293,7 @@ class ModelPool:
worker_type: WorkerType = WorkerType.REGULAR,
bootstrap_port: int | None = None,
ib_device: str | None = None,
instance_key: str | None = None,
) -> ModelInstance:
"""Launch a model instance.
@@ -303,6 +304,7 @@ class ModelPool:
worker_type: Worker type (REGULAR, PREFILL, or DECODE).
bootstrap_port: Bootstrap port for prefill workers in PD mode.
ib_device: InfiniBand device for PD disaggregation.
instance_key: Custom instance key, or None to auto-generate.
Returns:
The launched ModelInstance.
@@ -353,11 +355,15 @@ class ModelPool:
cmd.extend(["--disaggregation-ib-device", ib_device])
elif worker_type == WorkerType.DECODE:
cmd.extend(["--disaggregation-mode", "decode"])
# Base GPU ID 0 since CUDA_VISIBLE_DEVICES remaps the GPU
cmd.extend(["--base-gpu-id", "0"])
if ib_device:
cmd.extend(["--disaggregation-ib-device", ib_device])
# Build key based on worker type
if worker_type == WorkerType.REGULAR:
# Build key based on worker type (or use custom key)
if instance_key:
key = instance_key
elif worker_type == WorkerType.REGULAR:
key = f"{model_id}:{mode.value}"
else:
key = f"{model_id}:{mode.value}:{worker_type.value}"
@@ -560,7 +566,11 @@ class ModelPool:
return instance
def _evict_for_gpus(
self, required_gpus: int, exclude_model_id: str | None = None
self,
required_gpus: int,
exclude_model_id: str | None = None,
exclude_mode: ConnectionMode | None = None,
exclude_worker_types: set[WorkerType] | None = None,
) -> None:
"""Evict models until we have enough GPUs available.
@@ -570,29 +580,45 @@ class ModelPool:
Args:
required_gpus: Number of GPUs needed.
exclude_model_id: Model ID to exclude from eviction (test may need
multiple modes of the same model).
exclude_model_id: Model ID to exclude from eviction.
exclude_mode: Connection mode to exclude from eviction (optional).
exclude_worker_types: Worker types to exclude from eviction.
If None, falls back to excluding by model_id only (backward compatible).
"""
available = self.allocator.available_gpus()
if len(available) >= required_gpus:
return # Already have enough
# Sort by last_used descending (MRU eviction) - evict most recently used first
# Exclude instances of the same model_id (test may need multiple modes)
evictable = [
inst
for inst in self.instances.values()
if exclude_model_id is None or inst.model_id != exclude_model_id
]
evictable.sort(key=lambda x: x.last_used, reverse=True)
# Store (dict_key, instance) tuples to preserve the actual key for eviction
evictable: list[tuple[str, ModelInstance]] = []
for dict_key, inst in self.instances.items():
if exclude_worker_types is not None:
# Precise matching with worker types
# Must match model_id AND worker_type, mode is optional
if (
exclude_model_id is not None
and inst.model_id == exclude_model_id
and inst.worker_type in exclude_worker_types
):
# If mode is specified, also require mode match
if exclude_mode is None or inst.mode == exclude_mode:
continue
else:
# Backward compatible: exclude by model_id only
if exclude_model_id is not None and inst.model_id == exclude_model_id:
continue
evictable.append((dict_key, inst))
evictable.sort(key=lambda x: x[1].last_used, reverse=True)
freed_gpus = len(available)
for inst in evictable:
for dict_key, inst in evictable:
if freed_gpus >= required_gpus:
break
logger.info("Evicting model %s (MRU) to free GPUs", inst.key)
self._evict_instance(inst.key)
logger.info("Evicting model %s (MRU) to free GPUs", dict_key)
self._evict_instance(dict_key)
if inst.gpu_slot:
freed_gpus += len(inst.gpu_slot.gpu_ids)
@@ -608,7 +634,13 @@ class ModelPool:
spec = get_model_spec(model_id)
required_gpus = spec.get("tp", 1)
self._evict_for_gpus(required_gpus, exclude_model_id=model_id)
# Exclude REGULAR workers of same model from eviction (keep them)
# but allow evicting PD workers (PREFILL/DECODE) to free GPUs
self._evict_for_gpus(
required_gpus,
exclude_model_id=model_id,
exclude_worker_types={WorkerType.REGULAR},
)
available = self.allocator.available_gpus()
if len(available) < required_gpus:
@@ -682,6 +714,102 @@ class ModelPool:
if inst.model_id == model_id and inst.worker_type == worker_type
]
def launch_regular_workers(
self,
model_id: str,
num_workers: int,
mode: ConnectionMode = ConnectionMode.HTTP,
startup_timeout: int = DEFAULT_STARTUP_TIMEOUT,
allow_eviction: bool = True,
) -> list[ModelInstance]:
"""Launch multiple regular workers for load balancing.
Args:
model_id: Model identifier from MODEL_SPECS.
num_workers: Number of workers to launch.
mode: Connection mode (HTTP or GRPC).
startup_timeout: Timeout for workers to become healthy.
allow_eviction: If True, evict MRU models to free GPUs.
Returns:
List of ModelInstance objects.
"""
self._startup_timeout = startup_timeout
if model_id not in MODEL_SPECS:
raise ValueError(f"Unknown model: {model_id}")
spec = get_model_spec(model_id)
tp = spec.get("tp", 1)
required_gpus = num_workers * tp
# Check if we have enough GPUs
available = self.allocator.available_gpus()
if len(available) < required_gpus:
if allow_eviction:
logger.info(
"Need %d GPUs for %d workers, only %d available. Evicting MRU models...",
required_gpus,
num_workers,
len(available),
)
# Exclude REGULAR workers of same model/mode from eviction
self._evict_for_gpus(
required_gpus,
exclude_model_id=model_id,
exclude_mode=mode,
exclude_worker_types={WorkerType.REGULAR},
)
else:
logger.info(
"Need %d GPUs for %d workers, only %d available. "
"Skipping (eviction not allowed).",
required_gpus,
num_workers,
len(available),
)
return []
# Build allocation specs for all workers
allocation_specs = {}
for i in range(num_workers):
key = f"{model_id}:{mode.value}:{i}"
allocation_specs[key] = {
"model": spec["model"],
"memory_gb": spec.get("memory_gb", 16),
"tp": tp,
}
# Allocate GPU slots
slots = self.allocator.allocate_slots(allocation_specs)
slot_map = {slot.assigned_model: slot for slot in slots}
if not slots:
raise RuntimeError(
f"Failed to allocate GPU slots for {num_workers} workers after eviction. "
f"Need {required_gpus} GPUs."
)
instances: list[ModelInstance] = []
# Launch workers
for i in range(num_workers):
key = f"{model_id}:{mode.value}:{i}"
gpu_slot = slot_map.get(key)
instance = self._launch_model(
model_id=model_id,
mode=mode,
gpu_slot=gpu_slot,
worker_type=WorkerType.REGULAR,
instance_key=key,
)
instances.append(instance)
# Wait for all to be healthy
self._wait_all_healthy()
return instances
def launch_pd_workers(
self,
model_id: str,
@@ -728,7 +856,13 @@ class ModelPool:
required_gpus,
len(available),
)
self._evict_for_gpus(required_gpus, exclude_model_id=model_id)
# Exclude PD workers of same model/mode, but evict REGULAR workers
self._evict_for_gpus(
required_gpus,
exclude_model_id=model_id,
exclude_mode=mode,
exclude_worker_types={WorkerType.PREFILL, WorkerType.DECODE},
)
else:
logger.info(
"Need %d GPUs for PD workers, only %d available. "
@@ -781,6 +915,7 @@ class ModelPool:
worker_type=WorkerType.PREFILL,
bootstrap_port=bootstrap_port,
ib_device=ib_device,
instance_key=key,
)
prefill_instances.append(instance)
@@ -794,6 +929,7 @@ class ModelPool:
gpu_slot=gpu_slot,
worker_type=WorkerType.DECODE,
ib_device=ib_device,
instance_key=key,
)
decode_instances.append(instance)