minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[diffusion] CI: dynamic load-balanced partitioning for diffusion CI (#15528)

Browse Source 
Co-authored-by: daiweitao <dwti614707404@163.com>
Co-authored-by: SGLang CI <ci@sglang.ai>
This commit is contained in:
Prozac614
2026-04-12 13:02:43 +08:00
committed by GitHub
co-authored by daiweitao SGLang CI
parent d6c9d9116b
commit 45472d70cc
11 changed files with 1877 additions and 175 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.github/workflows/pr-test-multimodal-gen.yml
+108 -12
View File
@@ -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/
python/sglang/multimodal_gen/test/cli/test_generate_common.py
+1 -27
View File
@@ -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 = []
python/sglang/multimodal_gen/test/run_suite.py
+715 -121
View File
File diff suppressed because it is too large Load Diff
python/sglang/multimodal_gen/test/scripts/gen_diffusion_ci_outputs.py
+1 -1
View File
@@ -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:
python/sglang/multimodal_gen/test/server/perf_baselines.json
+20 -10
View File
@@ -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": {
python/sglang/multimodal_gen/test/server/test_server_common.py
+74 -2
View File
@@ -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"]:
python/sglang/multimodal_gen/test/server/testcase_configs.py
+6 -2
View File
@@ -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)
python/sglang/multimodal_gen/test/test_utils.py
+21
View File
@@ -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)
# ---------------------------------------------------------------------------
scripts/ci/utils/diffusion/compute_diffusion_partitions.py
Executable
+281
View File
@@ -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()
scripts/ci/utils/diffusion/diffusion_case_parser.py
Executable
+326
View File
@@ -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
scripts/ci/utils/diffusion/verify_diffusion_coverage.py
Executable
+324
View File
@@ -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()