inference job query reaper (#735)

.sqlx/query-19708db99f2b3c6283d2115a511e0eb1c99fab396ad0effd9eb2dd79926e9c93.json

@@ -0,0 +1,12 @@
+{
+  "db_name": "MySQL",
+  "query": "\nUPDATE generic_inference_jobs\nSET\n  status = 'cancelled_by_system',\n  retry_at = NULL,\n  internal_debugging_failure_reason = 'reap_stale_fakeyou_jobs'\nWHERE\n  job_type != 'fal_queue'\n  AND status IN (\n    'pending',\n    'started',\n    'attempt_failed'\n   )\n   AND created_at > NOW() - INTERVAL 30 MINUTE\nLIMIT 5000\n        ",
+  "describe": {
+    "columns": [],
+    "parameters": {
+      "Right": 0
+    },
+    "nullable": []
+  },
+  "hash": "19708db99f2b3c6283d2115a511e0eb1c99fab396ad0effd9eb2dd79926e9c93"
+}

_sql/migrations_squashed/generic_inference_jobs.sql

@@ -21,6 +21,8 @@ CREATE TABLE generic_inference_jobs (
   -- The new enum for the type of job, which will eventually replace `inference_category` and `maybe_model_type`.
   -- This isn't fully supported yet in the inference-job, but we'll start populating it and adding logic around it.
   -- We'll start this out as nullable, then migrate the old rows with a default column value.
+  --
+  -- Fal jobs have the type `fal_queue`.
   job_type VARCHAR(32) DEFAULT NULL,
 
   -- There is an index on this column.

crates/schema/database/mysql_queries/src/payloads/generic_inference_args/generic_inference_args.rs

@@ -1,3 +1,4 @@
+use log::{debug, warn};
 use crate::payloads::generic_inference_args::inner_payloads::f5_tts_payload::F5TTSPayload;
 use crate::payloads::generic_inference_args::inner_payloads::face_fusion_payload::FaceFusionPayload;
 use crate::payloads::generic_inference_args::inner_payloads::gptsovits_payload::GptSovitsPayload;
@@ -180,6 +181,19 @@ impl GenericInferenceArgs {
     Ok(serde_json::from_str(json)?)
   }
 
+  pub fn from_json_optional(json: Option<&str>) -> AnyhowResult<Option<Self>> {
+    match json {
+      None => Ok(None),
+      Some("" | "null" | "NULL" | "undefined" | "UNDEFINED") => {
+        // NB: Something inserted a few of these records (only in local dev).
+        // It's better not to stall the job queue than to assert correctness.
+        warn!("Received empty or null JSON `{json:?}` for GenericInferenceArgs, returning None.");
+        Ok(None)
+      },
+      Some(json_str) => Ok(Some(Self::from_json(json_str)?)),
+    }
+  }
+
   pub fn to_json(&self) -> AnyhowResult<String> {
     Ok(serde_json::to_string(self)?)
   }

crates/schema/database/mysql_queries/src/queries/generic_inference/job/list_available_generic_inference_jobs.rs

@@ -2,6 +2,7 @@ use std::collections::BTreeSet;
 
 use anyhow::anyhow;
 use chrono::Utc;
+use log::debug;
 use sqlx::MySqlPool;
 
 use enums::by_table::generic_inference_jobs::inference_category::InferenceCategory;
@@ -151,10 +152,8 @@ pub async fn list_available_generic_inference_jobs(
               .map(|token_type| InferenceInputSourceTokenType::from_str(token_type)
                   .map_err(|err| anyhow!("error: {:?}", err)))
               .transpose()?,
-          maybe_inference_args: record.maybe_inference_args
-              .as_deref()
-              .map(|args| GenericInferenceArgs::from_json(args))
-              .transpose()?,
+          maybe_inference_args: GenericInferenceArgs::from_json_optional(
+              record.maybe_inference_args.as_deref())?,
           maybe_raw_inference_text: record.maybe_raw_inference_text,
           status: JobStatusPlus::from_str(&record.status)
               .map_err(|err| anyhow!("JobStatus failure to parse: {:?}", err))?,

crates/schema/database/mysql_queries/src/queries/generic_inference/job/mod.rs

@@ -5,3 +5,4 @@ pub mod mark_generic_inference_job_completely_failed;
 pub mod mark_generic_inference_job_failure;
 pub mod mark_generic_inference_job_pending_and_grab_lock;
 pub mod mark_generic_inference_job_successfully_done;
+pub mod reap_stale_fakeyou_jobs;

crates/schema/database/mysql_queries/src/queries/generic_inference/job/reap_stale_fakeyou_jobs.rs

@@ -0,0 +1,52 @@
+use anyhow::anyhow;
+use log::info;
+use sqlx;
+use sqlx::pool::PoolConnection;
+use sqlx::{MySql, MySqlPool};
+
+use errors::AnyhowResult;
+
+/// Kill FakeYou jobs that are "stale".
+///
+/// - This leaves Fal jobs alone, as they do not use the job runner.
+/// - This kills jobs that are over a threshold (e.g. an hour) overdue.
+pub async fn reap_stale_fakeyou_jobs(mysql_pool: &MySqlPool) -> AnyhowResult<()>
+{
+  let mut connection = mysql_pool.acquire().await?;
+  reap_stale_fakeyou_jobs_using_connection(&mut connection).await
+}
+
+pub async fn reap_stale_fakeyou_jobs_using_connection(
+  mysql_connection: &mut PoolConnection<MySql>,
+) -> AnyhowResult<()> {
+  let query_result = sqlx::query!(
+        r#"
+UPDATE generic_inference_jobs
+SET
+  status = 'cancelled_by_system',
+  retry_at = NULL,
+  internal_debugging_failure_reason = 'reap_stale_fakeyou_jobs'
+WHERE
+  job_type != 'fal_queue'
+  AND status IN (
+    'pending',
+    'started',
+    'attempt_failed'
+   )
+   AND created_at > NOW() - INTERVAL 30 MINUTE
+LIMIT 5000
+        "#
+    )
+      .execute(&mut **mysql_connection)
+      .await;
+
+  match query_result {
+    Ok(result) => {
+      info!("Reaped stale FakeYou jobs: {} rows affected.", result.rows_affected());
+      Ok(())
+    },
+    Err(err) => {
+      Err(anyhow!("error with reaping stale jobs: {:?}", err))
+    }
+  }
+}

crates/service/job/inference_job/src/http_server/run_http_server.rs

@@ -40,6 +40,8 @@ pub fn run_http_server(args: CreateServerArgs) -> AnyhowResult<Server>
   // NB: We shouldn't be logging much as the /_status endpoint is all we aim to expose.
   let log_format = "[%{HOSTNAME}e] IP=[%{X-Forwarded-For}i] \"%r\" %s %b \"%{Referer}i\" \"%{User-Agent}i\" %T";
 
+  info!("HTTP server will bind to: {}", bind_address);
+
   let handle = HttpServer::new(move || {
     // NB: app_data being clone()'d below should all be safe (dependencies included)
     App::new()

crates/service/job/inference_job/src/main.rs

@@ -53,6 +53,7 @@ use crate::state::job_dependencies::{BucketDependencies, ClientDependencies, Dat
 use crate::state::job_specific_dependencies::JobSpecificDependencies;
 use crate::state::scoped_job_type_execution::ScopedJobTypeExecution;
 use crate::state::scoped_model_type_execution::ScopedModelTypeExecution;
+use crate::threads::reap_jobs_thread::reap_jobs_thread;
 use crate::util::filesystem::scoped_temp_dir_creator::ScopedTempDirCreator;
 use crate::util::instrumentation::JobInstruments;
 use crate::util::instrumentation::{init_otel_metrics_pipeline, JobInstrumentLabels};
@@ -61,6 +62,7 @@ use crate::util::model_weights_cache::model_weights_cache_directory::ModelWeight
 pub mod http_server;
 pub mod job;
 pub mod state;
+pub mod threads;
 pub mod util;
 
 // Buckets (shared config)
@@ -163,6 +165,8 @@ async fn main() -> AnyhowResult<()> {
       .connect(&db_connection_string)
       .await?;
 
+  let mysql_pool_2 = mysql_pool.clone();
+
   info!("Connected to MySQL.");
 
   let common_env = CommonEnv::read_from_env()?;
@@ -345,6 +349,14 @@ async fn main() -> AnyhowResult<()> {
 
   set_up_directories(&job_dependencies)?;
 
+  std::thread::spawn(move || {
+    let actix_runtime = actix_web::rt::System::new();
+
+    actix_runtime.block_on(reap_jobs_thread(mysql_pool_2));
+
+    warn!("Reap jobs thread terminated");
+  });
+
   std::thread::spawn(move || {
     let actix_runtime = actix_web::rt::System::new();
     let http_server_handle = launch_http_server(create_server_args);

crates/service/job/inference_job/src/threads/mod.rs

@@ -0,0 +1 @@
+pub mod reap_jobs_thread;

crates/service/job/inference_job/src/threads/reap_jobs_thread.rs

@@ -0,0 +1,33 @@
+use errors::AnyhowResult;
+use log::{error, info};
+use mysql_queries::queries::generic_inference::job::reap_stale_fakeyou_jobs::reap_stale_fakeyou_jobs;
+use sqlx::MySqlPool;
+use std::time::Duration;
+
+const WAIT_BETWEEN_REAPING_SECONDS: u64 = 60 * 2; // 2 minutes
+
+pub async fn reap_jobs_thread(
+  mysql_pool: MySqlPool,
+) -> ! {
+  info!("Database job reaping thread starting...");
+
+  loop {
+    tokio::time::sleep(Duration::from_secs(WAIT_BETWEEN_REAPING_SECONDS)).await;
+
+    info!("Reaping jobs from database...");
+
+    let result = do_reap_jobs(&mysql_pool).await;
+
+    if let Err(err) = result {
+      error!("Error reaping jobs: {:?}", err);
+      tokio::time::sleep(Duration::from_secs(300)).await;
+    }
+  }
+}
+
+pub async fn do_reap_jobs(
+  mysql_pool: &MySqlPool,
+) -> AnyhowResult<()> {
+  reap_stale_fakeyou_jobs(mysql_pool).await?;
+  Ok(())
+}