sgl-project--sglang/python/sglang/multimodal_gen/benchmarks/bench_offline_throughput.py

"""
Benchmark offline throughput for multimodal generation models (Image/Video Generation).

This script benchmarks generation throughput without running a server, using low-level APIs.
It provides detailed metrics on throughput, latency, and resource utilization.

# Usage Examples

## Text-to-Video with VBench dataset
python -m sglang.multimodal_gen.benchmarks.bench_offline_throughput \\
    --model-path Wan-AI/Wan2.1-T2V-1.3B-Diffusers \\
    --dataset vbench \\
    --num-prompts 20 \\
    --batch-size 1 \\
    --width 512 --height 512 --num-frames 16

## Random dataset for stress testing
python -m sglang.multimodal_gen.benchmarks.bench_offline_throughput \\
    --model-path Wan-AI/Wan2.1-T2V-1.3B-Diffusers \\
    --dataset random \\
    --num-prompts 100 \\
    --batch-size 1 \\
    --num-inference-steps 20 \\
    --output-file results.json
"""

import argparse
import dataclasses
import json
import time
from dataclasses import dataclass
from typing import Any, Dict, List, Optional, Tuple

import torch
from tqdm import tqdm

from sglang.multimodal_gen.benchmarks.datasets import RandomDataset, VBenchDataset
from sglang.multimodal_gen.runtime.entrypoints.diffusion_generator import DiffGenerator
from sglang.multimodal_gen.runtime.server_args import ServerArgs, set_global_server_args
from sglang.multimodal_gen.runtime.utils.logging_utils import (
    configure_logger,
    init_logger,
)
from sglang.multimodal_gen.test.test_utils import print_divider, print_value_formatted

logger = init_logger(__name__)


@dataclass
class BatchOutput:
    """Container for batch generation results."""

    latency: float = 0.0
    latency_per_sample: float = 0.0
    num_samples: int = 0
    total_frames: int = 0
    peak_memory_mb: float = 0.0
    success: bool = False
    error: str = ""


