[Feature] Profile the update weights from disk API of SGLang Diffusion

[zhaochenyang20]

Checklist

• If this is not a feature request but a general question, please start a discussion at https://github.com/sgl-project/sglang/discussions. Otherwise, it will be closed.
• Please use English. Otherwise, it will be closed.

Motivation

As we did in this comment:

#18306

We should profile the actual time breakdown in the update weights from disk.

Ideally speaking, 7B models' update should be within 1s (no considering save to disk time) in this https://github.com/zhaochenyang20/Awesome-ML-SYS-Tutorial/blob/main/sglang/latency-accelerate-for-weight-updates/readme.md

Related resources

No response

[Ratish1]

Hey @zhaochenyang20 , I will take up on this. Thanks

[Ratish1]

Hey @zhaochenyang20 ,

Code used

• Script: python/sglang/multimodal_gen/test/scripts/profile_update_weights_from_disk.py
• Command:
  python3 python/sglang/multimodal_gen/test/scripts/profile_update_weights_from_disk.py

Setup

• 1 GPU
• runs=5 (after 1 warmup), with_offload=false

Results (median / p95)

• update_weights_from_disk(flush_cache=True): 3.106s / 3.172s
• update_weights_from_disk(flush_cache=False): 3.036s / 3.272s
• target_modules=['text_encoder']: 1.048s / 1.084s
• target_modules=['vae']: 0.232s / 0.237s
• target_modules=['transformer']: 2.235s / 2.490s

Code:

"""Profile diffusion update_weights_from_disk API latency with external timing.

This script starts one diffusion server, runs selected API scenarios multiple
times, and writes a JSON report with per-run and aggregated latencies.
"""

from __future__ import annotations

import argparse
import json
import math
import os
import socket
import statistics
import sys
import time
from dataclasses import dataclass
from datetime import UTC, datetime
from pathlib import Path
from typing import Any

import requests

from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import maybe_download_model
from sglang.multimodal_gen.test.server.test_server_utils import ServerManager
from sglang.multimodal_gen.test.test_utils import get_dynamic_server_port


@dataclass
class Scenario:
    name: str
    endpoint: str
    payload: dict[str, Any]
    expected_status: int
    expected_success: bool | None


def _parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Profile diffusion update_weights_from_disk API latency."
    )
    parser.add_argument(
        "--serve-model",
        default="Qwen/Qwen-Image",
        help="Model loaded when server starts.",
    )
    parser.add_argument(
        "--runs",
        type=int,
        default=5,
        help="Measured runs per scenario.",
    )
    parser.add_argument(
        "--timeout",
        type=int,
        default=1800,
        help="HTTP timeout (seconds) for update/checksum requests.",
    )
    parser.add_argument(
        "--wait-secs",
        type=float,
        default=float(os.environ.get("SGLANG_TEST_WAIT_SECS", "1200")),
        help="Server startup wait deadline.",
    )
    parser.add_argument(
        "--num-gpus",
        type=int,
        default=1,
        help="Value for --num-gpus passed to sglang serve.",
    )
    parser.add_argument(
        "--port",
        type=int,
        default=0,
        help="Server port. If 0, auto-select dynamic port.",
    )
    parser.add_argument(
        "--with-offload",
        action="store_true",
        help="Enable --dit-layerwise-offload true for server startup.",
    )
    parser.add_argument(
        "--include-error-cases",
        action="store_true",
        help="Profile failure-path API calls (nonexistent model/module, missing path).",
    )
    parser.add_argument(
        "--include-checksum",
        action="store_true",
        help="Profile get_weights_checksum endpoint for transformer module.",
    )
    parser.add_argument(
        "--target-modules",
        default="text_encoder,vae,transformer",
        help=(
            "Comma-separated modules for dedicated target_modules timing "
            "scenarios (e.g. transformer,vae,text_encoder)."
        ),
    )
    parser.add_argument(
        "--output-json",
        default="update_weights_from_disk_profile_report.json",
        help="Path for JSON report output.",
    )
    return parser.parse_args()


def _round(x: float) -> float:
    return round(x, 3)


def _p95(values: list[float]) -> float:
    if not values:
        return 0.0
    sorted_values = sorted(values)
    idx = max(0, math.ceil(0.95 * len(sorted_values)) - 1)
    return sorted_values[idx]


def _summarize(latencies_s: list[float]) -> dict[str, float]:
    if not latencies_s:
        return {}
    return {
        "min_s": _round(min(latencies_s)),
        "max_s": _round(max(latencies_s)),
        "mean_s": _round(statistics.mean(latencies_s)),
        "median_s": _round(statistics.median(latencies_s)),
        "p95_s": _round(_p95(latencies_s)),
    }


