vllm-omni/tests/diffusion/test_diffusion_engine.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project

import asyncio
import queue
import threading
import time
from dataclasses import dataclass, field
from types import SimpleNamespace
from typing import Any

import pytest
import torch
from pytest_mock import MockerFixture

import vllm_omni.diffusion.diffusion_engine as diffusion_engine_module
from tests.helpers.mark import hardware_test
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
from vllm_omni.diffusion.diffusion_engine import (
    DiffusionEngine,
    DiffusionExecutionMode,
    _move_tensor_tree_to_cpu,
)
from vllm_omni.diffusion.request import OmniDiffusionRequest
from vllm_omni.diffusion.sched.interface import (
    CachedRequestData,
    NewRequestData,
    _AdmissionWaitDecision,
)
from vllm_omni.diffusion.sched.interface import (
    DiffusionSchedulerOutput as RealDiffusionSchedulerOutput,
)
from vllm_omni.inputs.data import OmniDiffusionSamplingParams

pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]


@dataclass
class DiffusionSchedulerOutput:
    step_id: int
    scheduled_new_reqs: list = field(default_factory=list)
    scheduled_cached_reqs: Any = None
    finished_req_ids: set = field(default_factory=set)
    num_running_reqs: int = 0
    num_waiting_reqs: int = 0

    @property
    def scheduled_request_ids(self):
        ids = [req.request_id for req in self.scheduled_new_reqs]
        if self.scheduled_cached_reqs:
            ids.extend(self.scheduled_cached_reqs.request_ids)
        return ids

    @property
    def is_empty(self):
        return len(self.scheduled_request_ids) == 0


class MockScheduler:
    def __init__(self):
        self._waiting_queue = []
        self._step_id = 0
        # Match RequestScheduler API used by request-batch admission wait.
        self.max_num_running_reqs = 5

    def add_request(self, request):
        self._waiting_queue.append(request)
        return request.request_id

    def has_requests(self):
        return len(self._waiting_queue) > 0

    def num_waiting_requests(self) -> int:
        return len(self._waiting_queue)

    def num_running_requests(self) -> int:
        return 0

    def get_admission_wait_decision(
        self,
        *,
        now: float,
        dp_concurrent: bool = False,
    ) -> _AdmissionWaitDecision:
        del dp_concurrent
        return _AdmissionWaitDecision(
            should_wait=True,
            deadline=now + 0.5,
            stable_window_s=0.05,
            max_batch=self.max_num_running_reqs,
        )

    def should_end_admission_wait(
        self,
        decision: _AdmissionWaitDecision,
        *,
        now: float,
        stable_since: float,
    ) -> bool:
        waiting = self.num_waiting_requests()
        return (
            waiting >= decision.max_batch
            or (waiting > 0 and now - stable_since >= decision.stable_window_s)
            or (decision.deadline is not None and now >= decision.deadline)
        )

    def schedule(self) -> DiffusionSchedulerOutput:
        if not self._waiting_queue:
            return DiffusionSchedulerOutput(step_id=self._step_id)

        batch = []
        while self._waiting_queue:
            req = self._waiting_queue.pop(0)
            batch.append(SimpleNamespace(request_id=req.request_id))

        output = DiffusionSchedulerOutput(step_id=self._step_id, scheduled_new_reqs=batch)
        self._step_id += 1
        return output

    def update_from_output(self, sched_output, runner_output):
        # assume all new req finished
        return [req.request_id for req in sched_output.scheduled_new_reqs]

    def get_request_state(self, request_id):
        return None

    def pop_request_state(self, request_id):
        return None

    def finish_requests(self, request_id, status):
        pass

    def close(self):
        pass

    def initialize(self, od_config):
        pass


def _make_admission_engine(pre_process_func) -> DiffusionEngine:
    engine = object.__new__(DiffusionEngine)
    engine.scheduler = MockScheduler()
    engine.pre_process_func = pre_process_func
    engine._out_streams = {}
    engine._rpc_lock = threading.RLock()
    engine._cv = threading.Condition(engine._rpc_lock)
    engine._closed = False
    return engine


