From 97a6f6e4ffe4dd1afe34bb044bd1e8748862e92f Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Wed, 19 Jul 2023 19:52:24 -0400 Subject: [PATCH] add health checker to download-job --- Cargo.lock | 1 + crates/service/job/download_job/Cargo.toml | 1 + .../download_job/src/job_loop/main_loop.rs | 14 +++- .../service/job/download_job/src/job_state.rs | 3 + .../so_vits_svc_model_check_command.rs | 7 +- crates/service/job/download_job/src/main.rs | 22 +++++- .../job/download_job/src/threads/mod.rs | 1 + .../src/threads/nvidia_smi_checker/mod.rs | 2 + .../nvidia_smi_health_check_status.rs | 28 +++++++ .../nvidia_smi_health_check_thread.rs | 79 +++++++++++++++++++ 10 files changed, 152 insertions(+), 6 deletions(-) create mode 100644 crates/service/job/download_job/src/threads/mod.rs create mode 100644 crates/service/job/download_job/src/threads/nvidia_smi_checker/mod.rs create mode 100644 crates/service/job/download_job/src/threads/nvidia_smi_checker/nvidia_smi_health_check_status.rs create mode 100644 crates/service/job/download_job/src/threads/nvidia_smi_checker/nvidia_smi_health_check_thread.rs diff --git a/Cargo.lock b/Cargo.lock index 4a318fa91f..b9a56207d4 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2380,6 +2380,7 @@ dependencies = [ "chrono", "clap", "cloud_storage", + "concurrency", "config", "container_common", "crockford", diff --git a/crates/service/job/download_job/Cargo.toml b/crates/service/job/download_job/Cargo.toml index 3d664a1534..ded8de06fd 100644 --- a/crates/service/job/download_job/Cargo.toml +++ b/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" } diff --git a/crates/service/job/download_job/src/job_loop/main_loop.rs b/crates/service/job/download_job/src/job_loop/main_loop.rs index 955eb10f98..4008d6254e 100644 --- a/crates/service/job/download_job/src/job_loop/main_loop.rs +++ b/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) -> 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(_) => {}, diff --git a/crates/service/job/download_job/src/job_state.rs b/crates/service/job/download_job/src/job_state.rs index e99b74ab1e..c760ee0097 100644 --- a/crates/service/job/download_job/src/job_state.rs +++ b/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 diff --git a/crates/service/job/download_job/src/job_types/voice_conversion/so_vits_svc/so_vits_svc_model_check_command.rs b/crates/service/job/download_job/src/job_types/voice_conversion/so_vits_svc/so_vits_svc_model_check_command.rs index f6c45296c4..f4b92557f0 100644 --- a/crates/service/job/download_job/src/job_types/voice_conversion/so_vits_svc/so_vits_svc_model_check_command.rs +++ b/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 diff --git a/crates/service/job/download_job/src/main.rs b/crates/service/job/download_job/src/main.rs index ac0b419682..2f09a72ea5 100644 --- a/crates/service/job/download_job/src/main.rs +++ b/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 { diff --git a/crates/service/job/download_job/src/threads/mod.rs b/crates/service/job/download_job/src/threads/mod.rs new file mode 100644 index 0000000000..49772f29b4 --- /dev/null +++ b/crates/service/job/download_job/src/threads/mod.rs @@ -0,0 +1 @@ +pub mod nvidia_smi_checker; diff --git a/crates/service/job/download_job/src/threads/nvidia_smi_checker/mod.rs b/crates/service/job/download_job/src/threads/nvidia_smi_checker/mod.rs new file mode 100644 index 0000000000..6af370de47 --- /dev/null +++ b/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; \ No newline at end of file diff --git a/crates/service/job/download_job/src/threads/nvidia_smi_checker/nvidia_smi_health_check_status.rs b/crates/service/job/download_job/src/threads/nvidia_smi_checker/nvidia_smi_health_check_status.rs new file mode 100644 index 0000000000..5b9d7aba10 --- /dev/null +++ b/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) + } +} diff --git a/crates/service/job/download_job/src/threads/nvidia_smi_checker/nvidia_smi_health_check_thread.rs b/crates/service/job/download_job/src/threads/nvidia_smi_checker/nvidia_smi_health_check_thread.rs new file mode 100644 index 0000000000..d80f6a4aca --- /dev/null +++ b/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 +}