From 2f17465a844eea2496a3d509d41ca683b705f92a Mon Sep 17 00:00:00 2001
From: Brandon Thomas <echelon@gmail.com>
Date: Thu, 3 Mar 2022 02:00:32 -0500
Subject: [PATCH] job sorts by priority during, but has starvation prevention +
 env config

---
 .../list_available_tts_inference_jobs.rs      | 58 +++++++++++++++++--
 .../storyteller_web/src/tts_inference_job.rs  | 30 +++++++++-
 2 files changed, 82 insertions(+), 6 deletions(-)

diff --git a/crates/lib/database_queries/src/queries/tts/tts_inference_jobs/list_available_tts_inference_jobs.rs b/crates/lib/database_queries/src/queries/tts/tts_inference_jobs/list_available_tts_inference_jobs.rs
index 20a54dc768..c70f536f30 100644
--- a/crates/lib/database_queries/src/queries/tts/tts_inference_jobs/list_available_tts_inference_jobs.rs
+++ b/crates/lib/database_queries/src/queries/tts/tts_inference_jobs/list_available_tts_inference_jobs.rs
@@ -36,12 +36,63 @@ pub struct AvailableTtsInferenceJob {
 }
 
 /// Query jobs that are ready to run
+/// If sorting by priority, be careful not to starve lower priority jobs
+///  (ie. if there's an issue with higher priorities.)
 pub async fn list_available_tts_inference_jobs(
   pool: &MySqlPool,
+  sort_by_priority: bool,
   num_records: u32
 ) -> AnyhowResult<Vec<AvailableTtsInferenceJob>> {
 
-  let job_records : Vec<AvailableTtsInferenceJobRawInternal> = sqlx::query_as!(
+  // NB: This query is awkwardly written twice because this is the only way the
+  // macro can statically type check the query, result types, and parameter bindings.
+  //
+  // The only difference here is the 'ORDER BY' clause !
+  let job_records : Vec<AvailableTtsInferenceJobRawInternal> = if sort_by_priority {
+    sqlx::query_as!(
+      AvailableTtsInferenceJobRawInternal,
+        r#"
+SELECT
+  id,
+  token AS inference_job_token,
+  uuid_idempotency_token,
+
+  model_token,
+  raw_inference_text,
+
+  creator_ip_address,
+  maybe_creator_user_token,
+  creator_set_visibility as `creator_set_visibility: crate::column_types::record_visibility::RecordVisibility`,
+
+  is_from_api,
+  is_for_twitch,
+
+  status as `status: crate::column_types::job_status::JobStatus`,
+  priority_level,
+  attempt_count,
+  failure_reason,
+
+  created_at,
+  updated_at,
+  retry_at
+FROM tts_inference_jobs
+WHERE
+  (
+    status IN ("pending", "attempt_failed")
+  )
+  AND
+  (
+    retry_at IS NULL
+    OR
+    retry_at < CURRENT_TIMESTAMP
+  )
+  ORDER BY priority_level DESC, id ASC
+  LIMIT ?
+        "#,
+      num_records,
+    ).fetch_all(pool).await?
+  } else {
+    sqlx::query_as!(
       AvailableTtsInferenceJobRawInternal,
         r#"
 SELECT
@@ -81,9 +132,8 @@ WHERE
   LIMIT ?
         "#,
       num_records,
-    )
-      .fetch_all(pool)
-      .await?;
+    ).fetch_all(pool).await?
+  };
 
   let job_records = job_records.into_iter()
       .map(|record : AvailableTtsInferenceJobRawInternal| {
diff --git a/crates/service/storyteller_web/src/tts_inference_job.rs b/crates/service/storyteller_web/src/tts_inference_job.rs
index 06106344b9..3183782df9 100644
--- a/crates/service/storyteller_web/src/tts_inference_job.rs
+++ b/crates/service/storyteller_web/src/tts_inference_job.rs
@@ -141,6 +141,11 @@ struct Inferencer {
 
   // Maximum number of synthesizer models to hold in memory.
   pub sidecar_max_synthesizer_models: usize,
+
+  // Typically we'll sort jobs by priority. Occasionally we introduce a chance for low
+  // priority jobs to run in the order they were enqueued.
+  // If this is set to "0", we no longer consider priority
+  pub low_priority_starvation_prevention_every_nth: usize,
 }
 
 #[tokio::main]
@@ -252,6 +257,10 @@ async fn main() -> AnyhowResult<()> {
   let sidecar_max_synthesizer_models = easyenv::get_env_num(
     "SIDECAR_MAX_SYNTHESIZER_MODELS", 3)?;
 
+  // Set to "0" to always treat low priority the same as high priority
+  let low_priority_starvation_prevention_every_nth= easyenv::get_env_num(
+    "LOW_PRIORITY_STARVATION_PREVENTION_EVERY_NTH", 3)?;
+
   let firehose_publisher = FirehosePublisher {
     mysql_pool: mysql_pool.clone(), // NB: MySqlPool is clone/send/sync safe
   };
@@ -308,6 +317,7 @@ async fn main() -> AnyhowResult<()> {
     job_batch_size: common_env.job_batch_size,
     no_op_logger_millis: common_env.no_op_logger_millis,
     sidecar_max_synthesizer_models,
+    low_priority_starvation_prevention_every_nth,
   };
 
   main_loop(inferencer).await;
@@ -325,11 +335,27 @@ async fn main_loop(inferencer: Inferencer) {
 
   let mut span_batch = Vec::new();
 
+  let mut sort_by_priority = true;
+  let mut sort_by_priority_count = 0;
+
   loop {
     let num_records = inferencer.job_batch_size;
 
-    let maybe_available_jobs
-        = list_available_tts_inference_jobs(&inferencer.mysql_pool, num_records).await;
+    // Don't completely starve low-priority jobs
+    if sort_by_priority_count >= inferencer.low_priority_starvation_prevention_every_nth {
+      sort_by_priority_count = 0;
+      sort_by_priority = false;
+    }
+
+    let maybe_available_jobs =
+        list_available_tts_inference_jobs(
+          &inferencer.mysql_pool,
+          sort_by_priority,
+          num_records
+        ).await;
+
+    sort_by_priority = true;
+    sort_by_priority_count += 1;
 
     let jobs = match maybe_available_jobs {
       Ok(jobs) => jobs,