| 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: |