chore: import upstream snapshot with attribution

python/sglang/multimodal_gen/test/scripts/gen_diffusion_ci_outputs.py

@@ -0,0 +1,226 @@
+#!/usr/bin/env python3
+"""
+Generate diffusion CI outputs for consistency testing.
+
+This script reuses the CI test code by calling run_suite.py with SGLANG_GEN_GT=1,
+ensuring that GT generation uses exactly the same code path as CI tests.
+
+Usage:
+    python gen_diffusion_ci_outputs.py --suite 1-gpu --partition-id 0 --total-partitions 2 --out-dir ./output
+    python gen_diffusion_ci_outputs.py --suite 1-gpu --case-ids qwen_image_t2i flux_image_t2i --out-dir ./output
+    python gen_diffusion_ci_outputs.py --suite 1-gpu-b200 --out-dir ./output
+"""
+
+import argparse
+import os
+import sys
+from pathlib import Path
+
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+from sglang.multimodal_gen.test.run_suite import (
+    SUITES,
+    PartitionItem,
+    _maybe_pin_update_weights_model_pair,
+    get_case_est_time,
+    get_suite_files_rel,
+    parse_partition_plan,
+    partition_items_by_lpt,
+)
+from sglang.multimodal_gen.test.runner.pytest_runner import (
+    collect_test_items,
+    run_pytest,
+)
+
+logger = init_logger(__name__)
+
+
+def main():
+    """Main entry point."""
+    parser = argparse.ArgumentParser(description="Generate diffusion CI outputs")
+    parser.add_argument(
+        "--suite",
+        type=str,
+        choices=list(SUITES.keys()),
+        required=True,
+        help="Test suite to run (choices: " + ", ".join(list(SUITES.keys())) + ")",
+    )
+    parser.add_argument(
+        "--partition-id",
+        type=int,
+        required=False,
+        help="Partition ID for matrix partitioning (0-based)",
+    )
+    parser.add_argument(
+        "--total-partitions",
+        type=int,
+        required=False,
+        help="Total number of partitions",
+    )
+    parser.add_argument(
+        "--out-dir",
+        type=str,
+        required=True,
+        help="Output directory for generated files",
+    )
+    parser.add_argument(
+        "--continue-on-error",
+        action="store_true",
+        help="Continue processing other cases if one fails",
+    )
+    parser.add_argument(
+        "--case-ids",
+        type=str,
+        nargs="*",
+        required=False,
+        help="Specific case IDs to run (space-separated). If provided, only these cases will be run.",
+    )
+    parser.add_argument(
+        "--partition-plan-json",
+        type=str,
+        required=False,
+        help="Full partition plan JSON for the current suite.",
+    )
+
+    args = parser.parse_args()
+
+    # Validate partition arguments
+    if args.partition_id is not None and args.total_partitions is not None:
+        if args.partition_id < 0 or args.partition_id >= args.total_partitions:
+            parser.error(f"partition-id must be in range [0, {args.total_partitions})")
+    elif args.partition_id is not None or args.total_partitions is not None:
+        parser.error(
+            "Both --partition-id and --total-partitions must be provided together"
+        )
+    if args.partition_plan_json and (
+        args.partition_id is None or args.total_partitions is None
+    ):
+        parser.error("--partition-plan-json requires partition-id and total-partitions")
+
+    # Create output directory
+    out_dir = Path(args.out_dir)
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    # Set environment variables for GT generation mode
+    os.environ["SGLANG_GEN_GT"] = "1"
+    os.environ["SGLANG_GT_OUTPUT_DIR"] = str(out_dir.absolute())
+    os.environ["SGLANG_SKIP_CONSISTENCY"] = (
+        "1"  # Skip consistency checks in GT gen mode
+    )
+
+    logger.info("GT generation mode enabled")
+    logger.info(f"Output directory: {out_dir}")
+
+    # Resolve test files path (same as run_suite.py)
+    current_file_path = Path(__file__).resolve()
+    test_root_dir = current_file_path.parent.parent  # scripts -> test
+    target_dir = test_root_dir / "server"
+
+    # GT generation only runs DiffusionTestCase parametrized cases. Standalone
+    # server tests such as disagg validate behavior but do not produce GT images.
+    suite_files_rel = get_suite_files_rel(args.suite, parametrized_only=True)
+
+    _maybe_pin_update_weights_model_pair(suite_files_rel)
+    suite_files_abs = []
+    for f_rel in suite_files_rel:
+        f_abs = target_dir / f_rel
+        if not f_abs.exists():
+            logger.warning(f"Test file {f_rel} not found in {target_dir}. Skipping.")
+            continue
+        suite_files_abs.append(str(f_abs))
+
+    if not suite_files_abs:
+        logger.error(f"No valid test files found for suite '{args.suite}'.")
+        sys.exit(1)
+
+    partition_id = args.partition_id if args.partition_id is not None else 0
+    total_partitions = args.total_partitions if args.total_partitions is not None else 1
+
+    selected_plan_case_ids = None
+    if args.partition_plan_json:
+        assignment = parse_partition_plan(
+            suite=args.suite,
+            partition_id=partition_id,
+            total_partitions=total_partitions,
+            plan_json=args.partition_plan_json,
+        )
+        selected_plan_case_ids = assignment.case_ids
+        if args.case_ids:
+            requested_case_ids = set(args.case_ids)
+            selected_plan_case_ids = [
+                case_id
+                for case_id in selected_plan_case_ids
+                if case_id in requested_case_ids
+            ]
+        if not selected_plan_case_ids:
+            logger.warning("No testcase cases assigned to this partition.")
+            sys.exit(0)
+
+    # Build pytest filter for case_ids if provided.
+    filter_expr = None
+    if selected_plan_case_ids is not None:
+        filters = [
+            f"test_diffusion_generation[{case_id}]"
+            for case_id in selected_plan_case_ids
+        ]
+        filter_expr = " or ".join(filters)
+        logger.info(f"Filtering by partition plan case IDs: {selected_plan_case_ids}")
+    elif args.case_ids:
+        # pytest parametrized test format: test_diffusion_generation[case_id]
+        filters = [f"test_diffusion_generation[{case_id}]" for case_id in args.case_ids]
+        filter_expr = " or ".join(filters)
+        logger.info(f"Filtering by case IDs: {args.case_ids}")
+
+    # Collect all test items and keep only testcase-based GT generators.
+    all_test_items = collect_test_items(suite_files_abs, filter_expr=filter_expr)
+    all_test_items = [
+        item for item in all_test_items if "test_diffusion_generation[" in item
+    ]
+
+    if not all_test_items:
+        logger.warning(f"No test items found for suite '{args.suite}'.")
+        sys.exit(0)
+
+    if selected_plan_case_ids is not None:
+        selected_case_id_set = set(selected_plan_case_ids)
+        my_items = [
+            item
+            for item in all_test_items
+            if item[item.index("[") + 1 : item.rindex("]")] in selected_case_id_set
+        ]
+    else:
+        # Partition by test items with the same LPT strategy used by CI partitioning.
+        partition_items = []
+        for item in all_test_items:
+            case_id = item[item.index("[") + 1 : item.rindex("]")]
+            partition_items.append(
+                PartitionItem(
+                    kind="case",
+                    item_id=item,
+                    est_time=get_case_est_time(case_id),
+                )
+            )
+
+        partitions = partition_items_by_lpt(partition_items, total_partitions)
+        my_items = [item.item_id for item in partitions[partition_id]]
+
+    logger.info(
+        f"Partition {partition_id}/{total_partitions}: "
+        f"running {len(my_items)} of {len(all_test_items)} test items"
+    )
+
+    if not my_items:
+        logger.warning("No items assigned to this partition. Exiting success.")
+        sys.exit(0)
+
+    # Run pytest with the specific test items (same as run_suite.py)
+    exit_code, _, _ = run_pytest(my_items)
+
+    if exit_code != 0:
+        if args.continue_on_error:
+            logger.warning(f"pytest exited with code {exit_code}")
+        else:
+            sys.exit(exit_code)
+
+
+if __name__ == "__main__":
+    main()

