[Test] Unify L2/L3 test layout, Buildkite steps, and test helpers (#2556)

Signed-off-by: wangyu <410167048@qq.com>
Signed-off-by: Huang, Zeyu <11222265+fhfuih@users.noreply.github.com>
Signed-off-by: wangyu <53896905+yenuo26@users.noreply.github.com>
Co-authored-by: Huang, Zeyu <11222265+fhfuih@users.noreply.github.com>

tests/diffusion/batching/test_diffusion_batching.py

@@ -0,0 +1,626 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+End-to-end test for diffusion batching via AsyncOmni.
+
+This test fires multiple concurrent ``AsyncOmni.generate()`` calls for a
+diffusion model and validates that every caller receives its correct
+individual result.  When the underlying diffusion stage is configured with
+``batch_size > 1`` (via stage config or ``StageDiffusionClient``), the
+requests will be batched internally.
+
+Even without explicit batching config this test is useful for verifying
+that concurrent async requests are handled correctly.
+
+Usage (standalone):
+
+    python tests/diffusion/batching/test_diffusion_batching.py \
+        --model <model_name_or_path> \
+        --num-prompts 8
+
+Or via pytest:
+
+    pytest tests/diffusion/batching/test_diffusion_batching.py -s
+"""
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+import sys
+import time
+import uuid
+from pathlib import Path
+
+import pytest
+import torch
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniRunner
+from vllm_omni.entrypoints.async_omni import AsyncOmni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.outputs import OmniRequestOutput
+from vllm_omni.platforms import current_omni_platform
+
+# ruff: noqa: E402
+REPO_ROOT = Path(__file__).resolve().parents[3]
+if str(REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(REPO_ROOT))
+
+
+# ------------------------------------------------------------------
+models = ["tiny-random/Qwen-Image"]
+
+# ------------------------------------------------------------------
+# Prompt fixtures
+# ------------------------------------------------------------------
+
+WARMUP_PROMPTS: list[dict[str, str]] = [
+    {"prompt": "a sunflower in a glass vase", "negative_prompt": "blurry"},
+    {"prompt": "a rocket launching into space", "negative_prompt": "low detail"},
+    {"prompt": "a small cottage in the snowy mountains", "negative_prompt": "foggy"},
+    {"prompt": "a colorful parrot sitting on a tree branch", "negative_prompt": "low contrast"},
+]
+
+TEST_PROMPTS: list[dict[str, str]] = [
+    {"prompt": "a cup of coffee on a table", "negative_prompt": "low resolution"},
+    {"prompt": "a toy dinosaur on a sandy beach", "negative_prompt": "cinematic, realistic"},
+    {"prompt": "a futuristic city skyline at sunset", "negative_prompt": "blurry, foggy"},
+    {"prompt": "a bowl of fresh strawberries", "negative_prompt": "low detail"},
+    {"prompt": "a medieval knight standing in the rain", "negative_prompt": "modern clothing"},
+    {"prompt": "a cat wearing sunglasses lounging in a garden", "negative_prompt": "dark lighting"},
+    {"prompt": "a spaceship flying above a volcano", "negative_prompt": "low contrast"},
+    {"prompt": "a watercolor painting of a mountain lake", "negative_prompt": "photo, realistic"},
+]
+
+
+# ------------------------------------------------------------------
+# Helpers
+# ------------------------------------------------------------------
+
+
+def _default_sampling_params(**overrides) -> OmniDiffusionSamplingParams:
+    defaults = dict(
+        num_inference_steps=2,
+        width=256,
+        height=256,
+        guidance_scale=0.0,
+    )
+    defaults.update(overrides)
+    return OmniDiffusionSamplingParams(**defaults)
+
+
+def _default_sync_sampling_params(**overrides) -> OmniDiffusionSamplingParams:
+    """Create sampling params for the synchronous Omni.generate() API."""
+    defaults = dict(
+        num_inference_steps=2,
+        width=256,
+        height=256,
+        guidance_scale=0.0,
+        generator=torch.Generator(current_omni_platform.device_type).manual_seed(42),
+    )
+    defaults.update(overrides)
+    return OmniDiffusionSamplingParams(**defaults)
+
+
+async def _collect_generate(omni: AsyncOmni, prompt, request_id, sampling_params_list) -> OmniRequestOutput:
+    """Consume the AsyncOmni.generate() async generator and return the last output."""
+    last_output: OmniRequestOutput | None = None
+    async for output in omni.generate(
+        prompt=prompt,
+        request_id=request_id,
+        sampling_params_list=sampling_params_list,
+    ):
+        last_output = output
+    if last_output is None:
+        raise RuntimeError(f"No output received for request {request_id}")
+    return last_output
+
+
+def _extract_images(output: OmniRequestOutput) -> list:
+    """Extract images from an OmniRequestOutput, handling both direct
+    and nested request_output structures."""
+    if output.images:
+        return output.images
+    # When the output comes from the orchestrator pipeline, images may be
+    # nested inside request_output.
+    inner = getattr(output, "request_output", None)
+    if inner is not None and hasattr(inner, "images") and inner.images:
+        return inner.images
+    return []
+
+
+# ------------------------------------------------------------------
+# Warm-up (async)
+# ------------------------------------------------------------------
+
+
+async def warmup(omni: AsyncOmni, prompts: list[dict[str, str]]) -> None:
+    """Warm-up: send prompts in parallel to pre-load the model."""
+    print(f"🔥 Warming up with {len(prompts)} prompts ...")
+    sp = _default_sampling_params(num_inference_steps=2)
+    start = time.perf_counter()
+
+    tasks = [
+        _collect_generate(
+            omni,
+            prompt=p,
+            request_id=f"warmup-{i}-{uuid.uuid4().hex[:8]}",
+            sampling_params_list=[sp],
+        )
+        for i, p in enumerate(prompts)
+    ]
+    await asyncio.gather(*tasks)
+
+    elapsed = time.perf_counter() - start
+    print(f"   Warm-up done in {elapsed:.2f}s\n")
+
+
+# ------------------------------------------------------------------
+# Single (sequential) benchmark
+# ------------------------------------------------------------------
+
+
+async def run_single(omni: AsyncOmni, prompts: list[dict[str, str]]) -> float:
+    """Run prompts one-by-one sequentially."""
+    print(f"🧩 Running SINGLE (sequential) mode – {len(prompts)} prompts ...")
+    sp = _default_sampling_params()
+    total_start = time.perf_counter()
+
+    for i, prompt in enumerate(prompts):
+        start = time.perf_counter()
+        result = await _collect_generate(
+            omni,
+            prompt=prompt,
+            request_id=f"single-{i}-{uuid.uuid4().hex[:8]}",
+            sampling_params_list=[sp],
+        )
+        elapsed = time.perf_counter() - start
+        images = _extract_images(result)
+        print(f"   prompt {i}: {elapsed:.2f}s  ({len(images)} images)")
+
+    total = time.perf_counter() - total_start
+    print(f"   ✅ Total single-mode: {total:.2f}s\n")
+    return total
+
+
+# ------------------------------------------------------------------
+# Batch (parallel) benchmark — concurrent individual requests
+# ------------------------------------------------------------------
+
+
+async def run_batch(
+    omni: AsyncOmni,
+    prompts: list[dict[str, str]],
+    label: str = "batch",
+) -> float:
+    """Send all prompts concurrently via asyncio.gather (one request per prompt)."""
+    print(f"⚙️  Running {label.upper()} mode – {len(prompts)} prompts concurrently ...")
+    sp = _default_sampling_params()
+    start = time.perf_counter()
+
+    tasks = [
+        _collect_generate(
+            omni,
+            prompt=p,
+            request_id=f"{label}-{i}-{uuid.uuid4().hex[:8]}",
+            sampling_params_list=[sp],
+        )
+        for i, p in enumerate(prompts)
+    ]
+    results = await asyncio.gather(*tasks)
+
+    elapsed = time.perf_counter() - start
+    for i, result in enumerate(results):
+        images = _extract_images(result)
+        print(f"   prompt {i}: {len(images)} images, request_id={result.request_id}")
+
+    print(f"   ✅ Total {label} mode: {elapsed:.2f}s\n")
+    return elapsed
+
+
+# ------------------------------------------------------------------
+# Explicit batch — single generate() call with list of prompts
+# ------------------------------------------------------------------
+
+
+async def run_batch_explicit(
+    omni: AsyncOmni,
+    prompts: list[dict[str, str]],
+    label: str = "batch_explicit",
+) -> float:
+    """Send all prompts as a single batch via generate(prompt=[...]).
+
+    Passes a *list* of prompts so that all are processed in **one**
+    ``DiffusionEngine.step()`` call.  The orchestrator detects
+    ``isinstance(prompt, list)`` and routes to the batch path.
+
+    A single ``OmniRequestOutput`` is yielded containing ALL generated
+    images combined.
+    """
+    print(f"⚡ Running {label.upper()} mode – {len(prompts)} prompts in ONE engine call ...")
+    sp = _default_sampling_params()
+    request_id = f"{label}-{uuid.uuid4().hex[:8]}"
+    start = time.perf_counter()
+
+    result: OmniRequestOutput | None = None
+    async for output in omni.generate(
+        prompt=prompts,
+        request_id=request_id,
+        sampling_params_list=[sp],
+    ):
+        result = output
+
+    elapsed = time.perf_counter() - start
+    if result is not None:
+        images = _extract_images(result)
+        print(f"   Got {len(images)} images total from batch, request_id={result.request_id}")
+    else:
+        print("   ⚠️  No output received from batch generate()")
+
+    print(f"   ✅ Total {label} mode: {elapsed:.2f}s\n")
+    return elapsed
+
+
+# ------------------------------------------------------------------
+# Async validation helpers
+# ------------------------------------------------------------------
+
+
+async def validate_batch_explicit(omni: AsyncOmni, prompts: list[dict[str, str]]) -> None:
+    """Validate generate(prompt=[...]) returns a single result with all images."""
+    print(f"🔍 Validating batch generate() correctness with {len(prompts)} prompts ...")
+    sp = _default_sampling_params()
+    request_id = f"validate-batch-{uuid.uuid4().hex[:8]}"
+
+    result: OmniRequestOutput | None = None
+    async for output in omni.generate(
+        prompt=prompts,
+        request_id=request_id,
+        sampling_params_list=[sp],
+    ):
+        result = output
+
+    assert result is not None, "No output received from batch generate()"
+    images = _extract_images(result)
+    # Batch mode returns ONE output with ALL images combined
+    assert len(images) == len(prompts), f"Expected {len(prompts)} images (one per prompt), got {len(images)}"
+    assert result.request_id == request_id, f"Expected request_id={request_id}, got {result.request_id}"
+    print(f"   ✅ Batch returned {len(images)} images with correct request_id.\n")
+
+
+async def validate_concurrent(omni: AsyncOmni, prompts: list[dict[str, str]]) -> None:
+    """Validate that every concurrent request receives a distinct result
+    with its own request_id."""
+    print(f"🔍 Validating concurrent correctness with {len(prompts)} prompts ...")
+    sp = _default_sampling_params()
+
+    request_ids = [f"validate-{i}-{uuid.uuid4().hex[:8]}" for i in range(len(prompts))]
+
+    tasks = [
+        _collect_generate(omni, prompt=p, request_id=rid, sampling_params_list=[sp])
+        for p, rid in zip(prompts, request_ids)
+    ]
+    results = await asyncio.gather(*tasks)
+
+    assert len(results) == len(prompts), f"Expected {len(prompts)} results, got {len(results)}"
+
+    returned_ids = [r.request_id for r in results]
+    for rid in request_ids:
+        assert rid in returned_ids, f"Missing request_id {rid} in results"
+
+    print("   ✅ All request_ids matched, results count correct.\n")
+
+
+# ------------------------------------------------------------------
+# Single vs Parallel comparison (CLI only)
+# ------------------------------------------------------------------
+
+
+async def compare_single_vs_parallel(
+    model: str,
+    prompts: list[dict[str, str]],
+    batch_size: int = 1,
+) -> None:
+    """Run the same prompts sequentially then in parallel and print a comparison."""
+
+    omni = AsyncOmni(model=model, diffusion_batch_size=batch_size)
+    try:
+        await warmup(omni, WARMUP_PROMPTS)
+        single_time = await run_single(omni, prompts)
+        parallel_time = await run_batch(omni, prompts, label="parallel")
+        explicit_time = await run_batch_explicit(omni, prompts, label="batch_explicit")
+    finally:
+        omni.shutdown()
+
+    speedup_parallel = single_time / parallel_time if parallel_time > 0 else float("inf")
+    speedup_explicit = single_time / explicit_time if explicit_time > 0 else float("inf")
+    print("=" * 60)
+    print(f"📊 Summary ({len(prompts)} prompts)")
+    print(f"   Sequential        : {single_time:.2f}s")
+    print(f"   Parallel (gather) : {parallel_time:.2f}s  ({speedup_parallel:.2f}x)")
+    print(f"   Explicit batch    : {explicit_time:.2f}s  ({speedup_explicit:.2f}x)")
+    print("=" * 60)
+
+
+# ------------------------------------------------------------------
+# CLI main entrypoint
+# ------------------------------------------------------------------
+
+
+async def main(model: str, num_prompts: int, mode: str, batch_size: int = 1) -> None:
+    prompts = (TEST_PROMPTS * ((num_prompts // len(TEST_PROMPTS)) + 1))[:num_prompts]
+
+    if mode == "compare":
+        await compare_single_vs_parallel(model, prompts, batch_size=batch_size)
+        return
+
+    omni = AsyncOmni(model=model, diffusion_batch_size=batch_size)
+    try:
+        await warmup(omni, WARMUP_PROMPTS)
+
+        if mode == "validate":
+            await validate_concurrent(omni, prompts)
+        elif mode == "validate_batch":
+            await validate_batch_explicit(omni, prompts)
+        elif mode == "batch":
+            await run_batch(omni, prompts, label="measurement")
+        elif mode == "batch_explicit":
+            await run_batch_explicit(omni, prompts)
+        elif mode == "single":
+            await run_single(omni, prompts)
+        else:
+            raise ValueError(f"Unknown mode: {mode}")
+    finally:
+        omni.shutdown()
+
+
+# ==================================================================
+# pytest test cases
+# ==================================================================
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "L4", "rocm": "MI325", "xpu": "B60"})
+@pytest.mark.parametrize("model_name", models)
+def test_diffusion_batching_sync_sequential(model_name: str):
+    """Test that synchronous Omni can generate images for multiple prompts
+    submitted sequentially (one at a time) and each returns a valid image."""
+    try:
+        with OmniRunner(model_name) as runner:
+            m = runner.omni
+            sp = _default_sync_sampling_params()
+            prompts = TEST_PROMPTS[:4]
+
+            for i, prompt in enumerate(prompts):
+                outputs = m.generate(prompt, sp)
+                first_output = outputs[0]
+                assert first_output.final_output_type == "image", (
+                    f"Expected 'image', got '{first_output.final_output_type}'"
+                )
+
+                # Images are surfaced both at top-level and inside request_output
+                images = _extract_images(first_output)
+                assert len(images) >= 1, f"Expected at least 1 image for prompt {i}, got {len(images)}"
+                assert images[0].width == 256
+                assert images[0].height == 256
+                print(f"   prompt {i}: OK ({len(images)} images)")
+    except Exception as e:
+        print(f"Test failed with error: {e}")
+        raise
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "L4", "rocm": "MI325", "xpu": "B60"})
+@pytest.mark.parametrize("model_name", models)
+def test_diffusion_batching_sync_multi_prompt(model_name: str):
+    """Test that synchronous Omni correctly handles a list of multiple
+    prompts submitted at once and returns one result per prompt.
+
+    Note: Omni.generate() iterates the list and submits each prompt
+    individually with its own request_id.  This tests concurrent request
+    handling at the diffusion stage, not the explicit list-batch path
+    (which is only available via AsyncOmni).
+    """
+    try:
+        with OmniRunner(model_name) as runner:
+            m = runner.omni
+            sp = _default_sync_sampling_params()
+            prompts = TEST_PROMPTS[:4]
+
+            outputs = m.generate(prompts, sp)
+            assert len(outputs) == len(prompts), f"Expected {len(prompts)} outputs, got {len(outputs)}"
+
+            for i, output in enumerate(outputs):
+                assert output.final_output_type == "image", (
+                    f"Output {i} final_output_type expected 'image', got '{output.final_output_type}'"
+                )
+                images = _extract_images(output)
+                assert images and len(images) >= 1, f"Expected at least 1 image for prompt {i}"
+                assert images[0].width == 256
+                assert images[0].height == 256
+                print(f"   prompt {i}: OK ({len(images)} images, request_id={output.request_id})")
+
+            # Verify all request_ids are distinct
+            request_ids = [o.request_id for o in outputs]
+            assert len(set(request_ids)) == len(request_ids), f"Duplicate request_ids found: {request_ids}"
+    except Exception as e:
+        print(f"Test failed with error: {e}")
+        raise
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "L4", "rocm": "MI325", "xpu": "B60"})
+@pytest.mark.parametrize("model_name", models)
+def test_diffusion_batching_async_concurrent(model_name: str):
+    """Test that AsyncOmni correctly handles multiple concurrent requests
+    fired via asyncio.gather. Each request_id must appear in the results."""
+
+    async def _inner():
+        omni = AsyncOmni(model=model_name, diffusion_batch_size=1)
+        try:
+            prompts = TEST_PROMPTS[:4]
+            sp = _default_sampling_params()
+            request_ids = [f"async-concurrent-{i}-{uuid.uuid4().hex[:8]}" for i in range(len(prompts))]
+
+            tasks = [
+                _collect_generate(omni, prompt=p, request_id=rid, sampling_params_list=[sp])
+                for p, rid in zip(prompts, request_ids)
+            ]
+            results = await asyncio.gather(*tasks)
+
+            assert len(results) == len(prompts), f"Expected {len(prompts)} results, got {len(results)}"
+
+            returned_ids = [r.request_id for r in results]
+            for rid in request_ids:
+                assert rid in returned_ids, f"Missing request_id {rid} in results"
+
+            for i, result in enumerate(results):
+                images = _extract_images(result)
+                assert len(images) >= 1, f"No images for prompt {i}"
+                assert images[0].width == 256
+                assert images[0].height == 256
+                print(f"   prompt {i}: OK ({len(images)} images, request_id={result.request_id})")
+        finally:
+            omni.shutdown()
+
+    asyncio.run(_inner())
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "L4", "rocm": "MI325", "xpu": "B60"})
+@pytest.mark.parametrize("model_name", models)
+def test_diffusion_batching_async_explicit_batch(model_name: str):
+    """Test that AsyncOmni batch mode (generate(prompt=[...])) dispatches
+    all prompts in a single engine call and returns a single combined result.
+
+    The list-prompt path routes through the orchestrator's
+    ``add_batch_request_async`` → ``AsyncOmni.generate_batch``
+    and yields ONE ``OmniRequestOutput`` with ALL images combined.
+    """
+
+    async def _inner():
+        omni = AsyncOmni(model=model_name, diffusion_batch_size=4)
+        try:
+            prompts = TEST_PROMPTS[:4]
+            sp = _default_sampling_params()
+            request_id = f"explicit-batch-{uuid.uuid4().hex[:8]}"
+
+            # Batch mode: pass list of prompts → single request_id
+            result: OmniRequestOutput | None = None
+            async for output in omni.generate(
+                prompt=prompts,
+                request_id=request_id,
+                sampling_params_list=[sp],
+            ):
+                result = output
+
+            assert result is not None, "No output received from batch generate()"
+
+            images = _extract_images(result)
+            # One image per prompt, all in a single output
+            assert len(images) == len(prompts), f"Expected {len(prompts)} images in combined output, got {len(images)}"
+            assert result.request_id == request_id, f"Expected request_id={request_id}, got {result.request_id}"
+            for i, img in enumerate(images):
+                assert img.width == 256, f"Image {i} width mismatch"
+                assert img.height == 256, f"Image {i} height mismatch"
+            print(f"   ✅ Batch returned {len(images)} images, request_id={result.request_id}")
+        finally:
+            omni.shutdown()
+
+    asyncio.run(_inner())
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "L4", "rocm": "MI325", "xpu": "B60"})
+@pytest.mark.parametrize("model_name", models)
+def test_diffusion_batching_num_outputs(model_name: str):
+    """Test that the diffusion model respects num_outputs_per_prompt and
+    generates the correct number of images per request."""
+    try:
+        with OmniRunner(model_name) as runner:
+            m = runner.omni
+            num_outputs = 2
+            sp = _default_sync_sampling_params(num_outputs_per_prompt=num_outputs)
+
+            outputs = m.generate(
+                "a photo of a cat sitting on a laptop keyboard",
+                sp,
+            )
+
+            first_output = outputs[0]
+            assert first_output.final_output_type == "image"
+            images = _extract_images(first_output)
+            assert images is not None and len(images) == num_outputs, (
+                f"Expected {num_outputs} images, got {len(images) if images else 0}"
+            )
+            for img in images:
+                assert img.width == 256
+                assert img.height == 256
+    except Exception as e:
+        print(f"Test failed with error: {e}")
+        raise
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "L4", "rocm": "MI325", "xpu": "B60"})
+@pytest.mark.parametrize("model_name", models)
+def test_diffusion_batching_distinct_results(model_name: str):
+    """Test that different prompts produce distinct images when batched,
+    ensuring the batching logic does not mix up results across requests."""
+    try:
+        with OmniRunner(model_name) as runner:
+            m = runner.omni
+            sp = _default_sync_sampling_params()
+            prompts = [
+                {"prompt": "a bright red apple on a white table", "negative_prompt": "blurry"},
+                {"prompt": "a blue ocean with white waves crashing", "negative_prompt": "blurry"},
+            ]
+
+            outputs = m.generate(prompts, sp)
+            assert len(outputs) == len(prompts), f"Expected {len(prompts)} outputs, got {len(outputs)}"
+
+            # Verify each output has a unique request_id
+            request_ids = [o.request_id for o in outputs]
+            assert len(set(request_ids)) == len(request_ids), f"Duplicate request_ids: {request_ids}"
+
+            # Verify each output has images
+            for i, output in enumerate(outputs):
+                images = _extract_images(output)
+                assert images and len(images) >= 1, f"No images for prompt {i}"
+                assert images[0].width == 256
+                assert images[0].height == 256
+    except Exception as e:
+        print(f"Test failed with error: {e}")
+        raise
+
+
+# ------------------------------------------------------------------
+# CLI
+# ------------------------------------------------------------------
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="E2E diffusion concurrent benchmark / validation")
+    parser.add_argument("--model", type=str, required=True, help="Model name or path")
+    parser.add_argument("--num-prompts", type=int, default=8, help="Number of prompts to run")
+    parser.add_argument("--batch-size", type=int, default=1, help="Diffusion batch size (1 = no batching)")
+    parser.add_argument(
+        "--mode",
+        choices=["batch", "batch_explicit", "single", "compare", "validate", "validate_batch"],
+        default="compare",
+        help=(
+            "Run mode: 'batch' (parallel gather), 'batch_explicit' (list-prompt batch API), "
+            "'single' (sequential), 'compare' (all three), "
+            "'validate' (concurrent correctness), 'validate_batch' (list-prompt correctness)"
+        ),
+    )
+    args = parser.parse_args()
+
+    asyncio.run(main(args.model, args.num_prompts, args.mode, batch_size=args.batch_size))