add health checker to download-job

Cargo.lock

@@ -2380,6 +2380,7 @@ dependencies = [
  "chrono",
  "clap",
  "cloud_storage",
+ "concurrency",
  "config",
  "container_common",
  "crockford",

crates/service/job/download_job/Cargo.toml

@@ -17,6 +17,7 @@ path = "src/main.rs"
 # Internal
 bootstrap = { path = "../../../lib/bootstrap" }
 cloud_storage = { path = "../../../lib/api_clients/cloud_storage" }
+concurrency = { path = "../../../lib/concurrency" }
 config = { path = "../../../lib/deprecated/config" }
 container_common = { path = "../../../lib/deprecated/container_common" }
 crockford = { path = "../../../lib/crockford" }

crates/service/job/download_job/src/job_loop/main_loop.rs

@@ -2,10 +2,11 @@ use container_common::anyhow_result::AnyhowResult;
 use crate::JobState;
 use crate::job_loop::process_single_job::process_single_job;
 use jobs_common::noop_logger::NoOpLogger;
-use log::warn;
+use log::{error, warn};
 use mysql_queries::queries::generic_download::job::list_available_generic_download_jobs::{AvailableDownloadJob, list_available_generic_download_jobs};
 use mysql_queries::queries::generic_download::job::mark_generic_download_job_failure::mark_generic_download_job_failure;
 use std::time::Duration;
+use anyhow::anyhow;
 
 // Job runner timeouts (guards MySQL)
 const START_TIMEOUT_MILLIS : u64 = 500;
@@ -17,6 +18,12 @@ pub async fn main_loop(job_state: JobState) {
   let mut noop_logger = NoOpLogger::new(job_state.no_op_logger_millis as i64);
 
   loop {
+    let gpu_is_missing = job_state.nvidia_smi_health_check_status.get_gpu_is_missing();
+    if gpu_is_missing {
+      error!("nvidia-smi health check failed; exiting program!");
+      panic!("nvidia-smi health check failed; exiting program!");
+    }
+
     let num_records = 1;
     let maybe_available_jobs = list_available_generic_download_jobs(&job_state.mysql_pool, num_records).await;
 
@@ -57,6 +64,11 @@ pub async fn main_loop(job_state: JobState) {
 
 async fn process_jobs(job_state: &JobState, jobs: Vec<AvailableDownloadJob>) -> AnyhowResult<()> {
   for job in jobs.into_iter() {
+    let gpu_is_missing = job_state.nvidia_smi_health_check_status.get_gpu_is_missing();
+    if gpu_is_missing {
+      return Err(anyhow!("nvidia-smi health check failed"));
+    }
+
     let result = process_single_job(job_state, &job).await;
     match result {
       Ok(_) => {},

crates/service/job/download_job/src/job_state.rs

@@ -8,6 +8,7 @@ use crate::job_types::vocoder::hifigan_tacotron::hifigan_model_check_command::Hi
 use crate::job_types::voice_conversion::rvc_v2::rvc_v2_model_check_command::RvcV2ModelCheckCommand;
 use crate::job_types::voice_conversion::so_vits_svc::so_vits_svc_model_check_command::SoVitsSvcModelCheckCommand;
 use crate::job_types::voice_conversion::softvc::softvc_model_check_command::SoftVcModelCheckCommand;
+use crate::threads::nvidia_smi_checker::nvidia_smi_health_check_status::NvidiaSmiHealthCheckStatus;
 use google_drive_common::google_drive_download_command::GoogleDriveDownloadCommand;
 use mysql_queries::common_inputs::container_environment_arg::ContainerEnvironmentArg;
 use mysql_queries::mediators::badge_granter::BadgeGranter;
@@ -29,6 +30,8 @@ pub struct JobState {
 
   pub bucket_path_unifier: BucketPathUnifier,
 
+  pub nvidia_smi_health_check_status: NvidiaSmiHealthCheckStatus,
+
   pub sidecar_configs: SidecarConfigs,
 
   // Root to store TTS results

crates/service/job/download_job/src/job_types/voice_conversion/so_vits_svc/so_vits_svc_model_check_command.rs

@@ -1,11 +1,10 @@
-use std::ffi::OsString;
+use anyhow::anyhow;
 use container_common::anyhow_result::AnyhowResult;
 use filesys::path_to_string::path_to_string;
 use log::info;
-use std::fs::OpenOptions;
+use std::ffi::OsString;
 use std::path::{Path, PathBuf};
-use anyhow::anyhow;
-use subprocess::{Popen, PopenConfig, Redirection};
+use subprocess::{Popen, PopenConfig};
 use subprocess_common::docker_options::{DockerFilesystemMount, DockerGpu, DockerOptions};
 
 /// This command is used to check tacotron for being a real model

crates/service/job/download_job/src/main.rs

@@ -20,6 +20,7 @@
 pub mod job_loop;
 pub mod job_state;
 pub mod job_types;
+pub mod threads;
 
 use bootstrap::bootstrap::{bootstrap, BootstrapArgs};
 use cloud_storage::bucket_client::BucketClient;
@@ -39,7 +40,7 @@ use crate::job_types::voice_conversion::so_vits_svc::so_vits_svc_model_check_com
 use crate::job_types::voice_conversion::softvc::softvc_model_check_command::SoftVcModelCheckCommand;
 use errors::AnyhowResult;
 use google_drive_common::google_drive_download_command::GoogleDriveDownloadCommand;
-use log::{info, warn};
+use log::info;
 use mysql_queries::common_inputs::container_environment_arg::ContainerEnvironmentArg;
 use mysql_queries::mediators::badge_granter::BadgeGranter;
 use mysql_queries::mediators::firehose_publisher::FirehosePublisher;
@@ -49,7 +50,10 @@ use sqlx::mysql::MySqlPoolOptions;
 use std::ops::Deref;
 use std::path::PathBuf;
 use std::time::Duration;
+use tokio::runtime::Runtime;
 use subprocess_common::docker_options::{DockerFilesystemMount, DockerGpu, DockerOptions};
+use crate::threads::nvidia_smi_checker::nvidia_smi_health_check_status::NvidiaSmiHealthCheckStatus;
+use crate::threads::nvidia_smi_checker::nvidia_smi_health_check_thread::nvidia_smi_health_check_thread;
 
 // Buckets
 const ENV_ACCESS_KEY : &'static str = "ACCESS_KEY";
@@ -317,6 +321,21 @@ async fn main() -> AnyhowResult<()> {
     firehose_publisher: firehose_publisher.clone(), // NB: Also safe
   };
 
+  let tokio_runtime = Runtime::new()?;
+
+  info!("Spawning nvidia-smi health checker.");
+
+  let nvidia_smi_health_check_status = NvidiaSmiHealthCheckStatus::new();
+  let nvidia_smi_health_check_status2 = nvidia_smi_health_check_status.clone();
+
+  tokio_runtime.spawn(async move {
+    nvidia_smi_health_check_thread(
+      nvidia_smi_health_check_status2,
+      easyenv::get_env_duration_seconds_or_default("NVIDIA_HEALTH_CHECK_TIMEOUT_SECONDS", Duration::from_secs(30)),
+    ).await;
+  });
+
+
   let job_state = JobState {
     download_temp_directory: temp_directory,
     mysql_pool,
@@ -324,6 +343,7 @@ async fn main() -> AnyhowResult<()> {
     bucket_client,
     bucket_path_unifier: BucketPathUnifier::default_paths(),
     bucket_root_tts_model_uploads: bucket_root.to_string(),
+    nvidia_smi_health_check_status,
     firehose_publisher,
     badge_granter,
     sidecar_configs: SidecarConfigs {

crates/service/job/download_job/src/threads/mod.rs

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

crates/service/job/download_job/src/threads/nvidia_smi_checker/mod.rs

@@ -0,0 +1,2 @@
+pub mod nvidia_smi_health_check_status;
+pub mod nvidia_smi_health_check_thread;

crates/service/job/download_job/src/threads/nvidia_smi_checker/nvidia_smi_health_check_status.rs

@@ -0,0 +1,28 @@
+use concurrency::relaxed_atomic_bool::RelaxedAtomicBool;
+
+#[derive(Clone)]
+pub struct NvidiaSmiHealthCheckStatus {
+  /// Occasionally the compute nodes lose the GPU.
+  /// When this happens, the GPU is gone and the node needs to be restarted.
+  /// (This might be a result of the so-vits-svc architecture doing weird things.)
+  gpu_is_missing: RelaxedAtomicBool,
+}
+
+impl NvidiaSmiHealthCheckStatus {
+
+  pub fn new() -> Self {
+    Self {
+      gpu_is_missing: RelaxedAtomicBool::new(false),
+    }
+  }
+
+  /// Read
+  pub fn get_gpu_is_missing(&self) -> bool {
+    self.gpu_is_missing.get()
+  }
+
+  /// Write
+  pub fn notify_gpu_missing(&self) {
+    self.gpu_is_missing.set(true)
+  }
+}

crates/service/job/download_job/src/threads/nvidia_smi_checker/nvidia_smi_health_check_thread.rs

@@ -0,0 +1,79 @@
+use crate::threads::nvidia_smi_checker::nvidia_smi_health_check_status::NvidiaSmiHealthCheckStatus;
+use log::{error, warn};
+use log::{debug, info};
+use std::thread;
+use std::time::Duration;
+use subprocess::{ExitStatus, Popen, PopenConfig};
+
+pub async fn nvidia_smi_health_check_thread(
+  health_check_status: NvidiaSmiHealthCheckStatus,
+  check_duration: Duration,
+) {
+  let mut consecutive_mystery_status_count = 0;
+
+  loop {
+    thread::sleep(check_duration);
+
+    debug!("Checking nvidia-smi health...");
+
+    let exit_status = run_nvidia_smi().await;
+
+    let mut unhealthy_exit = match exit_status {
+      ExitStatus::Exited(0) => {
+        // Healthy exit
+        consecutive_mystery_status_count = 0;
+        false
+      },
+      ExitStatus::Signaled(_) | ExitStatus::Other(_) | ExitStatus::Undetermined => {
+        // Mysterious exit codes (they might resolve)
+        consecutive_mystery_status_count += 1;
+        false
+      }
+      // Unhealthy exits
+      ExitStatus::Exited(1) => true,
+      ExitStatus::Exited(_other_exit_code) => true,
+    };
+
+    if consecutive_mystery_status_count > 5 {
+      unhealthy_exit = true;
+    }
+
+    if unhealthy_exit {
+      health_check_status.notify_gpu_missing();
+    }
+  }
+
+  error!("Should never happen: Nvidia SMI Health Checker Exits");
+}
+
+async fn run_nvidia_smi() ->ExitStatus {
+  let command_parts = [
+    "bash",
+    "-c",
+    "nvidia-smi",
+  ];
+
+  let config = PopenConfig::default();
+
+  let mut handle = match Popen::create(&command_parts, config) {
+    Ok(handle) => handle,
+    Err(err) => {
+      warn!("Error with running nvidia-smi: {:?}", err);
+      return ExitStatus::Undetermined
+    },
+  };
+
+  debug!("nvidia-smi health check pid: {:?}", handle.pid());
+
+  let exit_status = match handle.wait() {
+    Ok(exit_status) => exit_status,
+    Err(err) => {
+      warn!("Error with waiting for nvidia-smi: {:?}", err);
+      ExitStatus::Undetermined
+    }
+  };
+
+  debug!("nvidia-smi health check exit status: {:?}", exit_status);
+
+  exit_status
+}