1117 lines
38 KiB
Python
1117 lines
38 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
"""Unit tests for AR EngineCore vs diffusion worker pause/sleep routing."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from types import SimpleNamespace
|
|
from unittest.mock import AsyncMock, call
|
|
|
|
import pytest
|
|
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 CACHE_RESET_TIMEOUT_S, AsyncOmni
|
|
|
|
pytestmark = [pytest.mark.core_model]
|
|
|
|
|
|
def _make_omni(*, stage_types: list[str]) -> AsyncOmni:
|
|
omni = object.__new__(AsyncOmni)
|
|
omni._pause_cond = asyncio.Condition()
|
|
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
|
|
omni.event_resolver = SimpleNamespace(watch_task=lambda *a, **k: None, resolve=AsyncMock())
|
|
omni._final_output_handler = lambda: None
|
|
omni._clear_frontend_mm_cache = AsyncMock()
|
|
|
|
stage_configs = [SimpleNamespace(stage_type=stage_type) for stage_type in stage_types]
|
|
omni.engine = SimpleNamespace(
|
|
stage_configs=stage_configs,
|
|
stage_vllm_configs=[None] * len(stage_types),
|
|
collective_rpc_async=AsyncMock(return_value=[True]),
|
|
)
|
|
omni.collective_rpc = AsyncMock(return_value=[True])
|
|
return omni
|
|
|
|
|
|
@pytest.mark.cpu
|
|
def test_reset_mm_cache_routes_through_renderer(mocker):
|
|
async def run() -> None:
|
|
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
|
|
omni.input_processor = input_processor
|
|
|
|
await omni.reset_mm_cache()
|
|
|
|
renderer.clear_mm_cache_async.assert_awaited_once_with()
|
|
|
|
asyncio.run(run())
|
|
|
|
|
|
@pytest.mark.cpu
|
|
def test_split_stage_ids_by_type():
|
|
omni = _make_omni(stage_types=["llm", "diffusion", "llm"])
|
|
ar_ids, diff_ids = omni._split_stage_ids_by_type()
|
|
assert ar_ids == [0, 2]
|
|
assert diff_ids == [1]
|
|
|
|
|
|
@pytest.mark.cpu
|
|
def test_split_stage_ids_by_type_rejects_out_of_range():
|
|
omni = _make_omni(stage_types=["llm", "diffusion"])
|
|
with pytest.raises(ValueError, match=r"Invalid stage_ids \[2\].*0\.\.1"):
|
|
omni._split_stage_ids_by_type([0, 2])
|
|
with pytest.raises(ValueError, match=r"Invalid stage_ids \[-1\]"):
|
|
omni._split_stage_ids_by_type([-1])
|
|
|
|
|
|
@pytest.mark.cpu
|
|
def test_pause_generation_routes_ar_via_collective_rpc():
|
|
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="abort", clear_cache=True)
|
|
|
|
assert omni._paused is True
|
|
assert omni._hold_admission_until_resume is True
|
|
omni.collective_rpc.assert_awaited_once_with(
|
|
method="pause_scheduler",
|
|
args=(),
|
|
kwargs={"mode": "abort", "clear_cache": True},
|
|
stage_ids=[0],
|
|
)
|
|
|
|
asyncio.run(run())
|
|
|
|
|
|
@pytest.mark.cpu
|
|
def test_pause_generation_still_rpcs_when_already_paused():
|
|
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()
|
|
omni._paused = True
|
|
|
|
await omni.pause_generation(mode="abort", clear_cache=True, stage_ids=[0])
|
|
|
|
omni.collective_rpc.assert_awaited_once_with(
|
|
method="pause_scheduler",
|
|
args=(),
|
|
kwargs={"mode": "abort", "clear_cache": True},
|
|
stage_ids=[0],
|
|
)
|
|
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())
|
|
|
|
|
|
@pytest.mark.cpu
|
|
def test_resume_generation_resumes_ar_then_clears_frontend_pause():
|
|
async def run() -> None:
|
|
omni = _make_omni(stage_types=["llm", "diffusion"])
|
|
omni._paused = True
|
|
|
|
await omni.resume_generation()
|
|
|
|
omni.collective_rpc.assert_awaited_once_with(
|
|
method="resume_scheduler",
|
|
args=(),
|
|
kwargs=None,
|
|
stage_ids=[0],
|
|
)
|
|
assert omni._paused is False
|
|
assert omni._hold_admission_until_resume is False
|
|
|
|
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_release_kv_memory_requires_pause_and_wakes_only_ar_kv_allocations():
|
|
async def run() -> None:
|
|
omni = _make_omni(stage_types=["llm", "diffusion", "llm"])
|
|
with pytest.raises(RuntimeError, match="completed pause"):
|
|
await omni.release_kv_cache_memory()
|
|
omni.collective_rpc.assert_not_awaited()
|
|
|
|
omni._paused = True
|
|
await omni.release_kv_cache_memory()
|
|
omni._clear_frontend_mm_cache.assert_awaited_once_with()
|
|
omni.collective_rpc.assert_awaited_once_with(
|
|
method="release_kv_cache_memory",
|
|
args=(),
|
|
kwargs=None,
|
|
stage_ids=[0, 2],
|
|
timeout=CACHE_RESET_TIMEOUT_S,
|
|
)
|
|
assert omni._stage_sleeping_tags == {0: {"kv_cache"}, 2: {"kv_cache"}}
|
|
assert not omni._level2_sleeping
|
|
|
|
omni.collective_rpc.reset_mock()
|
|
await omni.wake_up(stage_ids=[0, 2], tags=["kv_cache"])
|
|
omni.collective_rpc.assert_awaited_once_with(
|
|
method="wake_up", args=(), kwargs={"tags": ["kv_cache"]}, stage_ids=[0, 2]
|
|
)
|
|
assert not omni._sleeping_tags
|
|
assert omni._paused
|
|
|
|
asyncio.run(run())
|
|
|
|
|
|
@pytest.mark.cpu
|
|
def test_release_kv_memory_skips_diffusion_and_reports_engine_failures():
|
|
async def run() -> None:
|
|
diffusion = _make_omni(stage_types=["diffusion"])
|
|
await diffusion.release_kv_cache_memory()
|
|
diffusion.collective_rpc.assert_not_awaited()
|
|
assert not diffusion._paused
|
|
|
|
omni = _make_omni(stage_types=["llm"])
|
|
omni._paused = True
|
|
omni.collective_rpc.return_value = [{"supported": False, "error": "requires all executor memory resident"}]
|
|
with pytest.raises(RuntimeError, match="all executor memory resident"):
|
|
await omni.release_kv_cache_memory()
|
|
assert not omni._sleeping_tags
|
|
|
|
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:
|
|
omni = _make_omni(stage_types=["llm", "diffusion"])
|
|
diffusion_ack = OmniACK(task_id="diff", status="SUCCESS", stage_id=1, rank=0)
|
|
omni._sleep_diffusion = AsyncMock(return_value=[diffusion_ack])
|
|
|
|
acks = await omni.sleep(stage_ids=[0, 1], level=1, mode="abort")
|
|
|
|
omni.collective_rpc.assert_awaited_once_with(
|
|
method="sleep",
|
|
args=(1, "abort"),
|
|
kwargs=None,
|
|
stage_ids=[0],
|
|
)
|
|
omni._sleep_diffusion.assert_awaited_once_with([1], 1)
|
|
assert {ack.stage_id for ack in acks} == {0, 1}
|
|
assert any(ack.metadata.get("path") == "engine_core" for ack in acks if ack.stage_id == 0)
|
|
assert CuMemTag.WEIGHTS.value in omni._sleeping_tags
|
|
assert CuMemTag.KV_CACHE.value in omni._sleeping_tags
|
|
# Sleep gates frontend admission for the trainer resume contract.
|
|
assert omni._paused is True
|
|
assert omni._hold_admission_until_resume is True
|
|
|
|
asyncio.run(run())
|
|
|
|
|
|
@pytest.mark.cpu
|
|
def test_sleep_blocks_admission_before_engine_core_rpc():
|
|
"""Problem 3: _paused must be set before awaiting sleep RPC."""
|
|
|
|
async def run() -> None:
|
|
omni = _make_omni(stage_types=["llm"])
|
|
paused_at_rpc: list[bool] = []
|
|
|
|
async def rpc_side_effect(**kwargs):
|
|
if kwargs.get("method") == "sleep":
|
|
paused_at_rpc.append(omni._paused)
|
|
return [True]
|
|
|
|
omni.collective_rpc = AsyncMock(side_effect=rpc_side_effect)
|
|
|
|
assert omni._paused is False
|
|
await omni.sleep(level=1, mode="abort")
|
|
|
|
assert paused_at_rpc == [True]
|
|
assert omni._paused is True
|
|
assert omni._hold_admission_until_resume is True
|
|
# Sleep must not also call pause_scheduler (EngineCore.sleep pauses).
|
|
assert all(c.kwargs.get("method") != "pause_scheduler" for c in omni.collective_rpc.await_args_list)
|
|
|
|
asyncio.run(run())
|
|
|
|
|
|
@pytest.mark.cpu
|
|
def test_sleep_waits_for_in_flight_generate_admission():
|
|
"""Sleep must not offload while generate() is still in add_request."""
|
|
|
|
async def run() -> None:
|
|
omni = _make_omni(stage_types=["llm"])
|
|
omni._admitting = 1
|
|
rpc_started = asyncio.Event()
|
|
|
|
async def rpc_side_effect(**kwargs):
|
|
if kwargs.get("method") == "sleep":
|
|
rpc_started.set()
|
|
return [True]
|
|
|
|
omni.collective_rpc = AsyncMock(side_effect=rpc_side_effect)
|
|
sleep_task = asyncio.create_task(omni.sleep(level=1, mode="abort"))
|
|
await asyncio.sleep(0.05)
|
|
assert not sleep_task.done()
|
|
assert not rpc_started.is_set()
|
|
await omni._release_generate_admission()
|
|
await asyncio.wait_for(sleep_task, timeout=1)
|
|
assert rpc_started.is_set()
|
|
assert omni._admitting == 0
|
|
|
|
asyncio.run(run())
|
|
|
|
|
|
@pytest.mark.cpu
|
|
def test_wake_up_routes_ar_via_collective_rpc_and_diffusion_to_worker_rpc():
|
|
async def run() -> None:
|
|
omni = _make_omni(stage_types=["llm", "diffusion"])
|
|
omni._paused = True
|
|
omni._hold_admission_until_resume = True
|
|
omni._sleeping_tags = {CuMemTag.WEIGHTS.value, CuMemTag.KV_CACHE.value}
|
|
diffusion_ack = OmniACK(task_id="diff", status="SUCCESS", stage_id=1, rank=0)
|
|
omni._wake_diffusion = AsyncMock(return_value=[diffusion_ack])
|
|
|
|
acks = await omni.wake_up(stage_ids=[0, 1])
|
|
|
|
omni.collective_rpc.assert_awaited_once()
|
|
wake_kwargs = omni.collective_rpc.await_args.kwargs
|
|
assert wake_kwargs["method"] == "wake_up"
|
|
assert wake_kwargs["stage_ids"] == [0]
|
|
assert set(wake_kwargs["kwargs"]["tags"]) == {
|
|
CuMemTag.WEIGHTS.value,
|
|
CuMemTag.KV_CACHE.value,
|
|
}
|
|
omni._wake_diffusion.assert_awaited_once()
|
|
assert {ack.stage_id for ack in acks} == {0, 1}
|
|
assert not omni._sleeping_tags
|
|
# Mixed/AR wake restores memory but does not resume frontend admission.
|
|
assert omni._paused is True
|
|
assert omni._hold_admission_until_resume is True
|
|
|
|
asyncio.run(run())
|
|
|
|
|
|
@pytest.mark.cpu
|
|
def test_wake_up_does_not_resume_frontend_admission():
|
|
async def run() -> None:
|
|
omni = _make_omni(stage_types=["llm"])
|
|
await omni.sleep(level=1, mode="abort")
|
|
assert omni._paused is True
|
|
assert omni._hold_admission_until_resume is True
|
|
|
|
await omni.wake_up()
|
|
|
|
assert omni._paused is True
|
|
assert omni._hold_admission_until_resume is True
|
|
assert not omni._sleeping_tags
|
|
# Explicit resume is required after AR sleep/wake.
|
|
await omni.resume_generation()
|
|
assert omni._paused is False
|
|
assert omni._hold_admission_until_resume is False
|
|
|
|
asyncio.run(run())
|
|
|
|
|
|
@pytest.mark.cpu
|
|
def test_sleep_level1_wake_without_tags_clears_all_sleeping_tags():
|
|
"""sleep(level=1) tracks WEIGHTS+KV; untagged wake_up must clear both."""
|
|
|
|
async def run() -> None:
|
|
omni = _make_omni(stage_types=["llm"])
|
|
|
|
await omni.sleep(level=1, mode="abort")
|
|
assert CuMemTag.WEIGHTS.value in omni._sleeping_tags
|
|
assert CuMemTag.KV_CACHE.value in omni._sleeping_tags
|
|
|
|
await omni.wake_up()
|
|
assert not omni._sleeping_tags
|
|
assert omni._paused is True
|
|
assert omni._hold_admission_until_resume is True
|
|
|
|
asyncio.run(run())
|
|
|
|
|
|
@pytest.mark.cpu
|
|
def test_sleep_diffusion_only_skips_engine_core_collective_rpc():
|
|
async def run() -> None:
|
|
omni = _make_omni(stage_types=["diffusion"])
|
|
omni._sleep_diffusion = AsyncMock(return_value=[OmniACK(task_id="d", status="SUCCESS", stage_id=0, rank=0)])
|
|
|
|
await omni.sleep(level=1)
|
|
|
|
# AR EngineCore sleep path must not run for diffusion-only.
|
|
assert call(method="sleep", args=(1, "abort"), kwargs=None, stage_ids=[0]) not in (
|
|
omni.collective_rpc.await_args_list
|
|
)
|
|
omni._sleep_diffusion.assert_awaited_once_with([0], 1)
|
|
assert omni._paused is True
|
|
assert omni._hold_admission_until_resume is False
|
|
|
|
asyncio.run(run())
|
|
|
|
|
|
@pytest.mark.cpu
|
|
def test_streaming_wait_for_first_chunk_does_not_block_sleep():
|
|
"""Waiting for the next client chunk must not hold an admission slot."""
|
|
|
|
async def run() -> None:
|
|
omni = _make_omni(stage_types=["llm"])
|
|
rpc_started = asyncio.Event()
|
|
first_chunk = asyncio.Event()
|
|
|
|
async def rpc_side_effect(**kwargs):
|
|
if kwargs.get("method") == "sleep":
|
|
rpc_started.set()
|
|
return [True]
|
|
|
|
omni.collective_rpc = AsyncMock(side_effect=rpc_side_effect)
|
|
|
|
async def wait_then_submit() -> None:
|
|
await first_chunk.wait()
|
|
await omni._submit_with_admission(asyncio.sleep(0))
|
|
|
|
wait_task = asyncio.create_task(wait_then_submit())
|
|
sleep_task = asyncio.create_task(omni.sleep(level=1, mode="abort"))
|
|
await asyncio.sleep(0.05)
|
|
assert not wait_task.done()
|
|
assert omni._admitting == 0
|
|
await asyncio.wait_for(sleep_task, timeout=1)
|
|
assert rpc_started.is_set()
|
|
|
|
wait_task.cancel()
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await wait_task
|
|
assert omni._admitting == 0
|
|
|
|
asyncio.run(run())
|
|
|
|
|
|
@pytest.mark.cpu
|
|
def test_streaming_in_flight_add_blocks_sleep():
|
|
"""Sleep must wait while a streaming ADD/update holds an admission slot."""
|
|
|
|
async def run() -> None:
|
|
omni = _make_omni(stage_types=["llm"])
|
|
add_started = asyncio.Event()
|
|
add_release = asyncio.Event()
|
|
rpc_started = asyncio.Event()
|
|
|
|
async def rpc_side_effect(**kwargs):
|
|
if kwargs.get("method") == "sleep":
|
|
rpc_started.set()
|
|
return [True]
|
|
|
|
omni.collective_rpc = AsyncMock(side_effect=rpc_side_effect)
|
|
|
|
async def slow_add() -> None:
|
|
add_started.set()
|
|
await add_release.wait()
|
|
|
|
add_task = asyncio.create_task(omni._submit_with_admission(slow_add()))
|
|
await add_started.wait()
|
|
sleep_task = asyncio.create_task(omni.sleep(level=1, mode="abort"))
|
|
await asyncio.sleep(0.05)
|
|
assert not sleep_task.done()
|
|
assert not rpc_started.is_set()
|
|
assert omni._admitting == 1
|
|
|
|
add_release.set()
|
|
await asyncio.wait_for(add_task, timeout=1)
|
|
await asyncio.wait_for(sleep_task, timeout=1)
|
|
assert rpc_started.is_set()
|
|
assert omni._admitting == 0
|
|
|
|
asyncio.run(run())
|
|
|
|
|
|
@pytest.mark.cpu
|
|
def test_wake_up_restores_admission_for_diffusion_only():
|
|
"""Pure diffusion must keep sleep → wake → generate (no resume)."""
|
|
|
|
async def run() -> None:
|
|
omni = _make_omni(stage_types=["diffusion"])
|
|
omni._sleep_diffusion = AsyncMock(return_value=[OmniACK(task_id="d", status="SUCCESS", stage_id=0, rank=0)])
|
|
omni._wake_diffusion = AsyncMock(return_value=[OmniACK(task_id="w", status="SUCCESS", stage_id=0, rank=0)])
|
|
|
|
await omni.sleep(level=1)
|
|
assert omni._paused is True
|
|
|
|
async def wait_for_admission() -> None:
|
|
async with omni._pause_cond:
|
|
await omni._pause_cond.wait_for(lambda: not omni._paused)
|
|
|
|
waiter = asyncio.create_task(wait_for_admission())
|
|
await omni.wake_up()
|
|
await asyncio.wait_for(waiter, timeout=1.0)
|
|
assert omni._paused is False
|
|
assert omni._hold_admission_until_resume is False
|
|
assert not omni._sleeping_tags
|
|
|
|
asyncio.run(run())
|
|
|
|
|
|
@pytest.mark.cpu
|
|
def test_pause_then_sleep_wake_keeps_admission_paused_for_diffusion():
|
|
"""Explicit pause_generation still requires resume after diffusion wake."""
|
|
|
|
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()
|
|
omni._sleep_diffusion = AsyncMock(return_value=[OmniACK(task_id="d", status="SUCCESS", stage_id=0, rank=0)])
|
|
omni._wake_diffusion = AsyncMock(return_value=[OmniACK(task_id="w", status="SUCCESS", stage_id=0, rank=0)])
|
|
|
|
await omni.pause_generation()
|
|
await omni.sleep(level=1)
|
|
await omni.wake_up()
|
|
|
|
assert omni._paused is True
|
|
assert omni._hold_admission_until_resume is True
|
|
await omni.resume_generation()
|
|
assert omni._paused is False
|
|
|
|
asyncio.run(run())
|
|
|
|
|
|
@pytest.mark.cpu
|
|
def test_partial_wake_does_not_skip_remaining_sleeping_stage():
|
|
"""sleep(stage_ids=[0]) then wake([0]) must not skip a later wake([1])."""
|
|
|
|
async def run() -> None:
|
|
omni = _make_omni(stage_types=["llm", "llm"])
|
|
await omni.sleep(stage_ids=[0, 1], level=1, mode="abort")
|
|
assert omni._stage_sleeping_tags.keys() == {0, 1}
|
|
|
|
await omni.wake_up(stage_ids=[0])
|
|
assert 0 not in omni._stage_sleeping_tags
|
|
assert 1 in omni._stage_sleeping_tags
|
|
assert CuMemTag.WEIGHTS.value in omni._sleeping_tags
|
|
|
|
omni.collective_rpc.reset_mock()
|
|
acks = await omni.wake_up(stage_ids=[1])
|
|
assert acks
|
|
omni.collective_rpc.assert_awaited_once()
|
|
wake_kwargs = omni.collective_rpc.await_args.kwargs
|
|
assert wake_kwargs["method"] == "wake_up"
|
|
assert wake_kwargs["stage_ids"] == [1]
|
|
assert not omni._sleeping_tags
|
|
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],
|
|
}
|
|
# 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])]
|
|
|
|
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())
|
|
|
|
|
|
@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())
|
|
|
|
|
|
_SLEEP_TAGS = {CuMemTag.WEIGHTS.value, CuMemTag.KV_CACHE.value}
|
|
|
|
# In-process stages return OmniACK, subprocess stages return its dict form,
|
|
# and StagePool returns an error dict when the RPC itself failed.
|
|
_FAILED_DIFFUSION_RESULTS = [
|
|
pytest.param(OmniACK(task_id="t", status="ERROR", error_msg="out of memory"), id="ack"),
|
|
pytest.param({"task_id": "t", "status": "ERROR", "error_msg": "out of memory"}, id="ack-dict"),
|
|
pytest.param({"supported": False, "error": "out of memory"}, id="rpc-error"),
|
|
]
|
|
|
|
|
|
@pytest.mark.cpu
|
|
@pytest.mark.parametrize("result", _FAILED_DIFFUSION_RESULTS)
|
|
def test_sleep_raises_when_diffusion_stage_fails(result):
|
|
async def run() -> None:
|
|
omni = _make_omni(stage_types=["diffusion"])
|
|
omni.collective_rpc = AsyncMock(return_value=[result])
|
|
|
|
with pytest.raises(RuntimeError, match="handle_sleep_task failed: out of memory"):
|
|
await omni.sleep(level=1)
|
|
|
|
# The stage may be partly asleep, so it stays on record for wake_up.
|
|
assert omni._sleeping_tags == _SLEEP_TAGS
|
|
|
|
asyncio.run(run())
|
|
|
|
|
|
@pytest.mark.cpu
|
|
@pytest.mark.parametrize("result", _FAILED_DIFFUSION_RESULTS)
|
|
def test_wake_up_raises_when_diffusion_stage_fails(result):
|
|
async def run() -> None:
|
|
omni = _make_omni(stage_types=["diffusion"])
|
|
await omni.sleep(level=1)
|
|
omni.collective_rpc = AsyncMock(return_value=[result])
|
|
|
|
with pytest.raises(RuntimeError, match="handle_wake_task failed: out of memory"):
|
|
await omni.wake_up()
|
|
|
|
assert omni._sleeping_tags == _SLEEP_TAGS
|
|
assert omni._paused is True
|
|
|
|
asyncio.run(run())
|
|
|
|
|
|
@pytest.mark.cpu
|
|
def test_sleep_records_all_stages_when_diffusion_stage_fails():
|
|
async def run() -> None:
|
|
omni = _make_omni(stage_types=["llm", "diffusion"])
|
|
omni._sleep_diffusion = AsyncMock(side_effect=RuntimeError("handle_sleep_task failed"))
|
|
|
|
with pytest.raises(RuntimeError):
|
|
await omni.sleep(level=1)
|
|
|
|
assert omni._stage_sleeping_tags == {0: _SLEEP_TAGS, 1: _SLEEP_TAGS}
|
|
|
|
asyncio.run(run())
|
|
|
|
|
|
@pytest.mark.cpu
|
|
def test_wake_up_clears_ar_stage_when_diffusion_stage_fails():
|
|
async def run() -> None:
|
|
omni = _make_omni(stage_types=["llm", "diffusion"])
|
|
await omni.sleep(level=1)
|
|
omni._wake_diffusion = AsyncMock(side_effect=RuntimeError("handle_wake_task failed"))
|
|
|
|
with pytest.raises(RuntimeError):
|
|
await omni.wake_up()
|
|
|
|
assert omni._stage_sleeping_tags == {1: _SLEEP_TAGS}
|
|
|
|
asyncio.run(run())
|
|
|
|
|
|
@pytest.mark.cpu
|
|
def test_sleep_level2_is_recorded_when_diffusion_stage_fails():
|
|
async def run() -> None:
|
|
omni = _make_omni(stage_types=["llm", "diffusion"])
|
|
omni._sleep_diffusion = AsyncMock(side_effect=RuntimeError("handle_sleep_task failed"))
|
|
|
|
with pytest.raises(RuntimeError):
|
|
await omni.sleep(level=2)
|
|
|
|
with pytest.raises(NotImplementedError):
|
|
await omni.wake_up(stage_ids=[0])
|
|
|
|
asyncio.run(run())
|
|
|
|
|
|
@pytest.mark.cpu
|
|
def test_wake_up_reaches_diffusion_stage_after_failed_sleep():
|
|
async def run() -> None:
|
|
omni = _make_omni(stage_types=["diffusion"])
|
|
omni.collective_rpc = AsyncMock(
|
|
return_value=[
|
|
[
|
|
OmniACK(task_id="t", status="SUCCESS", stage_id=0, rank=0),
|
|
OmniACK(task_id="t", status="ERROR", error_msg="out of memory"),
|
|
]
|
|
]
|
|
)
|
|
with pytest.raises(RuntimeError, match="out of memory"):
|
|
await omni.sleep(level=1)
|
|
|
|
omni.collective_rpc = AsyncMock(return_value=[OmniACK(task_id="t", status="SUCCESS", stage_id=0, rank=0)])
|
|
await omni.wake_up()
|
|
|
|
omni.collective_rpc.assert_awaited_once()
|
|
assert not omni._sleeping_tags
|
|
assert omni._paused is False
|
|
|
|
asyncio.run(run())
|
|
|
|
|
|
@pytest.mark.cpu
|
|
def test_wake_up_clears_diffusion_stages_that_woke_before_a_failure():
|
|
async def run() -> None:
|
|
omni = _make_omni(stage_types=["diffusion", "diffusion"])
|
|
await omni.sleep(level=1)
|
|
omni.collective_rpc = AsyncMock(
|
|
side_effect=[
|
|
[OmniACK(task_id="t", status="SUCCESS", stage_id=0, rank=0)],
|
|
[OmniACK(task_id="t", status="ERROR", error_msg="out of memory")],
|
|
]
|
|
)
|
|
|
|
with pytest.raises(RuntimeError, match="out of memory"):
|
|
await omni.wake_up()
|
|
|
|
assert omni._stage_sleeping_tags == {1: _SLEEP_TAGS}
|
|
|
|
asyncio.run(run())
|
|
|
|
|
|
@pytest.mark.cpu
|
|
def test_wake_up_settles_once_for_several_diffusion_stages(mocker):
|
|
async def run() -> None:
|
|
omni = _make_omni(stage_types=["diffusion", "diffusion"])
|
|
await omni.sleep(level=1)
|
|
settle = mocker.patch("vllm_omni.entrypoints.async_omni.asyncio.sleep", new=AsyncMock())
|
|
|
|
await omni.wake_up()
|
|
|
|
assert omni.collective_rpc.await_count == 3
|
|
settle.assert_awaited_once_with(0.1)
|
|
|
|
asyncio.run(run())
|
|
|
|
|
|
@pytest.mark.cpu
|
|
@pytest.mark.parametrize("result", [{"supported": False}, {"todo": "not supported yet"}], ids=["unsupported", "todo"])
|
|
def test_sleep_raises_on_rpc_result_without_error(result):
|
|
async def run() -> None:
|
|
omni = _make_omni(stage_types=["diffusion"])
|
|
omni.collective_rpc = AsyncMock(return_value=[result])
|
|
|
|
with pytest.raises(RuntimeError, match="handle_sleep_task failed"):
|
|
await omni.sleep(level=1)
|
|
|
|
asyncio.run(run())
|
|
|
|
|
|
@pytest.mark.cpu
|
|
def test_rpc_failure_result_is_not_resolved_as_ack():
|
|
async def run() -> None:
|
|
omni = _make_omni(stage_types=["diffusion"])
|
|
omni.event_resolver.resolve = AsyncMock()
|
|
omni.collective_rpc = AsyncMock(return_value=[{"supported": False, "error": "out of memory"}])
|
|
|
|
with pytest.raises(RuntimeError, match="out of memory"):
|
|
await omni.sleep(level=1)
|
|
|
|
omni.event_resolver.resolve.assert_not_awaited()
|
|
|
|
asyncio.run(run())
|
|
|
|
|
|
@pytest.mark.cpu
|
|
def test_sleep_raises_with_reason_of_rpc_error_envelope():
|
|
async def run() -> None:
|
|
omni = _make_omni(stage_types=["diffusion"])
|
|
omni.collective_rpc = AsyncMock(return_value=[{"error": True, "reason": "worker crashed"}])
|
|
|
|
with pytest.raises(RuntimeError, match="handle_sleep_task failed: worker crashed"):
|
|
await omni.sleep(level=1)
|
|
|
|
asyncio.run(run())
|
|
|
|
|
|
@pytest.mark.cpu
|
|
def test_failed_level2_sleep_can_still_be_woken():
|
|
async def run() -> None:
|
|
omni = _make_omni(stage_types=["diffusion"])
|
|
omni.collective_rpc = AsyncMock(return_value=[OmniACK(task_id="t", status="ERROR", error_msg="out of memory")])
|
|
with pytest.raises(RuntimeError, match="out of memory"):
|
|
await omni.sleep(level=2)
|
|
|
|
omni.collective_rpc = AsyncMock(return_value=[OmniACK(task_id="t", status="SUCCESS", stage_id=0, rank=0)])
|
|
await omni.wake_up()
|
|
|
|
omni.collective_rpc.assert_awaited_once()
|
|
assert omni._paused is False
|
|
|
|
asyncio.run(run())
|
|
|
|
|
|
@pytest.mark.cpu
|
|
def test_partial_level2_sleep_still_blocks_wake_up():
|
|
async def run() -> None:
|
|
omni = _make_omni(stage_types=["diffusion"])
|
|
omni.collective_rpc = AsyncMock(
|
|
return_value=[
|
|
[
|
|
OmniACK(task_id="t", status="SUCCESS", stage_id=0, rank=0),
|
|
OmniACK(task_id="t", status="ERROR", error_msg="out of memory"),
|
|
]
|
|
]
|
|
)
|
|
with pytest.raises(RuntimeError, match="out of memory"):
|
|
await omni.sleep(level=2)
|
|
|
|
with pytest.raises(NotImplementedError):
|
|
await omni.wake_up()
|
|
|
|
asyncio.run(run())
|
|
|
|
|
|
@pytest.mark.cpu
|
|
def test_is_sleeping_for_given_stages():
|
|
async def run() -> None:
|
|
omni = _make_omni(stage_types=["llm", "diffusion"])
|
|
await omni.sleep(stage_ids=[1], level=1)
|
|
|
|
assert await omni.is_sleeping(stage_ids=[1])
|
|
assert not await omni.is_sleeping(stage_ids=[0])
|
|
|
|
asyncio.run(run())
|