class _BatchCapablePipeline:
    supports_request_batch = True


class _SingleRequestPipeline:
    pass


class _SingleRequestOverridePipeline(_BatchCapablePipeline):
    def forward(self, req, prompt_ids=None):
        return DiffusionOutput(output=None)


def _make_request_mode_sched_output(*request_ids: str) -> RealDiffusionSchedulerOutput:
    new_reqs = [
        NewRequestData(
            request_id=request_id,
            req=OmniDiffusionRequest(
                prompt=f"prompt_{request_id}",
                sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
                request_id=request_id,
            ),
        )
        for request_id in request_ids
    ]
    return RealDiffusionSchedulerOutput(
        step_id=0,
        scheduled_new_reqs=new_reqs,
        scheduled_cached_reqs=CachedRequestData.make_empty(),
        finished_req_ids=set(),
        num_running_reqs=len(new_reqs),
        num_waiting_reqs=0,
    )


class TestRequestBatchCapability:
    pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]

    def test_supports_request_batch_uses_registered_model_class(self, monkeypatch: pytest.MonkeyPatch) -> None:
        od_config = SimpleNamespace(model_class_name="BatchPipeline", custom_pipeline_args=None)

        monkeypatch.setattr(
            diffusion_engine_module.DiffusionModelRegistry,
            "_try_load_model_cls",
            lambda model_class_name: _BatchCapablePipeline if model_class_name == "BatchPipeline" else None,
        )

        assert diffusion_engine_module.supports_request_batch(od_config) is True

    def test_engine_rejects_multi_seq_diffusers_using_adapter_capability(
        self,
        monkeypatch: pytest.MonkeyPatch,
        mocker: MockerFixture,
    ) -> None:
        od_config = SimpleNamespace(
            model_class_name="QwenImagePipeline",
            custom_pipeline_args=None,
            diffusion_load_format="diffusers",
            streaming_output=False,
            max_num_seqs=2,
        )
        registry_load = mocker.Mock(
            side_effect=lambda model_class_name: (
                _SingleRequestPipeline if model_class_name == "DiffusersAdapterPipeline" else _BatchCapablePipeline
            )
        )
        monkeypatch.setattr(
            diffusion_engine_module.DiffusionModelRegistry,
            "_try_load_model_cls",
            registry_load,
        )

        engine = DiffusionEngine.__new__(DiffusionEngine)
        with pytest.raises(ValueError, match="max_num_seqs=1"):
            engine._resolve_execution_mode(od_config)

        registry_load.assert_called_once_with("DiffusersAdapterPipeline")

    def test_engine_prefers_batch_capable_custom_pipeline_over_diffusers_adapter(
        self,
        monkeypatch: pytest.MonkeyPatch,
        mocker: MockerFixture,
    ) -> None:
        od_config = SimpleNamespace(
            model_class_name="QwenImagePipeline",
            custom_pipeline_args={"pipeline_class": _BatchCapablePipeline},
            diffusion_load_format="diffusers",
            streaming_output=False,
            max_num_seqs=2,
        )
        registry_load = mocker.Mock(return_value=_SingleRequestPipeline)
        monkeypatch.setattr(
            diffusion_engine_module.DiffusionModelRegistry,
            "_try_load_model_cls",
            registry_load,
        )

        engine = DiffusionEngine.__new__(DiffusionEngine)

        assert engine._resolve_execution_mode(od_config) is DiffusionExecutionMode.REQUEST_BATCH
        registry_load.assert_not_called()

    def test_supports_request_batch_uses_custom_pipeline_class(self, monkeypatch: pytest.MonkeyPatch) -> None:
        od_config = SimpleNamespace(
            model_class_name="SinglePipeline",
            custom_pipeline_args={"pipeline_class": _BatchCapablePipeline},
        )

        monkeypatch.setattr(
            diffusion_engine_module.DiffusionModelRegistry,
            "_try_load_model_cls",
            lambda model_class_name: _SingleRequestPipeline,
        )

        assert diffusion_engine_module.supports_request_batch(od_config) is True

    def test_supports_request_batch_uses_custom_pipeline_class_name(self, monkeypatch: pytest.MonkeyPatch) -> None:
        od_config = SimpleNamespace(
            model_class_name="SinglePipeline",
            custom_pipeline_args={"pipeline_class": "test.module.BatchPipeline"},
        )

        monkeypatch.setattr(
            diffusion_engine_module,
            "resolve_obj_by_qualname",
            lambda qualname: _BatchCapablePipeline if qualname == "test.module.BatchPipeline" else None,
        )
        monkeypatch.setattr(
            diffusion_engine_module.DiffusionModelRegistry,
            "_try_load_model_cls",
            lambda model_class_name: _SingleRequestPipeline,
        )

        assert diffusion_engine_module.supports_request_batch(od_config) is True

    def test_supports_request_batch_uses_only_explicit_pipeline_attribute_for_custom_override(
        self, monkeypatch: pytest.MonkeyPatch
    ) -> None:
        od_config = SimpleNamespace(
            model_class_name="BatchPipeline",
            custom_pipeline_args={"pipeline_class": _SingleRequestOverridePipeline},
        )
        monkeypatch.setattr(
            diffusion_engine_module.DiffusionModelRegistry,
            "_try_load_model_cls",
            lambda model_class_name: None,
        )

        assert diffusion_engine_module.supports_request_batch(od_config) is True

    def test_supports_request_batch_honors_explicit_false_on_custom_override(
        self, monkeypatch: pytest.MonkeyPatch
    ) -> None:
        class _ExplicitlyUnsupportedOverride(_BatchCapablePipeline):
            supports_request_batch = False

            def forward(self, req, prompt_ids=None):
                return DiffusionOutput(output=None)

        od_config = SimpleNamespace(
            model_class_name="BatchPipeline",
            custom_pipeline_args={"pipeline_class": _ExplicitlyUnsupportedOverride},
        )
        monkeypatch.setattr(
            diffusion_engine_module.DiffusionModelRegistry,
            "_try_load_model_cls",
            lambda model_class_name: None,
        )

        assert diffusion_engine_module.supports_request_batch(od_config) is False

    def test_supports_request_batch_rejects_invalid_custom_pipeline_class_name(
        self,
        monkeypatch: pytest.MonkeyPatch,
        mocker: MockerFixture,
    ) -> None:
        od_config = SimpleNamespace(
            model_class_name="BatchPipeline",
            custom_pipeline_args={"pipeline_class": "test.module.MissingPipeline"},
        )

        def fail_resolve(qualname):
            raise ImportError(qualname)

        monkeypatch.setattr(diffusion_engine_module, "resolve_obj_by_qualname", fail_resolve)
        registry_load = mocker.Mock(return_value=_BatchCapablePipeline)
        monkeypatch.setattr(
            diffusion_engine_module.DiffusionModelRegistry,
            "_try_load_model_cls",
            registry_load,
        )

        with pytest.raises(ValueError, match="Failed to resolve custom diffusion pipeline class"):
            diffusion_engine_module.supports_request_batch(od_config)
        registry_load.assert_not_called()

    def test_engine_uses_request_batch_mode_for_single_request_pipeline(
        self,
        monkeypatch: pytest.MonkeyPatch,
        mocker: MockerFixture,
    ) -> None:
        od_config = SimpleNamespace(
            model_class_name="SinglePipeline",
            custom_pipeline_args=None,
            streaming_output=False,
            max_num_seqs=1,
        )
        fake_executor = SimpleNamespace(
            execute_request=mocker.Mock(return_value="per-request"),
            execute_batch=mocker.Mock(return_value="batch"),
            execute_step=mocker.Mock(return_value="step"),
        )
        fake_executor_cls = mocker.Mock(return_value=fake_executor)

        monkeypatch.setattr(
            "vllm_omni.diffusion.diffusion_engine.get_diffusion_post_process_func",
            lambda *args, **kwargs: None,
        )
        monkeypatch.setattr(
            "vllm_omni.diffusion.diffusion_engine.get_diffusion_pre_process_func",
            lambda *args, **kwargs: None,
        )
        monkeypatch.setattr(
            "vllm_omni.diffusion.diffusion_engine.DiffusionExecutor.get_class",
            lambda *args, **kwargs: fake_executor_cls,
        )
        monkeypatch.setattr(
            diffusion_engine_module.DiffusionModelRegistry,
            "_try_load_model_cls",
            lambda model_class_name: _SingleRequestPipeline,
        )

        engine = DiffusionEngine(od_config)
        output = engine.execute_fn(_make_request_mode_sched_output("req-a"))

        assert engine.execution_mode == DiffusionExecutionMode.REQUEST_BATCH
        assert engine.supports_request_batch is False
        assert output == "batch"
        fake_executor.execute_batch.assert_called_once()
        fake_executor.execute_request.assert_not_called()

    def test_engine_rejects_multi_seq_request_batch_for_single_request_pipeline(
        self,
        monkeypatch: pytest.MonkeyPatch,
        mocker: MockerFixture,
    ) -> None:
        od_config = SimpleNamespace(
            model_class_name="SinglePipeline",
            custom_pipeline_args=None,
            streaming_output=False,
            max_num_seqs=2,
        )
        fake_executor_cls = mocker.Mock()

        monkeypatch.setattr(
            "vllm_omni.diffusion.diffusion_engine.get_diffusion_post_process_func",
            lambda *args, **kwargs: None,
        )
        monkeypatch.setattr(
            "vllm_omni.diffusion.diffusion_engine.get_diffusion_pre_process_func",
            lambda *args, **kwargs: None,
        )
        monkeypatch.setattr(
            "vllm_omni.diffusion.diffusion_engine.DiffusionExecutor.get_class",
            lambda *args, **kwargs: fake_executor_cls,
        )
        monkeypatch.setattr(
            diffusion_engine_module.DiffusionModelRegistry,
            "_try_load_model_cls",
            lambda model_class_name: _SingleRequestPipeline,
        )

        with pytest.raises(ValueError, match="max_num_seqs=1"):
            DiffusionEngine(od_config)
        fake_executor_cls.assert_not_called()

    def test_engine_allows_independent_dlo_dp_requests_for_single_request_pipeline(
        self,
        monkeypatch: pytest.MonkeyPatch,
        mocker: MockerFixture,
    ) -> None:
        od_config = SimpleNamespace(
            model_class_name="SinglePipeline",
            custom_pipeline_args=None,
            streaming_output=False,
            max_num_seqs=2,
            request_batch_max_wait_ms=0,
            parallel_config=SimpleNamespace(data_parallel_size=2),
            enable_distributed_layerwise_offload=True,
            dlo_use_allgather=True,
        )
        fake_executor = SimpleNamespace(
            execute_request=mocker.Mock(return_value="per-request"),
            execute_batch=mocker.Mock(return_value="batch"),
            execute_step=mocker.Mock(return_value="step"),
        )
        fake_executor_cls = mocker.Mock(return_value=fake_executor)

        monkeypatch.setattr(
            "vllm_omni.diffusion.diffusion_engine.get_diffusion_post_process_func",
            lambda *args, **kwargs: None,
        )
        monkeypatch.setattr(
            "vllm_omni.diffusion.diffusion_engine.get_diffusion_pre_process_func",
            lambda *args, **kwargs: None,
        )
        monkeypatch.setattr(
            "vllm_omni.diffusion.diffusion_engine.DiffusionExecutor.get_class",
            lambda *args, **kwargs: fake_executor_cls,
        )
        monkeypatch.setattr(
            diffusion_engine_module.DiffusionModelRegistry,
            "_try_load_model_cls",
            lambda model_class_name: _SingleRequestPipeline,
        )

        engine = DiffusionEngine(od_config)

        assert engine.execution_mode == DiffusionExecutionMode.REQUEST_BATCH
        assert engine.supports_request_batch is False
        assert engine.dp_concurrent is True
        assert engine.scheduler.max_num_running_reqs == 2
        assert od_config.request_batch_max_wait_ms == 0
        fake_executor_cls.assert_called_once_with(od_config)

    @pytest.mark.parametrize("request_ids", [("req-a",), ("req-a", "req-b")])
    def test_engine_enables_batch_dispatch_for_request_batch_pipeline(
        self,
        request_ids: tuple[str, ...],
        monkeypatch: pytest.MonkeyPatch,
        mocker: MockerFixture,
    ) -> None:
        od_config = SimpleNamespace(
            model_class_name="BatchPipeline",
            custom_pipeline_args=None,
            streaming_output=False,
            max_num_seqs=2,
        )
        fake_executor = SimpleNamespace(
            execute_request=mocker.Mock(return_value="per-request"),
            execute_batch=mocker.Mock(return_value="batch"),
            execute_step=mocker.Mock(return_value="step"),
        )
        fake_executor_cls = mocker.Mock(return_value=fake_executor)

        monkeypatch.setattr(
            "vllm_omni.diffusion.diffusion_engine.get_diffusion_post_process_func",
            lambda *args, **kwargs: None,
        )
        monkeypatch.setattr(
            "vllm_omni.diffusion.diffusion_engine.get_diffusion_pre_process_func",
            lambda *args, **kwargs: None,
        )
        monkeypatch.setattr(
            "vllm_omni.diffusion.diffusion_engine.DiffusionExecutor.get_class",
            lambda *args, **kwargs: fake_executor_cls,
        )
        monkeypatch.setattr(
            diffusion_engine_module.DiffusionModelRegistry,
            "_try_load_model_cls",
            lambda model_class_name: _BatchCapablePipeline,
        )

        engine = DiffusionEngine(od_config)
        output = engine.execute_fn(_make_request_mode_sched_output(*request_ids))

        assert engine.execution_mode == DiffusionExecutionMode.REQUEST_BATCH
        assert engine.supports_request_batch is True
        assert output == "batch"
        fake_executor.execute_batch.assert_called_once()
        fake_executor.execute_request.assert_not_called()

    def test_make_engine_runs_startup_warmup(self, monkeypatch: pytest.MonkeyPatch, mocker: MockerFixture) -> None:
        od_config = SimpleNamespace(
            model_class_name="SinglePipeline",
            custom_pipeline_args=None,
            engine_backend="default",
            streaming_output=False,
        )
        fake_executor = SimpleNamespace(
            execute_request=mocker.Mock(),
            execute_batch=mocker.Mock(),
            execute_step=mocker.Mock(),
        )
        fake_executor_cls = mocker.Mock(return_value=fake_executor)
        warmup = mocker.Mock()

        monkeypatch.setattr(
            "vllm_omni.diffusion.diffusion_engine.get_diffusion_post_process_func",
            lambda *args, **kwargs: None,
        )
        monkeypatch.setattr(
            "vllm_omni.diffusion.diffusion_engine.get_diffusion_pre_process_func",
            lambda *args, **kwargs: None,
        )
        monkeypatch.setattr(
            "vllm_omni.diffusion.diffusion_engine.DiffusionExecutor.get_class",
            lambda *args, **kwargs: fake_executor_cls,
        )
        monkeypatch.setattr(
            diffusion_engine_module.DiffusionModelRegistry,
            "_try_load_model_cls",
            lambda model_class_name: _SingleRequestPipeline,
        )
        monkeypatch.setattr(DiffusionEngine, "run_startup_warmup", warmup)

        engine = DiffusionEngine.make_engine(od_config)

        assert isinstance(engine, DiffusionEngine)
        warmup.assert_called_once_with()


