387 lines
14 KiB
Python
387 lines
14 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 the nvidia.py module structure 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 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__)
|
|
|
|
|
|
def _run_amd_smi(*args: str, timeout: int = 5) -> Optional[Any]:
|
|
"""Run amd-smi with the given arguments and return parsed JSON, or 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:
|
|
logger.warning("amd-smi query failed: %s", e)
|
|
return None
|
|
if result.returncode != 0 or not result.stdout.strip():
|
|
logger.warning("amd-smi returned code %d", result.returncode)
|
|
return None
|
|
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 (may be 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 we have seen), dict-shaped values with explicit units
|
|
(``{"value": 192, "unit": "GiB"}`` on newer releases), and plain
|
|
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 (including amd-smi) use binary units even
|
|
# when labeling them "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, which is 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 branch 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."""
|
|
# amd-smi metric 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 multiple keys in priority order.
|
|
# dict.get() returns "N/A" strings rather than falling through,
|
|
# so we must try each key and check if 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 to handle varying amd-smi output formats.
|
|
# Newer amd-smi versions may return {"value": 192, "unit": "GiB"}.
|
|
# Newer amd-smi uses "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 that is missing every
|
|
expected field (error response, unsupported card, hipless container).
|
|
In that case ``_extract_gpu_metrics`` produces a dict where every value
|
|
is ``None`` -- callers must surface this as ``available: False`` rather
|
|
than ``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 the
|
|
# .get() access with an isinstance check so a malformed scalar /
|
|
# string response from amd-smi 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 that should be 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 explicitly 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
|
|
# Filter out empty tokens after splitting. This tolerates minor
|
|
# typos like ``HIP_VISIBLE_DEVICES=",1"`` (leading comma, user
|
|
# clearly meant to narrow to 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):
|
|
# amd-smi returned a 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 return a dict with a "gpus" /
|
|
# "gpu" key wrapping the list. Guard non-dict / non-list envelopes
|
|
# (scalar / string fallbacks from malformed output) so the .get()
|
|
# access cannot raise AttributeError on an unexpected shape.
|
|
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 defensively: if amd-smi ever ships a
|
|
# scalar inside its "gpus" array (observed on some malformed
|
|
# output), _extract_gpu_metrics would raise AttributeError on
|
|
# the first .get() call.
|
|
if not isinstance(gpu_data, dict):
|
|
continue
|
|
# Use AMD-reported GPU ID when available, fall back to enumeration
|
|
# index. Newer amd-smi versions wrap scalars as ``{"value": 0,
|
|
# "unit": "none"}``, so route raw_id through ``_parse_numeric``
|
|
# which already 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.debug(
|
|
"amd-smi GPU id %r could not be parsed; falling back to "
|
|
"enumeration index %d",
|
|
raw_id,
|
|
fallback_idx,
|
|
)
|
|
idx = fallback_idx
|
|
else:
|
|
idx = int(parsed_id)
|
|
if idx not in visible_set:
|
|
continue
|
|
metrics = _extract_gpu_metrics(gpu_data)
|
|
if not _has_real_metrics(metrics):
|
|
# Skip ghost entries: an amd-smi response that decodes to a
|
|
# dict but contains 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",
|
|
}
|