vllm-omni/vllm_omni/engine/duplex/session/lease.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project

"""Generic engine lease primitives shared by optional session runtimes."""

from __future__ import annotations

from dataclasses import dataclass, field
from enum import Enum


class DuplexLeaseActivity(str, Enum):
    APPEND = "append"
    SIGNAL = "signal"
    PLAYBACK_ACK = "playback_ack"
    HEARTBEAT = "heartbeat"
    ATTACH = "attach"
    DETACH = "detach"
    RESUME = "resume"
    MODEL_OUTPUT = "model_output"


@dataclass(frozen=True)
class DuplexLeaseConfig:
    idle_ttl_s: float | None = 300.0
    disconnect_grace_s: float = 30.0

    def __post_init__(self) -> None:
        if self.idle_ttl_s is not None and self.idle_ttl_s <= 0:
            raise ValueError("idle_ttl_s must be positive or null")
        if self.disconnect_grace_s <= 0:
            raise ValueError("disconnect_grace_s must be positive")


@dataclass
class DuplexLeaseState:
    config: DuplexLeaseConfig
    generation: int
    last_activity: float
    detached_at: float | None = None
    active_operations: set[str] = field(default_factory=set)
    terminal_reason: str | None = None
    #: Control id of the resume that produced ``generation`` (``None`` for the
    #: open). A resume replayed with the same id is answered with the current
    #: generation instead of bumping it again: the caller lost the first
    #: answer (its RPC waiter was cancelled or timed out) and is finding out
    #: whether the resume landed.
    resumed_by: str | None = None

    @property
    def expires_at(self) -> float | None:
        if self.config.idle_ttl_s is None:
            return None
        return self.last_activity + self.config.idle_ttl_s

    def _require_open(self) -> None:
        if self.terminal_reason is not None:
            raise RuntimeError(f"duplex lease is terminal: {self.terminal_reason}")

    def touch(self, now: float, activity: DuplexLeaseActivity) -> None:
        self._require_open()
        if not isinstance(activity, DuplexLeaseActivity):
            raise TypeError("lease activity must be DuplexLeaseActivity")
        self.last_activity = now

    def detach(self, now: float) -> None:
        self.touch(now, DuplexLeaseActivity.DETACH)
        self.detached_at = now

    def resume(self, now: float, *, expected_generation: int, control_id: str | None = None) -> int:
        self._require_open()
        if control_id is not None and control_id == self.resumed_by:
            # The same resume asking again: it already produced this
            # generation, so answer without changing anything.
            return self.generation
        if expected_generation != self.generation:
            raise ValueError(f"duplex lease generation mismatch: expected {self.generation}, got {expected_generation}")
        self.generation += 1
        self.resumed_by = control_id
        self.detached_at = None
        self.last_activity = now
        return self.generation

    def begin_operation(self, now: float, operation_id: str) -> None:
        self._require_open()
        if not operation_id:
            raise ValueError("duplex lease operation_id must not be empty")
        if operation_id in self.active_operations:
            raise ValueError(f"duplex lease operation already active: {operation_id}")
        self.active_operations.add(operation_id)
        self.last_activity = now

    def end_operation(self, now: float, operation_id: str) -> None:
        self._require_open()
        if operation_id not in self.active_operations:
            raise KeyError(f"duplex lease operation is not active: {operation_id}")
        self.active_operations.remove(operation_id)
        self.last_activity = now

    def disconnect_grace_expired(self, now: float) -> bool:
        return (
            self.terminal_reason is None
            and self.detached_at is not None
            and now >= self.detached_at + self.config.disconnect_grace_s
        )

    def idle_expired(self, now: float) -> bool:
        expires_at = self.expires_at
        return (
            self.terminal_reason is None and not self.active_operations and expires_at is not None and now >= expires_at
        )

    def mark_terminal(self, reason: str) -> bool:
        if self.terminal_reason is not None:
            return False
        self.terminal_reason = reason
        self.generation += 1
        self.resumed_by = None
        self.detached_at = None
        self.active_operations.clear()
        return True


__all__ = [
    "DuplexLeaseActivity",
    "DuplexLeaseConfig",
    "DuplexLeaseState",
]