gpu_utils.py

"""
GPU Utilities for Topic Modeling System
Handles device detection and configuration for:
  - Multi-GPU CUDA (NVIDIA): RTX 5090 x2, etc.
  - MPS (Apple Silicon)
  - CPU fallback

Multi-GPU optimizations:
  - NCCL backend for fast GPU-to-GPU communication
  - Mixed precision (BF16 for RTX 5090, FP16 otherwise)
  - Per-GPU memory querying for dynamic batch sizing
"""
import os
import torch
import warnings
from typing import Optional, List


def setup_gpu_environment():
    """
    Setup optimal GPU environment settings.
    Call this BEFORE creating any models or loading data.
    Enables cuDNN auto-tuning, reproducibility, and NCCL for multi-GPU.
    """
    # Reproducibility
    torch.manual_seed(42)

    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(42)

        # cuDNN benchmark mode: lets cuDNN find the fastest conv algorithm
        # per input shape. Speeds up training when input sizes are fixed.
        torch.backends.cudnn.benchmark = True
        torch.backends.cudnn.enabled = True

        # Set NCCL environment for optimal GPU-to-GPU communication
        # NCCL is NVIDIA's library for collective operations (all-reduce, etc.)
        os.environ.setdefault("NCCL_DEBUG", "WARN")
        os.environ.setdefault("NCCL_SOCKET_IFNAME", "^lo,docker")   # Skip loopback
        os.environ.setdefault("NCCL_IB_DISABLE", "0")               # Allow InfiniBand
        os.environ.setdefault("NCCL_P2P_DISABLE", "0")              # Allow P2P transfers

    # Disable gradient computation globally — we are doing inference only
    torch.set_grad_enabled(False)

    print("✓ GPU environment configured")


def get_supported_precision(device_index: int = 0) -> str:
    """
    Return the best floating-point precision for the given CUDA device.
    RTX 5090 (Blackwell) and Ampere/Ada Lovelace support BF16 natively.
    BF16 has wider dynamic range than FP16 — preferred for NLP embeddings.

    Returns 'bf16', 'fp16', or 'fp32'.
    """
    if not torch.cuda.is_available():
        return "fp32"

    props = torch.cuda.get_device_properties(device_index)
    major = props.major  # CUDA compute capability major version

    if major >= 8:
        # Ampere (A100, RTX 3090), Ada Lovelace (RTX 4090),
        # Blackwell (RTX 5090) — all have native BF16 tensor cores
        return "bf16"
    elif major >= 7:
        # Volta (V100), Turing (RTX 2080) — FP16 tensor cores
        return "fp16"
    else:
        return "fp32"


class GPUManager:
    """
    Manages GPU device selection, memory reporting, and multi-GPU config.
    Detects all available CUDA GPUs and uses them all by default.
    """

    def __init__(self, force_cpu: bool = False):
        self.force_cpu = force_cpu
        self.device = self._detect_primary_device()
        self.n_gpus = self._count_gpus()
        self.device_name = self._get_device_name()
        self.precision = (
            get_supported_precision(0) if self.device == "cuda" else "fp32"
        )
        self._print_device_info()

    def _detect_primary_device(self) -> str:
        """Return 'cuda', 'mps', or 'cpu'."""
        if self.force_cpu:
            return "cpu"
        if torch.cuda.is_available():
            return "cuda"
        if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
            return "mps"
        return "cpu"

    def _count_gpus(self) -> int:
        """Return the number of usable CUDA GPUs (0 for MPS/CPU)."""
        if self.device == "cuda":
            return torch.cuda.device_count()
        return 0

    def _get_device_name(self) -> str:
        if self.device == "cuda":
            names = [torch.cuda.get_device_name(i) for i in range(self.n_gpus)]
            return " | ".join(names)
        if self.device == "mps":
            return "Apple Silicon GPU (Metal)"
        return "CPU"

    def _print_device_info(self):
        print("\n" + "=" * 80)
        print("GPU CONFIGURATION")
        print("=" * 80)
        print(f"Primary device : {self.device.upper()}")
        print(f"Device name(s) : {self.device_name}")

        if self.device == "cuda":
            print(f"GPU count      : {self.n_gpus}")
            print(f"Precision      : {self.precision.upper()}")
            print(f"CUDA version   : {torch.version.cuda}")
            total_vram = 0
            for i in range(self.n_gpus):
                props = torch.cuda.get_device_properties(i)
                mem_gb = props.total_memory / 1e9
                total_vram += mem_gb
                print(f"  GPU {i}: {props.name}  {mem_gb:.1f} GB VRAM")
            print(f"Total VRAM     : {total_vram:.1f} GB")
        elif self.device == "mps":
            print("MPS (Metal Performance Shaders) enabled")
        else:
            warnings.warn("⚠️  Running on CPU — this will be slow for large datasets.")

        print("=" * 80 + "\n")

    def get_device(self) -> str:
        """Return the primary device string used for single-GPU operations."""
        return self.device

    def get_all_cuda_devices(self) -> List[str]:
        """
        Return a list of all CUDA device strings, e.g. ['cuda:0', 'cuda:1'].
        Falls back to [self.device] for MPS/CPU.
        """
        if self.device == "cuda" and self.n_gpus > 1:
            return [f"cuda:{i}" for i in range(self.n_gpus)]
        return [self.device]

    def get_total_vram_gb(self) -> float:
        """Sum VRAM across all CUDA GPUs."""
        if self.device != "cuda":
            return 0.0
        return sum(
            torch.cuda.get_device_properties(i).total_memory / 1e9
            for i in range(self.n_gpus)
        )

    def clear_cache(self):
        """Free unused GPU memory on ALL devices."""
        if self.device == "cuda":
            for i in range(self.n_gpus):
                with torch.cuda.device(i):
                    torch.cuda.empty_cache()
            print(f"✓ CUDA cache cleared on {self.n_gpus} GPU(s)")
        elif self.device == "mps":
            print("✓ MPS cache management handled automatically")

    def print_memory_info(self):
        """Print current memory usage for each CUDA GPU."""
        if self.device != "cuda":
            return
        for i in range(self.n_gpus):
            props = torch.cuda.get_device_properties(i)
            allocated = torch.cuda.memory_allocated(i) / 1e9
            reserved  = torch.cuda.memory_reserved(i) / 1e9
            total     = props.total_memory / 1e9
            print(
                f"  GPU {i}: allocated {allocated:.2f} GB / "
                f"reserved {reserved:.2f} GB / total {total:.1f} GB"
            )

    def get_memory_info(self) -> dict:
        """Return per-GPU memory dict (CUDA only)."""
        if self.device != "cuda":
            return {"message": "Memory info only available for CUDA devices"}
        info = {}
        for i in range(self.n_gpus):
            allocated = torch.cuda.memory_allocated(i) / 1e9
            reserved  = torch.cuda.memory_reserved(i) / 1e9
            total     = torch.cuda.get_device_properties(i).total_memory / 1e9
            info[f"gpu_{i}"] = {
                "allocated": allocated,
                "reserved": reserved,
                "free": total - reserved,
                "total": total,
                "utilization": (allocated / total) * 100,
            }
        return info


