[Bugfix] Stop cancelled MiniMax-H3 requests at denoising boundaries (#7844)

Signed-off-by: Gao Han <hgaoaf@connect.ust.hk>

tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py

@@ -3494,3 +3494,48 @@ def test_long_video_shape_requires_explicit_opt_in(duration):
     assert frames >= duration * 24 and frames % 17 == 5
     assert video_t == (frames - 5) // 17 * 5 + 2
     assert audio_t == round(frames / 24 * 40)
+
+
+@pytest.mark.parametrize("preencode", [False, True])
+@pytest.mark.parametrize("cancel_phase", ["before_prepare", "prepare", "diffuse"])
+def test_request_cancellation_at_prepare_and_decode_boundaries(preencode, cancel_phase, monkeypatch):
+    from vllm_omni.diffusion.cancellation import RequestCancellationRegistry, request_cancellation_scope
+    from vllm_omni.diffusion.data import DiffusionRequestAbortedError
+    from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
+    from vllm_omni.platforms import current_omni_platform
+
+    synchronize = Mock()
+    monkeypatch.setattr(current_omni_platform, "synchronize", synchronize)
+    pipeline = object.__new__(MiniMaxH3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    registry = RequestCancellationRegistry()
+    signal = registry.create("request")
+
+    def prepare(*args):
+        if cancel_phase == "prepare":
+            registry.cancel(["request"])
+        return {"num_outputs": 1, "seed": 1101, "preencode_mp4": preencode}
+
+    def diffuse(**kwargs):
+        registry.cancel(["request"])
+        return torch.zeros(1), torch.zeros(1)
+
+    pipeline._prepare_request_inputs = Mock(side_effect=prepare)
+    pipeline._denoise_kwargs = Mock(return_value={})
+    pipeline.diffuse = Mock(side_effect=diffuse)
+    pipeline.decode = Mock()
+    pipeline.decode_to_mp4 = Mock()
+    try:
+        if cancel_phase == "before_prepare":
+            registry.cancel(["request"])
+        with request_cancellation_scope([signal]), pytest.raises(DiffusionRequestAbortedError):
+            pipeline.forward(_t2va_batch())
+        synchronize.assert_called_once_with()
+        if cancel_phase == "before_prepare":
+            pipeline._prepare_request_inputs.assert_not_called()
+        if cancel_phase != "diffuse":
+            pipeline.diffuse.assert_not_called()
+        pipeline.decode.assert_not_called()
+        pipeline.decode_to_mp4.assert_not_called()
+    finally:
+        registry.close()

tests/diffusion/models/minimax_h3/test_minimax_h3_step_execution.py

@@ -171,6 +171,41 @@ def test_step_execution_matches_request_mode_denoise_loop(num_steps, mocker):
     torch.testing.assert_close(state.extra[mod._STEP_AUDIO_ROWS], reference_audio)
 
 
+def test_request_mode_cancellation_stops_before_next_denoise_step(monkeypatch):
+    from vllm_omni.diffusion.cancellation import RequestCancellationRegistry, request_cancellation_scope
+    from vllm_omni.diffusion.data import DiffusionRequestAbortedError
+    from vllm_omni.diffusion.models.minimax_h3.denoise_loop import minimax_h3_denoise_loop
+    from vllm_omni.platforms import current_omni_platform
+
+    # This test runs real packing/Euler updates on CPU with the small DiT above.
+    monkeypatch.setattr(current_omni_platform, "synchronize", lambda: None)
+    branch, video_rows, audio_rows = _make_branch(text_len=9, latent_t=2, latent_h=4, latent_w=6, audio_t=3, seed=5)
+    registry = RequestCancellationRegistry()
+    signal = registry.create("request")
+    steps = []
+
+    def cancel_after_first_step(step, video, audio):
+        steps.append(step)
+        registry.cancel(["request"])
+
+    try:
+        with request_cancellation_scope([signal]), pytest.raises(DiffusionRequestAbortedError):
+            minimax_h3_denoise_loop(
+                model=_SegmentMeanModel(),
+                positive=branch,
+                initial_video_rows=video_rows,
+                initial_audio_rows=audio_rows,
+                keyframe_cond_rows=None,
+                sigmas_video=_sigmas(6, 12.0),
+                sigmas_audio=_sigmas(6, 3.0),
+                device=torch.device("cpu"),
+                on_step=cancel_after_first_step,
+            )
+    finally:
+        registry.close()
+    assert steps == [0]
+
+
 def test_step_execution_matches_request_mode_with_latent_edits():
     """Masked model rows, row timesteps, and scheduler math match both paths."""
     from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as mod

tests/diffusion/test_diffusion_engine_cleanup.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 import asyncio
 import queue
@@ -432,3 +432,68 @@ def test_finalize_aborted_request_without_async_output_skips_drop() -> None:
 
     assert output.aborted is True
     engine.executor.drop_output.assert_not_called()
+
+
+@pytest.mark.parametrize("scheduler_fails", [False, True])
+def test_fail_engine_releases_cancellation_signals_after_workers_stop(scheduler_fails):
+    from multiprocessing.shared_memory import SharedMemory
+
+    from vllm_omni.diffusion.cancellation import RequestCancellationRegistry
+
+    engine = _make_engine()
+    engine._fail_pending_rpcs = Mock()
+    registry = RequestCancellationRegistry()
+    engine._request_cancellations = registry
+    name = registry.create("running")
+
+    def shutdown():
+        # Cancellation is visible before shutdown, and late readers can attach.
+        reader = SharedMemory(name=name)
+        try:
+            assert reader.buf[0] == 1
+        finally:
+            reader.close()
+
+    engine.executor.shutdown.side_effect = shutdown
+    engine.scheduler.close = Mock(side_effect=RuntimeError("scheduler cleanup failed") if scheduler_fails else None)
+    try:
+        if scheduler_fails:
+            with pytest.raises(RuntimeError, match="scheduler cleanup failed"):
+                engine._fail_engine(RuntimeError("engine failed"))
+        else:
+            engine._fail_engine(RuntimeError("engine failed"))
+            engine.close()
+            engine.executor.shutdown.assert_called_once_with()
+        with pytest.raises(FileNotFoundError):
+            SharedMemory(name=name)
+        assert not registry._signals
+    finally:
+        registry.close()
+
+
+def test_close_retains_cancellation_signal_until_stuck_worker_stops():
+    from multiprocessing.shared_memory import SharedMemory
+
+    from vllm_omni.diffusion.cancellation import RequestCancellationRegistry
+
+    engine = _make_engine()
+    registry = RequestCancellationRegistry()
+    engine._request_cancellations = registry
+    name = registry.create("running")
+    engine.worker_thread = Mock()
+    engine.worker_thread.is_alive.side_effect = [True, True, False, False]
+    try:
+        engine.close()
+        engine.executor.shutdown.assert_not_called()
+        reader = SharedMemory(name=name)
+        try:
+            assert reader.buf[0] == 1
+        finally:
+            reader.close()
+        engine.close()
+        engine.executor.shutdown.assert_called_once_with()
+        with pytest.raises(FileNotFoundError):
+            SharedMemory(name=name)
+        assert not registry._signals
+    finally:
+        registry.close()

tests/diffusion/test_diffusion_model_runner.py

@@ -1411,3 +1411,44 @@ def test_vllm_set_forward_context_implementation(monkeypatch):
             ),
         ),
     ], ERROR_MESSAGE
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+@pytest.mark.parametrize("cancel_all", [False, True])
+def test_execute_model_batch_cancellation_preserves_live_peer(monkeypatch, cancel_all):
+    from vllm_omni.diffusion.cancellation import (
+        RequestCancellationRegistry,
+        check_request_cancellation,
+    )
+    from vllm_omni.platforms import current_omni_platform
+
+    monkeypatch.setattr(current_omni_platform, "synchronize", lambda: None)
+
+    class CancellableBatchPipeline(_BatchPipeline):
+        supports_request_cancellation = True
+
+        def forward(self, batch):
+            check_request_cancellation()
+            return super().forward(batch)
+
+    monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
+    monkeypatch.setattr(model_runner_module, "current_omni_platform", _fake_platform_for_peak_memory())
+    pipeline = CancellableBatchPipeline(outputs=[DiffusionOutput(output="a"), DiffusionOutput(output="b")])
+    runner = _make_batch_runner(pipeline)
+    sched = _make_scheduler_output(num_reqs=2)
+    registry = RequestCancellationRegistry()
+    try:
+        for entry in sched.scheduled_new_reqs:
+            entry.req.cancellation_signal = registry.create(entry.request_id)
+        registry.cancel(["req-0", "req-1"] if cancel_all else ["req-0"])
+        result = runner.execute_model_batch(sched, runner.od_config)
+        assert len(result.runner_outputs) == 2
+        assert [output.result.aborted for output in result.runner_outputs] == [cancel_all, cancel_all]
+        if not cancel_all:
+            assert result.runner_outputs[1].request_id == "req-1"
+            assert result.runner_outputs[1].result.output == "b"
+        else:
+            assert pipeline.last_batch is None
+    finally:
+        registry.close()