class TestDiffusionCompileConfig:
    pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]

    def test_config_defaults_to_dynamic_regional_compile(self) -> None:
        config = OmniDiffusionConfig(model="test")

        assert config.diffusion_compile_granularity == "regional"
        assert config.diffusion_compile_dynamic is True

    def test_from_kwargs_preserves_compile_controls(self) -> None:
        config = OmniDiffusionConfig.from_kwargs(
            model="test",
            diffusion_compile_granularity="full",
            diffusion_compile_dynamic=False,
        )

        assert config.diffusion_compile_granularity == "full"
        assert config.diffusion_compile_dynamic is False

    def test_config_rejects_invalid_compile_granularity(self) -> None:
        with pytest.raises(ValueError, match="diffusion_compile_granularity"):
            OmniDiffusionConfig(model="test", diffusion_compile_granularity="block")

    def test_config_rejects_non_boolean_compile_dynamic(self) -> None:
        with pytest.raises(TypeError, match="diffusion_compile_dynamic"):
            OmniDiffusionConfig(model="test", diffusion_compile_dynamic="false")

    @pytest.mark.parametrize(
        "kwargs, feature",
        [
            ({"parallel_config": {"ulysses_degree": 2}, "num_gpus": 2}, "sequence parallelism"),
            (
                {
                    "parallel_config": {"use_hsdp": True, "hsdp_shard_size": 2},
                    "num_gpus": 2,
                },
                "HSDP",
            ),
            ({"enable_cpu_offload": True}, "CPU offload"),
            ({"enable_layerwise_offload": True}, "layerwise offload"),
            (
                {
                    "diffusion_offload_config": {
                        "mode": "module",
                        "components": ["dit"],
                    }
                },
                "CPU offload",
            ),
            (
                {
                    "diffusion_offload_config": {
                        "mode": "layer",
                        "components": ["dit"],
                    }
                },
                "layerwise offload",
            ),
            (
                {
                    "diffusion_offload_config": {
                        "mode": "layer",
                        "components": ["dit"],
                        "layer_options": {"dit": {"weight_transfer": "allgather"}},
                    }
                },
                "distributed layerwise offload",
            ),
        ],
    )
    def test_full_compile_rejects_incompatible_features(self, kwargs, feature) -> None:
        with pytest.raises(ValueError, match=feature):
            OmniDiffusionConfig(
                model="test",
                diffusion_compile_granularity="full",
                **kwargs,
            )

    def test_raw_parallel_mapping_is_normalized_before_allgather_cache_validation(self) -> None:
        with pytest.raises(ValueError, match="rank-local cache decisions"):
            OmniDiffusionConfig(
                model="test",
                num_gpus=2,
                parallel_config={"data_parallel_size": 2},
                cache_backend="tea_cache",
                diffusion_offload_config={
                    "mode": "layer",
                    "components": ["dit"],
                    "layer_options": {"dit": {"weight_transfer": "allgather"}},
                },
            )

    def test_auto_dp_is_resolved_before_encoder_prompt_cache_validation(self) -> None:
        with pytest.raises(ValueError, match="rank-local cache hits"):
            OmniDiffusionConfig(
                model="test",
                num_gpus=2,
                parallel_config={"data_parallel_size": None},
                enable_prompt_embed_cache=True,
                diffusion_offload_config={
                    "mode": "layer",
                    "components": ["text_encoder"],
                    "layer_options": {"text_encoder": {"weight_transfer": "allgather"}},
                },
            )


