From c4693403d5306182f90e02c23831a8efc00a9a57 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Mon, 2 Mar 2026 21:25:56 +0100 Subject: [PATCH] feat(ml): wire hyperopt Prometheus metrics into training binaries Co-Authored-By: Claude Opus 4.6 --- crates/ml/examples/hyperopt_baseline_rl.rs | 15 +++++++ .../examples/hyperopt_baseline_supervised.rs | 45 +++++++++++++++++++ 2 files changed, 60 insertions(+) diff --git a/crates/ml/examples/hyperopt_baseline_rl.rs b/crates/ml/examples/hyperopt_baseline_rl.rs index 52a258dc1..fc5618044 100644 --- a/crates/ml/examples/hyperopt_baseline_rl.rs +++ b/crates/ml/examples/hyperopt_baseline_rl.rs @@ -166,6 +166,8 @@ fn run_dqn_hyperopt(args: &Args, parallel: usize, device: &candle_core::Device) .seed(args.seed) .build(); + training_metrics::set_hyperopt_trial("dqn", 0.0, args.trials as f64); + let start = Instant::now(); let result = if parallel > 1 { info!("Using parallel optimization ({} threads)", parallel); @@ -179,6 +181,9 @@ fn run_dqn_hyperopt(args: &Args, parallel: usize, device: &candle_core::Device) }; let elapsed = start.elapsed().as_secs_f64(); + training_metrics::set_hyperopt_trial("dqn", result.all_trials.len() as f64, args.trials as f64); + training_metrics::set_hyperopt_best_objective("dqn", result.best_objective); + info!("DQN hyperopt complete:"); info!(" Best objective: {:.6}", result.best_objective); info!(" Total trials: {}", result.all_trials.len()); @@ -230,6 +235,8 @@ fn run_ppo_hyperopt(args: &Args, parallel: usize, device: &candle_core::Device) .seed(args.seed) .build(); + training_metrics::set_hyperopt_trial("ppo", 0.0, args.trials as f64); + let start = Instant::now(); let result = if parallel > 1 { info!("Using parallel optimization ({} threads)", parallel); @@ -243,6 +250,9 @@ fn run_ppo_hyperopt(args: &Args, parallel: usize, device: &candle_core::Device) }; let elapsed = start.elapsed().as_secs_f64(); + training_metrics::set_hyperopt_trial("ppo", result.all_trials.len() as f64, args.trials as f64); + training_metrics::set_hyperopt_best_objective("ppo", result.best_objective); + info!("PPO hyperopt complete:"); info!(" Best objective: {:.6}", result.best_objective); info!(" Total trials: {}", result.all_trials.len()); @@ -281,6 +291,10 @@ fn main() -> Result<()> { let args = Args::parse(); + // Signal hyperopt mode active — will be cleared at exit + let hyperopt_model_label = args.model.clone(); + training_metrics::set_hyperopt_mode(&hyperopt_model_label, true); + info!("========================================"); info!(" Hyperopt Baseline Runner"); info!("========================================"); @@ -430,6 +444,7 @@ fn main() -> Result<()> { info!("========================================"); info!("{}", output_str); + training_metrics::set_hyperopt_mode(&hyperopt_model_label, false); training_metrics::set_active_workers(0.0); Ok(()) } diff --git a/crates/ml/examples/hyperopt_baseline_supervised.rs b/crates/ml/examples/hyperopt_baseline_supervised.rs index 742338b02..dfff07887 100644 --- a/crates/ml/examples/hyperopt_baseline_supervised.rs +++ b/crates/ml/examples/hyperopt_baseline_supervised.rs @@ -138,12 +138,17 @@ fn run_tft_hyperopt(args: &Args) -> Result { .seed(args.seed) .build(); + training_metrics::set_hyperopt_trial("tft", 0.0, args.trials as f64); + let start = Instant::now(); let result = optimizer .optimize(trainer) .context("TFT hyperopt optimization failed")?; let elapsed = start.elapsed().as_secs_f64(); + training_metrics::set_hyperopt_trial("tft", result.all_trials.len() as f64, args.trials as f64); + training_metrics::set_hyperopt_best_objective("tft", result.best_objective); + info!("TFT hyperopt complete:"); info!(" Best objective: {:.6}", result.best_objective); info!(" Total trials: {}", result.all_trials.len()); @@ -191,12 +196,17 @@ fn run_mamba2_hyperopt(args: &Args) -> Result { .seed(args.seed) .build(); + training_metrics::set_hyperopt_trial("mamba2", 0.0, args.trials as f64); + let start = Instant::now(); let result = optimizer .optimize(trainer) .context("Mamba2 hyperopt optimization failed")?; let elapsed = start.elapsed().as_secs_f64(); + training_metrics::set_hyperopt_trial("mamba2", result.all_trials.len() as f64, args.trials as f64); + training_metrics::set_hyperopt_best_objective("mamba2", result.best_objective); + info!("Mamba2 hyperopt complete:"); info!(" Best objective: {:.6}", result.best_objective); info!(" Total trials: {}", result.all_trials.len()); @@ -245,12 +255,17 @@ fn run_liquid_hyperopt(args: &Args) -> Result { .seed(args.seed) .build(); + training_metrics::set_hyperopt_trial("liquid", 0.0, args.trials as f64); + let start = Instant::now(); let result = optimizer .optimize(trainer) .context("Liquid hyperopt optimization failed")?; let elapsed = start.elapsed().as_secs_f64(); + training_metrics::set_hyperopt_trial("liquid", result.all_trials.len() as f64, args.trials as f64); + training_metrics::set_hyperopt_best_objective("liquid", result.best_objective); + info!("Liquid hyperopt complete:"); info!(" Best objective: {:.6}", result.best_objective); info!(" Total trials: {}", result.all_trials.len()); @@ -299,12 +314,17 @@ fn run_tggn_hyperopt(args: &Args) -> Result { .seed(args.seed) .build(); + training_metrics::set_hyperopt_trial("tggn", 0.0, args.trials as f64); + let start = Instant::now(); let result = optimizer .optimize(trainer) .context("TGGN hyperopt optimization failed")?; let elapsed = start.elapsed().as_secs_f64(); + training_metrics::set_hyperopt_trial("tggn", result.all_trials.len() as f64, args.trials as f64); + training_metrics::set_hyperopt_best_objective("tggn", result.best_objective); + info!("TGGN hyperopt complete:"); info!(" Best objective: {:.6}", result.best_objective); info!(" Total trials: {}", result.all_trials.len()); @@ -353,12 +373,17 @@ fn run_tlob_hyperopt(args: &Args) -> Result { .seed(args.seed) .build(); + training_metrics::set_hyperopt_trial("tlob", 0.0, args.trials as f64); + let start = Instant::now(); let result = optimizer .optimize(trainer) .context("TLOB hyperopt optimization failed")?; let elapsed = start.elapsed().as_secs_f64(); + training_metrics::set_hyperopt_trial("tlob", result.all_trials.len() as f64, args.trials as f64); + training_metrics::set_hyperopt_best_objective("tlob", result.best_objective); + info!("TLOB hyperopt complete:"); info!(" Best objective: {:.6}", result.best_objective); info!(" Total trials: {}", result.all_trials.len()); @@ -407,12 +432,17 @@ fn run_kan_hyperopt(args: &Args) -> Result { .seed(args.seed) .build(); + training_metrics::set_hyperopt_trial("kan", 0.0, args.trials as f64); + let start = Instant::now(); let result = optimizer .optimize(trainer) .context("KAN hyperopt optimization failed")?; let elapsed = start.elapsed().as_secs_f64(); + training_metrics::set_hyperopt_trial("kan", result.all_trials.len() as f64, args.trials as f64); + training_metrics::set_hyperopt_best_objective("kan", result.best_objective); + info!("KAN hyperopt complete:"); info!(" Best objective: {:.6}", result.best_objective); info!(" Total trials: {}", result.all_trials.len()); @@ -461,12 +491,17 @@ fn run_xlstm_hyperopt(args: &Args) -> Result { .seed(args.seed) .build(); + training_metrics::set_hyperopt_trial("xlstm", 0.0, args.trials as f64); + let start = Instant::now(); let result = optimizer .optimize(trainer) .context("xLSTM hyperopt optimization failed")?; let elapsed = start.elapsed().as_secs_f64(); + training_metrics::set_hyperopt_trial("xlstm", result.all_trials.len() as f64, args.trials as f64); + training_metrics::set_hyperopt_best_objective("xlstm", result.best_objective); + info!("xLSTM hyperopt complete:"); info!(" Best objective: {:.6}", result.best_objective); info!(" Total trials: {}", result.all_trials.len()); @@ -515,12 +550,17 @@ fn run_diffusion_hyperopt(args: &Args) -> Result { .seed(args.seed) .build(); + training_metrics::set_hyperopt_trial("diffusion", 0.0, args.trials as f64); + let start = Instant::now(); let result = optimizer .optimize(trainer) .context("Diffusion hyperopt optimization failed")?; let elapsed = start.elapsed().as_secs_f64(); + training_metrics::set_hyperopt_trial("diffusion", result.all_trials.len() as f64, args.trials as f64); + training_metrics::set_hyperopt_best_objective("diffusion", result.best_objective); + info!("Diffusion hyperopt complete:"); info!(" Best objective: {:.6}", result.best_objective); info!(" Total trials: {}", result.all_trials.len()); @@ -559,6 +599,10 @@ fn main() -> Result<()> { let args = Args::parse(); + // Signal hyperopt mode active — will be cleared at exit + let hyperopt_model_label = args.model.clone(); + training_metrics::set_hyperopt_mode(&hyperopt_model_label, true); + info!("========================================"); info!(" Hyperopt Baseline Supervised Runner"); info!("========================================"); @@ -652,6 +696,7 @@ fn main() -> Result<()> { info!("========================================"); info!("{}", output_str); + training_metrics::set_hyperopt_mode(&hyperopt_model_label, false); training_metrics::set_active_workers(0.0); Ok(()) }