@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 33fa95a0cf

[CI] Sample GPU utilization and memory alongside test timelines (#56541)

Signed-off-by: khluu <khluu000@gmail.com> Co-authored-by: Codex <noreply@openai.com>

Kevin H. Luu committed Sep 12, 2026 at 02:41 UTC 33fa95a0cf9ed6bf384a0a51667a01a4f0bb02d5
5 files changed +366 -5
.buildkite/scripts/ci-otel/README.md new
+43
@@ -0,0 +1,43 @@
1 +# CI timing and GPU samples
2 +
3 +The pipeline generator activates these helpers for trusted, traced CI commands.
4 +`ci_otel.sh` records command spans and `ci_pytest.sh` injects the pytest timing
5 +plugin. `ci_gpu.py` runs once per command, outside pytest workers, when
6 +`nvidia-smi` is available. It queries up to 16 container-visible NVIDIA devices
7 +once per second without importing PyTorch or creating a CUDA context.
8 +Set `CI_INFRA_GPU_SAMPLING=0` in the job environment to disable sampling.
9 +
10 +Samples contain device UUID, index, name, utilization percent, and used/total
11 +device memory in bytes. They are `ci.gpu.sample` OTLP events inside
12 +`ci.gpu.samples` spans, parented to the command span and authenticated using the
13 +existing job-scoped Buildkite OIDC identity. Batches upload every 30 seconds;
14 +the final batch is spooled for the shell's normal job-end upload. No host agent
15 +or long-lived dashboard token is needed. The dashboard companion change adds
16 +GPU controls alongside jobs, commands, and pytest rows in the Builds timeline.
17 +
18 +Device queries and periodic uploads each have a two-second timeout. Three
19 +consecutive query failures stop sampling. Failed periodic uploads are dropped;
20 +memory stays bounded, and the dashboard shows the missing intervals as gaps.
21 +Uploads can delay the next poll. The shell stops and joins the sampler before
22 +flushing traces, allows at most three seconds for shutdown, and preserves the
23 +original command exit status. Samples already uploaded survive an interrupted
24 +job; the last unflushed batch may be lost on forced termination.
25 +
26 +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.
33 +
34 +Run the helper tests without loading vLLM or requiring a GPU:
35 +
36 +```bash
37 +PATH="$PWD/.venv/bin:$PATH" .venv/bin/python -m pytest \
38 + .buildkite/scripts/ci-otel/tests/test_ci_otel.py -q
39 +```
40 +
41 +The suite uses a fake `nvidia-smi` to exercise the shell lifecycle, device-query
42 +failure handling, disable switch, event encoding, and bounded uploads. Real
43 +hardware sampling overhead still needs a canary run before broad deployment.
.buildkite/scripts/ci-otel/ci_gpu.py new
+145
@@ -0,0 +1,145 @@
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 io
10 +import math
11 +import os
12 +import secrets
13 +import signal
14 +import subprocess
15 +import sys
16 +import threading
17 +import time
18 +
19 +from ci_otel import Span, export_spans, record_spans
20 +
21 +INTERVAL_SECONDS = 1.0
22 +BATCH_SECONDS = 30
23 +MAX_DEVICES = 16
24 +QUERY_TIMEOUT_SECONDS = 2
25 +QUERY = "index,uuid,name,utilization.gpu,memory.used,memory.total,mig.mode.current"
26 +
27 +
28 +def _number(value: str, scale: int = 1, maximum: int | None = None) -> int | None:
29 + try:
30 + number = float(value)
31 + if not math.isfinite(number) or number < 0:
32 + return None
33 + if maximum is not None and number > maximum:
34 + return None
35 + return round(number * scale)
36 + except ValueError:
37 + return None
38 +
39 +
40 +def parse_samples(output: str, timestamp_ns: int) -> list[dict]:
41 + events = []
42 + for row in csv.reader(io.StringIO(output)):
43 + if len(row) != 7 or len(events) >= MAX_DEVICES:
44 + continue
45 + index, uuid, name, util, used, total, mig = (value.strip() for value in row)
46 + if not index.isdigit() or not uuid.startswith("GPU-"):
47 + continue
48 + attributes: dict[str, str | int | bool] = {
49 + "gpu.uuid": uuid,
50 + "gpu.index": int(index),
51 + "gpu.name": name,
52 + }
53 + # Parent-device memory on a MIG GPU is not the job's allocated memory.
54 + if mig.lower() == "enabled":
55 + attributes["gpu.status"] = "unsupported_mig"
56 + else:
57 + for key, value in (
58 + ("gpu.utilization", _number(util, maximum=100)),
59 + ("gpu.memory.used", _number(used, 1024**2)),
60 + ("gpu.memory.total", _number(total, 1024**2)),
61 + ):
62 + if value is not None:
63 + attributes[key] = value
64 + events.append(
65 + {"time_ns": timestamp_ns, "name": "ci.gpu.sample", "attributes": attributes}
66 + )
67 + return events
68 +
69 +
70 +def sample_span(events: list[dict]) -> Span:
71 + return Span(
72 + trace_id=os.environ["CI_INFRA_TRACE_ID"],
73 + span_id=secrets.token_hex(8),
74 + parent_span_id=os.environ["CI_INFRA_COMMAND_SPAN_ID"],
75 + name="ci.gpu.samples",
76 + start_ns=events[0]["time_ns"],
77 + end_ns=events[-1]["time_ns"],
78 + attributes={
79 + "ci.span.kind": "gpu-samples",
80 + "gpu.sample.interval_ms": int(INTERVAL_SECONDS * 1000),
81 + "gpu.scope": "container-visible-devices",
82 + },
83 + events=events,
84 + )
85 +
86 +
87 +def collect(parent_pid: int, stop: threading.Event) -> None:
88 + events: list[dict] = []
89 + batch_started = time.monotonic()
90 + failures = 0
91 + try:
92 + while not stop.is_set() and os.getppid() == parent_pid:
93 + started = time.monotonic()
94 + try:
95 + result = subprocess.run(
96 + [
97 + "nvidia-smi",
98 + f"--query-gpu={QUERY}",
99 + "--format=csv,noheader,nounits",
100 + ],
101 + capture_output=True,
102 + text=True,
103 + timeout=QUERY_TIMEOUT_SECONDS,
104 + check=True,
105 + )
106 + if stop.is_set():
107 + break
108 + samples = parse_samples(result.stdout, time.time_ns())
109 + if not samples:
110 + break
111 + events.extend(samples)
112 + failures = 0
113 + except (OSError, subprocess.SubprocessError):
114 + failures += 1
115 + if failures >= 3:
116 + print(
117 + "CI GPU sampling stopped: device query unavailable",
118 + file=sys.stderr,
119 + )
120 + break
121 + if events and time.monotonic() - batch_started >= BATCH_SECONDS:
122 + # Bounded memory and upload time even if the receiver is down.
123 + # Dropped batches appear as gaps, never as zero utilization.
124 + export_spans([sample_span(events)], timeout_seconds=2)
125 + events = []
126 + batch_started = time.monotonic()
127 + stop.wait(max(0, INTERVAL_SECONDS - (time.monotonic() - started)))
128 + finally:
129 + # The shell joins this process before the ordinary job-end spool flush.
130 + if events:
131 + record_spans([sample_span(events)])
132 +
133 +
134 +def main() -> None:
135 + stop = threading.Event()
136 + signal.signal(signal.SIGTERM, lambda *_: stop.set())
137 + signal.signal(signal.SIGINT, lambda *_: stop.set())
138 + try:
139 + collect(int(sys.argv[1]), stop)
140 + except Exception:
141 + print("CI GPU sampling skipped", file=sys.stderr)
142 +
143 +
144 +if __name__ == "__main__":
145 + main()
.buildkite/scripts/ci-otel/ci_otel.py
+25 -3
@@ -14,12 +14,13 @@ import sys
14 import time
15 import urllib.request
16 from collections.abc import Iterable
17 -from dataclasses import asdict, dataclass
17 +from dataclasses import asdict, dataclass, field
18 from pathlib import Path
19
20 ENDPOINT = os.getenv("CI_INFRA_OTEL_ENDPOINT", "https://ci.vllm.ai/api/otel/v1/traces")
21 AUDIENCE = os.getenv("CI_INFRA_OTEL_AUDIENCE", "https://ci.vllm.ai/api/otel")
22 MAX_BATCH_SIZE = 2_000
23 +MAX_BATCH_BYTES = 1024 * 1024
24 TEST_SPOOL_BATCH_SIZE = 100
25
26
@@ -56,6 +57,7 @@ class Span:
57 end_ns: int
58 attributes: dict[str, str | int | bool]
59 status_code: int = 1
60 + events: list[dict] = field(default_factory=list)
61
62
63 def _varint(value: int) -> bytes:
@@ -137,6 +139,12 @@ def _encode_span(span: Span) -> bytes:
139 )
140 for name, value in attributes.items():
141 encoded += _bytes_field(9, _attribute(name, value))
142 + for event in span.events:
143 + encoded_event = _fixed64_field(1, int(event["time_ns"]))
144 + encoded_event += _string_field(2, event["name"])
145 + for name, value in event["attributes"].items():
146 + encoded_event += _bytes_field(3, _attribute(name, value))
147 + encoded += _bytes_field(11, encoded_event)
148 return encoded + _bytes_field(15, _varint_field(3, span.status_code))
149
150
@@ -482,10 +490,24 @@ def export_spans(spans: list[Span], timeout_seconds: float = 3.0) -> bool:
490 try:
491 deadline = time.monotonic() + max(timeout_seconds, 0.1)
492 token = _oidc_token(deadline)
485 - for offset in range(0, len(spans), MAX_BATCH_SIZE):
493 + batches: list[list[Span]] = []
494 + batch: list[Span] = []
495 + size = 0
496 + for span in spans:
497 + span_size = len(_encode_span(span))
498 + if batch and (
499 + size + span_size > MAX_BATCH_BYTES or len(batch) >= MAX_BATCH_SIZE
500 + ):
501 + batches.append(batch)
502 + batch, size = [], 0
503 + batch.append(span)
504 + size += span_size
505 + if batch:
506 + batches.append(batch)
507 + for batch in batches:
508 request = urllib.request.Request(
509 ENDPOINT,
488 - data=encode_request(spans[offset : offset + MAX_BATCH_SIZE]),
510 + data=encode_request(batch),
511 method="POST",
512 headers={
513 "Authorization": f"Bearer {token}",
.buildkite/scripts/ci-otel/ci_otel.sh
+19
@@ -82,6 +82,13 @@ ci_otel_start() {
82 CI_INFRA_COMMAND_SPAN_ID="${_CI_INFRA_OTEL_COMMAND_SPAN_ID}"
83 export CI_INFRA_TRACE_ID CI_INFRA_COMMAND_SPAN_ID
84 _CI_INFRA_OTEL_ACTIVE=1
85 + _CI_INFRA_GPU_PID=""
86 + if [ "${CI_INFRA_GPU_SAMPLING:-1}" != "0" ] &&
87 + [ -f "${CI_INFRA_OTEL_DIR}/ci_gpu.py" ] &&
88 + command -v nvidia-smi >/dev/null 2>&1; then
89 + "${_CI_INFRA_OTEL_PYTHON}" "${CI_INFRA_OTEL_DIR}/ci_gpu.py" "$$" </dev/null &
90 + _CI_INFRA_GPU_PID=$!
91 + fi
92 }
93
94 ci_otel_finish() {
@@ -93,6 +100,18 @@ ci_otel_finish() {
100 "${_CI_INFRA_OTEL_PARENT_SPAN_ID}" "${_CI_INFRA_OTEL_START_NS}" \
101 "${_CI_INFRA_OTEL_COMMAND_INDEX}" "${_CI_INFRA_OTEL_COMMAND_STATUS}" \
102 "${_CI_INFRA_OTEL_COMMAND_LABEL}" || :
103 + if [ -n "${_CI_INFRA_GPU_PID:-}" ]; then
104 + kill -TERM "${_CI_INFRA_GPU_PID}" 2>/dev/null || :
105 + _CI_INFRA_GPU_WAIT=0
106 + while kill -0 "${_CI_INFRA_GPU_PID}" 2>/dev/null &&
107 + [ "${_CI_INFRA_GPU_WAIT}" -lt 30 ]; do
108 + sleep 0.1
109 + _CI_INFRA_GPU_WAIT=$((_CI_INFRA_GPU_WAIT + 1))
110 + done
111 + kill -KILL "${_CI_INFRA_GPU_PID}" 2>/dev/null || :
112 + wait "${_CI_INFRA_GPU_PID}" 2>/dev/null || :
113 + _CI_INFRA_GPU_PID=""
114 + fi
115 CI_INFRA_TRACE_ID=""
116 CI_INFRA_COMMAND_SPAN_ID=""
117 export CI_INFRA_TRACE_ID CI_INFRA_COMMAND_SPAN_ID
.buildkite/scripts/ci-otel/tests/test_ci_otel.py
+134 -2
@@ -8,6 +8,7 @@ import os
8 import shlex
9 import subprocess
10 import sys
11 +import threading
12 import time
13 from pathlib import Path
14
@@ -16,6 +17,7 @@ import pytest
17 SCRIPTS_DIR = Path(__file__).resolve().parents[1]
18 sys.path.insert(0, str(SCRIPTS_DIR))
19
20 +import ci_gpu # noqa: E402
21 import ci_otel # noqa: E402
22 from ci_otel import Span # noqa: E402
23
@@ -251,9 +253,15 @@ def test_export_failure_is_soft_and_bounded(monkeypatch):
253 assert time.monotonic() - started < 0.5
254
255
254 -def test_export_batches_with_one_oidc_token(monkeypatch):
256 +@pytest.mark.parametrize("byte_limited", [False, True])
257 +def test_export_batches_with_one_oidc_token(monkeypatch, byte_limited):
258 monkeypatch.setenv("BUILDKITE", "true")
256 - monkeypatch.setattr(ci_otel, "MAX_BATCH_SIZE", 2)
259 + if byte_limited:
260 + monkeypatch.setattr(
261 + ci_otel, "MAX_BATCH_BYTES", len(ci_otel._encode_span(_span())) * 2
262 + )
263 + else:
264 + monkeypatch.setattr(ci_otel, "MAX_BATCH_SIZE", 2)
265 token_calls = []
266 requests = []
267
@@ -786,3 +794,127 @@ def test_pytest_plugin_failure_cannot_fail_pytest(tmp_path):
794
795 assert result.returncode == 0, result.stderr
796 assert "1 passed" in result.stdout
797 +
798 +
799 +def test_gpu_samples_preserve_zero_and_omit_unsupported_metrics():
800 + events = ci_gpu.parse_samples(
801 + "0, GPU-a, H200, 0, 1024, 2048, Disabled\n"
802 + "1, GPU-b, H200, [N/A], [N/A], 2048, Disabled\n"
803 + "2, GPU-c, H200, 50, 1024, 2048, Enabled\n"
804 + "3, GPU-d, H200, NaN, -1, inf, Disabled\n",
805 + 1234,
806 + )
807 + assert len(events) == 4
808 + assert events[0]["attributes"]["gpu.utilization"] == 0
809 + assert events[0]["attributes"]["gpu.memory.used"] == 1024**3
810 + for event in events[1:]:
811 + assert "gpu.utilization" not in event["attributes"]
812 + assert "gpu.memory.used" not in event["attributes"]
813 + assert events[2]["attributes"]["gpu.status"] == "unsupported_mig"
814 + assert "gpu.memory.total" not in events[2]["attributes"]
815 +
816 +
817 +def test_gpu_batches_keep_command_identity_and_survive_spool_roundtrip(
818 + monkeypatch, tmp_path
819 +):
820 + monkeypatch.setenv("CI_INFRA_TRACE_ID", "01" * 16)
821 + monkeypatch.setenv("CI_INFRA_COMMAND_SPAN_ID", "02" * 8)
822 + monkeypatch.setenv("CI_INFRA_OTEL_SPOOL_DIR", str(tmp_path))
823 + span = ci_gpu.sample_span(
824 + ci_gpu.parse_samples("0, GPU-a, H200, 90, 1, 2, Disabled", 1234)
825 + )
826 + assert span.parent_span_id == "02" * 8
827 + assert span.trace_id == "01" * 16
828 + assert ci_otel.record_spans([span])
829 + assert ci_otel.load_spans() == [span]
830 + assert b"ci.gpu.sample" in ci_otel.encode_request([span])
831 +
832 +
833 +def test_gpu_query_failures_are_bounded_and_never_synthesize_idle_samples(monkeypatch):
834 + calls = []
835 +
836 + def query(*args, **kwargs):
837 + calls.append(kwargs["timeout"])
838 + raise subprocess.TimeoutExpired("nvidia-smi", kwargs["timeout"])
839 +
840 + monkeypatch.setattr(ci_gpu.subprocess, "run", query)
841 + monkeypatch.setattr(ci_gpu, "INTERVAL_SECONDS", 0)
842 + monkeypatch.setattr(
843 + ci_gpu, "record_spans", lambda _: pytest.fail("fabricated sample")
844 + )
845 + ci_gpu.collect(os.getppid(), threading.Event())
846 + assert calls == [2, 2, 2]
847 +
848 +
849 +def test_gpu_periodic_batches_are_bounded_even_when_upload_fails(monkeypatch):
850 + monkeypatch.setenv("CI_INFRA_TRACE_ID", "01" * 16)
851 + monkeypatch.setenv("CI_INFRA_COMMAND_SPAN_ID", "02" * 8)
852 + monkeypatch.setattr(ci_gpu, "BATCH_SECONDS", 0)
853 + monkeypatch.setattr(ci_gpu, "INTERVAL_SECONDS", 0)
854 + monkeypatch.setattr(
855 + ci_gpu.subprocess,
856 + "run",
857 + lambda *a, **k: subprocess.CompletedProcess(
858 + a, 0, "0, GPU-a, H200, 50, 1, 2, Disabled"
859 + ),
860 + )
861 + stop = threading.Event()
862 + batches = []
863 +
864 + def upload(spans, timeout_seconds):
865 + batches.append(spans)
866 + assert timeout_seconds == 2
867 + if len(batches) == 3:
868 + stop.set()
869 + return False
870 +
871 + monkeypatch.setattr(ci_gpu, "export_spans", upload)
872 + ci_gpu.collect(os.getppid(), stop)
873 + assert [len(batch[0].events) for batch in batches] == [1, 1, 1]
874 +
875 +
876 +@pytest.mark.parametrize("sampling", ["1", "0"])
877 +def test_gpu_sampler_stops_with_command_and_preserves_failure_status(
878 + tmp_path, sampling
879 +):
880 + bin_dir = tmp_path / "bin"
881 + bin_dir.mkdir()
882 + smi = bin_dir / "nvidia-smi"
883 + smi.write_text('#!/bin/sh\necho "0, GPU-a, H200, 70, 1024, 2048, Disabled"\n')
884 + smi.chmod(0o755)
885 + # Keep the spool for inspection, and avoid any network upload.
886 + shell = (
887 + f'. "{SCRIPTS_DIR / "ci_otel.sh"}"\n'
888 + "ci_otel_run 1 workload sh -c 'sleep 1.3; exit 7'\n"
889 + "exit $?"
890 + )
891 + result = subprocess.run(
892 + ["/bin/sh", "-c", shell],
893 + capture_output=True,
894 + text=True,
895 + timeout=8,
896 + env={
897 + **os.environ,
898 + "PATH": f"{bin_dir}:{os.environ['PATH']}",
899 + "BUILDKITE": "false",
900 + "CI_INFRA_GPU_SAMPLING": sampling,
901 + "CI_INFRA_OTEL_DIR": str(SCRIPTS_DIR),
902 + "CI_INFRA_OTEL_RUNTIME_DIR": str(tmp_path / "runtime"),
903 + "CI_INFRA_OTEL_SPOOL_DIR": str(tmp_path / "spans"),
904 + },
905 + )
906 + assert result.returncode == 7, result.stderr
907 + records = [
908 + json.loads(line)
909 + for path in (tmp_path / "spans").glob("spans-*.jsonl")
910 + for line in path.read_text().splitlines()
911 + ]
912 + command = next(record for record in records if record["name"] == "ci.command")
913 + batches = [record for record in records if record["name"] == "ci.gpu.samples"]
914 + assert bool(batches) == (sampling == "1")
915 + if batches:
916 + assert batches[0]["parent_span_id"] == command["span_id"]
917 + assert all(
918 + command["start_ns"] <= event["time_ns"] <= command["end_ns"]
919 + for event in batches[0]["events"]
920 + )