class TestRequestBatchAdmission:
    pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]

    def test_config_rejects_negative_request_batch_max_wait_ms(self) -> None:
        with pytest.raises(ValueError, match="request_batch_max_wait_ms"):
            OmniDiffusionConfig(model="test", request_batch_max_wait_ms=-1.0)

    @pytest.mark.parametrize("bad_wait", [float("nan"), float("inf"), float("-inf")])
    def test_config_rejects_nonfinite_request_batch_max_wait_ms(self, bad_wait: float) -> None:
        with pytest.raises(ValueError, match="request_batch_max_wait_ms"):
            OmniDiffusionConfig(model="test", request_batch_max_wait_ms=bad_wait)

    def test_config_normalizes_request_batch_max_wait_ms_to_float(self) -> None:
        config = OmniDiffusionConfig(model="test", request_batch_max_wait_ms=5)

        assert config.request_batch_max_wait_ms == 5.0
        assert isinstance(config.request_batch_max_wait_ms, float)

    def test_scheduler_exposes_waiting_and_running_counts(self) -> None:
        from vllm_omni.diffusion.sched import RequestScheduler

        od_config = SimpleNamespace(max_num_seqs=4)
        scheduler = RequestScheduler()
        scheduler.initialize(od_config)

        assert scheduler.num_waiting_requests() == 0
        assert scheduler.num_running_requests() == 0

        scheduler.add_request(
            OmniDiffusionRequest(
                prompt="prompt_a",
                sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
                request_id="req-a",
            )
        )
        scheduler.add_request(
            OmniDiffusionRequest(
                prompt="prompt_b",
                sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
                request_id="req-b",
            )
        )
        assert scheduler.num_waiting_requests() == 2
        assert scheduler.num_running_requests() == 0

        scheduler.schedule()
        assert scheduler.num_waiting_requests() == 0
        assert scheduler.num_running_requests() == 2

    def test_request_batch_admission_exits_early_when_waiting_queue_stable(self) -> None:
        from vllm_omni.diffusion.sched import RequestScheduler

        od_config = SimpleNamespace(
            max_num_seqs=32,
            request_batch_max_wait_ms=1000.0,
            step_execution=False,
        )
        scheduler = RequestScheduler()
        scheduler.initialize(od_config)
        for idx in range(2):
            scheduler.add_request(
                OmniDiffusionRequest(
                    prompt=f"prompt_{idx}",
                    sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
                    request_id=f"req-{idx}",
                )
            )

        engine = object.__new__(DiffusionEngine)
        engine.od_config = od_config
        engine.scheduler = scheduler
        engine.step_execution = False
        engine.supports_request_batch = True
        engine.stop_event = threading.Event()
        engine._rpc_lock = threading.RLock()
        engine._cv = threading.Condition(engine._rpc_lock)

        start = time.monotonic()
        with engine._cv:
            engine._wait_for_admission_if_needed_locked()
        waited_s = time.monotonic() - start

        # Stable-window exit (~50ms), not the full 1000ms deadline.
        assert waited_s < 0.5
        assert waited_s >= 0.04
        assert scheduler.num_waiting_requests() == 2
        assert scheduler.num_running_requests() == 0