def get_optimal_batch_size(
    device: str,
    num_documents: int,
    embedding_dim: int = 384,
    n_gpus: int = 1,
) -> int:
    """
    Calculate optimal batch size per GPU for embedding generation.

    Scales with total VRAM so that both GPUs are fully utilised.
    For multi-GPU, this returns the per-GPU batch size; multiply by n_gpus
    for the effective global batch size.

    Args:
        device        : 'cuda', 'mps', or 'cpu'
        num_documents : total number of documents (used for a sanity cap)
        embedding_dim : sentence-transformer output dim (default 384)
        n_gpus        : number of GPUs (scales the safe usage estimate)

    Returns:
        per-GPU batch size (int)
    """
    if device == "cuda":
        # Query the first GPU's VRAM; assume all GPUs are identical
        gpu_mem_gb = torch.cuda.get_device_properties(0).total_memory / 1e9

        # Use 70 % of each GPU's memory — leaves headroom for model weights
        # and BERTopic internals
        usable_mem_gb = gpu_mem_gb * 0.70

        # Empirical: sentence-transformers (MiniLM-L6) uses ~6-8 MB per sample
        # including activations and intermediate buffers
        mem_per_sample_mb = 8
        usable_mem_mb = usable_mem_gb * 1024

        batch_size = int(usable_mem_mb / mem_per_sample_mb)

        # Clamp to sensible hardware tiers
        if gpu_mem_gb < 8:        # RTX 1660 / 3060 class
            batch_size = max(64,  min(batch_size, 128))
        elif gpu_mem_gb < 16:     # RTX 3080 / 4070 class
            batch_size = max(128, min(batch_size, 256))
        elif gpu_mem_gb < 24:     # RTX 3090 / 4080 class
            batch_size = max(256, min(batch_size, 512))
        elif gpu_mem_gb < 36:     # RTX 4090 / 5080 (24 GB) class
            batch_size = max(512, min(batch_size, 1024))
        else:                     # RTX 5090 (32 GB) — plenty of room
            batch_size = max(512, min(batch_size, 1536))

        total_mem = gpu_mem_gb * n_gpus
        print(f"  Per-GPU batch size : {batch_size}")
        print(f"  Effective global   : {batch_size * n_gpus}  "
              f"({n_gpus} GPU × {batch_size})")
        print(f"  Total VRAM in use  : {total_mem:.1f} GB across {n_gpus} GPU(s)")

    elif device == "mps":
        # Apple Silicon: unified memory, conservative default
        batch_size = 128
        print(f"  MPS batch size: {batch_size}")

    else:
        # CPU: keep batches small to avoid RAM pressure
        batch_size = 32
        print(f"  CPU batch size: {batch_size}")

    return batch_size


# ---------------------------------------------------------------------------
# Module-level singleton
# ---------------------------------------------------------------------------
_gpu_manager: Optional[GPUManager] = None


def get_gpu_manager(force_cpu: bool = False) -> GPUManager:
    """Return (or create) the global GPUManager singleton."""
    global _gpu_manager
    if _gpu_manager is None or force_cpu:
        _gpu_manager = GPUManager(force_cpu=force_cpu)
    return _gpu_manager


# ---------------------------------------------------------------------------
# Self-test
# ---------------------------------------------------------------------------
if __name__ == "__main__":
    setup_gpu_environment()
    mgr = get_gpu_manager()

    print(f"Primary device : {mgr.get_device()}")
    print(f"All CUDA devs  : {mgr.get_all_cuda_devices()}")
    print(f"Total VRAM     : {mgr.get_total_vram_gb():.1f} GB")

    if mgr.device == "cuda":
        mgr.print_memory_info()

    bs = get_optimal_batch_size(
        device=mgr.get_device(),
        num_documents=2_000_000,
        n_gpus=mgr.n_gpus,
    )
    print(f"Recommended per-GPU batch size: {bs}")