Files
foxhunt/docs/plans/2026-02-27-training-deploy-plan.md
jgrusewski 2538d52324 docs: add training pipeline & model deployment implementation plan
11 tasks across 2 independent tracks:
- Track 1 (Tasks 1-5): Wire fxt train → gRPC → JobSpawner → K8sDispatcher → GPU Jobs
- Track 2 (Tasks 6-9): Wire promotion RPCs to real PromotionManager
- Shared (Tasks 10-11): Docker build + end-to-end validation

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-27 22:13:29 +01:00

22 KiB

Training Pipeline & Model Deployment Implementation Plan

For Claude: REQUIRED SUB-SKILL: Use superpowers:executing-plans to implement this plan task-by-task.

Goal: Wire the complete fxt CLI → training service → K8s GPU jobs pipeline and model promotion flow.

Architecture: Two independent tracks. Track 1 adds JobSpawner (PostgreSQL-backed batch queue) to the start_training gRPC path and adds an in-process queue consumer that dispatches to K8sDispatcher. Track 2 wires the stub promotion/completion RPCs to the real PromotionManager. Both tracks share S3 as the checkpoint handoff.

Tech Stack: Rust, tonic (gRPC), kube-rs, sqlx (PostgreSQL), Scaleway Kapsule (K8s), S3


Track 1: Training Pipeline

Task 1: Add JobSpawner to MLTrainingServiceImpl

Files:

  • Modify: services/ml_training_service/src/service.rs:53-82 (add job_spawner field)
  • Modify: services/ml_training_service/src/main.rs:394-424 (construct JobSpawner, pass to service)

Step 1: Write failing test

Add to services/ml_training_service/src/service.rs tests:

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_service_impl_has_job_spawner() {
        // Verify the struct has a job_spawner field
        // This is a compile-time check — if the field doesn't exist, this won't compile
        let _: fn(&MLTrainingServiceImpl) -> bool = |s| s.job_spawner.is_some();
    }
}

Step 2: Run test to verify it fails

Run: SQLX_OFFLINE=true cargo test -p ml_training_service --lib test_service_impl_has_job_spawner Expected: FAIL — no field job_spawner on type MLTrainingServiceImpl

Step 3: Add job_spawner field to MLTrainingServiceImpl

In service.rs, add to the struct (after line 60):

/// Database-backed job spawner for batch orchestration -- None when DB unavailable.
pub job_spawner: Option<Arc<crate::job_spawner::JobSpawner>>,

Update MLTrainingServiceImpl::new() to accept and store it:

pub fn new(
    orchestrator: Arc<TrainingOrchestrator>,
    tuning_manager: Arc<TuningManager>,
    _config: MLConfig,
    promotion_manager: Arc<crate::promotion_manager::PromotionManager>,
    k8s_dispatcher: Option<Arc<crate::k8s_dispatcher::K8sDispatcher>>,
    job_spawner: Option<Arc<crate::job_spawner::JobSpawner>>,
) -> Self {
    // ... existing code ...
    Self {
        orchestrator,
        tuning_handlers,
        batch_tuning_manager,
        promotion_manager,
        k8s_dispatcher,
        job_spawner,
    }
}

Step 4: Update main.rs to construct JobSpawner

In main.rs after the database initialization (line ~287), add:

// Initialize job spawner for batch training orchestration
let job_spawner = Arc::new(crate::job_spawner::JobSpawner::new(database.pool().clone()));
info!("Job spawner initialized for batch training orchestration");

Pass it to the service constructor:

let training_service = MLTrainingServiceImpl::new(
    Arc::clone(&orchestrator),
    Arc::clone(&tuning_manager),
    ml_config.clone(),
    promotion_manager,
    k8s_dispatcher,
    Some(job_spawner.clone()),
);

Step 5: Run test to verify it passes

Run: SQLX_OFFLINE=true cargo test -p ml_training_service --lib test_service_impl_has_job_spawner Expected: PASS

Step 6: Commit

git add services/ml_training_service/src/service.rs services/ml_training_service/src/main.rs
git commit -m "feat(training): add JobSpawner to MLTrainingServiceImpl"

Task 2: Wire start_training to JobSpawner

Files:

  • Modify: services/ml_training_service/src/service.rs:250-346 (start_training handler)

Step 1: Write failing test

