chore: import upstream snapshot with attribution
PR Test (NPU) / check-changes (push) Has been cancelled
PR Test (NPU) / pr-gate (push) Has been cancelled
PR Test (NPU) / set-image-config (push) Has been cancelled
PR Test (NPU) / stage-b-test-1-npu-a2 (0) (push) Has been cancelled
PR Test (NPU) / stage-b-test-1-npu-a2 (1) (push) Has been cancelled
PR Test (NPU) / stage-b-test-2-npu-a2 (0) (push) Has been cancelled
PR Test (NPU) / stage-b-test-2-npu-a2 (1) (push) Has been cancelled
PR Test (NPU) / stage-b-test-4-npu-a3 (push) Has been cancelled
PR Test (NPU) / stage-b-test-16-npu-a3 (push) Has been cancelled
PR Test (NPU) / multimodal-gen-test-1-npu-a3 (push) Has been cancelled
PR Test (NPU) / multimodal-gen-test-2-npu-a3 (push) Has been cancelled
PR Test (Arm64) / pr-gate (push) Has been cancelled
PR Test (Arm64) / check-changes (push) Has been cancelled
PR Test (Arm64) / build-test (push) Has been cancelled
PR Test (sgl-router) / gate (push) Has been cancelled
PR Test (sgl-router) / tier-1 — lint (push) Has been cancelled
PR Test (sgl-router) / tier-2 — build + test (push) Has been cancelled
PR Test (sgl-router) / tier-3 — docker (placeholder) (push) Has been cancelled
PR Test (sgl-router) / tier-3 — k8s integration (push) Has been cancelled
PR Test (sgl-router) / tier-3 — e2e (push) Has been cancelled
PR Test (sgl-router) / finish (push) Has been cancelled
PR Test (NPU) / single-node-poc (map[name:qwen3_6_27b_w8a8_1p_in64k_out1k_50ms runner:linux-aarch64-a3-2 test_case:test/registered/ascend/performance/qwen3_6_27b/test_npu_qwen3_6_27b_w8a8_1p_in64k_out1k_50ms.py test_type:perf]) (push) Has been cancelled
PR Test (NPU) / pr-test-npu-finish (push) Has been cancelled
PR Test (Xeon) / pr-gate (push) Has been cancelled
PR Test (Xeon) / check-changes (push) Has been cancelled
PR Test (Xeon) / build-test (, xeon-gnr, base-b-test-cpu) (push) Has been cancelled
PR Test (XPU) / check-changes (push) Has been cancelled
PR Test (XPU) / pr-gate (push) Has been cancelled
PR Test (XPU) / stage-a-test-1-gpu-xpu (push) Has been cancelled
PR Test (XPU) / wait-for-stage-a (push) Has been cancelled
PR Test (XPU) / stage-b-test-1-gpu-xpu (push) Has been cancelled
PR Test (XPU) / finish (push) Has been cancelled
CI Model Inventory / build-inventory (push) Has been cancelled
Lint / lint (push) Has been cancelled
PR Benchmark (SMG Components) / Benchmark Compilation Check (push) Has been cancelled
PR Benchmark (SMG Components) / Benchmark - Manual Policy (push) Has been cancelled
PR Benchmark (SMG Components) / Benchmark - Request Processing (push) Has been cancelled
PR Benchmark (SMG Components) / Benchmark Summary (push) Has been cancelled
PR Test (SMG) / build-wheel (push) Has been cancelled
Release SGLang Model Gateway to PyPI / build on windows (x86_64 - auto) (push) Has been cancelled
Release SGLang Model Gateway to PyPI / build on macos (x86_64 - auto) (push) Has been cancelled
PR Test (SMG) / python-unit-tests (push) Has been cancelled
PR Test (SMG) / unit-tests (push) Has been cancelled
PR Test (SMG) / benchmarks (push) Has been cancelled
PR Test (SMG) / chat-completions (push) Has been cancelled
PR Test (SMG) / chat-completions-4gpu (push) Has been cancelled
PR Test (SMG) / e2e (push) Has been cancelled
PR Test (SMG) / docker-build-test (push) Has been cancelled
PR Test (SMG) / k8s-integration (push) Has been cancelled
PR Test (SMG) / finish (push) Has been cancelled
PR Test (SMG) / summarize-benchmarks (push) Has been cancelled
Release SGLang Model Gateway Docker Image / publish (push) Has been cancelled
Release SGLang Model Gateway to PyPI / build on macos (aarch64 - auto) (push) Has been cancelled
Release SGLang Model Gateway to PyPI / build on linux (aarch64 - auto) (push) Has been cancelled
Release SGLang Model Gateway to PyPI / build on linux (x86_64 - auto) (push) Has been cancelled
Release SGLang Model Gateway to PyPI / build on linux (aarch64 - musllinux_1_1) (push) Has been cancelled
Release SGLang Model Gateway to PyPI / build on linux (x86_64 - musllinux_1_1) (push) Has been cancelled
Release SGLang Model Gateway to PyPI / Build SDist (push) Has been cancelled
Release SGLang Model Gateway to PyPI / Upload to PyPI (push) Has been cancelled
Release SGLang Kernels / build-cu129-matrix (aarch64, 12.9, 3.10, arm-kernel-build-node) (push) Has been cancelled
Release SGLang Kernels / build-cu129-matrix (x86_64, 12.9, 3.10, x64-kernel-build-node) (push) Has been cancelled
Release SGLang Kernels / release-cu129 (push) Has been cancelled
Release SGLang Kernels / build-cu130-matrix (aarch64, 13.0, 3.10, arm-kernel-build-node) (push) Has been cancelled
Release SGLang Kernels / build-cu130-matrix (x86_64, 13.0, 3.10, x64-kernel-build-node) (push) Has been cancelled
Release SGLang Kernels / release-cu130 (push) Has been cancelled
Release SGLang Kernels / build-rocm-matrix (3.10, 700) (push) Has been cancelled
Release SGLang Kernels / build-rocm-matrix (3.10, 720) (push) Has been cancelled
Release SGLang Kernels / release-rocm700 (push) Has been cancelled
Release SGLang Kernels / release-rocm720 (push) Has been cancelled
Release SGLang Kernels / build-musa43 (43, 3.10) (push) Has been cancelled
Release SGLang Kernels / release-musa43 (push) Has been cancelled
PR Test (NPU) / check-changes (push) Has been cancelled
PR Test (NPU) / pr-gate (push) Has been cancelled
PR Test (NPU) / set-image-config (push) Has been cancelled
PR Test (NPU) / stage-b-test-1-npu-a2 (0) (push) Has been cancelled
PR Test (NPU) / stage-b-test-1-npu-a2 (1) (push) Has been cancelled
PR Test (NPU) / stage-b-test-2-npu-a2 (0) (push) Has been cancelled
PR Test (NPU) / stage-b-test-2-npu-a2 (1) (push) Has been cancelled
PR Test (NPU) / stage-b-test-4-npu-a3 (push) Has been cancelled
PR Test (NPU) / stage-b-test-16-npu-a3 (push) Has been cancelled
PR Test (NPU) / multimodal-gen-test-1-npu-a3 (push) Has been cancelled
PR Test (NPU) / multimodal-gen-test-2-npu-a3 (push) Has been cancelled
PR Test (Arm64) / pr-gate (push) Has been cancelled
PR Test (Arm64) / check-changes (push) Has been cancelled
PR Test (Arm64) / build-test (push) Has been cancelled
PR Test (sgl-router) / gate (push) Has been cancelled
PR Test (sgl-router) / tier-1 — lint (push) Has been cancelled
PR Test (sgl-router) / tier-2 — build + test (push) Has been cancelled
PR Test (sgl-router) / tier-3 — docker (placeholder) (push) Has been cancelled
PR Test (sgl-router) / tier-3 — k8s integration (push) Has been cancelled
PR Test (sgl-router) / tier-3 — e2e (push) Has been cancelled
PR Test (sgl-router) / finish (push) Has been cancelled
PR Test (NPU) / single-node-poc (map[name:qwen3_6_27b_w8a8_1p_in64k_out1k_50ms runner:linux-aarch64-a3-2 test_case:test/registered/ascend/performance/qwen3_6_27b/test_npu_qwen3_6_27b_w8a8_1p_in64k_out1k_50ms.py test_type:perf]) (push) Has been cancelled
PR Test (NPU) / pr-test-npu-finish (push) Has been cancelled
PR Test (Xeon) / pr-gate (push) Has been cancelled
PR Test (Xeon) / check-changes (push) Has been cancelled
PR Test (Xeon) / build-test (, xeon-gnr, base-b-test-cpu) (push) Has been cancelled
PR Test (XPU) / check-changes (push) Has been cancelled
PR Test (XPU) / pr-gate (push) Has been cancelled
PR Test (XPU) / stage-a-test-1-gpu-xpu (push) Has been cancelled
PR Test (XPU) / wait-for-stage-a (push) Has been cancelled
PR Test (XPU) / stage-b-test-1-gpu-xpu (push) Has been cancelled
PR Test (XPU) / finish (push) Has been cancelled
CI Model Inventory / build-inventory (push) Has been cancelled
Lint / lint (push) Has been cancelled
PR Benchmark (SMG Components) / Benchmark Compilation Check (push) Has been cancelled
PR Benchmark (SMG Components) / Benchmark - Manual Policy (push) Has been cancelled
PR Benchmark (SMG Components) / Benchmark - Request Processing (push) Has been cancelled
PR Benchmark (SMG Components) / Benchmark Summary (push) Has been cancelled
PR Test (SMG) / build-wheel (push) Has been cancelled
Release SGLang Model Gateway to PyPI / build on windows (x86_64 - auto) (push) Has been cancelled
Release SGLang Model Gateway to PyPI / build on macos (x86_64 - auto) (push) Has been cancelled
PR Test (SMG) / python-unit-tests (push) Has been cancelled
PR Test (SMG) / unit-tests (push) Has been cancelled
PR Test (SMG) / benchmarks (push) Has been cancelled
PR Test (SMG) / chat-completions (push) Has been cancelled
PR Test (SMG) / chat-completions-4gpu (push) Has been cancelled
PR Test (SMG) / e2e (push) Has been cancelled
PR Test (SMG) / docker-build-test (push) Has been cancelled
PR Test (SMG) / k8s-integration (push) Has been cancelled
PR Test (SMG) / finish (push) Has been cancelled
PR Test (SMG) / summarize-benchmarks (push) Has been cancelled
Release SGLang Model Gateway Docker Image / publish (push) Has been cancelled
Release SGLang Model Gateway to PyPI / build on macos (aarch64 - auto) (push) Has been cancelled
Release SGLang Model Gateway to PyPI / build on linux (aarch64 - auto) (push) Has been cancelled
Release SGLang Model Gateway to PyPI / build on linux (x86_64 - auto) (push) Has been cancelled
Release SGLang Model Gateway to PyPI / build on linux (aarch64 - musllinux_1_1) (push) Has been cancelled
Release SGLang Model Gateway to PyPI / build on linux (x86_64 - musllinux_1_1) (push) Has been cancelled
Release SGLang Model Gateway to PyPI / Build SDist (push) Has been cancelled
Release SGLang Model Gateway to PyPI / Upload to PyPI (push) Has been cancelled
Release SGLang Kernels / build-cu129-matrix (aarch64, 12.9, 3.10, arm-kernel-build-node) (push) Has been cancelled
Release SGLang Kernels / build-cu129-matrix (x86_64, 12.9, 3.10, x64-kernel-build-node) (push) Has been cancelled
Release SGLang Kernels / release-cu129 (push) Has been cancelled
Release SGLang Kernels / build-cu130-matrix (aarch64, 13.0, 3.10, arm-kernel-build-node) (push) Has been cancelled
Release SGLang Kernels / build-cu130-matrix (x86_64, 13.0, 3.10, x64-kernel-build-node) (push) Has been cancelled
Release SGLang Kernels / release-cu130 (push) Has been cancelled
Release SGLang Kernels / build-rocm-matrix (3.10, 700) (push) Has been cancelled
Release SGLang Kernels / build-rocm-matrix (3.10, 720) (push) Has been cancelled
Release SGLang Kernels / release-rocm700 (push) Has been cancelled
Release SGLang Kernels / release-rocm720 (push) Has been cancelled
Release SGLang Kernels / build-musa43 (43, 3.10) (push) Has been cancelled
Release SGLang Kernels / release-musa43 (push) Has been cancelled
This commit is contained in:
Executable
+649
@@ -0,0 +1,649 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
CI Coverage Report Generator
|
||||
|
||||
Collects all CI test registrations from test/registered/ and generates
|
||||
a coverage report organized by folder, backend, and suite.
|
||||
|
||||
Usage:
|
||||
python scripts/ci/utils/ci_coverage_report.py [--output-format markdown|json]
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import glob
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
from collections import defaultdict
|
||||
from pathlib import Path
|
||||
|
||||
# Add the ci_register module path directly to avoid heavy sglang imports
|
||||
sys.path.insert(
|
||||
0,
|
||||
str(
|
||||
Path(__file__).parent.parent.parent.parent / "python" / "sglang" / "test" / "ci"
|
||||
),
|
||||
)
|
||||
|
||||
from ci_register import CIRegistry, HWBackend, ut_parse_one_file
|
||||
|
||||
# Display order for backend tables / sections. The list is sourced from
|
||||
# HWBackend so a newly-added enum member can never be silently dropped from
|
||||
# the report -- if the assert below fires, add the new backend name here in
|
||||
# the right display slot. Order isn't alphabetical: CUDA/AMD/NPU/CPU lead
|
||||
# (highest test volume historically), then accelerators that have been
|
||||
# wired into the registry more recently (XPU, MUSA).
|
||||
BACKEND_DISPLAY_ORDER = ("CUDA", "AMD", "NPU", "CPU", "XPU", "MUSA")
|
||||
assert set(BACKEND_DISPLAY_ORDER) == {
|
||||
b.name for b in HWBackend
|
||||
}, "BACKEND_DISPLAY_ORDER is out of sync with HWBackend"
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# multimodal_gen test coverage
|
||||
#
|
||||
# multimodal_gen tests live under python/sglang/multimodal_gen/test/ and use
|
||||
# their own run_suite.py / partitioning framework, NOT the register_*_ci()
|
||||
# registry used under test/registered/. To surface them in the daily
|
||||
# overview we synthesize CIRegistry records from file paths + filename
|
||||
# tokens, using the rules below. These are heuristics derived from how the
|
||||
# multimodal_gen workflows (pr-test-{musa,npu,amd}.yml, pr-test-multimodal-
|
||||
# gen.yml) currently invoke each file -- they MAY drift if a workflow
|
||||
# stops/starts running a directory.
|
||||
# --------------------------------------------------------------------------- #
|
||||
MULTIMODAL_GEN_TEST_DIR = "python/sglang/multimodal_gen/test"
|
||||
|
||||
# Subdirectory (relative to MULTIMODAL_GEN_TEST_DIR) -> backends those files
|
||||
# run on by default. Empty string is the top-level. Files whose tokenized
|
||||
# filename contains an explicit backend marker (see below) override this.
|
||||
_MM_GEN_SUBDIR_BACKENDS = {
|
||||
# Top-level helper-style tests (e.g. test_consistency_metrics.py).
|
||||
"": ("CUDA",),
|
||||
# server/ top-level: pr-test-multimodal-gen.yml drives CUDA, pr-test-amd
|
||||
# mirrors the same suite on AMD runners.
|
||||
"server": ("CUDA", "AMD"),
|
||||
"server/musa": ("MUSA",),
|
||||
"server/ascend": ("NPU",),
|
||||
"layers": ("CUDA",),
|
||||
# unit/ are portable CPU-style unit tests. pr-test-amd now runs the `unit`
|
||||
# suite on ROCm (multimodal-gen-unit-test-amd, both 7.0.0 and 7.2.0), so
|
||||
# they are AMD-covered too, not CUDA-only.
|
||||
"unit": ("CUDA", "AMD"),
|
||||
"cli": ("CUDA",),
|
||||
"manual": ("CUDA",),
|
||||
# Standalone server/CLI single-file tests (restructured out of server/).
|
||||
# Run on CUDA CI; AMD parity for these standalone files is TBD, so
|
||||
# CUDA-only for now (previously matched no rule and were dropped entirely).
|
||||
"single_test_file": ("CUDA",),
|
||||
"single_test_file/component_accuracy": ("CUDA",),
|
||||
# Nested unit suites run only on the CUDA lane today (they are not part of
|
||||
# the AMD `unit` suite that multimodal-gen-unit-test-amd executes).
|
||||
"unit/realtime": ("CUDA",),
|
||||
"unit/sana_wm": ("CUDA",),
|
||||
"unit/progressive_resolution": ("CUDA",),
|
||||
# musa-named unit layer kernels.
|
||||
"unit/musa/layers": ("MUSA",),
|
||||
}
|
||||
|
||||
# Filenames that match `test_*.py` by convention but contain no real tests
|
||||
# (utility / fixture modules). Skipped before classification.
|
||||
_MM_GEN_HELPER_FILENAMES = frozenset({"test_utils.py"})
|
||||
|
||||
# Filename token -> backend override. Tokenization is `stem.split("_")`, so a
|
||||
# file like `test_server_1_gpu_musa.py` yields tokens
|
||||
# {test, server, 1, gpu, musa} and matches `musa`. This correctly catches
|
||||
# both `test_musa_*.py` (musa-named kernels) and `test_*_musa*.py` (musa
|
||||
# variants of generic server tests). `nightly` is detected the same way and
|
||||
# flips the nightly flag without changing the backend.
|
||||
_MM_GEN_FILENAME_BACKEND_TOKENS = {
|
||||
"musa": ("MUSA",),
|
||||
"npu": ("NPU",),
|
||||
}
|
||||
|
||||
|
||||
def collect_all_tests(registered_dir: str) -> list[CIRegistry]:
|
||||
"""Collect all CI registrations from registered directory."""
|
||||
files = glob.glob(f"{registered_dir}/**/*.py", recursive=True)
|
||||
all_tests = []
|
||||
|
||||
for file in sorted(files):
|
||||
try:
|
||||
registries, _ = ut_parse_one_file(file)
|
||||
all_tests.extend(registries)
|
||||
except Exception as e:
|
||||
print(f"Warning: Failed to parse {file}: {e}", file=sys.stderr)
|
||||
|
||||
return all_tests
|
||||
|
||||
|
||||
def collect_multimodal_gen_tests(
|
||||
mm_gen_dir: str = MULTIMODAL_GEN_TEST_DIR,
|
||||
) -> list[CIRegistry]:
|
||||
"""Synthesize CIRegistry records for multimodal_gen tests.
|
||||
|
||||
multimodal_gen doesn't use register_*_ci(); see the module-level comment
|
||||
above MULTIMODAL_GEN_TEST_DIR for the rules. Returns one record per
|
||||
(file, backend) pair, matching the convention used for registered tests
|
||||
that target multiple backends.
|
||||
"""
|
||||
mm_gen_path = Path(mm_gen_dir)
|
||||
if not mm_gen_path.is_dir():
|
||||
return []
|
||||
|
||||
# Discover test files: anything matching test_*.py anywhere under the
|
||||
# tree. The csrc/ and apps/ subtrees aren't part of the test/ directory
|
||||
# so we don't need to exclude them.
|
||||
test_files = sorted(mm_gen_path.glob("**/test_*.py"))
|
||||
records: list[CIRegistry] = []
|
||||
|
||||
for file in test_files:
|
||||
rel = file.relative_to(mm_gen_path)
|
||||
subdir = "/".join(rel.parts[:-1]) # "" for top-level
|
||||
filename_only = rel.parts[-1]
|
||||
|
||||
if filename_only in _MM_GEN_HELPER_FILENAMES:
|
||||
continue # test_utils.py and similar -- helper modules, not tests
|
||||
|
||||
# Tokenize stem to look for explicit backend / nightly markers.
|
||||
stem_tokens = set(filename_only[:-3].split("_"))
|
||||
nightly = "nightly" in stem_tokens
|
||||
|
||||
backends: tuple[str, ...] = ()
|
||||
for token, override in _MM_GEN_FILENAME_BACKEND_TOKENS.items():
|
||||
if token in stem_tokens:
|
||||
backends = override
|
||||
break
|
||||
if not backends:
|
||||
backends = _MM_GEN_SUBDIR_BACKENDS.get(subdir, ())
|
||||
|
||||
if not backends:
|
||||
print(
|
||||
f"Warning: multimodal_gen file {file} matches no backend "
|
||||
f"rule (subdir={subdir!r}, tokens={sorted(stem_tokens)}); "
|
||||
f"add a rule to _MM_GEN_SUBDIR_BACKENDS.",
|
||||
file=sys.stderr,
|
||||
)
|
||||
continue
|
||||
|
||||
for backend_name in backends:
|
||||
records.append(
|
||||
CIRegistry(
|
||||
backend=HWBackend[backend_name],
|
||||
filename=str(file),
|
||||
# mm_gen does its own per-case partitioning -- file-level
|
||||
# estimates aren't available. Surfaced as 0 in the
|
||||
# by-suite section, which is correct (we don't know).
|
||||
est_time=0.0,
|
||||
suite=f"mm-gen-{backend_name.lower()}",
|
||||
nightly=nightly,
|
||||
disabled=None,
|
||||
)
|
||||
)
|
||||
|
||||
return records
|
||||
|
||||
|
||||
def get_folder_name(filename: str) -> str:
|
||||
"""Extract folder name from test filename.
|
||||
|
||||
Registered tests use test/registered/<folder>/test_*.py and map to
|
||||
<folder>. multimodal_gen tests live outside that tree and get a virtual
|
||||
`mm_gen/<subdir>` folder so they show up as their own rows in the
|
||||
Folder Summary table.
|
||||
"""
|
||||
parts = Path(filename).parts
|
||||
if "multimodal_gen" in parts:
|
||||
try:
|
||||
mg_idx = parts.index("multimodal_gen")
|
||||
test_idx = parts.index("test", mg_idx)
|
||||
sub_parts = parts[test_idx + 1 : -1] # strip 'test' anchor + file
|
||||
if sub_parts:
|
||||
return "mm_gen/" + "/".join(sub_parts)
|
||||
return "mm_gen"
|
||||
except ValueError:
|
||||
pass # fall through to default
|
||||
if "registered" in parts:
|
||||
idx = parts.index("registered")
|
||||
if idx + 1 < len(parts) - 1: # Has subfolder
|
||||
return parts[idx + 1]
|
||||
return "root"
|
||||
|
||||
|
||||
def get_test_basename(filename: str) -> str:
|
||||
"""Extract just the test file name from the path."""
|
||||
return Path(filename).name
|
||||
|
||||
|
||||
def organize_test_data(tests: list[CIRegistry]) -> dict:
|
||||
"""Organize tests into various groupings."""
|
||||
by_backend = defaultdict(list)
|
||||
by_folder = defaultdict(list)
|
||||
disabled_tests = []
|
||||
|
||||
for t in tests:
|
||||
by_backend[t.backend.name].append(t)
|
||||
by_folder[get_folder_name(t.filename)].append(t)
|
||||
if t.disabled:
|
||||
disabled_tests.append(t)
|
||||
|
||||
# Count unique test files (a file may be registered for multiple backends)
|
||||
unique_files = set(t.filename for t in tests)
|
||||
unique_enabled_files = set(t.filename for t in tests if not t.disabled)
|
||||
unique_disabled_files = set(t.filename for t in tests if t.disabled)
|
||||
|
||||
return {
|
||||
"total": len(tests),
|
||||
"total_unique_files": len(unique_files),
|
||||
"enabled": len(tests) - len(disabled_tests),
|
||||
"enabled_unique_files": len(unique_enabled_files),
|
||||
"disabled_count": len(disabled_tests),
|
||||
"disabled_unique_files": len(unique_disabled_files),
|
||||
"by_backend": by_backend,
|
||||
"by_folder": by_folder,
|
||||
"disabled_tests": disabled_tests,
|
||||
}
|
||||
|
||||
|
||||
def generate_summary_section(data: dict) -> str:
|
||||
"""Generate the summary/overview section."""
|
||||
lines = []
|
||||
lines.append("# CI Coverage Overview\n")
|
||||
lines.append(
|
||||
f"**Unique Test Files:** {data['total_unique_files']} ({data['enabled_unique_files']} enabled, {data['disabled_unique_files']} disabled)\n"
|
||||
)
|
||||
lines.append(
|
||||
f"**Total Registrations:** {data['total']} ({data['enabled']} enabled, {data['disabled_count']} disabled)\n"
|
||||
)
|
||||
lines.append(
|
||||
"*Note: A test file may be registered for multiple backends (e.g., CUDA + AMD), so total registrations > unique files.*\n"
|
||||
)
|
||||
|
||||
by_backend = data["by_backend"]
|
||||
by_folder = data["by_folder"]
|
||||
disabled_tests = data["disabled_tests"]
|
||||
|
||||
# Backend summary (collapsible)
|
||||
lines.append("<details>")
|
||||
lines.append("<summary><h2>Backend Summary</h2></summary>\n")
|
||||
lines.append("| Backend | Total | Enabled | Disabled | Per-Commit | Nightly |")
|
||||
lines.append("|---------|-------|---------|----------|------------|---------|")
|
||||
|
||||
for backend in BACKEND_DISPLAY_ORDER:
|
||||
backend_tests = by_backend.get(backend, [])
|
||||
if not backend_tests:
|
||||
continue
|
||||
b_total = len(backend_tests)
|
||||
b_disabled = sum(1 for t in backend_tests if t.disabled)
|
||||
b_enabled = b_total - b_disabled
|
||||
b_per_commit = sum(1 for t in backend_tests if not t.nightly and not t.disabled)
|
||||
b_nightly = sum(1 for t in backend_tests if t.nightly and not t.disabled)
|
||||
lines.append(
|
||||
f"| {backend} | {b_total} | {b_enabled} | {b_disabled} | {b_per_commit} | {b_nightly} |"
|
||||
)
|
||||
|
||||
lines.append("\n</details>\n")
|
||||
|
||||
# Folder summary (collapsible). Only show columns for backends that
|
||||
# have at least one registered test across the whole report -- otherwise
|
||||
# adding scaffolding for an unused backend would widen every row with a
|
||||
# column of zeros.
|
||||
active_backends = [b for b in BACKEND_DISPLAY_ORDER if by_backend.get(b)]
|
||||
lines.append("<details>")
|
||||
lines.append("<summary><h2>Folder Summary</h2></summary>\n")
|
||||
header_cells = ["Folder", *active_backends, "Total"]
|
||||
lines.append("| " + " | ".join(header_cells) + " |")
|
||||
lines.append("|" + "|".join(["-" * max(len(c), 3) for c in header_cells]) + "|")
|
||||
|
||||
for folder in sorted(by_folder.keys()):
|
||||
folder_tests = by_folder[folder]
|
||||
backend_counts = {b.name: 0 for b in HWBackend}
|
||||
for t in folder_tests:
|
||||
backend_counts[t.backend.name] += 1
|
||||
row = [folder] + [str(backend_counts[b]) for b in active_backends]
|
||||
row.append(str(len(folder_tests)))
|
||||
lines.append("| " + " | ".join(row) + " |")
|
||||
|
||||
lines.append("\n</details>\n")
|
||||
|
||||
# Disabled tests section (collapsible)
|
||||
if disabled_tests:
|
||||
lines.append("<details>")
|
||||
lines.append("<summary><h2>Disabled Tests</h2></summary>\n")
|
||||
lines.append("| File | Backend | Suite | Reason |")
|
||||
lines.append("|------|---------|-------|--------|")
|
||||
for t in sorted(disabled_tests, key=lambda x: (x.backend.name, x.filename)):
|
||||
test_name = get_test_basename(t.filename)
|
||||
reason = t.disabled[:50] + "..." if len(t.disabled) > 50 else t.disabled
|
||||
lines.append(
|
||||
f"| `{test_name}` | {t.backend.name} | {t.effective_suite} | {reason} |"
|
||||
)
|
||||
lines.append("\n</details>\n")
|
||||
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
def generate_by_folder_section(data: dict) -> str:
|
||||
"""Generate the 'All Tests by Folder' section."""
|
||||
lines = []
|
||||
by_folder = data["by_folder"]
|
||||
|
||||
lines.append("# All Tests by Folder\n")
|
||||
|
||||
for folder in sorted(by_folder.keys()):
|
||||
folder_tests = by_folder[folder]
|
||||
lines.append("<details>")
|
||||
lines.append(
|
||||
f"<summary><h2>{folder}/ ({len(folder_tests)} tests)</h2></summary>\n"
|
||||
)
|
||||
|
||||
# Group by backend within folder
|
||||
folder_by_backend = defaultdict(list)
|
||||
for t in folder_tests:
|
||||
folder_by_backend[t.backend.name].append(t)
|
||||
|
||||
for backend in BACKEND_DISPLAY_ORDER:
|
||||
backend_tests = folder_by_backend.get(backend, [])
|
||||
if not backend_tests:
|
||||
continue
|
||||
|
||||
lines.append(f"### {backend} ({len(backend_tests)} tests)\n")
|
||||
lines.append("| Test File | Suite | Est. Time | Status |")
|
||||
lines.append("|-----------|-------|-----------|--------|")
|
||||
|
||||
for t in sorted(backend_tests, key=lambda x: x.filename):
|
||||
test_name = get_test_basename(t.filename)
|
||||
status = (
|
||||
"Disabled"
|
||||
if t.disabled
|
||||
else ("Nightly" if t.nightly else "Per-Commit")
|
||||
)
|
||||
lines.append(
|
||||
f"| `{test_name}` | {t.effective_suite} | {t.est_time:.0f}s | {status} |"
|
||||
)
|
||||
|
||||
lines.append("")
|
||||
|
||||
lines.append("</details>\n")
|
||||
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
def generate_by_suite_section(data: dict) -> str:
|
||||
"""Generate the 'All Tests by Test Suite' section."""
|
||||
lines = []
|
||||
by_backend = data["by_backend"]
|
||||
|
||||
lines.append("# All Tests by Test Suite\n")
|
||||
|
||||
for backend in BACKEND_DISPLAY_ORDER:
|
||||
backend_tests = by_backend.get(backend, [])
|
||||
if not backend_tests:
|
||||
continue
|
||||
|
||||
b_total = len(backend_tests)
|
||||
b_disabled = sum(1 for t in backend_tests if t.disabled)
|
||||
b_enabled = b_total - b_disabled
|
||||
|
||||
lines.append("<details>")
|
||||
lines.append(
|
||||
f"<summary><h2>{backend} Backend ({b_enabled} enabled, {b_disabled} disabled)</h2></summary>\n"
|
||||
)
|
||||
|
||||
# Group by suite within backend
|
||||
backend_suites = defaultdict(list)
|
||||
for t in backend_tests:
|
||||
backend_suites[t.effective_suite].append(t)
|
||||
|
||||
for suite in sorted(backend_suites.keys()):
|
||||
suite_tests = backend_suites[suite]
|
||||
s_enabled = sum(1 for t in suite_tests if not t.disabled)
|
||||
s_disabled = sum(1 for t in suite_tests if t.disabled)
|
||||
s_est_time = sum(t.est_time for t in suite_tests if not t.disabled)
|
||||
is_nightly = any(t.nightly for t in suite_tests if not t.disabled)
|
||||
|
||||
suite_type = "Nightly" if is_nightly else "Per-Commit"
|
||||
lines.append("<details>")
|
||||
lines.append(
|
||||
f"<summary><h3>{suite} ({s_enabled} enabled, {s_disabled} disabled) - {suite_type}</h3></summary>\n"
|
||||
)
|
||||
lines.append(f"*Estimated total time: {s_est_time:.0f}s*\n")
|
||||
|
||||
lines.append("| Test File | Folder | Est. Time | Status |")
|
||||
lines.append("|-----------|--------|-----------|--------|")
|
||||
|
||||
for t in sorted(suite_tests, key=lambda x: x.filename):
|
||||
test_name = get_test_basename(t.filename)
|
||||
folder = get_folder_name(t.filename)
|
||||
if t.disabled:
|
||||
status = (
|
||||
f"Disabled: {t.disabled[:30]}..."
|
||||
if len(t.disabled) > 30
|
||||
else f"Disabled: {t.disabled}"
|
||||
)
|
||||
else:
|
||||
status = "Nightly" if t.nightly else "Per-Commit"
|
||||
lines.append(
|
||||
f"| `{test_name}` | {folder} | {t.est_time:.0f}s | {status} |"
|
||||
)
|
||||
|
||||
lines.append("\n</details>\n")
|
||||
|
||||
lines.append("</details>\n")
|
||||
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
def generate_markdown_report(tests: list[CIRegistry], section: str = "all") -> str:
|
||||
"""Generate markdown report for GitHub step summary."""
|
||||
data = organize_test_data(tests)
|
||||
|
||||
if section == "summary":
|
||||
return generate_summary_section(data)
|
||||
elif section == "by-folder":
|
||||
return generate_by_folder_section(data)
|
||||
elif section == "by-suite":
|
||||
return generate_by_suite_section(data)
|
||||
else: # "all"
|
||||
parts = [
|
||||
generate_summary_section(data),
|
||||
"---",
|
||||
generate_by_folder_section(data),
|
||||
"---",
|
||||
generate_by_suite_section(data),
|
||||
]
|
||||
return "\n".join(parts)
|
||||
|
||||
|
||||
def generate_json_report(tests: list[CIRegistry]) -> str:
|
||||
"""Generate JSON report with detailed test listings."""
|
||||
by_backend = defaultdict(list)
|
||||
by_folder = defaultdict(list)
|
||||
|
||||
for t in tests:
|
||||
by_backend[t.backend.name].append(t)
|
||||
by_folder[get_folder_name(t.filename)].append(t)
|
||||
|
||||
disabled_tests = [t for t in tests if t.disabled]
|
||||
|
||||
# Build structured data
|
||||
data = {
|
||||
"summary": {
|
||||
"total": len(tests),
|
||||
"enabled": len(tests) - len(disabled_tests),
|
||||
"disabled": len(disabled_tests),
|
||||
},
|
||||
"tests_by_folder": {},
|
||||
"tests_by_suite": {},
|
||||
"backend_summary": {},
|
||||
"folder_summary": {},
|
||||
"disabled_tests": [],
|
||||
}
|
||||
|
||||
# Section 1: Tests by Folder
|
||||
for folder in sorted(by_folder.keys()):
|
||||
folder_tests = by_folder[folder]
|
||||
folder_by_backend = defaultdict(list)
|
||||
for t in folder_tests:
|
||||
folder_by_backend[t.backend.name].append(t)
|
||||
|
||||
data["tests_by_folder"][folder] = {
|
||||
"total": len(folder_tests),
|
||||
"backends": {},
|
||||
}
|
||||
|
||||
for backend in BACKEND_DISPLAY_ORDER:
|
||||
backend_tests = folder_by_backend.get(backend, [])
|
||||
if backend_tests:
|
||||
data["tests_by_folder"][folder]["backends"][backend] = [
|
||||
{
|
||||
"filename": get_test_basename(t.filename),
|
||||
"suite": t.effective_suite,
|
||||
"est_time": t.est_time,
|
||||
"status": (
|
||||
"disabled"
|
||||
if t.disabled
|
||||
else ("nightly" if t.nightly else "per-commit")
|
||||
),
|
||||
}
|
||||
for t in sorted(backend_tests, key=lambda x: x.filename)
|
||||
]
|
||||
|
||||
# Section 2: Tests by Suite (Backend -> Suite)
|
||||
for backend in BACKEND_DISPLAY_ORDER:
|
||||
backend_tests = by_backend.get(backend, [])
|
||||
if not backend_tests:
|
||||
continue
|
||||
|
||||
backend_suites = defaultdict(list)
|
||||
for t in backend_tests:
|
||||
backend_suites[t.effective_suite].append(t)
|
||||
|
||||
data["tests_by_suite"][backend] = {
|
||||
"total": len(backend_tests),
|
||||
"enabled": sum(1 for t in backend_tests if not t.disabled),
|
||||
"disabled": sum(1 for t in backend_tests if t.disabled),
|
||||
"suites": {},
|
||||
}
|
||||
|
||||
for suite in sorted(backend_suites.keys()):
|
||||
suite_tests = backend_suites[suite]
|
||||
is_nightly = any(t.nightly for t in suite_tests if not t.disabled)
|
||||
|
||||
data["tests_by_suite"][backend]["suites"][suite] = {
|
||||
"total": len(suite_tests),
|
||||
"enabled": sum(1 for t in suite_tests if not t.disabled),
|
||||
"disabled": sum(1 for t in suite_tests if t.disabled),
|
||||
"est_time": sum(t.est_time for t in suite_tests if not t.disabled),
|
||||
"type": "nightly" if is_nightly else "per-commit",
|
||||
"tests": [
|
||||
{
|
||||
"filename": get_test_basename(t.filename),
|
||||
"folder": get_folder_name(t.filename),
|
||||
"est_time": t.est_time,
|
||||
"status": (
|
||||
"disabled"
|
||||
if t.disabled
|
||||
else ("nightly" if t.nightly else "per-commit")
|
||||
),
|
||||
"disabled_reason": t.disabled if t.disabled else None,
|
||||
}
|
||||
for t in sorted(suite_tests, key=lambda x: x.filename)
|
||||
],
|
||||
}
|
||||
|
||||
# Backend summary
|
||||
for backend in BACKEND_DISPLAY_ORDER:
|
||||
backend_tests = by_backend.get(backend, [])
|
||||
if backend_tests:
|
||||
data["backend_summary"][backend] = {
|
||||
"total": len(backend_tests),
|
||||
"enabled": sum(1 for t in backend_tests if not t.disabled),
|
||||
"disabled": sum(1 for t in backend_tests if t.disabled),
|
||||
"per_commit": sum(
|
||||
1 for t in backend_tests if not t.nightly and not t.disabled
|
||||
),
|
||||
"nightly": sum(
|
||||
1 for t in backend_tests if t.nightly and not t.disabled
|
||||
),
|
||||
}
|
||||
|
||||
# Folder summary -- one count per backend in HWBackend, in display
|
||||
# order, so every registered backend shows up regardless of whether
|
||||
# this folder has tests for it.
|
||||
for folder in sorted(by_folder.keys()):
|
||||
folder_tests = by_folder[folder]
|
||||
backend_counts = {b: 0 for b in BACKEND_DISPLAY_ORDER}
|
||||
for t in folder_tests:
|
||||
backend_counts[t.backend.name] += 1
|
||||
data["folder_summary"][folder] = {
|
||||
**backend_counts,
|
||||
"total": len(folder_tests),
|
||||
}
|
||||
|
||||
# Disabled tests
|
||||
for t in sorted(disabled_tests, key=lambda x: (x.backend.name, x.filename)):
|
||||
data["disabled_tests"].append(
|
||||
{
|
||||
"filename": get_test_basename(t.filename),
|
||||
"backend": t.backend.name,
|
||||
"suite": t.effective_suite,
|
||||
"reason": t.disabled,
|
||||
}
|
||||
)
|
||||
|
||||
return json.dumps(data, indent=2)
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Generate CI coverage report")
|
||||
parser.add_argument(
|
||||
"--output-format",
|
||||
choices=["markdown", "json"],
|
||||
default="markdown",
|
||||
help="Output format (default: markdown)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--section",
|
||||
choices=["all", "summary", "by-folder", "by-suite"],
|
||||
default="all",
|
||||
help="Which section to output (default: all). Only applies to markdown format.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--registered-dir",
|
||||
default="test/registered",
|
||||
help="Path to registered test directory",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--multimodal-gen-dir",
|
||||
default=MULTIMODAL_GEN_TEST_DIR,
|
||||
help=(
|
||||
"Path to multimodal_gen test directory. Pass empty string to "
|
||||
"skip multimodal_gen tests entirely."
|
||||
),
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
# Change to repo root if needed
|
||||
script_dir = Path(__file__).parent.parent
|
||||
repo_root = script_dir.parent.parent
|
||||
os.chdir(repo_root)
|
||||
|
||||
tests = collect_all_tests(args.registered_dir)
|
||||
if args.multimodal_gen_dir:
|
||||
tests.extend(collect_multimodal_gen_tests(args.multimodal_gen_dir))
|
||||
|
||||
if args.output_format == "markdown":
|
||||
report = generate_markdown_report(tests, section=args.section)
|
||||
else:
|
||||
report = generate_json_report(tests)
|
||||
|
||||
print(report)
|
||||
|
||||
# Write to GITHUB_STEP_SUMMARY if available
|
||||
summary_file = os.environ.get("GITHUB_STEP_SUMMARY")
|
||||
if summary_file and args.output_format == "markdown":
|
||||
with open(summary_file, "a") as f:
|
||||
f.write(report)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Executable
+146
@@ -0,0 +1,146 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Clean up stale HuggingFace cache artifacts from previous failed downloads.
|
||||
|
||||
This script removes incomplete marker files, temporary files, and lock files
|
||||
from the HuggingFace cache directory. These artifacts can accumulate from
|
||||
interrupted or failed downloads and may interfere with future downloads.
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import List
|
||||
|
||||
try:
|
||||
from huggingface_hub import constants
|
||||
|
||||
HF_HUB_AVAILABLE = True
|
||||
except ImportError:
|
||||
print("Warning: huggingface_hub not available")
|
||||
HF_HUB_AVAILABLE = False
|
||||
|
||||
|
||||
def get_hf_cache_dir() -> str:
|
||||
"""Get the HuggingFace cache directory."""
|
||||
if HF_HUB_AVAILABLE:
|
||||
return constants.HF_HUB_CACHE
|
||||
|
||||
# Fallback to environment variable or default
|
||||
hf_home = os.environ.get("HF_HOME", os.path.expanduser("~/.cache/huggingface"))
|
||||
return os.path.join(hf_home, "hub")
|
||||
|
||||
|
||||
def find_stale_artifacts(cache_dir: str) -> List[Path]:
|
||||
"""
|
||||
Find stale artifact files in the HuggingFace cache.
|
||||
|
||||
Args:
|
||||
cache_dir: HuggingFace cache directory
|
||||
|
||||
Returns:
|
||||
List of paths to stale artifact files
|
||||
"""
|
||||
cache_path = Path(cache_dir)
|
||||
|
||||
if not cache_path.exists():
|
||||
return []
|
||||
|
||||
# Patterns for stale files to clean up
|
||||
patterns = [
|
||||
"**/*.incomplete", # Incomplete download markers
|
||||
"**/*.tmp", # Temporary files
|
||||
"**/*.lock", # Lock files from interrupted downloads
|
||||
]
|
||||
|
||||
stale_files = []
|
||||
for pattern in patterns:
|
||||
stale_files.extend(cache_path.glob(pattern))
|
||||
|
||||
return stale_files
|
||||
|
||||
|
||||
def cleanup_artifacts(artifacts: List[Path]) -> tuple[int, int]:
|
||||
"""
|
||||
Remove stale artifact files.
|
||||
|
||||
Args:
|
||||
artifacts: List of file paths to remove
|
||||
|
||||
Returns:
|
||||
Tuple of (successful_removals, failed_removals)
|
||||
"""
|
||||
successful = 0
|
||||
failed = 0
|
||||
|
||||
for file_path in artifacts:
|
||||
try:
|
||||
file_path.unlink()
|
||||
print(f" Removed: {file_path}")
|
||||
successful += 1
|
||||
except Exception as e:
|
||||
print(f" Warning: Could not remove {file_path}: {e}")
|
||||
failed += 1
|
||||
|
||||
return successful, failed
|
||||
|
||||
|
||||
def main() -> int:
|
||||
"""
|
||||
Main cleanup logic.
|
||||
|
||||
Returns:
|
||||
Always returns 0 (cleanup is best-effort and should not fail CI)
|
||||
"""
|
||||
print("=" * 70)
|
||||
print("HuggingFace Cache Cleanup")
|
||||
print("=" * 70)
|
||||
|
||||
# Get cache directory
|
||||
cache_dir = get_hf_cache_dir()
|
||||
print(f"Cache directory: {cache_dir}")
|
||||
|
||||
if not os.path.exists(cache_dir):
|
||||
print("Cache directory does not exist - nothing to clean")
|
||||
return 0
|
||||
|
||||
print("-" * 70)
|
||||
|
||||
# Find stale artifacts
|
||||
print("Scanning for stale artifacts...")
|
||||
stale_artifacts = find_stale_artifacts(cache_dir)
|
||||
|
||||
if not stale_artifacts:
|
||||
print("✓ No stale cache artifacts found")
|
||||
return 0
|
||||
|
||||
# Clean up artifacts
|
||||
print(f"Found {len(stale_artifacts)} stale artifact(s) to remove:")
|
||||
successful, failed = cleanup_artifacts(stale_artifacts)
|
||||
|
||||
print("-" * 70)
|
||||
|
||||
# Summary
|
||||
if failed > 0:
|
||||
print(f"⚠ Cleaned up {successful} file(s), {failed} removal(s) failed")
|
||||
else:
|
||||
print(f"✓ Successfully cleaned up {successful} stale file(s)")
|
||||
|
||||
# Always return 0 - cleanup failures should not fail CI
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
try:
|
||||
exit_code = main()
|
||||
sys.exit(exit_code)
|
||||
except KeyboardInterrupt:
|
||||
print("\nInterrupted by user")
|
||||
sys.exit(0)
|
||||
except Exception as e:
|
||||
print(f"ERROR: Unexpected error during cleanup: {e}")
|
||||
import traceback
|
||||
|
||||
traceback.print_exc()
|
||||
# Still return 0 - cleanup failures should not fail CI
|
||||
sys.exit(0)
|
||||
@@ -0,0 +1,267 @@
|
||||
"""Sum est_time per per-commit suite and emit one $GITHUB_OUTPUT line
|
||||
keyed by suite name. Consumed by pr-test.yml stage jobs as
|
||||
`fromJson(needs.check-changes.outputs.partitions)['<suite>']`.
|
||||
|
||||
partitions={"base-b-test-1-gpu-small": {"size": 8, "arr": [0,...,7], "max_parallel": 2}, ...}
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import glob
|
||||
import importlib.util
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
from collections import defaultdict
|
||||
from datetime import datetime, timedelta
|
||||
|
||||
import yaml # PyYAML; preinstalled on ubuntu-latest GHA runners.
|
||||
|
||||
REPO_ROOT = os.path.dirname(
|
||||
os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
)
|
||||
|
||||
# Load ci_register.py directly: `import sglang.test...` pulls torch/orjson via
|
||||
# sglang.__init__ but check-changes runs on bare ubuntu-latest. ci_register
|
||||
# itself is stdlib-only (AST).
|
||||
_CI_REGISTER_PATH = os.path.join(
|
||||
REPO_ROOT, "python", "sglang", "test", "ci", "ci_register.py"
|
||||
)
|
||||
_spec = importlib.util.spec_from_file_location("ci_register", _CI_REGISTER_PATH)
|
||||
_ci_register = importlib.util.module_from_spec(_spec)
|
||||
_spec.loader.exec_module(_ci_register)
|
||||
collect_tests = _ci_register.collect_tests
|
||||
HWBackend = _ci_register.HWBackend
|
||||
|
||||
# pr-test-amd.yml / pr-test-npu.yml have their own dispatch.
|
||||
_TARGET_BACKENDS = {HWBackend.CUDA, HWBackend.CPU}
|
||||
|
||||
# base-a is the critical-path entry gate; pin its fanout to sanity-coverage
|
||||
# defaults instead of est_time. max_parallel = size (no throttle).
|
||||
_BASE_A_OVERRIDES = {
|
||||
"base-a-test-cpu": 8,
|
||||
"base-a-test-1-gpu-small": 1,
|
||||
}
|
||||
|
||||
_REUSABLE_STAGE_USES = "./.github/workflows/_pr-test-stage.yml"
|
||||
|
||||
|
||||
def load_run_timeouts(pr_test_yml_path: str) -> dict:
|
||||
"""Map `self_name -> run_timeout_minutes` from one pr-test*.yml. The input
|
||||
is required in `_pr-test-stage.yml` -- KeyError surfaces missing.
|
||||
Inline base-a-test-cpu is skipped (uses `_BASE_A_OVERRIDES`)."""
|
||||
with open(pr_test_yml_path) as f:
|
||||
wf = yaml.safe_load(f)
|
||||
timeouts = {}
|
||||
for job_id, job in (wf.get("jobs") or {}).items():
|
||||
if not isinstance(job, dict) or job.get("uses") != _REUSABLE_STAGE_USES:
|
||||
continue
|
||||
with_ = job.get("with") or {}
|
||||
suite = with_.get("self_name", job_id)
|
||||
timeouts[suite] = int(with_["run_timeout_minutes"])
|
||||
if not timeouts:
|
||||
raise RuntimeError(
|
||||
f"load_run_timeouts: no jobs matched uses={_REUSABLE_STAGE_USES!r} "
|
||||
f"in {pr_test_yml_path}. The reusable workflow path likely "
|
||||
"changed -- update _REUSABLE_STAGE_USES."
|
||||
)
|
||||
return timeouts
|
||||
|
||||
|
||||
def per_shard_target_seconds(suite: str, run_timeouts: dict) -> float:
|
||||
"""Per-shard wall budget = 0.75 * stage timeout. 0.75 is the inverse
|
||||
of LPT's 4/3 worst-case approximation ratio, so the most imbalanced
|
||||
LPT shard fills exactly the timeout."""
|
||||
return 0.75 * run_timeouts[suite] * 60
|
||||
|
||||
|
||||
def discover_files(repo_root: str) -> list[str]:
|
||||
test_dir = os.path.join(repo_root, "test")
|
||||
files = [
|
||||
f
|
||||
for f in glob.glob(
|
||||
os.path.join(test_dir, "registered", "**", "*.py"), recursive=True
|
||||
)
|
||||
if not f.endswith("/conftest.py") and not f.endswith("/__init__.py")
|
||||
]
|
||||
jit_kernel_dir = os.path.join(repo_root, "python", "sglang", "jit_kernel")
|
||||
files += glob.glob(
|
||||
os.path.join(jit_kernel_dir, "tests", "**", "test_*.py"), recursive=True
|
||||
)
|
||||
files += glob.glob(
|
||||
os.path.join(jit_kernel_dir, "benchmark", "**", "bench_*.py"), recursive=True
|
||||
)
|
||||
return files
|
||||
|
||||
|
||||
def load_partition_model(path):
|
||||
"""Read sglang-ci-stats' model.json; None on missing/unparsable.
|
||||
Cross-repo schema -- guard against non-dict top-level."""
|
||||
if not path or not os.path.exists(path):
|
||||
return None
|
||||
try:
|
||||
with open(path) as f:
|
||||
data = json.load(f)
|
||||
except (OSError, json.JSONDecodeError):
|
||||
return None
|
||||
return data if isinstance(data, dict) else None
|
||||
|
||||
|
||||
def compute_max_parallel(size: int) -> int:
|
||||
return max(size // 3, 1)
|
||||
|
||||
|
||||
def compute_partitions(
|
||||
tests, repo_root, run_timeouts, partition_model=None, full_parallel=False
|
||||
):
|
||||
"""Group per-commit tests by suite and emit partition metadata.
|
||||
|
||||
`run_timeouts`: `suite -> minutes` from `load_run_timeouts`.
|
||||
`partition_model`: optional sglang-ci-stats `model.json`; per-file
|
||||
`est` and per-suite `(coeff, bias)` each fall back independently to
|
||||
in-source `est_time` / `(1.0, 0.0)`.
|
||||
`full_parallel=True` lifts the matrix-fanout throttle.
|
||||
"""
|
||||
# Allowlist: stages pr-test.yml dispatches. Stress / weekly /
|
||||
# nightly-* live in test/registered/ but pr-test doesn't run them.
|
||||
dispatched_suites = set(run_timeouts) | set(_BASE_A_OVERRIDES)
|
||||
suite_tests = defaultdict(list)
|
||||
for t in tests:
|
||||
if t.backend not in _TARGET_BACKENDS:
|
||||
continue
|
||||
if t.nightly or t.disabled is not None:
|
||||
continue
|
||||
if t.effective_suite not in dispatched_suites:
|
||||
continue
|
||||
suite_tests[t.effective_suite].append(t)
|
||||
|
||||
est_table = (partition_model or {}).get("est", {})
|
||||
fit_table = (partition_model or {}).get("fit", {})
|
||||
|
||||
result = {}
|
||||
for suite, group in suite_tests.items():
|
||||
live_est = est_table.get(suite, {})
|
||||
total = 0.0
|
||||
for t in group:
|
||||
relpath = os.path.relpath(t.filename, repo_root)
|
||||
total += live_est.get(relpath, t.est_time)
|
||||
|
||||
fit = fit_table.get(suite) or {}
|
||||
coeff = fit.get("coeff", 1.0)
|
||||
bias = fit.get("bias", 0.0)
|
||||
|
||||
# Each shard pays `bias` once, so size >= coeff*total / (target-bias).
|
||||
if suite in _BASE_A_OVERRIDES:
|
||||
size = _BASE_A_OVERRIDES[suite]
|
||||
max_parallel = size
|
||||
else:
|
||||
target = per_shard_target_seconds(suite, run_timeouts)
|
||||
budget = target - bias
|
||||
if budget <= 0:
|
||||
raise RuntimeError(
|
||||
f"Suite {suite!r}: fit bias={bias}s >= target={target}s. "
|
||||
"Investigate the fit or raise the stage's run_timeout_minutes."
|
||||
)
|
||||
ideal_size = math.ceil(coeff * total / budget)
|
||||
# ideal_size > len(group) -> slowest single file alone exceeds
|
||||
# the per-shard budget; surface via raise instead of empty shard.
|
||||
if ideal_size > len(group):
|
||||
raise RuntimeError(
|
||||
f"Suite {suite!r}: needs {ideal_size} shards but has only "
|
||||
f"{len(group)} test file(s). target={target:.0f}s, "
|
||||
f"coeff={coeff}, bias={bias}s, total_est={total:.0f}s."
|
||||
)
|
||||
size = max(1, ideal_size)
|
||||
max_parallel = size if full_parallel else compute_max_parallel(size)
|
||||
result[suite] = {
|
||||
"size": size,
|
||||
"arr": list(range(size)),
|
||||
"max_parallel": max_parallel,
|
||||
}
|
||||
return result
|
||||
|
||||
|
||||
def format_fit_window(model: dict) -> str:
|
||||
"""Render the model's fit window as `[start, end)` for the step summary.
|
||||
|
||||
Surfacing the whole span keeps the lower-bound date from reading as a
|
||||
staleness marker (it trails today by fit_window_days)."""
|
||||
start = model.get("fit_window_start")
|
||||
days = model.get("fit_window_days")
|
||||
if not start or not isinstance(days, int):
|
||||
return f"fit_window_start={start}"
|
||||
end = (datetime.strptime(start[:10], "%Y-%m-%d") + timedelta(days=days)).strftime(
|
||||
"%Y-%m-%d"
|
||||
)
|
||||
return f"fit over [{start[:10]}, {end}) ({days}d window)"
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--repo-root", default=REPO_ROOT)
|
||||
parser.add_argument(
|
||||
"--output-format",
|
||||
choices=("gha", "json"),
|
||||
default="gha",
|
||||
help="`gha` emits `partitions=<json>` for $GITHUB_OUTPUT; `json` is raw",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--full-parallel",
|
||||
choices=("true", "false"),
|
||||
default="false",
|
||||
help="Lift the max_parallel throttle (set by schedule / `high priority`)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--partition-model-file",
|
||||
default=None,
|
||||
help="Path to sglang-ci-stats model.json (omit/missing -> static fallback)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--pr-test-yml",
|
||||
default=os.path.join(REPO_ROOT, ".github", "workflows", "pr-test.yml"),
|
||||
help="Path to pr-test*.yml; per-stage `run_timeout_minutes` is read from here.",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
files = discover_files(args.repo_root)
|
||||
# Warn-not-fail on unregistered files: run_suite.py catches this at
|
||||
# test-execution time with sanity_check=True; dispatch should keep going.
|
||||
all_tests = collect_tests(files, sanity_check=False)
|
||||
partition_model = load_partition_model(args.partition_model_file)
|
||||
run_timeouts = load_run_timeouts(args.pr_test_yml)
|
||||
|
||||
result = compute_partitions(
|
||||
all_tests,
|
||||
repo_root=args.repo_root,
|
||||
run_timeouts=run_timeouts,
|
||||
partition_model=partition_model,
|
||||
full_parallel=(args.full_parallel == "true"),
|
||||
)
|
||||
payload = json.dumps(result, separators=(",", ":"), sort_keys=True)
|
||||
if args.output_format == "gha":
|
||||
print(f"partitions={payload}")
|
||||
else:
|
||||
print(payload)
|
||||
|
||||
summary_path = os.environ.get("GITHUB_STEP_SUMMARY")
|
||||
if summary_path:
|
||||
with open(summary_path, "a") as f:
|
||||
f.write("## Partitions\n\n")
|
||||
if partition_model is None:
|
||||
src_note = "no live model -- static est_time + (coeff=1, bias=0)"
|
||||
else:
|
||||
src_note = (
|
||||
f"live model: {format_fit_window(partition_model)}, "
|
||||
f"`n_runs={partition_model.get('n_runs')}`"
|
||||
)
|
||||
f.write(
|
||||
f"`full_parallel={args.full_parallel}` "
|
||||
f"(`size//3` throttle is lifted when true); {src_note}\n\n"
|
||||
)
|
||||
f.write("| Suite | size | max_parallel |\n")
|
||||
f.write("|---|---:|---:|\n")
|
||||
for suite, info in sorted(result.items()):
|
||||
f.write(f"| `{suite}` | {info['size']} | {info['max_parallel']} |\n")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,191 @@
|
||||
{
|
||||
"_comment": "Per-model comparison config. Sampling params omitted where model defaults are correct — only override resolution, seed, and params that differ from defaults.",
|
||||
"test_image_url": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/cat.png",
|
||||
"cases": [
|
||||
{
|
||||
"id": "flux1_dev_t2i_1024",
|
||||
"model": "black-forest-labs/FLUX.1-dev",
|
||||
"task": "text-to-image",
|
||||
"prompt": "A futuristic cyberpunk city at night, neon lights reflecting on wet streets",
|
||||
"width": 1024,
|
||||
"height": 1024,
|
||||
"seed": 42,
|
||||
"num_gpus": 2,
|
||||
"frameworks": {
|
||||
"sglang": {
|
||||
"serve_args": "--warmup --dit-layerwise-offload false --tp-size 2",
|
||||
"extra_env": {}
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
"id": "flux2_dev_t2i_1024",
|
||||
"model": "black-forest-labs/FLUX.2-dev",
|
||||
"task": "text-to-image",
|
||||
"prompt": "A futuristic cyberpunk city at night, neon lights reflecting on wet streets",
|
||||
"width": 1024,
|
||||
"height": 1024,
|
||||
"seed": 42,
|
||||
"num_gpus": 2,
|
||||
"frameworks": {
|
||||
"sglang": {
|
||||
"serve_args": "--warmup --dit-layerwise-offload false --tp-size 2",
|
||||
"extra_env": {}
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
"id": "qwen_image_2512_t2i_1024",
|
||||
"model": "Qwen/Qwen-Image-2512",
|
||||
"task": "text-to-image",
|
||||
"prompt": "A futuristic cyberpunk city at night, neon lights reflecting on wet streets",
|
||||
"width": 1024,
|
||||
"height": 1024,
|
||||
"seed": 42,
|
||||
"num_gpus": 2,
|
||||
"frameworks": {
|
||||
"sglang": {
|
||||
"serve_args": "--warmup --tp-size 2",
|
||||
"extra_env": {}
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
"id": "qwen_image_edit_2511",
|
||||
"model": "Qwen/Qwen-Image-Edit-2511",
|
||||
"task": "image-edit",
|
||||
"prompt": "Make the cat wear a red hat",
|
||||
"reference_image": true,
|
||||
"width": 1024,
|
||||
"height": 1024,
|
||||
"seed": 42,
|
||||
"num_gpus": 2,
|
||||
"frameworks": {
|
||||
"sglang": {
|
||||
"serve_args": "--warmup --tp-size 2",
|
||||
"extra_env": {}
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
"id": "zimage_turbo_t2i_1024",
|
||||
"model": "Tongyi-MAI/Z-Image-Turbo",
|
||||
"task": "text-to-image",
|
||||
"prompt": "A futuristic cyberpunk city at night, neon lights reflecting on wet streets",
|
||||
"width": 1024,
|
||||
"height": 1024,
|
||||
"seed": 42,
|
||||
"num_gpus": 2,
|
||||
"frameworks": {
|
||||
"sglang": {
|
||||
"serve_args": "--warmup --tp-size 2",
|
||||
"extra_env": {}
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
"id": "wan22_t2v_a14b_720p",
|
||||
"model": "Wan-AI/Wan2.2-T2V-A14B-Diffusers",
|
||||
"task": "text-to-video",
|
||||
"prompt": "A cat and a dog baking a cake together in a kitchen.",
|
||||
"width": 1280,
|
||||
"height": 720,
|
||||
"num_frames": 81,
|
||||
"seed": 42,
|
||||
"num_gpus": 4,
|
||||
"frameworks": {
|
||||
"sglang": {
|
||||
"serve_args": "--warmup --enable-cfg-parallel --ulysses-degree 2 --text-encoder-cpu-offload --pin-cpu-memory",
|
||||
"extra_env": {}
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
"id": "wan22_ti2v_5b_720p",
|
||||
"model": "Wan-AI/Wan2.2-TI2V-5B-Diffusers",
|
||||
"task": "text-image-to-video",
|
||||
"prompt": "The cat starts walking slowly towards the camera.",
|
||||
"reference_image": true,
|
||||
"width": 1280,
|
||||
"height": 720,
|
||||
"num_frames": 81,
|
||||
"seed": 42,
|
||||
"num_gpus": 1,
|
||||
"frameworks": {
|
||||
"sglang": {
|
||||
"serve_args": "--warmup",
|
||||
"extra_env": {}
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
"id": "ltx2.3_twostage_ti2v_2gpus",
|
||||
"model": "Lightricks/LTX-2.3",
|
||||
"task": "text-image-to-video",
|
||||
"prompt": "The cat starts walking slowly towards the camera.",
|
||||
"reference_image": true,
|
||||
"width": 768,
|
||||
"height": 512,
|
||||
"num_frames": 121,
|
||||
"seed": 42,
|
||||
"num_gpus": 2,
|
||||
"frameworks": {
|
||||
"sglang": {
|
||||
"serve_args": "--warmup --pipeline-class-name LTX2TwoStagePipeline --cfg-parallel-size 2",
|
||||
"extra_env": {}
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
"id": "ideogram4_fp8_t2i_2gpu",
|
||||
"model": "ideogram-ai/ideogram-4-fp8",
|
||||
"task": "text-to-image",
|
||||
"prompt": "A futuristic cyberpunk city at night, neon lights reflecting on wet streets",
|
||||
"width": 1024,
|
||||
"height": 1024,
|
||||
"seed": 42,
|
||||
"num_gpus": 2,
|
||||
"frameworks": {
|
||||
"sglang": {
|
||||
"serve_args": "--warmup --tp-size 2 --attention-backend fa",
|
||||
"extra_env": {}
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
"id": "cosmos3_super_t2v_2gpu",
|
||||
"model": "nvidia/Cosmos3-Super",
|
||||
"task": "text-to-video",
|
||||
"prompt": "A cat and a dog baking a cake together in a kitchen.",
|
||||
"width": 1280,
|
||||
"height": 720,
|
||||
"num_frames": 81,
|
||||
"seed": 42,
|
||||
"num_gpus": 2,
|
||||
"frameworks": {
|
||||
"sglang": {
|
||||
"serve_args": "--warmup --tp-size 2",
|
||||
"extra_env": {"SGLANG_DISABLE_COSMOS3_GUARDRAILS": "1"}
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
"id": "wan22_i2v_a14b_720p",
|
||||
"model": "Wan-AI/Wan2.2-I2V-A14B-Diffusers",
|
||||
"task": "image-to-video",
|
||||
"prompt": "The cat starts walking slowly towards the camera.",
|
||||
"reference_image": true,
|
||||
"width": 1280,
|
||||
"height": 720,
|
||||
"num_frames": 81,
|
||||
"seed": 42,
|
||||
"num_gpus": 4,
|
||||
"frameworks": {
|
||||
"sglang": {
|
||||
"serve_args": "--warmup --enable-cfg-parallel --ulysses-degree 2 --text-encoder-cpu-offload --pin-cpu-memory",
|
||||
"extra_env": {}
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
+327
@@ -0,0 +1,327 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Compute dynamic partitions for diffusion CI tests.
|
||||
|
||||
This script runs on lightweight CI runners without sglang dependencies and uses
|
||||
AST parsing to extract parametrized cases plus standalone files from source.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import importlib.util
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
from diffusion_case_parser import (
|
||||
BASELINE_REL_PATH,
|
||||
RUN_SUITE_REL_PATH,
|
||||
DiffusionSuiteInfo,
|
||||
collect_diffusion_suites,
|
||||
resolve_case_config_path,
|
||||
)
|
||||
|
||||
|
||||
def _load_partitioning_helpers():
|
||||
repo_root = Path(__file__).resolve().parents[4]
|
||||
helper_path = repo_root / "python/sglang/multimodal_gen/test/partitioning.py"
|
||||
spec = importlib.util.spec_from_file_location(
|
||||
"diffusion_test_partitioning", helper_path
|
||||
)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
sys.modules[spec.name] = module
|
||||
spec.loader.exec_module(module)
|
||||
return module.PartitionItem, module.partition_items_by_lpt
|
||||
|
||||
|
||||
PartitionItem, partition_items_by_lpt = _load_partitioning_helpers()
|
||||
|
||||
SUITE_OUTPUT_NAMES = {"1-gpu": "1gpu", "2-gpu": "2gpu", "1-gpu-b200": "b200"}
|
||||
|
||||
USE_NPU_CONFIGS = os.getenv("USE_NPU_CONFIGS", "0").lower() in ("1", "true")
|
||||
|
||||
if USE_NPU_CONFIGS:
|
||||
SUITE_OUTPUT_NAMES = {"1-npu": "1npu", "2-npu": "2npu"}
|
||||
|
||||
DEFAULT_STANDALONE_EST_TIME_SECONDS = 300.0
|
||||
|
||||
|
||||
def validate_suite_case_coverage(suites: dict[str, DiffusionSuiteInfo]) -> None:
|
||||
"""
|
||||
Guardrail: dynamic diffusion suites must contain parametrized cases.
|
||||
"""
|
||||
suites_with_no_cases = []
|
||||
for suite_name in SUITE_OUTPUT_NAMES:
|
||||
suite_info = suites.get(suite_name)
|
||||
if suite_info is None:
|
||||
print(f"Error: Required suite '{suite_name}' not found in parsed suites.")
|
||||
sys.exit(1)
|
||||
if len(suite_info.cases) == 0:
|
||||
suites_with_no_cases.append(suite_name)
|
||||
|
||||
if suites_with_no_cases:
|
||||
joined = ", ".join(suites_with_no_cases)
|
||||
print(
|
||||
"Error: Parsed zero parametrized cases for diffusion suites: "
|
||||
f"{joined}. This usually means run_suite case imports changed but "
|
||||
"diffusion parser logic was not updated."
|
||||
)
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
def compute_partition_count(
|
||||
total_time_seconds: float,
|
||||
min_time_seconds: float,
|
||||
target_time_seconds: float,
|
||||
max_time_seconds: float,
|
||||
max_partitions: int,
|
||||
) -> int:
|
||||
if total_time_seconds <= 0:
|
||||
return 0
|
||||
|
||||
min_partition_count = max(1, math.ceil(total_time_seconds / max_time_seconds))
|
||||
max_partition_count = max(1, math.floor(total_time_seconds / min_time_seconds))
|
||||
|
||||
min_partition_count = min(min_partition_count, max_partitions)
|
||||
max_partition_count = min(max_partition_count, max_partitions)
|
||||
|
||||
if max_partition_count < min_partition_count:
|
||||
fallback_count = math.ceil(total_time_seconds / target_time_seconds)
|
||||
return max(1, min(fallback_count, max_partitions))
|
||||
|
||||
preferred_count = math.ceil(total_time_seconds / target_time_seconds)
|
||||
preferred_count = max(1, min(preferred_count, max_partitions))
|
||||
return max(min_partition_count, min(preferred_count, max_partition_count))
|
||||
|
||||
|
||||
def build_partition_items(
|
||||
suite_info: DiffusionSuiteInfo, include_standalone: bool = True
|
||||
) -> list[PartitionItem]:
|
||||
items = [
|
||||
PartitionItem(kind="case", item_id=case.case_id, est_time=case.est_time)
|
||||
for case in suite_info.cases
|
||||
]
|
||||
if not include_standalone:
|
||||
return items
|
||||
|
||||
items.extend(
|
||||
PartitionItem(
|
||||
kind="standalone",
|
||||
item_id=standalone_file,
|
||||
est_time=suite_info.standalone_est_times.get(
|
||||
standalone_file, DEFAULT_STANDALONE_EST_TIME_SECONDS
|
||||
),
|
||||
used_fallback_estimate=(
|
||||
standalone_file in suite_info.missing_standalone_estimates
|
||||
),
|
||||
)
|
||||
for standalone_file in suite_info.standalone_files
|
||||
)
|
||||
return items
|
||||
|
||||
|
||||
def build_matrix(partition_count: int) -> dict:
|
||||
if partition_count <= 0:
|
||||
return {"include": []}
|
||||
return {"include": [{"part": i} for i in range(partition_count)]}
|
||||
|
||||
|
||||
def build_partition_plan(
|
||||
suite_name: str,
|
||||
partitions: list[list[PartitionItem]],
|
||||
) -> dict:
|
||||
return {
|
||||
"suite": suite_name,
|
||||
"partition_count": len(partitions),
|
||||
"partitions": [
|
||||
{
|
||||
"part": idx,
|
||||
"case_ids": [item.item_id for item in partition if item.kind == "case"],
|
||||
"standalone_files": [
|
||||
item.item_id for item in partition if item.kind == "standalone"
|
||||
],
|
||||
"missing_standalone_estimates": [
|
||||
item.item_id
|
||||
for item in partition
|
||||
if item.kind == "standalone" and item.used_fallback_estimate
|
||||
],
|
||||
"estimated_time": round(sum(item.est_time for item in partition), 1),
|
||||
}
|
||||
for idx, partition in enumerate(partitions)
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
def output_github_value(name: str, value: dict) -> None:
|
||||
value_json = json.dumps(value, separators=(",", ":"))
|
||||
github_output = os.environ.get("GITHUB_OUTPUT")
|
||||
if github_output:
|
||||
with open(github_output, "a", encoding="utf-8") as f:
|
||||
f.write(f"{name}={value_json}\n")
|
||||
print(f"{name}={value_json}")
|
||||
|
||||
|
||||
def output_github_scalar(name: str, value: str) -> None:
|
||||
github_output = os.environ.get("GITHUB_OUTPUT")
|
||||
if github_output:
|
||||
with open(github_output, "a", encoding="utf-8") as f:
|
||||
f.write(f"{name}={value}\n")
|
||||
print(f"{name}={value}")
|
||||
|
||||
|
||||
def print_suite_summary(
|
||||
suite_name: str,
|
||||
suite_info: DiffusionSuiteInfo,
|
||||
partitions: list[list[PartitionItem]],
|
||||
include_standalone: bool = True,
|
||||
) -> None:
|
||||
total_time = sum(
|
||||
item.est_time
|
||||
for item in build_partition_items(
|
||||
suite_info, include_standalone=include_standalone
|
||||
)
|
||||
)
|
||||
print(f"{suite_name.upper()} suite:")
|
||||
print(f" Cases: {len(suite_info.cases)}")
|
||||
standalone_label = "Standalone files"
|
||||
if not include_standalone:
|
||||
standalone_label = "Standalone files ignored"
|
||||
print(f" {standalone_label}: {len(suite_info.standalone_files)}")
|
||||
print(
|
||||
f" Missing standalone estimates: {len(suite_info.missing_standalone_estimates)}"
|
||||
)
|
||||
if suite_info.missing_standalone_estimates:
|
||||
print(
|
||||
f" Fallback standalone estimate: "
|
||||
f"{DEFAULT_STANDALONE_EST_TIME_SECONDS:.1f}s"
|
||||
)
|
||||
for standalone_file in suite_info.missing_standalone_estimates:
|
||||
print(f" - {standalone_file}")
|
||||
print(f" Total estimated time: {total_time:.1f}s ({total_time/60:.1f} min)")
|
||||
print(f" Selected partitions: {len(partitions)}")
|
||||
print()
|
||||
|
||||
print(" Partition assignments:")
|
||||
for idx, partition in enumerate(partitions):
|
||||
partition_time = sum(item.est_time for item in partition)
|
||||
print(f" Partition {idx}:")
|
||||
print(
|
||||
f" Estimated time: {partition_time:.1f}s ({partition_time/60:.1f} min)"
|
||||
)
|
||||
for item in partition:
|
||||
fallback_suffix = (
|
||||
", fallback estimate"
|
||||
if item.kind == "standalone" and item.used_fallback_estimate
|
||||
else ""
|
||||
)
|
||||
print(
|
||||
f" - {item.kind}: {item.item_id} "
|
||||
f"({item.est_time:.1f}s{fallback_suffix})"
|
||||
)
|
||||
print()
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Compute diffusion test partitions for CI"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--min-time",
|
||||
type=float,
|
||||
default=1200.0,
|
||||
help="Minimum desired partition time in seconds (default: 1200 = 20 minutes)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--target-time",
|
||||
type=float,
|
||||
default=1800.0,
|
||||
help="Preferred partition time in seconds (default: 1800 = 30 minutes)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max-time",
|
||||
type=float,
|
||||
default=2400.0,
|
||||
help="Maximum desired partition time in seconds (default: 2400 = 40 minutes)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max-partitions",
|
||||
type=int,
|
||||
default=10,
|
||||
help="Maximum number of partitions (default: 10)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--parametrized-only",
|
||||
action="store_true",
|
||||
help="Only partition DiffusionTestCase parametrized cases.",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
script_dir = Path(__file__).resolve().parent
|
||||
repo_root = script_dir.parent.parent.parent.parent
|
||||
|
||||
baseline_path = repo_root / BASELINE_REL_PATH
|
||||
run_suite_path = repo_root / RUN_SUITE_REL_PATH
|
||||
|
||||
if not run_suite_path.exists():
|
||||
print(f"Error: Run suite not found: {run_suite_path}")
|
||||
sys.exit(1)
|
||||
try:
|
||||
if USE_NPU_CONFIGS:
|
||||
case_config_path = (
|
||||
repo_root
|
||||
/ "python/sglang/multimodal_gen/test/server/ascend/testcase_configs_npu.py"
|
||||
)
|
||||
else:
|
||||
case_config_path = resolve_case_config_path(repo_root, run_suite_path)
|
||||
except (RuntimeError, FileNotFoundError) as exc:
|
||||
print(f"Error: {exc}")
|
||||
sys.exit(1)
|
||||
|
||||
suites = collect_diffusion_suites(
|
||||
case_config_path,
|
||||
run_suite_path,
|
||||
baseline_path,
|
||||
)
|
||||
validate_suite_case_coverage(suites)
|
||||
|
||||
print("=== Diffusion Partition Computation ===")
|
||||
print(f"Min partition time: {args.min_time}s ({args.min_time/60:.1f} min)")
|
||||
print(f"Target partition time: {args.target_time}s ({args.target_time/60:.1f} min)")
|
||||
print(f"Max partition time: {args.max_time}s ({args.max_time/60:.1f} min)")
|
||||
print()
|
||||
|
||||
for suite_name, suite_info in suites.items():
|
||||
if suite_name not in SUITE_OUTPUT_NAMES:
|
||||
continue
|
||||
|
||||
items = build_partition_items(
|
||||
suite_info, include_standalone=not args.parametrized_only
|
||||
)
|
||||
total_time = sum(item.est_time for item in items)
|
||||
partition_count = compute_partition_count(
|
||||
total_time_seconds=total_time,
|
||||
min_time_seconds=args.min_time,
|
||||
target_time_seconds=args.target_time,
|
||||
max_time_seconds=args.max_time,
|
||||
max_partitions=args.max_partitions,
|
||||
)
|
||||
partitions = partition_items_by_lpt(items, partition_count)
|
||||
|
||||
print_suite_summary(
|
||||
suite_name,
|
||||
suite_info,
|
||||
partitions,
|
||||
include_standalone=not args.parametrized_only,
|
||||
)
|
||||
|
||||
output_name = SUITE_OUTPUT_NAMES[suite_name]
|
||||
output_github_value(f"matrix-{output_name}", build_matrix(partition_count))
|
||||
output_github_scalar(f"partition-count-{output_name}", str(partition_count))
|
||||
output_github_value(
|
||||
f"plan-{output_name}", build_partition_plan(suite_name, partitions)
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+517
@@ -0,0 +1,517 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
AST-based parser for diffusion test cases.
|
||||
|
||||
This module parses the diffusion case source and run_suite.py using AST to
|
||||
extract test case information without requiring sglang dependencies. The case
|
||||
source file is discovered from ONE_GPU_CASES/TWO_GPU_CASES imports in
|
||||
run_suite.py so CI keeps a single source of truth.
|
||||
|
||||
Usage:
|
||||
# From sibling scripts in this directory:
|
||||
from diffusion_case_parser import collect_diffusion_suites, resolve_case_config_path
|
||||
case_config_path = resolve_case_config_path(repo_root, run_suite_path)
|
||||
suites = collect_diffusion_suites(case_config_path, run_suite_path, baseline_path)
|
||||
"""
|
||||
|
||||
import ast
|
||||
import json
|
||||
import os
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Optional
|
||||
|
||||
# Mapping from list variable names to suite names
|
||||
CASE_LIST_TO_SUITE = {
|
||||
"ONE_GPU_CASES": "1-gpu",
|
||||
"ONE_GPU_CASES_A": "1-gpu",
|
||||
"ONE_GPU_CASES_B": "1-gpu",
|
||||
"ONE_GPU_CASES_C": "1-gpu-b200",
|
||||
"ONE_GPU_MODELOPT_FP8_CASES": "1-gpu",
|
||||
"ONE_GPU_MODELOPT_CASES": "1-gpu-b200",
|
||||
"ONE_GPU_B200_CASES": "1-gpu-b200",
|
||||
"TWO_GPU_CASES": "2-gpu",
|
||||
"TWO_GPU_CASES_A": "2-gpu",
|
||||
"TWO_GPU_CASES_B": "2-gpu",
|
||||
}
|
||||
|
||||
# Default estimated time for cases without baseline (5 minutes)
|
||||
DEFAULT_EST_TIME_SECONDS = 300.0
|
||||
|
||||
# Fixed overhead for server startup when estimated_full_test_time_s is not set
|
||||
STARTUP_OVERHEAD_SECONDS = 120.0
|
||||
|
||||
# Paths relative to repository root
|
||||
BASELINE_REL_PATH = "python/sglang/multimodal_gen/test/server/perf_baselines"
|
||||
BASELINE_PLATFORM_ORDER = ("h100", "b200", "5090")
|
||||
RUN_SUITE_REL_PATH = "python/sglang/multimodal_gen/test/run_suite.py"
|
||||
|
||||
USE_NPU_CONFIGS = os.getenv("USE_NPU_CONFIGS", "0").lower() in ("1", "true")
|
||||
|
||||
if USE_NPU_CONFIGS:
|
||||
BASELINE_REL_PATH = (
|
||||
"python/sglang/multimodal_gen/test/server/perf_baselines_npu.json"
|
||||
)
|
||||
CASE_LIST_TO_SUITE = {
|
||||
"ONE_NPU_CASES": "1-npu",
|
||||
"TWO_NPU_CASES": "2-npu",
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class DiffusionCaseInfo:
|
||||
"""Information about a single diffusion test case."""
|
||||
|
||||
case_id: str # e.g., "qwen_image_t2i"
|
||||
suite: str # "1-gpu" or "2-gpu"
|
||||
est_time: float # estimated time in seconds
|
||||
|
||||
|
||||
@dataclass
|
||||
class DiffusionSuiteInfo:
|
||||
"""Complete information for a test suite."""
|
||||
|
||||
suite: str # "1-gpu" or "2-gpu"
|
||||
cases: List[DiffusionCaseInfo] # parametrized test cases
|
||||
standalone_files: List[str] # standalone test files
|
||||
standalone_est_times: Dict[str, float] # standalone file -> estimated seconds
|
||||
missing_standalone_estimates: List[
|
||||
str
|
||||
] # standalone files without configured estimate
|
||||
|
||||
|
||||
class DiffusionTestCaseVisitor(ast.NodeVisitor):
|
||||
"""
|
||||
AST visitor to extract DiffusionTestCase definitions from the case config.
|
||||
|
||||
Parses assignments like:
|
||||
ONE_GPU_CASES_A: list[DiffusionTestCase] = [
|
||||
DiffusionTestCase("case_id", ...),
|
||||
...
|
||||
]
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self.cases: Dict[str, List[str]] = {} # list_name -> [case_id, ...]
|
||||
self.factory_case_ids: Dict[str, str] = {}
|
||||
|
||||
def visit_Module(self, node: ast.Module):
|
||||
for stmt in node.body:
|
||||
if not isinstance(stmt, ast.FunctionDef):
|
||||
continue
|
||||
case_id = self._extract_factory_case_id(stmt)
|
||||
if case_id:
|
||||
self.factory_case_ids[stmt.name] = case_id
|
||||
|
||||
self.generic_visit(node)
|
||||
|
||||
def visit_Expr(self, node: ast.Expr):
|
||||
"""Handle ``LIST.append(...)`` mutations at any nesting level.
|
||||
|
||||
Previously only module-top-level ``ast.Expr`` statements were scanned for
|
||||
``.append()`` calls, so cases registered under a platform guard such as
|
||||
``if not current_platform.is_hip(): ONE_GPU_CASES.append(...)`` (used by
|
||||
``hunyuan3d_shape_gen`` and ``turbo_wan2_1_t2v_1.3b``) were invisible to
|
||||
the partition planner and therefore never scheduled in CI. Visiting every
|
||||
``Expr`` lets ``generic_visit`` reach appends inside ``if``/``else`` blocks.
|
||||
"""
|
||||
self._process_expr(node.value)
|
||||
self.generic_visit(node)
|
||||
|
||||
def visit_Assign(self, node: ast.Assign):
|
||||
self._process_assignment(node.targets, node.value)
|
||||
self.generic_visit(node)
|
||||
|
||||
def visit_AnnAssign(self, node: ast.AnnAssign):
|
||||
if node.target and node.value:
|
||||
self._process_assignment([node.target], node.value)
|
||||
self.generic_visit(node)
|
||||
|
||||
def visit_AugAssign(self, node: ast.AugAssign):
|
||||
self._process_aug_assignment(node.target, node.op, node.value)
|
||||
self.generic_visit(node)
|
||||
|
||||
def _process_assignment(self, targets: List[ast.AST], value: ast.AST):
|
||||
"""Process an assignment to extract case IDs."""
|
||||
for target in targets:
|
||||
if isinstance(target, ast.Name):
|
||||
list_name = target.id
|
||||
case_ids = self._extract_case_ids(value)
|
||||
if case_ids is not None:
|
||||
self.cases[list_name] = case_ids
|
||||
|
||||
def _process_aug_assignment(self, target: ast.AST, op: ast.AST, value: ast.AST):
|
||||
"""Process `+=` style assignment to merge case lists."""
|
||||
if not isinstance(target, ast.Name) or not isinstance(op, ast.Add):
|
||||
return
|
||||
|
||||
if isinstance(value, ast.Name):
|
||||
target_suite = CASE_LIST_TO_SUITE.get(target.id)
|
||||
value_suite = CASE_LIST_TO_SUITE.get(value.id)
|
||||
if target_suite and value_suite and target_suite != value_suite:
|
||||
return
|
||||
|
||||
rhs_case_ids = self._extract_case_ids(value)
|
||||
if rhs_case_ids is None:
|
||||
return
|
||||
|
||||
lhs_case_ids = self.cases.get(target.id, [])
|
||||
self.cases[target.id] = [*lhs_case_ids, *rhs_case_ids]
|
||||
|
||||
def _process_expr(self, node: ast.AST):
|
||||
"""Process list mutation calls such as `ONE_GPU_CASES.append(...)`."""
|
||||
if not isinstance(node, ast.Call):
|
||||
return
|
||||
if not isinstance(node.func, ast.Attribute):
|
||||
return
|
||||
if node.func.attr != "append":
|
||||
return
|
||||
if not isinstance(node.func.value, ast.Name):
|
||||
return
|
||||
list_name = node.func.value.id
|
||||
if list_name not in CASE_LIST_TO_SUITE:
|
||||
return
|
||||
if len(node.args) != 1:
|
||||
return
|
||||
|
||||
case_id = self._extract_case_id_from_call(node.args[0])
|
||||
if case_id:
|
||||
self.cases.setdefault(list_name, []).append(case_id)
|
||||
|
||||
def _extract_case_ids(self, node: ast.AST) -> Optional[List[str]]:
|
||||
"""Extract case IDs from a supported expression."""
|
||||
if isinstance(node, ast.List):
|
||||
return self._extract_case_ids_from_list(node)
|
||||
|
||||
if isinstance(node, ast.Name):
|
||||
# Reference to a previously parsed list variable.
|
||||
if node.id not in self.cases:
|
||||
return None
|
||||
return list(self.cases[node.id])
|
||||
|
||||
if isinstance(node, ast.BinOp) and isinstance(node.op, ast.Add):
|
||||
left_ids = self._extract_case_ids(node.left)
|
||||
right_ids = self._extract_case_ids(node.right)
|
||||
if left_ids is None or right_ids is None:
|
||||
return None
|
||||
return [*left_ids, *right_ids]
|
||||
|
||||
return None
|
||||
|
||||
def _extract_case_ids_from_list(self, node: ast.List) -> List[str]:
|
||||
"""Extract case IDs from a literal list of DiffusionTestCase calls."""
|
||||
case_ids = []
|
||||
for elt in node.elts:
|
||||
if isinstance(elt, ast.Starred):
|
||||
starred_case_ids = self._extract_case_ids(elt.value)
|
||||
if starred_case_ids:
|
||||
case_ids.extend(starred_case_ids)
|
||||
continue
|
||||
case_id = self._extract_case_id_from_call(elt)
|
||||
if case_id:
|
||||
case_ids.append(case_id)
|
||||
return case_ids
|
||||
|
||||
def _extract_case_id_from_call(self, node: ast.AST) -> Optional[str]:
|
||||
"""Extract case_id from DiffusionTestCase(...) call."""
|
||||
if not isinstance(node, ast.Call):
|
||||
return None
|
||||
|
||||
# First positional argument is the case_id.
|
||||
if isinstance(node.func, ast.Name) and node.func.id in {
|
||||
"DiffusionTestCase",
|
||||
"_make_modelopt_ci_case",
|
||||
}:
|
||||
if node.args and isinstance(node.args[0], ast.Constant):
|
||||
return node.args[0].value
|
||||
if isinstance(node.func, ast.Name) and not node.args:
|
||||
return self.factory_case_ids.get(node.func.id)
|
||||
|
||||
return None
|
||||
|
||||
def _extract_factory_case_id(self, node: ast.FunctionDef) -> Optional[str]:
|
||||
for child in ast.walk(node):
|
||||
if not isinstance(child, ast.Return) or child.value is None:
|
||||
continue
|
||||
case_id = self._extract_case_id_from_call(child.value)
|
||||
if case_id:
|
||||
return case_id
|
||||
return None
|
||||
|
||||
|
||||
def resolve_case_config_path(repo_root: Path, run_suite_path: Path) -> Path:
|
||||
"""
|
||||
Resolve the diffusion case config path from run_suite imports.
|
||||
|
||||
run_suite.py must import BOTH ONE_GPU_CASES and TWO_GPU_CASES from the same
|
||||
module. That imported module is treated as the single source of truth.
|
||||
"""
|
||||
with open(run_suite_path, "r", encoding="utf-8") as f:
|
||||
content = f.read()
|
||||
|
||||
tree = ast.parse(content, filename=str(run_suite_path))
|
||||
one_gpu_module: Optional[str] = None
|
||||
two_gpu_module: Optional[str] = None
|
||||
|
||||
for node in ast.walk(tree):
|
||||
if not isinstance(node, ast.ImportFrom) or not node.module:
|
||||
continue
|
||||
imported_names = {alias.name for alias in node.names}
|
||||
if "ONE_GPU_CASES" in imported_names:
|
||||
one_gpu_module = node.module
|
||||
if "TWO_GPU_CASES" in imported_names:
|
||||
two_gpu_module = node.module
|
||||
|
||||
if one_gpu_module is None or two_gpu_module is None:
|
||||
raise RuntimeError(
|
||||
"run_suite.py must import BOTH ONE_GPU_CASES and TWO_GPU_CASES."
|
||||
)
|
||||
if one_gpu_module != two_gpu_module:
|
||||
raise RuntimeError(
|
||||
"run_suite.py imports ONE_GPU_CASES and TWO_GPU_CASES from different "
|
||||
f"modules: {one_gpu_module} vs {two_gpu_module}"
|
||||
)
|
||||
|
||||
rel_path = Path(*one_gpu_module.split(".")).with_suffix(".py")
|
||||
candidates = [repo_root / rel_path, repo_root / "python" / rel_path]
|
||||
case_config_path = next((path for path in candidates if path.exists()), None)
|
||||
if case_config_path is None:
|
||||
raise FileNotFoundError(
|
||||
"Resolved case config from run_suite does not exist. Checked: "
|
||||
+ ", ".join(str(path) for path in candidates)
|
||||
)
|
||||
return case_config_path
|
||||
|
||||
|
||||
class RunSuiteVisitor(ast.NodeVisitor):
|
||||
"""
|
||||
AST visitor to extract standalone metadata from run_suite.py.
|
||||
|
||||
Parses:
|
||||
STANDALONE_FILES = {
|
||||
"1-gpu": ["test_lora_format_adapter.py"],
|
||||
"2-gpu": [],
|
||||
}
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self.standalone_files: Dict[str, List[str]] = {}
|
||||
self.standalone_est_times: Dict[str, Dict[str, float]] = {}
|
||||
|
||||
def visit_Assign(self, node: ast.Assign):
|
||||
for target in node.targets:
|
||||
if isinstance(target, ast.Name) and target.id == "STANDALONE_FILES":
|
||||
self.standalone_files = self._extract_file_dict(node.value)
|
||||
if (
|
||||
isinstance(target, ast.Name)
|
||||
and target.id == "STANDALONE_FILE_EST_TIMES"
|
||||
):
|
||||
self.standalone_est_times = self._extract_est_time_dict(node.value)
|
||||
self.generic_visit(node)
|
||||
|
||||
def _extract_file_dict(self, node: ast.AST) -> Dict[str, List[str]]:
|
||||
"""Extract dictionary of suite -> file list."""
|
||||
result = {}
|
||||
if isinstance(node, ast.Dict):
|
||||
for key, value in zip(node.keys, node.values):
|
||||
if isinstance(key, ast.Constant) and isinstance(value, ast.List):
|
||||
suite = key.value
|
||||
files = [
|
||||
elt.value for elt in value.elts if isinstance(elt, ast.Constant)
|
||||
]
|
||||
result[suite] = files
|
||||
return result
|
||||
|
||||
def _extract_est_time_dict(self, node: ast.AST) -> Dict[str, Dict[str, float]]:
|
||||
"""Extract dictionary of suite -> standalone file -> estimated seconds."""
|
||||
result = {}
|
||||
if not isinstance(node, ast.Dict):
|
||||
return result
|
||||
|
||||
for key, value in zip(node.keys, node.values):
|
||||
if not isinstance(key, ast.Constant) or not isinstance(value, ast.Dict):
|
||||
continue
|
||||
|
||||
suite = key.value
|
||||
suite_est_times = {}
|
||||
for inner_key, inner_value in zip(value.keys, value.values):
|
||||
if not (
|
||||
isinstance(inner_key, ast.Constant)
|
||||
and isinstance(inner_value, ast.Constant)
|
||||
):
|
||||
continue
|
||||
suite_est_times[inner_key.value] = float(inner_value.value)
|
||||
result[suite] = suite_est_times
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def _iter_baseline_paths(baseline_path: Path) -> List[Path]:
|
||||
if baseline_path.is_file():
|
||||
return [baseline_path]
|
||||
if not baseline_path.is_dir():
|
||||
return []
|
||||
|
||||
ordered_paths = [
|
||||
baseline_path / f"{platform}.json" for platform in BASELINE_PLATFORM_ORDER
|
||||
]
|
||||
ordered_paths.extend(
|
||||
path
|
||||
for path in sorted(baseline_path.glob("*.json"))
|
||||
if path not in ordered_paths
|
||||
)
|
||||
return [path for path in ordered_paths if path.exists()]
|
||||
|
||||
|
||||
def load_baselines(baseline_path: Path) -> Dict[str, float]:
|
||||
"""
|
||||
Load performance baselines from a JSON file or platform baseline directory.
|
||||
|
||||
Returns:
|
||||
Dictionary mapping case_id to estimated time in seconds.
|
||||
"""
|
||||
baselines = {}
|
||||
for path in _iter_baseline_paths(baseline_path):
|
||||
with open(path, "r", encoding="utf-8") as f:
|
||||
data = json.load(f)
|
||||
|
||||
scenarios = data.get("scenarios", {})
|
||||
for case_id, scenario in scenarios.items():
|
||||
if scenario.get("estimated_full_test_time_s") is not None:
|
||||
est_time = scenario["estimated_full_test_time_s"]
|
||||
else:
|
||||
expected_e2e_ms = scenario.get("expected_e2e_ms", 0)
|
||||
est_time = expected_e2e_ms / 1000.0 + STARTUP_OVERHEAD_SECONDS
|
||||
baselines.setdefault(case_id, est_time)
|
||||
|
||||
return baselines
|
||||
|
||||
|
||||
def get_case_est_time(case_id: str, baselines: Dict[str, float]) -> float:
|
||||
"""Get estimated time for a case, with fallback to default."""
|
||||
return baselines.get(case_id, DEFAULT_EST_TIME_SECONDS)
|
||||
|
||||
|
||||
def parse_testcase_configs(config_path: Path) -> Dict[str, List[str]]:
|
||||
"""
|
||||
Parse a diffusion case config file to extract case IDs.
|
||||
|
||||
Returns:
|
||||
Dictionary mapping list name to case IDs.
|
||||
e.g., {"ONE_GPU_CASES_A": ["qwen_image_t2i", ...], ...}
|
||||
"""
|
||||
with open(config_path, "r", encoding="utf-8") as f:
|
||||
content = f.read()
|
||||
|
||||
tree = ast.parse(content, filename=str(config_path))
|
||||
visitor = DiffusionTestCaseVisitor()
|
||||
visitor.visit(tree)
|
||||
|
||||
return visitor.cases
|
||||
|
||||
|
||||
def parse_run_suite_standalone_data(
|
||||
run_suite_path: Path,
|
||||
) -> tuple[Dict[str, List[str]], Dict[str, Dict[str, float]]]:
|
||||
"""
|
||||
Parse run_suite.py to extract standalone file metadata.
|
||||
|
||||
Returns:
|
||||
Tuple of:
|
||||
- suite -> standalone file list
|
||||
- suite -> standalone file -> estimated seconds
|
||||
"""
|
||||
with open(run_suite_path, "r", encoding="utf-8") as f:
|
||||
content = f.read()
|
||||
|
||||
tree = ast.parse(content, filename=str(run_suite_path))
|
||||
visitor = RunSuiteVisitor()
|
||||
visitor.visit(tree)
|
||||
|
||||
return visitor.standalone_files, visitor.standalone_est_times
|
||||
|
||||
|
||||
def validate_standalone_est_times(
|
||||
standalone_files: Dict[str, List[str]],
|
||||
standalone_est_times: Dict[str, Dict[str, float]],
|
||||
) -> Dict[str, List[str]]:
|
||||
missing_by_suite = {}
|
||||
for suite, files in standalone_files.items():
|
||||
suite_est_times = standalone_est_times.get(suite, {})
|
||||
missing = [
|
||||
standalone_file
|
||||
for standalone_file in files
|
||||
if standalone_file not in suite_est_times
|
||||
]
|
||||
if missing:
|
||||
missing_by_suite[suite] = missing
|
||||
return missing_by_suite
|
||||
|
||||
|
||||
def collect_diffusion_suites(
|
||||
case_config_path: Path,
|
||||
run_suite_path: Path,
|
||||
baseline_path: Path,
|
||||
) -> Dict[str, DiffusionSuiteInfo]:
|
||||
"""
|
||||
Collect all diffusion test suite information using AST parsing.
|
||||
|
||||
Args:
|
||||
case_config_path: Path to case config (resolved from run_suite.py)
|
||||
run_suite_path: Path to run_suite.py
|
||||
baseline_path: Path to perf_baselines/ or a single baseline JSON file
|
||||
|
||||
Returns:
|
||||
Dictionary mapping suite name to DiffusionSuiteInfo.
|
||||
"""
|
||||
# Parse case IDs from the single source case config.
|
||||
case_lists = parse_testcase_configs(case_config_path)
|
||||
|
||||
# Parse standalone files from run_suite.py
|
||||
standalone_files, standalone_est_times = parse_run_suite_standalone_data(
|
||||
run_suite_path
|
||||
)
|
||||
missing_standalone_estimates = validate_standalone_est_times(
|
||||
standalone_files, standalone_est_times
|
||||
)
|
||||
|
||||
# Load baselines for time estimation
|
||||
baselines = load_baselines(baseline_path)
|
||||
|
||||
# Build suite info
|
||||
suites = {}
|
||||
for list_name, suite in CASE_LIST_TO_SUITE.items():
|
||||
case_ids = case_lists.get(list_name, [])
|
||||
cases = [
|
||||
DiffusionCaseInfo(
|
||||
case_id=cid,
|
||||
suite=suite,
|
||||
est_time=get_case_est_time(cid, baselines),
|
||||
)
|
||||
for cid in case_ids
|
||||
]
|
||||
|
||||
if suite not in suites:
|
||||
suites[suite] = DiffusionSuiteInfo(
|
||||
suite=suite,
|
||||
cases=[],
|
||||
standalone_files=standalone_files.get(suite, []),
|
||||
standalone_est_times=dict(standalone_est_times.get(suite, {})),
|
||||
missing_standalone_estimates=list(
|
||||
missing_standalone_estimates.get(suite, [])
|
||||
),
|
||||
)
|
||||
suites[suite].cases.extend(cases)
|
||||
|
||||
# Dedupe duplicated case IDs while preserving first-seen order.
|
||||
for suite_info in suites.values():
|
||||
seen_case_ids = set()
|
||||
deduped_cases = []
|
||||
for case in suite_info.cases:
|
||||
if case.case_id in seen_case_ids:
|
||||
continue
|
||||
seen_case_ids.add(case.case_id)
|
||||
deduped_cases.append(case)
|
||||
suite_info.cases = deduped_cases
|
||||
|
||||
return suites
|
||||
@@ -0,0 +1,835 @@
|
||||
"""Generate a Markdown dashboard for SGLang-Diffusion nightly benchmarks.
|
||||
|
||||
Reads current comparison results + historical data from sgl-project/ci-data repo
|
||||
and produces a Markdown report with tables and trend charts saved as PNG files.
|
||||
|
||||
Usage:
|
||||
python3 scripts/ci/utils/diffusion/generate_diffusion_dashboard.py \
|
||||
--results comparison-results.json \
|
||||
--output dashboard.md \
|
||||
--charts-dir comparison-charts/ \
|
||||
--history-dir history/ # optional, local history JSONs
|
||||
--fetch-history # fetch from GitHub API instead
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
from datetime import datetime, timezone
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# History fetching (from sgl-project/ci-data repo via GitHub API)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
CI_DATA_REPO_OWNER = "sgl-project"
|
||||
CI_DATA_REPO_NAME = "ci-data"
|
||||
CI_DATA_BRANCH = "main"
|
||||
HISTORY_PREFIX = "diffusion-comparisons"
|
||||
MAX_HISTORY_RUNS = 29
|
||||
|
||||
# Base URL for chart images pushed to sgl-project/ci-data
|
||||
CHARTS_RAW_BASE_URL = (
|
||||
f"https://raw.githubusercontent.com/{CI_DATA_REPO_OWNER}/{CI_DATA_REPO_NAME}"
|
||||
f"/{CI_DATA_BRANCH}/{HISTORY_PREFIX}/charts"
|
||||
)
|
||||
|
||||
|
||||
def _github_get(url: str, token: str) -> dict | list | None:
|
||||
"""Simple GET to GitHub API."""
|
||||
from urllib.error import HTTPError
|
||||
from urllib.request import Request, urlopen
|
||||
|
||||
headers = {
|
||||
"Accept": "application/vnd.github+json",
|
||||
"Authorization": f"Bearer {token}",
|
||||
"X-GitHub-Api-Version": "2022-11-28",
|
||||
}
|
||||
req = Request(url, headers=headers)
|
||||
try:
|
||||
with urlopen(req) as resp:
|
||||
return json.loads(resp.read().decode("utf-8"))
|
||||
except HTTPError as e:
|
||||
print(f" Warning: GitHub API request failed ({e.code}): {url}")
|
||||
return None
|
||||
except Exception as e:
|
||||
print(f" Warning: GitHub API request error: {e}")
|
||||
return None
|
||||
|
||||
|
||||
def fetch_history_from_github(token: str) -> list[dict]:
|
||||
"""Fetch recent comparison result JSONs from sgl-project/ci-data repo."""
|
||||
print("Fetching historical comparison data from GitHub...")
|
||||
url = (
|
||||
f"https://api.github.com/repos/{CI_DATA_REPO_OWNER}/{CI_DATA_REPO_NAME}"
|
||||
f"/contents/{HISTORY_PREFIX}?ref={CI_DATA_BRANCH}"
|
||||
)
|
||||
listing = _github_get(url, token)
|
||||
if not listing or not isinstance(listing, list):
|
||||
print(" No historical data found.")
|
||||
return []
|
||||
|
||||
# Filter JSON files and sort by name (date prefix) descending
|
||||
json_files = sorted(
|
||||
[f for f in listing if f["name"].endswith(".json")],
|
||||
key=lambda f: f["name"],
|
||||
reverse=True,
|
||||
)[:MAX_HISTORY_RUNS]
|
||||
|
||||
history = []
|
||||
for entry in json_files:
|
||||
raw_url = entry.get("download_url")
|
||||
if not raw_url:
|
||||
continue
|
||||
data = _github_get(raw_url, token)
|
||||
if data and isinstance(data, dict):
|
||||
history.append(data)
|
||||
print(f" Loaded {len(history)} historical run(s).")
|
||||
return history
|
||||
|
||||
|
||||
def load_history_from_dir(history_dir: str) -> list[dict]:
|
||||
"""Load historical JSONs from a local directory."""
|
||||
if not os.path.isdir(history_dir):
|
||||
return []
|
||||
files = sorted(
|
||||
[f for f in os.listdir(history_dir) if f.endswith(".json")],
|
||||
reverse=True,
|
||||
)[:MAX_HISTORY_RUNS]
|
||||
history = []
|
||||
for fname in files:
|
||||
try:
|
||||
with open(os.path.join(history_dir, fname)) as f:
|
||||
history.append(json.load(f))
|
||||
except Exception:
|
||||
pass
|
||||
return history
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Dashboard generation
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _fmt_latency(val: float | None) -> str:
|
||||
if val is None:
|
||||
return "N/A"
|
||||
return f"{val:.2f}"
|
||||
|
||||
|
||||
def _fmt_speedup(sglang_lat: float | None, other_lat: float | None) -> str:
|
||||
if sglang_lat is None or other_lat is None or sglang_lat <= 0:
|
||||
return "N/A"
|
||||
ratio = other_lat / sglang_lat
|
||||
return f"{ratio:.2f}x"
|
||||
|
||||
|
||||
def _short_date(ts: str) -> str:
|
||||
"""Extract short date from ISO timestamp."""
|
||||
try:
|
||||
dt = datetime.fromisoformat(ts.replace("Z", "+00:00"))
|
||||
return dt.strftime("%b %d")
|
||||
except Exception:
|
||||
return ts[:10]
|
||||
|
||||
|
||||
def _short_sha(sha: str) -> str:
|
||||
return sha[:7] if sha and sha != "unknown" else "?"
|
||||
|
||||
|
||||
def _assess_risk(
|
||||
cid: str,
|
||||
current_cases: dict[str, dict[str, float | None]],
|
||||
history: list[dict],
|
||||
other_frameworks: list[str],
|
||||
) -> tuple[str, str]:
|
||||
"""Assess risk for a given case, returning (emoji, reason).
|
||||
|
||||
Rules (checked in order):
|
||||
- N/A latency → ❌ broken
|
||||
- History exists: SGLang latency >5% vs avg of last 3 runs → ⚠️ regression
|
||||
- Competitor exists & SGLang slower → 🔴 competitive risk
|
||||
- SGLang faster than all competitors by >20% → 🟢 strong advantage
|
||||
- SGLang faster than all competitors by ≤20% → 🟡 moderate advantage
|
||||
- Default → ✅ stable
|
||||
"""
|
||||
sg_lat = current_cases.get(cid, {}).get("sglang")
|
||||
|
||||
# Broken: sglang latency is N/A
|
||||
if sg_lat is None:
|
||||
return "❌", f"{cid}: SGLang latency is N/A (broken)"
|
||||
|
||||
# Check regression against 3-run historical average
|
||||
if history:
|
||||
hist_lats: list[float] = []
|
||||
for run in history[:3]:
|
||||
run_cases = _extract_case_results(run)
|
||||
h_lat = run_cases.get(cid, {}).get("sglang")
|
||||
if h_lat is not None:
|
||||
hist_lats.append(h_lat)
|
||||
if hist_lats:
|
||||
avg_3 = sum(hist_lats) / len(hist_lats)
|
||||
if avg_3 > 0 and (sg_lat - avg_3) / avg_3 > 0.05:
|
||||
pct = (sg_lat - avg_3) / avg_3 * 100
|
||||
return (
|
||||
"⚠️",
|
||||
f"{cid}: SGLang regression +{pct:.1f}% vs 3-run avg "
|
||||
f"({sg_lat:.2f}s vs {avg_3:.2f}s)",
|
||||
)
|
||||
|
||||
# Check competitive risk
|
||||
if other_frameworks:
|
||||
competitor_lats: dict[str, float] = {}
|
||||
for ofw in other_frameworks:
|
||||
olat = current_cases.get(cid, {}).get(ofw)
|
||||
if olat is not None:
|
||||
competitor_lats[ofw] = olat
|
||||
|
||||
if competitor_lats:
|
||||
# SGLang slower than any competitor?
|
||||
for ofw, olat in competitor_lats.items():
|
||||
if sg_lat > olat:
|
||||
return (
|
||||
"🔴",
|
||||
f"{cid}: SGLang slower than {ofw} "
|
||||
f"({sg_lat:.2f}s vs {olat:.2f}s)",
|
||||
)
|
||||
|
||||
# SGLang faster — check margin
|
||||
min_competitor = min(competitor_lats.values())
|
||||
advantage = (min_competitor - sg_lat) / min_competitor
|
||||
if advantage > 0.20:
|
||||
return "🟢", ""
|
||||
else:
|
||||
return "🟡", ""
|
||||
|
||||
# Default: stable
|
||||
return "✅", ""
|
||||
|
||||
|
||||
def _trend_emoji(current: float | None, previous: float | None) -> str:
|
||||
if current is None or previous is None:
|
||||
return ""
|
||||
diff_pct = (current - previous) / previous * 100
|
||||
if diff_pct < -2:
|
||||
return " :arrow_down:" # faster (good)
|
||||
elif diff_pct > 2:
|
||||
return " :arrow_up:" # slower (bad)
|
||||
return " :left_right_arrow:"
|
||||
|
||||
|
||||
def _extract_case_results(run_data: dict) -> dict[str, dict[str, float | None]]:
|
||||
"""Extract {case_id: {framework: latency}} from a run."""
|
||||
mapping: dict[str, dict[str, float | None]] = {}
|
||||
for r in run_data.get("results", []):
|
||||
cid = r["case_id"]
|
||||
fw = r["framework"]
|
||||
if cid not in mapping:
|
||||
mapping[cid] = {}
|
||||
mapping[cid][fw] = r.get("latency_s")
|
||||
return mapping
|
||||
|
||||
|
||||
def _sanitize_filename(name: str) -> str:
|
||||
"""Sanitize a case ID to be a safe filename."""
|
||||
return name.replace("/", "_").replace(" ", "_").replace(":", "_")
|
||||
|
||||
|
||||
def generate_dashboard(
|
||||
current: dict,
|
||||
history: list[dict],
|
||||
charts_dir: str | None = None,
|
||||
) -> tuple[str, list[str]]:
|
||||
"""Generate full markdown dashboard.
|
||||
|
||||
Returns (markdown_string, alert_reasons) where alert_reasons is a list of
|
||||
human-readable strings for cases that need attention (empty if all is well).
|
||||
|
||||
If charts_dir is provided, saves chart PNGs as files to that directory
|
||||
and references them via raw.githubusercontent URLs. Otherwise, charts
|
||||
are omitted.
|
||||
|
||||
Returns the markdown string.
|
||||
"""
|
||||
lines: list[str] = []
|
||||
lines.append("# SGLang-Diffusion Nightly Performance Dashboard\n")
|
||||
ts = current.get("timestamp", datetime.now(timezone.utc).isoformat())
|
||||
sha = current.get("commit_sha", "unknown")
|
||||
lines.append(f"*Generated: {_short_date(ts)} | Commit: `{_short_sha(sha)}`*\n")
|
||||
|
||||
current_cases = _extract_case_results(current)
|
||||
case_ids = list(current_cases.keys())
|
||||
|
||||
# ---- Regression detection ----
|
||||
REGRESSION_THRESHOLD = 0.05 # 5%
|
||||
regressions: list[str] = []
|
||||
if history:
|
||||
prev_cases = _extract_case_results(history[0])
|
||||
for cid in case_ids:
|
||||
for fw in ("sglang", "vllm-omni"):
|
||||
cur = current_cases.get(cid, {}).get(fw)
|
||||
prev = prev_cases.get(cid, {}).get(fw)
|
||||
if cur and prev and prev > 0:
|
||||
pct = (cur - prev) / prev
|
||||
if pct > REGRESSION_THRESHOLD:
|
||||
regressions.append(
|
||||
f"**{cid}** ({fw}): {prev:.2f}s -> {cur:.2f}s "
|
||||
f"(+{pct*100:.1f}%)"
|
||||
)
|
||||
|
||||
if regressions:
|
||||
lines.append("> [!WARNING]\n> **Performance Regression Detected**\n>")
|
||||
for reg in regressions:
|
||||
lines.append(f"> - {reg}")
|
||||
lines.append("\n")
|
||||
|
||||
# Discover all frameworks present in results
|
||||
all_frameworks = []
|
||||
seen_fw = set()
|
||||
for r in current.get("results", []):
|
||||
fw = r["framework"]
|
||||
if fw not in seen_fw:
|
||||
all_frameworks.append(fw)
|
||||
seen_fw.add(fw)
|
||||
# Ensure sglang is first
|
||||
if "sglang" in all_frameworks:
|
||||
all_frameworks.remove("sglang")
|
||||
all_frameworks.insert(0, "sglang")
|
||||
other_frameworks = [fw for fw in all_frameworks if fw != "sglang"]
|
||||
|
||||
# ---- Section 1: SGLang-Diffusion performance (current run) ----
|
||||
lines.append("## SGLang-Diffusion Performance\n")
|
||||
|
||||
# Compute risk assessments for all cases
|
||||
risk_map: dict[str, tuple[str, str]] = {}
|
||||
for cid in case_ids:
|
||||
risk_map[cid] = _assess_risk(cid, current_cases, history, other_frameworks)
|
||||
|
||||
# Dynamic header
|
||||
header = "| Model | Risk |"
|
||||
sep = "|-------|------|"
|
||||
for fw in all_frameworks:
|
||||
header += f" {fw} (s) |"
|
||||
sep += "---------|"
|
||||
for ofw in other_frameworks:
|
||||
header += f" vs {ofw} |"
|
||||
sep += "---------|"
|
||||
lines.append(header)
|
||||
lines.append(sep)
|
||||
|
||||
# One row per case (deduplicated by case_id)
|
||||
seen_cases = set()
|
||||
for r in current.get("results", []):
|
||||
cid = r["case_id"]
|
||||
if cid in seen_cases:
|
||||
continue
|
||||
seen_cases.add(cid)
|
||||
|
||||
case_fws = current_cases.get(cid, {})
|
||||
sg_lat = case_fws.get("sglang")
|
||||
|
||||
risk_emoji, _ = risk_map.get(cid, ("✅", ""))
|
||||
row = f"| {r['model'].split('/')[-1]} | {risk_emoji} |"
|
||||
# Latency columns -- bold the fastest
|
||||
lats = {fw: case_fws.get(fw) for fw in all_frameworks}
|
||||
valid_lats = [v for v in lats.values() if v is not None]
|
||||
min_lat = min(valid_lats) if valid_lats else None
|
||||
for fw in all_frameworks:
|
||||
lat = lats[fw]
|
||||
if lat is not None and min_lat is not None and lat == min_lat:
|
||||
row += f" **{_fmt_latency(lat)}** |"
|
||||
else:
|
||||
row += f" {_fmt_latency(lat)} |"
|
||||
# Speedup columns
|
||||
for ofw in other_frameworks:
|
||||
row += f" {_fmt_speedup(sg_lat, case_fws.get(ofw))} |"
|
||||
lines.append(row)
|
||||
|
||||
# ---- Section 2: Speedup-over-time vs. other frameworks (rendered only when present) ----
|
||||
if history and other_frameworks:
|
||||
lines.append("\n## SGLang vs vLLM-Omni Speedup Over Time\n")
|
||||
|
||||
header = "| Date |"
|
||||
sep = "|------|"
|
||||
for cid in case_ids:
|
||||
header += f" {cid} |"
|
||||
sep += "---------|"
|
||||
lines.append(header)
|
||||
lines.append(sep)
|
||||
|
||||
all_runs = [current] + history
|
||||
for run in all_runs:
|
||||
run_cases = _extract_case_results(run)
|
||||
date = _short_date(run.get("timestamp", ""))
|
||||
row = f"| {date} |"
|
||||
for cid in case_ids:
|
||||
sg = run_cases.get(cid, {}).get("sglang")
|
||||
vl = run_cases.get(cid, {}).get("vllm-omni")
|
||||
row += f" {_fmt_speedup(sg, vl)} |"
|
||||
lines.append(row)
|
||||
|
||||
# ---- Section 4: Matplotlib Trend Charts (saved as PNG files) ----
|
||||
if history and charts_dir:
|
||||
all_runs = list(reversed([current] + history)) # chronological order
|
||||
|
||||
def _chart_label(run: dict) -> str:
|
||||
d = _short_date(run.get("timestamp", ""))
|
||||
s = _short_sha(run.get("commit_sha", ""))
|
||||
return f"{d}\n({s})"
|
||||
|
||||
try:
|
||||
import matplotlib
|
||||
|
||||
matplotlib.use("Agg")
|
||||
import matplotlib.pyplot as plt
|
||||
|
||||
os.makedirs(charts_dir, exist_ok=True)
|
||||
|
||||
# Per-case latency trend charts
|
||||
for cid in case_ids:
|
||||
labels = []
|
||||
sg_vals = []
|
||||
vl_vals = []
|
||||
for run in all_runs:
|
||||
run_cases = _extract_case_results(run)
|
||||
sg = run_cases.get(cid, {}).get("sglang")
|
||||
vl = run_cases.get(cid, {}).get("vllm-omni")
|
||||
if sg is None:
|
||||
continue
|
||||
labels.append(_chart_label(run))
|
||||
sg_vals.append(sg)
|
||||
vl_vals.append(vl)
|
||||
|
||||
if not sg_vals:
|
||||
continue
|
||||
|
||||
has_vl = any(v is not None for v in vl_vals)
|
||||
fig, ax = plt.subplots(figsize=(max(6, len(labels) * 1.2), 4))
|
||||
|
||||
# SGLang line
|
||||
ax.plot(
|
||||
range(len(sg_vals)),
|
||||
sg_vals,
|
||||
"o-",
|
||||
color="#2563eb",
|
||||
linewidth=2,
|
||||
markersize=6,
|
||||
label="SGLang",
|
||||
)
|
||||
for i, v in enumerate(sg_vals):
|
||||
ax.annotate(
|
||||
f"{v:.2f}s",
|
||||
(i, v),
|
||||
textcoords="offset points",
|
||||
xytext=(0, 10),
|
||||
ha="center",
|
||||
fontsize=8,
|
||||
fontweight="bold",
|
||||
color="#2563eb",
|
||||
)
|
||||
|
||||
# vLLM-Omni line (if data exists)
|
||||
if has_vl:
|
||||
vl_clean = [v if v is not None else float("nan") for v in vl_vals]
|
||||
ax.plot(
|
||||
range(len(vl_clean)),
|
||||
vl_clean,
|
||||
"s--",
|
||||
color="#dc2626",
|
||||
linewidth=2,
|
||||
markersize=5,
|
||||
label="vLLM-Omni",
|
||||
)
|
||||
for i, v in enumerate(vl_vals):
|
||||
if v is not None:
|
||||
ax.annotate(
|
||||
f"{v:.2f}s",
|
||||
(i, v),
|
||||
textcoords="offset points",
|
||||
xytext=(0, -14),
|
||||
ha="center",
|
||||
fontsize=8,
|
||||
color="#dc2626",
|
||||
)
|
||||
|
||||
ax.set_xticks(range(len(labels)))
|
||||
ax.set_xticklabels(labels, fontsize=7)
|
||||
ax.set_ylabel("Latency (s)")
|
||||
ax.set_title(f"Latency Trend -- {cid}", fontsize=11, fontweight="bold")
|
||||
ax.legend(loc="lower right", fontsize=8, framealpha=0.8)
|
||||
ax.grid(True, alpha=0.3)
|
||||
all_vals = sg_vals + [v for v in vl_vals if v is not None]
|
||||
y_min = min(all_vals)
|
||||
y_max = max(all_vals)
|
||||
y_range = y_max - y_min if y_max > y_min else max(y_max * 0.1, 0.1)
|
||||
ax.set_ylim(
|
||||
bottom=max(0, y_min - y_range * 0.3),
|
||||
top=y_max + y_range * 0.3,
|
||||
)
|
||||
|
||||
filename = f"latency_{_sanitize_filename(cid)}.png"
|
||||
chart_path = os.path.join(charts_dir, filename)
|
||||
fig.savefig(chart_path, format="png", dpi=120, bbox_inches="tight")
|
||||
plt.close(fig)
|
||||
print(f" Saved chart: {chart_path}")
|
||||
|
||||
chart_url = f"{CHARTS_RAW_BASE_URL}/{filename}"
|
||||
lines.append(f"\n### Latency Trend: {cid}\n")
|
||||
lines.append(f"\n")
|
||||
|
||||
# Speedup trend chart (only if multiple frameworks)
|
||||
if other_frameworks:
|
||||
fig, ax = plt.subplots(figsize=(max(6, len(all_runs) * 1.2), 4))
|
||||
colors = ["#2563eb", "#dc2626", "#16a34a", "#ea580c"]
|
||||
for ci_idx, cid in enumerate(case_ids):
|
||||
speedups = []
|
||||
run_labels = []
|
||||
for run in all_runs:
|
||||
run_cases = _extract_case_results(run)
|
||||
sg = run_cases.get(cid, {}).get("sglang")
|
||||
vl = run_cases.get(cid, {}).get("vllm-omni")
|
||||
if sg and vl and sg > 0:
|
||||
speedups.append(vl / sg)
|
||||
else:
|
||||
speedups.append(None)
|
||||
run_labels.append(_chart_label(run))
|
||||
clean = [v if v is not None else float("nan") for v in speedups]
|
||||
ax.plot(
|
||||
range(len(clean)),
|
||||
clean,
|
||||
"o-",
|
||||
color=colors[ci_idx % len(colors)],
|
||||
linewidth=2,
|
||||
markersize=5,
|
||||
label=cid,
|
||||
)
|
||||
|
||||
ax.set_xticks(range(len(run_labels)))
|
||||
ax.set_xticklabels(run_labels, fontsize=7)
|
||||
ax.set_ylabel("Speedup (x)")
|
||||
ax.set_title(
|
||||
"SGLang Speedup Over vLLM-Omni", fontsize=11, fontweight="bold"
|
||||
)
|
||||
ax.axhline(y=1.0, color="gray", linestyle=":", alpha=0.5)
|
||||
ax.legend(loc="upper left", fontsize=7)
|
||||
ax.grid(True, alpha=0.3)
|
||||
|
||||
filename = "speedup_trend.png"
|
||||
chart_path = os.path.join(charts_dir, filename)
|
||||
fig.savefig(chart_path, format="png", dpi=120, bbox_inches="tight")
|
||||
plt.close(fig)
|
||||
print(f" Saved chart: {chart_path}")
|
||||
|
||||
chart_url = f"{CHARTS_RAW_BASE_URL}/{filename}"
|
||||
lines.append("\n### Speedup Trend (SGLang vs vLLM-Omni)\n")
|
||||
lines.append(f"\n")
|
||||
|
||||
except ImportError:
|
||||
lines.append("\n*Charts unavailable (matplotlib not installed)*\n")
|
||||
|
||||
# ---- SGLang Performance Trend (raw data table, at the end) ----
|
||||
if history:
|
||||
lines.append(f"\n## SGLang Performance Trend (Last {len(history) + 1} Runs)\n")
|
||||
|
||||
header = "| Date | Commit |"
|
||||
sep = "|------|--------|"
|
||||
for cid in case_ids:
|
||||
header += f" {cid} (s) |"
|
||||
sep += "---------|"
|
||||
header += " Trend |"
|
||||
sep += "-------|"
|
||||
lines.append(header)
|
||||
lines.append(sep)
|
||||
|
||||
all_runs = [current] + history
|
||||
for i, run in enumerate(all_runs):
|
||||
run_cases = _extract_case_results(run)
|
||||
date = _short_date(run.get("timestamp", ""))
|
||||
sha_s = _short_sha(run.get("commit_sha", ""))
|
||||
row = f"| {date} | `{sha_s}` |"
|
||||
for cid in case_ids:
|
||||
lat = run_cases.get(cid, {}).get("sglang")
|
||||
row += f" {_fmt_latency(lat)} |"
|
||||
if i + 1 < len(all_runs):
|
||||
prev_cases = _extract_case_results(all_runs[i + 1])
|
||||
emojis = []
|
||||
for cid in case_ids:
|
||||
cur = run_cases.get(cid, {}).get("sglang")
|
||||
prev = prev_cases.get(cid, {}).get("sglang")
|
||||
emojis.append(_trend_emoji(cur, prev))
|
||||
row += " ".join(emojis) + " |"
|
||||
else:
|
||||
row += " -- |"
|
||||
lines.append(row)
|
||||
|
||||
# ---- Risk Notification ----
|
||||
alert_cases = [
|
||||
(cid, emoji, reason)
|
||||
for cid, (emoji, reason) in risk_map.items()
|
||||
if emoji in ("⚠️", "🔴", "❌")
|
||||
]
|
||||
if alert_cases:
|
||||
lines.append("\n> [!CAUTION]")
|
||||
lines.append("> **Action Required — Performance Alert**")
|
||||
lines.append(">")
|
||||
lines.append("> The following cases need attention:")
|
||||
for _cid, _emoji, reason in alert_cases:
|
||||
lines.append(f"> - {reason}")
|
||||
lines.append("")
|
||||
|
||||
# Footer
|
||||
lines.append("\n---")
|
||||
lines.append(
|
||||
"*Generated by `generate_diffusion_dashboard.py` in SGLang nightly CI.*"
|
||||
)
|
||||
|
||||
alert_reasons = [reason for _, _, reason in alert_cases]
|
||||
return "\n".join(lines) + "\n", alert_reasons
|
||||
|
||||
|
||||
ALERT_ASSIGNEES = ["mickqian", "bbuf", "yhyang201"]
|
||||
ALERT_LABEL = "perf-regression"
|
||||
|
||||
|
||||
ALERT_ISSUE_TITLE = "[Diffusion CI] Performance regression tracker"
|
||||
|
||||
|
||||
def _find_alert_issue(repo: str) -> tuple[str | None, bool]:
|
||||
"""Find the perf-regression tracker issue (open OR closed).
|
||||
|
||||
Returns (issue_number, is_open). Prefers an open issue; if none,
|
||||
returns the most recent closed one so it can be reopened.
|
||||
"""
|
||||
import subprocess
|
||||
|
||||
for state in ("open", "closed"):
|
||||
result = subprocess.run(
|
||||
[
|
||||
"gh",
|
||||
"issue",
|
||||
"list",
|
||||
"--repo",
|
||||
repo,
|
||||
"--label",
|
||||
ALERT_LABEL,
|
||||
"--state",
|
||||
state,
|
||||
"--json",
|
||||
"number",
|
||||
"--limit",
|
||||
"1",
|
||||
],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=30,
|
||||
)
|
||||
if result.returncode != 0 or not result.stdout.strip():
|
||||
continue
|
||||
issues = json.loads(result.stdout)
|
||||
if issues:
|
||||
return str(issues[0]["number"]), state == "open"
|
||||
return None, False
|
||||
|
||||
|
||||
def _create_alert_issue(alert_reasons: list[str]) -> None:
|
||||
"""Create or update the single perf-regression tracker issue.
|
||||
|
||||
Logic:
|
||||
- If an open issue exists → add a comment with the new alert.
|
||||
- If a closed issue exists → reopen it, then add a comment.
|
||||
- If no issue exists → create one.
|
||||
|
||||
This guarantees at most one tracker issue ever exists.
|
||||
|
||||
Uses `gh` (GitHub CLI) which is available in all GitHub Actions runners.
|
||||
Falls back silently outside CI.
|
||||
"""
|
||||
import subprocess
|
||||
|
||||
run_url = ""
|
||||
run_id = os.environ.get("GITHUB_RUN_ID", "")
|
||||
repo = os.environ.get("GITHUB_REPOSITORY", "sgl-project/sglang")
|
||||
server_url = os.environ.get("GITHUB_SERVER_URL", "https://github.com")
|
||||
if run_id:
|
||||
run_url = f"{server_url}/{repo}/actions/runs/{run_id}"
|
||||
|
||||
date = datetime.now(timezone.utc).strftime("%Y-%m-%d")
|
||||
|
||||
body_lines = [
|
||||
f"## Performance Alert — {date}",
|
||||
"",
|
||||
"The nightly diffusion benchmark detected the following issue(s):",
|
||||
"",
|
||||
]
|
||||
for reason in alert_reasons:
|
||||
body_lines.append(f"- {reason}")
|
||||
if run_url:
|
||||
body_lines += ["", f"**CI Run:** {run_url}"]
|
||||
body = "\n".join(body_lines)
|
||||
|
||||
try:
|
||||
existing, is_open = _find_alert_issue(repo)
|
||||
|
||||
if existing:
|
||||
# Reopen if closed
|
||||
if not is_open:
|
||||
subprocess.run(
|
||||
[
|
||||
"gh",
|
||||
"issue",
|
||||
"reopen",
|
||||
existing,
|
||||
"--repo",
|
||||
repo,
|
||||
],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=30,
|
||||
)
|
||||
print(f"Reopened alert issue #{existing}")
|
||||
|
||||
# Add comment
|
||||
result = subprocess.run(
|
||||
[
|
||||
"gh",
|
||||
"issue",
|
||||
"comment",
|
||||
existing,
|
||||
"--repo",
|
||||
repo,
|
||||
"--body",
|
||||
body,
|
||||
],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=30,
|
||||
)
|
||||
if result.returncode == 0:
|
||||
print(f"Commented on alert issue #{existing}")
|
||||
else:
|
||||
print(
|
||||
f"Warning: failed to comment on issue #{existing} "
|
||||
f"(rc={result.returncode}): {result.stderr.strip()}"
|
||||
)
|
||||
else:
|
||||
# Create a new issue
|
||||
cmd = [
|
||||
"gh",
|
||||
"issue",
|
||||
"create",
|
||||
"--repo",
|
||||
repo,
|
||||
"--title",
|
||||
ALERT_ISSUE_TITLE,
|
||||
"--body",
|
||||
body,
|
||||
"--label",
|
||||
ALERT_LABEL,
|
||||
]
|
||||
for user in ALERT_ASSIGNEES:
|
||||
cmd += ["--assignee", user]
|
||||
|
||||
result = subprocess.run(cmd, capture_output=True, text=True, timeout=30)
|
||||
if result.returncode == 0:
|
||||
print(f"Created alert issue: {result.stdout.strip()}")
|
||||
else:
|
||||
print(
|
||||
f"Warning: failed to create alert issue "
|
||||
f"(rc={result.returncode}): {result.stderr.strip()}"
|
||||
)
|
||||
except FileNotFoundError:
|
||||
print("Warning: `gh` CLI not found — skipping alert issue creation")
|
||||
except Exception as e:
|
||||
print(f"Warning: failed to create/update alert issue: {e}")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# CLI
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Generate SGLang-Diffusion nightly benchmark dashboard"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--results",
|
||||
required=True,
|
||||
help="Path to comparison-results.json from current run",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output",
|
||||
default="dashboard.md",
|
||||
help="Output markdown file path",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--charts-dir",
|
||||
default="comparison-charts",
|
||||
help="Directory to save chart PNG files (default: comparison-charts/)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--history-dir",
|
||||
default=None,
|
||||
help="Local directory containing historical comparison JSONs",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--fetch-history",
|
||||
action="store_true",
|
||||
help="Fetch history from ci-data GitHub repo",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--step-summary",
|
||||
action="store_true",
|
||||
help="Also write to $GITHUB_STEP_SUMMARY",
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# Load current results
|
||||
with open(args.results) as f:
|
||||
current = json.load(f)
|
||||
print(f"Loaded current results: {len(current.get('results', []))} entries")
|
||||
|
||||
# Load history
|
||||
history: list[dict] = []
|
||||
if args.fetch_history:
|
||||
token = os.environ.get("GH_PAT_FOR_NIGHTLY_CI_DATA") or os.environ.get(
|
||||
"GITHUB_TOKEN"
|
||||
)
|
||||
if token:
|
||||
history = fetch_history_from_github(token)
|
||||
else:
|
||||
print("Warning: No GitHub token available, skipping history fetch")
|
||||
elif args.history_dir:
|
||||
history = load_history_from_dir(args.history_dir)
|
||||
print(f"Loaded {len(history)} historical run(s) from {args.history_dir}")
|
||||
|
||||
# Generate dashboard
|
||||
markdown, alert_reasons = generate_dashboard(
|
||||
current, history, charts_dir=args.charts_dir
|
||||
)
|
||||
|
||||
# Write output
|
||||
os.makedirs(os.path.dirname(args.output) or ".", exist_ok=True)
|
||||
with open(args.output, "w") as f:
|
||||
f.write(markdown)
|
||||
print(f"Dashboard written to {args.output}")
|
||||
|
||||
# Write to GitHub Step Summary
|
||||
if args.step_summary:
|
||||
summary_file = os.environ.get("GITHUB_STEP_SUMMARY")
|
||||
if summary_file:
|
||||
with open(summary_file, "a") as f:
|
||||
f.write(markdown)
|
||||
print("Dashboard appended to $GITHUB_STEP_SUMMARY")
|
||||
else:
|
||||
print("Warning: $GITHUB_STEP_SUMMARY not set, skipping")
|
||||
|
||||
# Create GitHub Issue for performance alerts (so assignees get notified)
|
||||
if alert_reasons:
|
||||
_create_alert_issue(alert_reasons)
|
||||
else:
|
||||
print("No performance alerts — skipping issue creation.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,231 @@
|
||||
"""Publish SGLang-Diffusion nightly benchmark results to sgl-project/ci-data repo.
|
||||
|
||||
Pushes comparison-results.json, dashboard.md, and chart PNG files to the
|
||||
ci-data repository for historical tracking. Chart PNGs are stored under
|
||||
diffusion-comparisons/charts/ so they can be referenced via
|
||||
raw.githubusercontent URLs in the dashboard markdown (GitHub Step Summary
|
||||
blocks data: URIs).
|
||||
|
||||
Usage:
|
||||
python3 scripts/ci/utils/diffusion/publish_comparison_results.py \
|
||||
--results comparison-results.json \
|
||||
--dashboard dashboard.md \
|
||||
--charts-dir comparison-charts/
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
|
||||
# Reuse GitHub API helpers from publish_traces.
|
||||
# Support both direct script execution and package-style imports.
|
||||
if __package__:
|
||||
from ..publish_traces import (
|
||||
create_blobs,
|
||||
create_commit,
|
||||
create_tree,
|
||||
get_branch_sha,
|
||||
get_tree_sha,
|
||||
is_permission_error,
|
||||
is_rate_limit_error,
|
||||
update_branch_ref,
|
||||
verify_token_permissions,
|
||||
)
|
||||
else:
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
||||
from publish_traces import (
|
||||
create_blobs,
|
||||
create_commit,
|
||||
create_tree,
|
||||
get_branch_sha,
|
||||
get_tree_sha,
|
||||
is_permission_error,
|
||||
is_rate_limit_error,
|
||||
update_branch_ref,
|
||||
verify_token_permissions,
|
||||
)
|
||||
|
||||
# Repository configuration
|
||||
REPO_OWNER = "sgl-project"
|
||||
REPO_NAME = "ci-data"
|
||||
BRANCH = "main"
|
||||
STORAGE_PREFIX = "diffusion-comparisons"
|
||||
|
||||
|
||||
def _collect_chart_files(charts_dir: str) -> list[tuple[str, bytes]]:
|
||||
"""Collect PNG chart files from directory for upload."""
|
||||
files: list[tuple[str, bytes]] = []
|
||||
if not charts_dir or not os.path.isdir(charts_dir):
|
||||
return files
|
||||
|
||||
for entry in sorted(os.listdir(charts_dir)):
|
||||
if not entry.lower().endswith(".png"):
|
||||
continue
|
||||
full_path = os.path.join(charts_dir, entry)
|
||||
if not os.path.isfile(full_path):
|
||||
continue
|
||||
with open(full_path, "rb") as f:
|
||||
content = f.read()
|
||||
# Store charts under diffusion-comparisons/charts/
|
||||
repo_path = f"{STORAGE_PREFIX}/charts/{entry}"
|
||||
files.append((repo_path, content))
|
||||
|
||||
return files
|
||||
|
||||
|
||||
def publish_comparison(
|
||||
results_path: str,
|
||||
dashboard_path: str | None = None,
|
||||
charts_dir: str | None = None,
|
||||
) -> None:
|
||||
"""Publish comparison results, dashboard, and charts to ci-data repo."""
|
||||
token = os.environ.get("GH_PAT_FOR_NIGHTLY_CI_DATA") or os.environ.get(
|
||||
"GITHUB_TOKEN"
|
||||
)
|
||||
if not token:
|
||||
print("Error: GH_PAT_FOR_NIGHTLY_CI_DATA or GITHUB_TOKEN not set")
|
||||
sys.exit(1)
|
||||
|
||||
run_id = os.environ.get("GITHUB_RUN_ID", "local")
|
||||
run_number = os.environ.get("GITHUB_RUN_NUMBER", "0")
|
||||
|
||||
# Verify permissions
|
||||
perm = verify_token_permissions(REPO_OWNER, REPO_NAME, token)
|
||||
if perm == "rate_limited":
|
||||
print("Warning: Rate limited, skipping publish")
|
||||
return
|
||||
elif not perm:
|
||||
print("Error: Token permission verification failed")
|
||||
sys.exit(1)
|
||||
|
||||
# Prepare files to upload
|
||||
files_to_upload: list[tuple[str, bytes]] = []
|
||||
|
||||
# Results JSON: stored with date prefix for chronological ordering
|
||||
date_prefix = datetime.now(timezone.utc).strftime("%Y-%m-%d")
|
||||
results_target = f"{STORAGE_PREFIX}/{date_prefix}_{run_id}.json"
|
||||
with open(results_path, "rb") as f:
|
||||
files_to_upload.append((results_target, f.read()))
|
||||
|
||||
# Dashboard markdown: always overwrite latest
|
||||
if dashboard_path and os.path.exists(dashboard_path):
|
||||
dashboard_target = f"{STORAGE_PREFIX}/dashboard.md"
|
||||
with open(dashboard_path, "rb") as f:
|
||||
files_to_upload.append((dashboard_target, f.read()))
|
||||
|
||||
# Chart PNG files
|
||||
chart_files = _collect_chart_files(charts_dir)
|
||||
if chart_files:
|
||||
print(f"Found {len(chart_files)} chart PNG(s) to upload")
|
||||
files_to_upload.extend(chart_files)
|
||||
|
||||
print(f"Publishing {len(files_to_upload)} file(s) to {REPO_OWNER}/{REPO_NAME}")
|
||||
|
||||
# Create blobs
|
||||
try:
|
||||
tree_items = create_blobs(REPO_OWNER, REPO_NAME, files_to_upload, token)
|
||||
except Exception as e:
|
||||
if is_rate_limit_error(e):
|
||||
print("Warning: Rate limited during blob creation, skipping")
|
||||
return
|
||||
if is_permission_error(e):
|
||||
print(f"Error: No write permission to {REPO_OWNER}/{REPO_NAME}")
|
||||
sys.exit(1)
|
||||
raise
|
||||
|
||||
# Commit with retry (handle concurrent writes)
|
||||
max_retries = 5
|
||||
retry_delay = 5
|
||||
|
||||
for attempt in range(max_retries):
|
||||
try:
|
||||
branch_sha = get_branch_sha(REPO_OWNER, REPO_NAME, BRANCH, token)
|
||||
tree_sha = get_tree_sha(REPO_OWNER, REPO_NAME, branch_sha, token)
|
||||
|
||||
new_tree_sha = create_tree(
|
||||
REPO_OWNER, REPO_NAME, tree_sha, tree_items, token
|
||||
)
|
||||
|
||||
commit_msg = (
|
||||
f"Diffusion comparison results for run {run_id} (#{run_number})"
|
||||
)
|
||||
commit_sha = create_commit(
|
||||
REPO_OWNER, REPO_NAME, new_tree_sha, branch_sha, commit_msg, token
|
||||
)
|
||||
|
||||
update_branch_ref(REPO_OWNER, REPO_NAME, BRANCH, commit_sha, token)
|
||||
print(
|
||||
f"Successfully published comparison results (commit {commit_sha[:7]})"
|
||||
)
|
||||
return
|
||||
|
||||
except Exception as e:
|
||||
is_retryable = False
|
||||
if hasattr(e, "error_body"):
|
||||
body = getattr(e, "error_body", "")
|
||||
if "Update is not a fast forward" in body:
|
||||
is_retryable = True
|
||||
elif "Object does not exist" in body:
|
||||
is_retryable = True
|
||||
|
||||
from urllib.error import HTTPError
|
||||
|
||||
if isinstance(e, HTTPError) and e.code in [422, 500, 502, 503, 504]:
|
||||
is_retryable = True
|
||||
|
||||
if is_rate_limit_error(e):
|
||||
print("Warning: Rate limited, skipping publish")
|
||||
return
|
||||
|
||||
if is_permission_error(e):
|
||||
print(f"Error: No write permission to {REPO_OWNER}/{REPO_NAME}")
|
||||
sys.exit(1)
|
||||
|
||||
if is_retryable and attempt < max_retries - 1:
|
||||
print(
|
||||
f"Attempt {attempt + 1}/{max_retries} failed, retrying in {retry_delay}s..."
|
||||
)
|
||||
time.sleep(retry_delay)
|
||||
else:
|
||||
print(f"Failed to publish after {attempt + 1} attempts: {e}")
|
||||
raise
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Publish SGLang-Diffusion nightly benchmark results to ci-data"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--results",
|
||||
required=True,
|
||||
help="Path to comparison-results.json",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--dashboard",
|
||||
default=None,
|
||||
help="Path to dashboard.md (optional)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--charts-dir",
|
||||
default=None,
|
||||
help="Directory containing chart PNG files to upload (optional)",
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
if not os.path.exists(args.results):
|
||||
print(f"Error: Results file not found: {args.results}")
|
||||
sys.exit(1)
|
||||
|
||||
publish_comparison(
|
||||
results_path=args.results,
|
||||
dashboard_path=args.dashboard,
|
||||
charts_dir=args.charts_dir,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,507 @@
|
||||
"""
|
||||
Publish diffusion CI ground-truth images to sgl-project/ci-data
|
||||
via the GitHub API (same pattern as publish_traces.py).
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import base64
|
||||
import hashlib
|
||||
import io
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from urllib.error import HTTPError
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image, ImageFilter
|
||||
|
||||
# Reuse GitHub API helpers from publish_traces.
|
||||
# Support both direct script execution and package-style imports.
|
||||
if __package__:
|
||||
from ..publish_traces import (
|
||||
create_blobs,
|
||||
create_commit,
|
||||
create_tree,
|
||||
get_branch_sha,
|
||||
get_tree_sha,
|
||||
is_permission_error,
|
||||
is_rate_limit_error,
|
||||
make_github_request,
|
||||
update_branch_ref,
|
||||
verify_token_permissions,
|
||||
)
|
||||
else:
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
||||
from publish_traces import (
|
||||
create_blobs,
|
||||
create_commit,
|
||||
create_tree,
|
||||
get_branch_sha,
|
||||
get_tree_sha,
|
||||
is_permission_error,
|
||||
is_rate_limit_error,
|
||||
make_github_request,
|
||||
update_branch_ref,
|
||||
verify_token_permissions,
|
||||
)
|
||||
|
||||
REPO_OWNER = "sgl-project"
|
||||
REPO_NAME = "ci-data"
|
||||
BRANCH = "main"
|
||||
DEFAULT_TARGET_DIR = "diffusion-ci/consistency_gt/sglang_generated"
|
||||
|
||||
IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".webp"}
|
||||
QUALITY_MAX_SIDE = 256
|
||||
LOW_DETAIL_STD_THRESHOLD = 0.075
|
||||
LOW_DETAIL_ENTROPY_THRESHOLD = 0.55
|
||||
LOW_DETAIL_BLUR_RESIDUAL_THRESHOLD = 0.035
|
||||
LOW_DETAIL_GRADIENT_P95_THRESHOLD = 0.045
|
||||
RANDOM_NOISE_CORRELATION_THRESHOLD = 0.55
|
||||
RANDOM_NOISE_LOW_FREQUENCY_THRESHOLD = 0.20
|
||||
RANDOM_NOISE_BLUR_RESIDUAL_THRESHOLD = 0.045
|
||||
OLD_NEW_MIN_SSIM = 0.20
|
||||
OLD_NEW_MAX_MEAN_ABS_DIFF = 45.0
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ImageQualityMetrics:
|
||||
luminance_std: float
|
||||
entropy: float
|
||||
blur_residual: float
|
||||
gradient_p95: float
|
||||
neighbor_correlation: float
|
||||
low_frequency_ratio: float
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class OldNewMetrics:
|
||||
ssim: float
|
||||
mean_abs_diff: float
|
||||
|
||||
|
||||
def collect_images(source_dir, target_dir):
|
||||
"""Collect image files from source_dir and return list of (repo_path, content) tuples."""
|
||||
files = []
|
||||
for entry in sorted(os.listdir(source_dir)):
|
||||
ext = os.path.splitext(entry)[1].lower()
|
||||
if ext not in IMAGE_EXTENSIONS:
|
||||
continue
|
||||
full_path = os.path.join(source_dir, entry)
|
||||
if not os.path.isfile(full_path):
|
||||
continue
|
||||
with open(full_path, "rb") as f:
|
||||
content = f.read()
|
||||
repo_path = f"{target_dir}/{entry}"
|
||||
files.append((repo_path, content))
|
||||
return files
|
||||
|
||||
|
||||
def git_blob_sha(content):
|
||||
header = f"blob {len(content)}\0".encode()
|
||||
return hashlib.sha1(header + content).hexdigest()
|
||||
|
||||
|
||||
def get_remote_blob_shas(repo_owner, repo_name, target_dir, token):
|
||||
return {
|
||||
path: item["sha"]
|
||||
for path, item in get_remote_image_entries(
|
||||
repo_owner, repo_name, target_dir, token
|
||||
).items()
|
||||
}
|
||||
|
||||
|
||||
def get_remote_image_entries(repo_owner, repo_name, target_dir, token):
|
||||
url = (
|
||||
f"https://api.github.com/repos/{repo_owner}/{repo_name}/contents/"
|
||||
f"{target_dir}?ref={BRANCH}"
|
||||
)
|
||||
try:
|
||||
response = make_github_request(url, token)
|
||||
except HTTPError as e:
|
||||
if e.code == 404:
|
||||
return {}
|
||||
raise
|
||||
entries = json.loads(response)
|
||||
return {
|
||||
item["path"]: item
|
||||
for item in entries
|
||||
if item.get("type") == "file"
|
||||
and "sha" in item
|
||||
and os.path.splitext(item["path"])[1].lower() in IMAGE_EXTENSIONS
|
||||
}
|
||||
|
||||
|
||||
def filter_changed_files(files, remote_blob_shas):
|
||||
return [
|
||||
(path, content)
|
||||
for path, content in files
|
||||
if remote_blob_shas.get(path) != git_blob_sha(content)
|
||||
]
|
||||
|
||||
|
||||
def get_remote_blob_content(repo_owner, repo_name, blob_sha, token):
|
||||
url = f"https://api.github.com/repos/{repo_owner}/{repo_name}/git/blobs/{blob_sha}"
|
||||
response = make_github_request(url, token)
|
||||
blob = json.loads(response)
|
||||
if blob.get("encoding") != "base64":
|
||||
raise ValueError(
|
||||
f"Unexpected blob encoding for {blob_sha}: {blob.get('encoding')}"
|
||||
)
|
||||
return base64.b64decode(blob["content"])
|
||||
|
||||
|
||||
def _load_quality_image(content):
|
||||
with Image.open(io.BytesIO(content)) as image:
|
||||
image = image.convert("RGB")
|
||||
image.thumbnail((QUALITY_MAX_SIDE, QUALITY_MAX_SIDE), Image.Resampling.BICUBIC)
|
||||
return image.copy()
|
||||
|
||||
|
||||
def _image_to_rgb_array(image):
|
||||
return np.asarray(image, dtype=np.float32)
|
||||
|
||||
|
||||
def _luminance(rgb):
|
||||
return 0.299 * rgb[..., 0] + 0.587 * rgb[..., 1] + 0.114 * rgb[..., 2]
|
||||
|
||||
|
||||
def _neighbor_correlation(luma):
|
||||
def corr(a, b):
|
||||
a = a.ravel()
|
||||
b = b.ravel()
|
||||
if a.std() < 1e-6 or b.std() < 1e-6:
|
||||
return 1.0
|
||||
return float(np.corrcoef(a, b)[0, 1])
|
||||
|
||||
return (corr(luma[:, 1:], luma[:, :-1]) + corr(luma[1:, :], luma[:-1, :])) / 2
|
||||
|
||||
|
||||
def _low_frequency_ratio(luma):
|
||||
centered = luma - luma.mean()
|
||||
power = np.abs(np.fft.fftshift(np.fft.fft2(centered))) ** 2
|
||||
total_power = power.sum()
|
||||
if total_power < 1e-12:
|
||||
return 0.0
|
||||
|
||||
height, width = luma.shape
|
||||
y, x = np.ogrid[:height, :width]
|
||||
center_y = height // 2
|
||||
center_x = width // 2
|
||||
radius = np.sqrt((y - center_y) ** 2 + (x - center_x) ** 2)
|
||||
low_frequency_radius = min(height, width) * 0.08
|
||||
return float(power[radius <= low_frequency_radius].sum() / total_power)
|
||||
|
||||
|
||||
def compute_image_quality_metrics(content):
|
||||
image = _load_quality_image(content)
|
||||
rgb = _image_to_rgb_array(image)
|
||||
luma = _luminance(rgb) / 255.0
|
||||
|
||||
gradients = np.concatenate(
|
||||
[
|
||||
np.abs(np.diff(luma, axis=1)).ravel(),
|
||||
np.abs(np.diff(luma, axis=0)).ravel(),
|
||||
]
|
||||
)
|
||||
histogram, _ = np.histogram(luma, bins=32, range=(0, 1))
|
||||
probabilities = histogram / histogram.sum()
|
||||
nonzero_probabilities = probabilities[probabilities > 0]
|
||||
entropy = float(
|
||||
-(nonzero_probabilities * np.log2(nonzero_probabilities)).sum() / 5.0
|
||||
)
|
||||
blurred = _image_to_rgb_array(image.filter(ImageFilter.GaussianBlur(radius=3)))
|
||||
|
||||
return ImageQualityMetrics(
|
||||
luminance_std=float(luma.std()),
|
||||
entropy=entropy,
|
||||
blur_residual=float(np.mean(np.abs(rgb - blurred)) / 255.0),
|
||||
gradient_p95=float(np.percentile(gradients, 95)),
|
||||
neighbor_correlation=_neighbor_correlation(luma),
|
||||
low_frequency_ratio=_low_frequency_ratio(luma),
|
||||
)
|
||||
|
||||
|
||||
def get_quality_failure_reasons(metrics):
|
||||
reasons = []
|
||||
low_detail_static = (
|
||||
metrics.luminance_std < LOW_DETAIL_STD_THRESHOLD
|
||||
and metrics.entropy < LOW_DETAIL_ENTROPY_THRESHOLD
|
||||
and (
|
||||
metrics.blur_residual < LOW_DETAIL_BLUR_RESIDUAL_THRESHOLD
|
||||
or metrics.gradient_p95 < LOW_DETAIL_GRADIENT_P95_THRESHOLD
|
||||
)
|
||||
)
|
||||
high_frequency_noise = (
|
||||
metrics.neighbor_correlation < RANDOM_NOISE_CORRELATION_THRESHOLD
|
||||
and metrics.low_frequency_ratio < RANDOM_NOISE_LOW_FREQUENCY_THRESHOLD
|
||||
and metrics.blur_residual > RANDOM_NOISE_BLUR_RESIDUAL_THRESHOLD
|
||||
)
|
||||
if low_detail_static:
|
||||
reasons.append("low-contrast low-detail output")
|
||||
if high_frequency_noise:
|
||||
reasons.append("high-frequency random noise")
|
||||
return reasons
|
||||
|
||||
|
||||
def _resize_for_old_new_compare(content, size=None):
|
||||
with Image.open(io.BytesIO(content)) as image:
|
||||
image = image.convert("RGB")
|
||||
if size is None:
|
||||
image.thumbnail(
|
||||
(QUALITY_MAX_SIDE, QUALITY_MAX_SIDE), Image.Resampling.BICUBIC
|
||||
)
|
||||
else:
|
||||
image = image.resize(size, Image.Resampling.BICUBIC)
|
||||
return _image_to_rgb_array(image)
|
||||
|
||||
|
||||
def compute_old_new_metrics(old_content, new_content):
|
||||
old_rgb = _resize_for_old_new_compare(old_content)
|
||||
new_rgb = _resize_for_old_new_compare(
|
||||
new_content, size=(old_rgb.shape[1], old_rgb.shape[0])
|
||||
)
|
||||
old_luma = _luminance(old_rgb) / 255.0
|
||||
new_luma = _luminance(new_rgb) / 255.0
|
||||
|
||||
old_mean = old_luma.mean()
|
||||
new_mean = new_luma.mean()
|
||||
old_variance = old_luma.var()
|
||||
new_variance = new_luma.var()
|
||||
covariance = ((old_luma - old_mean) * (new_luma - new_mean)).mean()
|
||||
c1 = 0.01**2
|
||||
c2 = 0.03**2
|
||||
ssim = (
|
||||
(2 * old_mean * new_mean + c1)
|
||||
* (2 * covariance + c2)
|
||||
/ ((old_mean**2 + new_mean**2 + c1) * (old_variance + new_variance + c2))
|
||||
)
|
||||
|
||||
return OldNewMetrics(
|
||||
ssim=float(ssim),
|
||||
mean_abs_diff=float(np.abs(old_rgb - new_rgb).mean()),
|
||||
)
|
||||
|
||||
|
||||
def _format_quality_metrics(metrics):
|
||||
return (
|
||||
f"std={metrics.luminance_std:.4f}, entropy={metrics.entropy:.4f}, "
|
||||
f"blur_residual={metrics.blur_residual:.4f}, "
|
||||
f"gradient_p95={metrics.gradient_p95:.4f}, "
|
||||
f"neighbor_corr={metrics.neighbor_correlation:.4f}, "
|
||||
f"low_freq={metrics.low_frequency_ratio:.4f}"
|
||||
)
|
||||
|
||||
|
||||
def _format_old_new_metrics(metrics):
|
||||
return f"ssim={metrics.ssim:.4f}, mean_abs_diff={metrics.mean_abs_diff:.2f}"
|
||||
|
||||
|
||||
def validate_gt_files(files_to_upload, changed_files, remote_image_entries, token):
|
||||
failures = []
|
||||
for path, content in files_to_upload:
|
||||
quality_metrics = compute_image_quality_metrics(content)
|
||||
quality_reasons = get_quality_failure_reasons(quality_metrics)
|
||||
if quality_reasons:
|
||||
failures.append(
|
||||
f"{path}: {', '.join(quality_reasons)} "
|
||||
f"({_format_quality_metrics(quality_metrics)})"
|
||||
)
|
||||
|
||||
for path, content in changed_files:
|
||||
remote_entry = remote_image_entries.get(path)
|
||||
if not remote_entry:
|
||||
continue
|
||||
|
||||
old_content = get_remote_blob_content(
|
||||
REPO_OWNER, REPO_NAME, remote_entry["sha"], token
|
||||
)
|
||||
old_quality_metrics = compute_image_quality_metrics(old_content)
|
||||
old_quality_reasons = get_quality_failure_reasons(old_quality_metrics)
|
||||
if old_quality_reasons:
|
||||
print(
|
||||
f"Skipping old/new drift check for {path} because existing GT is "
|
||||
f"already suspicious: {', '.join(old_quality_reasons)} "
|
||||
f"({_format_quality_metrics(old_quality_metrics)})"
|
||||
)
|
||||
continue
|
||||
|
||||
old_new_metrics = compute_old_new_metrics(old_content, content)
|
||||
if (
|
||||
old_new_metrics.ssim < OLD_NEW_MIN_SSIM
|
||||
and old_new_metrics.mean_abs_diff > OLD_NEW_MAX_MEAN_ABS_DIFF
|
||||
):
|
||||
failures.append(
|
||||
f"{path}: changed too far from existing GT "
|
||||
f"({_format_old_new_metrics(old_new_metrics)})"
|
||||
)
|
||||
|
||||
if not failures:
|
||||
print(
|
||||
f"GT quality gate passed for {len(files_to_upload)} generated image(s) "
|
||||
f"and {len(changed_files)} changed image(s)."
|
||||
)
|
||||
return
|
||||
|
||||
print("GT quality gate failed; refusing to publish suspicious image updates:")
|
||||
for failure in failures:
|
||||
print(f" - {failure}")
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
def check_quality(source_dir, target_dir=None):
|
||||
target_dir = target_dir or DEFAULT_TARGET_DIR
|
||||
token = os.getenv("GITHUB_TOKEN")
|
||||
if not token:
|
||||
print("Error: GITHUB_TOKEN environment variable not set")
|
||||
sys.exit(1)
|
||||
|
||||
files_to_upload = collect_images(source_dir, target_dir)
|
||||
if not files_to_upload:
|
||||
print(f"No image files found in {source_dir}")
|
||||
return
|
||||
|
||||
remote_image_entries = get_remote_image_entries(
|
||||
REPO_OWNER, REPO_NAME, target_dir, token
|
||||
)
|
||||
remote_blob_shas = {
|
||||
path: item["sha"] for path, item in remote_image_entries.items()
|
||||
}
|
||||
changed_files = filter_changed_files(files_to_upload, remote_blob_shas)
|
||||
validate_gt_files(files_to_upload, changed_files, remote_image_entries, token)
|
||||
|
||||
|
||||
def publish(source_dir, target_dir=None):
|
||||
target_dir = target_dir or DEFAULT_TARGET_DIR
|
||||
token = os.getenv("GITHUB_TOKEN")
|
||||
if not token:
|
||||
print("Error: GITHUB_TOKEN environment variable not set")
|
||||
sys.exit(1)
|
||||
|
||||
files_to_upload = collect_images(source_dir, target_dir)
|
||||
if not files_to_upload:
|
||||
print(f"No image files found in {source_dir}")
|
||||
return
|
||||
|
||||
print(
|
||||
f"Found {len(files_to_upload)} image(s) to upload to {REPO_OWNER}/{REPO_NAME}/{target_dir}"
|
||||
)
|
||||
|
||||
# Verify token
|
||||
perm = verify_token_permissions(REPO_OWNER, REPO_NAME, token)
|
||||
if perm == "rate_limited":
|
||||
print("GitHub API rate-limited, skipping upload.")
|
||||
return
|
||||
if not perm:
|
||||
print("Token permission verification failed.")
|
||||
sys.exit(1)
|
||||
|
||||
# Commit with retry (handle concurrent pushes)
|
||||
max_retries = 5
|
||||
for attempt in range(max_retries):
|
||||
try:
|
||||
branch_sha = get_branch_sha(REPO_OWNER, REPO_NAME, BRANCH, token)
|
||||
tree_sha = get_tree_sha(REPO_OWNER, REPO_NAME, branch_sha, token)
|
||||
remote_image_entries = get_remote_image_entries(
|
||||
REPO_OWNER, REPO_NAME, target_dir, token
|
||||
)
|
||||
remote_blob_shas = {
|
||||
path: item["sha"] for path, item in remote_image_entries.items()
|
||||
}
|
||||
changed_files = filter_changed_files(files_to_upload, remote_blob_shas)
|
||||
validate_gt_files(
|
||||
files_to_upload, changed_files, remote_image_entries, token
|
||||
)
|
||||
if not changed_files:
|
||||
print("No image changes to publish.")
|
||||
return
|
||||
|
||||
try:
|
||||
tree_items = create_blobs(REPO_OWNER, REPO_NAME, changed_files, token)
|
||||
except Exception as e:
|
||||
if is_rate_limit_error(e):
|
||||
print("Rate-limited during blob creation, skipping.")
|
||||
return
|
||||
if is_permission_error(e):
|
||||
print(
|
||||
f"ERROR: Token lacks write permission to {REPO_OWNER}/{REPO_NAME}. "
|
||||
"Update GH_PAT_FOR_NIGHTLY_CI_DATA with a token that has contents:write."
|
||||
)
|
||||
sys.exit(1)
|
||||
raise
|
||||
|
||||
new_tree_sha = create_tree(
|
||||
REPO_OWNER, REPO_NAME, tree_sha, tree_items, token
|
||||
)
|
||||
if new_tree_sha == tree_sha:
|
||||
print("No tree changes to publish.")
|
||||
return
|
||||
|
||||
commit_msg = f"diffusion-ci: update images in {target_dir} ({len(changed_files)} files) [automated]"
|
||||
commit_sha = create_commit(
|
||||
REPO_OWNER, REPO_NAME, new_tree_sha, branch_sha, commit_msg, token
|
||||
)
|
||||
update_branch_ref(REPO_OWNER, REPO_NAME, BRANCH, commit_sha, token)
|
||||
print(
|
||||
f"Successfully pushed {len(changed_files)} changed images (commit {commit_sha[:10]})"
|
||||
)
|
||||
return
|
||||
except Exception as e:
|
||||
if is_rate_limit_error(e):
|
||||
print("Rate-limited, skipping.")
|
||||
return
|
||||
if is_permission_error(e):
|
||||
print(f"ERROR: permission denied to {REPO_OWNER}/{REPO_NAME}")
|
||||
sys.exit(1)
|
||||
|
||||
retryable = False
|
||||
if hasattr(e, "error_body"):
|
||||
if "Update is not a fast forward" in e.error_body:
|
||||
retryable = True
|
||||
elif "Object does not exist" in e.error_body:
|
||||
retryable = True
|
||||
|
||||
if isinstance(e, HTTPError) and e.code in [422, 500, 502, 503, 504]:
|
||||
retryable = True
|
||||
|
||||
if retryable and attempt < max_retries - 1:
|
||||
import time
|
||||
|
||||
wait = 2**attempt
|
||||
print(
|
||||
f"Attempt {attempt + 1}/{max_retries} failed, retrying in {wait}s..."
|
||||
)
|
||||
time.sleep(wait)
|
||||
else:
|
||||
print(f"Failed after {attempt + 1} attempts: {e}")
|
||||
raise
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Publish diffusion GT images to GitHub"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--source-dir", required=True, help="Directory containing GT images"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--target-dir",
|
||||
required=False,
|
||||
default=None,
|
||||
help=f"Target directory in the remote repo (default: {DEFAULT_TARGET_DIR})",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--check-only",
|
||||
action="store_true",
|
||||
help="Validate generated GT images without publishing them",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
if args.check_only:
|
||||
check_quality(args.source_dir, args.target_dir)
|
||||
else:
|
||||
publish(args.source_dir, args.target_dir)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
File diff suppressed because it is too large
Load Diff
+163
@@ -0,0 +1,163 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Collect and save diffusion performance metrics for artifact collection in CI.
|
||||
|
||||
This script reads diffusion test results from the pytest stash and saves them
|
||||
with metadata for the performance dashboard.
|
||||
|
||||
Usage:
|
||||
python3 scripts/ci/utils/diffusion/save_diffusion_metrics.py \
|
||||
--gpu-config 1-gpu-h100 \
|
||||
--run-id 12345678 \
|
||||
--output test/diffusion-metrics-1gpu.json \
|
||||
--results-json test/diffusion-results.json
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
from datetime import datetime, timezone
|
||||
|
||||
|
||||
def load_diffusion_results(results_file: str) -> list[dict]:
|
||||
"""Load diffusion performance results from JSON file."""
|
||||
if not os.path.exists(results_file):
|
||||
print(f"Warning: Results file not found: {results_file}")
|
||||
return []
|
||||
|
||||
try:
|
||||
with open(results_file, "r", encoding="utf-8") as f:
|
||||
data = json.load(f)
|
||||
return data if isinstance(data, list) else [data]
|
||||
except (json.JSONDecodeError, OSError) as e:
|
||||
print(f"Warning: Failed to parse {results_file}: {e}")
|
||||
return []
|
||||
|
||||
|
||||
def transform_diffusion_result(result: dict, gpu_config: str) -> dict:
|
||||
"""Transform a diffusion result to match dashboard expectations.
|
||||
|
||||
Dashboard expects:
|
||||
- Separate test_name, class_name
|
||||
- Numeric metrics in consistent units
|
||||
- Optional modality field
|
||||
"""
|
||||
return {
|
||||
"test_name": result.get("test_name"),
|
||||
"class_name": result.get("class_name"),
|
||||
"modality": result.get("modality", "image"),
|
||||
"e2e_ms": result.get("e2e_ms"),
|
||||
"avg_denoise_ms": result.get("avg_denoise_ms"),
|
||||
"median_denoise_ms": result.get("median_denoise_ms"),
|
||||
"stage_metrics": result.get("stage_metrics", {}),
|
||||
"sampled_steps": result.get("sampled_steps", {}),
|
||||
# Video-specific metrics (if present)
|
||||
"frames_per_second": result.get("frames_per_second"),
|
||||
"total_frames": result.get("total_frames"),
|
||||
"avg_frame_time_ms": result.get("avg_frame_time_ms"),
|
||||
}
|
||||
|
||||
|
||||
def group_results_by_class(results: list[dict], gpu_config: str) -> list[dict]:
|
||||
"""Group diffusion results by test class (suite).
|
||||
|
||||
Returns list with one entry per test class, containing all tests in that class.
|
||||
"""
|
||||
groups = {}
|
||||
|
||||
for result in results:
|
||||
class_name = result.get("class_name", "unknown")
|
||||
|
||||
if class_name not in groups:
|
||||
groups[class_name] = {
|
||||
"gpu_config": gpu_config,
|
||||
"test_suite": class_name,
|
||||
"tests": [],
|
||||
}
|
||||
|
||||
transformed = transform_diffusion_result(result, gpu_config)
|
||||
groups[class_name]["tests"].append(transformed)
|
||||
|
||||
return list(groups.values())
|
||||
|
||||
|
||||
def save_metrics(
|
||||
gpu_config: str,
|
||||
run_id: str,
|
||||
output_file: str,
|
||||
results_file: str,
|
||||
) -> bool:
|
||||
"""Collect diffusion metrics and save to output file."""
|
||||
timestamp = datetime.now(timezone.utc).isoformat()
|
||||
|
||||
# Load diffusion results
|
||||
raw_results = load_diffusion_results(results_file)
|
||||
print(f"Loaded {len(raw_results)} diffusion test result(s)")
|
||||
|
||||
# Group by test class
|
||||
grouped = group_results_by_class(raw_results, gpu_config)
|
||||
|
||||
# Create metrics structure
|
||||
metrics = {
|
||||
"run_id": run_id,
|
||||
"timestamp": timestamp,
|
||||
"gpu_config": gpu_config,
|
||||
"test_type": "diffusion",
|
||||
"results": grouped,
|
||||
}
|
||||
|
||||
# Ensure output directory exists and write output
|
||||
try:
|
||||
os.makedirs(os.path.dirname(output_file) or ".", exist_ok=True)
|
||||
with open(output_file, "w", encoding="utf-8") as f:
|
||||
json.dump(metrics, f, indent=2)
|
||||
|
||||
if not raw_results:
|
||||
print(f"Created empty metrics file: {output_file}")
|
||||
else:
|
||||
print(f"Saved diffusion metrics to: {output_file}")
|
||||
return True
|
||||
except OSError as e:
|
||||
print(f"Error writing metrics file: {e}")
|
||||
return False
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Collect diffusion performance metrics from test results"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--gpu-config",
|
||||
required=True,
|
||||
help="GPU configuration (e.g., 1-gpu-h100, 2-gpu-h100)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--run-id",
|
||||
required=True,
|
||||
help="GitHub Actions run ID",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output",
|
||||
required=True,
|
||||
help="Output file path for metrics JSON",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--results-json",
|
||||
required=True,
|
||||
help="Path to diffusion results JSON file",
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
success = save_metrics(
|
||||
gpu_config=args.gpu_config,
|
||||
run_id=args.run_id,
|
||||
output_file=args.output,
|
||||
results_file=args.results_json,
|
||||
)
|
||||
|
||||
sys.exit(0 if success else 1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+343
@@ -0,0 +1,343 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Verify 100% coverage of diffusion test cases.
|
||||
|
||||
This script checks that all expected test cases were executed across all partitions.
|
||||
Designed to run in the CI summary job after all partition jobs complete.
|
||||
|
||||
Usage:
|
||||
python scripts/ci/utils/diffusion/verify_diffusion_coverage.py --reports-dir <path>
|
||||
|
||||
Exit codes:
|
||||
0 - All cases executed (100% coverage)
|
||||
1 - Missing cases detected (coverage < 100%)
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
from diffusion_case_parser import (
|
||||
BASELINE_REL_PATH,
|
||||
RUN_SUITE_REL_PATH,
|
||||
collect_diffusion_suites,
|
||||
resolve_case_config_path,
|
||||
)
|
||||
|
||||
DYNAMIC_SUITES = {"1-gpu", "2-gpu"}
|
||||
|
||||
|
||||
def load_execution_reports(reports_dir: Path) -> list[dict]:
|
||||
"""Load all execution report JSON files from the given directory."""
|
||||
reports = []
|
||||
for json_file in reports_dir.glob("**/execution_report_*.json"):
|
||||
with open(json_file, "r", encoding="utf-8") as f:
|
||||
reports.append(json.load(f))
|
||||
return reports
|
||||
|
||||
|
||||
def get_expected_cases(repo_root: Path) -> dict[str, set[str]]:
|
||||
"""
|
||||
Get all expected cases from case config and run_suite.py.
|
||||
|
||||
Returns:
|
||||
Dictionary mapping suite name to set of expected case IDs.
|
||||
Standalone files are represented as "standalone:<filename>".
|
||||
"""
|
||||
baseline_path = repo_root / BASELINE_REL_PATH
|
||||
run_suite_path = repo_root / RUN_SUITE_REL_PATH
|
||||
case_config_path = resolve_case_config_path(repo_root, run_suite_path)
|
||||
|
||||
suites = collect_diffusion_suites(
|
||||
case_config_path,
|
||||
run_suite_path,
|
||||
baseline_path,
|
||||
)
|
||||
|
||||
expected = {}
|
||||
for suite_name, suite_info in suites.items():
|
||||
if suite_name not in DYNAMIC_SUITES:
|
||||
continue
|
||||
case_ids = set(case.case_id for case in suite_info.cases)
|
||||
# Add standalone files as special case IDs
|
||||
for standalone_file in suite_info.standalone_files:
|
||||
case_ids.add(f"standalone:{standalone_file}")
|
||||
expected[suite_name] = case_ids
|
||||
|
||||
empty_dynamic_suites = [
|
||||
suite_name
|
||||
for suite_name in DYNAMIC_SUITES
|
||||
if suite_name in expected
|
||||
and not any(
|
||||
not case_id.startswith("standalone:") for case_id in expected[suite_name]
|
||||
)
|
||||
]
|
||||
if empty_dynamic_suites:
|
||||
raise RuntimeError(
|
||||
"Parsed zero parametrized cases for diffusion suites: "
|
||||
+ ", ".join(sorted(empty_dynamic_suites))
|
||||
+ ". Refuse to pass coverage verification."
|
||||
)
|
||||
|
||||
return expected
|
||||
|
||||
|
||||
def collect_executed_cases(reports: list[dict]) -> dict[str, set[str]]:
|
||||
"""
|
||||
Collect all executed cases from execution reports.
|
||||
|
||||
Returns:
|
||||
Dictionary mapping suite name to set of executed case IDs.
|
||||
"""
|
||||
executed = {}
|
||||
for report in reports:
|
||||
suite = report["suite"]
|
||||
if suite not in executed:
|
||||
executed[suite] = set()
|
||||
|
||||
executed_cases = report.get("executed_cases", [])
|
||||
if executed_cases:
|
||||
executed[suite].update(executed_cases)
|
||||
elif report["is_standalone"]:
|
||||
standalone_file = report["standalone_file"]
|
||||
executed[suite].add(f"standalone:{standalone_file}")
|
||||
|
||||
return executed
|
||||
|
||||
|
||||
def collect_case_results(reports: list[dict]) -> dict[str, dict[str, str]]:
|
||||
"""
|
||||
Collect case results (pass/fail/error status) from execution reports.
|
||||
|
||||
Returns:
|
||||
Dictionary mapping suite name to {case_id: status} dictionary.
|
||||
"""
|
||||
results = {}
|
||||
for report in reports:
|
||||
suite = report["suite"]
|
||||
if suite not in results:
|
||||
results[suite] = {}
|
||||
|
||||
# Get case_results from report (empty dict for legacy reports)
|
||||
case_results = report.get("case_results", {})
|
||||
results[suite].update(case_results)
|
||||
|
||||
return results
|
||||
|
||||
|
||||
def collect_missing_standalone_estimates(reports: list[dict]) -> dict[str, set[str]]:
|
||||
missing_by_suite: dict[str, set[str]] = {}
|
||||
for report in reports:
|
||||
suite = report["suite"]
|
||||
missing = report.get("missing_standalone_estimates", [])
|
||||
if not missing:
|
||||
continue
|
||||
missing_by_suite.setdefault(suite, set()).update(missing)
|
||||
return missing_by_suite
|
||||
|
||||
|
||||
def collect_standalone_measurements(reports: list[dict]) -> dict[tuple[str, str], dict]:
|
||||
measurements: dict[tuple[str, str], dict] = {}
|
||||
for report in reports:
|
||||
for measurement in report.get("standalone_measurements", []):
|
||||
key = (measurement["suite"], measurement["standalone_file"])
|
||||
measurements[key] = measurement
|
||||
return measurements
|
||||
|
||||
|
||||
def print_missing_standalone_estimates_summary(
|
||||
missing_by_suite: dict[str, set[str]],
|
||||
measurements: dict[tuple[str, str], dict],
|
||||
) -> None:
|
||||
if not missing_by_suite:
|
||||
return
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print(
|
||||
"Add standalone estimate(s) to "
|
||||
"python/sglang/multimodal_gen/test/run_suite.py"
|
||||
)
|
||||
print("=" * 60)
|
||||
print("The following standalone file(s) used fallback estimate 300.0s.")
|
||||
print("Update STANDALONE_FILE_EST_TIMES with the measured runtime below:\n")
|
||||
|
||||
for suite in sorted(missing_by_suite):
|
||||
print(f'"{suite}": {{')
|
||||
for standalone_file in sorted(missing_by_suite[suite]):
|
||||
measurement = measurements.get((suite, standalone_file))
|
||||
measured_time = (
|
||||
measurement["measured_full_test_time_s"] if measurement else 300.0
|
||||
)
|
||||
print(f' "{standalone_file}": {measured_time:.1f},')
|
||||
print("}\n")
|
||||
|
||||
|
||||
def verify_coverage(
|
||||
expected: dict[str, set[str]],
|
||||
executed: dict[str, set[str]],
|
||||
) -> tuple[bool, dict[str, set[str]]]:
|
||||
"""
|
||||
Verify that all expected cases were executed.
|
||||
|
||||
Returns:
|
||||
Tuple of (is_complete, missing_cases_by_suite)
|
||||
"""
|
||||
missing = {}
|
||||
for suite, expected_cases in expected.items():
|
||||
executed_cases = executed.get(suite, set())
|
||||
suite_missing = expected_cases - executed_cases
|
||||
if suite_missing:
|
||||
missing[suite] = suite_missing
|
||||
|
||||
return len(missing) == 0, missing
|
||||
|
||||
|
||||
def print_results_summary(
|
||||
case_results: dict[str, dict[str, str]],
|
||||
) -> tuple[int, int, int]:
|
||||
"""
|
||||
Print test results summary and return counts.
|
||||
|
||||
Returns:
|
||||
Tuple of (passed_count, failed_count, error_count)
|
||||
"""
|
||||
# Check if we have any results data
|
||||
total_results = sum(len(results) for results in case_results.values())
|
||||
if total_results == 0:
|
||||
print("\nTest Results: No results data available (legacy reports)")
|
||||
return (0, 0, 0)
|
||||
|
||||
# Count by status
|
||||
passed_count = 0
|
||||
failed_count = 0
|
||||
error_count = 0
|
||||
failed_cases: dict[str, list[str]] = {}
|
||||
|
||||
for suite, results in case_results.items():
|
||||
for case_id, status in results.items():
|
||||
if status == "pass":
|
||||
passed_count += 1
|
||||
elif status == "fail":
|
||||
failed_count += 1
|
||||
if suite not in failed_cases:
|
||||
failed_cases[suite] = []
|
||||
failed_cases[suite].append(case_id)
|
||||
elif status == "error":
|
||||
error_count += 1
|
||||
if suite not in failed_cases:
|
||||
failed_cases[suite] = []
|
||||
failed_cases[suite].append(f"{case_id} (error)")
|
||||
|
||||
# Print summary
|
||||
total = passed_count + failed_count + error_count
|
||||
print("\n" + "=" * 60)
|
||||
print("Test Results Summary")
|
||||
print("=" * 60)
|
||||
print(f" Total executed: {total}")
|
||||
print(f" ✅ Passed: {passed_count}")
|
||||
print(f" ❌ Failed: {failed_count}")
|
||||
if error_count > 0:
|
||||
print(f" ⚠️ Errors: {error_count}")
|
||||
|
||||
# Print failed cases if any
|
||||
if failed_cases:
|
||||
print("\nFailed cases:")
|
||||
for suite, cases in sorted(failed_cases.items()):
|
||||
print(f" {suite}:")
|
||||
for case_id in sorted(cases):
|
||||
print(f" - {case_id}")
|
||||
|
||||
return (passed_count, failed_count, error_count)
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Verify 100% coverage of diffusion test cases"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--reports-dir",
|
||||
type=str,
|
||||
required=True,
|
||||
help="Directory containing execution report JSON files",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
# Determine repository root
|
||||
script_dir = Path(__file__).resolve().parent
|
||||
repo_root = script_dir.parent.parent.parent.parent
|
||||
|
||||
reports_dir = Path(args.reports_dir)
|
||||
|
||||
print("=" * 60)
|
||||
print("Diffusion CI Coverage Verification")
|
||||
print("=" * 60)
|
||||
|
||||
# Load execution reports
|
||||
reports = load_execution_reports(reports_dir)
|
||||
print(f"\nLoaded {len(reports)} execution reports")
|
||||
|
||||
if not reports:
|
||||
print("\nERROR: No execution reports found!")
|
||||
print(f"Expected reports in: {reports_dir}")
|
||||
sys.exit(1)
|
||||
|
||||
# Get expected cases
|
||||
try:
|
||||
expected = get_expected_cases(repo_root)
|
||||
except (RuntimeError, FileNotFoundError) as exc:
|
||||
print(f"\nERROR: {exc}")
|
||||
sys.exit(1)
|
||||
print("\nExpected cases by suite:")
|
||||
for suite, cases in expected.items():
|
||||
print(f" {suite}: {len(cases)} cases")
|
||||
|
||||
# Collect executed cases
|
||||
executed = collect_executed_cases(reports)
|
||||
print("\nExecuted cases by suite:")
|
||||
for suite, cases in executed.items():
|
||||
print(f" {suite}: {len(cases)} cases")
|
||||
|
||||
# Collect case results
|
||||
case_results = collect_case_results(reports)
|
||||
missing_standalone_estimates = collect_missing_standalone_estimates(reports)
|
||||
standalone_measurements = collect_standalone_measurements(reports)
|
||||
|
||||
# Verify coverage
|
||||
is_complete, missing = verify_coverage(expected, executed)
|
||||
|
||||
if is_complete:
|
||||
print("\n" + "=" * 60)
|
||||
print("✅ COVERAGE: 100% - All test cases executed")
|
||||
print("=" * 60)
|
||||
else:
|
||||
print("\n" + "=" * 60)
|
||||
print("❌ COVERAGE FAILURE: Missing test cases detected")
|
||||
print("=" * 60)
|
||||
for suite, cases in missing.items():
|
||||
print(f"\n{suite.upper()} suite - Missing {len(cases)} case(s):")
|
||||
for case_id in sorted(cases):
|
||||
print(f" - {case_id}")
|
||||
|
||||
# Print test results summary
|
||||
passed_count, failed_count, error_count = print_results_summary(case_results)
|
||||
print_missing_standalone_estimates_summary(
|
||||
missing_standalone_estimates, standalone_measurements
|
||||
)
|
||||
|
||||
# Exit with appropriate code
|
||||
if not is_complete:
|
||||
sys.exit(1)
|
||||
elif missing_standalone_estimates:
|
||||
sys.exit(1)
|
||||
elif failed_count > 0 or error_count > 0:
|
||||
print("\n" + "=" * 60)
|
||||
print("⚠️ WARNING: Some tests failed but coverage is complete")
|
||||
print("=" * 60)
|
||||
sys.exit(0) # Coverage is complete, failures are visible in results
|
||||
else:
|
||||
sys.exit(0)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,119 @@
|
||||
import argparse
|
||||
import datetime
|
||||
import json
|
||||
import sys
|
||||
|
||||
MOVING_TAGS = {"dev", "dev-cu12", "dev-cu13", "latest"}
|
||||
|
||||
|
||||
def render_tag_template(tag: str, version: str, date: str, short_sha: str) -> str:
|
||||
return (
|
||||
tag.replace("{version}", version)
|
||||
.replace("{date}", date)
|
||||
.replace("{short_sha}", short_sha)
|
||||
)
|
||||
|
||||
|
||||
def is_moving_tag(tag: str) -> bool:
|
||||
return tag in MOVING_TAGS or tag.startswith("latest-")
|
||||
|
||||
|
||||
def select_tag(
|
||||
tag_config: str, cuda: str, version: str, date: str, short_sha: str
|
||||
) -> str:
|
||||
entries = json.loads(tag_config)
|
||||
for entry in entries:
|
||||
if entry.get("cuda") != cuda:
|
||||
continue
|
||||
|
||||
tags = [
|
||||
render_tag_template(tag, version, date, short_sha)
|
||||
for tag in entry.get("tags", [])
|
||||
]
|
||||
if not tags:
|
||||
raise ValueError(f"No tags configured for CUDA variant {cuda}")
|
||||
|
||||
for tag in tags:
|
||||
if not is_moving_tag(tag):
|
||||
return tag
|
||||
return tags[0]
|
||||
|
||||
raise ValueError(f"CUDA variant {cuda} not found in tag_config")
|
||||
|
||||
|
||||
def build_arg_tokens(
|
||||
*,
|
||||
cuda: str,
|
||||
tag_config: str,
|
||||
image_repo: str,
|
||||
version: str,
|
||||
build_commit: str,
|
||||
build_url: str,
|
||||
date: str,
|
||||
short_sha: str,
|
||||
) -> list[str]:
|
||||
image_tag = select_tag(tag_config, cuda, version, date, short_sha)
|
||||
build_args = {
|
||||
"SGLANG_BUILD_COMMIT": build_commit,
|
||||
"SGLANG_BUILD_URL": build_url,
|
||||
"SGLANG_IMAGE_TAG": f"{image_repo}:{image_tag}",
|
||||
}
|
||||
|
||||
tokens = []
|
||||
for key, value in build_args.items():
|
||||
tokens.extend(["--build-arg", f"{key}={value}"])
|
||||
return tokens
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Emit docker build arguments for SGLang image metadata."
|
||||
)
|
||||
parser.add_argument("--cuda", required=True, help="CUDA variant from tag_config.")
|
||||
parser.add_argument("--tag-config", required=True, help="Docker tag JSON config.")
|
||||
parser.add_argument("--image-repo", required=True, help="Docker image repository.")
|
||||
parser.add_argument("--sgl-version", default="", help="SGLang release version.")
|
||||
parser.add_argument(
|
||||
"--build-commit",
|
||||
required=True,
|
||||
help="Commit checked out for the Docker build.",
|
||||
)
|
||||
parser.add_argument("--build-url", default="", help="CI run URL.")
|
||||
parser.add_argument(
|
||||
"--date",
|
||||
default=datetime.datetime.now(datetime.timezone.utc).strftime("%Y%m%d"),
|
||||
help="Date used for {date} tag templates.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--short-sha",
|
||||
default="",
|
||||
help="Short SHA used for {short_sha}; defaults to build commit prefix.",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main() -> int:
|
||||
args = parse_args()
|
||||
short_sha = args.short_sha or args.build_commit[:8]
|
||||
|
||||
try:
|
||||
tokens = build_arg_tokens(
|
||||
cuda=args.cuda,
|
||||
tag_config=args.tag_config,
|
||||
image_repo=args.image_repo,
|
||||
version=args.sgl_version,
|
||||
build_commit=args.build_commit,
|
||||
build_url=args.build_url,
|
||||
date=args.date,
|
||||
short_sha=short_sha,
|
||||
)
|
||||
except (json.JSONDecodeError, ValueError) as exc:
|
||||
print(f"error: {exc}", file=sys.stderr)
|
||||
return 1
|
||||
|
||||
print("\n".join(tokens))
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
Executable
+53
@@ -0,0 +1,53 @@
|
||||
#!/bin/bash
|
||||
# Ensure protoc is installed for router build (gRPC protobuf compilation).
|
||||
set -euxo pipefail
|
||||
|
||||
if command -v protoc >/dev/null 2>&1 && protoc --version >/dev/null 2>&1; then
|
||||
echo "protoc already installed: $(protoc --version)"
|
||||
exit 0
|
||||
fi
|
||||
|
||||
if command -v protoc >/dev/null 2>&1; then
|
||||
echo "protoc found but not runnable, reinstalling..."
|
||||
else
|
||||
echo "protoc not found, installing..."
|
||||
fi
|
||||
|
||||
ARCH=$(uname -m)
|
||||
|
||||
if command -v apt-get &> /dev/null; then
|
||||
# Ubuntu/Debian
|
||||
apt-get update || true # May fail due to unrelated broken packages
|
||||
PROTOC_APT_PACKAGES=(wget unzip)
|
||||
apt-get install -y --no-install-recommends "${PROTOC_APT_PACKAGES[@]}" || {
|
||||
echo "Warning: apt-get install failed, checking if required packages are available..."
|
||||
for pkg in "${PROTOC_APT_PACKAGES[@]}"; do
|
||||
if ! dpkg -l "$pkg" 2>/dev/null | grep -q "^ii"; then
|
||||
echo "ERROR: Required package $pkg is not installed and apt-get failed"
|
||||
exit 1
|
||||
fi
|
||||
done
|
||||
echo "All required packages are already installed, continuing..."
|
||||
}
|
||||
elif command -v yum &> /dev/null; then
|
||||
# RHEL/CentOS
|
||||
yum update -y
|
||||
yum install -y wget unzip
|
||||
else
|
||||
echo "ERROR: Neither apt-get nor yum found; cannot install protoc"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
if [ "$ARCH" = "aarch64" ] || [ "$ARCH" = "arm64" ]; then
|
||||
PROTOC_ARCH="aarch_64"
|
||||
else
|
||||
PROTOC_ARCH="x86_64"
|
||||
fi
|
||||
PROTOC_ZIP="protoc-32.0-linux-${PROTOC_ARCH}.zip"
|
||||
(
|
||||
cd /tmp
|
||||
wget "https://github.com/protocolbuffers/protobuf/releases/download/v32.0/${PROTOC_ZIP}"
|
||||
unzip -o "${PROTOC_ZIP}" -d /usr/local
|
||||
rm -f "${PROTOC_ZIP}"
|
||||
)
|
||||
protoc --version
|
||||
Executable
+24
@@ -0,0 +1,24 @@
|
||||
#!/bin/bash
|
||||
# Install protoc and a Rust toolchain (rustup/cargo). Required by setuptools-rust
|
||||
# to build the bundled native gRPC extension (rust/sglang-grpc) when installing
|
||||
# the main `sglang` wheel from source. Idempotent — both helpers no-op if
|
||||
# already installed.
|
||||
#
|
||||
# protoc installs system-wide (/usr/local) and apt deps, so it needs root.
|
||||
# rustup installs per-user under $HOME/.cargo, so it must run as the calling
|
||||
# user (running it under sudo would put cargo in /root/.cargo and the rest of
|
||||
# the job wouldn't find it).
|
||||
set -euxo pipefail
|
||||
|
||||
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
|
||||
if [ "$(id -u)" = "0" ]; then
|
||||
SUDO=""
|
||||
elif command -v sudo >/dev/null 2>&1; then
|
||||
SUDO="sudo"
|
||||
else
|
||||
SUDO=""
|
||||
fi
|
||||
|
||||
${SUDO} bash "${SCRIPT_DIR}/install_protoc.sh"
|
||||
bash "${SCRIPT_DIR}/install_rustup.sh"
|
||||
Executable
+60
@@ -0,0 +1,60 @@
|
||||
#!/bin/bash
|
||||
# Ensure a Rust toolchain (rustc/cargo) is installed for crates built from
|
||||
# source, e.g. the native gRPC extension bundled into the sglang wheel via
|
||||
# setuptools-rust. Minimum supported version is 1.85 (edition 2024).
|
||||
set -euxo pipefail
|
||||
|
||||
# Make cargo/rustc visible to the rest of this shell and to subsequent
|
||||
# GitHub Actions steps in the same job.
|
||||
export PATH="${CARGO_HOME:-$HOME/.cargo}/bin:${PATH}"
|
||||
if [ -n "${GITHUB_PATH:-}" ]; then
|
||||
# Self-heal if _runner_file_commands/ disappears mid-job on some self-hosted
|
||||
# runners; the runner reads this file by its registered UUID at step end, so
|
||||
# recreating the path keeps PATH propagation working for subsequent steps.
|
||||
mkdir -p "$(dirname "${GITHUB_PATH}")" 2>/dev/null || true
|
||||
echo "${CARGO_HOME:-$HOME/.cargo}/bin" >> "${GITHUB_PATH}" || true
|
||||
fi
|
||||
|
||||
if command -v cargo >/dev/null 2>&1 && command -v rustc >/dev/null 2>&1; then
|
||||
echo "rust already installed: $(rustc --version), $(cargo --version)"
|
||||
exit 0
|
||||
fi
|
||||
|
||||
echo "rust not found, installing via rustup..."
|
||||
|
||||
# rustup.rs requires curl — make sure it's present.
|
||||
if ! command -v curl >/dev/null 2>&1; then
|
||||
if command -v apt-get &> /dev/null; then
|
||||
apt-get update || true
|
||||
apt-get install -y --no-install-recommends curl ca-certificates
|
||||
elif command -v yum &> /dev/null; then
|
||||
yum install -y curl ca-certificates
|
||||
else
|
||||
echo "ERROR: curl is required to install rustup, but no supported package manager was found"
|
||||
exit 1
|
||||
fi
|
||||
fi
|
||||
|
||||
if [ -n "${RUSTUP_CACHE_URL:-}" ]; then
|
||||
# An in-cluster HTTP mirror is available (e.g. on NPU runners).
|
||||
export RUSTUP_DIST_SERVER="${RUSTUP_CACHE_URL}/rustup"
|
||||
export RUSTUP_UPDATE_ROOT="${RUSTUP_CACHE_URL}/rustup/rustup"
|
||||
case "$(uname -m)" in
|
||||
x86_64) RUSTUP_ARCH="x86_64-unknown-linux-gnu" ;;
|
||||
aarch64) RUSTUP_ARCH="aarch64-unknown-linux-gnu" ;;
|
||||
*) echo "ERROR: unsupported arch $(uname -m)"; exit 1 ;;
|
||||
esac
|
||||
RUSTUP_TMP="$(mktemp -d)"
|
||||
trap 'rm -rf "${RUSTUP_TMP}"' EXIT
|
||||
curl --retry 3 --retry-delay 2 -sSfL \
|
||||
"${RUSTUP_UPDATE_ROOT}/dist/${RUSTUP_ARCH}/rustup-init" \
|
||||
-o "${RUSTUP_TMP}/rustup-init"
|
||||
chmod +x "${RUSTUP_TMP}/rustup-init"
|
||||
"${RUSTUP_TMP}/rustup-init" -y --no-modify-path
|
||||
else
|
||||
curl --proto '=https' --tlsv1.2 --retry 3 --retry-delay 2 -sSf https://sh.rustup.rs \
|
||||
| sh -s -- -y --no-modify-path
|
||||
fi
|
||||
|
||||
rustc --version
|
||||
cargo --version
|
||||
Executable
+141
@@ -0,0 +1,141 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Merge per-partition metrics into a consolidated metrics file.
|
||||
|
||||
This script reads all per-partition metric JSON files and consolidates them
|
||||
into a single JSON file with run-level metadata.
|
||||
|
||||
Usage:
|
||||
python3 scripts/ci/utils/merge_metrics.py \
|
||||
--input-dir metrics/ \
|
||||
--output consolidated-metrics-12345678.json \
|
||||
--run-id 12345678 \
|
||||
--commit-sha abc123def456
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import glob
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
from datetime import datetime, timezone
|
||||
|
||||
|
||||
def find_partition_files(input_dir: str) -> list[str]:
|
||||
"""Find all partition metric files in the input directory."""
|
||||
patterns = [
|
||||
os.path.join(input_dir, "**/metrics-*.json"),
|
||||
os.path.join(input_dir, "**/diffusion-metrics-*.json"),
|
||||
os.path.join(input_dir, "**/comparison-metrics-*.json"),
|
||||
]
|
||||
files = set()
|
||||
for pattern in patterns:
|
||||
files.update(glob.glob(pattern, recursive=True))
|
||||
return list(files)
|
||||
|
||||
|
||||
def load_partition_metrics(filepath: str) -> dict | None:
|
||||
"""Load a partition metrics file."""
|
||||
try:
|
||||
with open(filepath, "r", encoding="utf-8") as f:
|
||||
return json.load(f)
|
||||
except (json.JSONDecodeError, OSError) as e:
|
||||
print(f"Warning: Failed to load {filepath}: {e}")
|
||||
return None
|
||||
|
||||
|
||||
def merge_metrics(
|
||||
input_dir: str,
|
||||
output_file: str,
|
||||
run_id: str,
|
||||
commit_sha: str,
|
||||
branch: str | None = None,
|
||||
) -> bool:
|
||||
"""Merge all partition metrics into a consolidated file."""
|
||||
run_date = datetime.now(timezone.utc).isoformat()
|
||||
|
||||
# Find all partition files
|
||||
partition_files = find_partition_files(input_dir)
|
||||
print(f"Found {len(partition_files)} partition file(s)")
|
||||
|
||||
all_results = []
|
||||
if not partition_files:
|
||||
print("No partition metrics files found")
|
||||
else:
|
||||
# Load all partition files
|
||||
for filepath in sorted(partition_files):
|
||||
print(f" Reading: {filepath}")
|
||||
metrics = load_partition_metrics(filepath)
|
||||
if metrics and "results" in metrics:
|
||||
all_results.extend(metrics["results"])
|
||||
print(f"Total results collected: {len(all_results)}")
|
||||
|
||||
# Create consolidated structure
|
||||
consolidated = {
|
||||
"run_id": run_id,
|
||||
"run_date": run_date,
|
||||
"commit_sha": commit_sha,
|
||||
"branch": branch,
|
||||
"results": all_results,
|
||||
}
|
||||
|
||||
# Ensure output directory exists and write output
|
||||
try:
|
||||
os.makedirs(os.path.dirname(output_file) or ".", exist_ok=True)
|
||||
with open(output_file, "w", encoding="utf-8") as f:
|
||||
json.dump(consolidated, f, indent=2)
|
||||
|
||||
if not partition_files:
|
||||
print(f"Created empty consolidated file: {output_file}")
|
||||
else:
|
||||
print(f"Saved consolidated metrics to: {output_file}")
|
||||
return True
|
||||
except OSError as e:
|
||||
print(f"Error writing consolidated file: {e}")
|
||||
return False
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Merge per-partition metrics into consolidated file"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--input-dir",
|
||||
required=True,
|
||||
help="Directory containing partition metric files",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output",
|
||||
required=True,
|
||||
help="Output file path for consolidated metrics JSON",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--run-id",
|
||||
required=True,
|
||||
help="GitHub Actions run ID",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--commit-sha",
|
||||
required=True,
|
||||
help="Git commit SHA",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--branch",
|
||||
default=None,
|
||||
help="Git branch name (optional)",
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
success = merge_metrics(
|
||||
input_dir=args.input_dir,
|
||||
output_file=args.output,
|
||||
run_id=args.run_id,
|
||||
commit_sha=args.commit_sha,
|
||||
branch=args.branch,
|
||||
)
|
||||
|
||||
sys.exit(0 if success else 1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+407
@@ -0,0 +1,407 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Pre-validate all cached HuggingFace models to provide detailed feedback.
|
||||
|
||||
This script runs once during CI initialization (in prepare_runner.sh) to:
|
||||
1. Scan snapshots in ~/.cache/huggingface/hub/ (with time/quantity limits)
|
||||
2. Validate completeness (config/tokenizer/weights)
|
||||
3. Output detailed failure reasons for debugging
|
||||
|
||||
NOTE: This script no longer writes shared validation markers. Each test run
|
||||
independently validates its cache using per-run markers to avoid cross-runner
|
||||
cache state pollution.
|
||||
"""
|
||||
|
||||
import glob
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
# Add python directory to path to import sglang modules
|
||||
REPO_ROOT = Path(__file__).parent.parent.parent.parent
|
||||
sys.path.insert(0, str(REPO_ROOT / "python"))
|
||||
|
||||
from sglang.srt.model_loader.ci_weight_validation import ( # noqa: E402
|
||||
_validate_diffusion_model,
|
||||
validate_cache_with_detailed_reason,
|
||||
)
|
||||
|
||||
# Limits to avoid spending too much time on validation
|
||||
MAX_VALIDATION_TIME_SECONDS = 300 # Max 5 minutes total
|
||||
|
||||
|
||||
def find_all_hf_snapshots():
|
||||
"""
|
||||
Find all HuggingFace snapshots in cache.
|
||||
|
||||
Returns:
|
||||
List of (model_name, snapshot_dir) tuples, sorted by mtime (newest first)
|
||||
"""
|
||||
hf_home = os.environ.get("HF_HOME", os.path.expanduser("~/.cache/huggingface"))
|
||||
hub_dir = os.path.join(hf_home, "hub")
|
||||
|
||||
if not os.path.isdir(hub_dir):
|
||||
print(f"HF hub directory not found: {hub_dir}")
|
||||
return []
|
||||
|
||||
snapshots = []
|
||||
|
||||
# Pattern: models--org--model/snapshots/hash
|
||||
for model_dir in glob.glob(os.path.join(hub_dir, "models--*")):
|
||||
# Extract model name from directory (models--org--model -> org/model)
|
||||
dir_name = os.path.basename(model_dir)
|
||||
if not dir_name.startswith("models--"):
|
||||
continue
|
||||
|
||||
# models--meta-llama--Llama-2-7b-hf -> meta-llama/Llama-2-7b-hf
|
||||
# Handle multi-part names: models--a--b--c -> a/b-c (join parts 1+ with /)
|
||||
parts = dir_name.split("--")
|
||||
if len(parts) < 3 or parts[0] != "models":
|
||||
# Invalid format, skip
|
||||
continue
|
||||
# Standard format: models--org--repo -> org/repo
|
||||
# Extended format: models--org--repo--extra -> org/repo-extra (join with -)
|
||||
model_name = parts[1] + "/" + "-".join(parts[2:])
|
||||
|
||||
snapshots_dir = os.path.join(model_dir, "snapshots")
|
||||
if not os.path.isdir(snapshots_dir):
|
||||
continue
|
||||
|
||||
# Find all snapshot hashes
|
||||
for snapshot_hash_dir in os.listdir(snapshots_dir):
|
||||
snapshot_path = os.path.join(snapshots_dir, snapshot_hash_dir)
|
||||
if os.path.isdir(snapshot_path):
|
||||
try:
|
||||
mtime = os.path.getmtime(snapshot_path)
|
||||
snapshots.append((model_name, snapshot_path, mtime))
|
||||
except OSError:
|
||||
continue
|
||||
|
||||
# Sort by mtime (newest first) - prioritize recently used models
|
||||
snapshots.sort(key=lambda x: x[2], reverse=True)
|
||||
|
||||
# Return without mtime
|
||||
return [(name, path) for name, path, _ in snapshots]
|
||||
|
||||
|
||||
def is_transformers_text_model(snapshot_dir):
|
||||
"""
|
||||
Check if a snapshot is a transformers text model.
|
||||
|
||||
Only excludes (returns False) for models with STRONG evidence of being
|
||||
diffusers/generation pipelines. Uses conservative heuristics to avoid
|
||||
false negatives on multimodal LLMs with tokenizers.
|
||||
|
||||
Args:
|
||||
snapshot_dir: Path to snapshot directory
|
||||
|
||||
Returns:
|
||||
True if this looks like a transformers text model, False otherwise (N/A)
|
||||
"""
|
||||
# Check for diffusers pipeline markers (strong evidence)
|
||||
diffusers_markers = [
|
||||
"model_index.json", # Diffusers pipeline config
|
||||
"scheduler", # Scheduler directory (diffusers)
|
||||
]
|
||||
if any(
|
||||
os.path.exists(os.path.join(snapshot_dir, marker))
|
||||
for marker in diffusers_markers
|
||||
):
|
||||
return False
|
||||
|
||||
config_path = os.path.join(snapshot_dir, "config.json")
|
||||
if not os.path.exists(config_path):
|
||||
# No config.json - likely not a transformers model
|
||||
return False
|
||||
|
||||
try:
|
||||
with open(config_path, "r", encoding="utf-8") as f:
|
||||
config = json.load(f)
|
||||
|
||||
# Check for explicit diffusers/generation model types (conservative keywords)
|
||||
model_type = config.get("_class_name") or config.get("model_type")
|
||||
if model_type:
|
||||
model_type_lower = str(model_type).lower()
|
||||
# Only exclude clear diffusion/generation models
|
||||
if any(
|
||||
keyword in model_type_lower
|
||||
for keyword in [
|
||||
"diffusion",
|
||||
"unet",
|
||||
"vae",
|
||||
"controlnet",
|
||||
"stable-diffusion",
|
||||
"latent-diffusion",
|
||||
]
|
||||
):
|
||||
return False
|
||||
|
||||
# Check architectures for explicit generation/diffusion classes
|
||||
architectures = config.get("architectures", [])
|
||||
if architectures:
|
||||
arch_str = " ".join(architectures).lower()
|
||||
# Conservative: only exclude obvious diffusion/generation architectures
|
||||
# Use word boundaries to avoid false positives (e.g., "dit" in "conditional")
|
||||
for keyword in [
|
||||
"diffusion",
|
||||
"unet2d",
|
||||
"unet3d",
|
||||
"vaedecoder", # More specific than "vae"
|
||||
"vaeencoder",
|
||||
"controlnet",
|
||||
"autoencoder",
|
||||
"ditmodel", # Diffusion Transformer - use more specific pattern
|
||||
"pixart", # PixArt diffusion model
|
||||
]:
|
||||
if keyword in arch_str:
|
||||
return False
|
||||
|
||||
# Check for standalone vision encoder/image processor (no text component)
|
||||
# Only if model name explicitly indicates non-text usage
|
||||
model_name = config.get("_name_or_path", "").lower()
|
||||
|
||||
if any(
|
||||
keyword in model_name
|
||||
for keyword in [
|
||||
"image-edit-", # Pure image editing (e.g., Qwen-Image-Edit)
|
||||
"-image-editing",
|
||||
"dit-", # DiT generation models
|
||||
"pixart-", # PixArt generation models
|
||||
]
|
||||
):
|
||||
# Additional check: does it have tokenizer? If yes, might be multimodal LLM
|
||||
has_tokenizer = any(
|
||||
os.path.exists(os.path.join(snapshot_dir, fname))
|
||||
for fname in ["tokenizer.json", "tokenizer.model", "tiktoken.model"]
|
||||
)
|
||||
if not has_tokenizer:
|
||||
# Image-edit model without tokenizer -> likely pure vision pipeline
|
||||
return False
|
||||
|
||||
# Default: assume it's a transformers text/multimodal model
|
||||
# Even if it lacks tokenizer, let validation report the actual error
|
||||
# (better false positive than false negative for text models)
|
||||
return True
|
||||
|
||||
except (json.JSONDecodeError, OSError, KeyError):
|
||||
# Can't parse config - assume it's transformers and let validation report failure
|
||||
return True
|
||||
|
||||
|
||||
def scan_weight_files(snapshot_dir):
|
||||
"""
|
||||
Scan for weight files in a snapshot.
|
||||
|
||||
Returns:
|
||||
List of weight file paths, or empty list if scan fails
|
||||
"""
|
||||
weight_files = []
|
||||
|
||||
# First, look for index files
|
||||
index_patterns = ["*.safetensors.index.json", "pytorch_model.bin.index.json"]
|
||||
index_files = []
|
||||
for pattern in index_patterns:
|
||||
index_files.extend(glob.glob(os.path.join(snapshot_dir, pattern)))
|
||||
|
||||
# If we have safetensors index, collect shards from it
|
||||
for index_file in index_files:
|
||||
if index_file.endswith(".safetensors.index.json"):
|
||||
try:
|
||||
with open(index_file, "r", encoding="utf-8") as f:
|
||||
index_data = json.load(f)
|
||||
weight_map = index_data.get("weight_map", {})
|
||||
for weight_file in set(weight_map.values()):
|
||||
weight_path = os.path.join(snapshot_dir, weight_file)
|
||||
if os.path.exists(weight_path):
|
||||
weight_files.append(weight_path)
|
||||
except Exception as e:
|
||||
print(
|
||||
f" Warning: Failed to parse index {os.path.basename(index_file)}: {e}"
|
||||
)
|
||||
|
||||
# If no index found or no shards from index, do recursive glob
|
||||
if not weight_files:
|
||||
matched = glob.glob(
|
||||
os.path.join(snapshot_dir, "**/*.safetensors"), recursive=True
|
||||
)
|
||||
MAX_WEIGHT_FILES = 1000
|
||||
if len(matched) > MAX_WEIGHT_FILES:
|
||||
print(
|
||||
f" Warning: Too many safetensors files ({len(matched)} > {MAX_WEIGHT_FILES})"
|
||||
)
|
||||
return []
|
||||
|
||||
for f in matched:
|
||||
if os.path.exists(f): # Filter out broken symlinks
|
||||
weight_files.append(f)
|
||||
|
||||
return weight_files
|
||||
|
||||
|
||||
def validate_snapshot(model_name, snapshot_dir, weight_files, validated_cache):
|
||||
"""
|
||||
Validate a snapshot and return detailed status.
|
||||
|
||||
Uses in-process cache to avoid duplicate validation within the same run.
|
||||
|
||||
Args:
|
||||
model_name: Model identifier
|
||||
snapshot_dir: Path to snapshot directory
|
||||
weight_files: List of weight files to validate
|
||||
validated_cache: Dict to track already-validated snapshots in this run
|
||||
|
||||
Returns:
|
||||
Tuple of (result, reason):
|
||||
- (True, None) if validation passed
|
||||
- (False, reason_str) if validation failed
|
||||
- (None, None) if skipped (already validated in this run)
|
||||
"""
|
||||
# Fast path: check in-process cache first
|
||||
if snapshot_dir in validated_cache:
|
||||
return None, None # Already validated in this run, skip
|
||||
|
||||
try:
|
||||
# Perform validation with detailed reason
|
||||
is_complete, reason = validate_cache_with_detailed_reason(
|
||||
snapshot_dir=snapshot_dir,
|
||||
weight_files=weight_files,
|
||||
model_name_or_path=model_name,
|
||||
)
|
||||
|
||||
# Cache result to avoid re-validation in this run
|
||||
validated_cache[snapshot_dir] = (is_complete, reason)
|
||||
|
||||
return is_complete, reason
|
||||
|
||||
except Exception as e:
|
||||
error_msg = f"Validation raised exception: {e}"
|
||||
return False, error_msg
|
||||
|
||||
|
||||
def main():
|
||||
start_time = time.time()
|
||||
|
||||
print("=" * 70)
|
||||
print("CI_OFFLINE: Pre-validating cached HuggingFace models")
|
||||
print("=" * 70)
|
||||
print(f"Max time: {MAX_VALIDATION_TIME_SECONDS}s")
|
||||
print()
|
||||
|
||||
print("Scanning HuggingFace cache for models...")
|
||||
snapshots = find_all_hf_snapshots()
|
||||
|
||||
if not snapshots:
|
||||
print("No cached models found, skipping validation")
|
||||
print("=" * 70)
|
||||
return
|
||||
|
||||
print(f"Found {len(snapshots)} snapshot(s) in cache")
|
||||
print()
|
||||
|
||||
validated_count = 0
|
||||
failed_count = 0
|
||||
skipped_count = 0
|
||||
processed_count = 0
|
||||
|
||||
# In-process cache to avoid re-validating same snapshot in this run
|
||||
validated_cache = {}
|
||||
|
||||
for model_name, snapshot_dir in snapshots:
|
||||
# Check time limit
|
||||
elapsed = time.time() - start_time
|
||||
if elapsed > MAX_VALIDATION_TIME_SECONDS:
|
||||
print()
|
||||
print(
|
||||
f"Time limit reached ({elapsed:.1f}s > {MAX_VALIDATION_TIME_SECONDS}s)"
|
||||
)
|
||||
print(
|
||||
f"Stopping validation, {len(snapshots) - processed_count} snapshots remaining"
|
||||
)
|
||||
break
|
||||
|
||||
snapshot_hash = os.path.basename(snapshot_dir)
|
||||
print(
|
||||
f"[{processed_count + 1}/{len(snapshots)}] {model_name} ({snapshot_hash[:8]}...)"
|
||||
)
|
||||
processed_count += 1
|
||||
|
||||
# Determine model type by checking for model_index.json (diffusers pipeline marker)
|
||||
model_index_path = os.path.join(snapshot_dir, "model_index.json")
|
||||
is_diffusion_model = os.path.exists(model_index_path)
|
||||
|
||||
if is_diffusion_model:
|
||||
# This is a diffusers pipeline - use diffusion validation
|
||||
try:
|
||||
is_valid, reason = _validate_diffusion_model(snapshot_dir)
|
||||
|
||||
if is_valid:
|
||||
print(" PASS (diffusion) - Cache complete & valid")
|
||||
validated_count += 1
|
||||
else:
|
||||
print(f" FAIL (diffusion) - {reason}")
|
||||
failed_count += 1
|
||||
|
||||
except Exception as e:
|
||||
print(f" FAIL (diffusion) - Validation raised exception: {e}")
|
||||
failed_count += 1
|
||||
|
||||
continue
|
||||
|
||||
# Transformers model - use standard validation
|
||||
# First check if this looks like a transformers text model
|
||||
if not is_transformers_text_model(snapshot_dir):
|
||||
# Not a recognized model type, skip
|
||||
print(
|
||||
" SKIP (unknown type) - Not a diffusers pipeline or transformers model"
|
||||
)
|
||||
skipped_count += 1
|
||||
continue
|
||||
|
||||
# Scan weight files
|
||||
weight_files = scan_weight_files(snapshot_dir)
|
||||
|
||||
if not weight_files:
|
||||
print(" SKIP (no weights) - empty or incomplete download")
|
||||
skipped_count += 1
|
||||
continue
|
||||
|
||||
# Validate
|
||||
try:
|
||||
result, reason = validate_snapshot(
|
||||
model_name, snapshot_dir, weight_files, validated_cache
|
||||
)
|
||||
|
||||
if result is True:
|
||||
print(" PASS - Cache complete & valid")
|
||||
validated_count += 1
|
||||
elif result is False:
|
||||
# Print detailed failure reason
|
||||
if reason:
|
||||
print(f" FAIL (incomplete) - {reason}")
|
||||
else:
|
||||
print(" FAIL (incomplete) - cache validation failed")
|
||||
failed_count += 1
|
||||
else: # None (skipped)
|
||||
print(" SKIP (already validated in this run)")
|
||||
skipped_count += 1
|
||||
|
||||
except Exception as e:
|
||||
print(f" FAIL (error) - Validation raised exception: {e}")
|
||||
failed_count += 1
|
||||
|
||||
elapsed_total = time.time() - start_time
|
||||
|
||||
print()
|
||||
print("=" * 70)
|
||||
print(f"Validation summary (completed in {elapsed_total:.1f}s):")
|
||||
print(f" PASS (complete & valid): {validated_count}")
|
||||
print(f" FAIL (incomplete/corrupted): {failed_count}")
|
||||
print(f" SKIP (no weights/duplicate): {skipped_count}")
|
||||
print(f" Total processed: {processed_count}/{len(snapshots)}")
|
||||
print("=" * 70)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,517 @@
|
||||
"""
|
||||
Publish performance traces to GitHub repository
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import base64
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
import warnings
|
||||
from urllib.error import HTTPError
|
||||
from urllib.request import Request, urlopen
|
||||
|
||||
|
||||
def is_rate_limit_error(e):
|
||||
"""Check if an exception is a GitHub rate limit error (not permission error)"""
|
||||
if not isinstance(e, HTTPError):
|
||||
return False
|
||||
if e.code == 429:
|
||||
return True
|
||||
if e.code == 403:
|
||||
# 403 can be rate limit OR permission error - check the message
|
||||
error_body = getattr(e, "error_body", "")
|
||||
if isinstance(error_body, str):
|
||||
# Rate limit errors contain specific phrases
|
||||
rate_limit_phrases = [
|
||||
"rate limit",
|
||||
"abuse detection",
|
||||
"secondary rate limit",
|
||||
]
|
||||
return any(phrase in error_body.lower() for phrase in rate_limit_phrases)
|
||||
return False
|
||||
|
||||
|
||||
def is_permission_error(e):
|
||||
"""Check if an exception is a GitHub permission error"""
|
||||
if not isinstance(e, HTTPError) or e.code != 403:
|
||||
return False
|
||||
error_body = getattr(e, "error_body", "")
|
||||
if isinstance(error_body, str):
|
||||
permission_phrases = [
|
||||
"resource not accessible",
|
||||
"must have push access",
|
||||
"permission",
|
||||
"denied",
|
||||
]
|
||||
return any(phrase in error_body.lower() for phrase in permission_phrases)
|
||||
return False
|
||||
|
||||
|
||||
def make_github_request(url, token, method="GET", data=None):
|
||||
"""Make authenticated request to GitHub API"""
|
||||
headers = {
|
||||
"Accept": "application/vnd.github+json",
|
||||
"Authorization": f"Bearer {token}",
|
||||
# "User-Agent": "sglang-ci",
|
||||
"X-GitHub-Api-Version": "2022-11-28",
|
||||
}
|
||||
|
||||
if data:
|
||||
headers["Content-Type"] = "application/json"
|
||||
data = json.dumps(data).encode("utf-8")
|
||||
|
||||
req = Request(url, data=data, headers=headers, method=method)
|
||||
|
||||
try:
|
||||
with urlopen(req) as response:
|
||||
return response.read().decode("utf-8")
|
||||
except HTTPError as e:
|
||||
print(f"GitHub API request failed: {e}")
|
||||
try:
|
||||
error_body = e.read().decode("utf-8")
|
||||
print(f"Error response body: {error_body}")
|
||||
e.error_body = error_body # Attach for later inspection
|
||||
except Exception:
|
||||
e.error_body = ""
|
||||
raise
|
||||
except Exception as e:
|
||||
print(f"GitHub API request failed with a non-HTTP error: {e}")
|
||||
raise
|
||||
|
||||
|
||||
def verify_token_permissions(repo_owner, repo_name, token):
|
||||
"""Verify that the token has necessary permissions for the repository"""
|
||||
print("Verifying token permissions...")
|
||||
|
||||
checks = [
|
||||
(
|
||||
f"https://api.github.com/repos/{repo_owner}/{repo_name}", # Check if we can access the repository
|
||||
"Repository access verified",
|
||||
),
|
||||
(
|
||||
f"https://api.github.com/repos/{repo_owner}/{repo_name}/contents", # Check if we can read the repository contents
|
||||
"Repository contents access verified",
|
||||
),
|
||||
]
|
||||
|
||||
for url, success_message in checks:
|
||||
try:
|
||||
response = make_github_request(url, token)
|
||||
if success_message == "Repository access verified":
|
||||
repo_data = json.loads(response)
|
||||
print(f"{success_message}: {repo_data['full_name']}")
|
||||
else:
|
||||
print(success_message)
|
||||
except Exception as e:
|
||||
if is_rate_limit_error(e):
|
||||
warnings.warn(
|
||||
"GitHub API rate limit exceeded during token verification."
|
||||
)
|
||||
return "rate_limited"
|
||||
print(f"Failed to verify permissions for {url}: {e}")
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
|
||||
def get_branch_sha(repo_owner, repo_name, branch, token):
|
||||
"""Get SHA of the branch head"""
|
||||
url = (
|
||||
f"https://api.github.com/repos/{repo_owner}/{repo_name}/git/refs/heads/{branch}"
|
||||
)
|
||||
response = make_github_request(url, token)
|
||||
data = json.loads(response)
|
||||
return data["object"]["sha"]
|
||||
|
||||
|
||||
def get_tree_sha(repo_owner, repo_name, commit_sha, token):
|
||||
"""Get tree SHA from commit"""
|
||||
url = f"https://api.github.com/repos/{repo_owner}/{repo_name}/git/commits/{commit_sha}"
|
||||
response = make_github_request(url, token)
|
||||
data = json.loads(response)
|
||||
return data["tree"]["sha"]
|
||||
|
||||
|
||||
def create_blob(repo_owner, repo_name, content, token, max_retries=3):
|
||||
"""Create a blob with file content"""
|
||||
url = f"https://api.github.com/repos/{repo_owner}/{repo_name}/git/blobs"
|
||||
|
||||
# Encode content as base64 for GitHub API
|
||||
content_b64 = base64.b64encode(content).decode("utf-8")
|
||||
|
||||
data = {"content": content_b64, "encoding": "base64"}
|
||||
|
||||
for attempt in range(max_retries):
|
||||
try:
|
||||
response = make_github_request(url, token, method="POST", data=data)
|
||||
return json.loads(response)["sha"]
|
||||
except Exception as e:
|
||||
# Don't retry on rate limit errors - fail fast
|
||||
if is_rate_limit_error(e):
|
||||
raise
|
||||
|
||||
if attempt < max_retries - 1:
|
||||
wait_time = 2**attempt # Exponential backoff: 1s, 2s, 4s
|
||||
print(
|
||||
f"Blob creation failed (attempt {attempt + 1}/{max_retries}), retrying in {wait_time}s..."
|
||||
)
|
||||
time.sleep(wait_time)
|
||||
else:
|
||||
raise
|
||||
|
||||
|
||||
def create_blobs(repo_owner, repo_name, files, token):
|
||||
"""Create blobs for all files and return tree items with blob SHAs"""
|
||||
tree_items = []
|
||||
for i, (file_path, content) in enumerate(files):
|
||||
# Create blob first to get SHA
|
||||
blob_sha = create_blob(repo_owner, repo_name, content, token)
|
||||
tree_items.append(
|
||||
{
|
||||
"path": file_path,
|
||||
"mode": "100644",
|
||||
"type": "blob",
|
||||
"sha": blob_sha,
|
||||
}
|
||||
)
|
||||
# Progress indicator for large uploads
|
||||
if (i + 1) % 10 == 0 or (i + 1) == len(files):
|
||||
print(f"Created {i + 1}/{len(files)} blobs...")
|
||||
return tree_items
|
||||
|
||||
|
||||
def create_tree(repo_owner, repo_name, base_tree_sha, tree_items, token, max_retries=3):
|
||||
"""Create a new tree from pre-created blob SHAs"""
|
||||
url = f"https://api.github.com/repos/{repo_owner}/{repo_name}/git/trees"
|
||||
|
||||
data = {"base_tree": base_tree_sha, "tree": tree_items}
|
||||
|
||||
for attempt in range(max_retries):
|
||||
try:
|
||||
response = make_github_request(url, token, method="POST", data=data)
|
||||
return json.loads(response)["sha"]
|
||||
except Exception as e:
|
||||
# Don't retry on rate limit errors - fail fast
|
||||
if is_rate_limit_error(e):
|
||||
raise
|
||||
|
||||
if attempt < max_retries - 1:
|
||||
wait_time = 2**attempt
|
||||
print(
|
||||
f"Tree creation failed (attempt {attempt + 1}/{max_retries}), retrying in {wait_time}s..."
|
||||
)
|
||||
time.sleep(wait_time)
|
||||
else:
|
||||
raise
|
||||
|
||||
|
||||
def create_commit(
|
||||
repo_owner, repo_name, tree_sha, parent_sha, message, token, max_retries=3
|
||||
):
|
||||
"""Create a new commit"""
|
||||
url = f"https://api.github.com/repos/{repo_owner}/{repo_name}/git/commits"
|
||||
|
||||
data = {"tree": tree_sha, "parents": [parent_sha], "message": message}
|
||||
|
||||
for attempt in range(max_retries):
|
||||
try:
|
||||
response = make_github_request(url, token, method="POST", data=data)
|
||||
commit_sha = json.loads(response)["sha"]
|
||||
|
||||
# Verify the commit was actually created
|
||||
verify_url = f"https://api.github.com/repos/{repo_owner}/{repo_name}/git/commits/{commit_sha}"
|
||||
verify_response = make_github_request(verify_url, token)
|
||||
verify_data = json.loads(verify_response)
|
||||
if verify_data["sha"] != commit_sha:
|
||||
raise Exception(
|
||||
f"Commit verification failed: expected {commit_sha}, got {verify_data['sha']}"
|
||||
)
|
||||
|
||||
return commit_sha
|
||||
except Exception as e:
|
||||
# Don't retry on rate limit errors - fail fast
|
||||
if is_rate_limit_error(e):
|
||||
raise
|
||||
|
||||
if attempt < max_retries - 1:
|
||||
wait_time = 2**attempt
|
||||
print(
|
||||
f"Commit creation failed (attempt {attempt + 1}/{max_retries}), retrying in {wait_time}s..."
|
||||
)
|
||||
time.sleep(wait_time)
|
||||
else:
|
||||
raise
|
||||
|
||||
|
||||
def update_branch_ref(repo_owner, repo_name, branch, commit_sha, token, max_retries=3):
|
||||
"""Update branch reference to point to new commit"""
|
||||
url = (
|
||||
f"https://api.github.com/repos/{repo_owner}/{repo_name}/git/refs/heads/{branch}"
|
||||
)
|
||||
|
||||
data = {"sha": commit_sha}
|
||||
|
||||
for attempt in range(max_retries):
|
||||
try:
|
||||
make_github_request(url, token, method="PATCH", data=data)
|
||||
return
|
||||
except HTTPError as e:
|
||||
# Don't retry on rate limit errors - fail fast
|
||||
if is_rate_limit_error(e):
|
||||
raise
|
||||
|
||||
# Check if this is an "Object does not exist" error
|
||||
is_object_not_exist = False
|
||||
if hasattr(e, "error_body"):
|
||||
try:
|
||||
error_data = json.loads(e.error_body)
|
||||
if "Object does not exist" in error_data.get("message", ""):
|
||||
is_object_not_exist = True
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
if is_object_not_exist and attempt < max_retries - 1:
|
||||
# This might be a transient consistency issue - wait and retry
|
||||
wait_time = 2**attempt
|
||||
print(
|
||||
f"Branch update failed with 'Object does not exist' (attempt {attempt + 1}/{max_retries}), waiting {wait_time}s for consistency..."
|
||||
)
|
||||
time.sleep(wait_time)
|
||||
else:
|
||||
raise
|
||||
except Exception as e:
|
||||
# Don't retry on rate limit errors - fail fast
|
||||
if is_rate_limit_error(e):
|
||||
raise
|
||||
|
||||
if attempt < max_retries - 1:
|
||||
wait_time = 2**attempt
|
||||
print(
|
||||
f"Branch update failed (attempt {attempt + 1}/{max_retries}), retrying in {wait_time}s..."
|
||||
)
|
||||
time.sleep(wait_time)
|
||||
else:
|
||||
raise
|
||||
|
||||
|
||||
def copy_trace_files(source_dir, target_base_path):
|
||||
"""Copy trace files and return list of files to upload.
|
||||
|
||||
Only uploads traces from TP rank 0 to avoid duplicated data across tensor parallel ranks.
|
||||
"""
|
||||
files_to_upload = []
|
||||
|
||||
if not os.path.exists(source_dir):
|
||||
print(f"Warning: Traces directory {source_dir} does not exist")
|
||||
return files_to_upload
|
||||
|
||||
# Walk through source directory and find .json.gz files
|
||||
for root, dirs, files in os.walk(source_dir):
|
||||
for file in files:
|
||||
if file.endswith(".json.gz"):
|
||||
|
||||
# Only upload TP rank 0 traces to avoid duplicates across tensor parallel ranks
|
||||
if "TP-" in file and "TP-0" not in file:
|
||||
continue
|
||||
|
||||
source_file = os.path.join(root, file)
|
||||
# Calculate relative path from source_dir
|
||||
rel_path = os.path.relpath(source_file, source_dir)
|
||||
target_path = f"{target_base_path}/{rel_path}"
|
||||
|
||||
# Read file content
|
||||
with open(source_file, "rb") as f:
|
||||
content = f.read()
|
||||
|
||||
files_to_upload.append((target_path, content))
|
||||
|
||||
return files_to_upload
|
||||
|
||||
|
||||
def publish_traces(traces_dir, run_id, run_number):
|
||||
"""Publish traces from a single directory to GitHub repository in a single commit"""
|
||||
target_base_path = f"traces/{run_id}"
|
||||
files_to_upload = copy_trace_files(traces_dir, target_base_path)
|
||||
|
||||
if not files_to_upload:
|
||||
print("No trace files found to upload")
|
||||
return
|
||||
|
||||
print(f"Found {len(files_to_upload)} files to upload")
|
||||
publish_traces_from_files(files_to_upload, run_id, run_number)
|
||||
|
||||
|
||||
def publish_traces_from_files(files_to_upload, run_id, run_number):
|
||||
"""Publish pre-collected trace files to GitHub repository in a single commit"""
|
||||
# Get environment variables
|
||||
token = os.getenv("GITHUB_TOKEN")
|
||||
if not token:
|
||||
print("Error: GITHUB_TOKEN environment variable not set")
|
||||
sys.exit(1)
|
||||
|
||||
# Repository configuration
|
||||
repo_owner = "sgl-project"
|
||||
repo_name = "ci-data"
|
||||
branch = "main"
|
||||
|
||||
# Verify token permissions before proceeding
|
||||
permission_check = verify_token_permissions(repo_owner, repo_name, token)
|
||||
if permission_check == "rate_limited":
|
||||
warnings.warn(
|
||||
"Skipping trace upload due to GitHub API rate limit. "
|
||||
"This is expected during high CI activity and does not indicate a test failure."
|
||||
)
|
||||
return
|
||||
elif not permission_check:
|
||||
print(
|
||||
"Token permission verification failed. Please check the token permissions."
|
||||
)
|
||||
sys.exit(1)
|
||||
|
||||
max_retries = 5
|
||||
retry_delay = 5 # seconds
|
||||
|
||||
# Create blobs once before retry loop to avoid re-uploading on failures
|
||||
try:
|
||||
tree_items = create_blobs(repo_owner, repo_name, files_to_upload, token)
|
||||
except Exception as e:
|
||||
# Check for rate limit errors during blob creation
|
||||
if is_rate_limit_error(e):
|
||||
warnings.warn(
|
||||
"GitHub API rate limit exceeded during blob creation. Skipping trace upload."
|
||||
)
|
||||
return
|
||||
# Check for permission errors - these should fail loudly
|
||||
if is_permission_error(e):
|
||||
print(
|
||||
f"ERROR: Token does not have write permission to {repo_owner}/{repo_name}. "
|
||||
"Please update the GH_PAT_FOR_NIGHTLY_CI_DATA secret with a token that has "
|
||||
"'contents: write' permission for the repository."
|
||||
)
|
||||
sys.exit(1)
|
||||
print(f"Failed to create blobs: {e}")
|
||||
raise
|
||||
|
||||
for attempt in range(max_retries):
|
||||
try:
|
||||
# Get current branch head
|
||||
branch_sha = get_branch_sha(repo_owner, repo_name, branch, token)
|
||||
print(f"Current branch head: {branch_sha}")
|
||||
|
||||
# Get current tree
|
||||
tree_sha = get_tree_sha(repo_owner, repo_name, branch_sha, token)
|
||||
print(f"Current tree SHA: {tree_sha}")
|
||||
|
||||
# Create new tree with pre-created blobs
|
||||
new_tree_sha = create_tree(
|
||||
repo_owner, repo_name, tree_sha, tree_items, token
|
||||
)
|
||||
print(f"Created new tree: {new_tree_sha}")
|
||||
|
||||
# Create commit
|
||||
commit_message = f"Nightly traces for run {run_id} at {run_number} ({len(files_to_upload)} files)"
|
||||
commit_sha = create_commit(
|
||||
repo_owner,
|
||||
repo_name,
|
||||
new_tree_sha,
|
||||
branch_sha,
|
||||
commit_message,
|
||||
token,
|
||||
)
|
||||
print(f"Created commit: {commit_sha}")
|
||||
|
||||
# Update branch reference
|
||||
update_branch_ref(repo_owner, repo_name, branch, commit_sha, token)
|
||||
print("Updated branch reference")
|
||||
|
||||
print("Successfully published all traces in a single commit")
|
||||
return
|
||||
|
||||
except Exception as e:
|
||||
# Check for retryable errors
|
||||
is_retryable = False
|
||||
error_type = "unknown"
|
||||
|
||||
if hasattr(e, "error_body"):
|
||||
if "Update is not a fast forward" in e.error_body:
|
||||
is_retryable = True
|
||||
error_type = "fast-forward conflict"
|
||||
elif "Object does not exist" in e.error_body:
|
||||
is_retryable = True
|
||||
error_type = "object consistency"
|
||||
|
||||
# Also retry on HTTP errors that might be transient
|
||||
if isinstance(e, HTTPError) and e.code in [422, 500, 502, 503, 504]:
|
||||
is_retryable = True
|
||||
error_type = f"HTTP {e.code}"
|
||||
|
||||
# Check for rate limit errors (non-fatal - just warn and skip)
|
||||
if is_rate_limit_error(e):
|
||||
warnings.warn("GitHub API rate limit exceeded. Skipping trace upload.")
|
||||
return
|
||||
|
||||
# Check for permission errors - these should fail loudly
|
||||
if is_permission_error(e):
|
||||
print(
|
||||
f"ERROR: Token does not have write permission to {repo_owner}/{repo_name}. "
|
||||
"Please update the GH_PAT_FOR_NIGHTLY_CI_DATA secret with a token that has "
|
||||
"'contents: write' permission for the repository."
|
||||
)
|
||||
sys.exit(1)
|
||||
|
||||
if is_retryable and attempt < max_retries - 1:
|
||||
print(
|
||||
f"Attempt {attempt + 1}/{max_retries} failed ({error_type}). Retrying in {retry_delay} seconds..."
|
||||
)
|
||||
time.sleep(retry_delay)
|
||||
else:
|
||||
print(f"Failed to publish traces after {attempt + 1} attempts: {e}")
|
||||
raise
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Publish performance traces to GitHub repository"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--traces-dir",
|
||||
type=str,
|
||||
action="append",
|
||||
dest="traces_dirs",
|
||||
required=True,
|
||||
help="Traces directory to publish (can be specified multiple times)",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
# Get environment variables
|
||||
run_id = os.getenv("GITHUB_RUN_ID", "test")
|
||||
run_number = os.getenv("GITHUB_RUN_NUMBER", "12345")
|
||||
|
||||
if not run_id or not run_number:
|
||||
print(
|
||||
"Error: GITHUB_RUN_ID and GITHUB_RUN_NUMBER environment variables must be set"
|
||||
)
|
||||
sys.exit(1)
|
||||
|
||||
# Collect trace files from all directories
|
||||
target_base_path = f"traces/{run_id}"
|
||||
all_files = []
|
||||
for traces_dir in args.traces_dirs:
|
||||
print(f"Processing traces from directory: {traces_dir}")
|
||||
files = copy_trace_files(traces_dir, target_base_path)
|
||||
all_files.extend(files)
|
||||
|
||||
if not all_files:
|
||||
print("No trace files found to upload across all directories")
|
||||
return
|
||||
|
||||
print(f"Found {len(all_files)} total files to upload")
|
||||
|
||||
# Publish all collected traces in a single commit
|
||||
publish_traces_from_files(all_files, run_id, run_number)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Executable
+1943
File diff suppressed because it is too large
Load Diff
+763
@@ -0,0 +1,763 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Runner Utilization Report
|
||||
|
||||
Analyzes GitHub Actions job data to calculate runner utilization metrics.
|
||||
Reports idle time, active time, and utilization percentage per runner label.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import random
|
||||
import subprocess
|
||||
import time
|
||||
from collections import Counter, defaultdict
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
from datetime import datetime, timedelta, timezone
|
||||
|
||||
# Labels to skip when grouping runners (GitHub default labels)
|
||||
DEFAULT_LABELS_TO_IGNORE = {"self-hosted", "Linux", "X64", "ARM64"}
|
||||
GITHUB_HOSTED_LABELS = {"ubuntu-latest", "ubuntu-22.04", "ubuntu-24.04"}
|
||||
|
||||
# Human-facing job outcome buckets, in display order, with emoji.
|
||||
STATUS_ORDER = ("pass", "fail", "cancel", "running", "queued")
|
||||
STATUS_EMOJI = {
|
||||
"pass": "✅",
|
||||
"fail": "❌",
|
||||
"cancel": "🚫",
|
||||
"running": "🔄",
|
||||
"queued": "⏳",
|
||||
}
|
||||
|
||||
|
||||
def format_status_counts(counts: dict) -> str:
|
||||
"""Compact per-label outcome summary, e.g. '✅120 ❌3 🔄2 ⏳4'."""
|
||||
parts = [f"{STATUS_EMOJI[s]}{counts[s]}" for s in STATUS_ORDER if counts.get(s)]
|
||||
return " ".join(parts) if parts else "—"
|
||||
|
||||
|
||||
def run_gh_command(args: list[str], max_retries: int = 10) -> dict:
|
||||
"""Run gh CLI command and return JSON result.
|
||||
|
||||
Retries on transient failures (5xx, secondary rate limits, network
|
||||
blips) with exponential backoff. The previous fail-fast behavior
|
||||
combined with `except Exception: return None` in the threadpool
|
||||
callers caused entire workflow runs to be silently dropped from
|
||||
the utilization numerator whenever GH API hiccuped, severely
|
||||
undercounting busy time on busy days.
|
||||
"""
|
||||
last_err = ""
|
||||
for attempt in range(max_retries):
|
||||
result = subprocess.run(
|
||||
["gh", "api"] + args,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
)
|
||||
if result.returncode == 0:
|
||||
return json.loads(result.stdout)
|
||||
last_err = result.stderr or "(no stderr)"
|
||||
# Detect retryable conditions: HTTP 5xx, secondary rate limit, abuse
|
||||
# detection, network resets. 4xx other than 429 are non-retryable.
|
||||
retryable = any(
|
||||
s in last_err
|
||||
for s in (
|
||||
"rate limit",
|
||||
"abuse",
|
||||
"Internal Server Error",
|
||||
"502",
|
||||
"503",
|
||||
"504",
|
||||
"Bad Gateway",
|
||||
"Gateway Time-out",
|
||||
"connection reset",
|
||||
"Connection reset",
|
||||
"EOF",
|
||||
"timeout",
|
||||
)
|
||||
)
|
||||
if not retryable:
|
||||
break
|
||||
# Exponential backoff with jitter, capped at 60s.
|
||||
delay = min(60, (2**attempt) + random.uniform(0, 1))
|
||||
time.sleep(delay)
|
||||
raise Exception(f"gh api failed after {max_retries} attempts: {last_err[:300]}")
|
||||
|
||||
|
||||
def get_workflow_runs(repo: str, hours: int = 24) -> list[dict]:
|
||||
"""Get workflow runs from the last N hours."""
|
||||
since = datetime.now(timezone.utc) - timedelta(hours=hours)
|
||||
|
||||
runs = []
|
||||
page = 1
|
||||
while True:
|
||||
data = run_gh_command(
|
||||
[
|
||||
f"repos/{repo}/actions/runs?per_page=100&page={page}",
|
||||
]
|
||||
)
|
||||
page_runs = data.get("workflow_runs", [])
|
||||
|
||||
# Filter by time
|
||||
for run in page_runs:
|
||||
created_at = parse_time(run.get("created_at"))
|
||||
if created_at and created_at >= since:
|
||||
runs.append(run)
|
||||
elif created_at and created_at < since:
|
||||
# Runs are ordered by created_at desc, so we can stop
|
||||
return runs
|
||||
|
||||
if len(page_runs) < 100:
|
||||
break
|
||||
page += 1
|
||||
if page > 50: # Safety limit (5000 runs)
|
||||
break
|
||||
return runs
|
||||
|
||||
|
||||
def get_jobs_for_run(repo: str, run_id: int) -> list[dict]:
|
||||
"""Get all jobs for a workflow run, including all retry attempts.
|
||||
|
||||
`filter=all` is required so that re-run attempts of the same job
|
||||
appear separately. Each attempt consumed host time on the runner
|
||||
pool, so for utilization we want them all summed in. The default
|
||||
(`filter=latest`) only returns the most recent attempt and silently
|
||||
hides time spent on prior retries.
|
||||
"""
|
||||
jobs = []
|
||||
page = 1
|
||||
while True:
|
||||
data = run_gh_command(
|
||||
[
|
||||
f"repos/{repo}/actions/runs/{run_id}/jobs"
|
||||
f"?per_page=100&page={page}&filter=all",
|
||||
]
|
||||
)
|
||||
jobs.extend(data.get("jobs", []))
|
||||
if len(data.get("jobs", [])) < 100:
|
||||
break
|
||||
page += 1
|
||||
if page > 20: # Safety limit (2000 jobs per run)
|
||||
break
|
||||
return jobs
|
||||
|
||||
|
||||
def get_runners(repo: str, online_only: bool = True) -> list[dict]:
|
||||
"""Get all self-hosted runners with pagination. Returns empty if no permission."""
|
||||
try:
|
||||
all_runners = []
|
||||
page = 1
|
||||
while True:
|
||||
data = run_gh_command(
|
||||
[f"repos/{repo}/actions/runners?per_page=100&page={page}"]
|
||||
)
|
||||
runners = data.get("runners", [])
|
||||
all_runners.extend(runners)
|
||||
if len(runners) < 100:
|
||||
break
|
||||
page += 1
|
||||
if page > 10: # Safety limit
|
||||
break
|
||||
if online_only:
|
||||
all_runners = [r for r in all_runners if r.get("status") == "online"]
|
||||
return all_runners
|
||||
except Exception as e:
|
||||
print(f"Warning: Cannot access runners API (need admin): {e}")
|
||||
return []
|
||||
|
||||
|
||||
def parse_time(time_str: str) -> datetime:
|
||||
"""Parse ISO timestamp to datetime."""
|
||||
if not time_str:
|
||||
return None
|
||||
return datetime.fromisoformat(time_str.replace("Z", "+00:00"))
|
||||
|
||||
|
||||
def classify_job(job: dict, now: datetime):
|
||||
"""Derive the queue-wait and busy interval for a single job.
|
||||
|
||||
Returns a job_info dict, or None when the job neither waited for nor
|
||||
occupied a runner (skipped / cancelled-before-start / missing data).
|
||||
|
||||
The queue wait runs from when the job entered the runner queue
|
||||
(`created_at`) until a runner picked it up (`started_at`) — or until
|
||||
`now` if it is still waiting.
|
||||
|
||||
GitHub API gotcha this exists to handle: a still-queued job reports
|
||||
status="queued", runner_name="" and `started_at` set to a PLACEHOLDER
|
||||
equal to `created_at` (not null). The previous code required both a
|
||||
runner_name and a `completed_at`, so every in-flight wait — the
|
||||
multi-hour 8-gpu jobs still sitting in the queue, i.e. the worst cases —
|
||||
was dropped, undercounting max/avg queue time. We therefore measure a
|
||||
queued job's wait against `now` rather than its bogus `started_at`, and
|
||||
don't require completion.
|
||||
"""
|
||||
status = job.get("status")
|
||||
runner_name = job.get("runner_name") or ""
|
||||
created_at = parse_time(job.get("created_at"))
|
||||
started_at = parse_time(job.get("started_at"))
|
||||
completed_at = parse_time(job.get("completed_at"))
|
||||
|
||||
if status == "queued":
|
||||
# Still waiting for a runner; ignore the placeholder started_at.
|
||||
queue_end, start, end = now, None, None
|
||||
elif status == "in_progress" and started_at is not None:
|
||||
# Running now: the wait is final and it still occupies the runner.
|
||||
queue_end, start, end = started_at, started_at, now
|
||||
elif (
|
||||
status == "completed"
|
||||
and started_at is not None
|
||||
and completed_at is not None
|
||||
and runner_name
|
||||
):
|
||||
queue_end, start, end = started_at, started_at, completed_at
|
||||
else:
|
||||
# Skipped, cancelled before start, or missing timestamps: never
|
||||
# waited for or occupied a runner.
|
||||
return None
|
||||
|
||||
if created_at is None:
|
||||
return None
|
||||
|
||||
queue_time = max(0.0, (queue_end - created_at).total_seconds())
|
||||
duration = (end - start).total_seconds() if start is not None else 0.0
|
||||
labels = [
|
||||
label
|
||||
for label in job.get("labels", [])
|
||||
if label not in DEFAULT_LABELS_TO_IGNORE | GITHUB_HOSTED_LABELS
|
||||
]
|
||||
|
||||
# Human-facing outcome bucket used by the report's status breakdown.
|
||||
if status == "queued":
|
||||
outcome = "queued"
|
||||
elif status == "in_progress":
|
||||
outcome = "running"
|
||||
else: # completed and actually ran
|
||||
outcome = {"success": "pass", "cancelled": "cancel"}.get(
|
||||
job.get("conclusion"), "fail"
|
||||
)
|
||||
|
||||
return {
|
||||
"start": start,
|
||||
"end": end,
|
||||
"created_at": created_at,
|
||||
"queue_end": queue_end,
|
||||
"duration": duration,
|
||||
"queue_time": queue_time,
|
||||
"job_name": job.get("name", ""),
|
||||
"runner_name": runner_name,
|
||||
"labels": labels,
|
||||
"status": outcome,
|
||||
"html_url": job.get("html_url", ""),
|
||||
}
|
||||
|
||||
|
||||
def calculate_concurrency_metrics(
|
||||
jobs: list,
|
||||
window_start: datetime,
|
||||
window_end: datetime,
|
||||
num_runners: int,
|
||||
) -> dict:
|
||||
"""Sweep-line algorithm: peak/avg concurrent, saturation time, peak queue."""
|
||||
if not jobs:
|
||||
return {
|
||||
"peak_concurrent": 0,
|
||||
"avg_concurrent": 0.0,
|
||||
"saturation_seconds": 0,
|
||||
"saturation_pct": 0.0,
|
||||
"peak_queue": 0,
|
||||
}
|
||||
window_seconds = (window_end - window_start).total_seconds()
|
||||
if window_seconds <= 0:
|
||||
return {
|
||||
"peak_concurrent": 0,
|
||||
"avg_concurrent": 0.0,
|
||||
"saturation_seconds": 0,
|
||||
"saturation_pct": 0.0,
|
||||
"peak_queue": 0,
|
||||
}
|
||||
running_events = []
|
||||
for job in jobs:
|
||||
start, end = job["start"], job["end"]
|
||||
# Still-queued jobs have no running interval yet (start/end are None).
|
||||
if start is None or end is None:
|
||||
continue
|
||||
if end < window_start or start > window_end:
|
||||
continue
|
||||
running_events.append((max(start, window_start), 1))
|
||||
running_events.append((min(end, window_end), -1))
|
||||
queue_events = []
|
||||
for job in jobs:
|
||||
created_at = job.get("created_at")
|
||||
# The wait ends when a runner picks the job up, or `now` if it is
|
||||
# still queued (queue_end was set to now upstream). Counting the
|
||||
# still-open waits is what makes peak_queue reflect the real backlog.
|
||||
queue_end = job.get("queue_end") or job["start"]
|
||||
if created_at and queue_end and created_at < queue_end:
|
||||
if queue_end < window_start or created_at > window_end:
|
||||
continue
|
||||
queue_events.append((max(created_at, window_start), 1))
|
||||
queue_events.append((min(queue_end, window_end), -1))
|
||||
running_events.sort(key=lambda e: (e[0], e[1] == 1))
|
||||
current_running = 0
|
||||
peak_running = 0
|
||||
prev_time = window_start
|
||||
total_running_seconds = 0.0
|
||||
saturation_seconds = 0.0
|
||||
for event_time, delta in running_events:
|
||||
td = (event_time - prev_time).total_seconds()
|
||||
if td > 0:
|
||||
total_running_seconds += current_running * td
|
||||
if current_running >= num_runners:
|
||||
saturation_seconds += td
|
||||
current_running += delta
|
||||
peak_running = max(peak_running, current_running)
|
||||
prev_time = event_time
|
||||
if prev_time < window_end:
|
||||
td = (window_end - prev_time).total_seconds()
|
||||
total_running_seconds += current_running * td
|
||||
if current_running >= num_runners:
|
||||
saturation_seconds += td
|
||||
queue_events.sort(key=lambda e: (e[0], e[1] == 1))
|
||||
current_queued = 0
|
||||
peak_queue = 0
|
||||
for _, delta in queue_events:
|
||||
current_queued += delta
|
||||
peak_queue = max(peak_queue, current_queued)
|
||||
avg_concurrent = total_running_seconds / window_seconds if window_seconds > 0 else 0
|
||||
return {
|
||||
"peak_concurrent": peak_running,
|
||||
"avg_concurrent": avg_concurrent,
|
||||
"saturation_seconds": saturation_seconds,
|
||||
"saturation_pct": (
|
||||
(saturation_seconds / window_seconds * 100) if window_seconds > 0 else 0
|
||||
),
|
||||
"peak_queue": peak_queue,
|
||||
}
|
||||
|
||||
|
||||
_NON_GPU_WORKFLOW_HINTS = (
|
||||
"lint",
|
||||
"deploy",
|
||||
"release",
|
||||
"publish",
|
||||
"docs",
|
||||
"doc",
|
||||
"mintlify",
|
||||
"runner utilization", # this very script
|
||||
"tag-and-rerun",
|
||||
"auto", # auto-merge etc.
|
||||
"label",
|
||||
"stale",
|
||||
"dependabot",
|
||||
"codeql",
|
||||
)
|
||||
|
||||
|
||||
def _likely_no_gpu_jobs(workflow_name: str) -> bool:
|
||||
"""Heuristic: skip per-run job-fetch for workflows that don't dispatch
|
||||
to self-hosted GPU runners. The GH API rate limit (~5000 req/hr per
|
||||
token) is the bottleneck on busy 24h windows where ~4000 workflow
|
||||
runs fire — but only a fraction of those (pr-test, nightly-test,
|
||||
pr-test-*kernel, etc.) actually run on GPU runners. Skipping the
|
||||
docs/lint/release runs cuts the API call budget by 2-4x.
|
||||
"""
|
||||
if not workflow_name:
|
||||
return False
|
||||
n = workflow_name.lower()
|
||||
return any(h in n for h in _NON_GPU_WORKFLOW_HINTS)
|
||||
|
||||
|
||||
def calculate_utilization(repo: str, hours: int = 24, runner_filter: str = None):
|
||||
"""Calculate runner utilization metrics."""
|
||||
|
||||
print(f"Fetching workflow runs from last {hours} hours...")
|
||||
all_runs = get_workflow_runs(repo, hours)
|
||||
runs = [r for r in all_runs if not _likely_no_gpu_jobs(r.get("name", ""))]
|
||||
skipped = len(all_runs) - len(runs)
|
||||
print(
|
||||
f"Found {len(all_runs)} workflow runs "
|
||||
f"({skipped} skipped as non-GPU: docs/lint/release/etc.)"
|
||||
)
|
||||
|
||||
# Try to get online runners from API
|
||||
print("Fetching online runners...")
|
||||
runners = get_runners(repo, online_only=True)
|
||||
|
||||
# Build label -> set of online runner names from API
|
||||
api_label_runners = defaultdict(set)
|
||||
if runners:
|
||||
for runner in runners:
|
||||
for label in runner.get("labels", []):
|
||||
label_name = label.get("name", "")
|
||||
if label_name not in DEFAULT_LABELS_TO_IGNORE:
|
||||
api_label_runners[label_name].add(runner["name"])
|
||||
print(f"Got {len(runners)} online runners from API")
|
||||
else:
|
||||
print("No runner API access, will use observed runners from job data")
|
||||
|
||||
# Track runners seen in jobs (for labels not in API or when API unavailable)
|
||||
job_label_runners = defaultdict(set)
|
||||
label_jobs = defaultdict(list) # label -> list of job_info
|
||||
# Per-host accumulation: each physical machine appears once regardless of
|
||||
# how many overlapping labels it advertises. This is what we use for the
|
||||
# "Per Host Utilization" section (the source-of-truth view).
|
||||
host_jobs = defaultdict(list) # runner_name -> list of job_info
|
||||
host_labels = defaultdict(set) # runner_name -> set of labels it ran jobs under
|
||||
|
||||
# Fetch jobs for all runs in parallel. Cap concurrency lower than the
|
||||
# GH API secondary rate-limit threshold to avoid bursts that silently
|
||||
# drop runs even with retries.
|
||||
total_runs = len(runs)
|
||||
print(f"Fetching jobs for {total_runs} runs in parallel...")
|
||||
|
||||
def fetch_jobs_for_run(run):
|
||||
"""Fetch jobs for a single run.
|
||||
|
||||
Returns (run_id, jobs, error_msg). `error_msg` is None on success.
|
||||
We surface failures rather than silently dropping the run so the
|
||||
caller can report how many runs' jobs are missing — silently
|
||||
dropping previously caused 4-gpu-b200 (and every other label) to
|
||||
report wildly different numbers depending on transient API hiccups.
|
||||
"""
|
||||
try:
|
||||
return (run["id"], get_jobs_for_run(repo, run["id"]), None)
|
||||
except Exception as e:
|
||||
return (run["id"], None, str(e)[:200])
|
||||
|
||||
all_jobs = []
|
||||
failed_runs = []
|
||||
# Concurrency=4 with longer retry budget keeps us well below the GH
|
||||
# API secondary rate-limit threshold (~10 req/s). On a 24h window
|
||||
# with ~1500 GPU-relevant runs (post-filter), this completes in ~5
|
||||
# min and almost never hits the rate limit.
|
||||
with ThreadPoolExecutor(max_workers=4) as executor:
|
||||
futures = [executor.submit(fetch_jobs_for_run, run) for run in runs]
|
||||
completed = 0
|
||||
for future in as_completed(futures):
|
||||
completed += 1
|
||||
if completed % 100 == 0:
|
||||
print(
|
||||
f"Fetched jobs for {completed}/{total_runs} runs "
|
||||
f"({len(failed_runs)} failed so far)..."
|
||||
)
|
||||
run_id, jobs, err = future.result()
|
||||
if err:
|
||||
failed_runs.append((run_id, err))
|
||||
elif jobs:
|
||||
all_jobs.extend(jobs)
|
||||
|
||||
print(f"Processing {len(all_jobs)} jobs...")
|
||||
if failed_runs:
|
||||
print(
|
||||
f"WARNING: {len(failed_runs)}/{total_runs} runs failed to fetch "
|
||||
f"after retries. Utilization will be undercounted. "
|
||||
f"First few errors:"
|
||||
)
|
||||
for rid, err in failed_runs[:5]:
|
||||
print(f" run {rid}: {err}")
|
||||
fetch_failure_pct = len(failed_runs) / total_runs * 100 if total_runs > 0 else 0
|
||||
|
||||
# `now` anchors the wait of jobs that are still queued or running. It is
|
||||
# captured once so every in-flight job is measured against a single
|
||||
# reference (matches window_end below to within processing time).
|
||||
now = datetime.now(timezone.utc)
|
||||
all_job_infos = [] # one entry per job (deduped across labels) for detail views
|
||||
for job in all_jobs:
|
||||
job_info = classify_job(job, now)
|
||||
if job_info is None:
|
||||
continue
|
||||
all_job_infos.append(job_info)
|
||||
runner_name = job_info["runner_name"]
|
||||
|
||||
# Per-host busy time only applies to jobs that actually occupied a
|
||||
# runner (ran or still running); a still-queued job has no host yet.
|
||||
if job_info["start"] is not None and runner_name:
|
||||
host_jobs[runner_name].append(job_info)
|
||||
|
||||
for label in job_info["labels"]:
|
||||
if runner_name:
|
||||
job_label_runners[label].add(runner_name)
|
||||
host_labels[runner_name].add(label)
|
||||
label_jobs[label].append(job_info)
|
||||
|
||||
# Merge API runners and job-observed runners
|
||||
# Prefer API count (online runners) when available
|
||||
# Include labels seen only on still-queued jobs (no online runner, no
|
||||
# completed job under them yet) so a fully-backed-up pool still reports.
|
||||
all_labels = (
|
||||
set(api_label_runners.keys())
|
||||
| set(job_label_runners.keys())
|
||||
| set(label_jobs.keys())
|
||||
)
|
||||
|
||||
# Filter labels if specified
|
||||
if runner_filter:
|
||||
all_labels = {lbl for lbl in all_labels if runner_filter in lbl}
|
||||
|
||||
print(f"Tracking {len(all_labels)} runner labels: {sorted(all_labels)}")
|
||||
|
||||
window_seconds = hours * 3600
|
||||
window_end = datetime.now(timezone.utc)
|
||||
window_start = window_end - timedelta(hours=hours)
|
||||
|
||||
# Per-host window-clamped busy time (each physical machine counted once).
|
||||
# This is the source of truth for how loaded each host actually is.
|
||||
host_busy_seconds = {}
|
||||
for host, jobs in host_jobs.items():
|
||||
busy = 0.0
|
||||
for j in jobs:
|
||||
cs = max(j["start"], window_start)
|
||||
ce = min(j["end"], window_end)
|
||||
if ce > cs:
|
||||
busy += (ce - cs).total_seconds()
|
||||
host_busy_seconds[host] = busy
|
||||
|
||||
results = []
|
||||
for label in sorted(all_labels):
|
||||
# Hosts to attribute to this label = union of currently-online
|
||||
# runners advertising the label PLUS hosts that actually ran a
|
||||
# job under it during the window. The union catches hosts that
|
||||
# went offline mid-window (their busy time is still real
|
||||
# capacity consumed) and hosts that came online late.
|
||||
hosts = api_label_runners.get(label, set()) | job_label_runners.get(
|
||||
label, set()
|
||||
)
|
||||
num_runners = len(hosts) if hosts else 1
|
||||
|
||||
# Pool busy time: sum of busy time across the hosts that could
|
||||
# serve this label, regardless of which sibling label actually
|
||||
# dispatched the job. This is the right denominator/numerator for
|
||||
# asking "how saturated is the underlying hardware that this
|
||||
# label depends on?" — sibling labels (e.g. `4-gpu-b200` and
|
||||
# `4-gpu-b200-low-disk`) compete for the same physical machines,
|
||||
# so their busy time should not be double-counted into separate
|
||||
# capacity buckets.
|
||||
active_seconds = sum(host_busy_seconds.get(h, 0.0) for h in hosts)
|
||||
capacity_seconds = num_runners * window_seconds
|
||||
utilization = (
|
||||
(active_seconds / capacity_seconds * 100) if capacity_seconds > 0 else 0
|
||||
)
|
||||
|
||||
# Job count + queue stats stay label-specific (only jobs that
|
||||
# were dispatched under THIS label).
|
||||
jobs = label_jobs.get(label, [])
|
||||
queue_times = [j["queue_time"] for j in jobs if j["queue_time"] > 0]
|
||||
avg_queue = sum(queue_times) / len(queue_times) if queue_times else 0
|
||||
max_queue = max(queue_times) if queue_times else 0
|
||||
# Outcome breakdown for this label (pass/fail/cancel/running/queued).
|
||||
status_counts = dict(Counter(j["status"] for j in jobs))
|
||||
|
||||
# Concurrency / saturation / queue-depth metrics. Use observed
|
||||
# peak as effective capacity if it's lower than the API count
|
||||
# (e.g. for autoscaling pools where most listeners sit idle).
|
||||
conc_initial = calculate_concurrency_metrics(
|
||||
jobs, window_start, window_end, num_runners
|
||||
)
|
||||
effective_runners = (
|
||||
min(num_runners, conc_initial["peak_concurrent"]) or num_runners
|
||||
)
|
||||
if effective_runners < num_runners and effective_runners > 0:
|
||||
conc = calculate_concurrency_metrics(
|
||||
jobs, window_start, window_end, effective_runners
|
||||
)
|
||||
else:
|
||||
conc = conc_initial
|
||||
|
||||
results.append(
|
||||
{
|
||||
"label": label,
|
||||
"num_runners": num_runners,
|
||||
"effective_runners": effective_runners,
|
||||
"num_jobs": len(jobs),
|
||||
"total_active_hours": active_seconds / 3600,
|
||||
"utilization_pct": utilization,
|
||||
"avg_queue_min": avg_queue / 60,
|
||||
"max_queue_min": max_queue / 60,
|
||||
"peak_concurrent": conc_initial["peak_concurrent"],
|
||||
"avg_concurrent": conc["avg_concurrent"],
|
||||
"saturation_hours": conc["saturation_seconds"] / 3600,
|
||||
"saturation_pct": conc["saturation_pct"],
|
||||
"peak_queue": conc["peak_queue"],
|
||||
"status_counts": status_counts,
|
||||
}
|
||||
)
|
||||
|
||||
# Per-job detail (deduped across labels), longest waits first, for the
|
||||
# links + status section of the report.
|
||||
longest_waits = sorted(all_job_infos, key=lambda j: j["queue_time"], reverse=True)
|
||||
return results, fetch_failure_pct, longest_waits
|
||||
|
||||
|
||||
def format_report(
|
||||
results: list[dict],
|
||||
hours: int,
|
||||
fetch_failure_pct: float = 0.0,
|
||||
longest_waits: list = None,
|
||||
top_n: int = 20,
|
||||
) -> str:
|
||||
"""One compact summary table — original schema, fixed columns.
|
||||
|
||||
Active (hrs) and Utilization now reflect the actual host pool's
|
||||
busy time (sum across all jobs on the hosts that advertise this
|
||||
label, regardless of which sibling label dispatched them). This
|
||||
makes the column meaningful for shared host pools — e.g.
|
||||
`4-gpu-b200` and `4-gpu-b200-low-disk` both consume the same
|
||||
physical hosts, so their utilization now reflects real hardware
|
||||
saturation instead of being divided across labels.
|
||||
"""
|
||||
lines = [
|
||||
"# Runner Utilization Report",
|
||||
"",
|
||||
f"**Time window:** Last {hours} hours · "
|
||||
f"**Generated:** {datetime.now(timezone.utc).strftime('%Y-%m-%d %H:%M UTC')}",
|
||||
"",
|
||||
]
|
||||
if fetch_failure_pct > 1.0:
|
||||
lines.append(
|
||||
f"⚠️ **Data completeness warning**: {fetch_failure_pct:.0f}% of "
|
||||
f"GPU-relevant workflow runs failed to fetch jobs after retries "
|
||||
f"(GH API rate limit). Active hours and utilization below are "
|
||||
f"under-counted by approximately this fraction."
|
||||
)
|
||||
lines.append("")
|
||||
lines.extend(
|
||||
[
|
||||
"| Label | Runners | Jobs | Active (hrs) | Utilization | Avg Queue | Max Queue | Status |",
|
||||
"|-------|---------|------|--------------|-------------|-----------|-----------|--------|",
|
||||
]
|
||||
)
|
||||
for r in results:
|
||||
bar = "█" * int(r["utilization_pct"] / 10) + "░" * (
|
||||
10 - int(r["utilization_pct"] / 10)
|
||||
)
|
||||
lines.append(
|
||||
f"| {r['label']} | {r['num_runners']} | {r['num_jobs']} | "
|
||||
f"{r['total_active_hours']:.1f} | "
|
||||
f"{r['utilization_pct']:.1f}% {bar} | "
|
||||
f"{r['avg_queue_min']:.1f}m | {r['max_queue_min']:.1f}m | "
|
||||
f"{format_status_counts(r.get('status_counts', {}))} |"
|
||||
)
|
||||
|
||||
# Longest queue waits — links to the actual jobs, with live status, so the
|
||||
# worst waits (including jobs still queued/running right now) are one click
|
||||
# away. This is the detail behind the Max Queue column.
|
||||
waits = [j for j in (longest_waits or []) if j.get("queue_time", 0) > 0][:top_n]
|
||||
if waits:
|
||||
lines.extend(
|
||||
[
|
||||
"",
|
||||
f"## Longest Queue Waits (top {len(waits)})",
|
||||
"",
|
||||
"| Wait | Status | Label | Job |",
|
||||
"|------|--------|-------|-----|",
|
||||
]
|
||||
)
|
||||
for j in waits:
|
||||
status = j.get("status", "")
|
||||
emoji = STATUS_EMOJI.get(status, "")
|
||||
label = ", ".join(j.get("labels", [])) or "—"
|
||||
name = j.get("job_name", "job")
|
||||
url = j.get("html_url", "")
|
||||
job_cell = f"[{name}]({url})" if url else name
|
||||
lines.append(
|
||||
f"| {j['queue_time'] / 60:.0f}m | {emoji} {status} | "
|
||||
f"{label} | {job_cell} |"
|
||||
)
|
||||
|
||||
# Concurrency Analysis section
|
||||
lines.extend(
|
||||
[
|
||||
"",
|
||||
"## Concurrency Analysis",
|
||||
"",
|
||||
"| Label | Runners (API/Effective) | Peak Concurrent | Avg Concurrent | Saturation Time | Peak Queue |",
|
||||
"|-------|-------------------------|-----------------|----------------|-----------------|------------|",
|
||||
]
|
||||
)
|
||||
for r in results:
|
||||
effective = r["effective_runners"]
|
||||
avg_pct = (r["avg_concurrent"] / effective * 100) if effective > 0 else 0
|
||||
runner_str = (
|
||||
f"{r['num_runners']}/{effective}"
|
||||
if effective != r["num_runners"]
|
||||
else str(r["num_runners"])
|
||||
)
|
||||
lines.append(
|
||||
f"| {r['label']} | {runner_str} | "
|
||||
f"{r['peak_concurrent']} | "
|
||||
f"{r['avg_concurrent']:.1f} ({avg_pct:.0f}%) | "
|
||||
f"{r['saturation_hours']:.1f}h ({r['saturation_pct']:.0f}%) | "
|
||||
f"{r['peak_queue']} jobs |"
|
||||
)
|
||||
|
||||
# Recommendations
|
||||
lines.extend(["", "## Recommendations", ""])
|
||||
has_recs = False
|
||||
for r in results:
|
||||
label = r["label"]
|
||||
sat_pct = r["saturation_pct"]
|
||||
peak_q = r["peak_queue"]
|
||||
effective = r["effective_runners"]
|
||||
avg_pct = (r["avg_concurrent"] / effective * 100) if effective > 0 else 0
|
||||
if sat_pct > 50 or peak_q > 5:
|
||||
lines.append(
|
||||
f"⚠️ **{label}**: High saturation ({sat_pct:.0f}%) "
|
||||
f"with queue buildup ({peak_q} jobs). Consider adding runners."
|
||||
)
|
||||
has_recs = True
|
||||
elif sat_pct > 20 or peak_q > 0:
|
||||
lines.append(
|
||||
f"📊 **{label}**: Moderate saturation ({sat_pct:.0f}%), "
|
||||
f"peak queue {peak_q} jobs. Monitor for trends."
|
||||
)
|
||||
has_recs = True
|
||||
elif avg_pct < 30 and r["num_jobs"] > 0:
|
||||
lines.append(
|
||||
f"💡 **{label}**: Low average utilization ({avg_pct:.0f}%). "
|
||||
f"Runner pool may be oversized."
|
||||
)
|
||||
has_recs = True
|
||||
else:
|
||||
lines.append(f"✓ **{label}**: Healthy utilization with minimal queueing.")
|
||||
if not has_recs and results:
|
||||
lines.append("All runner pools have healthy utilization.")
|
||||
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Generate runner utilization report")
|
||||
parser.add_argument("--repo", default="sgl-project/sglang", help="GitHub repo")
|
||||
parser.add_argument(
|
||||
"--hours", type=float, default=24, help="Time window in hours (fractional ok)"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--filter", type=str, help="Filter runner labels (e.g., '5090', 'h200')"
|
||||
)
|
||||
parser.add_argument("--output", type=str, help="Output file (default: stdout)")
|
||||
args = parser.parse_args()
|
||||
|
||||
results, fetch_failure_pct, longest_waits = calculate_utilization(
|
||||
args.repo, args.hours, args.filter
|
||||
)
|
||||
report = format_report(
|
||||
results, args.hours, fetch_failure_pct, longest_waits=longest_waits
|
||||
)
|
||||
|
||||
if args.output:
|
||||
with open(args.output, "w") as f:
|
||||
f.write(report)
|
||||
print(f"Report written to {args.output}")
|
||||
else:
|
||||
print(report)
|
||||
|
||||
# Also write to GITHUB_STEP_SUMMARY if available
|
||||
summary_file = os.environ.get("GITHUB_STEP_SUMMARY")
|
||||
if summary_file:
|
||||
with open(summary_file, "a") as f:
|
||||
f.write(report)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Executable
+245
@@ -0,0 +1,245 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Collect and save performance metrics from nightly benchmark results.
|
||||
|
||||
This script reads benchmark result JSON files from performance profile directories
|
||||
and saves them with metadata for artifact collection in CI.
|
||||
|
||||
Usage:
|
||||
python3 scripts/ci/utils/save_metrics.py \
|
||||
--gpu-config 8-gpu-h200 \
|
||||
--partition 0 \
|
||||
--run-id 12345678 \
|
||||
--output test/metrics-8gpu-h200-partition-0.json
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import glob
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
from datetime import datetime, timezone
|
||||
|
||||
|
||||
def find_result_files(search_dirs: list[str]) -> list[str]:
|
||||
"""Find all results_*.json files in the given directories."""
|
||||
result_files = set()
|
||||
for search_dir in search_dirs:
|
||||
if os.path.exists(search_dir):
|
||||
pattern = os.path.join(search_dir, "**/results_*.json")
|
||||
result_files.update(glob.glob(pattern, recursive=True))
|
||||
return list(result_files)
|
||||
|
||||
|
||||
def parse_result_file(filepath: str) -> list[dict]:
|
||||
"""Parse a benchmark result JSON file."""
|
||||
try:
|
||||
with open(filepath, "r", encoding="utf-8") as f:
|
||||
data = json.load(f)
|
||||
if isinstance(data, list):
|
||||
return data
|
||||
return [data]
|
||||
except (json.JSONDecodeError, OSError) as e:
|
||||
print(f"Warning: Failed to parse {filepath}: {e}")
|
||||
return []
|
||||
|
||||
|
||||
def transform_benchmark_result(result: dict, gpu_config: str, partition: int) -> dict:
|
||||
"""Transform a benchmark result to the metrics schema.
|
||||
|
||||
Note: input_len and output_len are preserved here for the flat benchmarks list,
|
||||
but are also used as grouping keys in benchmarks_by_io_len.
|
||||
"""
|
||||
# Handle None values safely for numeric conversions
|
||||
latency = result.get("latency")
|
||||
last_ttft = result.get("last_ttft")
|
||||
|
||||
return {
|
||||
"batch_size": result.get("batch_size"),
|
||||
"input_len": result.get("input_len"),
|
||||
"output_len": result.get("output_len"),
|
||||
"latency_ms": latency * 1000 if latency is not None else None,
|
||||
"input_throughput": result.get("input_throughput"),
|
||||
"output_throughput": result.get("output_throughput"),
|
||||
"overall_throughput": result.get("overall_throughput"),
|
||||
"ttft_ms": last_ttft * 1000 if last_ttft is not None else None,
|
||||
"acc_length": result.get("acc_length"),
|
||||
}
|
||||
|
||||
|
||||
def get_io_len_key(input_len: int, output_len: int) -> str:
|
||||
"""Generate a key for input/output length combination."""
|
||||
return f"{input_len}_{output_len}"
|
||||
|
||||
|
||||
def group_results_by_model(
|
||||
results: list[dict], gpu_config: str, partition: int
|
||||
) -> list[dict]:
|
||||
"""Group benchmark results by model, variant, and server_args.
|
||||
|
||||
Results are organized with two benchmark structures:
|
||||
- benchmarks: flat list of all benchmarks (for backward compatibility)
|
||||
- benchmarks_by_io_len: nested structure grouped by input/output length combinations
|
||||
"""
|
||||
groups = {}
|
||||
|
||||
for result in results:
|
||||
model_path = result.get("model_path", "unknown")
|
||||
run_name = result.get("run_name", "default")
|
||||
variant = run_name if run_name != "default" else None
|
||||
server_args = result.get("server_args")
|
||||
# Convert server_args list to tuple for use as dict key (lists are not hashable)
|
||||
server_args_key = tuple(server_args) if server_args else None
|
||||
|
||||
key = (model_path, variant, server_args_key)
|
||||
if key not in groups:
|
||||
groups[key] = {
|
||||
"gpu_config": gpu_config,
|
||||
"partition": partition,
|
||||
"model": model_path,
|
||||
"variant": variant,
|
||||
"server_args": server_args,
|
||||
"benchmarks": [],
|
||||
"benchmarks_by_io_len": {},
|
||||
}
|
||||
|
||||
transformed = transform_benchmark_result(result, gpu_config, partition)
|
||||
|
||||
# Add to flat benchmarks list (backward compatibility)
|
||||
groups[key]["benchmarks"].append(transformed)
|
||||
|
||||
# Add to nested benchmarks_by_io_len structure
|
||||
input_len = result.get("input_len")
|
||||
output_len = result.get("output_len")
|
||||
if input_len is not None and output_len is not None:
|
||||
io_key = get_io_len_key(input_len, output_len)
|
||||
if io_key not in groups[key]["benchmarks_by_io_len"]:
|
||||
groups[key]["benchmarks_by_io_len"][io_key] = {
|
||||
"input_len": input_len,
|
||||
"output_len": output_len,
|
||||
"benchmarks": [],
|
||||
}
|
||||
# For the nested structure, exclude input_len and output_len from individual benchmarks
|
||||
# since they're already in the parent
|
||||
nested_benchmark = {
|
||||
k: v
|
||||
for k, v in transformed.items()
|
||||
if k not in ("input_len", "output_len")
|
||||
}
|
||||
groups[key]["benchmarks_by_io_len"][io_key]["benchmarks"].append(
|
||||
nested_benchmark
|
||||
)
|
||||
|
||||
return list(groups.values())
|
||||
|
||||
|
||||
def save_metrics(
|
||||
gpu_config: str,
|
||||
partition: int,
|
||||
run_id: str,
|
||||
output_file: str,
|
||||
search_dirs: list[str],
|
||||
) -> bool:
|
||||
"""Collect metrics and save to output file."""
|
||||
timestamp = datetime.now(timezone.utc).isoformat()
|
||||
|
||||
# Find all result files
|
||||
result_files = find_result_files(search_dirs)
|
||||
print(f"Found {len(result_files)} result file(s)")
|
||||
|
||||
grouped = []
|
||||
if not result_files:
|
||||
print("No benchmark result files found")
|
||||
else:
|
||||
# Parse all result files
|
||||
all_results = []
|
||||
for filepath in sorted(result_files):
|
||||
print(f" Reading: {filepath}")
|
||||
results = parse_result_file(filepath)
|
||||
all_results.extend(results)
|
||||
print(f"Total benchmark results: {len(all_results)}")
|
||||
|
||||
# Group by model/variant
|
||||
grouped = group_results_by_model(all_results, gpu_config, partition)
|
||||
|
||||
# Create metrics structure
|
||||
metrics = {
|
||||
"run_id": run_id,
|
||||
"timestamp": timestamp,
|
||||
"gpu_config": gpu_config,
|
||||
"partition": partition,
|
||||
"results": grouped,
|
||||
}
|
||||
|
||||
# Ensure output directory exists and write output
|
||||
try:
|
||||
os.makedirs(os.path.dirname(output_file) or ".", exist_ok=True)
|
||||
with open(output_file, "w", encoding="utf-8") as f:
|
||||
json.dump(metrics, f, indent=2)
|
||||
|
||||
if not result_files:
|
||||
print(f"Created empty metrics file: {output_file}")
|
||||
else:
|
||||
print(f"Saved metrics to: {output_file}")
|
||||
return True
|
||||
except OSError as e:
|
||||
print(f"Error writing metrics file: {e}")
|
||||
return False
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Collect performance metrics from benchmark results"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--gpu-config",
|
||||
required=True,
|
||||
help="GPU configuration (e.g., 8-gpu-h200, 8-gpu-b200)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--partition",
|
||||
type=int,
|
||||
required=True,
|
||||
help="Partition number (0, 1, 2, etc.)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--run-id",
|
||||
required=True,
|
||||
help="GitHub Actions run ID",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output",
|
||||
required=True,
|
||||
help="Output file path for metrics JSON",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--search-dir",
|
||||
action="append",
|
||||
default=[],
|
||||
dest="search_dirs",
|
||||
help="Directory to search for result files (can be specified multiple times)",
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# Default search directories if none specified
|
||||
search_dirs = args.search_dirs or [
|
||||
"test/performance_profiles_8_gpu",
|
||||
"test/performance_profiles_text_models",
|
||||
"test/performance_profiles_vlms",
|
||||
"test",
|
||||
".",
|
||||
]
|
||||
|
||||
success = save_metrics(
|
||||
gpu_config=args.gpu_config,
|
||||
partition=args.partition,
|
||||
run_id=args.run_id,
|
||||
output_file=args.output,
|
||||
search_dirs=search_dirs,
|
||||
)
|
||||
|
||||
sys.exit(0 if success else 1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,195 @@
|
||||
"""Unit tests for runner_utilization_report.classify_job.
|
||||
|
||||
Pure-logic tests (no GitHub API, stdlib only) so they run in the
|
||||
runner-utilization workflow without installing dependencies:
|
||||
|
||||
python -m unittest discover -s scripts/ci/utils -p 'test_runner_utilization_report.py'
|
||||
|
||||
Regression guard for the queue-time underestimation bug: jobs still
|
||||
waiting in the runner queue (or still running) used to be dropped because
|
||||
the old code required a runner_name and a completed_at, so multi-hour
|
||||
8-gpu waits never showed up in max/avg queue time.
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import unittest
|
||||
from datetime import datetime, timedelta, timezone
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
import runner_utilization_report as rur # noqa: E402
|
||||
|
||||
NOW = datetime(2026, 5, 27, 21, 50, 56, tzinfo=timezone.utc)
|
||||
CREATED = NOW - timedelta(hours=4) # entered the queue 4h ago
|
||||
|
||||
|
||||
def _job(**kw):
|
||||
base = {
|
||||
"name": "base-c-test-8-gpu-h200 / base-c-test-8-gpu-h200 (3)",
|
||||
"status": "completed",
|
||||
"conclusion": "success",
|
||||
"runner_name": "h200-wk03",
|
||||
"labels": ["self-hosted", "X64", "8-gpu-h200"],
|
||||
"created_at": CREATED.isoformat().replace("+00:00", "Z"),
|
||||
"started_at": None,
|
||||
"completed_at": None,
|
||||
"html_url": "https://github.com/o/r/actions/runs/1/job/2",
|
||||
}
|
||||
base.update(kw)
|
||||
return base
|
||||
|
||||
|
||||
def _iso(dt):
|
||||
return dt.isoformat().replace("+00:00", "Z")
|
||||
|
||||
|
||||
class TestClassifyJob(unittest.TestCase):
|
||||
def test_queued_job_counts_ongoing_wait(self):
|
||||
"""The core bug: a still-queued job reports started_at == created_at
|
||||
(placeholder) and no completed_at. Its wait must be now - created_at,
|
||||
not 0, and it must not be dropped."""
|
||||
job = _job(
|
||||
status="queued",
|
||||
runner_name="",
|
||||
started_at=_iso(CREATED), # GitHub placeholder == created_at
|
||||
completed_at=None,
|
||||
)
|
||||
info = rur.classify_job(job, NOW)
|
||||
self.assertIsNotNone(info)
|
||||
self.assertAlmostEqual(info["queue_time"], 4 * 3600, delta=1)
|
||||
self.assertIsNone(info["start"]) # no runner occupied yet
|
||||
self.assertEqual(info["labels"], ["8-gpu-h200"]) # generic labels dropped
|
||||
|
||||
def test_in_progress_job_counts_final_wait(self):
|
||||
"""A running job's wait is final (started - created); old code dropped
|
||||
it for lacking completed_at."""
|
||||
started = CREATED + timedelta(hours=3)
|
||||
job = _job(status="in_progress", started_at=_iso(started), completed_at=None)
|
||||
info = rur.classify_job(job, NOW)
|
||||
self.assertIsNotNone(info)
|
||||
self.assertAlmostEqual(info["queue_time"], 3 * 3600, delta=1)
|
||||
self.assertEqual(info["start"], started)
|
||||
self.assertEqual(info["end"], NOW) # still occupying the runner
|
||||
|
||||
def test_completed_job_unchanged(self):
|
||||
started = CREATED + timedelta(minutes=30)
|
||||
completed = CREATED + timedelta(minutes=90)
|
||||
job = _job(started_at=_iso(started), completed_at=_iso(completed))
|
||||
info = rur.classify_job(job, NOW)
|
||||
self.assertAlmostEqual(info["queue_time"], 30 * 60, delta=1)
|
||||
self.assertAlmostEqual(info["duration"], 60 * 60, delta=1)
|
||||
self.assertEqual(info["end"], completed)
|
||||
|
||||
def test_skipped_job_dropped(self):
|
||||
"""Skipped / cancelled-before-start jobs never waited for a runner."""
|
||||
job = _job(status="completed", runner_name="", started_at=None)
|
||||
self.assertIsNone(rur.classify_job(job, NOW))
|
||||
|
||||
def test_queued_without_created_dropped(self):
|
||||
job = _job(status="queued", runner_name="", created_at=None, started_at=None)
|
||||
self.assertIsNone(rur.classify_job(job, NOW))
|
||||
|
||||
|
||||
class TestConcurrencyHandlesQueuedJobs(unittest.TestCase):
|
||||
def test_queued_job_does_not_crash_and_counts_in_peak_queue(self):
|
||||
window_start = NOW - timedelta(hours=24)
|
||||
queued = rur.classify_job(
|
||||
_job(status="queued", runner_name="", started_at=_iso(CREATED)), NOW
|
||||
)
|
||||
ran = rur.classify_job(
|
||||
_job(
|
||||
started_at=_iso(CREATED + timedelta(hours=1)),
|
||||
completed_at=_iso(CREATED + timedelta(hours=2)),
|
||||
),
|
||||
NOW,
|
||||
)
|
||||
conc = rur.calculate_concurrency_metrics(
|
||||
[queued, ran], window_start, NOW, num_runners=2
|
||||
)
|
||||
# Both jobs were waiting at CREATED before either started -> peak 2.
|
||||
self.assertEqual(conc["peak_queue"], 2)
|
||||
|
||||
|
||||
class TestStatusAndFormatting(unittest.TestCase):
|
||||
def test_status_mapping_and_url(self):
|
||||
for conclusion, expected in (
|
||||
("success", "pass"),
|
||||
("failure", "fail"),
|
||||
("timed_out", "fail"),
|
||||
("cancelled", "cancel"),
|
||||
):
|
||||
info = rur.classify_job(
|
||||
_job(
|
||||
conclusion=conclusion,
|
||||
started_at=_iso(CREATED + timedelta(minutes=5)),
|
||||
completed_at=_iso(CREATED + timedelta(minutes=10)),
|
||||
),
|
||||
NOW,
|
||||
)
|
||||
self.assertEqual(info["status"], expected)
|
||||
self.assertEqual(
|
||||
info["html_url"], "https://github.com/o/r/actions/runs/1/job/2"
|
||||
)
|
||||
|
||||
self.assertEqual(
|
||||
rur.classify_job(
|
||||
_job(status="queued", runner_name="", started_at=_iso(CREATED)), NOW
|
||||
)["status"],
|
||||
"queued",
|
||||
)
|
||||
self.assertEqual(
|
||||
rur.classify_job(
|
||||
_job(
|
||||
status="in_progress", started_at=_iso(CREATED + timedelta(hours=1))
|
||||
),
|
||||
NOW,
|
||||
)["status"],
|
||||
"running",
|
||||
)
|
||||
|
||||
def test_format_status_counts(self):
|
||||
self.assertEqual(rur.format_status_counts({}), "—")
|
||||
cell = rur.format_status_counts({"pass": 5, "queued": 2, "fail": 0})
|
||||
self.assertIn("✅5", cell)
|
||||
self.assertIn("⏳2", cell)
|
||||
self.assertNotIn("❌", cell) # zero counts omitted
|
||||
|
||||
def test_format_report_has_links_and_status(self):
|
||||
results = [
|
||||
{
|
||||
"label": "8-gpu-h200",
|
||||
"num_runners": 4,
|
||||
"effective_runners": 4,
|
||||
"num_jobs": 2,
|
||||
"total_active_hours": 1.0,
|
||||
"utilization_pct": 50.0,
|
||||
"avg_queue_min": 100.0,
|
||||
"max_queue_min": 264.0,
|
||||
"peak_concurrent": 1,
|
||||
"avg_concurrent": 0.5,
|
||||
"saturation_hours": 0.0,
|
||||
"saturation_pct": 0.0,
|
||||
"peak_queue": 2,
|
||||
"status_counts": {"pass": 1, "queued": 1},
|
||||
}
|
||||
]
|
||||
url = "https://github.com/sgl-project/sglang/actions/runs/1/job/2"
|
||||
waits = [
|
||||
{
|
||||
"queue_time": 264 * 60,
|
||||
"status": "queued",
|
||||
"labels": ["8-gpu-h200"],
|
||||
"job_name": "base-c-test-8-gpu-h200 (3)",
|
||||
"html_url": url,
|
||||
}
|
||||
]
|
||||
report = rur.format_report(results, 24, 0.0, longest_waits=waits)
|
||||
self.assertIn("| Status |", report) # new main-table column
|
||||
self.assertIn("Longest Queue Waits", report)
|
||||
self.assertIn(f"]({url})", report) # clickable job link
|
||||
self.assertIn("264m", report)
|
||||
self.assertIn("⏳", report)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Executable
+139
@@ -0,0 +1,139 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Update the per-batch status icon in a /rerun-test reply comment.
|
||||
|
||||
State machine for one batch line (anchored by a unique HTML-comment marker
|
||||
written by the slash-command handler):
|
||||
|
||||
dispatched ⏳ ... <!--rrt:i--> (handler, on dispatch)
|
||||
running 🔄 ... <!--rrt:i--> (start-beacon, on the test runner)
|
||||
done ✅/❌ ... <!--rrt:i:done--> (finalizer, after the test job)
|
||||
|
||||
The leading 🚀 on the line is the visual anchor that this came from a
|
||||
slash-command trigger and is preserved across all states.
|
||||
|
||||
Idempotency:
|
||||
- :done marker present -> no-op (covers reruns and start-after-finalizer race)
|
||||
- running and line already has 🔄 -> no-op
|
||||
- marker not found after retries -> warn and exit 0 so a single placeholder
|
||||
glitch does not amplify into N noisy job failures.
|
||||
|
||||
Concurrent updates against the same comment from different batches can race
|
||||
on the body (read-modify-write of the same field). The finalizer
|
||||
serializes itself via job-level concurrency; the beacon does not, because
|
||||
splitting it into its own job would re-queue the GPU runner. Worst case
|
||||
for the beacon race is one missed 🔄 flicker - comment stays consistent.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
import urllib.error
|
||||
import urllib.request
|
||||
|
||||
RETRY_DELAYS_SEC = [0, 5, 15]
|
||||
|
||||
STATUS_ICONS = {
|
||||
"running": "🔄",
|
||||
"success": "✅",
|
||||
"failure": "❌",
|
||||
}
|
||||
|
||||
TERMINAL_STATUSES = {"success", "failure"}
|
||||
|
||||
|
||||
def gh_request(method, url, token, body=None):
|
||||
headers = {
|
||||
"Authorization": f"Bearer {token}",
|
||||
"Accept": "application/vnd.github+json",
|
||||
"X-GitHub-Api-Version": "2022-11-28",
|
||||
}
|
||||
data = json.dumps(body).encode() if body is not None else None
|
||||
req = urllib.request.Request(url, data=data, method=method, headers=headers)
|
||||
if data is not None:
|
||||
req.add_header("Content-Type", "application/json")
|
||||
try:
|
||||
with urllib.request.urlopen(req, timeout=15) as resp:
|
||||
return resp.status, resp.read().decode()
|
||||
except urllib.error.HTTPError as e:
|
||||
return e.code, e.read().decode()
|
||||
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("--comment-id", required=True, type=int)
|
||||
ap.add_argument(
|
||||
"--marker", required=True, help="Per-batch marker, e.g. <!--rrt:0-->"
|
||||
)
|
||||
ap.add_argument(
|
||||
"--status",
|
||||
required=True,
|
||||
choices=list(STATUS_ICONS.keys()),
|
||||
)
|
||||
ap.add_argument("--repo", required=True, help="owner/repo")
|
||||
args = ap.parse_args()
|
||||
|
||||
token = os.environ.get("GITHUB_TOKEN")
|
||||
if not token:
|
||||
print("ERROR: GITHUB_TOKEN not set")
|
||||
return 1
|
||||
|
||||
icon = STATUS_ICONS[args.status]
|
||||
is_terminal = args.status in TERMINAL_STATUSES
|
||||
done_marker = args.marker.replace("-->", ":done-->")
|
||||
|
||||
url = f"https://api.github.com/repos/{args.repo}/issues/comments/{args.comment_id}"
|
||||
|
||||
body = None
|
||||
for attempt, delay in enumerate(RETRY_DELAYS_SEC):
|
||||
if delay:
|
||||
time.sleep(delay)
|
||||
status, text = gh_request("GET", url, token)
|
||||
if status != 200:
|
||||
print(f"GET failed: {status} {text}")
|
||||
return 1
|
||||
body = json.loads(text).get("body") or ""
|
||||
if done_marker in body:
|
||||
print(f"Marker {done_marker} already present; nothing to do.")
|
||||
return 0
|
||||
if args.marker in body:
|
||||
break
|
||||
print(
|
||||
f"Marker {args.marker} not found "
|
||||
f"(attempt {attempt + 1}/{len(RETRY_DELAYS_SEC)}); will retry."
|
||||
)
|
||||
else:
|
||||
print(
|
||||
f"WARNING: marker {args.marker} not found after "
|
||||
f"{len(RETRY_DELAYS_SEC)} attempts; skipping. "
|
||||
f"The handler may have failed to edit the placeholder comment."
|
||||
)
|
||||
return 0
|
||||
|
||||
new_lines = []
|
||||
for line in body.splitlines(keepends=True):
|
||||
if args.marker in line:
|
||||
if not is_terminal and icon in line:
|
||||
print(f"Line already has {icon}; nothing to do.")
|
||||
return 0
|
||||
for prior in ("⏳", "🔄"):
|
||||
if prior in line:
|
||||
line = line.replace(prior, icon, 1)
|
||||
break
|
||||
if is_terminal:
|
||||
line = line.replace(args.marker, done_marker)
|
||||
new_lines.append(line)
|
||||
|
||||
new_body = "".join(new_lines)
|
||||
status, text = gh_request("PATCH", url, token, body={"body": new_body})
|
||||
if status != 200:
|
||||
print(f"PATCH failed: {status} {text}")
|
||||
return 1
|
||||
print(f"Updated comment {args.comment_id}: {args.marker} -> {icon}")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
Executable
+1947
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user