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>
667 lines
21 KiB
Rust
667 lines
21 KiB
Rust
//! ML Pipeline Testing Harness
|
|
//!
|
|
//! Provides comprehensive testing capabilities for ML models including
|
|
//! inference testing, feature extraction, ensemble predictions, and
|
|
//! model performance monitoring.
|
|
|
|
use anyhow::Result;
|
|
use common::types::MarketTick;
|
|
use std::collections::HashMap;
|
|
use std::time::{Duration, Instant};
|
|
use tracing::{debug, info, warn};
|
|
|
|
/// ML model status information
|
|
#[derive(Debug, Clone)]
|
|
pub struct MLModelStatus {
|
|
pub mamba_available: bool,
|
|
pub dqn_available: bool,
|
|
pub ppo_available: bool,
|
|
pub tft_available: bool,
|
|
pub tlob_available: bool,
|
|
pub ensemble_available: bool,
|
|
}
|
|
|
|
impl MLModelStatus {
|
|
/// Check if any models are available
|
|
pub fn any_available(&self) -> bool {
|
|
self.mamba_available
|
|
|| self.dqn_available
|
|
|| self.ppo_available
|
|
|| self.tft_available
|
|
|| self.tlob_available
|
|
}
|
|
|
|
/// Get count of available models
|
|
pub fn available_count(&self) -> usize {
|
|
let mut count = 0;
|
|
if self.mamba_available {
|
|
count += 1;
|
|
}
|
|
if self.dqn_available {
|
|
count += 1;
|
|
}
|
|
if self.ppo_available {
|
|
count += 1;
|
|
}
|
|
if self.tft_available {
|
|
count += 1;
|
|
}
|
|
if self.tlob_available {
|
|
count += 1;
|
|
}
|
|
count
|
|
}
|
|
}
|
|
|
|
/// ML prediction result
|
|
#[derive(Debug, Clone)]
|
|
pub struct MLPrediction {
|
|
pub signal: f64, // Trading signal (-1.0 to 1.0)
|
|
pub confidence: f64, // Confidence (0.0 to 1.0)
|
|
pub model_name: String, // Name of the model that made the prediction
|
|
pub inference_time: Duration, // Time taken for inference
|
|
}
|
|
|
|
/// Ensemble prediction result
|
|
#[derive(Debug, Clone)]
|
|
pub struct EnsemblePrediction {
|
|
pub signal: f64, // Aggregated trading signal
|
|
pub confidence: f64, // Aggregated confidence
|
|
pub individual_predictions: Vec<MLPrediction>,
|
|
pub ensemble_method: String, // Method used for aggregation
|
|
pub total_inference_time: Duration,
|
|
pub prediction: PredictionType, // Discrete prediction type
|
|
pub signal_strength: f64, // Absolute signal strength
|
|
}
|
|
|
|
// Define PredictionType locally since tli crate is not available in e2e tests
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
|
pub enum PredictionType {
|
|
Buy,
|
|
Sell,
|
|
Hold,
|
|
StrongBuy,
|
|
StrongSell,
|
|
}
|
|
|
|
impl std::fmt::Display for PredictionType {
|
|
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
|
match self {
|
|
PredictionType::Buy => write!(f, "BUY"),
|
|
PredictionType::Sell => write!(f, "SELL"),
|
|
PredictionType::Hold => write!(f, "HOLD"),
|
|
PredictionType::StrongBuy => write!(f, "STRONG_BUY"),
|
|
PredictionType::StrongSell => write!(f, "STRONG_SELL"),
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Model performance metrics
|
|
#[derive(Debug, Clone)]
|
|
pub struct ModelMetrics {
|
|
pub inference_count: u64,
|
|
pub avg_latency_ms: f64,
|
|
pub error_rate: f64,
|
|
pub last_prediction_time: chrono::DateTime<chrono::Utc>,
|
|
}
|
|
|
|
/// Feature vector for ML models
|
|
#[derive(Debug, Clone)]
|
|
pub struct FeatureVector {
|
|
pub features: Vec<f64>,
|
|
pub feature_names: Vec<String>,
|
|
pub timestamp: chrono::DateTime<chrono::Utc>,
|
|
pub symbol: String,
|
|
}
|
|
|
|
/// ML Pipeline Testing Harness
|
|
#[derive(Debug)]
|
|
pub struct MLPipelineTestHarness {
|
|
model_status: MLModelStatus,
|
|
model_metrics: HashMap<String, ModelMetrics>,
|
|
feature_cache: HashMap<String, Vec<FeatureVector>>,
|
|
mock_mode: bool,
|
|
}
|
|
|
|
impl MLPipelineTestHarness {
|
|
/// Create a new ML pipeline test harness
|
|
pub async fn new() -> Result<Self> {
|
|
info!("🤖 Initializing ML pipeline test harness");
|
|
|
|
// Check if we're in mock mode (for CI/testing environments without GPU)
|
|
let mock_mode = std::env::var("ML_MOCK_MODE").unwrap_or_default() == "true";
|
|
|
|
if mock_mode {
|
|
info!("🎭 Running in mock mode - ML predictions will be simulated");
|
|
}
|
|
|
|
let model_status = Self::check_model_availability().await?;
|
|
|
|
info!("ML Models Status:");
|
|
info!(
|
|
" MAMBA: {}",
|
|
if model_status.mamba_available {
|
|
"✅"
|
|
} else {
|
|
"❌"
|
|
}
|
|
);
|
|
info!(
|
|
" DQN: {}",
|
|
if model_status.dqn_available {
|
|
"✅"
|
|
} else {
|
|
"❌"
|
|
}
|
|
);
|
|
info!(
|
|
" PPO: {}",
|
|
if model_status.ppo_available {
|
|
"✅"
|
|
} else {
|
|
"❌"
|
|
}
|
|
);
|
|
info!(
|
|
" TFT: {}",
|
|
if model_status.tft_available {
|
|
"✅"
|
|
} else {
|
|
"❌"
|
|
}
|
|
);
|
|
info!(
|
|
" TLOB: {}",
|
|
if model_status.tlob_available {
|
|
"✅"
|
|
} else {
|
|
"❌"
|
|
}
|
|
);
|
|
|
|
Ok(Self {
|
|
model_status,
|
|
model_metrics: HashMap::new(),
|
|
feature_cache: HashMap::new(),
|
|
mock_mode,
|
|
})
|
|
}
|
|
|
|
/// Check model health status
|
|
pub async fn check_models_health(&self) -> Result<MLModelStatus> {
|
|
Ok(self.model_status.clone())
|
|
}
|
|
|
|
/// Extract features from market data
|
|
pub async fn extract_features(
|
|
&mut self,
|
|
market_data: &[MarketTick],
|
|
) -> Result<Vec<FeatureVector>> {
|
|
debug!(
|
|
"🔧 Extracting features from {} market ticks",
|
|
market_data.len()
|
|
);
|
|
|
|
if market_data.is_empty() {
|
|
return Ok(Vec::new());
|
|
}
|
|
|
|
let mut features = Vec::new();
|
|
|
|
// Group by symbol for feature extraction
|
|
let mut symbol_data: HashMap<String, Vec<&MarketTick>> = HashMap::new();
|
|
for tick in market_data {
|
|
symbol_data
|
|
.entry(tick.symbol.to_string())
|
|
.or_default()
|
|
.push(tick);
|
|
}
|
|
|
|
for (symbol, ticks) in symbol_data {
|
|
let feature_vector = self.extract_features_for_symbol(&symbol, ticks).await?;
|
|
features.push(feature_vector);
|
|
}
|
|
|
|
// Cache features for later use
|
|
for feature in &features {
|
|
self.feature_cache
|
|
.entry(feature.symbol.clone())
|
|
.or_default()
|
|
.push(feature.clone());
|
|
}
|
|
|
|
debug!("✅ Extracted {} feature vectors", features.len());
|
|
Ok(features)
|
|
}
|
|
|
|
/// Extract features for a specific symbol
|
|
async fn extract_features_for_symbol(
|
|
&self,
|
|
symbol: &str,
|
|
ticks: Vec<&MarketTick>,
|
|
) -> Result<FeatureVector> {
|
|
if ticks.is_empty() {
|
|
return Err(anyhow::anyhow!("No ticks provided for feature extraction"));
|
|
}
|
|
|
|
// Sort by timestamp
|
|
let mut sorted_ticks = ticks;
|
|
sorted_ticks.sort_by_key(|t| t.timestamp);
|
|
|
|
let prices: Vec<f64> = sorted_ticks.iter().map(|t| t.price.as_f64()).collect();
|
|
let volumes: Vec<f64> = sorted_ticks.iter().map(|t| t.size.as_f64()).collect();
|
|
|
|
// Calculate technical indicators
|
|
let mut features = Vec::new();
|
|
let mut feature_names = Vec::new();
|
|
|
|
// Price-based features
|
|
if !prices.is_empty() {
|
|
let current_price = prices[prices.len() - 1];
|
|
let avg_price = prices.iter().sum::<f64>() / prices.len() as f64;
|
|
let price_std = Self::calculate_std(&prices);
|
|
|
|
features.push(current_price);
|
|
features.push(avg_price);
|
|
features.push(price_std);
|
|
features.push(current_price / avg_price - 1.0); // Price relative to average
|
|
|
|
feature_names.extend([
|
|
"current_price".to_string(),
|
|
"avg_price".to_string(),
|
|
"price_std".to_string(),
|
|
"price_rel_avg".to_string(),
|
|
]);
|
|
}
|
|
|
|
// Volume-based features
|
|
if !volumes.is_empty() {
|
|
let current_volume = volumes[volumes.len() - 1];
|
|
let avg_volume = volumes.iter().sum::<f64>() / volumes.len() as f64;
|
|
|
|
features.push(current_volume);
|
|
features.push(avg_volume);
|
|
features.push(current_volume / avg_volume.max(1.0)); // Volume relative to average
|
|
|
|
feature_names.extend([
|
|
"current_volume".to_string(),
|
|
"avg_volume".to_string(),
|
|
"volume_rel_avg".to_string(),
|
|
]);
|
|
}
|
|
|
|
// Returns-based features
|
|
if prices.len() >= 2 {
|
|
let returns: Vec<f64> = prices.windows(2).map(|w| (w[1] / w[0]) - 1.0).collect();
|
|
|
|
if !returns.is_empty() {
|
|
let last_return = returns[returns.len() - 1];
|
|
let avg_return = returns.iter().sum::<f64>() / returns.len() as f64;
|
|
let return_std = Self::calculate_std(&returns);
|
|
|
|
features.push(last_return);
|
|
features.push(avg_return);
|
|
features.push(return_std);
|
|
|
|
feature_names.extend([
|
|
"last_return".to_string(),
|
|
"avg_return".to_string(),
|
|
"return_std".to_string(),
|
|
]);
|
|
}
|
|
}
|
|
|
|
// Trend features
|
|
if prices.len() >= 5 {
|
|
let short_ma = prices[prices.len() - 5..].iter().sum::<f64>() / 5.0;
|
|
let long_ma = prices.iter().sum::<f64>() / prices.len() as f64;
|
|
let trend_strength = (short_ma / long_ma) - 1.0;
|
|
|
|
features.push(trend_strength);
|
|
feature_names.push("trend_strength".to_string());
|
|
}
|
|
|
|
Ok(FeatureVector {
|
|
features,
|
|
feature_names,
|
|
timestamp: chrono::Utc::now(),
|
|
symbol: symbol.to_string(),
|
|
})
|
|
}
|
|
|
|
/// Predict with MAMBA model
|
|
pub async fn predict_with_mamba(&mut self, features: &[FeatureVector]) -> Result<MLPrediction> {
|
|
self.predict_with_model("mamba", features).await
|
|
}
|
|
|
|
/// Predict with DQN model
|
|
pub async fn predict_with_dqn(&mut self, features: &[FeatureVector]) -> Result<MLPrediction> {
|
|
self.predict_with_model("dqn", features).await
|
|
}
|
|
|
|
/// Predict with TFT model
|
|
pub async fn predict_with_tft(&mut self, features: &[FeatureVector]) -> Result<MLPrediction> {
|
|
self.predict_with_model("tft", features).await
|
|
}
|
|
|
|
/// Predict with TLOB model
|
|
pub async fn predict_with_tlob(&mut self, features: &[FeatureVector]) -> Result<MLPrediction> {
|
|
self.predict_with_model("tlob", features).await
|
|
}
|
|
|
|
/// Generic model prediction
|
|
async fn predict_with_model(
|
|
&mut self,
|
|
model_name: &str,
|
|
features: &[FeatureVector],
|
|
) -> Result<MLPrediction> {
|
|
let start_time = Instant::now();
|
|
|
|
if features.is_empty() {
|
|
return Err(anyhow::anyhow!("No features provided for prediction"));
|
|
}
|
|
|
|
let prediction = if self.mock_mode {
|
|
self.mock_prediction(model_name, features).await?
|
|
} else {
|
|
self.real_prediction(model_name, features).await?
|
|
};
|
|
|
|
let inference_time = start_time.elapsed();
|
|
|
|
// Update model metrics
|
|
self.update_model_metrics(model_name, inference_time, true);
|
|
|
|
Ok(MLPrediction {
|
|
signal: prediction.0,
|
|
confidence: prediction.1,
|
|
model_name: model_name.to_string(),
|
|
inference_time,
|
|
})
|
|
}
|
|
|
|
/// Mock prediction for testing
|
|
async fn mock_prediction(
|
|
&self,
|
|
model_name: &str,
|
|
features: &[FeatureVector],
|
|
) -> Result<(f64, f64)> {
|
|
use rand::Rng;
|
|
let mut rng = rand::thread_rng();
|
|
|
|
// Simulate some processing time
|
|
tokio::time::sleep(Duration::from_millis(rng.gen_range(10..50))).await;
|
|
|
|
// Generate reasonable mock predictions based on features
|
|
let feature_sum: f64 = features.iter().flat_map(|f| &f.features).sum();
|
|
|
|
let normalized_sum = (feature_sum / 1000.0).tanh(); // Normalize to [-1, 1]
|
|
|
|
let signal = match model_name {
|
|
"mamba" => normalized_sum * 0.8 + rng.gen_range(-0.1..0.1),
|
|
"dqn" => normalized_sum * 0.6 + rng.gen_range(-0.2..0.2),
|
|
"tft" => normalized_sum * 0.9 + rng.gen_range(-0.05..0.05),
|
|
"tlob" => normalized_sum * 0.7 + rng.gen_range(-0.15..0.15),
|
|
_ => rng.gen_range(-0.5..0.5),
|
|
};
|
|
|
|
let confidence = rng.gen_range(0.6..0.95);
|
|
|
|
Ok((signal.clamp(-1.0, 1.0), confidence))
|
|
}
|
|
|
|
/// Real prediction (would integrate with actual ML models)
|
|
async fn real_prediction(
|
|
&self,
|
|
model_name: &str,
|
|
_features: &[FeatureVector],
|
|
) -> Result<(f64, f64)> {
|
|
// This would integrate with the actual ML models in the ml crate
|
|
warn!(
|
|
"Real ML prediction not implemented for {}, using mock",
|
|
model_name
|
|
);
|
|
|
|
// For now, fall back to mock prediction
|
|
self.mock_prediction(model_name, _features).await
|
|
}
|
|
|
|
/// Test ensemble prediction with raw feature array
|
|
pub async fn test_ensemble_prediction(
|
|
&mut self,
|
|
raw_features: Vec<f64>,
|
|
) -> Result<EnsemblePrediction> {
|
|
// Convert raw features to FeatureVector format
|
|
let feature_count = raw_features.len();
|
|
let feature_vector = FeatureVector {
|
|
features: raw_features,
|
|
feature_names: (0..feature_count)
|
|
.map(|i| format!("feature_{}", i))
|
|
.collect(),
|
|
timestamp: chrono::Utc::now(),
|
|
symbol: "TEST".to_string(),
|
|
};
|
|
|
|
self.predict_ensemble(&[feature_vector]).await
|
|
}
|
|
|
|
/// Ensemble prediction aggregating multiple models
|
|
pub async fn predict_ensemble(
|
|
&mut self,
|
|
features: &[FeatureVector],
|
|
) -> Result<EnsemblePrediction> {
|
|
let start_time = Instant::now();
|
|
let mut predictions = Vec::new();
|
|
|
|
// Get predictions from available models
|
|
if self.model_status.mamba_available {
|
|
match self.predict_with_mamba(features).await {
|
|
Ok(pred) => predictions.push(pred),
|
|
Err(e) => warn!("MAMBA prediction failed: {}", e),
|
|
}
|
|
}
|
|
|
|
if self.model_status.dqn_available {
|
|
match self.predict_with_dqn(features).await {
|
|
Ok(pred) => predictions.push(pred),
|
|
Err(e) => warn!("DQN prediction failed: {}", e),
|
|
}
|
|
}
|
|
|
|
if self.model_status.tft_available {
|
|
match self.predict_with_tft(features).await {
|
|
Ok(pred) => predictions.push(pred),
|
|
Err(e) => warn!("TFT prediction failed: {}", e),
|
|
}
|
|
}
|
|
|
|
if self.model_status.tlob_available {
|
|
match self.predict_with_tlob(features).await {
|
|
Ok(pred) => predictions.push(pred),
|
|
Err(e) => warn!("TLOB prediction failed: {}", e),
|
|
}
|
|
}
|
|
|
|
if predictions.is_empty() {
|
|
return Err(anyhow::anyhow!(
|
|
"No models available for ensemble prediction"
|
|
));
|
|
}
|
|
|
|
// Aggregate predictions using weighted average
|
|
let total_weight: f64 = predictions.iter().map(|p| p.confidence).sum();
|
|
let weighted_signal: f64 = predictions
|
|
.iter()
|
|
.map(|p| p.signal * p.confidence)
|
|
.sum::<f64>()
|
|
/ total_weight;
|
|
|
|
let avg_confidence: f64 =
|
|
predictions.iter().map(|p| p.confidence).sum::<f64>() / predictions.len() as f64;
|
|
|
|
let total_inference_time = start_time.elapsed();
|
|
|
|
// Convert signal to discrete prediction type
|
|
let prediction = if weighted_signal > 0.6 {
|
|
PredictionType::StrongBuy
|
|
} else if weighted_signal > 0.2 {
|
|
PredictionType::Buy
|
|
} else if weighted_signal < -0.6 {
|
|
PredictionType::StrongSell
|
|
} else if weighted_signal < -0.2 {
|
|
PredictionType::Sell
|
|
} else {
|
|
PredictionType::Hold
|
|
};
|
|
|
|
let signal_strength = weighted_signal.abs();
|
|
|
|
Ok(EnsemblePrediction {
|
|
signal: weighted_signal.clamp(-1.0, 1.0),
|
|
confidence: avg_confidence.clamp(0.0, 1.0),
|
|
individual_predictions: predictions,
|
|
ensemble_method: "weighted_average".to_string(),
|
|
total_inference_time,
|
|
prediction,
|
|
signal_strength,
|
|
})
|
|
}
|
|
|
|
/// Get model performance metrics
|
|
pub async fn get_model_metrics(&self) -> Result<HashMap<String, ModelMetrics>> {
|
|
Ok(self.model_metrics.clone())
|
|
}
|
|
|
|
/// Update model metrics
|
|
fn update_model_metrics(&mut self, model_name: &str, inference_time: Duration, success: bool) {
|
|
let metrics = self
|
|
.model_metrics
|
|
.entry(model_name.to_string())
|
|
.or_insert_with(|| ModelMetrics {
|
|
inference_count: 0,
|
|
avg_latency_ms: 0.0,
|
|
error_rate: 0.0,
|
|
last_prediction_time: chrono::Utc::now(),
|
|
});
|
|
|
|
metrics.inference_count += 1;
|
|
|
|
// Update average latency
|
|
let new_latency_ms = inference_time.as_millis() as f64;
|
|
if metrics.inference_count == 1 {
|
|
metrics.avg_latency_ms = new_latency_ms;
|
|
} else {
|
|
metrics.avg_latency_ms =
|
|
(metrics.avg_latency_ms * (metrics.inference_count - 1) as f64 + new_latency_ms)
|
|
/ metrics.inference_count as f64;
|
|
}
|
|
|
|
// Update error rate
|
|
if !success {
|
|
let error_count = (metrics.error_rate * (metrics.inference_count - 1) as f64) + 1.0;
|
|
metrics.error_rate = error_count / metrics.inference_count as f64;
|
|
} else {
|
|
let error_count = metrics.error_rate * (metrics.inference_count - 1) as f64;
|
|
metrics.error_rate = error_count / metrics.inference_count as f64;
|
|
}
|
|
|
|
metrics.last_prediction_time = chrono::Utc::now();
|
|
}
|
|
|
|
/// Disable a model for failover testing
|
|
pub async fn disable_model(&mut self, model_name: &str) -> Result<()> {
|
|
match model_name {
|
|
"mamba" => self.model_status.mamba_available = false,
|
|
"dqn" => self.model_status.dqn_available = false,
|
|
"ppo" => self.model_status.ppo_available = false,
|
|
"tft" => self.model_status.tft_available = false,
|
|
"tlob" => self.model_status.tlob_available = false,
|
|
_ => return Err(anyhow::anyhow!("Unknown model: {}", model_name)),
|
|
}
|
|
|
|
info!("🚫 Disabled model: {}", model_name);
|
|
Ok(())
|
|
}
|
|
|
|
/// Enable a model for failover testing
|
|
pub async fn enable_model(&mut self, model_name: &str) -> Result<()> {
|
|
match model_name {
|
|
"mamba" => self.model_status.mamba_available = true,
|
|
"dqn" => self.model_status.dqn_available = true,
|
|
"ppo" => self.model_status.ppo_available = true,
|
|
"tft" => self.model_status.tft_available = true,
|
|
"tlob" => self.model_status.tlob_available = true,
|
|
_ => return Err(anyhow::anyhow!("Unknown model: {}", model_name)),
|
|
}
|
|
|
|
info!("✅ Enabled model: {}", model_name);
|
|
Ok(())
|
|
}
|
|
|
|
/// Check model availability
|
|
async fn check_model_availability() -> Result<MLModelStatus> {
|
|
// In a real implementation, this would check for model files,
|
|
// GPU availability, etc. For testing, we'll assume models are available.
|
|
Ok(MLModelStatus {
|
|
mamba_available: true,
|
|
dqn_available: true,
|
|
ppo_available: true,
|
|
tft_available: true,
|
|
tlob_available: true,
|
|
ensemble_available: true,
|
|
})
|
|
}
|
|
|
|
/// Calculate standard deviation
|
|
fn calculate_std(values: &[f64]) -> f64 {
|
|
if values.len() < 2 {
|
|
return 0.0;
|
|
}
|
|
|
|
let mean = values.iter().sum::<f64>() / values.len() as f64;
|
|
let variance =
|
|
values.iter().map(|x| (x - mean).powi(2)).sum::<f64>() / (values.len() - 1) as f64;
|
|
|
|
variance.sqrt()
|
|
}
|
|
}
|
|
|
|
/// Type alias for compatibility with workflow code
|
|
pub type MLTestPipeline = MLPipelineTestHarness;
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[tokio::test]
|
|
async fn test_ml_harness_creation() {
|
|
let harness = MLPipelineTestHarness::new().await;
|
|
assert!(harness.is_ok());
|
|
|
|
let harness = harness.unwrap();
|
|
assert!(harness.model_status.any_available());
|
|
}
|
|
|
|
#[test]
|
|
fn test_model_status() {
|
|
let status = MLModelStatus {
|
|
mamba_available: true,
|
|
dqn_available: true,
|
|
ppo_available: false,
|
|
tft_available: false,
|
|
tlob_available: false,
|
|
ensemble_available: true,
|
|
};
|
|
|
|
assert!(status.any_available());
|
|
assert_eq!(status.available_count(), 2);
|
|
}
|
|
|
|
#[test]
|
|
fn test_std_calculation() {
|
|
let values = vec![1.0, 2.0, 3.0, 4.0, 5.0];
|
|
let std = MLPipelineTestHarness::calculate_std(&values);
|
|
assert!((std - 1.581).abs() < 0.01); // Approximately sqrt(2.5)
|
|
}
|
|
}
|