diff --git a/ml/src/gpu/capabilities.rs b/ml/src/gpu/capabilities.rs new file mode 100644 index 000000000..e69de29bb diff --git a/ml/src/gpu/memory_profile.rs b/ml/src/gpu/memory_profile.rs new file mode 100644 index 000000000..e69de29bb diff --git a/ml/src/gpu/mod.rs b/ml/src/gpu/mod.rs new file mode 100644 index 000000000..7d62b799c --- /dev/null +++ b/ml/src/gpu/mod.rs @@ -0,0 +1,98 @@ +//! GPU device management for ML training +//! +//! Provides centralized device selection, GPU capability detection, +//! and per-model batch size optimization. + +pub mod capabilities; +pub mod memory_profile; + +use candle_core::Device; +use serde::{Deserialize, Serialize}; + +use crate::MLError; + +/// Device configuration for ML training and inference. +/// +/// Use `Auto` for production — it detects CUDA if available, falls back to CPU. +/// Use `Cuda(id)` to pin to a specific GPU. Use `Cpu` to force CPU mode. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub enum DeviceConfig { + Cpu, + Cuda(usize), + Auto, +} + +impl DeviceConfig { + /// Resolve this config into a concrete candle `Device`. + /// + /// - `Cpu` → always `Device::Cpu` + /// - `Cuda(id)` → `Device::new_cuda(id)`, errors if unavailable + /// - `Auto` → CUDA device 0 if available, else CPU + pub fn resolve(&self) -> Result { + match self { + DeviceConfig::Cpu => Ok(Device::Cpu), + DeviceConfig::Cuda(id) => Device::new_cuda(*id).map_err(|e| { + MLError::ConfigurationError(format!( + "CUDA device {} required but unavailable: {}", + id, e + )) + }), + DeviceConfig::Auto => { + match Device::new_cuda(0) { + Ok(dev) => Ok(dev), + Err(_) => Ok(Device::Cpu), + } + } + } + } + + /// Returns true if this config will attempt to use a GPU. + pub fn is_gpu(&self) -> bool { + matches!(self, DeviceConfig::Cuda(_) | DeviceConfig::Auto) + } +} + +impl Default for DeviceConfig { + fn default() -> Self { + DeviceConfig::Auto + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_cpu_always_resolves() { + let device = DeviceConfig::Cpu.resolve().unwrap(); + assert!(!device.is_cuda()); + } + + #[test] + fn test_auto_resolves_without_error() { + let device = DeviceConfig::Auto.resolve().unwrap(); + let _ = device; + } + + #[test] + fn test_is_gpu() { + assert!(!DeviceConfig::Cpu.is_gpu()); + assert!(DeviceConfig::Cuda(0).is_gpu()); + assert!(DeviceConfig::Auto.is_gpu()); + } + + #[test] + fn test_default_is_auto() { + assert_eq!(DeviceConfig::default(), DeviceConfig::Auto); + } + + #[test] + fn test_serde_roundtrip() { + let configs = vec![DeviceConfig::Cpu, DeviceConfig::Cuda(0), DeviceConfig::Auto]; + for config in configs { + let json = serde_json::to_string(&config).unwrap(); + let deserialized: DeviceConfig = serde_json::from_str(&json).unwrap(); + assert_eq!(config, deserialized); + } + } +} diff --git a/ml/src/lib.rs b/ml/src/lib.rs index 761f95720..bcf2f460d 100644 --- a/ml/src/lib.rs +++ b/ml/src/lib.rs @@ -687,6 +687,7 @@ pub mod data_loaders; // Data loaders for ML training pub mod dqn; pub mod gradient_accumulation; // Gradient accumulation for mini-batch training pub mod gradient_utils; // Gradient utilities (clipping, monitoring) +pub mod gpu; pub mod diffusion; pub mod ensemble; pub mod evaluation; // DQN evaluation engine (backtest metrics, Sharpe ratio)