add routing tags to inference jobs

_sql/analytics/generic_inference/generic_inference_jobs.sql

@@ -45,7 +45,7 @@ from (
                  m.token = jobs.maybe_input_source_token
          where
              jobs.status != 'pending'
-        AND jobs.created_at > NOW() - INTERVAL 10 MINUTE
+        AND jobs.created_at > NOW() - INTERVAL 30 MINUTE
          order by id desc
              limit 5000
      ) as t

crates/lib/mysql_queries/sqlx-data.json

@@ -3927,6 +3927,16 @@
     },
     "query": "\nUPDATE users\nSET\n    profile_markdown = ?,\n    profile_rendered_html = ?,\n    discord_username = ?,\n    twitter_username = ?,\n    twitch_username = ?,\n    github_username = ?,\n    cashapp_username = ?,\n    website_url = ?,\n    version = version + 1\n\nWHERE users.token = ?\nLIMIT 1\n        "
   },
+  "3d01dbb4b3fa766a56d2f4d956ce0da7cd7e1c23c46ff1cbed578c9cd30c8e25": {
+    "describe": {
+      "columns": [],
+      "nullable": [],
+      "parameters": {
+        "Right": 15
+      }
+    },
+    "query": "\nINSERT INTO generic_inference_jobs\nSET\n  token = ?,\n  uuid_idempotency_token = ?,\n\n  inference_category = ?,\n  maybe_model_type = ?,\n  maybe_model_token = ?,\n\n  maybe_input_source_token = ?,\n  maybe_input_source_token_type = ?,\n\n  maybe_raw_inference_text = ?,\n\n  maybe_inference_args = ?,\n\n  maybe_creator_user_token = ?,\n  creator_ip_address = ?,\n  creator_set_visibility = ?,\n\n  priority_level = ?,\n  is_debug_request = ?,\n  maybe_routing_tag = ?,\n\n  status = \"pending\"\n        "
+  },
   "3e0387ef5d2569da801274c3ca80524e2b59b7eea577c1c96a28c186fb4fa0ea": {
     "describe": {
       "columns": [],
@@ -8419,16 +8429,6 @@
     },
     "query": "\nSELECT\n    jobs.token as `job_token: tokens::jobs::download::DownloadJobToken`,\n\n    jobs.status,\n    jobs.attempt_count,\n    jobs.on_success_downloaded_entity_token as maybe_downloaded_entity_token,\n    jobs.on_success_downloaded_entity_type as maybe_downloaded_entity_type,\n\n    jobs.created_at,\n    jobs.updated_at\n\nFROM generic_download_jobs as jobs\n\nWHERE jobs.token = ?\n        "
   },
