| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | """Dependency-free OTLP exporter and pytest timing plugin for CI.""" |
| 5 | |
| 6 | from __future__ import annotations |
| 7 | |
| 8 | import hashlib |
| 9 | import json |
| 10 | import os |
| 11 | import secrets |
| 12 | import struct |
| 13 | import sys |
| 14 | import time |
| 15 | import urllib.request |
| 16 | from collections.abc import Iterable |
| 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 | |
| 27 | def _parse_dockerfile_step(name: str) -> tuple[str, int, int, str] | None: |
| 28 | """Parse "[stage index/total] label" without regex.""" |
| 29 | if not name.startswith("["): |
| 30 | return None |
| 31 | close = name.find("]") |
| 32 | if close == -1: |
| 33 | return None |
| 34 | header = name[1:close] |
| 35 | label = name[close + 1 :].lstrip() |
| 36 | if not label: |
| 37 | return None |
| 38 | parts = header.rsplit(None, 1) |
| 39 | if len(parts) != 2: |
| 40 | return None |
| 41 | stage, step = parts |
| 42 | if "/" not in step: |
| 43 | return None |
| 44 | index_str, total_str = step.split("/", 1) |
| 45 | if not index_str.isdigit() or not total_str.isdigit(): |
| 46 | return None |
| 47 | return stage, int(index_str), int(total_str), label |
| 48 | |
| 49 | |
| 50 | @dataclass(frozen=True) |
| 51 | class Span: |
| 52 | trace_id: str |
| 53 | span_id: str |
| 54 | parent_span_id: str | None |
| 55 | name: str |
| 56 | start_ns: int |
| 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: |
| 64 | if value < 0: |
| 65 | value += 1 << 64 |
| 66 | encoded = bytearray() |
| 67 | while value > 0x7F: |
| 68 | encoded.append((value & 0x7F) | 0x80) |
| 69 | value >>= 7 |
| 70 | encoded.append(value) |
| 71 | return bytes(encoded) |
| 72 | |
| 73 | |
| 74 | def _field(number: int, wire_type: int) -> bytes: |
| 75 | return _varint((number << 3) | wire_type) |
| 76 | |
| 77 | |
| 78 | def _bytes_field(number: int, value: bytes) -> bytes: |
| 79 | return _field(number, 2) + _varint(len(value)) + value |
| 80 | |
| 81 | |
| 82 | def _string_field(number: int, value: str) -> bytes: |
| 83 | return _bytes_field(number, value.encode()) |
| 84 | |
| 85 | |
| 86 | def _varint_field(number: int, value: int) -> bytes: |
| 87 | return _field(number, 0) + _varint(value) |
| 88 | |
| 89 | |
| 90 | def _fixed64_field(number: int, value: int) -> bytes: |
| 91 | return _field(number, 1) + struct.pack("<Q", value) |
| 92 | |
| 93 | |
| 94 | def _attribute(name: str, value: str | int | bool) -> bytes: |
| 95 | if isinstance(value, bool): |
| 96 | encoded_value = _varint_field(2, int(value)) |
| 97 | elif isinstance(value, int): |
| 98 | encoded_value = _varint_field(3, value) |
| 99 | else: |
| 100 | encoded_value = _string_field(1, value) |
| 101 | return _string_field(1, name) + _bytes_field(2, encoded_value) |
| 102 | |
| 103 | |
| 104 | def _resource_attributes() -> dict[str, str | int | bool]: |
| 105 | build_number = os.getenv("BUILDKITE_BUILD_NUMBER", "0") |
| 106 | attributes: dict[str, str | int | bool] = { |
| 107 | "service.name": "vllm-ci", |
| 108 | "buildkite.organization.slug": os.getenv("BUILDKITE_ORGANIZATION_SLUG", ""), |
| 109 | "buildkite.pipeline.slug": os.getenv("BUILDKITE_PIPELINE_SLUG", ""), |
| 110 | "buildkite.build.id": os.getenv("BUILDKITE_BUILD_ID", ""), |
| 111 | "buildkite.build.number": int(build_number) if build_number.isdigit() else 0, |
| 112 | "buildkite.build.branch": os.getenv("BUILDKITE_BRANCH", ""), |
| 113 | "buildkite.build.commit": os.getenv("BUILDKITE_COMMIT", ""), |
| 114 | "buildkite.job.id": os.getenv("BUILDKITE_JOB_ID", ""), |
| 115 | "buildkite.job.label": os.getenv("BUILDKITE_LABEL", ""), |
| 116 | "buildkite.agent.queue": os.getenv("BUILDKITE_AGENT_META_DATA_QUEUE", ""), |
| 117 | } |
| 118 | return {name: value for name, value in attributes.items() if value != ""} |
| 119 | |
| 120 | |
| 121 | def _encode_span(span: Span) -> bytes: |
| 122 | encoded = b"".join( |
| 123 | ( |
| 124 | _bytes_field(1, bytes.fromhex(span.trace_id)), |
| 125 | _bytes_field(2, bytes.fromhex(span.span_id)), |
| 126 | _bytes_field(4, bytes.fromhex(span.parent_span_id)) |
| 127 | if span.parent_span_id |
| 128 | else b"", |
| 129 | _string_field(5, span.name), |
| 130 | _varint_field(6, 1), |
| 131 | _fixed64_field(7, span.start_ns), |
| 132 | _fixed64_field(8, span.end_ns), |
| 133 | ) |
| 134 | ) |
| 135 | attributes = dict(span.attributes) |
| 136 | if os.getenv("BUILDKITE_BUILD_URL") and os.getenv("BUILDKITE_JOB_ID"): |
| 137 | attributes["buildkite.job.web_url"] = ( |
| 138 | f"{os.environ['BUILDKITE_BUILD_URL']}#{os.environ['BUILDKITE_JOB_ID']}" |
| 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 | |
| 151 | def encode_request(spans: Iterable[Span]) -> bytes: |
| 152 | resource = b"".join( |
| 153 | _bytes_field(1, _attribute(name, value)) |
| 154 | for name, value in _resource_attributes().items() |
| 155 | ) |
| 156 | scope_spans = _bytes_field(1, _string_field(1, "vllm.ci")) |
| 157 | for span in spans: |
| 158 | scope_spans += _bytes_field(2, _encode_span(span)) |
| 159 | return _bytes_field(1, _bytes_field(1, resource) + _bytes_field(2, scope_spans)) |
| 160 | |
| 161 | |
| 162 | def new_context() -> tuple[str, str, str | None]: |
| 163 | parts = os.getenv("TRACEPARENT", "").lower().split("-") |
| 164 | valid = ( |
| 165 | len(parts) == 4 |
| 166 | and [len(part) for part in parts] == [2, 32, 16, 2] |
| 167 | and all(all(char in "0123456789abcdef" for char in part) for part in parts) |
| 168 | and parts[1] != "0" * 32 |
| 169 | and parts[2] != "0" * 16 |
| 170 | ) |
| 171 | if valid: |
| 172 | return parts[1], secrets.token_hex(8), parts[2] |
| 173 | build_id = os.getenv("BUILDKITE_BUILD_ID", "local") |
| 174 | trace_id = hashlib.sha256(build_id.encode()).hexdigest()[:32] |
| 175 | return trace_id, secrets.token_hex(8), None |
| 176 | |
| 177 | |
| 178 | def _spool_dir() -> Path | None: |
| 179 | value = os.getenv("CI_INFRA_OTEL_SPOOL_DIR") |
| 180 | return Path(value) if value else None |
| 181 | |
| 182 | |
| 183 | def record_spans(spans: Iterable[Span]) -> bool: |
| 184 | try: |
| 185 | spool_dir = _spool_dir() |
| 186 | records = [json.dumps(asdict(span), separators=(",", ":")) for span in spans] |
| 187 | if spool_dir is None or not records: |
| 188 | return False |
| 189 | spool_dir.mkdir(mode=0o700, parents=True, exist_ok=True) |
| 190 | with (spool_dir / f"spans-{os.getpid()}.jsonl").open( |
| 191 | "a", encoding="utf-8" |
| 192 | ) as output: |
| 193 | output.write("\n".join(records) + "\n") |
| 194 | return True |
| 195 | except Exception as error: |
| 196 | print(f"CI timing spool skipped: {error}", file=sys.stderr) |
| 197 | return False |
| 198 | |
| 199 | |
| 200 | def _read_spool() -> tuple[list[Span], list[Path]]: |
| 201 | spool_dir = _spool_dir() |
| 202 | if spool_dir is None or not spool_dir.is_dir(): |
| 203 | return [], [] |
| 204 | spans: list[Span] = [] |
| 205 | loaded_files: list[Path] = [] |
| 206 | for spool_file in sorted(spool_dir.glob("spans-*.jsonl")): |
| 207 | try: |
| 208 | spans.extend( |
| 209 | Span(**json.loads(record)) |
| 210 | for record in spool_file.read_text(encoding="utf-8").splitlines() |
| 211 | if record |
| 212 | ) |
| 213 | loaded_files.append(spool_file) |
| 214 | except Exception as error: |
| 215 | print( |
| 216 | f"CI timing spool ignored {spool_file.name}: {error}", file=sys.stderr |
| 217 | ) |
| 218 | return spans, loaded_files |
| 219 | |
| 220 | |
| 221 | def load_spans() -> list[Span]: |
| 222 | return _read_spool()[0] |
| 223 | |
| 224 | |
| 225 | def buildx_ref(metadata_path: str, target: str) -> str: |
| 226 | """Read the exact Buildx history reference emitted by bake.""" |
| 227 | metadata = json.loads(Path(metadata_path).read_text(encoding="utf-8")) |
| 228 | target_metadata = metadata.get(target, metadata) |
| 229 | reference = target_metadata.get("buildx.build.ref") |
| 230 | if not isinstance(reference, str) or not reference: |
| 231 | raise ValueError(f"Buildx metadata has no history reference for {target}") |
| 232 | return reference |
| 233 | |
| 234 | |
| 235 | def _buildkit_tags(span: dict) -> dict[str, str | int | bool]: |
| 236 | tags: dict[str, str | int | bool] = {} |
| 237 | for tag in span.get("tags", []): |
| 238 | key = tag.get("key") |
| 239 | value = tag.get("value") |
| 240 | if isinstance(key, str) and isinstance(value, (str, int, bool)): |
| 241 | tags[key] = value |
| 242 | return tags |
| 243 | |
| 244 | |
| 245 | def _buildkit_parent_ids(span: dict) -> list[str]: |
| 246 | return [ |
| 247 | reference["spanID"] |
| 248 | for reference in span.get("references", []) |
| 249 | if reference.get("refType") == "CHILD_OF" |
| 250 | and isinstance(reference.get("spanID"), str) |
| 251 | ] |
| 252 | |
| 253 | |
| 254 | def _buildkit_span_id(trace_id: str, value: str) -> str: |
| 255 | return hashlib.sha256(f"{trace_id}:{value}".encode()).hexdigest()[:16] |
| 256 | |
| 257 | |
| 258 | def buildkit_spans(payload: dict) -> list[Span]: |
| 259 | """Normalize Buildx's Jaeger JSON into dashboard-oriented OTLP spans.""" |
| 260 | trace_id, _, _ = new_context() |
| 261 | traces = payload.get("data", []) |
| 262 | native_trace_id = ( |
| 263 | traces[0].get("traceID", "buildkit") |
| 264 | if traces and isinstance(traces[0], dict) |
| 265 | else "buildkit" |
| 266 | ) |
| 267 | span_scope = ( |
| 268 | os.getenv("CI_INFRA_BUILDX_REF") |
| 269 | or os.getenv("BUILDKITE_JOB_ID") |
| 270 | or str(native_trace_id) |
| 271 | ) |
| 272 | steps: list[tuple[str, int, int, str, dict, int, int, int]] = [] |
| 273 | |
| 274 | for trace in traces: |
| 275 | raw_spans = [ |
| 276 | span |
| 277 | for span in trace.get("spans", []) |
| 278 | if isinstance(span, dict) and isinstance(span.get("spanID"), str) |
| 279 | ] |
| 280 | spans_by_id = {span["spanID"]: span for span in raw_spans} |
| 281 | children: dict[str, list[dict]] = {} |
| 282 | for span in raw_spans: |
| 283 | for parent_id in _buildkit_parent_ids(span): |
| 284 | if parent_id in spans_by_id: |
| 285 | children.setdefault(parent_id, []).append(span) |
| 286 | |
| 287 | def envelope( |
| 288 | span: dict, child_spans: dict[str, list[dict]] = children |
| 289 | ) -> tuple[int, int]: |
| 290 | start_us = int(span["startTime"]) |
| 291 | end_us = start_us + int(span["duration"]) |
| 292 | pending = list(child_spans.get(span["spanID"], [])) |
| 293 | seen: set[str] = set() |
| 294 | while pending: |
| 295 | child = pending.pop() |
| 296 | child_id = child["spanID"] |
| 297 | if child_id in seen: |
| 298 | continue |
| 299 | seen.add(child_id) |
| 300 | child_start = int(child["startTime"]) |
| 301 | start_us = min(start_us, child_start) |
| 302 | end_us = max(end_us, child_start + int(child["duration"])) |
| 303 | pending.extend(child_spans.get(child_id, [])) |
| 304 | return start_us * 1_000, end_us * 1_000 |
| 305 | |
| 306 | selected: dict[str, tuple[dict, int, int]] = {} |
| 307 | for span in raw_spans: |
| 308 | name = span.get("operationName") |
| 309 | if not isinstance(name, str) or name.startswith("cache request: "): |
| 310 | continue |
| 311 | step = _parse_dockerfile_step(name) |
| 312 | is_internal = name.startswith("[internal] ") |
| 313 | is_export = name.startswith(("exporting to ", "export to ")) |
| 314 | if not (step or is_internal or is_export): |
| 315 | continue |
| 316 | try: |
| 317 | start_ns, end_ns = envelope(span) |
| 318 | except (KeyError, TypeError, ValueError): |
| 319 | continue |
| 320 | vertex = _buildkit_tags(span).get("vertex") |
| 321 | key = str(vertex) if vertex else f"{name}:{span['spanID']}" |
| 322 | previous = selected.get(key) |
| 323 | if previous is None or end_ns - start_ns > previous[2] - previous[1]: |
| 324 | selected[key] = (span, start_ns, end_ns) |
| 325 | |
| 326 | for span, start_ns, end_ns in selected.values(): |
| 327 | name = span["operationName"] |
| 328 | step = _parse_dockerfile_step(name) |
| 329 | tags = _buildkit_tags(span) |
| 330 | failed = ( |
| 331 | tags.get("error") is True or tags.get("otel.status_code") == "ERROR" |
| 332 | ) |
| 333 | if step: |
| 334 | stage, step_index, step_total, label = step |
| 335 | else: |
| 336 | stage = "BuildKit internals" |
| 337 | label = name.removeprefix("[internal] ") |
| 338 | step_index = step_total = 0 |
| 339 | steps.append( |
| 340 | ( |
| 341 | stage, |
| 342 | step_index, |
| 343 | step_total, |
| 344 | label, |
| 345 | tags, |
| 346 | start_ns, |
| 347 | end_ns, |
| 348 | 2 if failed else 1, |
| 349 | ) |
| 350 | ) |
| 351 | |
| 352 | output: list[Span] = [] |
| 353 | build_ref = os.getenv("CI_INFRA_BUILDX_REF", "") |
| 354 | for stage in dict.fromkeys( |
| 355 | step[0] for step in steps if step[0] != "BuildKit internals" |
| 356 | ): |
| 357 | stage_steps = sorted( |
| 358 | (step for step in steps if step[0] == stage), key=lambda step: step[5] |
| 359 | ) |
| 360 | stage_id = _buildkit_span_id(trace_id, f"{span_scope}:stage:{stage}") |
| 361 | stage_start = min(step[5] for step in stage_steps) |
| 362 | stage_end = max(step[6] for step in stage_steps) |
| 363 | stage_failed = any(step[7] == 2 for step in stage_steps) |
| 364 | stage_attributes: dict[str, str | int | bool] = { |
| 365 | "ci.span.kind": "docker-stage", |
| 366 | "docker.stage": stage, |
| 367 | "docker.step.count": len(stage_steps), |
| 368 | } |
| 369 | if build_ref: |
| 370 | stage_attributes["docker.build.ref"] = build_ref |
| 371 | output.append( |
| 372 | Span( |
| 373 | trace_id=trace_id, |
| 374 | span_id=stage_id, |
| 375 | parent_span_id=None, |
| 376 | name="docker.stage", |
| 377 | start_ns=stage_start, |
| 378 | end_ns=stage_end, |
| 379 | attributes=stage_attributes, |
| 380 | status_code=2 if stage_failed else 1, |
| 381 | ) |
| 382 | ) |
| 383 | for ( |
| 384 | stage_name, |
| 385 | index, |
| 386 | total, |
| 387 | label, |
| 388 | tags, |
| 389 | start_ns, |
| 390 | end_ns, |
| 391 | status, |
| 392 | ) in stage_steps: |
| 393 | vertex = str(tags.get("vertex", "")) |
| 394 | attributes: dict[str, str | int | bool] = { |
| 395 | "ci.span.kind": "docker-instruction", |
| 396 | "docker.stage": stage_name, |
| 397 | "docker.step.label": label, |
| 398 | } |
| 399 | if index: |
| 400 | attributes["docker.step.index"] = index |
| 401 | attributes["docker.step.total"] = total |
| 402 | attributes["docker.instruction"] = label.split(None, 1)[0] |
| 403 | if vertex: |
| 404 | attributes["docker.vertex"] = vertex |
| 405 | output.append( |
| 406 | Span( |
| 407 | trace_id=trace_id, |
| 408 | span_id=_buildkit_span_id( |
| 409 | trace_id, |
| 410 | f"{span_scope}:step:{stage_name}:{index}:{vertex}:{label}", |
| 411 | ), |
| 412 | parent_span_id=stage_id, |
| 413 | name="docker.instruction", |
| 414 | start_ns=start_ns, |
| 415 | end_ns=end_ns, |
| 416 | attributes=attributes, |
| 417 | status_code=status, |
| 418 | ) |
| 419 | ) |
| 420 | for stage_name, index, total, label, tags, start_ns, end_ns, status in ( |
| 421 | step for step in steps if step[0] == "BuildKit internals" |
| 422 | ): |
| 423 | vertex = str(tags.get("vertex", "")) |
| 424 | attributes: dict[str, str | int | bool] = { |
| 425 | "ci.span.kind": "docker-internal", |
| 426 | "docker.stage": stage_name, |
| 427 | "docker.step.label": label, |
| 428 | } |
| 429 | if vertex: |
| 430 | attributes["docker.vertex"] = vertex |
| 431 | output.append( |
| 432 | Span( |
| 433 | trace_id=trace_id, |
| 434 | span_id=_buildkit_span_id( |
| 435 | trace_id, f"{span_scope}:internal:{vertex}:{label}" |
| 436 | ), |
| 437 | parent_span_id=None, |
| 438 | name="docker.internal", |
| 439 | start_ns=start_ns, |
| 440 | end_ns=end_ns, |
| 441 | attributes=attributes, |
| 442 | status_code=status, |
| 443 | ) |
| 444 | ) |
| 445 | return output |
| 446 | |
| 447 | |
| 448 | def record_buildkit_trace(trace_path: str) -> bool: |
| 449 | payload = json.loads(Path(trace_path).read_text(encoding="utf-8")) |
| 450 | return record_spans(buildkit_spans(payload)) |
| 451 | |
| 452 | |
| 453 | def _remaining_seconds(deadline: float) -> float: |
| 454 | remaining = deadline - time.monotonic() |
| 455 | if remaining <= 0: |
| 456 | raise TimeoutError("CI timing upload deadline exceeded") |
| 457 | return remaining |
| 458 | |
| 459 | |
| 460 | def _oidc_token(deadline: float) -> str: |
| 461 | access_token = os.getenv("BUILDKITE_AGENT_ACCESS_TOKEN", "") |
| 462 | job_id = os.getenv("BUILDKITE_JOB_ID", "") |
| 463 | if not access_token or not job_id: |
| 464 | raise RuntimeError("Buildkite job credentials are unavailable") |
| 465 | endpoint = os.getenv( |
| 466 | "BUILDKITE_AGENT_ENDPOINT", "https://agent.buildkite.com/v3" |
| 467 | ).rstrip("/") |
| 468 | request = urllib.request.Request( |
| 469 | f"{endpoint}/jobs/{job_id}/oidc/tokens", |
| 470 | data=json.dumps({"audience": AUDIENCE, "lifetime": 300}).encode(), |
| 471 | method="POST", |
| 472 | headers={ |
| 473 | "Authorization": f"Token {access_token}", |
| 474 | "Content-Type": "application/json", |
| 475 | "User-Agent": "vllm-ci-otel/1", |
| 476 | }, |
| 477 | ) |
| 478 | with urllib.request.urlopen( |
| 479 | request, timeout=_remaining_seconds(deadline) |
| 480 | ) as response: |
| 481 | token = json.loads(response.read()).get("token") |
| 482 | if not isinstance(token, str) or not token: |
| 483 | raise RuntimeError("Buildkite OIDC response did not contain a token") |
| 484 | return token |
| 485 | |
| 486 | |
| 487 | def export_spans(spans: list[Span], timeout_seconds: float = 3.0) -> bool: |
| 488 | if not spans or os.getenv("BUILDKITE") != "true": |
| 489 | return False |
| 490 | try: |
| 491 | deadline = time.monotonic() + max(timeout_seconds, 0.1) |
| 492 | token = _oidc_token(deadline) |
| 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, |
| 510 | data=encode_request(batch), |
| 511 | method="POST", |
| 512 | headers={ |
| 513 | "Authorization": f"Bearer {token}", |
| 514 | "Content-Type": "application/x-protobuf", |
| 515 | "User-Agent": "vllm-ci-otel/1", |
| 516 | }, |
| 517 | ) |
| 518 | with urllib.request.urlopen( |
| 519 | request, timeout=_remaining_seconds(deadline) |
| 520 | ) as response: |
| 521 | if response.status != 200: |
| 522 | raise RuntimeError(f"OTLP endpoint returned {response.status}") |
| 523 | return True |
| 524 | except Exception as error: |
| 525 | print(f"CI timing upload skipped: {error}", file=sys.stderr) |
| 526 | return False |
| 527 | |
| 528 | |
| 529 | def flush_spans(timeout_seconds: float = 3.0) -> bool: |
| 530 | spans, spool_files = _read_spool() |
| 531 | if not spool_files: |
| 532 | return True |
| 533 | if not spans or not export_spans(spans, timeout_seconds): |
| 534 | return False |
| 535 | for spool_file in spool_files: |
| 536 | try: |
| 537 | spool_file.unlink() |
| 538 | except FileNotFoundError: |
| 539 | pass |
| 540 | except Exception as error: |
| 541 | print(f"CI timing spool cleanup skipped: {error}", file=sys.stderr) |
| 542 | return True |
| 543 | |
| 544 | |
| 545 | _test_runs: dict[str, tuple[int, str]] = {} |
| 546 | _test_spans: list[Span] = [] |
| 547 | |
| 548 | |
| 549 | def _spool_test_spans() -> None: |
| 550 | if not _test_spans: |
| 551 | return |
| 552 | record_spans(_test_spans) |
| 553 | _test_spans.clear() |
| 554 | |
| 555 | |
| 556 | def _test_span(nodeid: str, start_ns: int, outcome: str) -> Span: |
| 557 | return Span( |
| 558 | trace_id=os.environ["CI_INFRA_TRACE_ID"], |
| 559 | span_id=secrets.token_hex(8), |
| 560 | parent_span_id=os.environ["CI_INFRA_COMMAND_SPAN_ID"], |
| 561 | name="pytest.test", |
| 562 | start_ns=start_ns, |
| 563 | end_ns=time.time_ns(), |
| 564 | attributes={ |
| 565 | "ci.span.kind": "test", |
| 566 | "test.nodeid": nodeid, |
| 567 | "test.file": nodeid.split("::", 1)[0], |
| 568 | "test.outcome": outcome, |
| 569 | }, |
| 570 | status_code=2 if outcome == "failed" else 1, |
| 571 | ) |
| 572 | |
| 573 | |
| 574 | def pytest_runtest_logstart(nodeid: str, location: tuple[str, int | None, str]): |
| 575 | try: |
| 576 | if os.getenv("CI_INFRA_TRACE_ID") and os.getenv("CI_INFRA_COMMAND_SPAN_ID"): |
| 577 | _test_runs[nodeid] = (time.time_ns(), "unknown") |
| 578 | except Exception: |
| 579 | pass |
| 580 | |
| 581 | |
| 582 | def pytest_runtest_logreport(report): |
| 583 | try: |
| 584 | if report.nodeid not in _test_runs: |
| 585 | return |
| 586 | start_ns, outcome = _test_runs[report.nodeid] |
| 587 | if report.failed: |
| 588 | outcome = "failed" |
| 589 | elif report.skipped and outcome != "failed": |
| 590 | outcome = "skipped" |
| 591 | elif report.when == "call" and outcome == "unknown": |
| 592 | outcome = "passed" |
| 593 | _test_runs[report.nodeid] = (start_ns, outcome) |
| 594 | except Exception: |
| 595 | pass |
| 596 | |
| 597 | |
| 598 | def pytest_runtest_logfinish(nodeid: str, location: tuple[str, int | None, str]): |
| 599 | try: |
| 600 | run = _test_runs.pop(nodeid, None) |
| 601 | if run: |
| 602 | _test_spans.append(_test_span(nodeid, *run)) |
| 603 | if len(_test_spans) >= TEST_SPOOL_BATCH_SIZE: |
| 604 | _spool_test_spans() |
| 605 | except Exception: |
| 606 | pass |
| 607 | |
| 608 | |
| 609 | def pytest_sessionfinish(session, exitstatus: int): |
| 610 | try: |
| 611 | for nodeid, run in list(_test_runs.items()): |
| 612 | _test_spans.append(_test_span(nodeid, *run)) |
| 613 | _test_runs.clear() |
| 614 | _spool_test_spans() |
| 615 | except Exception: |
| 616 | pass |
| 617 | |
| 618 | |
| 619 | def _record_command(arguments: list[str]) -> None: |
| 620 | trace_id, span_id, parent_span_id, start_ns, index, exit_code, label = arguments |
| 621 | exit_status = int(exit_code) |
| 622 | record_spans( |
| 623 | [ |
| 624 | Span( |
| 625 | trace_id=trace_id, |
| 626 | span_id=span_id, |
| 627 | parent_span_id=None if parent_span_id == "-" else parent_span_id, |
| 628 | name="ci.command", |
| 629 | start_ns=int(start_ns), |
| 630 | end_ns=time.time_ns(), |
| 631 | attributes={ |
| 632 | "ci.span.kind": "command", |
| 633 | "ci.command.index": int(index), |
| 634 | "ci.command.label": label, |
| 635 | "process.exit.code": exit_status, |
| 636 | }, |
| 637 | status_code=1 if exit_status == 0 else 2, |
| 638 | ) |
| 639 | ] |
| 640 | ) |
| 641 | |
| 642 | |
| 643 | def main(arguments: list[str] | None = None) -> int: |
| 644 | arguments = sys.argv[1:] if arguments is None else arguments |
| 645 | try: |
| 646 | if arguments == ["new-context"]: |
| 647 | trace_id, span_id, parent_span_id = new_context() |
| 648 | print(trace_id, span_id, parent_span_id or "-", time.time_ns()) |
| 649 | return 0 |
| 650 | if arguments == ["flush"]: |
| 651 | return 0 if flush_spans() else 1 |
| 652 | if len(arguments) == 3 and arguments[0] == "build-ref": |
| 653 | print(buildx_ref(arguments[1], arguments[2])) |
| 654 | return 0 |
| 655 | if len(arguments) == 2 and arguments[0] == "record-buildkit": |
| 656 | return 0 if record_buildkit_trace(arguments[1]) else 1 |
| 657 | if len(arguments) == 8 and arguments[0] == "record-command": |
| 658 | _record_command(arguments[1:]) |
| 659 | return 0 |
| 660 | print( |
| 661 | "usage: ci_otel.py {new-context|flush|build-ref ...|" |
| 662 | "record-buildkit ...|record-command ...}", |
| 663 | file=sys.stderr, |
| 664 | ) |
| 665 | return 2 |
| 666 | except Exception as error: |
| 667 | print(f"CI timing helper failed: {error}", file=sys.stderr) |
| 668 | return 1 |
| 669 | |
| 670 | |
| 671 | if __name__ == "__main__": |
| 672 | raise SystemExit(main()) |