add /v2/health/live,ready

llgtrt/src/routes/health_check.rs

@@ -1 +1,58 @@
-pub async fn route_health_check() {}
+use std::sync::Arc;
+
+use axum::{
+    extract::State,
+    http::{HeaderMap, StatusCode},
+    response::{IntoResponse, Response},
+    Json,
+};
+use serde_json::{json, Value};
+
+use crate::{
+    error::AppError,
+    routes::{completions, openai::CompletionCreateParams},
+    state::AppState,
+};
+
+pub async fn ready_check() {
+    log::info!("ready_check -> all good");
+}
+
+pub async fn live_check(
+    headers: HeaderMap,
+    State(app_state): State<Arc<AppState>>,
+) -> Result<Response, AppError> {
+    let req: CompletionCreateParams = serde_json::from_value(json!({
+        "model": "model",
+        "prompt": "Hi",
+        "max_tokens": 2
+    }))?;
+    let resp = completions::route_completions(headers, State(app_state), Json(req)).await?;
+    let status = resp.status();
+    let body = axum::body::to_bytes(resp.into_body(), 1024 * 1024).await?;
+    if status == StatusCode::OK {
+        let body: Value = serde_json::from_slice(&body)?;
+        log::debug!(
+            "Liveness check response: {}",
+            serde_json::to_string_pretty(&body)?
+        );
+        if body["choices"][0]["text"].as_str().is_some() {
+            log::info!("Liveness check complete");
+            Ok((StatusCode::OK, "Check complete").into_response())
+        } else {
+            log::error!(
+                "Liveness check failed: response body does not contain expected field; body: {}",
+                serde_json::to_string_pretty(&body)?
+            );
+            Ok((StatusCode::INTERNAL_SERVER_ERROR, "Check failed").into_response())
+        }
+
+    } else {
+        log::error!(
+            "Liveness check failed with status code: {}; body {}",
+            status,
+            String::from_utf8_lossy(&body)
+        );
+        Ok((status, body).into_response())
+    }
+}

llgtrt/src/startup.rs

@@ -201,7 +201,8 @@ pub async fn run_server(mut cli_config: CliConfig) -> anyhow::Result<()> {
     let app = Router::new()
         .route("/v1/completions", post(routes::route_completions))
         .route("/v1/chat/completions", post(routes::route_chat_completions))
-        .route("/health_check", get(routes::route_health_check))
+        .route("/v1/health/live", get(routes::live_check))
+        .route("/v1/health/ready", get(routes::ready_check))
         .route("/v1/run", post(routes::route_llguidance))
         .route("/guidance", post(routes::route_llguidance))
         .with_state(Arc::new(state))

scripts/test-infer.sh

@@ -86,6 +86,14 @@ curl -X POST "${TRT_API_BASE}chat/completions" \
      "temperature": 0.7
    }'
    ;;
+
+   live)
+    curl -v "${TRT_API_BASE}health/live"
+    ;;
+
+   ready)
+    curl -v "${TRT_API_BASE}health/ready"
+    ;;
 
   exn)
   curl -X POST "${TRT_API_BASE}chat/completions" \