@pytest.mark.core_model
@pytest.mark.diffusion
@pytest.mark.cpu
def test_move_tensor_tree_keeps_cpu_tensor_identity() -> None:
    tensor = torch.arange(8, dtype=torch.float32)

    moved = _move_tensor_tree_to_cpu(tensor)

    assert moved is tensor


@pytest.mark.core_model
@pytest.mark.diffusion
@pytest.mark.cpu
def test_move_tensor_tree_preserves_nested_structure_without_mutating_input() -> None:
    tensor = torch.arange(4, dtype=torch.float32)
    nested_tensor = torch.arange(6, dtype=torch.float32).reshape(2, 3)
    sentinel = object()
    payload = {
        "tensor": tensor,
        "list": [nested_tensor, sentinel],
        "tuple": ({"inner": tensor}, "metadata"),
        "scalar": 3,
    }

    moved = _move_tensor_tree_to_cpu(payload)

    assert moved is not payload
    assert set(moved) == {"tensor", "list", "tuple", "scalar"}
    assert moved["list"] is not payload["list"]
    assert moved["tuple"] is not payload["tuple"]
    assert moved["tuple"][0] is not payload["tuple"][0]
    assert moved["tensor"] is tensor
    assert moved["list"][0] is nested_tensor
    assert moved["list"][1] is sentinel
    assert moved["tuple"][0]["inner"] is tensor
    assert moved["tuple"][1] == "metadata"
    assert moved["scalar"] == 3
    assert payload["list"][0] is nested_tensor
    assert payload["list"][1] is sentinel
    assert payload["tuple"][0]["inner"] is tensor
    assert payload["tuple"][1] == "metadata"


