| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from __future__ import annotations |
| 5 | |
| 6 | import json |
| 7 | import os |
| 8 | import shlex |
| 9 | import shutil |
| 10 | import subprocess |
| 11 | import sys |
| 12 | import threading |
| 13 | import time |
| 14 | from pathlib import Path |
| 15 | from types import SimpleNamespace |
| 16 | |
| 17 | import pytest |
| 18 | |
| 19 | SCRIPTS_DIR = Path(__file__).resolve().parents[1] |
| 20 | sys.path.insert(0, str(SCRIPTS_DIR)) |
| 21 | |
| 22 | import ci_gpu # noqa: E402 |
| 23 | import ci_otel # noqa: E402 |
| 24 | from ci_otel import Span # noqa: E402 |
| 25 | |
| 26 | |
| 27 | def _span(**attributes) -> Span: |
| 28 | return Span( |
| 29 | trace_id="01" * 16, |
| 30 | span_id="02" * 8, |
| 31 | parent_span_id=None, |
| 32 | name="ci.command", |
| 33 | start_ns=100, |
| 34 | end_ns=200, |
| 35 | attributes=attributes, |
| 36 | ) |
| 37 | |
| 38 | |
| 39 | def _quoted(value: str) -> str: |
| 40 | return shlex.quote(value) |
| 41 | |
| 42 | |
| 43 | def test_context_continues_traceparent(monkeypatch): |
| 44 | monkeypatch.setenv( |
| 45 | "TRACEPARENT", |
| 46 | "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01", |
| 47 | ) |
| 48 | |
| 49 | trace_id, span_id, parent_span_id = ci_otel.new_context() |
| 50 | |
| 51 | assert trace_id == "4bf92f3577b34da6a3ce929d0e0e4736" |
| 52 | assert len(span_id) == 16 |
| 53 | assert parent_span_id == "00f067aa0ba902b7" |
| 54 | |
| 55 | |
| 56 | def test_otlp_payload_contains_dashboard_identity(monkeypatch): |
| 57 | monkeypatch.setenv("BUILDKITE_ORGANIZATION_SLUG", "vllm") |
| 58 | monkeypatch.setenv("BUILDKITE_PIPELINE_SLUG", "ci") |
| 59 | monkeypatch.setenv("BUILDKITE_BUILD_ID", "build-id") |
| 60 | monkeypatch.setenv("BUILDKITE_BUILD_NUMBER", "42") |
| 61 | monkeypatch.setenv("BUILDKITE_JOB_ID", "job-id") |
| 62 | |
| 63 | payload = ci_otel.encode_request([_span(**{"ci.span.kind": "command"})]) |
| 64 | |
| 65 | for value in (b"ci.command", b"ci.span.kind", b"buildkite.job.id", b"job-id"): |
| 66 | assert value in payload |
| 67 | |
| 68 | |
| 69 | def test_otlp_payload_accepts_negative_integer_attributes(): |
| 70 | payload = ci_otel.encode_request([_span(**{"process.exit.code": -1})]) |
| 71 | |
| 72 | assert b"process.exit.code" in payload |
| 73 | |
| 74 | |
| 75 | def test_spool_round_trip_is_local(monkeypatch, tmp_path): |
| 76 | monkeypatch.setenv("CI_INFRA_OTEL_SPOOL_DIR", str(tmp_path)) |
| 77 | monkeypatch.setattr( |
| 78 | ci_otel, |
| 79 | "_oidc_token", |
| 80 | lambda deadline: (_ for _ in ()).throw(AssertionError("unexpected upload")), |
| 81 | ) |
| 82 | span = _span(**{"ci.command.index": 1}) |
| 83 | |
| 84 | assert ci_otel.record_spans([span]) is True |
| 85 | assert ci_otel.load_spans() == [span] |
| 86 | |
| 87 | |
| 88 | def test_buildx_ref_supports_bake_metadata(tmp_path): |
| 89 | metadata = tmp_path / "metadata.json" |
| 90 | metadata.write_text( |
| 91 | json.dumps({"test-ci": {"buildx.build.ref": "builder/node/record-id"}}) |
| 92 | ) |
| 93 | |
| 94 | assert ci_otel.buildx_ref(str(metadata), "test-ci") == "builder/node/record-id" |
| 95 | |
| 96 | |
| 97 | def test_buildkit_trace_groups_stages_and_includes_descendant_time(monkeypatch): |
| 98 | monkeypatch.setenv("BUILDKITE_BUILD_ID", "build-id") |
| 99 | payload = { |
| 100 | "data": [ |
| 101 | { |
| 102 | "spans": [ |
| 103 | { |
| 104 | "spanID": "01" * 8, |
| 105 | "operationName": "[base 1/2] FROM alpine:3.21", |
| 106 | "startTime": 100, |
| 107 | "duration": 10, |
| 108 | "tags": [{"key": "vertex", "value": "sha256:from"}], |
| 109 | }, |
| 110 | { |
| 111 | "spanID": "02" * 8, |
| 112 | "operationName": "registry pull", |
| 113 | "references": [{"refType": "CHILD_OF", "spanID": "01" * 8}], |
| 114 | "startTime": 105, |
| 115 | "duration": 95, |
| 116 | }, |
| 117 | { |
| 118 | "spanID": "03" * 8, |
| 119 | "operationName": "[base 2/2] RUN uv pip install vllm", |
| 120 | "startTime": 210, |
| 121 | "duration": 40, |
| 122 | "tags": [{"key": "vertex", "value": "sha256:run"}], |
| 123 | }, |
| 124 | { |
| 125 | "spanID": "04" * 8, |
| 126 | "operationName": "[test 1/1] COPY --from=base /opt /opt", |
| 127 | "startTime": 260, |
| 128 | "duration": 20, |
| 129 | "tags": [{"key": "vertex", "value": "sha256:copy"}], |
| 130 | }, |
| 131 | { |
| 132 | "spanID": "05" * 8, |
| 133 | "operationName": "[internal] load build context", |
| 134 | "startTime": 90, |
| 135 | "duration": 30, |
| 136 | "tags": [{"key": "vertex", "value": "sha256:context"}], |
| 137 | }, |
| 138 | { |
| 139 | "spanID": "06" * 8, |
| 140 | "operationName": ( |
| 141 | "cache request: [base 2/2] RUN uv pip install vllm" |
| 142 | ), |
| 143 | "startTime": 205, |
| 144 | "duration": 2, |
| 145 | "tags": [{"key": "vertex", "value": "sha256:run"}], |
| 146 | }, |
| 147 | ] |
| 148 | } |
| 149 | ] |
| 150 | } |
| 151 | |
| 152 | spans = ci_otel.buildkit_spans(payload) |
| 153 | stages = { |
| 154 | span.attributes.get("docker.stage"): span |
| 155 | for span in spans |
| 156 | if span.name == "docker.stage" |
| 157 | } |
| 158 | instructions = [span for span in spans if span.name == "docker.instruction"] |
| 159 | |
| 160 | assert set(stages) == {"base", "test"} |
| 161 | assert stages["base"].start_ns == 100_000 |
| 162 | assert stages["base"].end_ns == 250_000 |
| 163 | assert len(instructions) == 3 |
| 164 | assert len([span for span in spans if span.name == "docker.internal"]) == 1 |
| 165 | from_span = next( |
| 166 | span for span in instructions if span.attributes["docker.instruction"] == "FROM" |
| 167 | ) |
| 168 | assert from_span.end_ns == 200_000 |
| 169 | assert from_span.parent_span_id == stages["base"].span_id |
| 170 | assert all(span.trace_id == stages["base"].trace_id for span in spans) |
| 171 | |
| 172 | |
| 173 | def test_record_buildkit_trace_spools_normalized_spans(monkeypatch, tmp_path): |
| 174 | trace_file = tmp_path / "trace.json" |
| 175 | trace_file.write_text( |
| 176 | json.dumps( |
| 177 | { |
| 178 | "data": [ |
| 179 | { |
| 180 | "spans": [ |
| 181 | { |
| 182 | "spanID": "01" * 8, |
| 183 | "operationName": "[base 1/1] RUN true", |
| 184 | "startTime": 100, |
| 185 | "duration": 20, |
| 186 | } |
| 187 | ] |
| 188 | } |
| 189 | ] |
| 190 | } |
| 191 | ) |
| 192 | ) |
| 193 | monkeypatch.setenv("CI_INFRA_OTEL_SPOOL_DIR", str(tmp_path / "spool")) |
| 194 | |
| 195 | assert ci_otel.record_buildkit_trace(str(trace_file)) is True |
| 196 | assert {span.attributes["ci.span.kind"] for span in ci_otel.load_spans()} == { |
| 197 | "docker-stage", |
| 198 | "docker-instruction", |
| 199 | } |
| 200 | |
| 201 | |
| 202 | def test_export_uses_buildkite_oidc_and_otlp(monkeypatch): |
| 203 | for name, value in { |
| 204 | "BUILDKITE": "true", |
| 205 | "BUILDKITE_AGENT_ACCESS_TOKEN": "agent-token", |
| 206 | "BUILDKITE_AGENT_ENDPOINT": "https://agent.example/v3", |
| 207 | "BUILDKITE_JOB_ID": "job-id", |
| 208 | }.items(): |
| 209 | monkeypatch.setenv(name, value) |
| 210 | requests = [] |
| 211 | |
| 212 | class Response: |
| 213 | status = 200 |
| 214 | |
| 215 | def __init__(self, body=b""): |
| 216 | self.body = body |
| 217 | |
| 218 | def read(self): |
| 219 | return self.body |
| 220 | |
| 221 | def __enter__(self): |
| 222 | return self |
| 223 | |
| 224 | def __exit__(self, *args): |
| 225 | return False |
| 226 | |
| 227 | def open_request(request, timeout): |
| 228 | requests.append((request, timeout)) |
| 229 | if request.full_url.endswith("/oidc/tokens"): |
| 230 | return Response(b'{"token":"oidc-token"}') |
| 231 | return Response() |
| 232 | |
| 233 | monkeypatch.setattr(ci_otel.urllib.request, "urlopen", open_request) |
| 234 | |
| 235 | assert ci_otel.export_spans([_span()], timeout_seconds=0.5) is True |
| 236 | oidc_request, upload_request = (request for request, _ in requests) |
| 237 | assert oidc_request.full_url.endswith("/jobs/job-id/oidc/tokens") |
| 238 | assert oidc_request.get_header("Authorization") == "Token agent-token" |
| 239 | assert json.loads(oidc_request.data)["audience"] == ci_otel.AUDIENCE |
| 240 | assert upload_request.get_header("Authorization") == "Bearer oidc-token" |
| 241 | assert upload_request.get_header("Content-type") == "application/x-protobuf" |
| 242 | assert all(0 < timeout <= 0.5 for _, timeout in requests) |
| 243 | |
| 244 | |
| 245 | def test_export_failure_is_soft_and_bounded(monkeypatch): |
| 246 | monkeypatch.setenv("BUILDKITE", "true") |
| 247 | monkeypatch.setattr( |
| 248 | ci_otel, |
| 249 | "_oidc_token", |
| 250 | lambda deadline: (_ for _ in ()).throw(TimeoutError("unavailable")), |
| 251 | ) |
| 252 | |
| 253 | started = time.monotonic() |
| 254 | assert ci_otel.export_spans([_span()], timeout_seconds=0.1) is False |
| 255 | assert time.monotonic() - started < 0.5 |
| 256 | |
| 257 | |
| 258 | @pytest.mark.parametrize("byte_limited", [False, True]) |
| 259 | def test_export_batches_with_one_oidc_token(monkeypatch, byte_limited): |
| 260 | monkeypatch.setenv("BUILDKITE", "true") |
| 261 | if byte_limited: |
| 262 | monkeypatch.setattr( |
| 263 | ci_otel, "MAX_BATCH_BYTES", len(ci_otel._encode_span(_span())) * 2 |
| 264 | ) |
| 265 | else: |
| 266 | monkeypatch.setattr(ci_otel, "MAX_BATCH_SIZE", 2) |
| 267 | token_calls = [] |
| 268 | requests = [] |
| 269 | |
| 270 | class Response: |
| 271 | status = 200 |
| 272 | |
| 273 | def __enter__(self): |
| 274 | return self |
| 275 | |
| 276 | def __exit__(self, *args): |
| 277 | return False |
| 278 | |
| 279 | def token(deadline): |
| 280 | token_calls.append(deadline) |
| 281 | return "token" |
| 282 | |
| 283 | def open_request(request, timeout): |
| 284 | requests.append(request) |
| 285 | return Response() |
| 286 | |
| 287 | monkeypatch.setattr(ci_otel, "_oidc_token", token) |
| 288 | monkeypatch.setattr(ci_otel.urllib.request, "urlopen", open_request) |
| 289 | |
| 290 | assert ci_otel.export_spans([_span() for _ in range(5)]) is True |
| 291 | assert len(token_calls) == 1 |
| 292 | assert len(requests) == 3 |
| 293 | |
| 294 | |
| 295 | def test_successful_flush_removes_spool_and_does_not_duplicate(monkeypatch, tmp_path): |
| 296 | monkeypatch.setenv("CI_INFRA_OTEL_SPOOL_DIR", str(tmp_path)) |
| 297 | uploaded = [] |
| 298 | assert ci_otel.record_spans([_span()]) is True |
| 299 | monkeypatch.setattr( |
| 300 | ci_otel, "export_spans", lambda spans, timeout: uploaded.extend(spans) or True |
| 301 | ) |
| 302 | |
| 303 | assert ci_otel.flush_spans() is True |
| 304 | assert len(uploaded) == 1 |
| 305 | assert list(tmp_path.glob("spans-*.jsonl")) == [] |
| 306 | assert ci_otel.flush_spans() is True |
| 307 | assert len(uploaded) == 1 |
| 308 | |
| 309 | |
| 310 | def test_failed_flush_preserves_spool(monkeypatch, tmp_path): |
| 311 | monkeypatch.setenv("CI_INFRA_OTEL_SPOOL_DIR", str(tmp_path)) |
| 312 | assert ci_otel.record_spans([_span()]) is True |
| 313 | monkeypatch.setattr(ci_otel, "export_spans", lambda spans, timeout: False) |
| 314 | |
| 315 | assert ci_otel.flush_spans() is False |
| 316 | assert len(list(tmp_path.glob("spans-*.jsonl"))) == 1 |
| 317 | |
| 318 | |
| 319 | def test_cli_rejects_bad_arguments_without_traceback(capsys): |
| 320 | assert ci_otel.main(["bogus"]) == 2 |
| 321 | assert ci_otel.main(["record-command", "too", "short"]) == 2 |
| 322 | assert ( |
| 323 | ci_otel.main(["record-command", "t", "s", "-", "not-an-int", "1", "0", "label"]) |
| 324 | == 1 |
| 325 | ) |
| 326 | |
| 327 | stderr = capsys.readouterr().err |
| 328 | assert "usage:" in stderr |
| 329 | assert "CI timing helper failed:" in stderr |
| 330 | assert "Traceback" not in stderr |
| 331 | |
| 332 | |
| 333 | def test_pytest_spans_are_spooled_incrementally(monkeypatch): |
| 334 | recorded = [] |
| 335 | monkeypatch.setenv("CI_INFRA_TRACE_ID", "01" * 16) |
| 336 | monkeypatch.setenv("CI_INFRA_COMMAND_SPAN_ID", "02" * 8) |
| 337 | monkeypatch.setattr(ci_otel, "TEST_SPOOL_BATCH_SIZE", 1) |
| 338 | monkeypatch.setattr( |
| 339 | ci_otel, "record_spans", lambda spans: recorded.extend(spans) or True |
| 340 | ) |
| 341 | ci_otel._test_runs.clear() |
| 342 | ci_otel._test_spans.clear() |
| 343 | ci_otel._test_runs["test_sample.py::test_one"] = (time.time_ns(), "passed") |
| 344 | |
| 345 | ci_otel.pytest_runtest_logfinish("test_sample.py::test_one", ("", 0, "")) |
| 346 | |
| 347 | assert len(recorded) == 1 |
| 348 | assert ci_otel._test_spans == [] |
| 349 | |
| 350 | |
| 351 | def test_shell_wrapper_records_commands_without_changing_shell_state(tmp_path): |
| 352 | first = f"ci_otel_start 1 {_quoted('export VALUE=ready')}" |
| 353 | second = f"ci_otel_start 2 {_quoted('check VALUE')}" |
| 354 | shell = ( |
| 355 | f'. "{SCRIPTS_DIR / "ci_otel.sh"}"; {first}; ' |
| 356 | f"export VALUE=ready; ci_otel_finish 0; {second}; " |
| 357 | 'test "$VALUE" = ready; ci_otel_finish 0' |
| 358 | ) |
| 359 | |
| 360 | result = subprocess.run( |
| 361 | ["/bin/sh", "-c", shell], |
| 362 | check=False, |
| 363 | capture_output=True, |
| 364 | text=True, |
| 365 | env={ |
| 366 | **os.environ, |
| 367 | "CI_INFRA_OTEL_DIR": str(SCRIPTS_DIR), |
| 368 | "CI_INFRA_OTEL_SPOOL_DIR": str(tmp_path), |
| 369 | }, |
| 370 | ) |
| 371 | |
| 372 | assert result.returncode == 0, result.stderr |
| 373 | records = "".join(path.read_text() for path in tmp_path.glob("spans-*.jsonl")) |
| 374 | assert "export VALUE=ready" in records |
| 375 | assert "check VALUE" in records |
| 376 | |
| 377 | |
| 378 | def test_shell_wrapper_preserves_failure_status(tmp_path): |
| 379 | shell = ( |
| 380 | f'. "{SCRIPTS_DIR / "ci_otel.sh"}"; ' |
| 381 | f"ci_otel_start 1 {_quoted('false')}; false; status=$?; " |
| 382 | 'ci_otel_finish "$status"; exit "$status"' |
| 383 | ) |
| 384 | |
| 385 | result = subprocess.run( |
| 386 | ["/bin/sh", "-c", shell], |
| 387 | check=False, |
| 388 | capture_output=True, |
| 389 | text=True, |
| 390 | env={ |
| 391 | **os.environ, |
| 392 | "CI_INFRA_OTEL_DIR": str(SCRIPTS_DIR), |
| 393 | "CI_INFRA_OTEL_SPOOL_DIR": str(tmp_path), |
| 394 | }, |
| 395 | ) |
| 396 | |
| 397 | assert result.returncode == 1 |
| 398 | |
| 399 | |
| 400 | def test_shell_wrapper_does_not_write_bytecode_to_helper_tree(tmp_path): |
| 401 | helper_dir = tmp_path / "ci-otel" |
| 402 | shutil.copytree( |
| 403 | SCRIPTS_DIR, |
| 404 | helper_dir, |
| 405 | ignore=shutil.ignore_patterns("__pycache__"), |
| 406 | ) |
| 407 | shell = ( |
| 408 | f'. "{helper_dir / "ci_otel.sh"}"; ' |
| 409 | f"ci_otel_start 1 {_quoted('true')}; ci_otel_finish 0" |
| 410 | ) |
| 411 | |
| 412 | result = subprocess.run( |
| 413 | ["/bin/sh", "-c", shell], |
| 414 | check=False, |
| 415 | capture_output=True, |
| 416 | text=True, |
| 417 | env={ |
| 418 | **os.environ, |
| 419 | "CI_INFRA_OTEL_DIR": str(helper_dir), |
| 420 | "CI_INFRA_GPU_SAMPLING": "0", |
| 421 | "CI_INFRA_OTEL_SPOOL_DIR": str(tmp_path / "spans"), |
| 422 | }, |
| 423 | ) |
| 424 | |
| 425 | assert result.returncode == 0, result.stderr |
| 426 | assert list(helper_dir.rglob("*.pyc")) == [] |
| 427 | |
| 428 | |
| 429 | def test_ci_otel_run_records_command_and_preserves_status(tmp_path): |
| 430 | shell = ( |
| 431 | f'. "{SCRIPTS_DIR / "ci_otel.sh"}"; ' |
| 432 | f"ci_otel_run 1 {_quoted('true')} true; " |
| 433 | f"ci_otel_run 2 {_quoted('false')} false; " |
| 434 | 'echo "status=$?"' |
| 435 | ) |
| 436 | |
| 437 | result = subprocess.run( |
| 438 | ["/bin/sh", "-c", shell], |
| 439 | check=False, |
| 440 | capture_output=True, |
| 441 | text=True, |
| 442 | env={ |
| 443 | **os.environ, |
| 444 | "CI_INFRA_OTEL_DIR": str(SCRIPTS_DIR), |
| 445 | "CI_INFRA_OTEL_SPOOL_DIR": str(tmp_path), |
| 446 | }, |
| 447 | ) |
| 448 | |
| 449 | assert result.returncode == 0, result.stderr |
| 450 | assert "status=1" in result.stdout |
| 451 | records = "".join(path.read_text() for path in tmp_path.glob("spans-*.jsonl")) |
| 452 | assert '"ci.command.index":1' in records |
| 453 | assert '"ci.command.index":2' in records |
| 454 | assert '"process.exit.code":0' in records |
| 455 | assert '"process.exit.code":1' in records |
| 456 | |
| 457 | |
| 458 | def test_ci_otel_run_fail_open_when_tracing_unavailable(tmp_path): |
| 459 | shell = ( |
| 460 | f'CI_INFRA_OTEL_DIR="{tmp_path / "missing"}"; export CI_INFRA_OTEL_DIR; ' |
| 461 | f'. "{SCRIPTS_DIR / "ci_otel.sh"}"; ' |
| 462 | f"ci_otel_run 1 {_quoted('echo ran')} echo ran" |
| 463 | ) |
| 464 | |
| 465 | result = subprocess.run( |
| 466 | ["/bin/sh", "-e", "-c", shell], |
| 467 | check=False, |
| 468 | capture_output=True, |
| 469 | text=True, |
| 470 | ) |
| 471 | |
| 472 | assert result.returncode == 0, result.stderr |
| 473 | assert result.stdout.strip() == "ran" |
| 474 | |
| 475 | |
| 476 | def test_ci_otel_run_handles_assignment_prefixed_commands(tmp_path): |
| 477 | shell = ( |
| 478 | f'. "{SCRIPTS_DIR / "ci_otel.sh"}"; ' |
| 479 | f"ci_otel_run 1 {_quoted('MY_VAR=hello env')} " |
| 480 | "MY_VAR=hello env" |
| 481 | ) |
| 482 | |
| 483 | result = subprocess.run( |
| 484 | ["/bin/sh", "-c", shell], |
| 485 | check=False, |
| 486 | capture_output=True, |
| 487 | text=True, |
| 488 | env={ |
| 489 | **os.environ, |
| 490 | "CI_INFRA_OTEL_DIR": str(SCRIPTS_DIR), |
| 491 | "CI_INFRA_OTEL_SPOOL_DIR": str(tmp_path), |
| 492 | }, |
| 493 | ) |
| 494 | |
| 495 | assert result.returncode == 0, result.stderr |
| 496 | assert "MY_VAR=hello" in result.stdout |
| 497 | |
| 498 | |
| 499 | def test_missing_helpers_do_not_block_the_test_command(tmp_path): |
| 500 | output = tmp_path / "ran" |
| 501 | shell = ( |
| 502 | f'CI_INFRA_OTEL_DIR="{tmp_path / "missing"}"; export CI_INFRA_OTEL_DIR; ' |
| 503 | f'. "{SCRIPTS_DIR / "ci_otel.sh"}"; printf ran > "{output}"' |
| 504 | ) |
| 505 | |
| 506 | result = subprocess.run( |
| 507 | ["/bin/sh", "-e", "-c", shell], |
| 508 | check=False, |
| 509 | capture_output=True, |
| 510 | text=True, |
| 511 | ) |
| 512 | |
| 513 | assert result.returncode == 0, result.stderr |
| 514 | assert output.read_text() == "ran" |
| 515 | |
| 516 | |
| 517 | def test_unset_helper_dir_does_not_block_the_test_command(): |
| 518 | shell = f'unset CI_INFRA_OTEL_DIR; . "{SCRIPTS_DIR / "ci_otel.sh"}"; echo ran' |
| 519 | |
| 520 | result = subprocess.run( |
| 521 | ["/bin/bash", "-e", "-c", shell], |
| 522 | check=False, |
| 523 | capture_output=True, |
| 524 | text=True, |
| 525 | ) |
| 526 | |
| 527 | assert result.returncode == 0, result.stderr |
| 528 | assert result.stdout.strip() == "ran" |
| 529 | |
| 530 | |
| 531 | def test_helper_dir_assignment_is_exported(): |
| 532 | shell = ( |
| 533 | "unset CI_INFRA_OTEL_DIR; " |
| 534 | f'CI_INFRA_OTEL_DIR="{SCRIPTS_DIR}" . "{SCRIPTS_DIR / "ci_otel.sh"}"; ' |
| 535 | f'/bin/sh -c \'test "$CI_INFRA_OTEL_DIR" = "{SCRIPTS_DIR}"\'' |
| 536 | ) |
| 537 | |
| 538 | result = subprocess.run( |
| 539 | ["/bin/sh", "-c", shell], check=False, capture_output=True, text=True |
| 540 | ) |
| 541 | |
| 542 | assert result.returncode == 0, result.stderr |
| 543 | |
| 544 | |
| 545 | @pytest.mark.parametrize("failure", ["new-context", "record-command"]) |
| 546 | def test_helper_failure_cannot_change_job_status(tmp_path, failure): |
| 547 | fake_bin = tmp_path / "bin" |
| 548 | fake_bin.mkdir() |
| 549 | fake_python = fake_bin / "python3" |
| 550 | fake_python.write_text( |
| 551 | "#!/bin/sh\n" |
| 552 | '[ "$1" = "-c" ] && exit 0\n' |
| 553 | '[ "$2" = "new-context" ] && {\n' |
| 554 | ' [ "$CI_OTEL_TEST_FAILURE" = "new-context" ] && exit 99\n' |
| 555 | ' echo "01010101010101010101010101010101 0202020202020202 - 1"\n' |
| 556 | " exit 0\n" |
| 557 | "}\n" |
| 558 | '[ "$2" = "record-command" ] && exit 99\n' |
| 559 | "exit 99\n" |
| 560 | ) |
| 561 | fake_python.chmod(0o755) |
| 562 | shell = ( |
| 563 | f'PATH="{fake_bin}:$PATH"; . "{SCRIPTS_DIR / "ci_otel.sh"}"; ' |
| 564 | f"ci_otel_start 1 {_quoted('true')}; ci_otel_finish 0; echo ran" |
| 565 | ) |
| 566 | |
| 567 | result = subprocess.run( |
| 568 | ["/bin/sh", "-e", "-c", shell], |
| 569 | check=False, |
| 570 | capture_output=True, |
| 571 | text=True, |
| 572 | env={ |
| 573 | **os.environ, |
| 574 | "CI_INFRA_OTEL_DIR": str(SCRIPTS_DIR), |
| 575 | "CI_INFRA_OTEL_RUNTIME_DIR": str(tmp_path / "runtime"), |
| 576 | "CI_OTEL_TEST_FAILURE": failure, |
| 577 | }, |
| 578 | ) |
| 579 | |
| 580 | assert result.returncode == 0, result.stderr |
| 581 | assert result.stdout.strip() == "ran" |
| 582 | |
| 583 | |
| 584 | def test_flush_failure_preserves_job_status(tmp_path): |
| 585 | fake_bin = tmp_path / "bin" |
| 586 | fake_bin.mkdir() |
| 587 | fake_python = fake_bin / "python3" |
| 588 | fake_python.write_text("#!/bin/sh\nexit 99\n") |
| 589 | fake_python.chmod(0o755) |
| 590 | |
| 591 | def run(status: int): |
| 592 | shell = ( |
| 593 | f'. "{SCRIPTS_DIR / "ci_otel.sh"}"; ' |
| 594 | f'PATH="{fake_bin}"; export PATH; exit {status}' |
| 595 | ) |
| 596 | return subprocess.run( |
| 597 | ["/bin/sh", "-c", shell], |
| 598 | check=False, |
| 599 | capture_output=True, |
| 600 | text=True, |
| 601 | env={**os.environ, "CI_INFRA_OTEL_DIR": str(SCRIPTS_DIR)}, |
| 602 | ) |
| 603 | |
| 604 | assert run(0).returncode == 0 |
| 605 | assert run(7).returncode == 7 |
| 606 | |
| 607 | |
| 608 | def test_shell_wrapper_uses_private_trace_ids(tmp_path): |
| 609 | shell = ( |
| 610 | f'. "{SCRIPTS_DIR / "ci_otel.sh"}"; ' |
| 611 | f"ci_otel_start 1 {_quoted('true')}; " |
| 612 | 'expected="$CI_INFRA_TRACE_ID"; ' |
| 613 | "CI_INFRA_TRACE_ID=corrupted; CI_INFRA_COMMAND_SPAN_ID=corrupted; " |
| 614 | 'ci_otel_finish 0; printf "%s" "$expected"' |
| 615 | ) |
| 616 | |
| 617 | result = subprocess.run( |
| 618 | ["/bin/sh", "-c", shell], |
| 619 | check=False, |
| 620 | capture_output=True, |
| 621 | text=True, |
| 622 | env={ |
| 623 | **os.environ, |
| 624 | "CI_INFRA_OTEL_DIR": str(SCRIPTS_DIR), |
| 625 | "CI_INFRA_OTEL_SPOOL_DIR": str(tmp_path), |
| 626 | }, |
| 627 | ) |
| 628 | |
| 629 | assert result.returncode == 0, result.stderr |
| 630 | record = json.loads(next(tmp_path.glob("spans-*.jsonl")).read_text()) |
| 631 | assert record["trace_id"] == result.stdout |
| 632 | assert record["span_id"] != "corrupted" |
| 633 | |
| 634 | |
| 635 | def test_shell_setup_is_idempotent(tmp_path): |
| 636 | runtime = tmp_path / "runtime" |
| 637 | shell = ( |
| 638 | f'. "{SCRIPTS_DIR / "ci_otel.sh"}"; ' |
| 639 | f'. "{SCRIPTS_DIR / "ci_otel.sh"}"; ' |
| 640 | 'printf "%s\n%s" "$PATH" "$CI_INFRA_OTEL_SHIM_PATHS"' |
| 641 | ) |
| 642 | |
| 643 | result = subprocess.run( |
| 644 | ["/bin/sh", "-c", shell], |
| 645 | check=False, |
| 646 | capture_output=True, |
| 647 | text=True, |
| 648 | env={ |
| 649 | **os.environ, |
| 650 | "CI_INFRA_OTEL_DIR": str(SCRIPTS_DIR), |
| 651 | "CI_INFRA_OTEL_RUNTIME_DIR": str(runtime), |
| 652 | }, |
| 653 | ) |
| 654 | |
| 655 | assert result.returncode == 0, result.stderr |
| 656 | path, shim_paths = result.stdout.splitlines() |
| 657 | assert path.split(":").count(str(runtime / "bin")) == 1 |
| 658 | assert shim_paths.split(":").count(str(runtime / "bin")) == 1 |
| 659 | |
| 660 | |
| 661 | def test_helper_owned_runtime_is_removed_on_exit(tmp_path): |
| 662 | runtime_file = tmp_path / "runtime" |
| 663 | shell = ( |
| 664 | f'. "{SCRIPTS_DIR / "ci_otel.sh"}"; ' |
| 665 | f'printf "%s" "$CI_INFRA_OTEL_RUNTIME_DIR" > "{runtime_file}"' |
| 666 | ) |
| 667 | |
| 668 | result = subprocess.run( |
| 669 | ["/bin/sh", "-c", shell], |
| 670 | check=False, |
| 671 | capture_output=True, |
| 672 | text=True, |
| 673 | env={**os.environ, "CI_INFRA_OTEL_DIR": str(SCRIPTS_DIR)}, |
| 674 | ) |
| 675 | |
| 676 | assert result.returncode == 0, result.stderr |
| 677 | assert not Path(runtime_file.read_text()).exists() |
| 678 | |
| 679 | |
| 680 | def test_caller_owned_runtime_is_preserved(tmp_path): |
| 681 | runtime = tmp_path / "runtime" |
| 682 | shell = f'. "{SCRIPTS_DIR / "ci_otel.sh"}"' |
| 683 | |
| 684 | result = subprocess.run( |
| 685 | ["/bin/sh", "-c", shell], |
| 686 | check=False, |
| 687 | capture_output=True, |
| 688 | text=True, |
| 689 | env={ |
| 690 | **os.environ, |
| 691 | "CI_INFRA_OTEL_DIR": str(SCRIPTS_DIR), |
| 692 | "CI_INFRA_OTEL_RUNTIME_DIR": str(runtime), |
| 693 | }, |
| 694 | ) |
| 695 | |
| 696 | assert result.returncode == 0, result.stderr |
| 697 | assert runtime.exists() |
| 698 | |
| 699 | |
| 700 | def test_pytest_shim_recovers_from_stale_recorded_path(tmp_path): |
| 701 | bin_dir = tmp_path / "bin" |
| 702 | bin_dir.mkdir() |
| 703 | real_pytest = bin_dir / "pytest" |
| 704 | real_pytest.write_text("#!/bin/sh\nprintf 'CURRENT-PYTEST'") |
| 705 | real_pytest.chmod(0o755) |
| 706 | |
| 707 | result = subprocess.run( |
| 708 | [str(SCRIPTS_DIR / "ci_pytest.sh"), "--version"], |
| 709 | check=False, |
| 710 | capture_output=True, |
| 711 | text=True, |
| 712 | env={ |
| 713 | **os.environ, |
| 714 | "PATH": str(bin_dir), |
| 715 | "CI_INFRA_OTEL_REAL_PYTEST": "/missing/pytest", |
| 716 | }, |
| 717 | ) |
| 718 | |
| 719 | assert result.returncode == 0, result.stderr |
| 720 | assert result.stdout == "CURRENT-PYTEST" |
| 721 | |
| 722 | |
| 723 | def test_path_change_after_source_uses_new_pytest(tmp_path): |
| 724 | old_bin = tmp_path / "old" |
| 725 | new_bin = tmp_path / "new" |
| 726 | runtime = tmp_path / "runtime" |
| 727 | old_bin.mkdir() |
| 728 | new_bin.mkdir() |
| 729 | for directory, output in ((old_bin, "OLD"), (new_bin, "NEW")): |
| 730 | executable = directory / "pytest" |
| 731 | executable.write_text(f"#!/bin/sh\nprintf '{output}'") |
| 732 | executable.chmod(0o755) |
| 733 | shell = ( |
| 734 | f'PATH="{old_bin}:$PATH"; . "{SCRIPTS_DIR / "ci_otel.sh"}"; ' |
| 735 | f'PATH="{new_bin}:$PATH"; hash -r; pytest' |
| 736 | ) |
| 737 | |
| 738 | result = subprocess.run( |
| 739 | ["/bin/sh", "-c", shell], |
| 740 | check=False, |
| 741 | capture_output=True, |
| 742 | text=True, |
| 743 | env={ |
| 744 | **os.environ, |
| 745 | "CI_INFRA_OTEL_DIR": str(SCRIPTS_DIR), |
| 746 | "CI_INFRA_OTEL_RUNTIME_DIR": str(runtime), |
| 747 | }, |
| 748 | ) |
| 749 | |
| 750 | assert result.returncode == 0, result.stderr |
| 751 | assert result.stdout == "NEW" |
| 752 | |
| 753 | |
| 754 | def test_pytest_shim_records_distinct_test_intervals_with_pythonpath_override( |
| 755 | tmp_path, |
| 756 | ): |
| 757 | test_file = tmp_path / "test_sample.py" |
| 758 | test_file.write_text( |
| 759 | "import time\n" |
| 760 | "def test_first(): time.sleep(0.02)\n" |
| 761 | "def test_second(): time.sleep(0.02)\n" |
| 762 | ) |
| 763 | workspace = tmp_path / "workspace" |
| 764 | workspace.mkdir() |
| 765 | bin_dir = tmp_path / "bin" |
| 766 | bin_dir.mkdir() |
| 767 | pytest_executable = bin_dir / "pytest" |
| 768 | pytest_executable.write_text( |
| 769 | f'#!/bin/sh\nexec {shlex.quote(sys.executable)} -m pytest "$@"\n' |
| 770 | ) |
| 771 | pytest_executable.chmod(0o755) |
| 772 | shell = ( |
| 773 | f'. "{SCRIPTS_DIR / "ci_otel.sh"}"\n' |
| 774 | f"ci_otel_start 1 {_quoted('pytest tests')}\n" |
| 775 | f"PYTHONPATH={shlex.quote(str(workspace))} " |
| 776 | f"pytest -q {shlex.quote(str(test_file))}\n" |
| 777 | "status=$?\n" |
| 778 | 'ci_otel_finish "$status"\n' |
| 779 | 'exit "$status"' |
| 780 | ) |
| 781 | |
| 782 | result = subprocess.run( |
| 783 | ["/bin/sh", "-c", shell], |
| 784 | check=False, |
| 785 | capture_output=True, |
| 786 | text=True, |
| 787 | env={ |
| 788 | **os.environ, |
| 789 | "CI_INFRA_OTEL_DIR": str(SCRIPTS_DIR), |
| 790 | "CI_INFRA_OTEL_SPOOL_DIR": str(tmp_path / "spans"), |
| 791 | "PATH": f"{bin_dir}:{os.environ['PATH']}", |
| 792 | "PYTEST_ADDOPTS": "", |
| 793 | }, |
| 794 | ) |
| 795 | |
| 796 | assert result.returncode == 0, result.stderr |
| 797 | records = [ |
| 798 | json.loads(line) |
| 799 | for path in (tmp_path / "spans").glob("spans-*.jsonl") |
| 800 | for line in path.read_text().splitlines() |
| 801 | ] |
| 802 | tests = [record for record in records if record["name"] == "pytest.test"] |
| 803 | assert len(tests) == 2 |
| 804 | assert tests[0]["start_ns"] < tests[1]["start_ns"] |
| 805 | assert tests[0]["end_ns"] <= tests[1]["start_ns"] |
| 806 | |
| 807 | |
| 808 | def test_pytest_plugin_failure_cannot_fail_pytest(tmp_path): |
| 809 | test_file = tmp_path / "test_sample.py" |
| 810 | test_file.write_text("def test_sample(): assert True\n") |
| 811 | |
| 812 | result = subprocess.run( |
| 813 | [sys.executable, "-m", "pytest", "-q", "-p", "ci_otel", str(test_file)], |
| 814 | check=False, |
| 815 | capture_output=True, |
| 816 | text=True, |
| 817 | env={ |
| 818 | **os.environ, |
| 819 | "PYTHONPATH": str(SCRIPTS_DIR), |
| 820 | "CI_INFRA_TRACE_ID": "01" * 16, |
| 821 | "CI_INFRA_COMMAND_SPAN_ID": "02" * 8, |
| 822 | "CI_INFRA_OTEL_SPOOL_DIR": "/dev/null/spans", |
| 823 | }, |
| 824 | ) |
| 825 | |
| 826 | assert result.returncode == 0, result.stderr |
| 827 | assert "1 passed" in result.stdout |
| 828 | |
| 829 | |
| 830 | def test_gpu_samples_preserve_zero_and_omit_unsupported_metrics(): |
| 831 | events = ci_gpu.parse_samples( |
| 832 | "0, GPU-a, H200, 0, 1024, 2048, Disabled\n" |
| 833 | "1, GPU-b, H200, [N/A], [N/A], 2048, Disabled\n" |
| 834 | "2, GPU-c, H200, 50, 1024, 2048, Enabled\n" |
| 835 | "3, GPU-d, H200, NaN, -1, inf, Disabled\n", |
| 836 | 1234, |
| 837 | ) |
| 838 | assert len(events) == 4 |
| 839 | assert events[0]["attributes"]["gpu.utilization"] == 0 |
| 840 | assert events[0]["attributes"]["gpu.memory.used"] == 1024**3 |
| 841 | for event in events[1:]: |
| 842 | assert "gpu.utilization" not in event["attributes"] |
| 843 | assert "gpu.memory.used" not in event["attributes"] |
| 844 | assert events[2]["attributes"]["gpu.status"] == "unsupported_mig" |
| 845 | assert "gpu.memory.total" not in events[2]["attributes"] |
| 846 | |
| 847 | |
| 848 | @pytest.mark.parametrize( |
| 849 | "cuda,nvidia,expected", |
| 850 | [ |
| 851 | ("MIG-a", "GPU-parent", ["MIG-a"]), |
| 852 | ("0", "MIG-a", ["MIG-a"]), |
| 853 | ("1,0", "MIG-a,MIG-b", ["MIG-b", "MIG-a"]), |
| 854 | (None, "MIG-a", ["MIG-a"]), |
| 855 | ("", "MIG-a", []), |
| 856 | ("-1", "MIG-a", []), |
| 857 | ("2", "MIG-a", []), |
| 858 | ("GPU-parent", "MIG-a", []), |
| 859 | ("0", "all", []), |
| 860 | ], |
| 861 | ) |
| 862 | def test_mig_assignment_obeys_cuda_visibility(monkeypatch, cuda, nvidia, expected): |
| 863 | monkeypatch.setenv("NVIDIA_VISIBLE_DEVICES", nvidia) |
| 864 | monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False) |
| 865 | if cuda is not None: |
| 866 | monkeypatch.setenv("CUDA_VISIBLE_DEVICES", cuda) |
| 867 | assert ci_gpu.visible_mig_devices() == expected |
| 868 | |
| 869 | |
| 870 | @pytest.mark.parametrize("used", [0, 3 * 1024**3]) |
| 871 | def test_mig_memory_uses_assigned_handle_without_parent_metrics(monkeypatch, used): |
| 872 | handles = [] |
| 873 | closed = [] |
| 874 | |
| 875 | def get_handle(uuid): |
| 876 | handles.append(uuid) |
| 877 | return "slice-handle" |
| 878 | |
| 879 | def get_memory(handle): |
| 880 | assert handle == "slice-handle" |
| 881 | return SimpleNamespace(used=used, total=16 * 1024**3) |
| 882 | |
| 883 | nvml = SimpleNamespace( |
| 884 | nvmlInit=lambda: None, |
| 885 | nvmlShutdown=lambda: closed.append(True), |
| 886 | nvmlDeviceGetHandleByUUID=get_handle, |
| 887 | nvmlDeviceIsMigDeviceHandle=lambda handle: handle == "slice-handle", |
| 888 | nvmlDeviceGetMemoryInfo=get_memory, |
| 889 | NVMLError=RuntimeError, |
| 890 | ) |
| 891 | monkeypatch.setattr(ci_gpu, "load_nvml", lambda: nvml) |
| 892 | samples = ci_gpu.query_mig_samples(["MIG-a"], 1234) |
| 893 | assert handles == ["MIG-a"] |
| 894 | assert closed == [True] |
| 895 | assert samples[0]["time_ns"] == 1234 |
| 896 | attributes = samples[0]["attributes"] |
| 897 | assert attributes["gpu.uuid"] == "MIG-a" |
| 898 | assert attributes["gpu.memory.used"] == used |
| 899 | assert attributes["gpu.memory.total"] == 16 * 1024**3 |
| 900 | assert "gpu.utilization" not in attributes |
| 901 | |
| 902 | |
| 903 | @pytest.mark.parametrize("is_mig", [True, False]) |
| 904 | def test_mig_memory_never_substitutes_parent_or_permission_failure(monkeypatch, is_mig): |
| 905 | def unavailable(handle): |
| 906 | if not is_mig: |
| 907 | pytest.fail("queried parent memory") |
| 908 | raise RuntimeError("Insufficient Permissions") |
| 909 | |
| 910 | nvml = SimpleNamespace( |
| 911 | nvmlInit=lambda: None, |
| 912 | nvmlShutdown=lambda: None, |
| 913 | nvmlDeviceGetHandleByUUID=lambda uuid: "handle", |
| 914 | nvmlDeviceIsMigDeviceHandle=lambda handle: is_mig, |
| 915 | nvmlDeviceGetMemoryInfo=unavailable, |
| 916 | NVMLError=RuntimeError, |
| 917 | ) |
| 918 | monkeypatch.setattr(ci_gpu, "load_nvml", lambda: nvml) |
| 919 | samples = ci_gpu.query_mig_samples(["MIG-a"], 1234) |
| 920 | if is_mig: |
| 921 | assert samples[0]["attributes"]["gpu.status"] == "unavailable_mig_memory" |
| 922 | assert "gpu.memory.used" not in samples[0]["attributes"] |
| 923 | else: |
| 924 | assert samples == [] |
| 925 | |
| 926 | |
| 927 | def test_gpu_batches_keep_command_identity_and_survive_spool_roundtrip( |
| 928 | monkeypatch, tmp_path |
| 929 | ): |
| 930 | monkeypatch.setenv("CI_INFRA_TRACE_ID", "01" * 16) |
| 931 | monkeypatch.setenv("CI_INFRA_COMMAND_SPAN_ID", "02" * 8) |
| 932 | monkeypatch.setenv("CI_INFRA_OTEL_SPOOL_DIR", str(tmp_path)) |
| 933 | span = ci_gpu.sample_span( |
| 934 | ci_gpu.parse_samples("0, GPU-a, H200, 90, 1, 2, Disabled", 1234) |
| 935 | ) |
| 936 | assert span.parent_span_id == "02" * 8 |
| 937 | assert span.trace_id == "01" * 16 |
| 938 | assert ci_otel.record_spans([span]) |
| 939 | assert ci_otel.load_spans() == [span] |
| 940 | assert b"ci.gpu.sample" in ci_otel.encode_request([span]) |
| 941 | |
| 942 | |
| 943 | @pytest.mark.parametrize("mig", [False, True]) |
| 944 | def test_gpu_query_failures_are_bounded_and_never_synthesize_idle_samples( |
| 945 | monkeypatch, mig |
| 946 | ): |
| 947 | monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "MIG-a" if mig else "0") |
| 948 | calls = [] |
| 949 | |
| 950 | def query(*args, **kwargs): |
| 951 | calls.append(kwargs["timeout"]) |
| 952 | raise subprocess.TimeoutExpired("nvidia-smi", kwargs["timeout"]) |
| 953 | |
| 954 | monkeypatch.setattr(ci_gpu.subprocess, "run", query) |
| 955 | monkeypatch.setattr(ci_gpu, "INTERVAL_SECONDS", 0) |
| 956 | monkeypatch.setattr( |
| 957 | ci_gpu, "record_spans", lambda _: pytest.fail("fabricated sample") |
| 958 | ) |
| 959 | ci_gpu.collect(os.getppid(), threading.Event()) |
| 960 | assert calls == [2, 2, 2] |
| 961 | |
| 962 | |
| 963 | def test_gpu_periodic_batches_are_bounded_even_when_upload_fails(monkeypatch): |
| 964 | monkeypatch.setenv("CI_INFRA_TRACE_ID", "01" * 16) |
| 965 | monkeypatch.setenv("CI_INFRA_COMMAND_SPAN_ID", "02" * 8) |
| 966 | monkeypatch.setattr(ci_gpu, "BATCH_SECONDS", 0) |
| 967 | monkeypatch.setattr(ci_gpu, "INTERVAL_SECONDS", 0) |
| 968 | monkeypatch.setattr( |
| 969 | ci_gpu.subprocess, |
| 970 | "run", |
| 971 | lambda *a, **k: subprocess.CompletedProcess( |
| 972 | a, 0, "0, GPU-a, H200, 50, 1, 2, Disabled" |
| 973 | ), |
| 974 | ) |
| 975 | stop = threading.Event() |
| 976 | batches = [] |
| 977 | |
| 978 | def upload(spans, timeout_seconds): |
| 979 | batches.append(spans) |
| 980 | assert timeout_seconds == 2 |
| 981 | if len(batches) == 3: |
| 982 | stop.set() |
| 983 | return False |
| 984 | |
| 985 | monkeypatch.setattr(ci_gpu, "export_spans", upload) |
| 986 | ci_gpu.collect(os.getppid(), stop) |
| 987 | assert [len(batch[0].events) for batch in batches] == [1, 1, 1] |
| 988 | |
| 989 | |
| 990 | def _write_executable(path: Path, body: str) -> None: |
| 991 | path.write_text(f"#!/bin/sh\n{body}") |
| 992 | path.chmod(0o755) |
| 993 | |
| 994 | |
| 995 | @pytest.mark.parametrize( |
| 996 | "workspace_owner,expect_sampler", |
| 997 | [("997 996", True), ("0 0", False)], |
| 998 | ) |
| 999 | def test_root_shell_runs_gpu_sampler_as_checkout_owner( |
| 1000 | tmp_path, workspace_owner, expect_sampler |
| 1001 | ): |
| 1002 | bin_dir = tmp_path / "bin" |
| 1003 | workspace = tmp_path / "workdir" |
| 1004 | runtime = tmp_path / "runtime" |
| 1005 | log = tmp_path / "calls.log" |
| 1006 | bin_dir.mkdir() |
| 1007 | workspace.mkdir() |
| 1008 | |
| 1009 | _write_executable( |
| 1010 | bin_dir / "id", |
| 1011 | '[ "$1" = "-u" ] && { echo 0; exit 0; }; exec /usr/bin/id "$@"\n', |
| 1012 | ) |
| 1013 | _write_executable( |
| 1014 | bin_dir / "stat", |
| 1015 | f"printf '%s\\n' '{workspace_owner}'\n", |
| 1016 | ) |
| 1017 | for command in ("chown", "chmod"): |
| 1018 | _write_executable(bin_dir / command, "exit 0\n") |
| 1019 | _write_executable( |
| 1020 | bin_dir / "setpriv", |
| 1021 | 'printf "setpriv %s\\n" "$*" >> "$CI_OTEL_TEST_LOG"\n' |
| 1022 | 'while [ "$#" -gt 0 ] && [ "$1" != "--" ]; do shift; done\n' |
| 1023 | '[ "$#" -gt 0 ] && shift\n' |
| 1024 | 'exec "$@"\n', |
| 1025 | ) |
| 1026 | _write_executable( |
| 1027 | bin_dir / "python3", |
| 1028 | '[ "$1" = "-c" ] && exit 0\n' |
| 1029 | '[ "$2" = "new-context" ] && {\n' |
| 1030 | ' echo "01010101010101010101010101010101 0202020202020202 - 1"\n' |
| 1031 | " exit 0\n" |
| 1032 | "}\n" |
| 1033 | 'case "$1" in *ci_gpu.py) echo gpu >> "$CI_OTEL_TEST_LOG";; esac\n' |
| 1034 | "exit 0\n", |
| 1035 | ) |
| 1036 | _write_executable(bin_dir / "nvidia-smi", "exit 0\n") |
| 1037 | shell = ( |
| 1038 | f'. "{SCRIPTS_DIR / "ci_otel.sh"}"; ' |
| 1039 | f"ci_otel_start 1 {_quoted('true')}; ci_otel_finish 0" |
| 1040 | ) |
| 1041 | |
| 1042 | result = subprocess.run( |
| 1043 | ["/bin/sh", "-c", shell], |
| 1044 | check=False, |
| 1045 | capture_output=True, |
| 1046 | text=True, |
| 1047 | env={ |
| 1048 | **os.environ, |
| 1049 | "PATH": f"{bin_dir}:{os.environ['PATH']}", |
| 1050 | "CI_INFRA_OTEL_DIR": str(SCRIPTS_DIR), |
| 1051 | "CI_INFRA_OTEL_RUNTIME_DIR": str(runtime), |
| 1052 | "CI_INFRA_OTEL_WORKSPACE_DIR": str(workspace), |
| 1053 | "CI_OTEL_TEST_LOG": str(log), |
| 1054 | }, |
| 1055 | ) |
| 1056 | |
| 1057 | assert result.returncode == 0, result.stderr |
| 1058 | calls = log.read_text() if log.exists() else "" |
| 1059 | if expect_sampler: |
| 1060 | expected = "--reuid=997 --regid=996 --clear-groups --no-new-privs --" |
| 1061 | assert calls.count("setpriv") == 2 |
| 1062 | assert calls.count(expected) == 2 |
| 1063 | assert "gpu" in calls |
| 1064 | else: |
| 1065 | assert calls == "" |
| 1066 | assert "non-root GPU sampler unavailable" in result.stderr |
| 1067 | |
| 1068 | |
| 1069 | @pytest.mark.parametrize("sampling", ["1", "0"]) |
| 1070 | def test_gpu_sampler_stops_with_command_and_preserves_failure_status( |
| 1071 | tmp_path, sampling |
| 1072 | ): |
| 1073 | bin_dir = tmp_path / "bin" |
| 1074 | bin_dir.mkdir() |
| 1075 | smi = bin_dir / "nvidia-smi" |
| 1076 | smi.write_text('#!/bin/sh\necho "0, GPU-a, H200, 70, 1024, 2048, Disabled"\n') |
| 1077 | smi.chmod(0o755) |
| 1078 | # Keep the spool for inspection, and avoid any network upload. |
| 1079 | shell = ( |
| 1080 | f'. "{SCRIPTS_DIR / "ci_otel.sh"}"\n' |
| 1081 | "ci_otel_run 1 workload sh -c 'sleep 1.3; exit 7'\n" |
| 1082 | "exit $?" |
| 1083 | ) |
| 1084 | result = subprocess.run( |
| 1085 | ["/bin/sh", "-c", shell], |
| 1086 | capture_output=True, |
| 1087 | text=True, |
| 1088 | timeout=8, |
| 1089 | env={ |
| 1090 | **os.environ, |
| 1091 | "PATH": f"{bin_dir}:{os.environ['PATH']}", |
| 1092 | "BUILDKITE": "false", |
| 1093 | "CI_INFRA_GPU_SAMPLING": sampling, |
| 1094 | "CI_INFRA_OTEL_DIR": str(SCRIPTS_DIR), |
| 1095 | "CI_INFRA_OTEL_RUNTIME_DIR": str(tmp_path / "runtime"), |
| 1096 | "CI_INFRA_OTEL_SPOOL_DIR": str(tmp_path / "spans"), |
| 1097 | }, |
| 1098 | ) |
| 1099 | assert result.returncode == 7, result.stderr |
| 1100 | records = [ |
| 1101 | json.loads(line) |
| 1102 | for path in (tmp_path / "spans").glob("spans-*.jsonl") |
| 1103 | for line in path.read_text().splitlines() |
| 1104 | ] |
| 1105 | command = next(record for record in records if record["name"] == "ci.command") |
| 1106 | batches = [record for record in records if record["name"] == "ci.gpu.samples"] |
| 1107 | assert bool(batches) == (sampling == "1") |
| 1108 | if batches: |
| 1109 | assert batches[0]["parent_span_id"] == command["span_id"] |
| 1110 | assert all( |
| 1111 | command["start_ns"] <= event["time_ns"] <= command["end_ns"] |
| 1112 | for event in batches[0]["events"] |
| 1113 | ) |
| 1114 | |
| 1115 | |
| 1116 | @pytest.mark.parametrize("is_mig", [True, False]) |
| 1117 | def test_cdi_discovery_uses_only_cuda_visible_uuids(monkeypatch, is_mig): |
| 1118 | identifiers = [] |
| 1119 | |
| 1120 | def get_count(output): |
| 1121 | output._obj.value = 1 |
| 1122 | return 0 |
| 1123 | |
| 1124 | def get_device(output, index): |
| 1125 | assert index == 0 |
| 1126 | output._obj.value = 7 |
| 1127 | return 0 |
| 1128 | |
| 1129 | def get_uuid(output, device): |
| 1130 | assert device.value == 7 |
| 1131 | output._obj[:] = bytes(range(16)) |
| 1132 | return 0 |
| 1133 | |
| 1134 | driver = SimpleNamespace( |
| 1135 | cuInit=lambda flags: 0, |
| 1136 | cuDeviceGetCount=get_count, |
| 1137 | cuDeviceGet=get_device, |
| 1138 | cuDeviceGetUuid_v2=get_uuid, |
| 1139 | ) |
| 1140 | monkeypatch.setattr(ci_gpu.ctypes, "CDLL", lambda _: driver) |
| 1141 | nvml = SimpleNamespace( |
| 1142 | nvmlInit=lambda: None, |
| 1143 | nvmlShutdown=lambda: None, |
| 1144 | nvmlDeviceGetHandleByUUID=lambda value: identifiers.append(value) or "handle", |
| 1145 | nvmlDeviceIsMigDeviceHandle=lambda handle: is_mig, |
| 1146 | NVMLError=RuntimeError, |
| 1147 | ) |
| 1148 | monkeypatch.setattr(ci_gpu, "load_nvml", lambda: nvml) |
| 1149 | expected = "MIG-00010203-0405-0607-0809-0a0b0c0d0e0f" |
| 1150 | assert ci_gpu.cuda_visible_mig_devices() == ([expected] if is_mig else []) |
| 1151 | assert identifiers == [expected] |
| 1152 | |
| 1153 | |
| 1154 | @pytest.mark.parametrize("discovery_fails", [False, True]) |
| 1155 | def test_cdi_sampling_switches_to_slice_or_keeps_unknown(monkeypatch, discovery_fails): |
| 1156 | monkeypatch.setenv("CI_INFRA_TRACE_ID", "01" * 16) |
| 1157 | monkeypatch.setenv("CI_INFRA_COMMAND_SPAN_ID", "02" * 8) |
| 1158 | monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False) |
| 1159 | monkeypatch.setenv("NVIDIA_VISIBLE_DEVICES", "void") |
| 1160 | monkeypatch.setattr(ci_gpu, "INTERVAL_SECONDS", 0) |
| 1161 | stop = threading.Event() |
| 1162 | commands = [] |
| 1163 | records = [] |
| 1164 | |
| 1165 | def query(command, **kwargs): |
| 1166 | assert kwargs["timeout"] == ci_gpu.QUERY_TIMEOUT_SECONDS |
| 1167 | commands.append(command) |
| 1168 | if "--discover-mig" in command: |
| 1169 | if discovery_fails: |
| 1170 | raise subprocess.TimeoutExpired(command, 2) |
| 1171 | return SimpleNamespace(stdout='["MIG-assigned"]') |
| 1172 | if len(commands) == 4: |
| 1173 | stop.set() |
| 1174 | if "--query-mig" in command: |
| 1175 | assert command[-1] == "MIG-assigned" |
| 1176 | return SimpleNamespace( |
| 1177 | stdout=json.dumps( |
| 1178 | [ |
| 1179 | { |
| 1180 | "time_ns": 2, |
| 1181 | "name": "ci.gpu.sample", |
| 1182 | "attributes": { |
| 1183 | "gpu.uuid": "MIG-assigned", |
| 1184 | "gpu.memory.used": 123, |
| 1185 | }, |
| 1186 | } |
| 1187 | ] |
| 1188 | ) |
| 1189 | ) |
| 1190 | return SimpleNamespace(stdout="0, GPU-parent, H200, 99, 100, 200, Enabled") |
| 1191 | |
| 1192 | monkeypatch.setattr(ci_gpu.subprocess, "run", query) |
| 1193 | monkeypatch.setattr(ci_gpu, "record_spans", lambda spans: records.extend(spans)) |
| 1194 | ci_gpu.collect(os.getppid(), stop) |
| 1195 | assert sum("--discover-mig" in command for command in commands) == 1 |
| 1196 | attributes = [event["attributes"] for span in records for event in span.events] |
| 1197 | if discovery_fails: |
| 1198 | assert all("gpu.memory.used" not in event for event in attributes) |
| 1199 | else: |
| 1200 | assert any(event.get("gpu.memory.used") == 123 for event in attributes) |
| 1201 | assert all(event["gpu.uuid"] == "MIG-assigned" for event in attributes) |
| 1202 | |
| 1203 | |
| 1204 | def test_bundled_nvml_load_does_not_import_vllm_or_torch(tmp_path): |
| 1205 | # Match test images: the collector is outside the installed vLLM package. |
| 1206 | helper = tmp_path / "ci_gpu_relocated.py" |
| 1207 | helper.write_text((SCRIPTS_DIR / "ci_gpu.py").read_text()) |
| 1208 | result = subprocess.run( |
| 1209 | [ |
| 1210 | sys.executable, |
| 1211 | "-c", |
| 1212 | "import sys; import ci_gpu_relocated as ci_gpu; " |
| 1213 | "binding = ci_gpu.load_nvml(); " |
| 1214 | "assert callable(binding.nvmlDeviceGetMemoryInfo); " |
| 1215 | "assert 'vllm' not in sys.modules; assert 'torch' not in sys.modules", |
| 1216 | ], |
| 1217 | env={**os.environ, "PYTHONPATH": f"{tmp_path}{os.pathsep}{SCRIPTS_DIR}"}, |
| 1218 | capture_output=True, |
| 1219 | text=True, |
| 1220 | timeout=5, |
| 1221 | ) |
| 1222 | assert result.returncode == 0, result.stderr |