From abc5ee6af12c38abec607035392cc3e8aec04659 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Fri, 27 Feb 2026 22:45:57 +0100 Subject: [PATCH] feat(training): add in-process queue consumer for K8s job dispatch Background tokio task polls JobSpawner.get_next_pending_job() every 5s, marks found jobs as Running, builds TrainingJobParams, and dispatches to K8s via K8sDispatcher. On dispatch failure the job is marked Failed. Co-Authored-By: Claude Opus 4.6 --- services/ml_training_service/src/lib.rs | 1 + services/ml_training_service/src/main.rs | 14 ++ .../ml_training_service/src/queue_consumer.rs | 181 ++++++++++++++++++ 3 files changed, 196 insertions(+) create mode 100644 services/ml_training_service/src/queue_consumer.rs diff --git a/services/ml_training_service/src/lib.rs b/services/ml_training_service/src/lib.rs index fcbf463b4..5e3a1163c 100644 --- a/services/ml_training_service/src/lib.rs +++ b/services/ml_training_service/src/lib.rs @@ -38,6 +38,7 @@ pub mod monitoring; pub mod optuna_persistence; pub mod orchestrator; pub mod promotion_manager; +pub mod queue_consumer; pub mod schema_types; pub mod service; pub mod simple_metrics; diff --git a/services/ml_training_service/src/main.rs b/services/ml_training_service/src/main.rs index 22b586437..26e0577cc 100644 --- a/services/ml_training_service/src/main.rs +++ b/services/ml_training_service/src/main.rs @@ -420,6 +420,20 @@ async fn serve(args: ServeArgs) -> Result<()> { )); info!("Job spawner initialized -- training jobs will be persisted to PostgreSQL"); + // Spawn queue consumer to poll pending jobs and dispatch to K8s + if let Some(ref dispatcher) = k8s_dispatcher { + let consumer_spawner = Arc::clone(&job_spawner); + let consumer_dispatcher = Arc::clone(dispatcher); + let consumer_config = + ml_training_service::queue_consumer::QueueConsumerConfig::default(); + tokio::spawn(ml_training_service::queue_consumer::run_queue_consumer( + consumer_spawner, + consumer_dispatcher, + consumer_config, + )); + info!("Queue consumer started -- polling for pending training jobs"); + } + // Create gRPC service let training_service = MLTrainingServiceImpl::new( Arc::clone(&orchestrator), diff --git a/services/ml_training_service/src/queue_consumer.rs b/services/ml_training_service/src/queue_consumer.rs new file mode 100644 index 000000000..d7ba3cfa0 --- /dev/null +++ b/services/ml_training_service/src/queue_consumer.rs @@ -0,0 +1,181 @@ +//! Queue Consumer — background loop that polls pending training jobs and dispatches them to K8s. +//! +//! Runs as a `tokio::spawn`-ed task inside the `ml_training_service` process. It checks +//! [`JobSpawner::get_next_pending_job`] every `poll_interval_secs` seconds and, when a job +//! is found, marks it as `Running`, builds [`TrainingJobParams`], and calls +//! [`K8sDispatcher::dispatch`]. On dispatch failure the job is marked `Failed`. +//! +//! The consumer is intentionally simple — one job at a time — because GPU training jobs +//! are heavyweight and we rely on K8s scheduling for parallelism. + +use std::sync::Arc; +use std::time::Duration; + +use tracing::{debug, error, info, warn}; + +use crate::job_spawner::JobSpawner; +use crate::k8s_dispatcher::{training_binary_for_model, K8sDispatcher, TrainingJobParams}; + +// --------------------------------------------------------------------------- +// Configuration +// --------------------------------------------------------------------------- + +/// Configuration for the queue consumer loop. +#[derive(Debug, Clone)] +pub struct QueueConsumerConfig { + /// How often (in seconds) to poll the database for pending jobs. + pub poll_interval_secs: u64, + /// Maximum number of concurrent K8s jobs the consumer will keep in flight. + /// Currently only `1` is supported — the field exists for forward compatibility. + pub max_concurrent_jobs: u32, + /// Default number of training epochs when the job config does not specify one. + pub default_epochs: u32, +} + +impl Default for QueueConsumerConfig { + fn default() -> Self { + Self { + poll_interval_secs: 5, + max_concurrent_jobs: 1, + default_epochs: 50, + } + } +} + +// --------------------------------------------------------------------------- +// Consumer loop +// --------------------------------------------------------------------------- + +/// Runs the queue consumer forever. +/// +/// This function is designed to be passed to [`tokio::spawn`]. It never returns under +/// normal operation — it only exits if the `JobSpawner` or `K8sDispatcher` is dropped +/// (which causes DB / K8s calls to fail permanently). +pub async fn run_queue_consumer( + spawner: Arc, + dispatcher: Arc, + config: QueueConsumerConfig, +) { + let poll_interval = Duration::from_secs(config.poll_interval_secs); + + info!( + poll_interval_secs = config.poll_interval_secs, + max_concurrent_jobs = config.max_concurrent_jobs, + "queue consumer started" + ); + + loop { + match spawner.get_next_pending_job().await { + Ok(Some(job)) => { + let job_id = job.id; + let model_type = job.model_type.clone(); + + info!( + job_id = %job_id, + model_type = %model_type, + batch_id = %job.batch_id, + "picked up pending job" + ); + + // Mark as Running *before* dispatching so no other consumer grabs it. + if let Err(e) = spawner.update_job_status(job_id, "Running").await { + warn!( + job_id = %job_id, + error = %e, + "failed to mark job as Running — skipping" + ); + tokio::time::sleep(poll_interval).await; + continue; + } + + // Extract parameters from config_json. + let symbol = job + .config_json + .get("asset") + .and_then(|v| v.as_str()) + .unwrap_or("ES.FUT") + .to_string(); + + let epochs = job + .config_json + .get("epochs") + .and_then(|v| v.as_u64()) + .map(|e| e as u32) + .unwrap_or(config.default_epochs); + + let binary = training_binary_for_model(&model_type).to_string(); + + let params = TrainingJobParams { + job_id, + model_type: model_type.clone(), + symbol, + epochs, + binary, + }; + + match dispatcher.dispatch(¶ms).await { + Ok(k8s_name) => { + info!( + job_id = %job_id, + k8s_name = %k8s_name, + "dispatched training job to K8s" + ); + } + Err(e) => { + error!( + job_id = %job_id, + error = %e, + "K8s dispatch failed — marking job as Failed" + ); + if let Err(update_err) = + spawner.update_job_status(job_id, "Failed").await + { + error!( + job_id = %job_id, + error = %update_err, + "failed to mark job as Failed after dispatch error" + ); + } + } + } + } + Ok(None) => { + debug!("no pending jobs — sleeping"); + } + Err(e) => { + warn!(error = %e, "error polling for pending jobs — will retry"); + } + } + + tokio::time::sleep(poll_interval).await; + } +} + +// --------------------------------------------------------------------------- +// Tests +// --------------------------------------------------------------------------- + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_queue_consumer_config_defaults() { + let config = QueueConsumerConfig::default(); + assert_eq!(config.poll_interval_secs, 5); + assert_eq!(config.max_concurrent_jobs, 1); + assert_eq!(config.default_epochs, 50); + } + + #[test] + fn test_queue_consumer_config_custom() { + let config = QueueConsumerConfig { + poll_interval_secs: 10, + max_concurrent_jobs: 3, + default_epochs: 100, + }; + assert_eq!(config.poll_interval_secs, 10); + assert_eq!(config.max_concurrent_jobs, 3); + assert_eq!(config.default_epochs, 100); + } +}