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