diff --git a/crates/service/job/inference_job/src/job/job_types/tts/process_single_tts_job.rs b/crates/service/job/inference_job/src/job/job_types/tts/process_single_tts_job.rs index fa117c65a5..7495535d87 100644 --- a/crates/service/job/inference_job/src/job/job_types/tts/process_single_tts_job.rs +++ b/crates/service/job/inference_job/src/job/job_types/tts/process_single_tts_job.rs @@ -1,4 +1,5 @@ use anyhow::anyhow; +use log::warn; use enums::by_table::generic_inference_jobs::inference_model_type::InferenceModelType; use enums::by_table::tts_models::tts_model_type::TtsModelType; @@ -93,21 +94,38 @@ async fn dispatch_fine_tuned_weights_model( }; match tts_model.tts_model_type { - TtsModelType::Tacotron2 => { - tacotron2_v2_early_fakeyou::process_job::process_job(ProcessJobArgs { - job_dependencies, - job, - tts_model: &tts_model, - raw_inference_text, - }).await - } TtsModelType::Vits => { vits::process_job::process_job(VitsProcessJobArgs { + job_dependencies, + job, + tts_model: &tts_model, + raw_inference_text, + }).await + } + TtsModelType::Tacotron2 => { + let result = tacotron2_v2_early_fakeyou::process_job::process_job(ProcessJobArgs { job_dependencies, job, tts_model: &tts_model, raw_inference_text, - }).await + }).await; + + if let Err(err) = result.as_ref() { + warn!("Error with TT2 inference. Requiring health check next iteration: {:?}", err); + job_dependencies + .job + .job_specific_dependencies + .maybe_tacotron2_dependencies + .as_ref() + .map(|deps| { + deps.sidecar.health_check_state.mark_maybe_needs_health_check(true) + .map_err(|err| { + warn!("Error with health check status lock: {:?}", err); + }) + }); + } + + result } } } diff --git a/crates/service/job/inference_job/src/job/job_types/tts/tacotron2_v2_early_fakeyou/health_check_state.rs b/crates/service/job/inference_job/src/job/job_types/tts/tacotron2_v2_early_fakeyou/health_check_state.rs new file mode 100644 index 0000000000..bd22c4cd67 --- /dev/null +++ b/crates/service/job/inference_job/src/job/job_types/tts/tacotron2_v2_early_fakeyou/health_check_state.rs @@ -0,0 +1,57 @@ +use std::sync::{Arc, LockResult, RwLock}; +use anyhow::anyhow; +use errors::AnyhowResult; + +#[derive(Clone)] +pub struct HealthCheckState { + inner: Arc>, +} + +struct InnerHealthCheckState { + checked_health_at_least_once: bool, + maybe_needs_health_check: bool, +} + +impl HealthCheckState { + pub fn new() -> Self { + Self { + inner: Arc::new(RwLock::new(InnerHealthCheckState { + checked_health_at_least_once: false, + maybe_needs_health_check: false, + })) + } + } + + //pub fn mark_first_health_check_done(&self) -> AnyhowResult<()> { + // match self.inner.write() { + // Err(err) => return Err(anyhow!("lock error: {err}")), + // Ok(mut lock) => { + // lock.checked_health_at_least_once = true; + // Ok(()) + // } + // } + //} + + pub fn mark_maybe_needs_health_check(&self, maybe_needs_health_check: bool) -> AnyhowResult<()> { + match self.inner.write() { + Err(err) => return Err(anyhow!("lock error: {err}")), + Ok(mut lock) => { + lock.maybe_needs_health_check = maybe_needs_health_check; + if !maybe_needs_health_check { + lock.checked_health_at_least_once = true; + } + Ok(()) + } + } + } + + pub fn needs_health_check(&self) -> AnyhowResult { + match self.inner.read() { + Err(err) => return Err(anyhow!("lock error: {err}")), + Ok(lock) => { + let needs_health_check = !lock.checked_health_at_least_once || lock.maybe_needs_health_check; + Ok(needs_health_check) + } + } + } +} diff --git a/crates/service/job/inference_job/src/job/job_types/tts/tacotron2_v2_early_fakeyou/health_check_trap.rs b/crates/service/job/inference_job/src/job/job_types/tts/tacotron2_v2_early_fakeyou/health_check_trap.rs new file mode 100644 index 0000000000..9a94d371bc --- /dev/null +++ b/crates/service/job/inference_job/src/job/job_types/tts/tacotron2_v2_early_fakeyou/health_check_trap.rs @@ -0,0 +1,45 @@ +use std::cmp::min; +use std::time::Duration; + +use log::{error, info}; + +use crate::job::job_types::tts::tacotron2_v2_early_fakeyou::tacotron2_sidecar_health_check_client::{HealthState, Tacotron2SidecarHealthCheckClient}; + +// Health check timeouts +const HEALTH_CHECK_START_TIMEOUT_MILLIS : u64 = 250; +const HEALTH_CHECK_INCREASE_TIMEOUT_MILLIS : u64 = 250; +const HEALTH_CHECK_MAX_TIMEOUT_MILLIS : u64 = 5000; + +pub async fn maybe_block_on_sidecar_health_check( + tts_sidecar_health_check_client: &Tacotron2SidecarHealthCheckClient +) { + let mut error_timeout_millis = HEALTH_CHECK_START_TIMEOUT_MILLIS; + + loop { + info!("Performing sidecar health check..."); + let maybe_health = + tts_sidecar_health_check_client.request_health_check().await; + + match maybe_health { + Ok(HealthState::Healthy) => { + return; + } + Ok(HealthState::Unhealthy) => { + error!("Sidecar health is reported as UNHEALTHY!"); + error_timeout_millis = increase_timeout(error_timeout_millis); + } + Err(err) => { + // NB: If we fail + error!("Sidecar health check failed for reason: {:?}", err); + error_timeout_millis = increase_timeout(error_timeout_millis); + } + } + + std::thread::sleep(Duration::from_millis(error_timeout_millis)); + } +} + +fn increase_timeout(current_timeout: u64) -> u64 { + let next_timeout = current_timeout + HEALTH_CHECK_INCREASE_TIMEOUT_MILLIS; + min(next_timeout, HEALTH_CHECK_MAX_TIMEOUT_MILLIS) +} diff --git a/crates/service/job/inference_job/src/job/job_types/tts/tacotron2_v2_early_fakeyou/mod.rs b/crates/service/job/inference_job/src/job/job_types/tts/tacotron2_v2_early_fakeyou/mod.rs index d8c0717997..2e03b29dc3 100644 --- a/crates/service/job/inference_job/src/job/job_types/tts/tacotron2_v2_early_fakeyou/mod.rs +++ b/crates/service/job/inference_job/src/job/job_types/tts/tacotron2_v2_early_fakeyou/mod.rs @@ -1,4 +1,6 @@ pub mod download_static_dependencies; +pub mod health_check_state; +pub mod health_check_trap; pub mod process_job; pub mod read_metadata_file; pub mod seconds_to_decoder_steps; diff --git a/crates/service/job/inference_job/src/job/job_types/tts/tacotron2_v2_early_fakeyou/process_job.rs b/crates/service/job/inference_job/src/job/job_types/tts/tacotron2_v2_early_fakeyou/process_job.rs index 61657b7218..95344784aa 100644 --- a/crates/service/job/inference_job/src/job/job_types/tts/tacotron2_v2_early_fakeyou/process_job.rs +++ b/crates/service/job/inference_job/src/job/job_types/tts/tacotron2_v2_early_fakeyou/process_job.rs @@ -24,6 +24,7 @@ use tts_common::text_pipelines::text_pipeline_type::TextPipelineType; use crate::job::job_loop::job_success_result::{JobSuccessResult, ResultEntity}; use crate::job::job_loop::process_single_job_error::ProcessSingleJobError; use crate::job::job_types::tts::tacotron2_v2_early_fakeyou::download_static_dependencies::download_static_dependencies; +use crate::job::job_types::tts::tacotron2_v2_early_fakeyou::health_check_trap::maybe_block_on_sidecar_health_check; use crate::job::job_types::tts::tacotron2_v2_early_fakeyou::seconds_to_decoder_steps::seconds_to_decoder_steps; use crate::job::job_types::tts::tacotron2_v2_early_fakeyou::tacotron2_inference_command::{InferenceArgs, MelMultiplyFactor}; use crate::job::job_types::tts::tacotron2_v2_early_fakeyou::vocoder_option::VocoderForInferenceOption; @@ -60,10 +61,20 @@ pub async fn process_job(args: ProcessJobArgs<'_>) -> Result