diff --git a/Cargo.lock b/Cargo.lock index 5d0e3dd8ab..ad01d3f8d2 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2541,6 +2541,8 @@ dependencies = [ "database_queries", "dotenv", "easyenv", + "enums", + "errors", "hostname", "hyper 0.14.14", "jobs_common", @@ -2555,6 +2557,7 @@ dependencies = [ "sqlx", "subprocess", "tempdir", + "tokens", "tokio 1.16.1", "tts_common", ] diff --git a/Dockerfile b/Dockerfile index ee90b49cae..709e9f751f 100644 --- a/Dockerfile +++ b/Dockerfile @@ -9,7 +9,7 @@ FROM ubuntu:jammy as rust-base # NB: This can be "stable" or another version. -ARG RUST_TOOLCHAIN="1.65.0" +ARG RUST_TOOLCHAIN="1.66.0" WORKDIR /tmp diff --git a/crates/_types/enums/src/workers/generic_inference_type.rs b/crates/_types/enums/src/workers/generic_inference_type.rs index e0d650849f..13920ad7d1 100644 --- a/crates/_types/enums/src/workers/generic_inference_type.rs +++ b/crates/_types/enums/src/workers/generic_inference_type.rs @@ -1,12 +1,14 @@ // TODO: Use macro-derived impls +use std::collections::BTreeSet; + /// Our "generic inference" pipeline supports a wide variety of ML models and other media. /// Each type of inference is identified by the following enum variants. /// These types are present in the HTTP API and database columns as serialized here. /// /// DO NOT CHANGE VALUES WITHOUT A MIGRATION STRATEGY. -#[derive(Clone, Copy, Eq, PartialEq, Debug, Deserialize, Serialize, sqlx::Type)] +#[derive(Clone, Copy, Eq, PartialEq, Debug, Hash, Ord, PartialOrd, Deserialize, Serialize, sqlx::Type)] pub enum GenericInferenceType { #[serde(rename = "text_to_speech")] #[sqlx(rename = "text_to_speech")] @@ -33,6 +35,15 @@ impl GenericInferenceType { _ => Err(format!("invalid value: {:?}", value)), } } + + pub fn all_variants() -> BTreeSet { + // NB: BTreeSet is sorted + // NB: BTreeSet::from() isn't const, but not worth using LazyStatic, etc. + BTreeSet::from([ + GenericInferenceType::TextToSpeech, + GenericInferenceType::VoiceConversion, + ]) + } } #[cfg(test)] @@ -57,4 +68,13 @@ mod tests { assert_eq!(GenericInferenceType::from_str("text_to_speech").unwrap(), GenericInferenceType::TextToSpeech); assert_eq!(GenericInferenceType::from_str("voice_conversion").unwrap(), GenericInferenceType::VoiceConversion); } + + #[test] + fn all_variants() { + let mut variants = GenericInferenceType::all_variants(); + assert_eq!(variants.len(), 2); + assert_eq!(variants.pop_first(), Some(GenericInferenceType::TextToSpeech)); + assert_eq!(variants.pop_first(), Some(GenericInferenceType::VoiceConversion)); + assert_eq!(variants.pop_first(), None); + } } diff --git a/crates/lib/database_queries/src/queries/generic_inference/job/list_available_generic_inference_jobs.rs b/crates/lib/database_queries/src/queries/generic_inference/job/list_available_generic_inference_jobs.rs index ef94050606..3eddb4876c 100644 --- a/crates/lib/database_queries/src/queries/generic_inference/job/list_available_generic_inference_jobs.rs +++ b/crates/lib/database_queries/src/queries/generic_inference/job/list_available_generic_inference_jobs.rs @@ -6,7 +6,10 @@ use crate::helpers::boolean_converters::i8_to_bool; use crate::queries::generic_inference::job::_keys::GenericInferenceJobId; use enums::core::visibility::Visibility; use enums::workers::generic_inference_type::GenericInferenceType; -use sqlx::MySqlPool; +use sqlx::mysql::MySqlArguments; +use sqlx::{MySql, MySqlPool}; +use std::collections::BTreeSet; +use std::future::Future; use std::path::Path; use tokens::jobs::inference::InferenceJobToken; @@ -41,56 +44,33 @@ pub struct AvailableInferenceJob { pub retry_at: Option>, } -pub async fn list_available_generic_inference_jobs(pool: &MySqlPool, num_records: u32) +pub struct ListAvailableGenericInferenceJobArgs<'a> { + pub num_records: u32, + pub is_debug_worker: bool, + pub sort_by_priority: bool, + pub maybe_scope_by_job_type: Option>, + pub mysql_pool: &'a MySqlPool, +} + +pub async fn list_available_generic_inference_jobs( + args: ListAvailableGenericInferenceJobArgs<'_>, +) -> AnyhowResult> { - let job_records = sqlx::query_as!( - AvailableInferenceJobRawInternal, - r#" -SELECT - id as `id: crate::queries::generic_inference::job::_keys::GenericInferenceJobId`, - token AS `inference_job_token: tokens::jobs::inference::InferenceJobToken`, + let ALL_TYPES = GenericInferenceType::all_variants(); - inference_type as `inference_type: enums::workers::generic_inference_type::GenericInferenceType`, - maybe_inference_args, - maybe_raw_inference_text, - maybe_model_token, + let inference_types = args.maybe_scope_by_job_type + .as_ref() + .map(|types| types.clone()) + .unwrap_or(ALL_TYPES); - maybe_creator_user_token, - creator_ip_address, - creator_set_visibility as `creator_set_visibility: enums::core::visibility::Visibility`, + let query = if args.sort_by_priority { + list_sorted_by_priority(args, inference_types).await + } else { + list_sorted_by_id(args, inference_types).await + }; - status as `status: crate::column_types::job_status::JobStatus`, - - - attempt_count, - priority_level, - is_from_premium_user, - is_from_api_user, - is_for_twitch, - is_debug_request, - - created_at, - updated_at, - retry_at -FROM generic_inference_jobs -WHERE - ( - status IN ("pending", "attempt_failed") - ) - AND - ( - retry_at IS NULL - OR - retry_at < CURRENT_TIMESTAMP - ) - ORDER BY id ASC - LIMIT ? - "#, - num_records, - ) - .fetch_all(pool) - .await?; + let job_records = query?; let job_records = job_records.into_iter() .map(|record : AvailableInferenceJobRawInternal| { @@ -121,6 +101,124 @@ WHERE Ok(job_records) } +async fn list_sorted_by_id(args: ListAvailableGenericInferenceJobArgs<'_>, inference_types: BTreeSet) -> Result, sqlx::Error> { + sqlx::query_as!( + AvailableInferenceJobRawInternal, + r#" +SELECT + id as `id: crate::queries::generic_inference::job::_keys::GenericInferenceJobId`, + token AS `inference_job_token: tokens::jobs::inference::InferenceJobToken`, + + inference_type as `inference_type: enums::workers::generic_inference_type::GenericInferenceType`, + maybe_inference_args, + maybe_raw_inference_text, + maybe_model_token, + + maybe_creator_user_token, + creator_ip_address, + creator_set_visibility as `creator_set_visibility: enums::core::visibility::Visibility`, + + status as `status: crate::column_types::job_status::JobStatus`, + + attempt_count, + priority_level, + is_from_premium_user, + is_from_api_user, + is_for_twitch, + is_debug_request, + + created_at, + updated_at, + retry_at +FROM generic_inference_jobs +WHERE + ( + inference_type IN (?) + ) + AND + ( + status IN ("pending", "attempt_failed") + ) + AND + ( + retry_at IS NULL + OR + retry_at < CURRENT_TIMESTAMP + ) + AND + ( + is_debug_request = ? + ) + ORDER BY id ASC + LIMIT ? + "#, + to_where_in_predicate(&inference_types), + args.is_debug_worker, + args.num_records, + ) + .fetch_all(args.mysql_pool) + .await +} + +async fn list_sorted_by_priority(args: ListAvailableGenericInferenceJobArgs<'_>, inference_types: BTreeSet) -> Result, sqlx::Error> { + sqlx::query_as!( + AvailableInferenceJobRawInternal, + r#" +SELECT + id as `id: crate::queries::generic_inference::job::_keys::GenericInferenceJobId`, + token AS `inference_job_token: tokens::jobs::inference::InferenceJobToken`, + + inference_type as `inference_type: enums::workers::generic_inference_type::GenericInferenceType`, + maybe_inference_args, + maybe_raw_inference_text, + maybe_model_token, + + maybe_creator_user_token, + creator_ip_address, + creator_set_visibility as `creator_set_visibility: enums::core::visibility::Visibility`, + + status as `status: crate::column_types::job_status::JobStatus`, + + attempt_count, + priority_level, + is_from_premium_user, + is_from_api_user, + is_for_twitch, + is_debug_request, + + created_at, + updated_at, + retry_at +FROM generic_inference_jobs +WHERE + ( + inference_type IN (?) + ) + AND + ( + status IN ("pending", "attempt_failed") + ) + AND + ( + retry_at IS NULL + OR + retry_at < CURRENT_TIMESTAMP + ) + AND + ( + is_debug_request = ? + ) + ORDER BY priority_level DESC, id ASC + LIMIT ? + "#, + to_where_in_predicate(&inference_types), + args.is_debug_worker, + args.num_records, + ) + .fetch_all(args.mysql_pool) + .await +} + #[derive(Debug)] struct AvailableInferenceJobRawInternal { pub id: GenericInferenceJobId, @@ -150,3 +248,40 @@ struct AvailableInferenceJobRawInternal { pub updated_at: chrono::DateTime, pub retry_at: Option>, } + +/// Return a comma-separated predicate, since SQLx does not yet support WHERE IN(?) for Vec, etc. +/// Issue: https://github.com/launchbadge/sqlx/issues/875 +fn to_where_in_predicate(types: &BTreeSet) -> String { + let mut vec = types.iter() + .map(|ty| ty.to_str()) // TODO: These strings might need to be quoted. + .collect::>(); + vec.sort(); // NB: For the benefit of tests. + vec.join(", ") +} + +#[cfg(test)] +mod tests { + use crate::queries::generic_inference::job::list_available_generic_inference_jobs::to_where_in_predicate; + use enums::workers::generic_inference_type::GenericInferenceType; + use std::collections::BTreeSet; + + #[test] + fn test_where_in_clause() { + // None + let types = BTreeSet::from([]); + assert_eq!(to_where_in_predicate(&types), "".to_string()); + + // Some + let types = BTreeSet::from([ + GenericInferenceType::VoiceConversion, + ]); + + assert_eq!(to_where_in_predicate(&types), "voice_conversion".to_string()); + // All + let types = BTreeSet::from([ + GenericInferenceType::TextToSpeech, + GenericInferenceType::VoiceConversion, + ]); + assert_eq!(to_where_in_predicate(&types), "text_to_speech, voice_conversion".to_string()); + } +} diff --git a/crates/lib/database_queries/src/queries/generic_inference/job/mod.rs b/crates/lib/database_queries/src/queries/generic_inference/job/mod.rs index 2a4c4e14a6..94cff4d755 100644 --- a/crates/lib/database_queries/src/queries/generic_inference/job/mod.rs +++ b/crates/lib/database_queries/src/queries/generic_inference/job/mod.rs @@ -1,6 +1,6 @@ pub mod _keys; -pub mod list_available_generic_inference_jobs; pub mod mark_generic_inference_job_completely_failed; -pub mod mark_generic_inference_job_successfully_done; 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 list_available_generic_inference_jobs; diff --git a/crates/service/job/download_job/src/job_steps/main_loop.rs b/crates/service/job/download_job/src/job_steps/main_loop.rs index 36d29753d8..e70395224e 100644 --- a/crates/service/job/download_job/src/job_steps/main_loop.rs +++ b/crates/service/job/download_job/src/job_steps/main_loop.rs @@ -7,6 +7,7 @@ use jobs_common::noop_logger::NoOpLogger; use log::warn; use std::time::Duration; +// Job runner timeouts (guards MySQL) const START_TIMEOUT_MILLIS : u64 = 500; const INCREASE_TIMEOUT_MILLIS : u64 = 1000; @@ -41,7 +42,7 @@ pub async fn main_loop(job_state: JobState) { match result { Ok(_) => {}, Err(e) => { - warn!("Error querying jobs: {:?}", e); + warn!("Error processing jobs: {:?}", e); std::thread::sleep(Duration::from_millis(error_timeout_millis)); error_timeout_millis += INCREASE_TIMEOUT_MILLIS; continue; diff --git a/crates/service/job/inference_job/Cargo.toml b/crates/service/job/inference_job/Cargo.toml index ad9a8c3559..3c54b2fb33 100644 --- a/crates/service/job/inference_job/Cargo.toml +++ b/crates/service/job/inference_job/Cargo.toml @@ -19,8 +19,11 @@ cloud_storage = { path = "../../../lib/cloud_storage" } config = { path = "../../../lib/config" } container_common = { path = "../../../lib/container_common" } database_queries = { path = "../../../lib/database_queries" } +enums = { path = "../../../_types/enums" } +errors = { path = "../../../lib/errors" } jobs_common = { path = "../../../lib/jobs_common" } memory_caching = { path = "../../../lib/memory_caching" } +tokens = { path = "../../../_types/tokens" } tts_common = { path = "../../../lib/tts_common" } # Vendored code diff --git a/crates/service/job/inference_job/src/job_steps/main_loop.rs b/crates/service/job/inference_job/src/job_steps/main_loop.rs index c164257073..4389049778 100644 --- a/crates/service/job/inference_job/src/job_steps/main_loop.rs +++ b/crates/service/job/inference_job/src/job_steps/main_loop.rs @@ -1,67 +1,38 @@ -use container_common::anyhow_result::AnyhowResult; -use container_common::collections::multiple_random_from_vec::multiple_random_from_vec; -use crate::caching::cache_miss_strategizer::CacheMissStrategy; -use crate::caching::virtual_lfu_cache::SyncVirtualLfuCache; use crate::job_steps::job_dependencies::JobDependencies; -use crate::job_steps::job_stats::JobStats; -use crate::job_steps::process_single_job::process_single_job; -use crate::job_steps::process_single_job_error::ProcessSingleJobError; -use crate::job_steps::sidecar_health_check_trap::maybe_block_on_sidecar_health_check; -use database_queries::queries::tts::tts_inference_jobs::list_available_tts_inference_jobs::{AvailableTtsInferenceJob, list_available_tts_inference_jobs, list_available_tts_inference_jobs_with_minimum_priority}; -use database_queries::queries::tts::tts_inference_jobs::mark_tts_inference_job_failure::mark_tts_inference_job_failure; -use database_queries::queries::tts::tts_inference_jobs::mark_tts_inference_job_permanently_dead::mark_tts_inference_job_permanently_dead; -use database_queries::queries::tts::tts_models::get_tts_model_for_inference::{get_tts_model_for_inference, TtsModelForInferenceError, TtsModelForInferenceRecord}; +use crate::job_steps::process_single_job_old::process_single_job_old; +use database_queries::queries::generic_inference::job::list_available_generic_inference_jobs::{AvailableInferenceJob, list_available_generic_inference_jobs, ListAvailableGenericInferenceJobArgs}; +use database_queries::queries::generic_inference::job::mark_generic_inference_job_failure::mark_generic_inference_job_failure; +use errors::AnyhowResult; use jobs_common::noop_logger::NoOpLogger; -use jobs_common::semi_persistent_cache_dir::SemiPersistentCacheDir; use log::{error, info, warn}; -use memory_caching::multi_item_ttl_cache::MultiItemTtlCache; -use newrelic_telemetry::Span; -use sqlx::MySqlPool; use std::time::Duration; +use crate::job_steps::process_single_job::process_single_job; // Job runner timeouts (guards MySQL) const START_TIMEOUT_MILLIS : u64 = 500; const INCREASE_TIMEOUT_MILLIS : u64 = 1000; -pub async fn main_loop(job_args: JobDependencies) { +pub async fn main_loop(job_dependencies: JobDependencies) { + let mut noop_logger = NoOpLogger::new(job_dependencies.no_op_logger_millis as i64); + let mut error_timeout_millis = START_TIMEOUT_MILLIS; - - let mut noop_logger = NoOpLogger::new(job_args.no_op_logger_millis as i64); - - let mut span_batch = Vec::new(); - let mut sort_by_priority = true; let mut sort_by_priority_count = 0; - let mut needs_health_check_at_start = true; // Run health check at startup. - loop { - let num_records = job_args.job_batch_size; - // Don't completely starve low-priority jobs - if sort_by_priority_count >= job_args.low_priority_starvation_prevention_every_nth { + if sort_by_priority_count >= job_dependencies.low_priority_starvation_prevention_every_nth { sort_by_priority_count = 0; sort_by_priority = false; } - let maybe_available_jobs = - if let Some(minimum_priority) = job_args.maybe_minimum_priority { - // Special high-priority workers - list_available_tts_inference_jobs_with_minimum_priority( - &job_args.mysql_pool, - minimum_priority, - num_records, - job_args.worker_details.is_debug_worker - ).await - } else { - // Normal path - list_available_tts_inference_jobs( - &job_args.mysql_pool, - sort_by_priority, - num_records, - job_args.worker_details.is_debug_worker - ).await - }; + let maybe_available_jobs = list_available_generic_inference_jobs(ListAvailableGenericInferenceJobArgs { + num_records: job_dependencies.job_batch_size, + is_debug_worker: false, // TODO + sort_by_priority, + maybe_scope_by_job_type: None, // TODO + mysql_pool: &job_dependencies.mysql_pool, + }).await; sort_by_priority = true; sort_by_priority_count += 1; @@ -69,7 +40,7 @@ pub async fn main_loop(job_args: JobDependencies) { let jobs = match maybe_available_jobs { Ok(jobs) => jobs, Err(e) => { - warn!("Error querying jobs: {:?}", e); + error!("Error querying jobs: {:?}", e); std::thread::sleep(Duration::from_millis(error_timeout_millis)); error_timeout_millis += INCREASE_TIMEOUT_MILLIS; continue; @@ -77,346 +48,52 @@ pub async fn main_loop(job_args: JobDependencies) { }; if jobs.is_empty() { - noop_logger.log_message_after_awhile("No TTS jobs picked up from database!"); - - std::thread::sleep(Duration::from_millis(job_args.job_batch_wait_millis)); + noop_logger.log_message_after_awhile("No jobs picked up from database!"); + std::thread::sleep(Duration::from_millis(job_dependencies.job_batch_wait_millis)); + error_timeout_millis = START_TIMEOUT_MILLIS; // reset continue; } info!("Queried {} jobs from database", jobs.len()); - let batch_result = process_jobs( - &job_args, - jobs, - needs_health_check_at_start, - ).await; + let batch_result = process_job_batch(&job_dependencies, jobs).await; - if needs_health_check_at_start { - needs_health_check_at_start = false; - } - - let mut spans = match batch_result { - Ok(spans) => spans, + match batch_result { + Ok(_) => {}, Err(e) => { - warn!("Error running job batch: {:?}", e); + error!("Error processing jobs: {:?}", e); std::thread::sleep(Duration::from_millis(error_timeout_millis)); error_timeout_millis += INCREASE_TIMEOUT_MILLIS; continue; } - }; - - if !job_args.newrelic_disabled { - span_batch.append(&mut spans); - - if span_batch.len() > 50 { - let spans_to_send = span_batch.split_off(0).into(); - job_args.newrelic_client.send_spans(spans_to_send).await; - } } error_timeout_millis = START_TIMEOUT_MILLIS; // reset - std::thread::sleep(Duration::from_millis(job_args.job_batch_wait_millis)); + std::thread::sleep(Duration::from_millis(job_dependencies.job_batch_wait_millis)); } } -/// Process a batch of jobs, returning the count of cold-cache skipped jobs. -async fn process_jobs( - inferencer: &JobDependencies, - jobs: Vec, - needs_health_check_at_start: bool, -) -> AnyhowResult> { - - if needs_health_check_at_start { - maybe_block_on_sidecar_health_check(&inferencer.http_clients.tts_sidecar_health_check_client).await; - } - - let mut maybe_sidecar_health_issue = false; - - let mut batch_spans = Vec::new(); +// TODO: A common interface/trait for each submodule (tts, webvc) to declare how to determine if the job is "ready". +// This probably returns a struct or enum with some measure of how many GB need to be downloaded. +async fn process_job_batch(job_state: &JobState, jobs: Vec) -> AnyhowResult<()> { for job in jobs.into_iter() { - let model_state_result = ModelState::query_model_and_check_filesystem( - &job, - &inferencer.mysql_pool, - &inferencer.caches.tts_model_record_cache, - &inferencer.semi_persistent_cache, - &inferencer.virtual_model_lfu, - ).await; - - let model_state = match model_state_result { - Ok(model_state) => model_state, - Err(e) => { - error!("TTS model fetch and state check error: {}, reason: {:?}", &job.model_token, &e); - - let (failure_reason, permanent_failure) = match e { - ModelStateError::ModelNotFound => ("model was not found", true), - ModelStateError::ModelDeleted => ("model has been deleted", true), - ModelStateError::CacheError { .. } => ("internal cache error", false), - ModelStateError::DatabaseError { .. } => ("unknown database error", false), - }; - - let internal_debugging_failure_reason = format!("model error: {:?}", e); - - let mut job_progress_reporter = inferencer - .job_progress_reporter - .new_tts_inference(&job.inference_job_token)?; - - job_progress_reporter.log_status(failure_reason)?; - - if permanent_failure { - warn!("Marking job permanently dead: {} because: {:?}", job.inference_job_token, &e); - - let _r = mark_tts_inference_job_permanently_dead( - &inferencer.mysql_pool, - job.id, - failure_reason, - &internal_debugging_failure_reason, - &inferencer.get_worker_name(), - ).await; - } else { - let _r = mark_tts_inference_job_failure( - &inferencer.mysql_pool, - &job, - failure_reason, - &internal_debugging_failure_reason, - inferencer.job_max_attempts, - &inferencer.get_worker_name(), - ).await; - } - - continue; - } - }; - - if !model_state.is_downloaded_to_filesystem || !model_state.is_in_memory_cache { - warn!("Model isn't ready: {} (downloaded = {}), (in memory = {})", - &model_state.model_record.model_token, - model_state.is_downloaded_to_filesystem, - model_state.is_in_memory_cache); - - let maybe_strategy = if !model_state.is_downloaded_to_filesystem { - inferencer.cache_miss_strategizers.disk_cache_miss(&model_state.model_record.model_token) - } else { - inferencer.cache_miss_strategizers.memory_cache_miss(&model_state.model_record.model_token) - }; - - match maybe_strategy { - Err(err) => { - warn!("Unable to process job: {:?}", err); - - let failure_reason = "cache error"; - let internal_debugging_failure_reason = format!("cache error: {:?}", err); - - let _r = mark_tts_inference_job_failure( - &inferencer.mysql_pool, - &job, - failure_reason, - &internal_debugging_failure_reason, - inferencer.job_max_attempts, - &inferencer.get_worker_name(), - ).await; - continue; - }, - Ok(CacheMissStrategy::WaitOrSkip) => { - // We're going to skip this for now. - // Maybe another worker has a warm cache and can continue. - warn!("Skipping TTS due to cold cache: {} ({})", - model_state.model_record.model_token, - model_state.model_record.title); - continue; - }, - Ok(CacheMissStrategy::Proceed) => {}, // We're going to go ahead... - } - } - - if maybe_sidecar_health_issue { - // Since we'll have a signal of the sidecar's health potentially being an issue, we don't - // need to background health check it from another thread. Instead we can react to the - // "potentially down" signal and block until it alleviates. - maybe_block_on_sidecar_health_check(&inferencer.http_clients.tts_sidecar_health_check_client).await; - maybe_sidecar_health_issue = false; - } - - let result = process_single_job(inferencer, &job, &model_state.model_record).await; + let result = process_single_job(job_state, &job).await; match result { - Ok((span1, span2)) => { - batch_spans.push(span1); - batch_spans.push(span2); - }, + Ok(_) => {}, Err(e) => { warn!("Failure to process job: {:?}", e); - - record_failure_and_maybe_slow_down(&inferencer.job_stats); - - maybe_sidecar_health_issue = true; - - let failure_reason = "failure processing job"; - let internal_debugging_failure_reason = format!("job error: {:?}", e); - - let _r = mark_tts_inference_job_failure( - &inferencer.mysql_pool, + let failure_reason = ""; + let _r = mark_generic_inference_job_failure( + &job_state.mysql_pool, &job, failure_reason, - &internal_debugging_failure_reason, - inferencer.job_max_attempts, - &inferencer.get_worker_name(), + job_state.job_max_attempts ).await; - - match e { - ProcessSingleJobError::Other(_) => {} // No-op - ProcessSingleJobError::FilesystemFull => { - // TODO: Refactor - we should stop processing all of these jobs as we'll lose out - // on this entire batch by attempting to clear the filesystem. This should be handled - // in the calling code. - delete_tts_synthesizers_from_cache(&inferencer.semi_persistent_cache)?; - } - } } } } - Ok(batch_spans) -} - -fn record_failure_and_maybe_slow_down(job_stats: &JobStats) { - let stats = match job_stats.increment_failure_count() { - Ok(stats) => stats, - Err(e) => { - warn!("Error recording stats and reacting to repeated failures: {:?}", e); - return; // Can't really do anything. - } - }; - - let seconds_timeout = match stats.consecutive_failure_count { - t if t > 100 => 180, - t if t > 50 => 60, - t if t > 20 => 30, - t if t > 10 => 10, - t if t > 5 => 5, - _ => return, // No timeout - }; - - info!("Slowing down {} seconds due to significant repeated failures: {:?}", - seconds_timeout, - stats); - - std::thread::sleep(Duration::from_secs(seconds_timeout)); -} - -/// Hack to delete locally cached TTS synthesizers to free up space from a full filesystem. -/// This is not intelligent and doesn't use any LRU/LFU mechanic. -/// This also relies on files not being read or written by concurrent workers while deleting. -fn delete_tts_synthesizers_from_cache(cache_dir: &SemiPersistentCacheDir) -> AnyhowResult<()> { - warn!("Deleting cached TTS synthesizers to free up disk space."); - - let tts_synthesizer_dir = cache_dir.tts_synthesizer_model_directory(); - - // TODO: When this is no longer sufficient, delete other types of locally-cached data. - let paths = std::fs::read_dir(tts_synthesizer_dir)? - .map(|res| res.map(|e| e.path())) - .collect::, std::io::Error>>()?; - - let models_to_delete = multiple_random_from_vec(&paths, 35); - - for model_to_delete in models_to_delete { - warn!("Deleting cached model file: {:?}", model_to_delete); - - let full_model_path = cache_dir.tts_synthesizer_model_path(model_to_delete); - std::fs::remove_file(full_model_path)?; - } - Ok(()) } - -/// We check both of these in one go so that we can reuse the ModelRecord later -/// without another DB query. -struct ModelState { - pub model_record: TtsModelForInferenceRecord, - pub is_downloaded_to_filesystem: bool, - pub is_in_memory_cache: bool, -} - -#[derive(Debug, Clone)] -enum ModelStateError { - ModelNotFound, - ModelDeleted, - CacheError { reason: String }, - DatabaseError { reason: String }, -} - -impl std::fmt::Display for ModelStateError { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self { - ModelStateError::ModelNotFound => write!(f, "ModelNotFound"), - ModelStateError::ModelDeleted => write!(f, "ModelDeleted"), - ModelStateError::CacheError { reason} => write!(f, "Cache error: {:?}", reason), - ModelStateError::DatabaseError { reason} => write!(f, "Database error: {:?}", reason), - } - } -} - -impl From for ModelStateError { - fn from(error: TtsModelForInferenceError) -> Self { - match error { - TtsModelForInferenceError::ModelNotFound => ModelStateError::ModelNotFound, - TtsModelForInferenceError::ModelDeleted => ModelStateError::ModelDeleted, - TtsModelForInferenceError::DatabaseError { reason } => ModelStateError::DatabaseError { reason } - } - } -} - -impl std::error::Error for ModelStateError {} - -impl ModelState { - /// Query the model details and see if the model file is on the filesystem in one go. - pub async fn query_model_and_check_filesystem( - job: &AvailableTtsInferenceJob, - mysql_pool: &MySqlPool, - tts_model_record_cache: &MultiItemTtlCache, - semi_persistent_cache: &SemiPersistentCacheDir, - virtual_cache: &SyncVirtualLfuCache, - ) -> Result { - // Many workers will be querying models constantly (n-many per batch). - // We can save on a lot of DB query volume by caching model records. - let maybe_cached_tts_model = - tts_model_record_cache.copy_without_bump_if_unexpired(job.model_token.clone()) - .ok() - .flatten(); - - let tts_model = match maybe_cached_tts_model { - Some(tts_model) => tts_model, - None => { - info!("Looking up TTS model record by token: {}", &job.model_token); - - let tts_model = get_tts_model_for_inference( - &mysql_pool, - &job.model_token - ).await?; - - tts_model_record_cache.store_copy(&job.model_token, &tts_model).ok(); - - tts_model - } - }; - - let tts_synthesizer_fs_path = semi_persistent_cache.tts_synthesizer_model_path( - &tts_model.model_token); - - let is_downloaded_to_filesystem = tts_synthesizer_fs_path.exists(); - - let path = tts_synthesizer_fs_path - .to_str() - .ok_or(ModelStateError::CacheError { reason: "could not make path".to_string() })? - .to_string(); - - let is_in_memory_cache = virtual_cache.in_cache(&path) - .map_err(|e| ModelStateError::CacheError { reason: format!("Model cache error: {:?}", e) })?; - - Ok(Self { - model_record: tts_model, - is_downloaded_to_filesystem, - is_in_memory_cache, - }) - } -} diff --git a/crates/service/job/inference_job/src/job_steps/main_loop_old.rs b/crates/service/job/inference_job/src/job_steps/main_loop_old.rs new file mode 100644 index 0000000000..c13727d237 --- /dev/null +++ b/crates/service/job/inference_job/src/job_steps/main_loop_old.rs @@ -0,0 +1,422 @@ +use container_common::anyhow_result::AnyhowResult; +use container_common::collections::multiple_random_from_vec::multiple_random_from_vec; +use crate::caching::cache_miss_strategizer::CacheMissStrategy; +use crate::caching::virtual_lfu_cache::SyncVirtualLfuCache; +use crate::job_steps::job_dependencies::JobDependencies; +use crate::job_steps::job_stats::JobStats; +use crate::job_steps::process_single_job_old::process_single_job_old; +use crate::job_steps::process_single_job_error::ProcessSingleJobError; +use crate::job_steps::sidecar_health_check_trap::maybe_block_on_sidecar_health_check; +use database_queries::queries::tts::tts_inference_jobs::list_available_tts_inference_jobs::{AvailableTtsInferenceJob, list_available_tts_inference_jobs, list_available_tts_inference_jobs_with_minimum_priority}; +use database_queries::queries::tts::tts_inference_jobs::mark_tts_inference_job_failure::mark_tts_inference_job_failure; +use database_queries::queries::tts::tts_inference_jobs::mark_tts_inference_job_permanently_dead::mark_tts_inference_job_permanently_dead; +use database_queries::queries::tts::tts_models::get_tts_model_for_inference::{get_tts_model_for_inference, TtsModelForInferenceError, TtsModelForInferenceRecord}; +use jobs_common::noop_logger::NoOpLogger; +use jobs_common::semi_persistent_cache_dir::SemiPersistentCacheDir; +use log::{error, info, warn}; +use memory_caching::multi_item_ttl_cache::MultiItemTtlCache; +use newrelic_telemetry::Span; +use sqlx::MySqlPool; +use std::time::Duration; + +// Job runner timeouts (guards MySQL) +const START_TIMEOUT_MILLIS : u64 = 500; +const INCREASE_TIMEOUT_MILLIS : u64 = 1000; + +pub async fn main_loop_old(job_args: JobDependencies) { + let mut error_timeout_millis = START_TIMEOUT_MILLIS; + + let mut noop_logger = NoOpLogger::new(job_args.no_op_logger_millis as i64); + + let mut span_batch = Vec::new(); + + let mut sort_by_priority = true; + let mut sort_by_priority_count = 0; + + let mut needs_health_check_at_start = true; // Run health check at startup. + + loop { + let num_records = job_args.job_batch_size; + + // Don't completely starve low-priority jobs + if sort_by_priority_count >= job_args.low_priority_starvation_prevention_every_nth { + sort_by_priority_count = 0; + sort_by_priority = false; + } + + let maybe_available_jobs = + if let Some(minimum_priority) = job_args.maybe_minimum_priority { + // Special high-priority workers + list_available_tts_inference_jobs_with_minimum_priority( + &job_args.mysql_pool, + minimum_priority, + num_records, + job_args.worker_details.is_debug_worker + ).await + } else { + // Normal path + list_available_tts_inference_jobs( + &job_args.mysql_pool, + sort_by_priority, + num_records, + job_args.worker_details.is_debug_worker + ).await + }; + + sort_by_priority = true; + sort_by_priority_count += 1; + + let jobs = match maybe_available_jobs { + Ok(jobs) => jobs, + Err(e) => { + warn!("Error querying jobs: {:?}", e); + std::thread::sleep(Duration::from_millis(error_timeout_millis)); + error_timeout_millis += INCREASE_TIMEOUT_MILLIS; + continue; + } + }; + + if jobs.is_empty() { + noop_logger.log_message_after_awhile("No inference jobs picked up from database!"); + + std::thread::sleep(Duration::from_millis(job_args.job_batch_wait_millis)); + continue; + } + + info!("Queried {} jobs from database", jobs.len()); + + let batch_result = process_jobs( + &job_args, + jobs, + needs_health_check_at_start, + ).await; + + if needs_health_check_at_start { + needs_health_check_at_start = false; + } + + let mut spans = match batch_result { + Ok(spans) => spans, + Err(e) => { + warn!("Error running job batch: {:?}", e); + std::thread::sleep(Duration::from_millis(error_timeout_millis)); + error_timeout_millis += INCREASE_TIMEOUT_MILLIS; + continue; + } + }; + + if !job_args.newrelic_disabled { + span_batch.append(&mut spans); + + if span_batch.len() > 50 { + let spans_to_send = span_batch.split_off(0).into(); + job_args.newrelic_client.send_spans(spans_to_send).await; + } + } + + error_timeout_millis = START_TIMEOUT_MILLIS; // reset + + std::thread::sleep(Duration::from_millis(job_args.job_batch_wait_millis)); + } +} + +/// Process a batch of jobs, returning the count of cold-cache skipped jobs. +async fn process_jobs( + inferencer: &JobDependencies, + jobs: Vec, + needs_health_check_at_start: bool, +) -> AnyhowResult> { + + if needs_health_check_at_start { + maybe_block_on_sidecar_health_check(&inferencer.http_clients.tts_sidecar_health_check_client).await; + } + + let mut maybe_sidecar_health_issue = false; + + let mut batch_spans = Vec::new(); + + for job in jobs.into_iter() { + let model_state_result = ModelState::query_model_and_check_filesystem( + &job, + &inferencer.mysql_pool, + &inferencer.caches.tts_model_record_cache, + &inferencer.semi_persistent_cache, + &inferencer.virtual_model_lfu, + ).await; + + let model_state = match model_state_result { + Ok(model_state) => model_state, + Err(e) => { + error!("TTS model fetch and state check error: {}, reason: {:?}", &job.model_token, &e); + + let (failure_reason, permanent_failure) = match e { + ModelStateError::ModelNotFound => ("model was not found", true), + ModelStateError::ModelDeleted => ("model has been deleted", true), + ModelStateError::CacheError { .. } => ("internal cache error", false), + ModelStateError::DatabaseError { .. } => ("unknown database error", false), + }; + + let internal_debugging_failure_reason = format!("model error: {:?}", e); + + let mut job_progress_reporter = inferencer + .job_progress_reporter + .new_tts_inference(&job.inference_job_token)?; + + job_progress_reporter.log_status(failure_reason)?; + + if permanent_failure { + warn!("Marking job permanently dead: {} because: {:?}", job.inference_job_token, &e); + + let _r = mark_tts_inference_job_permanently_dead( + &inferencer.mysql_pool, + job.id, + failure_reason, + &internal_debugging_failure_reason, + &inferencer.get_worker_name(), + ).await; + } else { + let _r = mark_tts_inference_job_failure( + &inferencer.mysql_pool, + &job, + failure_reason, + &internal_debugging_failure_reason, + inferencer.job_max_attempts, + &inferencer.get_worker_name(), + ).await; + } + + continue; + } + }; + + if !model_state.is_downloaded_to_filesystem || !model_state.is_in_memory_cache { + warn!("Model isn't ready: {} (downloaded = {}), (in memory = {})", + &model_state.model_record.model_token, + model_state.is_downloaded_to_filesystem, + model_state.is_in_memory_cache); + + let maybe_strategy = if !model_state.is_downloaded_to_filesystem { + inferencer.cache_miss_strategizers.disk_cache_miss(&model_state.model_record.model_token) + } else { + inferencer.cache_miss_strategizers.memory_cache_miss(&model_state.model_record.model_token) + }; + + match maybe_strategy { + Err(err) => { + warn!("Unable to process job: {:?}", err); + + let failure_reason = "cache error"; + let internal_debugging_failure_reason = format!("cache error: {:?}", err); + + let _r = mark_tts_inference_job_failure( + &inferencer.mysql_pool, + &job, + failure_reason, + &internal_debugging_failure_reason, + inferencer.job_max_attempts, + &inferencer.get_worker_name(), + ).await; + continue; + }, + Ok(CacheMissStrategy::WaitOrSkip) => { + // We're going to skip this for now. + // Maybe another worker has a warm cache and can continue. + warn!("Skipping TTS due to cold cache: {} ({})", + model_state.model_record.model_token, + model_state.model_record.title); + continue; + }, + Ok(CacheMissStrategy::Proceed) => {}, // We're going to go ahead... + } + } + + if maybe_sidecar_health_issue { + // Since we'll have a signal of the sidecar's health potentially being an issue, we don't + // need to background health check it from another thread. Instead we can react to the + // "potentially down" signal and block until it alleviates. + maybe_block_on_sidecar_health_check(&inferencer.http_clients.tts_sidecar_health_check_client).await; + maybe_sidecar_health_issue = false; + } + + let result = process_single_job_old(inferencer, &job, &model_state.model_record).await; + match result { + Ok((span1, span2)) => { + batch_spans.push(span1); + batch_spans.push(span2); + }, + Err(e) => { + warn!("Failure to process job: {:?}", e); + + record_failure_and_maybe_slow_down(&inferencer.job_stats); + + maybe_sidecar_health_issue = true; + + let failure_reason = "failure processing job"; + let internal_debugging_failure_reason = format!("job error: {:?}", e); + + let _r = mark_tts_inference_job_failure( + &inferencer.mysql_pool, + &job, + failure_reason, + &internal_debugging_failure_reason, + inferencer.job_max_attempts, + &inferencer.get_worker_name(), + ).await; + + match e { + ProcessSingleJobError::Other(_) => {} // No-op + ProcessSingleJobError::FilesystemFull => { + // TODO: Refactor - we should stop processing all of these jobs as we'll lose out + // on this entire batch by attempting to clear the filesystem. This should be handled + // in the calling code. + delete_tts_synthesizers_from_cache(&inferencer.semi_persistent_cache)?; + } + } + } + } + } + + Ok(batch_spans) +} + +fn record_failure_and_maybe_slow_down(job_stats: &JobStats) { + let stats = match job_stats.increment_failure_count() { + Ok(stats) => stats, + Err(e) => { + warn!("Error recording stats and reacting to repeated failures: {:?}", e); + return; // Can't really do anything. + } + }; + + let seconds_timeout = match stats.consecutive_failure_count { + t if t > 100 => 180, + t if t > 50 => 60, + t if t > 20 => 30, + t if t > 10 => 10, + t if t > 5 => 5, + _ => return, // No timeout + }; + + info!("Slowing down {} seconds due to significant repeated failures: {:?}", + seconds_timeout, + stats); + + std::thread::sleep(Duration::from_secs(seconds_timeout)); +} + +/// Hack to delete locally cached TTS synthesizers to free up space from a full filesystem. +/// This is not intelligent and doesn't use any LRU/LFU mechanic. +/// This also relies on files not being read or written by concurrent workers while deleting. +fn delete_tts_synthesizers_from_cache(cache_dir: &SemiPersistentCacheDir) -> AnyhowResult<()> { + warn!("Deleting cached TTS synthesizers to free up disk space."); + + let tts_synthesizer_dir = cache_dir.tts_synthesizer_model_directory(); + + // TODO: When this is no longer sufficient, delete other types of locally-cached data. + let paths = std::fs::read_dir(tts_synthesizer_dir)? + .map(|res| res.map(|e| e.path())) + .collect::, std::io::Error>>()?; + + let models_to_delete = multiple_random_from_vec(&paths, 35); + + for model_to_delete in models_to_delete { + warn!("Deleting cached model file: {:?}", model_to_delete); + + let full_model_path = cache_dir.tts_synthesizer_model_path(model_to_delete); + std::fs::remove_file(full_model_path)?; + } + + Ok(()) +} + +/// We check both of these in one go so that we can reuse the ModelRecord later +/// without another DB query. +struct ModelState { + pub model_record: TtsModelForInferenceRecord, + pub is_downloaded_to_filesystem: bool, + pub is_in_memory_cache: bool, +} + +#[derive(Debug, Clone)] +enum ModelStateError { + ModelNotFound, + ModelDeleted, + CacheError { reason: String }, + DatabaseError { reason: String }, +} + +impl std::fmt::Display for ModelStateError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + ModelStateError::ModelNotFound => write!(f, "ModelNotFound"), + ModelStateError::ModelDeleted => write!(f, "ModelDeleted"), + ModelStateError::CacheError { reason} => write!(f, "Cache error: {:?}", reason), + ModelStateError::DatabaseError { reason} => write!(f, "Database error: {:?}", reason), + } + } +} + +impl From for ModelStateError { + fn from(error: TtsModelForInferenceError) -> Self { + match error { + TtsModelForInferenceError::ModelNotFound => ModelStateError::ModelNotFound, + TtsModelForInferenceError::ModelDeleted => ModelStateError::ModelDeleted, + TtsModelForInferenceError::DatabaseError { reason } => ModelStateError::DatabaseError { reason } + } + } +} + +impl std::error::Error for ModelStateError {} + +impl ModelState { + /// Query the model details and see if the model file is on the filesystem in one go. + pub async fn query_model_and_check_filesystem( + job: &AvailableTtsInferenceJob, + mysql_pool: &MySqlPool, + tts_model_record_cache: &MultiItemTtlCache, + semi_persistent_cache: &SemiPersistentCacheDir, + virtual_cache: &SyncVirtualLfuCache, + ) -> Result { + // Many workers will be querying models constantly (n-many per batch). + // We can save on a lot of DB query volume by caching model records. + let maybe_cached_tts_model = + tts_model_record_cache.copy_without_bump_if_unexpired(job.model_token.clone()) + .ok() + .flatten(); + + let tts_model = match maybe_cached_tts_model { + Some(tts_model) => tts_model, + None => { + info!("Looking up TTS model record by token: {}", &job.model_token); + + let tts_model = get_tts_model_for_inference( + &mysql_pool, + &job.model_token + ).await?; + + tts_model_record_cache.store_copy(&job.model_token, &tts_model).ok(); + + tts_model + } + }; + + let tts_synthesizer_fs_path = semi_persistent_cache.tts_synthesizer_model_path( + &tts_model.model_token); + + let is_downloaded_to_filesystem = tts_synthesizer_fs_path.exists(); + + let path = tts_synthesizer_fs_path + .to_str() + .ok_or(ModelStateError::CacheError { reason: "could not make path".to_string() })? + .to_string(); + + let is_in_memory_cache = virtual_cache.in_cache(&path) + .map_err(|e| ModelStateError::CacheError { reason: format!("Model cache error: {:?}", e) })?; + + Ok(Self { + model_record: tts_model, + is_downloaded_to_filesystem, + is_in_memory_cache, + }) + } +} diff --git a/crates/service/job/inference_job/src/job_steps/mod.rs b/crates/service/job/inference_job/src/job_steps/mod.rs index 02569e6da9..e1f0d557b6 100644 --- a/crates/service/job/inference_job/src/job_steps/mod.rs +++ b/crates/service/job/inference_job/src/job_steps/mod.rs @@ -1,6 +1,9 @@ pub mod job_dependencies; pub mod job_stats; pub mod main_loop; +pub mod main_loop_old; pub mod process_single_job; pub mod process_single_job_error; +pub mod process_single_job_old; +pub mod select_job_batch; pub mod sidecar_health_check_trap; diff --git a/crates/service/job/inference_job/src/job_steps/process_single_job.rs b/crates/service/job/inference_job/src/job_steps/process_single_job.rs index 584ded0688..f91ac536f7 100644 --- a/crates/service/job/inference_job/src/job_steps/process_single_job.rs +++ b/crates/service/job/inference_job/src/job_steps/process_single_job.rs @@ -1,433 +1,120 @@ -// NB: Incrementally getting rid of build warnings... -#![forbid(unused_imports)] -#![forbid(unused_mut)] -#![forbid(unused_variables)] - use anyhow::anyhow; -use container_common::anyhow_result::AnyhowResult; -use container_common::filesystem::check_file_exists::check_file_exists; +use config::is_bad_download_url::is_bad_download_url; use container_common::filesystem::safe_delete_temp_directory::safe_delete_temp_directory; -use container_common::filesystem::safe_delete_temp_file::safe_delete_temp_file; -use container_common::hashing::hash_string_sha2::hash_string_sha2; -use container_common::token::random_uuid::generate_random_uuid; use crate::job_steps::job_dependencies::JobDependencies; -use crate::job_steps::process_single_job_error::ProcessSingleJobError; -use crate::job_types::tts::seconds_to_decoder_steps::seconds_to_decoder_steps; -use crate::util::download_file_from_bucket::maybe_download_file_from_bucket; -use database_queries::column_types::vocoder_type::VocoderType; -use database_queries::queries::tts::tts_inference_jobs::list_available_tts_inference_jobs::AvailableTtsInferenceJob; -use database_queries::queries::tts::tts_inference_jobs::mark_tts_inference_job_done::mark_tts_inference_job_done; -use database_queries::queries::tts::tts_inference_jobs::mark_tts_inference_job_pending_and_grab_lock::mark_tts_inference_job_pending_and_grab_lock; -use database_queries::queries::tts::tts_models::get_tts_model_for_inference::TtsModelForInferenceRecord; -use database_queries::queries::tts::tts_results::insert_tts_result::insert_tts_result; -use log::{warn, info, error}; -use newrelic_telemetry::Span; -use std::fs::File; -use std::io::Read; -use std::path::PathBuf; -use std::time::{SystemTime, UNIX_EPOCH, Instant}; +use database_queries::queries::generic_inference::job::list_available_generic_inference_jobs::AvailableInferenceJob; +use database_queries::queries::generic_inference::job::mark_generic_inference_job_pending_and_grab_lock::mark_generic_inference_job_pending_and_grab_lock; +use database_queries::queries::generic_inference::job::mark_generic_inference_job_successfully_done::mark_generic_inference_job_successfully_done; +use enums::workers::generic_inference_type::GenericInferenceType; +use errors::AnyhowResult; +use jobs_common::redis_job_status_logger::RedisJobStatusLogger; +use log::{info, warn}; use tempdir::TempDir; -use tts_common::clean_symbols::clean_symbols; -use tts_common::text_pipelines::guess_pipeline::guess_text_pipeline_heuristic; -use tts_common::text_pipelines::text_pipeline_type::TextPipelineType; -#[derive(Deserialize, Default)] -struct FileMetadata { - pub duration_millis: Option, - pub mimetype: Option, - pub file_size_bytes: u64, -} - -fn read_metadata_file(filename: &PathBuf) -> AnyhowResult { - let mut file = File::open(filename)?; - let mut buffer = String::new(); - file.read_to_string(&mut buffer)?; - Ok(serde_json::from_str(&buffer)?) -} - -pub async fn process_single_job( - job_args: &JobDependencies, - job: &AvailableTtsInferenceJob, - model_record: &TtsModelForInferenceRecord, -) -> Result<(Span, Span), ProcessSingleJobError> { - - let start = Instant::now(); - - let span_id = generate_random_uuid(); - let trace_id = generate_random_uuid(); - let maybe_user_token = job.maybe_creator_user_token.as_deref().unwrap_or(""); - - let mut job_iteration_span = Span::new(&span_id, &trace_id, get_timestamp_millis()) - .name("single job execution") - .attribute("user_token", maybe_user_token) - .service_name("tts-inference-job"); - - let span_id = generate_random_uuid(); - let trace_id = generate_random_uuid(); - - let created_at_timestamp = (job.created_at.timestamp() as u64) * 1000; - - let mut since_creation_span = Span::new(&span_id, &trace_id, created_at_timestamp) - .name("job since creation") - .attribute("user_token", maybe_user_token) - .service_name("tts-inference-job"); - - let mut job_progress_reporter = job_args - .job_progress_reporter - .new_tts_inference(&job.inference_job_token) - .map_err(|e| ProcessSingleJobError::Other(anyhow!(e)))?; +pub async fn process_single_job(job_dependencies: &JobDependencies, job: &AvailableInferenceJob) -> AnyhowResult<()> { + let mut redis = job_dependencies.redis_pool.get()?; + let mut redis_logger = RedisJobStatusLogger::new_generic_download(&mut redis, job.download_job_token.as_str()); // ==================== ATTEMPT TO GRAB JOB LOCK ==================== // - info!("Attempting to grab lock for job: {}", job.inference_job_token); - - let lock_acquired = - mark_tts_inference_job_pending_and_grab_lock(&job_args.mysql_pool, job.id) - .await - .map_err(|e| ProcessSingleJobError::Other(e))?; + let lock_acquired = mark_generic_inference_job_pending_and_grab_lock(&job_dependencies.mysql_pool, job.id).await?; if !lock_acquired { - warn!("Could not acquire job lock for: {:?}", &job.id); - let duration = start.elapsed(); - - since_creation_span.set_attribute("status", "failure"); - since_creation_span.set_duration(duration); - - job_iteration_span.set_attribute("status", "failure"); - job_iteration_span.set_duration(duration); - - return Ok((since_creation_span, job_iteration_span)); + warn!("Could not acquire job lock for: {}", &job.id.0); + return Ok(()) } - info!("Lock acquired for job: {}", job.inference_job_token); + info!("Beginning work on {:?} = {} | {}", job.download_type, job.download_job_token, job.download_url); - // ==================== CONFIRM OR DOWNLOAD WAVEGLOW VOCODER MODEL ==================== // + // ==================== SETUP TEMP DIRS ==================== // - let waveglow_vocoder_model_fs_path = { - let waveglow_vocoder_model_filename = job_args.waveglow_vocoder_model_filename.clone(); - let waveglow_vocoder_model_fs_path = job_args.semi_persistent_cache.tts_pretrained_vocoder_model_path(&waveglow_vocoder_model_filename); - let waveglow_vocoder_model_object_path = job_args.bucket_path_unifier.tts_pretrained_vocoders_path(&waveglow_vocoder_model_filename); + // TODO: Scoped temp dir - maybe_download_file_from_bucket( - "waveglow vocoder model", - &waveglow_vocoder_model_fs_path, - &waveglow_vocoder_model_object_path, - &job_args.private_bucket_client, - &mut job_progress_reporter, - "downloading vocoder (1 of 3)", - job.id.0, - &job_args.scoped_temp_dir_creator, - ).await?; + let temp_dir = format!("temp_{}", job.id.0); + let temp_dir = TempDir::new(&temp_dir)?; - waveglow_vocoder_model_fs_path - }; + // ==================== HANDLE DIFFERENT INFERENCE TYPES ==================== // - // ==================== CONFIRM OR DOWNLOAD HIFIGAN (NORMAL) VOCODER MODEL ==================== // + let mut entity_token : Option = None; + let mut entity_type : Option = None; - let pretrained_hifigan_vocoder_model_fs_path = { - let hifigan_vocoder_model_filename = job_args.hifigan_vocoder_model_filename.clone(); - let hifigan_vocoder_model_fs_path = job_args.semi_persistent_cache.tts_pretrained_vocoder_model_path(&hifigan_vocoder_model_filename); - let hifigan_vocoder_model_object_path = job_args.bucket_path_unifier.tts_pretrained_vocoders_path(&hifigan_vocoder_model_filename); - - maybe_download_file_from_bucket( - "hifigan vocoder model", - &hifigan_vocoder_model_fs_path, - &hifigan_vocoder_model_object_path, - &job_args.private_bucket_client, - &mut job_progress_reporter, - "downloading vocoder (2 of 3)", - job.id.0, - &job_args.scoped_temp_dir_creator, - ).await?; - - hifigan_vocoder_model_fs_path - }; - - // ==================== CONFIRM OR DOWNLOAD HIFIGAN (SUPERRES) VOCODER MODEL ==================== // - - let hifigan_superres_vocoder_model_fs_path = { - let hifigan_superres_vocoder_model_filename = job_args.hifigan_superres_vocoder_model_filename.clone(); - let hifigan_superres_vocoder_model_fs_path = job_args.semi_persistent_cache.tts_pretrained_vocoder_model_path(&hifigan_superres_vocoder_model_filename); - let hifigan_superres_vocoder_model_object_path = job_args.bucket_path_unifier.tts_pretrained_vocoders_path(&hifigan_superres_vocoder_model_filename); - - maybe_download_file_from_bucket( - "hifigan superres vocoder model", - &hifigan_superres_vocoder_model_fs_path, - &hifigan_superres_vocoder_model_object_path, - &job_args.private_bucket_client, - &mut job_progress_reporter, - "downloading vocoder (3 of 3)", - job.id.0, - &job_args.scoped_temp_dir_creator, - ).await?; - - hifigan_superres_vocoder_model_fs_path - }; - -// // ==================== CONFIRM OR DOWNLOAD OPTIONAL CUSTOM VOCODER MODEL ==================== // - - let custom_vocoder_fs_path = match &model_record.maybe_custom_vocoder { - None => None, - Some(vocoder) => { - let custom_vocoder_fs_path = job_args.semi_persistent_cache.custom_vocoder_model_path(&vocoder.vocoder_token); - let custom_vocoder_object_path = job_args.bucket_path_unifier.vocoder_path(&vocoder.vocoder_private_bucket_hash); - - maybe_download_file_from_bucket( - "custom vocoder", - &custom_vocoder_fs_path, - &custom_vocoder_object_path, - &job_args.private_bucket_client, - &mut job_progress_reporter, - "downloading user vocoder", - job.id.0, - &job_args.scoped_temp_dir_creator, - ).await?; - - Some(custom_vocoder_fs_path) + match job.inference_type { + GenericInferenceType::TextToSpeech => { + // TODO } - }; - - // ==================== CONFIRM OR DOWNLOAD TTS SYNTHESIZER MODEL ==================== // - - let tts_synthesizer_fs_path = { - let tts_synthesizer_fs_path = job_args.semi_persistent_cache.tts_synthesizer_model_path(&model_record.model_token); - let tts_synthesizer_object_path = job_args.bucket_path_unifier.tts_synthesizer_path(&model_record.private_bucket_hash); - - maybe_download_file_from_bucket( - "synthesizer", - &tts_synthesizer_fs_path, - &tts_synthesizer_object_path, - &job_args.private_bucket_client, - &mut job_progress_reporter, - "downloading synthesizer", - job.id.0, - &job_args.scoped_temp_dir_creator, - ).await?; - - tts_synthesizer_fs_path - }; - - // ==================== Preprocess text ==================== // - - let cleaned_inference_text = clean_symbols(&job.raw_inference_text); - - // ==================== WRITE TEXT TO FILE ==================== // - - info!("Creating tempdir for inference results."); - - let temp_dir = format!("temp_tts_inference_{}", job.id.0); - - // NB: TempDir exists until it goes out of scope, at which point it should delete from filesystem. - let temp_dir = TempDir::new(&temp_dir) - .map_err(|e| ProcessSingleJobError::from_io_error(e))?; - - let text_input_fs_path = temp_dir.path().join("inference_input.txt"); - - std::fs::write(&text_input_fs_path, &cleaned_inference_text) - .map_err(|e| ProcessSingleJobError::from_io_error(e))?; - - // ==================== RUN INFERENCE ==================== // - - job_progress_reporter.log_status("running inference") - .map_err(|e| ProcessSingleJobError::Other(e))?; - - // TODO: Fix this. - let maybe_unload_model_path = job_args - .virtual_model_lfu - .insert_returning_replaced(tts_synthesizer_fs_path.to_str().unwrap_or("")) - .map_err(|e| ProcessSingleJobError::Other(e))?; - - if let Some(model_path) = maybe_unload_model_path.as_deref() { - warn!("Remove model from LFU cache: {:?}", model_path); + GenericInferenceType::VoiceConversion => { + // TODO + } +// GenericDownloadType::HifiGan => { +// let results = process_hifigan_vocoder( +// job_state, +// job, +// &temp_dir, +// &download_filename, +// &mut redis_logger, +// ).await?; +// entity_token = results.entity_token.clone(); +// entity_type = results.entity_type.clone(); +// } +// GenericDownloadType::HifiGanRocketVc => { +// let results = process_hifigan_softvc_vocoder( +// job_state, +// job, +// &temp_dir, +// &download_filename, +// &mut redis_logger, +// ).await?; +// entity_token = results.entity_token.clone(); +// entity_type = results.entity_type.clone(); +// } +// GenericDownloadType::RocketVc => { +// let results = process_softvc_model( +// job_state, +// job, +// &temp_dir, +// &download_filename, +// &mut redis_logger, +// ).await?; +// entity_token = results.entity_token.clone(); +// entity_type = results.entity_type.clone(); +// } +// GenericDownloadType::Tacotron2 => { +// let results = process_tacotron_model( +// job_state, +// job, +// &temp_dir, +// &download_filename, +// &mut redis_logger, +// ).await?; +// entity_token = results.entity_token.clone(); +// entity_type = results.entity_type.clone(); +// } } - let output_audio_fs_path = temp_dir.path().join("output.wav"); - let output_metadata_fs_path = temp_dir.path().join("metadata.json"); - let output_spectrogram_fs_path = temp_dir.path().join("spectrogram.json"); - - info!("Running TTS inference..."); - - info!("Expected output audio filename: {:?}", &output_audio_fs_path); - info!("Expected output spectrogram filename: {:?}", &output_spectrogram_fs_path); - info!("Expected output metadata filename: {:?}", &output_metadata_fs_path); - - if let Some(model_path) = maybe_unload_model_path.as_deref() { - warn!("Unload model from sidecar: {:?}", &model_path); - } - - let mut pretrained_vocoder = VocoderType::HifiGanSuperResolution; - if let Some(default_vocoder) = model_record.maybe_default_pretrained_vocoder.as_deref() { - pretrained_vocoder = VocoderType::from_str(default_vocoder) - .map_err(|e| ProcessSingleJobError::Other(e))?; - } - - info!("With pretrained vocoder: {:?}", pretrained_vocoder); - - let text_pipeline_type_or_guess = model_record.text_pipeline_type - .as_deref() - .and_then(|pipeline_type| - TextPipelineType::from_str(pipeline_type).ok())// NB: If there's an error deserializing, turn it to None. - .unwrap_or_else(|| - guess_text_pipeline_heuristic(Some(model_record.created_at))); - - info!("With text pipeline type `{:?} ` (or guess: {:?})", - &model_record.text_pipeline_type, - &text_pipeline_type_or_guess); - - let hifigan_vocoder_model_fs_path_to_use = match custom_vocoder_fs_path { - None => { - info!("using pretrained HiFi-GAN vocoder"); - pretrained_hifigan_vocoder_model_fs_path - }, - Some(custom_vocoder_fs_path) => { - info!("using custom user-trained HiFi-GAN vocoder: {:?}", custom_vocoder_fs_path); - custom_vocoder_fs_path - }, - }; - - // NB: Tacotron operates on decoder steps. 1000 steps is the default and correlates to - // roughly 12 seconds max. Here we map seconds to decoder steps. - let max_decoder_steps = seconds_to_decoder_steps(job.max_duration_seconds); - - job_args.http_clients.tts_inference_sidecar_client.request_inference( - &cleaned_inference_text, - max_decoder_steps, - &tts_synthesizer_fs_path, - pretrained_vocoder, - &text_pipeline_type_or_guess.to_str(), - &hifigan_vocoder_model_fs_path_to_use, - &hifigan_superres_vocoder_model_fs_path, - &waveglow_vocoder_model_fs_path, - &output_audio_fs_path, - &output_spectrogram_fs_path, - &output_metadata_fs_path, - maybe_unload_model_path) - .await - .map_err(|e| ProcessSingleJobError::Other(e))?; - - // ==================== CHECK ALL FILES EXIST AND GET METADATA ==================== // - - info!("Checking that output files exist..."); - - check_file_exists(&output_audio_fs_path).map_err(|e| ProcessSingleJobError::Other(e))?; - check_file_exists(&output_spectrogram_fs_path).map_err(|e| ProcessSingleJobError::Other(e))?; - check_file_exists(&output_metadata_fs_path).map_err(|e| ProcessSingleJobError::Other(e))?; - - let file_metadata = read_metadata_file(&output_metadata_fs_path) - .map_err(|e| ProcessSingleJobError::Other(e))?; - - safe_delete_temp_file(&output_metadata_fs_path); - - // ==================== UPLOAD AUDIO TO BUCKET ==================== // - - job_progress_reporter.log_status("uploading result") - .map_err(|e| ProcessSingleJobError::Other(e))?; - - let audio_result_object_path = job_args.bucket_path_unifier.tts_inference_wav_audio_output_path( - &job.uuid_idempotency_token); // TODO: Don't use this! - - info!("Audio destination bucket path: {:?}", &audio_result_object_path); - - info!("Uploading audio..."); - - job_args.public_bucket_client.upload_filename_with_content_type( - &audio_result_object_path, - &output_audio_fs_path, - "audio/wav") - .await - .map_err(|e| ProcessSingleJobError::Other(e))?; - - safe_delete_temp_file(&output_audio_fs_path); - - // ==================== UPLOAD SPECTROGRAM TO BUCKETS ==================== // - - let spectrogram_result_object_path = job_args.bucket_path_unifier.tts_inference_spectrogram_output_path( - &job.uuid_idempotency_token); // TODO: Don't use this! - - info!("Spectrogram destination bucket path: {:?}", &spectrogram_result_object_path); - - info!("Uploading spectrogram..."); - - job_args.public_bucket_client.upload_filename_with_content_type( - &spectrogram_result_object_path, - &output_spectrogram_fs_path, - "application/json") - .await - .map_err(|e| ProcessSingleJobError::Other(e))?; - - safe_delete_temp_file(&output_spectrogram_fs_path); - - // ==================== DELETE DOWNLOADED FILE ==================== // - - // NB: We should be using a tempdir, but to make absolutely certain we don't overflow the disk... - safe_delete_temp_directory(&temp_dir); - - // ==================== SAVE RECORDS ==================== // - - let text_hash = hash_string_sha2(&cleaned_inference_text) - .map_err(|e| ProcessSingleJobError::Other(e))?; - - let worker_name = job_args.get_worker_name(); - - info!("Saving tts inference record..."); - - let (id, inference_result_token) = insert_tts_result( - &job_args.mysql_pool, - job, - &text_hash, - pretrained_vocoder, - &audio_result_object_path, - &spectrogram_result_object_path, - file_metadata.file_size_bytes, - file_metadata.duration_millis.unwrap_or(0), - job_args.worker_details.is_on_prem, - &worker_name, - job_args.worker_details.is_debug_worker) - .await - .map_err(|e| ProcessSingleJobError::Other(e))?; + // ===================================================== info!("Marking job complete..."); - mark_tts_inference_job_done( - &job_args.mysql_pool, - job.id, - true, - Some(&inference_result_token), - &worker_name) - .await - .map_err(|e| ProcessSingleJobError::Other(e))?; + mark_generic_inference_job_successfully_done( + &job_dependencies.mysql_pool, + job, + entity_type.as_deref(), + entity_token.as_deref(), + ).await?; - info!("TTS Done. Original text was: {}", &job.raw_inference_text); + info!("Saved model record: {} - {}", job.id.0, &job.inference_job_token); - job_args.firehose_publisher.tts_inference_finished( - job.maybe_creator_user_token.as_deref(), - &model_record.model_token, - &inference_result_token) + job_dependencies.firehose_publisher.publish_generic_download_finished( + &job.creator_user_token, + entity_token.as_deref()) .await .map_err(|e| { - error!("error publishing event: {:?}", e); - ProcessSingleJobError::Other(anyhow!("error publishing event")) + warn!("error publishing event: {:?}", e); + anyhow!("error publishing event") })?; - job_progress_reporter.log_status("done") - .map_err(|e| ProcessSingleJobError::Other(e))?; + redis_logger.log_status("done")?; - info!("Job {:?} complete success! Downloaded, ran inference, and uploaded. Saved model record: {}, Result Token: {}", - job.id, id, &inference_result_token); + info!("Job done: {}", job.id.0); - let duration = start.elapsed(); - - since_creation_span.set_attribute("status", "success"); - since_creation_span.set_duration(duration); - - job_iteration_span.set_attribute("status", "success"); - job_iteration_span.set_duration(duration); - - Ok((since_creation_span, job_iteration_span)) -} - -fn get_timestamp_millis() -> u64 { - SystemTime::now() - .duration_since(UNIX_EPOCH) - .map(|d| d.as_millis() as u64) - .unwrap_or(0) + Ok(()) } diff --git a/crates/service/job/inference_job/src/job_steps/process_single_job_old.rs b/crates/service/job/inference_job/src/job_steps/process_single_job_old.rs new file mode 100644 index 0000000000..049d443927 --- /dev/null +++ b/crates/service/job/inference_job/src/job_steps/process_single_job_old.rs @@ -0,0 +1,433 @@ +// NB: Incrementally getting rid of build warnings... +#![forbid(unused_imports)] +#![forbid(unused_mut)] +#![forbid(unused_variables)] + +use anyhow::anyhow; +use container_common::filesystem::check_file_exists::check_file_exists; +use container_common::filesystem::safe_delete_temp_directory::safe_delete_temp_directory; +use container_common::filesystem::safe_delete_temp_file::safe_delete_temp_file; +use container_common::hashing::hash_string_sha2::hash_string_sha2; +use container_common::token::random_uuid::generate_random_uuid; +use crate::job_steps::job_dependencies::JobDependencies; +use crate::job_steps::process_single_job_error::ProcessSingleJobError; +use crate::job_types::tts::seconds_to_decoder_steps::seconds_to_decoder_steps; +use crate::util::download_file_from_bucket::maybe_download_file_from_bucket; +use database_queries::column_types::vocoder_type::VocoderType; +use database_queries::queries::tts::tts_inference_jobs::list_available_tts_inference_jobs::AvailableTtsInferenceJob; +use database_queries::queries::tts::tts_inference_jobs::mark_tts_inference_job_done::mark_tts_inference_job_done; +use database_queries::queries::tts::tts_inference_jobs::mark_tts_inference_job_pending_and_grab_lock::mark_tts_inference_job_pending_and_grab_lock; +use database_queries::queries::tts::tts_models::get_tts_model_for_inference::TtsModelForInferenceRecord; +use database_queries::queries::tts::tts_results::insert_tts_result::insert_tts_result; +use errors::AnyhowResult; +use log::{warn, info, error}; +use newrelic_telemetry::Span; +use std::fs::File; +use std::io::Read; +use std::path::PathBuf; +use std::time::{SystemTime, UNIX_EPOCH, Instant}; +use tempdir::TempDir; +use tts_common::clean_symbols::clean_symbols; +use tts_common::text_pipelines::guess_pipeline::guess_text_pipeline_heuristic; +use tts_common::text_pipelines::text_pipeline_type::TextPipelineType; + +#[derive(Deserialize, Default)] +struct FileMetadata { + pub duration_millis: Option, + pub mimetype: Option, + pub file_size_bytes: u64, +} + +fn read_metadata_file(filename: &PathBuf) -> AnyhowResult { + let mut file = File::open(filename)?; + let mut buffer = String::new(); + file.read_to_string(&mut buffer)?; + Ok(serde_json::from_str(&buffer)?) +} + +pub async fn process_single_job_old( + job_args: &JobDependencies, + job: &AvailableTtsInferenceJob, + model_record: &TtsModelForInferenceRecord, +) -> Result<(Span, Span), ProcessSingleJobError> { + + let start = Instant::now(); + + let span_id = generate_random_uuid(); + let trace_id = generate_random_uuid(); + let maybe_user_token = job.maybe_creator_user_token.as_deref().unwrap_or(""); + + let mut job_iteration_span = Span::new(&span_id, &trace_id, get_timestamp_millis()) + .name("single job execution") + .attribute("user_token", maybe_user_token) + .service_name("tts-inference-job"); + + let span_id = generate_random_uuid(); + let trace_id = generate_random_uuid(); + + let created_at_timestamp = (job.created_at.timestamp() as u64) * 1000; + + let mut since_creation_span = Span::new(&span_id, &trace_id, created_at_timestamp) + .name("job since creation") + .attribute("user_token", maybe_user_token) + .service_name("tts-inference-job"); + + let mut job_progress_reporter = job_args + .job_progress_reporter + .new_tts_inference(&job.inference_job_token) + .map_err(|e| ProcessSingleJobError::Other(anyhow!(e)))?; + + // ==================== ATTEMPT TO GRAB JOB LOCK ==================== // + + info!("Attempting to grab lock for job: {}", job.inference_job_token); + + let lock_acquired = + mark_tts_inference_job_pending_and_grab_lock(&job_args.mysql_pool, job.id) + .await + .map_err(|e| ProcessSingleJobError::Other(e))?; + + if !lock_acquired { + warn!("Could not acquire job lock for: {:?}", &job.id); + let duration = start.elapsed(); + + since_creation_span.set_attribute("status", "failure"); + since_creation_span.set_duration(duration); + + job_iteration_span.set_attribute("status", "failure"); + job_iteration_span.set_duration(duration); + + return Ok((since_creation_span, job_iteration_span)); + } + + info!("Lock acquired for job: {}", job.inference_job_token); + + // ==================== CONFIRM OR DOWNLOAD WAVEGLOW VOCODER MODEL ==================== // + + let waveglow_vocoder_model_fs_path = { + let waveglow_vocoder_model_filename = job_args.waveglow_vocoder_model_filename.clone(); + let waveglow_vocoder_model_fs_path = job_args.semi_persistent_cache.tts_pretrained_vocoder_model_path(&waveglow_vocoder_model_filename); + let waveglow_vocoder_model_object_path = job_args.bucket_path_unifier.tts_pretrained_vocoders_path(&waveglow_vocoder_model_filename); + + maybe_download_file_from_bucket( + "waveglow vocoder model", + &waveglow_vocoder_model_fs_path, + &waveglow_vocoder_model_object_path, + &job_args.private_bucket_client, + &mut job_progress_reporter, + "downloading vocoder (1 of 3)", + job.id.0, + &job_args.scoped_temp_dir_creator, + ).await?; + + waveglow_vocoder_model_fs_path + }; + + // ==================== CONFIRM OR DOWNLOAD HIFIGAN (NORMAL) VOCODER MODEL ==================== // + + let pretrained_hifigan_vocoder_model_fs_path = { + let hifigan_vocoder_model_filename = job_args.hifigan_vocoder_model_filename.clone(); + let hifigan_vocoder_model_fs_path = job_args.semi_persistent_cache.tts_pretrained_vocoder_model_path(&hifigan_vocoder_model_filename); + let hifigan_vocoder_model_object_path = job_args.bucket_path_unifier.tts_pretrained_vocoders_path(&hifigan_vocoder_model_filename); + + maybe_download_file_from_bucket( + "hifigan vocoder model", + &hifigan_vocoder_model_fs_path, + &hifigan_vocoder_model_object_path, + &job_args.private_bucket_client, + &mut job_progress_reporter, + "downloading vocoder (2 of 3)", + job.id.0, + &job_args.scoped_temp_dir_creator, + ).await?; + + hifigan_vocoder_model_fs_path + }; + + // ==================== CONFIRM OR DOWNLOAD HIFIGAN (SUPERRES) VOCODER MODEL ==================== // + + let hifigan_superres_vocoder_model_fs_path = { + let hifigan_superres_vocoder_model_filename = job_args.hifigan_superres_vocoder_model_filename.clone(); + let hifigan_superres_vocoder_model_fs_path = job_args.semi_persistent_cache.tts_pretrained_vocoder_model_path(&hifigan_superres_vocoder_model_filename); + let hifigan_superres_vocoder_model_object_path = job_args.bucket_path_unifier.tts_pretrained_vocoders_path(&hifigan_superres_vocoder_model_filename); + + maybe_download_file_from_bucket( + "hifigan superres vocoder model", + &hifigan_superres_vocoder_model_fs_path, + &hifigan_superres_vocoder_model_object_path, + &job_args.private_bucket_client, + &mut job_progress_reporter, + "downloading vocoder (3 of 3)", + job.id.0, + &job_args.scoped_temp_dir_creator, + ).await?; + + hifigan_superres_vocoder_model_fs_path + }; + +// // ==================== CONFIRM OR DOWNLOAD OPTIONAL CUSTOM VOCODER MODEL ==================== // + + let custom_vocoder_fs_path = match &model_record.maybe_custom_vocoder { + None => None, + Some(vocoder) => { + let custom_vocoder_fs_path = job_args.semi_persistent_cache.custom_vocoder_model_path(&vocoder.vocoder_token); + let custom_vocoder_object_path = job_args.bucket_path_unifier.vocoder_path(&vocoder.vocoder_private_bucket_hash); + + maybe_download_file_from_bucket( + "custom vocoder", + &custom_vocoder_fs_path, + &custom_vocoder_object_path, + &job_args.private_bucket_client, + &mut job_progress_reporter, + "downloading user vocoder", + job.id.0, + &job_args.scoped_temp_dir_creator, + ).await?; + + Some(custom_vocoder_fs_path) + } + }; + + // ==================== CONFIRM OR DOWNLOAD TTS SYNTHESIZER MODEL ==================== // + + let tts_synthesizer_fs_path = { + let tts_synthesizer_fs_path = job_args.semi_persistent_cache.tts_synthesizer_model_path(&model_record.model_token); + let tts_synthesizer_object_path = job_args.bucket_path_unifier.tts_synthesizer_path(&model_record.private_bucket_hash); + + maybe_download_file_from_bucket( + "synthesizer", + &tts_synthesizer_fs_path, + &tts_synthesizer_object_path, + &job_args.private_bucket_client, + &mut job_progress_reporter, + "downloading synthesizer", + job.id.0, + &job_args.scoped_temp_dir_creator, + ).await?; + + tts_synthesizer_fs_path + }; + + // ==================== Preprocess text ==================== // + + let cleaned_inference_text = clean_symbols(&job.raw_inference_text); + + // ==================== WRITE TEXT TO FILE ==================== // + + info!("Creating tempdir for inference results."); + + let temp_dir = format!("temp_tts_inference_{}", job.id.0); + + // NB: TempDir exists until it goes out of scope, at which point it should delete from filesystem. + let temp_dir = TempDir::new(&temp_dir) + .map_err(|e| ProcessSingleJobError::from_io_error(e))?; + + let text_input_fs_path = temp_dir.path().join("inference_input.txt"); + + std::fs::write(&text_input_fs_path, &cleaned_inference_text) + .map_err(|e| ProcessSingleJobError::from_io_error(e))?; + + // ==================== RUN INFERENCE ==================== // + + job_progress_reporter.log_status("running inference") + .map_err(|e| ProcessSingleJobError::Other(e))?; + + // TODO: Fix this. + let maybe_unload_model_path = job_args + .virtual_model_lfu + .insert_returning_replaced(tts_synthesizer_fs_path.to_str().unwrap_or("")) + .map_err(|e| ProcessSingleJobError::Other(e))?; + + if let Some(model_path) = maybe_unload_model_path.as_deref() { + warn!("Remove model from LFU cache: {:?}", model_path); + } + + let output_audio_fs_path = temp_dir.path().join("output.wav"); + let output_metadata_fs_path = temp_dir.path().join("metadata.json"); + let output_spectrogram_fs_path = temp_dir.path().join("spectrogram.json"); + + info!("Running TTS inference..."); + + info!("Expected output audio filename: {:?}", &output_audio_fs_path); + info!("Expected output spectrogram filename: {:?}", &output_spectrogram_fs_path); + info!("Expected output metadata filename: {:?}", &output_metadata_fs_path); + + if let Some(model_path) = maybe_unload_model_path.as_deref() { + warn!("Unload model from sidecar: {:?}", &model_path); + } + + let mut pretrained_vocoder = VocoderType::HifiGanSuperResolution; + if let Some(default_vocoder) = model_record.maybe_default_pretrained_vocoder.as_deref() { + pretrained_vocoder = VocoderType::from_str(default_vocoder) + .map_err(|e| ProcessSingleJobError::Other(e))?; + } + + info!("With pretrained vocoder: {:?}", pretrained_vocoder); + + let text_pipeline_type_or_guess = model_record.text_pipeline_type + .as_deref() + .and_then(|pipeline_type| + TextPipelineType::from_str(pipeline_type).ok())// NB: If there's an error deserializing, turn it to None. + .unwrap_or_else(|| + guess_text_pipeline_heuristic(Some(model_record.created_at))); + + info!("With text pipeline type `{:?} ` (or guess: {:?})", + &model_record.text_pipeline_type, + &text_pipeline_type_or_guess); + + let hifigan_vocoder_model_fs_path_to_use = match custom_vocoder_fs_path { + None => { + info!("using pretrained HiFi-GAN vocoder"); + pretrained_hifigan_vocoder_model_fs_path + }, + Some(custom_vocoder_fs_path) => { + info!("using custom user-trained HiFi-GAN vocoder: {:?}", custom_vocoder_fs_path); + custom_vocoder_fs_path + }, + }; + + // NB: Tacotron operates on decoder steps. 1000 steps is the default and correlates to + // roughly 12 seconds max. Here we map seconds to decoder steps. + let max_decoder_steps = seconds_to_decoder_steps(job.max_duration_seconds); + + job_args.http_clients.tts_inference_sidecar_client.request_inference( + &cleaned_inference_text, + max_decoder_steps, + &tts_synthesizer_fs_path, + pretrained_vocoder, + &text_pipeline_type_or_guess.to_str(), + &hifigan_vocoder_model_fs_path_to_use, + &hifigan_superres_vocoder_model_fs_path, + &waveglow_vocoder_model_fs_path, + &output_audio_fs_path, + &output_spectrogram_fs_path, + &output_metadata_fs_path, + maybe_unload_model_path) + .await + .map_err(|e| ProcessSingleJobError::Other(e))?; + + // ==================== CHECK ALL FILES EXIST AND GET METADATA ==================== // + + info!("Checking that output files exist..."); + + check_file_exists(&output_audio_fs_path).map_err(|e| ProcessSingleJobError::Other(e))?; + check_file_exists(&output_spectrogram_fs_path).map_err(|e| ProcessSingleJobError::Other(e))?; + check_file_exists(&output_metadata_fs_path).map_err(|e| ProcessSingleJobError::Other(e))?; + + let file_metadata = read_metadata_file(&output_metadata_fs_path) + .map_err(|e| ProcessSingleJobError::Other(e))?; + + safe_delete_temp_file(&output_metadata_fs_path); + + // ==================== UPLOAD AUDIO TO BUCKET ==================== // + + job_progress_reporter.log_status("uploading result") + .map_err(|e| ProcessSingleJobError::Other(e))?; + + let audio_result_object_path = job_args.bucket_path_unifier.tts_inference_wav_audio_output_path( + &job.uuid_idempotency_token); // TODO: Don't use this! + + info!("Audio destination bucket path: {:?}", &audio_result_object_path); + + info!("Uploading audio..."); + + job_args.public_bucket_client.upload_filename_with_content_type( + &audio_result_object_path, + &output_audio_fs_path, + "audio/wav") + .await + .map_err(|e| ProcessSingleJobError::Other(e))?; + + safe_delete_temp_file(&output_audio_fs_path); + + // ==================== UPLOAD SPECTROGRAM TO BUCKETS ==================== // + + let spectrogram_result_object_path = job_args.bucket_path_unifier.tts_inference_spectrogram_output_path( + &job.uuid_idempotency_token); // TODO: Don't use this! + + info!("Spectrogram destination bucket path: {:?}", &spectrogram_result_object_path); + + info!("Uploading spectrogram..."); + + job_args.public_bucket_client.upload_filename_with_content_type( + &spectrogram_result_object_path, + &output_spectrogram_fs_path, + "application/json") + .await + .map_err(|e| ProcessSingleJobError::Other(e))?; + + safe_delete_temp_file(&output_spectrogram_fs_path); + + // ==================== DELETE DOWNLOADED FILE ==================== // + + // NB: We should be using a tempdir, but to make absolutely certain we don't overflow the disk... + safe_delete_temp_directory(&temp_dir); + + // ==================== SAVE RECORDS ==================== // + + let text_hash = hash_string_sha2(&cleaned_inference_text) + .map_err(|e| ProcessSingleJobError::Other(e))?; + + let worker_name = job_args.get_worker_name(); + + info!("Saving tts inference record..."); + + let (id, inference_result_token) = insert_tts_result( + &job_args.mysql_pool, + job, + &text_hash, + pretrained_vocoder, + &audio_result_object_path, + &spectrogram_result_object_path, + file_metadata.file_size_bytes, + file_metadata.duration_millis.unwrap_or(0), + job_args.worker_details.is_on_prem, + &worker_name, + job_args.worker_details.is_debug_worker) + .await + .map_err(|e| ProcessSingleJobError::Other(e))?; + + info!("Marking job complete..."); + mark_tts_inference_job_done( + &job_args.mysql_pool, + job.id, + true, + Some(&inference_result_token), + &worker_name) + .await + .map_err(|e| ProcessSingleJobError::Other(e))?; + + info!("TTS Done. Original text was: {}", &job.raw_inference_text); + + job_args.firehose_publisher.tts_inference_finished( + job.maybe_creator_user_token.as_deref(), + &model_record.model_token, + &inference_result_token) + .await + .map_err(|e| { + error!("error publishing event: {:?}", e); + ProcessSingleJobError::Other(anyhow!("error publishing event")) + })?; + + job_progress_reporter.log_status("done") + .map_err(|e| ProcessSingleJobError::Other(e))?; + + info!("Job {:?} complete success! Downloaded, ran inference, and uploaded. Saved model record: {}, Result Token: {}", + job.id, id, &inference_result_token); + + let duration = start.elapsed(); + + since_creation_span.set_attribute("status", "success"); + since_creation_span.set_duration(duration); + + job_iteration_span.set_attribute("status", "success"); + job_iteration_span.set_duration(duration); + + Ok((since_creation_span, job_iteration_span)) +} + +fn get_timestamp_millis() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|d| d.as_millis() as u64) + .unwrap_or(0) +} diff --git a/crates/service/job/inference_job/src/job_steps/select_job_batch.rs b/crates/service/job/inference_job/src/job_steps/select_job_batch.rs new file mode 100644 index 0000000000..5bb0fcdb1d --- /dev/null +++ b/crates/service/job/inference_job/src/job_steps/select_job_batch.rs @@ -0,0 +1,23 @@ +use database_queries::queries::generic_inference::job::list_available_generic_inference_jobs::{list_available_generic_inference_jobs, ListAvailableGenericInferenceJobArgs}; +use enums::workers::generic_inference_type::GenericInferenceType; +use sqlx::MySqlPool; +use std::collections::BTreeSet; + +pub struct SelectJobBatchArgs <'a> { + pub num_records: u32, + pub is_debug_worker: bool, + pub sort_by_priority: bool, + pub maybe_scope_by_job_type: Option>, + pub mysql_pool: &'a MySqlPool, +} + +pub async fn select_job_batch(args: SelectJobBatchArgs<'_>) { + let _results = list_available_generic_inference_jobs(ListAvailableGenericInferenceJobArgs { + num_records: args.num_records, + is_debug_worker: args.is_debug_worker, + sort_by_priority: args.sort_by_priority, + maybe_scope_by_job_type: args.maybe_scope_by_job_type, + mysql_pool: args.mysql_pool, + }).await; + +} diff --git a/crates/service/job/inference_job/src/main.rs b/crates/service/job/inference_job/src/main.rs index 841fdbfb22..5ace5dea1c 100644 --- a/crates/service/job/inference_job/src/main.rs +++ b/crates/service/job/inference_job/src/main.rs @@ -38,7 +38,7 @@ use crate::http_clients::tts_sidecar_health_check_client::TtsSidecarHealthCheckC use crate::job_steps::job_dependencies::JobWorkerDetails; use crate::job_steps::job_dependencies::{JobDependencies, JobCaches, JobHttpClients}; use crate::job_steps::job_stats::JobStats; -use crate::job_steps::main_loop::main_loop; +use crate::job_steps::main_loop_old::main_loop_old; use crate::script_execution::tacotron_inference_command::TacotronInferenceCommand; use crate::util::scoped_temp_dir_creator::ScopedTempDirCreator; use database_queries::mediators::firehose_publisher::FirehosePublisher; @@ -328,7 +328,7 @@ async fn main() -> AnyhowResult<()> { maybe_minimum_priority, }; - main_loop(inferencer).await; + main_loop_old(inferencer).await; Ok(()) }