# 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:

```rust
#[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):

```rust
/// 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:

```rust
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:

```rust
// 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:

```rust
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**

```bash
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**

```rust
#[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

```rust
// 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**

```bash
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`:

```rust
#[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**

```rust
//! 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:

```rust
// 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**

```bash
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**

```sql
-- 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**

```bash
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**

```bash
kubectl get pvc -n foxhunt training-output-pvc
```

If missing, apply it:

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

**Step 2: Verify both PVCs bound**

```bash
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**

```rust
#[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:

```rust
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**

```bash
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:

```rust
/// 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**

```rust
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**

```bash
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**

```rust
/// 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**

```rust
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**

```bash
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)**

```bash
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**

```bash
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**

```bash
# 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**

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

---

### Task 11: End-to-end validation

**Step 1: Verify training service is running**

```bash
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**

```bash
fxt train start tft ES.FUT --epochs 10
```

Expected: Job ID returned, status PENDING

**Step 3: Check job was dispatched**

```bash
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**

```bash
fxt model list
```

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