feat(ml_training_service): wire start_training handler to JobSpawner

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>
This commit is contained in:
jgrusewski
2026-02-27 22:37:46 +01:00
parent 3817b06f19
commit f3e485c2a1

View File

@@ -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"
);
}
}