10 tasks: completion markers, training-uploader crate, proto RPCs, K8s dispatcher, promotion manager, fxt CLI commands, infra updates. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
1074 lines
34 KiB
Markdown
1074 lines
34 KiB
Markdown
# On-Demand Training Dispatch Implementation Plan
|
|
|
|
> **For Claude:** REQUIRED SUB-SKILL: Use superpowers:executing-plans to implement this plan task-by-task.
|
|
|
|
**Goal:** Enable `ml_training_service` to dispatch GPU training jobs as K8s Jobs with a Rust sidecar uploader, S3 artifact storage, and model promotion with `fxt` CLI approval.
|
|
|
|
**Architecture:** The service creates K8s `batch/v1` Jobs programmatically via the `kube` crate. Each Job has a training container and a native sidecar (`training-uploader`) that watches for completion, uploads to S3, and calls back to the service via gRPC. The service compares metrics and queues models for promotion, approved via `fxt model approve`.
|
|
|
|
**Tech Stack:** Rust, tonic gRPC, kube-rs, object_store (S3), K8s 1.34 native sidecars, protobuf
|
|
|
|
---
|
|
|
|
## Task 1: Add DONE/FAILED marker + metrics.json to training binaries
|
|
|
|
The training binaries currently exit with `Ok(())` or `Err`. We need them to write a marker file and a metrics summary so the sidecar knows when training is done and what happened.
|
|
|
|
**Files:**
|
|
- Create: `crates/ml/examples/baseline_common/completion.rs`
|
|
- Modify: `crates/ml/examples/baseline_common/mod.rs`
|
|
- Modify: `crates/ml/examples/train_baseline_supervised.rs`
|
|
- Modify: `crates/ml/examples/train_baseline_rl.rs`
|
|
|
|
**Step 1: Create the completion module**
|
|
|
|
Create `crates/ml/examples/baseline_common/completion.rs`:
|
|
|
|
```rust
|
|
//! Training completion marker and metrics writer.
|
|
//!
|
|
//! Writes DONE/FAILED marker files and metrics.json to the output directory
|
|
//! so the training-uploader sidecar can detect completion and report results.
|
|
|
|
use std::path::Path;
|
|
use tracing::{error, info};
|
|
|
|
/// Metrics collected during training for the sidecar to report back.
|
|
#[derive(serde::Serialize)]
|
|
pub struct CompletionMetrics {
|
|
pub model: String,
|
|
pub symbol: String,
|
|
pub best_val_loss: Option<f64>,
|
|
pub sharpe_ratio: Option<f64>,
|
|
pub epochs_completed: usize,
|
|
pub folds_completed: usize,
|
|
}
|
|
|
|
/// Write a DONE marker and metrics.json on successful training completion.
|
|
pub fn write_success_marker(output_dir: &Path, metrics: &CompletionMetrics) {
|
|
let metrics_path = output_dir.join("metrics.json");
|
|
match serde_json::to_string_pretty(metrics) {
|
|
Ok(json) => {
|
|
if let Err(e) = std::fs::write(&metrics_path, &json) {
|
|
error!("Failed to write metrics.json: {}", e);
|
|
} else {
|
|
info!("Wrote metrics to {}", metrics_path.display());
|
|
}
|
|
}
|
|
Err(e) => error!("Failed to serialize metrics: {}", e),
|
|
}
|
|
|
|
let done_path = output_dir.join("DONE");
|
|
if let Err(e) = std::fs::write(&done_path, "") {
|
|
error!("Failed to write DONE marker: {}", e);
|
|
} else {
|
|
info!("Wrote DONE marker to {}", done_path.display());
|
|
}
|
|
}
|
|
|
|
/// Write a FAILED marker with error details.
|
|
pub fn write_failure_marker(output_dir: &Path, error_msg: &str) {
|
|
let failed_path = output_dir.join("FAILED");
|
|
if let Err(e) = std::fs::write(&failed_path, error_msg) {
|
|
error!("Failed to write FAILED marker: {}", e);
|
|
} else {
|
|
info!("Wrote FAILED marker to {}", failed_path.display());
|
|
}
|
|
}
|
|
```
|
|
|
|
**Step 2: Export from baseline_common/mod.rs**
|
|
|
|
Add to `crates/ml/examples/baseline_common/mod.rs`:
|
|
|
|
```rust
|
|
pub mod completion;
|
|
```
|
|
|
|
**Step 3: Wire into train_baseline_supervised.rs**
|
|
|
|
At the end of `main()`, after the training loop completes, call `write_success_marker` with collected metrics. Wrap the training body in a closure or match so failures call `write_failure_marker`.
|
|
|
|
Key pattern — the `--output-dir` arg already exists. After the fold loop:
|
|
|
|
```rust
|
|
use baseline_common::completion::{write_success_marker, write_failure_marker, CompletionMetrics};
|
|
|
|
// After fold loop, before Ok(())
|
|
let metrics = CompletionMetrics {
|
|
model: args.model.clone(),
|
|
symbol: args.symbol.clone(),
|
|
best_val_loss: fold_results.iter().map(|(_, l)| *l).reduce(f64::min),
|
|
sharpe_ratio: None, // supervised models don't compute sharpe during training
|
|
epochs_completed: args.epochs,
|
|
folds_completed: fold_results.len(),
|
|
};
|
|
write_success_marker(&args.output_dir, &metrics);
|
|
```
|
|
|
|
Add a top-level catch for errors:
|
|
|
|
```rust
|
|
fn main() -> Result<()> {
|
|
// ... existing setup ...
|
|
let result = run_training(&args);
|
|
if let Err(ref e) = result {
|
|
write_failure_marker(&args.output_dir, &format!("{:#}", e));
|
|
}
|
|
result
|
|
}
|
|
```
|
|
|
|
This requires extracting the training body into a `run_training()` function.
|
|
|
|
**Step 4: Wire into train_baseline_rl.rs**
|
|
|
|
Same pattern. The RL binary already has `--output-dir`. Add marker writes at success/failure points.
|
|
|
|
**Step 5: Verify**
|
|
|
|
Run: `SQLX_OFFLINE=true cargo check -p ml --example train_baseline_supervised --example train_baseline_rl`
|
|
Expected: compiles with zero errors
|
|
|
|
**Step 6: Commit**
|
|
|
|
```bash
|
|
git add crates/ml/examples/baseline_common/completion.rs crates/ml/examples/baseline_common/mod.rs \
|
|
crates/ml/examples/train_baseline_supervised.rs crates/ml/examples/train_baseline_rl.rs
|
|
git commit -m "feat(ml): add DONE/FAILED marker + metrics.json to training binaries"
|
|
```
|
|
|
|
---
|
|
|
|
## Task 2: Create training_uploader crate
|
|
|
|
New Rust binary that runs as a K8s native sidecar. Watches for DONE/FAILED marker, uploads model artifacts to S3, and reports back to `ml_training_service` via gRPC.
|
|
|
|
**Files:**
|
|
- Create: `crates/training_uploader/Cargo.toml`
|
|
- Create: `crates/training_uploader/src/main.rs`
|
|
- Modify: `Cargo.toml` (workspace members)
|
|
|
|
**Step 1: Create Cargo.toml**
|
|
|
|
Create `crates/training_uploader/Cargo.toml`:
|
|
|
|
```toml
|
|
[package]
|
|
name = "training_uploader"
|
|
version.workspace = true
|
|
edition.workspace = true
|
|
rust-version.workspace = true
|
|
authors.workspace = true
|
|
license.workspace = true
|
|
description = "K8s sidecar: uploads training artifacts to S3 and reports completion via gRPC"
|
|
|
|
[dependencies]
|
|
tokio.workspace = true
|
|
anyhow.workspace = true
|
|
tracing.workspace = true
|
|
tracing-subscriber.workspace = true
|
|
serde.workspace = true
|
|
serde_json.workspace = true
|
|
clap.workspace = true
|
|
object_store = { workspace = true, features = ["aws"] }
|
|
bytes.workspace = true
|
|
tonic.workspace = true
|
|
tonic-prost.workspace = true
|
|
prost.workspace = true
|
|
|
|
[build-dependencies]
|
|
tonic-prost-build.workspace = true
|
|
prost-build.workspace = true
|
|
```
|
|
|
|
**Step 2: Create build.rs**
|
|
|
|
Create `crates/training_uploader/build.rs`:
|
|
|
|
```rust
|
|
fn main() -> Result<(), Box<dyn std::error::Error>> {
|
|
tonic_prost_build::configure()
|
|
.build_server(false)
|
|
.build_client(true)
|
|
.compile_well_known_types(true)
|
|
.extern_path(".google.protobuf", "::prost_types")
|
|
.compile_protos(
|
|
&["../../services/ml_training_service/proto/ml_training.proto"],
|
|
&["../../services/ml_training_service/proto"],
|
|
)?;
|
|
println!("cargo:rerun-if-changed=../../services/ml_training_service/proto/ml_training.proto");
|
|
Ok(())
|
|
}
|
|
```
|
|
|
|
**Step 3: Create main.rs**
|
|
|
|
Create `crates/training_uploader/src/main.rs`. This is the core binary (~250 lines):
|
|
|
|
```rust
|
|
//! Training Uploader — K8s native sidecar
|
|
//!
|
|
//! Watches for DONE/FAILED marker in the training output directory,
|
|
//! uploads model artifacts to S3, and reports completion to ml_training_service.
|
|
|
|
#![allow(unused_crate_dependencies)]
|
|
|
|
use std::path::{Path, PathBuf};
|
|
use std::time::Duration;
|
|
|
|
use anyhow::{Context, Result};
|
|
use bytes::Bytes;
|
|
use clap::Parser;
|
|
use object_store::aws::AmazonS3Builder;
|
|
use object_store::{ObjectStore, PutPayload};
|
|
use tracing::{error, info, warn};
|
|
|
|
mod proto {
|
|
pub mod ml_training {
|
|
tonic::include_proto!("ml_training");
|
|
}
|
|
pub use ml_training::*;
|
|
}
|
|
|
|
#[derive(Parser, Debug)]
|
|
#[command(name = "training-uploader")]
|
|
struct Args {
|
|
/// Output directory to watch for DONE/FAILED marker
|
|
#[arg(long, default_value = "/output")]
|
|
output_dir: PathBuf,
|
|
|
|
/// S3 bucket name
|
|
#[arg(long, env = "S3_BUCKET")]
|
|
s3_bucket: String,
|
|
|
|
/// S3 endpoint URL (Scaleway)
|
|
#[arg(long, env = "S3_ENDPOINT")]
|
|
s3_endpoint: String,
|
|
|
|
/// S3 region
|
|
#[arg(long, env = "S3_REGION", default_value = "fr-par")]
|
|
s3_region: String,
|
|
|
|
/// S3 prefix path for this job's artifacts
|
|
#[arg(long, env = "S3_PREFIX")]
|
|
s3_prefix: String,
|
|
|
|
/// gRPC endpoint of ml_training_service
|
|
#[arg(long, env = "CALLBACK_ENDPOINT")]
|
|
callback_endpoint: String,
|
|
|
|
/// Job ID to report back
|
|
#[arg(long, env = "JOB_ID")]
|
|
job_id: String,
|
|
|
|
/// Poll interval in seconds
|
|
#[arg(long, default_value = "10")]
|
|
poll_interval: u64,
|
|
|
|
/// Maximum wait time in seconds (default: match Job activeDeadlineSeconds)
|
|
#[arg(long, default_value = "3600")]
|
|
max_wait: u64,
|
|
}
|
|
|
|
#[tokio::main]
|
|
async fn main() -> Result<()> {
|
|
tracing_subscriber::fmt()
|
|
.with_max_level(tracing::Level::INFO)
|
|
.with_target(false)
|
|
.init();
|
|
|
|
let args = Args::parse();
|
|
info!("Training uploader started for job {}", args.job_id);
|
|
info!("Watching: {}", args.output_dir.display());
|
|
info!("S3 target: s3://{}/{}", args.s3_bucket, args.s3_prefix);
|
|
|
|
// Wait for completion marker
|
|
let (success, error_msg) = wait_for_completion(&args.output_dir, args.poll_interval, args.max_wait).await;
|
|
|
|
// Read metrics if available
|
|
let metrics = read_metrics(&args.output_dir);
|
|
|
|
// Upload artifacts to S3 (even on failure — partial results may be useful)
|
|
let s3_path = match upload_artifacts(&args).await {
|
|
Ok(path) => {
|
|
info!("Uploaded artifacts to s3://{}/{}", args.s3_bucket, path);
|
|
path
|
|
}
|
|
Err(e) => {
|
|
error!("S3 upload failed: {:#}", e);
|
|
String::new()
|
|
}
|
|
};
|
|
|
|
// Report back to ml_training_service
|
|
if let Err(e) = report_completion(&args, success, &s3_path, &error_msg, &metrics).await {
|
|
error!("gRPC callback failed: {:#}", e);
|
|
// Don't fail the sidecar — the service can poll K8s Job status as fallback
|
|
}
|
|
|
|
info!("Uploader complete (success={})", success);
|
|
Ok(())
|
|
}
|
|
|
|
async fn wait_for_completion(dir: &Path, poll_secs: u64, max_wait: u64) -> (bool, String) {
|
|
let done_path = dir.join("DONE");
|
|
let failed_path = dir.join("FAILED");
|
|
let interval = Duration::from_secs(poll_secs);
|
|
let deadline = tokio::time::Instant::now() + Duration::from_secs(max_wait);
|
|
|
|
loop {
|
|
if done_path.exists() {
|
|
info!("DONE marker found");
|
|
return (true, String::new());
|
|
}
|
|
if failed_path.exists() {
|
|
let msg = std::fs::read_to_string(&failed_path).unwrap_or_default();
|
|
warn!("FAILED marker found: {}", msg);
|
|
return (false, msg);
|
|
}
|
|
if tokio::time::Instant::now() > deadline {
|
|
return (false, "Timeout waiting for training completion".to_owned());
|
|
}
|
|
tokio::time::sleep(interval).await;
|
|
}
|
|
}
|
|
|
|
fn read_metrics(dir: &Path) -> std::collections::HashMap<String, f64> {
|
|
let metrics_path = dir.join("metrics.json");
|
|
if !metrics_path.exists() {
|
|
return std::collections::HashMap::new();
|
|
}
|
|
match std::fs::read_to_string(&metrics_path) {
|
|
Ok(content) => serde_json::from_str::<serde_json::Value>(&content)
|
|
.ok()
|
|
.and_then(|v| v.as_object().cloned())
|
|
.map(|obj| {
|
|
obj.into_iter()
|
|
.filter_map(|(k, v)| v.as_f64().map(|f| (k, f)))
|
|
.collect()
|
|
})
|
|
.unwrap_or_default(),
|
|
Err(_) => std::collections::HashMap::new(),
|
|
}
|
|
}
|
|
|
|
async fn upload_artifacts(args: &Args) -> Result<String> {
|
|
let store = AmazonS3Builder::new()
|
|
.with_bucket_name(&args.s3_bucket)
|
|
.with_endpoint(&args.s3_endpoint)
|
|
.with_region(&args.s3_region)
|
|
.with_allow_http(false)
|
|
.build()
|
|
.context("Failed to build S3 client")?;
|
|
|
|
// Walk output directory and upload all files
|
|
let mut uploaded = 0usize;
|
|
upload_dir_recursive(&store, &args.output_dir, &args.s3_prefix, &mut uploaded).await?;
|
|
info!("Uploaded {} files to S3", uploaded);
|
|
Ok(args.s3_prefix.clone())
|
|
}
|
|
|
|
#[async_recursion::async_recursion]
|
|
async fn upload_dir_recursive(
|
|
store: &dyn ObjectStore,
|
|
local_dir: &Path,
|
|
s3_prefix: &str,
|
|
count: &mut usize,
|
|
) -> Result<()> {
|
|
let entries = std::fs::read_dir(local_dir)
|
|
.with_context(|| format!("Failed to read dir {}", local_dir.display()))?;
|
|
|
|
for entry in entries.flatten() {
|
|
let path = entry.path();
|
|
let name = entry.file_name().to_string_lossy().to_string();
|
|
|
|
if path.is_dir() {
|
|
let sub_prefix = format!("{}/{}", s3_prefix, name);
|
|
upload_dir_recursive(store, &path, &sub_prefix, count).await?;
|
|
} else {
|
|
let s3_key = format!("{}/{}", s3_prefix, name);
|
|
let data = std::fs::read(&path)
|
|
.with_context(|| format!("Failed to read {}", path.display()))?;
|
|
let s3_path = object_store::path::Path::from(s3_key.as_str());
|
|
store
|
|
.put(&s3_path, PutPayload::from(Bytes::from(data)))
|
|
.await
|
|
.with_context(|| format!("Failed to upload {}", s3_key))?;
|
|
*count += 1;
|
|
}
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
async fn report_completion(
|
|
args: &Args,
|
|
success: bool,
|
|
s3_path: &str,
|
|
error_msg: &str,
|
|
metrics: &std::collections::HashMap<String, f64>,
|
|
) -> Result<()> {
|
|
use proto::ml_training_service_client::MlTrainingServiceClient;
|
|
use proto::JobCompletionReport;
|
|
|
|
let mut client = MlTrainingServiceClient::connect(args.callback_endpoint.clone())
|
|
.await
|
|
.context("Failed to connect to ml_training_service")?;
|
|
|
|
let report = JobCompletionReport {
|
|
job_id: args.job_id.clone(),
|
|
s3_path: s3_path.to_owned(),
|
|
success,
|
|
error_message: error_msg.to_owned(),
|
|
metrics: metrics.iter().map(|(k, v)| (k.clone(), *v)).collect(),
|
|
};
|
|
|
|
let response = client.report_job_completion(report).await?;
|
|
let ack = response.into_inner();
|
|
info!(
|
|
"Service acknowledged: accepted={}, promotion_status={}",
|
|
ack.accepted, ack.promotion_status
|
|
);
|
|
Ok(())
|
|
}
|
|
```
|
|
|
|
Note: This references `JobCompletionReport` and `JobCompletionAck` proto messages that don't exist yet. They'll be added in Task 3 (proto update). The crate won't compile until then — that's expected.
|
|
|
|
**Step 4: Add to workspace**
|
|
|
|
In root `Cargo.toml`, add `"crates/training_uploader"` to the workspace members list.
|
|
|
|
**Step 5: Add async-recursion dependency**
|
|
|
|
Add `async-recursion = "1.0"` to `crates/training_uploader/Cargo.toml` dependencies.
|
|
|
|
**Step 6: Commit** (will not compile yet — proto messages needed from Task 3)
|
|
|
|
```bash
|
|
git add crates/training_uploader/ Cargo.toml
|
|
git commit -m "feat: add training-uploader sidecar crate (proto pending)"
|
|
```
|
|
|
|
---
|
|
|
|
## Task 3: Add new RPCs to ml_training.proto
|
|
|
|
Add the `ReportJobCompletion`, `ListPendingPromotions`, `ApprovePromotion`, and `RejectPromotion` RPCs.
|
|
|
|
**Files:**
|
|
- Modify: `services/ml_training_service/proto/ml_training.proto`
|
|
|
|
**Step 1: Add RPCs to the service definition**
|
|
|
|
In the `service MLTrainingService` block, add after the batch tuning RPCs:
|
|
|
|
```protobuf
|
|
// Job Completion Callback (called by training-uploader sidecar)
|
|
rpc ReportJobCompletion(JobCompletionReport) returns (JobCompletionAck);
|
|
|
|
// Model Promotion Management
|
|
rpc ListPendingPromotions(ListPendingPromotionsRequest) returns (ListPendingPromotionsResponse);
|
|
rpc ApprovePromotion(ApprovePromotionRequest) returns (ApprovePromotionResponse);
|
|
rpc RejectPromotion(RejectPromotionRequest) returns (RejectPromotionResponse);
|
|
```
|
|
|
|
**Step 2: Add the message definitions**
|
|
|
|
Add at the end of the proto file:
|
|
|
|
```protobuf
|
|
// --- Job Completion Callback (from training-uploader sidecar) ---
|
|
|
|
message JobCompletionReport {
|
|
string job_id = 1;
|
|
string s3_path = 2;
|
|
bool success = 3;
|
|
string error_message = 4;
|
|
map<string, double> metrics = 5;
|
|
}
|
|
|
|
message JobCompletionAck {
|
|
bool accepted = 1;
|
|
string promotion_status = 2; // "pending_promotion", "no_improvement", "registered", "error"
|
|
}
|
|
|
|
// --- Model Promotion ---
|
|
|
|
message ListPendingPromotionsRequest {}
|
|
|
|
message ListPendingPromotionsResponse {
|
|
repeated PendingPromotion promotions = 1;
|
|
}
|
|
|
|
message PendingPromotion {
|
|
string model_id = 1;
|
|
string model_type = 2;
|
|
string symbol = 3;
|
|
string s3_path = 4;
|
|
map<string, double> new_metrics = 5;
|
|
map<string, double> current_metrics = 6;
|
|
int64 trained_at = 7;
|
|
string job_id = 8;
|
|
}
|
|
|
|
message ApprovePromotionRequest {
|
|
string model_id = 1;
|
|
}
|
|
|
|
message ApprovePromotionResponse {
|
|
bool success = 1;
|
|
string message = 2;
|
|
}
|
|
|
|
message RejectPromotionRequest {
|
|
string model_id = 1;
|
|
string reason = 2;
|
|
}
|
|
|
|
message RejectPromotionResponse {
|
|
bool success = 1;
|
|
string message = 2;
|
|
}
|
|
```
|
|
|
|
**Step 3: Verify proto compiles**
|
|
|
|
Run: `SQLX_OFFLINE=true cargo check -p ml_training_service -p training_uploader`
|
|
Expected: compiles (the new RPCs won't have implementations yet — tonic generates default `unimplemented!()` stubs)
|
|
|
|
**Step 4: Commit**
|
|
|
|
```bash
|
|
git add services/ml_training_service/proto/ml_training.proto
|
|
git commit -m "feat(proto): add ReportJobCompletion and model promotion RPCs"
|
|
```
|
|
|
|
---
|
|
|
|
## Task 4: Implement K8s Job Dispatcher
|
|
|
|
New module in `ml_training_service` that creates K8s `batch/v1` Jobs using the `kube` crate.
|
|
|
|
**Files:**
|
|
- Create: `services/ml_training_service/src/k8s_dispatcher.rs`
|
|
- Modify: `services/ml_training_service/Cargo.toml`
|
|
- Modify: `services/ml_training_service/src/main.rs` (wire up)
|
|
|
|
**Step 1: Add kube dependencies**
|
|
|
|
Add to `services/ml_training_service/Cargo.toml` under `[dependencies]`:
|
|
|
|
```toml
|
|
# K8s API for job dispatch
|
|
kube = { version = "0.98", features = ["runtime", "client", "derive"] }
|
|
k8s-openapi = { version = "0.23", features = ["latest"] }
|
|
```
|
|
|
|
Also add to root `Cargo.toml` workspace dependencies if not already present.
|
|
|
|
**Step 2: Create k8s_dispatcher.rs**
|
|
|
|
Create `services/ml_training_service/src/k8s_dispatcher.rs`:
|
|
|
|
This module provides:
|
|
|
|
```rust
|
|
//! K8s Job Dispatcher
|
|
//!
|
|
//! Creates batch/v1 Jobs on the gpu-training pool for on-demand model training.
|
|
//! Each Job includes a training container and a native sidecar (training-uploader)
|
|
//! that handles S3 upload and gRPC completion callback.
|
|
|
|
use std::collections::BTreeMap;
|
|
|
|
use anyhow::{Context, Result};
|
|
use k8s_openapi::api::batch::v1::{Job, JobSpec};
|
|
use k8s_openapi::api::core::v1::{
|
|
Container, EnvVar, PodSpec, PodTemplateSpec, ResourceRequirements,
|
|
Volume, VolumeMount, PersistentVolumeClaimVolumeSource, EmptyDirVolumeSource,
|
|
};
|
|
use k8s_openapi::apimachinery::pkg::api::resource::Quantity;
|
|
use k8s_openapi::apimachinery::pkg::apis::meta::v1::ObjectMeta;
|
|
use kube::{Api, Client};
|
|
use tracing::{info, error};
|
|
use uuid::Uuid;
|
|
|
|
/// Configuration for the K8s job dispatcher
|
|
#[derive(Debug, Clone)]
|
|
pub struct DispatcherConfig {
|
|
pub namespace: String,
|
|
pub training_image: String,
|
|
pub uploader_image: String,
|
|
pub data_pvc: String,
|
|
pub output_pvc: String,
|
|
pub s3_bucket: String,
|
|
pub s3_endpoint: String,
|
|
pub s3_region: String,
|
|
pub callback_endpoint: String, // e.g. "http://ml-training-service.foxhunt:50053"
|
|
}
|
|
|
|
impl Default for DispatcherConfig {
|
|
fn default() -> Self {
|
|
Self {
|
|
namespace: std::env::var("K8S_NAMESPACE").unwrap_or_else(|_| "foxhunt".to_owned()),
|
|
training_image: std::env::var("TRAINING_IMAGE")
|
|
.unwrap_or_else(|_| "rg.fr-par.scw.cloud/foxhunt-ci/training:latest".to_owned()),
|
|
uploader_image: std::env::var("UPLOADER_IMAGE")
|
|
.unwrap_or_else(|_| "rg.fr-par.scw.cloud/foxhunt-ci/training:latest".to_owned()),
|
|
data_pvc: "training-data-pvc".to_owned(),
|
|
output_pvc: "training-output-pvc".to_owned(),
|
|
s3_bucket: std::env::var("S3_BUCKET").unwrap_or_else(|_| "foxhunt-models".to_owned()),
|
|
s3_endpoint: std::env::var("S3_ENDPOINT")
|
|
.unwrap_or_else(|_| "https://s3.fr-par.scw.cloud".to_owned()),
|
|
s3_region: "fr-par".to_owned(),
|
|
callback_endpoint: std::env::var("CALLBACK_ENDPOINT")
|
|
.unwrap_or_else(|_| "http://ml-training-service.foxhunt.svc.cluster.local:50053".to_owned()),
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Dispatch a training job to K8s
|
|
pub struct K8sDispatcher {
|
|
client: Client,
|
|
config: DispatcherConfig,
|
|
}
|
|
|
|
/// Parameters for a training job dispatch
|
|
#[derive(Debug, Clone)]
|
|
pub struct TrainingJobParams {
|
|
pub job_id: Uuid,
|
|
pub model_type: String, // e.g. "dqn", "tft", "mamba2"
|
|
pub symbol: String, // e.g. "ES.FUT"
|
|
pub epochs: u32,
|
|
pub binary: String, // e.g. "train_baseline_rl" or "train_baseline_supervised"
|
|
}
|
|
```
|
|
|
|
The `K8sDispatcher` should have:
|
|
|
|
- `async fn new() -> Result<Self>` — creates kube Client from in-cluster config
|
|
- `async fn dispatch(&self, params: &TrainingJobParams) -> Result<String>` — creates the Job, returns the K8s Job name
|
|
- `fn build_job_spec(&self, params: &TrainingJobParams) -> Job` — builds the Job manifest with training container + sidecar
|
|
- `async fn get_job_status(&self, job_name: &str) -> Result<Option<String>>` — polls Job status
|
|
- `async fn delete_job(&self, job_name: &str) -> Result<()>` — cleanup
|
|
|
|
The `build_job_spec` method should produce a Job matching the structure in `infra/k8s/training/job-template.yaml` but with the sidecar added:
|
|
|
|
- Main container: `training` — runs the specified binary with `--model`, `--symbol`, `--data-dir=/data`, `--output-dir=/output/{job_id}`
|
|
- Sidecar container: `uploader` — runs `training-uploader` with `--output-dir=/output/{job_id}`, S3 config, callback endpoint, job ID
|
|
- Sidecar uses `restartPolicy: Always` (K8s 1.28+ native sidecar)
|
|
- Shared volumes: `training-data` (PVC, RO), `output` (PVC, RW)
|
|
- Node selector: `gpu-training` pool
|
|
- GPU requests: 1 nvidia.com/gpu
|
|
- Labels: `foxhunt/job-type: training`, `foxhunt/model: {model}`, `foxhunt/job-id: {uuid}`
|
|
|
|
**Step 3: Verify**
|
|
|
|
Run: `SQLX_OFFLINE=true cargo check -p ml_training_service`
|
|
Expected: compiles
|
|
|
|
**Step 4: Commit**
|
|
|
|
```bash
|
|
git add services/ml_training_service/src/k8s_dispatcher.rs services/ml_training_service/Cargo.toml Cargo.toml
|
|
git commit -m "feat(service): add K8s Job dispatcher for GPU training"
|
|
```
|
|
|
|
---
|
|
|
|
## Task 5: Implement ReportJobCompletion + promotion logic
|
|
|
|
Wire the new `ReportJobCompletion` RPC in the service. On completion callback: register checkpoint, compare metrics, set promotion status.
|
|
|
|
**Files:**
|
|
- Create: `services/ml_training_service/src/promotion_manager.rs`
|
|
- Modify: `services/ml_training_service/src/service.rs`
|
|
|
|
**Step 1: Create promotion_manager.rs**
|
|
|
|
This module tracks models pending promotion and compares metrics:
|
|
|
|
```rust
|
|
//! Model Promotion Manager
|
|
//!
|
|
//! Tracks models awaiting operator approval for promotion to active status.
|
|
//! Compares new model metrics against currently active model to determine
|
|
//! if promotion should be offered.
|
|
|
|
use std::collections::HashMap;
|
|
use std::sync::Arc;
|
|
|
|
use chrono::{DateTime, Utc};
|
|
use tokio::sync::RwLock;
|
|
use tracing::info;
|
|
use uuid::Uuid;
|
|
|
|
#[derive(Debug, Clone)]
|
|
pub struct PendingModel {
|
|
pub model_id: Uuid,
|
|
pub model_type: String,
|
|
pub symbol: String,
|
|
pub s3_path: String,
|
|
pub new_metrics: HashMap<String, f64>,
|
|
pub current_metrics: HashMap<String, f64>,
|
|
pub trained_at: DateTime<Utc>,
|
|
pub job_id: String,
|
|
}
|
|
|
|
pub struct PromotionManager {
|
|
pending: Arc<RwLock<HashMap<Uuid, PendingModel>>>,
|
|
active_models: Arc<RwLock<HashMap<(String, String), ActiveModel>>>, // (model_type, symbol) -> active
|
|
}
|
|
|
|
#[derive(Debug, Clone)]
|
|
pub struct ActiveModel {
|
|
pub model_id: Uuid,
|
|
pub s3_path: String,
|
|
pub metrics: HashMap<String, f64>,
|
|
pub promoted_at: DateTime<Utc>,
|
|
}
|
|
```
|
|
|
|
Methods:
|
|
- `new()` — initialize empty
|
|
- `register_completion(job_id, model_type, symbol, s3_path, metrics) -> PromotionStatus` — compare against active, return status
|
|
- `list_pending() -> Vec<PendingModel>`
|
|
- `approve(model_id) -> Result<()>` — move to active, remove from pending
|
|
- `reject(model_id, reason) -> Result<()>` — remove from pending
|
|
- `get_active(model_type, symbol) -> Option<ActiveModel>`
|
|
|
|
Comparison logic: if new model has better `best_val_loss` (lower) or better `sharpe_ratio` (higher) than active, mark as `pending_promotion`. If no active model exists, auto-promote.
|
|
|
|
**Step 2: Implement the gRPC handlers in service.rs**
|
|
|
|
Add implementations for the 4 new RPCs in the `MlTrainingService` trait impl. Each delegates to `PromotionManager`.
|
|
|
|
**Step 3: Wire PromotionManager into MLTrainingServiceImpl**
|
|
|
|
Add `promotion_manager: Arc<PromotionManager>` field. Initialize in `main.rs`.
|
|
|
|
**Step 4: Verify**
|
|
|
|
Run: `SQLX_OFFLINE=true cargo check -p ml_training_service`
|
|
|
|
**Step 5: Commit**
|
|
|
|
```bash
|
|
git add services/ml_training_service/src/promotion_manager.rs services/ml_training_service/src/service.rs \
|
|
services/ml_training_service/src/main.rs
|
|
git commit -m "feat(service): implement ReportJobCompletion and model promotion"
|
|
```
|
|
|
|
---
|
|
|
|
## Task 6: Wire StartTraining to K8s dispatch
|
|
|
|
Modify the existing `StartTraining` RPC to dispatch a K8s Job instead of running training in-process.
|
|
|
|
**Files:**
|
|
- Modify: `services/ml_training_service/src/service.rs`
|
|
- Modify: `services/ml_training_service/src/main.rs`
|
|
|
|
**Step 1: Add K8sDispatcher to MLTrainingServiceImpl**
|
|
|
|
Add `dispatcher: Arc<K8sDispatcher>` field. Initialize in `main.rs` with `K8sDispatcher::new().await`.
|
|
|
|
**Step 2: Modify StartTraining handler**
|
|
|
|
The current `start_training` handler calls `self.orchestrator.submit_training_job()` which runs in-process. Change it to:
|
|
1. Determine the correct binary (RL vs supervised) from model_type
|
|
2. Create `TrainingJobParams`
|
|
3. Call `self.dispatcher.dispatch(¶ms).await`
|
|
4. Record the job in the orchestrator for status tracking
|
|
5. Return the job_id
|
|
|
|
The orchestrator's in-memory job tracking stays — it just no longer runs the training itself. Job status updates will come from the `ReportJobCompletion` callback instead.
|
|
|
|
**Step 3: Verify**
|
|
|
|
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/main.rs
|
|
git commit -m "feat(service): wire StartTraining to K8s Job dispatch"
|
|
```
|
|
|
|
---
|
|
|
|
## Task 7: Add fxt model subcommand
|
|
|
|
Add `fxt model list`, `fxt model approve`, `fxt model reject` commands.
|
|
|
|
**Files:**
|
|
- Create: `bin/fxt/src/commands/model/mod.rs`
|
|
- Create: `bin/fxt/src/commands/model/list.rs`
|
|
- Create: `bin/fxt/src/commands/model/approve.rs`
|
|
- Create: `bin/fxt/src/commands/model/reject.rs`
|
|
- Modify: `bin/fxt/src/commands/mod.rs`
|
|
- Modify: `bin/fxt/src/main.rs` (add Model subcommand to CLI)
|
|
|
|
**Step 1: Create model command module**
|
|
|
|
`bin/fxt/src/commands/model/mod.rs`:
|
|
|
|
```rust
|
|
pub mod approve;
|
|
pub mod list;
|
|
pub mod reject;
|
|
|
|
use anyhow::Result;
|
|
use clap::Subcommand;
|
|
|
|
#[derive(Subcommand, Debug)]
|
|
pub enum ModelCommand {
|
|
/// List models with active/pending status
|
|
List,
|
|
/// Approve a pending model promotion
|
|
Approve {
|
|
/// Model ID to approve
|
|
model_id: String,
|
|
},
|
|
/// Reject a pending model promotion
|
|
Reject {
|
|
/// Model ID to reject
|
|
model_id: String,
|
|
/// Reason for rejection
|
|
#[arg(long, default_value = "operator rejected")]
|
|
reason: String,
|
|
},
|
|
}
|
|
|
|
pub async fn execute_model_command(
|
|
command: ModelCommand,
|
|
api_gateway_url: &str,
|
|
jwt_token: &str,
|
|
) -> Result<()> {
|
|
match command {
|
|
ModelCommand::List => list::run(api_gateway_url, jwt_token).await,
|
|
ModelCommand::Approve { model_id } => approve::run(api_gateway_url, jwt_token, &model_id).await,
|
|
ModelCommand::Reject { model_id, reason } => reject::run(api_gateway_url, jwt_token, &model_id, &reason).await,
|
|
}
|
|
}
|
|
```
|
|
|
|
**Step 2: Implement list, approve, reject**
|
|
|
|
Each makes a gRPC call to the api_gateway (which proxies to ml_training_service). Follow the same pattern as the existing `train/list.rs` and `train/status.rs` commands.
|
|
|
|
- `list.rs` — calls `ListPendingPromotions`, displays table with model_type, symbol, metrics comparison, trained_at
|
|
- `approve.rs` — calls `ApprovePromotion`, prints success/failure
|
|
- `reject.rs` — calls `RejectPromotion`, prints success/failure
|
|
|
|
**Step 3: Wire into commands/mod.rs and main.rs**
|
|
|
|
Add `pub mod model;` to `commands/mod.rs`. Add `Model` variant to the main CLI enum in `main.rs`. Follow the pattern of the existing `Train` and `Tune` commands.
|
|
|
|
**Step 4: Verify**
|
|
|
|
Run: `SQLX_OFFLINE=true cargo check -p fxt`
|
|
|
|
**Step 5: Commit**
|
|
|
|
```bash
|
|
git add bin/fxt/src/commands/model/ bin/fxt/src/commands/mod.rs bin/fxt/src/main.rs
|
|
git commit -m "feat(fxt): add model list/approve/reject commands"
|
|
```
|
|
|
|
---
|
|
|
|
## Task 8: Add fxt train start subcommand
|
|
|
|
The existing `fxt train` has `list` and `status`. Add `start` to trigger on-demand training.
|
|
|
|
**Files:**
|
|
- Create: `bin/fxt/src/commands/train/start.rs`
|
|
- Modify: `bin/fxt/src/commands/train/mod.rs`
|
|
|
|
**Step 1: Create start.rs**
|
|
|
|
```rust
|
|
//! fxt train start — trigger on-demand training job
|
|
|
|
use anyhow::Result;
|
|
use clap::Args;
|
|
|
|
#[derive(Args, Debug)]
|
|
pub struct StartCommand {
|
|
/// Model type (dqn, ppo, tft, mamba2, tggn, tlob, liquid, kan, xlstm, diffusion)
|
|
pub model: String,
|
|
|
|
/// Symbol to train on (e.g. ES.FUT)
|
|
pub symbol: String,
|
|
|
|
/// Number of training epochs
|
|
#[arg(long, default_value = "50")]
|
|
pub epochs: u32,
|
|
}
|
|
```
|
|
|
|
The `run()` method calls `StartTraining` gRPC, prints the job ID and a status URL.
|
|
|
|
**Step 2: Wire into train/mod.rs**
|
|
|
|
Add `Start` variant to `TrainCommand` enum. Add `pub mod start;`.
|
|
|
|
**Step 3: Verify**
|
|
|
|
Run: `SQLX_OFFLINE=true cargo check -p fxt`
|
|
|
|
**Step 4: Commit**
|
|
|
|
```bash
|
|
git add bin/fxt/src/commands/train/start.rs bin/fxt/src/commands/train/mod.rs
|
|
git commit -m "feat(fxt): add train start command for on-demand training"
|
|
```
|
|
|
|
---
|
|
|
|
## Task 9: Update Dockerfile and K8s job template
|
|
|
|
Add `training-uploader` binary to the training Docker image and update the job template with the native sidecar.
|
|
|
|
**Files:**
|
|
- Modify: `infra/docker/Dockerfile.training` (add training-uploader to build)
|
|
- Modify: `infra/docker/Dockerfile.training-runtime` (copy uploader binary)
|
|
- Modify: `infra/k8s/training/job-template.yaml` (add sidecar)
|
|
- Modify: `.gitlab-ci.yml` (compile training_uploader in compile-services stage)
|
|
|
|
**Step 1: Update Dockerfile.training-runtime**
|
|
|
|
Add `COPY --from=build training_uploader /usr/local/bin/training-uploader` (or equivalent pattern matching the existing Dockerfile).
|
|
|
|
**Step 2: Update job-template.yaml**
|
|
|
|
Add the sidecar container alongside the training container:
|
|
|
|
```yaml
|
|
# Native sidecar (K8s 1.28+) — uploads artifacts on completion
|
|
initContainers:
|
|
- name: uploader
|
|
image: rg.fr-par.scw.cloud/foxhunt-ci/training:latest
|
|
restartPolicy: Always # Makes this a native sidecar
|
|
command: ["/usr/local/bin/training-uploader"]
|
|
env:
|
|
- name: JOB_ID
|
|
valueFrom:
|
|
fieldRef:
|
|
fieldPath: metadata.labels['foxhunt/job-id']
|
|
- name: S3_BUCKET
|
|
value: foxhunt-models
|
|
- name: S3_ENDPOINT
|
|
value: https://s3.fr-par.scw.cloud
|
|
- name: S3_PREFIX
|
|
value: "models/$(JOB_ID)"
|
|
- name: CALLBACK_ENDPOINT
|
|
value: http://ml-training-service.foxhunt.svc.cluster.local:50053
|
|
envFrom:
|
|
- secretRef:
|
|
name: s3-credentials
|
|
volumeMounts:
|
|
- name: output
|
|
mountPath: /output
|
|
readOnly: true
|
|
resources:
|
|
requests:
|
|
cpu: 100m
|
|
memory: 128Mi
|
|
limits:
|
|
cpu: 500m
|
|
memory: 256Mi
|
|
```
|
|
|
|
**Step 3: Update .gitlab-ci.yml compile-services stage**
|
|
|
|
Add `training_uploader` to the cargo build command that builds training binaries:
|
|
|
|
```yaml
|
|
- cargo build --release -p training_uploader
|
|
```
|
|
|
|
And copy it to build-out:
|
|
|
|
```yaml
|
|
cp target/release/training_uploader build-out/
|
|
strip build-out/training_uploader
|
|
```
|
|
|
|
**Step 4: Commit**
|
|
|
|
```bash
|
|
git add infra/docker/ infra/k8s/training/job-template.yaml .gitlab-ci.yml
|
|
git commit -m "infra: add training-uploader sidecar to Docker image and K8s job template"
|
|
```
|
|
|
|
---
|
|
|
|
## Task 10: Integration test
|
|
|
|
End-to-end test: verify the service can build a Job spec, the proto compiles, and the fxt commands parse correctly.
|
|
|
|
**Files:**
|
|
- Modify: `services/ml_training_service/src/k8s_dispatcher.rs` (add unit tests)
|
|
|
|
**Step 1: Add unit tests for Job spec building**
|
|
|
|
Test that `build_job_spec()` produces a valid Job with:
|
|
- Correct main container (binary, args, env)
|
|
- Sidecar container with correct env vars
|
|
- GPU resource requests
|
|
- Correct volume mounts
|
|
- Correct labels and node selector
|
|
|
|
```rust
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_build_job_spec() {
|
|
let config = DispatcherConfig::default();
|
|
let dispatcher = K8sDispatcher { client: todo!(), config };
|
|
// ... test job spec fields
|
|
}
|
|
}
|
|
```
|
|
|
|
Since `kube::Client` requires a cluster, extract `build_job_spec` as a pure function that takes `&DispatcherConfig` and `&TrainingJobParams` and returns `Job`. This makes it testable without a cluster.
|
|
|
|
**Step 2: Test fxt CLI parsing**
|
|
|
|
Run: `SQLX_OFFLINE=true cargo check -p fxt`
|
|
|
|
Verify `fxt train start --help` and `fxt model --help` parse correctly.
|
|
|
|
**Step 3: Full workspace check**
|
|
|
|
Run: `SQLX_OFFLINE=true cargo check --workspace`
|
|
Expected: zero errors
|
|
|
|
**Step 4: Commit**
|
|
|
|
```bash
|
|
git add services/ml_training_service/src/k8s_dispatcher.rs
|
|
git commit -m "test(service): add K8s dispatcher unit tests"
|
|
```
|
|
|
|
---
|
|
|
|
## Dependency Order
|
|
|
|
```
|
|
Task 1 (markers) ──┐
|
|
Task 2 (uploader) ──┼── Task 3 (proto) ── Task 5 (completion handler) ── Task 6 (wire StartTraining)
|
|
│ └── Task 4 (dispatcher) ──────────┘
|
|
│
|
|
Task 7 (fxt model) ──┤
|
|
Task 8 (fxt train start) ─┤
|
|
└── Task 9 (Docker/K8s) ── Task 10 (integration test)
|
|
```
|
|
|
|
Tasks 1, 2, 7, 8 can start in parallel. Task 3 unblocks 4 and 5. Task 9 is last infra change.
|