@pytest.mark.core_model
@pytest.mark.diffusion
@pytest.mark.cpu
def test_move_tensor_tree_returns_non_tensor_values_unchanged() -> None:
    value = object()

    moved = _move_tensor_tree_to_cpu(value)

    assert moved is value


@hardware_test(res={"cuda": "L4"}, num_cards=1)
def test_move_tensor_tree_moves_nested_cuda_tensors_to_cpu() -> None:
    tensor = torch.arange(8, dtype=torch.float32, device="cuda")
    other = torch.arange(4, dtype=torch.int64, device="cuda")
    payload = {"tensor": tensor, "items": [other, ("keep", tensor)]}

    moved = _move_tensor_tree_to_cpu(payload)

    assert moved["tensor"].device.type == "cpu"
    assert moved["items"][0].device.type == "cpu"
    assert moved["items"][1][1].device.type == "cpu"
    torch.testing.assert_close(moved["tensor"], tensor.cpu())
    torch.testing.assert_close(moved["items"][0], other.cpu())
    torch.testing.assert_close(moved["items"][1][1], tensor.cpu())
    assert moved["items"][1][0] == "keep"


async def _consume_final_output(generator):
    final_output = None
    async for output in generator:
        final_output = output
    if final_output is None:
        raise RuntimeError("Diffusion execution finished without output.")
    return final_output


