[CI] Add per-stage NVIDIA model inventory tool (#29447)
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
Kangyan-Zhou
co-authored by
Claude Opus 4.8
Claude Opus 4.8
parent
372a893744
commit
860244d4b4
Executable
+688
@@ -0,0 +1,688 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Build a per-stage model inventory for NVIDIA (CUDA) CI.
|
||||
|
||||
Emits a mapping `CI suite -> [HuggingFace model ids]` so the models a stage
|
||||
exercises can be pre-warmed into a runner cache. The mapping is produced by
|
||||
*static analysis* of the registered test files (no GPU, no sglang import), so
|
||||
it can run on a plain runner and stays fresh per commit.
|
||||
|
||||
How `suite -> files` is resolved
|
||||
Reuses the AST registry parser (`ut_parse_one_file` in ci_register.py, the
|
||||
same one `run_suite.py` uses): registered test files call
|
||||
`register_<backend>_ci(...)`; we group each file under its
|
||||
`effective_suite` for the requested backend (the property falls back to a
|
||||
legacy single-string `suite=` when `stage=`/`runner_config=` are unset).
|
||||
|
||||
How `file -> models` is resolved (best effort, recall-favoring)
|
||||
- A constant table built from `python/sglang/test/**/*.py` module-level
|
||||
assignments (`DEFAULT_MODEL_NAME_FOR_TEST = "meta-llama/..."`, including
|
||||
tuple/list values) plus each test file's own module-level constants.
|
||||
- Inline HuggingFace-id string literals in the file (f-string fragments are
|
||||
skipped: they are partial/dynamic and would yield truncated ids).
|
||||
- `ast.Name` references that resolve to a known model constant.
|
||||
Anything we cannot resolve is reported per suite as `unresolved_files`, and
|
||||
any file we cannot parse is reported in `parse_failures`, so recall gaps are
|
||||
visible rather than silent. A `--overrides` JSON file supplies models for
|
||||
dynamic cases and trims false positives.
|
||||
|
||||
How `runner label -> models` is aggregated
|
||||
Registration/prewarm decisions are made per GH runner *label* (a runner's
|
||||
`runs-on` tag), not per suite. Each suite's runner_config maps to a label
|
||||
via scripts/ci/runner_configs.yml (several configs can share one label,
|
||||
e.g. `4-gpu-h100` and `deepep-4-gpu-h100`), so `runner_labels` carries the
|
||||
per-label UNION -- the set a runner registered under that label must have
|
||||
cached before it takes jobs. Suites without a mappable runner_config are
|
||||
listed in `unmapped_suites`.
|
||||
|
||||
Usage:
|
||||
python3 scripts/ci/list_stage_models.py --backend cuda \
|
||||
--commit "$GITHUB_SHA" --output models-per-stage.json --markdown out.md
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import ast
|
||||
import glob
|
||||
import importlib.util
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import subprocess
|
||||
import sys
|
||||
from typing import Dict, List, Optional, Set, Tuple
|
||||
|
||||
# A HuggingFace repo id is `namespace/name`: exactly one slash, each side
|
||||
# starting alphanumeric and made of alnum plus `.`, `_`, `-`. Note `.` is kept
|
||||
# because real ids carry it (e.g. `RedHatAI/Llama-3.2-3B-quantized.w8a8`).
|
||||
MODEL_ID_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]*/[A-Za-z0-9][A-Za-z0-9._-]*$")
|
||||
|
||||
# `namespace/...` values that look like model ids but are MIME types or similar.
|
||||
_MIME_NAMESPACES = frozenset(
|
||||
{
|
||||
"application",
|
||||
"audio",
|
||||
"example",
|
||||
"font",
|
||||
"image",
|
||||
"message",
|
||||
"model",
|
||||
"multipart",
|
||||
"text",
|
||||
"video",
|
||||
}
|
||||
)
|
||||
|
||||
# Trailing extensions that mark a file path or weight file, not a model name.
|
||||
# (Real ids use suffixes like `-GGUF`/`.w8a8`, not these dotted extensions.)
|
||||
_FILE_EXTENSIONS = (
|
||||
".py",
|
||||
".json",
|
||||
".txt",
|
||||
".md",
|
||||
".rst",
|
||||
".yaml",
|
||||
".yml",
|
||||
".sh",
|
||||
".cu",
|
||||
".cuh",
|
||||
".cpp",
|
||||
".cc",
|
||||
".h",
|
||||
".hpp",
|
||||
".so",
|
||||
".png",
|
||||
".jpg",
|
||||
".jpeg",
|
||||
".gif",
|
||||
".csv",
|
||||
".log",
|
||||
".safetensors",
|
||||
".bin",
|
||||
".h5",
|
||||
".gguf",
|
||||
".pt",
|
||||
".pth",
|
||||
".onnx",
|
||||
)
|
||||
|
||||
# Non-test helper files under test/registered/ (skipped by basename, matching
|
||||
# scripts/ci/check_registered_tests.py). run_suite.py skips `cpu/utils.py` by
|
||||
# path; excluding every `utils.py` by basename is a superset that drops no
|
||||
# CUDA-registered test (the other `utils.py` registers CPU only).
|
||||
_NON_TEST_BASENAMES = frozenset({"conftest.py", "__init__.py", "utils.py"})
|
||||
|
||||
|
||||
def looks_like_model_id(value: str, deny: Optional[Set[str]] = None) -> bool:
|
||||
"""Heuristic: does ``value`` look like a HuggingFace repo id?
|
||||
|
||||
Recall-favoring (a few false positives are cheap for cache-warming) but
|
||||
drops the obvious non-models: MIME types, file paths, numeric ratios.
|
||||
"""
|
||||
if deny and value in deny:
|
||||
return False
|
||||
if not MODEL_ID_RE.match(value):
|
||||
return False
|
||||
if not any(c.isalpha() for c in value): # e.g. "2/3"
|
||||
return False
|
||||
namespace, name = value.split("/", 1)
|
||||
if len(namespace) < 2 or len(name) < 2: # e.g. "N/A"; real ids have longer parts
|
||||
return False
|
||||
if namespace.lower() in _MIME_NAMESPACES:
|
||||
return False
|
||||
if name.lower().endswith(_FILE_EXTENSIONS):
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def _string_values(node: ast.AST) -> List[str]:
|
||||
"""String constants directly held by ``node`` (a Constant, Tuple, or List)."""
|
||||
if isinstance(node, ast.Constant):
|
||||
return [node.value] if isinstance(node.value, str) else []
|
||||
if isinstance(node, (ast.Tuple, ast.List)):
|
||||
out: List[str] = []
|
||||
for elt in node.elts:
|
||||
if isinstance(elt, ast.Constant) and isinstance(elt.value, str):
|
||||
out.append(elt.value)
|
||||
return out
|
||||
return []
|
||||
|
||||
|
||||
def extract_constants_from_source(
|
||||
source: str, deny: Optional[Set[str]] = None
|
||||
) -> Dict[str, Set[str]]:
|
||||
"""Module-level ``NAME = "<model id>"`` (and tuple/list) assignments.
|
||||
|
||||
Handles both bare ``Assign`` and annotated ``AnnAssign`` (``NAME: str =
|
||||
...``). Returns ``{constant_name: {model_id, ...}}``; only model-shaped
|
||||
values are kept, so referencing a non-model constant later contributes
|
||||
nothing.
|
||||
"""
|
||||
table: Dict[str, Set[str]] = {}
|
||||
tree = ast.parse(source)
|
||||
for stmt in tree.body:
|
||||
if isinstance(stmt, ast.Assign):
|
||||
targets, value = stmt.targets, stmt.value
|
||||
elif isinstance(stmt, ast.AnnAssign) and stmt.value is not None:
|
||||
targets, value = [stmt.target], stmt.value
|
||||
else:
|
||||
continue
|
||||
models = {v for v in _string_values(value) if looks_like_model_id(v, deny)}
|
||||
if not models:
|
||||
continue
|
||||
for target in targets:
|
||||
if isinstance(target, ast.Name):
|
||||
table.setdefault(target.id, set()).update(models)
|
||||
return table
|
||||
|
||||
|
||||
def extract_models_from_source(
|
||||
source: str,
|
||||
const_table: Dict[str, Set[str]],
|
||||
deny: Optional[Set[str]] = None,
|
||||
) -> Set[str]:
|
||||
"""All model ids reachable from ``source``: inline literals + name refs.
|
||||
|
||||
``const_table`` is the merged (global + local) constant lookup. Name
|
||||
references resolve against it, so an imported ``DEFAULT_*`` constant is
|
||||
picked up even though its value is defined elsewhere. String fragments
|
||||
inside f-strings are skipped -- they are partial/dynamic and would yield
|
||||
truncated, non-existent ids (e.g. ``f"org/model-{ver}"``).
|
||||
"""
|
||||
tree = ast.parse(source)
|
||||
fstring_fragments = {
|
||||
id(part)
|
||||
for node in ast.walk(tree)
|
||||
if isinstance(node, ast.JoinedStr)
|
||||
for part in node.values
|
||||
if isinstance(part, ast.Constant)
|
||||
}
|
||||
found: Set[str] = set()
|
||||
for node in ast.walk(tree):
|
||||
if isinstance(node, ast.Constant) and isinstance(node.value, str):
|
||||
if id(node) in fstring_fragments:
|
||||
continue
|
||||
if looks_like_model_id(node.value, deny):
|
||||
found.add(node.value)
|
||||
elif isinstance(node, ast.Name) and node.id in const_table:
|
||||
found.update(m for m in const_table[node.id] if not (deny and m in deny))
|
||||
return found
|
||||
|
||||
|
||||
def build_global_constant_table(
|
||||
repo_root: str, deny: Optional[Set[str]] = None
|
||||
) -> Tuple[Dict[str, Set[str]], Dict[str, str]]:
|
||||
"""Constant table from every module under ``python/sglang/test/``.
|
||||
|
||||
These shared helpers (e.g. test_utils, lora_utils) define the ``DEFAULT_*``
|
||||
model constants test files reference by name. Returns ``(table, errors)``
|
||||
where ``errors`` maps any unparsable helper to its exception string -- a
|
||||
broken shared helper drops constants across many suites, so the gap must be
|
||||
surfaced rather than swallowed.
|
||||
"""
|
||||
table: Dict[str, Set[str]] = {}
|
||||
errors: Dict[str, str] = {}
|
||||
pattern = os.path.join(repo_root, "python", "sglang", "test", "**", "*.py")
|
||||
for path in glob.glob(pattern, recursive=True):
|
||||
try:
|
||||
with open(path, encoding="utf-8") as f:
|
||||
source = f.read()
|
||||
local = extract_constants_from_source(source, deny)
|
||||
except (OSError, SyntaxError) as exc:
|
||||
rel = os.path.relpath(path, repo_root)
|
||||
errors[rel] = f"{type(exc).__name__}: {exc}"
|
||||
print(
|
||||
f"WARNING: could not parse constant source {rel}; its model "
|
||||
f"constants are EXCLUDED: {errors[rel]}",
|
||||
file=sys.stderr,
|
||||
)
|
||||
continue
|
||||
for name, models in local.items():
|
||||
table.setdefault(name, set()).update(models)
|
||||
return table, errors
|
||||
|
||||
|
||||
def _load_ci_register(repo_root: str):
|
||||
"""Import ci_register.py by path, sidestepping the heavy `sglang` package."""
|
||||
spec = importlib.util.spec_from_file_location(
|
||||
"ci_register",
|
||||
os.path.join(repo_root, "python", "sglang", "test", "ci", "ci_register.py"),
|
||||
)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
def collect_suite_files(
|
||||
repo_root: str, backend_name: str, include_disabled: bool = False
|
||||
) -> Tuple[
|
||||
Dict[str, List[str]],
|
||||
Dict[str, bool],
|
||||
Dict[str, str],
|
||||
Dict[str, Optional[str]],
|
||||
]:
|
||||
"""Map ``effective_suite -> [relative test file]`` for one backend.
|
||||
|
||||
By default only enabled (`disabled is None`) registries are grouped, since a
|
||||
disabled suite does not run and thus needs no cache warming. Pass
|
||||
``include_disabled=True`` to also group disabled registries (useful to see
|
||||
what a suite *would* download once re-enabled). Returns the mapping,
|
||||
``{suite: is_nightly}``, ``{file: parse error}`` for files whose registry
|
||||
could not be parsed (their models are excluded -- surfaced, not silently
|
||||
dropped), and ``{suite: runner_config}`` (None for legacy single-string
|
||||
``suite=`` registrations, which carry no runner_config).
|
||||
"""
|
||||
ci_register = _load_ci_register(repo_root)
|
||||
backend = getattr(ci_register.HWBackend, backend_name.upper())
|
||||
|
||||
pattern = os.path.join(repo_root, "test", "registered", "**", "*.py")
|
||||
files = sorted(
|
||||
f
|
||||
for f in glob.glob(pattern, recursive=True)
|
||||
if os.path.basename(f) not in _NON_TEST_BASENAMES
|
||||
)
|
||||
|
||||
suite_files: Dict[str, List[str]] = {}
|
||||
suite_nightly: Dict[str, bool] = {}
|
||||
suite_runner_config: Dict[str, Optional[str]] = {}
|
||||
errors: Dict[str, str] = {}
|
||||
for path in files:
|
||||
rel = os.path.relpath(path, repo_root)
|
||||
# Narrow catch: SyntaxError (bad source), ValueError (malformed
|
||||
# registration, raised by RegistryVisitor), OSError (vanished file). A
|
||||
# broader failure (e.g. AttributeError from a parser API drift) should
|
||||
# crash loudly rather than silently empty the inventory.
|
||||
try:
|
||||
registries, _ = ci_register.ut_parse_one_file(path)
|
||||
except (SyntaxError, ValueError, OSError) as exc:
|
||||
errors[rel] = f"{type(exc).__name__}: {exc}"
|
||||
print(
|
||||
f"WARNING: could not parse {rel}; its models are EXCLUDED from "
|
||||
f"the inventory: {errors[rel]}",
|
||||
file=sys.stderr,
|
||||
)
|
||||
continue
|
||||
for r in registries:
|
||||
if r.backend != backend:
|
||||
continue
|
||||
if r.disabled is not None and not include_disabled:
|
||||
continue
|
||||
suite = r.effective_suite
|
||||
if suite is None:
|
||||
continue
|
||||
if rel not in suite_files.setdefault(suite, []):
|
||||
suite_files[suite].append(rel)
|
||||
suite_nightly[suite] = suite_nightly.get(suite, False) or bool(r.nightly)
|
||||
# Modern registrations name the suite `{stage}-test-{runner_config}`,
|
||||
# so every registry in a suite shares one runner_config; legacy
|
||||
# `suite=` registrations have none (stays None).
|
||||
if r.runner_config is not None:
|
||||
suite_runner_config[suite] = r.runner_config
|
||||
else:
|
||||
suite_runner_config.setdefault(suite, None)
|
||||
return suite_files, suite_nightly, errors, suite_runner_config
|
||||
|
||||
|
||||
def load_overrides(path: Optional[str]) -> Dict[str, object]:
|
||||
"""Read the overrides JSON: ``by_file``, ``by_suite``, ``deny``,
|
||||
``suite_labels`` (all optional).
|
||||
|
||||
``suite_labels`` maps a legacy ``suite=`` registration (which carries no
|
||||
runner_config) to the GH runner label(s) its dispatching workflow
|
||||
hardcodes in ``runs-on`` -- a LIST, since one suite can run on several
|
||||
labels. A present-but-null key is treated as its default, so a hand-edit
|
||||
like ``"deny": null`` does not blow up downstream iteration.
|
||||
"""
|
||||
overrides: Dict[str, object] = {
|
||||
"by_file": {},
|
||||
"by_suite": {},
|
||||
"deny": [],
|
||||
"suite_labels": {},
|
||||
}
|
||||
if not path:
|
||||
return overrides
|
||||
if not os.path.exists(path):
|
||||
raise FileNotFoundError(f"overrides file not found: {path}")
|
||||
with open(path, encoding="utf-8") as f:
|
||||
data = json.load(f)
|
||||
for key in ("by_file", "by_suite", "deny", "suite_labels"):
|
||||
if data.get(key) is not None:
|
||||
overrides[key] = data[key]
|
||||
return overrides
|
||||
|
||||
|
||||
# One runner_config entry in runner_configs.yml: a two-space-indented key with
|
||||
# a flow-style (inline `{...}`) mapping. This is the file's documented shape;
|
||||
# entries are matched line-by-line so the tool stays stdlib-only (the workflow
|
||||
# installs nothing, so PyYAML is not available).
|
||||
_RUNNER_CONFIG_LINE_RE = re.compile(r"^ ([A-Za-z0-9_-]+):\s*\{(.*)\}\s*$")
|
||||
_RUNS_ON_RE = re.compile(r"\bruns_on:\s*([^,}\s]+)")
|
||||
|
||||
# runner_configs.yml uses this placeholder for the dynamically-selected b200
|
||||
# runner label (resolved at workflow-load time by runner_configs.py --map).
|
||||
# The inventory keeps it literal unless --b200-runner substitutes it, so the
|
||||
# consumer can see the group is dynamic rather than silently guessing a label.
|
||||
B200_SENTINEL = "$b200_runner"
|
||||
|
||||
|
||||
def load_runner_labels(path: str) -> Dict[str, str]:
|
||||
"""Parse ``{runner_config: runs_on label}`` out of runner_configs.yml.
|
||||
|
||||
The mapping is what turns per-suite model sets into per-runner-LABEL sets:
|
||||
a runner is registered under a `runs_on` label (several runner_configs can
|
||||
share one, e.g. `4-gpu-h100` and `deepep-4-gpu-h100` both run on
|
||||
`4-gpu-h100`), so a runner's cache must cover the union of every suite
|
||||
that can land on its label. Raises ValueError on an entry without
|
||||
`runs_on` or a file with no entries at all -- a format drift must fail
|
||||
the workflow loudly, not silently empty the label aggregation.
|
||||
"""
|
||||
labels: Dict[str, str] = {}
|
||||
in_section = False
|
||||
with open(path, encoding="utf-8") as f:
|
||||
for line in f:
|
||||
if line.startswith("runner_configs:"):
|
||||
in_section = True
|
||||
continue
|
||||
if in_section and line.strip() and not line.startswith(" "):
|
||||
break # next top-level key
|
||||
if not in_section:
|
||||
continue
|
||||
m = _RUNNER_CONFIG_LINE_RE.match(line)
|
||||
if not m:
|
||||
continue
|
||||
name, body = m.group(1), m.group(2)
|
||||
runs_on = _RUNS_ON_RE.search(body)
|
||||
if not runs_on:
|
||||
raise ValueError(f"{path}: runner_config {name!r} has no runs_on field")
|
||||
labels[name] = runs_on.group(1)
|
||||
if not labels:
|
||||
raise ValueError(
|
||||
f"{path}: no runner_configs entries parsed -- format drift? "
|
||||
f"(expected two-space-indented `name: {{...}}` lines under "
|
||||
f"a `runner_configs:` key)"
|
||||
)
|
||||
return labels
|
||||
|
||||
|
||||
def build_inventory(
|
||||
repo_root: str,
|
||||
backend_name: str,
|
||||
overrides: Dict[str, object],
|
||||
commit: str,
|
||||
include_disabled: bool = False,
|
||||
b200_runner: Optional[str] = None,
|
||||
) -> Dict[str, object]:
|
||||
deny: Set[str] = set(overrides.get("deny", [])) # type: ignore[arg-type]
|
||||
by_file: Dict[str, List[str]] = overrides.get("by_file", {}) # type: ignore[assignment]
|
||||
by_suite: Dict[str, List[str]] = overrides.get("by_suite", {}) # type: ignore[assignment]
|
||||
suite_labels_override: Dict[str, List[str]] = overrides.get("suite_labels", {}) # type: ignore[assignment]
|
||||
|
||||
global_table, table_errors = build_global_constant_table(repo_root, deny)
|
||||
suite_files, suite_nightly, registry_errors, suite_runner_config = (
|
||||
collect_suite_files(repo_root, backend_name, include_disabled)
|
||||
)
|
||||
|
||||
runner_labels_map: Dict[str, str] = {}
|
||||
runner_configs_path = os.path.join(repo_root, "scripts", "ci", "runner_configs.yml")
|
||||
if os.path.exists(runner_configs_path):
|
||||
runner_labels_map = load_runner_labels(runner_configs_path)
|
||||
else:
|
||||
print(
|
||||
f"WARNING: {runner_configs_path} not found; every suite will be "
|
||||
f"reported as unmapped_suites (no per-runner-label aggregation).",
|
||||
file=sys.stderr,
|
||||
)
|
||||
|
||||
# Resolve models once per file (a file can belong to several suites).
|
||||
file_models: Dict[str, Set[str]] = {}
|
||||
extract_errors: Dict[str, str] = {}
|
||||
for files in suite_files.values():
|
||||
for rel in files:
|
||||
if rel in file_models:
|
||||
continue
|
||||
abs_path = os.path.join(repo_root, rel)
|
||||
try:
|
||||
with open(abs_path, encoding="utf-8") as f:
|
||||
source = f.read()
|
||||
local_table = extract_constants_from_source(source, deny)
|
||||
merged = dict(global_table)
|
||||
for name, models in local_table.items():
|
||||
merged.setdefault(name, set()).update(models)
|
||||
resolved = extract_models_from_source(source, merged, deny)
|
||||
except (OSError, SyntaxError) as exc:
|
||||
extract_errors[rel] = f"{type(exc).__name__}: {exc}"
|
||||
print(
|
||||
f"WARNING: could not extract models from {rel}; treating it "
|
||||
f"as unresolved: {extract_errors[rel]}",
|
||||
file=sys.stderr,
|
||||
)
|
||||
resolved = set()
|
||||
resolved.update(by_file.get(rel, []))
|
||||
file_models[rel] = resolved
|
||||
|
||||
suites: Dict[str, object] = {}
|
||||
all_models: Set[str] = set()
|
||||
for suite in sorted(suite_files):
|
||||
models: Set[str] = set(by_suite.get(suite, []))
|
||||
unresolved: List[str] = []
|
||||
for rel in suite_files[suite]:
|
||||
resolved = file_models.get(rel, set())
|
||||
if resolved:
|
||||
models.update(resolved)
|
||||
else:
|
||||
unresolved.append(rel)
|
||||
all_models.update(models)
|
||||
suites[suite] = {
|
||||
"nightly": suite_nightly.get(suite, False),
|
||||
"runner_config": suite_runner_config.get(suite),
|
||||
"models": sorted(models),
|
||||
"test_file_count": len(suite_files[suite]),
|
||||
"unresolved_files": sorted(unresolved),
|
||||
}
|
||||
|
||||
# Per-runner-LABEL aggregation: registration/prewarm decisions are made per
|
||||
# GH runner label (what a runner is registered with), and several suites --
|
||||
# via several runner_configs -- can route to one label. A runner's cache
|
||||
# must cover the UNION of every suite that can land on it. Label
|
||||
# resolution order: an explicit `suite_labels` override (legacy suites
|
||||
# whose runs-on lives hardcoded in their dispatching workflow; may name
|
||||
# several labels), else runner_config -> runner_configs.yml. Suites we
|
||||
# cannot map land in `unmapped_suites` -- visible, never silently
|
||||
# dropped, same contract as unresolved_files.
|
||||
label_models: Dict[str, Set[str]] = {}
|
||||
label_suites: Dict[str, List[str]] = {}
|
||||
unmapped_suites: List[str] = []
|
||||
for suite in sorted(suites):
|
||||
labels = suite_labels_override.get(suite)
|
||||
if labels is None:
|
||||
rc = suite_runner_config.get(suite)
|
||||
label = runner_labels_map.get(rc) if rc is not None else None
|
||||
labels = [label] if label is not None else []
|
||||
if not labels:
|
||||
unmapped_suites.append(suite)
|
||||
continue
|
||||
for label in labels:
|
||||
if label == B200_SENTINEL and b200_runner:
|
||||
label = b200_runner
|
||||
label_models.setdefault(label, set()).update(suites[suite]["models"])
|
||||
label_suites.setdefault(label, []).append(suite)
|
||||
runner_labels: Dict[str, object] = {
|
||||
label: {
|
||||
"models": sorted(label_models[label]),
|
||||
"suites": label_suites[label],
|
||||
}
|
||||
for label in sorted(label_models)
|
||||
}
|
||||
|
||||
parse_failures = {}
|
||||
parse_failures.update(table_errors)
|
||||
parse_failures.update(registry_errors)
|
||||
parse_failures.update(extract_errors)
|
||||
|
||||
return {
|
||||
"generated_at_commit": commit,
|
||||
"backend": backend_name.lower(),
|
||||
"suite_count": len(suites),
|
||||
"model_count": len(all_models),
|
||||
"runner_label_count": len(runner_labels),
|
||||
"all_models": sorted(all_models),
|
||||
"parse_failures": dict(sorted(parse_failures.items())),
|
||||
"runner_labels": runner_labels,
|
||||
"unmapped_suites": unmapped_suites,
|
||||
"suites": suites,
|
||||
}
|
||||
|
||||
|
||||
def render_markdown(inventory: Dict[str, object]) -> str:
|
||||
suites: Dict[str, dict] = inventory["suites"] # type: ignore[assignment]
|
||||
failures = inventory.get("parse_failures") or {}
|
||||
runner_labels: Dict[str, dict] = inventory.get("runner_labels") or {} # type: ignore[assignment]
|
||||
unmapped = inventory.get("unmapped_suites") or []
|
||||
lines = [
|
||||
f"## NVIDIA CI model inventory (`{inventory['backend']}`)",
|
||||
"",
|
||||
f"- Commit: `{inventory['generated_at_commit']}`",
|
||||
f"- Suites: **{inventory['suite_count']}**, "
|
||||
f"distinct models: **{inventory['model_count']}**, "
|
||||
f"runner labels: **{inventory.get('runner_label_count', 0)}**",
|
||||
]
|
||||
if failures:
|
||||
lines.append(
|
||||
f"- ⚠️ Unparsable files: **{len(failures)}** (see `parse_failures`)"
|
||||
)
|
||||
if unmapped:
|
||||
lines.append(
|
||||
f"- ⚠️ Suites with no runner label: **{len(unmapped)}** "
|
||||
f"({', '.join(f'`{s}`' for s in unmapped)})"
|
||||
)
|
||||
if runner_labels:
|
||||
lines += [
|
||||
"",
|
||||
"### Per runner label (prewarm a runner's cache with this union)",
|
||||
"",
|
||||
"| Runner label | Suites | Models |",
|
||||
"| --- | ---: | --- |",
|
||||
]
|
||||
for label in sorted(runner_labels):
|
||||
info = runner_labels[label]
|
||||
models = ", ".join(info["models"]) if info["models"] else "_(none)_"
|
||||
lines.append(f"| `{label}` | {len(info['suites'])} | {models} |")
|
||||
lines += [
|
||||
"",
|
||||
"### Per suite",
|
||||
"",
|
||||
"| Suite | Nightly | Models | Unresolved files |",
|
||||
"| --- | :---: | --- | ---: |",
|
||||
]
|
||||
for suite in sorted(suites):
|
||||
info = suites[suite]
|
||||
models = ", ".join(info["models"]) if info["models"] else "_(none)_"
|
||||
nightly = "✓" if info["nightly"] else ""
|
||||
lines.append(
|
||||
f"| `{suite}` | {nightly} | {models} | {len(info['unresolved_files'])} |"
|
||||
)
|
||||
return "\n".join(lines) + "\n"
|
||||
|
||||
|
||||
def resolve_commit(arg: Optional[str], repo_root: str) -> str:
|
||||
if arg:
|
||||
return arg
|
||||
env = os.environ.get("GITHUB_SHA")
|
||||
if env:
|
||||
return env
|
||||
try:
|
||||
return subprocess.check_output(
|
||||
["git", "-C", repo_root, "rev-parse", "HEAD"],
|
||||
text=True,
|
||||
stderr=subprocess.DEVNULL,
|
||||
).strip()
|
||||
except (OSError, subprocess.CalledProcessError):
|
||||
return "unknown"
|
||||
|
||||
|
||||
def main() -> int:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument(
|
||||
"--repo-root",
|
||||
default=".",
|
||||
help="Repository root (default: current directory).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--backend",
|
||||
default="cuda",
|
||||
help="Hardware backend name (default: cuda).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--overrides",
|
||||
default=os.path.join("scripts", "ci", "stage_models_overrides.json"),
|
||||
help="Path to the overrides JSON (relative paths resolve under "
|
||||
"--repo-root). A warning is printed if it does not exist.",
|
||||
)
|
||||
parser.add_argument("--commit", default=None, help="Commit sha to record.")
|
||||
parser.add_argument(
|
||||
"--include-disabled",
|
||||
action="store_true",
|
||||
help="Also include suites whose tests are currently disabled.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--b200-runner",
|
||||
default=None,
|
||||
help="Concrete runner label to substitute for the $b200_runner "
|
||||
"placeholder in runner_configs.yml (default: keep the placeholder "
|
||||
"as the runner_labels key, marking the group as dynamically routed).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output", default=None, help="Write JSON here (default: stdout)."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--markdown", default=None, help="Also write a Markdown summary here."
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
repo_root = os.path.abspath(args.repo_root)
|
||||
# Resolve a relative overrides path against repo_root (not cwd) so the flag
|
||||
# works when invoked from elsewhere with --repo-root.
|
||||
overrides_path = args.overrides
|
||||
if not os.path.isabs(overrides_path):
|
||||
overrides_path = os.path.join(repo_root, overrides_path)
|
||||
if not os.path.exists(overrides_path):
|
||||
print(
|
||||
f"WARNING: overrides file not found at {overrides_path}; "
|
||||
f"proceeding without overrides.",
|
||||
file=sys.stderr,
|
||||
)
|
||||
overrides_path = None
|
||||
overrides = load_overrides(overrides_path)
|
||||
commit = resolve_commit(args.commit, repo_root)
|
||||
|
||||
inventory = build_inventory(
|
||||
repo_root,
|
||||
args.backend,
|
||||
overrides,
|
||||
commit,
|
||||
args.include_disabled,
|
||||
b200_runner=args.b200_runner,
|
||||
)
|
||||
payload = json.dumps(inventory, indent=2, sort_keys=False) + "\n"
|
||||
|
||||
if args.output:
|
||||
with open(args.output, "w", encoding="utf-8") as f:
|
||||
f.write(payload)
|
||||
print(
|
||||
f"Wrote {args.output}: {inventory['suite_count']} suites, "
|
||||
f"{inventory['runner_label_count']} runner labels, "
|
||||
f"{inventory['model_count']} distinct models, "
|
||||
f"{len(inventory['parse_failures'])} parse failures, "
|
||||
f"{len(inventory['unmapped_suites'])} unmapped suites.",
|
||||
file=sys.stderr,
|
||||
)
|
||||
else:
|
||||
sys.stdout.write(payload)
|
||||
|
||||
if args.markdown:
|
||||
with open(args.markdown, "w", encoding="utf-8") as f:
|
||||
f.write(render_markdown(inventory))
|
||||
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
@@ -0,0 +1,31 @@
|
||||
{
|
||||
"_comment": "Manual overrides for list_stage_models.py. by_file/by_suite ADD models the static scan cannot see (models built dynamically, read from configs, passed via CLI args). deny REMOVES false-positive ids the heuristic mistakes for models. Keys in by_file are repo-relative test paths (e.g. test/registered/foo/test_bar.py). suite_labels maps legacy suite= registrations (no runner_config) to the GH runner label(s) their dispatching workflow hardcodes in runs-on -- a list, because one suite can run on several labels (nightly-8-gpu-common). $b200_runner is the dynamic-b200 placeholder from runner_configs.yml. Deliberately absent: nightly-4-gpu-gb300-* (run as k8s pods, not GHA runners) and nightly-2-gpu (registered but dispatched by no workflow); both stay visible in unmapped_suites.",
|
||||
"by_file": {},
|
||||
"by_suite": {},
|
||||
"suite_labels": {
|
||||
"base-b-kernel-benchmark-1-gpu-large": ["1-gpu-h100"],
|
||||
"base-b-kernel-unit-1-gpu-b200": ["$b200_runner"],
|
||||
"base-b-kernel-unit-1-gpu-large": ["1-gpu-h100"],
|
||||
"base-b-kernel-unit-8-gpu-h200": ["8-gpu-h200"],
|
||||
"nightly-1-gpu": ["1-gpu-h100"],
|
||||
"nightly-4-gpu": ["4-gpu-h100"],
|
||||
"nightly-4-gpu-b200": ["$b200_runner"],
|
||||
"nightly-8-gpu-b200": ["8-gpu-b200"],
|
||||
"nightly-8-gpu-common": ["8-gpu-h200", "8-gpu-b200"],
|
||||
"nightly-8-gpu-h200": ["8-gpu-h200"],
|
||||
"nightly-eval-text-2-gpu": ["2-gpu-h100"],
|
||||
"nightly-eval-vlm-2-gpu": ["2-gpu-h100"],
|
||||
"nightly-kernel-1-gpu": ["1-gpu-h100"],
|
||||
"nightly-kernel-8-gpu-h200": ["8-gpu-h200"],
|
||||
"nightly-perf-text-2-gpu": ["2-gpu-h100"],
|
||||
"nightly-perf-vlm-2-gpu": ["2-gpu-h100"],
|
||||
"nightly-precision-8-gpu-h200": ["8-gpu-h200"],
|
||||
"stress": ["8-gpu-h200"],
|
||||
"weekly-8-gpu-h200": ["8-gpu-h200"]
|
||||
},
|
||||
"deny": [
|
||||
"tok/req",
|
||||
"qwen/qwen3",
|
||||
"qwen/qwen3-vl"
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,602 @@
|
||||
"""Unit tests for list_stage_models extraction logic.
|
||||
|
||||
Pure-logic tests (stdlib only, no GPU, no sglang import) so they run in the
|
||||
ci-model-inventory workflow without installing dependencies:
|
||||
|
||||
python -m unittest discover -s scripts/ci -p 'test_list_stage_models.py'
|
||||
|
||||
Guards the recall/precision contract of the static model extractor: model-id
|
||||
shape filtering, constant-table resolution (incl. tuple values), inline +
|
||||
name-reference extraction, override merge semantics, and the suite->files +
|
||||
inventory assembly that drives cache-warming.
|
||||
"""
|
||||
|
||||
import os
|
||||
import shutil
|
||||
import sys
|
||||
import tempfile
|
||||
import unittest
|
||||
from unittest import mock
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
import list_stage_models as lsm # noqa: E402
|
||||
|
||||
# Repo root inferred from this file's location: <root>/scripts/ci/<this>.
|
||||
_REPO_ROOT = os.path.dirname(
|
||||
os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
)
|
||||
_REAL_CI_REGISTER = os.path.join(
|
||||
_REPO_ROOT, "python", "sglang", "test", "ci", "ci_register.py"
|
||||
)
|
||||
|
||||
|
||||
def _write(root, rel, content):
|
||||
path = os.path.join(root, rel)
|
||||
os.makedirs(os.path.dirname(path), exist_ok=True)
|
||||
with open(path, "w", encoding="utf-8") as f:
|
||||
f.write(content)
|
||||
|
||||
|
||||
def _make_fake_repo(root, registered, helpers=None):
|
||||
"""Build a temp repo: copy the real ci_register.py, write test + helper files.
|
||||
|
||||
``registered`` maps ``test/registered/...`` relpaths to file content;
|
||||
``helpers`` maps ``python/sglang/test/...`` relpaths to content.
|
||||
"""
|
||||
dst = os.path.join(root, "python", "sglang", "test", "ci", "ci_register.py")
|
||||
os.makedirs(os.path.dirname(dst), exist_ok=True)
|
||||
shutil.copy(_REAL_CI_REGISTER, dst)
|
||||
for rel, content in registered.items():
|
||||
_write(root, os.path.join("test", "registered", rel), content)
|
||||
for rel, content in (helpers or {}).items():
|
||||
_write(root, os.path.join("python", "sglang", "test", rel), content)
|
||||
|
||||
|
||||
class LooksLikeModelId(unittest.TestCase):
|
||||
def test_accepts_real_model_ids(self):
|
||||
for value in (
|
||||
"meta-llama/Llama-3.1-8B-Instruct",
|
||||
"RedHatAI/Llama-3.2-3B-quantized.w8a8", # dotted suffix is real
|
||||
"cross-encoder/ms-marco-MiniLM-L6-v2",
|
||||
"nvidia/DeepSeek-V3-0324-FP4",
|
||||
"lmsys/sglang-ci-dsv3-test",
|
||||
"Qwen/Qwen2-1.5B-Instruct-GGUF", # -GGUF suffix, not .gguf extension
|
||||
):
|
||||
self.assertTrue(lsm.looks_like_model_id(value), value)
|
||||
|
||||
def test_rejects_non_models(self):
|
||||
for value in (
|
||||
"text/plain", # MIME
|
||||
"application/json", # MIME
|
||||
"image/png", # MIME
|
||||
"Text/Plain", # MIME, case-insensitive
|
||||
"2/3", # numeric ratio, no letters
|
||||
"N/A", # single-char sides
|
||||
"configs/model.json", # path with file extension
|
||||
"org/weights.safetensors", # weight-file extension
|
||||
"org/model.bin", # weight-file extension
|
||||
"a/b/c", # too many slashes
|
||||
"/abs/path", # leading slash
|
||||
"./relative", # leading dot
|
||||
"just-a-string", # no slash
|
||||
):
|
||||
self.assertFalse(lsm.looks_like_model_id(value), value)
|
||||
|
||||
def test_deny_set_overrides(self):
|
||||
value = "org/looks-like-a-model"
|
||||
self.assertTrue(lsm.looks_like_model_id(value))
|
||||
self.assertFalse(lsm.looks_like_model_id(value, deny={value}))
|
||||
|
||||
|
||||
class ConstantTable(unittest.TestCase):
|
||||
def test_single_and_tuple_values(self):
|
||||
source = (
|
||||
'DEFAULT_MODEL = "meta-llama/Llama-3.1-8B-Instruct"\n'
|
||||
'PAIR = ("OPEA/Qwen2.5-0.5B-int4", "Intel/Qwen2-0.5B-int4")\n'
|
||||
'NOT_A_MODEL = "https://example.com/x"\n'
|
||||
"PORT = 30000\n"
|
||||
)
|
||||
table = lsm.extract_constants_from_source(source)
|
||||
self.assertEqual(table["DEFAULT_MODEL"], {"meta-llama/Llama-3.1-8B-Instruct"})
|
||||
self.assertEqual(
|
||||
table["PAIR"], {"OPEA/Qwen2.5-0.5B-int4", "Intel/Qwen2-0.5B-int4"}
|
||||
)
|
||||
self.assertNotIn("NOT_A_MODEL", table)
|
||||
self.assertNotIn("PORT", table)
|
||||
|
||||
def test_annotated_assignment(self):
|
||||
source = 'DEFAULT: str = "meta-llama/Llama-3.1-8B-Instruct"\n'
|
||||
table = lsm.extract_constants_from_source(source)
|
||||
self.assertEqual(table["DEFAULT"], {"meta-llama/Llama-3.1-8B-Instruct"})
|
||||
|
||||
def test_implicit_concatenation(self):
|
||||
# Python folds adjacent string literals at parse time into one Constant.
|
||||
source = 'M = (\n "meta-llama/"\n "Llama-3.1-8B-Instruct"\n)\n'
|
||||
table = lsm.extract_constants_from_source(source)
|
||||
self.assertEqual(table["M"], {"meta-llama/Llama-3.1-8B-Instruct"})
|
||||
|
||||
def test_deny_excludes_from_constant_table(self):
|
||||
source = 'BAD = "weird/thing"\nGOOD = "meta-llama/Llama-3.1-8B-Instruct"\n'
|
||||
table = lsm.extract_constants_from_source(source, deny={"weird/thing"})
|
||||
self.assertNotIn("BAD", table)
|
||||
self.assertIn("GOOD", table)
|
||||
|
||||
|
||||
class ExtractModels(unittest.TestCase):
|
||||
def test_inline_and_name_reference(self):
|
||||
const_table = {
|
||||
"DEFAULT_MODEL_NAME_FOR_TEST": {"meta-llama/Llama-3.1-8B-Instruct"}
|
||||
}
|
||||
source = (
|
||||
"from sglang.test.test_utils import DEFAULT_MODEL_NAME_FOR_TEST\n"
|
||||
"class T:\n"
|
||||
" model = DEFAULT_MODEL_NAME_FOR_TEST\n"
|
||||
' draft = "lmsys/sglang-EAGLE3-LLaMA3.1-Instruct-8B"\n'
|
||||
)
|
||||
models = lsm.extract_models_from_source(source, const_table)
|
||||
self.assertEqual(
|
||||
models,
|
||||
{
|
||||
"meta-llama/Llama-3.1-8B-Instruct",
|
||||
"lmsys/sglang-EAGLE3-LLaMA3.1-Instruct-8B",
|
||||
},
|
||||
)
|
||||
|
||||
def test_fstring_fragments_are_skipped(self):
|
||||
# An f-string with a placeholder yields a truncated, non-existent id;
|
||||
# it must NOT be picked up.
|
||||
source = 'msg = f"meta-llama/Llama-3.1-8B-Instruct-{suffix}"\n'
|
||||
self.assertEqual(lsm.extract_models_from_source(source, {}), set())
|
||||
|
||||
def test_fstring_name_placeholder_still_resolves(self):
|
||||
# The literal fragment is skipped, but a constant referenced inside the
|
||||
# f-string placeholder is still resolved.
|
||||
const_table = {"DRAFT": {"lmsys/sglang-EAGLE-llama2-chat-7B"}}
|
||||
source = 'path = f"prefix/{DRAFT}/topk"\n'
|
||||
self.assertEqual(
|
||||
lsm.extract_models_from_source(source, const_table),
|
||||
{"lmsys/sglang-EAGLE-llama2-chat-7B"},
|
||||
)
|
||||
|
||||
def test_model_less_file_yields_nothing(self):
|
||||
source = (
|
||||
"import torch\n"
|
||||
"def test_kernel():\n"
|
||||
" assert torch.cuda.is_available() or True\n"
|
||||
)
|
||||
self.assertEqual(lsm.extract_models_from_source(source, {}), set())
|
||||
|
||||
def test_deny_removes_name_resolved_model(self):
|
||||
# deny must apply to constant-resolved ids too, not just inline literals.
|
||||
const_table = {"M": {"weird/thing"}}
|
||||
source = "x = M\n"
|
||||
self.assertIn(
|
||||
"weird/thing", lsm.extract_models_from_source(source, const_table)
|
||||
)
|
||||
self.assertEqual(
|
||||
lsm.extract_models_from_source(source, const_table, deny={"weird/thing"}),
|
||||
set(),
|
||||
)
|
||||
|
||||
|
||||
class Overrides(unittest.TestCase):
|
||||
def test_load_defaults_when_missing(self):
|
||||
ov = lsm.load_overrides(None)
|
||||
self.assertEqual(
|
||||
ov, {"by_file": {}, "by_suite": {}, "deny": [], "suite_labels": {}}
|
||||
)
|
||||
|
||||
def test_null_values_fall_back_to_defaults(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
path = os.path.join(tmp, "ov.json")
|
||||
with open(path, "w", encoding="utf-8") as f:
|
||||
f.write('{"deny": null, "by_file": null}')
|
||||
ov = lsm.load_overrides(path)
|
||||
self.assertEqual(
|
||||
ov, {"by_file": {}, "by_suite": {}, "deny": [], "suite_labels": {}}
|
||||
)
|
||||
|
||||
|
||||
class CollectSuiteFiles(unittest.TestCase):
|
||||
REG = (
|
||||
"import unittest\n"
|
||||
"from sglang.test.ci.ci_register import register_cuda_ci, register_amd_ci\n"
|
||||
"{calls}\n"
|
||||
'MODEL = "{model}"\n'
|
||||
'if __name__ == "__main__":\n unittest.main()\n'
|
||||
)
|
||||
|
||||
def _repo(self, tmp):
|
||||
_make_fake_repo(
|
||||
tmp,
|
||||
registered={
|
||||
# enabled CUDA, two suites in one file -> dedupe per suite
|
||||
"a/test_a.py": self.REG.format(
|
||||
calls=(
|
||||
'register_cuda_ci(est_time=1, stage="base-x", '
|
||||
'runner_config="1-gpu")\n'
|
||||
'register_cuda_ci(est_time=1, suite="nightly-y", '
|
||||
"nightly=True)"
|
||||
),
|
||||
model="meta-llama/Llama-3.1-8B-Instruct",
|
||||
),
|
||||
# disabled CUDA -> excluded unless include_disabled
|
||||
"b/test_b.py": self.REG.format(
|
||||
calls=(
|
||||
'register_cuda_ci(est_time=1, stage="base-z", '
|
||||
'runner_config="1-gpu", disabled="flaky")'
|
||||
),
|
||||
model="Qwen/Qwen3-8B",
|
||||
),
|
||||
# AMD only -> never in CUDA inventory
|
||||
"c/test_c.py": self.REG.format(
|
||||
calls='register_amd_ci(est_time=1, suite="nightly-amd")',
|
||||
model="google/gemma-3-4b-it",
|
||||
),
|
||||
},
|
||||
)
|
||||
|
||||
def test_enabled_only_by_default(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
self._repo(tmp)
|
||||
suites, nightly, errors, runner_configs = lsm.collect_suite_files(
|
||||
tmp, "cuda"
|
||||
)
|
||||
self.assertEqual(set(suites), {"base-x-test-1-gpu", "nightly-y"})
|
||||
self.assertEqual(errors, {})
|
||||
self.assertTrue(nightly["nightly-y"])
|
||||
self.assertFalse(nightly["base-x-test-1-gpu"])
|
||||
# AMD suite never appears for the CUDA backend.
|
||||
self.assertNotIn("nightly-amd", suites)
|
||||
# Modern registrations carry their runner_config; legacy suite= has none.
|
||||
self.assertEqual(runner_configs["base-x-test-1-gpu"], "1-gpu")
|
||||
self.assertIsNone(runner_configs["nightly-y"])
|
||||
|
||||
def test_include_disabled(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
self._repo(tmp)
|
||||
suites, _, _, _ = lsm.collect_suite_files(
|
||||
tmp, "cuda", include_disabled=True
|
||||
)
|
||||
self.assertIn("base-z-test-1-gpu", suites)
|
||||
|
||||
def test_unparsable_registry_is_surfaced(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
_make_fake_repo(
|
||||
tmp,
|
||||
registered={
|
||||
# est_time missing -> RegistryVisitor raises ValueError
|
||||
"d/test_bad.py": (
|
||||
"from sglang.test.ci.ci_register import register_cuda_ci\n"
|
||||
'register_cuda_ci(suite="base-x")\n'
|
||||
),
|
||||
},
|
||||
)
|
||||
suites, _, errors, _ = lsm.collect_suite_files(tmp, "cuda")
|
||||
self.assertEqual(suites, {})
|
||||
self.assertIn("test/registered/d/test_bad.py", errors)
|
||||
|
||||
|
||||
class BuildInventory(unittest.TestCase):
|
||||
def _repo(self, tmp):
|
||||
_make_fake_repo(
|
||||
tmp,
|
||||
registered={
|
||||
# resolves a model via an imported constant
|
||||
"a/test_a.py": (
|
||||
"import unittest\n"
|
||||
"from sglang.test.ci.ci_register import register_cuda_ci\n"
|
||||
"from sglang.test.test_utils import DEFAULT_MODEL\n"
|
||||
'register_cuda_ci(est_time=1, stage="base-x", '
|
||||
'runner_config="1-gpu")\n'
|
||||
"MODEL = DEFAULT_MODEL\n"
|
||||
'if __name__ == "__main__":\n unittest.main()\n'
|
||||
),
|
||||
# model-less -> lands in unresolved_files (same suite as a)
|
||||
"b/test_b.py": (
|
||||
"import unittest\n"
|
||||
"from sglang.test.ci.ci_register import register_cuda_ci\n"
|
||||
'register_cuda_ci(est_time=1, stage="base-x", '
|
||||
'runner_config="1-gpu")\n'
|
||||
'if __name__ == "__main__":\n unittest.main()\n'
|
||||
),
|
||||
},
|
||||
helpers={"test_utils.py": 'DEFAULT_MODEL = "meta-llama/Llama-3.1-8B"\n'},
|
||||
)
|
||||
|
||||
def test_resolution_and_unresolved(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
self._repo(tmp)
|
||||
inv = lsm.build_inventory(tmp, "cuda", lsm.load_overrides(None), "sha123")
|
||||
self.assertEqual(inv["generated_at_commit"], "sha123")
|
||||
suite = inv["suites"]["base-x-test-1-gpu"]
|
||||
self.assertEqual(suite["models"], ["meta-llama/Llama-3.1-8B"])
|
||||
self.assertEqual(suite["test_file_count"], 2)
|
||||
self.assertEqual(suite["unresolved_files"], ["test/registered/b/test_b.py"])
|
||||
self.assertEqual(inv["model_count"], 1)
|
||||
self.assertEqual(inv["parse_failures"], {})
|
||||
|
||||
def test_by_file_override_adds_and_clears_unresolved(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
self._repo(tmp)
|
||||
overrides = {
|
||||
"by_file": {"test/registered/b/test_b.py": ["org/extra-model"]},
|
||||
"by_suite": {},
|
||||
"deny": [],
|
||||
}
|
||||
inv = lsm.build_inventory(tmp, "cuda", overrides, "sha")
|
||||
suite = inv["suites"]["base-x-test-1-gpu"]
|
||||
self.assertIn("org/extra-model", suite["models"])
|
||||
# by_file supplied a model for test_b -> no longer unresolved.
|
||||
self.assertEqual(suite["unresolved_files"], [])
|
||||
|
||||
def test_by_suite_override_adds_model(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
self._repo(tmp)
|
||||
overrides = {
|
||||
"by_file": {},
|
||||
"by_suite": {"base-x-test-1-gpu": ["org/suite-model"]},
|
||||
"deny": [],
|
||||
}
|
||||
inv = lsm.build_inventory(tmp, "cuda", overrides, "sha")
|
||||
self.assertIn(
|
||||
"org/suite-model", inv["suites"]["base-x-test-1-gpu"]["models"]
|
||||
)
|
||||
|
||||
|
||||
class RenderMarkdown(unittest.TestCase):
|
||||
def test_table_and_glyphs(self):
|
||||
inv = {
|
||||
"backend": "cuda",
|
||||
"generated_at_commit": "sha",
|
||||
"suite_count": 2,
|
||||
"model_count": 1,
|
||||
"parse_failures": {},
|
||||
"suites": {
|
||||
"n-suite": {
|
||||
"nightly": True,
|
||||
"models": ["org/model"],
|
||||
"test_file_count": 1,
|
||||
"unresolved_files": [],
|
||||
},
|
||||
"empty": {
|
||||
"nightly": False,
|
||||
"models": [],
|
||||
"test_file_count": 2,
|
||||
"unresolved_files": ["x.py", "y.py"],
|
||||
},
|
||||
},
|
||||
}
|
||||
md = lsm.render_markdown(inv)
|
||||
self.assertIn("| `n-suite` | ✓ | org/model | 0 |", md)
|
||||
self.assertIn("| `empty` | | _(none)_ | 2 |", md)
|
||||
|
||||
def test_parse_failures_line(self):
|
||||
inv = {
|
||||
"backend": "cuda",
|
||||
"generated_at_commit": "sha",
|
||||
"suite_count": 0,
|
||||
"model_count": 0,
|
||||
"parse_failures": {"x.py": "SyntaxError: bad"},
|
||||
"suites": {},
|
||||
}
|
||||
self.assertIn("Unparsable files", lsm.render_markdown(inv))
|
||||
|
||||
def test_runner_label_table_and_unmapped_note(self):
|
||||
inv = {
|
||||
"backend": "cuda",
|
||||
"generated_at_commit": "sha",
|
||||
"suite_count": 1,
|
||||
"model_count": 1,
|
||||
"runner_label_count": 1,
|
||||
"parse_failures": {},
|
||||
"runner_labels": {
|
||||
"1-gpu-h100": {"models": ["org/m"], "suites": ["s1", "s2"]}
|
||||
},
|
||||
"unmapped_suites": ["nightly-legacy"],
|
||||
"suites": {
|
||||
"s1": {"nightly": False, "models": ["org/m"], "unresolved_files": []}
|
||||
},
|
||||
}
|
||||
md = lsm.render_markdown(inv)
|
||||
self.assertIn("Per runner label", md)
|
||||
self.assertIn("| `1-gpu-h100` | 2 | org/m |", md)
|
||||
self.assertIn("no runner label: **1**", md)
|
||||
self.assertIn("`nightly-legacy`", md)
|
||||
|
||||
|
||||
_FAKE_RUNNER_CONFIGS_YML = """\
|
||||
# comment
|
||||
_anchors:
|
||||
default_install: &default scripts/ci/cuda/ci_install_dependency.sh
|
||||
|
||||
runner_configs:
|
||||
1-gpu: { install: *default, artifact_version: v4, runs_on: 1-gpu-h100 }
|
||||
deepep-1-gpu: { install: *default, artifact_version: v4, runs_on: 1-gpu-h100 }
|
||||
4-gpu-b200: { install: *default, artifact_version: v6, runs_on: $b200_runner }
|
||||
"""
|
||||
|
||||
|
||||
class LoadRunnerLabels(unittest.TestCase):
|
||||
def _load(self, content):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
path = os.path.join(tmp, "runner_configs.yml")
|
||||
with open(path, "w", encoding="utf-8") as f:
|
||||
f.write(content)
|
||||
return lsm.load_runner_labels(path)
|
||||
|
||||
def test_parses_flat_inline_maps(self):
|
||||
labels = self._load(_FAKE_RUNNER_CONFIGS_YML)
|
||||
self.assertEqual(
|
||||
labels,
|
||||
{
|
||||
"1-gpu": "1-gpu-h100",
|
||||
"deepep-1-gpu": "1-gpu-h100",
|
||||
"4-gpu-b200": lsm.B200_SENTINEL,
|
||||
},
|
||||
)
|
||||
|
||||
def test_missing_runs_on_is_loud(self):
|
||||
with self.assertRaises(ValueError):
|
||||
self._load("runner_configs:\n broken: { install: x }\n")
|
||||
|
||||
def test_empty_or_drifted_format_is_loud(self):
|
||||
with self.assertRaises(ValueError):
|
||||
self._load("something_else:\n a: 1\n")
|
||||
|
||||
def test_real_repo_file(self):
|
||||
"""Anchor the stdlib parser to the actual runner_configs.yml: format
|
||||
drift there must fail these tests (which the workflow runs before
|
||||
generating the inventory), not silently empty the label aggregation."""
|
||||
labels = lsm.load_runner_labels(
|
||||
os.path.join(_REPO_ROOT, "scripts", "ci", "runner_configs.yml")
|
||||
)
|
||||
# Two configs sharing a label is the reason the aggregation exists.
|
||||
self.assertEqual(labels["4-gpu-h100"], "4-gpu-h100")
|
||||
self.assertEqual(labels["deepep-4-gpu-h100"], "4-gpu-h100")
|
||||
self.assertEqual(labels["4-gpu-b200"], lsm.B200_SENTINEL)
|
||||
self.assertGreaterEqual(len(labels), 10)
|
||||
|
||||
|
||||
class RunnerLabelAggregation(unittest.TestCase):
|
||||
REG = (
|
||||
"import unittest\n"
|
||||
"from sglang.test.ci.ci_register import register_cuda_ci\n"
|
||||
"{calls}\n"
|
||||
'MODEL = "{model}"\n'
|
||||
'if __name__ == "__main__":\n unittest.main()\n'
|
||||
)
|
||||
|
||||
def _repo(self, tmp):
|
||||
_make_fake_repo(
|
||||
tmp,
|
||||
registered={
|
||||
# Two suites on runner_configs that share the 1-gpu-h100 label.
|
||||
"a/test_a.py": self.REG.format(
|
||||
calls=(
|
||||
'register_cuda_ci(est_time=1, stage="base-x", '
|
||||
'runner_config="1-gpu")'
|
||||
),
|
||||
model="meta-llama/Llama-3.1-8B-Instruct",
|
||||
),
|
||||
"b/test_b.py": self.REG.format(
|
||||
calls=(
|
||||
'register_cuda_ci(est_time=1, stage="base-y", '
|
||||
'runner_config="deepep-1-gpu")'
|
||||
),
|
||||
model="Qwen/Qwen3-8B",
|
||||
),
|
||||
# Sentinel-labeled config.
|
||||
"c/test_c.py": self.REG.format(
|
||||
calls=(
|
||||
'register_cuda_ci(est_time=1, stage="base-z", '
|
||||
'runner_config="4-gpu-b200")'
|
||||
),
|
||||
model="google/gemma-3-4b-it",
|
||||
),
|
||||
# Legacy suite= with no runner_config -> unmapped.
|
||||
"d/test_d.py": self.REG.format(
|
||||
calls='register_cuda_ci(est_time=1, suite="nightly-legacy")',
|
||||
model="openai/gpt-oss-20b",
|
||||
),
|
||||
},
|
||||
)
|
||||
_write(
|
||||
tmp,
|
||||
os.path.join("scripts", "ci", "runner_configs.yml"),
|
||||
_FAKE_RUNNER_CONFIGS_YML,
|
||||
)
|
||||
|
||||
def test_union_per_label_and_unmapped(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
self._repo(tmp)
|
||||
inv = lsm.build_inventory(tmp, "cuda", lsm.load_overrides(None), "sha")
|
||||
self.assertEqual(inv["runner_label_count"], 2)
|
||||
# Shared label carries the UNION of both suites' models.
|
||||
shared = inv["runner_labels"]["1-gpu-h100"]
|
||||
self.assertEqual(
|
||||
shared["models"],
|
||||
["Qwen/Qwen3-8B", "meta-llama/Llama-3.1-8B-Instruct"],
|
||||
)
|
||||
self.assertEqual(
|
||||
shared["suites"],
|
||||
["base-x-test-1-gpu", "base-y-test-deepep-1-gpu"],
|
||||
)
|
||||
# Sentinel stays literal without --b200-runner.
|
||||
self.assertIn(lsm.B200_SENTINEL, inv["runner_labels"])
|
||||
# Legacy suite is visible as unmapped, and still fully present
|
||||
# (with its models) in the per-suite section.
|
||||
self.assertEqual(inv["unmapped_suites"], ["nightly-legacy"])
|
||||
self.assertEqual(
|
||||
inv["suites"]["nightly-legacy"]["models"], ["openai/gpt-oss-20b"]
|
||||
)
|
||||
self.assertEqual(inv["suites"]["nightly-legacy"]["runner_config"], None)
|
||||
self.assertEqual(
|
||||
inv["suites"]["base-x-test-1-gpu"]["runner_config"], "1-gpu"
|
||||
)
|
||||
|
||||
def test_b200_runner_substitution(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
self._repo(tmp)
|
||||
inv = lsm.build_inventory(
|
||||
tmp,
|
||||
"cuda",
|
||||
lsm.load_overrides(None),
|
||||
"sha",
|
||||
b200_runner="4-gpu-b200-dyn",
|
||||
)
|
||||
self.assertNotIn(lsm.B200_SENTINEL, inv["runner_labels"])
|
||||
self.assertEqual(
|
||||
inv["runner_labels"]["4-gpu-b200-dyn"]["models"],
|
||||
["google/gemma-3-4b-it"],
|
||||
)
|
||||
|
||||
def test_suite_labels_override_maps_legacy_suite(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
self._repo(tmp)
|
||||
overrides = lsm.load_overrides(None)
|
||||
# Legacy suite dispatched by two workflows -> two labels; the
|
||||
# sentinel in an override is substituted like a yml-derived one.
|
||||
overrides["suite_labels"] = {
|
||||
"nightly-legacy": ["1-gpu-h100", "$b200_runner"]
|
||||
}
|
||||
inv = lsm.build_inventory(
|
||||
tmp, "cuda", overrides, "sha", b200_runner="b200-dyn"
|
||||
)
|
||||
self.assertEqual(inv["unmapped_suites"], [])
|
||||
self.assertIn(
|
||||
"openai/gpt-oss-20b", inv["runner_labels"]["1-gpu-h100"]["models"]
|
||||
)
|
||||
self.assertIn(
|
||||
"openai/gpt-oss-20b", inv["runner_labels"]["b200-dyn"]["models"]
|
||||
)
|
||||
self.assertIn(
|
||||
"nightly-legacy", inv["runner_labels"]["1-gpu-h100"]["suites"]
|
||||
)
|
||||
|
||||
def test_missing_yml_leaves_all_suites_unmapped(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
self._repo(tmp)
|
||||
os.remove(os.path.join(tmp, "scripts", "ci", "runner_configs.yml"))
|
||||
inv = lsm.build_inventory(tmp, "cuda", lsm.load_overrides(None), "sha")
|
||||
self.assertEqual(inv["runner_labels"], {})
|
||||
self.assertEqual(len(inv["unmapped_suites"]), 4)
|
||||
|
||||
|
||||
class ResolveCommit(unittest.TestCase):
|
||||
def test_explicit_arg_wins(self):
|
||||
self.assertEqual(lsm.resolve_commit("abc", "/nonexistent"), "abc")
|
||||
|
||||
def test_env_fallback(self):
|
||||
with mock.patch.dict(os.environ, {"GITHUB_SHA": "deadbeef"}, clear=True):
|
||||
self.assertEqual(lsm.resolve_commit(None, "/nonexistent"), "deadbeef")
|
||||
|
||||
def test_unknown_when_no_git(self):
|
||||
with tempfile.TemporaryDirectory() as tmp, mock.patch.dict(
|
||||
os.environ, {}, clear=True
|
||||
):
|
||||
self.assertEqual(lsm.resolve_commit(None, tmp), "unknown")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user