download job can now handle tacotron (untested)

crates/lib/database_queries/src/queries/tts/tts_download_jobs/tts_download_job_queries.rs

@@ -249,4 +249,4 @@ SET
   };
 
   Ok((record_id, model_token))
-}
+}

crates/lib/database_queries/src/queries/tts/tts_models/insert_tts_model_from_download_job.rs

@@ -0,0 +1,75 @@
+use anyhow::anyhow;
+use chrono::Utc;
+use container_common::anyhow_result::AnyhowResult;
+use crate::queries::generic_download::job::list_available_generic_download_jobs::AvailableDownloadJob;
+use crate::tokens::Tokens;
+use reusable_types::entity_visibility::EntityVisibility;
+use sqlx::MySqlPool;
+use std::path::Path;
+
+pub struct Args<'a, P: AsRef<Path>> {
+  pub title: &'a str,
+
+  pub original_download_url: &'a str,
+  pub original_filename: &'a str,
+  pub file_size_bytes: u64,
+
+  pub creator_user_token: &'a str,
+  pub creator_ip_address: &'a str,
+  pub creator_set_visibility: EntityVisibility,
+
+  pub private_bucket_hash: &'a str,
+  pub private_bucket_object_name: P,
+
+  pub mysql_pool: &'a MySqlPool,
+}
+
+
+pub async fn insert_tts_model_from_download_job<P: AsRef<Path>>(
+  args: Args<'_, P>,
+) -> AnyhowResult<(u64, String)> {
+
+  let model_token = Tokens::new_tts_model()?;
+
+  let private_bucket_object_name = &args.private_bucket_object_name
+      .as_ref()
+      .display()
+      .to_string();
+
+  let query_result = sqlx::query!(
+        r#"
+INSERT INTO tts_models
+SET
+  token = ?,
+  tts_model_type = "tacotron2",
+  title = ?,
+  description_markdown = '',
+  description_rendered_html = '',
+  creator_user_token = ?,
+  creator_ip_address_creation = ?,
+  creator_ip_address_last_update = ?,
+  original_download_url = ?,
+  private_bucket_hash = ?,
+  private_bucket_object_name = ?,
+  file_size_bytes = ?
+        "#,
+      &model_token,
+      args.title,
+      args.creator_user_token,
+      args.creator_ip_address,
+      args.creator_ip_address,
+      args.original_download_url,
+      args.private_bucket_hash,
+      private_bucket_object_name,
+      args.file_size_bytes
+    )
+      .execute(args.mysql_pool)
+      .await;
+
+  let record_id = match query_result {
+    Ok(res) => res.last_insert_id(),
+    Err(err) => return Err(anyhow!("Mysql error: {:?}", err)),
+  };
+
+  Ok((record_id, model_token))
+}

crates/lib/database_queries/src/queries/tts/tts_models/mod.rs

@@ -3,4 +3,5 @@ pub mod edit_tts_model_details;
 pub mod edit_tts_model_moderator_details;
 pub mod get_tts_model;
 pub mod get_tts_model_for_inference;
+pub mod insert_tts_model_from_download_job;
 pub mod list_tts_models;

crates/lib/reusable_types/src/generic_download_type.rs

@@ -12,6 +12,10 @@ pub enum GenericDownloadType {
   #[serde(rename = "melgan_vocodes")]
   #[sqlx(rename = "melgan_vocodes")]
   MelGanVocodes,
+
+  #[serde(rename = "tacotron2")]
+  #[sqlx(rename = "tacotron2")]
+  Tacotron2,
 }
 
 /// NB: Legacy API for older code.
@@ -20,6 +24,7 @@ impl GenericDownloadType {
     match self {
       Self::HifiGan => "hifigan",
       Self::MelGanVocodes => "melgan_vocodes",
+      Self::Tacotron2 => "tacotron2",
     }
   }
 
@@ -27,6 +32,7 @@ impl GenericDownloadType {
     match value {
       "hifigan" => Ok(Self::HifiGan),
       "melgan_vocodes" => Ok(Self::MelGanVocodes),
+      "tacotron2" => Ok(Self::Tacotron2),
       _ => Err(format!("invalid value: {:?}", value)),
     }
   }
