diff --git a/crates/lib/mysql_queries/src/queries/generic_inference/job/list_available_generic_inference_jobs.rs b/crates/lib/mysql_queries/src/queries/generic_inference/job/list_available_generic_inference_jobs.rs index 741cd5b95d..c9ce668d99 100644 --- a/crates/lib/mysql_queries/src/queries/generic_inference/job/list_available_generic_inference_jobs.rs +++ b/crates/lib/mysql_queries/src/queries/generic_inference/job/list_available_generic_inference_jobs.rs @@ -65,6 +65,11 @@ pub struct AvailableInferenceJob { pub created_at: chrono::DateTime, pub updated_at: chrono::DateTime, pub retry_at: Option>, + + /// "Keep alive" signals may have a race condition where Redis doesn't set the timestamp + /// in time. If so, we'll compare `created_at` versus `now` with some reasonable delta + /// (perhaps 30 sec). + pub database_clock: chrono::DateTime, } pub struct ListAvailableGenericInferenceJobArgs<'a> { @@ -129,6 +134,7 @@ pub async fn list_available_generic_inference_jobs( created_at: record.created_at, updated_at: record.updated_at, retry_at: record.retry_at, + database_clock: record.database_clock, }; Ok(record) }) @@ -187,7 +193,9 @@ SELECT created_at, updated_at, - retry_at + retry_at, + NOW() as database_clock + FROM generic_inference_jobs"#.to_string(); query.push_str(&format!(r#" @@ -271,7 +279,9 @@ SELECT created_at, updated_at, - retry_at + retry_at, + NOW() as database_clock + FROM generic_inference_jobs"#.to_string(); @@ -358,6 +368,8 @@ struct AvailableInferenceJobRawInternal { pub created_at: chrono::DateTime, pub updated_at: chrono::DateTime, pub retry_at: Option>, + + pub database_clock: chrono::DateTime, } /// Return a comma-separated predicate, since SQLx does not yet support WHERE IN(?) for Vec, etc. diff --git a/crates/service/job/inference_job/config/inference-job.development.env b/crates/service/job/inference_job/config/inference-job.development.env index 1c039ab567..27aa2b7b77 100644 --- a/crates/service/job/inference_job/config/inference-job.development.env +++ b/crates/service/job/inference_job/config/inference-job.development.env @@ -3,6 +3,9 @@ # NB: For my crapy internet connection BUCKET_TIMEOUT_SECONDS=6000 +# NB: To not step on storyteller-web (the main API gateway) +HTTP_BIND_ADDRESS=0.0.0.0:11223 + # NB: Disable if not testing Redis REDIS_FOR_KEEPALIVE_URL="redis://localhost/0" diff --git a/crates/service/job/inference_job/src/job/job_loop/process_single_job.rs b/crates/service/job/inference_job/src/job/job_loop/process_single_job.rs index 163631ba7f..587412f3fb 100644 --- a/crates/service/job/inference_job/src/job/job_loop/process_single_job.rs +++ b/crates/service/job/inference_job/src/job/job_loop/process_single_job.rs @@ -110,8 +110,19 @@ pub async fn process_single_job( let _ : Option = match redis.get(&keepalive_key) { Ok(None) => { - warn!("Job keepalive elapsed: {:?}", job.inference_job_token); - return Err(ProcessSingleJobError::KeepAliveElapsed) + // NB: There's a chance that we're racing the keepalive function. + // As a second check, we'll compare the database clock versus the `created_at`. + // If the delta is small, we'll allow it. + let delta = job.database_clock.signed_duration_since(job.created_at); + let delta = delta.num_seconds(); + + if delta < 60 && delta > -60 { + warn!("could not get redis keepalive, but time was within delta seconds: {}", delta); + None // Allow it + } else { + warn!("Job keepalive elapsed: {:?}", job.inference_job_token); + return Err(ProcessSingleJobError::KeepAliveElapsed) + } }, Ok(Some(value)) => Some(value), Err(e) => {