add TT2 HTTP sidecar health checking (hacky)

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
     }
   }
 }

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<RwLock<InnerHealthCheckState>>,
+}
+
+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<bool> {
+    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)
+      }
+    }
+  }
+}

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)
+}

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;

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<JobSuccessResult, P
       .as_ref()
       .ok_or_else(|| ProcessSingleJobError::JobSystemMisconfiguration(Some("missing Tacotron2 dependencies".to_string())))?;
 
-  // TODO(bt,2023-11-26): Need to stop the job before the sidecar spins up
-  //if needs_health_check_at_start {
-  //  maybe_block_on_sidecar_health_check(&inferencer.http_clients.tts_sidecar_health_check_client).await;
-  //}
+  // ==================== OPTIONAL SIDECAR HEALTH CHECK ==================== //
+
+  // TODO(bt,2023-11-28): Ideally we'd perform health checks before grabbing a lock on the job.
+  let maybe_needs_health_check =
+      model_dependencies.sidecar.use_sidecar_instead_of_shell &&
+          model_dependencies.sidecar.health_check_state.needs_health_check()
+              .map_err(|e| ProcessSingleJobError::Other(anyhow!(e)))?;
+
+  if maybe_needs_health_check {
+    maybe_block_on_sidecar_health_check(&model_dependencies.sidecar.health_check_client).await;
+
+    model_dependencies.sidecar.health_check_state.mark_maybe_needs_health_check(false)
+        .map_err(|e| ProcessSingleJobError::Other(anyhow!(e)))?;
+  }
 
   // ==================== CONFIRM OR DOWNLOAD STATIC DEPENDENCIES ==================== //
 

crates/service/job/inference_job/src/job/job_types/tts/tacotron2_v2_early_fakeyou/tacotron2_dependencies.rs

@@ -1,4 +1,5 @@
 use errors::AnyhowResult;
+use crate::job::job_types::tts::tacotron2_v2_early_fakeyou::health_check_state::HealthCheckState;
 
 use crate::job::job_types::tts::tacotron2_v2_early_fakeyou::tacotron2_inference_command::Tacotron2InferenceCommand;
 use crate::job::job_types::tts::tacotron2_v2_early_fakeyou::tacotron2_inference_sidecar_client::Tacotron2InferenceSidecarClient;
@@ -24,6 +25,7 @@ pub struct SidecarDeps {
   pub use_sidecar_instead_of_shell: bool,
   pub inference_client: Tacotron2InferenceSidecarClient,
   pub health_check_client: Tacotron2SidecarHealthCheckClient,
+  pub health_check_state: HealthCheckState,
 }
 
 impl Tacotron2Dependencies {
@@ -60,6 +62,7 @@ impl Tacotron2Dependencies {
         use_sidecar_instead_of_shell: true, // TODO: ENV VAR
         inference_client,
         health_check_client,
+        health_check_state: HealthCheckState::new(),
       }
     })
   }