8292e699e4Trim and tighten code comments and docstrings across studio/ Python. Comment-only: every changed file verified code-identical to main via AST/token comparison.
413 lines
16 KiB
Python
413 lines
16 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
|
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
|
|
|
"""AMD GPU monitoring via amd-smi.
|
|
|
|
Mirrors nvidia.py so hardware.py can swap backends based on IS_ROCM.
|
|
All functions return the same dict shapes as their nvidia.py counterparts.
|
|
"""
|
|
|
|
import json
|
|
import math
|
|
import os
|
|
import platform
|
|
import re
|
|
import subprocess
|
|
import sys
|
|
from typing import Any, Optional
|
|
|
|
from loggers import get_logger
|
|
from utils.native_path_leases import child_env_without_native_path_secret
|
|
from utils.subprocess_compat import windows_hidden_subprocess_kwargs
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
# amd-smi on Windows initialises the full ROCm runtime on first call, which
|
|
# can take 15-25 s on cold hardware. Linux is consistently < 2 s.
|
|
_AMD_SMI_DEFAULT_TIMEOUT = 30 if platform.system() == "Windows" else 10
|
|
|
|
# Circuit breaker: stop calling amd-smi after this many consecutive failures.
|
|
# On Windows, each failed call spawns a process that may show a UAC/DiskPart
|
|
# elevation prompt. Once we know amd-smi doesn't work, stop polling it.
|
|
_AMD_SMI_FAILURE_LIMIT = 3
|
|
_amd_smi_consecutive_failures = 0
|
|
_amd_smi_disabled = False
|
|
|
|
|
|
def _run_amd_smi(*args: str, timeout: int = _AMD_SMI_DEFAULT_TIMEOUT) -> Optional[Any]:
|
|
"""Run amd-smi with the given args and return parsed JSON, or None."""
|
|
global _amd_smi_consecutive_failures, _amd_smi_disabled
|
|
if _amd_smi_disabled:
|
|
return None
|
|
try:
|
|
result = subprocess.run(
|
|
["amd-smi", *args, "--json"],
|
|
capture_output = True,
|
|
text = True,
|
|
timeout = timeout,
|
|
env = child_env_without_native_path_secret(),
|
|
**windows_hidden_subprocess_kwargs(),
|
|
)
|
|
except (OSError, subprocess.TimeoutExpired) as e:
|
|
if isinstance(e, FileNotFoundError):
|
|
# amd-smi ships with Adrenalin, not the HIP SDK -- absence is
|
|
# expected on HIP SDK-only Windows setups. Log at debug only.
|
|
logger.debug("amd-smi not found (not in PATH): %s", e)
|
|
else:
|
|
logger.warning("amd-smi query failed: %s", e)
|
|
_amd_smi_consecutive_failures += 1
|
|
if _amd_smi_consecutive_failures >= _AMD_SMI_FAILURE_LIMIT:
|
|
logger.info(
|
|
"amd-smi not available (not installed; expected on HIP SDK-only systems); "
|
|
"GPU VRAM polling disabled"
|
|
)
|
|
_amd_smi_disabled = True
|
|
return None
|
|
if result.returncode != 0:
|
|
logger.warning("amd-smi returned code %d", result.returncode)
|
|
_amd_smi_consecutive_failures += 1
|
|
if _amd_smi_consecutive_failures >= _AMD_SMI_FAILURE_LIMIT:
|
|
logger.info(
|
|
"amd-smi not available (not installed; expected on HIP SDK-only systems); "
|
|
"GPU VRAM polling disabled"
|
|
)
|
|
_amd_smi_disabled = True
|
|
return None
|
|
if not result.stdout.strip():
|
|
# amd-smi exited 0 but produced no output (e.g. no GPUs visible on
|
|
# this query, or a version that emits nothing for --json). Not a tool
|
|
# failure, so don't count against the circuit breaker.
|
|
logger.debug("amd-smi exited 0 but returned no output")
|
|
return None
|
|
_amd_smi_consecutive_failures = 0 # reset on success
|
|
try:
|
|
return json.loads(result.stdout)
|
|
except json.JSONDecodeError:
|
|
logger.warning("Failed to parse amd-smi JSON output")
|
|
return None
|
|
|
|
|
|
def _parse_numeric(value: Any) -> Optional[float]:
|
|
"""Extract a numeric value from amd-smi output (str, int, float, or dict)."""
|
|
if value is None:
|
|
return None
|
|
# Newer amd-smi versions emit {"value": 10, "unit": "W"}
|
|
if isinstance(value, dict):
|
|
return _parse_numeric(value.get("value"))
|
|
if isinstance(value, (int, float)):
|
|
f = float(value)
|
|
return f if math.isfinite(f) else None
|
|
if isinstance(value, str):
|
|
# Strip units like "W", "C", "%", "MB", "MiB", "GB", "GiB" etc.
|
|
cleaned = re.sub(r"\s*[A-Za-z/%]+$", "", value.strip())
|
|
if not cleaned or cleaned.lower() in ("n/a", "none", "unknown"):
|
|
return None
|
|
try:
|
|
return float(cleaned)
|
|
except (ValueError, TypeError):
|
|
return None
|
|
return None
|
|
|
|
|
|
def _parse_memory_mb(value: Any) -> Optional[float]:
|
|
"""Parse a memory value from amd-smi output and return MB.
|
|
|
|
Handles bare numbers (assumed MB -- the amd-smi convention on every
|
|
version seen), dict values with explicit units (``{"value": 192,
|
|
"unit": "GiB"}`` on newer releases), and strings like ``"8192 MiB"``.
|
|
"""
|
|
unit = ""
|
|
raw_value = value
|
|
|
|
if isinstance(value, dict):
|
|
unit = str(value.get("unit", "")).strip().lower()
|
|
raw_value = value.get("value")
|
|
elif isinstance(value, str):
|
|
# Extract unit suffix from strings like "192 GiB" or "8192 MB"
|
|
m = re.match(r"^\s*([\d.]+)\s*([A-Za-z]+)\s*$", value.strip())
|
|
if m:
|
|
unit = m.group(2).lower()
|
|
|
|
num = _parse_numeric(raw_value if isinstance(value, dict) else value)
|
|
if num is None:
|
|
return None
|
|
|
|
# Unit conversion -- GPU tools (incl. amd-smi) use binary units even when
|
|
# labeled "GB" or "MB", so treat GB/GiB and MB/MiB the same.
|
|
if "gib" in unit or "gb" in unit:
|
|
return num * 1024
|
|
if "mib" in unit or "mb" in unit:
|
|
return num
|
|
if "kib" in unit or "kb" in unit:
|
|
return num / 1024
|
|
if unit in ("b", "byte", "bytes"):
|
|
# Plain bytes
|
|
return num / (1024 * 1024)
|
|
|
|
# No explicit unit -- default to MB, the amd-smi convention for bare
|
|
# numeric values. A previous heuristic assumed values above ~10M were
|
|
# bytes, but that misclassifies small VRAM allocations (e.g. 5 MB =
|
|
# 5,242,880 reported without a unit) as ~5 TB. Modern amd-smi always
|
|
# ships explicit units, so the heuristic only fired for legacy output
|
|
# where MB was already the convention.
|
|
return num
|
|
|
|
|
|
def _extract_gpu_metrics(gpu_data: dict) -> dict[str, Any]:
|
|
"""Extract standardized metrics from a single GPU's amd-smi data."""
|
|
# Output structure varies by version; try common paths
|
|
usage = gpu_data.get("usage", gpu_data.get("gpu_activity", {}))
|
|
if isinstance(usage, dict):
|
|
gpu_util = _parse_numeric(usage.get("gfx_activity", usage.get("gpu_use_percent")))
|
|
else:
|
|
gpu_util = _parse_numeric(usage)
|
|
|
|
# Temperature -- try keys in priority order. dict.get() returns "N/A"
|
|
# strings rather than falling through, so try each key and check it
|
|
# parses to a real number.
|
|
temp_data = gpu_data.get("temperature", {})
|
|
temp = None
|
|
if isinstance(temp_data, dict):
|
|
for temp_key in ("edge", "temperature_edge", "hotspot", "temperature_hotspot"):
|
|
temp = _parse_numeric(temp_data.get(temp_key))
|
|
if temp is not None:
|
|
break
|
|
else:
|
|
temp = _parse_numeric(temp_data)
|
|
|
|
# Power
|
|
power_data = gpu_data.get("power", {})
|
|
if isinstance(power_data, dict):
|
|
power_draw = _parse_numeric(
|
|
power_data.get(
|
|
"current_socket_power",
|
|
power_data.get("average_socket_power", power_data.get("socket_power")),
|
|
)
|
|
)
|
|
power_limit = _parse_numeric(power_data.get("power_cap", power_data.get("max_power_limit")))
|
|
else:
|
|
power_draw = None
|
|
power_limit = None
|
|
|
|
# VRAM -- unit-aware parsing for varying amd-smi output formats.
|
|
# Newer versions may return {"value": 192, "unit": "GiB"} and use
|
|
# "mem_usage" with "total_vram" / "used_vram" keys; older versions use
|
|
# "vram" or "fb_memory_usage" with "used" / "total".
|
|
vram_data = gpu_data.get(
|
|
"mem_usage",
|
|
gpu_data.get("vram", gpu_data.get("fb_memory_usage", {})),
|
|
)
|
|
if isinstance(vram_data, dict):
|
|
vram_used_mb = _parse_memory_mb(
|
|
vram_data.get("used_vram", vram_data.get("vram_used", vram_data.get("used")))
|
|
)
|
|
vram_total_mb = _parse_memory_mb(
|
|
vram_data.get("total_vram", vram_data.get("vram_total", vram_data.get("total")))
|
|
)
|
|
else:
|
|
vram_used_mb = None
|
|
vram_total_mb = None
|
|
|
|
# Build the standardized dict (same shape as nvidia._build_gpu_metrics)
|
|
vram_used_gb = round(vram_used_mb / 1024, 2) if vram_used_mb is not None else None
|
|
vram_total_gb = round(vram_total_mb / 1024, 2) if vram_total_mb is not None else None
|
|
vram_util = (
|
|
round((vram_used_mb / vram_total_mb) * 100, 1)
|
|
if vram_used_mb is not None and vram_total_mb is not None and vram_total_mb > 0
|
|
else None
|
|
)
|
|
power_util = (
|
|
round((power_draw / power_limit) * 100, 1)
|
|
if power_draw is not None and power_limit is not None and power_limit > 0
|
|
else None
|
|
)
|
|
|
|
return {
|
|
"gpu_utilization_pct": gpu_util,
|
|
"temperature_c": temp,
|
|
"vram_used_gb": vram_used_gb,
|
|
"vram_total_gb": vram_total_gb,
|
|
"vram_utilization_pct": vram_util,
|
|
"power_draw_w": power_draw,
|
|
"power_limit_w": power_limit,
|
|
"power_utilization_pct": power_util,
|
|
}
|
|
|
|
|
|
def _has_real_metrics(metrics: dict[str, Any]) -> bool:
|
|
"""Return True when ``metrics`` contains at least one non-None value.
|
|
|
|
``amd-smi`` can return a zero-exit JSON envelope missing every expected
|
|
field (error response, unsupported card, hipless container). Then
|
|
``_extract_gpu_metrics`` produces an all-``None`` dict -- callers must
|
|
surface this as ``available: False``, not ``available: True`` with empty data.
|
|
"""
|
|
return any(value is not None for value in metrics.values())
|
|
|
|
|
|
def get_physical_gpu_count() -> Optional[int]:
|
|
"""Return physical AMD GPU count via amd-smi, or None on failure."""
|
|
data = _run_amd_smi("list")
|
|
if data is None:
|
|
return None
|
|
if isinstance(data, list):
|
|
return len(data)
|
|
# Some versions return a dict with a "gpu" / "gpus" key. Guard .get()
|
|
# with isinstance so a malformed scalar/string response cannot raise
|
|
# AttributeError.
|
|
if not isinstance(data, dict):
|
|
return None
|
|
gpus = data.get("gpu", data.get("gpus", []))
|
|
if isinstance(gpus, list):
|
|
return len(gpus)
|
|
return None
|
|
|
|
|
|
def _first_visible_amd_gpu_id() -> Optional[str]:
|
|
"""Return the physical AMD GPU id treated as 'primary'.
|
|
|
|
Honours HIP_VISIBLE_DEVICES / ROCR_VISIBLE_DEVICES / CUDA_VISIBLE_DEVICES
|
|
in that order (HIP respects all three). Returns ``"0"`` when none are set,
|
|
and ``None`` when the env var narrows to zero GPUs ("" or "-1"), so callers
|
|
can short-circuit to "available: False".
|
|
"""
|
|
for env_name in (
|
|
"HIP_VISIBLE_DEVICES",
|
|
"ROCR_VISIBLE_DEVICES",
|
|
"CUDA_VISIBLE_DEVICES",
|
|
):
|
|
raw = os.environ.get(env_name)
|
|
if raw is None:
|
|
continue
|
|
raw = raw.strip()
|
|
if raw == "" or raw == "-1":
|
|
return None
|
|
# Drop empty tokens after splitting. Tolerates minor typos like
|
|
# ``HIP_VISIBLE_DEVICES=",1"`` (leading comma, clearly meant device 1)
|
|
# while still falling through to the next env var when every token is
|
|
# empty (e.g. ``,,,``).
|
|
tokens = [t.strip() for t in raw.split(",") if t.strip()]
|
|
if tokens:
|
|
return tokens[0]
|
|
return "0"
|
|
|
|
|
|
def get_primary_gpu_utilization() -> dict[str, Any]:
|
|
"""Return utilization metrics for the primary visible AMD GPU."""
|
|
gpu_idx = _first_visible_amd_gpu_id()
|
|
if gpu_idx is None:
|
|
return {"available": False}
|
|
data = _run_amd_smi("metric", "-g", gpu_idx)
|
|
if data is None:
|
|
return {"available": False}
|
|
|
|
# amd-smi may return:
|
|
# - a list of GPU dicts (older versions)
|
|
# - a dict with a "gpu_data" key wrapping a list (newer versions)
|
|
# - a single GPU dict (rare)
|
|
if isinstance(data, dict) and "gpu_data" in data:
|
|
data = data["gpu_data"]
|
|
if isinstance(data, list):
|
|
if len(data) == 0:
|
|
return {"available": False}
|
|
gpu_data = data[0]
|
|
else:
|
|
gpu_data = data
|
|
|
|
metrics = _extract_gpu_metrics(gpu_data)
|
|
if not _has_real_metrics(metrics):
|
|
# JSON envelope with no usable fields (error response or unsupported
|
|
# card). Surface as unavailable rather than available-with-empty-data
|
|
# so the UI does not render a ghost device.
|
|
return {"available": False}
|
|
metrics["available"] = True
|
|
return metrics
|
|
|
|
|
|
def get_visible_gpu_utilization(
|
|
parent_visible_ids: Optional[list[int]], parent_cuda_visible_devices: Optional[str] = None
|
|
) -> dict[str, Any]:
|
|
"""Return utilization metrics for visible AMD GPUs."""
|
|
if parent_visible_ids is None:
|
|
return {
|
|
"available": False,
|
|
"backend_cuda_visible_devices": parent_cuda_visible_devices,
|
|
"parent_visible_gpu_ids": [],
|
|
"devices": [],
|
|
"index_kind": "unresolved",
|
|
}
|
|
|
|
data = _run_amd_smi("metric")
|
|
if data is None:
|
|
return {
|
|
"available": False,
|
|
"backend_cuda_visible_devices": parent_cuda_visible_devices,
|
|
"parent_visible_gpu_ids": parent_visible_ids or [],
|
|
"devices": [],
|
|
"index_kind": "physical",
|
|
}
|
|
|
|
# Extract a device list from amd-smi's envelope. Newer versions return a
|
|
# JSON array directly; older versions wrap the list in a dict under
|
|
# "gpus" / "gpu". Guard non-dict / non-list envelopes (scalar/string
|
|
# fallbacks from malformed output) so .get() cannot raise AttributeError.
|
|
if isinstance(data, list):
|
|
gpu_list = data
|
|
elif isinstance(data, dict):
|
|
# Newer amd-smi wraps output in {"gpu_data": [...]}
|
|
gpu_list = data.get("gpu_data", data.get("gpus", data.get("gpu", [data])))
|
|
else:
|
|
gpu_list = [data]
|
|
visible_set = set(parent_visible_ids)
|
|
ordinal_map = {gpu_id: ordinal for ordinal, gpu_id in enumerate(parent_visible_ids)}
|
|
|
|
devices = []
|
|
for fallback_idx, gpu_data in enumerate(gpu_list):
|
|
# Skip non-dict entries: a scalar inside the "gpus" array (seen on
|
|
# some malformed output) would make _extract_gpu_metrics raise
|
|
# AttributeError on the first .get() call.
|
|
if not isinstance(gpu_data, dict):
|
|
continue
|
|
# Use AMD-reported GPU ID when available, else the enumeration index.
|
|
# Newer amd-smi wraps scalars as ``{"value": 0, "unit": "none"}``, so
|
|
# route raw_id through ``_parse_numeric`` which handles bare ints,
|
|
# floats, strings, and that dict shape uniformly.
|
|
raw_id = gpu_data.get("gpu", gpu_data.get("gpu_id", gpu_data.get("id", fallback_idx)))
|
|
parsed_id = _parse_numeric(raw_id)
|
|
if parsed_id is None:
|
|
logger.warning(
|
|
"amd-smi GPU id %r could not be parsed; falling back to enumeration index %d",
|
|
raw_id,
|
|
fallback_idx,
|
|
)
|
|
idx = fallback_idx
|
|
else:
|
|
rounded = round(parsed_id)
|
|
if rounded != parsed_id:
|
|
logger.warning(
|
|
"amd-smi GPU id %r parsed as non-integer %r; truncating to %d",
|
|
raw_id,
|
|
parsed_id,
|
|
rounded,
|
|
)
|
|
idx = int(rounded)
|
|
if idx not in visible_set:
|
|
continue
|
|
metrics = _extract_gpu_metrics(gpu_data)
|
|
if not _has_real_metrics(metrics):
|
|
# Skip ghost entries: a dict response with no usable fields
|
|
# (error envelope, etc.) would otherwise show up as a device row
|
|
# with all-None numbers in the UI.
|
|
continue
|
|
metrics["index"] = idx
|
|
metrics["index_kind"] = "physical"
|
|
metrics["visible_ordinal"] = ordinal_map.get(idx, len(devices))
|
|
devices.append(metrics)
|
|
|
|
return {
|
|
"available": len(devices) > 0,
|
|
"backend_cuda_visible_devices": parent_cuda_visible_devices,
|
|
"parent_visible_gpu_ids": parent_visible_ids or [],
|
|
"devices": devices,
|
|
"index_kind": "physical",
|
|
}
|