Yueqian LinandSy0307
e788ef6e76
Signed-off-by: Yueqian Lin <linyueqian@outlook.com> Signed-off-by: Sy0307 <24841857+Sy0307@users.noreply.github.com> Co-authored-by: Sy0307 <24841857+Sy0307@users.noreply.github.com>
102 lines
3.5 KiB
Python
102 lines
3.5 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Single-session lease teardown: a cancelled engine call must be drained.
|
|
|
|
Cancelling an asyncio task blocked on a worker-thread engine call returns
|
|
immediately while the thread keeps running (executor futures cannot be
|
|
interrupted). The server therefore tracks the in-flight future on the session
|
|
state and waits it out before releasing the capacity-one lease; otherwise a
|
|
fast reconnect could call ``open()`` on the shared engine while the previous
|
|
connection's final ``step()`` is still mutating it.
|
|
"""
|
|
|
|
import asyncio
|
|
import contextlib
|
|
import sys
|
|
import threading
|
|
import types
|
|
|
|
import pytest
|
|
|
|
try:
|
|
import sphn # noqa: F401
|
|
except ImportError:
|
|
# These lease tests never encode/decode Opus. Keep the optional serving
|
|
# dependency from hiding dependency-free teardown regressions.
|
|
sys.modules["sphn"] = types.ModuleType("sphn")
|
|
|
|
from vllm_omni.experimental.fullduplex.personaplex.config import PersonaPlexConfig
|
|
from vllm_omni.experimental.fullduplex.personaplex.serving.server import DuplexServer
|
|
from vllm_omni.experimental.fullduplex.personaplex.session import PersonaPlexServingSessionState
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cancelled_engine_call_is_drained_before_lease_release():
|
|
server = DuplexServer(PersonaPlexConfig()) # engine never loaded: stub calls only
|
|
state = PersonaPlexServingSessionState()
|
|
|
|
started = threading.Event()
|
|
release = threading.Event()
|
|
finished = threading.Event()
|
|
|
|
def slow_engine_op():
|
|
started.set()
|
|
release.wait(timeout=5.0)
|
|
finished.set()
|
|
return []
|
|
|
|
task = asyncio.create_task(server._engine_call(state, slow_engine_op))
|
|
await asyncio.to_thread(started.wait, 5.0)
|
|
|
|
# Cancel while the worker thread is mid-call: the await returns at once,
|
|
# the thread keeps running, and its completion event stays tracked on the
|
|
# state (the asyncio future itself reads as cancelled/done here, which is
|
|
# exactly why the drain must not rely on it).
|
|
task.cancel()
|
|
with contextlib.suppress(asyncio.CancelledError):
|
|
await task
|
|
assert state.inflight is not None
|
|
assert not state.inflight.is_set()
|
|
|
|
# The drain must genuinely wait for the worker, not just return.
|
|
drain = asyncio.create_task(server._drain_inflight(state))
|
|
await asyncio.sleep(0.05)
|
|
assert not drain.done()
|
|
|
|
release.set()
|
|
await asyncio.wait_for(drain, timeout=5.0)
|
|
assert finished.is_set()
|
|
assert state.inflight.is_set()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_completed_engine_call_clears_inflight():
|
|
server = DuplexServer(PersonaPlexConfig())
|
|
state = PersonaPlexServingSessionState()
|
|
result = await server._engine_call(state, lambda: "ok")
|
|
assert result == "ok"
|
|
assert state.inflight is None
|
|
await server._drain_inflight(state) # no-op on a clean state
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_drain_timeout_keeps_lease_until_worker_finishes():
|
|
server = DuplexServer(PersonaPlexConfig())
|
|
server._active = True
|
|
state = PersonaPlexServingSessionState()
|
|
state.inflight = threading.Event()
|
|
|
|
drained = await server._release_lease_after_drain(state, timeout=0.01)
|
|
|
|
assert drained is False
|
|
assert server._active is True
|
|
|
|
state.inflight.set()
|
|
for _ in range(100):
|
|
if not server._active:
|
|
break
|
|
await asyncio.sleep(0.01)
|
|
assert server._active is False
|