-  "7208d62c30125be3ff649e0bdf71cc090f1f52195053503f2a1383a713229684": {
-    "describe": {
-      "columns": [],
-      "nullable": [],
-      "parameters": {
-        "Right": 14
-      }
-    },
-    "query": "\nINSERT INTO generic_inference_jobs\nSET\n  token = ?,\n  uuid_idempotency_token = ?,\n\n  inference_category = ?,\n  maybe_model_type = ?,\n  maybe_model_token = ?,\n\n  maybe_input_source_token = ?,\n  maybe_input_source_token_type = ?,\n\n  maybe_raw_inference_text = ?,\n\n  maybe_inference_args = ?,\n\n  maybe_creator_user_token = ?,\n  creator_ip_address = ?,\n  creator_set_visibility = ?,\n\n  priority_level = ?,\n  is_debug_request = ?,\n\n  status = \"pending\"\n        "
-  },
   "73558104936aab6add2d9c789b368b5f155f1fe01d32e7314f2f120422866f54": {
     "describe": {
       "columns": [],

crates/lib/mysql_queries/src/queries/generic_inference/web/insert_generic_inference_job.rs

@@ -29,6 +29,7 @@ pub struct InsertGenericInferenceArgs<'a> {
 
   pub priority_level: u8,
   pub is_debug_request: bool,
+  pub maybe_routing_tag: Option<&'a str>,
 
   pub mysql_pool: &'a MySqlPool,
 }
@@ -39,6 +40,14 @@ pub async fn insert_generic_inference_job(args: InsertGenericInferenceArgs<'_>)
   let serialized_args_payload = serde_json::ser::to_string(&args.maybe_inference_args)
       .map_err(|_e| anyhow!("could not encode inference args"))?;
 
+  // The routing tag column is VARCHAR(32), so we should truncate.
+  let maybe_routing_tag = args.maybe_routing_tag
+      .map(|routing_tag| {
+        let mut routing_tag = routing_tag.trim().to_string();
+        routing_tag.truncate(32);
+        routing_tag
+      });
+
   let query = sqlx::query!(
         r#"
 INSERT INTO generic_inference_jobs
@@ -63,6 +72,7 @@ SET
 
   priority_level = ?,
   is_debug_request = ?,
+  maybe_routing_tag = ?,
 
   status = "pending"
         "#,
@@ -86,6 +96,7 @@ SET
         args.creator_set_visibility.to_str(),
         args.priority_level,
         args.is_debug_request,
+        maybe_routing_tag,
     );
 
   let query_result = query.execute(args.mysql_pool)

crates/service/job/inference_job/src/job/job_loop/process_single_job.rs

@@ -24,11 +24,28 @@ pub async fn process_single_job(
   //let mut redis = job_dependencies.redis_pool.get()?;
   //let mut redis_logger = RedisJobStatusLogger::new_generic_download(&mut redis, job.download_job_token.as_str());
 
+  let mut force_execution = false;
+
+  // Some jobs have "routing tags". These ensure that jobs only execute on certain hosts.
+  // This is typically for debugging or development.
+  if let Some(routing_tag) = job.maybe_routing_tag.as_deref() {
+    let routing_tag = routing_tag.to_lowercase();
+    let hostname = job_dependencies.container.hostname.to_ascii_lowercase();
+
+    if hostname.starts_with(&routing_tag) {
+      info!("Job has routing tag ({}) for execution on this host ({})", routing_tag, hostname);
+      force_execution = true;
+    } else {
+      info!("Job routing tag ({}) doesn't match hostname ({}); skipping...", routing_tag, hostname);
+      return Ok(ProcessSingleJobSuccessCase::JobSkippedForRoutingTagMismatch);
+    }
+  }
+
   let dependency_status = determine_dependency_status(job_dependencies, job)
       .await
       .map_err(|err| ProcessSingleJobError::Other(anyhow!("database or cache error: {:?}", err)))?;
 
-  if !dependency_status.models_already_on_filesystem {
+  if !force_execution && !dependency_status.models_already_on_filesystem {
     match dependency_status.maybe_model_token {
       None => {} // No model token, proceed
       Some(model_token) => {

crates/service/job/inference_job/src/job/job_loop/process_single_job_success_case.rs

@@ -9,6 +9,10 @@ pub enum ProcessSingleJobSuccessCase {
   /// up the job, or we might return to it later ourselves.
   JobTemporarilySkippedFilesAbsent,
 
+  /// The job has a routing tag that doesn't match this container.
+  /// This job will never execute on this container.
+  JobSkippedForRoutingTagMismatch,
+
   /// The lock for the job could not be obtained. Another worker might have it.
   /// If that workload fails, we could conceivably pick it up again in the future.
   LockNotObtained,

crates/service/web/storyteller_web/src/http_server/endpoints/tts/enqueue_infer_tts_handler/enqueue_infer_tts_handler.rs

@@ -54,6 +54,11 @@ const DEBUG_HEADER_NAME : &'static str = "enable-debug-mode";
 /// It'll be used to help test TT2 on the new system before it is rolled out for everyone.
 const NEW_JOB_SYSTEM_HEADER_NAME : &'static str = "new-job-system";
 
+/// The routing tag header can send workloads to particular k8s hosts.
+/// This is useful for catching the live logs or intercepting the job.
+/// NB: This is only for the new job system.
+const ROUTING_TAG_HEADER_NAME : &'static str = "routing-tag";
+
 const USER_FAKEYOU_USER_TOKEN : &'static str = "U:N5J8JXPW9BTYX";
 const USER_NEWS_STORY_USER_TOKEN : &'static str = "U:XAWRARC1N89X6";
 
@@ -224,6 +229,10 @@ pub async fn enqueue_infer_tts_handler(
   let use_new_job_system_for_request  = get_request_header_optional(&http_request, NEW_JOB_SYSTEM_HEADER_NAME)
       .is_some();
 
+  let maybe_routing_tag=
+      get_request_header_optional(&http_request, ROUTING_TAG_HEADER_NAME)
+          .map(|routing_tag| routing_tag.trim().to_string());
+
   // ==================== RATE LIMIT ==================== //
 
   if !disable_rate_limiter {
@@ -384,6 +393,7 @@ pub async fn enqueue_infer_tts_handler(
       creator_set_visibility: set_visibility,
       priority_level,
       is_debug_request,
+      maybe_routing_tag: maybe_routing_tag.as_deref(),
       mysql_pool: &server_state.mysql_pool,
     }).await;
 

crates/service/web/storyteller_web/src/http_server/endpoints/voice_conversion/inference/enqueue_voice_conversion_inference.rs

@@ -31,7 +31,11 @@ use tts_common::priority::FAKEYOU_INVESTOR_PRIORITY_LEVEL;
 
 /// Debug requests can get routed to special "debug-only" workers, which can
 /// be used to trial new code, run debugging, etc.
-const DEBUG_HEADER_NAME : &'static str = "enable_debug_mode";
+const DEBUG_HEADER_NAME : &'static str = "enable-debug-mode";
+
+/// The routing tag header can send workloads to particular k8s hosts.
+/// This is useful for catching the live logs or intercepting the job.
+const ROUTING_TAG_HEADER_NAME : &'static str = "routing-tag";
 
 #[derive(Deserialize)]
 pub struct EnqueueVoiceConversionInferenceRequest {
@@ -158,10 +162,15 @@ pub async fn enqueue_voice_conversion_inference_handler(
     }
   }
 
-  // ==================== DEBUG MODE ==================== //
+  // ==================== DEBUG MODE + ROUTING TAG ==================== //
 
-  let is_debug_request = get_request_header_optional(&http_request, DEBUG_HEADER_NAME)
-      .is_some();
+  let is_debug_request =
+      get_request_header_optional(&http_request, DEBUG_HEADER_NAME)
+          .is_some();
+
+  let maybe_routing_tag=
+      get_request_header_optional(&http_request, ROUTING_TAG_HEADER_NAME)
+          .map(|routing_tag| routing_tag.trim().to_string());
 
   // ==================== RATE LIMIT ==================== //
 
@@ -245,6 +254,7 @@ pub async fn enqueue_voice_conversion_inference_handler(
     creator_set_visibility: set_visibility,
     priority_level,
     is_debug_request,
+    maybe_routing_tag: maybe_routing_tag.as_deref(),
     mysql_pool: &server_state.mysql_pool,
   }).await;