Files
2026-07-13 13:17:40 +08:00

433 lines
17 KiB
Python

# MIT License
#
# Copyright (c) 2023 Advanced Micro Devices, Inc.
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
# Python bindings shim: wraps the amd-smi Python interface with the same
# smi_* function signatures previously provided by the rocm-smi ctypes binding.
import logging
from collections import namedtuple
# Defer amdsmi import so that importing this module never raises ModuleNotFoundError
# on systems without the amdsmi package (non-AMD machines, CI, etc.).
# smi_initialize() will raise ImportError if amdsmi is not installed.
try:
import amdsmi
from amdsmi import (
AmdSmiInitFlags,
AmdSmiMemoryType,
AmdSmiLibraryException,
)
_AMDSMI_AVAILABLE = True
except ImportError:
amdsmi = None
AmdSmiInitFlags = None
AmdSmiMemoryType = None
AmdSmiLibraryException = Exception # so except clauses below don't NameError
_AMDSMI_AVAILABLE = False
# Module-level handle list, indexed by integer device index.
# Populated during smi_initialize(), cleared on smi_shutdown().
_handles = []
# ── Namedtuple used by smi_get_compute_process_info_by_device ────────────────
# Keeps the same attribute interface (.process_id, .vram_usage) that
# gpu_providers.py expects on the returned objects.
ProcessInfo = namedtuple("ProcessInfo", ["process_id", "vram_usage"])
# ── Init / shutdown ───────────────────────────────────────────────────────────
def smi_initialize():
"""Initialize amd-smi and populate the device handle list."""
if not _AMDSMI_AVAILABLE:
raise ImportError(
"amdsmi package is not installed. "
"Install with: pip install amdsmi"
)
global _handles
amdsmi.amdsmi_init(AmdSmiInitFlags.INIT_AMD_GPUS)
_handles = amdsmi.amdsmi_get_processor_handles()
def smi_shutdown():
"""Shut down amd-smi and clear the handle list."""
global _handles
_handles = []
amdsmi.amdsmi_shut_down()
# ── Device enumeration ────────────────────────────────────────────────────────
def smi_get_device_count():
"""Return the number of AMD GPU devices."""
return len(_handles)
# ── Identity ──────────────────────────────────────────────────────────────────
def smi_get_device_id(dev):
"""Return the 16-bit device ID for GPU dev."""
try:
return amdsmi.amdsmi_get_gpu_id(_handles[dev])
except AmdSmiLibraryException as e:
logging.debug("smi_get_device_id(%d) failed: %s", dev, e)
return -1
def smi_get_device_name(dev):
"""Return the market name string for GPU dev (e.g. 'Instinct MI300X')."""
try:
return amdsmi.amdsmi_get_gpu_asic_info(_handles[dev])["market_name"]
except AmdSmiLibraryException as e:
logging.debug("smi_get_device_name(%d) failed: %s", dev, e)
return ""
def smi_get_device_unique_id(dev):
"""Return the UUID string for GPU dev (e.g. 'GPU-xxxxxxxx-xxxx-...')."""
try:
return amdsmi.amdsmi_get_gpu_device_uuid(_handles[dev])
except AmdSmiLibraryException as e:
logging.debug("smi_get_device_unique_id(%d) failed: %s", dev, e)
return ""
# ── Utilization ───────────────────────────────────────────────────────────────
def smi_get_device_utilization(dev):
"""Return GPU busy percent (0-100) for GPU dev."""
try:
return amdsmi.amdsmi_get_gpu_busy_percent(_handles[dev])
except AmdSmiLibraryException as e:
logging.debug("smi_get_device_utilization(%d) failed: %s", dev, e)
return -1
def smi_get_device_memory_busy(dev):
"""Return memory controller busy percent (0-100) for GPU dev."""
try:
activity = amdsmi.amdsmi_get_gpu_activity(_handles[dev])
val = activity.get("umc_activity", "N/A")
return val if val != "N/A" else -1
except AmdSmiLibraryException as e:
logging.debug("smi_get_device_memory_busy(%d) failed: %s", dev, e)
return -1
# ── Memory ────────────────────────────────────────────────────────────────────
def smi_get_device_memory_used(dev, type="VRAM"):
"""Return used VRAM for GPU dev in bytes."""
try:
return amdsmi.amdsmi_get_gpu_memory_usage(_handles[dev], AmdSmiMemoryType.VRAM)
except AmdSmiLibraryException as e:
logging.debug("smi_get_device_memory_used(%d) failed: %s", dev, e)
return -1
def smi_get_device_memory_total(dev, type="VRAM"):
"""Return total VRAM for GPU dev in bytes."""
try:
return amdsmi.amdsmi_get_gpu_memory_total(_handles[dev], AmdSmiMemoryType.VRAM)
except AmdSmiLibraryException as e:
logging.debug("smi_get_device_memory_total(%d) failed: %s", dev, e)
return -1
def smi_get_device_memory_reserved_pages(dev):
"""Return reserved/bad memory pages info for GPU dev, or -1 if unsupported."""
try:
pages = amdsmi.amdsmi_get_gpu_memory_reserved_pages(_handles[dev])
return (len(pages), pages)
except AmdSmiLibraryException as e:
logging.debug("smi_get_device_memory_reserved_pages(%d) failed: %s", dev, e)
return -1
# ── Power ─────────────────────────────────────────────────────────────────────
def smi_get_device_average_power(dev):
"""Return average socket power for GPU dev in watts (float), or -1."""
try:
info = amdsmi.amdsmi_get_power_info(_handles[dev])
val = info["average_socket_power"]
return float(val) if val != "N/A" else -1
except AmdSmiLibraryException as e:
logging.debug("smi_get_device_average_power(%d) failed: %s", dev, e)
return -1
# ── Processes ─────────────────────────────────────────────────────────────────
def smi_get_device_compute_process():
"""Return list of PIDs running compute workloads on any GPU."""
try:
procs = amdsmi.amdsmi_get_gpu_compute_process_info()
return [p["process_id"] for p in procs]
except AmdSmiLibraryException as e:
logging.debug("smi_get_device_compute_process() failed: %s", e)
return []
def smi_get_compute_process_info_by_device(device_id: int, proc_ids: list) -> list:
"""Return ProcessInfo(process_id, vram_usage) for compute processes on GPU device_id.
proc_ids is the compute-only PID list from smi_get_device_compute_process().
amdsmi_get_gpu_process_list returns all processes on the device (including graphics),
so filter to proc_ids to return only compute processes, matching the old API contract.
vram_usage is in bytes to match the contract expected by gpu_providers.py.
"""
try:
compute_pids = set(proc_ids)
entries = amdsmi.amdsmi_get_gpu_process_list(_handles[device_id])
result = []
for e in entries:
if e["pid"] not in compute_pids:
continue
# vram_mem is in bytes in amdsmi_get_gpu_process_list
vram_bytes = e.get("memory_usage", {}).get("vram_mem", 0)
result.append(ProcessInfo(
process_id=e["pid"],
vram_usage=vram_bytes,
))
return result
except AmdSmiLibraryException as e:
logging.debug(
"smi_get_compute_process_info_by_device(%d) failed: %s", device_id, e
)
return []
# ── Driver / kernel version ───────────────────────────────────────────────────
def smi_get_kernel_version():
"""Return the kernel driver version string."""
try:
if not _handles:
return ""
info = amdsmi.amdsmi_get_gpu_driver_info(_handles[0])
return info.get("driver_version", "")
except AmdSmiLibraryException as e:
logging.debug("smi_get_kernel_version() failed: %s", e)
return ""
# ── PCIe ─────────────────────────────────────────────────────────────────────
def smi_get_device_pcie_bandwidth(dev):
"""Return PCIe bandwidth info struct for GPU dev, or -1 if unsupported."""
try:
return amdsmi.amdsmi_get_gpu_pci_bandwidth(_handles[dev])
except AmdSmiLibraryException as e:
logging.debug("smi_get_device_pcie_bandwidth(%d) failed: %s", dev, e)
return -1
def smi_get_device_pci_id(dev):
"""Return the BDF ID (Bus/Device/Function) for GPU dev as an int."""
try:
return amdsmi.amdsmi_get_gpu_bdf_id(_handles[dev])
except AmdSmiLibraryException as e:
logging.debug("smi_get_device_pci_id(%d) failed: %s", dev, e)
return -1
def smi_get_device_topo_numa_affinity(dev):
"""Return NUMA affinity for GPU dev."""
try:
return amdsmi.amdsmi_get_gpu_topo_numa_affinity(_handles[dev])
except AmdSmiLibraryException as e:
logging.debug("smi_get_device_topo_numa_affinity(%d) failed: %s", dev, e)
return -1
def smi_get_device_pcie_throughput(dev):
"""Return PCIe throughput (sent + received) * max_packet_size for GPU dev."""
try:
info = amdsmi.amdsmi_get_gpu_pci_throughput(_handles[dev])
return (info["sent"] + info["received"]) * info["max_pkt_sz"]
except AmdSmiLibraryException as e:
logging.debug("smi_get_device_pcie_throughput(%d) failed: %s", dev, e)
return -1
def smi_get_device_pci_replay_counter(dev):
"""Return PCIe replay counter for GPU dev."""
try:
return amdsmi.amdsmi_get_gpu_pci_replay_counter(_handles[dev])
except AmdSmiLibraryException as e:
logging.debug("smi_get_device_pci_replay_counter(%d) failed: %s", dev, e)
return -1
# ── Compute partition ─────────────────────────────────────────────────────────
def smi_get_device_compute_partition(dev):
"""Return the current compute partition string for GPU dev."""
try:
return amdsmi.amdsmi_get_gpu_compute_partition(_handles[dev])
except AmdSmiLibraryException as e:
logging.debug("smi_get_device_compute_partition(%d) failed: %s", dev, e)
return ""
def smi_set_device_compute_partition(dev, partition):
"""Set the compute partition for GPU dev. partition is an AmdSmiComputePartitionType."""
try:
amdsmi.amdsmi_set_gpu_compute_partition(_handles[dev], partition)
return True
except AmdSmiLibraryException as e:
logging.debug("smi_set_device_compute_partition(%d) failed: %s", dev, e)
return False
def smi_reset_device_compute_partition(dev):
"""Reset the compute partition for GPU dev to its boot state."""
try:
amdsmi.amdsmi_set_gpu_accelerator_partition_profile(_handles[dev], 0)
return True
except AmdSmiLibraryException as e:
logging.debug("smi_reset_device_compute_partition(%d) failed: %s", dev, e)
return False
# ── Memory partition ──────────────────────────────────────────────────────────
def smi_get_device_memory_partition(dev):
"""Return the current memory partition string for GPU dev."""
try:
return amdsmi.amdsmi_get_gpu_memory_partition(_handles[dev])
except AmdSmiLibraryException as e:
logging.debug("smi_get_device_memory_partition(%d) failed: %s", dev, e)
return ""
def smi_set_device_memory_partition(dev, partition):
"""Set the memory partition for GPU dev. partition is an AmdSmiMemoryPartitionType."""
try:
amdsmi.amdsmi_set_gpu_memory_partition(_handles[dev], partition)
return True
except AmdSmiLibraryException as e:
logging.debug("smi_set_device_memory_partition(%d) failed: %s", dev, e)
return False
def smi_reset_device_memory_partition(dev):
"""Reset the memory partition for GPU dev to its boot state."""
try:
amdsmi.amdsmi_set_gpu_memory_partition_mode(_handles[dev], 0)
return True
except AmdSmiLibraryException as e:
logging.debug("smi_reset_device_memory_partition(%d) failed: %s", dev, e)
return False
# ── Hardware topology ─────────────────────────────────────────────────────────
def smi_get_device_topo_numa_node_number(dev):
"""Return the NUMA node number for GPU dev."""
try:
return amdsmi.amdsmi_topo_get_numa_node_number(_handles[dev])
except AmdSmiLibraryException as e:
logging.debug("smi_get_device_topo_numa_node_number(%d) failed: %s", dev, e)
return -1
def smi_get_device_topo_link_weight(dev_src, dev_dst):
"""Return the link weight between two GPUs."""
try:
return amdsmi.amdsmi_topo_get_link_weight(_handles[dev_src], _handles[dev_dst])
except AmdSmiLibraryException as e:
logging.debug(
"smi_get_device_topo_link_weight(%d, %d) failed: %s", dev_src, dev_dst, e
)
return -1
def smi_get_device_minmax_bandwidth(dev_src, dev_dst):
"""Return (min_bandwidth, max_bandwidth) between two GPUs, or -1."""
try:
result = amdsmi.amdsmi_get_minmax_bandwidth_between_processors(
_handles[dev_src], _handles[dev_dst]
)
return (result["min_bandwidth"], result["max_bandwidth"])
except AmdSmiLibraryException as e:
logging.debug(
"smi_get_device_minmax_bandwidth(%d, %d) failed: %s", dev_src, dev_dst, e
)
return -1
def smi_get_device_link_type(dev_src, dev_dst):
"""Return (hops, link_type) between two GPUs, or -1."""
try:
result = amdsmi.amdsmi_topo_get_link_type(_handles[dev_src], _handles[dev_dst])
return (result["hops"], result["type"])
except AmdSmiLibraryException as e:
logging.debug(
"smi_get_device_link_type(%d, %d) failed: %s", dev_src, dev_dst, e
)
return -1
def smi_is_device_p2p_accessible(dev_src, dev_dst):
"""Return True if two GPUs are P2P accessible, False if not, -1 on error."""
try:
return amdsmi.amdsmi_is_P2P_accessible(_handles[dev_src], _handles[dev_dst])
except AmdSmiLibraryException as e:
logging.debug(
"smi_is_device_p2p_accessible(%d, %d) failed: %s", dev_src, dev_dst, e
)
return -1
# ── XGMI ─────────────────────────────────────────────────────────────────────
def smi_get_device_xgmi_error_status(dev):
"""Return the XGMI error status (AmdSmiXgmiStatus int value) for GPU dev."""
try:
status = amdsmi.amdsmi_gpu_xgmi_error_status(_handles[dev])
return status.value
except AmdSmiLibraryException as e:
logging.debug("smi_get_device_xgmi_error_status(%d) failed: %s", dev, e)
return -1
def smi_reset_device_xgmi_error(dev):
"""Reset XGMI error status for GPU dev. Returns True on success."""
try:
amdsmi.amdsmi_reset_gpu_xgmi_error(_handles[dev])
return True
except AmdSmiLibraryException as e:
logging.debug("smi_reset_device_xgmi_error(%d) failed: %s", dev, e)
return False
def smi_get_device_xgmi_hive_id(dev):
"""Return the XGMI hive ID for GPU dev."""
try:
return amdsmi.amdsmi_get_xgmi_info(_handles[dev])["xgmi_hive_id"]
except AmdSmiLibraryException as e:
logging.debug("smi_get_device_xgmi_hive_id(%d) failed: %s", dev, e)
return -1