@pytest.mark.cpu
@pytest.mark.parametrize("entrypoint", ["add_request", "async_add_req_and_stream_response"])
def test_engine_admission_preprocesses_request_once(entrypoint: str) -> None:
    raw_request = OmniDiffusionRequest(
        prompt="raw",
        sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
        request_id="admission",
    )
    prepared_request = OmniDiffusionRequest(
        prompt="prepared",
        sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
        request_id="admission",
        prepared_layout=object(),
    )
    preprocess_calls = []

    def preprocess(request):
        preprocess_calls.append(request)
        return prepared_request

    engine = _make_admission_engine(preprocess)

    getattr(engine, entrypoint)(raw_request)

    assert preprocess_calls == [raw_request]
    assert engine.scheduler._waiting_queue == [prepared_request]


@pytest.mark.cpu
@pytest.mark.asyncio
async def test_async_add_req_and_stream_response():
    engine = object.__new__(DiffusionEngine)
    engine.scheduler = MockScheduler()
    engine._out_streams = {}
    engine.abort_queue: queue.Queue[str] = queue.Queue()
    engine._rpc_queue = queue.Queue()
    engine._rpc_lock = threading.RLock()
    engine._cv = threading.Condition(engine._rpc_lock)
    engine._init_lock = asyncio.Lock()
    engine._closed = False
    # Enable admission wait so concurrent adds land in one schedule wave;
    # otherwise the first request can execute alone (~1s) and the rest in a
    # second wave (~2s), failing the latency-spread assertion.
    engine.od_config = SimpleNamespace(
        streaming_output=False,
        request_batch_max_wait_ms=500.0,
    )
    engine._loop_started = False
    engine.main_loop = None
    engine.supports_request_batch = True
    engine.step_execution = False
    engine.execution_mode = DiffusionExecutionMode.REQUEST_BATCH

    def _finalize(rid, out, err=None, **kwargs):
        # Stream consumers stop on ``finished``; keep result_data for assertions.
        return SimpleNamespace(result_data=out.result.result_data, finished=True)

    engine._finalize_finished_request = _finalize

    def mock_execute_batch(sched_output):
        request_ids = sched_output.scheduled_request_ids

        time.sleep(1)

        class MockRunnerOutput:
            def __init__(self, ids):
                self.request_id = ids
                self.step_index = 0
                self.finished = True
                self._results = {rid: SimpleNamespace(result_data=f"data_{rid}") for rid in ids}

            def get_request_output(self, rid):
                return SimpleNamespace(result=self._results[rid], step_index=0, finished=True)

        return MockRunnerOutput(request_ids)

    engine.execute_fn = mock_execute_batch

    await engine._check_and_start_background_loop()

    async def run_task(rid):
        req = SimpleNamespace(request_id=rid)
        start = time.time()
        res = await _consume_final_output(engine.async_add_req_and_stream_response(req))
        return rid, res, time.time() - start

    task_ids = [f"req_{i}" for i in range(5)]
    tasks = [run_task(rid) for rid in task_ids]
    try:
        results = await asyncio.gather(*tasks)
    finally:
        with engine._cv:
            engine.stop_event.set()
            engine._cv.notify_all()
        engine.worker_thread.join(timeout=5)
    assert len(results) == 5
    for rid, res, elapsed in results:
        assert rid in res.result_data

    eps = 0.5
    latencies = [r[2] for r in results]
    assert max(latencies) - min(latencies) < eps