From 13d00f6edff008faad4db77b08c5bf724c48e0f0 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Fri, 16 Dec 2022 03:07:20 -0500 Subject: [PATCH] partially handle softvc model downloads --- .../src/bucket_path_unifier.rs | 26 ++++ .../service/job/download_job/src/job_state.rs | 2 + .../hifigan_softvc_model_check_command.rs | 12 +- .../job/download_job/src/job_types/mod.rs | 1 + .../download_job/src/job_types/softvc/mod.rs | 2 + .../job_types/softvc/process_softvc_model.rs | 140 ++++++++++++++++++ .../softvc/softvc_model_check_command.rs | 87 +++++++++++ 7 files changed, 264 insertions(+), 6 deletions(-) create mode 100644 crates/service/job/download_job/src/job_types/softvc/mod.rs create mode 100644 crates/service/job/download_job/src/job_types/softvc/process_softvc_model.rs create mode 100644 crates/service/job/download_job/src/job_types/softvc/softvc_model_check_command.rs diff --git a/crates/lib/storage_buckets_common/src/bucket_path_unifier.rs b/crates/lib/storage_buckets_common/src/bucket_path_unifier.rs index 45dbc96c1d..7ad5b8c9c3 100644 --- a/crates/lib/storage_buckets_common/src/bucket_path_unifier.rs +++ b/crates/lib/storage_buckets_common/src/bucket_path_unifier.rs @@ -16,6 +16,8 @@ pub struct BucketPathUnifier { pub w2l_inference_output_root: PathBuf, pub w2l_model_root: PathBuf, pub w2l_end_bump_root: PathBuf, + // VC + pub softvc_model_root: PathBuf, } impl BucketPathUnifier { @@ -33,6 +35,8 @@ impl BucketPathUnifier { w2l_inference_output_root: PathBuf::from("/w2l_inference_output"), w2l_model_root: PathBuf::from("/w2l_pretrained_models"), w2l_end_bump_root: PathBuf::from("/w2l_end_bumps"), + // VC + softvc_model_root: PathBuf::from("/user_uploaded_softvc_models"), } } @@ -169,6 +173,19 @@ impl BucketPathUnifier { .join(video_filename) } + // ==================== VOICE CONVERSION MODELS ==================== // + + pub fn softvc_model_path(&self, softvc_model_file_hash: &str) -> PathBuf { + let hashed_path = Self::hashed_directory_path(softvc_model_file_hash); + let model_filename = format!("{}.pt", &softvc_model_file_hash); + + self.softvc_model_root + .join(hashed_path) + .join(model_filename) + } + + // ==================== UTILITY ==================== // + pub fn hashed_directory_path(file_hash: &str) -> String { match file_hash.len() { 0 | 1=> "".to_string(), @@ -197,6 +214,8 @@ mod tests { w2l_inference_output_root: PathBuf::from("/test_path_w2l_output"), w2l_model_root: PathBuf::from("/test_path_w2l_pretrained_models"), w2l_end_bump_root: PathBuf::from("/test_path_w2l_end_bumps"), + // VOICE CONVERSION + softvc_model_root: PathBuf::from("/test_path_softvc_models"), } } @@ -297,6 +316,13 @@ mod tests { "/test_path_w2l_output/a/b/c/vocodes_video_TYPEABCDEF.mp4"); } + #[test] + fn test_softvc_model_path() { + let paths = get_instance(); + assert_eq!(paths.softvc_model_path("foobar").to_str().unwrap(), + "/test_path_softvc_models/f/o/o/foobar.pt"); + } + #[test] fn hashed_directory_path_length_zero() { assert_eq!(&BucketPathUnifier::hashed_directory_path(""), ""); diff --git a/crates/service/job/download_job/src/job_state.rs b/crates/service/job/download_job/src/job_state.rs index 82ded1e308..2ee2e40626 100644 --- a/crates/service/job/download_job/src/job_state.rs +++ b/crates/service/job/download_job/src/job_state.rs @@ -1,5 +1,6 @@ use crate::job_types::hifigan::hifigan_model_check_command::HifiGanModelCheckCommand; use crate::job_types::hifigan_softvc::hifigan_softvc_model_check_command::HifiGanSoftVcModelCheckCommand; +use crate::job_types::softvc::softvc_model_check_command::SoftVcModelCheckCommand; use crate::job_types::tacotron::tacotron_model_check_command::TacotronModelCheckCommand; use database_queries::mediators::badge_granter::BadgeGranter; use database_queries::mediators::firehose_publisher::FirehosePublisher; @@ -42,6 +43,7 @@ pub struct JobState { /// Configurations and interfaces to code deployed as sidecars or container mounts. pub struct SidecarConfigs { pub google_drive_downloader: GoogleDriveDownloadCommand, + pub softvc_model_check_command: SoftVcModelCheckCommand, pub tacotron_model_check_command: TacotronModelCheckCommand, pub hifigan_model_check_command: HifiGanModelCheckCommand, pub hifigan_softvc_model_check_command: HifiGanSoftVcModelCheckCommand, diff --git a/crates/service/job/download_job/src/job_types/hifigan_softvc/hifigan_softvc_model_check_command.rs b/crates/service/job/download_job/src/job_types/hifigan_softvc/hifigan_softvc_model_check_command.rs index 3cf34bb129..054b98f897 100644 --- a/crates/service/job/download_job/src/job_types/hifigan_softvc/hifigan_softvc_model_check_command.rs +++ b/crates/service/job/download_job/src/job_types/hifigan_softvc/hifigan_softvc_model_check_command.rs @@ -4,10 +4,10 @@ use std::path::Path; use subprocess::{Popen, PopenConfig}; use subprocess_common::docker_options::DockerOptions; -/// This command is used to check hifigan for being a real model +/// This command is used to check hifigan (softvc) for being a real model #[derive(Clone)] pub struct HifiGanSoftVcModelCheckCommand { - /// Where the HifiGan code lives + /// Where the HifiGan (softvc version) code lives hifigan_softvc_root_code_directory: String, /// eg. `source python/bin/activate` @@ -21,15 +21,15 @@ pub struct HifiGanSoftVcModelCheckCommand { impl HifiGanSoftVcModelCheckCommand { pub fn new( - hifigan_root_code_directory: &str, + hifigan_softvc_root_code_directory: &str, maybe_virtual_env_activation_command: Option<&str>, - hifigan_model_check_script_name: &str, + hifigan_softvc_model_check_script_name: &str, maybe_docker_options: Option, ) -> Self { Self { - hifigan_softvc_root_code_directory: hifigan_root_code_directory.to_string(), + hifigan_softvc_root_code_directory: hifigan_softvc_root_code_directory.to_string(), maybe_virtual_env_activation_command: maybe_virtual_env_activation_command.map(|s| s.to_string()), - hifigan_softvc_model_check_script_name: hifigan_model_check_script_name.to_string(), + hifigan_softvc_model_check_script_name: hifigan_softvc_model_check_script_name.to_string(), maybe_docker_options, } } diff --git a/crates/service/job/download_job/src/job_types/mod.rs b/crates/service/job/download_job/src/job_types/mod.rs index c4ca220216..8342328df3 100644 --- a/crates/service/job/download_job/src/job_types/mod.rs +++ b/crates/service/job/download_job/src/job_types/mod.rs @@ -1,4 +1,5 @@ pub mod hifigan; pub mod hifigan_softvc; +pub mod softvc; pub mod tacotron; pub mod talknet; \ No newline at end of file diff --git a/crates/service/job/download_job/src/job_types/softvc/mod.rs b/crates/service/job/download_job/src/job_types/softvc/mod.rs new file mode 100644 index 0000000000..87ade21b43 --- /dev/null +++ b/crates/service/job/download_job/src/job_types/softvc/mod.rs @@ -0,0 +1,2 @@ +pub mod softvc_model_check_command; +pub mod process_softvc_model; diff --git a/crates/service/job/download_job/src/job_types/softvc/process_softvc_model.rs b/crates/service/job/download_job/src/job_types/softvc/process_softvc_model.rs new file mode 100644 index 0000000000..cc9d84f130 --- /dev/null +++ b/crates/service/job/download_job/src/job_types/softvc/process_softvc_model.rs @@ -0,0 +1,140 @@ +use anyhow::anyhow; +use container_common::anyhow_result::AnyhowResult; +use container_common::filesystem::check_file_exists::check_file_exists; +use container_common::filesystem::safe_delete_temp_directory::safe_delete_temp_directory; +use container_common::filesystem::safe_delete_temp_file::safe_delete_temp_file; +use container_common::hashing::hash_file_sha2::hash_file_sha2; +use crate::JobState; +use crate::job_steps::job_results::JobResults; +use database_queries::queries::generic_download::job::list_available_generic_download_jobs::AvailableDownloadJob; +use database_queries::queries::tts::tts_models::insert_tts_model_from_download_job::insert_tts_model_from_download_job; +use database_queries::queries::tts::tts_models::insert_tts_model_from_download_job; +use jobs_common::redis_job_status_logger::RedisJobStatusLogger; +use log::{info, warn}; +use reusable_types::db::enums::entity_visibility::EntityVisibility; +use std::fs::File; +use std::io::Read; +use std::path::PathBuf; +use tempdir::TempDir; + +/// Returns the token of the entity. +pub async fn process_softvc_model<'a, 'b>( + job_state: &JobState, + job: &AvailableDownloadJob, + temp_dir: &TempDir, + download_filename: &str, + redis_logger: &'a mut RedisJobStatusLogger<'b>, +) -> AnyhowResult { + + // ==================== RUN MODEL CHECK ==================== // + + info!("Checking that softvc model is valid..."); + + redis_logger.log_status("checking rocket vc model")?; // NB: "rocket vc" is a codename + + let file_path = PathBuf::from(download_filename.clone()); + + let output_metadata_fs_path = temp_dir.path().join("metadata.json"); + + let model_check_result = job_state.sidecar_configs.softvc_model_check_command.execute( + &file_path, + &output_metadata_fs_path, + ); + + if let Err(e) = model_check_result { + safe_delete_temp_file(&file_path); + safe_delete_temp_directory(&temp_dir); + return Err(anyhow!("model check error: {:?}", e)); + } + + // ==================== CHECK ALL FILES EXIST AND GET METADATA ==================== // + + info!("Checking that metadata output file exists..."); + + check_file_exists(&output_metadata_fs_path)?; + + let file_metadata = match read_metadata_file(&output_metadata_fs_path) { + Ok(metadata) => metadata, + Err(e) => { + safe_delete_temp_file(&file_path); + safe_delete_temp_file(&output_metadata_fs_path); + safe_delete_temp_directory(&temp_dir); + return Err(e); + } + }; + + // ==================== UPLOAD MODEL FILE ==================== // + + info!("Uploading SoftVC model to GCS..."); + + let private_bucket_hash = hash_file_sha2(&download_filename)?; + + info!("File hash: {}", private_bucket_hash); + + let model_bucket_path = job_state.bucket_path_unifier.softvc_model_path(&private_bucket_hash); + + info!("Destination bucket path: {:?}", &model_bucket_path); + + redis_logger.log_status("uploading rocket vc model")?; // NB: "rocket vc" is a codename + + if let Err(e) = job_state.bucket_client.upload_filename(&model_bucket_path, &file_path).await { + safe_delete_temp_file(&output_metadata_fs_path); + safe_delete_temp_file(&file_path); + safe_delete_temp_directory(&temp_dir); + return Err(e); + } + + // ==================== DELETE DOWNLOADED FILE ==================== // + + // NB: We should be using a tempdir, but to make absolutely certain we don't overflow the disk... + safe_delete_temp_file(&output_metadata_fs_path); + safe_delete_temp_file(&file_path); + safe_delete_temp_directory(&temp_dir); + + // ==================== SAVE RECORDS ==================== // + + info!("Saving Soft VC record..."); + + // TODO FIX BELOW QUERIES FOR SOFTVC -- + // TODO FIX BELOW QUERIES FOR SOFTVC -- + // TODO FIX BELOW QUERIES FOR SOFTVC -- + // TODO FIX BELOW QUERIES FOR SOFTVC -- + + let (_id, model_token) = insert_tts_model_from_download_job(insert_tts_model_from_download_job::Args { + title: &job.title, + original_download_url: &job.download_url, + original_filename: &download_filename, + file_size_bytes: file_metadata.file_size_bytes, + creator_user_token: &job.creator_user_token, + creator_ip_address: &job.creator_ip_address, + creator_set_visibility: EntityVisibility::Public, // TODO: All models default to public at start + private_bucket_hash: &private_bucket_hash, + private_bucket_object_name: &model_bucket_path, + mysql_pool: &job_state.mysql_pool, + }).await?; + + job_state.badge_granter.maybe_grant_tts_model_uploads_badge(&job.creator_user_token) + .await + .map_err(|e| { + warn!("error maybe awarding badge: {:?}", e); + anyhow!("error maybe awarding badge") + })?; + + Ok(JobResults { + entity_token: Some(model_token), + entity_type: Some("tacotron2".to_string()), // NB: This may be different from `GenericDownloadType` in the future! + }) +} + +#[derive(Deserialize)] +struct FileMetadata { + pub file_size_bytes: u64, +} + +fn read_metadata_file(filename: &PathBuf) -> AnyhowResult { + let mut file = File::open(filename)?; + let mut buffer = String::new(); + file.read_to_string(&mut buffer)?; + Ok(serde_json::from_str(&buffer)?) +} + diff --git a/crates/service/job/download_job/src/job_types/softvc/softvc_model_check_command.rs b/crates/service/job/download_job/src/job_types/softvc/softvc_model_check_command.rs new file mode 100644 index 0000000000..ba91ba5976 --- /dev/null +++ b/crates/service/job/download_job/src/job_types/softvc/softvc_model_check_command.rs @@ -0,0 +1,87 @@ +use container_common::anyhow_result::AnyhowResult; +use log::info; +use std::fs::OpenOptions; +use std::path::Path; +use subprocess::{Popen, PopenConfig, Redirection}; +use subprocess_common::docker_options::DockerOptions; + +/// This command is used to check softvc for being a real model +#[derive(Clone)] +pub struct SoftVcModelCheckCommand { + /// Where the Tacotron code lives + softvc_root_code_directory: String, + + /// eg. `source python/bin/activate` + maybe_virtual_env_activation_command: Option, + + softvc_model_check_script_name: String, + + /// If this is run under Docker (eg. in development), these are the options. + maybe_docker_options: Option, +} + +impl SoftVcModelCheckCommand { + pub fn new( + softvc_root_code_directory: &str, + maybe_virtual_env_activation_command: Option<&str>, + softvc_model_check_script_name: &str, + maybe_docker_options: Option, + ) -> Self { + Self { + softvc_root_code_directory: softvc_root_code_directory.to_string(), + maybe_virtual_env_activation_command: maybe_virtual_env_activation_command.map(|s| s.to_string()), + softvc_model_check_script_name: softvc_model_check_script_name.to_string(), + maybe_docker_options, + } + } + + pub fn execute>( + &self, + synthesizer_checkpoint_path: P, + output_metadata_filename: P, + ) -> AnyhowResult<()> { + let mut command = String::new(); + + command.push_str("echo 'test'"); + command.push_str(" && "); + command.push_str(&format!("cd {}", self.softvc_root_code_directory)); + + if let Some(venv_command) = self.maybe_virtual_env_activation_command.as_deref() { + command.push_str(" && "); + command.push_str(venv_command); + command.push_str(" "); + } + + command.push_str(" && "); + command.push_str("python "); + command.push_str(&self.softvc_model_check_script_name); + command.push_str(" --synthesizer_checkpoint_path "); + command.push_str(&synthesizer_checkpoint_path.as_ref().display().to_string()); + command.push_str(" --output_metadata_filename "); + command.push_str(&output_metadata_filename.as_ref().display().to_string()); + + if let Some(docker_options) = self.maybe_docker_options.as_ref() { + command = docker_options.to_command_string(&command); + } + + info!("Command: {:?}", command); + + let command_parts = [ + "bash", + "-c", + &command + ]; + + let mut p = Popen::create(&command_parts, PopenConfig { + ..Default::default() + })?; + + info!("Pid : {:?}", p.pid()); + + let exit_status = p.wait()?; + + info!("Exit status: {:?}", exit_status); + + Ok(()) + } +}