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>
This commit is contained in:
jgrusewski
2026-02-27 22:13:29 +01:00
parent 93115c1fc5
commit 2538d52324

View File

@@ -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<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)