From 45472d70cc18c46f5c6f383bcc39d128b95a5480 Mon Sep 17 00:00:00 2001 From: Prozac614 Date: Sun, 12 Apr 2026 13:02:43 +0800 Subject: [PATCH] [diffusion] CI: dynamic load-balanced partitioning for diffusion CI (#15528) Co-authored-by: daiweitao Co-authored-by: SGLang CI --- .github/workflows/pr-test-multimodal-gen.yml | 120 ++- .../test/cli/test_generate_common.py | 28 +- .../sglang/multimodal_gen/test/run_suite.py | 836 +++++++++++++++--- .../test/scripts/gen_diffusion_ci_outputs.py | 2 +- .../test/server/perf_baselines.json | 30 +- .../test/server/test_server_common.py | 76 +- .../test/server/testcase_configs.py | 8 +- .../sglang/multimodal_gen/test/test_utils.py | 21 + .../diffusion/compute_diffusion_partitions.py | 281 ++++++ .../utils/diffusion/diffusion_case_parser.py | 326 +++++++ .../diffusion/verify_diffusion_coverage.py | 324 +++++++ 11 files changed, 1877 insertions(+), 175 deletions(-) create mode 100755 scripts/ci/utils/diffusion/compute_diffusion_partitions.py create mode 100755 scripts/ci/utils/diffusion/diffusion_case_parser.py create mode 100755 scripts/ci/utils/diffusion/verify_diffusion_coverage.py diff --git a/.github/workflows/pr-test-multimodal-gen.yml b/.github/workflows/pr-test-multimodal-gen.yml index d92b3a141..356b5254d 100644 --- a/.github/workflows/pr-test-multimodal-gen.yml +++ b/.github/workflows/pr-test-multimodal-gen.yml @@ -50,20 +50,57 @@ env: SKIP_STAGE_HEALTH_CHECK: ${{ inputs.skip_stage_health_check == 'true' }} jobs: - multimodal-gen-test-1-gpu: + compute-diffusion-partitions: if: | (inputs.target_stage == 'multimodal-gen-test-1-gpu') || + (inputs.target_stage == 'multimodal-gen-test-2-gpu') || ( !inputs.target_stage && - ((github.event_name == 'schedule' || inputs.test_parallel_dispatch == 'true') || (inputs.caller_needs_failure != 'true' && !cancelled())) && inputs.multimodal_gen == 'true' ) + runs-on: ubuntu-latest + outputs: + matrix-1gpu: ${{ steps.compute.outputs.matrix-1gpu }} + matrix-2gpu: ${{ steps.compute.outputs.matrix-2gpu }} + partition-count-1gpu: ${{ steps.compute.outputs['partition-count-1gpu'] }} + partition-count-2gpu: ${{ steps.compute.outputs['partition-count-2gpu'] }} + plan-1gpu: ${{ steps.compute.outputs.plan-1gpu }} + plan-2gpu: ${{ steps.compute.outputs.plan-2gpu }} + steps: + - name: Checkout code + uses: actions/checkout@v4 + with: + ref: ${{ inputs.pr_head_sha || inputs.git_ref || github.sha }} + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: '3.10' + + - name: Compute partitions + id: compute + run: | + python scripts/ci/utils/diffusion/compute_diffusion_partitions.py --min-time 1200 --target-time 1800 --max-time 2400 --max-partitions 10 + + multimodal-gen-test-1-gpu: + needs: compute-diffusion-partitions + if: | + always() && + needs.compute-diffusion-partitions.result == 'success' && + needs.compute-diffusion-partitions.outputs.matrix-1gpu != '{"include":[]}' && + ( + (inputs.target_stage == 'multimodal-gen-test-1-gpu') || + ( + !inputs.target_stage && + ((github.event_name == 'schedule' || inputs.test_parallel_dispatch == 'true') || (inputs.caller_needs_failure != 'true' && !cancelled())) && + inputs.multimodal_gen == 'true' + ) + ) runs-on: 1-gpu-h100 timeout-minutes: 240 strategy: fail-fast: false - matrix: - part: [0, 1] + matrix: ${{ fromJson(needs.compute-diffusion-partitions.outputs.matrix-1gpu) }} steps: - name: Checkout code uses: actions/checkout@v4 @@ -91,33 +128,48 @@ jobs: env: RUNAI_STREAMER_MEMORY_LIMIT: 0 CONTINUE_ON_ERROR_FLAG: ${{ inputs.continue_on_error == 'true' && '--continue-on-error' || '' }} + PARTITION_PLAN_JSON: ${{ needs.compute-diffusion-partitions.outputs.plan-1gpu }} run: | cd python python3 sglang/multimodal_gen/test/run_suite.py \ --suite 1-gpu \ --partition-id ${{ matrix.part }} \ - --total-partitions 2 \ + --total-partitions ${{ needs.compute-diffusion-partitions.outputs['partition-count-1gpu'] }} \ + --partition-plan-json "$PARTITION_PLAN_JSON" \ $CONTINUE_ON_ERROR_FLAG + - name: Upload execution report + if: always() + uses: actions/upload-artifact@v4 + with: + name: diffusion-report-1gpu-${{ matrix.part }} + path: python/sglang/multimodal_gen/test/execution_report_*.json + retention-days: 1 + - uses: ./.github/actions/upload-cuda-coredumps if: failure() with: artifact-suffix: ${{ matrix.part }} multimodal-gen-test-2-gpu: + needs: compute-diffusion-partitions if: | - (inputs.target_stage == 'multimodal-gen-test-2-gpu') || + always() && + needs.compute-diffusion-partitions.result == 'success' && + needs.compute-diffusion-partitions.outputs.matrix-2gpu != '{"include":[]}' && ( - !inputs.target_stage && - ((github.event_name == 'schedule' || inputs.test_parallel_dispatch == 'true') || (inputs.caller_needs_failure != 'true' && !cancelled())) && - inputs.multimodal_gen == 'true' + (inputs.target_stage == 'multimodal-gen-test-2-gpu') || + ( + !inputs.target_stage && + ((github.event_name == 'schedule' || inputs.test_parallel_dispatch == 'true') || (inputs.caller_needs_failure != 'true' && !cancelled())) && + inputs.multimodal_gen == 'true' + ) ) runs-on: 2-gpu-h100 timeout-minutes: 240 strategy: fail-fast: false - matrix: - part: [0, 1] + matrix: ${{ fromJson(needs.compute-diffusion-partitions.outputs.matrix-2gpu) }} steps: - name: Checkout code uses: actions/checkout@v4 @@ -146,14 +198,24 @@ jobs: env: RUNAI_STREAMER_MEMORY_LIMIT: 0 CONTINUE_ON_ERROR_FLAG: ${{ inputs.continue_on_error == 'true' && '--continue-on-error' || '' }} + PARTITION_PLAN_JSON: ${{ needs.compute-diffusion-partitions.outputs.plan-2gpu }} run: | cd python python3 sglang/multimodal_gen/test/run_suite.py \ --suite 2-gpu \ --partition-id ${{ matrix.part }} \ - --total-partitions 2 \ + --total-partitions ${{ needs.compute-diffusion-partitions.outputs['partition-count-2gpu'] }} \ + --partition-plan-json "$PARTITION_PLAN_JSON" \ $CONTINUE_ON_ERROR_FLAG + - name: Upload execution report + if: always() + uses: actions/upload-artifact@v4 + with: + name: diffusion-report-2gpu-${{ matrix.part }} + path: python/sglang/multimodal_gen/test/execution_report_*.json + retention-days: 1 + - uses: ./.github/actions/upload-cuda-coredumps if: failure() with: @@ -354,3 +416,37 @@ jobs: run: | cd python python3 sglang/multimodal_gen/test/run_suite.py --suite unit + + diffusion-coverage-check: + needs: [multimodal-gen-test-1-gpu, multimodal-gen-test-2-gpu] + if: | + always() && + inputs.multimodal_gen == 'true' && + ( + needs.multimodal-gen-test-1-gpu.result == 'success' || + needs.multimodal-gen-test-1-gpu.result == 'failure' || + needs.multimodal-gen-test-2-gpu.result == 'success' || + needs.multimodal-gen-test-2-gpu.result == 'failure' + ) + runs-on: ubuntu-latest + steps: + - name: Checkout code + uses: actions/checkout@v4 + with: + ref: ${{ inputs.pr_head_sha || inputs.git_ref || github.sha }} + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: '3.10' + + - name: Download all execution reports + uses: actions/download-artifact@v4 + with: + path: reports/ + pattern: diffusion-report-* + merge-multiple: true + + - name: Verify coverage + run: | + python scripts/ci/utils/diffusion/verify_diffusion_coverage.py --reports-dir reports/ diff --git a/python/sglang/multimodal_gen/test/cli/test_generate_common.py b/python/sglang/multimodal_gen/test/cli/test_generate_common.py index 196b82e03..abce50738 100644 --- a/python/sglang/multimodal_gen/test/cli/test_generate_common.py +++ b/python/sglang/multimodal_gen/test/cli/test_generate_common.py @@ -7,16 +7,13 @@ Common generate cli test, one test for image and video each import dataclasses import os import shlex -import subprocess -import sys import unittest -from typing import Optional from PIL import Image from sglang.multimodal_gen.configs.sample.sampling_params import DataType from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger -from sglang.multimodal_gen.test.test_utils import check_image_size +from sglang.multimodal_gen.test.test_utils import check_image_size, run_command logger = init_logger(__name__) @@ -28,29 +25,6 @@ class TestResult: succeed: bool -def run_command(command) -> Optional[float]: - """Runs a command and returns the execution time and status.""" - print(f"Running command: {shlex.join(command)}") - - with subprocess.Popen( - command, - stdout=subprocess.PIPE, - stderr=subprocess.STDOUT, - bufsize=0, - ) as process: - while True: - chunk = process.stdout.read(4096) - if not chunk: - break - sys.stdout.buffer.write(chunk) - sys.stdout.buffer.flush() - process.wait() - if process.returncode == 0: - return True - print(f"Command failed with exit code {process.returncode}") - return False - - class CLIBase(unittest.TestCase): model_path: str = None extra_args = [] diff --git a/python/sglang/multimodal_gen/test/run_suite.py b/python/sglang/multimodal_gen/test/run_suite.py index c32968e65..851c82a3b 100644 --- a/python/sglang/multimodal_gen/test/run_suite.py +++ b/python/sglang/multimodal_gen/test/run_suite.py @@ -1,26 +1,40 @@ """ Test runner for multimodal_gen that manages test suites and parallel execution. -Usage: - python3 run_suite.py --suite --partition-id --total-partitions - -Example: - python3 run_suite.py --suite 1-gpu --partition-id 0 --total-partitions 4 +For diffusion 1-gpu/2-gpu suites, cases are partitioned by estimated runtime +using LPT so each CI shard has a similar total runtime. """ import argparse +import copy +import json import os import random import subprocess import sys +import time +import xml.etree.ElementTree as ET +from dataclasses import dataclass from pathlib import Path import tabulate from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger +from sglang.multimodal_gen.test.server.testcase_configs import ( + BASELINE_CONFIG, + ONE_GPU_CASES_A, + ONE_GPU_CASES_B, + TWO_GPU_CASES_A, + TWO_GPU_CASES_B, + DiffusionTestCase, +) logger = init_logger(__name__) +DEFAULT_EST_TIME_SECONDS = 300.0 +STARTUP_OVERHEAD_SECONDS = 120.0 +DEFAULT_STANDALONE_EST_TIME_SECONDS = 300.0 + _UPDATE_WEIGHTS_FROM_DISK_TEST_FILE = "test_update_weights_from_disk.py" _UPDATE_WEIGHTS_MODEL_PAIR_ENV = "SGLANG_MMGEN_UPDATE_WEIGHTS_PAIR" _UPDATE_WEIGHTS_MODEL_PAIR_IDS = ( @@ -30,7 +44,6 @@ _UPDATE_WEIGHTS_MODEL_PAIR_IDS = ( def _discover_unit_tests() -> list[str]: - """Auto-discover all test_*.py files in the unit/ directory.""" unit_dir = Path(__file__).resolve().parent / "unit" if not unit_dir.is_dir(): return [] @@ -39,23 +52,8 @@ def _discover_unit_tests() -> list[str]: ) -SUITES = { - # no GPU required; safe to run on any CPU-only runner - # Auto-discovered from test/unit/test_*.py +FILE_SUITES = { "unit": _discover_unit_tests(), - "1-gpu": [ - "test_server_a.py", - "test_server_b.py", - # cli test - "../cli/test_generate_t2i_perf.py", - "test_update_weights_from_disk.py", - # add new 1-gpu test files here - ], - "2-gpu": [ - "test_server_2_gpu_a.py", - "test_server_2_gpu_b.py", - # add new 2-gpu test files here - ], "component-accuracy-1-gpu": [ "test_accuracy_1_gpu_a.py", "test_accuracy_1_gpu_b.py", @@ -70,21 +68,52 @@ SUITES = { } suites_ascend = { - "1-npu": [ - "ascend/test_server_1_npu.py", - # add new 1-npu test files here + "1-npu": ["ascend/test_server_1_npu.py"], + "2-npu": ["ascend/test_server_2_npu.py"], + "8-npu": ["ascend/test_server_8_npu.py"], +} +FILE_SUITES.update(suites_ascend) + +PARAMETRIZED_CASE_GROUPS = { + "1-gpu": [ + ("test_server_a.py", ONE_GPU_CASES_A), + ("test_server_b.py", ONE_GPU_CASES_B), ], - "2-npu": [ - "ascend/test_server_2_npu.py", - # add new 2-npu test files here - ], - "8-npu": [ - "ascend/test_server_8_npu.py", - # add new 8-npu test files here + "2-gpu": [ + ("test_server_2_gpu_a.py", TWO_GPU_CASES_A), + ("test_server_2_gpu_b.py", TWO_GPU_CASES_B), ], } -SUITES.update(suites_ascend) +STANDALONE_FILES = { + "1-gpu": [ + "../cli/test_generate_t2i_perf.py", + "test_update_weights_from_disk.py", + ], + "2-gpu": [], +} + +# New standalone files may omit an estimate once to learn the real CI runtime. +# CI will use a fallback estimate for sharding, run the test, then print a +# measured value that must be copied into STANDALONE_FILE_EST_TIMES. +STANDALONE_FILE_EST_TIMES = { + "1-gpu": { + "../cli/test_generate_t2i_perf.py": 240.0, + "test_update_weights_from_disk.py": 480.0, + }, + "2-gpu": {}, +} + +# Backward-compatible suite view for scripts that still operate on file lists. +SUITES = { + **FILE_SUITES, + **{ + suite: [filename for filename, _ in case_groups] + + STANDALONE_FILES.get(suite, []) + for suite, case_groups in PARAMETRIZED_CASE_GROUPS.items() + }, +} + STRICT_SUITES = {"unit"} COMPONENT_ACCURACY_SUITES = { "component-accuracy-1-gpu", @@ -92,14 +121,209 @@ COMPONENT_ACCURACY_SUITES = { } +@dataclass(frozen=True) +class PartitionAssignment: + case_ids: list[str] + standalone_files: list[str] + estimated_time: float | None = None + missing_standalone_estimates: list[str] | None = None + + +def get_case_est_time(case_id: str) -> float: + scenario = BASELINE_CONFIG.scenarios.get(case_id) + if scenario is None: + return DEFAULT_EST_TIME_SECONDS + if scenario.estimated_full_test_time_s is not None: + return scenario.estimated_full_test_time_s + return scenario.expected_e2e_ms / 1000.0 + STARTUP_OVERHEAD_SECONDS + + +def get_standalone_file_est_time( + suite: str, standalone_file: str +) -> tuple[float, bool]: + suite_est_times = STANDALONE_FILE_EST_TIMES.get(suite, {}) + if standalone_file not in suite_est_times: + return DEFAULT_STANDALONE_EST_TIME_SECONDS, True + return suite_est_times[standalone_file], False + + +def get_all_standalone_file_est_times() -> dict[str, dict[str, float]]: + return copy.deepcopy(STANDALONE_FILE_EST_TIMES) + + +def validate_standalone_file_est_times() -> dict[str, list[str]]: + missing_by_suite: dict[str, list[str]] = {} + for suite, standalone_files in STANDALONE_FILES.items(): + suite_est_times = STANDALONE_FILE_EST_TIMES.get(suite, {}) + missing = [ + standalone_file + for standalone_file in standalone_files + if standalone_file not in suite_est_times + ] + if missing: + missing_by_suite[suite] = missing + return missing_by_suite + + +def auto_partition( + cases: list[DiffusionTestCase], rank: int, size: int +) -> list[DiffusionTestCase]: + if not cases or size <= 0: + return [] + + sorted_cases = sorted(cases, key=lambda c: get_case_est_time(c.id), reverse=True) + partitions: list[list[DiffusionTestCase]] = [[] for _ in range(size)] + partition_sums = [0.0] * size + + for case in sorted_cases: + min_idx = partition_sums.index(min(partition_sums)) + partitions[min_idx].append(case) + partition_sums[min_idx] += get_case_est_time(case.id) + + return partitions[rank] if rank < size else [] + + +def _normalize_standalone_key(standalone_file: str) -> str: + return f"standalone:{standalone_file}" + + +def parse_partition_plan( + suite: str, + partition_id: int, + total_partitions: int, + plan_json: str, +) -> PartitionAssignment: + plan = json.loads(plan_json) + if plan.get("suite") != suite: + raise ValueError( + f"Partition plan suite mismatch: expected {suite!r}, " + f"got {plan.get('suite')!r}" + ) + + partition_count = plan.get("partition_count") + if partition_count != total_partitions: + raise ValueError( + f"Partition count mismatch for suite {suite!r}: " + f"plan={partition_count}, matrix={total_partitions}" + ) + + partitions = plan.get("partitions", []) + selected_partition = None + for partition in partitions: + if partition.get("part") == partition_id: + selected_partition = partition + break + + if selected_partition is None: + raise ValueError( + f"Partition {partition_id} not found in plan for suite {suite!r}" + ) + + return PartitionAssignment( + case_ids=list(selected_partition.get("case_ids", [])), + standalone_files=list(selected_partition.get("standalone_files", [])), + estimated_time=selected_partition.get("estimated_time"), + missing_standalone_estimates=list( + selected_partition.get("missing_standalone_estimates", []) + ), + ) + + +def _merge_execution_results( + executed_cases: list[str], + case_results: dict[str, str], + new_executed_cases: list[str], + new_case_results: dict[str, str], +) -> None: + executed_cases.extend( + case_id for case_id in new_executed_cases if case_id not in executed_cases + ) + case_results.update(new_case_results) + + +def _format_standalone_estimate_snippet( + suite: str, standalone_file: str, measured_full_test_time_s: float +) -> str: + return ( + f'"{suite}": {{\n' + f' "{standalone_file}": {measured_full_test_time_s:.1f},\n' + f"}}" + ) + + +def _print_missing_standalone_estimate_message( + suite: str, + standalone_file: str, + measured_full_test_time_s: float, +) -> None: + snippet = _format_standalone_estimate_snippet( + suite, standalone_file, measured_full_test_time_s + ) + logger.error( + f'\n{"=" * 60}\n' + f'Add standalone estimate for suite "{suite}" and file "{standalone_file}":\n\n' + f"File: python/sglang/multimodal_gen/test/run_suite.py\n\n" + f"Current partition used fallback estimate: " + f"{DEFAULT_STANDALONE_EST_TIME_SECONDS:.1f}s\n\n" + f"{snippet}\n" + f'{"=" * 60}\n' + ) + + +def _run_standalone_file( + suite: str, + standalone_rel: str, + target_dir: Path, + extra_filter: str | None = None, +) -> tuple[int, list[str], dict[str, str], dict]: + if standalone_rel == _UPDATE_WEIGHTS_FROM_DISK_TEST_FILE: + _maybe_pin_update_weights_model_pair([standalone_rel]) + + est_time, used_fallback_estimate = get_standalone_file_est_time( + suite, standalone_rel + ) + standalone_file = _resolve_suite_files(target_dir, [standalone_rel], strict=True)[0] + junit_xml_path = str( + target_dir / f"junit_results_{suite}_{Path(standalone_rel).stem}.xml" + ) + start_time = time.perf_counter() + exit_code, _, _ = run_pytest( + [standalone_file], + filter_expr=extra_filter, + junit_xml_path=junit_xml_path, + ) + measured_full_test_time_s = round(time.perf_counter() - start_time, 1) + standalone_key = _normalize_standalone_key(standalone_rel) + measurement = { + "suite": suite, + "standalone_file": standalone_rel, + "measured_full_test_time_s": measured_full_test_time_s, + "used_fallback_estimate": used_fallback_estimate, + "fallback_estimate_s": DEFAULT_STANDALONE_EST_TIME_SECONDS, + "had_configured_estimate": not used_fallback_estimate, + "configured_or_fallback_estimate_s": est_time, + } + if used_fallback_estimate: + _print_missing_standalone_estimate_message( + suite, standalone_rel, measured_full_test_time_s + ) + return ( + exit_code, + [standalone_key], + {standalone_key: "pass" if exit_code == 0 else "fail"}, + measurement, + ) + + def parse_args(): + suite_choices = sorted(set(FILE_SUITES) | set(PARAMETRIZED_CASE_GROUPS)) parser = argparse.ArgumentParser(description="Run multimodal_gen test suite") parser.add_argument( "--suite", type=str, required=True, - choices=list(SUITES.keys()), - help="The test suite to run (valid names are defined in SUITES)", + choices=suite_choices, + help="The test suite to run.", ) parser.add_argument( "--partition-id", @@ -130,13 +354,19 @@ def parse_args(): "--continue-on-error", action="store_true", default=False, - help="Continue running remaining tests even if one fails (for CI consistency; pytest already continues by default)", + help="Continue running remaining tests even if one fails.", + ) + parser.add_argument( + "--partition-plan-json", + type=str, + default=None, + help="Full partition plan JSON for the current suite.", ) return parser.parse_args() -def collect_test_items(files, filter_expr=None): - """Collect test item node IDs from the given files using pytest --collect-only.""" +def collect_test_items(files: list[str], filter_expr: str | None = None) -> list[str]: + """Collect test node IDs from the given files using pytest --collect-only.""" cmd = [sys.executable, "-m", "pytest", "--collect-only", "-q"] if filter_expr: cmd.extend(["-k", filter_expr]) @@ -146,14 +376,6 @@ def collect_test_items(files, filter_expr=None): print(f"Collecting tests from {len(files)} file(s){filter_note}") result = subprocess.run(cmd, capture_output=True, text=True) - # Check for collection errors - # pytest exit codes: - # 0: success - # 1: tests collected but some had errors during collection - # 2: test execution interrupted - # 3: internal error - # 4: command line usage error - # 5: no tests collected (may be expected with filters) if result.returncode not in (0, 5): error_msg = ( f"pytest --collect-only failed with exit code {result.returncode}\n" @@ -171,14 +393,10 @@ def collect_test_items(files, filter_expr=None): "No tests were collected (exit code 5). This may be expected with filters." ) - # Parse the output to extract test node IDs - # pytest -q outputs lines like: test_file.py::TestClass::test_method[param] test_items = [] for line in result.stdout.strip().split("\n"): line = line.strip() - # Skip empty lines and summary lines if line and "::" in line and not line.startswith(("=", "-", " ")): - # Handle lines that might have extra info after the test ID test_id = line.split()[0] if " " in line else line if "::" in test_id: test_items.append(test_id) @@ -187,6 +405,85 @@ def collect_test_items(files, filter_expr=None): return test_items +def parse_junit_xml_for_executed_cases(xml_path: str) -> list[str]: + if not Path(xml_path).exists(): + return [] + + executed_cases = [] + tree = ET.parse(xml_path) + root = tree.getroot() + + for testcase in root.iter("testcase"): + if testcase.find("skipped") is not None: + continue + + name = testcase.get("name", "") + if "[" in name and "]" in name: + case_id = name[name.index("[") + 1 : name.index("]")] + executed_cases.append(case_id) + + return executed_cases + + +def parse_junit_xml_for_case_results(xml_path: str) -> dict[str, str]: + if not Path(xml_path).exists(): + return {} + + case_results = {} + tree = ET.parse(xml_path) + root = tree.getroot() + + for testcase in root.iter("testcase"): + if testcase.find("skipped") is not None: + continue + + name = testcase.get("name", "") + if "[" not in name or "]" not in name: + continue + + case_id = name[name.index("[") + 1 : name.index("]")] + if testcase.find("failure") is not None: + case_results[case_id] = "fail" + elif testcase.find("error") is not None: + case_results[case_id] = "error" + else: + case_results[case_id] = "pass" + + return case_results + + +def write_execution_report( + suite: str, + partition_id: int, + total_partitions: int, + executed_cases: list[str], + is_standalone: bool = False, + standalone_file: str | None = None, + case_results: dict[str, str] | None = None, + missing_standalone_estimates: list[str] | None = None, + standalone_measurements: list[dict] | None = None, +) -> str: + report = { + "suite": suite, + "partition_id": partition_id, + "total_partitions": total_partitions, + "is_standalone": is_standalone, + "standalone_file": standalone_file, + "executed_cases": executed_cases, + "case_results": case_results or {}, + "missing_standalone_estimates": missing_standalone_estimates or [], + "standalone_measurements": standalone_measurements or [], + } + + report_filename = f"execution_report_{suite}_{partition_id}.json" + report_path = Path(__file__).parent / report_filename + with open(report_path, "w", encoding="utf-8") as f: + json.dump(report, f, indent=2) + + logger.info("Execution report written to: %s", report_path) + return str(report_path) + + def _run_pytest_attempt(cmd: list[str]) -> tuple[int, str]: process = subprocess.Popen( cmd, @@ -306,10 +603,17 @@ def _print_attempt_tail_summary( print("=" * 84, flush=True) -def run_pytest(files, filter_expr=None): +def run_pytest( + files: list[str], + filter_expr: str | None = None, + junit_xml_path: str | None = None, +) -> tuple[int, list[str], dict[str, str]]: if not files: print("No files to run.") - return 0 + return (0, [], {}) + + all_executed_cases: set[str] = set() + all_case_results: dict[str, str] = {} base_cmd = [ sys.executable, @@ -320,13 +624,13 @@ def run_pytest(files, filter_expr=None): "--tb=short", "--no-header", ] - + if junit_xml_path: + base_cmd.extend(["--junit-xml", junit_xml_path]) if filter_expr: base_cmd.extend(["-k", filter_expr]) max_retries = 6 attempt_reports = [] - for i in range(max_retries + 1): is_retry = i > 0 cmd = list(base_cmd) @@ -357,28 +661,33 @@ def run_pytest(files, filter_expr=None): } ) + if junit_xml_path: + all_executed_cases.update( + parse_junit_xml_for_executed_cases(junit_xml_path) + ) + all_case_results.update(parse_junit_xml_for_case_results(junit_xml_path)) + if returncode == 0: if is_retry: print(f"Recovered retryable failures on attempt {i + 1}.") _print_attempt_tail_summary(attempt_reports, len(files)) - return 0 - + return (0, list(all_executed_cases), all_case_results) if returncode == 5: print( "No tests collected (exit code 5). This is expected when filters " "deselect all tests in a partition. Treating as success." ) _print_attempt_tail_summary(attempt_reports, len(files)) - return 0 + return (0, list(all_executed_cases), all_case_results) if not retryable: _print_attempt_tail_summary(attempt_reports, len(files)) - return returncode + return (returncode, list(all_executed_cases), all_case_results) if i == max_retries: print(f"Max retry exceeded ({max_retries})") _print_attempt_tail_summary(attempt_reports, len(files)) - return returncode + return (returncode, list(all_executed_cases), all_case_results) print( f"Retryable failure detected on attempt {i + 1}. " @@ -386,7 +695,11 @@ def run_pytest(files, filter_expr=None): ) _print_attempt_tail_summary(attempt_reports, len(files)) - return attempt_reports[-1]["returncode"] + return ( + attempt_reports[-1]["returncode"], + list(all_executed_cases), + all_case_results, + ) def partition_test_files(files, partition_id, total_partitions): @@ -456,39 +769,351 @@ def _maybe_pin_update_weights_model_pair(suite_files_rel: list[str]) -> None: print(f"Selected {_UPDATE_WEIGHTS_MODEL_PAIR_ENV}={selected_pair} for this CI run") +def _resolve_suite_files( + target_dir: Path, suite_files_rel: list[str], strict: bool +) -> list[str]: + suite_files_abs = [] + for f_rel in suite_files_rel: + f_abs = target_dir / f_rel + if not f_abs.exists(): + msg = f"Test file {f_rel} not found in {target_dir}." + if strict: + print(f"Error: {msg}") + sys.exit(1) + print(f"Warning: {msg} Skipping.") + continue + suite_files_abs.append(str(f_abs)) + return suite_files_abs + + +def _run_file_suite(args, target_dir: Path) -> int: + suite_files_rel = FILE_SUITES[args.suite] + _maybe_pin_update_weights_model_pair(suite_files_rel) + suite_files_abs = _resolve_suite_files( + target_dir, suite_files_rel, args.suite in STRICT_SUITES + ) + + if not suite_files_abs: + print(f"No valid test files found for suite '{args.suite}'.") + return 1 if args.suite in STRICT_SUITES else 0 + + exit_code, _, _ = run_pytest( + suite_files_abs, + filter_expr=args.filter, + junit_xml_path=None, + ) + return exit_code + + +def _get_dynamic_suite_cases(suite: str) -> list[DiffusionTestCase]: + cases = [] + for _, case_group in PARAMETRIZED_CASE_GROUPS[suite]: + cases.extend(case_group) + return cases + + +def _get_parametrized_files_for_case_ids( + suite: str, case_ids: set[str], target_dir: Path +) -> list[str]: + files = [] + for filename, case_group in PARAMETRIZED_CASE_GROUPS[suite]: + if any(case.id in case_ids for case in case_group): + file_path = target_dir / filename + if file_path.exists(): + files.append(str(file_path)) + else: + logger.warning("Test file %s not found in %s", filename, target_dir) + return files + + +def _get_standalone_file(target_dir: Path, suite: str, index: int) -> str | None: + standalone_files = STANDALONE_FILES.get(suite, []) + if index < 0 or index >= len(standalone_files): + return None + file_path = target_dir / standalone_files[index] + if file_path.exists(): + return str(file_path) + logger.warning( + "Standalone test file %s not found in %s", standalone_files[index], target_dir + ) + return None + + +def _run_dynamic_suite(args, target_dir: Path) -> int: + if args.partition_plan_json: + assignment = parse_partition_plan( + suite=args.suite, + partition_id=args.partition_id, + total_partitions=args.total_partitions, + plan_json=args.partition_plan_json, + ) + + rows = [[args.suite, f"{args.partition_id + 1}/{args.total_partitions}"]] + print(tabulate.tabulate(rows, headers=["Suite", "Partition"], tablefmt="psql")) + + total_est_time = 0.0 + executed_cases: list[str] = [] + case_results: dict[str, str] = {} + missing_standalone_estimates: list[str] = [] + standalone_measurements: list[dict] = [] + overall_exit_code = 0 + + if assignment.case_ids: + case_id_set = set(assignment.case_ids) + total_est_time += sum( + get_case_est_time(case_id) for case_id in assignment.case_ids + ) + suite_files = _get_parametrized_files_for_case_ids( + args.suite, case_id_set, target_dir + ) + if not suite_files: + print( + f"No valid parametrized test files found for suite '{args.suite}'." + ) + return 0 + + partition_filter = " or ".join( + f"[{case_id}]" for case_id in assignment.case_ids + ) + filter_expr = ( + f"({partition_filter}) and ({args.filter})" + if args.filter + else partition_filter + ) + + print( + f"Running {len(assignment.case_ids)} parametrized cases with estimated total " + f"{sum(get_case_est_time(case_id) for case_id in assignment.case_ids):.1f}s:" + ) + for case_id in assignment.case_ids: + print(f" - case: {case_id} ({get_case_est_time(case_id):.1f}s)") + print(f"Test files: {[Path(f).name for f in suite_files]}") + print(f"Filter expression: {filter_expr}") + + junit_xml_path = str( + target_dir / f"junit_results_{args.suite}_{args.partition_id}.xml" + ) + exit_code, new_executed_cases, new_case_results = run_pytest( + suite_files, + filter_expr=filter_expr, + junit_xml_path=junit_xml_path, + ) + _merge_execution_results( + executed_cases, case_results, new_executed_cases, new_case_results + ) + if exit_code != 0 and overall_exit_code == 0: + overall_exit_code = exit_code + if exit_code != 0 and not args.continue_on_error: + write_execution_report( + suite=args.suite, + partition_id=args.partition_id, + total_partitions=args.total_partitions, + executed_cases=executed_cases, + is_standalone=False, + standalone_file=None, + case_results=case_results, + missing_standalone_estimates=missing_standalone_estimates, + standalone_measurements=standalone_measurements, + ) + return overall_exit_code + + if assignment.standalone_files: + standalone_estimate = sum( + get_standalone_file_est_time(args.suite, standalone_file)[0] + for standalone_file in assignment.standalone_files + ) + total_est_time += standalone_estimate + print( + f"Running {len(assignment.standalone_files)} standalone file(s) with estimated total " + f"{standalone_estimate:.1f}s:" + ) + for standalone_file in assignment.standalone_files: + est_time, used_fallback_estimate = get_standalone_file_est_time( + args.suite, standalone_file + ) + fallback_suffix = ( + f", fallback estimate {DEFAULT_STANDALONE_EST_TIME_SECONDS:.1f}s" + if used_fallback_estimate + else "" + ) + print( + f" - standalone: {standalone_file} " + f"({est_time:.1f}s{fallback_suffix})" + ) + + for standalone_file in assignment.standalone_files: + exit_code, new_executed_cases, new_case_results, measurement = ( + _run_standalone_file( + args.suite, + standalone_file, + target_dir, + extra_filter=args.filter, + ) + ) + if measurement["used_fallback_estimate"]: + missing_standalone_estimates.append(standalone_file) + standalone_measurements.append(measurement) + _merge_execution_results( + executed_cases, + case_results, + new_executed_cases, + new_case_results, + ) + if exit_code != 0 and overall_exit_code == 0: + overall_exit_code = exit_code + if exit_code != 0 and not args.continue_on_error: + break + + print(f"Partition estimated total time: {total_est_time:.1f}s") + write_execution_report( + suite=args.suite, + partition_id=args.partition_id, + total_partitions=args.total_partitions, + executed_cases=executed_cases, + is_standalone=False, + standalone_file=None, + case_results=case_results, + missing_standalone_estimates=missing_standalone_estimates, + standalone_measurements=standalone_measurements, + ) + return overall_exit_code + + all_cases = _get_dynamic_suite_cases(args.suite) + standalone_files = STANDALONE_FILES.get(args.suite, []) + parametrized_partitions = args.total_partitions - len(standalone_files) + + if parametrized_partitions < 0: + print( + f"Error: total_partitions ({args.total_partitions}) must be >= " + f"standalone files ({len(standalone_files)})" + ) + return 1 + + if args.partition_id < parametrized_partitions: + if not all_cases: + print(f"No cases found for suite '{args.suite}'.") + return 0 + + my_cases = auto_partition(all_cases, args.partition_id, parametrized_partitions) + if not my_cases: + print( + f"No cases assigned to partition {args.partition_id}. Exiting success." + ) + write_execution_report( + suite=args.suite, + partition_id=args.partition_id, + total_partitions=args.total_partitions, + executed_cases=[], + is_standalone=False, + missing_standalone_estimates=[], + standalone_measurements=[], + ) + return 0 + + case_ids = [case.id for case in my_cases] + case_id_set = set(case_ids) + total_est_time = sum(get_case_est_time(case.id) for case in my_cases) + suite_files = _get_parametrized_files_for_case_ids( + args.suite, case_id_set, target_dir + ) + + if not suite_files: + print(f"No valid parametrized test files found for suite '{args.suite}'.") + return 0 + + partition_filter = " or ".join(f"[{case_id}]" for case_id in case_ids) + filter_expr = ( + f"({partition_filter}) and ({args.filter})" + if args.filter + else partition_filter + ) + + rows = [[args.suite, f"{args.partition_id + 1}/{args.total_partitions}"]] + print(tabulate.tabulate(rows, headers=["Suite", "Partition"], tablefmt="psql")) + print( + f"Running {len(my_cases)} cases with estimated total " + f"{total_est_time:.1f}s:" + ) + for case in my_cases: + print(f" - {case.id} ({get_case_est_time(case.id):.1f}s)") + print(f"Test files: {[Path(f).name for f in suite_files]}") + print(f"Filter expression: {filter_expr}") + + junit_xml_path = str( + target_dir / f"junit_results_{args.suite}_{args.partition_id}.xml" + ) + exit_code, executed_cases, case_results = run_pytest( + suite_files, + filter_expr=filter_expr, + junit_xml_path=junit_xml_path, + ) + write_execution_report( + suite=args.suite, + partition_id=args.partition_id, + total_partitions=args.total_partitions, + executed_cases=executed_cases, + is_standalone=False, + case_results=case_results, + missing_standalone_estimates=[], + standalone_measurements=[], + ) + return exit_code + + standalone_idx = args.partition_id - parametrized_partitions + if standalone_idx >= len(standalone_files): + print( + f"ERROR: Standalone partition index {standalone_idx} exceeds available " + f"standalone files ({len(standalone_files)}) for suite '{args.suite}'." + ) + return 1 + + standalone_rel = standalone_files[standalone_idx] + print( + f"Suite: {args.suite} | Partition: {args.partition_id + 1}/{args.total_partitions} (standalone)" + ) + print(f"Running standalone test file: {Path(standalone_rel).name}") + exit_code, executed_cases, case_results, measurement = _run_standalone_file( + args.suite, + standalone_rel, + target_dir, + extra_filter=args.filter, + ) + write_execution_report( + suite=args.suite, + partition_id=args.partition_id, + total_partitions=args.total_partitions, + executed_cases=executed_cases, + is_standalone=True, + standalone_file=standalone_rel, + case_results=case_results, + missing_standalone_estimates=( + [standalone_rel] if measurement["used_fallback_estimate"] else [] + ), + standalone_measurements=[measurement], + ) + return exit_code + + def main(): args = parse_args() - - # 1. resolve base path - current_file_path = Path(__file__).resolve() - test_root_dir = current_file_path.parent + validate_standalone_file_est_times() + test_root_dir = Path(__file__).resolve().parent target_dir = test_root_dir / args.base_dir if not target_dir.exists(): print(f"Error: Target directory {target_dir} does not exist.") sys.exit(1) - # 2. get files from suite - suite_files_rel = SUITES[args.suite] - _maybe_pin_update_weights_model_pair(suite_files_rel) - - suite_files_abs = [] - for f_rel in suite_files_rel: - f_abs = target_dir / f_rel - if not f_abs.exists(): - msg = f"Test file {f_rel} not found in {target_dir}." - if args.suite in STRICT_SUITES: - print(f"Error: {msg}") - sys.exit(1) - print(f"Warning: {msg} Skipping.") - continue - suite_files_abs.append(str(f_abs)) - - if not suite_files_abs: - print(f"No valid test files found for suite '{args.suite}'.") - sys.exit(1 if args.suite in STRICT_SUITES else 0) - if args.suite in COMPONENT_ACCURACY_SUITES: + suite_files_rel = FILE_SUITES[args.suite] + suite_files_abs = _resolve_suite_files( + target_dir, suite_files_rel, args.suite in STRICT_SUITES + ) + + if not suite_files_abs: + print(f"No valid test files found for suite '{args.suite}'.") + sys.exit(1 if args.suite in STRICT_SUITES else 0) + my_files = partition_test_files( suite_files_abs, args.partition_id, args.total_partitions ) @@ -499,7 +1124,7 @@ def main(): headers = ["Suite", "Partition"] rows = [[args.suite, partition_info]] msg = tabulate.tabulate(rows, headers=headers, tablefmt="psql") + "\n" - msg += f"✅ Enabled {len(my_files)} file(s):\n" + msg += f"Enabled {len(my_files)} file(s):\n" for file_path in my_files: msg += f" - {file_path}\n" print(msg, flush=True) @@ -524,45 +1149,14 @@ def main(): ) msg = "\n" + tabulate.tabulate(rows, headers=headers, tablefmt="psql") + "\n" - msg += f"✅ Executed {len(my_files)} file(s):\n" + msg += f"Executed {len(my_files)} file(s):\n" for file_path in my_files: msg += f" - {file_path}\n" print(msg, flush=True) - - sys.exit(exit_code) - - # 3. collect all test items and partition by items (not files) - all_test_items = collect_test_items(suite_files_abs, filter_expr=args.filter) - - if not all_test_items: - print(f"No test items found for suite '{args.suite}'.") - sys.exit(0) - - # Partition by test items - my_items = [ - item - for i, item in enumerate(all_test_items) - if i % args.total_partitions == args.partition_id - ] - - # Print test info at beginning (similar to test/run_suite.py pretty_print_tests) - partition_info = f"{args.partition_id + 1}/{args.total_partitions} (0-based id={args.partition_id})" - headers = ["Suite", "Partition"] - rows = [[args.suite, partition_info]] - msg = tabulate.tabulate(rows, headers=headers, tablefmt="psql") + "\n" - msg += f"✅ Assigned {len(my_items)} test(s) from {len(suite_files_abs)} file(s):\n" - for f in suite_files_abs: - msg += f" - {os.path.basename(f)}\n" - print(msg, flush=True) - - if not my_items: - print("No items assigned to this partition. Exiting success.") - sys.exit(0) - - print(f"Running shard with {len(my_items)} assigned test item(s)") - - # 4. execute with the specific test items - exit_code = run_pytest(my_items) + elif args.suite in PARAMETRIZED_CASE_GROUPS: + exit_code = _run_dynamic_suite(args, target_dir) + else: + exit_code = _run_file_suite(args, target_dir) sys.exit(exit_code) diff --git a/python/sglang/multimodal_gen/test/scripts/gen_diffusion_ci_outputs.py b/python/sglang/multimodal_gen/test/scripts/gen_diffusion_ci_outputs.py index f36e803dd..b604074fa 100755 --- a/python/sglang/multimodal_gen/test/scripts/gen_diffusion_ci_outputs.py +++ b/python/sglang/multimodal_gen/test/scripts/gen_diffusion_ci_outputs.py @@ -148,7 +148,7 @@ def main(): sys.exit(0) # Run pytest with the specific test items (same as run_suite.py) - exit_code = run_pytest(my_items) + exit_code, _, _ = run_pytest(my_items) if exit_code != 0: if args.continue_on_error: diff --git a/python/sglang/multimodal_gen/test/server/perf_baselines.json b/python/sglang/multimodal_gen/test/server/perf_baselines.json index 2ff322ddb..b01abbb66 100644 --- a/python/sglang/multimodal_gen/test/server/perf_baselines.json +++ b/python/sglang/multimodal_gen/test/server/perf_baselines.json @@ -98,7 +98,8 @@ }, "expected_e2e_ms": 14959.11, "expected_avg_denoise_ms": 285.67, - "expected_median_denoise_ms": 286.1 + "expected_median_denoise_ms": 286.1, + "estimated_full_test_time_s": 129.1 }, "qwen_image_t2i_2_gpus": { "stages_ms": { @@ -228,7 +229,8 @@ }, "expected_e2e_ms": 7798.99, "expected_avg_denoise_ms": 150.77, - "expected_median_denoise_ms": 152.45 + "expected_median_denoise_ms": 152.45, + "estimated_full_test_time_s": 78.1 }, "flux_2_image_t2i": { "stages_ms": { @@ -294,7 +296,8 @@ }, "expected_e2e_ms": 329129.82, "expected_avg_denoise_ms": 489.43, - "expected_median_denoise_ms": 497.53 + "expected_median_denoise_ms": 497.53, + "estimated_full_test_time_s": 185.9 }, "flux_2_klein_image_t2i": { "stages_ms": { @@ -600,7 +603,8 @@ }, "expected_e2e_ms": 1292.92, "expected_avg_denoise_ms": 83.75, - "expected_median_denoise_ms": 93.58 + "expected_median_denoise_ms": 93.58, + "estimated_full_test_time_s": 49.8 }, "zimage_image_t2i_fp8": { "stages_ms": { @@ -624,7 +628,8 @@ }, "expected_e2e_ms": 1370.28, "expected_avg_denoise_ms": 85.97, - "expected_median_denoise_ms": 95.83 + "expected_median_denoise_ms": 95.83, + "estimated_full_test_time_s": 49.8 }, "zimage_image_t2i_multi_lora": { "stages_ms": { @@ -738,7 +743,8 @@ }, "expected_e2e_ms": 39706.9, "expected_avg_denoise_ms": 762.57, - "expected_median_denoise_ms": 765.44 + "expected_median_denoise_ms": 765.44, + "estimated_full_test_time_s": 178.8 }, "qwen_image_t2i_cache_dit_enabled": { "stages_ms": { @@ -1336,7 +1342,8 @@ }, "expected_e2e_ms": 42660.88, "expected_avg_denoise_ms": 788.2, - "expected_median_denoise_ms": 790.72 + "expected_median_denoise_ms": 790.72, + "estimated_full_test_time_s": 149.1 }, "fastwan2_2_ti2v_5b": { "stages_ms": { @@ -1355,7 +1362,8 @@ }, "expected_e2e_ms": 7722.91, "expected_avg_denoise_ms": 165.42, - "expected_median_denoise_ms": 165.66 + "expected_median_denoise_ms": 165.66, + "estimated_full_test_time_s": 75.3 }, "fast_hunyuan_video": { "stages_ms": { @@ -1434,7 +1442,8 @@ }, "expected_e2e_ms": 144621.32, "expected_avg_denoise_ms": 3434.22, - "expected_median_denoise_ms": 3428.99 + "expected_median_denoise_ms": 3428.99, + "estimated_full_test_time_s": 427.9 }, "turbo_wan2_2_i2v_a14b_2gpu": { "stages_ms": { @@ -2024,7 +2033,8 @@ }, "expected_e2e_ms": 24895.28, "expected_avg_denoise_ms": 596.59, - "expected_median_denoise_ms": 599.66 + "expected_median_denoise_ms": 599.66, + "estimated_full_test_time_s": 139.8 }, "fsdp-inference": { "stages_ms": { diff --git a/python/sglang/multimodal_gen/test/server/test_server_common.py b/python/sglang/multimodal_gen/test/server/test_server_common.py index c5ca8eeec..df4f3f3fc 100644 --- a/python/sglang/multimodal_gen/test/server/test_server_common.py +++ b/python/sglang/multimodal_gen/test/server/test_server_common.py @@ -8,6 +8,7 @@ If the actual run is significantly better than the baseline, the improved cases from __future__ import annotations import os +import time from pathlib import Path from typing import Any, Callable @@ -50,10 +51,15 @@ from sglang.multimodal_gen.test.test_utils import ( logger = init_logger(__name__) +# Track test cases missing estimated_full_test_time_s for time measurement output +_MISSING_ESTIMATED_TIME_CASES: set[str] = set() +_PENDING_BASELINE_DUMPS: dict[str, tuple["PerformanceSummary", bool]] = {} + @pytest.fixture def diffusion_server(case: DiffusionTestCase) -> ServerContext: """Start a diffusion server for a single case and tear it down afterwards.""" + _fixture_start_time = time.perf_counter() server_args = case.server_args # Skip ring attention tests on AMD/ROCm - Ring Attention requires Flash Attention @@ -118,10 +124,26 @@ def diffusion_server(case: DiffusionTestCase) -> ServerContext: env_vars["SGLANG_CACHE_DIT_ENABLED"] = "true" # start server + wait_deadline = float(os.environ.get("SGLANG_TEST_WAIT_SECS", "1200")) + logger.info( + "[server-test] Starting server for test case: %s\n" + " Model: %s\n" + " Port: %s\n" + " Wait deadline: %ss\n" + " Extra args: %s\n" + " Num GPUs: %s", + case.id, + server_args.model_path, + port, + wait_deadline, + extra_args, + server_args.num_gpus, + ) + manager = ServerManager( model=server_args.model_path, port=port, - wait_deadline=float(os.environ.get("SGLANG_TEST_WAIT_SECS", "1200")), + wait_deadline=wait_deadline, extra_args=extra_args, env_vars=env_vars, ) @@ -156,6 +178,38 @@ def diffusion_server(case: DiffusionTestCase) -> ServerContext: finally: ctx.cleanup() + _fixture_end_time = time.perf_counter() + _measured_full_time = _fixture_end_time - _fixture_start_time + is_baseline_generation_mode = os.environ.get("SGLANG_GEN_BASELINE", "0") == "1" + + pending_dump = _PENDING_BASELINE_DUMPS.pop(case.id, None) + if pending_dump is not None: + summary, missing_scenario = pending_dump + DiffusionServerBase()._dump_baseline_for_testcase( + case, + summary, + missing_scenario=missing_scenario, + measured_full_time=_measured_full_time, + ) + + scenario = BASELINE_CONFIG.scenarios.get(case.id) + needs_estimated_time = ( + scenario is None or scenario.estimated_full_test_time_s is None + ) + + if needs_estimated_time and not is_baseline_generation_mode: + _MISSING_ESTIMATED_TIME_CASES.add(case.id) + logger.error( + f'\n{"=" * 60}\n' + f'Add "estimated_full_test_time_s" to scenario "{case.id}":\n\n' + f"File: python/sglang/multimodal_gen/test/server/perf_baselines.json\n\n" + f' "{case.id}": {{\n' + f" ...\n" + f' "estimated_full_test_time_s": {_measured_full_time:.1f}\n' + f" }}\n" + f'{"=" * 60}\n' + ) + class DiffusionServerBase: """Performance tests for all diffusion models/scenarios. @@ -271,6 +325,16 @@ Consider updating perf_baselines.json with the snippets below: if not is_baseline_generation_mode: missing_scenario = True + # Check for missing estimated_full_test_time_s + missing_estimated_time = False + if ( + not missing_scenario + and not is_baseline_generation_mode + and scenario.estimated_full_test_time_s is None + ): + missing_estimated_time = True + _MISSING_ESTIMATED_TIME_CASES.add(case.id) + validator_name = case.server_args.custom_validator or "default" validator_class = VALIDATOR_REGISTRY.get(validator_name, PerformanceValidator) @@ -283,7 +347,11 @@ Consider updating perf_baselines.json with the snippets below: summary = validator.collect_metrics(perf_record) if case.run_perf_check: - if is_baseline_generation_mode or missing_scenario: + if is_baseline_generation_mode: + _PENDING_BASELINE_DUMPS[case.id] = (summary, missing_scenario) + return + + if missing_scenario: self._dump_baseline_for_testcase(case, summary, missing_scenario) if missing_scenario: pytest.fail( @@ -409,6 +477,7 @@ Consider updating perf_baselines.json with the snippets below: case: DiffusionTestCase, summary: "PerformanceSummary", missing_scenario: bool = False, + measured_full_time: float | None = None, ) -> None: """Dump performance metrics as a JSON scenario for baselines.""" import json @@ -426,6 +495,9 @@ Consider updating perf_baselines.json with the snippets below: "expected_median_denoise_ms": round(summary.median_denoise_ms, 2), } + if measured_full_time is not None: + baseline["estimated_full_test_time_s"] = round(measured_full_time, 1) + # Video-specific metrics if case.server_args.modality == "video": if "per_frame_generation" not in baseline["stages_ms"]: diff --git a/python/sglang/multimodal_gen/test/server/testcase_configs.py b/python/sglang/multimodal_gen/test/server/testcase_configs.py index e1c837691..00f9f8661 100644 --- a/python/sglang/multimodal_gen/test/server/testcase_configs.py +++ b/python/sglang/multimodal_gen/test/server/testcase_configs.py @@ -4,12 +4,13 @@ Configuration and data structures for diffusion performance tests. Usage: pytest python/sglang/multimodal_gen/test/server/test_server_a.py -# for a single testcase, look for the name of the testcases in DIFFUSION_CASES +# for a single testcase, look for the name of the testcase in ONE_GPU_CASES_A, +# ONE_GPU_CASES_B, ONE_GPU_CASES_C, TWO_GPU_CASES_A, or TWO_GPU_CASES_B pytest python/sglang/multimodal_gen/test/server/test_server_a.py -k qwen_image_t2i To add a new testcase: -1. add your testcase with case-id: `my_new_test_case_id` to DIFFUSION_CASES +1. add your testcase with case-id: `my_new_test_case_id` to the appropriate `*_CASES_*` list 2. run `SGLANG_GEN_BASELINE=1 pytest -s python/sglang/multimodal_gen/test/server/ -k my_new_test_case_id` 3. insert or override the corresponding scenario in `scenarios` section of perf_baselines.json with the output baseline of step-2 @@ -109,6 +110,7 @@ class ScenarioConfig: expected_e2e_ms: float expected_avg_denoise_ms: float expected_median_denoise_ms: float + estimated_full_test_time_s: float | None = None @dataclass @@ -140,6 +142,7 @@ class BaselineConfig: expected_e2e_ms=float(cfg["expected_e2e_ms"]), expected_avg_denoise_ms=float(cfg["expected_avg_denoise_ms"]), expected_median_denoise_ms=float(cfg["expected_median_denoise_ms"]), + estimated_full_test_time_s=cfg.get("estimated_full_test_time_s"), ) return cls( @@ -164,6 +167,7 @@ class BaselineConfig: expected_e2e_ms=float(cfg["expected_e2e_ms"]), expected_avg_denoise_ms=float(cfg["expected_avg_denoise_ms"]), expected_median_denoise_ms=float(cfg["expected_median_denoise_ms"]), + estimated_full_test_time_s=cfg.get("estimated_full_test_time_s"), ) self.scenarios.update(scenarios_new) diff --git a/python/sglang/multimodal_gen/test/test_utils.py b/python/sglang/multimodal_gen/test/test_utils.py index d25d6b18a..cda892d7c 100644 --- a/python/sglang/multimodal_gen/test/test_utils.py +++ b/python/sglang/multimodal_gen/test/test_utils.py @@ -5,6 +5,7 @@ import json import os import socket import subprocess +import sys import tempfile import time from dataclasses import dataclass @@ -193,6 +194,26 @@ def post_json( return httpx.post(urljoin(base_url, path), json=payload, timeout=timeout) +def run_command(command: list[str]) -> bool: + """Run a CLI command and return whether it succeeded.""" + print(f"Running command: {' '.join(command)}", flush=True) + with subprocess.Popen( + command, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + encoding="utf-8", + ) as process: + assert process.stdout is not None + for line in process.stdout: + sys.stdout.write(line) + process.wait() + if process.returncode == 0: + return True + print(f"Command failed with exit code {process.returncode}", flush=True) + return False + + # --------------------------------------------------------------------------- # GPU memory helpers (nvidia-smi) # --------------------------------------------------------------------------- diff --git a/scripts/ci/utils/diffusion/compute_diffusion_partitions.py b/scripts/ci/utils/diffusion/compute_diffusion_partitions.py new file mode 100755 index 000000000..a1af41027 --- /dev/null +++ b/scripts/ci/utils/diffusion/compute_diffusion_partitions.py @@ -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() diff --git a/scripts/ci/utils/diffusion/diffusion_case_parser.py b/scripts/ci/utils/diffusion/diffusion_case_parser.py new file mode 100755 index 000000000..0dabeadf3 --- /dev/null +++ b/scripts/ci/utils/diffusion/diffusion_case_parser.py @@ -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 diff --git a/scripts/ci/utils/diffusion/verify_diffusion_coverage.py b/scripts/ci/utils/diffusion/verify_diffusion_coverage.py new file mode 100755 index 000000000..5659fa828 --- /dev/null +++ b/scripts/ci/utils/diffusion/verify_diffusion_coverage.py @@ -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 + +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:". + """ + 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()