From 2538d5232467a468205c088a655f944ac0946458 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Fri, 27 Feb 2026 22:13:29 +0100 Subject: [PATCH] docs: add training pipeline & model deployment implementation plan MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- docs/plans/2026-02-27-training-deploy-plan.md | 738 ++++++++++++++++++ 1 file changed, 738 insertions(+) create mode 100644 docs/plans/2026-02-27-training-deploy-plan.md diff --git a/docs/plans/2026-02-27-training-deploy-plan.md b/docs/plans/2026-02-27-training-deploy-plan.md new file mode 100644 index 000000000..c29b74be6 --- /dev/null +++ b/docs/plans/2026-02-27-training-deploy-plan.md @@ -0,0 +1,738 @@ +# 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>, +``` + +Update `MLTrainingServiceImpl::new()` to accept and store it: + +```rust +pub fn new( + orchestrator: Arc, + tuning_manager: Arc, + _config: MLConfig, + promotion_manager: Arc, + k8s_dispatcher: Option>, + job_spawner: Option>, +) -> 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, + dispatcher: Arc, + 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(¶ms).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, +) -> Result, 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::() { + 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::() { + let _ = spawner.update_job_status(uuid, "Completed").await; + } + } + + // Register with promotion manager + let metrics: std::collections::HashMap = 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 { + self.pending.read().await.values().cloned().collect() +} +``` + +**Step 2: Wire the gRPC handler** + +```rust +async fn list_pending_promotions( + &self, + _request: Request, +) -> Result, 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 { + 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, +) -> Result, 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, +) -> Result, 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 +fxt model reject --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 +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)