CI: add server performance test for SGLang diffusion (#13091)

Co-authored-by: Mick <mickjagger19@icloud.com>

python/sglang/multimodal_gen/test/test_utils.py

@@ -1,5 +1,6 @@
 # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
 import dataclasses
+import json
 import os
 import shlex
 import socket
@@ -7,11 +8,13 @@ import subprocess
 import sys
 import time
 import unittest
-from typing import Optional
+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__)
@@ -52,11 +55,37 @@ def probe_port(host="127.0.0.1", port=30010, timeout=2.0) -> bool:
             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")
@@ -77,6 +106,113 @@ def check_image_size(ut, image, width, height):
     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