@@ -41,17 +47,20 @@ mod tests {
   fn test_serialization() {
     assert_serialization(GenericDownloadType::HifiGan, "hifigan");
     assert_serialization(GenericDownloadType::MelGanVocodes, "melgan_vocodes");
+    assert_serialization(GenericDownloadType::Tacotron2, "tacotron2");
   }
 
   #[test]
   fn to_str() {
     assert_eq!(GenericDownloadType::HifiGan.to_str(), "hifigan");
     assert_eq!(GenericDownloadType::MelGanVocodes.to_str(), "melgan_vocodes");
+    assert_eq!(GenericDownloadType::Tacotron2.to_str(), "tacotron2");
   }
 
   #[test]
   fn from_str() {
     assert_eq!(GenericDownloadType::from_str("hifigan").unwrap(), GenericDownloadType::HifiGan);
     assert_eq!(GenericDownloadType::from_str("melgan_vocodes").unwrap(), GenericDownloadType::MelGanVocodes);
+    assert_eq!(GenericDownloadType::from_str("tacotron2").unwrap(), GenericDownloadType::Tacotron2);
   }
 }

crates/service/download_job/src/job_state.rs

@@ -1,4 +1,5 @@
 use crate::job_types::hifigan::hifigan_model_check_command::HifiGanModelCheckCommand;
+use crate::job_types::tacotron::tacotron_model_check_command::TacotronModelCheckCommand;
 use database_queries::mediators::badge_granter::BadgeGranter;
 use database_queries::mediators::firehose_publisher::FirehosePublisher;
 use google_drive_common::google_drive_download_command::GoogleDriveDownloadCommand;
@@ -22,6 +23,7 @@ pub struct JobState {
 
   pub bucket_path_unifier: BucketPathUnifier,
 
+  pub tacotron_model_check_command: TacotronModelCheckCommand,
   pub hifigan_model_check_command: HifiGanModelCheckCommand,
 
   // Command to run

crates/service/download_job/src/job_steps/process_single_job.rs

@@ -4,6 +4,7 @@ use container_common::anyhow_result::AnyhowResult;
 use container_common::filesystem::safe_delete_temp_directory::safe_delete_temp_directory;
 use crate::JobState;
 use crate::job_types::hifigan::process_hifigan_vocoder::process_hifigan_vocoder;
+use crate::job_types::tacotron::process_tacotron_model::process_tacotron_model;
 use database_queries::queries::generic_download::job::list_available_generic_download_jobs::AvailableDownloadJob;
 use database_queries::queries::generic_download::job::mark_generic_download_job_done::mark_generic_download_job_done;
 use database_queries::queries::generic_download::job::mark_generic_download_job_pending_and_grab_lock::mark_generic_download_job_pending_and_grab_lock;
@@ -67,6 +68,17 @@ pub async fn process_single_job(job_state: &JobState, job: &AvailableDownloadJob
       entity_type = results.entity_type.clone();
     }
     GenericDownloadType::MelGanVocodes => {
+      let results = process_tacotron_model(
+        job_state,
+        job,
+        &temp_dir,
+        &download_filename,
+        &mut redis_logger,
+      ).await?;
+      entity_token = results.entity_token.clone();
+      entity_type = results.entity_type.clone();
+    }
+    GenericDownloadType::Tacotron2 => {
       // TODO
     }
   }

crates/service/download_job/src/job_types/mod.rs

@@ -1 +1,3 @@
-pub mod hifigan;
+pub mod hifigan;
+pub mod tacotron;
+pub mod talknet;

crates/service/download_job/src/job_types/tacotron/mod.rs

@@ -0,0 +1,2 @@
+pub mod tacotron_model_check_command;
+pub mod process_tacotron_model;

crates/service/download_job/src/job_types/tacotron/process_tacotron_model.rs

