dd62f3fcfd
Final cleanup: - 61 test files + 5 example files: candle imports replaced - 8 testing/integration files: migrated to cudarc/ml-core types - 3 services/trading_service test files: migrated - Root Cargo.toml: candle-core, candle-nn removed from [workspace.dependencies] - crates/ml/Cargo.toml: candle-nn dependency removed - testing/e2e/Cargo.toml: candle-core dependency removed Zero active candle_core/candle_nn/candle_optimisers code references remain. Zero candle dependency declarations in any Cargo.toml. Remaining "candle" strings are exclusively in doc comments. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
105 lines
3.2 KiB
Rust
105 lines
3.2 KiB
Rust
//! Device abstraction replacing `candle_core::Device`.
|
|
//!
|
|
//! Provides `MlDevice` enum with CPU and CUDA variants.
|
|
//! The CUDA variant holds `Arc<CudaDevice>` and `Arc<CudaStream>` for
|
|
//! direct cudarc interop.
|
|
|
|
use std::sync::Arc;
|
|
|
|
use crate::MLError;
|
|
|
|
#[cfg(feature = "cuda")]
|
|
use cudarc::driver::{CudaContext, CudaStream};
|
|
|
|
/// Device abstraction for the ML pipeline.
|
|
///
|
|
/// Replaces `candle_core::Device`. Only two variants: CPU (for checkpoint
|
|
/// serialization and preprocessing) and CUDA (for training and inference).
|
|
#[derive(Clone)]
|
|
pub enum MlDevice {
|
|
/// Host CPU — used for checkpoint I/O and lightweight preprocessing.
|
|
Cpu,
|
|
/// CUDA GPU with context + stream handles.
|
|
#[cfg(feature = "cuda")]
|
|
Cuda {
|
|
context: Arc<CudaContext>,
|
|
stream: Arc<CudaStream>,
|
|
},
|
|
}
|
|
|
|
impl std::fmt::Debug for MlDevice {
|
|
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
|
match self {
|
|
MlDevice::Cpu => write!(f, "MlDevice::Cpu"),
|
|
#[cfg(feature = "cuda")]
|
|
MlDevice::Cuda { .. } => write!(f, "MlDevice::Cuda"),
|
|
}
|
|
}
|
|
}
|
|
|
|
impl MlDevice {
|
|
/// Create a CUDA device with the given ordinal (0-indexed).
|
|
///
|
|
/// Initializes a `CudaDevice` and forks a non-default stream for all
|
|
/// subsequent operations.
|
|
#[cfg(feature = "cuda")]
|
|
pub fn cuda(ordinal: usize) -> Result<Self, MLError> {
|
|
let context = CudaContext::new(ordinal).map_err(|e| {
|
|
MLError::DeviceError(format!("Failed to open CUDA device {ordinal}: {e}"))
|
|
})?;
|
|
let stream = context.new_stream().map_err(|e| {
|
|
MLError::DeviceError(format!("Failed to create CUDA stream on device {ordinal}: {e}"))
|
|
})?;
|
|
Ok(MlDevice::Cuda {
|
|
context,
|
|
stream,
|
|
})
|
|
}
|
|
|
|
/// Returns `true` if this is a CUDA device.
|
|
pub fn is_cuda(&self) -> bool {
|
|
match self {
|
|
MlDevice::Cpu => false,
|
|
#[cfg(feature = "cuda")]
|
|
MlDevice::Cuda { .. } => true,
|
|
}
|
|
}
|
|
|
|
/// Returns `true` if this is the CPU device.
|
|
pub fn is_cpu(&self) -> bool {
|
|
matches!(self, MlDevice::Cpu)
|
|
}
|
|
|
|
/// Get the CUDA stream, or error if this is a CPU device.
|
|
#[cfg(feature = "cuda")]
|
|
pub fn cuda_stream(&self) -> Result<&Arc<CudaStream>, MLError> {
|
|
match self {
|
|
MlDevice::Cuda { stream, .. } => Ok(stream),
|
|
MlDevice::Cpu => Err(MLError::DeviceError(
|
|
"cuda_stream() called on CPU device".to_owned(),
|
|
)),
|
|
}
|
|
}
|
|
|
|
/// Get the CudaContext handle, or error if this is a CPU device.
|
|
#[cfg(feature = "cuda")]
|
|
pub fn cuda_context(&self) -> Result<&Arc<CudaContext>, MLError> {
|
|
match self {
|
|
MlDevice::Cuda { context, .. } => Ok(context),
|
|
MlDevice::Cpu => Err(MLError::DeviceError(
|
|
"cuda_context() called on CPU device".to_owned(),
|
|
)),
|
|
}
|
|
}
|
|
}
|
|
|
|
impl std::fmt::Display for MlDevice {
|
|
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
|
match self {
|
|
MlDevice::Cpu => write!(f, "cpu"),
|
|
#[cfg(feature = "cuda")]
|
|
MlDevice::Cuda { .. } => write!(f, "cuda"),
|
|
}
|
|
}
|
|
}
|