[Refactor][CI] Reorganize test tree, and move slow e2e to weekly (#4951)

Signed-off-by: wangyu <410167048@qq.com>

tests/diffusion/distributed/test_tensor_parallel.py

@@ -0,0 +1,185 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""DiT tensor-parallel parity and perf checks for diffusion models.
+
+Compares TP=2 against TP=1 on output correctness, latency, and peak memory.
+Uses ``Tongyi-MAI/Z-Image-Turbo`` as the parity model.
+"""
+
+from __future__ import annotations
+
+import os
+import time
+from pathlib import Path
+
+import numpy as np
+import pytest
+import torch
+from PIL import Image
+
+from tests.helpers.env import DeviceMemoryMonitor
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniRunner
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.outputs import OmniRequestOutput
+from vllm_omni.platforms import current_omni_platform
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
+PROMPT = "a photo of a cat sitting on a laptop keyboard"
+DEFAULT_MODEL = "Tongyi-MAI/Z-Image-Turbo"
+
+
+def _pil_to_float_rgb_tensor(img: Image.Image) -> torch.Tensor:
+    arr = np.asarray(img.convert("RGB"), dtype=np.float32) / 255.0
+    return torch.from_numpy(arr)
+
+
+def _diff_metrics(a: Image.Image, b: Image.Image) -> tuple[float, float]:
+    ta = _pil_to_float_rgb_tensor(a)
+    tb = _pil_to_float_rgb_tensor(b)
+    assert ta.shape == tb.shape, f"Image shapes differ: {ta.shape} vs {tb.shape}"
+    abs_diff = torch.abs(ta - tb)
+    p99_abs_diff = torch.quantile(abs_diff.flatten(), 0.99).item()
+    return abs_diff.mean().item(), p99_abs_diff
+
+
+def _extract_single_image(outputs) -> Image.Image:
+    first_output = outputs[0]
+    assert first_output.final_output_type == "image"
+    if not hasattr(first_output, "request_output") or not first_output.request_output:
+        raise ValueError("No request_output found in OmniRequestOutput")
+
+    req_out = first_output.request_output
+    if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
+        raise ValueError("Invalid request_output structure or missing 'images' key")
+
+    images = req_out.images
+    if images is None or len(images) != 1:
+        raise ValueError(f"Expected 1 image, got {0 if images is None else len(images)}")
+    return images[0]
+
+
+def _run_generate(
+    *,
+    tp_size: int,
+    height: int,
+    width: int,
+    num_inference_steps: int,
+    seed: int,
+    enforce_eager: bool,
+    num_requests: int = 4,
+) -> tuple[Image.Image, float, float]:
+    if num_requests < 2:
+        raise ValueError("num_requests must be >= 2 (1 warmup + >=1 timed)")
+
+    current_omni_platform.empty_cache()
+    device_index = current_omni_platform.current_device()
+    monitor = DeviceMemoryMonitor(device_index=device_index, interval=0.02)
+    monitor.start()
+    try:
+        with OmniRunner(
+            DEFAULT_MODEL,
+            parallel_config=DiffusionParallelConfig(tensor_parallel_size=tp_size),
+            enforce_eager=enforce_eager,
+        ) as runner:
+            gen = runner.omni.generate(
+                [PROMPT] * num_requests,
+                OmniDiffusionSamplingParams(
+                    height=height,
+                    width=width,
+                    num_inference_steps=num_inference_steps,
+                    guidance_scale=0.0,
+                    seed=seed,
+                    num_outputs_per_prompt=1,
+                ),
+                py_generator=True,
+            )
+
+            warmup_output = next(gen)
+
+            t_prev = time.perf_counter()
+            per_request_times_s: list[float] = []
+            last_output = warmup_output
+            for _ in range(num_requests - 1):
+                last_output = next(gen)
+                t_now = time.perf_counter()
+                per_request_times_s.append(t_now - t_prev)
+                t_prev = t_now
+
+            for _ in gen:
+                pass
+
+            median_time_s = float(np.median(per_request_times_s))
+            peak_memory_mb = monitor.peak_used_mb
+            return _extract_single_image([last_output]), median_time_s, peak_memory_mb
+    finally:
+        monitor.stop()
+
+
+@pytest.mark.advanced_model
+@pytest.mark.diffusion
+@pytest.mark.parallel
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2})
+def test_tensor_parallel_tp2(tmp_path: Path):
+    if current_omni_platform.is_npu():
+        pytest.skip("Tensor parallel parity test is only supported on CUDA and ROCm for now.")
+    if not current_omni_platform.is_available() or current_omni_platform.device_count() < 2:
+        pytest.skip("Tensor parallel TP=2 requires >= 2 devices.")
+
+    model_name = DEFAULT_MODEL
+    enforce_eager = False
+    height = 512
+    width = 512
+    num_inference_steps = 2
+    seed = 42
+
+    tp1_img, tp1_time_s, tp1_peak_mem = _run_generate(
+        tp_size=1,
+        height=height,
+        width=width,
+        num_inference_steps=num_inference_steps,
+        seed=seed,
+        enforce_eager=enforce_eager,
+    )
+    tp2_img, tp2_time_s, tp2_peak_mem = _run_generate(
+        tp_size=2,
+        height=height,
+        width=width,
+        num_inference_steps=num_inference_steps,
+        seed=seed,
+        enforce_eager=enforce_eager,
+    )
+
+    tp1_path = tmp_path / "tp1.png"
+    tp2_path = tmp_path / "tp2.png"
+    tp1_img.save(tp1_path)
+    tp2_img.save(tp2_path)
+
+    assert tp1_img.width == width and tp1_img.height == height
+    assert tp2_img.width == width and tp2_img.height == height
+
+    mean_abs_diff, p99_abs_diff = _diff_metrics(tp1_img, tp2_img)
+    mean_threshold = 3e-2
+    p99_threshold = 2.5e-1
+    print(
+        f"{model_name} TP image diff (TP=1 vs TP=2): "
+        f"mean_abs_diff={mean_abs_diff:.6e}, p99_abs_diff={p99_abs_diff:.6e}; "
+        f"thresholds: mean<={mean_threshold:.6e}, p99<={p99_threshold:.6e}; "
+        f"tp1_img={tp1_path}, tp2_img={tp2_path}"
+    )
+    assert mean_abs_diff <= mean_threshold and p99_abs_diff <= p99_threshold, (
+        f"Image diff exceeded threshold: mean_abs_diff={mean_abs_diff:.6e}, p99_abs_diff={p99_abs_diff:.6e} "
+        f"(thresholds: mean<={mean_threshold:.6e}, p99<={p99_threshold:.6e})"
+    )
+
+    print(f"TP perf (lower is better): tp1_time_s={tp1_time_s:.6f}, tp2_time_s={tp2_time_s:.6f}")
+    if not current_omni_platform.is_rocm():
+        assert tp2_time_s < tp1_time_s, f"Expected TP=2 to be faster than TP=1 (tp1={tp1_time_s}, tp2={tp2_time_s})"
+
+    print(f"TP peak memory (MB): tp1_peak_mem={tp1_peak_mem:.2f}, tp2_peak_mem={tp2_peak_mem:.2f}")
+    assert tp2_peak_mem < tp1_peak_mem, (
+        f"Expected TP=2 to use less peak memory than TP=1 (tp1={tp1_peak_mem}, tp2={tp2_peak_mem})"
+    )