diff --git a/services/ml_training_service/src/k8s_dispatcher.rs b/services/ml_training_service/src/k8s_dispatcher.rs index 4cf58cb80..efe7b3513 100644 --- a/services/ml_training_service/src/k8s_dispatcher.rs +++ b/services/ml_training_service/src/k8s_dispatcher.rs @@ -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", ¶ms.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] diff --git a/services/ml_training_service/src/service.rs b/services/ml_training_service/src/service.rs index 965268f71..f446a086a 100644 --- a/services/ml_training_service/src/service.rs +++ b/services/ml_training_service/src/service.rs @@ -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());