health check endpoint (not installed)

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",

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

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;

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

crates/lib/actix_helpers/src/response_serializers/mod.rs

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

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<SuccessAndFailureStats> {
+      // 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<SuccessAndFailureStats> {
       // NB: lock errors can't be moved between threads, so we change their type
       let mut lock = self.inner.write()

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"

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<Arc<HttpServerSharedState>>
+) -> Result<HttpResponse, HealthCheckError> {
+  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))
+  }
+}

crates/service/job/inference_job/src/http_server/endpoints/mod.rs

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

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

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;

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

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<VitsInferenceCommand> {
     maybe_nltk_data_cache.as_deref(),
     maybe_docker_options,
   ))
-}
+}