From 10c277d9cf48a5fccda14ef34849aac339d4fe3d Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Sat, 17 Jun 2023 07:49:27 -0700 Subject: [PATCH] health check endpoint (not installed) --- Cargo.lock | 4 + crates/lib/actix_helpers/Cargo.toml | 4 + crates/lib/actix_helpers/src/lib.rs | 1 + .../error_to_json_http_response.rs | 21 +++++ .../src/response_serializers/mod.rs | 1 + crates/lib/jobs_common/src/job_stats.rs | 15 +++- crates/service/job/inference_job/Cargo.toml | 4 + .../endpoints/health_check_handler.rs | 86 +++++++++++++++++++ .../src/http_server/endpoints/mod.rs | 1 + .../http_server/http_server_shared_state.rs | 6 ++ .../job/inference_job/src/http_server/mod.rs | 3 + .../src/http_server/run_http_server.rs | 51 +++++++++++ crates/service/job/inference_job/src/main.rs | 8 +- 13 files changed, 202 insertions(+), 3 deletions(-) create mode 100644 crates/lib/actix_helpers/src/response_serializers/error_to_json_http_response.rs create mode 100644 crates/lib/actix_helpers/src/response_serializers/mod.rs create mode 100644 crates/service/job/inference_job/src/http_server/endpoints/health_check_handler.rs create mode 100644 crates/service/job/inference_job/src/http_server/endpoints/mod.rs create mode 100644 crates/service/job/inference_job/src/http_server/http_server_shared_state.rs create mode 100644 crates/service/job/inference_job/src/http_server/mod.rs create mode 100644 crates/service/job/inference_job/src/http_server/run_http_server.rs diff --git a/Cargo.lock b/Cargo.lock index 2f2bde0002..6a6b3908ae 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -676,6 +676,8 @@ dependencies = [ "actix-web-actors", "errors", "futures-util", + "serde", + "serde_json", ] [[package]] @@ -3647,6 +3649,8 @@ dependencies = [ name = "inference-job" version = "0.0.1" dependencies = [ + "actix-web 4.0.0-beta.10", + "actix_helpers", "anyhow", "bootstrap", "buckets", diff --git a/crates/lib/actix_helpers/Cargo.toml b/crates/lib/actix_helpers/Cargo.toml index 4e203a60c2..1394e67edd 100644 --- a/crates/lib/actix_helpers/Cargo.toml +++ b/crates/lib/actix_helpers/Cargo.toml @@ -31,5 +31,9 @@ actix-web-actors = "4.0.0-beta.7" # Futures futures-util = "0.3.9" +# Serialization +serde = "1.0.125" +serde_json = "1.0.64" + [dev-dependencies] # None yet diff --git a/crates/lib/actix_helpers/src/lib.rs b/crates/lib/actix_helpers/src/lib.rs index fcf4b32b0a..f5d14811e3 100644 --- a/crates/lib/actix_helpers/src/lib.rs +++ b/crates/lib/actix_helpers/src/lib.rs @@ -20,4 +20,5 @@ pub mod extractors; pub mod middleware; +pub mod response_serializers; pub mod route_builder; diff --git a/crates/lib/actix_helpers/src/response_serializers/error_to_json_http_response.rs b/crates/lib/actix_helpers/src/response_serializers/error_to_json_http_response.rs new file mode 100644 index 0000000000..1ee91b1d7c --- /dev/null +++ b/crates/lib/actix_helpers/src/response_serializers/error_to_json_http_response.rs @@ -0,0 +1,21 @@ +use actix_web::HttpResponse; +use actix_web::HttpResponseBuilder; +use actix_web::error::ResponseError; +use actix_web::http::header::ContentType; +use serde::Serialize; + +/// Turn error responses into JSON HTTP responses with appropriate status code +pub fn error_to_json_http_response( + error: &T, +) -> HttpResponse + where T: ?Sized + Serialize + ResponseError, +{ + let json_body = match serde_json::to_string(error) { + Ok(json) => json, + Err(_) => "{}".to_string(), + }; + + HttpResponseBuilder::new(error.status_code()) + .content_type(ContentType::json()) + .body(json_body) +} diff --git a/crates/lib/actix_helpers/src/response_serializers/mod.rs b/crates/lib/actix_helpers/src/response_serializers/mod.rs new file mode 100644 index 0000000000..e03eb9b082 --- /dev/null +++ b/crates/lib/actix_helpers/src/response_serializers/mod.rs @@ -0,0 +1 @@ +pub mod error_to_json_http_response; \ No newline at end of file diff --git a/crates/lib/jobs_common/src/job_stats.rs b/crates/lib/jobs_common/src/job_stats.rs index b387ca679d..7701501d06 100644 --- a/crates/lib/jobs_common/src/job_stats.rs +++ b/crates/lib/jobs_common/src/job_stats.rs @@ -32,7 +32,20 @@ impl JobStats { inner: Arc::new(RwLock::new(JobStatsInner::default())), } } - + + pub fn get_status(&self) -> AnyhowResult { + // NB: lock errors can't be moved between threads, so we change their type + let mut lock = self.inner.read() + .map_err(|e| anyhow!("lock read error: {:?}", e))?; + + Ok(SuccessAndFailureStats { + total_success_count: lock.total_success_count, + total_failure_count: lock.total_failure_count, + consecutive_success_count: lock.consecutive_success_count, + consecutive_failure_count: lock.consecutive_failure_count, + }) + } + pub fn increment_failure_count(&self) -> AnyhowResult { // NB: lock errors can't be moved between threads, so we change their type let mut lock = self.inner.write() diff --git a/crates/service/job/inference_job/Cargo.toml b/crates/service/job/inference_job/Cargo.toml index a78dd447f7..9ccbe35f22 100644 --- a/crates/service/job/inference_job/Cargo.toml +++ b/crates/service/job/inference_job/Cargo.toml @@ -15,6 +15,7 @@ path = "src/main.rs" [dependencies] # Internal +actix_helpers = { path = "../../../lib/actix_helpers" } bootstrap = { path = "../../../lib/bootstrap" } buckets = { path = "../../../_types/buckets" } cloud_storage = { path = "../../../lib/api_clients/cloud_storage" } @@ -68,6 +69,9 @@ tokio = { version = "1.6.0", features = ["macros"] } # NB: Furthermore, the Tokio runtime is old, so I have to use an older Hyper. hyper = { version = "0.14.14" } +# HTTP Server +actix-web = "4.0.0-beta.10" + # Serialization serde = "1.0.125" serde_derive = "1.0.125" diff --git a/crates/service/job/inference_job/src/http_server/endpoints/health_check_handler.rs b/crates/service/job/inference_job/src/http_server/endpoints/health_check_handler.rs new file mode 100644 index 0000000000..1e642075e8 --- /dev/null +++ b/crates/service/job/inference_job/src/http_server/endpoints/health_check_handler.rs @@ -0,0 +1,86 @@ +use actix_helpers::response_serializers::error_to_json_http_response::error_to_json_http_response; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::http::header::ContentType; +use actix_web::{web, HttpResponse, HttpRequest}; +use crate::http_server::http_server_shared_state::HttpServerSharedState; +use log::error; +use std::sync::Arc; + +// =============== Success Response =============== + +#[derive(Serialize)] +pub struct HealthCheckResponse { + pub success: bool, + pub is_healthy: bool, + + pub consecutive_success_count: u64, + pub total_success_count: u64, + + pub consecutive_failure_count: u64, + pub total_failure_count: u64, +} + + +// =============== Error Response =============== + +#[derive(Debug, Serialize)] +pub enum HealthCheckError { + ServerError, +} + +impl ResponseError for HealthCheckError { + fn status_code(&self) -> StatusCode { + match *self { + HealthCheckError::ServerError => StatusCode::INTERNAL_SERVER_ERROR, + } + } + + fn error_response(&self) -> HttpResponse { + error_to_json_http_response(self) + } +} + +// NB: Not using derive_more::Display since Clion doesn't understand it. +impl std::fmt::Display for HealthCheckError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{:?}", self) + } +} + +// =============== Handler =============== + +pub async fn get_health_check_handler( + _http_request: HttpRequest, + server_state: web::Data> +) -> Result { + let job_stats = server_state.job_stats.get_status() + .map_err(|e| { + error!("Error serving health check status: {:?}", e); + HealthCheckError::ServerError + })?; + + let is_healthy = job_stats.consecutive_failure_count < 10; + + let response = HealthCheckResponse { + success: true, + is_healthy, + consecutive_failure_count: job_stats.consecutive_failure_count, + total_failure_count: job_stats.total_failure_count, + consecutive_success_count: job_stats.consecutive_success_count, + total_success_count: job_stats.total_success_count, + }; + + let body = serde_json::to_string(&response) + .map_err(|e| HealthCheckError::ServerError)?; + + if is_healthy { + Ok(HttpResponse::Ok() + .content_type(ContentType::json()) + .body(body)) + } else { + Ok(HttpResponse::InternalServerError() + .content_type(ContentType::json()) + .body(body)) + } +} diff --git a/crates/service/job/inference_job/src/http_server/endpoints/mod.rs b/crates/service/job/inference_job/src/http_server/endpoints/mod.rs new file mode 100644 index 0000000000..0f408270bf --- /dev/null +++ b/crates/service/job/inference_job/src/http_server/endpoints/mod.rs @@ -0,0 +1 @@ +pub mod health_check_handler; diff --git a/crates/service/job/inference_job/src/http_server/http_server_shared_state.rs b/crates/service/job/inference_job/src/http_server/http_server_shared_state.rs new file mode 100644 index 0000000000..7ab40d5e39 --- /dev/null +++ b/crates/service/job/inference_job/src/http_server/http_server_shared_state.rs @@ -0,0 +1,6 @@ +use jobs_common::job_stats::JobStats; + +#[derive(Clone)] +pub struct HttpServerSharedState { + pub job_stats: JobStats, +} diff --git a/crates/service/job/inference_job/src/http_server/mod.rs b/crates/service/job/inference_job/src/http_server/mod.rs new file mode 100644 index 0000000000..5d888824cc --- /dev/null +++ b/crates/service/job/inference_job/src/http_server/mod.rs @@ -0,0 +1,3 @@ +pub mod endpoints; +pub mod http_server_shared_state; +pub mod run_http_server; diff --git a/crates/service/job/inference_job/src/http_server/run_http_server.rs b/crates/service/job/inference_job/src/http_server/run_http_server.rs new file mode 100644 index 0000000000..55752d37f1 --- /dev/null +++ b/crates/service/job/inference_job/src/http_server/run_http_server.rs @@ -0,0 +1,51 @@ +use actix_web::middleware::Logger; +use actix_web::{App, HttpResponse, HttpServer, middleware, web}; +use crate::http_server::endpoints::health_check_handler::get_health_check_handler; +use crate::http_server::http_server_shared_state::HttpServerSharedState; +use crate::job_dependencies::JobDependencies; +use errors::AnyhowResult; +use log::info; +use std::sync::Arc; + +const DEFAULT_BIND_ADDRESS : &'static str = "0.0.0.0:12345"; +const DEFAULT_NUM_WORKERS : usize = 4; + +pub async fn run_http_server(job_dependencies: &JobDependencies) -> AnyhowResult<()> +{ + // HTTP server args + let bind_address = easyenv::get_env_string_or_default("HTTP_BIND_ADDRESS", DEFAULT_BIND_ADDRESS); + let num_workers = easyenv::get_env_num("HTTP_NUM_WORKERS", DEFAULT_NUM_WORKERS)?; + let hostname = job_dependencies.container.hostname.clone(); + + let server_state = HttpServerSharedState { + job_stats: job_dependencies.job_stats.clone(), + }; + + let server_state_arc = web::Data::new(Arc::new(server_state)); + + info!("Starting HTTP service (just for k8s health checking)."); + + // NB: We shouldn't be logging much as the /_status endpoint is all we aim to expose. + let log_format = "[%{HOSTNAME}e] IP=[%{X-Forwarded-For}i] \"%r\" %s %b \"%{Referer}i\" \"%{User-Agent}i\" %T"; + + HttpServer::new(move || { + // NB: app_data being clone()'d below should all be safe (dependencies included) + App::new() + .app_data(server_state_arc.clone()) + .wrap(Logger::new(&log_format) + .exclude("/_status") + ) + .wrap(middleware::Compress::default()) + .service( + web::resource("/_status") + .route(web::post().to(get_health_check_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + }) + .bind(bind_address)? + .workers(num_workers) + .run() + .await?; + + Ok(()) +} diff --git a/crates/service/job/inference_job/src/main.rs b/crates/service/job/inference_job/src/main.rs index 5970a0b617..ff4d774479 100644 --- a/crates/service/job/inference_job/src/main.rs +++ b/crates/service/job/inference_job/src/main.rs @@ -17,9 +17,10 @@ #[macro_use] extern crate serde_derive; +pub mod http_server; pub mod job; -pub mod util; pub mod job_dependencies; +pub mod util; use bootstrap::bootstrap::{bootstrap, BootstrapArgs}; use clap::{App, Arg}; @@ -54,6 +55,7 @@ use sqlx::mysql::MySqlPoolOptions; use std::path::PathBuf; use std::time::Duration; use subprocess_common::docker_options::{DockerEnvVar, DockerFilesystemMount, DockerGpu, DockerOptions}; +use crate::http_server::run_http_server::run_http_server; // Buckets (shared config) const ENV_ACCESS_KEY : &'static str = "ACCESS_KEY"; @@ -296,6 +298,8 @@ async fn main() -> AnyhowResult<()> { }, }; + //run_http_server(&job_dependencies).await?; + main_loop(job_dependencies).await; Ok(()) @@ -356,4 +360,4 @@ fn vits_inference_command() -> AnyhowResult { maybe_nltk_data_cache.as_deref(), maybe_docker_options, )) -} \ No newline at end of file +}