filesystem path tool, tests pass again

src/buckets/bucket_path_unifier.rs

@@ -1,15 +1,35 @@
 use crate::buckets::bucket_paths::hash_to_bucket_path;
 use std::path::PathBuf;
+use crate::util::anyhow_result::AnyhowResult;
 
 /// This is designed to make it centrally configurable where
 /// different types of objects are stored.
 pub struct BucketPathUnifier {
-  pub user_uploaded_audio_for_w2l_root: PathBuf,
   pub user_uploaded_w2l_templates_root: PathBuf,
+  pub user_uploaded_audio_for_w2l_root: PathBuf,
+  pub tts_inference_output_root: PathBuf,
+  pub w2l_inference_output_root: PathBuf,
 }
 
 impl BucketPathUnifier {
 
+  // TODO
+  //pub fn from_env_vars() -> AnyhowResult<Self> {
+  //  Ok(Self {
+  //    user_uploaded_audio_for_w2l_root: ,
+  //    user_uploaded_w2l_templates_root: ,
+  //  })
+  //}
+
+  pub fn default_paths() -> Self {
+    Self {
+      user_uploaded_audio_for_w2l_root: PathBuf::from("/user_uploaded_w2l_audio"),
+      user_uploaded_w2l_templates_root: PathBuf::from("/user_uploaded_w2l_templates"),
+      tts_inference_output_root: PathBuf::from("/tts_inference_output"),
+      w2l_inference_output_root: PathBuf::from("/w2l_inference_output"),
+    }
+  }
+
   pub fn end_bump_video_for_w2l_path(&self, end_bump_filename: &str) -> String {
     "".to_string()
   }
@@ -34,10 +54,20 @@ impl BucketPathUnifier {
       .join(audio_uuid)
   }
 
-  pub fn pretrained_w2l_model_path(&self, model_name: &str) -> String {
-    "".to_string()
+  // W2L inference output videos
+  pub fn w2l_inference_video_output_path(&self, w2l_inference_token: &str) -> PathBuf {
+    let hashed_path = Self::hashed_directory_path(w2l_inference_token);
+    let video_filename = format!("{}_video.mp4", &w2l_inference_token);
+
+    self.w2l_inference_output_root
+      .join(hashed_path)
+      .join(video_filename)
   }
 
+  //pub fn pretrained_w2l_model_path(&self, model_name: &str) -> String {
+  //  "".to_string()
+  //}
+
   pub fn hashed_directory_path(file_hash: &str) -> String {
     match file_hash.len() {
       0 | 1=> "".to_string(),
@@ -55,8 +85,10 @@ mod tests {
 
   fn get_instance() -> BucketPathUnifier {
     BucketPathUnifier {
-      user_uploaded_audio_for_w2l_root: PathBuf::from("/test_path_w2l_audio"),
       user_uploaded_w2l_templates_root: PathBuf::from("/test_path_w2l_templates"),
+      user_uploaded_audio_for_w2l_root: PathBuf::from("/test_path_w2l_audio"),
+      tts_inference_output_root: PathBuf::from("/test_path_tts_output"),
+      w2l_inference_output_root: PathBuf::from("/test_path_w2l_output"),
     }
   }
 
@@ -74,6 +106,13 @@ mod tests {
                "/test_path_w2l_templates/f/o/o/foobar_detected_faces.pickle");
   }
 
+  #[test]
+  fn test_w2l_inference_video_output_path() {
+    let paths = get_instance();
+    assert_eq!(paths.w2l_inference_video_output_path("foobar").to_str().unwrap(),
+               "/test_path_w2l_output/f/o/o/foobar_video.mp4");
+  }
+
   #[test]
   fn hashed_directory_path_length_zero() {
     assert_eq!(&BucketPathUnifier::hashed_directory_path(""), "");

src/util/mod.rs

@@ -1,6 +1,7 @@
 pub mod anyhow_result;
 pub mod email_to_gravatar;
 pub mod filesystem;
-pub mod random_prefix_crockford_token;
 pub mod random_crockford_token;
-pub mod random_uuid;
+pub mod random_prefix_crockford_token;
+pub mod random_uuid;
+pub mod semi_persistent_cache_dir;

src/util/semi_persistent_cache_dir.rs

@@ -0,0 +1,68 @@
+use std::path::{PathBuf, Path};
+use std::fs;
+use crate::util::anyhow_result::AnyhowResult;
+
+/// These are for files on the worker filesystems
+pub struct SemiPersistentCacheDir {
+  cache_root: PathBuf,
+  tts_model_root: PathBuf,
+  w2l_model_root: PathBuf,
+  w2l_face_templates_root: PathBuf,
+  video_asset_root: PathBuf, // end bump, etc.
+}
+
+impl SemiPersistentCacheDir {
+
+
+  /// Everything is rooted at `/file_cache`.
+  pub fn default_paths() -> Self {
+    // NB: This is the root of the filesystem
+    Self::configured_root("/file_cache/")
+  }
+
+  pub fn configured_root(root_path: &str) -> Self {
+    let cache_root = PathBuf::from(root_path);
+    Self {
+      cache_root: cache_root.clone(),
+      tts_model_root: cache_root.join("tts/models/"),
+      w2l_model_root: cache_root.join("w2l/models/"),
+      w2l_face_templates_root: cache_root.join("w2l/face_templates/"),
+      video_asset_root: cache_root.join("static_video_assets/"),
+    }
+  }
+
+  /// We cache TTS models here.
+  pub fn tts_model_path(&self) -> &Path {
+    &self.tts_model_root
+  }
+
+  /// We cache W2L models here.
+  pub fn w2l_model_path(&self) -> &Path {
+    &self.w2l_model_root
+  }
+
+  /// Face templates go here.
+  pub fn w2l_face_templates_path(&self) -> &Path {
+    &self.w2l_face_templates_root
+  }
+
+  /// End bump, etc.
+  pub fn video_asset_path(&self) -> &Path {
+    &self.video_asset_root
+  }
+
+  pub fn create_w2l_model_path(&self) -> AnyhowResult<()> {
+    let _ = fs::create_dir_all(self.w2l_model_path())?;
+    Ok(())
+  }
+
+  pub fn create_w2l_face_templates_path(&self) -> AnyhowResult<()> {
+    let _ = fs::create_dir_all(self.w2l_face_templates_path())?;
+    Ok(())
+  }
+
+  pub fn create_video_asset_path(&self) -> AnyhowResult<()> {
+    let _ = fs::create_dir_all(self.video_asset_path())?;
+    Ok(())
+  }
+}

src/w2l_inference_job.rs

@@ -15,6 +15,7 @@ pub mod util;
 use anyhow::anyhow;
 use chrono::Utc;
 use crate::buckets::bucket_client::BucketClient;
+use crate::buckets::bucket_path_unifier::BucketPathUnifier;
 use crate::buckets::bucket_paths::hash_to_bucket_path;
 use crate::buckets::file_hashing::get_file_hash;
 use crate::job_queries::w2l_inference_job_queries::W2lInferenceJobRecord;
@@ -22,11 +23,16 @@ use crate::job_queries::w2l_inference_job_queries::insert_w2l_result;
 use crate::job_queries::w2l_inference_job_queries::mark_w2l_inference_job_done;
 use crate::job_queries::w2l_inference_job_queries::mark_w2l_inference_job_failure;
 use crate::job_queries::w2l_inference_job_queries::query_w2l_inference_job_records;
+use crate::script_execution::ffmpeg_generate_preview_image_command::FfmpegGeneratePreviewImageCommand;
+use crate::script_execution::ffmpeg_generate_preview_video_command::FfmpegGeneratePreviewVideoCommand;
 use crate::script_execution::google_drive_download_command::GoogleDriveDownloadCommand;
+use crate::script_execution::imagemagick_generate_preview_image_command::ImagemagickGeneratePreviewImageCommand;
+use crate::script_execution::wav2lip_inference_command::Wav2LipInferenceCommand;
 use crate::util::anyhow_result::AnyhowResult;
 use crate::util::filesystem::check_directory_exists;
 use crate::util::filesystem::check_file_exists;
 use crate::util::random_crockford_token::random_crockford_token;
+use crate::util::semi_persistent_cache_dir::SemiPersistentCacheDir;
 use data_encoding::{HEXUPPER, HEXLOWER, HEXLOWER_PERMISSIVE};
 use log::{warn, info};
 use ring::digest::{Context, Digest, SHA256};
@@ -38,10 +44,6 @@ use std::path::{PathBuf, Path};
 use std::process::Command;
 use std::time::Duration;
 use tempdir::TempDir;
-use crate::script_execution::wav2lip_inference_command::Wav2LipInferenceCommand;
-use crate::script_execution::ffmpeg_generate_preview_video_command::FfmpegGeneratePreviewVideoCommand;
-use crate::script_execution::ffmpeg_generate_preview_image_command::FfmpegGeneratePreviewImageCommand;
-use crate::script_execution::imagemagick_generate_preview_image_command::ImagemagickGeneratePreviewImageCommand;
 
 // Buckets (shared config)
 const ENV_ACCESS_KEY : &'static str = "ACCESS_KEY";
@@ -53,10 +55,8 @@ const ENV_PRIVATE_BUCKET_NAME : &'static str = "W2L_PRIVATE_DOWNLOAD_BUCKET_NAME
 // Buckets (public data)
 const ENV_PUBLIC_BUCKET_NAME : &'static str = "W2L_PUBLIC_DOWNLOAD_BUCKET_NAME";
 
-// Various bucket roots
-const ENV_DOWNLOAD_BUCKET_ROOT : &'static str = "W2L_DOWNLOAD_BUCKET_ROOT";
-const ENV_AUDIO_UPLOADS_BUCKET_ROOT : &'static str = "AUDIO_UPLOADS_BUCKET_ROOT";
-const ENV_INFERENCE_OUTPUT_BUCKET_ROOT : &'static str = "W2L_INFERENCE_OUTPUT_BUCKET_ROOT";
+// Where models and other assets get downloaded to.
+const ENV_SEMIPERSISTENT_CACHE_DIR : &'static str = "SEMIPERSISTENT_CACHE_DIR";
 
 // Python code
 const ENV_CODE_DIRECTORY : &'static str = "W2L_CODE_DIRECTORY";
@@ -74,6 +74,9 @@ struct Inferencer {
   pub private_bucket_client: BucketClient,
   pub public_bucket_client: BucketClient,
 
+  pub bucket_path_unifier: BucketPathUnifier,
+  pub semi_persistent_cache: SemiPersistentCacheDir,
+
   pub google_drive_downloader: GoogleDriveDownloadCommand,
   pub w2l_inference: Wav2LipInferenceCommand,
   //pub ffmpeg_image_preview_generator: FfmpegGeneratePreviewImageCommand,
@@ -82,13 +85,6 @@ struct Inferencer {
 
   // Command to run
   pub inference_script: String,
-
-  // Root to retrieve W2L templates (public and private)
-  pub w2l_template_uploads_bucket_root: String,
-  // Root to retrieve uploaded audio (public and private)
-  pub audio_uploads_bucket_root: String,
-  // Root to upload inference output
-  pub inference_output_bucket_root: String,
 }
 
 #[tokio::main]
@@ -117,11 +113,6 @@ async fn main() -> AnyhowResult<()> {
   let private_bucket_name = easyenv::get_env_string_required(ENV_PRIVATE_BUCKET_NAME)?;
   let public_bucket_name = easyenv::get_env_string_required(ENV_PUBLIC_BUCKET_NAME)?;
 
-  // Bucket roots
-  let w2l_template_bucket_root = easyenv::get_env_string_required(ENV_DOWNLOAD_BUCKET_ROOT)?;
-  let audio_bucket_root = easyenv::get_env_string_required(ENV_AUDIO_UPLOADS_BUCKET_ROOT)?;
-  let inference_output_bucket_root = easyenv::get_env_string_required(ENV_INFERENCE_OUTPUT_BUCKET_ROOT)?;
-
   let private_bucket_client = BucketClient::create(
     &access_key,
     &secret_key,
@@ -176,6 +167,17 @@ async fn main() -> AnyhowResult<()> {
 
   let inference_script = "TODO".to_string();
 
+  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_w2l_model_path()?;
+  semi_persistent_cache.create_w2l_face_templates_path()?;
+  semi_persistent_cache.create_w2l_model_path()?;
+
   let inferencer = Inferencer {
     download_temp_directory: temp_directory,
     mysql_pool,
@@ -187,9 +189,8 @@ async fn main() -> AnyhowResult<()> {
     //ffmpeg_video_preview_generator: FfmpegGeneratePreviewVideoCommand {},
     //imagemagick_image_preview_generator: ImagemagickGeneratePreviewImageCommand {},
     w2l_inference: w2l_inference_command,
-    w2l_template_uploads_bucket_root: w2l_template_bucket_root.to_string(),
-    audio_uploads_bucket_root: audio_bucket_root.to_string(),
-    inference_output_bucket_root: inference_output_bucket_root.to_string(),
+    bucket_path_unifier: BucketPathUnifier::default_paths(),
+    semi_persistent_cache,
   };
 
   main_loop(inferencer).await;
@@ -286,19 +287,17 @@ fn read_metadata_file(filename: &str) -> AnyhowResult<FileMetadata> {
 
 async fn process_job(inferencer: &Inferencer, job: &W2lInferenceJobRecord) -> AnyhowResult<()> {
 
-  // TODO: 1. Check if w2l model is downloaded and download it
-  // TODO: 2. Check if w2l template faces are downloaded and download it
-  // TODO: 3. Download audio
+  // TODO 1. Mark Processing
+  //
+  // TODO: 2. Check if w2l model is downloaded / download it to a stable cache location
+  // TODO: 3. Check if w2l template faces are downloaded and download it
+  // TODO: 4. Download user audio
 
+  // TODO: 5. Process Inference
 
-  // TODO: 1. Mark processing.
-  // TODO: 2. Download. (DONE)
-  // TODO: 3. Process template with face detection (DONE)
-  // TODO: 4. Determine if picture or video
-  // TODO: 5. Take a screenshot/gif
-  // TODO: 6. Upload all (partially done).
-  // TODO: 7. Save record. (DONE)
-  // TODO: 8. Mark job done. (DONE)
+  // TODO 6. Upload result
+  // TODO 7. Save record
+  // TODO 8. Mark job done
 
   let temp_dir = format!("temp_{}", job.id);
   let temp_dir = TempDir::new(&temp_dir)?;
@@ -354,9 +353,11 @@ async fn process_job(inferencer: &Inferencer, job: &W2lInferenceJobRecord) -> An
   info!("File hash: {}", private_bucket_hash);
 
   // Full path to video/image
-  let full_object_path = hash_to_bucket_path(
-    &private_bucket_hash,
-    Some(&inferencer.w2l_template_uploads_bucket_root))?;
+  //let full_object_path = hash_to_bucket_path(
+  //  &private_bucket_hash,
+  //  Some(&inferencer.w2l_template_uploads_bucket_root))?;
+
+  let full_object_path = "TODO";
 
   // ==================== GENERATE VIDEO PREVIEWS ==================== //