From b128cdb3f5da9e7dc5e81638ef8cf8f46525eb28 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Mon, 2 Jan 2023 06:34:35 -0500 Subject: [PATCH] inference job handlers --- .../src/job_steps/process_single_job.rs | 4 +- .../src/job_steps/process_single_job_old.rs | 42 +-- .../src/job_types/tts/handle_tts_inference.rs | 344 ++++++++++++++++++ .../inference_job/src/job_types/tts/mod.rs | 2 + .../src/job_types/tts/read_metadata_file.rs | 18 + 5 files changed, 379 insertions(+), 31 deletions(-) create mode 100644 crates/service/job/inference_job/src/job_types/tts/handle_tts_inference.rs create mode 100644 crates/service/job/inference_job/src/job_types/tts/read_metadata_file.rs diff --git a/crates/service/job/inference_job/src/job_steps/process_single_job.rs b/crates/service/job/inference_job/src/job_steps/process_single_job.rs index f91ac536f7..abcdb2631a 100644 --- a/crates/service/job/inference_job/src/job_steps/process_single_job.rs +++ b/crates/service/job/inference_job/src/job_steps/process_single_job.rs @@ -28,10 +28,8 @@ pub async fn process_single_job(job_dependencies: &JobDependencies, job: &Availa // ==================== SETUP TEMP DIRS ==================== // - // TODO: Scoped temp dir - let temp_dir = format!("temp_{}", job.id.0); - let temp_dir = TempDir::new(&temp_dir)?; + let temp_dir = job_dependencies.scoped_temp_dir_creator.new_tempdir(&temp_dir)?; // ==================== HANDLE DIFFERENT INFERENCE TYPES ==================== // diff --git a/crates/service/job/inference_job/src/job_steps/process_single_job_old.rs b/crates/service/job/inference_job/src/job_steps/process_single_job_old.rs index 049d443927..7b9bde648b 100644 --- a/crates/service/job/inference_job/src/job_steps/process_single_job_old.rs +++ b/crates/service/job/inference_job/src/job_steps/process_single_job_old.rs @@ -31,20 +31,6 @@ use tts_common::clean_symbols::clean_symbols; use tts_common::text_pipelines::guess_pipeline::guess_text_pipeline_heuristic; use tts_common::text_pipelines::text_pipeline_type::TextPipelineType; -#[derive(Deserialize, Default)] -struct FileMetadata { - pub duration_millis: Option, - pub mimetype: Option, - 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)?) -} - pub async fn process_single_job_old( job_args: &JobDependencies, job: &AvailableTtsInferenceJob, @@ -53,24 +39,24 @@ pub async fn process_single_job_old( let start = Instant::now(); - let span_id = generate_random_uuid(); - let trace_id = generate_random_uuid(); - let maybe_user_token = job.maybe_creator_user_token.as_deref().unwrap_or(""); + //let span_id = generate_random_uuid(); + //let trace_id = generate_random_uuid(); + //let maybe_user_token = job.maybe_creator_user_token.as_deref().unwrap_or(""); - let mut job_iteration_span = Span::new(&span_id, &trace_id, get_timestamp_millis()) - .name("single job execution") - .attribute("user_token", maybe_user_token) - .service_name("tts-inference-job"); + //let mut job_iteration_span = Span::new(&span_id, &trace_id, get_timestamp_millis()) + // .name("single job execution") + // .attribute("user_token", maybe_user_token) + // .service_name("tts-inference-job"); - let span_id = generate_random_uuid(); - let trace_id = generate_random_uuid(); + //let span_id = generate_random_uuid(); + //let trace_id = generate_random_uuid(); - let created_at_timestamp = (job.created_at.timestamp() as u64) * 1000; + //let created_at_timestamp = (job.created_at.timestamp() as u64) * 1000; - let mut since_creation_span = Span::new(&span_id, &trace_id, created_at_timestamp) - .name("job since creation") - .attribute("user_token", maybe_user_token) - .service_name("tts-inference-job"); + //let mut since_creation_span = Span::new(&span_id, &trace_id, created_at_timestamp) + // .name("job since creation") + // .attribute("user_token", maybe_user_token) + // .service_name("tts-inference-job"); let mut job_progress_reporter = job_args .job_progress_reporter diff --git a/crates/service/job/inference_job/src/job_types/tts/handle_tts_inference.rs b/crates/service/job/inference_job/src/job_types/tts/handle_tts_inference.rs new file mode 100644 index 0000000000..6e4a882ff7 --- /dev/null +++ b/crates/service/job/inference_job/src/job_types/tts/handle_tts_inference.rs @@ -0,0 +1,344 @@ +use anyhow::anyhow; +use log::{info, warn}; +use tempdir::TempDir; +use crate::job_steps::job_dependencies::JobDependencies; +use crate::job_steps::process_single_job_error::ProcessSingleJobError; +use database_queries::queries::generic_inference::job::list_available_generic_inference_jobs::AvailableInferenceJob; +use errors::AnyhowResult; +use tts_common::clean_symbols::clean_symbols; +use crate::util::download_file_from_bucket::maybe_download_file_from_bucket; + +pub async fn handle_tts_inference( + job_dependencies: &JobDependencies, + job: &AvailableInferenceJob, +) -> AnyhowResult<()> { + + let mut job_progress_reporter = job_dependencies + .job_progress_reporter + .new_tts_inference(job.inference_job_token.as_str()) + .map_err(|e| ProcessSingleJobError::Other(anyhow!(e)))?; + + // ==================== CONFIRM OR DOWNLOAD WAVEGLOW VOCODER MODEL ==================== // + + let waveglow_vocoder_model_fs_path = { + let waveglow_vocoder_model_filename = job_dependencies.waveglow_vocoder_model_filename.clone(); + let waveglow_vocoder_model_fs_path = job_dependencies.semi_persistent_cache.tts_pretrained_vocoder_model_path(&waveglow_vocoder_model_filename); + let waveglow_vocoder_model_object_path = job_dependencies.bucket_path_unifier.tts_pretrained_vocoders_path(&waveglow_vocoder_model_filename); + + maybe_download_file_from_bucket( + "waveglow vocoder model", + &waveglow_vocoder_model_fs_path, + &waveglow_vocoder_model_object_path, + &job_dependencies.private_bucket_client, + &mut job_progress_reporter, + "downloading vocoder (1 of 3)", + job.id.0, + &job_dependencies.scoped_temp_dir_creator, + ).await?; + + waveglow_vocoder_model_fs_path + }; + + // ==================== CONFIRM OR DOWNLOAD HIFIGAN (NORMAL) VOCODER MODEL ==================== // + + let pretrained_hifigan_vocoder_model_fs_path = { + let hifigan_vocoder_model_filename = job_dependencies.hifigan_vocoder_model_filename.clone(); + let hifigan_vocoder_model_fs_path = job_dependencies.semi_persistent_cache.tts_pretrained_vocoder_model_path(&hifigan_vocoder_model_filename); + let hifigan_vocoder_model_object_path = job_dependencies.bucket_path_unifier.tts_pretrained_vocoders_path(&hifigan_vocoder_model_filename); + + maybe_download_file_from_bucket( + "hifigan vocoder model", + &hifigan_vocoder_model_fs_path, + &hifigan_vocoder_model_object_path, + &job_dependencies.private_bucket_client, + &mut job_progress_reporter, + "downloading vocoder (2 of 3)", + job.id.0, + &job_dependencies.scoped_temp_dir_creator, + ).await?; + + hifigan_vocoder_model_fs_path + }; + + // ==================== CONFIRM OR DOWNLOAD HIFIGAN (SUPERRES) VOCODER MODEL ==================== // + + let hifigan_superres_vocoder_model_fs_path = { + let hifigan_superres_vocoder_model_filename = job_dependencies.hifigan_superres_vocoder_model_filename.clone(); + let hifigan_superres_vocoder_model_fs_path = job_dependencies.semi_persistent_cache.tts_pretrained_vocoder_model_path(&hifigan_superres_vocoder_model_filename); + let hifigan_superres_vocoder_model_object_path = job_dependencies.bucket_path_unifier.tts_pretrained_vocoders_path(&hifigan_superres_vocoder_model_filename); + + maybe_download_file_from_bucket( + "hifigan superres vocoder model", + &hifigan_superres_vocoder_model_fs_path, + &hifigan_superres_vocoder_model_object_path, + &job_dependencies.private_bucket_client, + &mut job_progress_reporter, + "downloading vocoder (3 of 3)", + job.id.0, + &job_dependencies.scoped_temp_dir_creator, + ).await?; + + hifigan_superres_vocoder_model_fs_path + }; + +// // ==================== CONFIRM OR DOWNLOAD OPTIONAL CUSTOM VOCODER MODEL ==================== // + + let custom_vocoder_fs_path = match &model_record.maybe_custom_vocoder { + None => None, + Some(vocoder) => { + let custom_vocoder_fs_path = job_dependencies.semi_persistent_cache.custom_vocoder_model_path(&vocoder.vocoder_token); + let custom_vocoder_object_path = job_dependencies.bucket_path_unifier.vocoder_path(&vocoder.vocoder_private_bucket_hash); + + maybe_download_file_from_bucket( + "custom vocoder", + &custom_vocoder_fs_path, + &custom_vocoder_object_path, + &job_dependencies.private_bucket_client, + &mut job_progress_reporter, + "downloading user vocoder", + job.id.0, + &job_dependencies.scoped_temp_dir_creator, + ).await?; + + Some(custom_vocoder_fs_path) + } + }; + + // ==================== CONFIRM OR DOWNLOAD TTS SYNTHESIZER MODEL ==================== // + + let tts_synthesizer_fs_path = { + let tts_synthesizer_fs_path = job_dependencies.semi_persistent_cache.tts_synthesizer_model_path(&model_record.model_token); + let tts_synthesizer_object_path = job_dependencies.bucket_path_unifier.tts_synthesizer_path(&model_record.private_bucket_hash); + + maybe_download_file_from_bucket( + "synthesizer", + &tts_synthesizer_fs_path, + &tts_synthesizer_object_path, + &job_dependencies.private_bucket_client, + &mut job_progress_reporter, + "downloading synthesizer", + job.id.0, + &job_dependencies.scoped_temp_dir_creator, + ).await?; + + tts_synthesizer_fs_path + }; + + // ==================== Preprocess text ==================== // + + let cleaned_inference_text = clean_symbols(&job.raw_inference_text); + + // ==================== WRITE TEXT TO FILE ==================== // + + info!("Creating tempdir for inference results."); + + let temp_dir = format!("temp_tts_inference_{}", job.id.0); + + // NB: TempDir exists until it goes out of scope, at which point it should delete from filesystem. + let temp_dir = TempDir::new(&temp_dir) + .map_err(|e| ProcessSingleJobError::from_io_error(e))?; + + let text_input_fs_path = temp_dir.path().join("inference_input.txt"); + + std::fs::write(&text_input_fs_path, &cleaned_inference_text) + .map_err(|e| ProcessSingleJobError::from_io_error(e))?; + + // ==================== RUN INFERENCE ==================== // + + job_progress_reporter.log_status("running inference") + .map_err(|e| ProcessSingleJobError::Other(e))?; + + // TODO: Fix this. + let maybe_unload_model_path = job_dependencies + .virtual_model_lfu + .insert_returning_replaced(tts_synthesizer_fs_path.to_str().unwrap_or("")) + .map_err(|e| ProcessSingleJobError::Other(e))?; + + if let Some(model_path) = maybe_unload_model_path.as_deref() { + warn!("Remove model from LFU cache: {:?}", model_path); + } + + let output_audio_fs_path = temp_dir.path().join("output.wav"); + let output_metadata_fs_path = temp_dir.path().join("metadata.json"); + let output_spectrogram_fs_path = temp_dir.path().join("spectrogram.json"); + + info!("Running TTS inference..."); + + info!("Expected output audio filename: {:?}", &output_audio_fs_path); + info!("Expected output spectrogram filename: {:?}", &output_spectrogram_fs_path); + info!("Expected output metadata filename: {:?}", &output_metadata_fs_path); + + if let Some(model_path) = maybe_unload_model_path.as_deref() { + warn!("Unload model from sidecar: {:?}", &model_path); + } + + let mut pretrained_vocoder = VocoderType::HifiGanSuperResolution; + if let Some(default_vocoder) = model_record.maybe_default_pretrained_vocoder.as_deref() { + pretrained_vocoder = VocoderType::from_str(default_vocoder) + .map_err(|e| ProcessSingleJobError::Other(e))?; + } + + info!("With pretrained vocoder: {:?}", pretrained_vocoder); + + let text_pipeline_type_or_guess = model_record.text_pipeline_type + .as_deref() + .and_then(|pipeline_type| + TextPipelineType::from_str(pipeline_type).ok())// NB: If there's an error deserializing, turn it to None. + .unwrap_or_else(|| + guess_text_pipeline_heuristic(Some(model_record.created_at))); + + info!("With text pipeline type `{:?} ` (or guess: {:?})", + &model_record.text_pipeline_type, + &text_pipeline_type_or_guess); + + let hifigan_vocoder_model_fs_path_to_use = match custom_vocoder_fs_path { + None => { + info!("using pretrained HiFi-GAN vocoder"); + pretrained_hifigan_vocoder_model_fs_path + }, + Some(custom_vocoder_fs_path) => { + info!("using custom user-trained HiFi-GAN vocoder: {:?}", custom_vocoder_fs_path); + custom_vocoder_fs_path + }, + }; + + // NB: Tacotron operates on decoder steps. 1000 steps is the default and correlates to + // roughly 12 seconds max. Here we map seconds to decoder steps. + let max_decoder_steps = seconds_to_decoder_steps(job.max_duration_seconds); + + job_dependencies.http_clients.tts_inference_sidecar_client.request_inference( + &cleaned_inference_text, + max_decoder_steps, + &tts_synthesizer_fs_path, + pretrained_vocoder, + &text_pipeline_type_or_guess.to_str(), + &hifigan_vocoder_model_fs_path_to_use, + &hifigan_superres_vocoder_model_fs_path, + &waveglow_vocoder_model_fs_path, + &output_audio_fs_path, + &output_spectrogram_fs_path, + &output_metadata_fs_path, + maybe_unload_model_path) + .await + .map_err(|e| ProcessSingleJobError::Other(e))?; + + // ==================== CHECK ALL FILES EXIST AND GET METADATA ==================== // + + info!("Checking that output files exist..."); + + check_file_exists(&output_audio_fs_path).map_err(|e| ProcessSingleJobError::Other(e))?; + check_file_exists(&output_spectrogram_fs_path).map_err(|e| ProcessSingleJobError::Other(e))?; + check_file_exists(&output_metadata_fs_path).map_err(|e| ProcessSingleJobError::Other(e))?; + + let file_metadata = read_metadata_file(&output_metadata_fs_path) + .map_err(|e| ProcessSingleJobError::Other(e))?; + + safe_delete_temp_file(&output_metadata_fs_path); + + // ==================== UPLOAD AUDIO TO BUCKET ==================== // + + job_progress_reporter.log_status("uploading result") + .map_err(|e| ProcessSingleJobError::Other(e))?; + + let audio_result_object_path = job_dependencies.bucket_path_unifier.tts_inference_wav_audio_output_path( + &job.uuid_idempotency_token); // TODO: Don't use this! + + info!("Audio destination bucket path: {:?}", &audio_result_object_path); + + info!("Uploading audio..."); + + job_dependencies.public_bucket_client.upload_filename_with_content_type( + &audio_result_object_path, + &output_audio_fs_path, + "audio/wav") + .await + .map_err(|e| ProcessSingleJobError::Other(e))?; + + safe_delete_temp_file(&output_audio_fs_path); + + // ==================== UPLOAD SPECTROGRAM TO BUCKETS ==================== // + + let spectrogram_result_object_path = job_dependencies.bucket_path_unifier.tts_inference_spectrogram_output_path( + &job.uuid_idempotency_token); // TODO: Don't use this! + + info!("Spectrogram destination bucket path: {:?}", &spectrogram_result_object_path); + + info!("Uploading spectrogram..."); + + job_dependencies.public_bucket_client.upload_filename_with_content_type( + &spectrogram_result_object_path, + &output_spectrogram_fs_path, + "application/json") + .await + .map_err(|e| ProcessSingleJobError::Other(e))?; + + safe_delete_temp_file(&output_spectrogram_fs_path); + + // ==================== DELETE DOWNLOADED FILE ==================== // + + // NB: We should be using a tempdir, but to make absolutely certain we don't overflow the disk... + safe_delete_temp_directory(&temp_dir); + + // ==================== SAVE RECORDS ==================== // + + let text_hash = hash_string_sha2(&cleaned_inference_text) + .map_err(|e| ProcessSingleJobError::Other(e))?; + + let worker_name = job_dependencies.get_worker_name(); + + info!("Saving tts inference record..."); + + let (id, inference_result_token) = insert_tts_result( + &job_dependencies.mysql_pool, + job, + &text_hash, + pretrained_vocoder, + &audio_result_object_path, + &spectrogram_result_object_path, + file_metadata.file_size_bytes, + file_metadata.duration_millis.unwrap_or(0), + job_dependencies.worker_details.is_on_prem, + &worker_name, + job_dependencies.worker_details.is_debug_worker) + .await + .map_err(|e| ProcessSingleJobError::Other(e))?; + + info!("Marking job complete..."); + mark_tts_inference_job_done( + &job_dependencies.mysql_pool, + job.id, + true, + Some(&inference_result_token), + &worker_name) + .await + .map_err(|e| ProcessSingleJobError::Other(e))?; + + info!("TTS Done. Original text was: {}", &job.raw_inference_text); + + job_dependencies.firehose_publisher.tts_inference_finished( + job.maybe_creator_user_token.as_deref(), + &model_record.model_token, + &inference_result_token) + .await + .map_err(|e| { + error!("error publishing event: {:?}", e); + ProcessSingleJobError::Other(anyhow!("error publishing event")) + })?; + + job_progress_reporter.log_status("done") + .map_err(|e| ProcessSingleJobError::Other(e))?; + + info!("Job {:?} complete success! Downloaded, ran inference, and uploaded. Saved model record: {}, Result Token: {}", + job.id, id, &inference_result_token); + + let duration = start.elapsed(); + + since_creation_span.set_attribute("status", "success"); + since_creation_span.set_duration(duration); + + job_iteration_span.set_attribute("status", "success"); + job_iteration_span.set_duration(duration); + + Ok((since_creation_span, job_iteration_span)) + Ok(()) +} \ No newline at end of file diff --git a/crates/service/job/inference_job/src/job_types/tts/mod.rs b/crates/service/job/inference_job/src/job_types/tts/mod.rs index 5e9ab82e7f..673e15b371 100644 --- a/crates/service/job/inference_job/src/job_types/tts/mod.rs +++ b/crates/service/job/inference_job/src/job_types/tts/mod.rs @@ -1 +1,3 @@ +pub mod handle_tts_inference; +pub mod read_metadata_file; pub mod seconds_to_decoder_steps; diff --git a/crates/service/job/inference_job/src/job_types/tts/read_metadata_file.rs b/crates/service/job/inference_job/src/job_types/tts/read_metadata_file.rs new file mode 100644 index 0000000000..514813d9c3 --- /dev/null +++ b/crates/service/job/inference_job/src/job_types/tts/read_metadata_file.rs @@ -0,0 +1,18 @@ +use errors::AnyhowResult; +use std::fs::File; +use std::io::Read; +use std::path::PathBuf; + +#[derive(Deserialize, Default)] +pub struct FileMetadata { + pub duration_millis: Option, + pub mimetype: Option, + pub file_size_bytes: u64, +} + +pub 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)?) +}