utils/gpu_utils.py

"""
Multi-GPU device management for Indic Heritage Studio v2.

Provides:
  - GPUInfo: per-device info (name, VRAM, free/used)
  - list_gpus(): snapshot of all available GPUs
  - assign_pipeline_device(): round-robin or pinned assignment
  - VRAMGuard: context manager that warns on low VRAM
  - shard_workload(): split a list of items across N GPUs for batch jobs
"""
from __future__ import annotations

import logging
from contextlib import contextmanager
from dataclasses import dataclass
from typing import Iterable, List, Sequence

log = logging.getLogger(__name__)


@dataclass
class GPUInfo:
    index: int
    name: str
    vram_total_gb: float
    vram_free_gb: float
    vram_used_gb: float

    @property
    def free_ratio(self) -> float:
        return self.vram_free_gb / max(self.vram_total_gb, 1e-6)


def list_gpus() -> List[GPUInfo]:
    """Return a snapshot of every available CUDA/ROCm device."""
    try:
        import torch
        if not torch.cuda.is_available():
            return []
        gpus = []
        for i in range(torch.cuda.device_count()):
            free, total = torch.cuda.mem_get_info(i)
            used = total - free
            gpus.append(GPUInfo(
                index=i,
                name=torch.cuda.get_device_name(i),
                vram_total_gb=round(total / 1e9, 2),
                vram_free_gb=round(free / 1e9, 2),
                vram_used_gb=round(used / 1e9, 2),
            ))
        return gpus
    except Exception as exc:
        log.warning("list_gpus failed: %s", exc)
        return []


def shard_workload(items: Sequence, n_shards: int) -> List[List]:
    """Split a sequence into n_shards contiguous chunks (last shard gets remainder)."""
    if n_shards <= 0:
        return [list(items)]
    n = len(items)
    base = n // n_shards
    extra = n % n_shards
    shards = []
    start = 0
    for i in range(n_shards):
        size = base + (1 if i < extra else 0)
        shards.append(list(items[start:start + size]))
        start += size
    return shards


@contextmanager
def VRAMGuard(device: str, min_free_gb: float = 2.0, label: str = "pipeline"):
    """Warn if VRAM on a device drops below min_free_gb during the block.

    Properly propagates exceptions from the wrapped block (unlike the
    previous version which raised 'generator didn't stop after throw()').
    """
    import torch
    try:
        idx = int(device.split(":")[-1]) if ":" in device else 0
    except Exception:
        idx = 0

    try:
        free_before, _ = torch.cuda.mem_get_info(idx)
    except Exception:
        free_before = 0

    try:
        yield
    finally:
        # Always check VRAM, even if the wrapped block raised
        try:
            free_after, _ = torch.cuda.mem_get_info(idx)
            if free_before > 0:
                leaked_gb = (free_before - free_after) / 1e9
                if leaked_gb > 1.0:
                    log.warning(
                        "[%s] %.2f GB VRAM leaked on %s (free %.2f → %.2f GB). "
                        "Possible missing cleanup.",
                        label, leaked_gb, device,
                        free_before / 1e9, free_after / 1e9,
                    )
                if free_after / 1e9 < min_free_gb:
                    log.warning(
                        "[%s] Low VRAM on %s: only %.2f GB free",
                        label, device, free_after / 1e9,
                    )
        except Exception as exc:
            log.debug("VRAMGuard post-check skipped: %s", exc)


def assign_pipeline_device(pipeline_name: str) -> str:
    """Return the cuda device assigned to a named pipeline."""
    from config.settings import settings
    return settings.get_pipeline_device(pipeline_name)


def get_batch_worker_devices() -> List[str]:
    """Devices available for batch parallel workers."""
    from config.settings import settings
    return settings.get_batch_worker_devices()


def move_model_to_device(model, device: str):
    """Move a model to a device and return it (no-op if already there)."""
    try:
        return model.to(device)
    except Exception as exc:
        log.warning("Failed to move model to %s: %s", device, exc)
        return model


def empty_cache_all() -> None:
    """Empty cache across all visible GPUs."""
    try:
        import torch
        if not torch.cuda.is_available():
            return
        for i in range(torch.cuda.device_count()):
            with torch.cuda.device(i):
                torch.cuda.empty_cache()
    except Exception:
        pass