diff --git a/src/tts_inference_job.rs b/src/tts_inference_job.rs index aded992df9..20178b84ad 100644 --- a/src/tts_inference_job.rs +++ b/src/tts_inference_job.rs @@ -19,7 +19,7 @@ pub mod script_execution; pub mod shared_constants; pub mod util; -use anyhow::anyhow; +use anyhow::{anyhow, Error}; use chrono::{Utc, DateTime, TimeZone}; use crate::clients::tts_inference_sidecar_client::TtsInferenceSidecarClient; use crate::common_env::CommonEnv; @@ -42,6 +42,9 @@ use crate::util::filesystem::check_directory_exists; use crate::util::filesystem::check_file_exists; use crate::util::hashing::hash_file_sha2::hash_file_sha2; use crate::util::hashing::hash_string_sha2::hash_string_sha2; +use crate::util::jobs::cache_miss_strategizer::CacheMissStrategizer; +use crate::util::jobs::cache_miss_strategizer::CacheMissStrategy; +use crate::util::jobs::cache_miss_strategizer_multi::SyncMultiCacheMissStrategizer; use crate::util::jobs::virtual_lfu_cache::SyncVirtualLfuCache; use crate::util::noop_logger::NoOpLogger; use crate::util::random_crockford_token::random_crockford_token; @@ -98,6 +101,7 @@ struct Inferencer { // Keep tabs of which models to hold in the sidecar memory with this virtual LRU cache pub virtual_model_lfu: SyncVirtualLfuCache, + pub cache_miss_strategizers: SyncMultiCacheMissStrategizer, // Command to run pub tts_vocoder_model_filename: String, @@ -220,6 +224,23 @@ async fn main() -> AnyhowResult<()> { let virtual_lfu_cache = SyncVirtualLfuCache::new(2)?; + let cache_miss_strategizers = { + let in_memory_strategizer = CacheMissStrategizer::new( + chrono::Duration::milliseconds(5_000), + chrono::Duration::milliseconds(60_000), + ); + + let on_disk_strategizer = CacheMissStrategizer::new( + chrono::Duration::milliseconds(20_000), + chrono::Duration::milliseconds(120_000), + ); + + SyncMultiCacheMissStrategizer::new( + in_memory_strategizer, + on_disk_strategizer, + ) + }; + let inferencer = Inferencer { download_temp_directory: temp_directory, mysql_pool, @@ -231,6 +252,7 @@ async fn main() -> AnyhowResult<()> { tts_inference_command, tts_inference_sidecar_client, virtual_model_lfu: virtual_lfu_cache, + cache_miss_strategizers, bucket_path_unifier: BucketPathUnifier::default_paths(), semi_persistent_cache, firehose_publisher, @@ -255,8 +277,6 @@ async fn main_loop(inferencer: Inferencer) { let mut noop_logger = NoOpLogger::new(inferencer.no_op_logger_millis as i64); - let mut cold_cache_okay = false; - loop { let num_records = inferencer.job_batch_size; @@ -285,28 +305,20 @@ async fn main_loop(inferencer: Inferencer) { let batch_result = process_jobs( &inferencer, jobs, - cold_cache_okay).await; + ).await; - let cold_cache_skipped_count = match batch_result { - Ok(cold_cache_skipped_count) => cold_cache_skipped_count, + match batch_result { + Ok(_) => {}, Err(e) => { warn!("Error querying jobs: {:?}", e); std::thread::sleep(Duration::from_millis(error_timeout_millis)); error_timeout_millis += INCREASE_TIMEOUT_MILLIS; continue; } - }; + } error_timeout_millis = START_TIMEOUT_MILLIS; // reset - // We mostly want to ignore a cold cache unless other workers - // aren't picking up the slack. - if cold_cache_okay { - cold_cache_okay = false; - } else if cold_cache_skipped_count > 0 { - cold_cache_okay = true; - } - std::thread::sleep(Duration::from_millis(inferencer.job_batch_wait_millis)); } } @@ -315,9 +327,7 @@ async fn main_loop(inferencer: Inferencer) { async fn process_jobs( inferencer: &Inferencer, jobs: Vec, - cold_cache_okay: bool -) -> AnyhowResult { - let mut skipped_for_cold_cache = 0; +) -> AnyhowResult<()> { for job in jobs.into_iter() { let model_state_result = ModelState::query_model_and_check_filesystem( @@ -344,18 +354,37 @@ async fn process_jobs( }; if !model_state.is_downloaded_to_filesystem || !model_state.is_in_memory_cache { - warn!("Model isn't in cache (downloaded = {}), (in memory = {})", + warn!("Model isn't ready (downloaded = {}), (in memory = {})", model_state.is_downloaded_to_filesystem, model_state.is_in_memory_cache); - if !cold_cache_okay { - // 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); + 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) + }; - skipped_for_cold_cache += 1; - continue; + match maybe_strategy { + Err(err) => { + warn!("Failure to process job: {:?}", err); + let failure_reason = ""; + let _r = mark_tts_inference_job_failure( + &inferencer.mysql_pool, + &job, + failure_reason, + inferencer.job_max_attempts + ).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... } } @@ -375,7 +404,7 @@ async fn process_jobs( } } - Ok(skipped_for_cold_cache) + Ok(()) } #[derive(Deserialize, Default)] diff --git a/src/util/jobs/cache_miss_strategizer.rs b/src/util/jobs/cache_miss_strategizer.rs index b8ce6d1ecd..333fdbbda7 100644 --- a/src/util/jobs/cache_miss_strategizer.rs +++ b/src/util/jobs/cache_miss_strategizer.rs @@ -1,6 +1,7 @@ use anyhow::anyhow; use chrono::{DateTime, Utc, Duration}; use std::collections::HashMap; +use std::hash::Hash; /// How to handle cache miss events, as dictated by the `CacheMissStrategizer`. #[derive(Eq, PartialEq, Copy, Clone, Debug)] @@ -24,7 +25,7 @@ struct CacheMissLog { /// /// Keep track of a cache and tell the caller when to proceed with work. /// If it isn't time yet, skip or wait and let another consumer handle it. -pub struct CacheMissStrategizer { +pub struct CacheMissStrategizer { /// How long it takes after the first cache miss to do work. max_cold_duration: Duration, @@ -34,7 +35,7 @@ pub struct CacheMissStrategizer { forget_duration: Duration, /// Log of all cache misses. - cache_miss_log: HashMap, + cache_miss_log: HashMap, /// Clock can be replaced for testing. get_time_function: Box DateTime>, @@ -56,7 +57,7 @@ impl CacheMissLog { } } -impl CacheMissStrategizer { +impl CacheMissStrategizer { pub fn new(max_cold_duration: Duration, forget_duration: Duration) -> Self { Self { max_cold_duration, @@ -81,7 +82,7 @@ impl CacheMissStrategizer { // NB: Not threadsafe due to multiple operations against hashes! #[must_use] - pub fn cache_miss(&mut self, id: i64) -> CacheMissStrategy { + pub fn cache_miss(&mut self, id: T) -> CacheMissStrategy { let now : DateTime = (self.get_time_function)(); if let Some((_id, cache_miss_log)) = self.cache_miss_log.get_key_value(&id) { diff --git a/src/util/jobs/cache_miss_strategizer_multi.rs b/src/util/jobs/cache_miss_strategizer_multi.rs index 4c91f051a7..554be524cd 100644 --- a/src/util/jobs/cache_miss_strategizer_multi.rs +++ b/src/util/jobs/cache_miss_strategizer_multi.rs @@ -7,8 +7,8 @@ use std::sync::{RwLock, Arc}; /// Keep track of multiple caches, each with different time penalties. /// (This isn't strictly necessary) pub struct MultiCacheMissStrategizer { - in_memory_log: CacheMissStrategizer, - on_disk_log: CacheMissStrategizer, + in_memory_log: CacheMissStrategizer, + on_disk_log: CacheMissStrategizer, // ... future caches ? } @@ -19,7 +19,10 @@ pub struct SyncMultiCacheMissStrategizer { } impl SyncMultiCacheMissStrategizer { - pub fn new(in_memory_log: CacheMissStrategizer, on_disk_log: CacheMissStrategizer) -> Self { + pub fn new( + in_memory_log: CacheMissStrategizer, + on_disk_log: CacheMissStrategizer, + ) -> Self { let multi_cache_log = MultiCacheMissStrategizer::new( in_memory_log, on_disk_log, @@ -30,26 +33,29 @@ impl SyncMultiCacheMissStrategizer { } #[must_use] - pub fn memory_cache_miss(&self, id: i64) -> AnyhowResult { + pub fn memory_cache_miss(&self, token: &str) -> AnyhowResult { let result = self.multi_cache_log .write() - .map(|mut cache_logs| cache_logs.memory_cache_miss(id)) + .map(|mut cache_logs| cache_logs.memory_cache_miss(token)) .map_err(|err| anyhow!("mutex error: {:?}", err))?; Ok(result) } #[must_use] - pub fn disk_cache_miss(&self, id: i64) -> AnyhowResult { + pub fn disk_cache_miss(&self, token: &str) -> AnyhowResult { let result = self.multi_cache_log .write() - .map(|mut cache_logs| cache_logs.disk_cache_miss(id)) + .map(|mut cache_logs| cache_logs.disk_cache_miss(token)) .map_err(|err| anyhow!("mutex error: {:?}", err))?; Ok(result) } } impl MultiCacheMissStrategizer { - pub fn new(in_memory_log: CacheMissStrategizer, on_disk_log: CacheMissStrategizer) -> Self { + pub fn new( + in_memory_log: CacheMissStrategizer, + on_disk_log: CacheMissStrategizer, + ) -> Self { Self { in_memory_log, on_disk_log, @@ -57,12 +63,12 @@ impl MultiCacheMissStrategizer { } #[must_use] - pub fn memory_cache_miss(&mut self, id: i64) -> CacheMissStrategy { - self.in_memory_log.cache_miss(id) + pub fn memory_cache_miss(&mut self, token: &str) -> CacheMissStrategy { + self.in_memory_log.cache_miss(token.to_string()) } #[must_use] - pub fn disk_cache_miss(&mut self, id: i64) -> CacheMissStrategy { - self.on_disk_log.cache_miss(id) + pub fn disk_cache_miss(&mut self, token: &str) -> CacheMissStrategy { + self.on_disk_log.cache_miss(token.to_string()) } }