Spaces:
Runtime error
Runtime error
File size: 2,746 Bytes
102f4d2 | 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 | """AMD ROCm device management and monitoring."""
import os
import subprocess
from loguru import logger
def get_device() -> str:
"""Return 'cuda' (ROCm uses cuda device name in PyTorch) or 'cpu'."""
try:
import torch
if torch.cuda.is_available():
device_name = torch.cuda.get_device_name(0)
logger.info(f"GPU detected: {device_name}")
return "cuda"
except ImportError:
pass
logger.warning("No GPU available, falling back to CPU")
return "cpu"
def get_vram_gb() -> float:
"""Return available VRAM in GB."""
try:
import torch
if torch.cuda.is_available():
total = torch.cuda.get_device_properties(0).total_memory
return round(total / 1024**3, 1)
except Exception:
pass
return 0.0
def get_gpu_utilization() -> dict:
"""Return GPU utilization stats via rocm-smi."""
try:
result = subprocess.run(
["rocm-smi", "--showuse", "--showmemuse", "--csv"],
capture_output=True, text=True, timeout=5
)
if result.returncode == 0:
lines = result.stdout.strip().split("\n")
if len(lines) >= 2:
headers = lines[0].split(",")
values = lines[1].split(",")
return dict(zip(headers, values))
except (FileNotFoundError, subprocess.TimeoutExpired):
pass
# Fallback: PyTorch memory stats
try:
import torch
if torch.cuda.is_available():
allocated = torch.cuda.memory_allocated(0) / 1024**3
reserved = torch.cuda.memory_reserved(0) / 1024**3
total = torch.cuda.get_device_properties(0).total_memory / 1024**3
return {
"vram_used_gb": round(allocated, 2),
"vram_reserved_gb": round(reserved, 2),
"vram_total_gb": round(total, 2),
"vram_pct": round(allocated / total * 100, 1) if total > 0 else 0,
}
except Exception:
pass
return {}
def get_optimal_batch_size(model_type: str = "whisper") -> int:
"""Return optimal batch size based on available VRAM."""
vram = get_vram_gb()
if model_type == "whisper":
if vram >= 48:
return 32
elif vram >= 24:
return 16
elif vram >= 16:
return 8
return 4
elif model_type == "vision":
if vram >= 80:
return 8
elif vram >= 48:
return 4
return 1
return 1
def log_gpu_status():
stats = get_gpu_utilization()
if stats:
logger.info(f"GPU stats: {stats}")
else:
logger.info(f"GPU: {get_device()} | VRAM: {get_vram_gb()} GB")
|