[MISC] Add full support of AMD GPU to unit test and benchmark infra.

pyproject.toml

@@ -85,6 +85,9 @@ dev = [
     # Note that 'pytest-rerunfailures' is incompatible with 'pytest-forked'
     # - 16.0 is causing pytest-xdist to crash in case of failure or skipped tests
     # "pytest-rerunfailures!=16.0",
+    # GPU queries for the test infrastructure (tests/gpu_info.py); non-macOS so the CUDA device check runs on
+    # Linux and Windows. AMD's 'amdsmi' ships with ROCm, not PyPI, so the AMD backend is import-gated instead.
+    "nvidia-ml-py; platform_system != 'Darwin'",
     "setproctitle", # allows renaming the test processes on the cluster
     "syrupy",
     "huggingface_hub[hf_xet]",

tests/conftest.py

@@ -2,7 +2,6 @@
 import gc
 import logging
 import os
-import re
 import shutil
 import subprocess
 import sys
@@ -20,6 +19,8 @@
 from PIL import Image
 from syrupy.extensions.image import PNGImageSnapshotExtension
 
+from tests.gpu_info import detect_gpu_backend
+
 # Mock tkinter module for backward compatibility because it is a hard dependency for old Genesis versions
 has_tkinter = False
 try:
@@ -102,12 +103,14 @@ def _skip_reason(reason):
 
 
 def is_mem_monitoring_supported():
-    try:
-        assert sys.platform.startswith("linux")
-        subprocess.check_output(["nvidia-smi"], stderr=subprocess.STDOUT, timeout=10)
+    if not sys.platform.startswith("linux"):
+        return False, "mem-monitoring only supported on linux"
+
+    backend = detect_gpu_backend()
+    if backend is not None:
         return True, None
-    except Exception as exc:  # platform or nvidia-smi unavailable
-        return False, exc
+
+    return False, "no supported GPU backend detected"
 
 
 def pytest_make_parametrize_id(config, val, argname):
@@ -246,43 +249,34 @@ def _get_gpu_indices():
         return tuple(map(int, cuda_visible_devices.split(",")))
 
     if sys.platform == "linux":
-        nvidia_gpu_interface_path = "/proc/driver/nvidia/gpus/"
-        try:
-            return tuple(range(len(os.listdir(nvidia_gpu_interface_path))))
-        except FileNotFoundError:
-            warnings.warn(
-                f"'{nvidia_gpu_interface_path}' is not available. Multi-GPU support will be disabled. This is expected "
-                "on WSL2 where the NVIDIA proc interface is not mounted.",
-                stacklevel=2,
-            )
+        backend = detect_gpu_backend()
+        if backend is not None:
+            device_count = backend.get_device_count()
+            if device_count > 0:
+                return tuple(range(device_count))
+
+        warnings.warn(
+            "No GPU backend detected (neither NVML nor AMD SMI); multi-GPU support will be disabled.",
+            stacklevel=2,
+        )
 
     return (0,)
 
 
 def _torch_get_gpu_idx(device):
-    if sys.platform == "linux":
-        import torch
+    # The caller only invokes this for a CUDA device, so torch is using this GPU and its identity must be
+    # confirmable. Returns the resolved physical device index, or -1 when it cannot be confirmed (no GPU
+    # management library, or a UUID unknown to it), which the caller turns into a hard error rather than
+    # letting an unverified device through.
+    import torch
 
-        device_property = torch.cuda.get_device_properties(device)
-        device_uuid = str(device_property.uuid)
+    device_uuid = str(torch.cuda.get_device_properties(device).uuid)
 
-        nvidia_gpu_interface_path = "/proc/driver/nvidia/gpus/"
-        try:
-            for device_idx, device_path in enumerate(os.listdir(nvidia_gpu_interface_path)):
-                with open(os.path.join(nvidia_gpu_interface_path, device_path, "information"), "r") as f:
-                    device_info = f.read()
-                if re.search(rf"GPU UUID:\s+GPU-{device_uuid}", device_info):
-                    return device_idx
-            else:
-                return -1
-        except FileNotFoundError:
-            warnings.warn(
-                f"'{nvidia_gpu_interface_path}' is not available. Multi-GPU support will be disabled. This is expected "
-                "on WSL2 where the NVIDIA proc interface is not mounted.",
-                stacklevel=2,
-            )
+    backend = detect_gpu_backend()
+    if backend is None:
+        return -1
 
-    return 0
+    return backend.get_device_index_from_uuid(device_uuid)
 
 
 def _get_egl_index(gpu_index):
@@ -339,34 +333,12 @@ def pytest_xdist_auto_num_workers(config):
     else:
         # Cannot rely on 'torch' because this would force loading devices before configuring CUDA device visibility
         devices_vram_memory = None
-        try:
-            result = subprocess.run(
-                ["nvidia-smi", "--query-gpu=memory.total", "--format=csv,noheader,nounits"],
-                stdout=subprocess.PIPE,
-                stderr=subprocess.PIPE,
-                check=True,
-                text=True,
-            )
-            devices_vram_memory = tuple(int(e.strip()) for e in result.stdout.splitlines())
-        except ValueError:
-            # Unknown VRAM. Assuming unbounded.
-            vram_memory = float("inf")
-        except (FileNotFoundError, subprocess.CalledProcessError):
-            try:
-                result = subprocess.run(
-                    ["rocm-smi", "--showmeminfo", "vram", "-d", "0-255"],
-                    stdout=subprocess.PIPE,
-                    stderr=subprocess.PIPE,
-                    check=True,
-                    text=True,
-                )
-                devices_vram_memory = tuple(
-                    int(m.group(1)) for m in re.finditer(r"VRAM Total:\s+(\d+)\s*MiB", result.stdout)
-                )
-            except (FileNotFoundError, subprocess.CalledProcessError):
-                pass
-        if devices_vram_memory is not None:
-            assert len(set(devices_vram_memory)) == 1, "Heterogeneous Nvidia GPU devices not supported."
+        backend = detect_gpu_backend()
+        if backend is not None:
+            devices_vram_memory = backend.get_device_vram_mib()
+
+        if devices_vram_memory:
+            assert len(set(devices_vram_memory)) == 1, "Heterogeneous GPU devices not supported."
             num_gpus = len(devices_vram_memory)
             vram_memory = sum(devices_vram_memory) / 1024
         else:
@@ -836,6 +808,9 @@ def _RigidSimStaticConfig_init(self, *args, **kwargs):
             monkeypatch.setattr(RigidSimStaticConfig, "__init__", _RigidSimStaticConfig_init)
 
         if gs.backend != gs.cpu and gs.device.index is not None:
+            # The device torch selected must be one this worker is allowed to use. Anything else - including a
+            # -1 meaning the device could not be confirmed - fails hard rather than letting an unverified device
+            # through, on every platform.
             device_idx = _torch_get_gpu_idx(gs.device.index)
             if device_idx not in _get_gpu_indices():
                 raise RuntimeError(f"Invalid CUDA GPU device, got {device_idx}, not in {_get_gpu_indices()}.")

tests/gpu_info.py

@@ -0,0 +1,175 @@
+"""Cross-vendor GPU information for the test infrastructure.
+
+Each backend queries the GPUs through the vendor management library - NVIDIA Management Library (NVML) via
+nvidia-ml-py for NVIDIA, AMD SMI (amdsmi) for AMD - rather than parsing command-line tools or reading the
+driver proc/sysfs interface. The management libraries reach the devices through the same driver path as the
+compute runtime, so they report the GPUs actually allocated to the current process even inside a container
+whose proc/sysfs interface does not reflect that allocation (e.g. an attached or namespaced container). The
+management libraries are optional dependencies, so they are imported lazily and each backend is gated on
+is_available().
+"""
+
+import warnings
+from abc import ABC, abstractmethod
+
+
+class GpuBackend(ABC):
+    """Vendor-specific accessor for the GPUs visible to the current process."""
+
+    @classmethod
+    @abstractmethod
+    def is_available(cls) -> bool:
+        """Whether this backend's management library loads and its GPUs are usable in the current process."""
+
+    @abstractmethod
+    def get_device_count(self) -> int:
+        """Number of GPUs visible to the current process."""
+
+    @abstractmethod
+    def get_device_vram_mib(self) -> tuple[int, ...]:
+        """Total VRAM in MiB of each visible GPU, ordered by device index."""
+
+    @abstractmethod
+    def get_device_index_from_uuid(self, device_uuid: str) -> int:
+        """Device index of the GPU whose UUID matches, or -1 if no visible GPU does."""
+
+    @abstractmethod
+    def get_per_process_vram_mib(self) -> dict[int, int]:
+        """VRAM in MiB used across the visible GPUs by each process, keyed by process id."""
+
+
+class NvidiaBackend(GpuBackend):
+    """NVIDIA backend backed by the NVIDIA Management Library (NVML) through nvidia-ml-py."""
+
+    @classmethod
+    def is_available(cls) -> bool:
+        try:
+            import pynvml
+        except ImportError:
+            return False
+        # Validate the same calls __init__ relies on, so a driver that loads but cannot enumerate reports the
+        # backend as unavailable here instead of raising when it is later constructed.
+        try:
+            pynvml.nvmlInit()
+            pynvml.nvmlDeviceGetCount()
+        except pynvml.NVMLError:
+            return False
+        return True
+
+    def __init__(self):
+        import pynvml
+
+        # NVML reference-counts initialization and the test worker is short-lived, so the matching shutdown is
+        # left to process exit; the handles stay valid for the lifetime of this backend.
+        pynvml.nvmlInit()
+        self._pynvml = pynvml
+        self._handles = tuple(pynvml.nvmlDeviceGetHandleByIndex(i) for i in range(pynvml.nvmlDeviceGetCount()))
+
+    def get_device_count(self) -> int:
+        return len(self._handles)
+
+    def get_device_vram_mib(self) -> tuple[int, ...]:
+        return tuple(self._pynvml.nvmlDeviceGetMemoryInfo(handle).total >> 20 for handle in self._handles)
+
+    def get_device_index_from_uuid(self, device_uuid: str) -> int:
+        target = device_uuid.replace("-", "").lower()
+        for index, handle in enumerate(self._handles):
+            uuid = self._pynvml.nvmlDeviceGetUUID(handle)
+            if isinstance(uuid, bytes):
+                uuid = uuid.decode()
+            # NVML reports the UUID as 'GPU-<uuid>' while torch reports the bare UUID.
+            if uuid.removeprefix("GPU-").replace("-", "").lower() == target:
+                return index
+        return -1
+
+    def get_per_process_vram_mib(self) -> dict[int, int]:
+        usage: dict[int, int] = {}
+        for handle in self._handles:
+            # Genesis test workers use the GPU both for compute (Quadrants) and for rendering (EGL/OpenGL), which
+            # the driver reports as separate process lists.
+            for get_processes in (
+                self._pynvml.nvmlDeviceGetComputeRunningProcesses,
+                self._pynvml.nvmlDeviceGetGraphicsRunningProcesses,
+            ):
+                for proc in get_processes(handle):
+                    # usedGpuMemory is None when the driver cannot attribute memory to the process.
+                    if proc.usedGpuMemory is not None:
+                        usage[proc.pid] = usage.get(proc.pid, 0) + (proc.usedGpuMemory >> 20)
+        return usage
+
+
+class AmdBackend(GpuBackend):
+    """AMD backend backed by AMD SMI (amdsmi), the management library shipped with ROCm."""
+
+    @classmethod
+    def is_available(cls) -> bool:
+        try:
+            import amdsmi
+        except ImportError:
+            return False
+        # Validate the same calls __init__ relies on, so a ROCm setup that loads but cannot enumerate (e.g. a
+        # container without /dev/kfd access) reports the backend as unavailable here instead of raising when it
+        # is later constructed.
+        try:
+            amdsmi.amdsmi_init()
+            amdsmi.amdsmi_get_processor_handles()
+        except amdsmi.AmdSmiException:
+            return False
+        return True
+
+    def __init__(self):
+        import amdsmi
+
+        # Shutdown is left to process exit, mirroring the NVIDIA backend.
+        amdsmi.amdsmi_init()
+        self._amdsmi = amdsmi
+        self._handles = tuple(amdsmi.amdsmi_get_processor_handles())
+
+    def get_device_count(self) -> int:
+        return len(self._handles)
+
+    def get_device_vram_mib(self) -> tuple[int, ...]:
+        return tuple(
+            self._amdsmi.amdsmi_get_gpu_memory_total(handle, self._amdsmi.AmdSmiMemoryType.VRAM) >> 20
+            for handle in self._handles
+        )
+
+    def get_device_index_from_uuid(self, device_uuid: str) -> int:
+        target = device_uuid.replace("-", "").lower()
+        for index, handle in enumerate(self._handles):
+            uuid = self._amdsmi.amdsmi_get_gpu_device_uuid(handle)
+            if uuid.replace("-", "").lower() == target:
+                return index
+        return -1
+
+    def get_per_process_vram_mib(self) -> dict[int, int]:
+        usage: dict[int, int] = {}
+        for handle in self._handles:
+            for proc in self._amdsmi.amdsmi_get_gpu_process_list(handle):
+                # amdsmi_get_gpu_process_list returns process info dicts on recent ROCm and opaque handles on
+                # older ones, which must be resolved to a dict through amdsmi_get_gpu_process_info.
+                info = proc if isinstance(proc, dict) else self._amdsmi.amdsmi_get_gpu_process_info(handle, proc)
+                mem = info.get("memory_usage", {}).get("vram_mem") or info.get("mem")
+                if mem is not None:
+                    pid = int(info["pid"])
+                    usage[pid] = usage.get(pid, 0) + (int(mem) >> 20)
+        return usage
+
+
+def detect_gpu_backend() -> GpuBackend | None:
+    """Return a fresh instance of the first available GPU backend, or None when no GPU backend is usable.
+
+    A new instance is built on every call so the management library is initialized in the calling process,
+    which keeps the backend valid in forked children (pytest-forked runs every test in one) where the parent's
+    session and device handles would be stale.
+    """
+    backend_classes = (NvidiaBackend, AmdBackend)
+    available_backends = [backend_cls for backend_cls in backend_classes if backend_cls.is_available()]
+
+    if not available_backends:
+        return None
+
+    if len(available_backends) > 1:
+        warnings.warn("Multiple GPU backends were detected on the current system; using the first one.", stacklevel=2)
+
+    return available_backends[0]()

tests/monitor_test_mem.py

@@ -1,10 +1,14 @@
-from collections import defaultdict
-import subprocess
-import time
-import os
 import argparse
-import psutil
+import os
 import re
+import time
+from collections import defaultdict
+
+import psutil
+
+# This module is launched as a standalone script ('python tests/monitor_test_mem.py'), so its own directory is
+# on sys.path and 'gpu_info' is imported as a top-level module rather than through the 'tests' package.
+from gpu_info import detect_gpu_backend
 
 
 CHECK_INTERVAL = 2.0
@@ -48,26 +52,12 @@ def parse_test_name(test_name: str) -> dict[str, str]:
 
 
 def get_cuda_usage() -> dict[int, int]:
-    output = subprocess.check_output(["nvidia-smi"]).decode("utf-8")
-    section = 0
-    subsec = 0
-    res = {}
-    for line in output.split("\n"):
-        if line.startswith("|============"):
-            section += 1
-            subsec = 0
-            continue
-        if line.startswith("+-------"):
-            subsec += 1
-            continue
-        if section == 2 and subsec == 0:
-            if "No running processes" in line:
-                continue
-            split_line = line.split()
-            pid = int(split_line[4])
-            mem = int(split_line[-2].split("MiB")[0])
-            res[pid] = mem
-    return res
+    """VRAM in MiB used on the GPUs by each process, keyed by process id."""
+    backend = detect_gpu_backend()
+    if backend is not None:
+        return backend.get_per_process_vram_mib()
+
+    raise RuntimeError("No supported GPU backend available.")
 
 
 def get_test_name_by_pid() -> dict[int, str]: