[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>
Kevin H. Luu committed
Sep 13, 2026 at 21:30 UTC
dc36fcce902a63eab06c1b93a5c4a5ee178a0c56
3 files changed
+371
-12
.buildkite/scripts/ci-otel/README.md
+11
-3
@@ -27,9 +27,17 @@ These are device readings during each test's time interval, not process-level
27
kernel attribution. Other processes and overlapping tests can contribute to
28
the readings. Subsecond tests may have no sample. Missing readings never become
29
zero utilization. MIG parent-device memory and utilization are omitted because
30
-they cannot be attributed to the job's partition. The initial implementation
31
-covers NVIDIA commands in the existing tracing rollout; CPU jobs and AMD
32
-mirrors do not collect GPU samples.
30
+they cannot be attributed to the job's partition. When the environment identifies
31
+an assigned MIG UUID, the collector queries that instance's used/total memory
32
+through vLLM's standalone bundled NVML binding in a bounded subprocess.
33
+The binding is loaded directly to avoid vLLM/Torch import side effects. Numeric
34
+`CUDA_VISIBLE_DEVICES` indices are resolved against explicit MIG UUIDs in
35
+`NVIDIA_VISIBLE_DEVICES`. For CDI assignments without an explicit UUID, a bounded
36
+helper initializes the CUDA driver and reads only CUDA-visible device UUIDs,
37
+without creating CUDA contexts or importing Torch. NVML verifies each MIG handle.
38
+Discovery or permission failures remain unknown; NVML sibling instances are never
39
+enumerated. MIG utilization stays absent.
40
+CPU jobs and AMD mirrors do not collect GPU samples.
41
42
Run the helper tests without loading vLLM or requiring a GPU:
43
.buildkite/scripts/ci-otel/ci_gpu.py
+166
-8
@@ -6,7 +6,10 @@
6
from __future__ import annotations
7
8
import csv
9
+import ctypes
10
+import importlib.util
11
import io
12
+import json
13
import math
14
import os
15
import secrets
@@ -15,6 +18,8 @@ import subprocess
18
import sys
19
import threading
20
import time
21
+import uuid as uuid_module
22
+from pathlib import Path
23
24
from ci_otel import Span, export_spans, record_spans
25
@@ -67,6 +72,118 @@ def parse_samples(output: str, timestamp_ns: int) -> list[dict]:
72
return events
73
74
75
+def visible_mig_devices() -> list[str]:
76
+ """Resolve explicit MIG assignments without enumerating sibling instances."""
77
+ nvidia = os.environ.get("NVIDIA_VISIBLE_DEVICES", "").split(",")
78
+ nvidia = [value.strip() for value in nvidia]
79
+ cuda = os.environ.get("CUDA_VISIBLE_DEVICES")
80
+ devices = nvidia if cuda is None else [value.strip() for value in cuda.split(",")]
81
+ if devices and all(value.isdigit() for value in devices):
82
+ if not all(value.startswith("MIG-") for value in nvidia):
83
+ return []
84
+ try:
85
+ devices = [nvidia[int(value)] for value in devices]
86
+ except IndexError:
87
+ return []
88
+ if not devices or not all(value.startswith("MIG-") for value in devices):
89
+ return []
90
+ return list(dict.fromkeys(devices))[:MAX_DEVICES]
91
+
92
+
93
+def load_nvml():
94
+ """Load the standalone bundled binding without importing vLLM or Torch."""
95
+ # Test images install vLLM as a wheel and keep only tests/helpers in the
96
+ # checkout. Finding the top-level package does not execute its __init__.
97
+ package = importlib.util.find_spec("vllm")
98
+ if package is None or package.origin is None:
99
+ raise ImportError("Installed vLLM package unavailable")
100
+ path = Path(package.origin).parent / "third_party/pynvml.py"
101
+ spec = importlib.util.spec_from_file_location("ci_nvml", path)
102
+ if spec is None or spec.loader is None:
103
+ raise ImportError("Bundled NVML binding unavailable")
104
+ module = importlib.util.module_from_spec(spec)
105
+ sys.modules[spec.name] = module
106
+ spec.loader.exec_module(module)
107
+ return module
108
+
109
+
110
+def cuda_visible_mig_devices() -> list[str]:
111
+ """Discover CDI assignments in a subprocess without creating CUDA contexts."""
112
+ pynvml = load_nvml()
113
+
114
+ driver = ctypes.CDLL("libcuda.so.1")
115
+ count = ctypes.c_int()
116
+ if driver.cuInit(0) or driver.cuDeviceGetCount(ctypes.byref(count)):
117
+ return []
118
+ devices = []
119
+ pynvml.nvmlInit()
120
+ try:
121
+ for index in range(min(count.value, MAX_DEVICES)):
122
+ device = ctypes.c_int()
123
+ identifier = (ctypes.c_ubyte * 16)()
124
+ if driver.cuDeviceGet(ctypes.byref(device), index):
125
+ continue
126
+ # The v2 entry point returns the compute instance UUID under MIG.
127
+ if driver.cuDeviceGetUuid_v2(ctypes.byref(identifier), device):
128
+ continue
129
+ uuid = "MIG-" + str(uuid_module.UUID(bytes=bytes(identifier)))
130
+ try:
131
+ handle = pynvml.nvmlDeviceGetHandleByUUID(uuid)
132
+ if pynvml.nvmlDeviceIsMigDeviceHandle(handle):
133
+ devices.append(uuid)
134
+ except pynvml.NVMLError:
135
+ continue
136
+ finally:
137
+ pynvml.nvmlShutdown()
138
+ return devices
139
+
140
+
141
+def query_mig_samples(devices: list[str], timestamp_ns: int) -> list[dict]:
142
+ # Use the binding already shipped with vLLM; never initialize CUDA/Torch.
143
+ pynvml = load_nvml()
144
+
145
+ events = []
146
+ pynvml.nvmlInit()
147
+ try:
148
+ for index, uuid in enumerate(devices[:MAX_DEVICES]):
149
+ if not uuid.startswith("MIG-"):
150
+ continue
151
+ attributes: dict[str, str | int | bool] = {
152
+ "gpu.uuid": uuid,
153
+ "gpu.index": index,
154
+ "gpu.name": "NVIDIA MIG",
155
+ "gpu.status": "unavailable_mig_memory",
156
+ }
157
+ try:
158
+ handle = pynvml.nvmlDeviceGetHandleByUUID(uuid)
159
+ if not pynvml.nvmlDeviceIsMigDeviceHandle(handle):
160
+ continue
161
+ memory = pynvml.nvmlDeviceGetMemoryInfo(handle)
162
+ used = _number(str(memory.used))
163
+ total = _number(str(memory.total))
164
+ if used is not None and total and used <= total:
165
+ attributes.update(
166
+ {
167
+ "gpu.memory.used": used,
168
+ "gpu.memory.total": total,
169
+ "gpu.status": "mig_memory_only",
170
+ }
171
+ )
172
+ # NVML does not support per-MIG utilization on these devices.
173
+ except pynvml.NVMLError:
174
+ pass
175
+ events.append(
176
+ {
177
+ "time_ns": timestamp_ns,
178
+ "name": "ci.gpu.sample",
179
+ "attributes": attributes,
180
+ }
181
+ )
182
+ finally:
183
+ pynvml.nvmlShutdown()
184
+ return events
185
+
186
+
187
def sample_span(events: list[dict]) -> Span:
188
return Span(
189
trace_id=os.environ["CI_INFRA_TRACE_ID"],
@@ -88,16 +205,23 @@ def collect(parent_pid: int, stop: threading.Event) -> None:
205
events: list[dict] = []
206
batch_started = time.monotonic()
207
failures = 0
208
+ discovery_attempted = False
209
+ mig_devices = visible_mig_devices()
210
+ query = (
211
+ [sys.executable, __file__, "--query-mig", *mig_devices]
212
+ if mig_devices
213
+ else [
214
+ "nvidia-smi",
215
+ f"--query-gpu={QUERY}",
216
+ "--format=csv,noheader,nounits",
217
+ ]
218
+ )
219
try:
220
while not stop.is_set() and os.getppid() == parent_pid:
221
started = time.monotonic()
222
try:
223
result = subprocess.run(
96
- [
97
- "nvidia-smi",
98
- f"--query-gpu={QUERY}",
99
- "--format=csv,noheader,nounits",
100
- ],
224
+ query,
225
capture_output=True,
226
text=True,
227
timeout=QUERY_TIMEOUT_SECONDS,
@@ -105,12 +229,41 @@ def collect(parent_pid: int, stop: threading.Event) -> None:
229
)
230
if stop.is_set():
231
break
108
- samples = parse_samples(result.stdout, time.time_ns())
232
+ samples = (
233
+ json.loads(result.stdout)
234
+ if mig_devices
235
+ else parse_samples(result.stdout, time.time_ns())
236
+ )
237
+ if (
238
+ not mig_devices
239
+ and not discovery_attempted
240
+ and any(
241
+ sample["attributes"].get("gpu.status") == "unsupported_mig"
242
+ for sample in samples
243
+ )
244
+ ):
245
+ discovery_attempted = True
246
+ # CDI can assign a MIG device without exporting its UUID.
247
+ # Only CUDA-visible devices are inspected, never NVML siblings.
248
+ try:
249
+ discovered = subprocess.run(
250
+ [sys.executable, __file__, "--discover-mig"],
251
+ capture_output=True,
252
+ text=True,
253
+ timeout=QUERY_TIMEOUT_SECONDS,
254
+ check=True,
255
+ )
256
+ mig_devices = json.loads(discovered.stdout)
257
+ except (OSError, subprocess.SubprocessError, json.JSONDecodeError):
258
+ mig_devices = []
259
+ if mig_devices:
260
+ query = [sys.executable, __file__, "--query-mig", *mig_devices]
261
+ continue # Do not create an extra parent-device series.
262
if not samples:
263
break
264
events.extend(samples)
265
failures = 0
113
- except (OSError, subprocess.SubprocessError):
266
+ except (OSError, subprocess.SubprocessError, json.JSONDecodeError):
267
failures += 1
268
if failures >= 3:
269
print(
@@ -142,4 +295,9 @@ def main() -> None:
295
296
297
if __name__ == "__main__":
145
- main()
298
+ if len(sys.argv) > 1 and sys.argv[1] == "--discover-mig":
299
+ print(json.dumps(cuda_visible_mig_devices()))
300
+ elif len(sys.argv) > 1 and sys.argv[1] == "--query-mig":
301
+ print(json.dumps(query_mig_samples(sys.argv[2:], time.time_ns())))
302
+ else:
303
+ main()
.buildkite/scripts/ci-otel/tests/test_ci_otel.py
+194
-1
@@ -11,6 +11,7 @@ import sys
11
import threading
12
import time
13
from pathlib import Path
14
+from types import SimpleNamespace
15
16
import pytest
17
@@ -814,6 +815,85 @@ def test_gpu_samples_preserve_zero_and_omit_unsupported_metrics():
815
assert "gpu.memory.total" not in events[2]["attributes"]
816
817
818
+@pytest.mark.parametrize(
819
+ "cuda,nvidia,expected",
820
+ [
821
+ ("MIG-a", "GPU-parent", ["MIG-a"]),
822
+ ("0", "MIG-a", ["MIG-a"]),
823
+ ("1,0", "MIG-a,MIG-b", ["MIG-b", "MIG-a"]),
824
+ (None, "MIG-a", ["MIG-a"]),
825
+ ("", "MIG-a", []),
826
+ ("-1", "MIG-a", []),
827
+ ("2", "MIG-a", []),
828
+ ("GPU-parent", "MIG-a", []),
829
+ ("0", "all", []),
830
+ ],
831
+)
832
+def test_mig_assignment_obeys_cuda_visibility(monkeypatch, cuda, nvidia, expected):
833
+ monkeypatch.setenv("NVIDIA_VISIBLE_DEVICES", nvidia)
834
+ monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False)
835
+ if cuda is not None:
836
+ monkeypatch.setenv("CUDA_VISIBLE_DEVICES", cuda)
837
+ assert ci_gpu.visible_mig_devices() == expected
838
+
839
+
840
+@pytest.mark.parametrize("used", [0, 3 * 1024**3])
841
+def test_mig_memory_uses_assigned_handle_without_parent_metrics(monkeypatch, used):
842
+ handles = []
843
+ closed = []
844
+
845
+ def get_handle(uuid):
846
+ handles.append(uuid)
847
+ return "slice-handle"
848
+
849
+ def get_memory(handle):
850
+ assert handle == "slice-handle"
851
+ return SimpleNamespace(used=used, total=16 * 1024**3)
852
+
853
+ nvml = SimpleNamespace(
854
+ nvmlInit=lambda: None,
855
+ nvmlShutdown=lambda: closed.append(True),
856
+ nvmlDeviceGetHandleByUUID=get_handle,
857
+ nvmlDeviceIsMigDeviceHandle=lambda handle: handle == "slice-handle",
858
+ nvmlDeviceGetMemoryInfo=get_memory,
859
+ NVMLError=RuntimeError,
860
+ )
861
+ monkeypatch.setattr(ci_gpu, "load_nvml", lambda: nvml)
862
+ samples = ci_gpu.query_mig_samples(["MIG-a"], 1234)
863
+ assert handles == ["MIG-a"]
864
+ assert closed == [True]
865
+ assert samples[0]["time_ns"] == 1234
866
+ attributes = samples[0]["attributes"]
867
+ assert attributes["gpu.uuid"] == "MIG-a"
868
+ assert attributes["gpu.memory.used"] == used
869
+ assert attributes["gpu.memory.total"] == 16 * 1024**3
870
+ assert "gpu.utilization" not in attributes
871
+
872
+
873
+@pytest.mark.parametrize("is_mig", [True, False])
874
+def test_mig_memory_never_substitutes_parent_or_permission_failure(monkeypatch, is_mig):
875
+ def unavailable(handle):
876
+ if not is_mig:
877
+ pytest.fail("queried parent memory")
878
+ raise RuntimeError("Insufficient Permissions")
879
+
880
+ nvml = SimpleNamespace(
881
+ nvmlInit=lambda: None,
882
+ nvmlShutdown=lambda: None,
883
+ nvmlDeviceGetHandleByUUID=lambda uuid: "handle",
884
+ nvmlDeviceIsMigDeviceHandle=lambda handle: is_mig,
885
+ nvmlDeviceGetMemoryInfo=unavailable,
886
+ NVMLError=RuntimeError,
887
+ )
888
+ monkeypatch.setattr(ci_gpu, "load_nvml", lambda: nvml)
889
+ samples = ci_gpu.query_mig_samples(["MIG-a"], 1234)
890
+ if is_mig:
891
+ assert samples[0]["attributes"]["gpu.status"] == "unavailable_mig_memory"
892
+ assert "gpu.memory.used" not in samples[0]["attributes"]
893
+ else:
894
+ assert samples == []
895
+
896
+
897
def test_gpu_batches_keep_command_identity_and_survive_spool_roundtrip(
898
monkeypatch, tmp_path
899
):
@@ -830,7 +910,11 @@ def test_gpu_batches_keep_command_identity_and_survive_spool_roundtrip(
910
assert b"ci.gpu.sample" in ci_otel.encode_request([span])
911
912
833
-def test_gpu_query_failures_are_bounded_and_never_synthesize_idle_samples(monkeypatch):
913
+@pytest.mark.parametrize("mig", [False, True])
914
+def test_gpu_query_failures_are_bounded_and_never_synthesize_idle_samples(
915
+ monkeypatch, mig
916
+):
917
+ monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "MIG-a" if mig else "0")
918
calls = []
919
920
def query(*args, **kwargs):
@@ -918,3 +1002,112 @@ def test_gpu_sampler_stops_with_command_and_preserves_failure_status(
1002
command["start_ns"] <= event["time_ns"] <= command["end_ns"]
1003
for event in batches[0]["events"]
1004
)
1005
+
1006
+
1007
+@pytest.mark.parametrize("is_mig", [True, False])
1008
+def test_cdi_discovery_uses_only_cuda_visible_uuids(monkeypatch, is_mig):
1009
+ identifiers = []
1010
+
1011
+ def get_count(output):
1012
+ output._obj.value = 1
1013
+ return 0
1014
+
1015
+ def get_device(output, index):
1016
+ assert index == 0
1017
+ output._obj.value = 7
1018
+ return 0
1019
+
1020
+ def get_uuid(output, device):
1021
+ assert device.value == 7
1022
+ output._obj[:] = bytes(range(16))
1023
+ return 0
1024
+
1025
+ driver = SimpleNamespace(
1026
+ cuInit=lambda flags: 0,
1027
+ cuDeviceGetCount=get_count,
1028
+ cuDeviceGet=get_device,
1029
+ cuDeviceGetUuid_v2=get_uuid,
1030
+ )
1031
+ monkeypatch.setattr(ci_gpu.ctypes, "CDLL", lambda _: driver)
1032
+ nvml = SimpleNamespace(
1033
+ nvmlInit=lambda: None,
1034
+ nvmlShutdown=lambda: None,
1035
+ nvmlDeviceGetHandleByUUID=lambda value: identifiers.append(value) or "handle",
1036
+ nvmlDeviceIsMigDeviceHandle=lambda handle: is_mig,
1037
+ NVMLError=RuntimeError,
1038
+ )
1039
+ monkeypatch.setattr(ci_gpu, "load_nvml", lambda: nvml)
1040
+ expected = "MIG-00010203-0405-0607-0809-0a0b0c0d0e0f"
1041
+ assert ci_gpu.cuda_visible_mig_devices() == ([expected] if is_mig else [])
1042
+ assert identifiers == [expected]
1043
+
1044
+
1045
+@pytest.mark.parametrize("discovery_fails", [False, True])
1046
+def test_cdi_sampling_switches_to_slice_or_keeps_unknown(monkeypatch, discovery_fails):
1047
+ monkeypatch.setenv("CI_INFRA_TRACE_ID", "01" * 16)
1048
+ monkeypatch.setenv("CI_INFRA_COMMAND_SPAN_ID", "02" * 8)
1049
+ monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False)
1050
+ monkeypatch.setenv("NVIDIA_VISIBLE_DEVICES", "void")
1051
+ monkeypatch.setattr(ci_gpu, "INTERVAL_SECONDS", 0)
1052
+ stop = threading.Event()
1053
+ commands = []
1054
+ records = []
1055
+
1056
+ def query(command, **kwargs):
1057
+ assert kwargs["timeout"] == ci_gpu.QUERY_TIMEOUT_SECONDS
1058
+ commands.append(command)
1059
+ if "--discover-mig" in command:
1060
+ if discovery_fails:
1061
+ raise subprocess.TimeoutExpired(command, 2)
1062
+ return SimpleNamespace(stdout='["MIG-assigned"]')
1063
+ if len(commands) == 4:
1064
+ stop.set()
1065
+ if "--query-mig" in command:
1066
+ assert command[-1] == "MIG-assigned"
1067
+ return SimpleNamespace(
1068
+ stdout=json.dumps(
1069
+ [
1070
+ {
1071
+ "time_ns": 2,
1072
+ "name": "ci.gpu.sample",
1073
+ "attributes": {
1074
+ "gpu.uuid": "MIG-assigned",
1075
+ "gpu.memory.used": 123,
1076
+ },
1077
+ }
1078
+ ]
1079
+ )
1080
+ )
1081
+ return SimpleNamespace(stdout="0, GPU-parent, H200, 99, 100, 200, Enabled")
1082
+
1083
+ monkeypatch.setattr(ci_gpu.subprocess, "run", query)
1084
+ monkeypatch.setattr(ci_gpu, "record_spans", lambda spans: records.extend(spans))
1085
+ ci_gpu.collect(os.getppid(), stop)
1086
+ assert sum("--discover-mig" in command for command in commands) == 1
1087
+ attributes = [event["attributes"] for span in records for event in span.events]
1088
+ if discovery_fails:
1089
+ assert all("gpu.memory.used" not in event for event in attributes)
1090
+ else:
1091
+ assert any(event.get("gpu.memory.used") == 123 for event in attributes)
1092
+ assert all(event["gpu.uuid"] == "MIG-assigned" for event in attributes)
1093
+
1094
+
1095
+def test_bundled_nvml_load_does_not_import_vllm_or_torch(tmp_path):
1096
+ # Match test images: the collector is outside the installed vLLM package.
1097
+ helper = tmp_path / "ci_gpu_relocated.py"
1098
+ helper.write_text((SCRIPTS_DIR / "ci_gpu.py").read_text())
1099
+ result = subprocess.run(
1100
+ [
1101
+ sys.executable,
1102
+ "-c",
1103
+ "import sys; import ci_gpu_relocated as ci_gpu; "
1104
+ "binding = ci_gpu.load_nvml(); "
1105
+ "assert callable(binding.nvmlDeviceGetMemoryInfo); "
1106
+ "assert 'vllm' not in sys.modules; assert 'torch' not in sys.modules",
1107
+ ],
1108
+ env={**os.environ, "PYTHONPATH": f"{tmp_path}{os.pathsep}{SCRIPTS_DIR}"},
1109
+ capture_output=True,
1110
+ text=True,
1111
+ timeout=5,
1112
+ )
1113
+ assert result.returncode == 0, result.stderr