diff --git a/tests/engine/test_orchestrator_error_handling.py b/tests/engine/test_orchestrator_error_handling.py index 21a26279d..f05c5f132 100644 --- a/tests/engine/test_orchestrator_error_handling.py +++ b/tests/engine/test_orchestrator_error_handling.py @@ -1206,3 +1206,110 @@ async def test_rpc_failure_capture_is_limited_to_pause_and_resume() -> None: finally: for q in queues: q.close() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("method", ["reset_mm_cache", "reset_encoder_cache", "reset_prefix_cache"]) +async def test_cache_reset_rpc_failure_is_nonfatal_and_visits_all_replicas(method): + clients = [FakeStageClient(stage_type="llm") for _ in range(2)] + setattr(clients[0], f"{method}_async", AsyncMock(side_effect=RuntimeError("reset failed"))) + ok = AsyncMock(return_value=True) + setattr(clients[1], f"{method}_async", ok) + orchestrator, queues = _build_bare_orchestrator(_build_stage_pools([clients])) + try: + await orchestrator._handle_collective_rpc( + CollectiveRPCRequestMessage(rpc_id="reset", method=method, args=(), kwargs={}, stage_ids=[0], timeout=1.0) + ) + result = queues[2].async_q.get_nowait() + assert result.results[0]["supported"] is False + assert "reset failed" in result.results[0]["error"] + assert result.results[1] is True + ok.assert_awaited_once_with() + # Errors before StagePool dispatch (e.g. a missing route) must also + # be serialized at the orchestrator boundary, then allow another RPC. + orchestrator.stage_pools[0].collective_rpc = AsyncMock(side_effect=RuntimeError("route failed")) + await orchestrator._handle_collective_rpc( + CollectiveRPCRequestMessage(rpc_id="route", method=method, args=(), kwargs={}, stage_ids=[0]) + ) + assert all("route failed" in r["error"] for r in queues[2].async_q.get_nowait().results) + orchestrator.stage_pools[0].collective_rpc = AsyncMock(return_value=True) + await orchestrator._handle_collective_rpc( + CollectiveRPCRequestMessage(rpc_id="next", method=method, args=(), kwargs={}, stage_ids=[0]) + ) + assert queues[2].async_q.get_nowait().results == [True, True] + finally: + for q in queues: + q.close() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "method, kwargs, clears", + [ + ("reset_mm_cache", {}, True), + ("sleep", {}, True), + ("pause_scheduler", {"clear_cache": True}, True), + ("pause_scheduler", {"clear_cache": False}, False), + ("reset_prefix_cache", {}, False), + ("reset_encoder_cache", {}, False), + ], +) +async def test_control_rpc_clears_only_selected_downstream_sender(method, kwargs, clears): + pools = _build_stage_pools([[FakeStageClient(stage_type="llm") for _ in range(2)] for _ in range(2)]) + orchestrator, queues = _build_bare_orchestrator(pools) + events = [] + renderers = [SimpleNamespace(clear_mm_cache_async=AsyncMock()) for _ in range(2)] + renderers[1].clear_mm_cache_async.side_effect = lambda: events.append("sender") + orchestrator._stage_input_processors = {i: SimpleNamespace(renderer=r) for i, r in enumerate(renderers)} + + async def rpc(**kwargs): + events.append("receiver") + return True + + pools[0].collective_rpc = AsyncMock() + pools[1].collective_rpc = AsyncMock(side_effect=rpc) + try: + await orchestrator._handle_collective_rpc( + CollectiveRPCRequestMessage( + rpc_id="reset", method=method, args=(), kwargs=kwargs, stage_ids=[1], timeout=1.0 + ) + ) + renderers[0].clear_mm_cache_async.assert_not_awaited() + pools[0].collective_rpc.assert_not_awaited() + assert renderers[1].clear_mm_cache_async.await_count == int(clears) + assert events == (["sender"] if clears else []) + ["receiver", "receiver"] + assert queues[2].async_q.get_nowait().results == [True, True] + finally: + for q in queues: + q.close() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("failure", ["error", "timeout"]) +async def test_downstream_sender_clear_failure_is_nonfatal(failure): + pools = _build_stage_pools([[FakeStageClient(stage_type="llm")]]) + orchestrator, queues = _build_bare_orchestrator(pools) + + async def clear(): + if failure == "timeout": + await asyncio.Event().wait() + raise RuntimeError("sender cache failed") + + renderer = SimpleNamespace(clear_mm_cache_async=AsyncMock(side_effect=clear)) + orchestrator._stage_input_processors[0] = SimpleNamespace(renderer=renderer) + pools[0].collective_rpc = AsyncMock(return_value=None) + message = CollectiveRPCRequestMessage( + rpc_id="reset", method="reset_mm_cache", args=(), kwargs={}, stage_ids=[0], timeout=0.01 + ) + try: + await orchestrator._handle_collective_rpc(message) + result = queues[2].async_q.get_nowait().results[0] + assert result["supported"] is False + assert ("TimeoutError" if failure == "timeout" else "sender cache failed") in result["error"] + pools[0].collective_rpc.assert_not_awaited() + renderer.clear_mm_cache_async = AsyncMock() + await orchestrator._handle_collective_rpc(message) + assert queues[2].async_q.get_nowait().results == [None] + finally: + for q in queues: + q.close() diff --git a/tests/engine/test_stage_pool_collective_rpc.py b/tests/engine/test_stage_pool_collective_rpc.py index e8a228827..281b4b504 100644 --- a/tests/engine/test_stage_pool_collective_rpc.py +++ b/tests/engine/test_stage_pool_collective_rpc.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Unit tests for StagePool.collective_rpc EngineCore control dispatch.""" from __future__ import annotations @@ -40,14 +43,14 @@ def test_collective_rpc_normalizes_none_args_on_control_helper(): def test_collective_rpc_unrelated_async_helper_uses_collective_path(): async def run() -> None: other = AsyncMock(return_value="should-not-run") - pool, client = _make_pool(reset_prefix_cache_async=other) + pool, client = _make_pool(is_sleeping_async=other) - result = await pool.collective_rpc(0, "reset_prefix_cache", timeout=1.5, args=("x",), kwargs={"k": 1}) + result = await pool.collective_rpc(0, "is_sleeping", timeout=1.5, args=("x",), kwargs={"k": 1}) assert result == {"via": "collective"} other.assert_not_awaited() client.collective_rpc_async.assert_awaited_once_with( - method="reset_prefix_cache", + method="is_sleeping", timeout=1.5, args=("x",), kwargs={"k": 1}, @@ -63,7 +66,7 @@ def test_collective_rpc_control_helper_honors_timeout(): await asyncio.sleep(1.0) return "slept" - pool, _client = _make_pool(sleep_async=slow_sleep) + pool, _client = _make_pool(sleep_async=AsyncMock(side_effect=slow_sleep)) with pytest.raises(asyncio.TimeoutError): await pool.collective_rpc(0, "sleep", timeout=0.01, args=(1,)) @@ -80,11 +83,38 @@ def test_collective_rpc_control_method_reraises_worker_error(): asyncio.run(run()) +@pytest.mark.cpu +def test_collective_rpc_reset_caches_use_control_helpers(): + async def run() -> None: + reset_prefix = AsyncMock(return_value=True) + reset_encoder = AsyncMock(return_value=None) + reset_mm = AsyncMock(return_value=None) + pool, client = _make_pool( + reset_prefix_cache_async=reset_prefix, + reset_encoder_cache_async=reset_encoder, + reset_mm_cache_async=reset_mm, + ) + + result = await pool.collective_rpc(0, "reset_prefix_cache", args=(True, True), timeout=2.0) + assert result is True + reset_prefix.assert_awaited_once_with(True, True) + + await pool.collective_rpc(0, "reset_encoder_cache", timeout=2.0) + reset_encoder.assert_awaited_once() + + await pool.collective_rpc(0, "reset_mm_cache", timeout=2.0) + reset_mm.assert_awaited_once() + + client.collective_rpc_async.assert_not_awaited() + + asyncio.run(run()) + + @pytest.mark.cpu def test_collective_rpc_non_control_method_still_returns_error_dict(): async def run() -> None: pool, _client = _make_pool(collective_rpc_async=AsyncMock(side_effect=RuntimeError("probe failed"))) - result = await pool.collective_rpc(0, "reset_prefix_cache") + result = await pool.collective_rpc(0, "is_sleeping") assert result["supported"] is False assert "probe failed" in result["error"] @@ -122,3 +152,28 @@ def test_abort_requests_does_not_commit_op_state_when_engine_abort_fails(): abort.assert_awaited_once() asyncio.run(run()) + + +@pytest.mark.cpu +@pytest.mark.parametrize("method", ["reset_prefix_cache", "reset_encoder_cache", "reset_mm_cache"]) +@pytest.mark.parametrize("failure", ["error", "timeout", "missing"]) +def test_cache_reset_failure_is_serialized_and_pool_remains_usable(method, failure): + async def run(): + async def reset(): + if failure == "timeout": + await asyncio.Event().wait() + raise RuntimeError("cache reset failed") + + pool, client = _make_pool() + if failure != "missing": + setattr(client, f"{method}_async", AsyncMock(side_effect=reset)) + result = await pool.collective_rpc(0, method, timeout=0.01) + assert result["supported"] is False + assert result["error"] + if failure == "timeout": + assert "timed out" in result["error"] + client.collective_rpc_async.assert_not_awaited() + # A subsequent control request still completes on the same pool. + assert await pool.collective_rpc(0, "is_sleeping") == {"via": "collective"} + + asyncio.run(run()) diff --git a/tests/entrypoints/test_async_omni_cache_reset.py b/tests/entrypoints/test_async_omni_cache_reset.py new file mode 100644 index 000000000..eb0d610c6 --- /dev/null +++ b/tests/entrypoints/test_async_omni_cache_reset.py @@ -0,0 +1,77 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Real EngineCore cache invalidation through the AsyncOmni control plane.""" + +import pytest +from PIL import Image +from vllm import SamplingParams + +from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config +from vllm_omni.entrypoints.async_omni import AsyncOmni + + +@pytest.mark.advanced_model +@pytest.mark.omni +@pytest.mark.cuda +@pytest.mark.asyncio +async def test_cache_resets_invalidate_real_engine_and_allow_repeated_multimodal_inputs(): + deploy_config = modify_stage_config( + get_deploy_config_path("ci/qwen2_5_omni_thinker_only.yaml"), + updates={ + "stages": { + 0: { + "enable_prefix_caching": True, + "mm_processor_cache_gb": 0.1, + "max_model_len": 2048, + "max_num_batched_tokens": 2048, + "num_gpu_blocks_override": 256, + "enforce_eager": True, + } + } + }, + ) + engine = AsyncOmni( + model="Qwen/Qwen2.5-Omni-7B", + deploy_config=deploy_config, + stage_init_timeout=600, + ) + + async def generate(prompt, request_id): + final = None + async for output in engine.generate( + prompt, + request_id=request_id, + sampling_params_list=[SamplingParams(max_tokens=4, temperature=0)], + ): + final = output + assert final is not None and final.finished + assert final.outputs and final.outputs[0].token_ids + return final + + try: + prompt = "Explain how a cache works. " * 40 + cold = await generate(prompt, "cold") + warm = await generate(prompt, "warm") + assert cold.num_cached_tokens == 0 + assert warm.num_cached_tokens > 0 + + assert await engine.reset_prefix_cache(reset_running_requests=True, reset_connector=True) + await engine.reset_encoder_cache() + await engine.reset_mm_cache() + reset = await generate(prompt, "reset") + assert reset.num_cached_tokens == 0 + assert reset.outputs[0].token_ids == cold.outputs[0].token_ids + + image_prompt = { + "prompt": "<|im_start|>user\n<|vision_start|><|IMAGE|><|vision_end|>Describe this image.<|im_end|>\n<|im_start|>assistant\n", + "multi_modal_data": {"image": Image.new("RGB", (56, 56), "red")}, + } + before = await generate(image_prompt, "image-before") + await engine.reset_mm_cache() + await engine.reset_encoder_cache() + assert await engine.reset_prefix_cache() + after = await generate(image_prompt, "image-after") + assert after.outputs[0].token_ids == before.outputs[0].token_ids + finally: + engine.shutdown() diff --git a/tests/entrypoints/test_async_omni_pause_sleep_routing.py b/tests/entrypoints/test_async_omni_pause_sleep_routing.py index 887cf8885..e4b229f90 100644 --- a/tests/entrypoints/test_async_omni_pause_sleep_routing.py +++ b/tests/entrypoints/test_async_omni_pause_sleep_routing.py @@ -14,7 +14,7 @@ from vllm.renderers import BaseRenderer from vllm.v1.engine.input_processor import InputProcessor from vllm_omni.diffusion.data import CuMemTag, OmniACK -from vllm_omni.entrypoints.async_omni import AsyncOmni +from vllm_omni.entrypoints.async_omni import CACHE_RESET_TIMEOUT_S, AsyncOmni pytestmark = [pytest.mark.core_model] @@ -31,7 +31,7 @@ def _make_omni(*, stage_types: list[str]) -> AsyncOmni: omni._level2_sleeping = False omni.event_resolver = SimpleNamespace(watch_task=lambda *a, **k: None, resolve=AsyncMock()) omni._final_output_handler = lambda: None - omni.reset_mm_cache = AsyncMock() + omni._clear_frontend_mm_cache = AsyncMock() stage_configs = [SimpleNamespace(stage_type=stage_type) for stage_type in stage_types] omni.engine = SimpleNamespace( @@ -46,7 +46,8 @@ def _make_omni(*, stage_types: list[str]) -> AsyncOmni: @pytest.mark.cpu def test_reset_mm_cache_routes_through_renderer(mocker): async def run() -> None: - omni = object.__new__(AsyncOmni) + omni = _make_omni(stage_types=["llm"]) + del omni._clear_frontend_mm_cache renderer = mocker.create_autospec(BaseRenderer, instance=True) input_processor = mocker.create_autospec(InputProcessor, instance=True) input_processor.renderer = renderer @@ -115,9 +116,10 @@ def test_pause_generation_still_rpcs_when_already_paused(): kwargs={"mode": "abort", "clear_cache": True}, stage_ids=[0], ) - omni.reset_prefix_cache.assert_awaited_once() - omni.reset_mm_cache.assert_awaited_once() - omni.reset_encoder_cache.assert_awaited_once() + omni.reset_prefix_cache.assert_not_awaited() + omni.reset_mm_cache.assert_not_awaited() + omni.reset_encoder_cache.assert_not_awaited() + omni._clear_frontend_mm_cache.assert_awaited_once_with() asyncio.run(run()) @@ -142,6 +144,84 @@ def test_resume_generation_resumes_ar_then_clears_frontend_pause(): asyncio.run(run()) +@pytest.mark.cpu +def test_reset_prefix_cache_forwards_to_ar_stages(): + async def run() -> None: + omni = _make_omni(stage_types=["llm", "diffusion"]) + + result = await omni.reset_prefix_cache(reset_running_requests=True, reset_connector=True) + + assert result is True + omni.collective_rpc.assert_awaited_once_with( + method="reset_prefix_cache", + args=(True, True), + kwargs=None, + stage_ids=[0], + timeout=CACHE_RESET_TIMEOUT_S, + ) + + asyncio.run(run()) + + +@pytest.mark.cpu +def test_reset_prefix_cache_diffusion_only_skips_rpc(): + async def run() -> None: + omni = _make_omni(stage_types=["diffusion"]) + + result = await omni.reset_prefix_cache() + + assert result is True + omni.collective_rpc.assert_not_awaited() + + asyncio.run(run()) + + +@pytest.mark.cpu +def test_reset_encoder_cache_forwards_to_ar_stages(): + async def run() -> None: + omni = _make_omni(stage_types=["llm", "diffusion"]) + + await omni.reset_encoder_cache() + + omni.collective_rpc.assert_awaited_once_with( + method="reset_encoder_cache", + args=(), + kwargs=None, + stage_ids=[0], + timeout=CACHE_RESET_TIMEOUT_S, + ) + + asyncio.run(run()) + + +@pytest.mark.cpu +def test_reset_mm_cache_clears_frontend_then_forwards_to_ar_stages(): + async def run() -> None: + omni = _make_omni(stage_types=["llm", "diffusion"]) + calls = [] + renderer = SimpleNamespace(clear_mm_cache_async=AsyncMock(side_effect=lambda: calls.append("frontend"))) + omni.input_processor = SimpleNamespace(renderer=renderer) + del omni._clear_frontend_mm_cache + + async def rpc(**kwargs): + calls.append("engine") + return [None] + + omni.collective_rpc.side_effect = rpc + await omni.reset_mm_cache() + + assert calls == ["frontend", "engine"] + omni.collective_rpc.assert_awaited_once_with( + method="reset_mm_cache", + args=(), + kwargs=None, + stage_ids=[0], + timeout=CACHE_RESET_TIMEOUT_S, + ) + + asyncio.run(run()) + + @pytest.mark.cpu def test_sleep_routes_ar_via_collective_rpc_and_diffusion_to_worker_rpc(): async def run() -> None: @@ -539,8 +619,9 @@ def test_mixed_engine_keep_sends_engine_core_and_diffusion_pauses(): "kwargs": {"mode": "keep"}, "stage_ids": [1], } - # The frontend cache reset still runs after the backend pauses. - omni.reset_prefix_cache.assert_awaited_once() + # Only P0 needs clearing here: pause_scheduler already resets P1. + omni.reset_prefix_cache.assert_not_awaited() + omni._clear_frontend_mm_cache.assert_awaited_once_with() await omni.resume_generation() assert _rpc_methods(omni)[2:] == [("resume_scheduler", [0]), ("resume_scheduler", [1])] @@ -579,9 +660,11 @@ def test_keep_then_sleep_wake_resume_keeps_admission_paused_until_resume(): 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] + 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) @@ -683,3 +766,63 @@ def test_two_ar_stages_partial_resume_keeps_admission_paused(): assert await omni.is_paused() is False asyncio.run(run()) + + +@pytest.mark.cpu +@pytest.mark.parametrize("method", ["reset_prefix_cache", "reset_encoder_cache", "reset_mm_cache"]) +@pytest.mark.parametrize("stage_ids, targets", [(None, [0, 2]), ([2], [2]), ([1], []), ([], [])]) +def test_cache_resets_respect_stage_selection(method, stage_ids, targets): + async def run(): + omni = _make_omni(stage_types=["llm", "diffusion", "llm"]) + await getattr(omni, method)(stage_ids=stage_ids, timeout=2.0) + if targets: + assert omni.collective_rpc.await_args.kwargs["stage_ids"] == targets + assert omni.collective_rpc.await_args.kwargs["timeout"] == 2.0 + else: + omni.collective_rpc.assert_not_awaited() + assert omni._clear_frontend_mm_cache.await_count == int(method == "reset_mm_cache" and 0 in targets) + + asyncio.run(run()) + + +@pytest.mark.cpu +@pytest.mark.parametrize("method", ["reset_prefix_cache", "reset_encoder_cache", "reset_mm_cache"]) +@pytest.mark.parametrize("failure", [{"todo": "unsupported"}, {"supported": False}, {"error": "reset failed"}]) +def test_cache_resets_reject_failure_from_any_replica(method, failure): + async def run(): + omni = _make_omni(stage_types=["llm"]) + # Exercise the actual collective_rpc too: a single stage can have + # multiple replicas, so result indexes are not stage IDs. + del omni.collective_rpc + omni.engine.collective_rpc_async.return_value = [True, failure] + with pytest.raises(RuntimeError, match=method): + await getattr(omni, method)() + + asyncio.run(run()) + + +@pytest.mark.cpu +def test_prefix_reset_returns_false_if_any_replica_refuses(): + async def run(): + omni = _make_omni(stage_types=["llm"]) + omni.collective_rpc.return_value = [True, [True, False]] + assert await omni.reset_prefix_cache() is False + + asyncio.run(run()) + + +@pytest.mark.cpu +@pytest.mark.parametrize("operation", ["pause_generation", "sleep"]) +@pytest.mark.parametrize("stage_ids", [[0], [2]]) +def test_pause_and_sleep_do_not_reset_unselected_stages_or_reset_twice(operation, stage_ids): + async def run(): + omni = _make_omni(stage_types=["llm", "diffusion", "llm"]) + kwargs = {"clear_cache": True} if operation == "pause_generation" else {} + await getattr(omni, operation)(stage_ids=stage_ids, **kwargs) + assert omni.collective_rpc.await_count == 1 + rpc = omni.collective_rpc.await_args.kwargs + assert rpc["method"] == ("pause_scheduler" if operation == "pause_generation" else "sleep") + assert rpc["stage_ids"] == stage_ids + assert omni._clear_frontend_mm_cache.await_count == int(0 in stage_ids) + + asyncio.run(run()) diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index deec36199..c5a242732 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -704,8 +704,21 @@ class OrchestratorBase: results: list[Any] = [] stage_ids: list[int] = [] for pool in target_pools: + clear_sender = pool.stage_type != "diffusion" and ( + method in ("reset_mm_cache", "sleep") + or (method == "pause_scheduler" and kwargs.get("clear_cache", args[1] if len(args) > 1 else True)) + ) for replica_id in pool.live_replica_ids(): try: + if clear_sender: + processor = self._stage_input_processors.get(pool.stage_id) + if processor is not None: + clear = processor.renderer.clear_mm_cache_async() + if timeout is None: + await clear + else: + await asyncio.wait_for(clear, timeout=timeout) + clear_sender = False stage_result = await pool.collective_rpc( replica_id=replica_id, method=method, @@ -714,13 +727,14 @@ class OrchestratorBase: kwargs=kwargs, ) except Exception as exc: - if method not in ("pause_scheduler", "resume_scheduler"): + if ( + method not in ("pause_scheduler", "resume_scheduler") + and method not in StagePool._CACHE_RESET_METHODS + ): 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. + # Administrative failures must reach the caller without + # terminating unrelated stages. _engine_core_rpc raises + # on this result; remaining replicas still receive the RPC. 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/engine/stage_pool.py b/vllm_omni/engine/stage_pool.py index cb79501cb..222cd3b5d 100644 --- a/vllm_omni/engine/stage_pool.py +++ b/vllm_omni/engine/stage_pool.py @@ -101,12 +101,16 @@ class StagePool: DISPATCH_RETRY_INTERVAL_S: float = 0.1 # Only these EngineCore helpers may skip collective_rpc_async. A generic # ``{method}_async`` on AsyncMPClient must not silently drop timeout. + _CACHE_RESET_METHODS = frozenset({"reset_prefix_cache", "reset_encoder_cache", "reset_mm_cache"}) _ENGINE_CORE_CONTROL_ASYNC_METHODS = frozenset( { "pause_scheduler", "resume_scheduler", "sleep", "wake_up", + "reset_prefix_cache", + "reset_encoder_cache", + "reset_mm_cache", } ) @@ -1330,6 +1334,8 @@ class StagePool: if timeout is not None: return await asyncio.wait_for(result, timeout=timeout) return await result + if method in self._CACHE_RESET_METHODS: + return {"supported": False, "error": f"EngineCore helper {method}_async is unavailable"} return await client.collective_rpc_async( method=method, @@ -1344,7 +1350,7 @@ class StagePool: replica_id, method, ) - if method in self._ENGINE_CORE_CONTROL_ASYNC_METHODS: + if method in self._ENGINE_CORE_CONTROL_ASYNC_METHODS and method not in self._CACHE_RESET_METHODS: raise if isinstance(exc, TimeoutError): error = f"{type(exc).__name__}: {method} timed out after {timeout}s" diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index 438cf08c8..1c27def99 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/vllm_omni/entrypoints/async_omni.py @@ -46,6 +46,8 @@ if TYPE_CHECKING: logger = init_logger(__name__) +CACHE_RESET_TIMEOUT_S = 60.0 + class AsyncOmni(AsyncOmniBase, EngineClient): """Asynchronous unified entry point for multi-stage pipelines using AsyncOmniEngine. @@ -535,17 +537,15 @@ class AsyncOmni(AsyncOmniBase, EngineClient): stage_ids=stage_ids, ) - unsupported_stage_ids: list[int] = [] - effective_stage_ids = stage_ids or list(range(len(results))) - for index, result in enumerate(results): - if isinstance(result, dict) and result.get("todo"): - unsupported_stage_ids.append(effective_stage_ids[index]) - - if unsupported_stage_ids: + unsupported_results = [ + index for index, result in enumerate(results) if isinstance(result, dict) and result.get("todo") + ] + if unsupported_results: logger.warning( - "[AsyncOmni] collective_rpc(%s) has TODO support on stage(s): %s", + "[AsyncOmni] collective_rpc(%s) has TODO support in replica result(s) %s (requested stages: %s)", method, - unsupported_stage_ids, + unsupported_results, + stage_ids, ) return results @@ -557,23 +557,28 @@ class AsyncOmni(AsyncOmniBase, EngineClient): stage_ids: list[int], args: tuple[Any, ...] = (), kwargs: dict[str, Any] | None = None, + timeout: float | None = None, ) -> list[Any]: - """Call an engine control helper via collective_rpc (orchestrator loop). - - StagePool resolves ``{method}_async`` on the AR client when present - (vLLM AsyncMPClient convention); diffusion stages answer the same - method names inside DiffusionEngine. Raises if any replica reports - failure. - """ + """Call EngineCore helpers and reject failures from any replica.""" + timeout_kwargs = {"timeout": timeout} if timeout is not None else {} results = await self.collective_rpc( method=method, args=args, kwargs=kwargs, stage_ids=stage_ids, + **timeout_kwargs, ) + + def check_result(result: Any) -> None: + if isinstance(result, (list, tuple)): + for item in result: + check_result(item) + elif isinstance(result, dict): + if result.get("error") or result.get("todo") or result.get("supported") is False: + raise RuntimeError(f"{method} failed: {result}") + for result in results: - if isinstance(result, dict) and result.get("error"): - raise RuntimeError(f"{method} failed: {result['error']}") + check_result(result) return results @staticmethod @@ -768,13 +773,8 @@ class AsyncOmni(AsyncOmniBase, EngineClient): # Frontend / sender-side cache clear (P0). EngineCore.pause_scheduler # already clears AR-side caches when clear_cache=True. - if clear_cache: - await self.reset_prefix_cache( - reset_running_requests=not wait_for_inflight_requests, - reset_connector=True, - ) - await self.reset_mm_cache() - await self.reset_encoder_cache() + if clear_cache and 0 in ar_stage_ids: + await self._clear_frontend_mm_cache() async def resume_generation(self, stage_ids: list[int] | None = None) -> None: """Resume generation after :meth:`pause_generation`.""" @@ -835,34 +835,63 @@ class AsyncOmni(AsyncOmniBase, EngineClient): """ return await self.collective_rpc(method="profile", args=(False, None), stage_ids=stages) - async def reset_mm_cache(self) -> None: - """Reset the frontend (P0) multimodal processor cache. - - ``EngineCore.sleep(level>=1)`` already clears the P1 receiver cache. - Clearing P0 avoids hash-only follow-up requests after that reset. - """ + async def _clear_frontend_mm_cache(self) -> None: + """Clear P0 through the renderer's serialized multimodal executor.""" renderer = self.renderer if renderer is not None: await renderer.clear_mm_cache_async() - async def reset_encoder_cache(self) -> None: - """Reset the encoder cache for all stages. + async def reset_mm_cache( + self, *, stage_ids: list[int] | None = None, timeout: float = CACHE_RESET_TIMEOUT_S + ) -> None: + """Reset sender and receiver MM caches on selected AR stages. - TODO: Forward to Orchestrator process via message. + By default all AR stages are selected; diffusion stages are skipped. + Stage 0's sender is the frontend renderer. Downstream senders are + cleared by the orchestrator before resetting their engine cores. + Call while generation is paused to avoid racing new inputs. """ - logger.warning("[AsyncOmni] reset_encoder_cache not yet supported with Orchestrator process") + ar_stage_ids, _ = self._split_stage_ids_by_type(stage_ids) + if 0 in ar_stage_ids: + await asyncio.wait_for(self._clear_frontend_mm_cache(), timeout=timeout) + if ar_stage_ids: + await self._engine_core_rpc("reset_mm_cache", stage_ids=ar_stage_ids, timeout=timeout) + + async def reset_encoder_cache( + self, *, stage_ids: list[int] | None = None, timeout: float = CACHE_RESET_TIMEOUT_S + ) -> None: + """Reset encoder caches on selected AR stages (all by default). + + Diffusion stages are skipped. RPC failures are raised to the caller. + """ + ar_stage_ids, _ = self._split_stage_ids_by_type(stage_ids) + if ar_stage_ids: + await self._engine_core_rpc("reset_encoder_cache", stage_ids=ar_stage_ids, timeout=timeout) async def reset_prefix_cache( self, reset_running_requests: bool = False, reset_connector: bool = False, + *, + stage_ids: list[int] | None = None, + timeout: float = CACHE_RESET_TIMEOUT_S, ) -> bool: - """Reset the prefix cache for all stages. + """Reset prefix caches on selected AR stages (all by default). - TODO: Forward to Orchestrator process via message. + Diffusion stages are skipped. Return False if a stage cannot reset + its cache; unsupported operations and RPC failures raise instead of + silently retaining KV computed under previous model weights. """ - logger.warning("[AsyncOmni] reset_prefix_cache not yet supported with Orchestrator process") - return True + ar_stage_ids, _ = self._split_stage_ids_by_type(stage_ids) + if not ar_stage_ids: + return True + results = await self._engine_core_rpc( + "reset_prefix_cache", + stage_ids=ar_stage_ids, + args=(reset_running_requests, reset_connector), + timeout=timeout, + ) + return all(self._coerce_stage_bool(result) for result in results) async def sleep( self, stage_ids: list[int] | None = None, level: int = 2, mode: PauseMode = "abort" @@ -895,11 +924,12 @@ class AsyncOmni(AsyncOmniBase, EngineClient): self._paused = True await self._pause_cond.wait_for(lambda: getattr(self, "_admitting", 0) == 0) - # P0 sender cache must drop hashes before EngineCore.sleep clears P1. - await self.reset_mm_cache() + ar_stage_ids, diffusion_stage_ids = self._split_stage_ids_by_type(stage_ids) + # EngineCore.sleep resets receiver caches itself; only clear P0 here. + if 0 in ar_stage_ids: + await self._clear_frontend_mm_cache() self._final_output_handler() - ar_stage_ids, diffusion_stage_ids = self._split_stage_ids_by_type(stage_ids) final_acks: list[OmniACK] = [] if ar_stage_ids: self._hold_admission_until_resume = True