diff --git a/crates/ml/src/ensemble/mod.rs b/crates/ml/src/ensemble/mod.rs index 37c64d559..88630f784 100644 --- a/crates/ml/src/ensemble/mod.rs +++ b/crates/ml/src/ensemble/mod.rs @@ -26,6 +26,7 @@ pub mod weight_optimizer; pub mod gate_optimizer; pub mod model_adapter; pub mod cuda_streams; +pub mod stream_ensemble; // Re-export key types that are used across ensemble modules pub use model_adapter::{EnsembleModelAdapter, build_production_strategy}; diff --git a/crates/ml/src/ensemble/stream_ensemble.rs b/crates/ml/src/ensemble/stream_ensemble.rs new file mode 100644 index 000000000..6a48673b0 --- /dev/null +++ b/crates/ml/src/ensemble/stream_ensemble.rs @@ -0,0 +1,679 @@ +//! Stream-aware ensemble with per-model CUDA device handles. +//! +//! Each model runs on its own CUDA stream for true GPU-level parallelism. +//! Falls back to standard rayon parallelism on CPU. +//! +//! # Design (Option C — per-model CudaDevice handles) +//! +//! On CUDA: each model adapter gets its own `CudaDevice::new(0)` handle so the +//! CUDA driver assigns a unique default stream. `predict()` runs all models via +//! rayon with per-device streams, syncs all, then aggregates on the primary device. +//! +//! On CPU: all adapters share the same device, rayon provides parallelism (same +//! behaviour as [`InferenceEnsemble`](super::inference_ensemble::InferenceEnsemble)). + +use rayon::prelude::*; + +use candle_core::Tensor; + +use crate::ensemble::cuda_streams::CudaStreamPool; +use crate::ensemble::inference_adapter::{ + EnsemblePrediction, FeatureVector, ModelInferenceAdapter, PredictionMeta, RawPrediction, +}; +use crate::{MLError, MLResult}; + +/// Ensemble that runs models on separate CUDA streams for true +/// GPU-level parallelism. Falls back to rayon on CPU. +/// +/// Uses [`CudaStreamPool`] for stream management and synchronisation. +#[allow(missing_debug_implementations)] +pub struct StreamAwareEnsemble { + /// Model adapters (each may have its own device handle) + adapters: Vec>, + /// Weights per model (normalised, indexed by adapter position) + weights: Vec, + /// Stream pool for CUDA synchronisation (no-op on CPU) + stream_pool: CudaStreamPool, +} + +impl StreamAwareEnsemble { + /// Create a new stream-aware ensemble. + /// + /// `weights` are normalised so they sum to 1.0. If the provided weights + /// slice is shorter than `adapters`, missing entries default to equal + /// share. If `weights` is empty, all models get equal weight. + pub fn new( + adapters: Vec>, + weights: Vec, + device: &candle_core::Device, + ) -> MLResult { + let n = adapters.len(); + let stream_pool = CudaStreamPool::new(device, n)?; + + let raw_weights = if weights.is_empty() || n == 0 { + vec![1.0_f64; n] + } else { + let mut w = weights; + w.resize(n, 1.0); + w + }; + + let weight_sum: f64 = raw_weights.iter().sum(); + let normalized = if weight_sum > 0.0 { + raw_weights.iter().map(|w| w / weight_sum).collect() + } else { + vec![1.0 / n.max(1) as f64; n] + }; + + Ok(Self { + adapters, + weights: normalized, + stream_pool, + }) + } + + /// Run ensemble prediction with stream-level parallelism. + /// + /// 1. Runs all ready adapters in parallel via rayon (each on its own + /// CUDA stream if GPU, or rayon thread if CPU). + /// 2. Synchronises all CUDA streams (no-op on CPU). + /// 3. Aggregates predictions via confidence-weighted voting with + /// GPU-tensor and CPU-scalar paths (same algorithm as + /// [`InferenceEnsemble`]). + pub fn predict(&self, features: &FeatureVector) -> MLResult { + // Parallel inference via rayon — each adapter on its own stream if CUDA + let predictions: Vec<(usize, String, RawPrediction)> = self + .adapters + .par_iter() + .enumerate() + .filter(|(_, a)| a.is_ready()) + .filter_map(|(i, adapter)| { + let name = adapter.model_name().to_string(); + match adapter.predict_raw(features) { + Ok(pred) => { + if !pred.direction_scalar.is_finite() || !pred.confidence.is_finite() { + tracing::warn!( + model = %name, + direction = %pred.direction_scalar, + confidence = %pred.confidence, + "StreamAwareEnsemble: model returned NaN/Inf, skipping" + ); + None + } else { + Some((i, name, pred)) + } + } + Err(e) => { + tracing::warn!( + model = %name, + error = %e, + "StreamAwareEnsemble: model prediction failed, skipping" + ); + None + } + } + }) + .collect(); + + if predictions.is_empty() { + return Ok(EnsemblePrediction { + model_name: "STREAM_ENSEMBLE(empty)".to_owned(), + direction: 0.0, + confidence: 0.0, + metadata: PredictionMeta::default(), + }); + } + + // Synchronise all CUDA streams before aggregation (no-op on CPU) + self.stream_pool.sync_all()?; + + // Partition into GPU-tensor vs CPU-scalar predictions + let (gpu_preds, cpu_preds): (Vec<_>, Vec<_>) = predictions + .into_iter() + .partition(|(_, _, p)| p.tensor.is_some()); + + let mut model_names: Vec = Vec::new(); + let mut total_confidence_sum = 0.0_f64; + let mut total_count: usize = 0; + + // --- GPU path: stack tensors, sigmoid, weighted-sum, extract scalar --- + let gpu_result = if !gpu_preds.is_empty() { + match self.aggregate_gpu(&gpu_preds) { + Ok((direction, count)) => { + for (_, name, pred) in &gpu_preds { + model_names.push(name.clone()); + total_confidence_sum += pred.confidence.clamp(0.0, 1.0); + } + total_count += gpu_preds.len(); + Some((direction, count)) + } + Err(e) => { + tracing::warn!( + error = %e, + "GPU aggregation failed, falling back to CPU for {} models", + gpu_preds.len() + ); + let fallback = self.aggregate_cpu(&gpu_preds); + for (_, name, pred) in &gpu_preds { + model_names.push(name.clone()); + total_confidence_sum += pred.confidence.clamp(0.0, 1.0); + } + total_count += gpu_preds.len(); + fallback + } + } + } else { + None + }; + + // --- CPU path: scalar weighted average --- + let cpu_result = if !cpu_preds.is_empty() { + let result = self.aggregate_cpu(&cpu_preds); + for (_, name, pred) in &cpu_preds { + model_names.push(name.clone()); + total_confidence_sum += pred.confidence.clamp(0.0, 1.0); + } + total_count += cpu_preds.len(); + result + } else { + None + }; + + // Merge GPU and CPU directions by model-count-weighted average + let direction = match (gpu_result, cpu_result) { + (Some((gpu_dir, gpu_n)), Some((cpu_dir, cpu_n))) => { + let total_n = (gpu_n + cpu_n) as f64; + (gpu_dir * gpu_n as f64 + cpu_dir * cpu_n as f64) / total_n + } + (Some((dir, _)), None) | (None, Some((dir, _))) => dir, + (None, None) => 0.0, + }; + + let avg_confidence = if total_count > 0 { + total_confidence_sum / total_count as f64 + } else { + 0.0 + }; + + let ensemble_name = format!("STREAM_ENSEMBLE({})", model_names.join("+")); + + Ok(EnsemblePrediction { + model_name: ensemble_name, + direction, + confidence: avg_confidence, + metadata: PredictionMeta::default(), + }) + } + + /// Number of adapters in the ensemble. + pub fn adapter_count(&self) -> usize { + self.adapters.len() + } + + /// Number of adapters whose `is_ready()` returns true. + pub fn ready_count(&self) -> usize { + self.adapters.iter().filter(|a| a.is_ready()).count() + } + + /// Whether the underlying stream pool is on a CUDA device. + pub fn is_cuda(&self) -> bool { + self.stream_pool.is_cuda() + } + + // ----------------------------------------------------------------------- + // Private aggregation helpers + // ----------------------------------------------------------------------- + + /// GPU aggregation: stack tensors, sigmoid, confidence-weighted sum, extract. + fn aggregate_gpu( + &self, + preds: &[(usize, String, RawPrediction)], + ) -> Result<(f64, usize), MLError> { + let tensors: Vec = preds + .iter() + .filter_map(|(_, _, p)| p.tensor.as_ref().cloned()) + .collect(); + + if tensors.is_empty() { + return Err(MLError::InferenceError( + "No GPU tensors available for aggregation".to_owned(), + )); + } + + let stacked = Tensor::stack(&tensors, 0).map_err(|e| { + MLError::TensorOperationError(format!("Failed to stack GPU tensors: {e}")) + })?; + + let sigmoided = candle_nn::ops::sigmoid(&stacked).map_err(|e| { + MLError::TensorOperationError(format!("Sigmoid failed: {e}")) + })?; + + // Build confidence * model-weight vector on same device + let weights_f32: Vec = preds + .iter() + .map(|(idx, _, p)| { + let conf = p.confidence.clamp(0.0, 1.0) as f32; + let w = self.weights.get(*idx).copied().unwrap_or(1.0) as f32; + w * conf + }) + .collect(); + + let device = stacked.device(); + let n = weights_f32.len(); + + let weight_t = Tensor::from_vec(weights_f32, n, device).map_err(|e| { + MLError::TensorCreationError { + operation: "stream_ensemble_weight_tensor".to_owned(), + reason: format!("{e}"), + } + })?; + + let weight_sum = weight_t.sum_all().map_err(|e| { + MLError::TensorOperationError(format!("Weight sum failed: {e}")) + })?; + + let weight_sum_val = weight_sum.to_scalar::().map_err(|e| { + MLError::TensorOperationError(format!("Weight sum extraction failed: {e}")) + })?; + + if weight_sum_val.abs() < f32::EPSILON { + return Ok((0.0, preds.len())); + } + + let normalized = weight_t.broadcast_div(&weight_sum).map_err(|e| { + MLError::TensorOperationError(format!("Weight normalisation failed: {e}")) + })?; + + let weighted = sigmoided.mul(&normalized).map_err(|e| { + MLError::TensorOperationError(format!("Weighted multiply failed: {e}")) + })?; + + let result = weighted.sum_all().map_err(|e| { + MLError::TensorOperationError(format!("Sum extraction failed: {e}")) + })?; + + let direction_f32 = result.to_scalar::().map_err(|e| { + MLError::TensorOperationError(format!("Scalar extraction failed: {e}")) + })?; + + // Sigmoid output is [0,1]; remap to [-1,1] direction space + let direction = (f64::from(direction_f32) * 2.0) - 1.0; + + Ok((direction, preds.len())) + } + + /// CPU-side scalar weighted average. + fn aggregate_cpu( + &self, + preds: &[(usize, String, RawPrediction)], + ) -> Option<(f64, usize)> { + if preds.is_empty() { + return None; + } + + let mut weighted_direction_sum = 0.0_f64; + let mut weight_confidence_sum = 0.0_f64; + + for (idx, _, pred) in preds { + let confidence = pred.confidence.clamp(0.0, 1.0); + let w = self.weights.get(*idx).copied().unwrap_or(1.0); + let wc = w * confidence; + weighted_direction_sum += pred.direction_scalar * wc; + weight_confidence_sum += wc; + } + + let direction = if weight_confidence_sum.abs() < f64::EPSILON { + 0.0 + } else { + weighted_direction_sum / weight_confidence_sum + }; + + Some((direction, preds.len())) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use candle_core::Device; + + /// A simple test adapter with configurable direction, confidence, and readiness. + struct DummyAdapter { + name: String, + direction: f64, + confidence: f64, + ready: bool, + } + + impl ModelInferenceAdapter for DummyAdapter { + fn model_name(&self) -> &str { + &self.name + } + + fn predict(&self, _features: &FeatureVector) -> MLResult { + Ok(EnsemblePrediction { + model_name: self.name.clone(), + direction: self.direction, + confidence: self.confidence, + metadata: PredictionMeta::default(), + }) + } + + fn is_ready(&self) -> bool { + self.ready + } + } + + /// An adapter that always fails prediction. + struct FailingAdapter { + name: String, + } + + impl ModelInferenceAdapter for FailingAdapter { + fn model_name(&self) -> &str { + &self.name + } + + fn predict(&self, _features: &FeatureVector) -> MLResult { + Err(MLError::InferenceError("intentional test failure".to_owned())) + } + + fn is_ready(&self) -> bool { + true + } + } + + fn make_features() -> FeatureVector { + FeatureVector { + values: vec![0.0; 51], + timestamp: 1_700_000_000_000_000, + } + } + + #[test] + fn test_stream_ensemble_empty() { + let ensemble = StreamAwareEnsemble::new(vec![], vec![], &Device::Cpu) + .expect("empty ensemble should create"); + assert_eq!(ensemble.adapter_count(), 0); + assert_eq!(ensemble.ready_count(), 0); + assert!(!ensemble.is_cuda()); + + let pred = ensemble + .predict(&make_features()) + .expect("empty predict should succeed"); + assert_eq!(pred.direction, 0.0); + assert_eq!(pred.confidence, 0.0); + assert_eq!(pred.model_name, "STREAM_ENSEMBLE(empty)"); + } + + #[test] + fn test_stream_ensemble_single_model() { + let adapters: Vec> = vec![Box::new(DummyAdapter { + name: "DQN".to_owned(), + direction: 0.7, + confidence: 0.9, + ready: true, + })]; + + let ensemble = StreamAwareEnsemble::new(adapters, vec![1.0], &Device::Cpu) + .expect("single model should create"); + assert_eq!(ensemble.adapter_count(), 1); + assert_eq!(ensemble.ready_count(), 1); + + let pred = ensemble + .predict(&make_features()) + .expect("predict should succeed"); + + // Single model: direction comes straight through the CPU path + assert!( + (pred.direction - 0.7).abs() < 1e-9, + "expected direction 0.7, got {}", + pred.direction + ); + assert!( + (pred.confidence - 0.9).abs() < 1e-9, + "expected confidence 0.9, got {}", + pred.confidence + ); + assert!(pred.model_name.contains("DQN")); + } + + #[test] + fn test_stream_ensemble_two_models_equal_weight() { + // Model A: bullish (dir=1.0, conf=0.8) + // Model B: bearish (dir=-1.0, conf=0.6) + // CPU path: weighted_direction = (1.0*0.5*0.8 + (-1.0)*0.5*0.6) / (0.5*0.8 + 0.5*0.6) + // = (0.4 - 0.3) / (0.4 + 0.3) = 0.1 / 0.7 ~ 0.1429 + let adapters: Vec> = vec![ + Box::new(DummyAdapter { + name: "A".to_owned(), + direction: 1.0, + confidence: 0.8, + ready: true, + }), + Box::new(DummyAdapter { + name: "B".to_owned(), + direction: -1.0, + confidence: 0.6, + ready: true, + }), + ]; + + let ensemble = StreamAwareEnsemble::new(adapters, vec![1.0, 1.0], &Device::Cpu) + .expect("two-model ensemble should create"); + + let pred = ensemble + .predict(&make_features()) + .expect("predict should succeed"); + + // Net direction should be positive (bullish model has higher confidence) + assert!( + pred.direction > 0.0, + "direction should be positive, got {}", + pred.direction + ); + assert!( + pred.direction < 0.5, + "direction should be < 0.5, got {}", + pred.direction + ); + // Confidence = average = (0.8 + 0.6) / 2 = 0.7 + assert!( + (pred.confidence - 0.7).abs() < 1e-9, + "expected confidence 0.7, got {}", + pred.confidence + ); + assert!(pred.model_name.contains('A')); + assert!(pred.model_name.contains('B')); + } + + #[test] + fn test_stream_ensemble_skips_unready_models() { + let adapters: Vec> = vec![ + Box::new(DummyAdapter { + name: "Ready".to_owned(), + direction: 1.0, + confidence: 0.9, + ready: true, + }), + Box::new(DummyAdapter { + name: "NotReady".to_owned(), + direction: -1.0, + confidence: 0.9, + ready: false, + }), + ]; + + let ensemble = StreamAwareEnsemble::new(adapters, vec![1.0, 1.0], &Device::Cpu) + .expect("ensemble should create"); + assert_eq!(ensemble.ready_count(), 1); + + let pred = ensemble + .predict(&make_features()) + .expect("predict should succeed"); + + assert!(pred.model_name.contains("Ready")); + assert!(!pred.model_name.contains("NotReady")); + assert!(pred.direction > 0.5); + } + + #[test] + fn test_stream_ensemble_skips_nan_predictions() { + let adapters: Vec> = vec![ + Box::new(DummyAdapter { + name: "Valid".to_owned(), + direction: 0.8, + confidence: 0.7, + ready: true, + }), + Box::new(DummyAdapter { + name: "NaN_Model".to_owned(), + direction: f64::NAN, + confidence: 0.9, + ready: true, + }), + ]; + + let ensemble = StreamAwareEnsemble::new(adapters, vec![1.0, 1.0], &Device::Cpu) + .expect("ensemble should create"); + + let pred = ensemble + .predict(&make_features()) + .expect("predict should succeed"); + + assert!(pred.model_name.contains("Valid")); + assert!(!pred.model_name.contains("NaN_Model")); + assert!( + (pred.direction - 0.8).abs() < 1e-9, + "expected direction 0.8, got {}", + pred.direction + ); + } + + #[test] + fn test_stream_ensemble_skips_failing_models() { + let adapters: Vec> = vec![ + Box::new(DummyAdapter { + name: "Good".to_owned(), + direction: 0.5, + confidence: 0.8, + ready: true, + }), + Box::new(FailingAdapter { + name: "Broken".to_owned(), + }), + ]; + + let ensemble = StreamAwareEnsemble::new(adapters, vec![1.0, 1.0], &Device::Cpu) + .expect("ensemble should create"); + + let pred = ensemble + .predict(&make_features()) + .expect("predict should succeed despite one failing model"); + + assert!(pred.model_name.contains("Good")); + assert!(!pred.model_name.contains("Broken")); + } + + #[test] + fn test_stream_ensemble_weight_normalisation() { + // Weights [3.0, 1.0] should normalise to [0.75, 0.25] + let adapters: Vec> = vec![ + Box::new(DummyAdapter { + name: "Heavy".to_owned(), + direction: 1.0, + confidence: 0.8, + ready: true, + }), + Box::new(DummyAdapter { + name: "Light".to_owned(), + direction: -1.0, + confidence: 0.8, + ready: true, + }), + ]; + + let ensemble = StreamAwareEnsemble::new(adapters, vec![3.0, 1.0], &Device::Cpu) + .expect("ensemble should create"); + + let pred = ensemble + .predict(&make_features()) + .expect("predict should succeed"); + + // Heavy (0.75) bullish should dominate over Light (0.25) bearish + assert!( + pred.direction > 0.3, + "heavy bullish model should dominate, got {}", + pred.direction + ); + } + + #[test] + fn test_stream_ensemble_all_models_fail_returns_neutral() { + let adapters: Vec> = vec![ + Box::new(FailingAdapter { + name: "Fail1".to_owned(), + }), + Box::new(FailingAdapter { + name: "Fail2".to_owned(), + }), + ]; + + let ensemble = StreamAwareEnsemble::new(adapters, vec![1.0, 1.0], &Device::Cpu) + .expect("ensemble should create"); + + let pred = ensemble + .predict(&make_features()) + .expect("predict should return neutral when all fail"); + + assert_eq!(pred.direction, 0.0); + assert_eq!(pred.confidence, 0.0); + assert_eq!(pred.model_name, "STREAM_ENSEMBLE(empty)"); + } + + #[test] + fn test_stream_ensemble_missing_weights_default() { + // Provide fewer weights than adapters — missing ones should default + let adapters: Vec> = vec![ + Box::new(DummyAdapter { + name: "A".to_owned(), + direction: 0.5, + confidence: 0.8, + ready: true, + }), + Box::new(DummyAdapter { + name: "B".to_owned(), + direction: 0.5, + confidence: 0.8, + ready: true, + }), + Box::new(DummyAdapter { + name: "C".to_owned(), + direction: 0.5, + confidence: 0.8, + ready: true, + }), + ]; + + let ensemble = StreamAwareEnsemble::new(adapters, vec![1.0], &Device::Cpu) + .expect("missing weights should use defaults"); + + let pred = ensemble + .predict(&make_features()) + .expect("predict should succeed"); + + // All models agree on 0.5, so direction should be ~0.5 + assert!( + (pred.direction - 0.5).abs() < 1e-9, + "expected direction 0.5, got {}", + pred.direction + ); + } + + #[test] + fn test_stream_ensemble_cpu_not_cuda() { + let ensemble = StreamAwareEnsemble::new(vec![], vec![], &Device::Cpu) + .expect("cpu ensemble should create"); + assert!(!ensemble.is_cuda()); + } +}