rework where temp storage lives: download vs workspace

crates/lib/jobs_common/src/semi_persistent_cache_dir.rs

@@ -29,8 +29,8 @@ impl SemiPersistentCacheDir {
     Self::configured_root("/file_cache/")
   }
 
-  pub fn configured_root(root_path: &str) -> Self {
-    let cache_root = PathBuf::from(root_path);
+  pub fn configured_root<P: AsRef<Path>>(root_path: P) -> Self {
+    let cache_root = root_path.as_ref().to_path_buf();
     Self {
       cache_root: cache_root.clone(),
 

crates/service/job/inference_job/src/job/job_loop/determine_dependency_status.rs

@@ -42,6 +42,7 @@ pub async fn determine_dependency_status(job_dependencies: &JobDependencies, job
         TtsModelType::Tacotron2 => {
           maybe_model_token = Some(model.model_token.to_string());
           Some(job_dependencies
+              .fs
               .semi_persistent_cache
               .tts_synthesizer_model_path(model.model_token.as_str()))
         }
@@ -58,6 +59,7 @@ pub async fn determine_dependency_status(job_dependencies: &JobDependencies, job
         VoiceConversionModelType::SoVitsSvc => {
           maybe_model_token = Some(model.token.to_string());
           Some(job_dependencies
+              .fs
               .semi_persistent_cache
               .voice_conversion_model_path(model.token.as_str()))
         }

crates/service/job/inference_job/src/job/job_loop/main_loop.rs

@@ -123,7 +123,7 @@ async fn process_job_batch(job_dependencies: &JobDependencies, jobs: Vec<Availab
           // Post failure handling
           ProcessSingleJobError::FilesystemFull => {
             warn!("Clearing full filesystem...");
-            clear_full_filesystem(&job_dependencies.semi_persistent_cache)?;
+            clear_full_filesystem(&job_dependencies.fs.semi_persistent_cache)?;
           }
           // No-op
           ProcessSingleJobError::Other(_) => {}

crates/service/job/inference_job/src/job/job_loop/process_single_job.rs

@@ -124,7 +124,7 @@ pub async fn process_single_job(
   // ==================== SETUP TEMP DIRS ==================== //
 
   let temp_dir = format!("temp_{}", job.id.0);
-  let temp_dir = job_dependencies.scoped_temp_dir_creator.new_tempdir(&temp_dir)
+  let temp_dir = job_dependencies.fs.scoped_temp_dir_creator.new_tempdir(&temp_dir)
       .map_err(|err| ProcessSingleJobError::Other(anyhow!("filesystem error: {:?}", err)))?;
 
   let _p = temp_dir.path(); // TODO: Just so the build doesn't complain about unused. Remove.

crates/service/job/inference_job/src/job/job_types/tts/tacotron2_v2_early_fakeyou/process_job.rs

@@ -53,7 +53,7 @@ pub async fn process_job(args: ProcessJobArgs<'_>) -> Result<JobSuccessResult, P
 
   let waveglow_vocoder_model_fs_path = {
     let waveglow_vocoder_model_filename = args.job_dependencies.job_type_details.tacotron2_old_vocodes.waveglow_vocoder_model_filename.clone();
-    let waveglow_vocoder_model_fs_path = args.job_dependencies.semi_persistent_cache.tts_pretrained_vocoder_model_path(&waveglow_vocoder_model_filename);
+    let waveglow_vocoder_model_fs_path = args.job_dependencies.fs.semi_persistent_cache.tts_pretrained_vocoder_model_path(&waveglow_vocoder_model_filename);
     let waveglow_vocoder_model_object_path = args.job_dependencies.bucket_path_unifier.tts_pretrained_vocoders_path(&waveglow_vocoder_model_filename);
 
     maybe_download_file_from_bucket(
@@ -64,7 +64,7 @@ pub async fn process_job(args: ProcessJobArgs<'_>) -> Result<JobSuccessResult, P
       &mut job_progress_reporter,
       "downloading vocoder (1 of 3)",
       job.id.0,
-      &args.job_dependencies.scoped_temp_dir_creator,
+      &args.job_dependencies.fs.scoped_temp_dir_creator,
     ).await?;
 
     waveglow_vocoder_model_fs_path
@@ -74,7 +74,7 @@ pub async fn process_job(args: ProcessJobArgs<'_>) -> Result<JobSuccessResult, P
 
   let pretrained_hifigan_vocoder_model_fs_path = {
     let hifigan_vocoder_model_filename = args.job_dependencies.job_type_details.tacotron2_old_vocodes.hifigan_vocoder_model_filename.clone();
-    let hifigan_vocoder_model_fs_path = args.job_dependencies.semi_persistent_cache.tts_pretrained_vocoder_model_path(&hifigan_vocoder_model_filename);
+    let hifigan_vocoder_model_fs_path = args.job_dependencies.fs.semi_persistent_cache.tts_pretrained_vocoder_model_path(&hifigan_vocoder_model_filename);
     let hifigan_vocoder_model_object_path = args.job_dependencies.bucket_path_unifier.tts_pretrained_vocoders_path(&hifigan_vocoder_model_filename);
 
     maybe_download_file_from_bucket(
@@ -85,7 +85,7 @@ pub async fn process_job(args: ProcessJobArgs<'_>) -> Result<JobSuccessResult, P
       &mut job_progress_reporter,
       "downloading vocoder (2 of 3)",
       job.id.0,
-      &args.job_dependencies.scoped_temp_dir_creator,
+      &args.job_dependencies.fs.scoped_temp_dir_creator,
     ).await?;
 
     hifigan_vocoder_model_fs_path
@@ -95,7 +95,7 @@ pub async fn process_job(args: ProcessJobArgs<'_>) -> Result<JobSuccessResult, P
 
   let hifigan_superres_vocoder_model_fs_path = {
     let hifigan_superres_vocoder_model_filename = args.job_dependencies.job_type_details.tacotron2_old_vocodes.hifigan_superres_vocoder_model_filename.clone();
-    let hifigan_superres_vocoder_model_fs_path = args.job_dependencies.semi_persistent_cache.tts_pretrained_vocoder_model_path(&hifigan_superres_vocoder_model_filename);
+    let hifigan_superres_vocoder_model_fs_path = args.job_dependencies.fs.semi_persistent_cache.tts_pretrained_vocoder_model_path(&hifigan_superres_vocoder_model_filename);
     let hifigan_superres_vocoder_model_object_path = args.job_dependencies.bucket_path_unifier.tts_pretrained_vocoders_path(&hifigan_superres_vocoder_model_filename);
 
     maybe_download_file_from_bucket(
@@ -106,7 +106,7 @@ pub async fn process_job(args: ProcessJobArgs<'_>) -> Result<JobSuccessResult, P
       &mut job_progress_reporter,
       "downloading vocoder (3 of 3)",
       job.id.0,
-      &args.job_dependencies.scoped_temp_dir_creator,
+      &args.job_dependencies.fs.scoped_temp_dir_creator,
     ).await?;
 
     hifigan_superres_vocoder_model_fs_path
@@ -117,7 +117,7 @@ pub async fn process_job(args: ProcessJobArgs<'_>) -> Result<JobSuccessResult, P
   let custom_vocoder_fs_path = match &tts_model.maybe_custom_vocoder {
     None => None,
     Some(vocoder) => {
-      let custom_vocoder_fs_path = args.job_dependencies.semi_persistent_cache.custom_vocoder_model_path(&vocoder.vocoder_token);
+      let custom_vocoder_fs_path = args.job_dependencies.fs.semi_persistent_cache.custom_vocoder_model_path(&vocoder.vocoder_token);
       let custom_vocoder_object_path  = args.job_dependencies.bucket_path_unifier.vocoder_path(&vocoder.vocoder_private_bucket_hash);
 
       maybe_download_file_from_bucket(
@@ -128,7 +128,7 @@ pub async fn process_job(args: ProcessJobArgs<'_>) -> Result<JobSuccessResult, P
         &mut job_progress_reporter,
         "downloading user vocoder",
         job.id.0,
-        &args.job_dependencies.scoped_temp_dir_creator,
+        &args.job_dependencies.fs.scoped_temp_dir_creator,
       ).await?;
 
       Some(custom_vocoder_fs_path)
@@ -138,7 +138,7 @@ pub async fn process_job(args: ProcessJobArgs<'_>) -> Result<JobSuccessResult, P
   // ==================== CONFIRM OR DOWNLOAD TTS SYNTHESIZER MODEL ==================== //
 
   let tts_synthesizer_fs_path = {
-    let tts_synthesizer_fs_path = args.job_dependencies.semi_persistent_cache.tts_synthesizer_model_path(tts_model.model_token.as_str());
+    let tts_synthesizer_fs_path = args.job_dependencies.fs.semi_persistent_cache.tts_synthesizer_model_path(tts_model.model_token.as_str());
     let tts_synthesizer_object_path  = args.job_dependencies.bucket_path_unifier.tts_synthesizer_path(&tts_model.private_bucket_hash);
 
     maybe_download_file_from_bucket(
@@ -149,7 +149,7 @@ pub async fn process_job(args: ProcessJobArgs<'_>) -> Result<JobSuccessResult, P
       &mut job_progress_reporter,
       "downloading synthesizer",
       job.id.0,
-      &args.job_dependencies.scoped_temp_dir_creator,
+      &args.job_dependencies.fs.scoped_temp_dir_creator,
     ).await?;
 
     tts_synthesizer_fs_path

crates/service/job/inference_job/src/job/job_types/tts/vits/process_job.rs

@@ -59,7 +59,7 @@ pub async fn process_job(args: VitsProcessJobArgs<'_>) -> Result<JobSuccessResul
   // ==================== CONFIRM OR DOWNLOAD VITS SYNTHESIZER MODEL ==================== //
 
   let vits_traced_synthesizer_fs_path = {
-    let vits_traced_synthesizer_fs_path = args.job_dependencies.semi_persistent_cache.tts_synthesizer_model_path(tts_model.model_token.as_str());
+    let vits_traced_synthesizer_fs_path = args.job_dependencies.fs.semi_persistent_cache.tts_synthesizer_model_path(tts_model.model_token.as_str());
 
     // NB: We're using traced models, not the original model files.
     // We generate these at time of upload from the original model files.
@@ -74,7 +74,7 @@ pub async fn process_job(args: VitsProcessJobArgs<'_>) -> Result<JobSuccessResul
       &mut job_progress_reporter,
       "downloading synthesizer",
       job.id.0,
-      &args.job_dependencies.scoped_temp_dir_creator,
+      &args.job_dependencies.fs.scoped_temp_dir_creator,
     ).await?;
 
     vits_traced_synthesizer_fs_path

crates/service/job/inference_job/src/job/job_types/vc/so_vits_svc/process_job.rs

@@ -57,9 +57,9 @@ pub async fn process_job(args: SoVitsSvcProcessJobArgs<'_>) -> Result<JobSuccess
   // ==================== CONFIRM OR DOWNLOAD SO-VITS-SVC SYNTHESIZER MODEL ==================== //
 
   let so_vits_svc_fs_path = {
-    let so_vits_svc_fs_path = args.job_dependencies.semi_persistent_cache.voice_conversion_model_path(vc_model.token.as_str());
+    let so_vits_svc_fs_path = args.job_dependencies.fs.semi_persistent_cache.voice_conversion_model_path(vc_model.token.as_str());
 
-    create_dir_all_if_missing(args.job_dependencies.semi_persistent_cache.voice_conversion_model_directory())
+    create_dir_all_if_missing(args.job_dependencies.fs.semi_persistent_cache.voice_conversion_model_directory())
         .map_err(|e| {
           error!("could not create model storage directory: {:?}", e);
           ProcessSingleJobError::from_io_error(e)
@@ -75,7 +75,7 @@ pub async fn process_job(args: SoVitsSvcProcessJobArgs<'_>) -> Result<JobSuccess
       &mut job_progress_reporter,
       "downloading so-vits-svc model",
       job.id.0,
-      &args.job_dependencies.scoped_temp_dir_creator,
+      &args.job_dependencies.fs.scoped_temp_dir_creator,
     ).await?;
 
     so_vits_svc_fs_path
@@ -126,7 +126,7 @@ pub async fn process_job(args: SoVitsSvcProcessJobArgs<'_>) -> Result<JobSuccess
       &mut job_progress_reporter,
       "downloading",
       job.id.0,
-      &args.job_dependencies.scoped_temp_dir_creator,
+      &args.job_dependencies.fs.scoped_temp_dir_creator,
     ).await?;
 
     original_media_upload_fs_path

crates/service/job/inference_job/src/job_dependencies.rs

@@ -21,8 +21,8 @@ use std::path::PathBuf;
 use memory_caching::ttl_key_counter::TtlKeyCounter;
 
 pub struct JobDependencies {
-  pub download_temp_directory: PathBuf,
-  pub scoped_temp_dir_creator: ScopedTempDirCreator,
+  /// Filesystem info and utils
+  pub fs: FileSystemDetails,
 
   pub mysql_pool: MySqlPool,
 
@@ -43,7 +43,6 @@ pub struct JobDependencies {
   pub firehose_publisher: FirehosePublisher,
 
   pub bucket_path_unifier: BucketPathUnifier,
-  pub semi_persistent_cache: SemiPersistentCacheDir,
 
   pub job_stats: JobStats,
 
@@ -90,6 +89,24 @@ pub struct JobDependencies {
   pub container_db: ContainerEnvironmentArg, // Same info, but for database.
 }
 
+pub struct FileSystemDetails {
+  /// Temporary directory for storing downloads
+  /// (In prod, typically model files from GCS / NFS PVC mount)
+  pub temp_directory_downloads: PathBuf,
+
+  /// Temporary directory for storing work
+  /// (In prod, typically python inference outputs / emptyDir mount)
+  pub temp_directory_work: PathBuf,
+
+  // TODO: Rename
+  /// Creates workspaces for scratch files
+  pub scoped_temp_dir_creator: ScopedTempDirCreator,
+
+  // TODO: Rename
+  /// Organizes downloaded files
+  pub semi_persistent_cache: SemiPersistentCacheDir,
+}
+
 pub struct JobWorkerDetails {
   // Debug workers only process special debug requests. They're silent otherwise.
   // Non-debug workers ignore debug requests. This is so we can deploy special code

crates/service/job/inference_job/src/main.rs

@@ -35,14 +35,16 @@ use crate::job::job_loop::main_loop::main_loop;
 use crate::job::job_types::tts::tacotron2_v2_early_fakeyou::tacotron2_inference_command::Tacotron2InferenceCommand;
 use crate::job::job_types::tts::vits::vits_inference_command::VitsInferenceCommand;
 use crate::job::job_types::vc::so_vits_svc::so_vits_svc_inference_command::SoVitsSvcInferenceCommand;
-use crate::job_dependencies::{JobCaches, JobDependencies, JobTypeDetails, JobWorkerDetails, SoVitsSvcDetails, Tacotron2VocodesDetails, VitsDetails};
+use crate::job_dependencies::{FileSystemDetails, JobCaches, JobDependencies, JobTypeDetails, JobWorkerDetails, SoVitsSvcDetails, Tacotron2VocodesDetails, VitsDetails};
 use crate::util::scoped_temp_dir_creator::ScopedTempDirCreator;
+use filesys::create_dir_all_if_missing::create_dir_all_if_missing;
 use jobs_common::job_progress_reporter::job_progress_reporter::JobProgressReporterBuilder;
 use jobs_common::job_progress_reporter::noop_job_progress_reporter::NoOpJobProgressReporterBuilder;
 use jobs_common::job_progress_reporter::redis_job_progress_reporter::RedisJobProgressReporterBuilder;
 use jobs_common::semi_persistent_cache_dir::SemiPersistentCacheDir;
 use log::{error, info, warn};
 use memory_caching::multi_item_ttl_cache::MultiItemTtlCache;
+use memory_caching::ttl_key_counter::TtlKeyCounter;
 use mysql_queries::common_inputs::container_environment_arg::ContainerEnvironmentArg;
 use mysql_queries::mediators::firehose_publisher::FirehosePublisher;
 use newrelic_telemetry::ClientBuilder;
@@ -51,7 +53,6 @@ use r2d2_redis::r2d2;
 use sqlx::mysql::MySqlPoolOptions;
 use std::path::PathBuf;
 use std::time::Duration;
-use memory_caching::ttl_key_counter::TtlKeyCounter;
 use subprocess_common::docker_options::{DockerEnvVar, DockerFilesystemMount, DockerGpu, DockerOptions};
 
 // Buckets (shared config)
@@ -63,14 +64,9 @@ const ENV_REGION_NAME : &'static str = "REGION_NAME";
 const ENV_PRIVATE_BUCKET_NAME : &'static str = "PRIVATE_BUCKET_NAME";
 const ENV_PUBLIC_BUCKET_NAME : &'static str = "PUBLIC_BUCKET_NAME";
 
-// Where models and other assets get downloaded to.
-const ENV_SEMIPERSISTENT_CACHE_DIR : &'static str = "SEMIPERSISTENT_CACHE_DIR";
-
 // HTTP sidecar
 const ENV_TTS_INFERENCE_SIDECAR_HOSTNAME: &'static str = "TTS_INFERENCE_SIDECAR_HOSTNAME";
 
-const DEFAULT_TEMP_DIR: &'static str = "/tmp";
-
 #[tokio::main]
 async fn main() -> AnyhowResult<()> {
 
@@ -116,13 +112,35 @@ async fn main() -> AnyhowResult<()> {
     Some(bucket_timeout),
   )?;
 
-  let temp_directory = easyenv::get_env_string_or_default(
-    "DOWNLOAD_TEMP_DIR",
-    DEFAULT_TEMP_DIR);
+  // Where we download models and resources to (typically a shared NFS volume in prod).
+  let temp_directory_downloads = easyenv::get_env_pathbuf_or_default(
+    "TEMP_DIR_DOWNLOADS",
+    PathBuf::from("/tmp/downloads")
+  );
 
-  let temp_directory = PathBuf::from(temp_directory);
+  // Where we store scratch files for workloads (temporary processing).
+  let temp_directory_work = easyenv::get_env_pathbuf_or_default(
+    "TEMP_DIR_WORK",
+    PathBuf::from("/tmp/work")
+  );
 
-  check_directory_exists(&temp_directory)?;
+  if !container_environment.server_environment.is_deployed_in_production() {
+    warn!("Creating temporary directories for non-production / development only!");
+    create_dir_all_if_missing(&temp_directory_downloads)?;
+    create_dir_all_if_missing(&temp_directory_work)?;
+  }
+
+  check_directory_exists(&temp_directory_downloads)?;
+  check_directory_exists(&temp_directory_work)?;
+
+  let semi_persistent_cache =
+      SemiPersistentCacheDir::configured_root(&temp_directory_downloads);
+
+  info!("Creating pod semi-persistent cache dirs...");
+  semi_persistent_cache.create_custom_vocoder_model_path()?;
+  semi_persistent_cache.create_tts_pretrained_vocoder_model_path()?;
+  semi_persistent_cache.create_tts_synthesizer_model_path()?;
+  semi_persistent_cache.create_voice_conversion_model_path()?;
 
   let db_connection_string =
       easyenv::get_env_string_or_default(
@@ -138,18 +156,6 @@ async fn main() -> AnyhowResult<()> {
 
   let common_env = CommonEnv::read_from_env()?;
 
-  let persistent_cache_path = easyenv::get_env_string_or_default(
-    ENV_SEMIPERSISTENT_CACHE_DIR,
-    "/tmp");
-
-  let semi_persistent_cache =
-      SemiPersistentCacheDir::configured_root(&persistent_cache_path);
-
-  info!("Creating pod semi-persistent cache dirs...");
-  semi_persistent_cache.create_tts_synthesizer_model_path()?;
-  semi_persistent_cache.create_tts_pretrained_vocoder_model_path()?;
-  semi_persistent_cache.create_custom_vocoder_model_path()?;
-
   let waveglow_vocoder_model_filename = easyenv::get_env_string_or_default(
     "TTS_WAVEGLOW_VOCODER_MODEL_FILENAME", "waveglow.pth");
 
@@ -217,8 +223,12 @@ async fn main() -> AnyhowResult<()> {
       };
 
   let job_dependencies = JobDependencies {
-    scoped_temp_dir_creator: ScopedTempDirCreator::for_directory(&temp_directory),
-    download_temp_directory: temp_directory,
+    fs: FileSystemDetails {
+      temp_directory_downloads: temp_directory_downloads.clone(),
+      temp_directory_work,
+      scoped_temp_dir_creator: ScopedTempDirCreator::for_directory(&temp_directory_downloads),
+      semi_persistent_cache,
+    },
     mysql_pool,
     maybe_redis_pool: None, // TODO(bt, 2023-01-11): See note in JobDependencies
     maybe_keepalive_redis_pool,
@@ -254,7 +264,6 @@ async fn main() -> AnyhowResult<()> {
     cold_filesystem_cache_starvation_threshold:
       easyenv::get_env_num("COLD_FILESYSTEM_CACHE_STARVATION_THRESHOLD", 3)?,
     bucket_path_unifier: BucketPathUnifier::default_paths(),
-    semi_persistent_cache,
     firehose_publisher,
     job_batch_wait_millis: common_env.job_batch_wait_millis,
     job_max_attempts: common_env.job_max_attempts as u16,