From 06a875e6fc75f6987606e0636295e21801aa6ebf Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Fri, 6 Mar 2026 00:32:09 +0100 Subject: [PATCH] fix(metrics): push training metrics to pushgateway before pod exit Ephemeral Argo workflow pods terminate after training completes, causing Prometheus to lose all scraped metrics. Add push_to_gateway() to POST final metrics to the existing pushgateway service so they persist on the Grafana training dashboard after pod completion. Co-Authored-By: Claude Opus 4.6 --- crates/common/src/metrics/server.rs | 60 ++++++++++++++++++- crates/ml/examples/evaluate_baseline.rs | 5 ++ crates/ml/examples/evaluate_supervised.rs | 6 ++ crates/ml/examples/hyperopt_baseline_rl.rs | 6 ++ .../examples/hyperopt_baseline_supervised.rs | 6 ++ crates/ml/examples/train_baseline_rl.rs | 5 ++ .../ml/examples/train_baseline_supervised.rs | 5 ++ 7 files changed, 92 insertions(+), 1 deletion(-) diff --git a/crates/common/src/metrics/server.rs b/crates/common/src/metrics/server.rs index 8e47f29b1..a30165b85 100644 --- a/crates/common/src/metrics/server.rs +++ b/crates/common/src/metrics/server.rs @@ -5,7 +5,7 @@ //! HTTP server use this instead. use std::io::{BufRead, BufReader, Read as IoRead, Write as IoWrite}; -use std::net::TcpListener; +use std::net::{TcpListener, TcpStream}; use super::gather_metrics; @@ -69,3 +69,61 @@ pub fn start_metrics_server(port: u16) { } }); } + +/// Push all registered metrics to a Prometheus Pushgateway. +/// +/// Uses a plain TCP connection to POST the text exposition format. +/// The `gateway` address should be `host:port` (e.g. `pushgateway.foxhunt.svc.cluster.local:9091`). +/// `job` is the grouping key (typically the binary name like `hyperopt_baseline_rl`). +/// +/// Falls back to the `PUSHGATEWAY_URL` env var if `gateway` is `None`. +/// Returns `Ok(())` on success, logs and returns error on failure. +pub fn push_to_gateway(gateway: Option<&str>, job: &str) -> Result<(), String> { + let addr = match gateway { + Some(a) => a.to_string(), + None => std::env::var("PUSHGATEWAY_URL") + .unwrap_or_else(|_| "pushgateway.foxhunt.svc.cluster.local:9091".to_string()), + }; + + let body = gather_metrics(); + if body.is_empty() { + tracing::debug!("No metrics to push to gateway"); + return Ok(()); + } + + let path = format!("/metrics/job/{job}"); + let request = format!( + "POST {path} HTTP/1.1\r\n\ + Host: {addr}\r\n\ + Content-Type: text/plain; version=0.0.4\r\n\ + Content-Length: {}\r\n\ + Connection: close\r\n\r\n\ + {body}", + body.len(), + ); + + let mut stream = TcpStream::connect(&addr).map_err(|e| { + let msg = format!("Failed to connect to pushgateway at {addr}: {e}"); + tracing::warn!("{}", msg); + msg + })?; + + stream.write_all(request.as_bytes()).map_err(|e| { + let msg = format!("Failed to send metrics to pushgateway: {e}"); + tracing::warn!("{}", msg); + msg + })?; + + // Read response status line + let mut response = String::new(); + let _ = BufReader::new(&stream).read_line(&mut response); + + if response.contains("200") || response.contains("202") { + tracing::info!("Pushed metrics to pushgateway ({addr}) for job={job}"); + Ok(()) + } else { + let msg = format!("Pushgateway returned unexpected response: {}", response.trim()); + tracing::warn!("{}", msg); + Err(msg) + } +} diff --git a/crates/ml/examples/evaluate_baseline.rs b/crates/ml/examples/evaluate_baseline.rs index 176c38580..e7a415e39 100644 --- a/crates/ml/examples/evaluate_baseline.rs +++ b/crates/ml/examples/evaluate_baseline.rs @@ -1219,5 +1219,10 @@ fn main() -> Result<()> { tm::set_active_workers(0.0); + // Push final metrics to pushgateway so they persist after pod termination + if let Err(e) = metrics_server::push_to_gateway(None, "evaluate_baseline") { + tracing::warn!("Failed to push metrics to gateway (non-fatal): {e}"); + } + Ok(()) } diff --git a/crates/ml/examples/evaluate_supervised.rs b/crates/ml/examples/evaluate_supervised.rs index acec57b4d..6285a98fb 100644 --- a/crates/ml/examples/evaluate_supervised.rs +++ b/crates/ml/examples/evaluate_supervised.rs @@ -836,5 +836,11 @@ fn main() -> Result<()> { info!(" Total fold evaluations: {}", report.folds.len()); tm::set_active_workers(0.0); + + // Push final metrics to pushgateway so they persist after pod termination + if let Err(e) = metrics_server::push_to_gateway(None, "evaluate_supervised") { + tracing::warn!("Failed to push metrics to gateway (non-fatal): {e}"); + } + Ok(()) } diff --git a/crates/ml/examples/hyperopt_baseline_rl.rs b/crates/ml/examples/hyperopt_baseline_rl.rs index 4ef6340f3..1d4fb4fdb 100644 --- a/crates/ml/examples/hyperopt_baseline_rl.rs +++ b/crates/ml/examples/hyperopt_baseline_rl.rs @@ -483,5 +483,11 @@ fn main() -> Result<()> { training_metrics::set_hyperopt_mode(&hyperopt_model_label, false); training_metrics::set_active_workers(0.0); + + // Push final metrics to pushgateway so they persist after pod termination + if let Err(e) = metrics_server::push_to_gateway(None, "hyperopt_baseline_rl") { + warn!("Failed to push metrics to gateway (non-fatal): {e}"); + } + Ok(()) } diff --git a/crates/ml/examples/hyperopt_baseline_supervised.rs b/crates/ml/examples/hyperopt_baseline_supervised.rs index 9a1e4afc6..3fba18855 100644 --- a/crates/ml/examples/hyperopt_baseline_supervised.rs +++ b/crates/ml/examples/hyperopt_baseline_supervised.rs @@ -714,5 +714,11 @@ fn main() -> Result<()> { training_metrics::set_hyperopt_mode(&hyperopt_model_label, false); training_metrics::set_active_workers(0.0); + + // Push final metrics to pushgateway so they persist after pod termination + if let Err(e) = metrics_server::push_to_gateway(None, "hyperopt_baseline_supervised") { + warn!("Failed to push metrics to gateway (non-fatal): {e}"); + } + Ok(()) } diff --git a/crates/ml/examples/train_baseline_rl.rs b/crates/ml/examples/train_baseline_rl.rs index 9a691febd..87315885c 100644 --- a/crates/ml/examples/train_baseline_rl.rs +++ b/crates/ml/examples/train_baseline_rl.rs @@ -911,6 +911,11 @@ fn main() -> Result<()> { let result = run_training(&args); metrics::set_active_workers(0.0); + // Push final metrics to pushgateway so they persist after pod termination + if let Err(e) = metrics_server::push_to_gateway(None, "train_baseline_rl") { + tracing::warn!("Failed to push metrics to gateway (non-fatal): {e}"); + } + match result { Ok(results) => { for training_result in &results { diff --git a/crates/ml/examples/train_baseline_supervised.rs b/crates/ml/examples/train_baseline_supervised.rs index 91fdd0b6b..3d0548512 100644 --- a/crates/ml/examples/train_baseline_supervised.rs +++ b/crates/ml/examples/train_baseline_supervised.rs @@ -822,6 +822,11 @@ fn main() -> Result<()> { let result = run_training(&args); metrics::set_active_workers(0.0); + // Push final metrics to pushgateway so they persist after pod termination + if let Err(e) = metrics_server::push_to_gateway(None, "train_baseline_supervised") { + tracing::warn!("Failed to push metrics to gateway (non-fatal): {e}"); + } + match result { Ok(results) => { for training_result in &results {