diff --git a/crates/common/src/metrics/server.rs b/crates/common/src/metrics/server.rs index f9d79bc4b..8e47f29b1 100644 --- a/crates/common/src/metrics/server.rs +++ b/crates/common/src/metrics/server.rs @@ -4,7 +4,7 @@ //! existing HTTP stack. Training binaries and CLI tools that don't have an //! HTTP server use this instead. -use std::io::{BufRead, BufReader, Write as IoWrite}; +use std::io::{BufRead, BufReader, Read as IoRead, Write as IoWrite}; use std::net::TcpListener; use super::gather_metrics; @@ -14,6 +14,11 @@ use super::gather_metrics; /// Responds to `GET /metrics` with the global registry output in Prometheus /// text exposition format. Any other request gets a 404. The thread is /// detached and dies with the process. +/// +/// Safety hardening: +/// - 5-second read timeout prevents slow-loris connections +/// - 8KB request line limit prevents memory abuse +/// - Correct Content-Type charset for Prometheus scrapers pub fn start_metrics_server(port: u16) { std::thread::spawn(move || { let addr = format!("0.0.0.0:{port}"); @@ -33,21 +38,25 @@ pub fn start_metrics_server(port: u16) { continue; }; - // Read the HTTP request line (e.g. "GET /metrics HTTP/1.1") - let mut reader = BufReader::new(&stream); + // 5-second read timeout prevents slow-loris connections + _ = stream.set_read_timeout(Some(std::time::Duration::from_secs(5))); + + // Read the HTTP request line, limited to 8KB to prevent memory abuse let mut request_line = String::new(); - if reader.read_line(&mut request_line).is_err() { + if BufReader::new((&stream).take(8192)) + .read_line(&mut request_line) + .is_err() + { continue; } let is_metrics = request_line.starts_with("GET /metrics"); - drop(reader); if is_metrics { let body = gather_metrics(); let response = format!( "HTTP/1.1 200 OK\r\n\ - Content-Type: text/plain; version=0.0.4\r\n\ + Content-Type: text/plain; version=0.0.4; charset=utf-8\r\n\ Content-Length: {}\r\n\r\n\ {}", body.len(), diff --git a/crates/ml/src/ensemble/mod.rs b/crates/ml/src/ensemble/mod.rs index b6e053a1f..90912d3ec 100644 --- a/crates/ml/src/ensemble/mod.rs +++ b/crates/ml/src/ensemble/mod.rs @@ -24,8 +24,10 @@ pub mod adapters; pub mod conviction_gates; pub mod weight_optimizer; pub mod gate_optimizer; +pub mod model_adapter; // Re-export key types that are used across ensemble modules +pub use model_adapter::{EnsembleModelAdapter, build_production_strategy}; pub use ab_testing::{ ABGroup, ABMetricsTracker, ABTestConfig, ABTestResults, ABTestRouter, GroupMetrics, Recommendation, StatisticalTestResult, diff --git a/crates/ml/src/ensemble/model_adapter.rs b/crates/ml/src/ensemble/model_adapter.rs new file mode 100644 index 000000000..60880ff19 --- /dev/null +++ b/crates/ml/src/ensemble/model_adapter.rs @@ -0,0 +1,128 @@ +//! Bridge between ml model registry and common::ml_strategy::MLModelAdapter trait. +//! +//! EnsembleModelAdapter wraps a model ID and implements MLModelAdapter so it can +//! be injected into SharedMLStrategy. When real checkpoint loading is wired, this +//! adapter will delegate to the loaded model for inference. + +use anyhow::Result; +use chrono::Utc; +use common::ml_strategy::{MLModelAdapter, MLPrediction, SharedMLStrategy}; +use crate::features::production_adapter::ProductionFeatureExtractorAdapter; + +/// Adapter that will delegate predict() to a loaded model checkpoint. +/// +/// Currently returns neutral predictions with zero confidence (filtered out +/// by the confidence threshold) until checkpoint loading is production-ready. +#[derive(Debug)] +pub struct EnsembleModelAdapter { + model_id: String, +} + +impl EnsembleModelAdapter { + pub fn new(model_id: impl Into) -> Self { + Self { + model_id: model_id.into(), + } + } +} + +impl MLModelAdapter for EnsembleModelAdapter { + fn predict(&self, features: &[f64]) -> Result { + // TODO: Wire to real model inference via checkpoint loading when + // the model registry is production-ready. For now, return neutral + // prediction so the ensemble pipeline is fully wired end-to-end. + let _ = features; + Ok(MLPrediction { + model_id: self.model_id.clone(), + prediction_value: 0.5, + confidence: 0.0, // Zero confidence = filtered out by threshold + features: vec![], + timestamp: Utc::now(), + inference_latency_us: 0, + }) + } + + fn model_id(&self) -> &str { + &self.model_id + } + + fn validate_prediction(&mut self, _prediction: &MLPrediction, _actual_outcome: bool) { + // Performance tracking handled by SharedMLStrategy + } +} + +/// Build a production-ready SharedMLStrategy with model adapters for each +/// model in the 10-model ensemble. +/// +/// Returns a strategy with: +/// - ProductionFeatureExtractorAdapter (51 features currently) +/// - One EnsembleModelAdapter per known model type +/// +/// When no checkpoints are loaded, models return neutral predictions with +/// zero confidence, which get filtered out by the confidence threshold. +pub fn build_production_strategy(min_confidence_threshold: f64) -> SharedMLStrategy { + let extractor = Box::new(ProductionFeatureExtractorAdapter::new()); + + let model_ids = [ + "dqn", "ppo", "tft", "mamba2", "tggn", "tlob", "liquid", "kan", "xlstm", "diffusion", + ]; + let models: Vec> = model_ids + .iter() + .map(|&id| Box::new(EnsembleModelAdapter::new(id)) as Box) + .collect(); + + SharedMLStrategy::new(extractor, models, min_confidence_threshold) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_ensemble_model_adapter_neutral_prediction() { + let adapter = EnsembleModelAdapter::new("dqn"); + let features = vec![0.1; 51]; + let prediction = adapter.predict(&features).unwrap(); + + assert_eq!(prediction.model_id, "dqn"); + assert_eq!(prediction.prediction_value, 0.5); + assert_eq!(prediction.confidence, 0.0); + } + + #[test] + fn test_build_production_strategy() { + let strategy = build_production_strategy(0.6); + assert_eq!(strategy.min_confidence_threshold(), 0.6); + } + + #[tokio::test] + async fn test_production_strategy_filters_neutral_predictions() { + let strategy = build_production_strategy(0.5); + let predictions = strategy + .get_ensemble_prediction(100.0, 1000.0, Utc::now()) + .await + .unwrap(); + + // All 10 adapters return confidence=0.0, so threshold=0.5 filters them all + assert!( + predictions.is_empty(), + "Neutral predictions (confidence=0.0) should be filtered by threshold=0.5" + ); + } + + #[tokio::test] + async fn test_production_strategy_zero_threshold_keeps_all() { + let strategy = build_production_strategy(0.0); + let predictions = strategy + .get_ensemble_prediction(100.0, 1000.0, Utc::now()) + .await + .unwrap(); + + // With threshold=0.0, all 10 neutral predictions pass through + assert_eq!( + predictions.len(), + 10, + "Zero threshold should keep all 10 model predictions" + ); + } +} diff --git a/services/backtesting_service/src/ml_strategy_engine.rs b/services/backtesting_service/src/ml_strategy_engine.rs index ad8073522..1813508dd 100644 --- a/services/backtesting_service/src/ml_strategy_engine.rs +++ b/services/backtesting_service/src/ml_strategy_engine.rs @@ -20,8 +20,6 @@ use config::structures::BacktestingStrategyConfig; // Import shared ML strategy (ONE SINGLE SYSTEM) use common::ml_strategy::{MLPrediction as CommonMLPrediction, SharedMLStrategy}; -// Import ProductionFeatureExtractorAdapter for 225-feature extraction -use ml::features::production_adapter::ProductionFeatureExtractorAdapter; // Import UnifiedFeatureExtractor (256 features, production system) use ml::features::unified::{FeatureExtractionConfig, UnifiedFeatureExtractor}; @@ -119,12 +117,7 @@ impl MLPoweredStrategy { pub fn new(name: String, _lookback_periods: usize) -> Result { // Use shared ML strategy (ONE SINGLE SYSTEM) with ProductionFeatureExtractorAdapter (225 features) let min_confidence_threshold = 0.6; - let production_extractor = Box::new(ProductionFeatureExtractorAdapter::new()); - let strategy = Arc::new(SharedMLStrategy::new( - production_extractor, - vec![], // Models injected via EnsembleModelAdapter once registry is available - min_confidence_threshold, - )); + let strategy = Arc::new(ml::ensemble::build_production_strategy(min_confidence_threshold)); // Initialize UnifiedFeatureExtractor (256 features) let feature_config = FeatureExtractionConfig::default();