| 957 | return cls._instance |
| 958 | |
| 959 | class DeviceManager: |
| 960 | def __init__(self): |
| 961 | self.available_devices = self._detect_devices() |
| 962 | self.device_stats = {device: {'memory_used': 0, 'active_models': 0} for device in self.available_devices} |
| 963 | |
| 964 | def _detect_devices(self) -> List[str]: |
| 965 | devices = ['cpu'] |
| 966 | if torch.cuda.is_available(): |
| 967 | devices.extend([f'cuda:{i}' for i in range(torch.cuda.device_count())]) |
| 968 | if torch.backends.mps.is_available(): |
| 969 | devices.append('mps') |
| 970 | return devices |
| 971 | |
| 972 | def get_optimal_device(self, model_size: int = 0) -> str: |
| 973 | if not self.available_devices: |
| 974 | return 'cpu' |
| 975 | |
| 976 | # Prefer CUDA devices if available |
| 977 | cuda_devices = [d for d in self.available_devices if 'cuda' in d] |
| 978 | if cuda_devices: |
| 979 | # Find CUDA device with most free memory |
| 980 | max_free_memory = 0 |
| 981 | optimal_device = cuda_devices[0] |
| 982 | |
| 983 | for device in cuda_devices: |
| 984 | idx = int(device.split(':')[1]) |
| 985 | free_memory = torch.cuda.get_device_properties(idx).total_memory - torch.cuda.memory_allocated(idx) |
| 986 | if free_memory > max_free_memory: |
| 987 | max_free_memory = free_memory |
| 988 | optimal_device = device |
| 989 | |
| 990 | return optimal_device |
| 991 | |
| 992 | # Fall back to MPS if available |
| 993 | if 'mps' in self.available_devices: |
| 994 | return 'mps' |
| 995 | |
| 996 | return 'cpu' |
| 997 | |
| 998 | def track_device_usage(self, device: str, memory_delta: int): |
| 999 | if device in self.device_stats: |
| 1000 | self.device_stats[device]['memory_used'] += memory_delta |
| 1001 | |
| 1002 | class ModelManager: |
| 1003 | def __init__(self, cache_manager: CacheManager, device_manager: DeviceManager): |