fix(monitoring): report NVML metrics for the GPUs visible to CUDA

CHANGELOG.md

@@ -20,6 +20,7 @@ title: "Changelog"
 - fix: save Torch-TensorRT AOT models whose inputs include non-tensor values
 - fix: `inspect()` accepts `MinMaxRandomDataset`, directly or through `DataLoaderFactory`, and no longer replaces the factory's dataset
 - fix: AOT tuning writes and clears module caches under `config.cache_dir` instead of always using `AITUNE_CACHE_DIR`
+- fix: hardware metrics report `cuda:N` GPU memory, utilization, and power for the N-th GPU visible through `CUDA_VISIBLE_DEVICES` instead of the N-th GPU in the machine
 - breaking change: concrete strategy constructors now require explicit backend candidates; use `resolve_strategy` for dynamic defaults
 - breaking change: TorchInductor JIT, Torch-TensorRT JIT, and TorchAO checkpoints now require disk-backed warmup samples; legacy checkpoints with inline sample data are no longer supported
 

aitune/utils/monitoring/hardware_metrics/metrics.py

@@ -3,6 +3,7 @@
 """Metrics providers for hardware metrics."""
 
 import atexit
+import os
 from abc import ABC, abstractmethod
 from typing import Any
 
@@ -14,6 +15,7 @@
     nvmlDeviceGetMemoryInfo,
     nvmlDeviceGetPowerUsage,
     nvmlDeviceGetUtilizationRates,
+    nvmlDeviceGetUUID,
     nvmlInit,
     nvmlShutdown,
 )
@@ -53,13 +55,58 @@ def get_metrics(self) -> dict[str, Any]:
         }
 
 
+def _resolve_visible_device(entry: str, uuids: list[str], by_uuid: bool) -> int | None:
+    """Returns the NVML index of one CUDA_VISIBLE_DEVICES entry, or None if CUDA would stop parsing there."""
+    if by_uuid:
+        matches = [index for index, uuid in enumerate(uuids) if entry.startswith("GPU-") and uuid.startswith(entry)]
+        return matches[0] if len(matches) == 1 else None
+    if entry.isdecimal() and int(entry) < len(uuids):
+        return int(entry)
+    return None
+
+
+def _visible_device_indices() -> list[int]:
+    """Returns the NVML indices of the GPUs visible to CUDA, in CUDA device order.
+
+    NVML enumerates every GPU in the machine and ignores CUDA_VISIBLE_DEVICES, while
+    ``cuda:N`` is the N-th visible GPU. Entries are device indices or (unique prefixes of)
+    ``GPU-`` UUIDs, and the first entry decides which. Like CUDA, parsing stops at the first
+    invalid entry, and a repeated entry hides every GPU. Index entries assume that CUDA
+    enumerates GPUs in PCI bus order, as NVML does; this holds for identical GPUs or with
+    ``CUDA_DEVICE_ORDER=PCI_BUS_ID``. UUID entries are always exact.
+    """
+    count = nvmlDeviceGetCount()
+    visible = os.environ.get("CUDA_VISIBLE_DEVICES")
+    if visible is None:
+        return list(range(count))
+    uuids = []
+    for index in range(count):
+        uuid = nvmlDeviceGetUUID(nvmlDeviceGetHandleByIndex(index))
+        uuids.append(uuid.decode() if isinstance(uuid, bytes) else uuid)
+    entries = [entry.strip() for entry in visible.split(",")]
+    by_uuid = entries[0].startswith("GPU-")
+    indices = []
+    for entry in entries:
+        index = _resolve_visible_device(entry, uuids, by_uuid)
+        if index is None:
+            break
+        if index in indices:
+            return []
+        indices.append(index)
+    return indices
+
+
 class NvmlMetricProvider(AbstractHardwareMetricProvider):
-    """Nvml metric provider."""
+    """Nvml metric provider.
+
+    Metrics are keyed by CUDA device (``cuda:N``), so only GPUs visible through
+    CUDA_VISIBLE_DEVICES are reported.
+    """
 
     def __init__(self):
         """Initializes the NvmlMetricProvider."""
         nvmlInit()
-        self.handles = [nvmlDeviceGetHandleByIndex(i) for i in range(nvmlDeviceGetCount())]
+        self.handles = [nvmlDeviceGetHandleByIndex(i) for i in _visible_device_indices()]
         self._closed = False
         atexit.register(self.close)
 

tests/unit/utils/monitoring/hardware_metrics/test_metrics.py

@@ -3,15 +3,95 @@
 
 """Tests for hardware metric providers."""
 
+from types import SimpleNamespace
+
 import pytest
 import torch
 
