From cfafc0d13e7e5033b3d2dcbb014eb6aacd36bdf5 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sun, 1 Mar 2026 15:01:33 +0100 Subject: [PATCH] feat(training): add Prometheus metrics export to training binaries Create shared baseline_common/metrics.rs module that registers all 18 dashboard-expected metrics (11 gauges, 5 counters, 2 histograms) and spawns a lightweight HTTP metrics server on port 9094. Instrument train_baseline_supervised and train_baseline_rl with: - Epoch progress, training/validation loss gauges - Checkpoint save timing, size, and failure counters - NaN/gradient explosion detection counters - Data loading latency histograms - Active workers lifecycle (1 on start, 0 on exit) Co-Authored-By: Claude Opus 4.6 --- crates/ml/examples/baseline_common/metrics.rs | 258 ++++++++++++++++++ crates/ml/examples/baseline_common/mod.rs | 1 + crates/ml/examples/train_baseline_rl.rs | 69 ++++- .../ml/examples/train_baseline_supervised.rs | 46 +++- 4 files changed, 371 insertions(+), 3 deletions(-) create mode 100644 crates/ml/examples/baseline_common/metrics.rs diff --git a/crates/ml/examples/baseline_common/metrics.rs b/crates/ml/examples/baseline_common/metrics.rs new file mode 100644 index 000000000..c4a3ecad3 --- /dev/null +++ b/crates/ml/examples/baseline_common/metrics.rs @@ -0,0 +1,258 @@ +//! Prometheus metrics for training binaries. +//! +//! Registers 18 dashboard-expected metrics and exposes them via a lightweight +//! HTTP endpoint on `/metrics` (default port 9094, matching K8s annotations). + +use std::io::{BufRead, BufReader, Write as IoWrite}; +use std::net::TcpListener; + +/// Register all 18 training metrics with the global Prometheus registry. +/// +/// Safe to call multiple times — `common::metrics` silently ignores duplicate +/// registrations. +pub fn init_training_metrics() { + use common::metrics::{ + register_counter_vec, register_gauge, register_gauge_vec, register_histogram_vec, + }; + + let mf = &["model", "fold"]; + let m = &["model"]; + + // Gauges (model + fold) + let _ = register_gauge_vec("foxhunt_training_current_epoch", "Current epoch number", mf); + let _ = register_gauge_vec("foxhunt_training_epoch_loss", "Training loss at end of epoch", mf); + let _ = register_gauge_vec( + "foxhunt_training_validation_loss", + "Validation loss at end of epoch", + mf, + ); + let _ = register_gauge_vec( + "foxhunt_training_batches_per_second", + "Training throughput in batches per second", + mf, + ); + let _ = register_gauge_vec( + "foxhunt_training_batches_processed", + "Total batches processed this run", + mf, + ); + let _ = register_gauge_vec( + "foxhunt_training_iteration_seconds", + "Last epoch wall time in seconds", + mf, + ); + let _ = register_gauge_vec( + "foxhunt_training_eval_accuracy", + "Eval accuracy (supervised models)", + mf, + ); + let _ = register_gauge_vec( + "foxhunt_training_eval_precision", + "Eval precision (supervised models)", + mf, + ); + let _ = register_gauge_vec( + "foxhunt_training_eval_recall", + "Eval recall (supervised models)", + mf, + ); + let _ = register_gauge_vec( + "foxhunt_training_eval_f1", + "Eval F1 score (supervised models)", + mf, + ); + let _ = register_gauge_vec( + "foxhunt_training_checkpoint_size_bytes", + "Size of last checkpoint file in bytes", + mf, + ); + + // Counters (model + fold) + let _ = register_counter_vec( + "foxhunt_training_checkpoint_saves_total", + "Successful checkpoint saves", + mf, + ); + let _ = register_counter_vec( + "foxhunt_training_checkpoint_failures_total", + "Failed checkpoint saves", + mf, + ); + let _ = register_counter_vec( + "foxhunt_training_nan_detected_total", + "NaN loss events detected", + mf, + ); + let _ = register_counter_vec( + "foxhunt_training_gradient_explosion_total", + "Gradient explosion events", + mf, + ); + let _ = register_counter_vec( + "foxhunt_training_feature_errors_total", + "Feature extraction failures", + mf, + ); + + // Histograms (model only) + let _ = register_histogram_vec( + "foxhunt_training_checkpoint_duration_seconds", + "Checkpoint save latency", + m, + vec![0.1, 0.5, 1.0, 2.0, 5.0, 10.0, 30.0], + ); + let _ = register_histogram_vec( + "foxhunt_training_data_load_seconds", + "Data loading latency per fold", + m, + vec![0.5, 1.0, 5.0, 10.0, 30.0, 60.0, 120.0, 300.0], + ); + + // Special gauge (no labels) + let _ = register_gauge("foxhunt_training_active_workers", "Number of active training workers"); +} + +/// Spawn a background daemon thread serving Prometheus metrics over HTTP. +/// +/// Responds to `GET /metrics` with the global registry output. Any other +/// request gets a 404. The thread is detached and dies with the process. +pub fn start_metrics_server(port: u16) { + std::thread::spawn(move || { + let addr = format!("0.0.0.0:{port}"); + let listener = match TcpListener::bind(&addr) { + Ok(l) => { + tracing::info!("Prometheus metrics server listening on {addr}"); + l + } + Err(e) => { + tracing::warn!("Failed to bind metrics server on {addr}: {e}"); + return; + } + }; + + for stream in listener.incoming() { + let Ok(mut stream) = stream else { + continue; + }; + + // Read request line (e.g. "GET /metrics HTTP/1.1") + let mut reader = BufReader::new(&stream); + let mut request_line = String::new(); + if reader.read_line(&mut request_line).is_err() { + continue; + } + + let is_metrics = request_line.starts_with("GET /metrics"); + drop(reader); + + if is_metrics { + let body = common::metrics::gather_metrics(); + let response = format!( + "HTTP/1.1 200 OK\r\nContent-Type: text/plain; version=0.0.4\r\nContent-Length: {}\r\n\r\n{}", + body.len(), + body + ); + let _ = stream.write_all(response.as_bytes()); + } else { + let _ = stream.write_all(b"HTTP/1.1 404 Not Found\r\nContent-Length: 0\r\n\r\n"); + } + } + }); +} + +// --------------------------------------------------------------------------- +// Helper functions wrapping common::metrics +// --------------------------------------------------------------------------- + +pub fn set_epoch(model: &str, fold: &str, epoch: f64) { + common::metrics::set_gauge_vec("foxhunt_training_current_epoch", &[model, fold], epoch); +} + +pub fn set_epoch_loss(model: &str, fold: &str, loss: f64) { + common::metrics::set_gauge_vec("foxhunt_training_epoch_loss", &[model, fold], loss); +} + +pub fn set_validation_loss(model: &str, fold: &str, loss: f64) { + common::metrics::set_gauge_vec("foxhunt_training_validation_loss", &[model, fold], loss); +} + +pub fn set_batches_per_second(model: &str, fold: &str, bps: f64) { + common::metrics::set_gauge_vec("foxhunt_training_batches_per_second", &[model, fold], bps); +} + +pub fn set_batches_processed(model: &str, fold: &str, total: f64) { + common::metrics::set_gauge_vec("foxhunt_training_batches_processed", &[model, fold], total); +} + +pub fn set_iteration_seconds(model: &str, fold: &str, secs: f64) { + common::metrics::set_gauge_vec("foxhunt_training_iteration_seconds", &[model, fold], secs); +} + +pub fn set_eval_metrics(model: &str, fold: &str, accuracy: f64, precision: f64, recall: f64, f1: f64) { + common::metrics::set_gauge_vec("foxhunt_training_eval_accuracy", &[model, fold], accuracy); + common::metrics::set_gauge_vec("foxhunt_training_eval_precision", &[model, fold], precision); + common::metrics::set_gauge_vec("foxhunt_training_eval_recall", &[model, fold], recall); + common::metrics::set_gauge_vec("foxhunt_training_eval_f1", &[model, fold], f1); +} + +pub fn record_checkpoint_save(model: &str, fold: &str, duration_secs: f64, size_bytes: f64) { + common::metrics::increment_counter_vec( + "foxhunt_training_checkpoint_saves_total", + &[model, fold], + 1.0, + ); + common::metrics::set_gauge_vec( + "foxhunt_training_checkpoint_size_bytes", + &[model, fold], + size_bytes, + ); + common::metrics::observe_histogram_vec( + "foxhunt_training_checkpoint_duration_seconds", + &[model], + duration_secs, + ); +} + +pub fn record_checkpoint_failure(model: &str, fold: &str) { + common::metrics::increment_counter_vec( + "foxhunt_training_checkpoint_failures_total", + &[model, fold], + 1.0, + ); +} + +pub fn record_nan_detected(model: &str, fold: &str) { + common::metrics::increment_counter_vec( + "foxhunt_training_nan_detected_total", + &[model, fold], + 1.0, + ); +} + +pub fn record_gradient_explosion(model: &str, fold: &str) { + common::metrics::increment_counter_vec( + "foxhunt_training_gradient_explosion_total", + &[model, fold], + 1.0, + ); +} + +pub fn record_feature_error(model: &str, fold: &str) { + common::metrics::increment_counter_vec( + "foxhunt_training_feature_errors_total", + &[model, fold], + 1.0, + ); +} + +pub fn record_data_load(model: &str, duration_secs: f64) { + common::metrics::observe_histogram_vec( + "foxhunt_training_data_load_seconds", + &[model], + duration_secs, + ); +} + +pub fn set_active_workers(value: f64) { + common::metrics::set_gauge("foxhunt_training_active_workers", value); +} diff --git a/crates/ml/examples/baseline_common/mod.rs b/crates/ml/examples/baseline_common/mod.rs index 5b60d647c..62c83a9e5 100644 --- a/crates/ml/examples/baseline_common/mod.rs +++ b/crates/ml/examples/baseline_common/mod.rs @@ -6,6 +6,7 @@ #[allow(dead_code)] pub mod completion; +pub mod metrics; use std::path::{Path, PathBuf}; diff --git a/crates/ml/examples/train_baseline_rl.rs b/crates/ml/examples/train_baseline_rl.rs index 5cff268a5..55440d1c5 100644 --- a/crates/ml/examples/train_baseline_rl.rs +++ b/crates/ml/examples/train_baseline_rl.rs @@ -34,6 +34,7 @@ use ml::dqn::{DQNConfig, Experience, DQN}; #[allow(unreachable_pub)] mod baseline_common; use baseline_common::completion::{write_failure_marker, write_success_marker, CompletionMetrics}; +use baseline_common::metrics; use baseline_common::{load_all_bars, spread_cost_bps}; use ml::features::extraction::extract_ml_features; use ml::ppo::ppo::{PPOConfig, PPO}; @@ -281,8 +282,10 @@ fn train_dqn_fold( let mut best_val_loss = f64::MAX; let mut epochs_without_improvement = 0_usize; let mut rng = rand::thread_rng(); + let fold_str = fold.to_string(); for epoch in 0..args.epochs { + let epoch_start = std::time::Instant::now(); // --- Training pass --- let mut epoch_loss = 0.0_f64; let mut epoch_steps = 0_usize; @@ -338,6 +341,9 @@ fn train_dqn_fold( // Train step (returns (loss, grad_norm)) if let Ok((loss, _grad_norm)) = dqn.train_step(None) { + if loss.is_nan() || loss.is_infinite() { + metrics::record_nan_detected("dqn", &fold_str); + } epoch_loss += loss as f64; epoch_steps += 1; } @@ -352,6 +358,21 @@ fn train_dqn_fold( // --- Validation pass --- let val_loss = evaluate_dqn_validation(&mut dqn, val_features, val_bars, args); + // Prometheus metrics + let epoch_elapsed = epoch_start.elapsed(); + metrics::set_epoch("dqn", &fold_str, (epoch + 1) as f64); + metrics::set_epoch_loss("dqn", &fold_str, avg_train_loss); + metrics::set_validation_loss("dqn", &fold_str, val_loss); + metrics::set_batches_processed("dqn", &fold_str, epoch_steps as f64); + metrics::set_iteration_seconds("dqn", &fold_str, epoch_elapsed.as_secs_f64()); + if epoch_elapsed.as_secs_f64() > 0.001 { + metrics::set_batches_per_second( + "dqn", + &fold_str, + epoch_steps as f64 / epoch_elapsed.as_secs_f64(), + ); + } + // Decay epsilon at epoch level let current_eps = dqn.get_epsilon(); let new_eps = (current_eps * 0.995_f32).max(0.05); @@ -374,9 +395,18 @@ fn train_dqn_fold( // Save best checkpoint let ckpt_path = output_dir.join(format!("dqn_fold{}_best.safetensors", fold)); + let ckpt_start = std::time::Instant::now(); if let Err(e) = dqn.get_q_network_vars().save(ckpt_path.to_string_lossy().as_ref()) { + metrics::record_checkpoint_failure("dqn", &fold_str); warn!(" [DQN] Failed to save checkpoint: {}", e); } else { + let ckpt_size = ckpt_path.metadata().map(|m| m.len() as f64).unwrap_or(0.0); + metrics::record_checkpoint_save( + "dqn", + &fold_str, + ckpt_start.elapsed().as_secs_f64(), + ckpt_size, + ); info!(" [DQN] Saved best checkpoint: {}", ckpt_path.display()); } } else { @@ -508,6 +538,7 @@ fn train_ppo_fold( let mut best_val_loss = f64::MAX; let mut epochs_without_improvement = 0_usize; + let fold_str = fold.to_string(); let n_train = train_features.len(); if n_train < 2 { @@ -516,6 +547,7 @@ fn train_ppo_fold( } for epoch in 0..args.epochs { + let epoch_start = std::time::Instant::now(); // --- Collect trajectory --- let trajectory = collect_ppo_trajectory( &ppo, @@ -551,6 +583,16 @@ fn train_ppo_fold( // --- Validation pass --- let val_loss = evaluate_ppo_validation(&ppo, val_features, val_bars, args.tx_cost_bps, args.tick_size, args.spread_ticks); + // Prometheus metrics + let epoch_elapsed = epoch_start.elapsed(); + metrics::set_epoch("ppo", &fold_str, (epoch + 1) as f64); + metrics::set_epoch_loss("ppo", &fold_str, policy_loss as f64); + metrics::set_validation_loss("ppo", &fold_str, val_loss); + metrics::set_iteration_seconds("ppo", &fold_str, epoch_elapsed.as_secs_f64()); + if policy_loss.is_nan() || value_loss.is_nan() { + metrics::record_nan_detected("ppo", &fold_str); + } + info!( " [PPO] Fold {} Epoch {}/{} -- policy_loss={:.6} value_loss={:.6} val_metric={:.6}", fold, @@ -571,13 +613,22 @@ fn train_ppo_fold( let critic_path = output_dir.join(format!("ppo_fold{}_critic.safetensors", fold)); let meta_path = output_dir.join(format!("ppo_fold{}_meta.json", fold)); + let ckpt_start = std::time::Instant::now(); if let Err(e) = ppo.save_checkpoint( &actor_path.to_string_lossy(), &critic_path.to_string_lossy(), &meta_path.to_string_lossy(), ) { + metrics::record_checkpoint_failure("ppo", &fold_str); warn!(" [PPO] Failed to save checkpoint: {}", e); } else { + let ckpt_size = actor_path.metadata().map(|m| m.len() as f64).unwrap_or(0.0); + metrics::record_checkpoint_save( + "ppo", + &fold_str, + ckpt_start.elapsed().as_secs_f64(), + ckpt_size, + ); info!(" [PPO] Saved best checkpoint: {}", actor_path.display()); } } else { @@ -770,6 +821,7 @@ fn run_training(args: &Args) -> Result> { // 1. Load all OHLCV bars from DBN files info!("Step 1/5: Loading OHLCV bars from DBN files..."); + let data_load_start = std::time::Instant::now(); let bars = load_all_bars(&args.data_dir, &args.symbol)?; if bars.is_empty() { anyhow::bail!("No bars loaded from {}", args.data_dir.display()); @@ -811,6 +863,14 @@ fn run_training(args: &Args) -> Result> { } info!(" Generated {} walk-forward folds", windows.len()); + // Record data loading + feature extraction time + if train_dqn { + metrics::record_data_load("dqn", data_load_start.elapsed().as_secs_f64()); + } + if train_ppo { + metrics::record_data_load("ppo", data_load_start.elapsed().as_secs_f64()); + } + // Create output directory std::fs::create_dir_all(&args.output_dir) .with_context(|| format!("Failed to create output dir: {}", args.output_dir.display()))?; @@ -978,6 +1038,10 @@ fn main() -> Result<()> { ) .init(); + metrics::init_training_metrics(); + metrics::start_metrics_server(9094); + metrics::set_active_workers(1.0); + let args = Args::parse(); // Ensure output directory exists before training so markers can always be written. @@ -985,7 +1049,10 @@ fn main() -> Result<()> { error!("Failed to create output dir {}: {}", args.output_dir.display(), e); } - match run_training(&args) { + let result = run_training(&args); + metrics::set_active_workers(0.0); + + match result { Ok(results) => { for result in &results { let best_val = result diff --git a/crates/ml/examples/train_baseline_supervised.rs b/crates/ml/examples/train_baseline_supervised.rs index a4254eef8..50d99eb97 100644 --- a/crates/ml/examples/train_baseline_supervised.rs +++ b/crates/ml/examples/train_baseline_supervised.rs @@ -42,6 +42,7 @@ use ml::walk_forward::{generate_walk_forward_windows, NormStats, WalkForwardConf #[allow(unreachable_pub)] mod baseline_common; use baseline_common::completion::{write_failure_marker, write_success_marker, CompletionMetrics}; +use baseline_common::metrics; use baseline_common::{load_all_bars, spread_cost_bps}; // Model adapter imports -- using verified paths from existing examples @@ -459,6 +460,8 @@ fn run_training_epoch( adapter: &mut dyn UnifiedTrainable, train_pairs: &[(Tensor, Tensor)], batch_size: usize, + model_name: &str, + fold_str: &str, ) -> Result { let mut epoch_loss_sum = 0.0_f64; let mut epoch_steps = 0_usize; @@ -497,6 +500,10 @@ fn run_training_epoch( .to_scalar::() .map_err(|e| anyhow::anyhow!("loss to_scalar failed: {}", e))?; + if loss_val.is_nan() || loss_val.is_infinite() { + metrics::record_nan_detected(model_name, fold_str); + } + adapter .backward(&loss) .map_err(|e| anyhow::anyhow!("backward failed: {}", e))?; @@ -509,6 +516,8 @@ fn run_training_epoch( batch_start = batch_end; } + metrics::set_batches_processed(model_name, fold_str, epoch_steps as f64); + Ok(if epoch_steps > 0 { epoch_loss_sum / epoch_steps as f64 } else { @@ -539,17 +548,24 @@ fn train_fold( let mut best_val_loss = f64::MAX; let mut epochs_without_improvement = 0_usize; + let fold_str = fold.to_string(); for epoch in 0..args.epochs { let epoch_start = Instant::now(); let avg_train_loss = - run_training_epoch(adapter.as_mut(), train_pairs, args.batch_size)?; + run_training_epoch(adapter.as_mut(), train_pairs, args.batch_size, model_name, &fold_str)?; let val_loss = adapter .validate(val_pairs) .map_err(|e| anyhow::anyhow!("validation failed: {}", e))?; let elapsed = epoch_start.elapsed(); + // Prometheus metrics + metrics::set_epoch(model_name, &fold_str, (epoch + 1) as f64); + metrics::set_epoch_loss(model_name, &fold_str, avg_train_loss); + metrics::set_validation_loss(model_name, &fold_str, val_loss); + metrics::set_iteration_seconds(model_name, &fold_str, elapsed.as_secs_f64()); + info!( " Fold {} Epoch {}/{}: train_loss={:.6}, val_loss={:.6}, lr={:.2e}, time={:.1}s", fold, @@ -566,9 +582,22 @@ fn train_fold( epochs_without_improvement = 0; let ckpt_path = output_dir.join(format!("{}_fold{}_best", model_name, fold)); let ckpt_str = ckpt_path.to_str().unwrap_or("checkpoint"); + let ckpt_start = Instant::now(); if let Err(e) = adapter.save_checkpoint(ckpt_str) { + metrics::record_checkpoint_failure(model_name, &fold_str); warn!(" Failed to save checkpoint: {}", e); } else { + let ckpt_size = ckpt_path + .with_extension("safetensors") + .metadata() + .map(|m| m.len() as f64) + .unwrap_or(0.0); + metrics::record_checkpoint_save( + model_name, + &fold_str, + ckpt_start.elapsed().as_secs_f64(), + ckpt_size, + ); info!( " [{}] New best val_loss={:.6}, checkpoint saved", model_name, val_loss @@ -637,6 +666,7 @@ fn run_training(args: &Args) -> Result> { args.symbol, args.data_dir.display() ); + let data_load_start = Instant::now(); let all_bars = load_all_bars(&args.data_dir, &args.symbol)?; info!("Loaded {} bars", all_bars.len()); @@ -664,6 +694,11 @@ fn run_training(args: &Args) -> Result> { let mut all_results = Vec::new(); + // Record data loading time (shared across all models) + for mn in &models_to_train { + metrics::record_data_load(mn, data_load_start.elapsed().as_secs_f64()); + } + // Train each model for model_name in &models_to_train { info!("=== Training model: {} ===", model_name); @@ -769,6 +804,10 @@ fn main() -> Result<()> { ) .init(); + metrics::init_training_metrics(); + metrics::start_metrics_server(9094); + metrics::set_active_workers(1.0); + let args = Args::parse(); // Ensure output directory exists before training so markers can always be written. @@ -776,7 +815,10 @@ fn main() -> Result<()> { error!("Failed to create output dir {}: {}", args.output_dir.display(), e); } - match run_training(&args) { + let result = run_training(&args); + metrics::set_active_workers(0.0); + + match result { Ok(results) => { for result in &results { let best_val = result