Spaces:
Running on Zero
Running on Zero
File size: 4,931 Bytes
36333c5 d63e724 36333c5 d63e724 36333c5 eb808a5 36333c5 eb808a5 36333c5 993f563 d63e724 36333c5 d63e724 36333c5 993f563 36333c5 | 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 148 149 150 151 152 153 154 155 | """Accelerator selection, inference contexts, and memory cleanup."""
from __future__ import annotations
import gc
import logging
from contextlib import ExitStack, contextmanager
from typing import Any, Iterator
from config import Settings
logger = logging.getLogger(__name__)
def import_torch() -> Any:
"""Import torch lazily so startup never loads a CUDA context."""
try:
import torch
except ImportError as exc: # pragma: no cover - required in deployment
raise RuntimeError("PyTorch is not installed") from exc
return torch
def detect_device(configured: str = "auto") -> str:
"""Select the requested accelerator, falling back safely to CPU."""
try:
torch = import_torch()
except RuntimeError:
if configured in {"auto", "cpu"}:
return "cpu"
raise
if configured != "auto":
if configured.startswith("cuda"):
if not torch.cuda.is_available():
return "cpu"
if ":" in configured:
device_index = int(configured.split(":", 1)[1])
if device_index >= torch.cuda.device_count():
return "cpu"
if configured == "mps":
mps = getattr(torch.backends, "mps", None)
return "mps" if mps is not None and mps.is_available() else "cpu"
return configured
if torch.cuda.is_available():
return "cuda"
mps = getattr(torch.backends, "mps", None)
if mps is not None and mps.is_available():
return "mps"
return "cpu"
def preferred_dtype(device: str, mixed_precision: bool = True) -> Any:
"""Choose an inference dtype supported by the selected device."""
torch = import_torch()
if not mixed_precision or device == "cpu":
return torch.float32
if device.startswith("cuda"):
return torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
return torch.float16 if device == "mps" else torch.float32
def gpu_available() -> bool:
"""Return scheduled ZeroGPU or directly attached CUDA availability."""
if zerogpu_enabled():
return True
try:
torch = import_torch()
except RuntimeError:
return False
return bool(torch.cuda.is_available())
def zerogpu_enabled() -> bool:
"""Return whether this process is running in a Hugging Face ZeroGPU Space."""
try:
from spaces.config import Config
except ImportError:
return False
return bool(Config.zero_gpu)
def initialize_zerogpu() -> None:
"""Run the scheduler startup hook when Gradio is mounted into FastAPI."""
if not zerogpu_enabled():
return
from spaces.zero import startup
startup()
@contextmanager
def inference_context(settings: Settings, device: str) -> Iterator[None]:
"""Disable autograd and enable CUDA mixed precision when configured."""
torch = import_torch()
with ExitStack() as stack:
stack.enter_context(torch.no_grad())
stack.enter_context(torch.inference_mode())
if settings.mixed_precision and device.startswith("cuda"):
stack.enter_context(
torch.autocast(device_type="cuda", dtype=preferred_dtype(device, True))
)
yield
def cleanup_memory() -> None:
"""Collect Python garbage and return cached accelerator memory."""
gc.collect()
try:
torch = import_torch()
except RuntimeError:
return
# Cache cleanup must never initialize CUDA. ZeroGPU deliberately reports CUDA
# as available in the parent process while rejecting low-level initialization;
# a real CUDA context exists only inside the scheduled worker.
try:
if torch.cuda.is_initialized():
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
except Exception:
logger.debug("CUDA cache cleanup skipped", exc_info=True)
mps_backend = getattr(getattr(torch, "backends", None), "mps", None)
mps = getattr(torch, "mps", None)
try:
if (
mps_backend is not None
and mps_backend.is_available()
and mps is not None
and hasattr(mps, "empty_cache")
):
mps.empty_cache()
except Exception:
logger.debug("MPS cache cleanup skipped", exc_info=True)
def memory_stats() -> dict[str, float]:
"""Return process and CUDA memory consumption in MiB."""
stats: dict[str, float] = {}
try:
import psutil
stats["rss_mb"] = round(psutil.Process().memory_info().rss / 1024**2, 2)
except ImportError:
pass
try:
torch = import_torch()
except RuntimeError:
return stats
if torch.cuda.is_initialized():
stats.update(
cuda_allocated_mb=round(torch.cuda.memory_allocated() / 1024**2, 2),
cuda_reserved_mb=round(torch.cuda.memory_reserved() / 1024**2, 2),
)
return stats
|