foxhunt/services/ml_training_service/src/queue_consumer.rs

//! 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};

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<JobSpawner>,
    dispatcher: Arc<K8sDispatcher>,
    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 {
                    error!(
                        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,
                    extra_args: vec![],
                };

                match dispatcher.dispatch(&params).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) => {
                debug!(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);
    }
}