diff --git a/docs/features/sleep_mode.md b/docs/features/sleep_mode.md index b4eec162d..9a740519f 100644 --- a/docs/features/sleep_mode.md +++ b/docs/features/sleep_mode.md @@ -9,9 +9,11 @@ This feature is inherited from [vLLM’s Sleep Mode](https://blog.vllm.ai/2025/1 ## 1. Feature Documentation ### Overview + Omni Sleep Mode provides a mechanism to "sleep" specific model stages. When a stage enters sleep, its physical VRAM is reclaimed by the system, while the process state is preserved for rapid "wake-up" without full re-initialization. ### Sleep Levels + We support two levels of hibernation to balance recovery speed and memory efficiency: | Level | Name | Mechanism | Recovery Speed | Memory Freed | @@ -29,15 +31,67 @@ Omni Sleep Mode is optimized for high-performance computing backends: * **Huawei NPU**: Supported via Ascend memory scavenging. ### Hardware Requirements + * **Memory Considerations**: System RAM must be sufficient to hold offloaded weights during sleep. * **TP Support**: Tensor Parallel groups synchronize sleep/wake transitions across all workers. ---- +### Quiescing a busy diffusion stage: `pause_generation(mode="keep")` +`sleep()` on a diffusion stage is a worker-level RPC: it offloads memory but does not stop the +stage's scheduler, so a stage that still has queued requests must be made quiet first. Aborting +discards the batch that is running. `pause_generation(mode="keep")` keeps it instead: + +* The batch that is already running finishes on the weights it started with and its outputs are + delivered normally. +* Requests that have not started stay queued and run after `resume_generation()`, on whatever + weights are loaded by then. A request never continues denoising across a weight change. +* The call returns only after every replica has acknowledged: the batch that was running has + finished executing, its asynchronous output copies are done, and the device is synchronized; its + outputs are then delivered on the normal path. On a non-preemptible backend that acknowledgement + arrives at the next batch boundary, so expect it to take as long as the remaining batch. +* Control RPCs (`list_loras`, `sleep`, `wake_up`, `resume_generation`) keep working while paused. + +Recommended order for a weight update on a diffusion stage: + +```python +await engine.pause_generation(mode="keep", clear_cache=False) +await engine.sleep(level=1) # optional: free memory for the trainer +await engine.wake_up() # admission stays closed until resume +# install new weights here +await engine.resume_generation() +``` + +Support matrix and limits: + +| Configuration | `mode="keep"` | +| :--- | :--- | +| Request-level execution (`step_execution=False`) with the in-process (`uni`) executor | Supported | +| Request-level execution with the multi-process (`mp`) executor, including asynchronous output | Supported; the acknowledgement waits for the output copies | +| Request-level execution with asynchronous KV prefetch | Supported; the acknowledgement joins any prefetch still in flight | +| Step-level execution (`step_execution=True` / streaming output) | Not supported: the call raises `NotImplementedError` | +| `mode="abort"` / `mode="wait"` / `wait_for_inflight_requests=True` on a diffusion stage | Unchanged: frontend admission is paused only | + +Behavior notes: + +* Callers that already passed `mode="keep"` for a diffusion stage used to get an immediate return + that only closed frontend admission. They now also stop the backend and wait for the ACK. +* A timeout or error from `pause_generation` means the pause did **not** complete; the scheduler + stays closed and the caller must retry the pause or call `resume_generation()` explicitly. Do not + sleep or mutate weights after a failed pause. +* `resume_generation(stage_ids=...)` that leaves another paused stage closed (an AR stage, or a + keep-paused diffusion stage) keeps frontend admission closed as well, so new requests cannot + queue on a stage that will not schedule them. Admission reopens once every stage paused by + `pause_generation` has been resumed. +* `clear_cache` has no diffusion-specific effect. +* Without a pause call nothing changes: no barrier runs and batching, abort and output ordering are + as before. + +--- ## 2. Usage Examples ### Python API Example + You can programmatically control the lifecycle of stages using the `AsyncOmni` engine. ```python @@ -65,6 +119,7 @@ if __name__ == "__main__": ``` ### server command Example + Start the server with sleep mode enabled: The first method @@ -92,9 +147,6 @@ python3 -m vllm_omni.entrypoints.openai.api_server \ ``` - - - ### Test Scenarios & Commands #### Scenario 1: LLM Engine Sleep @@ -113,9 +165,8 @@ curl -X POST http://localhost:8000/v1/omni/sleep \ Tip: Open a new terminal and run rocm-smi or nvidia-smi or to observe the immediate drop in VRAM usage. - - #### Scenario 2: Diffusion Sleep + Objective: Verify VRAM reclamation for Stage 1 (Diffusion). Trigger sleep (Level 1 or Level 2) via client: @@ -128,9 +179,8 @@ curl -X POST http://localhost:8000/v1/omni/sleep \ ``` - - #### Scenario 3: Multi-Stage Coordinated Stress Test + Objective: Test concurrent sleep and rapid wake-up across multiple stages. Concurrent Sleep (Stage 0 & 1): @@ -143,7 +193,6 @@ curl -X POST http://localhost:8000/v1/omni/sleep \ ``` - Rapid Wake-up: ``` @@ -154,8 +203,8 @@ curl -X POST http://localhost:8000/v1/omni/wakeup \ ``` - #### Scenario 4: Full Lifecycle Memory Audit & Functional Integrity + Objective: Audit the complete flow from Sleep to Wake-up followed by an Inference validation. Check Initial State: Observe baseline VRAM usage. @@ -196,20 +245,16 @@ curl -X POST http://localhost:8000/v1/images/generations \ ``` - - - ## 3. API Reference - ### Methods | Method | Arguments | Return Type | Description | | :--- | :--- | :--- | :--- | | **sleep** | `stage_ids: List[int], level: int` | `List[OmniACK]` | Triggers hibernation for specified stages. | | **wake_up** | `stage_ids: List[int]` | `List[OmniACK]` | Reloads weights and re-maps memory. | - - +| **pause_generation** | `mode: str, stage_ids: List[int]` | `None` | Stops admission; with `mode="keep"` also stops diffusion schedulers and returns after their ACK. | +| **resume_generation** | `stage_ids: List[int]` | `None` | Reopens paused schedulers, then admission. | ### OmniACK Dataclass Fields @@ -233,7 +278,8 @@ The metadata field is a dynamic dictionary containing hardware-specific telemetr } ``` -#### Core Utility: +#### Core Utility + **VRAM Reclamation Audit (total_freed_gib)**: Converts raw freed_bytes into human-readable GiB. It serves as the primary metric to verify that Level 2 sleep has successfully purged model weights from VRAM. **Residual & Fragmentation Monitoring (rank_residual_gib)**: Reports the remaining VRAM footprint after memory de-mapping. A low residual value (e.g., 2.07 GiB) confirms a successful "clean" state, ensuring the device is ready for high-memory co-located tasks like training or diffusion pipelines. diff --git a/tests/diffusion/distributed/test_pause_keep_two_ranks.py b/tests/diffusion/distributed/test_pause_keep_two_ranks.py new file mode 100644 index 000000000..6ebd4743f --- /dev/null +++ b/tests/diffusion/distributed/test_pause_keep_two_ranks.py @@ -0,0 +1,81 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +""" +Two-rank ``pause_generation(mode="keep")``: every rank runs the pause barrier +before the ACK, and a barrier failure on a non-zero rank fails the pause +without aborting or dropping the running batch. +""" + +from __future__ import annotations + +import asyncio +from collections import defaultdict +from contextlib import ExitStack + +import pytest + +from tests.e2e.features.custom_pipeline.test_diffusion_pause_keep import ( + CUSTOM_PIPELINE_CLASS, + MODEL, + WORKER_EXTENSION_CLASS, + _generate, + _pause_round, + _wait_for, +) +from tests.helpers.mark import hardware_test +from vllm_omni.entrypoints.async_omni import AsyncOmni + + +def _all_true(results) -> bool: + flat = [] + stack = list(results) + while stack: + item = stack.pop() + if isinstance(item, list): + stack.extend(item) + else: + flat.append(item) + return bool(flat) and all(item is True for item in flat) + + +@pytest.mark.core_model +@pytest.mark.diffusion +@hardware_test(res={"cuda": "L4"}, num_cards=2) +@pytest.mark.asyncio +async def test_keep_pause_two_ranks_ack_and_rank_failure(): + with ExitStack() as after: + engine = AsyncOmni( + model=MODEL, + custom_pipeline_args={"pipeline_class": CUSTOM_PIPELINE_CLASS}, + worker_extension_cls=WORKER_EXTENSION_CLASS, + enforce_eager=True, + enable_sleep_mode=True, + num_gpus=2, + parallel_config={"ulysses_degree": 2, "ulysses_mode": "advanced_uaa"}, + ) + after.callback(engine.shutdown) + events: dict[str, dict[str, float]] = defaultdict(dict) + assert _all_true(await engine.collective_rpc(method="start_counting_synchronize_device")) + + pause_latencies = [await _pause_round(engine, events, round_index) for round_index in range(3)] + # The result is the all-rank AND, so True means both ranks ran the barrier. + assert _all_true(await engine.collective_rpc(method="synchronize_device_seen")) + + assert _all_true(await engine.collective_rpc(method="set_synchronize_device_failure", args=(1, True))) + a = asyncio.create_task(_generate(engine, "a-fail", events)) + await _wait_for(lambda: "started" in events["a-fail"]) + with pytest.raises(RuntimeError, match="pause_scheduler failed"): + await engine.pause_generation(mode="keep", clear_cache=False) + assert await engine.is_paused() + a_output = await asyncio.wait_for(a, 60.0) + assert a_output.images, "a failed barrier must not abort the running batch" + + assert _all_true(await engine.collective_rpc(method="set_synchronize_device_failure", args=(1, False))) + await engine.resume_generation() + assert not await engine.is_paused() + recovered = await asyncio.wait_for(_generate(engine, "after-recovery", events), 60.0) + assert recovered.images + + assert not engine.request_states + print(f"[sp2] rounds=3 pause_latency_s min/max={min(pause_latencies):.3f}/{max(pause_latencies):.3f}") diff --git a/tests/diffusion/test_async_output_worker.py b/tests/diffusion/test_async_output_worker.py index 8a81f30d4..9465f94b3 100644 --- a/tests/diffusion/test_async_output_worker.py +++ b/tests/diffusion/test_async_output_worker.py @@ -497,3 +497,99 @@ class TestWorkerProcShutdown: result_mq.shutdown.assert_called_once_with() assert proc.mq is None assert proc.result_mq is None + + +class TestSynchronizeDeviceBarrier: + """The pause barrier must not report a quiet device while outputs still pack.""" + + @staticmethod + def _rpc(timeout): + return { + "method": "synchronize_device", + "args": (), + "kwargs": {"timeout": timeout}, + "output_rank": 0, + "exec_all_ranks": False, + "collect_rank_status": False, + } + + def test_rejects_undrained_output(self, mocker): + proc = _make_worker_proc(step_execution=False) + proc.worker = MagicMock() + mocker.patch.object(proc, "drain_async_outputs", return_value=False) + + with pytest.raises(TimeoutError, match="did not drain"): + proc._execute_rpc(self._rpc(0.01)) + + proc.worker.execute_method.assert_not_called() + + @pytest.mark.parametrize("timeout", [0.5, None]) + def test_drains_with_the_call_timeout_then_synchronizes(self, mocker, timeout): + proc = _make_worker_proc(step_execution=False) + proc.worker = MagicMock() + call_order: list = [] + mocker.patch.object( + proc, + "drain_async_outputs", + side_effect=lambda **kw: call_order.append(("drain", kw)) or True, + ) + proc.worker.execute_method.side_effect = lambda *a, **kw: call_order.append(("sync", a, kw)) + + _, should_reply = proc._execute_rpc(self._rpc(timeout)) + + assert call_order == [ + ("drain", {"timeout": timeout}), + ("sync", ("synchronize_device",), {"timeout": timeout}), + ] + assert should_reply is True + + def test_completes_only_after_output_ready_is_published(self, mocker): + proc = _make_worker_proc(step_execution=False) + proc.worker = MagicMock() + + mock_platform = mocker.MagicMock() + mock_platform.record_device_event.return_value = MagicMock() + mocker.patch( + "vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform", + mock_platform, + ) + mock_torch = mocker.MagicMock() + mock_torch.accelerator.current_accelerator.return_value.type = "cpu" + mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.torch", mock_torch) + release = threading.Event() + mocker.patch( + "vllm_omni.diffusion.worker.diffusion_worker.pack_diffusion_output_shm", + side_effect=lambda *a, **kw: release.wait(timeout=5.0), + ) + + proc._return_result(DiffusionOutput(output="data"), rpc_id="1") + loop_thread = threading.Thread(target=proc._async_output_loop, daemon=True) + loop_thread.start() + + barrier_done = threading.Event() + + def run_barrier(): + proc._execute_rpc(self._rpc(5.0)) + barrier_done.set() + + barrier_thread = threading.Thread(target=run_barrier, daemon=True) + barrier_thread.start() + try: + assert not barrier_done.wait(timeout=0.2) + proc.worker.execute_method.assert_not_called() + + release.set() + assert barrier_done.wait(timeout=5.0) + + kinds = [ + call.args[0].kind + for call in proc.result_mq.enqueue.call_args_list + if isinstance(call.args[0], AsyncDiffusionOutput) + ] + assert kinds == [AsyncOutputKind.COMPUTE_DONE, AsyncOutputKind.OUTPUT_READY] + proc.worker.execute_method.assert_called_once_with("synchronize_device", timeout=5.0) + finally: + release.set() + proc._async_output_queue.put(None) + loop_thread.join(timeout=2.0) + barrier_thread.join(timeout=2.0) diff --git a/tests/diffusion/test_diffusion_engine_rpc_routing.py b/tests/diffusion/test_diffusion_engine_rpc_routing.py index bb64e5f01..d26f20c1d 100644 --- a/tests/diffusion/test_diffusion_engine_rpc_routing.py +++ b/tests/diffusion/test_diffusion_engine_rpc_routing.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 """Regression tests for the busy-loop-based RPC routing in DiffusionEngine. These tests guard the invariants established when ``collective_rpc`` was @@ -60,6 +60,12 @@ class _ConcurrencyTrackingExecutor: self.is_failed = False self._closed = False self.od_config = SimpleNamespace() + # Pause tests: a call whose gate is registered (by request id, or by + # method name for ``synchronize_device``) blocks inside the executor + # until the gate is set; ``version`` tags each call so a test can tell + # which "weights" a request ran on. + self.gates: dict[str, threading.Event] = {} + self.version = "v1" def collective_rpc( self, @@ -81,16 +87,25 @@ class _ConcurrencyTrackingExecutor: "kwargs": kwargs, "unique_reply_rank": unique_reply_rank, "thread": threading.get_ident(), + "version": self.version, } ) try: if self.rpc_delay: time.sleep(self.rpc_delay) + if method == "synchronize_device": + gate = self.gates.get(method) + if gate is not None: + gate.wait(timeout=5.0) + return [None] # Distinguish per-request execution from raw RPC by method name. if method in {"execute_model", "execute_stepwise", "generate"}: # args[0] is the request-like object for execute_model. req = args[0] if args else None tag = req.request_id if req is not None and hasattr(req, "request_id") else "unknown" + gate = self.gates.get(tag) + if gate is not None: + gate.wait(timeout=5.0) return DiffusionOutput(error=f"result_for_{tag}", finished=True) tag = args[0] if args else method return DiffusionOutput(error=f"rpc_result_for_{tag}") @@ -133,21 +148,27 @@ def _make_request(tag: str): def _make_engine_with_loop( - loop: asyncio.AbstractEventLoop, + loop: asyncio.AbstractEventLoop | None, rpc_delay: float = 0.0, + *, + scheduler: RequestScheduler | None = None, + start_loop: bool = True, ): """Construct a ``DiffusionEngine`` skeleton with a real busy loop. The engine is wired with a fake executor that asserts no concurrent - calls and a real ``RequestScheduler``. + calls and a real ``RequestScheduler``. ``start_loop=False`` leaves the + engine in its bootstrap state (no busy-loop thread). """ engine = DiffusionEngine.__new__(DiffusionEngine) engine._closed = False engine.od_config = SimpleNamespace(streaming_output=False) engine.executor = _ConcurrencyTrackingExecutor(rpc_delay=rpc_delay) - sched = RequestScheduler() - sched.initialize(SimpleNamespace(max_num_seqs=1, request_batch_max_wait_ms=0.0)) + sched = scheduler + if sched is None: + sched = RequestScheduler() + sched.initialize(SimpleNamespace(max_num_seqs=1, request_batch_max_wait_ms=0.0)) engine.scheduler = sched engine.step_execution = False engine.supports_request_batch = False @@ -160,9 +181,14 @@ def _make_engine_with_loop( engine._closed = False engine.abort_queue = queue.Queue() engine._rpc_queue = queue.Queue() + engine._scheduling_paused = False engine.main_loop = loop engine.stop_event = threading.Event() + if not start_loop: + engine._loop_started = False + engine.worker_thread = None + return engine engine._loop_started = True engine.worker_thread = threading.Thread(target=engine._busy_loop, daemon=True) engine.worker_thread.start() @@ -170,6 +196,10 @@ def _make_engine_with_loop( def _stop_engine(engine: DiffusionEngine) -> None: + if engine.worker_thread is None: + return + for gate in engine.executor.gates.values(): + gate.set() with engine._cv: engine.stop_event.set() engine._cv.notify_all() @@ -565,3 +595,319 @@ async def test_rpc_and_request_results_do_not_swap(): assert req_res.error == "result_for_req1" # Belt and braces: never overlapped on the executor. assert engine.executor.max_concurrent == 1 + + +# ───────────────────────── scheduler pause (keep) ───────────────────────── + + +class _GatedAdmissionScheduler(RequestScheduler): + """Batching wait that only ends when the test opens admission.""" + + def __init__(self) -> None: + super().__init__() + self.admission_open = threading.Event() + + def get_admission_wait_decision(self, *, now: float, dp_concurrent: bool = False): + decision = super().get_admission_wait_decision(now=now, dp_concurrent=dp_concurrent) + return decision.__class__(should_wait=True, deadline=None, max_batch=self.max_num_running_reqs) + + def should_end_admission_wait(self, decision, *, now: float, stable_since: float) -> bool: + return self.admission_open.is_set() + + +def _pause(engine: DiffusionEngine, timeout: float | None = None): + return engine.async_collective_rpc("pause_scheduler", timeout=timeout, kwargs={"mode": "keep"}) + + +def _resume(engine: DiffusionEngine): + return engine.async_collective_rpc("resume_scheduler") + + +def _methods(engine: DiffusionEngine) -> list[str]: + return [call["method"] for call in engine.executor.calls] + + +def _executed(engine: DiffusionEngine) -> list[tuple[str, str]]: + return [ + (call["args"][0].request_id, call["version"]) + for call in engine.executor.calls + if call["method"] == "execute_model" + ] + + +def _submit(engine: DiffusionEngine, tag: str) -> asyncio.Task: + return asyncio.create_task(_consume_final_output(engine.async_add_req_and_stream_response(_make_request(tag)))) + + +async def _wait_until(predicate, timeout: float = 3.0) -> None: + deadline = time.monotonic() + timeout + while not predicate(): + if time.monotonic() > deadline: + raise AssertionError("condition not reached in time") + await asyncio.sleep(0.005) + + +@pytest.mark.asyncio +async def test_keep_pause_retains_waiting_request_and_serves_control_rpcs(): + loop = asyncio.get_running_loop() + engine = _make_engine_with_loop(loop) + try: + assert await _pause(engine) is None + queued = _submit(engine, "queued") + result = await engine.async_collective_rpc("list_loras", unique_reply_rank=0) + assert result.error == "rpc_result_for_list_loras" + await asyncio.sleep(0.05) + assert engine.scheduler.num_waiting_requests() == 1 + assert _executed(engine) == [] + assert not queued.done() + + await _resume(engine) + assert (await asyncio.wait_for(queued, 3.0)).error == "result_for_queued" + assert _methods(engine) == ["synchronize_device", "list_loras", "execute_model"] + finally: + _stop_engine(engine) + + +@pytest.mark.asyncio +async def test_keep_pause_acks_after_running_batch_and_freezes_queue(): + loop = asyncio.get_running_loop() + engine = _make_engine_with_loop(loop) + gate = engine.executor.gates["A"] = threading.Event() + try: + a = _submit(engine, "A") + await _wait_until(lambda: _executed(engine) == [("A", "v1")]) + b = _submit(engine, "B") + await _wait_until(lambda: engine.scheduler.num_waiting_requests() == 1) + + pause = asyncio.create_task(_pause(engine)) + await asyncio.sleep(0.1) + assert not pause.done() + + gate.set() + await asyncio.wait_for(pause, 3.0) + assert _methods(engine) == ["execute_model", "synchronize_device"] + assert (await asyncio.wait_for(a, 3.0)).error == "result_for_A" + assert engine.scheduler.num_waiting_requests() == 1 + assert not b.done() + + engine.executor.version = "v2" + await _resume(engine) + assert (await asyncio.wait_for(b, 3.0)).error == "result_for_B" + assert _executed(engine) == [("A", "v1"), ("B", "v2")] + finally: + _stop_engine(engine) + + +@pytest.mark.asyncio +async def test_keep_pause_ack_does_not_release_later_rpcs_before_the_batch_is_emitted(): + """A pause ACK must not open the queue to RPCs that touch device memory + (sleep) before the batch that just ran has been delivered. + """ + loop = asyncio.get_running_loop() + engine = _make_engine_with_loop(loop) + barrier_gate = engine.executor.gates["synchronize_device"] = threading.Event() + methods_at_emit: list[list[str]] = [] + emit_outputs = engine._emit_outputs + + def _recording_emit(*args, **kwargs): + methods_at_emit.append(_methods(engine)) + return emit_outputs(*args, **kwargs) + + engine._emit_outputs = _recording_emit + try: + gate = engine.executor.gates["A"] = threading.Event() + a = _submit(engine, "A") + await _wait_until(lambda: _executed(engine) == [("A", "v1")]) + pause = asyncio.create_task(_pause(engine)) + await _wait_until(lambda: engine._scheduling_paused) + gate.set() + # The barrier holds the busy loop inside the post-execute RPC drain, + # so the sleep is queued before the pause ACK is published. + await _wait_until(lambda: "synchronize_device" in _methods(engine)) + sleep_rpc = asyncio.create_task(engine.async_collective_rpc("handle_sleep_task", args=("t",))) + await _wait_until(lambda: not engine._rpc_queue.empty()) + barrier_gate.set() + + await asyncio.wait_for(pause, 3.0) + assert (await asyncio.wait_for(a, 3.0)).error == "result_for_A" + assert (await asyncio.wait_for(sleep_rpc, 3.0)).error == "rpc_result_for_t" + finally: + _stop_engine(engine) + + assert _methods(engine) == ["execute_model", "synchronize_device", "handle_sleep_task"] + assert methods_at_emit and "handle_sleep_task" not in methods_at_emit[0] + + +@pytest.mark.asyncio +async def test_timed_out_pause_still_ends_the_drain_before_the_batch_is_emitted(): + """A pause whose caller already gave up is still a drain barrier: its + cancelled task must not let a later RPC run before the batch is emitted. + """ + loop = asyncio.get_running_loop() + engine = _make_engine_with_loop(loop) + methods_at_emit: list[list[str]] = [] + emit_outputs = engine._emit_outputs + + def _recording_emit(*args, **kwargs): + methods_at_emit.append(_methods(engine)) + return emit_outputs(*args, **kwargs) + + engine._emit_outputs = _recording_emit + try: + gate = engine.executor.gates["A"] = threading.Event() + a = _submit(engine, "A") + await _wait_until(lambda: _executed(engine) == [("A", "v1")]) + + # The caller times out while A is still executing: the pause future is + # cancelled, but its task stays queued ahead of the sleep. + with pytest.raises(TimeoutError): + await _pause(engine, timeout=0.1) + sleep_rpc = asyncio.create_task(engine.async_collective_rpc("handle_sleep_task", args=("t",))) + await _wait_until(lambda: engine._rpc_queue.qsize() == 2) + gate.set() + + assert (await asyncio.wait_for(a, 3.0)).error == "result_for_A" + assert (await asyncio.wait_for(sleep_rpc, 3.0)).error == "rpc_result_for_t" + finally: + _stop_engine(engine) + + assert _methods(engine) == ["execute_model", "handle_sleep_task"] + assert methods_at_emit and "handle_sleep_task" not in methods_at_emit[0] + + +@pytest.mark.asyncio +async def test_pause_rejects_unsupported_modes_without_closing_the_gate(): + loop = asyncio.get_running_loop() + engine = _make_engine_with_loop(loop) + try: + for kwargs in ({"mode": "abort"}, {"mode": "wait"}, {}): + with pytest.raises(ValueError): + await engine.async_collective_rpc("pause_scheduler", kwargs=kwargs) + engine.execution_mode = DiffusionExecutionMode.STEP_BATCH + with pytest.raises(NotImplementedError): + await _pause(engine) + engine.execution_mode = DiffusionExecutionMode.REQUEST_BATCH + + assert (await asyncio.wait_for(_submit(engine, "still-runs"), 3.0)).error == "result_for_still-runs" + assert _methods(engine) == ["execute_model"] + finally: + _stop_engine(engine) + + +@pytest.mark.asyncio +async def test_keep_pause_during_batching_wait_holds_the_batch(): + loop = asyncio.get_running_loop() + sched = _GatedAdmissionScheduler() + sched.initialize(SimpleNamespace(max_num_seqs=4, request_batch_max_wait_ms=1000.0)) + engine = _make_engine_with_loop(loop, scheduler=sched) + try: + waiting = _submit(engine, "W") + await _wait_until(lambda: engine.scheduler.num_waiting_requests() == 1) + await asyncio.sleep(0.05) + + await asyncio.wait_for(_pause(engine), 3.0) + assert _methods(engine) == ["synchronize_device"] + assert engine.scheduler.num_waiting_requests() == 1 + + await _resume(engine) + await asyncio.sleep(0.05) + assert _executed(engine) == [] + sched.admission_open.set() + assert (await asyncio.wait_for(waiting, 3.0)).error == "result_for_W" + finally: + sched.admission_open.set() + _stop_engine(engine) + + +@pytest.mark.asyncio +async def test_pause_timeout_leaves_gate_closed_until_explicit_resume(): + loop = asyncio.get_running_loop() + engine = _make_engine_with_loop(loop) + gate = engine.executor.gates["A"] = threading.Event() + try: + a = _submit(engine, "A") + await _wait_until(lambda: _executed(engine) == [("A", "v1")]) + b = _submit(engine, "B") + await _wait_until(lambda: engine.scheduler.num_waiting_requests() == 1) + + with pytest.raises(TimeoutError): + await _pause(engine, timeout=0.1) + gate.set() + assert (await asyncio.wait_for(a, 3.0)).error == "result_for_A" + await asyncio.sleep(0.1) + assert _executed(engine) == [("A", "v1")] + assert not b.done() + + await _resume(engine) + assert (await asyncio.wait_for(b, 3.0)).error == "result_for_B" + assert "synchronize_device" not in _methods(engine) + finally: + _stop_engine(engine) + + +@pytest.mark.asyncio +async def test_repeated_pause_and_resume_are_idempotent(): + loop = asyncio.get_running_loop() + engine = _make_engine_with_loop(loop) + try: + await _resume(engine) + await _pause(engine) + await _pause(engine) + queued = _submit(engine, "Q") + await asyncio.sleep(0.05) + assert _executed(engine) == [] + await _resume(engine) + await _resume(engine) + assert (await asyncio.wait_for(queued, 3.0)).error == "result_for_Q" + assert _methods(engine) == ["synchronize_device", "synchronize_device", "execute_model"] + finally: + _stop_engine(engine) + + +@pytest.mark.asyncio +async def test_abort_of_waiting_request_while_paused_delivers_terminal(): + loop = asyncio.get_running_loop() + engine = _make_engine_with_loop(loop) + try: + await _pause(engine) + b = _submit(engine, "B") + await _wait_until(lambda: engine.scheduler.num_waiting_requests() == 1) + engine.abort("B") + + result = await asyncio.wait_for(b, 3.0) + assert result.aborted is True + assert engine.scheduler.get_request_state("B") is None + assert _executed(engine) == [] + + await _resume(engine) + assert (await asyncio.wait_for(_submit(engine, "C"), 3.0)).error == "result_for_C" + assert _executed(engine) == [("C", "v1")] + assert engine.worker_thread.is_alive() + finally: + _stop_engine(engine) + + +def test_pause_and_resume_before_loop_starts_use_engine_local_path(): + engine = _make_engine_with_loop(None, start_loop=False) + + assert engine.collective_rpc("pause_scheduler", kwargs={"mode": "keep"}) is None + assert _methods(engine) == ["synchronize_device"] + + assert engine.collective_rpc("resume_scheduler") is None + assert engine.add_req_and_wait_for_response(_make_request("X")).error == "result_for_X" + assert _executed(engine) == [("X", "v1")] + + +@pytest.mark.asyncio +async def test_requests_without_pause_never_trigger_the_barrier(): + loop = asyncio.get_running_loop() + engine = _make_engine_with_loop(loop) + try: + first = await asyncio.wait_for(_submit(engine, "r1"), 3.0) + await engine.async_collective_rpc("ping", args=("p",), unique_reply_rank=0) + second = await asyncio.wait_for(_submit(engine, "r2"), 3.0) + finally: + _stop_engine(engine) + + assert (first.error, second.error) == ("result_for_r1", "result_for_r2") + assert _methods(engine) == ["execute_model", "ping", "execute_model"] diff --git a/tests/diffusion/test_diffusion_scheduler.py b/tests/diffusion/test_diffusion_scheduler.py index 38da458b0..c0009c30f 100644 --- a/tests/diffusion/test_diffusion_scheduler.py +++ b/tests/diffusion/test_diffusion_scheduler.py @@ -2092,3 +2092,19 @@ class TestStepScheduler: with pytest.raises(ValueError): self.scheduler.add_request(request) + + +class TestPendingFinishedRequestIds: + def test_reports_finished_ids_that_still_hold_state_without_clearing_them(self): + sched = RequestScheduler() + sched.initialize(SimpleNamespace(max_num_seqs=1, request_batch_max_wait_ms=0.0)) + sched.add_request(_make_request("a")) + sched.add_request(_make_request("b")) + + sched.finish_requests("a", DiffusionRequestStatus.FINISHED_ABORTED) + assert sched.pending_finished_request_ids() == {"a"} + + sched.pop_request_state("a") + assert sched.pending_finished_request_ids() == set() + # The next wave still ships the id to the worker for its own cleanup. + assert sched.schedule().finished_req_ids == {"a"} diff --git a/tests/diffusion/test_diffusion_worker_synchronize_device.py b/tests/diffusion/test_diffusion_worker_synchronize_device.py new file mode 100644 index 000000000..6e4d3f4bf --- /dev/null +++ b/tests/diffusion/test_diffusion_worker_synchronize_device.py @@ -0,0 +1,53 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Unit tests for the device step of the diffusion pause barrier.""" + +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest + +from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _make_worker(runner: object) -> DiffusionWorker: + worker = object.__new__(DiffusionWorker) + worker.model_runner = runner + worker.rank = 0 + return worker + + +def test_synchronize_device_waits_for_prefetch_before_device_sync(mocker): + platform = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform") + order: list = [] + manager = MagicMock() + manager.wait_prefetch.side_effect = lambda timeout: order.append(("prefetch", timeout)) or True + platform.synchronize.side_effect = lambda: order.append("sync") + worker = _make_worker(SimpleNamespace(kv_transfer_manager=manager)) + + worker.synchronize_device(timeout=2.5) + + assert order == [("prefetch", 2.5), "sync"] + + +def test_synchronize_device_fails_when_prefetch_is_still_outstanding(mocker): + platform = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform") + manager = MagicMock() + manager.wait_prefetch.return_value = False + worker = _make_worker(SimpleNamespace(kv_transfer_manager=manager)) + + with pytest.raises(TimeoutError, match="prefetch"): + worker.synchronize_device(timeout=0.01) + + platform.synchronize.assert_not_called() + + +def test_synchronize_device_without_transfer_manager_only_syncs(mocker): + platform = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform") + worker = _make_worker(SimpleNamespace(kv_transfer_manager=None)) + + worker.synchronize_device(timeout=None) + + platform.synchronize.assert_called_once_with() diff --git a/tests/distributed/omni_connectors/test_kv_async_prefetch.py b/tests/distributed/omni_connectors/test_kv_async_prefetch.py index 6ef5a3f8c..b596e97bb 100644 --- a/tests/distributed/omni_connectors/test_kv_async_prefetch.py +++ b/tests/distributed/omni_connectors/test_kv_async_prefetch.py @@ -5,6 +5,7 @@ mock connector, dedup, per-call sender_info isolation, role classification, and miss/abort paths. GPU H2D and D2D paths need CUDA (integration tests). """ +import threading from types import SimpleNamespace import pytest @@ -354,3 +355,41 @@ def test_consume_then_apply_attaches_payload(): receiver._record_stream_for_prefetched(data) receiver.apply_kv_cache_to_request(req, data) assert req.past_key_values is not None + + +# --------------------------------------------------------------------------- # +# wait_prefetch: the pause barrier's view of outstanding background loads +# --------------------------------------------------------------------------- # + + +def test_wait_prefetch_is_true_when_nothing_outstanding(): + _, receiver, _ = _make_sender_receiver() + assert receiver.wait_prefetch(timeout=0.01) is True + + +def test_wait_prefetch_is_true_when_prefetch_disabled(): + _, receiver, _ = _make_sender_receiver(async_prefetch=False) + assert receiver.wait_prefetch(timeout=0.01) is True + + +def test_wait_prefetch_reports_a_running_prefetch_and_keeps_its_payload(): + sender, receiver, connector = _make_sender_receiver() + _seed_payload(sender, "rid-w") + gate = threading.Event() + real_get = connector.get + + def slow_get(*args, **kwargs): + gate.wait(timeout=5.0) + return real_get(*args, **kwargs) + + connector.get = slow_get + receiver.start_prefetch({"request_id": "rid-w", "kv_sender_info": _SENDER_INFO}) + try: + assert receiver.wait_prefetch(timeout=0.05) is False + gate.set() + assert receiver.wait_prefetch(timeout=5.0) is True + data, _ = receiver.consume_prefetched_kv(_req("rid-w")) + assert data is not None + finally: + gate.set() + receiver.shutdown_prefetch() diff --git a/tests/e2e/features/custom_pipeline/test_diffusion_pause_keep.py b/tests/e2e/features/custom_pipeline/test_diffusion_pause_keep.py new file mode 100644 index 000000000..cd8231473 --- /dev/null +++ b/tests/e2e/features/custom_pipeline/test_diffusion_pause_keep.py @@ -0,0 +1,183 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +""" +Lifecycle test for ``AsyncOmni.pause_generation(mode="keep")`` on a diffusion stage. + +Per round, request A is started, request B is queued behind it, and the pause +is issued while A is still running. The pause must return only after A has +been delivered, B must stay queued through ``sleep`` / ``wake_up``, and B may +start only after ``resume_generation``. + +A second test follows the order the docs recommend for a weight update: sleep +is issued as soon as the pause returns, without awaiting the request that was +running, which must still be delivered intact. + +Both single-GPU executor backends are covered explicitly: ``uni`` (in-process +worker) and ``mp`` (worker process with the asynchronous D2H output thread). + +Usage: + pytest tests/e2e/features/custom_pipeline/test_diffusion_pause_keep.py -v -s +""" + +from __future__ import annotations + +import asyncio +import os +import time +from collections import defaultdict +from contextlib import ExitStack + +import pytest + +from tests.helpers.mark import hardware_test +from vllm_omni.diffusion.data import is_diffusion_request_started_output +from vllm_omni.entrypoints.async_omni import AsyncOmni +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +MODEL = "tiny-random/Qwen-Image" +CUSTOM_PIPELINE_CLASS = "tests.e2e.features.helpers.custom_pipeline.QwenImagePipelineWithLogProbForTest" +WORKER_EXTENSION_CLASS = "tests.e2e.features.helpers.custom_pipeline.vLLMOmniColocateWorkerExtensionForTest" + +# CI runs a couple of rounds; set PAUSE_KEEP_ROUNDS=10 for a soak run. +ROUNDS = int(os.environ.get("PAUSE_KEEP_ROUNDS", "2")) +# Enough denoising steps that the pause lands while A is still executing. +NUM_INFERENCE_STEPS = 40 + + +def _sampling_params() -> OmniDiffusionSamplingParams: + return OmniDiffusionSamplingParams( + num_inference_steps=NUM_INFERENCE_STEPS, + guidance_scale=0.0, + height=256, + width=256, + seed=42, + emit_request_lifecycle=True, + ) + + +async def _generate(engine: AsyncOmni, request_id: str, events: dict[str, dict[str, float]]): + final = None + async for output in engine.generate( + prompt={"prompt_ids": list(range(50))}, + request_id=request_id, + sampling_params_list=[_sampling_params()], + output_modalities=["image"], + ): + if is_diffusion_request_started_output(output): + events[request_id]["started"] = time.monotonic() + continue + final = output + events[request_id]["finished"] = time.monotonic() + return final + + +async def _wait_for(predicate, timeout: float = 30.0) -> None: + deadline = time.monotonic() + timeout + while not predicate(): + assert time.monotonic() < deadline, "condition not reached in time" + await asyncio.sleep(0.005) + + +async def _pause_round(engine: AsyncOmni, events: dict[str, dict[str, float]], round_index: int) -> float: + """A runs, B queues, keep-pause lands while A runs; B starts only after resume. + + Returns the pause latency. + """ + a_id, b_id = f"a{round_index}", f"b{round_index}" + a = asyncio.create_task(_generate(engine, a_id, events)) + await _wait_for(lambda: "started" in events[a_id]) + b = asyncio.create_task(_generate(engine, b_id, events)) + await asyncio.sleep(0.05) + + assert not a.done(), "A finished before the pause was issued; raise NUM_INFERENCE_STEPS" + pause_start = time.monotonic() + await engine.pause_generation(mode="keep", clear_cache=False) + pause_latency = time.monotonic() - pause_start + assert await engine.is_paused() + + a_output = await asyncio.wait_for(a, 60.0) + assert a_output.images, "A must be delivered, not aborted, by a keep pause" + assert "started" not in events[b_id], "B must not start while paused" + + # Control RPCs and the sleep cycle stay available on the paused stage. + assert await engine.list_loras() == [] + await engine.sleep(level=1) + assert await engine.is_sleeping() + await engine.wake_up() + assert not await engine.is_sleeping() + assert await engine.is_paused() + assert not b.done() and "started" not in events[b_id] + + resume_at = time.monotonic() + await engine.resume_generation() + b_output = await asyncio.wait_for(b, 60.0) + assert b_output.images, "B must run after resume" + assert events[b_id]["started"] >= resume_at + return pause_latency + + +@pytest.mark.core_model +@pytest.mark.diffusion +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.asyncio +@pytest.mark.parametrize("backend", ["uni", "mp"]) +async def test_keep_pause_sleep_wake_resume_lifecycle(backend: str): + with ExitStack() as after: + engine = AsyncOmni( + model=MODEL, + custom_pipeline_args={"pipeline_class": CUSTOM_PIPELINE_CLASS}, + worker_extension_cls=WORKER_EXTENSION_CLASS, + enforce_eager=True, + enable_sleep_mode=True, + distributed_executor_backend=backend, + ) + after.callback(engine.shutdown) + events: dict[str, dict[str, float]] = defaultdict(dict) + + pause_latencies = [await _pause_round(engine, events, round_index) for round_index in range(ROUNDS)] + + assert not engine.request_states, "every request must be finalized" + print( + f"[{backend}] rounds={ROUNDS} steps={NUM_INFERENCE_STEPS} " + f"pause_latency_s min/median/max=" + f"{min(pause_latencies):.3f}/{sorted(pause_latencies)[ROUNDS // 2]:.3f}/{max(pause_latencies):.3f}" + ) + + +@pytest.mark.core_model +@pytest.mark.diffusion +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.asyncio +@pytest.mark.parametrize("backend", ["uni", "mp"]) +async def test_sleep_right_after_keep_pause_keeps_the_running_batch(backend: str): + """The documented weight-update order: sleep is issued as soon as the pause + returns, with the request that was running still un-awaited. + """ + with ExitStack() as after: + engine = AsyncOmni( + model=MODEL, + custom_pipeline_args={"pipeline_class": CUSTOM_PIPELINE_CLASS}, + worker_extension_cls=WORKER_EXTENSION_CLASS, + enforce_eager=True, + enable_sleep_mode=True, + distributed_executor_backend=backend, + ) + after.callback(engine.shutdown) + events: dict[str, dict[str, float]] = defaultdict(dict) + + a = asyncio.create_task(_generate(engine, "a", events)) + await _wait_for(lambda: "started" in events["a"]) + await asyncio.sleep(0.05) + assert not a.done(), "A finished before the pause was issued; raise NUM_INFERENCE_STEPS" + + await engine.pause_generation(mode="keep", clear_cache=False) + await engine.sleep(level=1) + + a_output = await asyncio.wait_for(a, 60.0) + assert a_output.images, "A must survive a sleep issued right after the pause ACK" + + await engine.wake_up() + await engine.resume_generation() + b_output = await asyncio.wait_for(_generate(engine, "b", events), 120.0) + assert b_output.images, "the stage must generate again after wake and resume" diff --git a/tests/e2e/features/helpers/custom_pipeline.py b/tests/e2e/features/helpers/custom_pipeline.py index 63cb21056..3f1b48563 100644 --- a/tests/e2e/features/helpers/custom_pipeline.py +++ b/tests/e2e/features/helpers/custom_pipeline.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 """Test doubles for custom-pipeline E2E and RLHF integration. @@ -290,6 +290,36 @@ class vLLMOmniColocateWorkerExtensionForTest(CustomPipelineWorkerExtension): """Return a stable identifier for assertions in unit tests.""" return "vllm-omni-colocate-worker-extension-for-test" + def start_counting_synchronize_device(self) -> bool: + """Count pause-barrier runs on this rank (test-only).""" + self._synchronize_device_calls = 0 + inner = type(self).synchronize_device + + def _counted(timeout=None): + self._synchronize_device_calls += 1 + return inner(self, timeout=timeout) + + self.synchronize_device = _counted + return True + + def synchronize_device_seen(self) -> bool: + """True from every rank only when each rank ran the barrier at least once.""" + return getattr(self, "_synchronize_device_calls", 0) > 0 + + def set_synchronize_device_failure(self, rank: int, fail: bool) -> bool: + """Make (or stop making) the pause barrier fail on ``rank`` (test-only).""" + if self.rank != rank: + return True + if fail: + + def _fail(timeout=None): + raise RuntimeError(f"injected barrier failure on rank {rank}") + + self.synchronize_device = _fail + else: + self.__dict__.pop("synchronize_device", None) + return True + # --------------------------------------------------------------------------- # Pipeline layer diff --git a/tests/engine/test_orchestrator_error_handling.py b/tests/engine/test_orchestrator_error_handling.py index b8949786e..21a26279d 100644 --- a/tests/engine/test_orchestrator_error_handling.py +++ b/tests/engine/test_orchestrator_error_handling.py @@ -13,6 +13,7 @@ import asyncio import queue import time from types import SimpleNamespace +from unittest.mock import AsyncMock import janus import pytest @@ -22,6 +23,7 @@ from vllm.v1.serial_utils import MsgpackEncoder from vllm_omni.engine.messages import ( AddCompanionRequestMessage, + CollectiveRPCRequestMessage, EngineQueueMessage, ErrorMessage, ShutdownRequestMessage, @@ -1143,3 +1145,64 @@ async def test_diffusion_client_error_output_propagates_non_400_status( finally: orchestrator_fixture.request_sync_q.put_nowait(ShutdownRequestMessage()) orchestrator_fixture.thread.join(timeout=5) + + +@pytest.mark.asyncio +async def test_control_rpc_failure_is_reported_to_the_caller_not_fatal() -> None: + """A control RPC that fails on one replica must come back as that + replica's error result instead of raising out of _request_handler. + """ + + class _FailingClient(FakeStageClient): + async def collective_rpc_async(self, *args, **kwargs): + raise TimeoutError("outputs did not drain") + + ok = FakeStageClient(stage_type="diffusion") + ok.collective_rpc_async = AsyncMock(return_value=None) + orchestrator, queues = _build_bare_orchestrator( + _build_stage_pools([[ok], [_FailingClient(stage_type="diffusion")]]) + ) + try: + await orchestrator._handle_collective_rpc( + CollectiveRPCRequestMessage( + rpc_id="rpc-1", + method="pause_scheduler", + args=(), + kwargs={"mode": "keep"}, + stage_ids=[0, 1], + ) + ) + + result = queues[2].async_q.get_nowait() + assert result.rpc_id == "rpc-1" + assert result.stage_ids == [0, 1] + assert result.results[0] is None + assert result.results[1]["supported"] is False + assert "outputs did not drain" in result.results[1]["error"] + finally: + for q in queues: + q.close() + + +@pytest.mark.asyncio +async def test_rpc_failure_capture_is_limited_to_pause_and_resume() -> None: + """Only pause/resume failures are reported as a result; every other RPC + keeps propagating the way it does today. + """ + orchestrator, queues = _build_bare_orchestrator(_build_stage_pools([[FakeStageClient(stage_type="diffusion")]])) + orchestrator.stage_pools[0].collective_rpc = AsyncMock(side_effect=TimeoutError("worker died")) + + def _msg(rpc_id: str, method: str) -> CollectiveRPCRequestMessage: + return CollectiveRPCRequestMessage(rpc_id=rpc_id, method=method, args=(), kwargs={}, stage_ids=[0]) + + try: + await orchestrator._handle_collective_rpc(_msg("rpc-resume", "resume_scheduler")) + result = queues[2].async_q.get_nowait() + assert result.results[0]["supported"] is False + assert "worker died" in result.results[0]["error"] + + with pytest.raises(TimeoutError): + await orchestrator._handle_collective_rpc(_msg("rpc-sleep", "sleep")) + finally: + for q in queues: + q.close() diff --git a/tests/entrypoints/test_async_omni_pause_sleep_routing.py b/tests/entrypoints/test_async_omni_pause_sleep_routing.py index 38ec40163..d3a60e5a6 100644 --- a/tests/entrypoints/test_async_omni_pause_sleep_routing.py +++ b/tests/entrypoints/test_async_omni_pause_sleep_routing.py @@ -25,6 +25,7 @@ def _make_omni(*, stage_types: list[str]) -> AsyncOmni: omni._paused = False omni._admitting = 0 omni._hold_admission_until_resume = False + omni._paused_stage_ids = set() omni._sleeping_tags = set() omni._stage_sleeping_tags = {} omni._level2_sleeping = False @@ -460,3 +461,225 @@ def test_partial_wake_does_not_skip_remaining_sleeping_stage(): assert not omni._stage_sleeping_tags asyncio.run(run()) + + +# ───────────────────── diffusion pause_generation(mode="keep") ───────────────────── + + +def _rpc_methods(omni: AsyncOmni) -> list[tuple[str, list[int] | None]]: + return [(call.kwargs["method"], call.kwargs.get("stage_ids")) for call in omni.collective_rpc.await_args_list] + + +@pytest.mark.cpu +def test_diffusion_keep_routes_to_backend(): + async def run() -> None: + omni = _make_omni(stage_types=["diffusion"]) + + await omni.pause_generation(mode="keep", clear_cache=False) + + omni.collective_rpc.assert_awaited_once_with( + method="pause_scheduler", + args=(), + kwargs={"mode": "keep"}, + stage_ids=[0], + ) + assert await omni.is_paused() is True + assert omni._hold_admission_until_resume is True + + asyncio.run(run()) + + +@pytest.mark.cpu +@pytest.mark.parametrize( + "pause_kwargs", + [ + {}, + {"mode": "abort"}, + {"mode": "wait"}, + {"mode": "keep", "wait_for_inflight_requests": True}, + ], +) +def test_diffusion_non_keep_modes_pause_frontend_only(pause_kwargs): + async def run() -> None: + omni = _make_omni(stage_types=["diffusion"]) + omni.reset_prefix_cache = AsyncMock(return_value=True) + omni.reset_mm_cache = AsyncMock() + omni.reset_encoder_cache = AsyncMock() + + await omni.pause_generation(**pause_kwargs) + assert await omni.is_paused() is True + omni.collective_rpc.assert_not_awaited() + + await omni.resume_generation() + omni.collective_rpc.assert_not_awaited() + assert await omni.is_paused() is False + + asyncio.run(run()) + + +@pytest.mark.cpu +def test_mixed_engine_keep_sends_engine_core_and_diffusion_pauses(): + async def run() -> None: + omni = _make_omni(stage_types=["llm", "diffusion"]) + omni.reset_prefix_cache = AsyncMock(return_value=True) + omni.reset_mm_cache = AsyncMock() + omni.reset_encoder_cache = AsyncMock() + + await omni.pause_generation(mode="keep", clear_cache=True) + + assert omni.collective_rpc.await_args_list[0].kwargs == { + "method": "pause_scheduler", + "args": (), + "kwargs": {"mode": "keep", "clear_cache": True}, + "stage_ids": [0], + } + assert omni.collective_rpc.await_args_list[1].kwargs == { + "method": "pause_scheduler", + "args": (), + "kwargs": {"mode": "keep"}, + "stage_ids": [1], + } + # The frontend cache reset still runs after the backend pauses. + omni.reset_prefix_cache.assert_awaited_once() + + await omni.resume_generation() + assert _rpc_methods(omni)[2:] == [("resume_scheduler", [0]), ("resume_scheduler", [1])] + + asyncio.run(run()) + + +@pytest.mark.cpu +def test_keep_targets_only_requested_diffusion_stages(): + async def run() -> None: + omni = _make_omni(stage_types=["diffusion", "diffusion"]) + + await omni.pause_generation(mode="keep", clear_cache=False, stage_ids=[1]) + assert _rpc_methods(omni) == [("pause_scheduler", [1])] + + await omni.resume_generation(stage_ids=[0]) + assert _rpc_methods(omni) == [("pause_scheduler", [1])] + # Stage 1's scheduler is still closed: admission must stay closed too, + # otherwise new requests queue on a stage that will not schedule them. + assert await omni.is_paused() is True + + await omni.resume_generation(stage_ids=[1]) + assert _rpc_methods(omni) == [("pause_scheduler", [1]), ("resume_scheduler", [1])] + assert await omni.is_paused() is False + + await omni.resume_generation() + assert len(_rpc_methods(omni)) == 2 + + asyncio.run(run()) + + +@pytest.mark.cpu +def test_keep_then_sleep_wake_resume_keeps_admission_paused_until_resume(): + async def run() -> None: + omni = _make_omni(stage_types=["diffusion"]) + omni.reset_prefix_cache = AsyncMock(return_value=True) + omni.reset_encoder_cache = AsyncMock() + omni.collective_rpc = AsyncMock( + side_effect=lambda **kw: [OmniACK(task_id="t", status="SUCCESS", stage_id=0, rank=0)] + if kw["method"] in {"handle_sleep_task", "handle_wake_task"} + else [None] + ) + + await omni.pause_generation(mode="keep", clear_cache=False) + await omni.sleep(level=1) + await omni.wake_up() + assert await omni.is_paused() is True + + await omni.resume_generation() + assert await omni.is_paused() is False + assert [m for m, _ in _rpc_methods(omni)] == [ + "pause_scheduler", + "handle_sleep_task", + "handle_wake_task", + "resume_scheduler", + ] + + asyncio.run(run()) + + +@pytest.mark.cpu +@pytest.mark.parametrize( + "failure", + [ + AsyncMock(side_effect=NotImplementedError("request-level execution only")), + AsyncMock(return_value=[{"supported": False, "error": "TimeoutError: outputs did not drain"}]), + ], +) +def test_keep_failure_surfaces_and_leaves_explicit_resume_possible(failure): + async def run() -> None: + omni = _make_omni(stage_types=["diffusion"]) + omni.collective_rpc = failure + + with pytest.raises((NotImplementedError, RuntimeError)): + await omni.pause_generation(mode="keep", clear_cache=False) + + assert await omni.is_paused() is True + assert omni._paused_stage_ids == {0} + + omni.collective_rpc = AsyncMock(return_value=[None]) + await omni.resume_generation() + omni.collective_rpc.assert_awaited_once_with( + method="resume_scheduler", + args=(), + kwargs=None, + stage_ids=[0], + ) + assert await omni.is_paused() is False + assert omni._paused_stage_ids == set() + + asyncio.run(run()) + + +@pytest.mark.cpu +def test_resume_failure_keeps_admission_paused_and_target_recorded(): + async def run() -> None: + omni = _make_omni(stage_types=["diffusion"]) + await omni.pause_generation(mode="keep", clear_cache=False) + + omni.collective_rpc = AsyncMock(return_value=[{"error": "replica gone"}]) + with pytest.raises(RuntimeError, match="resume_scheduler failed"): + await omni.resume_generation() + + assert await omni.is_paused() is True + assert omni._paused_stage_ids == {0} + + asyncio.run(run()) + + +@pytest.mark.cpu +@pytest.mark.parametrize("first, second", [([1], [0]), ([0], [1])]) +def test_mixed_engine_partial_resume_keeps_admission_paused(first, second): + async def run() -> None: + omni = _make_omni(stage_types=["llm", "diffusion"]) + await omni.pause_generation(mode="keep", clear_cache=False) + + await omni.resume_generation(stage_ids=first) + assert await omni.is_paused() is True + + await omni.resume_generation(stage_ids=second) + assert await omni.is_paused() is False + assert _rpc_methods(omni)[2:] == [("resume_scheduler", first), ("resume_scheduler", second)] + + asyncio.run(run()) + + +@pytest.mark.cpu +def test_two_ar_stages_partial_resume_keeps_admission_paused(): + async def run() -> None: + omni = _make_omni(stage_types=["llm", "llm"]) + omni.reset_prefix_cache = AsyncMock(return_value=True) + omni.reset_mm_cache = AsyncMock() + omni.reset_encoder_cache = AsyncMock() + await omni.pause_generation(mode="abort") + + await omni.resume_generation(stage_ids=[0]) + assert await omni.is_paused() is True + + await omni.resume_generation(stage_ids=[1]) + assert await omni.is_paused() is False + + asyncio.run(run()) diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index 5e3082f45..6440c2fe8 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/vllm_omni/diffusion/diffusion_engine.py @@ -228,9 +228,10 @@ class DiffusionEngine: #: remains a pure explicit user override (never mutated by engines). default_diffusion_model_runner_cls: str | None = None - # Class-level default so tests using object.__new__ (without __init__) - # don't hit AttributeError when _busy_loop accesses self.dp_concurrent. + # Class-level defaults so tests using object.__new__ (without __init__) + # don't hit AttributeError when _busy_loop accesses them. dp_concurrent: bool = False + _scheduling_paused: bool = False def __init__( self, @@ -387,6 +388,8 @@ class DiffusionEngine: self._shutdown_complete = False self.abort_queue: queue.Queue[str] = queue.Queue() self._rpc_queue: queue.Queue[_RpcTask] = queue.Queue() + # pause_scheduler(mode="keep"): no new batch is scheduled while set. + self._scheduling_paused = False # Copied onto the existing output metrics payload so queue monitoring # reuses the normal diffusion result path without additional IPC. self._scheduler_num_waiting_reqs = 0 @@ -584,10 +587,16 @@ class DiffusionEngine: while not self.stop_event.is_set(): self._process_aborts_queue() self._process_rpc_queue() + if self._scheduling_paused: + # No wave runs while paused, so requests finished by an abort + # would otherwise wait for the resume to surface. + with self._cv: + pending_finished = self.scheduler.pending_finished_request_ids() + self._emit_finished_outputs(pending_finished, None) with self._cv: while ( - not self.scheduler.has_requests() + (self._scheduling_paused or not self.scheduler.has_requests()) and self._rpc_queue.empty() and self.abort_queue.empty() and not self.stop_event.is_set() @@ -597,11 +606,14 @@ class DiffusionEngine: if self.stop_event.is_set(): break - if not self.scheduler.has_requests(): + if self._scheduling_paused or not self.scheduler.has_requests(): # Only RPC / abort work pending; loop back to drain it. continue self._wait_for_admission_if_needed_locked() + if self._scheduling_paused: + # The gate closed while the batching wait released the lock. + continue try: sched_output = self.scheduler.schedule() @@ -693,7 +705,7 @@ class DiffusionEngine: last_waiting = -1 stable_since = start - while not self.stop_event.is_set(): + while not self.stop_event.is_set() and not self._scheduling_paused: waiting = self.scheduler.num_waiting_requests() now = time.monotonic() @@ -733,19 +745,34 @@ class DiffusionEngine: except queue.Empty: return - fut = task.future - if fut.cancelled() or fut.done(): - continue + self._run_rpc_task(task) - remaining: float | None = None - if task.deadline is not None: - remaining = task.deadline - time.monotonic() - if remaining <= 0: - if not fut.done(): - fut.set_exception(TimeoutError(f"RPC call to {task.method} timed out before execution.")) - continue + if task.method == "pause_scheduler": + # A dequeued pause ends this drain whatever became of it, so + # the batch that just ran is delivered before anything the + # caller queued behind the pause. The rest of the queue is + # picked up on the next pass of the busy loop. + return - try: + def _run_rpc_task(self, task: _RpcTask) -> None: + fut = task.future + if fut.cancelled() or fut.done(): + return + + remaining: float | None = None + if task.deadline is not None: + remaining = task.deadline - time.monotonic() + if remaining <= 0: + if not fut.done(): + fut.set_exception(TimeoutError(f"RPC call to {task.method} timed out before execution.")) + return + + try: + if task.method == "pause_scheduler": + result = self._run_pause_barrier(remaining) + elif task.method == "resume_scheduler": + result = self._open_scheduling_gate() + else: result = self.executor.collective_rpc( method=task.method, timeout=remaining, @@ -753,16 +780,16 @@ class DiffusionEngine: kwargs=task.kwargs, unique_reply_rank=task.unique_reply_rank, ) - except BaseException as exc: # noqa: BLE001 - propagate to caller - # The future may have been cancelled (e.g. by a sync timeout - # or asyncio cancellation) while the executor call was - # running. Setting state on a cancelled/done future raises - # InvalidStateError, which would kill the busy loop. - if not fut.done(): - fut.set_exception(exc) - else: - if not fut.done(): - fut.set_result(result) + except BaseException as exc: # noqa: BLE001 - propagate to caller + # The future may have been cancelled (e.g. by a sync timeout + # or asyncio cancellation) while the executor call was + # running. Setting state on a cancelled/done future raises + # InvalidStateError, which would kill the busy loop. + if not fut.done(): + fut.set_exception(exc) + else: + if not fut.done(): + fut.set_result(result) def _fail_pending_rpcs(self, exc: BaseException) -> None: while True: @@ -773,6 +800,30 @@ class DiffusionEngine: if not task.future.done(): task.future.set_exception(exc) + def _check_pause_request(self, kwargs: dict | None) -> None: + mode = (kwargs or {}).get("mode") + if mode != "keep": + raise ValueError(f"DiffusionEngine pause supports mode='keep' only, got {mode!r}.") + if self.execution_mode != DiffusionExecutionMode.REQUEST_BATCH: + raise NotImplementedError("DiffusionEngine pause supports request-level execution only.") + + def _run_pause_barrier(self, timeout: float | None) -> None: + """Return once every worker has finished the device work of the batch that ran.""" + self.executor.collective_rpc("synchronize_device", timeout=timeout, kwargs={"timeout": timeout}) + + def _open_scheduling_gate(self) -> None: + with self._cv: + self._scheduling_paused = False + self._cv.notify_all() + + def _run_engine_control(self, method: str, timeout: float | None, kwargs: dict | None) -> None: + """Bootstrap-path dispatch for engine-local control methods.""" + if method == "pause_scheduler": + self._check_pause_request(kwargs) + self._scheduling_paused = True + return self._run_pause_barrier(timeout) + return self._open_scheduling_gate() + def _remove_diffusion_kv_requests(self, request_ids: Iterable[str]) -> None: """Clear terminal Worker rows while Scheduler owns the allocations.""" @@ -1309,6 +1360,9 @@ class DiffusionEngine: unique_reply_rank=unique_reply_rank, ) with self._cv: + if method == "pause_scheduler": + self._check_pause_request(kwargs) + self._scheduling_paused = True self._rpc_queue.put(task) self._cv.notify_all() return task @@ -1351,6 +1405,8 @@ class DiffusionEngine: # between the outer check and acquiring the lock, in which # case we should use the queued path for proper ordering. if not self._loop_started: + if method in ("pause_scheduler", "resume_scheduler"): + return self._run_engine_control(method, timeout, kwargs) return self.executor.collective_rpc( method=method, timeout=timeout, diff --git a/vllm_omni/diffusion/sched/base_scheduler.py b/vllm_omni/diffusion/sched/base_scheduler.py index 4fd695dc5..f0ca6c3ac 100644 --- a/vllm_omni/diffusion/sched/base_scheduler.py +++ b/vllm_omni/diffusion/sched/base_scheduler.py @@ -343,6 +343,10 @@ class BaseScheduler(ABC): def num_running_requests(self) -> int: return len(self._running) + def pending_finished_request_ids(self) -> set[str]: + """Finished requests whose state the engine has not consumed yet.""" + return {request_id for request_id in self._finished_req_ids if request_id in self._request_states} + def get_admission_wait_decision( self, *, diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py index 4eadf33db..f4574b35e 100644 --- a/vllm_omni/diffusion/worker/diffusion_worker.py +++ b/vllm_omni/diffusion/worker/diffusion_worker.py @@ -877,6 +877,19 @@ class DiffusionWorker: logger.info(f"[Worker {self.rank}] Wake-up complete.") return True + def synchronize_device(self, timeout: float | None = None) -> None: + """Wait until this rank has no device work left. + + A KV prefetch still receiving on its background thread has queued no + device work yet, so it is joined first. ``timeout`` bounds that join + (and the multi-process worker's output drain before this call); the + device wait itself is unbounded. + """ + manager = getattr(self.model_runner, "kv_transfer_manager", None) + if manager is not None and not manager.wait_prefetch(timeout=timeout): + raise TimeoutError("Diffusion KV prefetch did not finish before pause") + current_omni_platform.synchronize() + def handle_sleep_task(self, task: OmniSleepTask | dict) -> OmniACK | None: from vllm_omni.platforms import current_omni_platform @@ -1214,11 +1227,13 @@ class WorkerProc: self._async_output_pending = max(0, self._async_output_pending - 1) self._async_output_done.notify_all() - def drain_async_outputs(self, timeout: float = _ASYNC_OUTPUT_DRAIN_TIMEOUT_S) -> bool: + def drain_async_outputs(self, timeout: float | None = None) -> bool: """Block until background D2H/SHM packing has no work left. Returns False if outputs are still in flight when ``timeout`` expires. """ + if timeout is None: + timeout = _ASYNC_OUTPUT_DRAIN_TIMEOUT_S with self._async_output_done: if self._async_output_pending == 0: return True @@ -1334,6 +1349,8 @@ class WorkerProc: } try: + if method == "synchronize_device" and not self.drain_async_outputs(timeout=kwargs.get("timeout")): + raise TimeoutError("Diffusion async outputs did not drain before pause") if method in _MEMORY_RELEASING_METHODS: self.drain_async_outputs() # Use execute_method from WorkerWrapperBase for consistent method resolution diff --git a/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py b/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py index 6ab0b36a2..77731b013 100644 --- a/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py +++ b/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py @@ -8,6 +8,7 @@ import struct import time from collections.abc import Callable from concurrent.futures import ThreadPoolExecutor +from concurrent.futures import wait as wait_futures from dataclasses import asdict, dataclass from typing import Any @@ -1346,6 +1347,16 @@ class OmniKVTransferManager: logger.exception("KV load failed for %s; falling back to sync receive", request_id) return None, 0 + def wait_prefetch(self, timeout: float | None = None) -> bool: + """Block until no background prefetch is running; payloads stay consumable. + + Returns False if a prefetch is still in flight when ``timeout`` expires. + """ + futures = list(self._prefetch_futures.values()) + if not futures: + return True + return not wait_futures(futures, timeout=timeout).not_done + def _discard_future(self, request_id: str) -> None: """Cancel an unstarted prefetch or attach a callback to drop a running one.""" fut = self._prefetch_futures.pop(request_id, None) diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index abcd92874..351f0204c 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -696,13 +696,23 @@ class OrchestratorBase: stage_ids: list[int] = [] for pool in target_pools: for replica_id in pool.live_replica_ids(): - stage_result = await pool.collective_rpc( - replica_id=replica_id, - method=method, - timeout=timeout, - args=args, - kwargs=kwargs, - ) + try: + stage_result = await pool.collective_rpc( + replica_id=replica_id, + method=method, + timeout=timeout, + args=args, + kwargs=kwargs, + ) + except Exception as exc: + if method not in ("pause_scheduler", "resume_scheduler"): + raise + # A pause or resume that fails on one replica leaves the + # engine usable, so report it as this replica's result + # rather than out of the request handler; the caller + # reaches both through _engine_core_rpc, which raises on + # the error result and can then retry or resume. + stage_result = {"supported": False, "error": f"{type(exc).__name__}: {exc}"} stage_ids.append(pool.stage_id) results.append(stage_result) diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index 1291e68f0..ba4e75a11 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/vllm_omni/entrypoints/async_omni.py @@ -95,6 +95,10 @@ class AsyncOmni(AsyncOmniBase, EngineClient): # sleep uses _paused as a temporary admission gate and clears it # on wake so sleep → wake → generate keeps working. self._hold_admission_until_resume: bool = False + # Stages whose scheduler pause_generation closed (AR in any mode, + # diffusion with mode="keep"); admission stays closed until + # resume_generation has reopened all of them. + self._paused_stage_ids: set[int] = set() self._sleeping_tags: set[str] = set() self._stage_sleeping_tags: dict[int, set[str]] = {} self._level2_sleeping: bool = False @@ -542,10 +546,12 @@ class AsyncOmni(AsyncOmniBase, EngineClient): args: tuple[Any, ...] = (), kwargs: dict[str, Any] | None = None, ) -> list[Any]: - """Call an AR EngineCore helper via collective_rpc (orchestrator loop). + """Call an engine control helper via collective_rpc (orchestrator loop). StagePool resolves ``{method}_async`` on the AR client when present - (vLLM AsyncMPClient convention). Raises if any replica reports failure. + (vLLM AsyncMPClient convention); diffusion stages answer the same + method names inside DiffusionEngine. Raises if any replica reports + failure. """ results = await self.collective_rpc( method=method, @@ -695,8 +701,13 @@ class AsyncOmni(AsyncOmniBase, EngineClient): 1. Stop frontend admission (``_paused``). 2. For AR/LLM stages, call EngineCore.pause_scheduler via the Orchestrator loop (abort/wait/keep + optional cache clear). - 3. Diffusion stages have no EngineCore scheduler — only frontend - admission is paused for them. + 3. For diffusion stages, ``mode="keep"`` pauses the DiffusionEngine + scheduler and returns once the batch that was running has finished + on every worker; that batch is delivered before any control RPC + issued after this call runs, so the documented pause -> sleep order + is safe. Queued requests stay queued until + :meth:`resume_generation`. Other modes pause frontend admission + only. Note: ``sleep()`` already pauses the AR scheduler internally (same as vLLM EngineCore.sleep). Call this API when you need pause *without* @@ -711,7 +722,12 @@ class AsyncOmni(AsyncOmniBase, EngineClient): self._paused = True self._hold_admission_until_resume = True - ar_stage_ids, _diffusion_stage_ids = self._split_stage_ids_by_type(stage_ids) + ar_stage_ids, diffusion_stage_ids = self._split_stage_ids_by_type(stage_ids) + if mode != "keep": + diffusion_stage_ids = [] + # Recorded before the RPCs so a failed or cancelled pause can still + # be undone with an explicit resume_generation. + self._paused_stage_ids.update(ar_stage_ids, diffusion_stage_ids) if ar_stage_ids: logger.info( "[%s] Pausing AR stage(s) %s via EngineCore.pause_scheduler(mode=%s)", @@ -726,6 +742,17 @@ class AsyncOmni(AsyncOmniBase, EngineClient): stage_ids=ar_stage_ids, kwargs={"mode": mode, "clear_cache": clear_cache}, ) + if diffusion_stage_ids: + logger.info( + "[%s] Pausing diffusion stage(s) %s via DiffusionEngine pause_scheduler(mode=keep)", + self._name, + diffusion_stage_ids, + ) + await self._engine_core_rpc( + "pause_scheduler", + stage_ids=diffusion_stage_ids, + kwargs={"mode": "keep"}, + ) # Frontend / sender-side cache clear (P0). EngineCore.pause_scheduler # already clears AR-side caches when clear_cache=True. @@ -739,10 +766,26 @@ class AsyncOmni(AsyncOmniBase, EngineClient): async def resume_generation(self, stage_ids: list[int] | None = None) -> None: """Resume generation after :meth:`pause_generation`.""" - ar_stage_ids, _diffusion_stage_ids = self._split_stage_ids_by_type(stage_ids) + ar_stage_ids, diffusion_stage_ids = self._split_stage_ids_by_type(stage_ids) if ar_stage_ids: logger.info("[%s] Resuming AR stage(s) %s via EngineCore", self._name, ar_stage_ids) await self._engine_core_rpc("resume_scheduler", stage_ids=ar_stage_ids) + self._paused_stage_ids.difference_update(ar_stage_ids) + diffusion_stage_ids = [sid for sid in diffusion_stage_ids if sid in self._paused_stage_ids] + if diffusion_stage_ids: + logger.info("[%s] Resuming diffusion stage(s) %s via DiffusionEngine", self._name, diffusion_stage_ids) + await self._engine_core_rpc("resume_scheduler", stage_ids=diffusion_stage_ids) + self._paused_stage_ids.difference_update(diffusion_stage_ids) + + if self._paused_stage_ids: + # Reopening admission now would let new requests queue on a stage + # whose scheduler is still closed. + logger.info( + "[%s] Admission stays paused: stage(s) %s are still paused", + self._name, + sorted(self._paused_stage_ids), + ) + return async with self._pause_cond: self._paused = False @@ -816,8 +859,9 @@ class AsyncOmni(AsyncOmniBase, EngineClient): AR/LLM stages use EngineCore.sleep (pause scheduler, wait idle, then offload/discard memory) — matching vLLM AsyncLLM.sleep. - Diffusion stages keep the existing worker-level handle_sleep_task RPC - because StageDiffusionProc does not expose EngineCore.pause_scheduler. + Diffusion stages keep the worker-level handle_sleep_task RPC, which + does not stop the DiffusionEngine scheduler; quiesce a busy diffusion + stage first with ``pause_generation(mode="keep")`` (or abort it). Frontend admission is blocked at the start of this call (``_paused``) so pipelined :meth:`generate` cannot race into stages while sleep is