@dataclass
class BenchArgs:
    """Benchmark configuration for multimodal generation."""

    # Diffusion Model Configuration
    num_inference_steps: int = 20
    guidance_scale: float = 7.5
    seed: int = 42
    disable_safety_checker: bool = False

    # Output Configuration
    width: int = 32
    height: int = 32
    num_frames: int = 1
    fps: int = 24

    # Dataset & Benchmark
    dataset: str = "random"
    dataset_path: str = ""
    task_name: str = "unknown"
    num_prompts: int = 10
    num_outputs_per_prompt: int = 1
    batch_size: int = 1
    random_request_config: str = None
    random_request_seed: int = 42

    # Benchmark Execution
    skip_warmup: bool = False
    output_file: str = ""
    disable_tqdm: bool = False

    # Profiling
    profile: bool = False
    num_profiled_timesteps: int = 5
    profile_all_stages: bool = False

    @staticmethod
    def add_cli_args(parser: argparse.ArgumentParser):
        """Add benchmark-specific CLI arguments."""
        # Diffusion Model Configuration
        parser.add_argument(
            "--num-inference-steps",
            type=int,
            default=20,
            help="Number of denoising steps",
        )
        parser.add_argument(
            "--guidance-scale",
            type=float,
            default=7.5,
            help="Classifier-free guidance scale",
        )
        parser.add_argument("--seed", type=int, default=42, help="Random seed")
        parser.add_argument(
            "--disable-safety-checker",
            action="store_true",
            help="Disable NSFW detection",
        )

        # Output Configuration
        parser.add_argument("--width", type=int, default=32, help="Image/video width")
        parser.add_argument("--height", type=int, default=32, help="Image/video height")
        parser.add_argument(
            "--num-frames", type=int, default=1, help="Number of frames for video"
        )
        parser.add_argument("--fps", type=int, default=24, help="FPS for video")

        # Dataset & Benchmark
        parser.add_argument(
            "--dataset",
            type=str,
            default="random",
            choices=["vbench", "random"],
            help="Dataset to use",
        )
        parser.add_argument(
            "--dataset-path",
            type=str,
            default="",
            help="Path to dataset (prompts file or image directory)",
        )
        parser.add_argument(
            "--task-name",
            type=str,
            default="unknown",
            help="Task name for benchmark identification",
        )
        parser.add_argument(
            "--num-prompts",
            type=int,
            default=10,
            help="Total number of prompts to benchmark",
        )
        parser.add_argument(
            "--num-outputs-per-prompt",
            type=int,
            default=1,
            help="Number of generated outputs requested per prompt",
        )
        parser.add_argument(
            "--batch-size",
            type=int,
            default=1,
            help="Batch size per generation call (currently only bs=1 is supported)",
        )

        parser.add_argument(
            "--random-request-config",
            type=str,
            default=None,
            help=(
                "JSON string defining random request profiles. "
                "Each profile may contain: width, height, num_inference_steps, etc. "
                "The 'weight' field controls sampling probability (relative weight)."
            ),
        )
        parser.add_argument(
            "--random-request-seed",
            type=int,
            default=42,
            help="Random seed for sampling request profiles (default: 42).",
        )

        # Benchmark Execution
        parser.add_argument(
            "--skip-warmup", action="store_true", help="Skip warmup batch"
        )
        parser.add_argument(
            "--output-file",
            type=str,
            default="",
            help="Output JSON file for results (append mode)",
        )
        parser.add_argument(
            "--disable-tqdm",
            action="store_true",
            help="Disable progress bar",
        )
        parser.add_argument(
            "--profile",
            action="store_true",
            help=(
                "Enable PyTorch profiler for diffusion generation. "
                "Set SGLANG_DIFFUSION_TORCH_PROFILER_DIR to control trace output directory."
            ),
        )
        parser.add_argument(
            "--num-profiled-timesteps",
            type=int,
            default=5,
            help=(
                "Number of denoising timesteps to profile after warmup. "
                "Use -1 to profile all denoising timesteps."
            ),
        )
        parser.add_argument(
            "--profile-all-stages",
            action="store_true",
            help="Profile all diffusion pipeline stages instead of only denoising steps.",
        )

    @classmethod
    def from_cli_args(cls, args: argparse.Namespace):
        """Create BenchArgs from parsed CLI arguments."""
        attrs = [attr.name for attr in dataclasses.fields(cls)]
        return cls(**{attr: getattr(args, attr) for attr in attrs})


def initialize_engine(server_args: ServerArgs) -> DiffGenerator:
    """Initialize diffusion pipeline engine."""
    logger.info("Initializing engine...")
    engine = DiffGenerator.from_server_args(server_args, local_mode=True)
    logger.info("Engine initialized successfully")
    return engine


def generate_batch(
    engine: DiffGenerator,
    bench_args: BenchArgs,
    prompts: List[str],
    user_sampling_params: List[Dict[str, Any]],
) -> BatchOutput:
    """Generate batch of images/videos synchronously."""
    assert len(user_sampling_params) == len(prompts), (
        f"user_sampling_params length ({len(user_sampling_params)}) must match "
        f"prompts length ({len(prompts)})"
    )

    output = BatchOutput()
    start_time = time.perf_counter()

    torch.get_device_module().reset_peak_memory_stats()

    for prompt, params in zip(prompts, user_sampling_params):
        try:
            sampling_params_kwargs = dict(params)
            sampling_params_kwargs["prompt"] = prompt
            result = engine.generate(sampling_params_kwargs=sampling_params_kwargs)

            if result is not None:
                if isinstance(result, list):
                    output.total_frames += len(result)
                else:
                    output.total_frames += 1
            output.num_samples += 1
        except Exception as e:
            logger.error(f"Generation failed for prompt '{prompt[:50]}...': {e}")
            output.error = str(e)

    output.latency = time.perf_counter() - start_time
    output.latency_per_sample = output.latency / len(prompts) if prompts else 0.0
    output.success = output.num_samples > 0
    output.peak_memory_mb = torch.get_device_module().max_memory_allocated() / (
        1024 * 1024
    )

    logger.debug(
        f"Batch generated: {output.num_samples}/{len(prompts)} samples in {output.latency:.2f}s"
    )

    return output


