9c3d741a08
Move 17 library crates into crates/, CLI binary into bin/fxt, consolidate 10 test crates into testing/, split config crate from deployment config files. Root directory reduced from 38+ to ~17 directories. All Cargo.toml paths and build.rs proto refs updated. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
100 lines
2.8 KiB
Rust
100 lines
2.8 KiB
Rust
//! 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-detect: CUDA device 0 if available, else CPU (silent fallback -- never errors).
|
|
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<Device, MLError> {
|
|
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);
|
|
}
|
|
}
|
|
}
|