request.py

"""Per-request resource lifecycle.

Tracks every resource acquired during a generation (LoRAs including turbo,
VAE residency, keyframe images) and tears them all down in `_RequestContext.release()`,
even on failure. The authoritative record of what this task applied is
`_RequestContext.applied` / the parent `GenerationTask.applied_loras` — never
the global `*_merged` sets as request truth.

Intended to be used under `_MODEL_MUTATION_LOCK` for the mutate+generate
critical section when physical merge/unmerge is required.
"""
from __future__ import annotations

import time

from config import NSFW_LORA_CONFIGS, PLACEMENT
from lora.apply import _remove_custom, _remove_lora, _remove_nsfw, _remove_turbo
import state as _state
from state import (
    _CUSTOM_LoRA,
    _CUSTOM_merged,
    _LORA_ACCESS,
    _LORA_MAX_ENTRIES,
    _NSFW_LoRA,
    _NSFW_merged,
    _log,
)
from task import _task_log


def _touch_lora(key):
    """Update the last-access timestamp for a LoRA entry."""
    _LORA_ACCESS[key] = time.time()


def _evict_stale_loras():
    """Remove the oldest LoRA entries when either store exceeds _LORA_MAX_ENTRIES."""
    for store, merged, _scale_fn in [
        (_CUSTOM_LoRA, _CUSTOM_merged, None),
        (_NSFW_LoRA, _NSFW_merged, None),
    ]:
        if len(store) <= _LORA_MAX_ENTRIES:
            continue
        sorted_keys = sorted(store.keys(), key=lambda k: _LORA_ACCESS.get(k, 0))
        n_evict = len(store) - _LORA_MAX_ENTRIES
        for key in sorted_keys[:n_evict]:
            try:
                scale = NSFW_LORA_CONFIGS[key]["scale"] if key in NSFW_LORA_CONFIGS else 1.0
                _remove_lora(key, scale, store, merged)
                store.pop(key, None)
                _LORA_ACCESS.pop(key, None)
                _log(f"evicted stale LoRA {key}")
            except Exception as _e:
                _log(f"evict failed for {key}: {_e}", "ERROR")


class _RequestContext:
    """Per-request (task-scoped) resource lifecycle manager."""

    def __init__(self, task=None):
        self.task = task
        self._custom: list[tuple[tuple[str, str], float]] = []   # (key, scale)
        self._nsfw: list[tuple[str, float]] = []                 # (key, scale)
        self._turbo: bool = False
        self._vae_on_gpu = False
        self._pil_images: list = []                               # keyframe PILs to del
        self.applied: list[str] = []                              # human-readable names of applied adapters

    # ── Acquire helpers ────────────────────────────────────────────────────────

    def acquire_vae(self):
        """Move the two VAEs to GPU for this request (pack mode)."""
        if PLACEMENT == "pack" and _state.PIPE is not None:
            _state.PIPE.vae.to("cuda")
            _state.PIPE.audio_vae.to("cuda")
            self._vae_on_gpu = True

    def acquire_turbo(self):
        """Track that turbo LoRA was merged (will be unmerged on release)."""
        self._turbo = True
        self.applied.append("turbo")
        if self.task is not None:
            self.task.applied_loras.append("turbo")

    def acquire_custom_lora(self, key, scale):
        """Track a custom LoRA that was merged (will be unmerged on release)."""
        self._custom.append((key, scale))
        label = f"custom:{key[0]}/{key[1]}"
        self.applied.append(label)
        if self.task is not None:
            self.task.applied_loras.append(label)
        _touch_lora(key)

    def acquire_nsfw_lora(self, key, scale):
        """Track an NSFW LoRA that was merged (will be unmerged on release)."""
        self._nsfw.append((key, scale))
        label = f"nsfw:{key}"
        self.applied.append(label)
        if self.task is not None:
            self.task.applied_loras.append(label)
        _touch_lora(key)

    def track_pil(self, img):
        """Register a PIL image for cleanup on release."""
        if img is not None:
            self._pil_images.append(img)

    # ── Release ────────────────────────────────────────────────────────────────

    def release(self):
        """Release ALL acquired resources.  Safe to call multiple times.
        Unmerge order is reverse of application so baseline is restored.
        """
        # 1. Unmerge LoRAs in reverse application order — each independently.
        for key, scale in reversed(self._custom):
            try:
                _remove_custom(key, scale)
            except Exception as _e:
                _task_log(self.task, f"custom LoRA unmerge failed for {key}: {_e}", "ERROR")
        self._custom.clear()
        for key, scale in reversed(self._nsfw):
            try:
                _remove_nsfw(key, scale)
            except Exception as _e:
                _task_log(self.task, f"NSFW LoRA unmerge failed for {key}: {_e}", "ERROR")
        self._nsfw.clear()
        if self._turbo:
            try:
                _remove_turbo()
            except Exception as _e:
                _task_log(self.task, f"turbo LoRA unmerge failed: {_e}", "ERROR")
            self._turbo = False

        # 2. Move VAEs back to CPU to free ~10 GB GPU between requests.
        if self._vae_on_gpu and _state.PIPE is not None:
            try:
                _state.PIPE.vae.to("cpu")
                _state.PIPE.audio_vae.to("cpu")
            except Exception as _e:
                _task_log(self.task, f"VAE CPU offload failed: {_e}", "ERROR")
            self._vae_on_gpu = False

        # 3. Drop PIL images.
        for img in self._pil_images:
            try:
                img.close()
            except Exception:
                pass
        self._pil_images.clear()

        # 4. Evict stale LoRAs from the stores (LRU, bucket-cached so re-load is fast).
        try:
            _evict_stale_loras()
        except Exception as _e:
            _task_log(self.task, f"LoRA eviction error: {_e}", "ERROR")

        # 5. Aggressive CUDA cleanup: del + gc + empty_cache.
        import gc as _gc
        _gc.collect()
        try:
            import torch as _torch
            _torch.cuda.empty_cache()
        except Exception:
            pass