225 lines
8.2 KiB
Python
225 lines
8.2 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import FrozenInstanceError
|
|
|
|
import pytest
|
|
|
|
from vllm_omni.experimental.fullduplex.engine.duplex_runtime import DuplexInputMode
|
|
from vllm_omni.experimental.fullduplex.engine.duplex_session import DuplexSessionRuntimeManager
|
|
from vllm_omni.experimental.fullduplex.engine.lease import (
|
|
DuplexLeaseActivity,
|
|
DuplexLeaseConfig,
|
|
)
|
|
from vllm_omni.experimental.fullduplex.engine.messages import DuplexFence
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
|
|
|
|
|
class FakeMonotonicClock:
|
|
def __init__(self, value: float) -> None:
|
|
self.value = value
|
|
|
|
def __call__(self) -> float:
|
|
return self.value
|
|
|
|
def advance(self, seconds: float) -> None:
|
|
self.value += seconds
|
|
|
|
|
|
def _lease_config(*, idle_ttl_s: float | None = 300.0) -> DuplexLeaseConfig:
|
|
return DuplexLeaseConfig(
|
|
idle_ttl_s=idle_ttl_s,
|
|
disconnect_grace_s=30.0,
|
|
)
|
|
|
|
|
|
def test_open_touch_detach_and_idle_expiry_use_monotonic_time() -> None:
|
|
clock = FakeMonotonicClock(100.0)
|
|
manager = DuplexSessionRuntimeManager(clock=clock)
|
|
fence = DuplexFence("sid-expiry")
|
|
session = manager.open_session(fence, lease_config=_lease_config())
|
|
|
|
assert session.lease.last_activity == 100.0
|
|
clock.advance(10.0)
|
|
session.touch(fence, DuplexLeaseActivity.HEARTBEAT)
|
|
session.detach(fence)
|
|
clock.advance(29.0)
|
|
assert session.lease.disconnect_grace_expired(clock()) is False
|
|
assert manager.collect_expired() == []
|
|
|
|
clock.advance(272.0)
|
|
expired = manager.collect_expired()
|
|
|
|
assert [item.session_id for item in expired] == ["sid-expiry"]
|
|
assert expired[0].reason == "disconnect_grace_expired"
|
|
assert manager.get("sid-expiry") is session
|
|
manager.finalize_close_session(session)
|
|
assert manager.get("sid-expiry") is None
|
|
|
|
|
|
def test_detach_grace_and_resume_advance_token_independent_lease_generation() -> None:
|
|
clock = FakeMonotonicClock(10.0)
|
|
manager = DuplexSessionRuntimeManager(clock=clock)
|
|
fence = DuplexFence("sid-resume")
|
|
session = manager.open_session(fence, lease_config=_lease_config())
|
|
|
|
session.detach(fence)
|
|
clock.advance(29.0)
|
|
assert session.lease.disconnect_grace_expired(clock()) is False
|
|
clock.advance(2.0)
|
|
assert session.lease.disconnect_grace_expired(clock()) is True
|
|
|
|
generation = session.resume(fence, expected_lease_generation=0)
|
|
|
|
assert generation == 1
|
|
assert session.lease.detached_at is None
|
|
assert session.lease.disconnect_grace_expired(clock()) is False
|
|
with pytest.raises(ValueError, match="lease generation mismatch"):
|
|
session.resume(fence, expected_lease_generation=0)
|
|
|
|
|
|
def test_active_operation_prevents_mid_transaction_expiry() -> None:
|
|
clock = FakeMonotonicClock(0.0)
|
|
manager = DuplexSessionRuntimeManager(clock=clock)
|
|
fence = DuplexFence("sid-operation")
|
|
session = manager.open_session(fence, lease_config=_lease_config(idle_ttl_s=5.0))
|
|
|
|
session.begin_operation(fence, "append-1")
|
|
clock.advance(10.0)
|
|
assert manager.collect_expired() == []
|
|
|
|
session.end_operation(fence, "append-1")
|
|
assert manager.collect_expired() == []
|
|
clock.advance(6.0)
|
|
assert [item.session_id for item in manager.collect_expired()] == ["sid-operation"]
|
|
|
|
|
|
def test_sessions_have_independent_activity_deadlines_and_resources() -> None:
|
|
clock = FakeMonotonicClock(0.0)
|
|
manager = DuplexSessionRuntimeManager(clock=clock)
|
|
fence_a = DuplexFence("sid-a")
|
|
fence_b = DuplexFence("sid-b")
|
|
session_a = manager.open_session(fence_a, lease_config=_lease_config(idle_ttl_s=10.0))
|
|
session_b = manager.open_session(fence_b, lease_config=_lease_config(idle_ttl_s=10.0))
|
|
session_a.reserve_stage_request(0, "req-a-reserved", fence=fence_a)
|
|
session_a.bind_stage_request(1, "req-a-submitted", fence=fence_a)
|
|
session_b.bind_stage_request(0, "req-b", fence=fence_b)
|
|
|
|
clock.advance(6.0)
|
|
session_b.touch(fence_b, DuplexLeaseActivity.MODEL_OUTPUT)
|
|
clock.advance(5.0)
|
|
expired = manager.collect_expired()
|
|
|
|
assert len(expired) == 1
|
|
assert expired[0].session_id == "sid-a"
|
|
assert expired[0].reserved_request_ids == ("req-a-reserved",)
|
|
assert expired[0].submitted_request_ids == ("req-a-submitted",)
|
|
assert manager.get("sid-b") is session_b
|
|
assert session_b.resource_request_ids() == ["req-b"]
|
|
|
|
|
|
def test_close_and_reaper_select_exactly_one_terminal_transition() -> None:
|
|
clock = FakeMonotonicClock(0.0)
|
|
manager = DuplexSessionRuntimeManager(clock=clock)
|
|
fence = DuplexFence("sid-race")
|
|
manager.open_session(fence, lease_config=_lease_config(idle_ttl_s=1.0))
|
|
clock.advance(2.0)
|
|
|
|
assert len(manager.collect_expired()) == 1
|
|
assert manager.close_session(fence, reason="explicit_close") is None
|
|
assert manager.collect_expired() == []
|
|
|
|
fence_2 = DuplexFence("sid-race-2")
|
|
manager.open_session(fence_2, lease_config=_lease_config(idle_ttl_s=1.0))
|
|
assert manager.close_session(fence_2, reason="explicit_close") is not None
|
|
clock.advance(2.0)
|
|
assert manager.collect_expired() == []
|
|
|
|
|
|
def test_stale_fence_cannot_touch_detach_or_resume_lease() -> None:
|
|
clock = FakeMonotonicClock(0.0)
|
|
manager = DuplexSessionRuntimeManager(clock=clock)
|
|
current = DuplexFence("sid-fence", epoch=1)
|
|
stale = DuplexFence("sid-fence", epoch=0)
|
|
session = manager.open_session(current, lease_config=_lease_config())
|
|
|
|
with pytest.raises(RuntimeError, match="fence mismatch"):
|
|
session.touch(stale, DuplexLeaseActivity.APPEND)
|
|
with pytest.raises(RuntimeError, match="fence mismatch"):
|
|
session.detach(stale)
|
|
with pytest.raises(RuntimeError, match="fence mismatch"):
|
|
session.resume(stale, expected_lease_generation=0)
|
|
|
|
|
|
def test_disabled_idle_expiry_never_collects_session() -> None:
|
|
clock = FakeMonotonicClock(0.0)
|
|
manager = DuplexSessionRuntimeManager(clock=clock)
|
|
manager.open_session(DuplexFence("sid-no-expiry"), lease_config=_lease_config(idle_ttl_s=None))
|
|
|
|
clock.advance(1_000_000.0)
|
|
|
|
assert manager.collect_expired() == []
|
|
|
|
|
|
def test_detached_session_expires_at_disconnect_grace_when_idle_ttl_is_disabled() -> None:
|
|
clock = FakeMonotonicClock(0.0)
|
|
manager = DuplexSessionRuntimeManager(clock=clock)
|
|
fence = DuplexFence("sid-disconnect-grace")
|
|
session = manager.open_session(fence, lease_config=_lease_config(idle_ttl_s=None))
|
|
|
|
session.detach(fence)
|
|
clock.advance(29.0)
|
|
assert manager.collect_expired() == []
|
|
|
|
clock.advance(1.0)
|
|
expired = manager.collect_expired()
|
|
|
|
assert [item.session_id for item in expired] == [fence.session_id]
|
|
assert expired[0].reason == "disconnect_grace_expired"
|
|
|
|
|
|
def test_runtime_manager_enforces_server_owned_session_admission() -> None:
|
|
manager = DuplexSessionRuntimeManager(max_sessions=2)
|
|
manager.open_session(DuplexFence("sid-admission-a"))
|
|
manager.open_session(DuplexFence("sid-admission-b"))
|
|
|
|
with pytest.raises(RuntimeError, match="duplex_session_capacity_exhausted"):
|
|
manager.open_session(DuplexFence("sid-admission-c"))
|
|
|
|
|
|
def test_completed_append_cache_is_bounded() -> None:
|
|
manager = DuplexSessionRuntimeManager(completed_append_limit=2)
|
|
fence = DuplexFence("sid-completed-cache")
|
|
session = manager.open_session(fence)
|
|
|
|
for index in range(3):
|
|
session.record_completed_append(
|
|
f"operation-{index}",
|
|
fence=fence,
|
|
mode=DuplexInputMode.TURN_COMMIT_ONLY,
|
|
final=True,
|
|
stage_results=[{"index": index}],
|
|
)
|
|
|
|
assert list(session.completed_appends) == ["operation-1", "operation-2"]
|
|
|
|
session.accept_fence(DuplexFence(fence.session_id, epoch=1))
|
|
assert session.completed_appends == {}
|
|
|
|
|
|
def test_lease_config_and_expiry_record_are_immutable() -> None:
|
|
config = _lease_config()
|
|
with pytest.raises(FrozenInstanceError):
|
|
config.idle_ttl_s = 1.0 # type: ignore[misc]
|
|
|
|
clock = FakeMonotonicClock(0.0)
|
|
manager = DuplexSessionRuntimeManager(clock=clock)
|
|
manager.open_session(DuplexFence("sid-record"), lease_config=_lease_config(idle_ttl_s=1.0))
|
|
clock.advance(2.0)
|
|
record = manager.collect_expired()[0]
|
|
with pytest.raises(FrozenInstanceError):
|
|
record.reason = "changed" # type: ignore[misc]
|