"""
Backend manager for selecting and configuring computational backends.
This module provides utilities for automatically detecting available backends
and selecting the appropriate one based on hardware and user preferences.
"""
import warnings
from typing import Optional
from .base import Backend
from .numpy_backend import NumPyBackend
# Try to import CuPy backend
try:
from .cupy_backend import CUPY_AVAILABLE, CuPyBackend
except ImportError:
CUPY_AVAILABLE = False
CuPyBackend = None
# Try to import Metal backend
try:
from .metal_backend import METAL_AVAILABLE, MetalBackend
except ImportError:
METAL_AVAILABLE = False
MetalBackend = None
# Global backend instance
_CURRENT_BACKEND: Optional[Backend] = None
[docs]
def list_available_backends() -> list[str]:
"""
List all available backends on this system.
Returns
-------
List[str]
List of backend names (e.g., ['numpy', 'cupy', 'metal']).
"""
available = ["numpy"] # NumPy is always available
if CUPY_AVAILABLE:
try:
# Try to initialize CuPy to verify CUDA is working
import cupy as cp
cp.cuda.Device(0).use()
available.append("cupy")
except Exception:
pass
if METAL_AVAILABLE:
try:
# Try to initialize Metal to verify it's working
import platform
import Metal
if platform.system() == "Darwin":
devices = Metal.MTLCopyAllDevices()
if len(devices) > 0:
available.append("metal")
except Exception:
pass
return available
def get_backend_info() -> dict[str, any]:
"""
Get information about available backends and their capabilities.
Returns
-------
Dict[str, any]
Dictionary with backend information.
"""
info = {
"available_backends": list_available_backends(),
"current_backend": _CURRENT_BACKEND.name if _CURRENT_BACKEND else None,
"numpy_available": True,
"cupy_available": CUPY_AVAILABLE,
"metal_available": METAL_AVAILABLE,
}
if CUPY_AVAILABLE:
try:
import cupy as cp
info["cuda_version"] = cp.cuda.runtime.runtimeGetVersion()
info["num_devices"] = cp.cuda.runtime.getDeviceCount()
# Get info for each device
devices = []
for i in range(info["num_devices"]):
device = cp.cuda.Device(i)
try:
device_name = device.name
except AttributeError:
device_name = f"GPU:{i}"
devices.append(
{
"id": i,
"name": device_name,
"compute_capability": device.compute_capability,
"total_memory_mb": device.mem_info[1] / (1024**2),
}
)
info["cuda_devices"] = devices
except Exception as e:
info["cuda_error"] = str(e)
if METAL_AVAILABLE:
try:
import platform
import Metal
if platform.system() == "Darwin":
devices = Metal.MTLCopyAllDevices()
info["num_metal_devices"] = len(devices)
# Get info for each Metal device
metal_devices = []
for i, device in enumerate(devices):
metal_devices.append(
{
"id": i,
"name": device.name(),
"max_buffer_size": device.maxBufferLength(),
"recommended_max_working_set_size": device.recommendedMaxWorkingSetSize(),
"is_low_power": device.isLowPower(),
"is_headless": device.isHeadless(),
}
)
info["metal_devices"] = metal_devices
except Exception as e:
info["metal_error"] = str(e)
return info
[docs]
def set_backend(backend: str, device_id: int = 0) -> Backend:
"""
Set the global backend for computations.
Parameters
----------
backend : str
Backend name ('numpy', 'cupy', or 'metal').
device_id : int, optional
Device ID for GPU backends. Default is 0.
Returns
-------
Backend
The initialized backend instance.
Raises
------
ValueError
If the requested backend is not available.
"""
global _CURRENT_BACKEND
backend = backend.lower()
if backend == "numpy":
_CURRENT_BACKEND = NumPyBackend()
elif backend == "cupy":
if not CUPY_AVAILABLE:
raise ValueError(
"CuPy backend requested but CuPy is not available. "
"Install with: pip install cupy-cuda12x"
)
_CURRENT_BACKEND = CuPyBackend(device_id=device_id)
elif backend == "metal":
if not METAL_AVAILABLE:
raise ValueError(
"Metal backend requested but Metal is not available. "
"Metal backend requires macOS with Metal framework."
)
_CURRENT_BACKEND = MetalBackend(device_id=device_id)
else:
available = list_available_backends()
raise ValueError(
f"Unknown backend '{backend}'. Available backends: {available}"
)
return _CURRENT_BACKEND
[docs]
def get_backend(backend: Optional[str] = None, device_id: int = 0) -> Backend:
"""
Get a backend instance.
If no backend is specified, returns the current global backend or
automatically selects the best available backend (GPU preferred).
Parameters
----------
backend : str, optional
Backend name ('numpy', 'cupy', or 'metal'). If None, uses current or auto-detects.
device_id : int, optional
Device ID for GPU backends. Default is 0.
Returns
-------
Backend
The backend instance.
"""
global _CURRENT_BACKEND
# If specific backend requested, set and return it
if backend is not None:
return set_backend(backend, device_id)
# If we have a current backend, return it
if _CURRENT_BACKEND is not None:
return _CURRENT_BACKEND
# Auto-detect best backend
available = list_available_backends()
# Prefer GPU if available (Metal on macOS, CUDA otherwise)
import platform
if platform.system() == "Darwin" and "metal" in available:
warnings.warn(
"No backend specified. Auto-selecting Metal (GPU) backend. "
"Set explicitly with set_backend() or get_backend(backend='numpy')",
UserWarning,
stacklevel=2,
)
return set_backend("metal", device_id)
elif "cupy" in available:
warnings.warn(
"No backend specified. Auto-selecting CuPy (GPU) backend. "
"Set explicitly with set_backend() or get_backend(backend='numpy')",
UserWarning,
stacklevel=2,
)
return set_backend("cupy", device_id)
else:
# Fall back to NumPy
return set_backend("numpy")
def auto_select_backend() -> Backend:
"""
Automatically select the best available backend.
Prefers GPU (Metal on macOS, CuPy otherwise) if available, otherwise uses CPU (NumPy).
Returns
-------
Backend
The selected backend instance.
"""
return get_backend()
# Initialize with NumPy backend by default
_CURRENT_BACKEND = NumPyBackend()