277 lines
10 KiB
Python
277 lines
10 KiB
Python
# 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()
|