From f3e485c2a1cc2d81c1a4e3df221148fa0b170430 Mon Sep 17 00:00:00 2001
From: jgrusewski <jgrusewski@localhost>
Date: Fri, 27 Feb 2026 22:37:46 +0100
Subject: [PATCH] feat(ml_training_service): wire start_training handler to
 JobSpawner
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Add JobSpawner as the highest-priority dispatch path in start_training.
When job_spawner is available, training requests are persisted to
PostgreSQL via spawn_batch() and a batch ID is returned immediately.
The queue consumer (Task 3) will later poll for pending jobs and
dispatch them to K8s.

Priority order: JobSpawner (DB) → K8sDispatcher (direct) → orchestrator (in-process).

Also adds test_start_training_model_binary_mapping unit test.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 services/ml_training_service/src/service.rs | 64 +++++++++++++++++----
 1 file changed, 53 insertions(+), 11 deletions(-)

diff --git a/services/ml_training_service/src/service.rs b/services/ml_training_service/src/service.rs
index 09b457132..2186cf0ba 100644
--- a/services/ml_training_service/src/service.rs
+++ b/services/ml_training_service/src/service.rs
@@ -266,18 +266,48 @@ impl MlTrainingService for MLTrainingServiceImpl {
         // Generate a job ID
         let job_id = Uuid::new_v4();
 
-        // Try K8s dispatch first, fall back to in-process orchestrator
-        if let Some(ref dispatcher) = self.k8s_dispatcher {
-            // Extract symbol from data_source file_path, 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()),
-                    _ => None,
-                })
-                .unwrap_or_else(|| "ES.FUT".to_string());
+        // Extract symbol from data_source file_path, 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()),
+                _ => None,
+            })
+            .unwrap_or_else(|| "ES.FUT".to_string());
 
+        // Priority 1: Batch flow — persist job to DB, queue consumer dispatches to K8s
+        if let Some(ref spawner) = self.job_spawner {
+            let model = common::model_types::ModelType::from_str(&req.model_type).ok_or_else(
+                || Status::invalid_argument(format!("Unknown model type: {}", req.model_type)),
+            )?;
+
+            let asset = crate::job_spawner::Asset {
+                symbol: symbol.clone(),
+                data_file: std::path::PathBuf::from(format!(
+                    "/data/futures-baseline/{}",
+                    symbol
+                )),
+            };
+
+            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 K8s/orchestrator: {}", e);
+                    // Fall through to K8s dispatcher or in-process orchestrator
+                }
+            }
+        }
+
+        // Priority 2: Direct K8s dispatch
+        if let Some(ref dispatcher) = self.k8s_dispatcher {
             let epochs = req
                 .hyperparameters
                 .as_ref()
@@ -1592,4 +1622,16 @@ mod tests {
     fn test_service_impl_has_job_spawner() {
         let _: fn(&MLTrainingServiceImpl) -> bool = |s| s.job_spawner.is_some();
     }
+
+    #[test]
+    fn test_start_training_model_binary_mapping() {
+        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");
+        assert_eq!(training_binary_for_model("ppo"), "train_baseline_rl");
+        assert_eq!(
+            training_binary_for_model("mamba2"),
+            "train_baseline_supervised"
+        );
+    }
 }