[smg][ci]: migrate benchmarks to e2e_test/benchmarks/, use parent conftest (#16597)
This commit is contained in:
@@ -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")
|
||||
@@ -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)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user