web/services/gpu.py

"""Gestione della VRAM: famiglie ML da scaricare e primitive CUDA (best effort).

Isolate dal `ServiceRegistry` perché sono dettagli di basso livello legati a
torch/CUDA: il registry decide COSA liberare (engine in cache, provider), qui ci
sono il COME (svuotare la cache CUDA) e il catalogo delle famiglie ML coinvolte.
"""

from __future__ import annotations

import logging

# Famiglie ML "pesanti" da liberare dalla VRAM quando si isola l'LLM-layer
# (il quale, però, HA bisogno della famiglia 'llm' → non la si scarica qui).
GPU_FAMILIES = ("ner", "ner_comuni", "zero_shot", "embedding", "baseline")
# TUTTE le famiglie ML (incl. l'LLM generativo): liberate prima di caricare il
# GIUDICE, che è un client LLM separato → nessun modello deve restare in VRAM.
ALL_ML_FAMILIES = ("ner", "ner_comuni", "zero_shot", "embedding", "baseline", "llm")


def log_vram(when: str, logger: logging.Logger) -> None:
    """Log dell'occupazione VRAM (best effort, solo se torch+cuda presenti)."""
    try:
        import torch
        if torch.cuda.is_available():
            alloc = torch.cuda.memory_allocated() / (1024 ** 3)
            reserved = torch.cuda.memory_reserved() / (1024 ** 3)
            logger.info(
                "[VRAM] %s: allocata %.2f GB, riservata %.2f GB",
                when, alloc, reserved,
            )
    except Exception:  # noqa: BLE001 - best effort
        pass


def release_cuda_cache() -> None:
    """Forza il GC e svuota la cache CUDA, se disponibile (best effort)."""
    import gc
    gc.collect()
    try:
        import torch
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
            torch.cuda.synchronize()
    except Exception:  # noqa: BLE001 - best effort
        pass