Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 29 additions & 0 deletions src/server/routes/availability_tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -159,6 +159,25 @@ async fn post(app: axum::Router, path: &str, body: Value) -> (StatusCode, Value)
(status, body)
}

async fn get(app: axum::Router, path: &str) -> (StatusCode, Value) {
let response = app
.oneshot(
Request::builder()
.method(Method::GET)
.uri(path)
.body(Body::empty())
.expect("request builds"),
)
.await
.expect("route responds");
let status = response.status();
let bytes = axum::body::to_bytes(response.into_body(), usize::MAX)
.await
.expect("body reads");
let body = serde_json::from_slice(&bytes).expect("health body is JSON");
(status, body)
}

#[tokio::test]
async fn embedding_only_server_returns_501_on_all_generation_routes() {
let cases = [
Expand Down Expand Up @@ -231,6 +250,16 @@ async fn audio_only_server_names_the_served_routes() {
assert!(message.contains("/v1/audio/translations"), "{body}");
}

#[tokio::test]
async fn audio_only_server_reports_healthy_after_its_worker_is_ready() {
let state = state_with(ModelProvider::chat_unavailable_for_route_tests())
.with_audio_model(Some(Arc::new(SideAudioProvider)));
let (status, body) = get(create_app(state), "/health").await;

assert_eq!(status, StatusCode::OK, "{body}");
assert_eq!(body, json!({ "status": "ok" }));
}

#[tokio::test]
async fn exited_chat_worker_returns_503_without_channel_details() {
let state = state_with(ModelProvider::exited_chat_worker_for_route_tests());
Expand Down
9 changes: 9 additions & 0 deletions src/server/routes/health.rs
Original file line number Diff line number Diff line change
Expand Up @@ -60,13 +60,22 @@ fn side_model_ready(state: &AppState) -> bool {
///
/// `200 {"status": "ok"}` once the serving worker is up; the b10621 loading
/// envelope before that. Never reports saturation (see module docs).
///
/// Used by: `mlxcel-server` and `mlxcel serve`.
pub async fn health_check(State(state): State<AppState>) -> Response {
// A server started with b10621's `--embeddings` or `--reranking` has no
// chat worker to be "loaded", and reporting it unhealthy forever would
// make the mode unusable behind any container probe (#1452). Readiness
// there is whether the worker the mode selected came up.
let ready = if state.config.embedding_serving_mode.blocks_generation() {
side_model_ready(&state)
} else if state.model_provider.is_chat_unavailable() && state.audio_model.is_some() {
// Whisper and Kokoro checkpoints are primary audio-only models: their
// dedicated worker has already loaded before the provider is installed,
// while the deliberately unused chat worker records a terminal state.
// Treat that combination as ready so the documented `-m <checkpoint>`
// audio server can pass container health checks.
true
} else {
state.model_provider.is_loaded()
};
Expand Down