diff --git a/crates/service/tts_inference_job/src/job_steps/download_file_from_bucket.rs b/crates/service/tts_inference_job/src/job_steps/download_file_from_bucket.rs index f8d3386183..fd6f345f23 100644 --- a/crates/service/tts_inference_job/src/job_steps/download_file_from_bucket.rs +++ b/crates/service/tts_inference_job/src/job_steps/download_file_from_bucket.rs @@ -1,9 +1,8 @@ use std::path::Path; use log::{info, warn}; -use tempdir::TempDir; use jobs_common::redis_job_status_logger::RedisJobStatusLogger; use storage_buckets_common::bucket_client::BucketClient; -use crate::ProcessSingleJobError; +use crate::{ProcessSingleJobError, ScopedTempDirCreator}; // TODO(bt, 2022-07-15): Make a concrete type for bucket paths @@ -15,6 +14,7 @@ pub async fn maybe_download_file_from_bucket( redis_logger: &mut RedisJobStatusLogger<'_>, redis_status_update_description: &str, job_id: i64, + scoped_tempdir_creator: &ScopedTempDirCreator, ) -> Result<(), ProcessSingleJobError> { if file_path.exists() { @@ -32,7 +32,7 @@ pub async fn maybe_download_file_from_bucket( let temp_dir = format!("temp_download_{}", job_id); // NB: TempDir exists until it goes out of scope, at which point it should delete from filesystem. - let temp_dir = TempDir::new(&temp_dir) + let temp_dir = scoped_tempdir_creator.new_tempdir(&temp_dir) .map_err(|e| ProcessSingleJobError::from_io_error(e))?; let temp_path = temp_dir.path().join("download.part"); 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 16df80db56..eed8afad35 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 @@ -14,9 +14,12 @@ use memory_caching::multi_item_ttl_cache::MultiItemTtlCache; use storage_buckets_common::bucket_client::BucketClient; use storage_buckets_common::bucket_path_unifier::BucketPathUnifier; use crate::http_clients::tts_sidecar_health_check_client::TtsSidecarHealthCheckClient; +use crate::util::scoped_temp_dir_creator::ScopedTempDirCreator; pub struct JobArgs { pub download_temp_directory: PathBuf, + pub scoped_temp_dir_creator: ScopedTempDirCreator, + pub mysql_pool: MySqlPool, pub redis_pool: r2d2::Pool, diff --git a/crates/service/tts_inference_job/src/job_steps/process_single_job.rs b/crates/service/tts_inference_job/src/job_steps/process_single_job.rs index f7370360f1..83c12a0d41 100644 --- a/crates/service/tts_inference_job/src/job_steps/process_single_job.rs +++ b/crates/service/tts_inference_job/src/job_steps/process_single_job.rs @@ -113,6 +113,7 @@ pub async fn process_single_job( &mut redis_logger, "downloading vocoder (1 of 3)", job.id.0, + &inferencer.scoped_temp_dir_creator, ).await?; waveglow_vocoder_model_fs_path @@ -133,6 +134,7 @@ pub async fn process_single_job( &mut redis_logger, "downloading vocoder (2 of 3)", job.id.0, + &inferencer.scoped_temp_dir_creator, ).await?; hifigan_vocoder_model_fs_path @@ -153,6 +155,7 @@ pub async fn process_single_job( &mut redis_logger, "downloading vocoder (3 of 3)", job.id.0, + &inferencer.scoped_temp_dir_creator, ).await?; hifigan_superres_vocoder_model_fs_path @@ -191,6 +194,7 @@ pub async fn process_single_job( &mut redis_logger, "downloading synthesizer", job.id.0, + &inferencer.scoped_temp_dir_creator, ).await?; tts_synthesizer_fs_path diff --git a/crates/service/tts_inference_job/src/main.rs b/crates/service/tts_inference_job/src/main.rs index 25a4389457..36958fb6b1 100644 --- a/crates/service/tts_inference_job/src/main.rs +++ b/crates/service/tts_inference_job/src/main.rs @@ -18,6 +18,7 @@ pub mod caching; pub mod http_clients; pub mod job_steps; pub mod script_execution; +pub mod util; use clap::{App, Arg}; use config::common_env::CommonEnv; @@ -59,6 +60,7 @@ use std::path::PathBuf; use std::time::Duration; use storage_buckets_common::bucket_client::BucketClient; use storage_buckets_common::bucket_path_unifier::BucketPathUnifier; +use crate::util::scoped_temp_dir_creator::ScopedTempDirCreator; // Buckets (shared config) const ENV_ACCESS_KEY : &'static str = "ACCESS_KEY"; @@ -271,6 +273,7 @@ async fn main() -> AnyhowResult<()> { info!("Is debug worker? {}", is_debug_worker); let inferencer = JobArgs { + scoped_temp_dir_creator: ScopedTempDirCreator::for_directory(&temp_directory), download_temp_directory: temp_directory, mysql_pool, redis_pool, diff --git a/crates/service/tts_inference_job/src/util/mod.rs b/crates/service/tts_inference_job/src/util/mod.rs new file mode 100644 index 0000000000..db617bc657 --- /dev/null +++ b/crates/service/tts_inference_job/src/util/mod.rs @@ -0,0 +1 @@ +pub mod scoped_temp_dir_creator; \ No newline at end of file diff --git a/crates/service/tts_inference_job/src/util/scoped_temp_dir_creator.rs b/crates/service/tts_inference_job/src/util/scoped_temp_dir_creator.rs new file mode 100644 index 0000000000..6bbfc03427 --- /dev/null +++ b/crates/service/tts_inference_job/src/util/scoped_temp_dir_creator.rs @@ -0,0 +1,20 @@ +use std::path::{Path, PathBuf}; +use tempdir::TempDir; + +#[derive(Clone)] +pub struct ScopedTempDirCreator { + base_dir: PathBuf, +} + +impl ScopedTempDirCreator { + + pub fn for_directory>(base_dir: P) -> Self { + Self { + base_dir: PathBuf::from(&base_dir.as_ref()), + } + } + + pub fn new_tempdir(&self, name: &str) -> std::io::Result { + TempDir::new_in(&self.base_dir, name) + } +} \ No newline at end of file