bge_m3_embedding_server/handler/
health.rs

1// Copyright (c) 2026 J. Patrick Fulton
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! `GET /health` and `GET /health/deep` handlers.
16//!
17//! `/health` returns lightweight readiness status from in-memory atomics.
18//! `/health/deep` runs a tiny canary inference (batch=1, seq≈8 tokens) and
19//! returns `503` if the actual embedding pipeline is broken.
20
21use std::sync::Arc;
22use std::sync::atomic::Ordering;
23use std::time::Duration;
24
25use axum::{Json, extract::State, http::StatusCode, response::IntoResponse};
26
27use crate::state::{AppState, ProbeStatus};
28
29/// Fixed canary text for `/health/deep`. Tokenises to ~8 tokens with the
30/// BGE-M3 `SentencePiece` vocabulary, meeting the "batch=1, seq=8" goal.
31const DEEP_HEALTH_CANARY: &str = "embedding service canary health check ok";
32
33/// Timeout for the canary embed call in `/health/deep`.
34///
35/// If the worker pool is so overloaded or broken that even the canary cannot
36/// complete within this budget, the endpoint returns 503. 30 s is generous
37/// for a single batch-1 request but avoids masking a genuinely hung session.
38const DEEP_HEALTH_TIMEOUT: Duration = Duration::from_secs(30);
39
40/// Handles `GET /health` — returns readiness status, worker counts, and tuning diagnostics.
41///
42/// Returns `503` while models are loading or if all workers have exited; returns
43/// `200 ok` (or `200 warn` when fewer workers are live than configured) with the
44/// current cost-model coefficients and probe status in the `tuning` block.
45pub async fn health(State(state): State<Arc<AppState>>) -> impl IntoResponse {
46    let ready = state.ready.load(Ordering::Acquire);
47    let live = state.pool.live_worker_count();
48    let loaded = state.pool.loaded_worker_count();
49    let total = state.total_workers;
50
51    if !ready {
52        return (
53            StatusCode::SERVICE_UNAVAILABLE,
54            Json(serde_json::json!({"status": "loading"})),
55        )
56            .into_response();
57    }
58
59    if live == 0 {
60        return (
61            StatusCode::SERVICE_UNAVAILABLE,
62            Json(serde_json::json!({
63                "status": "fail",
64                "workers": { "live": live, "total": total }
65            })),
66        )
67            .into_response();
68    }
69
70    if loaded == 0 {
71        return (
72            StatusCode::OK,
73            Json(serde_json::json!({
74                "status": "idle",
75                "workers": { "live": live, "total": total }
76            })),
77        )
78            .into_response();
79    }
80
81    let status = if live < total { "warn" } else { "ok" };
82
83    // Read the live cost model and probe status atomically.
84    let cm = state.cost_model.load();
85    let probe_status = ProbeStatus::from_u8(state.probe_status.load(Ordering::Acquire)).as_str();
86
87    let mut tuning = serde_json::json!({
88        "a_bytes_per_token": cm.a,
89        "b_bytes_per_token_sq": cm.b,
90        "max_workspace_bytes": cm.max_workspace_bytes,
91        "probe_status": probe_status,
92    });
93
94    // Add static memory fields when available (written before probe starts).
95    if let Some(ti) = state.tuning.get() {
96        tuning["memory_source"] = serde_json::Value::String(ti.memory_source.clone());
97        tuning["available_bytes"] =
98            serde_json::Value::Number(serde_json::Number::from(ti.available_bytes));
99        tuning["model_rss_bytes_per_worker"] =
100            serde_json::Value::Number(serde_json::Number::from(ti.model_rss_bytes_per_worker));
101    }
102
103    let body = serde_json::json!({
104        "status": status,
105        "workers": { "live": live, "total": total },
106        "max_seq_length": state.max_seq_length,
107        "tuning": tuning,
108    });
109
110    (StatusCode::OK, Json(body)).into_response()
111}
112
113/// Handles `GET /health/deep` — runs a tiny canary inference and returns
114/// `503` if the actual embedding pipeline is broken.
115///
116/// Unlike `GET /health` (which reads only in-memory atomics), this handler
117/// submits a real `embed_dense` call through the worker pool and exercises
118/// the full tokenize → ORT `session.run()` → projection path, including
119/// `TensorRT` engine dispatch on GPU builds. It is the strongest available
120/// liveness signal because it catches the silent-failure mode where
121/// `/health` returned `200 ok` while every real
122/// embedding request returned `500` due to a broken TRT CUDA context.
123///
124/// # Response codes
125///
126/// | Code | Condition |
127/// |---|---|
128/// | `200 ok` | Server is ready and the canary embed succeeded |
129/// | `503 loading` | Server is still loading models |
130/// | `503 fail` | Canary embed failed or timed out |
131///
132/// # ECS and ALB configuration
133///
134/// Point both `ECS healthCheck.command` and the ALB target-group health check
135/// at `/health/deep`. The 30-second inference timeout ensures the health check
136/// never hangs indefinitely; keep the ECS `healthCheckGracePeriodSeconds`
137/// large enough to cover TRT cold-start (≥ 10 800 s for a full 24-shape grid).
138pub async fn health_deep(State(state): State<Arc<AppState>>) -> impl IntoResponse {
139    if !state.ready.load(Ordering::Acquire) {
140        return (
141            StatusCode::SERVICE_UNAVAILABLE,
142            Json(serde_json::json!({"status": "loading"})),
143        )
144            .into_response();
145    }
146
147    let canary = tokio::time::timeout(
148        DEEP_HEALTH_TIMEOUT,
149        state.pool.dense(vec![DEEP_HEALTH_CANARY.to_string()]),
150    )
151    .await;
152
153    match canary {
154        Ok(Ok(_)) => (StatusCode::OK, Json(serde_json::json!({"status": "ok"}))).into_response(),
155        Ok(Err(e)) => {
156            tracing::warn!(error = %e, "health/deep: canary embed failed");
157            (
158                StatusCode::SERVICE_UNAVAILABLE,
159                Json(serde_json::json!({"status": "fail", "error": e.to_string()})),
160            )
161                .into_response()
162        }
163        Err(_elapsed) => {
164            tracing::warn!(
165                timeout_secs = DEEP_HEALTH_TIMEOUT.as_secs(),
166                "health/deep: canary embed timed out"
167            );
168            (
169                StatusCode::SERVICE_UNAVAILABLE,
170                Json(serde_json::json!({"status": "fail", "error": "canary embed timed out"})),
171            )
172                .into_response()
173        }
174    }
175}
bge_m3_embedding_server/handler/health.rs

bge_m3_embedding_server/handler/
health.rs
