| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | """Best-effort, container-visible NVIDIA device samples for command timelines.""" |
| 5 | |
| 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 |
| 16 | import signal |
| 17 | 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 | |
| 26 | INTERVAL_SECONDS = 1.0 |
| 27 | BATCH_SECONDS = 30 |
| 28 | MAX_DEVICES = 16 |
| 29 | QUERY_TIMEOUT_SECONDS = 2 |
| 30 | QUERY = "index,uuid,name,utilization.gpu,memory.used,memory.total,mig.mode.current" |
| 31 | |
| 32 | |
| 33 | def _number(value: str, scale: int = 1, maximum: int | None = None) -> int | None: |
| 34 | try: |
| 35 | number = float(value) |
| 36 | if not math.isfinite(number) or number < 0: |
| 37 | return None |
| 38 | if maximum is not None and number > maximum: |
| 39 | return None |
| 40 | return round(number * scale) |
| 41 | except ValueError: |
| 42 | return None |
| 43 | |
| 44 | |
| 45 | def parse_samples(output: str, timestamp_ns: int) -> list[dict]: |
| 46 | events = [] |
| 47 | for row in csv.reader(io.StringIO(output)): |
| 48 | if len(row) != 7 or len(events) >= MAX_DEVICES: |
| 49 | continue |
| 50 | index, uuid, name, util, used, total, mig = (value.strip() for value in row) |
| 51 | if not index.isdigit() or not uuid.startswith("GPU-"): |
| 52 | continue |
| 53 | attributes: dict[str, str | int | bool] = { |
| 54 | "gpu.uuid": uuid, |
| 55 | "gpu.index": int(index), |
| 56 | "gpu.name": name, |
| 57 | } |
| 58 | # Parent-device memory on a MIG GPU is not the job's allocated memory. |
| 59 | if mig.lower() == "enabled": |
| 60 | attributes["gpu.status"] = "unsupported_mig" |
| 61 | else: |
| 62 | for key, value in ( |
| 63 | ("gpu.utilization", _number(util, maximum=100)), |
| 64 | ("gpu.memory.used", _number(used, 1024**2)), |
| 65 | ("gpu.memory.total", _number(total, 1024**2)), |
| 66 | ): |
| 67 | if value is not None: |
| 68 | attributes[key] = value |
| 69 | events.append( |
| 70 | {"time_ns": timestamp_ns, "name": "ci.gpu.sample", "attributes": attributes} |
| 71 | ) |
| 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"], |
| 190 | span_id=secrets.token_hex(8), |
| 191 | parent_span_id=os.environ["CI_INFRA_COMMAND_SPAN_ID"], |
| 192 | name="ci.gpu.samples", |
| 193 | start_ns=events[0]["time_ns"], |
| 194 | end_ns=events[-1]["time_ns"], |
| 195 | attributes={ |
| 196 | "ci.span.kind": "gpu-samples", |
| 197 | "gpu.sample.interval_ms": int(INTERVAL_SECONDS * 1000), |
| 198 | "gpu.scope": "container-visible-devices", |
| 199 | }, |
| 200 | events=events, |
| 201 | ) |
| 202 | |
| 203 | |
| 204 | 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( |
| 224 | query, |
| 225 | capture_output=True, |
| 226 | text=True, |
| 227 | timeout=QUERY_TIMEOUT_SECONDS, |
| 228 | check=True, |
| 229 | ) |
| 230 | if stop.is_set(): |
| 231 | break |
| 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 |
| 266 | except (OSError, subprocess.SubprocessError, json.JSONDecodeError): |
| 267 | failures += 1 |
| 268 | if failures >= 3: |
| 269 | print( |
| 270 | "CI GPU sampling stopped: device query unavailable", |
| 271 | file=sys.stderr, |
| 272 | ) |
| 273 | break |
| 274 | if events and time.monotonic() - batch_started >= BATCH_SECONDS: |
| 275 | # Bounded memory and upload time even if the receiver is down. |
| 276 | # Dropped batches appear as gaps, never as zero utilization. |
| 277 | export_spans([sample_span(events)], timeout_seconds=2) |
| 278 | events = [] |
| 279 | batch_started = time.monotonic() |
| 280 | stop.wait(max(0, INTERVAL_SECONDS - (time.monotonic() - started))) |
| 281 | finally: |
| 282 | # The shell joins this process before the ordinary job-end spool flush. |
| 283 | if events: |
| 284 | record_spans([sample_span(events)]) |
| 285 | |
| 286 | |
| 287 | def main() -> None: |
| 288 | stop = threading.Event() |
| 289 | signal.signal(signal.SIGTERM, lambda *_: stop.set()) |
| 290 | signal.signal(signal.SIGINT, lambda *_: stop.set()) |
| 291 | try: |
| 292 | collect(int(sys.argv[1]), stop) |
| 293 | except Exception: |
| 294 | print("CI GPU sampling skipped", file=sys.stderr) |
| 295 | |
| 296 | |
| 297 | if __name__ == "__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() |