File size: 5,026 Bytes
094b8e7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 | # -*- 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)
|