diff --git a/.buildkite/scripts/ci-otel/README.md b/.buildkite/scripts/ci-otel/README.md index 3ad552b5db..318cec5f21 100644 --- a/.buildkite/scripts/ci-otel/README.md +++ b/.buildkite/scripts/ci-otel/README.md @@ -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: diff --git a/.buildkite/scripts/ci-otel/ci_gpu.py b/.buildkite/scripts/ci-otel/ci_gpu.py index 304d738b88..0c77f5170b 100644 --- a/.buildkite/scripts/ci-otel/ci_gpu.py +++ b/.buildkite/scripts/ci-otel/ci_gpu.py @@ -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() diff --git a/.buildkite/scripts/ci-otel/tests/test_ci_otel.py b/.buildkite/scripts/ci-otel/tests/test_ci_otel.py index f636e7ed18..ba54f9f273 100644 --- a/.buildkite/scripts/ci-otel/tests/test_ci_otel.py +++ b/.buildkite/scripts/ci-otel/tests/test_ci_otel.py @@ -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