File size: 5,626 Bytes
8ae78b0 | 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 156 157 158 159 160 161 | import torch
import time
import functools
import logging
import os
import psutil
import gc
try:
from app.utils.logging_utils import setup_logger
except ImportError:
# Try relative imports for running from project root
from behavior_backend.app.utils.logging_utils import setup_logger
# Configure logging
logger = setup_logger(__name__)
def get_system_memory_info():
"""
Get system memory information.
Returns:
dict: Memory information
"""
memory = psutil.virtual_memory()
return {
"total": memory.total / (1024 ** 3), # GB
"available": memory.available / (1024 ** 3), # GB
"percent_used": memory.percent,
"process_usage": psutil.Process(os.getpid()).memory_info().rss / (1024 ** 3) # GB
}
def log_memory_usage(message=""):
"""
Log current memory usage.
Args:
message: Optional message to include in the log
"""
mem_info = get_system_memory_info()
logger.info(f"Memory usage {message}: "
f"Total: {mem_info['total']:.2f}GB, "
f"Available: {mem_info['available']:.2f}GB, "
f"Used: {mem_info['percent_used']}%, "
f"Process: {mem_info['process_usage']:.2f}GB")
def get_available_device():
"""
Determine the best available device with proper error handling.
Returns:
str: 'cuda', 'mps', or 'cpu' depending on availability
"""
logger.info("=== GPU DETECTION ===")
# Check available memory first
mem_info = get_system_memory_info()
if mem_info['available'] < 2.0: # Less than 2GB available
logger.warning(f"Low system memory: {mem_info['available']:.2f}GB available. Forcing CPU usage.")
return "cpu"
# First try CUDA (NVIDIA GPUs)
if torch.cuda.is_available():
try:
# Simplified CUDA test with better error handling
logger.info("CUDA detected - attempting verification")
# Use a smaller and simpler operation
test_tensor = torch.tensor([1.0], device="cuda")
test_tensor = test_tensor + 1.0 # Simple operation
result = test_tensor.item() # Get the value back to validate operation
# If we get here, the CUDA operation worked
test_tensor = test_tensor.cpu() # Move back to CPU to free CUDA memory
torch.cuda.empty_cache() # Clear CUDA cache
logger.info(f" NVIDIA GPU (CUDA) detected and verified working (test result: {result})")
return "cuda"
except Exception as e:
logger.warning(f"CUDA detected but test failed: {e}")
torch.cuda.empty_cache() # Clear CUDA cache
# Then try MPS (Apple Silicon)
if hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
try:
# Test MPS with a small operation
test_tensor = torch.zeros(1).to('mps')
test_tensor = test_tensor + 1
test_tensor.cpu() # Move back to CPU to free MPS memory
logger.info(" Apple Silicon GPU (MPS) detected and verified working")
return "mps"
except Exception as e:
logger.warning(f" MPS detected but test failed: {e}")
# Fall back to CPU
logger.info(" No GPU detected or all GPU tests failed, using CPU")
return "cpu"
def run_on_device(func):
"""
Decorator to run a function on the best available device.
Args:
func: The function to decorate
Returns:
A wrapped function that runs on the best available device
"""
@functools.wraps(func)
def wrapper(*args, **kwargs):
# Log memory before operation
log_memory_usage(f"before {func.__name__}")
# Force garbage collection before operation
gc.collect()
# Get device if not already specified
device = get_available_device()
# Add device to kwargs if not already present
if 'device' not in kwargs:
kwargs['device'] = device
try:
start_time = time.time()
result = func(*args, **kwargs)
end_time = time.time()
logger.debug(f"Function {func.__name__} ran on {device} in {end_time - start_time:.4f} seconds")
return result
except Exception as e:
# Check if this is the SparseMPS error
if "SparseMPS" in str(e) and device == "mps":
logger.warning(f"MPS error detected: {e}")
logger.warning("Falling back to CPU for this operation")
# Update device to CPU and retry
kwargs['device'] = 'cpu'
# Force garbage collection before retry
gc.collect()
start_time = time.time()
result = func(*args, **kwargs)
end_time = time.time()
logger.debug(f"Function {func.__name__} ran on CPU (fallback) in {end_time - start_time:.4f} seconds")
return result
else:
# Re-raise other exceptions
raise
finally:
# Force garbage collection after operation
gc.collect()
if device == 'cuda':
torch.cuda.empty_cache()
# Log memory after operation
log_memory_usage(f"after {func.__name__}")
return wrapper
# Initialize device once at module level
device = get_available_device() |