# 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 ( DIFFUSION_REQUEST_LIFECYCLE_KEY, DIFFUSION_REQUEST_STARTED, 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, ) @pytest.mark.cpu def test_request_started_output_is_emitted_only_for_opted_in_requests() -> None: sched_output = _make_request_mode_sched_output("tracked", "untracked") sched_output.scheduled_new_reqs[0].req.sampling_params.emit_request_lifecycle = True engine = object.__new__(DiffusionEngine) emitted = [] engine._put_output = lambda request_id, output: emitted.append((request_id, output)) engine._emit_request_started_outputs(sched_output) assert len(emitted) == 1 request_id, output = emitted[0] assert request_id == "tracked" assert output.request_started is True assert output.finished is False [formatted] = engine.postprocess_output(sched_output.scheduled_new_reqs[0].req, output) assert formatted.custom_output == { DIFFUSION_REQUEST_LIFECYCLE_KEY: DIFFUSION_REQUEST_STARTED, } assert formatted.finished is False 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