# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Tests for step-level diffusion execution across runner / worker / executor / engine.""" import contextlib import os import queue import threading from contextlib import contextmanager from types import SimpleNamespace import pytest import torch from pytest_mock import MockerFixture import vllm_omni.diffusion.worker.diffusion_model_runner as model_runner_module from tests.helpers.mark import hardware_test from vllm_omni.config.config_factory import StageConfigFactory from vllm_omni.diffusion.data import DiffusionOutput from vllm_omni.diffusion.diffusion_engine import DiffusionEngine from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode from vllm_omni.diffusion.diffusion_kv.metadata import DiffusionKVMetadata from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin from vllm_omni.diffusion.distributed.comm import RingComm, SeqAllToAll4D from vllm_omni.diffusion.distributed.parallel_state import ( destroy_distributed_env, get_sp_group, init_distributed_environment, initialize_model_parallel, ) from vllm_omni.diffusion.executor.multiproc_executor import MultiprocDiffusionExecutor from vllm_omni.diffusion.ipc import ( pack_diffusion_output_shm, unpack_diffusion_output_shm, ) from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import ( DiffusionPipelineProfilerMixin, ) from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.sched import StepScheduler from vllm_omni.diffusion.sched.interface import ( CachedRequestData, DiffusionSchedulerOutput, NewRequestData, ) from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker from vllm_omni.diffusion.worker.input_batch import InputBatch from vllm_omni.diffusion.worker.utils import RunnerOutput, StepRequestState from vllm_omni.errors import OmniClientError from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.platforms import current_omni_platform pytestmark = [pytest.mark.core_model, pytest.mark.diffusion] # --------------------------------------------------------------------------- # Helpers & fixtures # --------------------------------------------------------------------------- @contextmanager def _noop_forward_context(*args, **kwargs): del args, kwargs yield def _update_environment_variables(envs_dict: dict[str, str]) -> None: for key, value in envs_dict.items(): os.environ[key] = value class _StepPipeline: """Minimal pipeline stub that supports step-wise execution.""" supports_step_execution = True def __init__(self): self.prepare_calls = 0 self.denoise_calls = 0 self.scheduler_calls = 0 self.decode_calls = 0 def prepare_encode(self, state, **kwargs): del kwargs self.prepare_calls += 1 state.timesteps = [torch.tensor(10), torch.tensor(5)] state.latents = torch.tensor([0.0]) state.prompt_embeds = torch.tensor([[0.0, 0.0], [1.0, 1.0]]) return state def denoise_step(self, input_batch, **kwargs): self.denoise_calls += 1 return torch.full_like(input_batch.prompt_embeds, fill_value=0.5) def step_scheduler(self, state, noise_pred, **kwargs): del noise_pred, kwargs self.scheduler_calls += 1 state.step_index += 1 def post_decode(self, state, **kwargs): del kwargs self.decode_calls += 1 return DiffusionOutput(output=torch.tensor([state.step_index], dtype=torch.float32)) class _ChunkedStepPipeline(_StepPipeline): """Streaming stub: two chunks of two denoise steps each, one decode per chunk.""" supports_chunk_step_grouping = False def prepare_encode(self, state, **kwargs): del kwargs self.prepare_calls += 1 state.timesteps = [torch.tensor(10), torch.tensor(5), torch.tensor(10), torch.tensor(5)] state.latents = torch.tensor([0.0]) state.prompt_embeds = torch.tensor([[0.0, 0.0], [1.0, 1.0]]) state.chunk_num_steps = 2 state.total_chunks = 2 state.chunk_index = 0 state.step_in_chunk = 0 return state def step_scheduler(self, state, noise_pred, **kwargs): del noise_pred, kwargs self.scheduler_calls += 1 state.step_in_chunk += 1 state.step_index = state.step_in_chunk def post_decode(self, state, **kwargs): del kwargs self.decode_calls += 1 result = DiffusionOutput(output=torch.tensor([float(state.chunk_index)])) state.chunk_index += 1 state.step_in_chunk = 0 return result class _BatchDependentChunkPipeline(_ChunkedStepPipeline): """Streaming stub whose noise depends on the batch's latents and timesteps, so a stale batch shows.""" def __init__(self): super().__init__() self.seen: list[tuple[float, float]] = [] def prepare_encode(self, state, **kwargs): super().prepare_encode(state, **kwargs) state.latents = torch.tensor([1.0]) return state def denoise_step(self, input_batch, **kwargs): self.denoise_calls += 1 latent = float(input_batch.latents.reshape(-1)[0]) timestep = float(input_batch.timesteps.reshape(-1)[0]) self.seen.append((latent, timestep)) return torch.full((input_batch.latents.shape[0], 1), latent + timestep) def step_scheduler(self, state, noise_pred, **kwargs): del kwargs self.scheduler_calls += 1 state.latents = state.latents + noise_pred.reshape(-1)[0] state.step_in_chunk += 1 state.step_index = state.step_in_chunk def post_decode(self, state, **kwargs): self.final_latent = state.latents.reshape(-1)[0].clone() return super().post_decode(state, **kwargs) class _FailsMidChunkPipeline(_ChunkedStepPipeline): """Streaming stub whose second denoise step of the first chunk raises.""" def step_scheduler(self, state, noise_pred, **kwargs): super().step_scheduler(state, noise_pred, **kwargs) if state.chunk_index == 0 and state.step_in_chunk == 2: raise RuntimeError("boom mid chunk") class _PerRequestErrorStepPipeline(_StepPipeline): def prepare_encode(self, state, **kwargs): if state.prompt == "fail-prepare": raise OmniClientError( "invalid request input", status_code=422, error_type="UnprocessableEntityError", ) return super().prepare_encode(state, **kwargs) def denoise_step(self, input_batch, **kwargs): del kwargs self.denoise_calls += 1 return torch.ones_like(input_batch.latents) class _ProfilingStepPipeline(_StepPipeline): enable_diffusion_pipeline_profiler = True def __init__(self): super().__init__() self._stage_durations: dict[str, float] = {} @property def stage_durations(self) -> dict[str, float]: return dict(self._stage_durations) def clear_profiler_records(self) -> None: self._stage_durations.clear() def prepare_encode(self, state, **kwargs): result = super().prepare_encode(state, **kwargs) self._stage_durations["QwenImagePipeline.text_encoder.forward"] = 1.0 return result def denoise_step(self, input_batch, **kwargs): result = super().denoise_step(input_batch, **kwargs) self._stage_durations["QwenImagePipeline.diffuse"] = 2.0 return result def post_decode(self, state, **kwargs): result = super().post_decode(state, **kwargs) self._stage_durations["QwenImagePipeline.vae.decode"] = 3.0 return result class _AutoDenoiseProfilerPipeline(DiffusionPipelineProfilerMixin): _PROFILER_TARGETS: list[str] = [] def __init__(self): self.setup_diffusion_pipeline_profiler( enable_diffusion_pipeline_profiler=True, ) def forward(self): return None def denoise_step(self): return "ok" class _InterruptingStepPipeline(_StepPipeline): interrupt = True def denoise_step(self, state, **kwargs): del state, kwargs self.denoise_calls += 1 return None def step_scheduler(self, state, noise_pred, **kwargs): del state, noise_pred, kwargs raise AssertionError("step_scheduler should not run after interrupt") def post_decode(self, state, **kwargs): del state, kwargs raise AssertionError("post_decode should not run after interrupt") class _FakePeakMemoryPlatform: def __init__(self, reserved_mb: list[float]): self._reserved_mb = reserved_mb self.reset_calls = 0 def reset_peak_memory_stats(self): self.reset_calls += 1 def max_memory_reserved(self): index = min(self.reset_calls - 1, len(self._reserved_mb) - 1) return int(self._reserved_mb[index] * 1024**2) def max_memory_allocated(self): index = min(self.reset_calls - 1, len(self._reserved_mb) - 1) return int((self._reserved_mb[index] - 100) * 1024**2) def is_available(self) -> bool: return True class _IdentityNoiseTransformer(torch.nn.Module): def forward(self, x: torch.Tensor, **kwargs): del kwargs return (x,) class _AdditiveScheduler: def step(self, noise_pred: torch.Tensor, t: torch.Tensor, latents: torch.Tensor, return_dict: bool = False): del t, return_dict return (latents + noise_pred,) class _DistributedStepPipeline(CFGParallelMixin): supports_step_execution = True def __init__(self, mode: str, device: torch.device): self.mode = mode self.device = device self._interrupt = False self.scheduler = _AdditiveScheduler() self.transformer = _IdentityNoiseTransformer() @property def interrupt(self): return self._interrupt def prepare_encode(self, state, **kwargs): del kwargs state.timesteps = [torch.tensor(1.0, device=self.device)] state.latents = torch.ones((1, 1), device=self.device) state.step_index = 0 state.scheduler = self.scheduler state.do_true_cfg = self.mode == "cfg" state.prompt_embeds = torch.tensor([[0.0, 0.0], [1.0, 1.0]]) return state def denoise_step(self, state, **kwargs): del kwargs if self.mode == "ulysses": sp_group = get_sp_group().ulysses_group seq_world_size = torch.distributed.get_world_size(sp_group) input_tensor = torch.randn(1, 2, 2 * seq_world_size, 2, device=self.device) original = input_tensor.clone() intermediate = SeqAllToAll4D.apply(sp_group, input_tensor, 2, 1, False) output = SeqAllToAll4D.apply(sp_group, intermediate, 1, 2, False) torch.testing.assert_close(output, original, rtol=1e-5, atol=1e-5) return torch.ones_like(state.latents) if self.mode == "ring": ring_group = get_sp_group().ring_group rank = torch.distributed.get_rank(ring_group) world_size = torch.distributed.get_world_size(ring_group) comm = RingComm(ring_group) input_tensor = torch.full((1, 2, 2), float(rank + 1), device=self.device) recv_tensor = comm.send_recv(input_tensor) comm.commit() comm.wait() expected = torch.full_like(recv_tensor, float(((rank - 1) % world_size) + 1)) torch.testing.assert_close(recv_tensor, expected, rtol=1e-5, atol=1e-5) return torch.ones_like(state.latents) positive_kwargs = {"x": state.latents + 1} negative_kwargs = {"x": state.latents - 1} return self.predict_noise_maybe_with_cfg( do_true_cfg=True, true_cfg_scale=1.0, positive_kwargs=positive_kwargs, negative_kwargs=negative_kwargs, cfg_normalize=False, ) def step_scheduler(self, state, noise_pred, **kwargs): del kwargs if self.mode == "cfg": state.latents = self.scheduler_step_maybe_with_cfg( noise_pred, state.current_timestep, state.latents, do_true_cfg=True, per_request_scheduler=state.scheduler, ) else: state.latents = state.latents + noise_pred state.step_index += 1 def post_decode(self, state, **kwargs): del kwargs return DiffusionOutput(output=state.latents.detach().cpu()) def _make_step_request(num_inference_steps: int = 2): return OmniDiffusionRequest( prompt="a prompt", request_id="req-1", sampling_params=OmniDiffusionSamplingParams(num_inference_steps=num_inference_steps), ) def _assert_aborted_output(output: DiffusionOutput, request_id: str) -> None: assert output.output is None assert output.error is None assert output.aborted is True assert output.abort_message == f"Request {request_id} aborted." def _make_engine_request(req_id: str = "req-1", num_inference_steps: int = 2) -> OmniDiffusionRequest: return OmniDiffusionRequest( prompt=f"prompt-{req_id}", sampling_params=OmniDiffusionSamplingParams(num_inference_steps=num_inference_steps), request_id=req_id, ) def _make_vllm_config(): @contextlib.contextmanager def set_priority(*args, **kwargs): yield return SimpleNamespace( kernel_config=SimpleNamespace(ir_op_priority=SimpleNamespace(set_priority=set_priority)), compilation_config=SimpleNamespace(ir_enable_torch_wrap=True), ) def _make_runner( diffusion_kv_mode: DiffusionKVCacheMode = DiffusionKVCacheMode.DENSE_LEGACY, runner_cls: type[DiffusionModelRunner] = DiffusionModelRunner, ): runner = object.__new__(runner_cls) runner.vllm_config = _make_vllm_config() runner.od_config = SimpleNamespace( cache_backend=None, diffusion_kv_mode=diffusion_kv_mode, parallel_config=SimpleNamespace(use_hsdp=False), streaming_output=False, ) runner.device = torch.device("cpu") runner.pipeline = _StepPipeline() runner.cache_backend = None runner.offload_backend = None runner.state_cache = {} runner.kv_transfer_manager = SimpleNamespace( receive_multi_kv_cache_distributed=lambda req, cfg_kv_collect_func=None, target_device=None: None ) return runner def _make_distributed_runner(mode: str, device: torch.device): runner = object.__new__(DiffusionModelRunner) runner.vllm_config = _make_vllm_config() runner.od_config = SimpleNamespace( cache_backend=None, diffusion_kv_mode=DiffusionKVCacheMode.DENSE_LEGACY, parallel_config=SimpleNamespace(use_hsdp=False), streaming_output=False, ) runner.device = device runner.pipeline = _DistributedStepPipeline(mode=mode, device=device) runner.cache_backend = None runner.offload_backend = None runner.state_cache = {} runner.kv_transfer_manager = SimpleNamespace( receive_multi_kv_cache_distributed=lambda req, cfg_kv_collect_func=None, target_device=None: None ) return runner def _make_scheduler_output(req, request_id="req-1", step_id=0, finished_req_ids=None): req.request_id = request_id return DiffusionSchedulerOutput( step_id=step_id, scheduled_new_reqs=[NewRequestData(request_id=request_id, req=req)], scheduled_cached_reqs=CachedRequestData.make_empty(), finished_req_ids=set() if finished_req_ids is None else set(finished_req_ids), num_running_reqs=1, num_waiting_reqs=0, ) def _make_batch_scheduler_output(reqs, *, step_id=0, finished_req_ids=None): """Scheduler output for a homogeneous batch (one NewRequestData per req).""" new_reqs = [NewRequestData(request_id=r.request_id, req=r) for r in reqs] return DiffusionSchedulerOutput( step_id=step_id, scheduled_new_reqs=new_reqs, scheduled_cached_reqs=CachedRequestData.make_empty(), finished_req_ids=set() if finished_req_ids is None else set(finished_req_ids), num_running_reqs=len(new_reqs), num_waiting_reqs=0, ) def _make_input_batch_state(request_id: str, latent_value: float) -> StepRequestState: state = StepRequestState( request_id=request_id, sampling=SimpleNamespace(), prompt=None, ) state.latents = torch.tensor([[latent_value]]) state.timesteps = torch.tensor([1.0]) return state def _make_cached_scheduler_output(request_id="req-1", step_id=1, finished_req_ids=None): return DiffusionSchedulerOutput( step_id=step_id, scheduled_new_reqs=[], scheduled_cached_reqs=CachedRequestData(request_ids=[request_id]), finished_req_ids=set() if finished_req_ids is None else set(finished_req_ids), num_running_reqs=1, num_waiting_reqs=0, ) def _make_engine(scheduler, execute_fn=None) -> DiffusionEngine: engine = object.__new__(DiffusionEngine) engine.od_config = SimpleNamespace(model_class_name="QwenImagePipeline", streaming_output=False) engine.pre_process_func = None engine.post_process_func = None engine.scheduler = scheduler engine.execute_fn = execute_fn engine._rpc_lock = threading.RLock() engine._cv = threading.Condition(engine._rpc_lock) engine._closed = False engine.abort_queue = queue.Queue() return engine def _expected_output_for_mode(mode: str) -> torch.Tensor: if mode == "cfg": return torch.tensor([[3.0]]) return torch.tensor([[2.0]]) def _distributed_step_worker(local_rank: int, world_size: int, mode: str, master_port: str): device = torch.device(f"{current_omni_platform.device_type}:{local_rank}") current_omni_platform.set_device(device) _update_environment_variables( { "RANK": str(local_rank), "LOCAL_RANK": str(local_rank), "WORLD_SIZE": str(world_size), "MASTER_ADDR": "localhost", "MASTER_PORT": master_port, } ) model_runner_module.set_forward_context = _noop_forward_context try: init_distributed_environment() if mode == "ulysses": initialize_model_parallel(ulysses_degree=world_size) elif mode == "ring": initialize_model_parallel(ring_degree=world_size) elif mode == "cfg": initialize_model_parallel(cfg_parallel_size=world_size) else: raise ValueError(f"Unsupported distributed test mode: {mode}") runner = _make_distributed_runner(mode, device) result = DiffusionModelRunner.execute_stepwise( runner, _make_scheduler_output(_make_step_request(num_inference_steps=1), step_id=0), ) output = result.get_request_output("req-1") assert output.finished is True assert output.result is not None torch.testing.assert_close(output.result.output, _expected_output_for_mode(mode), rtol=1e-5, atol=1e-5) assert "req-1" not in runner.state_cache finally: destroy_distributed_env() # --------------------------------------------------------------------------- # Runner / Worker # --------------------------------------------------------------------------- @pytest.mark.cpu def test_input_batch_cached_repack_refreshes_state_references_without_prompt_embeds(): first_state = _make_input_batch_state("req-1", 1.0) batch = InputBatch.make_batch([first_state]) assert batch.prompt_embeds is None replacement_state = _make_input_batch_state("req-1", 2.0) repacked = InputBatch.make_batch([replacement_state], cached_batch=batch) assert repacked is batch assert repacked.states[0] is replacement_state assert repacked.prompt_embeds is None torch.testing.assert_close(repacked.latents, replacement_state.latents) @pytest.mark.cpu def test_input_batch_cached_repack_keeps_static_prompt_fields_for_same_composition(): first_state = _make_input_batch_state("req-1", 1.0) first_state.prompt_embeds = torch.ones(1, 2, 3) first_state.prompt_embeds_mask = torch.ones(1, 2, dtype=torch.bool) batch = InputBatch.make_batch([first_state]) replacement_state = _make_input_batch_state("req-1", 2.0) replacement_state.prompt_embeds = torch.full((1, 2, 3), 2.0) replacement_state.prompt_embeds_mask = torch.ones(1, 2, dtype=torch.bool) repacked = InputBatch.make_batch([replacement_state], cached_batch=batch) assert repacked.states[0] is replacement_state torch.testing.assert_close(repacked.latents, replacement_state.latents) torch.testing.assert_close(repacked.prompt_embeds, torch.ones(1, 2, 3)) @pytest.mark.cpu def test_make_batch_identity_mapping_never_reads_the_device_tensor_back(monkeypatch): """The default index mapping is built on the host; reading it back from the device would sync every step.""" def forbidden(self, *args, **kwargs): raise AssertionError("make_batch read a tensor back to the host") monkeypatch.setattr(torch.Tensor, "tolist", forbidden) monkeypatch.setattr(torch.Tensor, "cpu", forbidden) states = [_make_input_batch_state("req-1", 1.0), _make_input_batch_state("req-2", 2.0)] batch = InputBatch.make_batch(states) repacked = InputBatch.make_batch(states, cached_batch=batch) assert repacked is batch assert list(batch.idx_mapping_np) == [0, 1] assert batch.idx_mapping.dtype is torch.int32 and batch.idx_mapping.shape == (2,) assert batch.request_ids == ["req-1", "req-2"] @pytest.mark.cpu def test_make_batch_explicit_mapping_still_selects_and_orders_states(): states = [_make_input_batch_state("req-1", 1.0), _make_input_batch_state("req-2", 2.0)] batch = InputBatch.make_batch(states, idx_mapping=torch.tensor([1, 0])) assert batch.request_ids == ["req-2", "req-1"] assert list(batch.idx_mapping_np) == [1, 0] torch.testing.assert_close(batch.latents, torch.tensor([[2.0], [1.0]])) @pytest.mark.cpu def test_step_profiler_reports_denoise_step_as_diffuse(monkeypatch): monkeypatch.setattr(current_omni_platform, "synchronize", lambda: None) pipeline = _AutoDenoiseProfilerPipeline() assert pipeline.denoise_step() == "ok" assert any(key.endswith(".diffuse") for key in pipeline.stage_durations) assert not any(key.endswith(".denoise_step") for key in pipeline.stage_durations) @pytest.mark.cpu class TestRunner: """DiffusionModelRunner.execute_stepwise""" @pytest.fixture(autouse=True) def mock_platform_memory(self, monkeypatch): monkeypatch.setattr(model_runner_module.current_omni_platform, "is_available", lambda: True) monkeypatch.setattr(model_runner_module.current_omni_platform, "reset_peak_memory_stats", lambda: None) monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_reserved", lambda: 0) monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_allocated", lambda: 0) @pytest.mark.parametrize( ("diffusion_kv_mode", "should_cleanup"), [ (DiffusionKVCacheMode.DENSE_LEGACY, False), (DiffusionKVCacheMode.PAGED_SCHEDULER, True), ], ) def test_scheduler_finished_cleanup_only_runs_for_paged_kv( self, mocker: MockerFixture, diffusion_kv_mode: DiffusionKVCacheMode, should_cleanup: bool, ): runner = _make_runner(diffusion_kv_mode) cleanup = mocker.patch.object(runner, "remove_diffusion_kv_requests") expected_output = object() mocker.patch.object(runner, "_execute_stepwise_core", return_value=expected_output) scheduler_output = DiffusionSchedulerOutput( step_id=1, scheduled_new_reqs=[], scheduled_cached_reqs=CachedRequestData.make_empty(), finished_req_ids={"finished-request"}, num_running_reqs=0, num_waiting_reqs=0, ) output = DiffusionModelRunner.execute_stepwise(runner, scheduler_output) assert output is expected_output if should_cleanup: cleanup.assert_called_once_with(["finished-request"]) else: cleanup.assert_not_called() @pytest.mark.parametrize( ("diffusion_kv_mode", "should_cleanup"), [ (DiffusionKVCacheMode.DENSE_LEGACY, False), (DiffusionKVCacheMode.PAGED_SCHEDULER, True), ], ) def test_terminal_cleanup_only_runs_for_paged_kv( self, monkeypatch: pytest.MonkeyPatch, mocker: MockerFixture, diffusion_kv_mode: DiffusionKVCacheMode, should_cleanup: bool, ): runner = _make_runner(diffusion_kv_mode) cleanup = mocker.patch.object(runner, "remove_diffusion_kv_requests") install = mocker.patch.object(runner, "install_diffusion_kv_metadata") req = _make_step_request() scheduler_output = _make_scheduler_output(req, step_id=0) if should_cleanup: scheduler_output.scheduled_new_reqs[0].diffusion_kv_metadata = DiffusionKVMetadata( request_id=req.request_id, allocation_generation=1, sequences=(), ) monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context) first_output = DiffusionModelRunner.execute_stepwise(runner, scheduler_output) assert first_output.get_request_output(req.request_id).finished is False cleanup.assert_not_called() output = DiffusionModelRunner.execute_stepwise(runner, _make_cached_scheduler_output(step_id=1)) assert output.get_request_output(req.request_id).finished is True if should_cleanup: install.assert_called_once_with(scheduler_output.scheduled_new_reqs[0].diffusion_kv_metadata) cleanup.assert_called_once_with([req.request_id]) else: install.assert_not_called() cleanup.assert_not_called() def test_completes_request_and_clears_state(self, monkeypatch): runner = _make_runner() req = _make_step_request() monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context) result = DiffusionModelRunner.execute_stepwise(runner, _make_scheduler_output(req, step_id=0)) first = result.get_request_output("req-1") assert first.request_id == "req-1" assert first.step_index == 1 assert first.finished is False assert first.result is None assert "req-1" in runner.state_cache result = DiffusionModelRunner.execute_stepwise(runner, _make_cached_scheduler_output(step_id=1)) second = result.get_request_output("req-1") assert second.request_id == "req-1" assert second.step_index == 2 assert second.finished is True assert second.result is not None assert second.result.error is None assert torch.equal(second.result.output, torch.tensor([2.0])) assert "req-1" not in runner.state_cache assert runner.pipeline.prepare_calls == 1 assert runner.pipeline.denoise_calls == 2 assert runner.pipeline.scheduler_calls == 2 assert runner.pipeline.decode_calls == 1 @staticmethod def _make_grouping_runner(pipeline, *, grouping: bool): """An AR runner without KV (so it inherits the stepwise entry) over a streaming stub pipeline.""" from vllm_omni.experimental.ar_diffusion.runner import ARDiffusionModelRunner runner = _make_runner(runner_cls=ARDiffusionModelRunner) runner.od_config.streaming_output = True runner.kv_cache = None runner.pipeline = pipeline runner.pipeline.supports_chunk_step_grouping = grouping return runner @staticmethod def _drive_to_completion(runner, num_steps=4): outputs = [runner.execute_stepwise(_make_scheduler_output(_make_step_request(num_steps)))] step = 1 while not outputs[-1].get_request_output("req-1").finished: outputs.append(runner.execute_stepwise(_make_cached_scheduler_output(step_id=step))) step += 1 return [o.get_request_output("req-1") for o in outputs] @pytest.mark.parametrize("grouping", [False, True]) def test_ar_runner_runs_every_step_of_a_chunk_in_one_call_when_the_pipeline_declares_it( self, monkeypatch, grouping ): """With the capability declared, one AR runner call drives a whole chunk; the outputs are the same either way.""" runner = self._make_grouping_runner(_ChunkedStepPipeline(), grouping=grouping) monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context) denoise_calls_after_each_call = [] original = runner.execute_stepwise def counting(scheduler_output): output = original(scheduler_output) denoise_calls_after_each_call.append(runner.pipeline.denoise_calls) return output runner.execute_stepwise = counting per_call = self._drive_to_completion(runner) # Chunk outputs, and the steps they were emitted at, are identical. emitted = [(o.step_index, float(o.result.output[0])) for o in per_call if o.result is not None] assert emitted == [(2, 0.0), (2, 1.0)] assert per_call[-1].finished is True assert runner.pipeline.denoise_calls == 4 and runner.pipeline.scheduler_calls == 4 assert runner.pipeline.decode_calls == 2 assert runner.pipeline.prepare_calls == 1 # One call per chunk with the capability, one per step without; a # grouped call never crosses into the next chunk. assert denoise_calls_after_each_call == ([2, 4] if grouping else [1, 2, 3, 4]) assert "req-1" not in runner.state_cache def test_ar_runner_continues_a_grouped_chunk_as_a_cached_request(self, monkeypatch): """The first step admits the new request; every further step of the chunk is a cached-request continuation.""" runner = self._make_grouping_runner(_ChunkedStepPipeline(), grouping=True) monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context) seen = [] original_core = DiffusionModelRunner._execute_stepwise_core def recording_core(self_, scheduler_output, **kwargs): seen.append(scheduler_output) return original_core(self_, scheduler_output, **kwargs) monkeypatch.setattr(DiffusionModelRunner, "_execute_stepwise_core", recording_core) first = runner.execute_stepwise(_make_scheduler_output(_make_step_request(4), finished_req_ids={"old"})) assert first.get_request_output("req-1").result is not None assert runner.pipeline.prepare_calls == 1 assert len(seen) == 2 assert [n.request_id for n in seen[0].scheduled_new_reqs] == ["req-1"] continuation = seen[1] assert continuation.scheduled_new_reqs == [] assert continuation.scheduled_cached_reqs.request_ids == ["req-1"] assert continuation.finished_req_ids == set() assert continuation.kv_prefetch_job is None and continuation.kv_connector_metadata is None assert continuation.step_id == seen[0].step_id def test_ar_runner_grouped_steps_see_the_same_inputs_as_per_step_calls(self, monkeypatch): """Every grouped step sees the latents and timestep step_scheduler just advanced, exactly as per-step calls do.""" finals = {} seen = {} for grouping in (False, True): runner = self._make_grouping_runner(_BatchDependentChunkPipeline(), grouping=grouping) monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context) self._drive_to_completion(runner) finals[grouping] = float(runner.pipeline.final_latent) seen[grouping] = runner.pipeline.seen # 1 -> +(1+10)=12 -> +(12+5)=29 -> +(29+10)=68 -> +(68+5)=141 in both modes. assert seen[False] == [(1.0, 10.0), (12.0, 5.0), (29.0, 10.0), (68.0, 5.0)] assert seen[True] == seen[False] assert finals[True] == finals[False] == 141.0 def test_ar_runner_grouped_chunk_stops_at_a_failing_step_and_drops_the_state(self, monkeypatch): """A failure mid-chunk ends the grouped call with that step's error; nothing runs after it.""" runner = self._make_grouping_runner(_FailsMidChunkPipeline(), grouping=True) monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context) output = runner.execute_stepwise(_make_scheduler_output(_make_step_request(4))) failed = output.get_request_output("req-1") assert failed.finished is True assert failed.result is not None and "boom mid chunk" in failed.result.error assert runner.pipeline.denoise_calls == 2 and runner.pipeline.decode_calls == 0 assert "req-1" not in runner.state_cache def test_ar_runner_fails_a_chunk_that_never_advances(self, monkeypatch): pipeline = _ChunkedStepPipeline() monkeypatch.setattr(pipeline, "step_scheduler", lambda *args, **kwargs: None) runner = self._make_grouping_runner(pipeline, grouping=True) monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context) output = runner.execute_stepwise(_make_scheduler_output(_make_step_request(4))) failed = output.get_request_output("req-1") assert failed.finished assert "within 2 denoise steps" in failed.result.error assert pipeline.denoise_calls == 2 assert "req-1" not in runner.state_cache def test_base_runner_never_groups_and_the_ar_policy_needs_capability_streaming_and_one_request(self, monkeypatch): """The shared runner is one step per call; the AR runner groups only a lone streaming request that declares it.""" base = _make_runner() base.od_config.streaming_output = True base.pipeline = _ChunkedStepPipeline() base.pipeline.supports_chunk_step_grouping = True monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context) first = DiffusionModelRunner.execute_stepwise(base, _make_scheduler_output(_make_step_request(4))) assert first.get_request_output("req-1").result is None and base.pipeline.denoise_calls == 1 ar = self._make_grouping_runner(_ChunkedStepPipeline(), grouping=False) one = _make_cached_scheduler_output() two = _make_batch_scheduler_output([_make_step_request(4), _make_step_request(4)]) two.scheduled_new_reqs[1].request_id = two.scheduled_new_reqs[1].req.request_id = "req-2" assert ar._groups_chunk_steps(one) is False # no capability declared ar.pipeline.supports_chunk_step_grouping = True assert ar._groups_chunk_steps(one) is True assert ar._groups_chunk_steps(two) is False # not for a batch ar.od_config.streaming_output = False assert ar._groups_chunk_steps(one) is False # not without streaming output def test_stepwise_output_includes_stage_and_peak_metrics(self, monkeypatch): runner = _make_runner() runner.pipeline = _ProfilingStepPipeline() req = _make_step_request() reset_calls = [] monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context) monkeypatch.setattr( model_runner_module.current_omni_platform, "is_available", lambda: True, ) monkeypatch.setattr( model_runner_module.current_omni_platform, "reset_peak_memory_stats", lambda: reset_calls.append(True), ) monkeypatch.setattr( model_runner_module.current_omni_platform, "max_memory_reserved", lambda: 2 * 1024**2, ) monkeypatch.setattr( model_runner_module.current_omni_platform, "max_memory_allocated", lambda: 1024**2, ) DiffusionModelRunner.execute_stepwise( runner, _make_scheduler_output(req, step_id=0), ) result = DiffusionModelRunner.execute_stepwise( runner, _make_cached_scheduler_output(step_id=1), ) output = result.get_request_output("req-1") assert output.finished is True assert output.result is not None assert output.result.peak_memory_mb == 2 assert output.result.stage_durations == { "QwenImagePipeline.text_encoder.forward": 1.0, "QwenImagePipeline.diffuse": 4.0, "QwenImagePipeline.vae.decode": 3.0, } assert reset_calls == [True] def test_carries_peak_memory_across_stepwise_request_lifecycle(self, monkeypatch): runner = _make_runner() req = _make_step_request() fake_platform = _FakePeakMemoryPlatform([1500.0, 1200.0]) monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context) monkeypatch.setattr(model_runner_module, "current_omni_platform", fake_platform) first = DiffusionModelRunner.execute_stepwise(runner, _make_scheduler_output(req, step_id=0)) first_output = first.get_request_output("req-1") assert first_output.finished is False assert first_output.result is None second = DiffusionModelRunner.execute_stepwise(runner, _make_cached_scheduler_output(step_id=1)) second_output = second.get_request_output("req-1") assert second_output.finished is True assert second_output.result is not None assert second_output.result.peak_memory_mb == pytest.approx(1500.0) def test_rejects_multi_request_step_batch(self): runner = _make_runner() req_1 = _make_step_request() req_2 = _make_step_request() req_2.request_id = "req-2" scheduler_output = DiffusionSchedulerOutput( step_id=0, scheduled_new_reqs=[ NewRequestData(request_id="req-1", req=req_1), NewRequestData(request_id="req-2", req=req_2), ], scheduled_cached_reqs=CachedRequestData.make_empty(), finished_req_ids=set(), num_running_reqs=2, num_waiting_reqs=0, ) result = DiffusionModelRunner.execute_stepwise(runner, scheduler_output) assert len(result) == 2 def test_prepare_error_isolated_from_healthy_request(self, monkeypatch): runner = _make_runner() runner.pipeline = _PerRequestErrorStepPipeline() bad_req = _make_step_request() bad_req.request_id = "req-bad" bad_req.prompt = "fail-prepare" good_req = _make_step_request() good_req.request_id = "req-good" good_req.prompt = "valid" monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context) result = DiffusionModelRunner.execute_stepwise( runner, _make_batch_scheduler_output([bad_req, good_req]), ) bad_output = result.get_request_output("req-bad") assert bad_output.finished is True assert bad_output.result is not None assert bad_output.result.error == "invalid request input" assert bad_output.result.error_status_code == 422 assert bad_output.result.error_type == "UnprocessableEntityError" good_output = result.get_request_output("req-good") assert good_output.finished is False assert good_output.step_index == 1 assert good_output.result is None assert set(runner.state_cache) == {"req-good"} assert runner.pipeline.denoise_calls == 1 assert runner.pipeline.scheduler_calls == 1 def test_prepare_error_on_other_rank_is_isolated_via_all_reduce(self, monkeypatch): """P1-B: rank-0-only pipeline work (e.g. H3 reference-video prep) can raise on rank 0 while non-zero ranks return cleanly and then hang on a downstream ``dist.broadcast``. Beyond the pipeline-side guard, the runner all-reduces a per-request failure flag across the DiT group so a rank that ``prepare_encode`` returned cleanly on still skips the request and does not enter the denoise step batch. """ runner = _make_runner() runner.pipeline = _StepPipeline() # locally always succeeds monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context) # Simulate a peer DiT rank having reported failure. We patch the # helper directly (rather than ``torch.distributed`` broadly) because # ``execute_stepwise`` also queries ``torch.distributed.get_rank``. monkeypatch.setattr( model_runner_module, "_dit_any_rank_failed", lambda local_failed: True, ) req = _make_step_request(num_inference_steps=1) result = DiffusionModelRunner.execute_stepwise( runner, _make_scheduler_output(req), ) output = result.get_request_output("req-1") assert output.finished is True assert output.result is not None assert output.result.error is not None assert "another DiT rank" in output.result.error # A rank that skipped the request must not have driven a denoise step. assert runner.pipeline.denoise_calls == 0 assert runner.state_cache == {} def test_peer_prepare_encode_failure_skips_interaction_session_initialization(self, monkeypatch): """One rank's failure in ``_prepare_batch_inputs -> prepare_encode`` must be handled before all ranks initialize interaction sessions. ``maybe_prepare_initial_session`` calls ``synchronized_monotonic_time()`` (a broadcast). If this rank's ``prepare_encode`` succeeded while a peer failed, entering that broadcast while the peer enters the failure all-reduce would hang. """ from unittest.mock import MagicMock from vllm_omni.diffusion.interaction.types import ChunkMediaSpec, InteractionChunkMetadata class _InteractionStepPipeline(_StepPipeline): def __init__(self): super().__init__() self.prepare_next_chunk_calls = 0 def peek_chunk_media(self, state): del state return ChunkMediaSpec(num_media_frames=8, fps=16.0, num_latent_frames=8) def apply_interaction_at_chunk_boundary(self, state): del state def prepare_next_chunk(self, state): del state self.prepare_next_chunk_calls += 1 runner = _make_runner() pipeline = _InteractionStepPipeline() runner.pipeline = pipeline coordinator = MagicMock() coordinator.maybe_prepare_initial_session.return_value = InteractionChunkMetadata( started_event_ids=[], active_event_ids=[], completed_event_ids=[], ) runner._interaction_coordinator = coordinator monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context) # Peer failed during prepare_encode; this rank's local encode succeeds. monkeypatch.setattr( model_runner_module, "_dit_any_rank_failed", lambda local_failed: True, ) req = _make_step_request(num_inference_steps=1) result = DiffusionModelRunner.execute_stepwise(runner, _make_scheduler_output(req)) output = result.get_request_output("req-1") assert output.finished is True assert output.result is not None assert "another DiT rank" in (output.result.error or "") assert pipeline.prepare_calls == 1 coordinator.maybe_prepare_initial_session.assert_not_called() assert pipeline.prepare_next_chunk_calls == 0 assert pipeline.denoise_calls == 0 assert runner.state_cache == {} def test_receives_kv_payload_before_prepare_encode(self, monkeypatch): runner = _make_runner() captured: dict[str, object] = {} kv_payload = object() class _CapturingStepPipeline(_StepPipeline): def prepare_encode(self, state, **kwargs): captured["past_key_values"] = getattr(state.sampling, "past_key_values", None) return super().prepare_encode(state, **kwargs) class _KVTransferManager: def receive_multi_kv_cache_distributed(self, req, cfg_kv_collect_func=None, target_device=None): captured["cfg_kv_collect_func"] = cfg_kv_collect_func captured["target_device"] = target_device req.sampling_params.past_key_values = kv_payload runner.pipeline = _CapturingStepPipeline() runner.pipeline.device = torch.device("cpu") runner.od_config.cfg_kv_collect_func = "collect-cfg" runner.kv_transfer_manager = _KVTransferManager() monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context) req = _make_step_request() DiffusionModelRunner.execute_stepwise(runner, _make_scheduler_output(req, step_id=0)) assert captured["past_key_values"] is kv_payload assert captured["cfg_kv_collect_func"] == "collect-cfg" assert captured["target_device"] == torch.device("cpu") assert getattr(req.sampling_params, "past_key_values", None) is None def test_rejects_missing_cached_state(self): runner = _make_runner() with pytest.raises(ValueError, match="Missing cached state"): DiffusionModelRunner.execute_stepwise(runner, _make_cached_scheduler_output(request_id="req-missing")) def test_interrupt_marks_request_finished_and_clears_state(self, monkeypatch): runner = _make_runner() runner.pipeline = _InterruptingStepPipeline() req = _make_step_request() monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context) result = DiffusionModelRunner.execute_stepwise(runner, _make_scheduler_output(req, step_id=0)) output = result.get_request_output("req-1") assert output.request_id == "req-1" assert output.step_index == 0 assert output.finished is True assert output.result is not None assert output.result.error == "stepwise denoise interrupted" assert "req-1" not in runner.state_cache assert runner.pipeline.prepare_calls == 1 assert runner.pipeline.denoise_calls == 1 assert runner.pipeline.scheduler_calls == 0 assert runner.pipeline.decode_calls == 0 def test_load_model_rejects_unsupported_step_execution(self, monkeypatch): class _RequestOnlyPipeline: pass class _FakeLoader: def __init__(self, *args, **kwargs): del args, kwargs def load_model(self, **kwargs): del kwargs return _RequestOnlyPipeline() class _FakeProfiler: consumed_memory = 0 def __enter__(self): return self def __exit__(self, exc_type, exc, tb): del exc_type, exc, tb return False runner = object.__new__(DiffusionModelRunner) runner.vllm_config = _make_vllm_config() runner.od_config = SimpleNamespace( enable_cpu_offload=False, enable_layerwise_offload=False, enforce_eager=True, cache_backend=None, cache_config=None, step_execution=True, model_class_name="RequestOnlyPipeline", parallel_config=SimpleNamespace(use_hsdp=False), streaming_output=False, ) runner.device = torch.device("cpu") runner.pipeline = None runner.cache_backend = None runner.offload_backend = None runner.state_cache = {} runner.kv_transfer_manager = SimpleNamespace() monkeypatch.setattr(model_runner_module, "DiffusersPipelineLoader", _FakeLoader) monkeypatch.setattr(model_runner_module, "DeviceMemoryProfiler", _FakeProfiler) monkeypatch.setattr( model_runner_module, "enable_offload_backend", lambda _config, pipeline, **_kwargs: (pipeline, None), ) monkeypatch.setattr(model_runner_module, "get_cache_backend", lambda *args, **kwargs: None) with pytest.raises(ValueError, match="RequestOnlyPipeline"): DiffusionModelRunner.load_model(runner) class _RecordingLoRAManager: def __init__(self) -> None: self.calls: list[tuple[object | None, float]] = [] def set_active_adapter(self, adapter, scale: float = 1.0) -> None: self.calls.append((adapter, scale)) def _make_step_worker(lora_manager=None, *, expected_output=None): """Build a bare DiffusionWorker primed for execute_stepwise tests.""" worker = object.__new__(DiffusionWorker) worker.lora_manager = lora_manager worker._step_lora_state = {} output = expected_output if expected_output is not None else RunnerOutput(request_id="req-1") worker.model_runner = SimpleNamespace(execute_stepwise=lambda arg: output) return worker @pytest.mark.cpu class TestWorker: """DiffusionWorker.execute_stepwise""" def test_delegates_to_model_runner(self): expected = RunnerOutput(request_id="req-1", step_index=1, finished=False, result=None) worker = _make_step_worker(expected_output=expected) scheduler_output = _make_scheduler_output(_make_engine_request("req-1"), request_id="req-1") output = DiffusionWorker.execute_stepwise(worker, scheduler_output) assert output is expected def test_deactivates_lora_when_request_has_no_adapter(self): manager = _RecordingLoRAManager() worker = _make_step_worker(lora_manager=manager) scheduler_output = _make_scheduler_output(_make_engine_request("req-1"), request_id="req-1") DiffusionWorker.execute_stepwise(worker, scheduler_output) assert manager.calls == [(None, 1.0)] def test_activates_lora_for_step_requests(self): from vllm_omni.lora.request import LoRARequest lora_request = LoRARequest(lora_name="adapter", lora_int_id=7, lora_path="/tmp/lora") request = _make_engine_request("req-1") request.sampling_params.lora_request = lora_request request.sampling_params.lora_scale = 0.75 manager = _RecordingLoRAManager() worker = _make_step_worker(lora_manager=manager) scheduler_output = _make_scheduler_output(request, request_id="req-1") DiffusionWorker.execute_stepwise(worker, scheduler_output) assert manager.calls == [(lora_request, 0.75)] def test_recovers_lora_for_cached_step_requests(self): from vllm_omni.lora.request import LoRARequest lora_request = LoRARequest(lora_name="adapter", lora_int_id=11, lora_path="/tmp/lora") request = _make_engine_request("req-1") request.sampling_params.lora_request = lora_request request.sampling_params.lora_scale = 0.5 manager = _RecordingLoRAManager() worker = _make_step_worker(lora_manager=manager) first = _make_scheduler_output(request, request_id="req-1") second = _make_cached_scheduler_output(request_id="req-1", step_id=1) DiffusionWorker.execute_stepwise(worker, first) DiffusionWorker.execute_stepwise(worker, second) assert manager.calls == [(lora_request, 0.5), (lora_request, 0.5)] def test_activates_single_lora_for_homogeneous_batch(self): """Multiple requests sharing the same LoRA → exactly one activation, and every request id is registered in ``_step_lora_state``.""" from vllm_omni.lora.request import LoRARequest lora_request = LoRARequest(lora_name="adapter", lora_int_id=9, lora_path="/tmp/lora") reqs = [] for rid in ("req-1", "req-2", "req-3"): r = _make_engine_request(rid) r.sampling_params.lora_request = lora_request r.sampling_params.lora_scale = 0.6 reqs.append(r) manager = _RecordingLoRAManager() worker = _make_step_worker(lora_manager=manager) scheduler_output = _make_batch_scheduler_output(reqs) DiffusionWorker.execute_stepwise(worker, scheduler_output) assert manager.calls == [(lora_request, 0.6)] assert set(worker._step_lora_state) == {"req-1", "req-2", "req-3"} for entry in worker._step_lora_state.values(): assert entry == (lora_request, 0.6) def test_evicts_step_lora_state_for_finished_requests(self): from vllm_omni.lora.request import LoRARequest lora_request = LoRARequest(lora_name="adapter", lora_int_id=3, lora_path="/tmp/lora") finishing = _make_engine_request("req-1") finishing.sampling_params.lora_request = lora_request next_request = _make_engine_request("req-2") next_request.sampling_params.lora_request = lora_request worker = _make_step_worker(lora_manager=_RecordingLoRAManager()) first = _make_scheduler_output(finishing, request_id="req-1") next_batch = _make_scheduler_output( next_request, request_id="req-2", step_id=1, finished_req_ids={"req-1"}, ) DiffusionWorker.execute_stepwise(worker, first) assert "req-1" in worker._step_lora_state DiffusionWorker.execute_stepwise(worker, next_batch) assert "req-1" not in worker._step_lora_state assert worker._step_lora_state == {"req-2": (lora_request, 1.0)} @pytest.mark.cpu class TestExecutor: """MultiprocDiffusionExecutor.execute_step""" def test_execute_step_passes_through_runner_output(self, mocker: MockerFixture): executor = object.__new__(MultiprocDiffusionExecutor) executor.od_config = SimpleNamespace(streaming_output=False) executor._ensure_open = lambda: None expected = RunnerOutput(request_id="req-step", step_index=1, finished=False, result=None) executor.collective_rpc = mocker.Mock(return_value=expected) request = _make_engine_request("req-step", num_inference_steps=2) scheduler_output = _make_scheduler_output(request, request_id="req-step") output = MultiprocDiffusionExecutor.execute_step(executor, scheduler_output) assert output is expected @pytest.mark.cpu class TestEngine: """Step-execution paths in DiffusionEngine.add_req_and_wait_for_response""" @pytest.mark.parametrize( ("execute_fn", "expected_error"), [ ( lambda _: RunnerOutput( request_id="req-error", step_index=1, finished=True, result=DiffusionOutput(error="boom"), ), "boom", ), ( lambda _: (_ for _ in ()).throw(RuntimeError("gpu on fire")), "gpu on fire", ), ], ) def test_step_engine_returns_error(self, execute_fn, expected_error, mocker: MockerFixture): scheduler = StepScheduler() scheduler.initialize(SimpleNamespace()) engine = _make_engine(scheduler, execute_fn=execute_fn) output = engine.add_req_and_wait_for_response(_make_engine_request("req-error", num_inference_steps=2)) assert output.output is None assert expected_error in output.error def test_step_execution_completes(self, mocker: MockerFixture): scheduler = StepScheduler() scheduler.initialize(SimpleNamespace()) engine = _make_engine(scheduler) request = _make_engine_request("req-step", num_inference_steps=2) call_count = {"n": 0} def execute_fn(_): call_count["n"] += 1 finished = call_count["n"] == 2 return RunnerOutput( request_id="req-step", step_index=call_count["n"], finished=finished, result=(DiffusionOutput(output=torch.tensor([2.0])) if finished else None), ) engine.execute_fn = execute_fn output = engine.add_req_and_wait_for_response(request) assert call_count["n"] == 2 assert output.error is None assert torch.equal(output.output, torch.tensor([2.0])) def test_step_abort_stops_rescheduling_after_first_step(self, mocker: MockerFixture): scheduler = StepScheduler() scheduler.initialize(SimpleNamespace()) engine = _make_engine(scheduler) request = _make_engine_request("req-stop", num_inference_steps=4) step = {"n": 0} def execute_fn(_): step["n"] += 1 engine.abort("req-stop") return RunnerOutput( request_id="req-stop", step_index=1, finished=False, result=None, ) engine.execute_fn = execute_fn output = engine.add_req_and_wait_for_response(request) assert step["n"] == 1 _assert_aborted_output(output, "req-stop") def test_step_abort_after_reschedule_returns_aborted_output(self, mocker: MockerFixture): scheduler = StepScheduler() scheduler.initialize(SimpleNamespace()) engine = _make_engine(scheduler) request = _make_engine_request("req-mid", num_inference_steps=4) step = {"n": 0} def execute_fn(sched_output): step["n"] += 1 if step["n"] == 2: assert sched_output == _make_cached_scheduler_output("req-mid", step_id=1) engine.abort("req-mid") return RunnerOutput( request_id="req-mid", step_index=step["n"], finished=False, result=None, ) engine.execute_fn = execute_fn output = engine.add_req_and_wait_for_response(request) assert step["n"] == 2 _assert_aborted_output(output, "req-mid") def test_finished_step_without_result_returns_error(self, mocker: MockerFixture): scheduler = StepScheduler() scheduler.initialize(SimpleNamespace()) engine = _make_engine( scheduler, execute_fn=lambda _: RunnerOutput( request_id="req-missing", step_index=1, finished=True, result=None, ), ) output = engine.add_req_and_wait_for_response(_make_engine_request("req-missing", num_inference_steps=1)) assert output.output is None assert output.error == "Diffusion execution finished without a final output." @pytest.mark.cpu class TestIPC: def test_pack_unpack_runner_output_shm(self): tensor = torch.zeros(300_000, dtype=torch.float32) output = RunnerOutput(request_id="req-1", finished=True, result=DiffusionOutput(output=tensor)) packed = pack_diffusion_output_shm(output) assert isinstance(packed.result.output, dict) assert packed.result.output["__tensor_shm__"] is True unpacked = unpack_diffusion_output_shm(packed) assert isinstance(unpacked.result.output, torch.Tensor) torch.testing.assert_close(unpacked.result.output, tensor) @pytest.mark.cpu class TestSupportedPipelines: """Step-execution protocol checks for supported pipelines.""" def test_default_stage_config_includes_step_execution(self): stage_cfg = StageConfigFactory.create_default_diffusion( { "step_execution": True, } )[0] assert stage_cfg["engine_args"]["step_execution"] is True def test_qwen_image_supports_step_execution(self): from vllm_omni.diffusion.models.interface import SupportsStepExecution, supports_step_execution from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image import QwenImagePipeline # Avoid loading model weights; protocol membership depends on the class contract. pipeline = object.__new__(QwenImagePipeline) assert pipeline.supports_step_execution is True assert supports_step_execution(pipeline) is True assert isinstance(pipeline, SupportsStepExecution) is True @hardware_test( res={"cuda": "L4"}, num_cards=2, ) def test_execute_stepwise_with_ulysses_parallel(): world_size = 2 if current_omni_platform.get_device_count() < world_size: pytest.skip(f"Test requires {world_size} devices") torch.multiprocessing.spawn( _distributed_step_worker, args=(world_size, "ulysses", "29540"), nprocs=world_size, ) @hardware_test( res={"cuda": "L4"}, num_cards=2, ) def test_execute_stepwise_with_ring_parallel(): world_size = 2 if current_omni_platform.get_device_count() < world_size: pytest.skip(f"Test requires {world_size} devices") torch.multiprocessing.spawn( _distributed_step_worker, args=(world_size, "ring", "29541"), nprocs=world_size, ) @hardware_test( res={"cuda": "L4"}, num_cards=2, ) def test_execute_stepwise_with_cfg_parallel(): world_size = 2 if current_omni_platform.get_device_count() < world_size: pytest.skip(f"Test requires {world_size} devices") torch.multiprocessing.spawn( _distributed_step_worker, args=(world_size, "cfg", "29542"), nprocs=world_size, )