gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(ml_training_service): align K8s job spec with production runtime

Browse Source 
- Replace separate training/uploader images with single runtime_image
- Add fetch-binaries initContainer (rclone from S3 binaries bucket)
- Switch to emptyDir for output and binaries (no output PVC needed)
- Add Cilium CNI toleration for fresh scale-from-zero nodes
- Extract symbol from file_path last component in data_source
- Increase active_deadline_seconds to 6 hours for hyperopt jobs

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-02-28 20:58:37 +01:00
parent 14d405f112
commit b5fca4d158
2 changed files with 192 additions and 108 deletions
Download Patch File Download Diff File Expand all files Collapse all files

290
services/ml_training_service/src/k8s_dispatcher.rs
View File

@@ -9,9 +9,9 @@ use std::collections::BTreeMap;
use anyhow::{Context, Result};
use k8s_openapi::api::batch::v1::{Job, JobSpec};
use k8s_openapi::api::core::v1::{
Container, EnvFromSource, EnvVar, LocalObjectReference, PersistentVolumeClaimVolumeSource,
PodSpec, PodTemplateSpec, ResourceRequirements, SecretEnvSource, Toleration, Volume,
VolumeMount,
Container, EmptyDirVolumeSource, EnvFromSource, EnvVar, LocalObjectReference,
PersistentVolumeClaimVolumeSource, PodSpec, PodTemplateSpec, ResourceRequirements,
SecretEnvSource, SecurityContext, Toleration, Volume, VolumeMount,
};
use k8s_openapi::apimachinery::pkg::api::resource::Quantity;
use k8s_openapi::apimachinery::pkg::apis::meta::v1::ObjectMeta;
@@ -30,13 +30,12 @@ use uuid::Uuid;
#[derive(Debug, Clone)]
pub struct DispatcherConfig {
pub namespace: String,
pub training_image: String,
pub uploader_image: String,
pub runtime_image: String,
pub data_pvc: String,
pub output_pvc: String,
pub s3_bucket: String,
pub s3_endpoint: String,
pub s3_region: String,
pub binaries_bucket: String,
pub callback_endpoint: String,
}
@@ -51,19 +50,17 @@ impl DispatcherConfig {
pub fn from_env() -> Self {
Self {
namespace: std::env::var("K8S_NAMESPACE").unwrap_or_else(|_| "foxhunt".to_string()),
training_image: std::env::var("TRAINING_IMAGE").unwrap_or_else(|_| {
"rg.fr-par.scw.cloud/foxhunt-ci/training:latest".to_string()
}),
uploader_image: std::env::var("UPLOADER_IMAGE").unwrap_or_else(|_| {
"rg.fr-par.scw.cloud/foxhunt-ci/training:latest".to_string()
runtime_image: std::env::var("TRAINING_RUNTIME_IMAGE").unwrap_or_else(|_| {
"rg.fr-par.scw.cloud/foxhunt-ci/foxhunt-training-runtime:latest".to_string()
}),
data_pvc: "training-data-pvc".to_string(),
output_pvc: "training-output-pvc".to_string(),
s3_bucket: std::env::var("S3_BUCKET")
.unwrap_or_else(|_| "foxhunt-models".to_string()),
s3_endpoint: std::env::var("S3_ENDPOINT")
.unwrap_or_else(|_| "https://s3.fr-par.scw.cloud".to_string()),
s3_region: "fr-par".to_string(),
binaries_bucket: std::env::var("BINARIES_BUCKET")
.unwrap_or_else(|_| "foxhunt-binaries".to_string()),
callback_endpoint: std::env::var("CALLBACK_ENDPOINT").unwrap_or_else(|_| {
"http://ml-training-service.foxhunt.svc.cluster.local:50053".to_string()
}),
@@ -187,8 +184,11 @@ impl K8sDispatcher {
/// Builds a complete `batch/v1` Job manifest without touching the cluster.
///
/// This is a pure function, suitable for unit testing without a Tokio runtime
/// or K8s client.
/// Matches the production `infra/k8s/training/job-template.yaml`:
/// - fetch-binaries initContainer (rclone from S3)
/// - uploader native sidecar (K8s 1.28+)
/// - main training container with GPU
/// - emptyDir for output and binaries, PVC for training data
pub fn build_job_spec(config: &DispatcherConfig, params: &TrainingJobParams) -> Job {
let raw_name = format!("training-{}-{}", params.model_type, params.job_id);
// K8s names are limited to 63 characters and must be lowercase DNS-compatible.
@@ -212,18 +212,83 @@ pub fn build_job_spec(config: &DispatcherConfig, params: &TrainingJobParams) ->
params.job_id.to_string(),
);
// -- Sidecar (native K8s 1.28+) ------------------------------------------
// -- fetch-binaries initContainer (runs first, downloads from S3) ---------
let fetch_binaries = Container {
name: "fetch-binaries".to_string(),
image: Some(config.runtime_image.clone()),
security_context: Some(SecurityContext {
run_as_user: Some(0),
..Default::default()
}),
command: Some(vec!["/bin/sh".to_string(), "-c".to_string()]),
args: Some(vec![format!(
"set -e\n\
export RCLONE_S3_PROVIDER=Scaleway RCLONE_S3_ENDPOINT=s3.fr-par.scw.cloud RCLONE_S3_REGION=fr-par\n\
rclone sync :s3:{}/latest/training/ /binaries/\n\
chmod +x /binaries/*\n\
ls -lh /binaries/\n\
echo \"Fetched training binaries from S3\"",
config.binaries_bucket
)]),
env: Some(vec![
EnvVar {
name: "RCLONE_S3_ACCESS_KEY_ID".to_string(),
value_from: Some(k8s_openapi::api::core::v1::EnvVarSource {
secret_key_ref: Some(k8s_openapi::api::core::v1::SecretKeySelector {
name: "s3-credentials".to_string(),
key: "access-key".to_string(),
optional: Some(false),
}),
..Default::default()
}),
..Default::default()
},
EnvVar {
name: "RCLONE_S3_SECRET_ACCESS_KEY".to_string(),
value_from: Some(k8s_openapi::api::core::v1::EnvVarSource {
secret_key_ref: Some(k8s_openapi::api::core::v1::SecretKeySelector {
name: "s3-credentials".to_string(),
key: "secret-key".to_string(),
optional: Some(false),
}),
..Default::default()
}),
..Default::default()
},
]),
volume_mounts: Some(vec![VolumeMount {
name: "binaries".to_string(),
mount_path: "/binaries".to_string(),
..Default::default()
}]),
resources: Some(ResourceRequirements {
requests: Some(BTreeMap::from([
("cpu".to_string(), Quantity("100m".to_string())),
("memory".to_string(), Quantity("64Mi".to_string())),
])),
limits: Some(BTreeMap::from([
("cpu".to_string(), Quantity("500m".to_string())),
("memory".to_string(), Quantity("256Mi".to_string())),
])),
..Default::default()
}),
..Default::default()
};
// -- Uploader sidecar (native K8s 1.28+) ----------------------------------
let uploader = Container {
name: "uploader".to_string(),
image: Some(config.uploader_image.clone()),
image: Some(config.runtime_image.clone()),
restart_policy: Some("Always".to_string()), // native sidecar
command: Some(vec!["/usr/local/bin/training_uploader".to_string()]),
command: Some(vec!["/binaries/training_uploader".to_string()]),
env: Some(vec![
env_var("JOB_ID", &params.job_id.to_string()),
env_var("S3_BUCKET", &config.s3_bucket),
env_var("S3_ENDPOINT", &config.s3_endpoint),
env_var("S3_REGION", &config.s3_region),
env_var("S3_PREFIX", &format!("models/{}", params.job_id)),
env_var("CALLBACK_ENDPOINT", &config.callback_endpoint),
env_var("RUST_LOG", "info"),
]),
env_from: Some(vec![EnvFromSource {
secret_ref: Some(SecretEnvSource {
@@ -232,12 +297,20 @@ pub fn build_job_spec(config: &DispatcherConfig, params: &TrainingJobParams) ->
}),
..Default::default()
}]),
volume_mounts: Some(vec![VolumeMount {
name: "output".to_string(),
mount_path: "/output".to_string(),
read_only: Some(false),
..Default::default()
}]),
volume_mounts: Some(vec![
VolumeMount {
name: "output".to_string(),
mount_path: "/output".to_string(),
read_only: Some(true),
..Default::default()
},
VolumeMount {
name: "binaries".to_string(),
mount_path: "/binaries".to_string(),
read_only: Some(true),
..Default::default()
},
]),
resources: Some(ResourceRequirements {
requests: Some(BTreeMap::from([
("cpu".to_string(), Quantity("100m".to_string())),
@@ -255,15 +328,12 @@ pub fn build_job_spec(config: &DispatcherConfig, params: &TrainingJobParams) ->
// -- Main training container ----------------------------------------------
let training = Container {
name: "training".to_string(),
image: Some(config.training_image.clone()),
command: Some(vec![format!("/usr/local/bin/{}", params.binary)]),
image: Some(config.runtime_image.clone()),
command: Some(vec![format!("/binaries/{}", params.binary)]),
args: Some(vec![
"--model".to_string(),
params.model_type.clone(),
"--symbol".to_string(),
params.symbol.clone(),
"--data-dir=/data".to_string(),
format!("--output-dir=/output/{}", params.job_id),
format!("--symbol={}", params.symbol),
"--data-dir=/data/futures-baseline".to_string(),
"--output-dir=/output".to_string(),
format!("--epochs={}", params.epochs),
]),
env: Some(vec![
@@ -280,7 +350,12 @@ pub fn build_job_spec(config: &DispatcherConfig, params: &TrainingJobParams) ->
VolumeMount {
name: "output".to_string(),
mount_path: "/output".to_string(),
read_only: Some(false),
..Default::default()
},
VolumeMount {
name: "binaries".to_string(),
mount_path: "/binaries".to_string(),
read_only: Some(true),
..Default::default()
},
]),
@@ -312,9 +387,17 @@ pub fn build_job_spec(config: &DispatcherConfig, params: &TrainingJobParams) ->
},
Volume {
name: "output".to_string(),
persistent_volume_claim: Some(PersistentVolumeClaimVolumeSource {
claim_name: config.output_pvc.clone(),
read_only: Some(false),
empty_dir: Some(EmptyDirVolumeSource {
size_limit: Some(Quantity("2Gi".to_string())),
..Default::default()
}),
..Default::default()
},
Volume {
name: "binaries".to_string(),
empty_dir: Some(EmptyDirVolumeSource {
size_limit: Some(Quantity("500Mi".to_string())),
..Default::default()
}),
..Default::default()
},
@@ -329,16 +412,26 @@ pub fn build_job_spec(config: &DispatcherConfig, params: &TrainingJobParams) ->
// -- Pod spec --------------------------------------------------------------
let pod_spec = PodSpec {
init_containers: Some(vec![uploader]),
// fetch-binaries runs first (sequential), then uploader starts as native sidecar
init_containers: Some(vec![fetch_binaries, uploader]),
containers: vec![training],
volumes: Some(volumes),
node_selector: Some(node_selector),
tolerations: Some(vec![Toleration {
key: Some("nvidia.com/gpu".to_string()),
operator: Some("Exists".to_string()),
effect: Some("NoSchedule".to_string()),
..Default::default()
}]),
tolerations: Some(vec![
Toleration {
key: Some("nvidia.com/gpu".to_string()),
operator: Some("Exists".to_string()),
effect: Some("NoSchedule".to_string()),
..Default::default()
},
// Cilium CNI takes ~30s on fresh scale-from-zero nodes
Toleration {
key: Some("node.cilium.io/agent-not-ready".to_string()),
operator: Some("Exists".to_string()),
effect: Some("NoSchedule".to_string()),
..Default::default()
},
]),
image_pull_secrets: Some(vec![LocalObjectReference {
name: "scw-registry".to_string(),
}]),
@@ -356,7 +449,7 @@ pub fn build_job_spec(config: &DispatcherConfig, params: &TrainingJobParams) ->
},
spec: Some(JobSpec {
backoff_limit: Some(1),
active_deadline_seconds: Some(3600),
active_deadline_seconds: Some(21600), // 6 hours
ttl_seconds_after_finished: Some(600),
template: PodTemplateSpec {
metadata: Some(ObjectMeta {
@@ -410,13 +503,12 @@ mod tests {
fn test_config() -> DispatcherConfig {
DispatcherConfig {
namespace: "foxhunt-test".to_string(),
training_image: "registry.example.com/training:v1".to_string(),
uploader_image: "registry.example.com/uploader:v1".to_string(),
runtime_image: "registry.example.com/foxhunt-training-runtime:v1".to_string(),
data_pvc: "training-data-pvc".to_string(),
output_pvc: "training-output-pvc".to_string(),
s3_bucket: "test-bucket".to_string(),
s3_endpoint: "https://s3.test.example.com".to_string(),
s3_region: "us-test-1".to_string(),
binaries_bucket: "test-binaries".to_string(),
callback_endpoint: "http://callback.test.svc:50053".to_string(),
}
}
@@ -481,7 +573,7 @@ mod tests {
// -- Job spec --
let spec = job.spec.as_ref().unwrap();
assert_eq!(spec.backoff_limit, Some(1));
assert_eq!(spec.active_deadline_seconds, Some(3600));
assert_eq!(spec.active_deadline_seconds, Some(21600)); // 6 hours
assert_eq!(spec.ttl_seconds_after_finished, Some(600));
// -- Pod spec --
@@ -494,11 +586,11 @@ mod tests {
"gpu-training"
);
// Tolerations
// Tolerations (GPU + Cilium)
let tols = pod_spec.tolerations.as_ref().unwrap();
assert_eq!(tols.len(), 1);
assert_eq!(tols.len(), 2);
assert_eq!(tols[0].key.as_deref().unwrap(), "nvidia.com/gpu");
assert_eq!(tols[0].effect.as_deref().unwrap(), "NoSchedule");
assert_eq!(tols[1].key.as_deref().unwrap(), "node.cilium.io/agent-not-ready");
// Image pull secrets
let ips = pod_spec.image_pull_secrets.as_ref().unwrap();
@@ -507,10 +599,21 @@ mod tests {
// Restart policy
assert_eq!(pod_spec.restart_policy.as_deref().unwrap(), "Never");
// -- Sidecar (init container with restart_policy: Always) --
// -- Init containers: fetch-binaries + uploader sidecar --
let init_containers = pod_spec.init_containers.as_ref().unwrap();
assert_eq!(init_containers.len(), 1);
let sidecar = &init_containers[0];
assert_eq!(init_containers.len(), 2);
// fetch-binaries (sequential init, runs first)
let fetcher = &init_containers[0];
assert_eq!(fetcher.name, "fetch-binaries");
assert!(fetcher.restart_policy.is_none(), "fetcher is not a sidecar");
assert_eq!(
fetcher.security_context.as_ref().unwrap().run_as_user,
Some(0)
);
// uploader (native sidecar)
let sidecar = &init_containers[1];
assert_eq!(sidecar.name, "uploader");
assert_eq!(
sidecar.restart_policy.as_deref().unwrap(),
@@ -519,7 +622,7 @@ mod tests {
);
assert_eq!(
sidecar.image.as_deref().unwrap(),
"registry.example.com/uploader:v1"
"registry.example.com/foxhunt-training-runtime:v1"
);
// Sidecar env
@@ -531,27 +634,11 @@ mod tests {
assert!(env_names.contains(&"S3_PREFIX"));
assert!(env_names.contains(&"CALLBACK_ENDPOINT"));
// Sidecar envFrom (s3-credentials secret)
let env_from = sidecar.env_from.as_ref().unwrap();
assert_eq!(env_from.len(), 1);
let secret_name = &env_from[0]
.secret_ref
.as_ref()
.unwrap()
.name;
assert_eq!(secret_name, "s3-credentials");
// Sidecar volume mounts
// Sidecar volume mounts (output + binaries)
let sidecar_mounts = sidecar.volume_mounts.as_ref().unwrap();
assert_eq!(sidecar_mounts.len(), 1);
assert_eq!(sidecar_mounts[0].name, "output");
assert_eq!(sidecar_mounts[0].mount_path, "/output");
// Sidecar resources (lightweight)
let sidecar_res = sidecar.resources.as_ref().unwrap();
let sidecar_req = sidecar_res.requests.as_ref().unwrap();
assert_eq!(sidecar_req.get("cpu").unwrap().0, "100m");
assert_eq!(sidecar_req.get("memory").unwrap().0, "128Mi");
assert_eq!(sidecar_mounts.len(), 2);
assert!(sidecar_mounts.iter().any(|m| m.name == "output"));
assert!(sidecar_mounts.iter().any(|m| m.name == "binaries"));
// -- Main container --
let containers = &pod_spec.containers;
@@ -560,20 +647,17 @@ mod tests {
assert_eq!(main.name, "training");
assert_eq!(
main.image.as_deref().unwrap(),
"registry.example.com/training:v1"
"registry.example.com/foxhunt-training-runtime:v1"
);
// Command and args
// Command uses /binaries/ path
let cmd = main.command.as_ref().unwrap();
assert_eq!(cmd[0], "/usr/local/bin/train_baseline_supervised");
assert_eq!(cmd[0], "/binaries/train_baseline_supervised");
let args = main.args.as_ref().unwrap();
assert!(args.contains(&"--model".to_string()));
assert!(args.contains(&"tft".to_string()));
assert!(args.contains(&"--symbol".to_string()));
assert!(args.contains(&"ES.FUT".to_string()));
assert!(args.contains(&"--data-dir=/data".to_string()));
assert!(args.iter().any(|a| a.starts_with("--output-dir=/output/")));
assert!(args.contains(&"--symbol=ES.FUT".to_string()));
assert!(args.contains(&"--data-dir=/data/futures-baseline".to_string()));
assert!(args.contains(&"--output-dir=/output".to_string()));
assert!(args.contains(&"--epochs=50".to_string()));
// GPU resources
@@ -583,34 +667,28 @@ mod tests {
assert_eq!(requests.get("cpu").unwrap().0, "4");
assert_eq!(requests.get("memory").unwrap().0, "16Gi");
let limits = res.limits.as_ref().unwrap();
assert_eq!(limits.get("nvidia.com/gpu").unwrap().0, "1");
assert_eq!(limits.get("cpu").unwrap().0, "8");
assert_eq!(limits.get("memory").unwrap().0, "32Gi");
// Volume mounts
// Volume mounts (training-data + output + binaries)
let mounts = main.volume_mounts.as_ref().unwrap();
assert_eq!(mounts.len(), 2);
let data_mount = mounts.iter().find(|m| m.name == "training-data").unwrap();
assert_eq!(data_mount.mount_path, "/data");
assert_eq!(data_mount.read_only, Some(true));
let output_mount = mounts.iter().find(|m| m.name == "output").unwrap();
assert_eq!(output_mount.mount_path, "/output");
assert_eq!(mounts.len(), 3);
assert!(mounts.iter().any(|m| m.name == "training-data" && m.mount_path == "/data"));
assert!(mounts.iter().any(|m| m.name == "output" && m.mount_path == "/output"));
assert!(mounts.iter().any(|m| m.name == "binaries" && m.mount_path == "/binaries"));
// -- Volumes --
// -- Volumes (PVC + 2 emptyDir) --
let volumes = pod_spec.volumes.as_ref().unwrap();
assert_eq!(volumes.len(), 2);
let data_vol = volumes
.iter()
.find(|v| v.name == "training-data")
.unwrap();
let data_pvc = data_vol.persistent_volume_claim.as_ref().unwrap();
assert_eq!(data_pvc.claim_name, "training-data-pvc");
assert_eq!(data_pvc.read_only, Some(true));
assert_eq!(volumes.len(), 3);
let data_vol = volumes.iter().find(|v| v.name == "training-data").unwrap();
assert_eq!(
data_vol.persistent_volume_claim.as_ref().unwrap().claim_name,
"training-data-pvc"
);
let output_vol = volumes.iter().find(|v| v.name == "output").unwrap();
let output_pvc = output_vol.persistent_volume_claim.as_ref().unwrap();
assert_eq!(output_pvc.claim_name, "training-output-pvc");
assert!(output_vol.empty_dir.is_some(), "output should be emptyDir");
let binaries_vol = volumes.iter().find(|v| v.name == "binaries").unwrap();
assert!(binaries_vol.empty_dir.is_some(), "binaries should be emptyDir");
}
#[test]

10
services/ml_training_service/src/service.rs
View File

@@ -266,12 +266,18 @@ impl MlTrainingService for MLTrainingServiceImpl {
// Generate a job ID
let job_id = Uuid::new_v4();
// Extract symbol from data_source file_path, or use default
// Extract symbol from data_source file_path (take last path component), or use default
let symbol = req
.data_source
.as_ref()
.and_then(|ds| match &ds.source {
Some(proto::data_source::Source::FilePath(p)) => Some(p.clone()),
Some(proto::data_source::Source::FilePath(p)) => {
// file_path may be "data/cache/futures-baseline/ES.FUT" — extract just the symbol
std::path::Path::new(p)
.file_name()
.and_then(|f| f.to_str())
.map(String::from)
}
_ => None,
})
.unwrap_or_else(|| "ES.FUT".to_string());