diff --git a/python/sglang/multimodal_gen/test/server/diffusion_config.py b/python/sglang/multimodal_gen/test/server/diffusion_config.py new file mode 100644 index 000000000..16381c608 --- /dev/null +++ b/python/sglang/multimodal_gen/test/server/diffusion_config.py @@ -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")) diff --git a/python/sglang/multimodal_gen/test/server/diffusion_server.py b/python/sglang/multimodal_gen/test/server/diffusion_server.py new file mode 100644 index 000000000..fabfe0e60 --- /dev/null +++ b/python/sglang/multimodal_gen/test/server/diffusion_server.py @@ -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, +} diff --git a/python/sglang/multimodal_gen/test/server/perf_baselines.json b/python/sglang/multimodal_gen/test/server/perf_baselines.json index f05fbf7cf..473343c9d 100644 --- a/python/sglang/multimodal_gen/test/server/perf_baselines.json +++ b/python/sglang/multimodal_gen/test/server/perf_baselines.json @@ -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 + } + } +} diff --git a/python/sglang/multimodal_gen/test/server/test_server_performance.py b/python/sglang/multimodal_gen/test/server/test_server_performance.py index 61df5550a..d9ac0a7ae 100644 --- a/python/sglang/multimodal_gen/test/server/test_server_performance.py +++ b/python/sglang/multimodal_gen/test/server/test_server_performance.py @@ -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) diff --git a/python/sglang/multimodal_gen/test/test_utils.py b/python/sglang/multimodal_gen/test/test_utils.py index 0e947dae0..a923e5697 100644 --- a/python/sglang/multimodal_gen/test/test_utils.py +++ b/python/sglang/multimodal_gen/test/test_utils.py @@ -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", + )