refactor: support amd detection

Signed-off-by: thxCode <[email protected]>

Author: thxCode

gpustack_runtime/cmds/detector.py

@@ -114,14 +114,12 @@ def format_devices_table(devs: Devices) -> str:
     # Device rows
     device_lines = []
     for dev in devs:
-        row_data = [
-            dev.index,
+        row_data: list[str] = [
+            str(dev.index),
             dev.name if dev.name else "N/A",
-            f"{dev.memory_used}MiB / {dev.memory}MiB"
-            if dev.memory and dev.memory_used
-            else "N/A",
-            f"{dev.cores_utilization}%" if dev.cores_utilization else "N/A",
-            f"{dev.temperature}C" if dev.temperature else "N/A",
+            f"{dev.memory_used}MiB / {dev.memory}MiB",
+            f"{dev.cores_utilization}%",
+            f"{dev.temperature}C" if dev.temperature is not None else "N/A",
             dev.compute_capability if dev.compute_capability else "N/A",
         ]
 

gpustack_runtime/detector/__init__.py

@@ -10,12 +10,14 @@
     supported_backends,
     supported_manufacturers,
 )
+from .amd import AMDDetector
 from .ascend import AscendDetector
 from .nvidia import NVIDIADetector
 
 detectors: list[Detector] = [
-    NVIDIADetector(),
+    AMDDetector(),
     AscendDetector(),
+    NVIDIADetector(),
 ]
 
 

gpustack_runtime/detector/__types__.py

@@ -149,34 +149,34 @@ class Device:
     """
     UUID of the device.
     """
-    driver_version: str = ""
+    driver_version: str | None = None
     """
     Driver version of the device.
     """
     driver_version_tuple: list[int | str] | None = None
     """
     Driver version tuple of the device.
-    None if `driver_version` is blank.
+    None if `driver_version` is missed.
     """
-    runtime_version: str = ""
+    runtime_version: str | None = None
     """
     Runtime version of the device.
     """
     runtime_version_tuple: list[int | str] | None = None
     """
     Runtime version tuple of the device.
-    None if `runtime_version` is blank.
+    None if `runtime_version` is missed.
     """
-    compute_capability: str = ""
+    compute_capability: str | None = None
     """
     Compute capability of the device.
     """
     compute_capability_tuple: list[int | str] | None = None
     """
     Compute capability tuple of the device.
-    None if `compute_capability` is blank.
+    None if `compute_capability` is missed.
     """
-    cores: int = 0
+    cores: int | None = None
     """
     Total cores of the device.
     """
@@ -196,15 +196,15 @@ class Device:
     """
     Memory utilization of the device in percentage.
     """
-    temperature: int = 0
+    temperature: int | None = None
     """
     Temperature of the device in Celsius.
     """
-    power: int = 0
+    power: int | None = None
     """
     Power consumption of the device in Watts.
     """
-    power_used: int = 0
+    power_used: int | None = None
     """
     Used power of the device in Watts.
     """

gpustack_runtime/detector/__utils__.py

@@ -11,7 +11,7 @@ class PCIDevice:
     """
     Vendor ID of the PCI device.
     """
-    path: Path
+    path: str
     """
     Path to the PCI device in sysfs.
     """