#[tokio::test]
async fn test_start_training_creates_batch_job() {
    // This test validates the flow: start_training -> JobSpawner::spawn_batch
    // It will be an integration test that requires a database
    // For now, test the model-to-binary mapping which is pure
    use crate::k8s_dispatcher::training_binary_for_model;
    assert_eq!(training_binary_for_model("tft"), "train_baseline_supervised");
    assert_eq!(training_binary_for_model("dqn"), "train_baseline_rl");
}

Step 2: Modify start_training handler

In the start_training method, after the K8s dispatcher attempt (line ~322), add the JobSpawner path. The logic should be:

  1. If job_spawner is Some AND k8s_dispatcher is Some → use batch flow (JobSpawner creates DB records, queue consumer dispatches)
  2. If only k8s_dispatcher → dispatch directly (current behavior)
  3. Fallback → in-process orchestrator
// After the existing K8s dispatcher block, before the orchestrator fallback:
if let Some(ref spawner) = self.job_spawner {
    let asset = crate::job_spawner::Asset {
        symbol: symbol.clone(),
        data_file: std::path::PathBuf::from(format!("/data/futures-baseline/{}", symbol)),
    };
    let model = common::model_types::ModelType::from_str(&req.model_type)
        .map_err(|e| Status::invalid_argument(format!("Invalid model type: {}", e)))?;

    match spawner.spawn_batch(vec![asset], vec![model]).await {
        Ok(batch) => {
            info!(batch_id = %batch.batch_id, "Batch job created in database");
            return Ok(Response::new(StartTrainingResponse {
                job_id: batch.batch_id.to_string(),
                status: ProtoTrainingStatus::Pending as i32,
                message: format!("Batch training job queued: {}", batch.batch_id),
            }));
        }
        Err(e) => {
            warn!("JobSpawner failed, falling back to orchestrator: {}", e);
        }
    }
}

Step 3: Run tests

Run: SQLX_OFFLINE=true cargo check -p ml_training_service Expected: Compiles (may need to add from_str import or adjust ModelType conversion)

Step 4: Commit

git add services/ml_training_service/src/service.rs
git commit -m "feat(training): wire start_training gRPC handler to JobSpawner"

Task 3: Add queue consumer loop

Files:

  • Create: services/ml_training_service/src/queue_consumer.rs
  • Modify: services/ml_training_service/src/lib.rs (add module)
  • Modify: services/ml_training_service/src/main.rs (spawn consumer)

Step 1: Write failing test

In queue_consumer.rs:

#[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);
    }
}

Step 2: Implement queue consumer

//! Queue Consumer — polls JobSpawner for pending jobs and dispatches to K8s.
//!
//! Runs as a tokio::spawn'd task inside the ml_training_service process.

use std::sync::Arc;
use std::time::Duration;

use tracing::{error, info, warn};

use crate::job_spawner::JobSpawner;
use crate::k8s_dispatcher::{K8sDispatcher, TrainingJobParams, training_binary_for_model};

/// Configuration for the queue consumer.
#[derive(Debug, Clone)]
pub struct QueueConsumerConfig {
    /// Seconds between polls for pending jobs.
    pub poll_interval_secs: u64,
    /// Maximum concurrent K8s jobs to dispatch (currently sequential).
    pub max_concurrent_jobs: u32,
}

impl Default for QueueConsumerConfig {
    fn default() -> Self {
        Self {
            poll_interval_secs: 5,
            max_concurrent_jobs: 1,
        }
    }
}

