minke.yu/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 19:58:04 +08:00
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
python/sglang/multimodal_gen/test/server/diffusion_config.py
+219
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"))
python/sglang/multimodal_gen/test/server/diffusion_server.py
+420
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,
}
python/sglang/multimodal_gen/test/server/perf_baselines.json
+142 -82
View File
@@ -1,82 +1,142 @@
{ {
"metadata": { "metadata": {
"model": "Qwen/Qwen-Image", "model": "Diffusion Server",
"hardware": "CI H100 80GB pool", "hardware": "CI H100 80GB pool",
"description": "Reference numbers captured from the CI diffusion server baseline run" "description": "Reference numbers captured from the CI diffusion server baseline run"
}, },
"tolerances": { "tolerances": {
"e2e": 0.25, "e2e": 0.25,
"stage": 0.3, "stage": 0.3,
"denoise_step": 0.1, "denoise_step": 0.1,
"denoise_agg": 0.1 "denoise_agg": 0.1
}, },
"sampling": { "sampling": {
"step_fractions": [ "step_fractions": [
0.0, 0.0,
0.2, 0.2,
0.4, 0.4,
0.6, 0.6,
0.8, 0.8,
1.0 1.0
], ],
"warmup_requests": { "warmup_requests": {
"text": 1, "text": 1,
"image_edit": 0 "image_edit": 0
} }
}, },
"scenarios": { "scenarios": {
"text_to_image": { "text_to_image": {
"notes": "Single-image generation using the default prompt", "notes": "Single-image generation using the default prompt",
"expected_e2e_ms": 74500.0, "expected_e2e_ms": 74500.0,
"expected_avg_denoise_ms": 422.42, "expected_avg_denoise_ms": 422.42,
"expected_median_denoise_ms": 410.62, "expected_median_denoise_ms": 410.62,
"stages_ms": { "stages_ms": {
"InputValidationStage": 0.1, "InputValidationStage": 0.1,
"TextEncodingStage": 834.2, "TextEncodingStage": 834.2,
"ConditioningStage": 0.1, "ConditioningStage": 0.1,
"TimestepPreparationStage": 10.6, "TimestepPreparationStage": 10.6,
"LatentPreparationStage": 5.2, "LatentPreparationStage": 9.0,
"DenoisingStage": 21202.6, "DenoisingStage": 21202.6,
"DecodingStage": 476.12 "DecodingStage": 476.12
}, },
"denoise_step_ms": { "denoise_step_ms": {
"0": 1077.77, "1": 345.13, "2": 413.8, "3": 405.49, "4": 408.14, "5": 409.06, "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, "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, "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, "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, "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, "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, "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, "42": 411.09, "43": 410.87, "44": 411.37, "45": 411.68, "46": 411.0, "47": 410.09,
"48": 412.72, "49": 410.42 "48": 412.72, "49": 410.42
} }
}, },
"image_edit": { "image_edit": {
"notes": "single uploaded reference image, Qwen/Qwen-Image-Edit", "notes": "single uploaded reference image, Qwen/Qwen-Image-Edit",
"expected_e2e_ms": 138500.0, "expected_e2e_ms": 138500.0,
"expected_avg_denoise_ms": 720.0, "expected_avg_denoise_ms": 720.0,
"expected_median_denoise_ms": 718.0, "expected_median_denoise_ms": 718.0,
"stages_ms": { "stages_ms": {
"InputValidationStage": 14, "InputValidationStage": 23,
"ImageEncodingStage": 1400.0, "ImageEncodingStage": 990.0,
"ImageVAEEncodingStage": 252.76, "ImageVAEEncodingStage": 340.0,
"ConditioningStage": 0.13, "ConditioningStage": 0.13,
"TimestepPreparationStage": 13.78, "TimestepPreparationStage": 13.78,
"LatentPreparationStage": 9.18, "LatentPreparationStage": 10.0,
"DenoisingStage": 36000.0, "DenoisingStage": 36000.0,
"DecodingStage": 645 "DecodingStage": 645
}, },
"denoise_step_ms": { "denoise_step_ms": {
"0": 720.0, "1": 720.0, "2": 720.0, "3": 720.0, "4": 720.0, "5": 720.0, "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, "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, "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, "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, "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, "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, "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, "42": 720.0, "43": 720.0, "44": 720.0, "45": 720.0, "46": 720.0, "47": 720.0,
"48": 720.0, "49": 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
}
}
}
python/sglang/multimodal_gen/test/server/test_server_performance.py
+361 -401
View File
@@ -1,34 +1,39 @@
# Server-based diffusion performance test: """
# - Launches an sglang diffusion server via the CLI. Config-driven diffusion performance test with pytest parametrization.
# - Issues an OpenAI-compatible Images API request. Adding a new model/scenario = adding one DiffusionCase entry in diffusion_config.py.
# - Extracts all performance metrics from performance.log (no stdout parsing). """
# - Verifies E2E, stage-level, and denoising-step latencies with configurable buffers.
from __future__ import annotations from __future__ import annotations
import base64
import json
import os import os
import statistics
import subprocess
import tempfile
import time import time
from pathlib import Path from pathlib import Path
from typing import Any, Sequence from typing import Any, Callable
import pytest import pytest
from openai import OpenAI 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.runtime.utils.logging_utils import init_logger
from sglang.multimodal_gen.test.server.conftest import _GLOBAL_PERF_RESULTS 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 ( from sglang.multimodal_gen.test.test_utils import (
get_dynamic_server_port, get_dynamic_server_port,
is_jpeg,
is_png,
prepare_perf_log,
read_perf_records, read_perf_records,
sample_step_indices, validate_image,
validate_openai_video,
wait_for_perf_record, wait_for_perf_record,
wait_for_stage_metrics, wait_for_stage_metrics,
) )
@@ -36,261 +41,74 @@ from sglang.multimodal_gen.test.test_utils import (
logger = init_logger(__name__) logger = init_logger(__name__)
_BASELINE_PATH = Path(__file__).with_name("perf_baselines.json") @pytest.fixture(params=DIFFUSION_CASES, ids=lambda c: c.id)
with _BASELINE_PATH.open("r", encoding="utf-8") as _fh: def case(request) -> DiffusionCase:
_BASELINE_CONFIG = json.load(_fh) """Provide a DiffusionCase for each test."""
return request.param
_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"]
def _tolerance_from_env(var_name: str, default: float) -> float: @pytest.fixture
override = os.environ.get(var_name) def diffusion_server(case: DiffusionCase) -> ServerContext:
if override is not None: """Start a diffusion server for a single case and tear it down afterwards."""
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__))
default_port = get_dynamic_server_port() default_port = get_dynamic_server_port()
port = int(os.environ.get("SGLANG_TEST_SERVER_PORT", default_port)) port = int(os.environ.get("SGLANG_TEST_SERVER_PORT", default_port))
port = getattr(cls, "SERVER_PORT", port)
model = getattr(cls, "MODEL_PATH") # start server
wait_deadline = float(os.environ.get("SGLANG_TEST_WAIT_SECS", "1200")) manager = ServerManager(
serve_extra_args = os.environ.get("SGLANG_TEST_SERVE_ARGS", "") model=case.model_path,
port=port,
safe_model_name = model.replace("/", "_") wait_deadline=float(os.environ.get("SGLANG_TEST_WAIT_SECS", "1200")),
stdout_path = ( extra_args=os.environ.get("SGLANG_TEST_SERVE_ARGS", ""),
Path(tempfile.gettempdir()) / f"sgl_server_{port}_{safe_model_name}.log"
) )
stdout_path.unlink(missing_ok=True) ctx = manager.start()
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)
if case.startup_grace_seconds > 0:
logger.info( logger.info(
"[server-test] Waiting for server to initialize... elapsed=%ss", "[server-test] Waiting %.1fs for %s to settle",
int(time.time() - start), case.startup_grace_seconds,
case.id,
) )
time.sleep(5) time.sleep(case.startup_grace_seconds)
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)
try: 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: except Exception as exc:
logger.error("Warm-up requests failed: %s", exc) logger.error("Warm-up failed for %s: %s", case.id, exc)
kill_process_tree(process.pid) ctx.cleanup()
raise raise
yield ctx
try: try:
kill_process_tree(process.pid) yield ctx
except Exception: finally:
pass ctx.cleanup()
try:
stdout_fh.flush()
stdout_fh.close()
except Exception:
pass
@pytest.mark.usefixtures("diffusion_server") class TestDiffusionPerformance:
class DiffusionPerfTestBase: """Performance tests for all diffusion models/scenarios.
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
STAGE_EXPECTATIONS: dict This single test class runs against all cases defined in DIFFUSION_CASES.
STEP_EXPECTATIONS: dict Each case gets its own server instance via the parametrized fixture.
EXPECTED_E2E_MS: float """
EXPECTED_AVG_DENOISE_MS: float
EXPECTED_MEDIAN_DENOISE_MS: float
_perf_results: list[dict[str, Any]] = [] _perf_results: list[dict[str, Any]] = []
@@ -304,187 +122,329 @@ class DiffusionPerfTestBase:
result["class_name"] = cls.__name__ result["class_name"] = cls.__name__
_GLOBAL_PERF_RESULTS.append(result) _GLOBAL_PERF_RESULTS.append(result)
def _client(self) -> OpenAI: def _client(self, ctx: ServerContext) -> OpenAI:
"""Get OpenAI client for the server."""
return OpenAI( return OpenAI(
api_key="sglang-anything", 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: def _run_and_collect(
return self.server_ctx["perf_log_path"] self,
ctx: ServerContext,
def _record_result(self, test_name: str, summary: dict[str, Any]) -> None: case: DiffusionCase,
if not summary: generate_fn: Callable[[], None],
return ) -> tuple[dict, dict]:
entry = {"test_name": test_name, **summary} """Run generation and collect performance records."""
self.__class__._perf_results.append(entry) log_path = ctx.perf_log_path
def _run_and_collect_records(self, generate_fn) -> tuple[dict, dict]:
log_path = self._perf_log_path()
prev_len = len(read_perf_records(log_path)) prev_len = len(read_perf_records(log_path))
generate_fn() generate_fn()
perf_record, _ = wait_for_perf_record( perf_record, _ = wait_for_perf_record(
"total_inference_time", "total_inference_time",
prev_len, prev_len,
log_path, log_path,
) )
scenario = BASELINE_CONFIG.scenarios[case.scenario_name]
stage_metrics, _ = wait_for_stage_metrics( stage_metrics, _ = wait_for_stage_metrics(
perf_record.get("request_id", ""), perf_record.get("request_id", ""),
prev_len, prev_len,
len(self.STAGE_EXPECTATIONS), len(scenario.stages_ms),
log_path, log_path,
) )
return perf_record, stage_metrics return perf_record, stage_metrics
def _generate_image(self): def _generate_for_case(
client = self._client() self,
result = client.images.generate( ctx: ServerContext,
model=self.MODEL_PATH, case: DiffusionCase,
prompt=self.PROMPT, ) -> Callable[[], None]:
n=1, """Return appropriate generation function for the case."""
size=self.OUTPUT_SIZE, client = self._client(ctx)
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_image_edit(self): def _create_and_download_video(
if not self.IMAGE_EDIT_PROMPT: *,
pytest.skip("Image edit prompt not configured") model: str,
if not self.IMAGE_EDIT_PATH.exists(): size: str,
pytest.skip(f"Image edit file missing: {self.IMAGE_EDIT_PATH}") prompt: str | None = None,
client = self._client() seconds: int | None = None,
with self.IMAGE_EDIT_PATH.open("rb") as fh: input_reference: Any | None = None,
result = client.images.edit( ) -> bytes:
model=self.MODEL_PATH, """
image=fh, Create a video job via /v1/videos, poll until completion,
prompt=self.IMAGE_EDIT_PROMPT, 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, n=1,
size=self.OUTPUT_SIZE, size=case.output_size,
response_format="b64_json", response_format="b64_json",
) )
image_bytes = base64.b64decode(result.data[0].b64_json) validate_image(result.data[0].b64_json)
assert is_png(image_bytes) or is_jpeg(
image_bytes
), "Edited image must be PNG or JPEG"
def _assert_metrics(self, perf_record: dict, stage_metrics: dict): def generate_image_edit():
e2e_ms = float(perf_record.get("total_duration_ms", 0.0)) """TI2I: Text + Image ? Image edit."""
assert e2e_ms > 0, "E2E duration missing from perf log" if not case.image_edit_prompt or not case.image_edit_path:
e2e_upper = self.EXPECTED_E2E_MS * (1 + E2E_TOLERANCE_RATIO) pytest.skip(f"{case.id}: no edit config")
assert (
e2e_ms <= e2e_upper
), f"E2E time {e2e_ms:.2f}ms exceeds allowed {e2e_upper:.2f}ms"
steps = [ # Handle URL or local path
step if case.is_image_url():
for step in perf_record.get("steps", []) or [] image_path = download_image_from_url(str(case.image_edit_path))
if step.get("name") == "denoising_step_guided" and "duration_ms" in step else:
] image_path = Path(case.image_edit_path)
assert steps, "Denoising step timings missing from perf log" if not image_path.exists():
pytest.skip(f"{case.id}: file missing: {image_path}")
durations = [float(step["duration_ms"]) for step in steps] with image_path.open("rb") as fh:
avg_duration = sum(durations) / len(durations) result = client.images.edit(
median_duration = statistics.median(durations) 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) # VIDEO MODE
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"
avg_per_step = { def generate_video():
int(step.get("index")): float(step["duration_ms"]) """T2V: Text ? Video."""
for step in steps if not case.prompt:
if step.get("index") is not None 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" # video-specific metrics
for stage, expected in self.STAGE_EXPECTATIONS.items(): if summary.frames_per_second:
actual = stage_metrics.get(stage) result.update(
assert actual is not None, f"Stage {stage} timing missing" {
upper_bound = expected * (1 + STAGE_TOLERANCE_RATIO) "frames_per_second": summary.frames_per_second,
assert ( "total_frames": summary.total_frames,
actual <= upper_bound "avg_frame_time_ms": summary.avg_frame_time_ms,
), f"Stage {stage} took {actual:.2f}ms > allowed {upper_bound:.2f}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( logger.info(
"[Perf] %s/%s: E2E %.2f ms; Avg denoise %.2f ms; Median %.2f ms", "[BASELINE] %s expected_avg_denoise_ms = %.2f",
self.__class__.__name__, case.id,
perf_record.get("test_name", "test"), summary.avg_denoise_ms,
e2e_ms, )
avg_duration, logger.info(
median_duration, "[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 { # Only log video-specific metrics when they exist
"e2e_ms": e2e_ms, if summary.frames_per_second is not None:
"avg_denoise_ms": avg_duration, logger.info(
"median_denoise_ms": median_duration, "[BASELINE] %s frames_per_second = %.2f",
"stage_metrics": stage_metrics, case.id,
"sampled_steps": sampled_steps, 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): Pytest will execute this test once per case in DIFFUSION_CASES,
"""Performance tests for the Qwen/Qwen-image model.""" with test IDs like:
- test_diffusion_perf[qwen_image_text]
MODEL_PATH = "Qwen/Qwen-Image" - test_diffusion_perf[qwen_image_edit]
STARTUP_GRACE_SECONDS = 30.0 - etc.
WARMUP_IMAGE_EDIT_REQUESTS = 0 """
STAGE_EXPECTATIONS = _TEXT_SCENARIO["stages_ms"] generate_fn = self._generate_for_case(diffusion_server, case)
STEP_EXPECTATIONS = { perf_record, stage_metrics = self._run_and_collect(
int(k): v for k, v in _TEXT_SCENARIO["denoise_step_ms"].items() diffusion_server,
} case,
EXPECTED_E2E_MS = float(_TEXT_SCENARIO["expected_e2e_ms"]) generate_fn,
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
) )
summary = self._assert_metrics(perf_record, stage_metrics) self._validate_and_record(case, perf_record, stage_metrics)
self._record_result("image_edit", summary)
python/sglang/multimodal_gen/test/test_utils.py
+443 -415
View File
@@ -1,415 +1,443 @@
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
import dataclasses import base64
import json import dataclasses
import os import json
import shlex import os
import socket import shlex
import subprocess import socket
import sys import subprocess
import time import sys
import unittest import time
from pathlib import Path import unittest
from typing import Optional, Sequence from pathlib import Path
from typing import Optional, Sequence
from PIL import Image
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.configs.sample.base import DataType
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger 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__)
logger = init_logger(__name__)
def run_command(command) -> Optional[float]:
"""Runs a command and returns the execution time and status.""" def run_command(command) -> Optional[float]:
print(f"Running command: {shlex.join(command)}") """Runs a command and returns the execution time and status."""
print(f"Running command: {shlex.join(command)}")
duration = None
with subprocess.Popen( duration = None
command, with subprocess.Popen(
stdout=subprocess.PIPE, command,
stderr=subprocess.STDOUT, stdout=subprocess.PIPE,
text=True, stderr=subprocess.STDOUT,
encoding="utf-8", text=True,
) as process: encoding="utf-8",
for line in process.stdout: ) as process:
sys.stdout.write(line) for line in process.stdout:
if "Pixel data generated" in line: sys.stdout.write(line)
words = line.split(" ") if "Pixel data generated" in line:
duration = float(words[-2]) words = line.split(" ")
duration = float(words[-2])
if process.returncode == 0:
return duration if process.returncode == 0:
else: return duration
print(f"Command failed with exit code {process.returncode}") else:
return None 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: def probe_port(host="127.0.0.1", port=30010, timeout=2.0) -> bool:
s.settimeout(timeout) with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
try: s.settimeout(timeout)
s.connect((host, port)) try:
return True s.connect((host, port))
except OSError: return True
return False except OSError:
return False
def is_in_ci() -> bool:
return get_bool_env_var("SGLANG_IS_IN_CI") 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") def get_dynamic_server_port() -> int:
if not cuda_devices: cuda_devices = os.environ.get("CUDA_VISIBLE_DEVICES", "0")
cuda_devices = "0" if not cuda_devices:
try: cuda_devices = "0"
first_device_id = int(cuda_devices.split(",")[0].strip()[0]) try:
except (ValueError, IndexError): first_device_id = int(cuda_devices.split(",")[0].strip()[0])
first_device_id = 0 except (ValueError, IndexError):
first_device_id = 0
if is_in_ci():
base_port = 10000 + first_device_id * 2000 if is_in_ci():
else: base_port = 10000 + first_device_id * 2000
base_port = 20000 + first_device_id * 1000 else:
base_port = 20000 + first_device_id * 1000
return base_port + 1000
return base_port + 1000
def is_mp4(data):
idx = data.find(b"ftyp") def is_mp4(data):
return 0 <= idx <= 32 idx = data.find(b"ftyp")
return 0 <= idx <= 32
def is_jpeg(data: bytes) -> bool:
# JPEG files start with: FF D8 FF def is_jpeg(data: bytes) -> bool:
return data.startswith(b"\xff\xd8\xff") # 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 def is_png(data):
return data.startswith(b"\x89PNG\r\n\x1a\n") # 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 def wait_for_port(host="127.0.0.1", port=30010, deadline=300.0, interval=0.5):
last_err = None end = time.time() + deadline
while time.time() < end: last_err = None
if probe_port(host, port, timeout=interval): while time.time() < end:
return True if probe_port(host, port, timeout=interval):
time.sleep(interval) return True
raise TimeoutError(f"Port {host}:{port} not ready. Last error: {last_err}") 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 def check_image_size(ut, image, width, height):
ut.assertEqual(image.size, (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.""" def get_perf_log_dir(start_file: Path) -> Path:
this_file = start_file.resolve() """Mirror runtime/utils/performance_logger.py behaviour for locating logs."""
root_logs = this_file.parents[3] / "logs" this_file = start_file.resolve()
fallback = this_file.parents[2] / "logs" root_logs = this_file.parents[3] / "logs"
return root_logs if root_logs.exists() or not fallback.exists() else fallback 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) def _ensure_log_path(log_dir: Path) -> Path:
return log_dir / "performance.log" 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.""" def clear_perf_log(log_dir: Path) -> Path:
log_path = _ensure_log_path(log_dir) """Delete the perf log file so tests can watch for fresh entries."""
if log_path.exists(): log_path = _ensure_log_path(log_dir)
log_path.unlink() if log_path.exists():
logger.info("[server-test] Monitoring perf log at %s", log_path.as_posix()) log_path.unlink()
return log_path 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.""" def prepare_perf_log(start_file: Path) -> tuple[Path, Path]:
log_dir = get_perf_log_dir(start_file) """Convenience helper to resolve and clear the perf log in one call."""
log_path = clear_perf_log(log_dir) log_dir = get_perf_log_dir(start_file)
return log_dir, log_path 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(): def read_perf_records(log_path: Path) -> list[dict]:
return [] if not log_path.exists():
records: list[dict] = [] return []
with log_path.open("r", encoding="utf-8") as fh: records: list[dict] = []
for line in fh: with log_path.open("r", encoding="utf-8") as fh:
line = line.strip() for line in fh:
if not line: line = line.strip()
continue if not line:
try: continue
records.append(json.loads(line)) try:
except json.JSONDecodeError: records.append(json.loads(line))
continue except json.JSONDecodeError:
return records continue
return records
def wait_for_perf_record(
tag: str, def wait_for_perf_record(
prev_len: int, tag: str,
log_path: Path, prev_len: int,
timeout: float = 120.0, log_path: Path,
) -> tuple[dict, int]: timeout: float = 120.0,
deadline = time.time() + timeout ) -> tuple[dict, int]:
while time.time() < deadline: deadline = time.time() + timeout
records = read_perf_records(log_path) while time.time() < deadline:
if len(records) > prev_len: records = read_perf_records(log_path)
for rec in records[prev_len:]: if len(records) > prev_len:
if rec.get("tag") == tag: for rec in records[prev_len:]:
return rec, len(records) if rec.get("tag") == tag:
time.sleep(0.5) return rec, len(records)
raise AssertionError( time.sleep(0.5)
f"Timeout waiting for perf log entry '{tag}' (start_len={prev_len})" raise AssertionError(
) f"Timeout waiting for perf log entry '{tag}' (start_len={prev_len})"
)
def wait_for_stage_metrics(
request_id: str, def wait_for_stage_metrics(
prev_len: int, request_id: str,
expected_count: int, prev_len: int,
log_path: Path, expected_count: int,
timeout: float = 120.0, log_path: Path,
) -> tuple[dict[str, float], int]: timeout: float = 120.0,
deadline = time.time() + timeout ) -> tuple[dict[str, float], int]:
metrics: dict[str, float] = {} deadline = time.time() + timeout
while time.time() < deadline: metrics: dict[str, float] = {}
records = read_perf_records(log_path) while time.time() < deadline:
for rec in records[prev_len:]: records = read_perf_records(log_path)
if ( for rec in records[prev_len:]:
rec.get("tag") == "pipeline_stage_metric" if (
and rec.get("request_id") == request_id rec.get("tag") == "pipeline_stage_metric"
): and rec.get("request_id") == request_id
stage = rec.get("stage") ):
duration = rec.get("duration_ms") stage = rec.get("stage")
if stage is not None and duration is not None: duration = rec.get("duration_ms")
metrics[str(stage)] = float(duration) if stage is not None and duration is not None:
if len(metrics) >= expected_count: metrics[str(stage)] = float(duration)
return metrics, len(records) if len(metrics) >= expected_count:
time.sleep(0.5) return metrics, len(records)
raise AssertionError( time.sleep(0.5)
f"Timeout waiting for stage metrics for request {request_id} " raise AssertionError(
f"(collected={len(metrics)} expected={expected_count})" 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] def sample_step_indices(
) -> list[int]: step_map: dict[int, float], fractions: Sequence[float]
if not step_map: ) -> list[int]:
return [] if not step_map:
max_idx = max(step_map.keys()) return []
indices = set() max_idx = max(step_map.keys())
for fraction in fractions: indices = set()
idx = min(max_idx, max(0, int(round(fraction * max_idx)))) for fraction in fractions:
if idx in step_map: idx = min(max_idx, max(0, int(round(fraction * max_idx))))
indices.add(idx) if idx in step_map:
return sorted(indices) indices.add(idx)
return sorted(indices)
@dataclasses.dataclass
class TestResult: def validate_image(b64_json: str) -> None:
name: str """Decode and validate that image is PNG or JPEG."""
key: str image_bytes = base64.b64decode(b64_json)
duration: Optional[float] assert is_png(image_bytes) or is_jpeg(image_bytes), "Image must be PNG or JPEG"
succeed: bool
@property def validate_video(b64_json: str) -> None:
def duration_str(self): """Decode and validate that video is a valid format."""
return f"{self.duration:.4f}" if self.duration else "NA" 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"
class TestCLIBase(unittest.TestCase): )
model_path: str = None is_webm = video_bytes[:4] == b"\x1a\x45\xdf\xa3"
extra_args = [] assert is_mp4 or is_webm, "Video must be MP4 or WebM"
data_type: DataType = None
# tested on h100
thresholds = {} def validate_openai_video(video_bytes: bytes) -> None:
"""Validate that video is MP4 or WebM by magic bytes."""
width: int = 720 is_mp4 = (
height: int = 720 video_bytes.startswith(b"\x00\x00\x00\x18")
output_path: str = "test_outputs" or video_bytes.startswith(b"\x00\x00\x00\x1c")
or video_bytes[4:8] == b"ftyp"
base_command = [ )
"sglang", is_webm = video_bytes.startswith(b"\x1a\x45\xdf\xa3")
"generate", assert is_mp4 or is_webm, "Video must be MP4 or WebM"
"--text-encoder-cpu-offload",
"--pin-cpu-memory",
"--prompt", @dataclasses.dataclass
"A curious raccoon", class TestResult:
"--save-output", name: str
"--log-level=debug", key: str
f"--width={width}", duration: Optional[float]
f"--height={height}", succeed: bool
f"--output-path={output_path}",
] @property
def duration_str(self):
results = [] return f"{self.duration:.4f}" if self.duration else "NA"
@classmethod
def setUpClass(cls): class TestCLIBase(unittest.TestCase):
cls.results = [] model_path: str = None
extra_args = []
def _run_command(self, name: str, model_path: str, test_key: str = "", args=[]): data_type: DataType = None
command = ( # tested on h100
self.base_command thresholds = {}
+ [f"--model-path={model_path}"]
+ shlex.split(args or "") width: int = 720
+ ["--output-file-name", f"{name}"] height: int = 720
+ self.extra_args output_path: str = "test_outputs"
)
duration = run_command(command) base_command = [
status = "Success" if duration else "Failed" "sglang",
succeed = duration is not None "generate",
"--text-encoder-cpu-offload",
duration = float(duration) if succeed else None "--pin-cpu-memory",
self.results.append(TestResult(name, test_key, duration, succeed)) "--prompt",
"A curious raccoon",
return name, duration, status "--save-output",
"--log-level=debug",
f"--width={width}",
class TestGenerateBase(TestCLIBase): f"--height={height}",
model_path: str = None f"--output-path={output_path}",
extra_args = [] ]
data_type: DataType = None
# tested on h100 results = []
thresholds = {}
@classmethod
width: int = 720 def setUpClass(cls):
height: int = 720 cls.results = []
output_path: str = "test_outputs"
image_path: str | None = None def _run_command(self, name: str, model_path: str, test_key: str = "", args=[]):
prompt: str | None = "A curious raccoon" command = (
self.base_command
base_command = [ + [f"--model-path={model_path}"]
"sglang", + shlex.split(args or "")
"generate", + ["--output-file-name", f"{name}"]
# "--text-encoder-cpu-offload", + self.extra_args
# "--pin-cpu-memory", )
f"--prompt", duration = run_command(command)
f"{prompt}", status = "Success" if duration else "Failed"
"--save-output", succeed = duration is not None
"--log-level=debug",
f"--width={width}", duration = float(duration) if succeed else None
f"--height={height}", self.results.append(TestResult(name, test_key, duration, succeed))
f"--output-path={output_path}",
] return name, duration, status
results: list[TestResult] = []
class TestGenerateBase(TestCLIBase):
@classmethod model_path: str = None
def setUpClass(cls): extra_args = []
cls.results = [] data_type: DataType = None
# tested on h100
@classmethod thresholds = {}
def tearDownClass(cls):
# Print markdown table width: int = 720
print("\n## Test Results\n") height: int = 720
print("| Test Case | Duration | Status |") output_path: str = "test_outputs"
print("|--------------------------------|----------|---------|") image_path: str | None = None
test_keys = ["test_single_gpu", "test_cfg_parallel", "test_usp", "test_mixed"] prompt: str | None = "A curious raccoon"
test_key_to_order = {
test_key: order for order, test_key in enumerate(test_keys) base_command = [
} "sglang",
"generate",
ordered_results: list[TestResult] = [None] * len(test_keys) # "--text-encoder-cpu-offload",
for result in cls.results: # "--pin-cpu-memory",
order = test_key_to_order[result.key] f"--prompt",
ordered_results[order] = result f"{prompt}",
"--save-output",
for result in ordered_results: "--log-level=debug",
if not result: f"--width={width}",
continue f"--height={height}",
status = ( f"--output-path={output_path}",
"Succeed" ]
if (
result.succeed results: list[TestResult] = []
and float(result.duration) <= float(cls.thresholds[result.key])
) @classmethod
else "Failed" def setUpClass(cls):
) cls.results = []
print(f"| {result.name:<30} | {result.duration_str:<8} | {status:<7} |")
print() @classmethod
durations = [result.duration_str for result in cls.results] def tearDownClass(cls):
print(" | ".join([""] + durations + [""])) # Print markdown table
print("\n## Test Results\n")
def _run_test(self, name: str, args, model_path: str, test_key: str): print("| Test Case | Duration | Status |")
time_threshold = self.thresholds[test_key] print("|--------------------------------|----------|---------|")
name, duration, status = self._run_command( test_keys = ["test_single_gpu", "test_cfg_parallel", "test_usp", "test_mixed"]
name, args=args, model_path=model_path, test_key=test_key test_key_to_order = {
) test_key: order for order, test_key in enumerate(test_keys)
self.verify(status, name, duration, time_threshold) }
def verify(self, status, name, duration, time_threshold): ordered_results: list[TestResult] = [None] * len(test_keys)
print("-" * 80) for result in cls.results:
print("\n" * 3) order = test_key_to_order[result.key]
ordered_results[order] = result
# test task status
self.assertEqual(status, "Success", f"{name} command failed") for result in ordered_results:
self.assertIsNotNone(duration, f"Could not parse duration for {name}") if not result:
self.assertLessEqual( continue
duration, status = (
time_threshold, "Succeed"
f"{name} failed with {duration:.4f}s > {time_threshold}s", if (
) result.succeed
and float(result.duration) <= float(cls.thresholds[result.key])
# test output file )
path = os.path.join( else "Failed"
self.output_path, f"{name}.{self.data_type.get_default_extension()}" )
) print(f"| {result.name:<30} | {result.duration_str:<8} | {status:<7} |")
self.assertTrue(os.path.exists(path), f"Output file not exist for {path}") print()
if self.data_type == DataType.IMAGE: durations = [result.duration_str for result in cls.results]
with Image.open(path) as image: print(" | ".join([""] + durations + [""]))
check_image_size(self, image, self.width, self.height)
logger.info(f"{name} passed in {duration:.4f}s (threshold: {time_threshold}s)") def _run_test(self, name: str, args, model_path: str, test_key: str):
time_threshold = self.thresholds[test_key]
def model_name(self): name, duration, status = self._run_command(
return self.model_path.split("/")[-1] name, args=args, model_path=model_path, test_key=test_key
)
def test_single_gpu(self): self.verify(status, name, duration, time_threshold)
"""single gpu"""
self._run_test( def verify(self, status, name, duration, time_threshold):
name=f"{self.model_name()}_single_gpu", print("-" * 80)
args=None, print("\n" * 3)
model_path=self.model_path,
test_key="test_single_gpu", # test task status
) self.assertEqual(status, "Success", f"{name} command failed")
self.assertIsNotNone(duration, f"Could not parse duration for {name}")
def test_cfg_parallel(self): self.assertLessEqual(
"""cfg parallel""" duration,
if self.data_type == DataType.IMAGE: time_threshold,
return f"{name} failed with {duration:.4f}s > {time_threshold}s",
self._run_test( )
name=f"{self.model_name()}_cfg_parallel",
args="--num-gpus 2 --enable-cfg-parallel", # test output file
model_path=self.model_path, path = os.path.join(
test_key="test_cfg_parallel", self.output_path, f"{name}.{self.data_type.get_default_extension()}"
) )
self.assertTrue(os.path.exists(path), f"Output file not exist for {path}")
def test_usp(self): if self.data_type == DataType.IMAGE:
"""usp""" with Image.open(path) as image:
if self.data_type == DataType.IMAGE: check_image_size(self, image, self.width, self.height)
return logger.info(f"{name} passed in {duration:.4f}s (threshold: {time_threshold}s)")
self._run_test(
name=f"{self.model_name()}_usp", def model_name(self):
args="--num-gpus 4 --ulysses-degree=2 --ring-degree=2", return self.model_path.split("/")[-1]
model_path=self.model_path,
test_key="test_usp", def test_single_gpu(self):
) """single gpu"""
self._run_test(
def test_mixed(self): name=f"{self.model_name()}_single_gpu",
"""mixed""" args=None,
if self.data_type == DataType.IMAGE: model_path=self.model_path,
return test_key="test_single_gpu",
self._run_test( )
name=f"{self.model_name()}_mixed",
args="--num-gpus 4 --ulysses-degree=2 --ring-degree=1 --enable-cfg-parallel", def test_cfg_parallel(self):
model_path=self.model_path, """cfg parallel"""
test_key="test_mixed", 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",
)