File size: 4,550 Bytes
15d68eb | 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 142 143 144 145 146 147 | """
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
|