def _timed_post(
    base_url: str,
    endpoint: str,
    payload: dict[str, Any],
    timeout: int,
) -> tuple[dict[str, Any], int, float]:
    t0 = time.perf_counter()
    response = requests.post(
        f"{base_url}{endpoint}",
        json=payload,
        timeout=timeout,
    )
    elapsed_s = time.perf_counter() - t0
    try:
        body = response.json()
    except Exception:
        body = {"raw_text": response.text}
    return body, response.status_code, elapsed_s


def _run_scenario(
    base_url: str,
    scenario: Scenario,
    runs: int,
    timeout: int,
) -> dict[str, Any]:
    records: list[dict[str, Any]] = []
    latencies_s: list[float] = []

    print(f"\n=== {scenario.name} ===")
    for idx in range(runs):
        body, status_code, latency_s = _timed_post(
            base_url=base_url,
            endpoint=scenario.endpoint,
            payload=scenario.payload,
            timeout=timeout,
        )
        success = body.get("success")
        print(
            f"[{idx + 1}/{runs}] status={status_code} success={success} "
            f"latency={latency_s:.3f}s"
        )

        if status_code != scenario.expected_status:
            raise RuntimeError(
                f"{scenario.name}: expected status {scenario.expected_status}, "
                f"got {status_code}, body={body}"
            )
        if (
            scenario.expected_success is not None
            and bool(success) != scenario.expected_success
        ):
            raise RuntimeError(
                f"{scenario.name}: expected success={scenario.expected_success}, "
                f"got {success}, body={body}"
            )

        records.append(
            {
                "status_code": status_code,
                "success": success,
                "latency_s": _round(latency_s),
                "message": body.get("message"),
            }
        )
        latencies_s.append(latency_s)

    summary = _summarize(latencies_s)
    print(
        f"{scenario.name} summary: median={summary.get('median_s', 0):.3f}s "
        f"p95={summary.get('p95_s', 0):.3f}s"
    )
    return {
        "name": scenario.name,
        "endpoint": scenario.endpoint,
        "payload": scenario.payload,
        "expected_status": scenario.expected_status,
        "expected_success": scenario.expected_success,
        "runs": records,
        "summary": summary,
    }


def _profile_checksum(
    base_url: str,
    runs: int,
    timeout: int,
) -> dict[str, Any]:
    scenario = Scenario(
        name="get_weights_checksum(module_names=['transformer'])",
        endpoint="/get_weights_checksum",
        payload={"module_names": ["transformer"]},
        expected_status=200,
        expected_success=None,
    )
    records: list[dict[str, Any]] = []
    latencies_s: list[float] = []

    print(f"\n=== {scenario.name} ===")
    for idx in range(runs):
        body, status_code, latency_s = _timed_post(
            base_url=base_url,
            endpoint=scenario.endpoint,
            payload=scenario.payload,
            timeout=timeout,
        )
        keys = sorted(body.keys())
        print(f"[{idx + 1}/{runs}] status={status_code} keys={keys} latency={latency_s:.3f}s")
        if status_code != scenario.expected_status:
            raise RuntimeError(
                f"{scenario.name}: expected status {scenario.expected_status}, "
                f"got {status_code}, body={body}"
            )
        records.append(
            {
                "status_code": status_code,
                "latency_s": _round(latency_s),
                "keys": keys,
            }
        )
        latencies_s.append(latency_s)

    summary = _summarize(latencies_s)
    print(
        f"{scenario.name} summary: median={summary.get('median_s', 0):.3f}s "
        f"p95={summary.get('p95_s', 0):.3f}s"
    )
    return {
        "name": scenario.name,
        "endpoint": scenario.endpoint,
        "payload": scenario.payload,
        "runs": records,
        "summary": summary,
    }


