diff --git a/docs/serving/videos_api.md b/docs/serving/videos_api.md index da0fa1fd3..053d2590b 100644 --- a/docs/serving/videos_api.md +++ b/docs/serving/videos_api.md @@ -120,7 +120,18 @@ request for execution. `DELETE /v1/videos/{video_id}` issues a bounded engine abort (`VLLM_OMNI_ABORT_TIMEOUT`, default 2s), then cancels the frontend task. Cancellation cleanup is also bounded and best-effort: it confirms -the abort was queued, and the current request batch may still drain. +the abort was submitted. In request execution mode, MiniMax-H3 checks +cancellation at model boundaries, +including before and after input preparation, after each denoising step, and +before decode, and skips the remaining generation when cancelled. Input +preparation or decoding already in progress may finish before the next boundary. +Successful steps do not force device synchronization for cancellation. +Parallel ranks agree before stopping; independent requests sharing a distributed +AllGather offload wave can stop that wave early only when all its requests are +cancelled, so a cancelled request cannot strand its live peers in a collective. +In step execution mode, cancellation is handled by the scheduler between steps; +the request-mode component-boundary checks do not apply within an active step. +Other pipelines may still drain their current request batch. The job is then re-read so a completed save is not orphaned. ### Synchronous Response diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py index 11266cff3..2c776d458 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py +++ b/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() diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_step_execution.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_step_execution.py index 00fdb4cdc..5bba75f93 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_step_execution.py +++ b/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 diff --git a/tests/diffusion/test_diffusion_engine_cleanup.py b/tests/diffusion/test_diffusion_engine_cleanup.py index 0d593fe73..3f138e721 100644 --- a/tests/diffusion/test_diffusion_engine_cleanup.py +++ b/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() diff --git a/tests/diffusion/test_diffusion_model_runner.py b/tests/diffusion/test_diffusion_model_runner.py index b914061d0..86a1e5030 100644 --- a/tests/diffusion/test_diffusion_model_runner.py +++ b/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() diff --git a/tests/diffusion/test_request_cancellation.py b/tests/diffusion/test_request_cancellation.py new file mode 100644 index 000000000..c4ce3785c --- /dev/null +++ b/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() diff --git a/vllm_omni/diffusion/cancellation.py b/vllm_omni/diffusion/cancellation.py new file mode 100644 index 000000000..bf0633aec --- /dev/null +++ b/vllm_omni/diffusion/cancellation.py @@ -0,0 +1,141 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Cooperative cancellation for local, full-forward diffusion workers. + +The engine owns each signal until the worker has returned. Unlike an executor +RPC, writing the signal does not wait behind the forward being cancelled. +Opted-in pipelines call ``check_request_cancellation`` at safe boundaries. +""" + +from __future__ import annotations + +import threading +from collections.abc import Iterable, Iterator, Sequence +from contextlib import contextmanager +from contextvars import ContextVar +from multiprocessing.shared_memory import SharedMemory + + +class RequestCancellationRegistry: + """Engine-owned signals; scheduler mutation remains on the engine thread.""" + + def __init__(self) -> None: + self._signals: dict[str, SharedMemory] = {} + self._lock = threading.Lock() + + def create(self, request_id: str) -> str: + with self._lock: + if request_id in self._signals: + raise ValueError(f"Duplicate cancellation request: {request_id}") + signal = SharedMemory(create=True, size=1) + signal.buf[0] = 0 + self._signals[request_id] = signal + return signal.name + + def cancel(self, request_ids: Iterable[str]) -> None: + with self._lock: + for request_id in request_ids: + signal = self._signals.get(request_id) + if signal is not None: + signal.buf[0] = 1 + + def cancel_all(self) -> None: + with self._lock: + for signal in self._signals.values(): + signal.buf[0] = 1 + + @staticmethod + def _dispose(signal: SharedMemory) -> None: + try: + signal.unlink() + except FileNotFoundError: + # A process-level resource tracker may already have removed it. + pass + finally: + signal.close() + + def finish(self, request_id: str) -> None: + """Release only after execution returns, including aborted execution.""" + with self._lock: + signal = self._signals.pop(request_id, None) + if signal is not None: + self._dispose(signal) + + def close(self) -> None: + """Release remaining signals after the executor has shut down.""" + with self._lock: + for signal in self._signals.values(): + self._dispose(signal) + self._signals.clear() + + +_current_signals: ContextVar[tuple[SharedMemory | None, ...] | None] = ContextVar( + "diffusion_request_cancellation", default=None +) + + +@contextmanager +def request_cancellation_scope(signal_names: Sequence[str | None], *, enabled: bool = True) -> Iterator[None]: + """Attach worker readers without taking ownership of the engine's names.""" + signals: list[SharedMemory | None] = [] + try: + for name in signal_names if enabled else (): + signals.append(SharedMemory(name=name) if name is not None else None) + # An opted-in rank without a local signal must still join every + # checkpoint collective and vote "not cancelled". None means only + # that this execution does not use cooperative cancellation at all. + token = _current_signals.set(tuple(signals) if enabled else None) + try: + yield + finally: + _current_signals.reset(token) + finally: + for signal in signals: + if signal is not None: + signal.close() + + +def check_request_cancellation(*, synchronize: bool = False) -> None: + """Stop a cancelled execution wave without stranding a peer's collectives. + + ``synchronize=True`` drains queued device work only when cancellation has + been requested locally, before re-reading the flags. Every abort drains queued + device work before unwinding. Successful steps retain asynchronous execution + when no local cancellation is observed. + Independent requests coupled by an AllGather offload wave must all be cancelled + before the wave can exit; cancelling one must not abort its live peers. + """ + signals = _current_signals.get() + if signals is None: + return + + import torch + + locally_requested = any(signal is not None and signal.buf[0] for signal in signals) + device_drained = False + if synchronize and locally_requested: + from vllm_omni.platforms import current_omni_platform + + current_omni_platform.synchronize() + device_drained = True + + cancelled = bool(signals) and all(signal is not None and signal.buf[0] for signal in signals) + if torch.distributed.is_initialized(): + from vllm_omni.diffusion.distributed.parallel_state import get_world_group + + group = get_world_group() + if group.world_size > 1: + # Reading a shared byte alone is insufficient: ranks can observe a + # concurrent write on different steps. Make one collective decision. + flag = torch.tensor([int(cancelled)], dtype=torch.int32, device="cpu") + torch.distributed.all_reduce(flag, op=torch.distributed.ReduceOp.MIN, group=group.cpu_group) + cancelled = bool(flag.item()) + + if cancelled: + from vllm_omni.diffusion.data import DiffusionRequestAbortedError + + if not device_drained: + from vllm_omni.platforms import current_omni_platform + + current_omni_platform.synchronize() + raise DiffusionRequestAbortedError("Request cancelled at a model execution boundary") diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index babfac099..e2562ad12 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/vllm_omni/diffusion/diffusion_engine.py @@ -25,6 +25,7 @@ from vllm.utils.import_utils import resolve_obj_by_qualname from vllm.v1.engine.exceptions import EngineDeadError from vllm.v1.kv_cache_interface import KVCacheConfig +from vllm_omni.diffusion.cancellation import RequestCancellationRegistry from vllm_omni.diffusion.data import ( DIFFUSION_REQUEST_LIFECYCLE_KEY, DIFFUSION_REQUEST_STARTED, @@ -180,6 +181,19 @@ def _max_num_seqs(od_config: OmniDiffusionConfig) -> int: return 1 +def supports_request_cancellation(od_config: OmniDiffusionConfig) -> bool: + """Whether the local pipeline checks cooperative cancellation boundaries.""" + model_cls = _resolve_custom_pipeline_cls(getattr(od_config, "custom_pipeline_args", None)) + if model_cls is None: + name = ( + "DiffusersAdapterPipeline" + if uses_diffusers_adapter(od_config) + else getattr(od_config, "model_class_name", None) + ) + model_cls = DiffusionModelRegistry._try_load_model_cls(name) + return getattr(model_cls, "supports_request_cancellation", False) is True + + def _uses_dlo_dp_concurrency(od_config: OmniDiffusionConfig) -> bool: parallel_config = getattr(od_config, "parallel_config", None) dp_size = getattr(parallel_config, "data_parallel_size", 1) @@ -230,6 +244,8 @@ class DiffusionEngine: # don't hit AttributeError when _busy_loop accesses them. dp_concurrent: bool = False _scheduling_paused: bool = False + # Disabled until runtime initialization resolves the pipeline capability. + _request_cancellations: RequestCancellationRegistry | None = None def __init__( self, @@ -393,6 +409,9 @@ class DiffusionEngine: self._shutting_down = False self._shutdown_complete = False self.abort_queue: queue.Queue[str] = queue.Queue() + self._request_cancellations = ( + RequestCancellationRegistry() if supports_request_cancellation(self.od_config) else None + ) self._rpc_queue: queue.Queue[_RpcTask] = queue.Queue() # pause_scheduler(mode="keep"): no new batch is scheduled while set. self._scheduling_paused = False @@ -913,6 +932,8 @@ class DiffusionEngine: # path, and repeating executor.shutdown() is unsafe. self._shutting_down = True self._closed = True + if self._request_cancellations is not None: + self._request_cancellations.cancel_all() if self.stop_event is not None: self.stop_event.set() streams = list(self._out_streams.values()) @@ -925,7 +946,13 @@ class DiffusionEngine: # If Worker shutdown fails, retain the Scheduler reservations. A # remote producer may still be writing into those allocations. self.executor.shutdown() - self.scheduler.close() + try: + self.scheduler.close() + finally: + # Workers are down, so readers cannot attach or use the signals. + # Release them even if scheduler cleanup itself fails. + if self._request_cancellations is not None: + self._request_cancellations.close() self._shutdown_complete = True def _emit_finished_outputs( @@ -1106,7 +1133,15 @@ class DiffusionEngine: if self._closed: raise RuntimeError("DiffusionEngine is closed.") queue: asyncio.Queue[DiffusionOutput] = asyncio.Queue() - request_id = self.scheduler.add_request(request) + if self._request_cancellations is not None: + request.cancellation_signal = self._request_cancellations.create(request.request_id) + try: + request_id = self.scheduler.add_request(request) + except BaseException: + if self._request_cancellations is not None: + self._request_cancellations.finish(request.request_id) + request.cancellation_signal = None + raise self._out_streams[request_id] = queue self._cv.notify_all() @@ -1517,6 +1552,8 @@ class DiffusionEngine: return if not self._closed: self._closed = True + if self._request_cancellations is not None: + self._request_cancellations.cancel_all() if self.stop_event is not None: self.stop_event.set() pending_streams = list(self._out_streams.values()) @@ -1532,6 +1569,9 @@ class DiffusionEngine: if worker_thread.is_alive(): worker_thread.join(timeout=10) if worker_thread.is_alive(): + # Keep cancellation names available: an in-flight worker may + # not have attached its readers yet. A later close releases + # them after execution and executor shutdown have completed. logger.warning( "Worker thread did not terminate within 10s; scheduler and executor shutdown will be deferred." ) @@ -1550,6 +1590,8 @@ class DiffusionEngine: else: self.scheduler.close() self.executor.shutdown() + if self._request_cancellations is not None: + self._request_cancellations.close() self._shutdown_complete = True def abort(self, request_id: str | Iterable[str]) -> None: @@ -1558,6 +1600,10 @@ class DiffusionEngine: with self._cv: if self._closed: return + if self._request_cancellations is not None: + # Do not queue this behind the full-forward executor call. + # Scheduler state is still mutated only by the busy loop. + self._request_cancellations.cancel(request_ids) for req_id in request_ids: self.abort_queue.put(req_id) self._cv.notify_all() @@ -1601,6 +1647,8 @@ class DiffusionEngine: state = self.scheduler.get_request_state(request_id) popped_state = self.scheduler.pop_request_state(request_id) state = state or popped_state + if self._request_cancellations is not None: + self._request_cancellations.finish(request_id) if state is None: raise RuntimeError(f"Diffusion scheduler lost state for request {request_id}.") diff --git a/vllm_omni/diffusion/models/minimax_h3/denoise_loop.py b/vllm_omni/diffusion/models/minimax_h3/denoise_loop.py index 1bd429b4d..1b358a070 100644 --- a/vllm_omni/diffusion/models/minimax_h3/denoise_loop.py +++ b/vllm_omni/diffusion/models/minimax_h3/denoise_loop.py @@ -16,6 +16,7 @@ from typing import Any import torch from vllm_omni.diffusion.attention.backends.abstract import VideoTokenLayout, VideoTokenSpan +from vllm_omni.diffusion.cancellation import check_request_cancellation from vllm_omni.diffusion.forward_context import ( set_forward_context_denoise_step_idx, set_forward_context_denoise_timestep, @@ -366,6 +367,7 @@ def minimax_h3_denoise_loop( num_steps = len(sigmas_video) - 1 for step in range(num_steps): + check_request_cancellation() step_cm = step_profiler(step) if step_profiler is not None else nullcontext() with step_cm: s_v, s_v_next = sigmas_video[step], sigmas_video[step + 1] @@ -450,6 +452,7 @@ def minimax_h3_denoise_loop( audio_rows[~audio_update] = audio_anchor # per-step audio ref reset if on_step is not None: on_step(step, video_rows, audio_rows) + check_request_cancellation(synchronize=True) minimax_h3_publish_denoise_progress(None, None, None) return video_rows, audio_rows diff --git a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py index be9d359d0..98722b932 100644 --- a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py +++ b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py @@ -27,6 +27,7 @@ from vllm_omni.diffusion.cache.cachedit import ( CacheDiTBackend, RequestScopedCacheDiTRuntime, ) +from vllm_omni.diffusion.cancellation import check_request_cancellation from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig from vllm_omni.diffusion.distributed.parallel_state import get_world_group, init_world_group from vllm_omni.diffusion.distributed.utils import get_local_device @@ -604,6 +605,7 @@ class MiniMaxH3Pipeline( """CFG-distilled joint video/audio generation for MiniMax H3.""" supports_step_execution: ClassVar[bool] = True + supports_request_cancellation: ClassVar[bool] = True _dit_modules: ClassVar[list[str]] = ["transformer", "transformers_ref"] _encoder_modules: ClassVar[list[str]] = ["text_encoder"] @@ -2619,15 +2621,18 @@ class MiniMaxH3Pipeline( def forward(self, request: DiffusionRequestBatch) -> DiffusionOutput: if len(request.prompts) != 1: raise OmniClientError("MiniMax H3 supports one request at a time") + check_request_cancellation() context = self._prepare_request_inputs( request.prompts[0], request.sampling_params, ) + check_request_cancellation() denoise_kwargs = self._denoise_kwargs(context) num_outputs = context["num_outputs"] videos = [] audios = [] for output_seed in _minimax_h3_output_seeds(context["seed"], num_outputs): + check_request_cancellation() output_kwargs = {**denoise_kwargs, "seed": output_seed} if context.get("continuation") is None: video_latent, audio_latent = self.diffuse(**output_kwargs) @@ -2640,6 +2645,7 @@ class MiniMaxH3Pipeline( overlap_frames=overlap_frames, text_conditioning=context.get("continuation_text_conditioning"), ) + check_request_cancellation() if context["preencode_mp4"]: videos.append( self.decode_to_mp4( diff --git a/vllm_omni/diffusion/request.py b/vllm_omni/diffusion/request.py index b606c39d7..b4125dcde 100644 --- a/vllm_omni/diffusion/request.py +++ b/vllm_omni/diffusion/request.py @@ -5,7 +5,7 @@ from __future__ import annotations import random -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType @@ -81,6 +81,9 @@ class OmniDiffusionRequest: kv_recv_ms: float = 0.0 # Time spent waiting for initial admission by the diffusion scheduler. scheduler_queue_wait_ms: float | None = None + # Engine-owned shared-memory signal for cooperative full-forward workers. + # Only the engine creates/unlinks it; workers attach readers until return. + cancellation_signal: str | None = field(default=None, repr=False) def __post_init__(self): """Initialize dependent fields after dataclass initialization.""" diff --git a/vllm_omni/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py index 6e4f3813d..50e801ae7 100644 --- a/vllm_omni/diffusion/worker/diffusion_model_runner.py +++ b/vllm_omni/diffusion/worker/diffusion_model_runner.py @@ -31,8 +31,9 @@ from vllm_omni.diffusion.cache.prompt_embed_cache import ( resolve_prompt_embed_cache_config, ) from vllm_omni.diffusion.cache.selector import get_cache_backend +from vllm_omni.diffusion.cancellation import check_request_cancellation, request_cancellation_scope from vllm_omni.diffusion.compile import regionally_compile -from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.data import DiffusionOutput, DiffusionRequestAbortedError, OmniDiffusionConfig from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode from vllm_omni.diffusion.diffusion_kv.metadata import DiffusionKVMetadata from vllm_omni.diffusion.diffusion_kv.model_runner_backend import DiffusionKVModelRunnerBackend @@ -853,15 +854,29 @@ class DiffusionModelRunner(DiffusionStagePayloadMixin): in_diffusion_kv_memory_profile=in_diffusion_kv_memory_profile, ), paged_kv_context, + request_cancellation_scope( + [getattr(req, "cancellation_signal", None) for req in reqs], + enabled=getattr(self.pipeline, "supports_request_cancellation", False) is True, + ), ): with record_function(record_name): - raw_outputs = self.pipeline.forward(batch) - outputs = _normalize_pipeline_outputs( - raw_outputs, - expected_count=len(reqs), - allow_single_output=allow_single_output, - pipeline_name=type(self.pipeline).__name__, - ) + try: + check_request_cancellation() + raw_outputs = self.pipeline.forward(batch) + outputs = _normalize_pipeline_outputs( + raw_outputs, + expected_count=len(reqs), + allow_single_output=allow_single_output, + pipeline_name=type(self.pipeline).__name__, + ) + except DiffusionRequestAbortedError as exc: + # The checkpoint aborts only a fully cancelled wave; + # a mixed batch must keep running for its live peers. + logger.info( + "Stopped cancelled diffusion request(s) %s at a model execution boundary", + [req.request_id for req in reqs], + ) + outputs = [DiffusionOutput(aborted=True, abort_message=str(exc)) for _ in reqs] with record_function("prepare_output_for_transport"): outputs = [ self._prepare_output_for_transport(output, req.sampling_params)