[CI] Collect GPU memory telemetry for MIG slices (#56747)

Signed-off-by: Kevin Luu <51931015+khluu@users.noreply.github.com> Signed-off-by: khluu <khluu000@gmail.com> Co-authored-by: OpenAI Codex <codex@openai.com>
2026-09-13 21:30:29 -07:00
parent 214248c7d8
commit dc36fcce90
3 changed files with 371 additions and 12 deletions
@@ -27,9 +27,17 @@ These are device readings during each test's time interval, not process-level
 kernel attribution. Other processes and overlapping tests can contribute to
 the readings. Subsecond tests may have no sample. Missing readings never become
 zero utilization. MIG parent-device memory and utilization are omitted because
-they cannot be attributed to the job's partition. The initial implementation
-covers NVIDIA commands in the existing tracing rollout; CPU jobs and AMD
-mirrors do not collect GPU samples.
+they cannot be attributed to the job's partition. When the environment identifies
+an assigned MIG UUID, the collector queries that instance's used/total memory
+through vLLM's standalone bundled NVML binding in a bounded subprocess.
+The binding is loaded directly to avoid vLLM/Torch import side effects. Numeric
+`CUDA_VISIBLE_DEVICES` indices are resolved against explicit MIG UUIDs in
+`NVIDIA_VISIBLE_DEVICES`. For CDI assignments without an explicit UUID, a bounded
+helper initializes the CUDA driver and reads only CUDA-visible device UUIDs,
+without creating CUDA contexts or importing Torch. NVML verifies each MIG handle.
+Discovery or permission failures remain unknown; NVML sibling instances are never
+enumerated. MIG utilization stays absent.
+CPU jobs and AMD mirrors do not collect GPU samples.

 Run the helper tests without loading vLLM or requiring a GPU:

@@ -6,7 +6,10 @@
 from __future__ import annotations

 import csv
+import ctypes
+import importlib.util
 import io
+import json
 import math
 import os
 import secrets
@@ -15,6 +18,8 @@ import subprocess
 import sys
 import threading
 import time
+import uuid as uuid_module
+from pathlib import Path

 from ci_otel import Span, export_spans, record_spans

@@ -67,6 +72,118 @@ def parse_samples(output: str, timestamp_ns: int) -> list[dict]:
    return events


+def visible_mig_devices() -> list[str]:
+    """Resolve explicit MIG assignments without enumerating sibling instances."""
+    nvidia = os.environ.get("NVIDIA_VISIBLE_DEVICES", "").split(",")
+    nvidia = [value.strip() for value in nvidia]
+    cuda = os.environ.get("CUDA_VISIBLE_DEVICES")
+    devices = nvidia if cuda is None else [value.strip() for value in cuda.split(",")]
+    if devices and all(value.isdigit() for value in devices):
+        if not all(value.startswith("MIG-") for value in nvidia):
+            return []
+        try:
+            devices = [nvidia[int(value)] for value in devices]
+        except IndexError:
+            return []
+    if not devices or not all(value.startswith("MIG-") for value in devices):
+        return []
+    return list(dict.fromkeys(devices))[:MAX_DEVICES]
+
+
+def load_nvml():
+    """Load the standalone bundled binding without importing vLLM or Torch."""
+    # Test images install vLLM as a wheel and keep only tests/helpers in the
+    # checkout. Finding the top-level package does not execute its __init__.
+    package = importlib.util.find_spec("vllm")
+    if package is None or package.origin is None:
+        raise ImportError("Installed vLLM package unavailable")
+    path = Path(package.origin).parent / "third_party/pynvml.py"
+    spec = importlib.util.spec_from_file_location("ci_nvml", path)
+    if spec is None or spec.loader is None:
+        raise ImportError("Bundled NVML binding unavailable")
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[spec.name] = module
+    spec.loader.exec_module(module)
+    return module
+
+
+def cuda_visible_mig_devices() -> list[str]:
+    """Discover CDI assignments in a subprocess without creating CUDA contexts."""
+    pynvml = load_nvml()
+
+    driver = ctypes.CDLL("libcuda.so.1")
+    count = ctypes.c_int()
+    if driver.cuInit(0) or driver.cuDeviceGetCount(ctypes.byref(count)):
+        return []
+    devices = []
+    pynvml.nvmlInit()
+    try:
+        for index in range(min(count.value, MAX_DEVICES)):
+            device = ctypes.c_int()
+            identifier = (ctypes.c_ubyte * 16)()
+            if driver.cuDeviceGet(ctypes.byref(device), index):
+                continue
+            # The v2 entry point returns the compute instance UUID under MIG.
+            if driver.cuDeviceGetUuid_v2(ctypes.byref(identifier), device):
+                continue
+            uuid = "MIG-" + str(uuid_module.UUID(bytes=bytes(identifier)))
+            try:
+                handle = pynvml.nvmlDeviceGetHandleByUUID(uuid)
+                if pynvml.nvmlDeviceIsMigDeviceHandle(handle):
+                    devices.append(uuid)
+            except pynvml.NVMLError:
+                continue
+    finally:
+        pynvml.nvmlShutdown()
+    return devices
+
+
+def query_mig_samples(devices: list[str], timestamp_ns: int) -> list[dict]:
+    # Use the binding already shipped with vLLM; never initialize CUDA/Torch.
+    pynvml = load_nvml()
+
+    events = []
+    pynvml.nvmlInit()
+    try:
+        for index, uuid in enumerate(devices[:MAX_DEVICES]):
+            if not uuid.startswith("MIG-"):
+                continue
+            attributes: dict[str, str | int | bool] = {
+                "gpu.uuid": uuid,
+                "gpu.index": index,
+                "gpu.name": "NVIDIA MIG",
+                "gpu.status": "unavailable_mig_memory",
+            }
+            try:
+                handle = pynvml.nvmlDeviceGetHandleByUUID(uuid)
+                if not pynvml.nvmlDeviceIsMigDeviceHandle(handle):
+                    continue
+                memory = pynvml.nvmlDeviceGetMemoryInfo(handle)
+                used = _number(str(memory.used))
+                total = _number(str(memory.total))
+                if used is not None and total and used <= total:
+                    attributes.update(
+                        {
+                            "gpu.memory.used": used,
+                            "gpu.memory.total": total,
+                            "gpu.status": "mig_memory_only",
+                        }
+                    )
+                # NVML does not support per-MIG utilization on these devices.
+            except pynvml.NVMLError:
+                pass
+            events.append(
+                {
+                    "time_ns": timestamp_ns,
+                    "name": "ci.gpu.sample",
+                    "attributes": attributes,
+                }
+            )
+    finally:
+        pynvml.nvmlShutdown()
+    return events
+
+
 def sample_span(events: list[dict]) -> Span:
    return Span(
        trace_id=os.environ["CI_INFRA_TRACE_ID"],
@@ -88,16 +205,23 @@ def collect(parent_pid: int, stop: threading.Event) -> None:
    events: list[dict] = []
    batch_started = time.monotonic()
    failures = 0
+    discovery_attempted = False
+    mig_devices = visible_mig_devices()
+    query = (
+        [sys.executable, __file__, "--query-mig", *mig_devices]
+        if mig_devices
+        else [
+            "nvidia-smi",
+            f"--query-gpu={QUERY}",
+            "--format=csv,noheader,nounits",
+        ]
+    )
    try:
        while not stop.is_set() and os.getppid() == parent_pid:
            started = time.monotonic()
            try:
                result = subprocess.run(
-                    [
-                        "nvidia-smi",
-                        f"--query-gpu={QUERY}",
-                        "--format=csv,noheader,nounits",
-                    ],
+                    query,
                    capture_output=True,
                    text=True,
                    timeout=QUERY_TIMEOUT_SECONDS,
@@ -105,12 +229,41 @@ def collect(parent_pid: int, stop: threading.Event) -> None:
                )
                if stop.is_set():
                    break
-                samples = parse_samples(result.stdout, time.time_ns())
+                samples = (
+                    json.loads(result.stdout)
+                    if mig_devices
+                    else parse_samples(result.stdout, time.time_ns())
+                )
+                if (
+                    not mig_devices
+                    and not discovery_attempted
+                    and any(
+                        sample["attributes"].get("gpu.status") == "unsupported_mig"
+                        for sample in samples
+                    )
+                ):
+                    discovery_attempted = True
+                    # CDI can assign a MIG device without exporting its UUID.
+                    # Only CUDA-visible devices are inspected, never NVML siblings.
+                    try:
+                        discovered = subprocess.run(
+                            [sys.executable, __file__, "--discover-mig"],
+                            capture_output=True,
+                            text=True,
+                            timeout=QUERY_TIMEOUT_SECONDS,
+                            check=True,
+                        )
+                        mig_devices = json.loads(discovered.stdout)
+                    except (OSError, subprocess.SubprocessError, json.JSONDecodeError):
+                        mig_devices = []
+                    if mig_devices:
+                        query = [sys.executable, __file__, "--query-mig", *mig_devices]
+                        continue  # Do not create an extra parent-device series.
                if not samples:
                    break
                events.extend(samples)
                failures = 0
-            except (OSError, subprocess.SubprocessError):
+            except (OSError, subprocess.SubprocessError, json.JSONDecodeError):
                failures += 1
                if failures >= 3:
                    print(
@@ -142,4 +295,9 @@ def main() -> None:


 if __name__ == "__main__":
-    main()
+    if len(sys.argv) > 1 and sys.argv[1] == "--discover-mig":
+        print(json.dumps(cuda_visible_mig_devices()))
+    elif len(sys.argv) > 1 and sys.argv[1] == "--query-mig":
+        print(json.dumps(query_mig_samples(sys.argv[2:], time.time_ns())))
+    else:
+        main()
@@ -11,6 +11,7 @@ import sys
 import threading
 import time
 from pathlib import Path
+from types import SimpleNamespace

 import pytest

@@ -814,6 +815,85 @@ def test_gpu_samples_preserve_zero_and_omit_unsupported_metrics():
    assert "gpu.memory.total" not in events[2]["attributes"]


+@pytest.mark.parametrize(
+    "cuda,nvidia,expected",
+    [
+        ("MIG-a", "GPU-parent", ["MIG-a"]),
+        ("0", "MIG-a", ["MIG-a"]),
+        ("1,0", "MIG-a,MIG-b", ["MIG-b", "MIG-a"]),
+        (None, "MIG-a", ["MIG-a"]),
+        ("", "MIG-a", []),
+        ("-1", "MIG-a", []),
+        ("2", "MIG-a", []),
+        ("GPU-parent", "MIG-a", []),
+        ("0", "all", []),
+    ],
+)
+def test_mig_assignment_obeys_cuda_visibility(monkeypatch, cuda, nvidia, expected):
+    monkeypatch.setenv("NVIDIA_VISIBLE_DEVICES", nvidia)
+    monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False)
+    if cuda is not None:
+        monkeypatch.setenv("CUDA_VISIBLE_DEVICES", cuda)
+    assert ci_gpu.visible_mig_devices() == expected
+
+
+@pytest.mark.parametrize("used", [0, 3 * 1024**3])
+def test_mig_memory_uses_assigned_handle_without_parent_metrics(monkeypatch, used):
+    handles = []
+    closed = []
+
+    def get_handle(uuid):
+        handles.append(uuid)
+        return "slice-handle"
+
+    def get_memory(handle):
+        assert handle == "slice-handle"
+        return SimpleNamespace(used=used, total=16 * 1024**3)
+
+    nvml = SimpleNamespace(
+        nvmlInit=lambda: None,
+        nvmlShutdown=lambda: closed.append(True),
+        nvmlDeviceGetHandleByUUID=get_handle,
+        nvmlDeviceIsMigDeviceHandle=lambda handle: handle == "slice-handle",
+        nvmlDeviceGetMemoryInfo=get_memory,
+        NVMLError=RuntimeError,
+    )
+    monkeypatch.setattr(ci_gpu, "load_nvml", lambda: nvml)
+    samples = ci_gpu.query_mig_samples(["MIG-a"], 1234)
+    assert handles == ["MIG-a"]
+    assert closed == [True]
+    assert samples[0]["time_ns"] == 1234
+    attributes = samples[0]["attributes"]
+    assert attributes["gpu.uuid"] == "MIG-a"
+    assert attributes["gpu.memory.used"] == used
+    assert attributes["gpu.memory.total"] == 16 * 1024**3
+    assert "gpu.utilization" not in attributes
+
+
+@pytest.mark.parametrize("is_mig", [True, False])
+def test_mig_memory_never_substitutes_parent_or_permission_failure(monkeypatch, is_mig):
+    def unavailable(handle):
+        if not is_mig:
+            pytest.fail("queried parent memory")
+        raise RuntimeError("Insufficient Permissions")
+
+    nvml = SimpleNamespace(
+        nvmlInit=lambda: None,
+        nvmlShutdown=lambda: None,
+        nvmlDeviceGetHandleByUUID=lambda uuid: "handle",
+        nvmlDeviceIsMigDeviceHandle=lambda handle: is_mig,
+        nvmlDeviceGetMemoryInfo=unavailable,
+        NVMLError=RuntimeError,
+    )
+    monkeypatch.setattr(ci_gpu, "load_nvml", lambda: nvml)
+    samples = ci_gpu.query_mig_samples(["MIG-a"], 1234)
+    if is_mig:
+        assert samples[0]["attributes"]["gpu.status"] == "unavailable_mig_memory"
+        assert "gpu.memory.used" not in samples[0]["attributes"]
+    else:
+        assert samples == []
+
+
 def test_gpu_batches_keep_command_identity_and_survive_spool_roundtrip(
    monkeypatch, tmp_path
 ):
@@ -830,7 +910,11 @@ def test_gpu_batches_keep_command_identity_and_survive_spool_roundtrip(
    assert b"ci.gpu.sample" in ci_otel.encode_request([span])


-def test_gpu_query_failures_are_bounded_and_never_synthesize_idle_samples(monkeypatch):
+@pytest.mark.parametrize("mig", [False, True])
+def test_gpu_query_failures_are_bounded_and_never_synthesize_idle_samples(
+    monkeypatch, mig
+):
+    monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "MIG-a" if mig else "0")
    calls = []

    def query(*args, **kwargs):
@@ -918,3 +1002,112 @@ def test_gpu_sampler_stops_with_command_and_preserves_failure_status(
            command["start_ns"] <= event["time_ns"] <= command["end_ns"]
            for event in batches[0]["events"]
        )
+
+
+@pytest.mark.parametrize("is_mig", [True, False])
+def test_cdi_discovery_uses_only_cuda_visible_uuids(monkeypatch, is_mig):
+    identifiers = []
+
+    def get_count(output):
+        output._obj.value = 1
+        return 0
+
+    def get_device(output, index):
+        assert index == 0
+        output._obj.value = 7
+        return 0
+
+    def get_uuid(output, device):
+        assert device.value == 7
+        output._obj[:] = bytes(range(16))
+        return 0
+
+    driver = SimpleNamespace(
+        cuInit=lambda flags: 0,
+        cuDeviceGetCount=get_count,
+        cuDeviceGet=get_device,
+        cuDeviceGetUuid_v2=get_uuid,
+    )
+    monkeypatch.setattr(ci_gpu.ctypes, "CDLL", lambda _: driver)
+    nvml = SimpleNamespace(
+        nvmlInit=lambda: None,
+        nvmlShutdown=lambda: None,
+        nvmlDeviceGetHandleByUUID=lambda value: identifiers.append(value) or "handle",
+        nvmlDeviceIsMigDeviceHandle=lambda handle: is_mig,
+        NVMLError=RuntimeError,
+    )
+    monkeypatch.setattr(ci_gpu, "load_nvml", lambda: nvml)
+    expected = "MIG-00010203-0405-0607-0809-0a0b0c0d0e0f"
+    assert ci_gpu.cuda_visible_mig_devices() == ([expected] if is_mig else [])
+    assert identifiers == [expected]
+
+
+@pytest.mark.parametrize("discovery_fails", [False, True])
+def test_cdi_sampling_switches_to_slice_or_keeps_unknown(monkeypatch, discovery_fails):
+    monkeypatch.setenv("CI_INFRA_TRACE_ID", "01" * 16)
+    monkeypatch.setenv("CI_INFRA_COMMAND_SPAN_ID", "02" * 8)
+    monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False)
+    monkeypatch.setenv("NVIDIA_VISIBLE_DEVICES", "void")
+    monkeypatch.setattr(ci_gpu, "INTERVAL_SECONDS", 0)
+    stop = threading.Event()
+    commands = []
+    records = []
+
+    def query(command, **kwargs):
+        assert kwargs["timeout"] == ci_gpu.QUERY_TIMEOUT_SECONDS
+        commands.append(command)
+        if "--discover-mig" in command:
+            if discovery_fails:
+                raise subprocess.TimeoutExpired(command, 2)
+            return SimpleNamespace(stdout='["MIG-assigned"]')
+        if len(commands) == 4:
+            stop.set()
+        if "--query-mig" in command:
+            assert command[-1] == "MIG-assigned"
+            return SimpleNamespace(
+                stdout=json.dumps(
+                    [
+                        {
+                            "time_ns": 2,
+                            "name": "ci.gpu.sample",
+                            "attributes": {
+                                "gpu.uuid": "MIG-assigned",
+                                "gpu.memory.used": 123,
+                            },
+                        }
+                    ]
+                )
+            )
+        return SimpleNamespace(stdout="0, GPU-parent, H200, 99, 100, 200, Enabled")
+
+    monkeypatch.setattr(ci_gpu.subprocess, "run", query)
+    monkeypatch.setattr(ci_gpu, "record_spans", lambda spans: records.extend(spans))
+    ci_gpu.collect(os.getppid(), stop)
+    assert sum("--discover-mig" in command for command in commands) == 1
+    attributes = [event["attributes"] for span in records for event in span.events]
+    if discovery_fails:
+        assert all("gpu.memory.used" not in event for event in attributes)
+    else:
+        assert any(event.get("gpu.memory.used") == 123 for event in attributes)
+        assert all(event["gpu.uuid"] == "MIG-assigned" for event in attributes)
+
+
+def test_bundled_nvml_load_does_not_import_vllm_or_torch(tmp_path):
+    # Match test images: the collector is outside the installed vLLM package.
+    helper = tmp_path / "ci_gpu_relocated.py"
+    helper.write_text((SCRIPTS_DIR / "ci_gpu.py").read_text())
+    result = subprocess.run(
+        [
+            sys.executable,
+            "-c",
+            "import sys; import ci_gpu_relocated as ci_gpu; "
+            "binding = ci_gpu.load_nvml(); "
+            "assert callable(binding.nvmlDeviceGetMemoryInfo); "
+            "assert 'vllm' not in sys.modules; assert 'torch' not in sys.modules",
+        ],
+        env={**os.environ, "PYTHONPATH": f"{tmp_path}{os.pathsep}{SCRIPTS_DIR}"},
+        capture_output=True,
+        text=True,
+        timeout=5,
+    )
+    assert result.returncode == 0, result.stderr