@@ -80,7 +80,7 @@ def get_pci_devices(
         pci_devices.append(
             PCIDevice(
                 vendor=dev_vendor,
-                path=dev_path,
+                path=str(dev_path),
                 address=dev_address,
                 class_=dev_class,
                 config=dev_config,
@@ -92,7 +92,7 @@ def get_pci_devices(
 
 @dataclass
 class DeviceFile:
-    path: Path
+    path: str
     """
     Path to the device file.
     """
@@ -102,7 +102,7 @@ class DeviceFile:
     """
 
 
-def get_device_files(pattern: str, directory: Path = Path("/dev")) -> list[DeviceFile]:
+def get_device_files(pattern: str, directory: Path | str = "/dev") -> list[DeviceFile]:
     r"""
     Get device files with the given pattern.
 
@@ -123,24 +123,27 @@ def get_device_files(pattern: str, directory: Path = Path("/dev")) -> list[Devic
         msg = "Pattern must include a regex group for the number, e.g nvidia(?P<number>\\d+)."
         raise ValueError(msg)
 
+    if isinstance(directory, str):
+        directory = Path(directory)
+
     device_files = []
     if not directory.exists():
         return device_files
 
     regex = re.compile(f"^{directory!s}/{pattern}$")
-    for path in directory.iterdir():
-        matched = regex.match(str(path))
+    for file_path in directory.iterdir():
+        matched = regex.match(str(file_path))
         if not matched:
             continue
-        number = matched.group("number")
+        file_number = matched.group("number")
         try:
-            number = int(number)
+            file_number = int(file_number)
         except ValueError:
-            number = None
+            file_number = None
         device_files.append(
             DeviceFile(
-                path=path,
-                number=number,
+                path=str(file_path),
+                number=file_number,
             ),
         )
 

gpustack_runtime/detector/amd.py

@@ -0,0 +1,178 @@
+from __future__ import annotations
+
+import contextlib
+import logging
+from functools import lru_cache
+
+from .. import envs
+from .__types__ import Detector, Device, Devices, ManufacturerEnum
+from .__utils__ import PCIDevice, get_device_files, get_pci_devices
+
+logger = logging.getLogger(__name__)
+
+
+class AMDDetector(Detector):
+    """
+    Detect AMD GPUs.
+    """
+
+    @staticmethod
+    @lru_cache
+    def is_supported() -> bool:
+        """
+        Check if the AMD detector is supported.
+
+        Returns:
+            True if supported, False otherwise.
+
+        """
+        supported = False
+        if envs.GPUSTACK_RUNTIME_DETECT.lower() not in ("auto", "amd"):
+            logger.debug("AMD detection is disabled by environment variable")
+            return supported
+
+        pci_devs = AMDDetector.detect_pci_devices()
+        if not pci_devs:
+            logger.debug("No AMD PCI devices found")
+            return supported
+
+        try:
+            import amdsmi as pyamdsmi  # noqa: PLC0415
+        except ImportError:
+            if logger.isEnabledFor(logging.DEBUG):
+                logger.exception("amdsmi module is not installed")
+            return supported
+
+        try:
+            pyamdsmi.amdsmi_init()
+            pyamdsmi.amdsmi_shut_down()
+            supported = True
+        except pyamdsmi.AmdSmiException:
+            if logger.isEnabledFor(logging.DEBUG):
+                logger.exception("Failed to initialize AMD SMI")
+
+        return supported
+
+    @staticmethod
+    @lru_cache
+    def detect_pci_devices() -> dict[str, PCIDevice] | None:
+        pci_devs = get_pci_devices(vendor="0x1002")
+        if not pci_devs:
+            return None
+        return {dev.address: dev for dev in pci_devs}
+
+    def __init__(self):
+        super().__init__(ManufacturerEnum.AMD)
+
+    def detect(self) -> Devices | None:
+        """
+        Detect AMD GPUs.
+
+        Returns:
+            A list of detected AMD GPU devices,
+            or None if detection fails.
+
+        """
+        if not self.is_supported():
+            return None
+
+        try:
+            import amdsmi as pyamdsmi  # noqa: PLC0415
+        except ImportError:
+            if logger.isEnabledFor(logging.DEBUG):
+                logger.exception("amdsmi module is not installed")
+            return None
+
+        ret: Devices = []
+
+        try:
+            pyamdsmi.amdsmi_init()
+
+            sys_runtime_ver = None
+            sys_runtime_ver_t = None
+
+            devs = pyamdsmi.amdsmi_get_processor_handles()
+            dev_files = None
+            for dev_idx, dev in enumerate(devs):
+                dev_index = dev_idx
+                if envs.GPUSTACK_RUNTIME_DETECT_PHYSICAL_INDEX_PRIORITY:
+                    if dev_files is None:
+                        dev_files = get_device_files(
+                            pattern=r"card(?P<number>\d+)",
+                            directory="/dev/dri",
+                        )
+                    if len(dev_files) > dev_idx:
+                        dev_file = dev_files[dev_idx]
+                        if dev_file.number is not None:
+                            dev_index = dev_file.number - 1
+
+                dev_uuid = pyamdsmi.amdsmi_get_gpu_device_uuid(dev)
+
+                dev_gpu_driver_info = pyamdsmi.amdsmi_get_gpu_driver_info(dev)
+                dev_driver_ver = dev_gpu_driver_info.get("driver_version")
+                dev_driver_ver_t = [
+                    int(v) if v.isdigit() else v for v in dev_driver_ver.split(".")
+                ]
+
+                dev_gpu_board_info = pyamdsmi.amdsmi_get_gpu_board_info(dev)
+                dev_name = "AMD " + dev_gpu_board_info.get("product_name")
+
+                dev_gpu_metrics_info = pyamdsmi.amdsmi_get_gpu_metrics_info(dev)
+                dev_cores_util = dev_gpu_metrics_info.get("average_gfx_activity", 0)
+                dev_gpu_vram_usage = pyamdsmi.amdsmi_get_gpu_vram_usage(dev)
+                dev_mem = dev_gpu_vram_usage.get("vram_total")
+                dev_mem_used = dev_gpu_vram_usage.get("vram_used")
+                dev_temp = dev_gpu_metrics_info.get("temperature_hotspot", 0)
+
+                dev_power_info = pyamdsmi.amdsmi_get_power_info(dev)
+                dev_power = dev_power_info.get("power_limit", 0) // 1000000  # uW to W
+                dev_power_used = (
+                    dev_power_info.get("current_socket_power")
+                    if dev_power_info.get("current_socket_power", "N/A") != "N/A"
+                    else dev_power_info.get("average_socket_power", 0)
+                )
+
+                dev_compute_partition = None
+                with contextlib.suppress(pyamdsmi.AmdSmiException):
+                    dev_compute_partition = pyamdsmi.amdsmi_get_gpu_compute_partition(
+                        dev,
+                    )
+
+                dev_appendix = {
+                    "vgpu": dev_compute_partition is not None,
+                }
+
+                ret.append(
+                    Device(
+                        manufacturer=self.manufacturer,
+                        index=dev_index,
+                        name=dev_name,
+                        uuid=dev_uuid,
+                        driver_version=dev_driver_ver,
+                        driver_version_tuple=dev_driver_ver_t,
+                        runtime_version=sys_runtime_ver,
+                        runtime_version_tuple=sys_runtime_ver_t,
+                        cores_utilization=dev_cores_util,
+                        memory=dev_mem,
+                        memory_used=dev_mem_used,
+                        memory_utilization=(
+                            (dev_mem_used * 100 // dev_mem) if dev_mem > 0 else 0
+                        ),
+                        temperature=dev_temp,
+                        power=dev_power,
+                        power_used=dev_power_used,
+                        appendix=dev_appendix,
+                    ),
+                )
+        except pyamdsmi.AmdSmiException:
+            if logger.isEnabledFor(logging.DEBUG):
+                logger.exception("Failed to fetch devices")
+            raise
+        except Exception:
+            if logger.isEnabledFor(logging.DEBUG):
+                logger.exception("Failed to process devices fetching")
+            raise
+        finally:
+            pyamdsmi.amdsmi_shut_down()
+
+        return ret