def calculate_metrics(
    outputs: List[BatchOutput],
    total_duration: float,
    resolution: Tuple[int, int, int],
    num_requests: int,
    all_sampling_params: Optional[List[Dict[str, Any]]] = None,
) -> Dict[str, Any]:
    """Calculate generation-specific throughput metrics."""
    successful = [o for o in outputs if o.success]
    num_success = sum(o.num_samples for o in successful)
    total_frames = sum(o.total_frames for o in successful)
    peak_memory = max((o.peak_memory_mb for o in outputs), default=0)

    width, height, frames = resolution
    if all_sampling_params:
        total_pixels = sum(
            p.get("width", width)
            * p.get("height", height)
            * p.get("num_frames", frames)
            for p in all_sampling_params[:num_success]
        )
    else:
        total_pixels = num_success * width * height * frames

    metrics = {
        "num_requests": num_requests,
        "successful_requests": num_success,
        "failed_requests": num_requests - num_success,
        "total_duration_seconds": total_duration,
        "total_frames_generated": total_frames,
        "total_pixels_generated": total_pixels,
        "images_per_second": num_success / total_duration if total_duration > 0 else 0,
        "frames_per_second": total_frames / total_duration if total_duration > 0 else 0,
        "megapixels_per_second": (
            total_pixels / (total_duration * 1e6) if total_duration > 0 else 0
        ),
        "requests_per_second": (
            num_success / total_duration if total_duration > 0 else 0
        ),
        "latency_per_request_seconds": (
            total_duration / num_success if num_success > 0 else 0
        ),
        "peak_memory_mb": peak_memory,
    }

    return metrics


def throughput_test(
    server_args: ServerArgs,
    bench_args: BenchArgs,
) -> Dict[str, Any]:
    """Main throughput benchmark function."""
    configure_logger(server_args=server_args)
    logger.info("Starting offline throughput benchmark...")

    engine = initialize_engine(server_args)

    if bench_args.random_request_config and bench_args.dataset != "random":
        raise ValueError(
            "--random-request-config can only be used with --dataset random"
        )

    if bench_args.num_outputs_per_prompt != 1:
        raise ValueError(
            "bench_offline_throughput currently supports only --num-outputs-per-prompt 1"
        )

    logger.info(f"Loading {bench_args.dataset} dataset...")
    if bench_args.dataset == "vbench":
        bench_args.task_name = str(engine.server_args.pipeline_config.task_type)
        dataset = VBenchDataset(bench_args)
    elif bench_args.dataset == "random":
        dataset = RandomDataset(bench_args)
    else:
        raise ValueError(f"Unknown dataset: {bench_args.dataset}")

    _sampling_params = {
        "guidance_scale": bench_args.guidance_scale,
        "num_inference_steps": bench_args.num_inference_steps,
        "height": bench_args.height,
        "width": bench_args.width,
        "num_frames": bench_args.num_frames,
        "num_outputs_per_prompt": bench_args.num_outputs_per_prompt,
        "seed": bench_args.seed,
        "profile": bench_args.profile,
        "num_profiled_timesteps": bench_args.num_profiled_timesteps,
        "profile_all_stages": bench_args.profile_all_stages,
    }
    if bench_args.disable_safety_checker:
        _sampling_params["safety_checker"] = None

    total_count = min(bench_args.num_prompts, len(dataset))
    all_prompts = [dataset[i].prompt for i in range(total_count)]

    if bench_args.random_request_config:
        all_sampling_params = []
        for i in range(total_count):
            params = dict(_sampling_params)
            params.update(dataset.get_sampling_params(i))
            all_sampling_params.append(params)
    else:
        all_sampling_params = [_sampling_params] * total_count

    if not bench_args.skip_warmup:
        logger.info("Running warmup batch...")
        warmup_count = min(bench_args.batch_size, total_count)
        warmup_prompts = all_prompts[:warmup_count]
        warmup_sampling_params = [
            {**p, "profile": False} for p in all_sampling_params[:warmup_count]
        ]
        generate_batch(engine, bench_args, warmup_prompts, warmup_sampling_params)

    logger.info(f"Running benchmark with {bench_args.num_prompts} prompts...")
    outputs: List[BatchOutput] = []

    start_time = time.perf_counter()

    num_batches = (total_count + bench_args.batch_size - 1) // bench_args.batch_size
    pbar = tqdm(
        total=num_batches,
        disable=bench_args.disable_tqdm,
        desc="Benchmark",
    )

    for batch_start in range(0, total_count, bench_args.batch_size):
        batch_end = min(batch_start + bench_args.batch_size, total_count)
        batch_prompts = all_prompts[batch_start:batch_end]
        batch_sampling_params = all_sampling_params[batch_start:batch_end]

        batch_output = generate_batch(
            engine, bench_args, batch_prompts, batch_sampling_params
        )
        outputs.append(batch_output)

        pbar.update(1)

    pbar.close()
    total_duration = time.perf_counter() - start_time

    resolution = (bench_args.width, bench_args.height, bench_args.num_frames)
    metrics = calculate_metrics(
        outputs,
        total_duration,
        resolution=resolution,
        num_requests=total_count,
        all_sampling_params=all_sampling_params,
    )

    display_results(
        metrics,
        bench_args,
        model_path=server_args.model_path,
    )

    if bench_args.output_file:
        save_results(metrics, bench_args, server_args)

    return metrics


