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

[diffusion] refactor and added tests for Flux, T2V, TI2V, I2V(#13344)

Browse Source
This commit is contained in:
Adarsh Shirawalmath
2025-11-16 17:28:04 +05:30
committed by GitHub
parent efc5d8f5ed
commit d724670873
5 changed files with 1585 additions and 898 deletions
Download Patch File Download Diff File Expand all files Collapse all files

219
python/sglang/multimodal_gen/test/server/diffusion_config.py Normal file
View File

@@ -0,0 +1,219 @@
"""
Configuration and data structures for diffusion performance tests.
"""
from __future__ import annotations
import json
import os
from dataclasses import dataclass
from pathlib import Path
from typing import Sequence
@dataclass
class ToleranceConfig:
"""Tolerance ratios for performance validation."""
e2e: float
stage: float
denoise_step: float
denoise_agg: float
@dataclass
class ScenarioConfig:
"""Expected performance metrics for a test scenario."""
stages_ms: dict[str, float]
denoise_step_ms: dict[int, float]
expected_e2e_ms: float
expected_avg_denoise_ms: float
expected_median_denoise_ms: float
@dataclass
class BaselineConfig:
"""Full baseline configuration."""
scenarios: dict[str, ScenarioConfig]
step_fractions: Sequence[float]
warmup_defaults: dict[str, int]
tolerances: ToleranceConfig
@classmethod
def load(cls, path: Path) -> BaselineConfig:
"""Load baseline configuration from JSON file."""
with path.open("r", encoding="utf-8") as fh:
data = json.load(fh)
tol_data = data["tolerances"]
tolerances = ToleranceConfig(
e2e=float(os.getenv("SGLANG_E2E_TOLERANCE", tol_data["e2e"])),
stage=float(os.getenv("SGLANG_STAGE_TIME_TOLERANCE", tol_data["stage"])),
denoise_step=float(
os.getenv("SGLANG_DENOISE_STEP_TOLERANCE", tol_data["denoise_step"])
),
denoise_agg=float(
os.getenv("SGLANG_DENOISE_AGG_TOLERANCE", tol_data["denoise_agg"])
),
)
scenarios = {}
for name, cfg in data["scenarios"].items():
scenarios[name] = ScenarioConfig(
stages_ms=cfg["stages_ms"],
denoise_step_ms={int(k): v for k, v in cfg["denoise_step_ms"].items()},
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"]),
)
return cls(
scenarios=scenarios,
step_fractions=tuple(data["sampling"]["step_fractions"]),
warmup_defaults=data["sampling"].get("warmup_requests", {}),
tolerances=tolerances,
)
@dataclass(frozen=True)
class DiffusionCase:
"""Configuration for a single model/scenario test case."""
id: str # pytest test id
model_path: str # HF repo or local path
scenario_name: str # key into BASELINE_CONFIG.scenarios
modality: str = "image" # "image" or "video" or "3d"
prompt: str | None = None # text prompt for generation
output_size: str = "1024x1024" # output image dimensions (or video resolution)
num_frames: int | None = None # for video: number of frames
fps: int | None = None # for video: frames per second
warmup_text: int = 1 # number of text-to-image/video warmups
warmup_edit: int = 0 # number of image/video-edit warmups
image_edit_prompt: str | None = None # prompt for editing
image_edit_path: Path | str | None = (
None # input image/video for editing (Path or URL)
)
startup_grace_seconds: float = 0.0 # wait time after server starts
custom_validator: str | None = None # optional custom validator name
seconds: int = 4 # for video: duration in seconds
def is_image_url(self) -> bool:
"""Check if image_edit_path is a URL."""
if self.image_edit_path is None:
return False
return isinstance(self.image_edit_path, str) and (
self.image_edit_path.startswith("http://")
or self.image_edit_path.startswith("https://")
)
@dataclass
class PerformanceSummary:
"""Summary of performance metrics."""
e2e_ms: float
avg_denoise_ms: float
median_denoise_ms: float
stage_metrics: dict[str, float]
sampled_steps: dict[int, float]
frames_per_second: float | None = None
total_frames: int | None = None
avg_frame_time_ms: float | None = None
# Common paths
IMAGE_INPUT_FILE = Path(__file__).resolve().parents[1] / "test_files" / "girl.jpg"
# All test cases with clean default values
# To test different models, simply add more DiffusionCase entries
DIFFUSION_CASES: list[DiffusionCase] = [
# === Text to Image (T2I) ===
DiffusionCase(
id="qwen_image_t2i",
model_path="Qwen/Qwen-Image",
scenario_name="text_to_image",
modality="image",
prompt="A futuristic cityscape at sunset with flying cars",
output_size="1024x1024",
warmup_text=1,
warmup_edit=0,
startup_grace_seconds=30.0,
),
DiffusionCase(
id="flux_image_t2i",
model_path="black-forest-labs/FLUX.1-dev",
scenario_name="text_to_image",
modality="image",
prompt="A futuristic cityscape at sunset with flying cars",
output_size="1024x1024",
warmup_text=1,
warmup_edit=0,
startup_grace_seconds=30.0,
),
# === Text and Image to Image (TI2I) ===
DiffusionCase(
id="qwen_image_edit_ti2i",
model_path="Qwen/Qwen-Image-Edit",
scenario_name="image_edit",
modality="image",
prompt=None, # not used for editing
output_size="1024x1536",
warmup_text=0,
warmup_edit=1,
image_edit_prompt="Convert 2D style to 3D style",
image_edit_path="https://github.com/lm-sys/lm-sys.github.io/releases/download/test/TI2I_Qwen_Image_Edit_Input.jpg",
startup_grace_seconds=30.0,
),
# === Text to Video (T2V) ===
DiffusionCase(
id="fastwan2_1_t2v",
model_path="Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
scenario_name="text_to_video",
modality="video",
prompt="A curious raccoon",
output_size="848x480",
seconds=4,
warmup_text=0, # warmups only for image gen models
warmup_edit=0,
startup_grace_seconds=30.0,
custom_validator="video",
),
# # === Image to Video (I2V) ===
# DiffusionCase(
# id="wan2_1_i2v_480p",
# model_path="Wan-AI/Wan2.1-I2V-14B-Diffusers",
# scenario_name="image_to_video",
# modality="video",
# prompt="generate", # passing in something since failing if no prompt is passed
# warmup_text=0, # warmups only for image gen models
# warmup_edit=0,
# output_size="1024x1536",
# image_edit_prompt="generate",
# image_edit_path="https://github.com/lm-sys/lm-sys.github.io/releases/download/test/TI2I_Qwen_Image_Edit_Input.jpg",
# startup_grace_seconds=30.0,
# custom_validator="video",
# seconds=4,
# ),
# === Text and Image to Video (TI2V) ===
DiffusionCase(
id="wan2_2_ti2v_5b",
model_path="Wan-AI/Wan2.2-TI2V-5B-Diffusers",
scenario_name="text_image_to_video",
modality="video",
prompt="Animate this image",
output_size="832x1104",
warmup_text=0, # warmups only for image gen models
warmup_edit=0,
image_edit_prompt="Add dynamic motion to the scene",
image_edit_path="https://github.com/lm-sys/lm-sys.github.io/releases/download/test/TI2I_Qwen_Image_Edit_Input.jpg",
startup_grace_seconds=30.0,
custom_validator="video",
seconds=4,
),
]
# Load global configuration
BASELINE_CONFIG = BaselineConfig.load(Path(__file__).with_name("perf_baselines.json"))

420
python/sglang/multimodal_gen/test/server/diffusion_server.py Normal file
View File

@@ -0,0 +1,420 @@
"""
Server management and performance validation for diffusion tests.
"""
from __future__ import annotations
import os
import statistics
import subprocess
import tempfile
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Sequence
from urllib.request import urlopen
from openai import OpenAI
from sglang.multimodal_gen.runtime.utils.common import kill_process_tree
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
from sglang.multimodal_gen.test.server.diffusion_config import (
PerformanceSummary,
ScenarioConfig,
ToleranceConfig,
)
from sglang.multimodal_gen.test.test_utils import (
prepare_perf_log,
sample_step_indices,
validate_image,
)
logger = init_logger(__name__)
def download_image_from_url(url: str) -> Path:
"""Download an image from a URL to a temporary file.
Args:
url: The URL of the image to download
Returns:
Path to the downloaded temporary file
"""
logger.info(f"Downloading image from URL: {url}")
# Determine file extension from URL
ext = ".jpg" # default
if url.lower().endswith((".png", ".jpeg", ".jpg", ".webp", ".gif")):
ext = url[url.rfind(".") :]
# Create temporary file
temp_file = (
Path(tempfile.gettempdir()) / f"diffusion_test_image_{int(time.time())}{ext}"
)
try:
with urlopen(url, timeout=30) as response:
temp_file.write_bytes(response.read())
logger.info(f"Downloaded image to: {temp_file}")
return temp_file
except Exception as e:
logger.error(f"Failed to download image from {url}: {e}")
raise
@dataclass
class ServerContext:
"""Context for a running diffusion server."""
port: int
process: subprocess.Popen
model: str
stdout_file: Path
perf_log_path: Path
log_dir: Path
_stdout_fh: Any = field(repr=False)
def cleanup(self) -> None:
"""Clean up server resources."""
try:
kill_process_tree(self.process.pid)
except Exception:
pass
try:
self._stdout_fh.flush()
self._stdout_fh.close()
except Exception:
pass
class ServerManager:
"""Manages diffusion server lifecycle."""
def __init__(
self,
model: str,
port: int,
wait_deadline: float = 1200.0,
extra_args: str = "",
):
self.model = model
self.port = port
self.wait_deadline = wait_deadline
self.extra_args = extra_args
def start(self) -> ServerContext:
"""Start the diffusion server and wait for readiness."""
log_dir, perf_log_path = prepare_perf_log(Path(__file__))
safe_model_name = self.model.replace("/", "_")
stdout_path = (
Path(tempfile.gettempdir())
/ f"sgl_server_{self.port}_{safe_model_name}.log"
)
stdout_path.unlink(missing_ok=True)
command = [
"sglang",
"serve",
"--model-path",
self.model,
"--port",
str(self.port),
"--log-level=debug",
]
if self.extra_args.strip():
command.extend(self.extra_args.strip().split())
env = os.environ.copy()
env["SGL_DIFFUSION_STAGE_LOGGING"] = "1"
env["SGLANG_PERF_LOG_DIR"] = log_dir.as_posix()
stdout_fh = stdout_path.open("w", encoding="utf-8", buffering=1)
process = subprocess.Popen(
command,
stdout=stdout_fh,
stderr=subprocess.STDOUT,
text=True,
bufsize=1,
env=env,
)
logger.info(
"[server-test] Starting server pid=%s, model=%s, log=%s",
process.pid,
self.model,
stdout_path,
)
self._wait_for_ready(process, stdout_path)
return ServerContext(
port=self.port,
process=process,
model=self.model,
stdout_file=stdout_path,
perf_log_path=perf_log_path,
log_dir=log_dir,
_stdout_fh=stdout_fh,
)
def _wait_for_ready(self, process: subprocess.Popen, stdout_path: Path) -> None:
"""Wait for server to become ready."""
start = time.time()
ready_message = "Application startup complete."
while time.time() - start < self.wait_deadline:
if process.poll() is not None:
tail = self._get_log_tail(stdout_path)
raise RuntimeError(
f"Server exited early (code {process.returncode}).\n{tail}"
)
if stdout_path.exists():
try:
content = stdout_path.read_text(encoding="utf-8", errors="ignore")
if ready_message in content:
logger.info("[server-test] Server ready")
return
except Exception as e:
logger.debug("Could not read log yet: %s", e)
elapsed = int(time.time() - start)
logger.info("[server-test] Waiting for server... elapsed=%ss", elapsed)
time.sleep(5)
tail = self._get_log_tail(stdout_path)
raise TimeoutError(f"Server not ready within {self.wait_deadline}s.\n{tail}")
@staticmethod
def _get_log_tail(path: Path, lines: int = 200) -> str:
"""Get the last N lines from a log file."""
try:
content = path.read_text(encoding="utf-8", errors="ignore")
return "\n".join(content.splitlines()[-lines:])
except Exception:
return ""
class WarmupRunner:
"""Handles warmup requests for a server."""
def __init__(
self,
port: int,
model: str,
prompt: str,
output_size: str,
):
self.client = OpenAI(
api_key="sglang-anything",
base_url=f"http://localhost:{port}/v1",
)
self.model = model
self.prompt = prompt
self.output_size = output_size
def run_text_warmups(self, count: int) -> None:
"""Run text-to-image warmup requests."""
if count <= 0:
return
logger.info("[server-test] Running %s text warm-up(s)", count)
for _ in range(count):
result = self.client.images.generate(
model=self.model,
prompt=self.prompt,
n=1,
size=self.output_size,
response_format="b64_json",
)
validate_image(result.data[0].b64_json)
def run_edit_warmups(
self,
count: int,
edit_prompt: str,
image_path: Path,
) -> None:
"""Run image-edit warmup requests."""
if count <= 0:
return
if not image_path.exists():
logger.warning(
"[server-test] Skipping edit warmup: image missing at %s", image_path
)
return
logger.info("[server-test] Running %s edit warm-up(s)", count)
for _ in range(count):
with image_path.open("rb") as fh:
result = self.client.images.edit(
model=self.model,
image=fh,
prompt=edit_prompt,
n=1,
size=self.output_size,
response_format="b64_json",
)
validate_image(result.data[0].b64_json)
class PerformanceValidator:
"""Validates performance metrics against expectations."""
def __init__(
self,
scenario: ScenarioConfig,
tolerances: ToleranceConfig,
step_fractions: Sequence[float],
):
self.scenario = scenario
self.tolerances = tolerances
self.step_fractions = step_fractions
def validate(
self,
perf_record: dict,
stage_metrics: dict,
) -> PerformanceSummary:
"""Validate all performance metrics and return summary."""
self._validate_e2e(perf_record)
avg_denoise, median_denoise = self._validate_denoise_agg(perf_record)
sampled_steps = self._validate_denoise_steps(perf_record)
self._validate_stages(stage_metrics)
return PerformanceSummary(
e2e_ms=float(perf_record["total_duration_ms"]),
avg_denoise_ms=avg_denoise,
median_denoise_ms=median_denoise,
stage_metrics=stage_metrics,
sampled_steps=sampled_steps,
)
def _validate_e2e(self, perf_record: dict) -> None:
"""Validate end-to-end performance."""
e2e_ms = float(perf_record.get("total_duration_ms", 0.0))
assert e2e_ms > 0, "E2E duration missing"
upper = self.scenario.expected_e2e_ms * (1 + self.tolerances.e2e)
assert e2e_ms <= upper, f"E2E {e2e_ms:.2f}ms exceeds {upper:.2f}ms"
def _validate_denoise_agg(self, perf_record: dict) -> tuple[float, float]:
"""Validate aggregate denoising metrics."""
steps = [
s
for s in perf_record.get("steps", []) or []
if s.get("name") == "denoising_step_guided" and "duration_ms" in s
]
assert steps, "Denoising step timings missing"
durations = [float(s["duration_ms"]) for s in steps]
avg = sum(durations) / len(durations)
median = statistics.median(durations)
avg_upper = self.scenario.expected_avg_denoise_ms * (
1 + self.tolerances.denoise_agg
)
med_upper = self.scenario.expected_median_denoise_ms * (
1 + self.tolerances.denoise_agg
)
assert avg <= avg_upper, f"Avg denoise {avg:.2f}ms exceeds {avg_upper:.2f}ms"
assert (
median <= med_upper
), f"Median denoise {median:.2f}ms exceeds {med_upper:.2f}ms"
return avg, median
def _validate_denoise_steps(self, perf_record: dict) -> dict[int, float]:
"""Validate individual denoising steps."""
steps = [
s
for s in perf_record.get("steps", []) or []
if s.get("name") == "denoising_step_guided" and "duration_ms" in s
]
per_step = {
int(s["index"]): float(s["duration_ms"])
for s in steps
if s.get("index") is not None
}
sample_indices = sample_step_indices(per_step, self.step_fractions)
sampled = {idx: per_step[idx] for idx in sample_indices}
for idx in sample_indices:
expected = self.scenario.denoise_step_ms.get(idx)
if expected is None:
continue
actual = per_step[idx]
upper = expected * (1 + self.tolerances.denoise_step)
assert actual <= upper, f"Step {idx}: {actual:.2f}ms > {upper:.2f}ms"
return sampled
def _validate_stages(self, stage_metrics: dict) -> None:
"""Validate stage-level metrics."""
assert stage_metrics, "Stage metrics missing"
for stage, expected in self.scenario.stages_ms.items():
actual = stage_metrics.get(stage)
assert actual is not None, f"Stage {stage} timing missing"
upper = expected * (1 + self.tolerances.stage)
assert actual <= upper, f"Stage {stage}: {actual:.2f}ms > {upper:.2f}ms"
class VideoPerformanceValidator(PerformanceValidator):
"""Extended validator for video diffusion with frame-level metrics."""
def validate(
self,
perf_record: dict,
stage_metrics: dict,
num_frames: int | None = None,
) -> PerformanceSummary:
"""Validate video metrics including frame generation rates."""
summary = super().validate(perf_record, stage_metrics)
if num_frames and summary.e2e_ms > 0:
summary.total_frames = num_frames
summary.avg_frame_time_ms = summary.e2e_ms / num_frames
summary.frames_per_second = 1000.0 / summary.avg_frame_time_ms
self._validate_frame_rate(summary)
return summary
def _validate_frame_rate(self, summary: PerformanceSummary) -> None:
"""Validate frame generation performance."""
expected_frame_time = self.scenario.stages_ms.get("per_frame_generation")
if expected_frame_time and summary.avg_frame_time_ms:
upper = expected_frame_time * (1 + self.tolerances.stage)
assert (
summary.avg_frame_time_ms <= upper
), f"Avg frame time {summary.avg_frame_time_ms:.2f}ms exceeds {upper:.2f}ms"
def _validate_stages(self, stage_metrics: dict) -> None:
"""Validate video-specific stages."""
assert stage_metrics, "Stage metrics missing"
for stage, expected in self.scenario.stages_ms.items():
if stage == "per_frame_generation":
continue
actual = stage_metrics.get(stage)
assert actual is not None, f"Stage {stage} timing missing"
upper = expected * (1 + self.tolerances.stage)
assert actual <= upper, f"Stage {stage}: {actual:.2f}ms > {upper:.2f}ms"
# Registry of validators by name
VALIDATOR_REGISTRY = {
"default": PerformanceValidator,
"video": VideoPerformanceValidator,
}

224
python/sglang/multimodal_gen/test/server/perf_baselines.json
View File

@@ -1,82 +1,142 @@
{
"metadata": {
"model": "Qwen/Qwen-Image",
"hardware": "CI H100 80GB pool",
"description": "Reference numbers captured from the CI diffusion server baseline run"
},
"tolerances": {
"e2e": 0.25,
"stage": 0.3,
"denoise_step": 0.1,
"denoise_agg": 0.1
},
"sampling": {
"step_fractions": [
0.0,
0.2,
0.4,
0.6,
0.8,
1.0
],
"warmup_requests": {
"text": 1,
"image_edit": 0
}
},
"scenarios": {
"text_to_image": {
"notes": "Single-image generation using the default prompt",
"expected_e2e_ms": 74500.0,
"expected_avg_denoise_ms": 422.42,
"expected_median_denoise_ms": 410.62,
"stages_ms": {
"InputValidationStage": 0.1,
"TextEncodingStage": 834.2,
"ConditioningStage": 0.1,
"TimestepPreparationStage": 10.6,
"LatentPreparationStage": 5.2,
"DenoisingStage": 21202.6,
"DecodingStage": 476.12
},
"denoise_step_ms": {
"0": 1077.77, "1": 345.13, "2": 413.8, "3": 405.49, "4": 408.14, "5": 409.06,
"6": 408.85, "7": 410.53, "8": 407.51, "9": 409.44, "10": 408.65, "11": 410.14,
"12": 411.74, "13": 409.59, "14": 409.17, "15": 410.78, "16": 410.66, "17": 410.58,
"18": 411.27, "19": 410.51, "20": 409.03, "21": 410.16, "22": 409.42, "23": 411.03,
"24": 410.18, "25": 409.72, "26": 410.26, "27": 410.21, "28": 410.71, "29": 410.76,
"30": 411.06, "31": 410.1, "32": 410.55, "33": 410.77, "34": 410.74, "35": 411.75,
"36": 410.78, "37": 411.56, "38": 410.85, "39": 411.08, "40": 411.12, "41": 411.1,
"42": 411.09, "43": 410.87, "44": 411.37, "45": 411.68, "46": 411.0, "47": 410.09,
"48": 412.72, "49": 410.42
}
},
"image_edit": {
"notes": "single uploaded reference image, Qwen/Qwen-Image-Edit",
"expected_e2e_ms": 138500.0,
"expected_avg_denoise_ms": 720.0,
"expected_median_denoise_ms": 718.0,
"stages_ms": {
"InputValidationStage": 14,
"ImageEncodingStage": 1400.0,
"ImageVAEEncodingStage": 252.76,
"ConditioningStage": 0.13,
"TimestepPreparationStage": 13.78,
"LatentPreparationStage": 9.18,
"DenoisingStage": 36000.0,
"DecodingStage": 645
},
"denoise_step_ms": {
"0": 720.0, "1": 720.0, "2": 720.0, "3": 720.0, "4": 720.0, "5": 720.0,
"6": 720.0, "7": 720.0, "8": 720.0, "9": 720.0, "10": 720.0, "11": 720.0,
"12": 720.0, "13": 720.0, "14": 720.0, "15": 720.0, "16": 720.0, "17": 720.0,
"18": 720.0, "19": 720.0, "20": 720.0, "21": 720.0, "22": 720.0, "23": 720.0,
"24": 720.0, "25": 720.0, "26": 720.0, "27": 720.0, "28": 720.0, "29": 720.0,
"30": 720.0, "31": 720.0, "32": 720.0, "33": 720.0, "34": 720.0, "35": 720.0,
"36": 720.0, "37": 720.0, "38": 720.0, "39": 720.0, "40": 720.0, "41": 720.0,
"42": 720.0, "43": 720.0, "44": 720.0, "45": 720.0, "46": 720.0, "47": 720.0,
"48": 720.0, "49": 720.0
}
}
}
}
{
"metadata": {
"model": "Diffusion Server",
"hardware": "CI H100 80GB pool",
"description": "Reference numbers captured from the CI diffusion server baseline run"
},
"tolerances": {
"e2e": 0.25,
"stage": 0.3,
"denoise_step": 0.1,
"denoise_agg": 0.1
},
"sampling": {
"step_fractions": [
0.0,
0.2,
0.4,
0.6,
0.8,
1.0
],
"warmup_requests": {
"text": 1,
"image_edit": 0
}
},
"scenarios": {
"text_to_image": {
"notes": "Single-image generation using the default prompt",
"expected_e2e_ms": 74500.0,
"expected_avg_denoise_ms": 422.42,
"expected_median_denoise_ms": 410.62,
"stages_ms": {
"InputValidationStage": 0.1,
"TextEncodingStage": 834.2,
"ConditioningStage": 0.1,
"TimestepPreparationStage": 10.6,
"LatentPreparationStage": 9.0,
"DenoisingStage": 21202.6,
"DecodingStage": 476.12
},
"denoise_step_ms": {
"0": 1077.77, "1": 345.13, "2": 413.8, "3": 405.49, "4": 408.14, "5": 409.06,
"6": 408.85, "7": 410.53, "8": 407.51, "9": 409.44, "10": 408.65, "11": 410.14,
"12": 411.74, "13": 409.59, "14": 409.17, "15": 410.78, "16": 410.66, "17": 410.58,
"18": 411.27, "19": 410.51, "20": 409.03, "21": 410.16, "22": 409.42, "23": 411.03,
"24": 410.18, "25": 409.72, "26": 410.26, "27": 410.21, "28": 410.71, "29": 410.76,
"30": 411.06, "31": 410.1, "32": 410.55, "33": 410.77, "34": 410.74, "35": 411.75,
"36": 410.78, "37": 411.56, "38": 410.85, "39": 411.08, "40": 411.12, "41": 411.1,
"42": 411.09, "43": 410.87, "44": 411.37, "45": 411.68, "46": 411.0, "47": 410.09,
"48": 412.72, "49": 410.42
}
},
"image_edit": {
"notes": "single uploaded reference image, Qwen/Qwen-Image-Edit",
"expected_e2e_ms": 138500.0,
"expected_avg_denoise_ms": 720.0,
"expected_median_denoise_ms": 718.0,
"stages_ms": {
"InputValidationStage": 23,
"ImageEncodingStage": 990.0,
"ImageVAEEncodingStage": 340.0,
"ConditioningStage": 0.13,
"TimestepPreparationStage": 13.78,
"LatentPreparationStage": 10.0,
"DenoisingStage": 36000.0,
"DecodingStage": 645
},
"denoise_step_ms": {
"0": 720.0, "1": 720.0, "2": 720.0, "3": 720.0, "4": 720.0, "5": 720.0,
"6": 720.0, "7": 720.0, "8": 720.0, "9": 720.0, "10": 720.0, "11": 720.0,
"12": 720.0, "13": 720.0, "14": 720.0, "15": 720.0, "16": 720.0, "17": 720.0,
"18": 720.0, "19": 720.0, "20": 720.0, "21": 720.0, "22": 720.0, "23": 720.0,
"24": 720.0, "25": 720.0, "26": 720.0, "27": 720.0, "28": 720.0, "29": 720.0,
"30": 720.0, "31": 720.0, "32": 720.0, "33": 720.0, "34": 720.0, "35": 720.0,
"36": 720.0, "37": 720.0, "38": 720.0, "39": 720.0, "40": 720.0, "41": 720.0,
"42": 720.0, "43": 720.0, "44": 720.0, "45": 720.0, "46": 720.0, "47": 720.0,
"48": 720.0, "49": 720.0
}
},
"text_to_video": {
"notes": "Single-video generation using the default prompt",
"expected_e2e_ms": 95616.59,
"expected_avg_denoise_ms": 1798.77,
"expected_median_denoise_ms": 1786.78,
"stages_ms": {
"InputValidationStage": 1.03,
"TextEncodingStage": 3450.0,
"ConditioningStage": 1.0,
"TimestepPreparationStage": 6.0,
"LatentPreparationStage": 15.0,
"DenoisingStage": 90100.0,
"DecodingStage": 3650.0
},
"denoise_step_ms": {
"0": 3500.0, "10": 1800.0, "20": 1800.0, "29": 1800.0, "39": 1800.0, "49": 1800.0
},
"frames_per_second": 0.51,
"total_frames": 49,
"avg_frame_time_ms": 1951.36
},
"image_to_video": {
"notes": "Image-to-Video generation baseline placeholder: TODO(bug)",
"expected_e2e_ms": 1000000000.0,
"expected_avg_denoise_ms": 1000000000.0,
"expected_median_denoise_ms": 1000000000.0,
"stages_ms": {},
"denoise_step_ms": {},
"frames_per_second": null,
"total_frames": null,
"avg_frame_time_ms": null
},
"text_image_to_video": {
"notes": "Text-and-Image-to-Video generation baseline for Wan2.2-TI2V-5B",
"expected_e2e_ms": 178300.0,
"expected_avg_denoise_ms": 3250.0,
"expected_median_denoise_ms": 3260.0,
"stages_ms": {
"InputValidationStage": 80.0,
"TextEncodingStage": 3000.0,
"ConditioningStage": 1.0,
"TimestepPreparationStage": 6.0,
"LatentPreparationStage": 30.0,
"DenoisingStage": 162900.0,
"DecodingStage": 13500.0
},
"denoise_step_ms": {
"0": 3700.0,
"10": 3300.0,
"20": 3300.0,
"29": 3300.0,
"39": 3300.0,
"49": 3300.0
},
"frames_per_second": null,
"total_frames": null,
"avg_frame_time_ms": null
}
}
}

762
python/sglang/multimodal_gen/test/server/test_server_performance.py
View File

@@ -1,34 +1,39 @@
# Server-based diffusion performance test:
# - Launches an sglang diffusion server via the CLI.
# - Issues an OpenAI-compatible Images API request.
# - Extracts all performance metrics from performance.log (no stdout parsing).
# - Verifies E2E, stage-level, and denoising-step latencies with configurable buffers.
"""
Config-driven diffusion performance test with pytest parametrization.
Adding a new model/scenario = adding one DiffusionCase entry in diffusion_config.py.
"""
from __future__ import annotations
import base64
import json
import os
import statistics
import subprocess
import tempfile
import time
from pathlib import Path
from typing import Any, Sequence
from typing import Any, Callable
import pytest
from openai import OpenAI
from sglang.multimodal_gen.runtime.utils.common import kill_process_tree
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
from sglang.multimodal_gen.test.server.conftest import _GLOBAL_PERF_RESULTS
from sglang.multimodal_gen.test.server.diffusion_config import (
BASELINE_CONFIG,
DIFFUSION_CASES,
DiffusionCase,
)
from sglang.multimodal_gen.test.server.diffusion_server import (
VALIDATOR_REGISTRY,
PerformanceValidator,
ServerContext,
ServerManager,
VideoPerformanceValidator,
WarmupRunner,
download_image_from_url,
)
from sglang.multimodal_gen.test.test_utils import (
get_dynamic_server_port,
is_jpeg,
is_png,
prepare_perf_log,
read_perf_records,
sample_step_indices,
validate_image,
validate_openai_video,
wait_for_perf_record,
wait_for_stage_metrics,
)
@@ -36,261 +41,74 @@ from sglang.multimodal_gen.test.test_utils import (
logger = init_logger(__name__)
_BASELINE_PATH = Path(__file__).with_name("perf_baselines.json")
with _BASELINE_PATH.open("r", encoding="utf-8") as _fh:
_BASELINE_CONFIG = json.load(_fh)
_SCENARIOS = _BASELINE_CONFIG["scenarios"]
_TEXT_SCENARIO = _SCENARIOS["text_to_image"]
_IMAGE_EDIT_SCENARIO = _SCENARIOS["image_edit"]
STEP_SAMPLE_FRACTIONS: Sequence[float] = tuple(
_BASELINE_CONFIG["sampling"]["step_fractions"]
)
_WARMUP_DEFAULTS = _BASELINE_CONFIG["sampling"].get("warmup_requests", {})
_DEFAULT_WARMUP_TEXT = int(_WARMUP_DEFAULTS.get("text", 1))
_DEFAULT_WARMUP_EDIT = int(_WARMUP_DEFAULTS.get("image_edit", 0))
_TOLERANCES = _BASELINE_CONFIG["tolerances"]
@pytest.fixture(params=DIFFUSION_CASES, ids=lambda c: c.id)
def case(request) -> DiffusionCase:
"""Provide a DiffusionCase for each test."""
return request.param
def _tolerance_from_env(var_name: str, default: float) -> float:
override = os.environ.get(var_name)
if override is not None:
return float(override)
return float(default)
E2E_TOLERANCE_RATIO = _tolerance_from_env("SGLANG_E2E_TOLERANCE", _TOLERANCES["e2e"])
STAGE_TOLERANCE_RATIO = _tolerance_from_env(
"SGLANG_STAGE_TIME_TOLERANCE", _TOLERANCES["stage"]
)
DENOISE_STEP_TOLERANCE_RATIO = _tolerance_from_env(
"SGLANG_DENOISE_STEP_TOLERANCE", _TOLERANCES["denoise_step"]
)
DENOISE_AGG_TOLERANCE_RATIO = _tolerance_from_env(
"SGLANG_DENOISE_AGG_TOLERANCE", _TOLERANCES["denoise_agg"]
)
def _decode_and_validate_image(b64_json: str) -> None:
image_bytes = base64.b64decode(b64_json)
assert is_png(image_bytes) or is_jpeg(
image_bytes
), "Warm-up image must be PNG or JPEG"
def _run_warmup_requests(cls, port: int) -> None:
warmup_text_requests = int(getattr(cls, "WARMUP_TEXT_REQUESTS", 1))
warmup_edit_requests = int(getattr(cls, "WARMUP_IMAGE_EDIT_REQUESTS", 0))
if warmup_text_requests <= 0 and warmup_edit_requests <= 0:
return
client = OpenAI(
api_key="sglang-anything",
base_url=f"http://localhost:{port}/v1",
)
prompt = getattr(cls, "PROMPT", "A colorful raccoon icon")
output_size = getattr(cls, "OUTPUT_SIZE", "1024x1024")
logger.info(
"[server-test] Running %s text warm-up(s) and %s edit warm-up(s)",
warmup_text_requests,
warmup_edit_requests,
)
for _ in range(warmup_text_requests):
result = client.images.generate(
model=getattr(cls, "MODEL_PATH"),
prompt=prompt,
n=1,
size=output_size,
response_format="b64_json",
)
_decode_and_validate_image(result.data[0].b64_json)
if warmup_edit_requests > 0:
edit_prompt = getattr(cls, "IMAGE_EDIT_PROMPT", None)
edit_path: Path | None = getattr(cls, "IMAGE_EDIT_PATH", None)
if not edit_prompt or not edit_path or not edit_path.exists():
logger.warning(
"[server-test] Skipping image-edit warm-up: prompt=%s path=%s exists=%s",
bool(edit_prompt),
edit_path,
edit_path.exists() if edit_path else False,
)
return
for _ in range(warmup_edit_requests):
with edit_path.open("rb") as fh:
result = client.images.edit(
model=getattr(cls, "MODEL_PATH"),
image=fh,
prompt=edit_prompt,
n=1,
size=output_size,
response_format="b64_json",
)
_decode_and_validate_image(result.data[0].b64_json)
@pytest.fixture(scope="class")
def diffusion_server(request):
cls = request.cls
log_dir, perf_log_path = prepare_perf_log(Path(__file__))
@pytest.fixture
def diffusion_server(case: DiffusionCase) -> ServerContext:
"""Start a diffusion server for a single case and tear it down afterwards."""
default_port = get_dynamic_server_port()
port = int(os.environ.get("SGLANG_TEST_SERVER_PORT", default_port))
port = getattr(cls, "SERVER_PORT", port)
model = getattr(cls, "MODEL_PATH")
wait_deadline = float(os.environ.get("SGLANG_TEST_WAIT_SECS", "1200"))
serve_extra_args = os.environ.get("SGLANG_TEST_SERVE_ARGS", "")
safe_model_name = model.replace("/", "_")
stdout_path = (
Path(tempfile.gettempdir()) / f"sgl_server_{port}_{safe_model_name}.log"
# start server
manager = ServerManager(
model=case.model_path,
port=port,
wait_deadline=float(os.environ.get("SGLANG_TEST_WAIT_SECS", "1200")),
extra_args=os.environ.get("SGLANG_TEST_SERVE_ARGS", ""),
)
stdout_path.unlink(missing_ok=True)
base_command = [
"sglang",
"serve",
"--model-path",
model,
"--port",
str(port),
"--log-level=debug",
]
if serve_extra_args.strip():
base_command += serve_extra_args.strip().split()
env = os.environ.copy()
env["SGL_DIFFUSION_STAGE_LOGGING"] = "1"
env["SGLANG_PERF_LOG_DIR"] = log_dir.as_posix()
stdout_fh = stdout_path.open("w", encoding="utf-8", buffering=1)
process = subprocess.Popen(
base_command,
stdout=stdout_fh,
stderr=subprocess.STDOUT,
text=True,
bufsize=1,
env=env,
)
logger.info(
"[server-test] Starting diffusion server pid=%s, model=%s, log=%s",
process.pid,
model,
stdout_path.as_posix(),
)
start = time.time()
server_ready_message = "Application startup complete."
server_ready = False
while time.time() - start < wait_deadline:
if process.poll() is not None:
tail = ""
try:
tail = "\n".join(
stdout_path.read_text(
encoding="utf-8", errors="ignore"
).splitlines()[-200:]
)
except Exception:
pass
raise RuntimeError(
f"Server exited early (code {process.returncode}). Last logs:\n{tail}"
)
if stdout_path.exists():
try:
log_content = stdout_path.read_text(encoding="utf-8", errors="ignore")
if server_ready_message in log_content:
logger.info("[server-test] Server is fully loaded and ready.")
server_ready = True
break
except Exception as e:
logger.debug("Could not read server log file yet: %s", e)
ctx = manager.start()
if case.startup_grace_seconds > 0:
logger.info(
"[server-test] Waiting for server to initialize... elapsed=%ss",
int(time.time() - start),
"[server-test] Waiting %.1fs for %s to settle",
case.startup_grace_seconds,
case.id,
)
time.sleep(5)
if not server_ready:
tail = ""
try:
tail = "\n".join(
stdout_path.read_text(encoding="utf-8", errors="ignore").splitlines()[
-200:
]
)
except Exception:
pass
raise TimeoutError(
f"Server did not become ready within {wait_deadline}s. Last logs:\n{tail}"
)
ctx = {
"port": port,
"stdout_file": stdout_path,
"process": process,
"model": model,
"fh": stdout_fh,
"perf_log_path": perf_log_path,
"log_dir": log_dir,
}
request.cls.server_ctx = ctx
request.cls.perf_log_path = perf_log_path
grace = float(getattr(cls, "STARTUP_GRACE_SECONDS", 0.0) or 0.0)
if grace > 0:
logger.info(
"[server-test] Waiting %.1fs before warm-ups to let model settle", grace
)
time.sleep(grace)
time.sleep(case.startup_grace_seconds)
try:
_run_warmup_requests(cls, port)
warmup = WarmupRunner(
port=ctx.port,
model=case.model_path,
prompt=case.prompt or "A colorful raccoon icon",
output_size=case.output_size,
)
warmup.run_text_warmups(case.warmup_text)
if case.warmup_edit > 0 and case.image_edit_prompt and case.image_edit_path:
# Handle URL or local path
image_path = case.image_edit_path
if case.is_image_url():
image_path = download_image_from_url(str(case.image_edit_path))
else:
image_path = Path(case.image_edit_path)
warmup.run_edit_warmups(
count=case.warmup_edit,
edit_prompt=case.image_edit_prompt,
image_path=image_path,
)
except Exception as exc:
logger.error("Warm-up requests failed: %s", exc)
kill_process_tree(process.pid)
logger.error("Warm-up failed for %s: %s", case.id, exc)
ctx.cleanup()
raise
yield ctx
try:
kill_process_tree(process.pid)
except Exception:
pass
try:
stdout_fh.flush()
stdout_fh.close()
except Exception:
pass
yield ctx
finally:
ctx.cleanup()
@pytest.mark.usefixtures("diffusion_server")
class DiffusionPerfTestBase:
MODEL_PATH: str
# SERVER_PORT = int(os.environ.get("SGLANG_TEST_SERVER_PORT", "30100"))
PROMPT = "A Logo With Bold Large Text: SGL Diffusion"
IMAGE_EDIT_PROMPT: str | None = None
IMAGE_EDIT_PATH = Path(__file__).resolve().parents[1] / "test_files" / "girl.jpg"
OUTPUT_SIZE = "1024x1024"
WARMUP_TEXT_REQUESTS = _DEFAULT_WARMUP_TEXT
WARMUP_IMAGE_EDIT_REQUESTS = _DEFAULT_WARMUP_EDIT
STARTUP_GRACE_SECONDS = 0.0
class TestDiffusionPerformance:
"""Performance tests for all diffusion models/scenarios.
STAGE_EXPECTATIONS: dict
STEP_EXPECTATIONS: dict
EXPECTED_E2E_MS: float
EXPECTED_AVG_DENOISE_MS: float
EXPECTED_MEDIAN_DENOISE_MS: float
This single test class runs against all cases defined in DIFFUSION_CASES.
Each case gets its own server instance via the parametrized fixture.
"""
_perf_results: list[dict[str, Any]] = []
@@ -304,187 +122,329 @@ class DiffusionPerfTestBase:
result["class_name"] = cls.__name__
_GLOBAL_PERF_RESULTS.append(result)
def _client(self) -> OpenAI:
def _client(self, ctx: ServerContext) -> OpenAI:
"""Get OpenAI client for the server."""
return OpenAI(
api_key="sglang-anything",
base_url=f"http://localhost:{self.server_ctx['port']}/v1",
base_url=f"http://localhost:{ctx.port}/v1",
)
def _perf_log_path(self) -> Path:
return self.server_ctx["perf_log_path"]
def _record_result(self, test_name: str, summary: dict[str, Any]) -> None:
if not summary:
return
entry = {"test_name": test_name, **summary}
self.__class__._perf_results.append(entry)
def _run_and_collect_records(self, generate_fn) -> tuple[dict, dict]:
log_path = self._perf_log_path()
def _run_and_collect(
self,
ctx: ServerContext,
case: DiffusionCase,
generate_fn: Callable[[], None],
) -> tuple[dict, dict]:
"""Run generation and collect performance records."""
log_path = ctx.perf_log_path
prev_len = len(read_perf_records(log_path))
generate_fn()
perf_record, _ = wait_for_perf_record(
"total_inference_time",
prev_len,
log_path,
)
scenario = BASELINE_CONFIG.scenarios[case.scenario_name]
stage_metrics, _ = wait_for_stage_metrics(
perf_record.get("request_id", ""),
prev_len,
len(self.STAGE_EXPECTATIONS),
len(scenario.stages_ms),
log_path,
)
return perf_record, stage_metrics
def _generate_image(self):
client = self._client()
result = client.images.generate(
model=self.MODEL_PATH,
prompt=self.PROMPT,
n=1,
size=self.OUTPUT_SIZE,
response_format="b64_json",
)
image_bytes = base64.b64decode(result.data[0].b64_json)
assert is_png(image_bytes) or is_jpeg(
image_bytes
), "Generated image must be PNG or JPEG"
def _generate_for_case(
self,
ctx: ServerContext,
case: DiffusionCase,
) -> Callable[[], None]:
"""Return appropriate generation function for the case."""
client = self._client(ctx)
def _generate_image_edit(self):
if not self.IMAGE_EDIT_PROMPT:
pytest.skip("Image edit prompt not configured")
if not self.IMAGE_EDIT_PATH.exists():
pytest.skip(f"Image edit file missing: {self.IMAGE_EDIT_PATH}")
client = self._client()
with self.IMAGE_EDIT_PATH.open("rb") as fh:
result = client.images.edit(
model=self.MODEL_PATH,
image=fh,
prompt=self.IMAGE_EDIT_PROMPT,
def _create_and_download_video(
*,
model: str,
size: str,
prompt: str | None = None,
seconds: int | None = None,
input_reference: Any | None = None,
) -> bytes:
"""
Create a video job via /v1/videos, poll until completion,
then download the binary content and validate it.
"""
create_kwargs: dict[str, Any] = {
"model": model,
"size": size,
}
if prompt is not None:
create_kwargs["prompt"] = prompt
if seconds is not None:
create_kwargs["seconds"] = seconds
if input_reference is not None:
create_kwargs["input_reference"] = input_reference # triggers multipart
# create video job
job = client.videos.create(**create_kwargs) # type: ignore[attr-defined]
video_id = job.id
deadline = time.time() + 600
while True:
page = client.videos.list() # type: ignore[attr-defined]
item = next((v for v in page.data if v.id == video_id), None)
if item and getattr(item, "status", None) == "completed":
break
if time.time() > deadline:
pytest.fail(
f"{case.id}: video job {video_id} did not complete in time"
)
time.sleep(5)
# download video
resp = client.videos.download_content(video_id=video_id) # type: ignore[attr-defined]
content = resp.read()
validate_openai_video(content)
return content
# for all tests, seconds = case.seconds or fallback 4 seconds
video_seconds = case.seconds or 4
# -------------------------
# IMAGE MODE
# -------------------------
def generate_image():
"""T2I: Text to Image generation."""
if not case.prompt:
pytest.skip(f"{case.id}: no text prompt configured")
result = client.images.generate(
model=case.model_path,
prompt=case.prompt,
n=1,
size=self.OUTPUT_SIZE,
size=case.output_size,
response_format="b64_json",
)
image_bytes = base64.b64decode(result.data[0].b64_json)
assert is_png(image_bytes) or is_jpeg(
image_bytes
), "Edited image must be PNG or JPEG"
validate_image(result.data[0].b64_json)
def _assert_metrics(self, perf_record: dict, stage_metrics: dict):
e2e_ms = float(perf_record.get("total_duration_ms", 0.0))
assert e2e_ms > 0, "E2E duration missing from perf log"
e2e_upper = self.EXPECTED_E2E_MS * (1 + E2E_TOLERANCE_RATIO)
assert (
e2e_ms <= e2e_upper
), f"E2E time {e2e_ms:.2f}ms exceeds allowed {e2e_upper:.2f}ms"
def generate_image_edit():
"""TI2I: Text + Image ? Image edit."""
if not case.image_edit_prompt or not case.image_edit_path:
pytest.skip(f"{case.id}: no edit config")
steps = [
step
for step in perf_record.get("steps", []) or []
if step.get("name") == "denoising_step_guided" and "duration_ms" in step
]
assert steps, "Denoising step timings missing from perf log"
# Handle URL or local path
if case.is_image_url():
image_path = download_image_from_url(str(case.image_edit_path))
else:
image_path = Path(case.image_edit_path)
if not image_path.exists():
pytest.skip(f"{case.id}: file missing: {image_path}")
durations = [float(step["duration_ms"]) for step in steps]
avg_duration = sum(durations) / len(durations)
median_duration = statistics.median(durations)
with image_path.open("rb") as fh:
result = client.images.edit(
model=case.model_path,
image=fh,
prompt=case.image_edit_prompt,
n=1,
size=case.output_size,
response_format="b64_json",
)
validate_image(result.data[0].b64_json)
avg_upper = self.EXPECTED_AVG_DENOISE_MS * (1 + DENOISE_AGG_TOLERANCE_RATIO)
med_upper = self.EXPECTED_MEDIAN_DENOISE_MS * (1 + DENOISE_AGG_TOLERANCE_RATIO)
assert (
avg_duration <= avg_upper
), f"Avg denoise {avg_duration:.2f}ms exceeds {avg_upper:.2f}ms"
assert (
median_duration <= med_upper
), f"Median denoise {median_duration:.2f}ms exceeds {med_upper:.2f}ms"
# -------------------------
# VIDEO MODE
# -------------------------
avg_per_step = {
int(step.get("index")): float(step["duration_ms"])
for step in steps
if step.get("index") is not None
def generate_video():
"""T2V: Text ? Video."""
if not case.prompt:
pytest.skip(f"{case.id}: no text prompt configured")
_create_and_download_video(
model=case.model_path,
prompt=case.prompt,
size=case.output_size,
seconds=video_seconds,
)
def generate_image_to_video():
"""I2V: Image ? Video (optional prompt)."""
if not case.image_edit_path:
pytest.skip(f"{case.id}: no input image configured")
# Handle URL or local path
if case.is_image_url():
image_path = download_image_from_url(str(case.image_edit_path))
else:
image_path = Path(case.image_edit_path)
if not image_path.exists():
pytest.skip(f"{case.id}: file missing: {image_path}")
with image_path.open("rb") as fh:
_create_and_download_video(
model=case.model_path,
prompt=case.image_edit_prompt,
size=case.output_size,
seconds=video_seconds,
input_reference=fh,
)
def generate_text_image_to_video():
"""TI2V: Text + Image ? Video."""
if not case.image_edit_prompt or not case.image_edit_path:
pytest.skip(f"{case.id}: no edit config")
# Handle URL or local path
if case.is_image_url():
image_path = download_image_from_url(str(case.image_edit_path))
else:
image_path = Path(case.image_edit_path)
if not image_path.exists():
pytest.skip(f"{case.id}: file missing: {image_path}")
with image_path.open("rb") as fh:
_create_and_download_video(
model=case.model_path,
prompt=case.image_edit_prompt,
size=case.output_size,
seconds=video_seconds,
input_reference=fh,
)
if case.modality == "video":
if case.image_edit_path and case.image_edit_prompt:
return generate_text_image_to_video
elif case.image_edit_path:
return generate_image_to_video
else:
return generate_video
# Image modality
if case.image_edit_prompt and case.image_edit_path:
return generate_image_edit
return generate_image
def _validate_and_record(
self,
case: DiffusionCase,
perf_record: dict,
stage_metrics: dict,
) -> None:
"""Validate metrics and record results."""
scenario = BASELINE_CONFIG.scenarios[case.scenario_name]
validator_name = case.custom_validator or "default"
validator_class = VALIDATOR_REGISTRY.get(validator_name, PerformanceValidator)
validator = validator_class(
scenario=scenario,
tolerances=BASELINE_CONFIG.tolerances,
step_fractions=BASELINE_CONFIG.step_fractions,
)
if isinstance(validator, VideoPerformanceValidator):
summary = validator.validate(perf_record, stage_metrics, case.num_frames)
else:
summary = validator.validate(perf_record, stage_metrics)
if case.modality == "video" and summary.frames_per_second:
logger.info(
"[Perf] %s: E2E %.2f ms; Avg %.2f ms; FPS %.2f; Frames %d",
case.id,
summary.e2e_ms,
summary.avg_denoise_ms,
summary.frames_per_second,
summary.total_frames or 0,
)
else:
logger.info(
"[Perf] %s: E2E %.2f ms; Avg %.2f ms; Median %.2f ms",
case.id,
summary.e2e_ms,
summary.avg_denoise_ms,
summary.median_denoise_ms,
)
result = {
"test_name": case.id,
"modality": case.modality,
"e2e_ms": summary.e2e_ms,
"avg_denoise_ms": summary.avg_denoise_ms,
"median_denoise_ms": summary.median_denoise_ms,
"stage_metrics": summary.stage_metrics,
"sampled_steps": summary.sampled_steps,
}
sample_indices = sample_step_indices(avg_per_step, STEP_SAMPLE_FRACTIONS)
sampled_steps = {idx: avg_per_step[idx] for idx in sample_indices}
for idx in sample_indices:
expected = self.STEP_EXPECTATIONS.get(idx)
if expected is None:
continue
actual = avg_per_step[idx]
upper_bound = expected * (1 + DENOISE_STEP_TOLERANCE_RATIO)
assert (
actual <= upper_bound
), f"Denoise step {idx} took {actual:.2f}ms > allowed {upper_bound:.2f}ms"
assert stage_metrics, "Stage metrics missing from performance log"
for stage, expected in self.STAGE_EXPECTATIONS.items():
actual = stage_metrics.get(stage)
assert actual is not None, f"Stage {stage} timing missing"
upper_bound = expected * (1 + STAGE_TOLERANCE_RATIO)
assert (
actual <= upper_bound
), f"Stage {stage} took {actual:.2f}ms > allowed {upper_bound:.2f}ms"
# video-specific metrics
if summary.frames_per_second:
result.update(
{
"frames_per_second": summary.frames_per_second,
"total_frames": summary.total_frames,
"avg_frame_time_ms": summary.avg_frame_time_ms,
}
)
# Log to pytest console during the run for immediate feedback
self.__class__._perf_results.append(result)
logger.info("[BASELINE] %s expected_e2e_ms = %.2f", case.id, summary.e2e_ms)
logger.info(
"[Perf] %s/%s: E2E %.2f ms; Avg denoise %.2f ms; Median %.2f ms",
self.__class__.__name__,
perf_record.get("test_name", "test"),
e2e_ms,
avg_duration,
median_duration,
"[BASELINE] %s expected_avg_denoise_ms = %.2f",
case.id,
summary.avg_denoise_ms,
)
logger.info(
"[BASELINE] %s expected_median_denoise_ms = %.2f",
case.id,
summary.median_denoise_ms,
)
logger.info("[BASELINE] %s stages_ms = %r", case.id, summary.stage_metrics)
logger.info(
"[BASELINE] %s denoise_step_ms = %r", case.id, summary.sampled_steps
)
return {
"e2e_ms": e2e_ms,
"avg_denoise_ms": avg_duration,
"median_denoise_ms": median_duration,
"stage_metrics": stage_metrics,
"sampled_steps": sampled_steps,
}
# Only log video-specific metrics when they exist
if summary.frames_per_second is not None:
logger.info(
"[BASELINE] %s frames_per_second = %.2f",
case.id,
summary.frames_per_second,
)
if summary.total_frames is not None:
logger.info(
"[BASELINE] %s total_frames = %d", case.id, summary.total_frames
)
if summary.avg_frame_time_ms is not None:
logger.info(
"[BASELINE] %s avg_frame_time_ms = %.2f",
case.id,
summary.avg_frame_time_ms,
)
def test_diffusion_perf(
self,
case: DiffusionCase,
diffusion_server: ServerContext,
):
"""Single parametrized test that runs for all cases.
class TestQwenImageGeneration(DiffusionPerfTestBase):
"""Performance tests for the Qwen/Qwen-image model."""
MODEL_PATH = "Qwen/Qwen-Image"
STARTUP_GRACE_SECONDS = 30.0
WARMUP_IMAGE_EDIT_REQUESTS = 0
STAGE_EXPECTATIONS = _TEXT_SCENARIO["stages_ms"]
STEP_EXPECTATIONS = {
int(k): v for k, v in _TEXT_SCENARIO["denoise_step_ms"].items()
}
EXPECTED_E2E_MS = float(_TEXT_SCENARIO["expected_e2e_ms"])
EXPECTED_AVG_DENOISE_MS = float(_TEXT_SCENARIO["expected_avg_denoise_ms"])
EXPECTED_MEDIAN_DENOISE_MS = float(_TEXT_SCENARIO["expected_median_denoise_ms"])
def test_text_to_image_performance(self):
perf_record, stage_metrics = self._run_and_collect_records(self._generate_image)
summary = self._assert_metrics(perf_record, stage_metrics)
self._record_result("text_to_image", summary)
class TestQwenImageEdit(DiffusionPerfTestBase):
"""Performance tests for the Qwen/Qwen-Image-Edit model."""
MODEL_PATH = "Qwen/Qwen-Image-Edit"
IMAGE_EDIT_PROMPT = "Convert 2D style to 3D style"
OUTPUT_SIZE = "1024x1536"
STARTUP_GRACE_SECONDS = 30.0
WARMUP_TEXT_REQUESTS = 0
WARMUP_IMAGE_EDIT_REQUESTS = 1
STAGE_EXPECTATIONS = _IMAGE_EDIT_SCENARIO["stages_ms"]
STEP_EXPECTATIONS = {
int(k): v for k, v in _IMAGE_EDIT_SCENARIO["denoise_step_ms"].items()
}
EXPECTED_E2E_MS = float(_IMAGE_EDIT_SCENARIO["expected_e2e_ms"])
EXPECTED_AVG_DENOISE_MS = float(_IMAGE_EDIT_SCENARIO["expected_avg_denoise_ms"])
EXPECTED_MEDIAN_DENOISE_MS = float(
_IMAGE_EDIT_SCENARIO["expected_median_denoise_ms"]
)
def test_image_edit_performance(self):
perf_record, stage_metrics = self._run_and_collect_records(
self._generate_image_edit
Pytest will execute this test once per case in DIFFUSION_CASES,
with test IDs like:
- test_diffusion_perf[qwen_image_text]
- test_diffusion_perf[qwen_image_edit]
- etc.
"""
generate_fn = self._generate_for_case(diffusion_server, case)
perf_record, stage_metrics = self._run_and_collect(
diffusion_server,
case,
generate_fn,
)
summary = self._assert_metrics(perf_record, stage_metrics)
self._record_result("image_edit", summary)
self._validate_and_record(case, perf_record, stage_metrics)

858
python/sglang/multimodal_gen/test/test_utils.py
View File

@@ -1,415 +1,443 @@
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
import dataclasses
import json
import os
import shlex
import socket
import subprocess
import sys
import time
import unittest
from pathlib import Path
from typing import Optional, Sequence
from PIL import Image
from sglang.multimodal_gen.configs.sample.base import DataType
from sglang.multimodal_gen.runtime.utils.common import get_bool_env_var
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
logger = init_logger(__name__)
def run_command(command) -> Optional[float]:
"""Runs a command and returns the execution time and status."""
print(f"Running command: {shlex.join(command)}")
duration = None
with subprocess.Popen(
command,
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT,
text=True,
encoding="utf-8",
) as process:
for line in process.stdout:
sys.stdout.write(line)
if "Pixel data generated" in line:
words = line.split(" ")
duration = float(words[-2])
if process.returncode == 0:
return duration
else:
print(f"Command failed with exit code {process.returncode}")
return None
def probe_port(host="127.0.0.1", port=30010, timeout=2.0) -> bool:
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
s.settimeout(timeout)
try:
s.connect((host, port))
return True
except OSError:
return False
def is_in_ci() -> bool:
return get_bool_env_var("SGLANG_IS_IN_CI")
def get_dynamic_server_port() -> int:
cuda_devices = os.environ.get("CUDA_VISIBLE_DEVICES", "0")
if not cuda_devices:
cuda_devices = "0"
try:
first_device_id = int(cuda_devices.split(",")[0].strip()[0])
except (ValueError, IndexError):
first_device_id = 0
if is_in_ci():
base_port = 10000 + first_device_id * 2000
else:
base_port = 20000 + first_device_id * 1000
return base_port + 1000
def is_mp4(data):
idx = data.find(b"ftyp")
return 0 <= idx <= 32
def is_jpeg(data: bytes) -> bool:
# JPEG files start with: FF D8 FF
return data.startswith(b"\xff\xd8\xff")
def is_png(data):
# PNG files start with: 89 50 4E 47 0D 0A 1A 0A
return data.startswith(b"\x89PNG\r\n\x1a\n")
def wait_for_port(host="127.0.0.1", port=30010, deadline=300.0, interval=0.5):
end = time.time() + deadline
last_err = None
while time.time() < end:
if probe_port(host, port, timeout=interval):
return True
time.sleep(interval)
raise TimeoutError(f"Port {host}:{port} not ready. Last error: {last_err}")
def check_image_size(ut, image, width, height):
# check image size
ut.assertEqual(image.size, (width, height))
def get_perf_log_dir(start_file: Path) -> Path:
"""Mirror runtime/utils/performance_logger.py behaviour for locating logs."""
this_file = start_file.resolve()
root_logs = this_file.parents[3] / "logs"
fallback = this_file.parents[2] / "logs"
return root_logs if root_logs.exists() or not fallback.exists() else fallback
def _ensure_log_path(log_dir: Path) -> Path:
log_dir.mkdir(parents=True, exist_ok=True)
return log_dir / "performance.log"
def clear_perf_log(log_dir: Path) -> Path:
"""Delete the perf log file so tests can watch for fresh entries."""
log_path = _ensure_log_path(log_dir)
if log_path.exists():
log_path.unlink()
logger.info("[server-test] Monitoring perf log at %s", log_path.as_posix())
return log_path
def prepare_perf_log(start_file: Path) -> tuple[Path, Path]:
"""Convenience helper to resolve and clear the perf log in one call."""
log_dir = get_perf_log_dir(start_file)
log_path = clear_perf_log(log_dir)
return log_dir, log_path
def read_perf_records(log_path: Path) -> list[dict]:
if not log_path.exists():
return []
records: list[dict] = []
with log_path.open("r", encoding="utf-8") as fh:
for line in fh:
line = line.strip()
if not line:
continue
try:
records.append(json.loads(line))
except json.JSONDecodeError:
continue
return records
def wait_for_perf_record(
tag: str,
prev_len: int,
log_path: Path,
timeout: float = 120.0,
) -> tuple[dict, int]:
deadline = time.time() + timeout
while time.time() < deadline:
records = read_perf_records(log_path)
if len(records) > prev_len:
for rec in records[prev_len:]:
if rec.get("tag") == tag:
return rec, len(records)
time.sleep(0.5)
raise AssertionError(
f"Timeout waiting for perf log entry '{tag}' (start_len={prev_len})"
)
def wait_for_stage_metrics(
request_id: str,
prev_len: int,
expected_count: int,
log_path: Path,
timeout: float = 120.0,
) -> tuple[dict[str, float], int]:
deadline = time.time() + timeout
metrics: dict[str, float] = {}
while time.time() < deadline:
records = read_perf_records(log_path)
for rec in records[prev_len:]:
if (
rec.get("tag") == "pipeline_stage_metric"
and rec.get("request_id") == request_id
):
stage = rec.get("stage")
duration = rec.get("duration_ms")
if stage is not None and duration is not None:
metrics[str(stage)] = float(duration)
if len(metrics) >= expected_count:
return metrics, len(records)
time.sleep(0.5)
raise AssertionError(
f"Timeout waiting for stage metrics for request {request_id} "
f"(collected={len(metrics)} expected={expected_count})"
)
def sample_step_indices(
step_map: dict[int, float], fractions: Sequence[float]
) -> list[int]:
if not step_map:
return []
max_idx = max(step_map.keys())
indices = set()
for fraction in fractions:
idx = min(max_idx, max(0, int(round(fraction * max_idx))))
if idx in step_map:
indices.add(idx)
return sorted(indices)
@dataclasses.dataclass
class TestResult:
name: str
key: str
duration: Optional[float]
succeed: bool
@property
def duration_str(self):
return f"{self.duration:.4f}" if self.duration else "NA"
class TestCLIBase(unittest.TestCase):
model_path: str = None
extra_args = []
data_type: DataType = None
# tested on h100
thresholds = {}
width: int = 720
height: int = 720
output_path: str = "test_outputs"
base_command = [
"sglang",
"generate",
"--text-encoder-cpu-offload",
"--pin-cpu-memory",
"--prompt",
"A curious raccoon",
"--save-output",
"--log-level=debug",
f"--width={width}",
f"--height={height}",
f"--output-path={output_path}",
]
results = []
@classmethod
def setUpClass(cls):
cls.results = []
def _run_command(self, name: str, model_path: str, test_key: str = "", args=[]):
command = (
self.base_command
+ [f"--model-path={model_path}"]
+ shlex.split(args or "")
+ ["--output-file-name", f"{name}"]
+ self.extra_args
)
duration = run_command(command)
status = "Success" if duration else "Failed"
succeed = duration is not None
duration = float(duration) if succeed else None
self.results.append(TestResult(name, test_key, duration, succeed))
return name, duration, status
class TestGenerateBase(TestCLIBase):
model_path: str = None
extra_args = []
data_type: DataType = None
# tested on h100
thresholds = {}
width: int = 720
height: int = 720
output_path: str = "test_outputs"
image_path: str | None = None
prompt: str | None = "A curious raccoon"
base_command = [
"sglang",
"generate",
# "--text-encoder-cpu-offload",
# "--pin-cpu-memory",
f"--prompt",
f"{prompt}",
"--save-output",
"--log-level=debug",
f"--width={width}",
f"--height={height}",
f"--output-path={output_path}",
]
results: list[TestResult] = []
@classmethod
def setUpClass(cls):
cls.results = []
@classmethod
def tearDownClass(cls):
# Print markdown table
print("\n## Test Results\n")
print("| Test Case | Duration | Status |")
print("|--------------------------------|----------|---------|")
test_keys = ["test_single_gpu", "test_cfg_parallel", "test_usp", "test_mixed"]
test_key_to_order = {
test_key: order for order, test_key in enumerate(test_keys)
}
ordered_results: list[TestResult] = [None] * len(test_keys)
for result in cls.results:
order = test_key_to_order[result.key]
ordered_results[order] = result
for result in ordered_results:
if not result:
continue
status = (
"Succeed"
if (
result.succeed
and float(result.duration) <= float(cls.thresholds[result.key])
)
else "Failed"
)
print(f"| {result.name:<30} | {result.duration_str:<8} | {status:<7} |")
print()
durations = [result.duration_str for result in cls.results]
print(" | ".join([""] + durations + [""]))
def _run_test(self, name: str, args, model_path: str, test_key: str):
time_threshold = self.thresholds[test_key]
name, duration, status = self._run_command(
name, args=args, model_path=model_path, test_key=test_key
)
self.verify(status, name, duration, time_threshold)
def verify(self, status, name, duration, time_threshold):
print("-" * 80)
print("\n" * 3)
# test task status
self.assertEqual(status, "Success", f"{name} command failed")
self.assertIsNotNone(duration, f"Could not parse duration for {name}")
self.assertLessEqual(
duration,
time_threshold,
f"{name} failed with {duration:.4f}s > {time_threshold}s",
)
# test output file
path = os.path.join(
self.output_path, f"{name}.{self.data_type.get_default_extension()}"
)
self.assertTrue(os.path.exists(path), f"Output file not exist for {path}")
if self.data_type == DataType.IMAGE:
with Image.open(path) as image:
check_image_size(self, image, self.width, self.height)
logger.info(f"{name} passed in {duration:.4f}s (threshold: {time_threshold}s)")
def model_name(self):
return self.model_path.split("/")[-1]
def test_single_gpu(self):
"""single gpu"""
self._run_test(
name=f"{self.model_name()}_single_gpu",
args=None,
model_path=self.model_path,
test_key="test_single_gpu",
)
def test_cfg_parallel(self):
"""cfg parallel"""
if self.data_type == DataType.IMAGE:
return
self._run_test(
name=f"{self.model_name()}_cfg_parallel",
args="--num-gpus 2 --enable-cfg-parallel",
model_path=self.model_path,
test_key="test_cfg_parallel",
)
def test_usp(self):
"""usp"""
if self.data_type == DataType.IMAGE:
return
self._run_test(
name=f"{self.model_name()}_usp",
args="--num-gpus 4 --ulysses-degree=2 --ring-degree=2",
model_path=self.model_path,
test_key="test_usp",
)
def test_mixed(self):
"""mixed"""
if self.data_type == DataType.IMAGE:
return
self._run_test(
name=f"{self.model_name()}_mixed",
args="--num-gpus 4 --ulysses-degree=2 --ring-degree=1 --enable-cfg-parallel",
model_path=self.model_path,
test_key="test_mixed",
)
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
import base64
import dataclasses
import json
import os
import shlex
import socket
import subprocess
import sys
import time
import unittest
from pathlib import Path
from typing import Optional, Sequence
from PIL import Image
from sglang.multimodal_gen.configs.sample.base import DataType
from sglang.multimodal_gen.runtime.utils.common import get_bool_env_var
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
logger = init_logger(__name__)
def run_command(command) -> Optional[float]:
"""Runs a command and returns the execution time and status."""
print(f"Running command: {shlex.join(command)}")
duration = None
with subprocess.Popen(
command,
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT,
text=True,
encoding="utf-8",
) as process:
for line in process.stdout:
sys.stdout.write(line)
if "Pixel data generated" in line:
words = line.split(" ")
duration = float(words[-2])
if process.returncode == 0:
return duration
else:
print(f"Command failed with exit code {process.returncode}")
return None
def probe_port(host="127.0.0.1", port=30010, timeout=2.0) -> bool:
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
s.settimeout(timeout)
try:
s.connect((host, port))
return True
except OSError:
return False
def is_in_ci() -> bool:
return get_bool_env_var("SGLANG_IS_IN_CI")
def get_dynamic_server_port() -> int:
cuda_devices = os.environ.get("CUDA_VISIBLE_DEVICES", "0")
if not cuda_devices:
cuda_devices = "0"
try:
first_device_id = int(cuda_devices.split(",")[0].strip()[0])
except (ValueError, IndexError):
first_device_id = 0
if is_in_ci():
base_port = 10000 + first_device_id * 2000
else:
base_port = 20000 + first_device_id * 1000
return base_port + 1000
def is_mp4(data):
idx = data.find(b"ftyp")
return 0 <= idx <= 32
def is_jpeg(data: bytes) -> bool:
# JPEG files start with: FF D8 FF
return data.startswith(b"\xff\xd8\xff")
def is_png(data):
# PNG files start with: 89 50 4E 47 0D 0A 1A 0A
return data.startswith(b"\x89PNG\r\n\x1a\n")
def wait_for_port(host="127.0.0.1", port=30010, deadline=300.0, interval=0.5):
end = time.time() + deadline
last_err = None
while time.time() < end:
if probe_port(host, port, timeout=interval):
return True
time.sleep(interval)
raise TimeoutError(f"Port {host}:{port} not ready. Last error: {last_err}")
def check_image_size(ut, image, width, height):
# check image size
ut.assertEqual(image.size, (width, height))
def get_perf_log_dir(start_file: Path) -> Path:
"""Mirror runtime/utils/performance_logger.py behaviour for locating logs."""
this_file = start_file.resolve()
root_logs = this_file.parents[3] / "logs"
fallback = this_file.parents[2] / "logs"
return root_logs if root_logs.exists() or not fallback.exists() else fallback
def _ensure_log_path(log_dir: Path) -> Path:
log_dir.mkdir(parents=True, exist_ok=True)
return log_dir / "performance.log"
def clear_perf_log(log_dir: Path) -> Path:
"""Delete the perf log file so tests can watch for fresh entries."""
log_path = _ensure_log_path(log_dir)
if log_path.exists():
log_path.unlink()
logger.info("[server-test] Monitoring perf log at %s", log_path.as_posix())
return log_path
def prepare_perf_log(start_file: Path) -> tuple[Path, Path]:
"""Convenience helper to resolve and clear the perf log in one call."""
log_dir = get_perf_log_dir(start_file)
log_path = clear_perf_log(log_dir)
return log_dir, log_path
def read_perf_records(log_path: Path) -> list[dict]:
if not log_path.exists():
return []
records: list[dict] = []
with log_path.open("r", encoding="utf-8") as fh:
for line in fh:
line = line.strip()
if not line:
continue
try:
records.append(json.loads(line))
except json.JSONDecodeError:
continue
return records
def wait_for_perf_record(
tag: str,
prev_len: int,
log_path: Path,
timeout: float = 120.0,
) -> tuple[dict, int]:
deadline = time.time() + timeout
while time.time() < deadline:
records = read_perf_records(log_path)
if len(records) > prev_len:
for rec in records[prev_len:]:
if rec.get("tag") == tag:
return rec, len(records)
time.sleep(0.5)
raise AssertionError(
f"Timeout waiting for perf log entry '{tag}' (start_len={prev_len})"
)
def wait_for_stage_metrics(
request_id: str,
prev_len: int,
expected_count: int,
log_path: Path,
timeout: float = 120.0,
) -> tuple[dict[str, float], int]:
deadline = time.time() + timeout
metrics: dict[str, float] = {}
while time.time() < deadline:
records = read_perf_records(log_path)
for rec in records[prev_len:]:
if (
rec.get("tag") == "pipeline_stage_metric"
and rec.get("request_id") == request_id
):
stage = rec.get("stage")
duration = rec.get("duration_ms")
if stage is not None and duration is not None:
metrics[str(stage)] = float(duration)
if len(metrics) >= expected_count:
return metrics, len(records)
time.sleep(0.5)
raise AssertionError(
f"Timeout waiting for stage metrics for request {request_id} "
f"(collected={len(metrics)} expected={expected_count})"
)
def sample_step_indices(
step_map: dict[int, float], fractions: Sequence[float]
) -> list[int]:
if not step_map:
return []
max_idx = max(step_map.keys())
indices = set()
for fraction in fractions:
idx = min(max_idx, max(0, int(round(fraction * max_idx))))
if idx in step_map:
indices.add(idx)
return sorted(indices)
def validate_image(b64_json: str) -> None:
"""Decode and validate that image is PNG or JPEG."""
image_bytes = base64.b64decode(b64_json)
assert is_png(image_bytes) or is_jpeg(image_bytes), "Image must be PNG or JPEG"
def validate_video(b64_json: str) -> None:
"""Decode and validate that video is a valid format."""
video_bytes = base64.b64decode(b64_json)
is_mp4 = (
video_bytes[:4] == b"\x00\x00\x00\x18" or video_bytes[:4] == b"\x00\x00\x00\x1c"
)
is_webm = video_bytes[:4] == b"\x1a\x45\xdf\xa3"
assert is_mp4 or is_webm, "Video must be MP4 or WebM"
def validate_openai_video(video_bytes: bytes) -> None:
"""Validate that video is MP4 or WebM by magic bytes."""
is_mp4 = (
video_bytes.startswith(b"\x00\x00\x00\x18")
or video_bytes.startswith(b"\x00\x00\x00\x1c")
or video_bytes[4:8] == b"ftyp"
)
is_webm = video_bytes.startswith(b"\x1a\x45\xdf\xa3")
assert is_mp4 or is_webm, "Video must be MP4 or WebM"
@dataclasses.dataclass
class TestResult:
name: str
key: str
duration: Optional[float]
succeed: bool
@property
def duration_str(self):
return f"{self.duration:.4f}" if self.duration else "NA"
class TestCLIBase(unittest.TestCase):
model_path: str = None
extra_args = []
data_type: DataType = None
# tested on h100
thresholds = {}
width: int = 720
height: int = 720
output_path: str = "test_outputs"
base_command = [
"sglang",
"generate",
"--text-encoder-cpu-offload",
"--pin-cpu-memory",
"--prompt",
"A curious raccoon",
"--save-output",
"--log-level=debug",
f"--width={width}",
f"--height={height}",
f"--output-path={output_path}",
]
results = []
@classmethod
def setUpClass(cls):
cls.results = []
def _run_command(self, name: str, model_path: str, test_key: str = "", args=[]):
command = (
self.base_command
+ [f"--model-path={model_path}"]
+ shlex.split(args or "")
+ ["--output-file-name", f"{name}"]
+ self.extra_args
)
duration = run_command(command)
status = "Success" if duration else "Failed"
succeed = duration is not None
duration = float(duration) if succeed else None
self.results.append(TestResult(name, test_key, duration, succeed))
return name, duration, status
class TestGenerateBase(TestCLIBase):
model_path: str = None
extra_args = []
data_type: DataType = None
# tested on h100
thresholds = {}
width: int = 720
height: int = 720
output_path: str = "test_outputs"
image_path: str | None = None
prompt: str | None = "A curious raccoon"
base_command = [
"sglang",
"generate",
# "--text-encoder-cpu-offload",
# "--pin-cpu-memory",
f"--prompt",
f"{prompt}",
"--save-output",
"--log-level=debug",
f"--width={width}",
f"--height={height}",
f"--output-path={output_path}",
]
results: list[TestResult] = []
@classmethod
def setUpClass(cls):
cls.results = []
@classmethod
def tearDownClass(cls):
# Print markdown table
print("\n## Test Results\n")
print("| Test Case | Duration | Status |")
print("|--------------------------------|----------|---------|")
test_keys = ["test_single_gpu", "test_cfg_parallel", "test_usp", "test_mixed"]
test_key_to_order = {
test_key: order for order, test_key in enumerate(test_keys)
}
ordered_results: list[TestResult] = [None] * len(test_keys)
for result in cls.results:
order = test_key_to_order[result.key]
ordered_results[order] = result
for result in ordered_results:
if not result:
continue
status = (
"Succeed"
if (
result.succeed
and float(result.duration) <= float(cls.thresholds[result.key])
)
else "Failed"
)
print(f"| {result.name:<30} | {result.duration_str:<8} | {status:<7} |")
print()
durations = [result.duration_str for result in cls.results]
print(" | ".join([""] + durations + [""]))
def _run_test(self, name: str, args, model_path: str, test_key: str):
time_threshold = self.thresholds[test_key]
name, duration, status = self._run_command(
name, args=args, model_path=model_path, test_key=test_key
)
self.verify(status, name, duration, time_threshold)
def verify(self, status, name, duration, time_threshold):
print("-" * 80)
print("\n" * 3)
# test task status
self.assertEqual(status, "Success", f"{name} command failed")
self.assertIsNotNone(duration, f"Could not parse duration for {name}")
self.assertLessEqual(
duration,
time_threshold,
f"{name} failed with {duration:.4f}s > {time_threshold}s",
)
# test output file
path = os.path.join(
self.output_path, f"{name}.{self.data_type.get_default_extension()}"
)
self.assertTrue(os.path.exists(path), f"Output file not exist for {path}")
if self.data_type == DataType.IMAGE:
with Image.open(path) as image:
check_image_size(self, image, self.width, self.height)
logger.info(f"{name} passed in {duration:.4f}s (threshold: {time_threshold}s)")
def model_name(self):
return self.model_path.split("/")[-1]
def test_single_gpu(self):
"""single gpu"""
self._run_test(
name=f"{self.model_name()}_single_gpu",
args=None,
model_path=self.model_path,
test_key="test_single_gpu",
)
def test_cfg_parallel(self):
"""cfg parallel"""
if self.data_type == DataType.IMAGE:
return
self._run_test(
name=f"{self.model_name()}_cfg_parallel",
args="--num-gpus 2 --enable-cfg-parallel",
model_path=self.model_path,
test_key="test_cfg_parallel",
)
def test_usp(self):
"""usp"""
if self.data_type == DataType.IMAGE:
return
self._run_test(
name=f"{self.model_name()}_usp",
args="--num-gpus 4 --ulysses-degree=2 --ring-degree=2",
model_path=self.model_path,
test_key="test_usp",
)
def test_mixed(self):
"""mixed"""
if self.data_type == DataType.IMAGE:
return
self._run_test(
name=f"{self.model_name()}_mixed",
args="--num-gpus 4 --ulysses-degree=2 --ring-degree=1 --enable-cfg-parallel",
model_path=self.model_path,
test_key="test_mixed",
)