diff --git a/src/job_queries/w2l_inference_job_queries.rs b/src/job_queries/w2l_inference_job_queries.rs index 21b488765d..35b9350010 100644 --- a/src/job_queries/w2l_inference_job_queries.rs +++ b/src/job_queries/w2l_inference_job_queries.rs @@ -129,11 +129,12 @@ WHERE id = ? Ok(()) } -pub async fn mark_w2l_inference_job_done(pool: &MySqlPool, - job: &W2lInferenceJobRecord, - success: bool) - -> AnyhowResult<()> -{ +pub async fn mark_w2l_inference_job_done( + pool: &MySqlPool, + job: &W2lInferenceJobRecord, + success: bool, + maybe_result_token: Option<&str> +) -> AnyhowResult<()> { let status = if success { "complete_success" } else { "complete_failure" }; let query_result = sqlx::query!( @@ -141,11 +142,13 @@ pub async fn mark_w2l_inference_job_done(pool: &MySqlPool, UPDATE w2l_inference_jobs SET status = ?, + on_success_result_token = ?, failure_reason = NULL, retry_at = NULL WHERE id = ? "#, status, + maybe_result_token, job.id, ) .execute(pool) diff --git a/src/tts_inference_job.rs b/src/tts_inference_job.rs index 9a37d0745a..d4b27cca80 100644 --- a/src/tts_inference_job.rs +++ b/src/tts_inference_job.rs @@ -444,7 +444,12 @@ async fn process_job(inferencer: &Inferencer, job: &TtsInferenceJobRecord) -> An .await?; info!("Marking job complete..."); - mark_tts_inference_job_done(&inferencer.mysql_pool, job, true, Some(&inference_result_token)).await?; + mark_tts_inference_job_done( + &inferencer.mysql_pool, + job, + true, + Some(&inference_result_token) + ).await?; inferencer.firehose_publisher.tts_inference_finished( job.maybe_creator_user_token.as_deref(), diff --git a/src/w2l_inference_job.rs b/src/w2l_inference_job.rs index 36bb4b2b60..425a263335 100644 --- a/src/w2l_inference_job.rs +++ b/src/w2l_inference_job.rs @@ -531,7 +531,12 @@ async fn process_job(inferencer: &Inferencer, job: &W2lInferenceJobRecord) -> An .await?; info!("Marking job complete..."); - mark_w2l_inference_job_done(&inferencer.mysql_pool, job, true).await?; + mark_w2l_inference_job_done( + &inferencer.mysql_pool, + job, + true, + Some(&inference_result_token) + ).await?; inferencer.firehose_publisher.w2l_inference_finished( job.maybe_creator_user_token.as_deref(),