def display_results(
    metrics: Dict[str, Any],
    bench_args: BenchArgs,
    model_path: str,
):
    """Display benchmark results in console."""
    print(
        "\n{s:{c}^{n}}".format(s=" Offline Throughput Benchmark Result ", n=110, c="=")
    )
    print_value_formatted("Model:", model_path)
    print_value_formatted("Dataset:", bench_args.dataset)
    print_value_formatted(
        "Resolution:",
        f"{bench_args.width}x{bench_args.height}x{bench_args.num_frames}",
    )
    print_value_formatted("Num Inference Steps:", bench_args.num_inference_steps)
    print_divider(75)
    print_value_formatted("Total Requests:", metrics["num_requests"])
    print_value_formatted("Successful Requests:", metrics["successful_requests"])
    print_value_formatted("Failed Requests:", metrics["failed_requests"])
    print_value_formatted(
        "Total Duration (seconds):", metrics["total_duration_seconds"]
    )
    print_divider(75)
    print_value_formatted("Frames Generated:", metrics["total_frames_generated"])
    print_value_formatted(
        "Megapixels Generated:", metrics["total_pixels_generated"] / 1e6
    )
    print_divider(75)
    print_value_formatted(
        "Frame Throughput (frames/sec):", metrics["frames_per_second"]
    )
    print_value_formatted("MP Throughput (MP/sec):", metrics["megapixels_per_second"])
    print_value_formatted("Requests Per Second:", metrics["requests_per_second"])
    print_value_formatted(
        "Latency Per Request (sec):", metrics["latency_per_request_seconds"]
    )
    print_value_formatted("Peak Memory (MB):", metrics["peak_memory_mb"])
    print_divider(110, "=")


def save_results(
    metrics: Dict[str, Any],
    bench_args: BenchArgs,
    server_args: ServerArgs,
):
    """Save benchmark results to JSON file."""
    result = {
        "metadata": {
            "timestamp": time.strftime("%Y-%m-%dT%H:%M:%S"),
            "model_path": server_args.model_path,
            "task_type": bench_args.task_name,
            "backend": "engine",
        },
        "configuration": {
            "num_inference_steps": bench_args.num_inference_steps,
            "guidance_scale": bench_args.guidance_scale,
            "seed": bench_args.seed,
            "batch_size": bench_args.batch_size,
            "num_prompts": bench_args.num_prompts,
            "resolution": f"{bench_args.width}x{bench_args.height}x{bench_args.num_frames}",
            "dataset": bench_args.dataset,
        },
        "results": metrics,
    }

    with open(bench_args.output_file, "a") as f:
        f.write(json.dumps(result) + "\n")

    logger.info(f"Results saved to {bench_args.output_file}")


def main():
    """Main entry point."""
    parser = argparse.ArgumentParser(
        description="Offline throughput benchmark for multimodal generation models"
    )

    ServerArgs.add_cli_args(parser)
    BenchArgs.add_cli_args(parser)

    args, unknown_args = parser.parse_known_args()

    server_args = ServerArgs.from_cli_args(args, unknown_args)
    bench_args = BenchArgs.from_cli_args(args)

    set_global_server_args(server_args)

    result = throughput_test(server_args, bench_args)

    return result


if __name__ == "__main__":
    main()