101 lines
3.5 KiB
Python
101 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
|