398 lines
13 KiB
Rust
398 lines
13 KiB
Rust
//! Lightweight synchronous ensemble coordinator
|
|
//!
|
|
//! Wraps N [`ModelInferenceAdapter`] instances and aggregates their
|
|
//! predictions via confidence-weighted voting. Models that are not
|
|
//! ready or that fail inference are gracefully skipped with warnings.
|
|
|
|
use std::collections::HashMap;
|
|
|
|
use crate::ensemble::inference_adapter::{
|
|
EnsemblePrediction, FeatureVector, ModelInferenceAdapter, PredictionMeta,
|
|
};
|
|
use crate::{MLError, MLResult};
|
|
|
|
/// Synchronous coordinator that aggregates predictions from multiple
|
|
/// [`ModelInferenceAdapter`] instances using confidence-weighted voting.
|
|
///
|
|
/// # Aggregation formula
|
|
///
|
|
/// ```text
|
|
/// weighted_direction = sum(direction_i * weight_i * confidence_i)
|
|
/// / sum(weight_i * confidence_i)
|
|
/// ```
|
|
///
|
|
/// where `weight_i` defaults to 1.0 and can be overridden via
|
|
/// [`set_weight`](InferenceEnsemble::set_weight).
|
|
#[allow(missing_debug_implementations)]
|
|
pub struct InferenceEnsemble {
|
|
adapters: Vec<Box<dyn ModelInferenceAdapter>>,
|
|
weights: HashMap<String, f64>,
|
|
}
|
|
|
|
impl InferenceEnsemble {
|
|
/// Create a new ensemble from a vector of adapters.
|
|
///
|
|
/// All adapters start with a default weight of 1.0.
|
|
pub fn new(adapters: Vec<Box<dyn ModelInferenceAdapter>>) -> Self {
|
|
Self {
|
|
adapters,
|
|
weights: HashMap::new(),
|
|
}
|
|
}
|
|
|
|
/// Set a custom weight for a model identified by name.
|
|
///
|
|
/// Models without an explicit weight use 1.0.
|
|
pub fn set_weight(&mut self, model_name: &str, weight: f64) {
|
|
self.weights.insert(model_name.to_string(), weight);
|
|
}
|
|
|
|
/// Return the number of adapters whose `is_ready()` returns true.
|
|
pub fn ready_count(&self) -> usize {
|
|
self.adapters.iter().filter(|a| a.is_ready()).count()
|
|
}
|
|
|
|
/// Run inference across all ready adapters and aggregate via
|
|
/// confidence-weighted voting.
|
|
///
|
|
/// Returns [`MLError::InferenceError`] if no models are ready.
|
|
pub fn predict(&self, features: &FeatureVector) -> MLResult<EnsemblePrediction> {
|
|
let ready_adapters: Vec<&Box<dyn ModelInferenceAdapter>> =
|
|
self.adapters.iter().filter(|a| a.is_ready()).collect();
|
|
|
|
if ready_adapters.is_empty() {
|
|
return Err(MLError::InferenceError(
|
|
"No models are ready for inference".to_string(),
|
|
));
|
|
}
|
|
|
|
let mut weighted_direction_sum = 0.0_f64;
|
|
let mut weight_confidence_sum = 0.0_f64;
|
|
let mut confidence_sum = 0.0_f64;
|
|
let mut successful_count = 0_usize;
|
|
let mut model_names: Vec<String> = Vec::new();
|
|
|
|
for adapter in &ready_adapters {
|
|
let model_name = adapter.model_name().to_string();
|
|
match adapter.predict(features) {
|
|
Ok(pred) => {
|
|
// Circuit breaker: skip NaN/Inf predictions
|
|
if !pred.direction.is_finite() || !pred.confidence.is_finite() {
|
|
tracing::warn!(
|
|
model = %model_name,
|
|
direction = %pred.direction,
|
|
confidence = %pred.confidence,
|
|
"Model returned NaN/Inf prediction, skipping (circuit breaker)"
|
|
);
|
|
continue;
|
|
}
|
|
// Clamp confidence to valid range
|
|
let confidence = pred.confidence.clamp(0.0, 1.0);
|
|
let w = self
|
|
.weights
|
|
.get(&model_name)
|
|
.copied()
|
|
.unwrap_or(1.0);
|
|
let wc = w * confidence;
|
|
weighted_direction_sum += pred.direction * wc;
|
|
weight_confidence_sum += wc;
|
|
confidence_sum += confidence;
|
|
successful_count += 1;
|
|
model_names.push(model_name);
|
|
}
|
|
Err(e) => {
|
|
tracing::warn!(
|
|
model = %model_name,
|
|
error = %e,
|
|
"Model prediction failed, skipping"
|
|
);
|
|
}
|
|
}
|
|
}
|
|
|
|
if successful_count == 0 {
|
|
return Err(MLError::InferenceError(
|
|
"All ready models failed during inference".to_string(),
|
|
));
|
|
}
|
|
|
|
let direction = if weight_confidence_sum.abs() < f64::EPSILON {
|
|
0.0
|
|
} else {
|
|
weighted_direction_sum / weight_confidence_sum
|
|
};
|
|
|
|
let avg_confidence = confidence_sum / successful_count as f64;
|
|
|
|
let ensemble_name = format!("ENSEMBLE({})", model_names.join("+"));
|
|
|
|
Ok(EnsemblePrediction {
|
|
model_name: ensemble_name,
|
|
direction,
|
|
confidence: avg_confidence,
|
|
metadata: PredictionMeta::default(),
|
|
})
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
/// 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
|
|
}
|
|
}
|
|
|
|
fn make_features() -> FeatureVector {
|
|
FeatureVector {
|
|
values: vec![0.0; 51],
|
|
timestamp: 1_700_000_000_000_000,
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_ensemble_equal_weight_aggregation() {
|
|
// Model A: bullish (dir=1.0, conf=0.8)
|
|
// Model B: bearish (dir=-1.0, conf=0.6)
|
|
// weighted_direction = (1.0*1.0*0.8 + (-1.0)*1.0*0.6) / (1.0*0.8 + 1.0*0.6)
|
|
// = (0.8 - 0.6) / 1.4
|
|
// = 0.2 / 1.4
|
|
// ~ 0.1429
|
|
let adapters: Vec<Box<dyn ModelInferenceAdapter>> = vec![
|
|
Box::new(DummyAdapter {
|
|
name: "A".to_string(),
|
|
direction: 1.0,
|
|
confidence: 0.8,
|
|
ready: true,
|
|
}),
|
|
Box::new(DummyAdapter {
|
|
name: "B".to_string(),
|
|
direction: -1.0,
|
|
confidence: 0.6,
|
|
ready: true,
|
|
}),
|
|
];
|
|
|
|
let ensemble = InferenceEnsemble::new(adapters);
|
|
let features = make_features();
|
|
let pred = ensemble.predict(&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
|
|
);
|
|
// But not too strong since the bearish model partially cancels
|
|
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,
|
|
"confidence should be 0.7, got {}",
|
|
pred.confidence
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_ensemble_skips_unready_models() {
|
|
let adapters: Vec<Box<dyn ModelInferenceAdapter>> = vec![
|
|
Box::new(DummyAdapter {
|
|
name: "Ready".to_string(),
|
|
direction: 1.0,
|
|
confidence: 0.9,
|
|
ready: true,
|
|
}),
|
|
Box::new(DummyAdapter {
|
|
name: "NotReady".to_string(),
|
|
direction: -1.0,
|
|
confidence: 0.9,
|
|
ready: false,
|
|
}),
|
|
];
|
|
|
|
let ensemble = InferenceEnsemble::new(adapters);
|
|
assert_eq!(ensemble.ready_count(), 1);
|
|
|
|
let features = make_features();
|
|
let pred = ensemble.predict(&features).expect("predict should succeed");
|
|
|
|
// Only the ready (bullish) model contributes
|
|
assert!(
|
|
pred.direction > 0.5,
|
|
"direction should be > 0.5 with only bullish model, got {}",
|
|
pred.direction
|
|
);
|
|
// Model name should only include the ready model
|
|
assert!(
|
|
pred.model_name.contains("Ready"),
|
|
"model_name should contain 'Ready', got {}",
|
|
pred.model_name
|
|
);
|
|
assert!(
|
|
!pred.model_name.contains("NotReady"),
|
|
"model_name should NOT contain 'NotReady', got {}",
|
|
pred.model_name
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_ensemble_custom_weights() {
|
|
// Model A: bullish, weight 0.75
|
|
// Model B: bearish, weight 0.25
|
|
// weighted_direction = (1.0*0.75*0.8 + (-1.0)*0.25*0.8) / (0.75*0.8 + 0.25*0.8)
|
|
// = (0.6 - 0.2) / (0.6 + 0.2)
|
|
// = 0.4 / 0.8
|
|
// = 0.5
|
|
let adapters: Vec<Box<dyn ModelInferenceAdapter>> = vec![
|
|
Box::new(DummyAdapter {
|
|
name: "Heavy".to_string(),
|
|
direction: 1.0,
|
|
confidence: 0.8,
|
|
ready: true,
|
|
}),
|
|
Box::new(DummyAdapter {
|
|
name: "Light".to_string(),
|
|
direction: -1.0,
|
|
confidence: 0.8,
|
|
ready: true,
|
|
}),
|
|
];
|
|
|
|
let mut ensemble = InferenceEnsemble::new(adapters);
|
|
ensemble.set_weight("Heavy", 0.75);
|
|
ensemble.set_weight("Light", 0.25);
|
|
|
|
let features = make_features();
|
|
let pred = ensemble.predict(&features).expect("predict should succeed");
|
|
|
|
// The heavier-weighted bullish model should dominate
|
|
assert!(
|
|
pred.direction > 0.0,
|
|
"direction should be positive (heavy bullish dominates), got {}",
|
|
pred.direction
|
|
);
|
|
// With equal confidences, the 3:1 weight ratio strongly favors bullish
|
|
// Expected: 0.5, but any positive value > 0.3 confirms dominance
|
|
assert!(
|
|
pred.direction > 0.3,
|
|
"direction should be > 0.3 showing heavy model dominance, got {}",
|
|
pred.direction
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_ensemble_filters_nan_predictions() {
|
|
// NaN model should be skipped; the valid model's prediction stands alone.
|
|
let adapters: Vec<Box<dyn ModelInferenceAdapter>> = vec![
|
|
Box::new(DummyAdapter {
|
|
name: "Valid".to_string(),
|
|
direction: 0.8,
|
|
confidence: 0.7,
|
|
ready: true,
|
|
}),
|
|
Box::new(DummyAdapter {
|
|
name: "NaN_Model".to_string(),
|
|
direction: f64::NAN,
|
|
confidence: 0.9,
|
|
ready: true,
|
|
}),
|
|
];
|
|
|
|
let ensemble = InferenceEnsemble::new(adapters);
|
|
let features = make_features();
|
|
let pred = ensemble.predict(&features).expect("predict should succeed");
|
|
|
|
// Only the Valid model should contribute
|
|
assert!(
|
|
pred.model_name.contains("Valid"),
|
|
"model_name should contain 'Valid', got {}",
|
|
pred.model_name
|
|
);
|
|
assert!(
|
|
!pred.model_name.contains("NaN_Model"),
|
|
"model_name should NOT contain 'NaN_Model', got {}",
|
|
pred.model_name
|
|
);
|
|
// direction should come entirely from the Valid model
|
|
assert!(
|
|
(pred.direction - 0.8).abs() < 1e-9,
|
|
"direction should be 0.8, got {}",
|
|
pred.direction
|
|
);
|
|
// confidence should be from the single valid model
|
|
assert!(
|
|
(pred.confidence - 0.7).abs() < 1e-9,
|
|
"confidence should be 0.7, got {}",
|
|
pred.confidence
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_ensemble_filters_extreme_confidence() {
|
|
// Model with confidence 5.0 should be clamped to 1.0
|
|
let adapters: Vec<Box<dyn ModelInferenceAdapter>> = vec![
|
|
Box::new(DummyAdapter {
|
|
name: "Normal".to_string(),
|
|
direction: 1.0,
|
|
confidence: 0.6,
|
|
ready: true,
|
|
}),
|
|
Box::new(DummyAdapter {
|
|
name: "Extreme".to_string(),
|
|
direction: 1.0,
|
|
confidence: 5.0,
|
|
ready: true,
|
|
}),
|
|
];
|
|
|
|
let ensemble = InferenceEnsemble::new(adapters);
|
|
let features = make_features();
|
|
let pred = ensemble.predict(&features).expect("predict should succeed");
|
|
|
|
// Both models should contribute (extreme confidence is clamped, not skipped)
|
|
assert!(
|
|
pred.model_name.contains("Normal"),
|
|
"model_name should contain 'Normal', got {}",
|
|
pred.model_name
|
|
);
|
|
assert!(
|
|
pred.model_name.contains("Extreme"),
|
|
"model_name should contain 'Extreme', got {}",
|
|
pred.model_name
|
|
);
|
|
// Confidence should be clamped: avg of 0.6 and 1.0 = 0.8
|
|
assert!(
|
|
(pred.confidence - 0.8).abs() < 1e-9,
|
|
"confidence should be 0.8 (avg of 0.6 and clamped 1.0), got {}",
|
|
pred.confidence
|
|
);
|
|
}
|
|
}
|