From a4f0aec7065b33f30c15cfbcbe9276f8aca140ef Mon Sep 17 00:00:00 2001
From: Michal Moskal <michal@moskal.me>
Date: Mon, 25 Nov 2024 22:51:15 +0000
Subject: [PATCH] rename /health/live to /health/model and have /health/live
 just check canEnqueueRequests

---
 llgtrt/src/async_exec.rs          |  4 ++++
 llgtrt/src/routes/health_check.rs | 15 +++++++++++++--
 llgtrt/src/startup.rs             |  1 +
 3 files changed, 18 insertions(+), 2 deletions(-)

diff --git a/llgtrt/src/async_exec.rs b/llgtrt/src/async_exec.rs
index fa7f75f..c74bb8f 100644
--- a/llgtrt/src/async_exec.rs
+++ b/llgtrt/src/async_exec.rs
@@ -419,6 +419,10 @@ impl AsyncExecutor {
         Ok((res, tok_env, chat_builder))
     }
 
+    pub fn can_enqueue_request(&self) -> bool {
+        self.executor.can_enqueue_request()
+    }
+
     pub fn add_request(
         &mut self,
         init: RequestInit,
diff --git a/llgtrt/src/routes/health_check.rs b/llgtrt/src/routes/health_check.rs
index 46920a7..34f9689 100644
--- a/llgtrt/src/routes/health_check.rs
+++ b/llgtrt/src/routes/health_check.rs
@@ -9,6 +9,7 @@ use axum::{
 use serde_json::{json, Value};
 
 use crate::{
+    async_exec::AsyncExecutor,
     error::AppError,
     routes::{completions, openai::CompletionCreateParams},
     state::AppState,
@@ -18,7 +19,18 @@ pub async fn ready_check() {
     log::info!("ready_check -> all good");
 }
 
-pub async fn live_check(
+pub async fn live_check() -> Result<Response, AppError> {
+    let can_enq = AsyncExecutor::lock().can_enqueue_request();
+    if can_enq {
+        log::info!("live_check -> all good");
+        Ok((StatusCode::OK, "Check complete").into_response())
+    } else {
+        log::error!("live_check -> failed");
+        Ok((StatusCode::INTERNAL_SERVER_ERROR, "Check failed").into_response())
+    }
+}
+
+pub async fn model_check(
     headers: HeaderMap,
     State(app_state): State<Arc<AppState>>,
 ) -> Result<Response, AppError> {
@@ -46,7 +58,6 @@ pub async fn live_check(
             );
             Ok((StatusCode::INTERNAL_SERVER_ERROR, "Check failed").into_response())
         }
-        
     } else {
         log::error!(
             "Liveness check failed with status code: {}; body {}",
diff --git a/llgtrt/src/startup.rs b/llgtrt/src/startup.rs
index 4646301..810f67d 100644
--- a/llgtrt/src/startup.rs
+++ b/llgtrt/src/startup.rs
@@ -202,6 +202,7 @@ pub async fn run_server(mut cli_config: CliConfig) -> anyhow::Result<()> {
         .route("/v1/completions", post(routes::route_completions))
         .route("/v1/chat/completions", post(routes::route_chat_completions))
         .route("/v1/health/live", get(routes::live_check))
+        .route("/v1/health/model", get(routes::model_check))
         .route("/v1/health/ready", get(routes::ready_check))
         .route("/v1/run", post(routes::route_llguidance))
         .route("/guidance", post(routes::route_llguidance))