python/sglang/multimodal_gen/test/scripts/gen_perf_baselines.py

@@ -0,0 +1,211 @@
+import argparse
+import inspect
+import json
+import os
+import re
+import sys
+from pathlib import Path
+
+from openai import OpenAI
+
+from sglang.multimodal_gen.test.server.test_server_utils import (
+    ServerManager,
+    get_generate_fn,
+)
+from sglang.multimodal_gen.test.server.testcase_configs import (
+    BASELINE_CONFIG,
+    DiffusionTestCase,
+)
+from sglang.multimodal_gen.test.test_utils import (
+    get_dynamic_server_port,
+    wait_for_req_perf_record,
+)
+
+
+def _all_cases() -> list[DiffusionTestCase]:
+    import sglang.multimodal_gen.test.server.testcase_configs as cfg
+
+    cases: list[DiffusionTestCase] = []
+    for _, v in inspect.getmembers(cfg):
+        if isinstance(v, list) and v and isinstance(v[0], DiffusionTestCase):
+            cases.extend(v)
+
+    seen: set[str] = set()
+    out: list[DiffusionTestCase] = []
+    for c in cases:
+        if c.id not in seen:
+            seen.add(c.id)
+            out.append(c)
+    return out
+
+
+def _baseline_path() -> Path:
+    import sglang.multimodal_gen.test.server.testcase_configs as cfg
+
+    return cfg.get_perf_baseline_path()
+
+
+def _openai_client(port: int) -> OpenAI:
+    return OpenAI(api_key="sglang-anything", base_url=f"http://localhost:{port}/v1")
+
+
+def _build_server_extra_args(case: DiffusionTestCase) -> str:
+    server_args = case.server_args
+    a = os.environ.get("SGLANG_TEST_SERVE_ARGS", "")
+    a += f" --num-gpus {server_args.num_gpus}"
+    if server_args.tp_size is not None:
+        a += f" --tp-size {server_args.tp_size}"
+    if server_args.ulysses_degree is not None:
+        a += f" --ulysses-degree {server_args.ulysses_degree}"
+    if server_args.dit_layerwise_offload:
+        a += " --dit-layerwise-offload true"
+    if server_args.dit_offload_prefetch_size:
+        a += f" --dit-offload-prefetch-size {server_args.dit_offload_prefetch_size}"
+    if server_args.text_encoder_cpu_offload:
+        a += " --text-encoder-cpu-offload"
+    if server_args.ring_degree is not None:
+        a += f" --ring-degree {server_args.ring_degree}"
+    if server_args.lora_path:
+        a += f" --lora-path {server_args.lora_path}"
+
+    # default warmup
+    a += " --warmup"
+
+    for extra_arg in server_args.extras:
+        a += f" {extra_arg}"
+    return a
+
+
+def _build_env_vars(case: DiffusionTestCase) -> dict[str, str]:
+    if case.server_args.enable_cache_dit:
+        return {"SGLANG_CACHE_DIT_ENABLED": "true"}
+    return {}
+
+
+def _torch_cleanup() -> None:
+    try:
+        import gc
+
+        gc.collect()
+    except Exception:
+        pass
+    try:
+        import torch
+
+        if torch.get_device_module().is_available():
+            torch.get_device_module().synchronize()
+            torch.get_device_module().empty_cache()
+    except Exception:
+        pass
+
+
+def _run_case(case: DiffusionTestCase) -> dict:
+    default_port = get_dynamic_server_port()
+    port = int(os.environ.get("SGLANG_TEST_SERVER_PORT", default_port))
+    mgr = ServerManager(
+        model=case.server_args.model_path,
+        port=port,
+        wait_deadline=float(os.environ.get("SGLANG_TEST_WAIT_SECS", "1200")),
+        extra_args=_build_server_extra_args(case),
+        env_vars=_build_env_vars(case),
+    )
+    ctx = mgr.start()
+    try:
+        sp = case.sampling_params
+        client = _openai_client(ctx.port)
+        gen = get_generate_fn(
+            model_path=case.server_args.model_path,
+            modality=case.server_args.modality,
+            sampling_params=sp,
+        )
+        rid, _ = gen(case.id, client)
+        rec = wait_for_req_perf_record(
+            rid,
+            ctx.perf_log_path,
+            timeout=float(os.environ.get("SGLANG_PERF_TIMEOUT", "300")),
+        )
+        if rec is None:
+            raise RuntimeError(f"missing perf record: {case.id}")
+        from sglang.multimodal_gen.test.server.testcase_configs import (
+            PerformanceSummary,
+        )
+
+        perf = PerformanceSummary.from_req_perf_record(
+            rec, BASELINE_CONFIG.step_fractions
+        )
+        if case.server_args.modality == "video" and sp.num_frames and sp.num_frames > 0:
+            if "per_frame_generation" not in perf.stage_metrics:
+                perf.stage_metrics["per_frame_generation"] = perf.e2e_ms / sp.num_frames
+
+        return {
+            "stages_ms": {k: round(v, 2) for k, v in perf.stage_metrics.items()},
+            "denoise_step_ms": {
+                str(k): round(v, 2) for k, v in perf.all_denoise_steps.items()
+            },
+            "expected_e2e_ms": round(perf.e2e_ms, 2),
+            "expected_avg_denoise_ms": round(perf.avg_denoise_ms, 2),
+            "expected_median_denoise_ms": round(perf.median_denoise_ms, 2),
+        }
+    finally:
+        ctx.cleanup()
+
+
+def main() -> int:
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--baseline", default="")
+    ap.add_argument("--out", default="")
+    ap.add_argument("--match", default="")
+    ap.add_argument("--case", action="append", default=[])
+    ap.add_argument("--all-from-baseline", action="store_true")
+    ap.add_argument("--timeout", type=float, default=300.0)
+    args = ap.parse_args()
+
+    os.environ.setdefault("SGLANG_GEN_BASELINE", "1")
+    os.environ["SGLANG_PERF_TIMEOUT"] = str(args.timeout)
+
+    baseline_path = Path(args.baseline) if args.baseline else _baseline_path()
+    out_path = Path(args.out) if args.out else baseline_path
+    data = json.loads(baseline_path.read_text(encoding="utf-8"))
+    scenarios = data.setdefault("scenarios", {})
+
+    ids = set(args.case) if args.case else None
+    pat = re.compile(args.match) if args.match else None
+    if args.all_from_baseline:
+        ids = set(scenarios.keys())
+        pat = None
+
+    all_cases = _all_cases()
+    cases = []
+    for c in all_cases:
+        if ids and c.id not in ids:
+            continue
+        if pat and not pat.search(c.id):
+            continue
+        cases.append(c)
+
+    if args.all_from_baseline and ids:
+        case_ids = {c.id for c in all_cases}
+        missing = sorted([i for i in ids if i not in case_ids])
+        if missing:
+            sys.stderr.write(f"missing cases in testcase_configs.py: {len(missing)}\n")
+
+    if not cases:
+        return 0
+
+    for c in cases:
+        prev = scenarios.get(c.id, {})
+        note = prev.get("notes")
+        baseline = _run_case(c)
+        if note is not None:
+            baseline["notes"] = note
+        scenarios[c.id] = baseline
+        sys.stdout.write(f"{c.id}\n")
+        sys.stdout.flush()
+        _torch_cleanup()
+
+    out_path.write_text(json.dumps(data, indent=4) + "\n", encoding="utf-8")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())