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
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)
# ---------------------------------------------------------------------------