[diffusion] CI: dynamic load-balanced partitioning for diffusion CI (#15528)
Co-authored-by: daiweitao <dwti614707404@163.com> Co-authored-by: SGLang CI <ci@sglang.ai>
This commit is contained in:
Prozac614
co-authored by
daiweitao
SGLang CI
daiweitao
SGLang CI
parent
d6c9d9116b
commit
45472d70cc
+281
@@ -0,0 +1,281 @@
|
||||
#!/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 json
|
||||
import math
|
||||
import os
|
||||
import sys
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
|
||||
from diffusion_case_parser import (
|
||||
BASELINE_REL_PATH,
|
||||
RUN_SUITE_REL_PATH,
|
||||
TESTCASE_CONFIG_REL_PATH,
|
||||
DiffusionSuiteInfo,
|
||||
collect_diffusion_suites,
|
||||
)
|
||||
|
||||
SUITE_OUTPUT_NAMES = {
|
||||
"1-gpu": "1gpu",
|
||||
"2-gpu": "2gpu",
|
||||
}
|
||||
DEFAULT_STANDALONE_EST_TIME_SECONDS = 300.0
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class PartitionItem:
|
||||
kind: str
|
||||
item_id: str
|
||||
est_time: float
|
||||
used_fallback_estimate: bool = False
|
||||
|
||||
|
||||
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) -> list[PartitionItem]:
|
||||
items = [
|
||||
PartitionItem(kind="case", item_id=case.case_id, est_time=case.est_time)
|
||||
for case in suite_info.cases
|
||||
]
|
||||
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 lpt_partition(
|
||||
items: list[PartitionItem], num_partitions: int
|
||||
) -> list[list[PartitionItem]]:
|
||||
if not items or num_partitions <= 0:
|
||||
return []
|
||||
|
||||
sorted_items = sorted(
|
||||
items,
|
||||
key=lambda item: (-item.est_time, item.kind, item.item_id),
|
||||
)
|
||||
partitions: list[list[PartitionItem]] = [[] for _ in range(num_partitions)]
|
||||
partition_sums = [0.0] * num_partitions
|
||||
|
||||
for item in sorted_items:
|
||||
min_idx = partition_sums.index(min(partition_sums))
|
||||
partitions[min_idx].append(item)
|
||||
partition_sums[min_idx] += item.est_time
|
||||
|
||||
return partitions
|
||||
|
||||
|
||||
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]],
|
||||
) -> None:
|
||||
total_time = sum(item.est_time for item in build_partition_items(suite_info))
|
||||
print(f"{suite_name.upper()} suite:")
|
||||
print(f" Cases: {len(suite_info.cases)}")
|
||||
print(f" Standalone files: {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)",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
script_dir = Path(__file__).resolve().parent
|
||||
repo_root = script_dir.parent.parent.parent.parent
|
||||
|
||||
testcase_config_path = repo_root / TESTCASE_CONFIG_REL_PATH
|
||||
baseline_path = repo_root / BASELINE_REL_PATH
|
||||
run_suite_path = repo_root / RUN_SUITE_REL_PATH
|
||||
|
||||
if not testcase_config_path.exists():
|
||||
print(f"Error: Testcase config not found: {testcase_config_path}")
|
||||
sys.exit(1)
|
||||
if not run_suite_path.exists():
|
||||
print(f"Error: Run suite not found: {run_suite_path}")
|
||||
sys.exit(1)
|
||||
|
||||
suites = collect_diffusion_suites(
|
||||
testcase_config_path,
|
||||
run_suite_path,
|
||||
baseline_path,
|
||||
)
|
||||
|
||||
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)
|
||||
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 = lpt_partition(items, partition_count)
|
||||
|
||||
print_suite_summary(suite_name, suite_info, partitions)
|
||||
|
||||
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()
|
||||
+326
@@ -0,0 +1,326 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
AST-based parser for diffusion test cases.
|
||||
|
||||
This module parses testcase_configs.py and run_suite.py using AST
|
||||
to extract test case information without requiring sglang dependencies.
|
||||
Designed to run on lightweight CI runners (ubuntu-latest).
|
||||
|
||||
Usage:
|
||||
# From sibling scripts in this directory:
|
||||
from diffusion_case_parser import collect_diffusion_suites
|
||||
suites = collect_diffusion_suites(testcase_config_path, run_suite_path, baseline_path)
|
||||
"""
|
||||
|
||||
import ast
|
||||
import json
|
||||
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_A": "1-gpu",
|
||||
"ONE_GPU_CASES_B": "1-gpu",
|
||||
"ONE_GPU_CASES_C": "1-gpu-b200",
|
||||
"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
|
||||
TESTCASE_CONFIG_REL_PATH = (
|
||||
"python/sglang/multimodal_gen/test/server/testcase_configs.py"
|
||||
)
|
||||
BASELINE_REL_PATH = "python/sglang/multimodal_gen/test/server/perf_baselines.json"
|
||||
RUN_SUITE_REL_PATH = "python/sglang/multimodal_gen/test/run_suite.py"
|
||||
|
||||
|
||||
@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 testcase_configs.py.
|
||||
|
||||
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, ...]
|
||||
|
||||
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 _process_assignment(self, targets: List[ast.AST], value: ast.AST):
|
||||
"""Process an assignment to extract case IDs if it's a known list."""
|
||||
for target in targets:
|
||||
if isinstance(target, ast.Name) and target.id in CASE_LIST_TO_SUITE:
|
||||
list_name = target.id
|
||||
case_ids = self._extract_case_ids_from_list(value)
|
||||
self.cases[list_name] = case_ids
|
||||
|
||||
def _extract_case_ids_from_list(self, node: ast.AST) -> List[str]:
|
||||
"""Extract case IDs from a list of DiffusionTestCase calls."""
|
||||
case_ids = []
|
||||
if isinstance(node, ast.List):
|
||||
for elt in node.elts:
|
||||
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
|
||||
|
||||
# Check if it's a DiffusionTestCase call
|
||||
if isinstance(node.func, ast.Name) and node.func.id == "DiffusionTestCase":
|
||||
# First positional argument is the case_id
|
||||
if node.args and isinstance(node.args[0], ast.Constant):
|
||||
return node.args[0].value
|
||||
|
||||
return None
|
||||
|
||||
|
||||
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 load_baselines(baseline_path: Path) -> Dict[str, float]:
|
||||
"""
|
||||
Load performance baselines from JSON file.
|
||||
|
||||
Returns:
|
||||
Dictionary mapping case_id to estimated time in seconds.
|
||||
"""
|
||||
if not baseline_path.exists():
|
||||
return {}
|
||||
|
||||
with open(baseline_path, "r", encoding="utf-8") as f:
|
||||
data = json.load(f)
|
||||
|
||||
baselines = {}
|
||||
scenarios = data.get("scenarios", {})
|
||||
|
||||
for case_id, scenario in scenarios.items():
|
||||
if scenario.get("estimated_full_test_time_s") is not None:
|
||||
baselines[case_id] = scenario["estimated_full_test_time_s"]
|
||||
else:
|
||||
expected_e2e_ms = scenario.get("expected_e2e_ms", 0)
|
||||
baselines[case_id] = expected_e2e_ms / 1000.0 + STARTUP_OVERHEAD_SECONDS
|
||||
|
||||
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 testcase_configs.py 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(
|
||||
testcase_config_path: Path,
|
||||
run_suite_path: Path,
|
||||
baseline_path: Path,
|
||||
) -> Dict[str, DiffusionSuiteInfo]:
|
||||
"""
|
||||
Collect all diffusion test suite information using AST parsing.
|
||||
|
||||
Args:
|
||||
testcase_config_path: Path to testcase_configs.py
|
||||
run_suite_path: Path to run_suite.py
|
||||
baseline_path: Path to perf_baselines.json
|
||||
|
||||
Returns:
|
||||
Dictionary mapping suite name to DiffusionSuiteInfo.
|
||||
"""
|
||||
# Parse case IDs from testcase_configs.py
|
||||
case_lists = parse_testcase_configs(testcase_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)
|
||||
|
||||
return suites
|
||||
+324
@@ -0,0 +1,324 @@
|
||||
#!/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,
|
||||
TESTCASE_CONFIG_REL_PATH,
|
||||
collect_diffusion_suites,
|
||||
)
|
||||
|
||||
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 testcase_configs.py and run_suite.py.
|
||||
|
||||
Returns:
|
||||
Dictionary mapping suite name to set of expected case IDs.
|
||||
Standalone files are represented as "standalone:<filename>".
|
||||
"""
|
||||
testcase_config_path = repo_root / TESTCASE_CONFIG_REL_PATH
|
||||
baseline_path = repo_root / BASELINE_REL_PATH
|
||||
run_suite_path = repo_root / RUN_SUITE_REL_PATH
|
||||
|
||||
suites = collect_diffusion_suites(
|
||||
testcase_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
|
||||
|
||||
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
|
||||
expected = get_expected_cases(repo_root)
|
||||
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()
|
||||
Reference in New Issue
Block a user