@@ -0,0 +1,136 @@
+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::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_tacotron_model<'a, 'b>(
+  job_state: &JobState,
+  job: &AvailableDownloadJob,
+  temp_dir: &TempDir,
+  download_filename: &str,
+  redis_logger: &'a mut RedisJobStatusLogger<'b>,
+) -> AnyhowResult<JobResults> {
+
+  // ==================== RUN MODEL CHECK ==================== //
+
+  info!("Checking that model is valid...");
+
+  redis_logger.log_status("checking tacotron model")?;
+
+  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.tacotron_model_check_command.execute(
+    &file_path,
+    &output_metadata_fs_path,
+    false
+  );
+
+  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 Tacotron TTS 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.tts_synthesizer_path(&private_bucket_hash);
+
+  info!("Destination bucket path: {:?}", &model_bucket_path);
+
+  redis_logger.log_status("uploading tacotron TTS model")?;
+
+  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 TTS model record...");
+
+  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<FileMetadata> {
+  let mut file = File::open(filename)?;
+  let mut buffer = String::new();
+  file.read_to_string(&mut buffer)?;
+  Ok(serde_json::from_str(&buffer)?)
+}
+

crates/service/download_job/src/job_types/tacotron/tacotron_model_check_command.rs

@@ -0,0 +1,128 @@
+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 tacotron for being a real model
+#[derive(Clone)]
+pub struct TacotronModelCheckCommand {
+  /// Where the Tacotron code lives
+  tacotron_root_code_directory: String,
+  
+  /// eg. `source python/bin/activate`
+  virtual_env_activation_command: String,
+  
+  tacotron_model_check_script_name: String,
+
+  /// If this is run under Docker (eg. in development), these are the options.
+  maybe_docker_options: Option<DockerOptions>,
+}
+
+impl TacotronModelCheckCommand {
+  pub fn new(
+    tacotron_root_code_directory: &str,
+    virtual_env_activation_command: &str,
+    tacotron_model_check_script_name: &str,
+    maybe_docker_options: Option<DockerOptions>,
+  ) -> Self {
+    Self {
+      tacotron_root_code_directory: tacotron_root_code_directory.to_string(),
+      virtual_env_activation_command: virtual_env_activation_command.to_string(),
+      tacotron_model_check_script_name: tacotron_model_check_script_name.to_string(),
+      maybe_docker_options,
+    }
+  }
+
+  pub fn execute<P: AsRef<Path>>(
+    &self,
+    synthesizer_checkpoint_path: P,
+    output_metadata_filename: P,
+    spawn_process: bool
+  ) -> AnyhowResult<()> {
+    let mut command = String::new();
+
+    command.push_str("echo 'test'");
+    command.push_str(" && ");
+    command.push_str(&format!("cd {}", self.tacotron_root_code_directory));
+    command.push_str(" && ");
+    command.push_str(&self.virtual_env_activation_command);
+    command.push_str(" && ");
+    command.push_str("python ");
+    command.push_str(&self.tacotron_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
+    ];
+
+    if spawn_process {
+      // NB: This forks and returns immediately.
+      //let _child_pid = command_builder.spawn()?;
+
+      let stdout_file = OpenOptions::new()
+          .read(true)
+          .write(true)
+          .create(true)
+          .truncate(true)
+          .open("/tmp/tacotron_upload_stdout.txt")?;
+
+      let stderr_file = OpenOptions::new()
+          .read(true)
+          .write(true)
+          .create(true)
+          .truncate(true)
+          .open("/tmp/tacotron_upload_stderr.txt")?;
+
+      let mut p = Popen::create(&command_parts, PopenConfig {
+        //stdout: Redirection::Pipe,
+        //stderr: Redirection::Pipe,
+        stdout: Redirection::File(stdout_file),
+        stderr: Redirection::File(stderr_file),
+        ..Default::default()
+      })?;
+
+      info!("Pid : {:?}", p.pid());
+
+      p.detach();
+
+    } else {
+      // NB: This is a blocking call.
+      /*let output = command_builder.output()?;
+
+      info!("Output status: {}", output.status);
+      info!("Stdout: {:?}", String::from_utf8(output.stdout));
+      error!("Stderr: {:?}", String::from_utf8(output.stderr));
+
+      if !output.status.success() {
+        bail!("Bad error code: {:?}", output.status);
+      }*/
+
+      let mut p = Popen::create(&command_parts, PopenConfig {
+        //stdout: Redirection::Pipe,
+        //stderr: Redirection::Pipe,
+        ..Default::default()
+      })?;
+
+      info!("Pid : {:?}", p.pid());
+
+      let exit_status = p.wait()?;
+
+      info!("Exit status: {:?}", exit_status);
+    }
+
+    Ok(())
+  }
+}

crates/service/download_job/src/job_types/talknet/mod.rs

@@ -0,0 +1 @@
+pub mod talknet_model_check_command;

crates/service/download_job/src/job_types/talknet/talknet_model_check_command.rs

@@ -0,0 +1,111 @@
+use container_common::anyhow_result::AnyhowResult;
+use log::info;
+use subprocess::{Popen, PopenConfig, Redirection};
+use std::fs::OpenOptions;
+use std::path::Path;
+
+/// This command is used to check talknet for being a real model
+#[derive(Clone)]
+pub struct TalknetModelCheckCommand {
+  talknet_directory: String,
+  check_model_script_name: String,
+}
+
+impl TalknetModelCheckCommand {
+  pub fn new(
+    talknet_directory: &str,
+    check_model_script_name: &str,
+  ) -> Self {
+    Self {
+      talknet_directory: talknet_directory.to_string(),
+      check_model_script_name: check_model_script_name.to_string(),
+    }
+  }
+
+  pub fn execute<P: AsRef<Path>>(
+    &self,
+    synthesizer_checkpoint_path: P,
+    output_metadata_filename: P,
+    spawn_process: bool
+  ) -> AnyhowResult<()> {
+    let mut command = String::new();
+
+    command.push_str("echo 'test'");
+    command.push_str(" && ");
+    command.push_str(&format!("cd {}", self.talknet_directory));
+    command.push_str(" && ");
+    command.push_str("source python/bin/activate");
+    command.push_str(" && ");
+    command.push_str("python ");
+    command.push_str(&self.check_model_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());
+
+    info!("Command: {:?}", command);
+
+    let command_parts = [
+      "bash",
+      "-c",
+      &command
+    ];
+
+    if spawn_process {
+      // NB: This forks and returns immediately.
+      //let _child_pid = command_builder.spawn()?;
+
+      let stdout_file = OpenOptions::new()
+          .read(true)
+          .write(true)
+          .create(true)
+          .truncate(true)
+          .open("/tmp/talknet_upload_stdout.txt")?;
+
+      let stderr_file = OpenOptions::new()
+          .read(true)
+          .write(true)
+          .create(true)
+          .truncate(true)
+          .open("/tmp/talknet_upload_stderr.txt")?;
+
+      let mut p = Popen::create(&command_parts, PopenConfig {
+        //stdout: Redirection::Pipe,
+        //stderr: Redirection::Pipe,
+        stdout: Redirection::File(stdout_file),
+        stderr: Redirection::File(stderr_file),
+        ..Default::default()
+      })?;
+
+      info!("Pid : {:?}", p.pid());
+
+      p.detach();
+
+    } else {
+      // NB: This is a blocking call.
+      /*let output = command_builder.output()?;
+
+      info!("Output status: {}", output.status);
+      info!("Stdout: {:?}", String::from_utf8(output.stdout));
+      error!("Stderr: {:?}", String::from_utf8(output.stderr));
+
+      if !output.status.success() {
+        bail!("Bad error code: {:?}", output.status);
+      }*/
+
+      let mut p = Popen::create(&command_parts, PopenConfig {
+        //stdout: Redirection::Pipe,
+        //stderr: Redirection::Pipe,
+        ..Default::default()
+      })?;
+
+      info!("Pid : {:?}", p.pid());
+
+      let exit_status = p.wait()?;
+
+      info!("Exit status: {:?}", exit_status);
+    }
+
+    Ok(())
+  }
+}

crates/service/download_job/src/main.rs

@@ -42,6 +42,7 @@ use std::time::Duration;
 use storage_buckets_common::bucket_client::BucketClient;
 use storage_buckets_common::bucket_path_unifier::BucketPathUnifier;
 use subprocess_common::docker_options::{DockerFilesystemMount, DockerGpu, DockerOptions};
+use crate::job_types::tacotron::tacotron_model_check_command::TacotronModelCheckCommand;
 
 // Buckets
 const ENV_ACCESS_KEY : &'static str = "ACCESS_KEY";
@@ -94,6 +95,7 @@ async fn main() -> AnyhowResult<()> {
     "DOWNLOAD_SCRIPT",
     "./scripts/download_internet_file.py");
 
+  // TODO/FIXME: Cannot be deployed as currently written.
   //let google_drive_downloader = GoogleDriveDownloadCommand::new(&download_script);
   let google_drive_downloader = GoogleDriveDownloadCommand::new_local_dev_docker(
     "./download_internet_file.py",
@@ -143,6 +145,19 @@ async fn main() -> AnyhowResult<()> {
     firehose_publisher: firehose_publisher.clone(), // NB: Also safe
   };
 
+  // TODO/FIXME: Cannot be deployed as currently written.
+  let tacotron_model_check_command = TacotronModelCheckCommand::new(
+    "/models/tts",
+    "source python/bin/activate",
+    "./vocodes_model_check_tacotron.py",
+    Some(DockerOptions {
+      image_name: "5642d0fd7fc1".to_string(), // storyteller-ml
+      maybe_bind_mount: Some(DockerFilesystemMount::tmp_to_tmp()),
+      maybe_gpu: Some(DockerGpu::All),
+    })
+  );
+
+  // TODO/FIXME: Cannot be deployed as currently written.
   let hifigan_model_check_command= HifiGanModelCheckCommand::new(
     //&easyenv::get_env_string_required("HIFIGAN_ROOT_CODE_DIRECTORY")?,
     "/models/tts",
@@ -172,6 +187,7 @@ async fn main() -> AnyhowResult<()> {
     bucket_root_tts_model_uploads: bucket_root.to_string(),
     firehose_publisher,
     badge_granter,
+    tacotron_model_check_command,
     hifigan_model_check_command,
     job_batch_wait_millis: common_env.job_batch_wait_millis,
     job_max_attempts: common_env.job_max_attempts as i32,