From dbd7e949b7bf9fc18908c47d29f88d77528755c5 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Sat, 6 Aug 2022 22:31:20 -0400 Subject: [PATCH] WIP generic-download-job --- .../src/mediators/badge_granter.rs | 89 +++++- .../src/queries/generic_download/job/_keys.rs | 5 + .../list_available_generic_download_jobs.rs | 123 ++++++++ .../job/mark_generic_download_job_done.rs | 39 +++ .../job/mark_generic_download_job_failure.rs | 45 +++ ...eric_download_job_pending_and_grab_lock.rs | 84 ++++++ .../src/queries/generic_download/job/mod.rs | 5 + .../src/queries/generic_download/mod.rs | 1 + .../tts_download_job_queries.rs | 4 - .../generic_download_job/src/job_state.rs | 41 +++ .../generic_download_job/src/job_steps/mod.rs | 2 + .../src/job_steps/process_hifigan_vocoder.rs | 118 ++++++++ .../src/job_steps/process_single_job.rs | 94 ++++++ .../service/generic_download_job/src/main.rs | 277 ++---------------- 14 files changed, 674 insertions(+), 253 deletions(-) create mode 100644 crates/lib/database_queries/src/queries/generic_download/job/_keys.rs create mode 100644 crates/lib/database_queries/src/queries/generic_download/job/list_available_generic_download_jobs.rs create mode 100644 crates/lib/database_queries/src/queries/generic_download/job/mark_generic_download_job_done.rs create mode 100644 crates/lib/database_queries/src/queries/generic_download/job/mark_generic_download_job_failure.rs create mode 100644 crates/lib/database_queries/src/queries/generic_download/job/mark_generic_download_job_pending_and_grab_lock.rs create mode 100644 crates/lib/database_queries/src/queries/generic_download/job/mod.rs create mode 100644 crates/service/generic_download_job/src/job_state.rs create mode 100644 crates/service/generic_download_job/src/job_steps/mod.rs create mode 100644 crates/service/generic_download_job/src/job_steps/process_hifigan_vocoder.rs create mode 100644 crates/service/generic_download_job/src/job_steps/process_single_job.rs diff --git a/crates/lib/database_queries/src/mediators/badge_granter.rs b/crates/lib/database_queries/src/mediators/badge_granter.rs index d771744beb..d3f199d60e 100644 --- a/crates/lib/database_queries/src/mediators/badge_granter.rs +++ b/crates/lib/database_queries/src/mediators/badge_granter.rs @@ -12,7 +12,7 @@ pub enum UserBadge { // Granted for early vocodes users. EarlyUser, - // Granted for uploading models + // Granted for uploading tts models TtsModelUploader1, TtsModelUploader5, TtsModelUploader10, @@ -23,6 +23,17 @@ pub enum UserBadge { TtsModelUploader200, TtsModelUploader250, + // Granted for vocoder models + VocoderModelUploader1, + VocoderModelUploader5, + VocoderModelUploader10, + VocoderModelUploader20, + VocoderModelUploader50, + VocoderModelUploader100, + VocoderModelUploader150, + VocoderModelUploader200, + VocoderModelUploader250, + // Granted for uploading templates W2lTemplateUploader1, W2lTemplateUploader10, @@ -63,6 +74,15 @@ impl UserBadge { UserBadge::TtsModelUploader150 => "tts_model_uploader_150", UserBadge::TtsModelUploader200 => "tts_model_uploader_200", UserBadge::TtsModelUploader250 => "tts_model_uploader_250", + UserBadge::VocoderModelUploader1 => "vocoder_model_uploader_1", + UserBadge::VocoderModelUploader5 => "vocoder_model_uploader_5", + UserBadge::VocoderModelUploader10 => "vocoder_model_uploader_10", + UserBadge::VocoderModelUploader20 => "vocoder_model_uploader_20", + UserBadge::VocoderModelUploader50 => "vocoder_model_uploader_50", + UserBadge::VocoderModelUploader100 => "vocoder_model_uploader_100", + UserBadge::VocoderModelUploader150 => "vocoder_model_uploader_150", + UserBadge::VocoderModelUploader200 => "vocoder_model_uploader_200", + UserBadge::VocoderModelUploader250 => "vocoder_model_uploader_250", UserBadge::W2lTemplateUploader1 => "w2l_template_uploader_1", UserBadge::W2lTemplateUploader10 => "w2l_template_uploader_10", UserBadge::W2lTemplateUploader50 => "w2l_template_uploader_50", @@ -206,6 +226,51 @@ impl BadgeGranter { Ok(()) } + /// This needs to be called *after* successful upload. + pub async fn maybe_grant_vocoder_model_uploads_badge(&self, user_token: &str) -> AnyhowResult<()> { + let count = self.count_vocoder_models_uploaded(user_token).await?; + + let mut maybe_badge = None; + + if count >= 250 { + maybe_badge = Some(UserBadge::VocoderModelUploader250); + } else if count >= 200 { + maybe_badge = Some(UserBadge::VocoderModelUploader200); + } else if count >= 150 { + maybe_badge = Some(UserBadge::VocoderModelUploader150); + } else if count >= 100 { + maybe_badge = Some(UserBadge::VocoderModelUploader100); + } else if count >= 50 { + maybe_badge = Some(UserBadge::VocoderModelUploader50); + } else if count >= 20 { + maybe_badge = Some(UserBadge::VocoderModelUploader20); + } else if count >= 10 { + maybe_badge = Some(UserBadge::VocoderModelUploader10); + } else if count >= 5 { + maybe_badge = Some(UserBadge::VocoderModelUploader5); + } else if count >= 1 { + maybe_badge = Some(UserBadge::VocoderModelUploader1); + } + + let badge = match maybe_badge { + Some(badge) => badge, + None => return Ok(()), + }; + + if self.has_badge(user_token, badge).await? { + return Ok(()) + } + + let _record_id = self.insert( + badge, + user_token, + ).await?; + + self.firehose_publisher.publish_user_badge_granted(user_token, badge.to_db_value()) + .await?; + + Ok(()) + } // ======================================================================= @@ -267,6 +332,28 @@ LIMIT 1 self.handle_count_query(maybe_result) } + async fn count_vocoder_models_uploaded( + &self, + user_token: &str, + ) -> AnyhowResult { + // NB: This could get expensive! + let maybe_result = sqlx::query_as!( + CountRecord, + r#" +SELECT count(*) as count +FROM vocoder_models +WHERE + creator_user_token = ? +LIMIT 1 + "#, + user_token + ) + .fetch_one(&self.mysql_pool) + .await; + + self.handle_count_query(maybe_result) + } + async fn count_w2l_templates_uploaded( &self, user_token: &str, diff --git a/crates/lib/database_queries/src/queries/generic_download/job/_keys.rs b/crates/lib/database_queries/src/queries/generic_download/job/_keys.rs new file mode 100644 index 0000000000..a2ade0e9f9 --- /dev/null +++ b/crates/lib/database_queries/src/queries/generic_download/job/_keys.rs @@ -0,0 +1,5 @@ +//! Primary and foreign keys + +/// Since jobs are internal, we'll use IDs as primary keys. +#[derive(Clone, Copy, PartialOrd, PartialEq, Eq, Ord, Debug)] +pub struct GenericDownloadJobId(pub i64); diff --git a/crates/lib/database_queries/src/queries/generic_download/job/list_available_generic_download_jobs.rs b/crates/lib/database_queries/src/queries/generic_download/job/list_available_generic_download_jobs.rs new file mode 100644 index 0000000000..2e83e136b7 --- /dev/null +++ b/crates/lib/database_queries/src/queries/generic_download/job/list_available_generic_download_jobs.rs @@ -0,0 +1,123 @@ +use anyhow::anyhow; +use chrono::Utc; +use container_common::anyhow_result::AnyhowResult; +use crate::column_types::job_status::JobStatus; +use crate::queries::generic_download::job::_keys::GenericDownloadJobId; +use crate::tokens::Tokens; +use reusable_types::entity_visibility::EntityVisibility; +use reusable_types::generic_download_type::GenericDownloadType; +use sqlx::MySqlPool; +use std::path::Path; + +/// table: generic_download_jobs +#[derive(Debug)] +pub struct AvailableDownloadJob { + pub id: GenericDownloadJobId, + pub download_job_token: String, + + pub creator_user_token: String, + pub creator_ip_address: String, + pub creator_set_visibility: EntityVisibility, + + pub download_type: GenericDownloadType, + pub download_url: Option, + pub title: String, + + pub status: JobStatus, + pub attempt_count: i32, + pub failure_reason: Option, + + pub created_at: chrono::DateTime, + pub updated_at: chrono::DateTime, + pub retry_at: Option>, +} + +pub async fn list_available_generic_download_jobs(pool: &MySqlPool, num_records: u32) + -> AnyhowResult> +{ + let job_records = sqlx::query_as!( + AvailableDownloadJobRawInternal, + r#" +SELECT + id, + token AS download_job_token, + + creator_user_token, + creator_ip_address, + creator_set_visibility as `creator_set_visibility: reusable_types::entity_visibility::EntityVisibility`, + + download_type as `download_type: reusable_types::generic_download_type::GenericDownloadType`, + download_url, + title, + + status as `status: crate::column_types::job_status::JobStatus`, + attempt_count, + failure_reason, + + created_at, + updated_at, + retry_at +FROM generic_download_jobs +WHERE + ( + status IN ("pending", "attempt_failed") + ) + AND + ( + retry_at IS NULL + OR + retry_at < CURRENT_TIMESTAMP + ) + ORDER BY id ASC + LIMIT ? + "#, + num_records, + ) + .fetch_all(pool) + .await?; + + let job_records = job_records.into_iter() + .map(|record : AvailableDownloadJobRawInternal| { + AvailableDownloadJob { + id: GenericDownloadJobId(record.id), + download_job_token: record.download_job_token, + creator_ip_address: record.creator_ip_address, + creator_user_token: record.creator_user_token, + creator_set_visibility: record.creator_set_visibility, + download_type: record.download_type, + download_url: record.download_url, + title: record.title, + status: record.status, + attempt_count: record.attempt_count, + failure_reason: record.failure_reason, + created_at: record.created_at, + updated_at: record.updated_at, + retry_at: record.retry_at, + } + }) + .collect::>(); + + Ok(job_records) +} + +#[derive(Debug)] +struct AvailableDownloadJobRawInternal { + pub id: i64, + pub download_job_token: String, + + pub creator_user_token: String, + pub creator_ip_address: String, + pub creator_set_visibility: EntityVisibility, + + pub download_type: GenericDownloadType, + pub download_url: Option, + pub title: String, + + pub status: JobStatus, + pub attempt_count: i32, + pub failure_reason: Option, + + pub created_at: chrono::DateTime, + pub updated_at: chrono::DateTime, + pub retry_at: Option>, +} diff --git a/crates/lib/database_queries/src/queries/generic_download/job/mark_generic_download_job_done.rs b/crates/lib/database_queries/src/queries/generic_download/job/mark_generic_download_job_done.rs new file mode 100644 index 0000000000..fdc6de7e67 --- /dev/null +++ b/crates/lib/database_queries/src/queries/generic_download/job/mark_generic_download_job_done.rs @@ -0,0 +1,39 @@ +use anyhow::anyhow; +use container_common::anyhow_result::AnyhowResult; +use crate::queries::generic_download::job::list_available_generic_download_jobs::AvailableDownloadJob; +use sqlx::MySqlPool; + +pub async fn mark_generic_download_job_done( + pool: &MySqlPool, + job: &AvailableDownloadJob, + success: bool, + maybe_entity_token: Option<&str>, + maybe_entity_type: Option<&str>, +) -> AnyhowResult<()> +{ + let status = if success { "complete_success" } else { "complete_failure" }; + + let query_result = sqlx::query!( + r#" +UPDATE generic_download_jobs +SET + status = ?, + on_success_downloaded_entity_token = ?, + on_success_downloaded_entity_type = ?, + failure_reason = NULL, + retry_at = NULL +WHERE id = ? + "#, + status, + maybe_entity_token, + maybe_entity_type, + job.id.0, + ) + .execute(pool) + .await; + + match query_result { + Err(err) => Err(anyhow!("error with query: {:?}", err)), + Ok(_r) => Ok(()), + } +} diff --git a/crates/lib/database_queries/src/queries/generic_download/job/mark_generic_download_job_failure.rs b/crates/lib/database_queries/src/queries/generic_download/job/mark_generic_download_job_failure.rs new file mode 100644 index 0000000000..b6de45c92c --- /dev/null +++ b/crates/lib/database_queries/src/queries/generic_download/job/mark_generic_download_job_failure.rs @@ -0,0 +1,45 @@ +use anyhow::anyhow; +use container_common::anyhow_result::AnyhowResult; +use crate::queries::generic_download::job::list_available_generic_download_jobs::AvailableDownloadJob; +use sqlx::MySqlPool; +use sqlx; + +/// Mark a single download job failure. The job may be re-run. +pub async fn mark_generic_download_job_failure( + pool: &MySqlPool, + job: &AvailableDownloadJob, + failure_reason: &str, + max_attempts: i32 +) -> AnyhowResult<()> { + + // statuses: "attempt_failed", "complete_failure", "dead" + let mut next_status = "attempt_failed"; + + if job.attempt_count >= max_attempts { + // NB: Job attempt count is incremented at start + next_status = "dead"; + } + + let query_result = sqlx::query!( + r#" +UPDATE generic_download_jobs +SET + status = ?, + failure_reason = ?, + retry_at = NOW() + interval 2 minute +WHERE id = ? + "#, + next_status, + failure_reason.to_string(), + job.id.0, + ) + .execute(pool) + .await; + + match query_result { + Err(err) => { + Err(anyhow!("error with query: {:?}", err)) + }, + Ok(_r) => Ok(()), + } +} diff --git a/crates/lib/database_queries/src/queries/generic_download/job/mark_generic_download_job_pending_and_grab_lock.rs b/crates/lib/database_queries/src/queries/generic_download/job/mark_generic_download_job_pending_and_grab_lock.rs new file mode 100644 index 0000000000..532ea770f9 --- /dev/null +++ b/crates/lib/database_queries/src/queries/generic_download/job/mark_generic_download_job_pending_and_grab_lock.rs @@ -0,0 +1,84 @@ +use anyhow::anyhow; +use container_common::anyhow_result::AnyhowResult; +use sqlx::MySqlPool; +use sqlx; +use crate::queries::generic_download::job::_keys::GenericDownloadJobId; + +// TODO(2022-02-09): It's unclear how stuck "pending" jobs don't become perma-dead +// I think this logic is flawed and the job system needs to pick up stuck "pending" jobs. + +pub struct GenericDownloadJobLockRecord { + id: i64, + status: String, +} + +pub async fn mark_generic_download_job_pending_and_grab_lock( + pool: &MySqlPool, + job_id: GenericDownloadJobId +) -> AnyhowResult { + + // NB: We use transactions and "SELECT ... FOR UPDATE" to simulate mutexes. + let mut transaction = pool.begin().await?; + + let maybe_record = sqlx::query_as!( + TtsInferenceLockRecord, + r#" +SELECT + id, + status +FROM generic_download_jobs +WHERE id = ? +FOR UPDATE + "#, + job_id.0, + ) + .fetch_one(&mut transaction) + .await; + + let record : GenericDownloadJobLockRecord = match maybe_record { + Ok(record) => record, + Err(err) => { + match err { + sqlx::Error::RowNotFound => { + return Err(anyhow!("could not job")); + }, + _ => { + return Err(anyhow!("query error")); + } + } + } + }; + + let can_transact = match record.status.as_ref() { + "pending" => true, // It's okay for us to take the lock. + "attempt_failed" => true, // We can retry. + "started" => false, // Job in progress (another job beat us, and we can't take the lock) + "complete_success" => false, // Job already complete + "complete_failure" => false, // Job already complete (permanently dead; no need to retry) + "dead" => false, // Job already complete (permanently dead; retries exhausted) + _ => false, // Future-proof + }; + + if !can_transact { + transaction.rollback().await?; + return Ok(false); + } + + let _acquire_lock = sqlx::query!( + r#" +UPDATE generic_download_jobs +SET + status = 'started', + attempt_count = attempt_count + 1, + retry_at = NOW() + interval 2 minute +WHERE id = ? + "#, + job_id.0, + ) + .execute(&mut transaction) + .await?; + + transaction.commit().await?; + + Ok(true) +} diff --git a/crates/lib/database_queries/src/queries/generic_download/job/mod.rs b/crates/lib/database_queries/src/queries/generic_download/job/mod.rs new file mode 100644 index 0000000000..ac425ffd26 --- /dev/null +++ b/crates/lib/database_queries/src/queries/generic_download/job/mod.rs @@ -0,0 +1,5 @@ +pub mod _keys; +pub mod list_available_generic_download_jobs; +pub mod mark_generic_download_job_done; +pub mod mark_generic_download_job_failure; +pub mod mark_generic_download_job_pending_and_grab_lock; \ No newline at end of file diff --git a/crates/lib/database_queries/src/queries/generic_download/mod.rs b/crates/lib/database_queries/src/queries/generic_download/mod.rs index 4def6e37b1..4feea997bc 100644 --- a/crates/lib/database_queries/src/queries/generic_download/mod.rs +++ b/crates/lib/database_queries/src/queries/generic_download/mod.rs @@ -1,2 +1,3 @@ pub mod get_generic_download_job_status; pub mod insert_generic_download_job; +pub mod job; diff --git a/crates/lib/database_queries/src/queries/tts/tts_download_jobs/tts_download_job_queries.rs b/crates/lib/database_queries/src/queries/tts/tts_download_jobs/tts_download_job_queries.rs index 558f800cfd..e99c1e1317 100644 --- a/crates/lib/database_queries/src/queries/tts/tts_download_jobs/tts_download_job_queries.rs +++ b/crates/lib/database_queries/src/queries/tts/tts_download_jobs/tts_download_job_queries.rs @@ -1,7 +1,3 @@ -//! NB: This seems required for sqlx to generate the cached queries. -//! Sqlx's prepare needs a *single* binary to work against, so we need to -//! include these in the main binary to generate all the queries. - use anyhow::anyhow; use chrono::Utc; use container_common::anyhow_result::AnyhowResult; diff --git a/crates/service/generic_download_job/src/job_state.rs b/crates/service/generic_download_job/src/job_state.rs new file mode 100644 index 0000000000..c92c1817fd --- /dev/null +++ b/crates/service/generic_download_job/src/job_state.rs @@ -0,0 +1,41 @@ +use crate::HifiGanModelCheckCommand; +use database_queries::mediators::badge_granter::BadgeGranter; +use database_queries::mediators::firehose_publisher::FirehosePublisher; +use google_drive_common::google_drive_download_command::GoogleDriveDownloadCommand; +use r2d2_redis::RedisConnectionManager; +use r2d2_redis::r2d2; +use sqlx::MySqlPool; +use std::path::PathBuf; +use storage_buckets_common::bucket_client::BucketClient; +use storage_buckets_common::bucket_path_unifier::BucketPathUnifier; + +pub struct JobState { + pub download_temp_directory: PathBuf, + pub mysql_pool: MySqlPool, + + pub redis_pool: r2d2::Pool, + + pub bucket_client: BucketClient, + pub firehose_publisher: FirehosePublisher, + pub badge_granter: BadgeGranter, + pub google_drive_downloader: GoogleDriveDownloadCommand, + + pub bucket_path_unifier: BucketPathUnifier, + + pub hifigan_model_check_command: HifiGanModelCheckCommand, + + // Command to run + pub download_script: String, + // Root to store TTS results + pub bucket_root_tts_model_uploads: String, + + // Sleep between batches + pub job_batch_wait_millis: u64, + + // How long to wait between log lines + pub no_op_logger_millis: u64, + + // Max job attempts before failure. + // NB: This is an i32 so we don't need to convert to db column type. + pub job_max_attempts: i32, +} diff --git a/crates/service/generic_download_job/src/job_steps/mod.rs b/crates/service/generic_download_job/src/job_steps/mod.rs new file mode 100644 index 0000000000..c1accdd4bb --- /dev/null +++ b/crates/service/generic_download_job/src/job_steps/mod.rs @@ -0,0 +1,2 @@ +pub mod process_single_job; +pub mod process_hifigan_vocoder; \ No newline at end of file diff --git a/crates/service/generic_download_job/src/job_steps/process_hifigan_vocoder.rs b/crates/service/generic_download_job/src/job_steps/process_hifigan_vocoder.rs new file mode 100644 index 0000000000..7ff26e4e68 --- /dev/null +++ b/crates/service/generic_download_job/src/job_steps/process_hifigan_vocoder.rs @@ -0,0 +1,118 @@ +use std::path::PathBuf; +use log::info; +use tempdir::TempDir; +use container_common::anyhow_result::AnyhowResult; +use database_queries::queries::generic_download::job::list_available_generic_download_jobs::AvailableDownloadJob; +use crate::JobState; + +pub async fn process_hifigan_vocoder( + job_state: &JobState, + job: &AvailableDownloadJob, + temp_dir: &TempDir +) -> AnyhowResult<()> { + + // ==================== RUN MODEL CHECK ==================== // + + info!("Checking that model is valid..."); + + redis_logger.log_status("checking model")?; + + let file_path = PathBuf::from(download_filename.clone()); + + let output_metadata_fs_path = temp_dir.path().join("metadata.json"); + + match model_type { + TtsModelType::Tacotron2 => { + let result = job_state.tacotron_tts_check.execute( + &file_path, + &output_metadata_fs_path, + false, + ); + + if let Err(e) = result { + safe_delete_temp_file(&file_path); + safe_delete_temp_directory(&temp_dir); + } + }, + TtsModelType::Talknet => { + let result = job_state.talknet_tts_check.execute( + &file_path, + &output_metadata_fs_path, + false, + ); + + if let Err(e) = result { + safe_delete_temp_file(&file_path); + safe_delete_temp_directory(&temp_dir); + } + }, + } + + // ==================== 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 model to GCS..."); + + let private_bucket_hash = hash_file_sha2(&download_filename)?; + + info!("File hash: {}", private_bucket_hash); + + let synthesizer_model_bucket_path = match model_type { + TtsModelType::Tacotron2 => job_state.bucket_path_unifier.tts_synthesizer_path(&private_bucket_hash), + TtsModelType::Talknet => job_state.bucket_path_unifier.tts_zipped_synthesizer_path(&private_bucket_hash), + }; + + info!("Destination bucket path: {:?}", &synthesizer_model_bucket_path); + + redis_logger.log_status("uploading model")?; + + if let Err(e) = job_state.bucket_client.upload_filename(&synthesizer_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 model record..."); + let (id, model_token) = insert_tts_model( + &job_state.mysql_pool, + job, + &private_bucket_hash, + synthesizer_model_bucket_path, + file_metadata.file_size_bytes) + .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(()) +} diff --git a/crates/service/generic_download_job/src/job_steps/process_single_job.rs b/crates/service/generic_download_job/src/job_steps/process_single_job.rs new file mode 100644 index 0000000000..18718ad7b8 --- /dev/null +++ b/crates/service/generic_download_job/src/job_steps/process_single_job.rs @@ -0,0 +1,94 @@ +use anyhow::anyhow; +use config::bad_urls::is_bad_tts_model_download_url; +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 database_queries::queries::generic_download::job::list_available_generic_download_jobs::AvailableDownloadJob; +use database_queries::queries::generic_download::job::mark_generic_download_job_pending_and_grab_lock::mark_generic_download_job_pending_and_grab_lock; +use jobs_common::redis_job_status_logger::RedisJobStatusLogger; +use log::{info, warn}; +use std::path::PathBuf; +use tempdir::TempDir; +use database_queries::queries::generic_download::job::mark_generic_download_job_done::mark_generic_download_job_done; +use crate::job_steps::process_hifigan_vocoder::process_hifigan_vocoder; + +pub async fn process_single_job(job_state: &JobState, job: &AvailableDownloadJob) -> AnyhowResult<()> { + let mut redis = job_state.redis_pool.get()?; + let mut redis_logger = RedisJobStatusLogger::new_tts_download( + &mut redis, + &job.token); + + // ==================== ATTEMPT TO GRAB JOB LOCK ==================== // + + let lock_acquired = mark_generic_download_job_pending_and_grab_lock(&job_state.mysql_pool, job.id).await?; + + if !lock_acquired { + warn!("Could not acquire job lock for: {}", &job.id.0); + return Ok(()) + } + + // ==================== SETUP TEMP DIRS ==================== // + + let temp_dir = format!("temp_{}", job.id.0); + let temp_dir = TempDir::new(&temp_dir)?; + + // ==================== DOWNLOAD MODEL FILE ==================== // + + info!("Calling downloader..."); + + redis_logger.log_status("downloading model")?; + + let download_url = job.download_url.as_ref() + .map(|c| c.to_string()) + .unwrap_or("".to_string()); + + if is_bad_tts_model_download_url(&download_url)? { + warn!("Bad download URL: `{}`", &download_url); + return Err(anyhow!("Bad download URL: `{}`", &download_url)); + } + + let download_filename = match job_state.google_drive_downloader.download_file(&download_url, &temp_dir).await { + Ok(filename) => filename, + Err(e) => { + safe_delete_temp_directory(&temp_dir); + return Err(e); + } + }; + + // ==================== HANDLE DIFFERENT DOWNLOAD TYPES ==================== // + + match job.download_type { + GenericDownloadType::HifiGan => { + process_hifigan_vocoder(job_state, job, &temp_dir).await?; + } + } + + // ===================================================== + + info!("Marking job complete..."); + mark_generic_download_job_done( + &job_state.mysql_pool, + job, + true, + Some(&entity_token), + Some(&entity_type), + ).await?; + + info!("Saved model record: {}", id); + + job_state.firehose_publisher.publish_generic_download_finished(&job.creator_user_token, &model_token) + .await + .map_err(|e| { + warn!("error publishing event: {:?}", e); + anyhow!("error publishing event") + })?; + + redis_logger.log_status("done")?; + + info!("Job done: {}", job.id.0); + + Ok(()) +} diff --git a/crates/service/generic_download_job/src/main.rs b/crates/service/generic_download_job/src/main.rs index 837e0a30cc..5ac646f1b6 100644 --- a/crates/service/generic_download_job/src/main.rs +++ b/crates/service/generic_download_job/src/main.rs @@ -10,6 +10,8 @@ #[macro_use] extern crate serde_derive; +pub mod job_state; +pub mod job_steps; pub mod script_execution; use anyhow::anyhow; @@ -24,16 +26,20 @@ 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::job_state::JobState; +use crate::job_steps::process_single_job::process_single_job; +use crate::script_execution::hifigan_model_check_command::HifiGanModelCheckCommand; use crate::script_execution::tacotron_model_check_command::TacotronModelCheckCommand; use crate::script_execution::talknet_model_check_command::TalknetModelCheckCommand; use database_queries::column_types::tts_model_type::TtsModelType; use database_queries::mediators::badge_granter::BadgeGranter; use database_queries::mediators::firehose_publisher::FirehosePublisher; +use database_queries::queries::generic_download::job::list_available_generic_download_jobs::{AvailableDownloadJob, list_available_generic_download_jobs}; +use database_queries::queries::generic_download::job::mark_generic_download_job_failure::mark_generic_download_job_failure; use database_queries::queries::tts::tts_download_jobs::tts_download_job_queries::TtsUploadJobRecord; use database_queries::queries::tts::tts_download_jobs::tts_download_job_queries::grab_job_lock_and_mark_pending; use database_queries::queries::tts::tts_download_jobs::tts_download_job_queries::insert_tts_model; use database_queries::queries::tts::tts_download_jobs::tts_download_job_queries::mark_tts_upload_job_done; -use database_queries::queries::tts::tts_download_jobs::tts_download_job_queries::mark_tts_upload_job_failure; use database_queries::queries::tts::tts_download_jobs::tts_download_job_queries::query_tts_upload_job_records; use google_drive_common::google_drive_download_command::GoogleDriveDownloadCommand; use jobs_common::noop_logger::NoOpLogger; @@ -53,7 +59,6 @@ use std::time::Duration; use storage_buckets_common::bucket_client::BucketClient; use storage_buckets_common::bucket_path_unifier::BucketPathUnifier; use tempdir::TempDir; -use crate::script_execution::hifigan_model_check_command::HifiGanModelCheckCommand; // Buckets const ENV_ACCESS_KEY : &'static str = "ACCESS_KEY"; @@ -64,37 +69,6 @@ const ENV_BUCKET_ROOT : &'static str = "TTS_DOWNLOAD_BUCKET_ROOT"; const DEFAULT_TEMP_DIR: &'static str = "/tmp"; -struct Downloader { - pub download_temp_directory: PathBuf, - pub mysql_pool: MySqlPool, - - pub redis_pool: r2d2::Pool, - - pub bucket_client: BucketClient, - pub firehose_publisher: FirehosePublisher, - pub badge_granter: BadgeGranter, - pub google_drive_downloader: GoogleDriveDownloadCommand, - - pub bucket_path_unifier: BucketPathUnifier, - - pub hifigan_model_check_command: HifiGanModelCheckCommand, - - // Command to run - pub download_script: String, - // Root to store TTS results - pub bucket_root_tts_model_uploads: String, - - // Sleep between batches - pub job_batch_wait_millis: u64, - - // How long to wait between log lines - pub no_op_logger_millis: u64, - - // Max job attempts before failure. - // NB: This is an i32 so we don't need to convert to db column type. - pub job_max_attempts: i32, -} - #[tokio::main] async fn main() -> AnyhowResult<()> { easyenv::init_all_with_default_logging(Some(DEFAULT_RUST_LOG)); @@ -106,7 +80,7 @@ async fn main() -> AnyhowResult<()> { let server_hostname = hostname::get() .ok() .and_then(|h| h.into_string().ok()) - .unwrap_or("tts-download-job".to_string()); + .unwrap_or("generic-download-job".to_string()); info!("Hostname: {}", &server_hostname); @@ -170,22 +144,17 @@ async fn main() -> AnyhowResult<()> { firehose_publisher: firehose_publisher.clone(), // NB: Also safe }; - let hifigan_root_code_directory = easyenv::get_env_string_required("HIFIGAN_ROOT_CODE_DIRECTORY")?; - let hifigan_virtual_env_activation_command = easyenv::get_env_string_or_default( - "HIFIGAN_VIRTUAL_ENV_ACTIVATION_COMMAND", - "source python-tacotron/bin/activate"); - - let hifigan_model_check_script_name = easyenv::get_env_string_or_default( - "HIFIGAN_MODEL_CHECK_SCRIPT_NAME", - "vocodes_model_check_hifigan.py"); - let hifigan_model_check_command= HifiGanModelCheckCommand::new( - &hifigan_root_code_directory, - &hifigan_virtual_env_activation_command, - &hifigan_model_check_script_name, + &easyenv::get_env_string_required("HIFIGAN_ROOT_CODE_DIRECTORY")?, + &easyenv::get_env_string_or_default( + "HIFIGAN_VIRTUAL_ENV_ACTIVATION_COMMAND", + "source python-tacotron/bin/activate"), + &easyenv::get_env_string_or_default( + "HIFIGAN_MODEL_CHECK_SCRIPT_NAME", + "vocodes_model_check_hifigan.py"), ); - let downloader = Downloader { + let job_state = JobState { download_temp_directory: temp_directory, mysql_pool, redis_pool, @@ -202,7 +171,7 @@ async fn main() -> AnyhowResult<()> { no_op_logger_millis: common_env.no_op_logger_millis, }; - main_loop(downloader).await; + main_loop(job_state).await; Ok(()) } @@ -210,16 +179,16 @@ async fn main() -> AnyhowResult<()> { const START_TIMEOUT_MILLIS : u64 = 500; const INCREASE_TIMEOUT_MILLIS : u64 = 1000; -async fn main_loop(downloader: Downloader) { +async fn main_loop(job_state: JobState) { let mut error_timeout_millis = START_TIMEOUT_MILLIS; - let mut noop_logger = NoOpLogger::new(downloader.no_op_logger_millis as i64); + let mut noop_logger = NoOpLogger::new(job_state.no_op_logger_millis as i64); loop { let num_records = 1; - let query_result = query_tts_upload_job_records(&downloader.mysql_pool, num_records).await; + let maybe_available_jobs = list_available_generic_download_jobs(&job_state.mysql_pool, num_records).await; - let jobs = match query_result { + let jobs = match maybe_available_jobs { Ok(jobs) => jobs, Err(e) => { warn!("Error querying jobs: {:?}", e); @@ -232,11 +201,11 @@ async fn main_loop(downloader: Downloader) { if jobs.is_empty() { noop_logger.log_after_awhile(); - std::thread::sleep(Duration::from_millis(downloader.job_batch_wait_millis)); + std::thread::sleep(Duration::from_millis(job_state.job_batch_wait_millis)); continue; } - let result = process_jobs(&downloader, jobs).await; + let result = process_jobs(&job_state, jobs).await; match result { Ok(_) => {}, @@ -250,23 +219,23 @@ async fn main_loop(downloader: Downloader) { error_timeout_millis = START_TIMEOUT_MILLIS; // reset - std::thread::sleep(Duration::from_millis(downloader.job_batch_wait_millis)); + std::thread::sleep(Duration::from_millis(job_state.job_batch_wait_millis)); } } -async fn process_jobs(downloader: &Downloader, jobs: Vec) -> AnyhowResult<()> { +async fn process_jobs(job_state: &JobState, jobs: Vec) -> AnyhowResult<()> { for job in jobs.into_iter() { - let result = process_job(downloader, &job).await; + let result = process_single_job(job_state, &job).await; match result { Ok(_) => {}, Err(e) => { warn!("Failure to process job: {:?}", e); let failure_reason = ""; - let _r = mark_tts_upload_job_failure( - &downloader.mysql_pool, + let _r = mark_generic_download_job_failure( + &job_state.mysql_pool, &job, failure_reason, - downloader.job_max_attempts + job_state.job_max_attempts ).await; } } @@ -287,191 +256,3 @@ fn read_metadata_file(filename: &PathBuf) -> AnyhowResult { Ok(serde_json::from_str(&buffer)?) } -async fn process_job(downloader: &Downloader, job: &TtsUploadJobRecord) -> AnyhowResult<()> { - // TODO: 1. Mark processing. (DONE) - // TODO: 2. Download. (DONE) - // TODO: 3. Upload. (DONE) - // TODO: 4. Save record. (DONE) - // TODO: 5. Mark job done. (DONE) - - let mut redis = downloader.redis_pool.get()?; - let mut redis_logger = RedisJobStatusLogger::new_tts_download( - &mut redis, - &job.token); - - // ==================== ATTEMPT TO GRAB JOB LOCK ==================== // - - let lock_acquired = grab_job_lock_and_mark_pending(&downloader.mysql_pool, job).await?; - - if !lock_acquired { - warn!("Could not acquire job lock for: {}", &job.id); - return Ok(()) - } - - // ==================== SETUP TEMP DIRS ==================== // - - let temp_dir = format!("temp_{}", job.id); - let temp_dir = TempDir::new(&temp_dir)?; - - // ==================== DOWNLOAD MODEL FILE ==================== // - - info!("Calling downloader..."); - - redis_logger.log_status("downloading model")?; - - let download_url = job.download_url.as_ref() - .map(|c| c.to_string()) - .unwrap_or("".to_string()); - - if is_bad_tts_model_download_url(&download_url)? { - warn!("Bad download URL: `{}`", &download_url); - return Err(anyhow!("Bad download URL: `{}`", &download_url)); - } - - let download_filename = match downloader.google_drive_downloader.download_file(&download_url, &temp_dir).await { - Ok(filename) => filename, - Err(e) => { - safe_delete_temp_directory(&temp_dir); - return Err(e); - } - }; - - let model_type = if download_filename.to_lowercase().ends_with("zip") { - TtsModelType::Talknet - } else { - TtsModelType::Tacotron2 - }; - - info!("Uploaded model type: {:?}", model_type); - - if model_type == TtsModelType::Talknet { - // TODO: Finish supporting TalkNet - warn!("Unsupported model type"); - return Err(anyhow!("unsupported model type")); - } - - // ==================== RUN MODEL CHECK ==================== // - - info!("Checking that model is valid..."); - - redis_logger.log_status("checking model")?; - - let file_path = PathBuf::from(download_filename.clone()); - - let output_metadata_fs_path = temp_dir.path().join("metadata.json"); - - match model_type { - TtsModelType::Tacotron2 => { - let result = downloader.tacotron_tts_check.execute( - &file_path, - &output_metadata_fs_path, - false, - ); - - if let Err(e) = result { - safe_delete_temp_file(&file_path); - safe_delete_temp_directory(&temp_dir); - } - }, - TtsModelType::Talknet => { - let result = downloader.talknet_tts_check.execute( - &file_path, - &output_metadata_fs_path, - false, - ); - - if let Err(e) = result { - safe_delete_temp_file(&file_path); - safe_delete_temp_directory(&temp_dir); - } - }, - } - - // ==================== 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 model to GCS..."); - - let private_bucket_hash = hash_file_sha2(&download_filename)?; - - info!("File hash: {}", private_bucket_hash); - - let synthesizer_model_bucket_path = match model_type { - TtsModelType::Tacotron2 => downloader.bucket_path_unifier.tts_synthesizer_path(&private_bucket_hash), - TtsModelType::Talknet => downloader.bucket_path_unifier.tts_zipped_synthesizer_path(&private_bucket_hash), - }; - - info!("Destination bucket path: {:?}", &synthesizer_model_bucket_path); - - redis_logger.log_status("uploading model")?; - - if let Err(e) = downloader.bucket_client.upload_filename(&synthesizer_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 model record..."); - let (id, model_token) = insert_tts_model( - &downloader.mysql_pool, - job, - &private_bucket_hash, - synthesizer_model_bucket_path, - file_metadata.file_size_bytes) - .await?; - - info!("Marking job complete..."); - mark_tts_upload_job_done( - &downloader.mysql_pool, - job, - true, - Some(&model_token) - ).await?; - - info!("Saved model record: {}", id); - - downloader.firehose_publisher.publish_tts_model_upload_finished(&job.creator_user_token, &model_token) - .await - .map_err(|e| { - warn!("error publishing event: {:?}", e); - anyhow!("error publishing event") - })?; - - downloader.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") - })?; - - redis_logger.log_status("done")?; - - info!("Job done: {}", job.id); - - Ok(()) -}