From 9b85b9efd596ed8a69a22dbfaaa01446cac577c3 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sun, 1 Mar 2026 02:36:05 +0100 Subject: [PATCH] feat(ml): add Pushgateway metrics client for training jobs Training jobs are short-lived K8s Jobs that can't be reliably scraped by Prometheus. This adds a TrainingMetricsPusher that pushes epoch-level metrics (loss, accuracy, throughput, gradient health, checkpoints) to the Pushgateway so they persist for Prometheus to scrape. Co-Authored-By: Claude Opus 4.6 --- crates/ml/src/training.rs | 1 + crates/ml/src/training/push_metrics.rs | 211 +++++++++++++++++++++++++ 2 files changed, 212 insertions(+) create mode 100644 crates/ml/src/training/push_metrics.rs diff --git a/crates/ml/src/training.rs b/crates/ml/src/training.rs index 410b6be92..4a3c79519 100644 --- a/crates/ml/src/training.rs +++ b/crates/ml/src/training.rs @@ -12,6 +12,7 @@ // Sub-modules for specialized training components pub mod orchestrator; +pub mod push_metrics; pub mod unified_data_loader; pub mod unified_trainer; // NEW: Unified training trait for all models // NEW: Model-agnostic training orchestrator diff --git a/crates/ml/src/training/push_metrics.rs b/crates/ml/src/training/push_metrics.rs new file mode 100644 index 000000000..01d8d166c --- /dev/null +++ b/crates/ml/src/training/push_metrics.rs @@ -0,0 +1,211 @@ +//! Prometheus Pushgateway integration for training job metrics +//! +//! Training jobs are short-lived K8s Jobs. They push epoch-level metrics +//! to the Pushgateway so Prometheus can scrape them persistently. + +use std::fmt::Write; + +/// Pushgateway client for training metrics +#[derive(Debug)] +pub struct TrainingMetricsPusher { + pushgateway_url: String, + job_id: String, + model_name: String, + client: reqwest::Client, +} + +impl TrainingMetricsPusher { + /// Create a new pusher. Falls back to in-cluster Pushgateway if `PUSHGATEWAY_URL` not set. + pub fn new(job_id: &str, model_name: &str) -> Self { + let pushgateway_url = std::env::var("PUSHGATEWAY_URL") + .unwrap_or_else(|_| "http://pushgateway.foxhunt.svc.cluster.local:9091".to_string()); + Self { + pushgateway_url, + job_id: job_id.to_string(), + model_name: model_name.to_string(), + client: reqwest::Client::new(), + } + } + + /// Push current training state to Pushgateway. + /// Non-fatal: logs errors but never panics. + pub async fn push(&self, state: &TrainingState) { + let mut body = String::new(); + + // Epoch progress + let _ = writeln!(body, "# HELP foxhunt_training_current_epoch Current training epoch"); + let _ = writeln!(body, "# TYPE foxhunt_training_current_epoch gauge"); + let _ = writeln!(body, "foxhunt_training_current_epoch {}", state.epoch); + + // Loss + let _ = writeln!( + body, + "# HELP foxhunt_training_epoch_loss Training loss value" + ); + let _ = writeln!(body, "# TYPE foxhunt_training_epoch_loss gauge"); + let _ = writeln!( + body, + "foxhunt_training_epoch_loss{{split=\"train\"}} {}", + state.train_loss + ); + if let Some(val_loss) = state.val_loss { + let _ = writeln!( + body, + "foxhunt_training_epoch_loss{{split=\"val\"}} {}", + val_loss + ); + } + + // Accuracy + if let Some(accuracy) = state.accuracy { + let _ = writeln!( + body, + "# HELP foxhunt_training_eval_accuracy Model evaluation accuracy" + ); + let _ = writeln!(body, "# TYPE foxhunt_training_eval_accuracy gauge"); + let _ = writeln!(body, "foxhunt_training_eval_accuracy {}", accuracy); + } + + // Batches processed + let _ = writeln!( + body, + "# HELP foxhunt_training_batches_processed Total batches processed" + ); + let _ = writeln!(body, "# TYPE foxhunt_training_batches_processed counter"); + let _ = writeln!( + body, + "foxhunt_training_batches_processed {}", + state.batches_processed + ); + + // Training speed + if state.batches_per_second > 0.0 { + let _ = writeln!( + body, + "# HELP foxhunt_training_batches_per_second Training throughput" + ); + let _ = writeln!(body, "# TYPE foxhunt_training_batches_per_second gauge"); + let _ = writeln!( + body, + "foxhunt_training_batches_per_second {}", + state.batches_per_second + ); + } + + // Learning rate + let _ = writeln!( + body, + "# HELP foxhunt_training_learning_rate Current learning rate" + ); + let _ = writeln!(body, "# TYPE foxhunt_training_learning_rate gauge"); + let _ = writeln!( + body, + "foxhunt_training_learning_rate {}", + state.learning_rate + ); + + // NaN / gradient events + if state.nan_count > 0 { + let _ = writeln!( + body, + "# HELP foxhunt_training_nan_detected_total NaN gradient events" + ); + let _ = writeln!(body, "# TYPE foxhunt_training_nan_detected_total counter"); + let _ = writeln!( + body, + "foxhunt_training_nan_detected_total {}", + state.nan_count + ); + } + if state.gradient_explosion_count > 0 { + let _ = writeln!( + body, + "# HELP foxhunt_training_gradient_explosion_total Gradient clipping events" + ); + let _ = writeln!( + body, + "# TYPE foxhunt_training_gradient_explosion_total counter" + ); + let _ = writeln!( + body, + "foxhunt_training_gradient_explosion_total {}", + state.gradient_explosion_count + ); + } + + // Checkpoint saves + let _ = writeln!( + body, + "# HELP foxhunt_training_checkpoint_saves_total Checkpoints saved" + ); + let _ = writeln!( + body, + "# TYPE foxhunt_training_checkpoint_saves_total counter" + ); + let _ = writeln!( + body, + "foxhunt_training_checkpoint_saves_total {}", + state.checkpoint_saves + ); + + let url = format!( + "{}/metrics/job/{}/model/{}", + self.pushgateway_url, self.job_id, self.model_name + ); + + match self.client.put(&url).body(body).send().await { + Ok(resp) if resp.status().is_success() => {} + Ok(resp) => { + eprintln!( + "Pushgateway returned {}: {}", + resp.status(), + resp.text().await.unwrap_or_default() + ); + } + Err(e) => { + eprintln!("Failed to push metrics to Pushgateway: {e}"); + } + } + } + + /// Delete metrics for this job from Pushgateway (call on training completion). + pub async fn cleanup(&self) { + let url = format!( + "{}/metrics/job/{}/model/{}", + self.pushgateway_url, self.job_id, self.model_name + ); + let _ = self.client.delete(&url).send().await; + } +} + +/// Current training state to push +#[derive(Debug)] +pub struct TrainingState { + pub epoch: u64, + pub train_loss: f64, + pub val_loss: Option, + pub accuracy: Option, + pub batches_processed: u64, + pub batches_per_second: f64, + pub learning_rate: f64, + pub nan_count: u64, + pub gradient_explosion_count: u64, + pub checkpoint_saves: u64, +} + +impl Default for TrainingState { + fn default() -> Self { + Self { + epoch: 0, + train_loss: 0.0, + val_loss: None, + accuracy: None, + batches_processed: 0, + batches_per_second: 0.0, + learning_rate: 0.001, + nan_count: 0, + gradient_explosion_count: 0, + checkpoint_saves: 0, + } + } +}