+import aitune.utils.monitoring.hardware_metrics.metrics as metrics_module
 from aitune.utils.monitoring.hardware_metrics.metrics import (
     HostMetricProvider,
     NvmlMetricProvider,
     TorchMetricProvider,
 )
 
+_FAKE_UUIDS = [f"GPU-{index}a{index}b0000-0000-0000-0000-00000000000{index}" for index in range(8)]
+
+
+@pytest.fixture
+def fake_nvml(monkeypatch):
+    """Replaces NVML with an 8-GPU fake whose handles are the NVML indices."""
+    monkeypatch.setattr(metrics_module, "nvmlInit", lambda: None)
+    monkeypatch.setattr(metrics_module, "nvmlShutdown", lambda: None)
+    monkeypatch.setattr(metrics_module, "nvmlDeviceGetCount", lambda: len(_FAKE_UUIDS))
+    monkeypatch.setattr(metrics_module, "nvmlDeviceGetHandleByIndex", lambda index: index)
+    monkeypatch.setattr(metrics_module, "nvmlDeviceGetUUID", lambda handle: _FAKE_UUIDS[handle])
+    monkeypatch.setattr(
+        metrics_module,
+        "nvmlDeviceGetMemoryInfo",
+        lambda handle: SimpleNamespace(used=handle, free=100 - handle, total=100),
+    )
+    monkeypatch.setattr(metrics_module, "nvmlDeviceGetUtilizationRates", lambda handle: SimpleNamespace(gpu=handle))
+    monkeypatch.setattr(metrics_module, "nvmlDeviceGetPowerUsage", lambda handle: 1000 * handle)
+
+
+@pytest.mark.parametrize(
+    "visible_devices, expected",
+    [
+        (None, list(range(8))),
+        ("4", [4]),
+        ("6,4", [6, 4]),
+        (" 5 , 7 ", [5, 7]),
+        (_FAKE_UUIDS[3], [3]),
+        (f"GPU-6a6,{_FAKE_UUIDS[2]}", [6, 2]),
+        ("4,9,5", [4]),
+        ("4,4", []),
+        (f"4,{_FAKE_UUIDS[5]}", [4]),
+        (f"{_FAKE_UUIDS[6]},4", [6]),
+        ("4,GPU-", [4]),
+        ("4,MIG-1234", [4]),
+        ("-1", []),
+        ("", []),
+    ],
+)
+def test_visible_device_indices_follow_cuda_visible_devices(fake_nvml, monkeypatch, visible_devices, expected):
+    """Test that NVML devices are selected and ordered the way CUDA parses CUDA_VISIBLE_DEVICES."""
+    if visible_devices is None:
+        monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False)
+    else:
+        monkeypatch.setenv("CUDA_VISIBLE_DEVICES", visible_devices)
+    assert metrics_module._visible_device_indices() == expected
+
+
+def test_visible_device_indices_accept_bytes_uuids(fake_nvml, monkeypatch):
+    """Test that UUIDs returned as bytes by older NVML bindings are matched."""
+    monkeypatch.setattr(metrics_module, "nvmlDeviceGetUUID", lambda handle: _FAKE_UUIDS[handle].encode())
+    monkeypatch.setenv("CUDA_VISIBLE_DEVICES", _FAKE_UUIDS[1])
+    assert metrics_module._visible_device_indices() == [1]
+
+
+def test_nvml_metric_provider_keys_metrics_by_cuda_device(fake_nvml, monkeypatch):
+    """Test that cuda:N metrics describe the N-th visible GPU, not the N-th GPU in the machine."""
+    monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "6,4")
+    provider = NvmlMetricProvider()
+    try:
+        metrics = provider.get_metrics()
+    finally:
+        provider.close()
+
+    assert metrics == {
+        "cuda:0_memory_used": 6,
+        "cuda:0_memory_free": 94,
+        "cuda:0_memory_total": 100,
+        "cuda:0_utilization": 6,
+        "cuda:0_power_usage_milliwatts": 6000,
+        "cuda:1_memory_used": 4,
+        "cuda:1_memory_free": 96,
+        "cuda:1_memory_total": 100,
+        "cuda:1_utilization": 4,
+        "cuda:1_power_usage_milliwatts": 4000,
+    }
+
 
 def test_host_metric_provider_get_metrics_has_expected_keys():
     """Test that HostMetricProvider.get_metrics returns expected keys."""
@@ -49,6 +129,7 @@
     metrics = provider.get_metrics()
 
     num_gpus = len(provider.handles)
+    assert num_gpus == torch.cuda.device_count()
     for i in range(num_gpus):
         assert f"cuda:{i}_memory_used" in metrics
         assert f"cuda:{i}_memory_free" in metrics