From fe3dd7d66bcc255bec998e0e12533a4e3e37cfc0 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Tue, 2 May 2023 21:31:39 -0700 Subject: [PATCH] add routing tags to inference jobs --- .../generic_inference_jobs.sql | 2 +- crates/lib/mysql_queries/sqlx-data.json | 20 +++++++++---------- .../web/insert_generic_inference_job.rs | 11 ++++++++++ .../src/job/job_loop/process_single_job.rs | 19 +++++++++++++++++- .../process_single_job_success_case.rs | 4 ++++ .../enqueue_infer_tts_handler.rs | 10 ++++++++++ .../enqueue_voice_conversion_inference.rs | 18 +++++++++++++---- 7 files changed, 68 insertions(+), 16 deletions(-) diff --git a/_sql/analytics/generic_inference/generic_inference_jobs.sql b/_sql/analytics/generic_inference/generic_inference_jobs.sql index 1e5b720d21..6d66149cfe 100644 --- a/_sql/analytics/generic_inference/generic_inference_jobs.sql +++ b/_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 diff --git a/crates/lib/mysql_queries/sqlx-data.json b/crates/lib/mysql_queries/sqlx-data.json index 0978d67f19..4b209cd867 100644 --- a/crates/lib/mysql_queries/sqlx-data.json +++ b/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": [], diff --git a/crates/lib/mysql_queries/src/queries/generic_inference/web/insert_generic_inference_job.rs b/crates/lib/mysql_queries/src/queries/generic_inference/web/insert_generic_inference_job.rs index 57003933b0..b0bce60312 100644 --- a/crates/lib/mysql_queries/src/queries/generic_inference/web/insert_generic_inference_job.rs +++ b/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) 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 135a5ba791..1a549fe006 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 @@ -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) => { diff --git a/crates/service/job/inference_job/src/job/job_loop/process_single_job_success_case.rs b/crates/service/job/inference_job/src/job/job_loop/process_single_job_success_case.rs index a2ebda3eb3..7dcb4de802 100644 --- a/crates/service/job/inference_job/src/job/job_loop/process_single_job_success_case.rs +++ b/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, diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/tts/enqueue_infer_tts_handler/enqueue_infer_tts_handler.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/tts/enqueue_infer_tts_handler/enqueue_infer_tts_handler.rs index 619ed54914..3f64d8da91 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/tts/enqueue_infer_tts_handler/enqueue_infer_tts_handler.rs +++ b/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; diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/voice_conversion/inference/enqueue_voice_conversion_inference.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/voice_conversion/inference/enqueue_voice_conversion_inference.rs index a510125e9d..e86c34f57f 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/voice_conversion/inference/enqueue_voice_conversion_inference.rs +++ b/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;