diff --git a/pyproject.toml b/pyproject.toml index dee2ce8279..f0c3560a01 100644 --- a/pyproject.toml +++ b/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]", diff --git a/tests/conftest.py b/tests/conftest.py index 4049c4af81..e30add6199 100644 --- a/tests/conftest.py +++ b/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()}.") diff --git a/tests/gpu_info.py b/tests/gpu_info.py new file mode 100644 index 0000000000..501324e2c4 --- /dev/null +++ b/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-' 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]() diff --git a/tests/monitor_test_mem.py b/tests/monitor_test_mem.py index 2c98dc233d..531dfebd43 100644 --- a/tests/monitor_test_mem.py +++ b/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]: