feat(serving): wire report_job_completion to PromotionManager

Replace the stub report_job_completion gRPC handler with a real
implementation that:
- Validates job_id as UUID upfront
- On failure: updates DB status to Failed (best-effort), returns "failed"
- On success: updates DB status to Completed, looks up model_type and
  symbol from child_jobs table, registers with PromotionManager, and
  returns the actual promotion status string

Also adds JobSpawner::get_job_by_id() to fetch individual child jobs
from PostgreSQL, and two new tests covering promotion status mapping
and the PromotionManager integration path.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-02-27 22:56:39 +01:00
parent b4a5b8d235
commit 8324b98f72
2 changed files with 156 additions and 6 deletions

View File

@@ -503,6 +503,33 @@ impl JobSpawner {
Ok(())
}
/// Fetch a single child job by its ID
///
/// # Arguments
///
/// * `job_id` - Child job identifier
///
/// # Returns
///
/// * `Ok(Some(ChildJob))` - The child job
/// * `Ok(None)` - Job not found
/// * `Err` - Database error
pub async fn get_job_by_id(&self, job_id: Uuid) -> Result<Option<ChildJob>> {
let job = sqlx::query_as::<_, ChildJob>(
r#"
SELECT id, batch_id, model_type, status, created_at, config_json
FROM child_jobs
WHERE id = $1
"#,
)
.bind(job_id)
.fetch_optional(&self.db_pool)
.await
.context("Failed to fetch job by id")?;
Ok(job)
}
/// Get all child jobs for a batch
///
/// # Arguments

View File

@@ -1141,14 +1141,92 @@ impl MlTrainingService for MLTrainingServiceImpl {
"received job completion report"
);
// TODO(task-6): wire to promotion pipeline
// Parse the job UUID from the report
let job_uuid = uuid::Uuid::parse_str(&report.job_id).map_err(|e| {
Status::invalid_argument(format!("invalid job_id UUID: {}", e))
})?;
// ── Failed job path ──────────────────────────────────────────────
if !report.success {
// Best-effort DB status update; log warning if it fails
if let Some(ref spawner) = self.job_spawner {
if let Err(e) = spawner.update_job_status(job_uuid, "Failed").await {
warn!(job_id = %report.job_id, error = %e, "failed to update job status to Failed");
}
}
return Ok(Response::new(JobCompletionAck {
accepted: true,
promotion_status: "failed".to_string(),
}));
}
// ── Successful job path ──────────────────────────────────────────
// 1. Update DB status to Completed
if let Some(ref spawner) = self.job_spawner {
if let Err(e) = spawner.update_job_status(job_uuid, "Completed").await {
warn!(job_id = %report.job_id, error = %e, "failed to update job status to Completed");
}
}
// 2. Look up model_type and symbol from the child_jobs table
let (model_type, symbol) = match self.job_spawner {
Some(ref spawner) => {
match spawner.get_job_by_id(job_uuid).await {
Ok(Some(child_job)) => {
// Symbol is stored in config_json.asset
let sym = child_job
.config_json
.get("asset")
.and_then(|v| v.as_str())
.unwrap_or("UNKNOWN")
.to_string();
(child_job.model_type, sym)
}
Ok(None) => {
warn!(job_id = %report.job_id, "child job not found in DB, using defaults");
("UNKNOWN".to_string(), "UNKNOWN".to_string())
}
Err(e) => {
warn!(job_id = %report.job_id, error = %e, "failed to look up child job");
("UNKNOWN".to_string(), "UNKNOWN".to_string())
}
}
}
None => {
// No DB — cannot resolve model_type/symbol
debug!(job_id = %report.job_id, "no JobSpawner, cannot resolve model metadata");
("UNKNOWN".to_string(), "UNKNOWN".to_string())
}
};
// 3. Register with PromotionManager
let metrics: std::collections::HashMap<String, f64> = report.metrics.into_iter().collect();
let status = self
.promotion_manager
.register_completion(&report.job_id, &model_type, &symbol, &report.s3_path, metrics)
.await;
let promotion_status = match status {
crate::promotion_manager::PromotionStatus::PendingPromotion => "pending_promotion",
crate::promotion_manager::PromotionStatus::NoImprovement => "no_improvement",
crate::promotion_manager::PromotionStatus::Registered => "registered",
crate::promotion_manager::PromotionStatus::Error(ref msg) => {
warn!(job_id = %report.job_id, error = %msg, "promotion registration error");
"error"
}
};
info!(
job_id = %report.job_id,
model_type = %model_type,
symbol = %symbol,
promotion_status = %promotion_status,
"job completion processed"
);
Ok(Response::new(JobCompletionAck {
accepted: true,
promotion_status: if report.success {
"pending_promotion".to_string()
} else {
"failed".to_string()
},
promotion_status: promotion_status.to_string(),
}))
}
@@ -1634,4 +1712,49 @@ mod tests {
"train_baseline_supervised"
);
}
#[test]
fn test_promotion_status_string_mapping() {
// Verify the string representations match what gRPC clients expect
use crate::promotion_manager::PromotionStatus;
let cases: Vec<(PromotionStatus, &str)> = vec![
(PromotionStatus::PendingPromotion, "pending_promotion"),
(PromotionStatus::NoImprovement, "no_improvement"),
(PromotionStatus::Registered, "registered"),
(PromotionStatus::Error("test".to_string()), "error"),
];
for (status, expected) in cases {
let result = match status {
PromotionStatus::PendingPromotion => "pending_promotion",
PromotionStatus::NoImprovement => "no_improvement",
PromotionStatus::Registered => "registered",
PromotionStatus::Error(_) => "error",
};
assert_eq!(result, expected);
}
}
#[tokio::test]
async fn test_report_job_completion_registers_with_promotion_manager() {
// Integration test: verify that a successful completion report
// flows through to the PromotionManager and returns the correct status
use crate::promotion_manager::PromotionManager;
let pm = PromotionManager::new();
let metrics = HashMap::from([
("best_val_loss".to_string(), 0.05),
("sharpe_ratio".to_string(), 1.5),
]);
// First registration for a type+symbol returns Registered
let status = pm
.register_completion("job-1", "DQN", "ES.FUT", "s3://b/m.bin", metrics)
.await;
assert_eq!(
status,
crate::promotion_manager::PromotionStatus::Registered
);
}
}