[Bugfix] Stop cancelled MiniMax-H3 requests at denoising boundaries (#7844)
Signed-off-by: Gao Han <hgaoaf@connect.ust.hk>
This commit is contained in:
Gao Han authored and GitHub committed
1 parent
b308e196c9
commit
69de153f6e
12 files changed
+702
-13
No files matched your search
@@ -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()
|
||||
Reference in new issue
Block a user