From b9fed6a9339f3ef9fe6d89837f73751be2bdf11c Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Thu, 14 Jul 2022 20:03:32 -0400 Subject: [PATCH] Cache model lookups, removing some query load from DB --- .../tts_models/get_tts_model_for_inference.rs | 1 + .../src/job_steps/job_args.rs | 9 +++++ crates/service/tts_inference_job/src/main.rs | 38 ++++++++++++++++--- 3 files changed, 42 insertions(+), 6 deletions(-) diff --git a/crates/lib/database_queries/src/queries/tts/tts_models/get_tts_model_for_inference.rs b/crates/lib/database_queries/src/queries/tts/tts_models/get_tts_model_for_inference.rs index f5134cbb51..c12910499a 100644 --- a/crates/lib/database_queries/src/queries/tts/tts_models/get_tts_model_for_inference.rs +++ b/crates/lib/database_queries/src/queries/tts/tts_models/get_tts_model_for_inference.rs @@ -6,6 +6,7 @@ use sqlx; // TODO: Can probably just reuse another query. +#[derive(Clone)] pub struct TtsModelForInferenceRecord { pub model_token: String, pub tts_model_type: String, diff --git a/crates/service/tts_inference_job/src/job_steps/job_args.rs b/crates/service/tts_inference_job/src/job_steps/job_args.rs index cb680b8454..7ac49c896e 100644 --- a/crates/service/tts_inference_job/src/job_steps/job_args.rs +++ b/crates/service/tts_inference_job/src/job_steps/job_args.rs @@ -9,6 +9,8 @@ use r2d2_redis::RedisConnectionManager; use r2d2_redis::r2d2; use sqlx::MySqlPool; use std::path::PathBuf; +use database_queries::queries::tts::tts_models::get_tts_model_for_inference::TtsModelForInferenceRecord; +use memory_caching::multi_item_ttl_cache::MultiItemTtlCache; use storage_buckets_common::bucket_client::BucketClient; use storage_buckets_common::bucket_path_unifier::BucketPathUnifier; @@ -39,6 +41,9 @@ pub struct JobArgs { pub virtual_model_lfu: SyncVirtualLfuCache, pub cache_miss_strategizers: SyncMultiCacheMissStrategizer, + // In-process cache of database lookup records, etc. + pub caches: JobCaches, + // Waveglow vocoder filename pub waveglow_vocoder_model_filename: String, @@ -86,3 +91,7 @@ pub struct JobWorkerDetails { // Hostname of the worker. pub worker_hostname: String, } + +pub struct JobCaches { + pub tts_model_record_cache: MultiItemTtlCache, +} diff --git a/crates/service/tts_inference_job/src/main.rs b/crates/service/tts_inference_job/src/main.rs index db12006f71..28dd68ba78 100644 --- a/crates/service/tts_inference_job/src/main.rs +++ b/crates/service/tts_inference_job/src/main.rs @@ -31,7 +31,7 @@ use crate::caching::cache_miss_strategizer::CacheMissStrategy; use crate::caching::cache_miss_strategizer_multi::SyncMultiCacheMissStrategizer; use crate::caching::virtual_lfu_cache::SyncVirtualLfuCache; use crate::http_clients::tts_inference_sidecar_client::TtsInferenceSidecarClient; -use crate::job_steps::job_args::JobArgs; +use crate::job_steps::job_args::{JobArgs, JobCaches}; use crate::job_steps::job_args::JobWorkerDetails; use crate::job_steps::process_single_job::process_single_job; use crate::job_steps::process_single_job_error::ProcessSingleJobError; @@ -54,6 +54,7 @@ use sqlx::mysql::MySqlPoolOptions; use std::ops::Deref; use std::path::PathBuf; use std::time::Duration; +use memory_caching::multi_item_ttl_cache::MultiItemTtlCache; use storage_buckets_common::bucket_client::BucketClient; use storage_buckets_common::bucket_path_unifier::BucketPathUnifier; @@ -242,6 +243,10 @@ async fn main() -> AnyhowResult<()> { ) }; + let model_cache_duration = std::time::Duration::from_millis( + easyenv::get_env_num("TTS_MODEL_RECORD_CACHE_MILLIS", 300_000)?, // Five minutes + ); + let license_key = easyenv::get_env_string_required("NEWRELIC_API_KEY")?; let newrelic_disabled = easyenv::get_env_bool_or_default("IS_NEWRELIC_DISABLED", false); @@ -275,6 +280,9 @@ async fn main() -> AnyhowResult<()> { worker_hostname: server_hostname.clone(), is_debug_worker, }, + caches: JobCaches { + tts_model_record_cache: MultiItemTtlCache::create_with_duration(model_cache_duration), + }, virtual_model_lfu: virtual_lfu_cache, cache_miss_strategizers, bucket_path_unifier: BucketPathUnifier::default_paths(), @@ -402,6 +410,7 @@ async fn process_jobs( 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; @@ -583,15 +592,32 @@ impl ModelState { 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 { - info!("Looking up TTS model by token: {}", &job.model_token); + // 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 = get_tts_model_for_inference( - &mysql_pool, - &job.model_token - ).await?; + 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);