/// Starts the queue consumer loop. Call via `tokio::spawn(run_queue_consumer(...))`.
pub async fn run_queue_consumer(
    spawner: Arc<JobSpawner>,
    dispatcher: Arc<K8sDispatcher>,
    config: QueueConsumerConfig,
) {
    info!(
        poll_interval = config.poll_interval_secs,
        "Queue consumer started"
    );

    let interval = Duration::from_secs(config.poll_interval_secs);

    loop {
        match spawner.get_next_pending_job().await {
            Ok(Some(job)) => {
                let job_id = job.id;
                info!(job_id = %job_id, model = %job.model_type, "Dispatching pending job");

                // Mark as running
                if let Err(e) = spawner.update_job_status(job_id, "Running").await {
                    error!(job_id = %job_id, "Failed to mark job as Running: {}", e);
                    continue;
                }

                // Extract symbol from config_json
                let symbol = job.config_json
                    .get("asset")
                    .and_then(|v| v.as_str())
                    .unwrap_or("ES.FUT")
                    .to_string();

                let binary = training_binary_for_model(&job.model_type);

                let params = TrainingJobParams {
                    job_id,
                    model_type: job.model_type.clone(),
                    symbol,
                    epochs: 50, // TODO: extract from config_json
                    binary: binary.to_string(),
                };

                match dispatcher.dispatch(&params).await {
                    Ok(k8s_name) => {
                        info!(
                            job_id = %job_id,
                            k8s_name = %k8s_name,
                            "K8s job dispatched successfully"
                        );
                    }
                    Err(e) => {
                        error!(job_id = %job_id, "K8s dispatch failed: {}", e);
                        if let Err(e2) = spawner.update_job_status(job_id, "Failed").await {
                            error!(job_id = %job_id, "Failed to mark job as Failed: {}", e2);
                        }
                    }
                }
            }
            Ok(None) => {
                // No pending jobs — sleep and retry
            }
            Err(e) => {
                warn!("Error polling for pending jobs: {}", e);
            }
        }

        tokio::time::sleep(interval).await;
    }
}

Step 3: Add module to lib.rs

Add pub mod queue_consumer; to services/ml_training_service/src/lib.rs.

Step 4: Spawn consumer in main.rs

In main.rs, after the service is constructed (line ~424), add:

// Spawn queue consumer if both JobSpawner and K8sDispatcher are available
if let Some(ref dispatcher) = k8s_dispatcher {
    let consumer_spawner = job_spawner.clone();
    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");
}

Step 5: Run tests

Run: SQLX_OFFLINE=true cargo test -p ml_training_service --lib test_queue_consumer Expected: PASS

Step 6: Commit

git add services/ml_training_service/src/queue_consumer.rs services/ml_training_service/src/lib.rs services/ml_training_service/src/main.rs
git commit -m "feat(training): add in-process queue consumer for K8s job dispatch"

Task 4: Expand child_jobs model_type constraint

Files:

  • Create: migrations/049_expand_child_jobs_model_types.sql

The existing chk_child_job_model_type constraint only allows 5 model types (DQN, PPO, MAMBA-2, TFT, TLOB). We support 10. Fix it.

Step 1: Write migration

-- Migration 049: Expand child_jobs model_type constraint to all 10 supported models
ALTER TABLE child_jobs DROP CONSTRAINT IF EXISTS chk_child_job_model_type;
ALTER TABLE child_jobs ADD CONSTRAINT chk_child_job_model_type CHECK (
    model_type IN ('DQN', 'PPO', 'MAMBA-2', 'TFT', 'TLOB', 'TGGN', 'LIQUID', 'KAN', 'XLSTM', 'DIFFUSION')
);

Step 2: Commit

git add migrations/049_expand_child_jobs_model_types.sql
git commit -m "fix(db): expand child_jobs model_type constraint to all 10 models"

Task 5: Deploy training-output-pvc

Files:

  • Verify: infra/k8s/training/training-output-pvc.yaml (already exists)

Step 1: Check if PVC exists in cluster

kubectl get pvc -n foxhunt training-output-pvc

If missing, apply it:

kubectl apply -f infra/k8s/training/training-output-pvc.yaml

Step 2: Verify both PVCs bound

kubectl get pvc -n foxhunt | grep training

Expected: Both training-data-pvc and training-output-pvc in Bound state.

Step 3: Commit (no code change needed if YAML exists)


Track 2: Model Serving & Validation

Task 6: Wire report_job_completion to PromotionManager

Files:

  • Modify: services/ml_training_service/src/service.rs:1097-1118 (report_job_completion handler)

Step 1: Write failing test

#[tokio::test]
async fn test_promotion_manager_register() {
    let pm = crate::promotion_manager::PromotionManager::new();
    let status = pm.register_completion(
        "test-job-1",
        "TFT",
        "ES.FUT",
        "s3://foxhunt-models/test/checkpoint.safetensors",
        &[("sharpe".to_string(), 1.5)].into_iter().collect(),
    ).await;
    // First model for this (type, symbol) pair should be Registered
    assert_eq!(status, crate::promotion_manager::PromotionStatus::Registered);
}

Step 2: Wire report_job_completion to PromotionManager

Replace the stub implementation:

