src/aurora/dados/cache.py

# -*- coding: utf-8 -*-
"""Cache em RAM e disco com orçamento explícito (item 1: otimização de RAM
e armazenamento). RAM: LRU por bytes estimados; Disco: arquivos .pt nomeados
por hash, com índice JSON e limites de total.

Política: prioridade RAM para tensores quentes (lote corrente), disco para
blobs frios (embeddings pré-calculados, amostras do corpus).
"""
from __future__ import annotations

import hashlib
import json
import os
from collections import OrderedDict

import torch


def _estimar_bytes(obj) -> int:
    if isinstance(obj, torch.Tensor):
        return obj.element_size() * obj.numel()
    if isinstance(obj, (list, tuple)):
        return sum(_estimar_bytes(o) for o in obj) + 64
    if isinstance(obj, dict):
        return sum(_estimar_bytes(k) + _estimar_bytes(v) for k, v in obj.items()) + 64
    return len(str(obj).encode("utf-8")) + 48


def _chave_hash(chave: str) -> str:
    return hashlib.sha1(chave.encode("utf-8")).hexdigest()[:16]


class CacheRAM:
    """LRU com limite de bytes. Evict: menos-recentemente-usado primeiro."""

    def __init__(self, limite_mb: float = 256.0):
        self.limite = int(limite_mb * 1024 * 1024)
        self._uso = 0
        self._mapa: OrderedDict[str, tuple] = OrderedDict()  # chave -> (valor, bytes)

    def get(self, chave: str):
        item = self._mapa.get(chave)
        if item is None:
            return None
        self._mapa.move_to_end(chave)  # toca LRU
        return item[0]

    def put(self, chave: str, valor) -> None:
        if chave in self._mapa:
            self._uso -= self._mapa[chave][1]
            del self._mapa[chave]
        n_bytes = _estimar_bytes(valor)
        if n_bytes > self.limite:  # indivisível e maior que tudo: descarta
            return
        while self._uso + n_bytes > self.limite and self._mapa:
            _, (_, b) = self._mapa.popitem(last=False)
            self._uso -= b
        self._mapa[chave] = (valor, n_bytes)
        self._uso += n_bytes

    @property
    def uso_mb(self) -> float:
        return self._uso / 1024 / 1024

    def __len__(self):
        return len(self._mapa)


class CacheDisco:
    """Blobs frios em disco com índice JSON e limite de armazenamento."""

    def __init__(self, diretorio: str, limite_mb: float = 2048.0):
        self.dir = diretorio
        os.makedirs(diretorio, exist_ok=True)
        self.indice_path = os.path.join(diretorio, "indice.json")
        self.limite = limite_mb * 1024 * 1024
        self.indice: dict[str, dict] = self._carregar_indice()

    def _carregar_indice(self) -> dict:
        if os.path.exists(self.indice_path):
            with open(self.indice_path, "r", encoding="utf-8") as f:
                return json.load(f)
        return {}

    def _salvar_indice(self) -> None:
        with open(self.indice_path, "w", encoding="utf-8") as f:
            json.dump(self.indice, f, ensure_ascii=False)

    def _evict_se_necessario(self, novo_bytes: int) -> None:
        total = sum(m["bytes"] for m in self.indice.values())
        while total + novo_bytes > self.limite and self.indice:
            chave_mais_velha = min(self.indice, key=lambda k: self.indice[k]["t"])
            meta = self.indice.pop(chave_mais_velha)
            try:
                os.remove(meta["arquivo"])
            except FileNotFoundError:
                pass
            total -= meta["bytes"]

    def put(self, chave: str, tensor: torch.Tensor) -> None:
        h = _chave_hash(chave)
        caminho = os.path.join(self.dir, h + ".pt")
        torch.save(tensor.cpu(), caminho)
        n_bytes = os.path.getsize(caminho)
        self._evict_se_necessario(n_bytes)
        self.indice[chave] = {"arquivo": caminho, "bytes": n_bytes, "t": os.path.getmtime(caminho)}
        self._salvar_indice()

    def get(self, chave: str) -> torch.Tensor | None:
        meta = self.indice.get(chave)
        if meta is None or not os.path.exists(meta["arquivo"]):
            return None
        return torch.load(meta["arquivo"], map_location="cpu", weights_only=True)


class CacheDuplo:
    """Fachada: RAM primeiro, disco como respaldo. Métricas p/ telemetria."""

    def __init__(self, limite_ram_mb: float, dir_disco: str, limite_disco_mb: float = 2048.0):
        self.ram = CacheRAM(limite_ram_mb)
        self.disco = CacheDisco(dir_disco, limite_disco_mb)
        self.estatisticas = {"hits_ram": 0, "hits_disco": 0, "misses": 0}

    def get(self, chave: str):
        v = self.ram.get(chave)
        if v is not None:
            self.estatisticas["hits_ram"] += 1
            return v
        v = self.disco.get(chave)
        if v is not None:
            self.estatisticas["hits_disco"] += 1
            self.ram.put(chave, v)  # promove
            return v
        self.estatisticas["misses"] += 1
        return None

    def put(self, chave: str, valor, persistir: bool = False) -> None:
        self.ram.put(chave, valor)
        if persistir and isinstance(valor, torch.Tensor):
            self.disco.put(chave, valor)