rpa_vision_v3/core/gpu/device_policy.py

"""Résolution de device paramétrable (auto/cuda/cpu) avec garde-fou VRAM.

Permet de basculer les étages CPU-par-défaut de la cascade vision (OCR docTR,
EasyOCR, YOLO/SoM) vers le GPU local **quand la VRAM est libre**, SANS jamais
hardcoder cuda. La politique anti-concurrence VRAM (tout sur CPU) datait d'une
époque où les VLM tournaient sur la RTX 5070 locale ; ils tournent désormais
sur un DGX distant (tunnel SSH `:11434`), libérant ~9 Go localement.

Logique de garde-fou inspirée de `core/embedding/clip_embedder.py` (lignes
~65-82) : `torch.cuda.is_available()` + `torch.cuda.mem_get_info()`.

Contraintes :
    - JAMAIS de hardcode cuda ;
    - aucun appel réseau ;
    - import-safe : aucun chargement de modèle, aucune allocation GPU à l'import ;
    - fallback CPU propre partout (jamais de crash si pas de GPU).

Override global : variable d'environnement `RPA_VISION_DEVICE` ∈ {cpu, cuda, auto}.
"""

from __future__ import annotations

import logging
import os
from typing import Optional

import torch

logger = logging.getLogger(__name__)

_GB = 1024 ** 3

# Valeurs reconnues pour l'argument `requested` et l'override env.
_VALID = {"cpu", "cuda", "auto"}

# Garde-fous par défaut (Go).
DEFAULT_MIN_FREE_GB = 2.0   # VRAM libre minimale pour autoriser cuda
DEFAULT_MAX_TOTAL_GB = 6.0  # plafond d'usage VRAM total après bascule
# Au-delà de ce total VRAM, on considère une grosse carte (data-center) ou une
# mémoire UNIFIÉE (DGX GB10 : ~121 Go partagés CPU+GPU). Dans ce cas `used`
# (= total - free) inclut la RAM système → le plafond fixe `max_total_gb` (pensé
# pour la RTX 12 Go dédiés) devient un faux positif qui force CPU à tort. On ne
# l'applique donc QUE sous ce seuil ; au-dessus, seul `free ≥ min_free_gb` décide.
DEFAULT_LARGE_VRAM_GB = 24.0


def _env_override() -> Optional[str]:
    """Lit l'override `RPA_VISION_DEVICE` s'il est présent et valide.

    Retourne None si absent ou invalide (on retombe alors sur `requested`).
    """
    raw = os.getenv("RPA_VISION_DEVICE", "").strip().lower()
    if not raw:
        return None
    if raw in _VALID:
        return raw
    logger.warning(
        "RPA_VISION_DEVICE='%s' invalide (attendu cpu/cuda/auto) — ignoré",
        raw,
    )
    return None


def _cuda_available() -> bool:
    """`torch.cuda.is_available()` protégé contre toute exception driver."""
    try:
        return bool(torch.cuda.is_available())
    except Exception as e:  # pragma: no cover - dépend du driver
        logger.debug("torch.cuda.is_available a levé : %s — CPU", e)
        return False


def _free_total_gb() -> Optional[tuple[float, float]]:
    """VRAM (libre, totale) en Go via mem_get_info, ou None si indisponible."""
    try:
        free_bytes, total_bytes = torch.cuda.mem_get_info()
        return free_bytes / _GB, total_bytes / _GB
    except Exception as e:  # pragma: no cover - dépend du driver
        logger.debug("torch.cuda.mem_get_info a levé : %s", e)
        return None


def resolve_device(
    requested: str = "auto",
    min_free_gb: float = DEFAULT_MIN_FREE_GB,
    max_total_gb: float = DEFAULT_MAX_TOTAL_GB,
) -> str:
    """Résout le device effectif ("cuda" ou "cpu") selon la politique VRAM.

    Args:
        requested: "cpu", "cuda" ou "auto" (défaut). L'override env
            `RPA_VISION_DEVICE` prime sur cet argument s'il est présent/valide.
        min_free_gb: VRAM libre minimale (Go) pour autoriser cuda en mode auto.
        max_total_gb: plafond d'usage VRAM total (Go). Si basculer cuda ferait
            dépasser ce plafond (used = total - free), on reste CPU. Garde-fou
            contre la saturation quand d'autres process occupent déjà le GPU.

    Returns:
        "cuda" ou "cpu". Toujours "cpu" en cas de doute (fallback propre).

    Politique :
        - "cpu"  → "cpu" ;
        - "cuda" → "cuda" si cuda dispo, sinon "cpu" (fallback loggé) ;
        - "auto" → "cuda" si cuda dispo ET free ≥ min_free_gb ET
                   used ≤ max_total_gb, sinon "cpu".
    """
    effective = _env_override() or (requested or "auto").strip().lower()
    if effective not in _VALID:
        logger.warning(
            "device demandé '%s' invalide (attendu cpu/cuda/auto) — auto",
            effective,
        )
        effective = "auto"

    if effective == "cpu":
        return "cpu"

    if not _cuda_available():
        if effective == "cuda":
            logger.info("device=cuda demandé mais CUDA indisponible — fallback CPU")
        return "cpu"

    if effective == "cuda":
        # Demande explicite : on respecte sans appliquer le garde-fou VRAM
        # (l'appelant assume). CUDA est dispo → cuda.
        return "cuda"

    # effective == "auto" : garde-fou VRAM.
    mem = _free_total_gb()
    if mem is None:
        logger.info("auto: mem_get_info indisponible — CPU par prudence")
        return "cpu"

    free_gb, total_gb = mem
    used_gb = total_gb - free_gb

    if free_gb < min_free_gb:
        logger.info(
            "auto: VRAM libre %.1f Go < seuil %.1f Go — CPU",
            free_gb, min_free_gb,
        )
        return "cpu"

    # Plafond d'usage : seulement sur carte dédiée "petite" (type RTX). Sur grosse
    # mémoire / mémoire unifiée (GB10), `used` inclut la RAM système → non pertinent.
    if total_gb <= DEFAULT_LARGE_VRAM_GB and used_gb > max_total_gb:
        logger.info(
            "auto: usage VRAM %.1f Go > plafond %.1f Go (carte %.1f Go) — CPU",
            used_gb, max_total_gb, total_gb,
        )
        return "cpu"

    if total_gb > DEFAULT_LARGE_VRAM_GB:
        logger.info(
            "auto: grosse mémoire/unifiée %.1f Go, libre %.1f Go — CUDA (plafond ignoré)",
            total_gb, free_gb,
        )
        return "cuda"

    logger.info(
        "auto: VRAM libre %.1f Go (usage %.1f/%.1f Go) — CUDA",
        free_gb, used_gb, total_gb,
    )
    return "cuda"