mirror of https://github.com/razor-ai/soup.git
216 lines
7.6 KiB
Python
216 lines
7.6 KiB
Python
"""v0.43.0 Part C — Profiling extras.
|
||
|
||
- `memory_snapshot_context`: torch.cuda.memory._record_memory_history wrapper
|
||
that writes a `<run_id>.snapshot.pickle` to the profiles dir under cwd.
|
||
- `enable_detect_anomaly`: thin context manager around
|
||
`torch.autograd.set_detect_anomaly(True)`.
|
||
- `nccl_bandwidth_check`: scaffold that returns the expected upper-bound
|
||
bandwidth for a (gpu_pair, link) tuple. Live measurement deferred — we
|
||
expose the reference table so `soup doctor --nccl` can warn when measured
|
||
perf is well below expectation.
|
||
|
||
Containment, redaction, and exception-narrowing follow v0.34.0 `crash.py` /
|
||
v0.34.0 `profiling.py` policies.
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import contextlib
|
||
import math
|
||
import os
|
||
from dataclasses import dataclass
|
||
from types import MappingProxyType
|
||
from typing import Iterator, Mapping
|
||
|
||
from soup_cli.utils.paths import is_under_cwd
|
||
|
||
# Reference NVLink/PCIe bandwidth ceilings (GB/s, unidirectional).
|
||
# Source: NVIDIA published topology specs for H100/A100/V100.
|
||
_BANDWIDTH_REFERENCE: Mapping[tuple[str, str], float] = MappingProxyType({
|
||
("h100", "nvlink"): 450.0, # NVLink 4 (18 links × 25 GB/s)
|
||
("h100", "pcie"): 64.0, # PCIe Gen5 x16
|
||
("a100", "nvlink"): 300.0, # NVLink 3 (12 links × 25 GB/s)
|
||
("a100", "pcie"): 32.0, # PCIe Gen4 x16
|
||
("v100", "nvlink"): 150.0, # NVLink 2 (6 links × 25 GB/s)
|
||
("v100", "pcie"): 16.0, # PCIe Gen3 x16
|
||
("rtx4090", "pcie"): 64.0,
|
||
("rtx3090", "pcie"): 32.0,
|
||
})
|
||
|
||
|
||
def _validate_run_id(run_id: object) -> str:
|
||
if not isinstance(run_id, str):
|
||
raise ValueError("run_id must be a string")
|
||
if not run_id:
|
||
raise ValueError("run_id must not be empty")
|
||
if "\x00" in run_id:
|
||
raise ValueError("run_id must not contain null bytes")
|
||
if run_id in {".", ".."}:
|
||
raise ValueError(f"run_id contains forbidden token '{run_id}'")
|
||
for ch in ("/", "\\"):
|
||
if ch in run_id:
|
||
raise ValueError(f"run_id must not contain path separator '{ch}'")
|
||
return run_id
|
||
|
||
|
||
def resolve_snapshot_path(run_id: str, *, base_dir: str = "profiles") -> str:
|
||
"""Return realpath of `<cwd>/<base_dir>/<run_id>.snapshot.pickle`.
|
||
|
||
Rejects values that escape cwd. Mirrors v0.34.0 `resolve_trace_path`.
|
||
"""
|
||
rid = _validate_run_id(run_id)
|
||
if not isinstance(base_dir, str) or not base_dir:
|
||
raise ValueError("base_dir must be a non-empty string")
|
||
if "\x00" in base_dir:
|
||
raise ValueError("base_dir must not contain null bytes")
|
||
# Reject path separators and `..` components in base_dir so an absolute
|
||
# / parent-traversing path can never sneak through realpath on Windows
|
||
# short-name systems (security review fix).
|
||
if base_dir in {".", ".."}:
|
||
raise ValueError("base_dir must not be '.' or '..'")
|
||
if os.path.isabs(base_dir):
|
||
raise ValueError("base_dir must be a relative path under cwd")
|
||
parts = [p for p in base_dir.replace("\\", "/").split("/") if p]
|
||
if any(p == ".." for p in parts):
|
||
raise ValueError("base_dir must not contain '..' segments")
|
||
target = os.path.realpath(os.path.join(base_dir, f"{rid}.snapshot.pickle"))
|
||
if not is_under_cwd(target):
|
||
raise ValueError("snapshot path must stay under cwd")
|
||
return target
|
||
|
||
|
||
@contextlib.contextmanager
|
||
def memory_snapshot_context(
|
||
run_id: str,
|
||
*,
|
||
base_dir: str = "profiles",
|
||
max_entries: int = 100_000,
|
||
) -> Iterator[str | None]:
|
||
"""Record CUDA memory history; on exit, dump pickle + stop recording.
|
||
|
||
Yields the snapshot path on success, or None when torch is missing /
|
||
CUDA is unavailable / the recording API is missing. Never raises through
|
||
the context exit when torch failures are missing-dep style — those map
|
||
to a yielded None.
|
||
"""
|
||
if isinstance(max_entries, bool) or not isinstance(max_entries, int):
|
||
raise ValueError("max_entries must be a positive int")
|
||
if max_entries < 1 or max_entries > 10_000_000:
|
||
raise ValueError("max_entries must be in [1, 10_000_000]")
|
||
|
||
path = resolve_snapshot_path(run_id, base_dir=base_dir)
|
||
try:
|
||
import torch # type: ignore[import-not-found]
|
||
except ImportError:
|
||
yield None
|
||
return
|
||
cuda = getattr(torch, "cuda", None)
|
||
if cuda is None or not cuda.is_available():
|
||
yield None
|
||
return
|
||
record = getattr(getattr(cuda, "memory", None), "_record_memory_history", None)
|
||
dump = getattr(getattr(cuda, "memory", None), "_dump_snapshot", None)
|
||
if record is None or dump is None:
|
||
yield None
|
||
return
|
||
os.makedirs(os.path.dirname(path), exist_ok=True)
|
||
# Narrow the RuntimeError catch to only the entry call so a user-body
|
||
# RuntimeError cannot trigger a double-yield in this generator
|
||
# (review fix — the previous wide `except RuntimeError` would have
|
||
# raised "generator already executing" on user-body failures).
|
||
try:
|
||
record(max_entries=max_entries)
|
||
except RuntimeError:
|
||
yield None
|
||
return
|
||
try:
|
||
yield path
|
||
finally:
|
||
try:
|
||
dump(path)
|
||
finally:
|
||
try:
|
||
record(enabled=None) # type: ignore[arg-type]
|
||
except (TypeError, ValueError):
|
||
pass
|
||
|
||
|
||
@contextlib.contextmanager
|
||
def detect_anomaly_context() -> Iterator[bool]:
|
||
"""torch.autograd.set_detect_anomaly(True) wrapper.
|
||
|
||
Yields True when activated, False when torch is missing.
|
||
"""
|
||
try:
|
||
import torch # type: ignore[import-not-found]
|
||
except ImportError:
|
||
yield False
|
||
return
|
||
set_detect = getattr(getattr(torch, "autograd", None), "set_detect_anomaly", None)
|
||
if set_detect is None:
|
||
yield False
|
||
return
|
||
set_detect(True)
|
||
try:
|
||
yield True
|
||
finally:
|
||
set_detect(False)
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class BandwidthExpectation:
|
||
gpu: str
|
||
link: str
|
||
expected_gb_per_sec: float
|
||
|
||
|
||
def expected_bandwidth(gpu: str, link: str) -> float | None:
|
||
"""Return reference bandwidth (GB/s) for `(gpu, link)`, or None."""
|
||
if not isinstance(gpu, str) or not isinstance(link, str):
|
||
return None
|
||
return _BANDWIDTH_REFERENCE.get((gpu.lower(), link.lower()))
|
||
|
||
|
||
def nccl_bandwidth_check(
|
||
*, gpu: str, link: str, measured_gb_per_sec: float
|
||
) -> dict:
|
||
"""Compare measured bandwidth against the reference table.
|
||
|
||
Returns a dict with `expected`, `measured`, `ratio`, and `status`.
|
||
|
||
- status="OK" : ratio >= 0.80
|
||
- status="MINOR" : 0.50 <= ratio < 0.80
|
||
- status="MAJOR" : ratio < 0.50 (silent degradation likely)
|
||
- status="UNKNOWN": no reference entry for (gpu, link)
|
||
"""
|
||
if (
|
||
isinstance(measured_gb_per_sec, bool)
|
||
or not isinstance(measured_gb_per_sec, (int, float))
|
||
):
|
||
raise ValueError("measured_gb_per_sec must be a number")
|
||
if not math.isfinite(float(measured_gb_per_sec)):
|
||
raise ValueError("measured_gb_per_sec must be finite")
|
||
if measured_gb_per_sec < 0:
|
||
raise ValueError("measured_gb_per_sec must be >= 0")
|
||
|
||
expected = expected_bandwidth(gpu, link)
|
||
if expected is None or expected <= 0:
|
||
return {
|
||
"expected_gb_per_sec": None,
|
||
"measured_gb_per_sec": float(measured_gb_per_sec),
|
||
"ratio": None,
|
||
"status": "UNKNOWN",
|
||
}
|
||
ratio = float(measured_gb_per_sec) / expected
|
||
if ratio >= 0.80:
|
||
status = "OK"
|
||
elif ratio >= 0.50:
|
||
status = "MINOR"
|
||
else:
|
||
status = "MAJOR"
|
||
return {
|
||
"expected_gb_per_sec": expected,
|
||
"measured_gb_per_sec": float(measured_gb_per_sec),
|
||
"ratio": round(ratio, 4),
|
||
"status": status,
|
||
}
|