feat(ml): add StreamAwareEnsemble with per-model CUDA streams
Stream-aware ensemble that runs models on separate CUDA streams for true GPU-level parallelism. Falls back to rayon on CPU. Uses CudaStreamPool for synchronization. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -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};
|
||||
|
||||
679
crates/ml/src/ensemble/stream_ensemble.rs
Normal file
679
crates/ml/src/ensemble/stream_ensemble.rs
Normal file
@@ -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<Box<dyn ModelInferenceAdapter>>,
|
||||
/// Weights per model (normalised, indexed by adapter position)
|
||||
weights: Vec<f64>,
|
||||
/// 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<Box<dyn ModelInferenceAdapter>>,
|
||||
weights: Vec<f64>,
|
||||
device: &candle_core::Device,
|
||||
) -> MLResult<Self> {
|
||||
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<EnsemblePrediction> {
|
||||
// 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<String> = 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<Tensor> = 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<f32> = 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::<f32>().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::<f32>().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<EnsemblePrediction> {
|
||||
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<EnsemblePrediction> {
|
||||
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<Box<dyn ModelInferenceAdapter>> = 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<Box<dyn ModelInferenceAdapter>> = 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<Box<dyn ModelInferenceAdapter>> = 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<Box<dyn ModelInferenceAdapter>> = 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<Box<dyn ModelInferenceAdapter>> = 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<Box<dyn ModelInferenceAdapter>> = 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<Box<dyn ModelInferenceAdapter>> = 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<Box<dyn ModelInferenceAdapter>> = 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());
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user