async fn report_job_completion(
    &self,
    request: Request<JobCompletionReport>,
) -> Result<Response<JobCompletionAck>, Status> {
    let report = request.into_inner();
    info!(
        job_id = %report.job_id,
        model_type = %report.model_type,
        success = report.success,
        "Job completion report received"
    );

    if !report.success {
        // Update job status to Failed if we have a spawner
        if let Some(ref spawner) = self.job_spawner {
            if let Ok(uuid) = report.job_id.parse::<uuid::Uuid>() {
                let _ = spawner.update_job_status(uuid, "Failed").await;
            }
        }
        return Ok(Response::new(JobCompletionAck {
            acknowledged: true,
            promotion_status: "failed".to_string(),
        }));
    }

    // Update job status to Completed
    if let Some(ref spawner) = self.job_spawner {
        if let Ok(uuid) = report.job_id.parse::<uuid::Uuid>() {
            let _ = spawner.update_job_status(uuid, "Completed").await;
        }
    }

    // Register with promotion manager
    let metrics: std::collections::HashMap<String, f64> = report.metrics.clone();
    let status = self.promotion_manager.register_completion(
        &report.job_id,
        &report.model_type,
        &report.symbol,
        &report.s3_path,
        &metrics,
    ).await;

    let promotion_str = match status {
        crate::promotion_manager::PromotionStatus::PendingPromotion => "pending_promotion",
        crate::promotion_manager::PromotionStatus::NoImprovement => "no_improvement",
        crate::promotion_manager::PromotionStatus::Registered => "registered",
        crate::promotion_manager::PromotionStatus::Error(ref e) => {
            warn!("Promotion registration error: {}", e);
            "error"
        }
    };

    Ok(Response::new(JobCompletionAck {
        acknowledged: true,
        promotion_status: promotion_str.to_string(),
    }))
}

Step 3: Run tests

Run: SQLX_OFFLINE=true cargo check -p ml_training_service

Step 4: Commit

git add services/ml_training_service/src/service.rs
git commit -m "feat(serving): wire report_job_completion to PromotionManager"

Task 7: Wire list_pending_promotions

Files:

  • Modify: services/ml_training_service/src/service.rs:1120-1128
  • Modify: services/ml_training_service/src/promotion_manager.rs (add list_pending method if needed)

Step 1: Check if PromotionManager has list_pending

Read promotion_manager.rs to verify list_pending() exists. If not, add:

/// List all models pending promotion.
pub async fn list_pending(&self) -> Vec<PendingModel> {
    self.pending.read().await.values().cloned().collect()
}

Step 2: Wire the gRPC handler

async fn list_pending_promotions(
    &self,
    _request: Request<ListPendingPromotionsRequest>,
) -> Result<Response<ListPendingPromotionsResponse>, Status> {
    let pending = self.promotion_manager.list_pending().await;

    let promotions = pending.iter().map(|p| proto::PendingPromotion {
        model_id: p.model_id.clone(),
        model_type: p.model_type.clone(),
        symbol: p.symbol.clone(),
        s3_path: p.s3_path.clone(),
        job_id: p.job_id.clone(),
        new_metrics: p.new_metrics.clone(),
        current_metrics: p.current_metrics.clone(),
    }).collect();

    Ok(Response::new(ListPendingPromotionsResponse { promotions }))
}

Step 3: Check proto definition has PendingPromotion message

If the proto doesn't have the PendingPromotion message, this will need adjustment to match whatever fields the proto actually defines. Check the generated code.

Step 4: Commit

git add services/ml_training_service/src/service.rs services/ml_training_service/src/promotion_manager.rs
git commit -m "feat(serving): wire list_pending_promotions to PromotionManager"

Task 8: Wire approve_promotion and reject_promotion

Files:

  • Modify: services/ml_training_service/src/service.rs:1130-1155
  • Modify: services/ml_training_service/src/promotion_manager.rs (add approve/reject methods)

Step 1: Add approve/reject to PromotionManager

/// Approve a pending promotion — moves model to active.
pub async fn approve(&self, model_id: &str) -> Result<ActiveModel, String> {
    let mut pending = self.pending.write().await;
    let model = pending.remove(model_id)
        .ok_or_else(|| format!("No pending model: {}", model_id))?;

    let active = ActiveModel {
        model_id: model.model_id.clone(),
        s3_path: model.s3_path.clone(),
        metrics: model.new_metrics.clone(),
        promoted_at: chrono::Utc::now(),
    };

    let key = (model.model_type.clone(), model.symbol.clone());
    self.active_models.write().await.insert(key, active.clone());

    info!(model_id = %model_id, "Model promoted to active");
    Ok(active)
}

