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:
@@ -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"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user