tests/diffusion/test_request_cancellation.py

@@ -0,0 +1,276 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Cancellation transport, ownership and collective agreement regressions."""
+
+import multiprocessing as mp
+import queue
+import threading
+from concurrent.futures import ThreadPoolExecutor
+from multiprocessing.shared_memory import SharedMemory
+from unittest.mock import patch
+
+import pytest
+
+from vllm_omni.diffusion.cancellation import (
+    RequestCancellationRegistry,
+    check_request_cancellation,
+    request_cancellation_scope,
+)
+from vllm_omni.diffusion.data import DiffusionRequestAbortedError, OmniDiffusionConfig
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+@pytest.fixture
+def registry(monkeypatch):
+    from vllm_omni.platforms import current_omni_platform
+
+    # These tests exercise the cancellation protocol without device work.
+    monkeypatch.setattr(current_omni_platform, "synchronize", lambda: None)
+    signals = RequestCancellationRegistry()
+    yield signals
+    signals.close()
+
+
+def test_cancel_is_request_scoped_and_id_reuse_is_fresh(registry):
+    first = registry.create("first")
+    other = registry.create("other")
+    registry.cancel(["missing", "first"])
+    with pytest.raises(ValueError, match="Duplicate cancellation request"):
+        registry.create("first")
+    with request_cancellation_scope([first]):
+        with pytest.raises(DiffusionRequestAbortedError):
+            check_request_cancellation()
+    with request_cancellation_scope([other]):
+        check_request_cancellation()
+    registry.finish("first")
+    registry.finish("first")
+    with pytest.raises(FileNotFoundError):
+        SharedMemory(name=first)
+    replacement = registry.create("first")
+    with request_cancellation_scope([replacement]):
+        check_request_cancellation()
+
+
+def test_scope_restores_outer_state_and_does_not_leak_to_threads(registry):
+    signal = registry.create("request")
+    registry.cancel_all()
+    with request_cancellation_scope([signal]):
+        with request_cancellation_scope([None]):
+            check_request_cancellation()
+        with ThreadPoolExecutor(max_workers=1) as executor:
+            executor.submit(check_request_cancellation).result(timeout=10)
+        with pytest.raises(DiffusionRequestAbortedError):
+            check_request_cancellation()
+    check_request_cancellation()
+
+
+def test_mixed_batch_keeps_uncancelled_request_running(registry):
+    first = registry.create("first")
+    other = registry.create("other")
+    with request_cancellation_scope([first, other]):
+        registry.cancel(["first"])
+        check_request_cancellation()
+        registry.cancel(["other"])
+        with pytest.raises(DiffusionRequestAbortedError):
+            check_request_cancellation()
+
+
+def test_checkpoint_observes_cancel_during_device_wait(registry, monkeypatch):
+    from vllm_omni.platforms import current_omni_platform
+
+    first = registry.create("first")
+    signal = registry.create("request")
+    registry.cancel(["first"])
+    # A local cancellation triggers the device wait. A peer can be cancelled
+    # during that wait; re-read before deciding whether the whole batch stops.
+    monkeypatch.setattr(current_omni_platform, "synchronize", lambda: registry.cancel(["request"]))
+    with request_cancellation_scope([first, signal]):
+        with pytest.raises(DiffusionRequestAbortedError):
+            check_request_cancellation(synchronize=True)
+
+
+@pytest.mark.parametrize("synchronize_early", [False, True])
+def test_checkpoint_drains_once_before_abort_only(registry, monkeypatch, synchronize_early):
+    from unittest.mock import Mock
+
+    from vllm_omni.platforms import current_omni_platform
+
+    synchronize = Mock()
+    monkeypatch.setattr(current_omni_platform, "synchronize", synchronize)
+    signal = registry.create("request")
+    with request_cancellation_scope([signal]):
+        check_request_cancellation(synchronize=synchronize_early)
+        synchronize.assert_not_called()
+        # Simulate work queued after a clean checkpoint. Cancellation at the
+        # next checkpoint must drain it before exception unwinding drops refs.
+        pending_work = [object()]
+        synchronize.side_effect = pending_work.clear
+        registry.cancel(["request"])
+        with pytest.raises(DiffusionRequestAbortedError):
+            try:
+                check_request_cancellation(synchronize=synchronize_early)
+            finally:
+                assert not pending_work
+        synchronize.assert_called_once_with()
+
+
+def test_close_removes_all_owned_signals(registry):
+    names = [registry.create(str(i)) for i in range(3)]
+    registry.cancel_all()
+    registry.close()
+    registry.close()
+    for name in names:
+        with pytest.raises(FileNotFoundError):
+            SharedMemory(name=name)
+
+
+def _admission_engine(registry):
+    from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
+    from vllm_omni.diffusion.sched.request_scheduler import RequestScheduler
+
+    engine = object.__new__(DiffusionEngine)
+    engine.scheduler = RequestScheduler()
+    engine.scheduler.initialize(OmniDiffusionConfig())
+    engine._request_cancellations = registry
+    engine._cv = threading.Condition(threading.RLock())
+    engine._closed = False
+    engine._out_streams = {}
+    engine.abort_queue = queue.Queue()
+    return engine
+
+
+def _request(request_id):
+    from vllm_omni.diffusion.request import OmniDiffusionRequest
+    from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+    return OmniDiffusionRequest(
+        prompt="A fox in the snow.", sampling_params=OmniDiffusionSamplingParams(), request_id=request_id
+    )
+
+
+def test_engine_signals_worker_before_busy_loop_drains_abort(registry):
+    from vllm_omni.diffusion.sched.interface import DiffusionRequestStatus
+
+    engine = _admission_engine(registry)
+    request = _request("running")
+    engine._add_prepared_request(request)
+    engine.scheduler.schedule()
+    signal = request.cancellation_signal
+    assert signal is not None
+    with request_cancellation_scope([signal]):
+        engine.abort("running")
+        # The busy loop is still occupied by execution. Only the signal has
+        # changed; scheduler mutation must stay on that loop's thread.
+        assert engine.scheduler.get_request_state("running").status == DiffusionRequestStatus.RUNNING
+        with pytest.raises(DiffusionRequestAbortedError):
+            check_request_cancellation()
+    engine._process_aborts_queue()
+    assert engine._finalize_finished_request("running").aborted
+    assert engine.scheduler.get_request_state("running") is None
+    with pytest.raises(FileNotFoundError):
+        SharedMemory(name=signal)
+
+
+def test_rejected_admission_releases_signal(registry, monkeypatch):
+    engine = _admission_engine(registry)
+    allocated_names = []
+
+    def reject(request):
+        allocated_names.append(request.cancellation_signal)
+        raise ValueError("admission rejected")
+
+    monkeypatch.setattr(engine.scheduler, "add_request", reject)
+    request = _request("rejected")
+    with pytest.raises(ValueError, match="admission rejected"):
+        engine._add_prepared_request(request)
+    assert request.cancellation_signal is None
+    with pytest.raises(FileNotFoundError):
+        SharedMemory(name=allocated_names[0])
+    registry.create("rejected")
+
+
+def _collective_worker(rank, name, rendezvous, ready, first_check, second_check, results):
+    import torch
+    import torch.distributed as dist
+
+    from vllm_omni.diffusion.distributed import parallel_state
+    from vllm_omni.platforms import current_omni_platform
+
+    dist.init_process_group("gloo", init_method=f"file://{rendezvous}", rank=rank, world_size=2)
+    # Exercise a real CPU collective without constructing GPU model groups.
+    group = object.__new__(parallel_state.GroupCoordinator)
+    group.world_size = 2
+    group.cpu_group = dist.group.WORLD
+    parallel_state._WORLD = group
+    try:
+        with (
+            torch.inference_mode(),
+            patch.object(current_omni_platform, "synchronize") as synchronize,
+            request_cancellation_scope([name]),
+        ):
+            dist.barrier()
+            if rank == 0:
+                ready.set()
+            assert first_check.wait(120)
+            check_request_cancellation()
+            results.put((rank, "partial_wave_kept_running"))
+            # A replica missing its token (or all local requests) still votes
+            # false. Both ranks must reach each collective without deadlocking.
+            for absent in ([None], []):
+                with request_cancellation_scope([name] if rank == 0 else absent):
+                    check_request_cancellation()
+            results.put((rank, "missing_signal_kept_running"))
+            synchronize.assert_not_called()
+            assert second_check.wait(120)
+            try:
+                check_request_cancellation()
+            except DiffusionRequestAbortedError:
+                synchronize.assert_called_once_with()
+                results.put((rank, "fully_cancelled_wave_stopped"))
+            else:
+                raise AssertionError("All cancelled ranks must stop")
+    finally:
+        dist.destroy_process_group()
+
+
+def test_spawned_ranks_agree_before_leaving_collective_wave(registry, tmp_path):
+    ctx = mp.get_context("spawn")
+    names = [registry.create(str(rank)) for rank in range(2)]
+    ready, first_check, second_check = ctx.Event(), ctx.Event(), ctx.Event()
+    results = ctx.Queue()
+    workers = [
+        ctx.Process(
+            target=_collective_worker,
+            args=(rank, names[rank], str(tmp_path / "rendezvous"), ready, first_check, second_check, results),
+        )
+        for rank in range(2)
+    ]
+    for worker in workers:
+        worker.start()
+    try:
+        assert ready.wait(120)
+        registry.cancel(["0"])
+        first_check.set()
+        assert {results.get(timeout=120) for _ in range(2 * len(workers))} == {
+            (0, "partial_wave_kept_running"),
+            (1, "partial_wave_kept_running"),
+            (0, "missing_signal_kept_running"),
+            (1, "missing_signal_kept_running"),
+        }
+        registry.cancel(["1"])
+        second_check.set()
+        assert {results.get(timeout=120) for _ in workers} == {
+            (0, "fully_cancelled_wave_stopped"),
+            (1, "fully_cancelled_wave_stopped"),
+        }
+        for worker in workers:
+            worker.join(timeout=30)
+            assert worker.exitcode == 0
+    finally:
+        for worker in workers:
+            if worker.is_alive():
+                worker.terminate()
+            worker.join(timeout=30)
+        results.close()
+        results.join_thread()