/// Reject a pending promotion — removes from pending list.
pub async fn reject(&self, model_id: &str, reason: &str) -> Result<(), String> {
    let mut pending = self.pending.write().await;
    if pending.remove(model_id).is_none() {
        return Err(format!("No pending model: {}", model_id));
    }
    info!(model_id = %model_id, reason = %reason, "Model promotion rejected");
    Ok(())
}

Step 2: Wire gRPC handlers

async fn approve_promotion(
    &self,
    request: Request<ApprovePromotionRequest>,
) -> Result<Response<ApprovePromotionResponse>, Status> {
    let req = request.into_inner();
    info!(model_id = %req.model_id, "Promotion approval requested");

    match self.promotion_manager.approve(&req.model_id).await {
        Ok(active) => Ok(Response::new(ApprovePromotionResponse {
            success: true,
            message: format!("Model {} promoted. S3: {}", active.model_id, active.s3_path),
        })),
        Err(e) => Ok(Response::new(ApprovePromotionResponse {
            success: false,
            message: e,
        })),
    }
}

async fn reject_promotion(
    &self,
    request: Request<RejectPromotionRequest>,
) -> Result<Response<RejectPromotionResponse>, Status> {
    let req = request.into_inner();
    info!(model_id = %req.model_id, reason = %req.reason, "Promotion rejection requested");

    match self.promotion_manager.reject(&req.model_id, &req.reason).await {
        Ok(()) => Ok(Response::new(RejectPromotionResponse {
            success: true,
            message: format!("Model {} promotion rejected", req.model_id),
        })),
        Err(e) => Ok(Response::new(RejectPromotionResponse {
            success: false,
            message: e,
        })),
    }
}

Step 3: Run tests

Run: SQLX_OFFLINE=true cargo check -p ml_training_service

Step 4: Commit

git add services/ml_training_service/src/service.rs services/ml_training_service/src/promotion_manager.rs
git commit -m "feat(serving): wire approve/reject promotion to PromotionManager"

Task 9: Verify fxt model commands work end-to-end

Files:

  • Read: bin/fxt/src/commands/model/list.rs
  • Read: bin/fxt/src/commands/model/approve.rs
  • Read: bin/fxt/src/commands/model/reject.rs

Step 1: Verify fxt model commands compile

Run: SQLX_OFFLINE=true cargo check -p fxt

Step 2: Test locally (if training service is running)

fxt model list
fxt model approve <model-id>
fxt model reject <model-id> --reason "test"

Step 3: Commit (if any fixes needed)


Task 10: Build and push Docker images

Files:

  • Read: infra/docker/Dockerfile.training (training binaries image)
  • Read: K8s deployment for ml-training-service

Step 1: Build training image

docker build -f infra/docker/Dockerfile.training -t rg.fr-par.scw.cloud/foxhunt-ci/training:latest .
docker push rg.fr-par.scw.cloud/foxhunt-ci/training:latest

Step 2: Build and push ml-training-service image

# This would use the service Dockerfile
docker build -f infra/docker/Dockerfile.ml-training-service -t rg.fr-par.scw.cloud/foxhunt-ci/ml-training-service:latest .
docker push rg.fr-par.scw.cloud/foxhunt-ci/ml-training-service:latest

Step 3: Redeploy to K8s

kubectl rollout restart deployment/ml-training-service -n foxhunt

Task 11: End-to-end validation

Step 1: Verify training service is running

kubectl get pods -n foxhunt -l app=ml-training-service
kubectl logs -n foxhunt deployment/ml-training-service --tail=50

Look for: "Queue consumer started" and "K8s dispatcher initialized"

Step 2: Trigger a training job

fxt train start tft ES.FUT --epochs 10

Expected: Job ID returned, status PENDING

Step 3: Check job was dispatched

fxt train status <job-id>
kubectl get jobs -n foxhunt -l foxhunt/job-type=training

Expected: K8s Job created on gpu-training pool

Step 4: Check model promotion after completion

fxt model list

Expected: New model appears as pending (or registered if first for this type/symbol)