lightseekorg--tokenspeed/python/tokenspeed/runtime/engine/memory_occupation.py

# Copyright (c) 2026 LightSeek Foundation
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

"""GPU-memory data plane for sleep/wake (release/resume_memory_occupation).

Pairs with the control-plane :class:`PauseController`. A *release* pauses and
drains the scheduler (delegated to the pause controller) and only then frees GPU
memory via the torch_memory_saver adapter; a *resume* re-maps memory and, when
the KV region returns, repairs the KV cache. Tags (``weights``, ``kv_cache``)
are freed/restored independently so the RL loop can resume weights, push fresh
weights, then resume the KV cache.

This module knows nothing about ZMQ beyond a ``send_func`` reply socket, and
nothing about scheduling beyond delegating drain to the pause controller.
"""

from __future__ import annotations

from collections.abc import Callable

from tokenspeed.runtime.engine.io_struct import (
    IsSleepingReqInput,
    IsSleepingReqOutput,
    ReleaseMemoryOccupationReqInput,
    ReleaseMemoryOccupationReqOutput,
    ResumeMemoryOccupationReqInput,
    ResumeMemoryOccupationReqOutput,
)
from tokenspeed.runtime.engine.pause import PauseController

VALID_TAGS = ("weights", "kv_cache")


def _normalize_tags(tags: list[str] | None) -> tuple[list[str] | None, str | None]:
    """Return ``(ordered_tags, error)``. ``None`` ⇒ all tags. The order is fixed
    (``weights`` before ``kv_cache``) for determinism; an unknown tag yields an
    error and ``None`` tags."""
    if tags is None:
        return list(VALID_TAGS), None
    bad = [t for t in tags if t not in VALID_TAGS]
    if bad:
        return None, f"invalid tags: {bad!r}; valid: {list(VALID_TAGS)}"
    return [t for t in VALID_TAGS if t in tags], None


class MemoryOccupationController:
    """Owns the data-plane release/resume for one scheduler event loop."""

    def __init__(
        self,
        *,
        send_func,
        pause_controller: PauseController,
        adapter,
        enabled: bool,
        reset_caches_fn: Callable[[], None],
        kv_repair_fn: Callable[[], None],
        kv_cache_release_allowed: bool = True,
    ) -> None:
        self._send = send_func
        self._pause = pause_controller
        self._adapter = adapter
        self._enabled = enabled
        self._reset_caches = reset_caches_fn
        self._kv_repair = kv_repair_fn
        # Whether releasing the ``kv_cache`` region is safe in this engine. False
        # when prefix caching is on but the scheduler exposes no prefix-cache
        # reset: discarded KV would leave stale cache entries that survive wake
        # and produce silently wrong hits, so we reject the release up front
        # rather than free KV we cannot invalidate. A declared capability (set
        # at construction), not a runtime duck-typed probe.
        self._kv_cache_release_allowed = kv_cache_release_allowed
        self.released_tags: set[str] = set()

    @property
    def is_sleeping(self) -> bool:
        """True while any GPU-memory region is released (data-plane sleep)."""
        return bool(self.released_tags)

    # -- control-request handlers (driven by the request handler) -------------

    def handle_release(self, req: ReleaseMemoryOccupationReqInput) -> None:
        if not self._enabled:
            self._fail_release("memory saver not enabled (--enable-memory-saver)")
            return
        tags, err = _normalize_tags(req.tags)
        if err is not None:
            self._fail_release(err)
            return
        if "kv_cache" in tags and not self._kv_cache_release_allowed:
            self._fail_release(
                "cannot release kv_cache: prefix caching is enabled but the "
                "scheduler has no prefix-cache reset, so discarded KV would "
                "leave stale cache entries after wake"
            )
            return
        already = [t for t in tags if t in self.released_tags]
        if already:
            self._fail_release(f"tags already released: {already!r}")
            return
        if self._pause.is_drain_pending:
            self._fail_release("a pause or release is already in progress")
            return
        # Defer the actual free until the scheduler drains. wait-mode: let
        # in-flight requests finish (the RL rollout is idle between steps).
        self._pause.request_drain(
            abort_inflight=False,
            on_drained=lambda: self._finish_release(tags),
            on_cancelled=lambda: self._fail_release("resumed before release drained"),
        )

    def _finish_release(self, tags: list[str]) -> None:
        if "kv_cache" in tags:
            # KV is discarded; any prefix-cache entry pointing at it is stale.
            self._reset_caches()
        for tag in tags:
            self._adapter.pause(tag=tag)
            self.released_tags.add(tag)
        self._pause.set_released(True)
        self._send.send_pyobj(ReleaseMemoryOccupationReqOutput(success=True))

    def _fail_release(self, message: str) -> None:
        self._send.send_pyobj(
            ReleaseMemoryOccupationReqOutput(success=False, message=message)
        )

    def handle_resume(self, req: ResumeMemoryOccupationReqInput) -> None:
        # ``None`` on resume means "wake exactly what is asleep" — NOT all valid
        # tags. Expanding to all (as on release) would make a bare resume after a
        # partial release (e.g. weights-only) fail on the never-released tag,
        # stranding the released region. Resolve against ``released_tags`` so a
        # partial release round-trips with a default resume_memory_occupation().
        if req.tags is None:
            tags = [t for t in VALID_TAGS if t in self.released_tags]
        else:
            tags, err = _normalize_tags(req.tags)
            if err is not None:
                self._send.send_pyobj(
                    ResumeMemoryOccupationReqOutput(success=False, message=err)
                )
                return
        not_released = [t for t in tags if t not in self.released_tags]
        if not_released:
            self._send.send_pyobj(
                ResumeMemoryOccupationReqOutput(
                    success=False, message=f"tags not released: {not_released!r}"
                )
            )
            return
        for tag in tags:
            self._adapter.resume(tag=tag)
            self.released_tags.discard(tag)
        if "kv_cache" in tags:
            self._kv_repair()
        if not self.released_tags:
            # Fully awake → resume scheduling.
            self._pause.set_released(False)
        self._send.send_pyobj(ResumeMemoryOccupationReqOutput(success=True))

    def handle_is_sleeping(self, req: IsSleepingReqInput) -> None:
        self._send.send_pyobj(IsSleepingReqOutput(is_sleeping=self.is_sleeping))