vllm-omni/vllm_omni/experimental/world_models/session_state/manager.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Session memory manager (RFC #4480).

``SessionStateManager`` maps ``session_id -> {name: StateObject}`` and is the
single authority for session lifecycle and eviction. It evicts by session count
(LRU), matching the per-model session caches it adapts (e.g. DreamZero's), so
an adapted model keeps the same eviction behaviour. A byte budget is recorded
for observability; enforcing it is left to a byte-budget eviction planner (see
RFC #4480).
"""

from __future__ import annotations

import logging
import os
import threading
from collections import OrderedDict
from collections.abc import Callable
from contextlib import AbstractContextManager
from typing import TypeVar

from vllm_omni.experimental.world_models.session_state.accounting import (
    SeenStorages,
    device_bytes,
    merge_bytes,
)
from vllm_omni.experimental.world_models.session_state.base import StateObject

logger = logging.getLogger(__name__)

# Matches DreamZero's MAX_DREAMZERO_SESSIONS cap so an adapted DreamZero
# evicts identically.
DEFAULT_MAX_SESSIONS = 64

M = TypeVar("M", bound=StateObject)


class SessionState:
    """The named ``StateObject`` collection for one session."""

    def __init__(self) -> None:
        self._objects: dict[str, StateObject] = {}
        # Session-scoped scalar/tensor metadata that is not itself a typed
        # StateObject (e.g. counters, cached conditioning tensors). Persists
        # across the per-call adapters that read/write it.
        # ``object`` (not ``Any``) so every read site must narrow explicitly.
        self.attrs: dict[str, object] = {}

    def get(self, name: str) -> StateObject | None:
        return self._objects.get(name)

    def put(self, name: str, obj: StateObject) -> StateObject:
        self._objects[name] = obj
        return obj

    def get_or_create(self, name: str, factory: Callable[[], M], cls: type[M]) -> M:
        """Return the resident object of type ``cls`` under ``name``, building
        it via ``factory`` when absent, wrong-typed, or no longer resident."""
        obj = self._objects.get(name)
        if not isinstance(obj, cls) or not obj.resident:
            obj = factory()
            self._objects[name] = obj
        return obj

    def names(self) -> list[str]:
        return list(self._objects)

    def reset(self) -> None:
        for obj in self._objects.values():
            obj.reset()
        # Clear session-scoped metadata too, so a reset leaves nothing stale
        # (otherwise language/current_start_frame/clip_feas/ys would survive).
        self.attrs.clear()

    def nbytes_by_device(self, seen: SeenStorages | None = None) -> dict[str, int]:
        """Bytes this session holds, keyed by device.

        Both buckets are walked. A ``StateObject`` and a session attribute are
        different in what the *manager* may do with them -- release, accumulate,
        stage and commit versus merely carry -- but not in whether their bytes
        occupy memory, so the accounting covers both. One ``seen`` set spans the
        whole session, so a tensor aliased between the two buckets is counted
        once.
        """
        shared: SeenStorages = set() if seen is None else seen
        totals: dict[str, int] = {}
        for obj in self._objects.values():
            merge_bytes(totals, obj.nbytes_by_device(shared))
        merge_bytes(totals, device_bytes(self.attrs, shared))
        return totals

    @property
    def nbytes(self) -> int:
        """Total bytes held, summed over every device -- see
        ``StateObject.nbytes`` on why a budget wants the per-device split."""
        return sum(self.nbytes_by_device().values())


class SessionStateManager:
    """Owns per-session state and arbitrates the (count-based) LRU.

    The manager lives beside ``DiffusionRequestState`` (owned by the pipeline),
    not inside its ``extra`` dict: the manager is cross-request and long-lived,
    whereas ``DiffusionRequestState`` is per-request and transient.
    """

    def __init__(
        self,
        max_sessions: int = DEFAULT_MAX_SESSIONS,
        byte_budget: int | None = None,
        lock_factory: Callable[[], AbstractContextManager[object]] = threading.Lock,
    ) -> None:
        if max_sessions <= 0:
            raise ValueError(f"max_sessions must be positive, got {max_sessions}")
        self.max_sessions = max_sessions
        # Recorded for observability; enforcement is left to an eviction
        # planner (see RFC #4480). Not scoped to a device: which pool a budget
        # applies to is a question for whoever enforces it, so read
        # ``nbytes_by_device()`` and pick the pool at the point of enforcement.
        self.byte_budget = byte_budget
        self._sessions: OrderedDict[str, SessionState] = OrderedDict()
        # The locking strategy is injected rather than hard-coded: callers
        # that run the manager from worker threads keep the default
        # ``threading.Lock``, while an async or single-threaded owner can pass
        # e.g. ``contextlib.nullcontext`` (or a custom strategy) without any
        # call-site change -- the manager only ever does ``with self._lock:``.
        # Heavier operations (paged allocation, eviction planning) keep this
        # seam.
        self._lock = lock_factory()
        self.hits = 0
        self.misses = 0
        self.evictions = 0

    @staticmethod
    def _key(session_id: str | None) -> str:
        return str(session_id or "default")

    def get_or_create_session(self, session_id: str | None) -> SessionState:
        key = self._key(session_id)
        with self._lock:
            session = self._sessions.get(key)
            if session is None:
                self.misses += 1
                session = SessionState()
                self._sessions[key] = session
                while len(self._sessions) > self.max_sessions:
                    # Drop the oldest from the table only; do not free its
                    # buffers. An adapter still using the session keeps its own
                    # reference, so its state survives (matching a model's own
                    # per-session state object, which the caller holds even
                    # after it leaves the table). Unreferenced sessions are
                    # GC-reclaimed.
                    self._sessions.popitem(last=False)
                    self.evictions += 1
            else:
                self.hits += 1
                self._sessions.move_to_end(key)
            return session

    def drop_session(self, session_id: str | None) -> bool:
        """Release one session explicitly and free its buffers.

        This is the end-of-session signal a model's engine raises on close,
        eviction, or reset -- distinct from the count-based LRU in
        ``get_or_create_session``, which only drops the *table* entry and lets an
        adapter still holding the session keep it alive. Here the session is
        finished, so its ``StateObject`` buffers and carried attributes are reset
        to release their memory. Returns whether a session was present.
        """
        key = self._key(session_id)
        with self._lock:
            session = self._sessions.pop(key, None)
            if session is None:
                return False
            session.reset()
            return True

    def __len__(self) -> int:
        with self._lock:
            return len(self._sessions)

    def __contains__(self, session_id: str | None) -> bool:
        with self._lock:
            return self._key(session_id) in self._sessions

    def nbytes_by_device(self) -> dict[str, int]:
        """Bytes held across every live session, keyed by device."""
        with self._lock:
            return self._nbytes_by_device_locked()

    def _nbytes_by_device_locked(self) -> dict[str, int]:
        # One ``seen`` set across all sessions: sessions can legitimately share
        # a storage (a conditioning tensor broadcast to several of them), and
        # the total is what the process holds, not the sum of what each session
        # would hold alone.
        shared: SeenStorages = set()
        totals: dict[str, int] = {}
        for session in self._sessions.values():
            merge_bytes(totals, session.nbytes_by_device(shared))
        return totals

    def stats(self) -> dict[str, int]:
        """Counters plus per-device byte totals.

        Byte keys are prefixed ``nbytes:`` and named by device -- ``nbytes:cpu``,
        ``nbytes:cuda:0`` -- because host and device memory are different
        resources and their sum answers no question. ``total_nbytes`` is kept
        for callers that only want a single number.
        """
        with self._lock:
            by_device = self._nbytes_by_device_locked()
            stats: dict[str, int] = {
                "sessions": len(self._sessions),
                "max_sessions": self.max_sessions,
                "total_nbytes": sum(by_device.values()),
                "hits": self.hits,
                "misses": self.misses,
                "evictions": self.evictions,
            }
            for device, count in sorted(by_device.items()):
                stats[f"nbytes:{device}"] = count
            return stats


def resolve_session_state_config(
    enable: bool | None = None,
    max_sessions: int | None = None,
) -> tuple[bool, int]:
    """Combine explicit args with env-var overrides.

    Environment variables (for quick enablement without touching config files):

        ``OMNI_DIFFUSION_SESSION_STATE_MANAGER`` (``1``/``0``/``true``/``false``)
        ``OMNI_DIFFUSION_SESSION_STATE_MANAGER_MAX_SESSIONS`` (positive int)

    Precedence: a *set* environment variable wins over the value passed in
    (i.e. over ``OmniDiffusionConfig.enable_session_state_manager``), in both
    directions -- ``...SESSION_MEMORY_MANAGER=0`` force-disables the manager
    even when the config enables it. Unset or unparsable values leave the
    passed-in value untouched (an unparsable ``MAX_SESSIONS`` logs a warning
    and is ignored). This mirrors how other omni feature toggles allow quick
    A/B flips on a deployed config.
    """
    env_enable = os.environ.get("OMNI_DIFFUSION_SESSION_STATE_MANAGER")
    if env_enable is not None:
        parsed = env_enable.strip().lower()
        if parsed in ("1", "true", "yes", "on"):
            enable = True
        elif parsed in ("0", "false", "no", "off"):
            enable = False

    env_size = os.environ.get("OMNI_DIFFUSION_SESSION_STATE_MANAGER_MAX_SESSIONS")
    if env_size is not None:
        try:
            env_size_int = int(env_size)
            if env_size_int > 0:
                max_sessions = env_size_int
        except ValueError:
            logger.warning(
                "Ignoring non-integer OMNI_DIFFUSION_SESSION_STATE_MANAGER_MAX_SESSIONS=%r.",
                env_size,
            )

    return bool(enable), int(max_sessions or DEFAULT_MAX_SESSIONS)