def main() -> int:
    args = _parse_args()
    # RL-style profiling path: update from disk using the same model path.
    update_model = args.serve_model
    port = args.port or get_dynamic_server_port()
    base_url = f"http://localhost:{port}"

    report: dict[str, Any] = {
        "metadata": {
            "timestamp_utc": datetime.now(UTC).isoformat(),
            "hostname": socket.gethostname(),
            "serve_model": args.serve_model,
            "update_model": args.serve_model,
            "runs": args.runs,
            "timeout_s": args.timeout,
            "wait_secs": args.wait_secs,
            "num_gpus": args.num_gpus,
            "with_offload": args.with_offload,
            "base_url": base_url,
            "target_modules": [
                module.strip()
                for module in args.target_modules.split(",")
                if module.strip()
            ],
        },
        "prepare": {},
        "scenarios": [],
    }

    maybe_download_times: dict[str, float] = {}
    t0 = time.perf_counter()
    maybe_download_model(args.serve_model)
    maybe_download_times[args.serve_model] = _round(time.perf_counter() - t0)
    print(
        f"maybe_download_model({args.serve_model}): "
        f"{maybe_download_times[args.serve_model]:.3f}s"
    )
    report["prepare"]["maybe_download_model_s"] = maybe_download_times

    extra_args = f"--num-gpus {args.num_gpus}"
    if args.with_offload:
        extra_args += " --dit-layerwise-offload true"
    manager = ServerManager(
        model=args.serve_model,
        port=port,
        wait_deadline=args.wait_secs,
        extra_args=extra_args,
    )

    t0 = time.perf_counter()
    ctx = manager.start()
    server_start_s = _round(time.perf_counter() - t0)
    report["prepare"]["server_start_s"] = server_start_s
    print(f"server_start: {server_start_s:.3f}s")

    try:
        warmup_scenario = Scenario(
            name="warmup_update_weights_from_disk",
            endpoint="/update_weights_from_disk",
            payload={"model_path": update_model, "flush_cache": True},
            expected_status=200,
            expected_success=True,
        )
        warmup_body, warmup_status, warmup_latency_s = _timed_post(
            base_url=base_url,
            endpoint=warmup_scenario.endpoint,
            payload=warmup_scenario.payload,
            timeout=args.timeout,
        )
        if warmup_status != 200 or not warmup_body.get("success", False):
            raise RuntimeError(
                f"Warmup failed: status={warmup_status}, body={warmup_body}"
            )
        report["prepare"]["warmup"] = {
            "status_code": warmup_status,
            "latency_s": _round(warmup_latency_s),
        }
        print(f"warmup_update_weights_from_disk: {warmup_latency_s:.3f}s")

        scenarios: list[Scenario] = [
            Scenario(
                name="update_weights_from_disk(flush_cache=True)",
                endpoint="/update_weights_from_disk",
                payload={"model_path": update_model, "flush_cache": True},
                expected_status=200,
                expected_success=True,
            ),
            Scenario(
                name="update_weights_from_disk(flush_cache=False)",
                endpoint="/update_weights_from_disk",
                payload={"model_path": update_model, "flush_cache": False},
                expected_status=200,
                expected_success=True,
            ),
        ]

        target_modules = [
            module.strip()
            for module in args.target_modules.split(",")
            if module.strip()
        ]
        for target_module in target_modules:
            scenarios.append(
                Scenario(
                    name=f"update_weights_from_disk(target_modules=['{target_module}'])",
                    endpoint="/update_weights_from_disk",
                    payload={
                        "model_path": update_model,
                        "flush_cache": True,
                        "target_modules": [target_module],
                    },
                    expected_status=200,
                    expected_success=True,
                )
            )

        if args.include_error_cases:
            scenarios.extend(
                [
                    Scenario(
                        name="update_weights_from_disk(nonexistent_model)",
                        endpoint="/update_weights_from_disk",
                        payload={
                            "model_path": "/nonexistent/path/to/model",
                            "flush_cache": True,
                        },
                        expected_status=400,
                        expected_success=False,
                    ),
                    Scenario(
                        name="update_weights_from_disk(missing_model_path)",
                        endpoint="/update_weights_from_disk",
                        payload={},
                        expected_status=400,
                        expected_success=False,
                    ),
                    Scenario(
                        name="update_weights_from_disk(nonexistent_module)",
                        endpoint="/update_weights_from_disk",
                        payload={
                            "model_path": update_model,
                            "flush_cache": True,
                            "target_modules": ["nonexistent_module"],
                        },
                        expected_status=400,
                        expected_success=False,
                    ),
                ]
            )

        for scenario in scenarios:
            scenario_report = _run_scenario(
                base_url=base_url,
                scenario=scenario,
                runs=args.runs,
                timeout=args.timeout,
            )
            report["scenarios"].append(scenario_report)

        if args.include_checksum:
            checksum_report = _profile_checksum(
                base_url=base_url,
                runs=args.runs,
                timeout=args.timeout,
            )
            report["scenarios"].append(checksum_report)
    finally:
        ctx.cleanup()

    out_path = Path(args.output_json)
    out_path.parent.mkdir(parents=True, exist_ok=True)
    out_path.write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8")
    print(f"\nreport written: {out_path}")

    print("\n=== Summary (median_s / p95_s) ===")
    for scenario in report["scenarios"]:
        summary = scenario["summary"]
        print(
            f"{scenario['name']}: "
            f"{summary.get('median_s', 0):.3f}s / {summary.get('p95_s', 0):.3f}s"
        )
    return 0


if __name__ == "__main__":
    sys.exit(main())

[zhaochenyang20]

Thanks ratish in providing this results:

update_weights_from_disk(flush_cache=True): 3.106s / 3.172s
update_weights_from_disk(flush_cache=False): 3.036s / 3.272s
target_modules=['text_encoder']: 1.048s / 1.084s
target_modules=['vae']: 0.232s / 0.237s
target_modules=['transformer']: 2.235s / 2.490s