From 3c0b4432be91d4a5fda590f3fb3236cd087882c5 Mon Sep 17 00:00:00 2001 From: Sahil Singh Date: Mon, 8 Jun 2026 14:58:37 +0530 Subject: [PATCH] [Rust Frontend] Add /pause, /resume, /is_paused endpoints (#44499) Signed-off-by: Sahil Singh Co-authored-by: Bugen Zhao --- rust/src/engine-core-client/src/client.rs | 32 +++++ rust/src/server/src/routes.rs | 4 + rust/src/server/src/routes/pause.rs | 90 ++++++++++++ rust/src/server/src/routes/tests.rs | 158 ++++++++++++++++++++++ 4 files changed, 284 insertions(+) create mode 100644 rust/src/server/src/routes/pause.rs diff --git a/rust/src/engine-core-client/src/client.rs b/rust/src/engine-core-client/src/client.rs index 2a8c3c7418..8f5daba79e 100644 --- a/rust/src/engine-core-client/src/client.rs +++ b/rust/src/engine-core-client/src/client.rs @@ -692,6 +692,38 @@ impl EngineCoreClient { Ok(()) } + /// Pause the scheduler so generation can be halted + pub async fn pause_scheduler(&self, mode: &str, clear_cache: bool) -> Result<()> { + self.call_utility::<(), _>("pause_scheduler", (mode, clear_cache)).await?; + Ok(()) + } + + /// Resume the scheduler after a pause + pub async fn resume_scheduler(&self) -> Result<()> { + self.call_utility::<(), _>("resume_scheduler", ()).await?; + Ok(()) + } + + /// Return whether the scheduler is currently in any pause state. + pub async fn is_scheduler_paused(&self) -> Result { + let results: Vec = self.call_utility("is_scheduler_paused", ()).await?; + // `engine_count >= 1` is enforced during startup handshake, so `results` + // is normally non-empty; fall back to a fail-loud error rather than + // indexing in case that invariant is ever bypassed. + let first = *results.first().ok_or_else(|| Error::InconsistentUtilityResults { + method: "is_scheduler_paused".to_string(), + values: "[]".to_string(), + })?; + if results.iter().all(|&v| v == first) { + Ok(first) + } else { + Err(Error::InconsistentUtilityResults { + method: "is_scheduler_paused".to_string(), + values: format!("{results:?}"), + }) + } + } + /// Shut down local client tasks and close transport state. pub async fn shutdown(self) -> Result<()> { let Self { diff --git a/rust/src/server/src/routes.rs b/rust/src/server/src/routes.rs index a0473c783a..e1ad2e310c 100644 --- a/rust/src/server/src/routes.rs +++ b/rust/src/server/src/routes.rs @@ -6,6 +6,7 @@ mod load; mod lora; mod metrics; pub(crate) mod openai; +mod pause; mod server_info; mod sleep; mod version; @@ -90,6 +91,9 @@ fn build_router_with_options( .route("/sleep", post(sleep::sleep)) .route("/wake_up", post(sleep::wake_up)) .route("/is_sleeping", get(sleep::is_sleeping)) + .route("/pause", post(pause::pause)) + .route("/resume", post(pause::resume)) + .route("/is_paused", get(pause::is_paused)) .route("/server_info", get(server_info::server_info)) } diff --git a/rust/src/server/src/routes/pause.rs b/rust/src/server/src/routes/pause.rs new file mode 100644 index 0000000000..ef55a6f5e4 --- /dev/null +++ b/rust/src/server/src/routes/pause.rs @@ -0,0 +1,90 @@ +use std::sync::Arc; + +use axum::Json; +use axum::extract::{Query, State}; +use serde::{Deserialize, Serialize}; + +use crate::error::ApiError; +use crate::state::AppState; +use crate::utils::utility_call_error; + +#[derive(Debug, Deserialize)] +pub(crate) struct PauseParams { + #[serde(default = "default_pause_mode")] + mode: String, + #[serde(default = "default_clear_cache")] + clear_cache: bool, +} + +#[derive(Serialize)] +pub(crate) struct StatusResponse { + status: &'static str, +} + +#[derive(Serialize)] +pub(crate) struct IsPausedResponse { + is_paused: bool, +} + +/// Pause modes accepted by the engine (`PauseMode` in the Python frontend). +const VALID_PAUSE_MODES: [&str; 3] = ["abort", "wait", "keep"]; + +fn default_pause_mode() -> String { + "abort".to_string() +} + +const fn default_clear_cache() -> bool { + true +} + +// TODO: the Python frontend also accepts the deprecated +// `wait_for_inflight_requests` flag (equivalent to `mode="wait"`); it is +// intentionally omitted here in favor of the `mode` parameter. + +/// Pause the scheduler so generation can be halted (e.g. for weight updates). +pub async fn pause( + State(state): State>, + Query(params): Query, +) -> Result, ApiError> { + if !VALID_PAUSE_MODES.contains(¶ms.mode.as_str()) { + return Err(ApiError::invalid_request( + format!( + "Invalid pause mode '{}'; expected one of: abort, wait, keep", + params.mode + ), + Some("mode"), + )); + } + + state + .engine_core_client() + .pause_scheduler(¶ms.mode, params.clear_cache) + .await + .map_err(|error| utility_call_error("pause", error))?; + + Ok(Json(StatusResponse { status: "paused" })) +} + +/// Resume the scheduler after a pause. +pub async fn resume(State(state): State>) -> Result, ApiError> { + state + .engine_core_client() + .resume_scheduler() + .await + .map_err(|error| utility_call_error("resume", error))?; + + Ok(Json(StatusResponse { status: "resumed" })) +} + +/// Return whether the scheduler is currently paused. +pub async fn is_paused( + State(state): State>, +) -> Result, ApiError> { + let is_paused = state + .engine_core_client() + .is_scheduler_paused() + .await + .map_err(|error| utility_call_error("is_paused", error))?; + + Ok(Json(IsPausedResponse { is_paused })) +} diff --git a/rust/src/server/src/routes/tests.rs b/rust/src/server/src/routes/tests.rs index 0a269f4ce1..9f9288958e 100644 --- a/rust/src/server/src/routes/tests.rs +++ b/rust/src/server/src/routes/tests.rs @@ -4228,6 +4228,161 @@ async fn is_sleeping_route_returns_json_payload() { ); } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn pause_route_uses_python_compatible_default_query_values() { + let (app, engine_task) = test_admin_app_with_engine_script(|dealer, push| { + boxed_test_future(async move { + let utility = recv_engine_message(dealer).await; + assert_eq!(utility[0].as_ref(), &[0x03]); + + let payload = decode_value(&utility[1]).expect("decode utility payload"); + let array = payload.as_array().expect("utility payload array"); + let call_id = array[1].as_u64().expect("call id"); + + assert_eq!(array[2], Value::from("pause_scheduler")); + assert_eq!( + array[3], + Value::Array(vec![Value::from("abort"), Value::from(true)]) + ); + + send_outputs(push, utility_outputs(call_id, utility_none_result())).await; + }) + }) + .await; + + let response = app + .clone() + .call( + Request::builder() + .method("POST") + .uri("/pause") + .body(Body::empty()) + .expect("build request"), + ) + .await + .expect("call app"); + + assert_eq!(response.status(), StatusCode::OK); + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body"); + engine_task.await.expect("mock engine task"); + + assert_eq!( + serde_json::from_slice::(&body).expect("decode json"), + json!({ "status": "paused" }) + ); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn pause_route_rejects_invalid_mode() { + let (app, engine_task) = + test_admin_app_with_engine_script(|_dealer, _push| boxed_test_future(async move {})).await; + + let response = app + .clone() + .call( + Request::builder() + .method("POST") + .uri("/pause?mode=banana") + .body(Body::empty()) + .expect("build request"), + ) + .await + .expect("call app"); + + assert_eq!(response.status(), StatusCode::BAD_REQUEST); + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body"); + let json: serde_json::Value = serde_json::from_slice(&body).expect("decode json"); + assert_eq!(json["error"]["type"], "invalid_request_error"); + assert_eq!(json["error"]["param"], "mode"); + engine_task.abort_and_join().await; +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn resume_route_sends_no_args() { + let (app, engine_task) = test_admin_app_with_engine_script(|dealer, push| { + boxed_test_future(async move { + let utility = recv_engine_message(dealer).await; + assert_eq!(utility[0].as_ref(), &[0x03]); + + let payload = decode_value(&utility[1]).expect("decode utility payload"); + let array = payload.as_array().expect("utility payload array"); + let call_id = array[1].as_u64().expect("call id"); + + assert_eq!(array[2], Value::from("resume_scheduler")); + assert_eq!(array[3], Value::Array(Vec::new())); + + send_outputs(push, utility_outputs(call_id, utility_none_result())).await; + }) + }) + .await; + + let response = app + .clone() + .call( + Request::builder() + .method("POST") + .uri("/resume") + .body(Body::empty()) + .expect("build request"), + ) + .await + .expect("call app"); + + assert_eq!(response.status(), StatusCode::OK); + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body"); + engine_task.await.expect("mock engine task"); + + assert_eq!( + serde_json::from_slice::(&body).expect("decode json"), + json!({ "status": "resumed" }) + ); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn is_paused_route_returns_json_payload() { + let (app, engine_task) = test_admin_app_with_engine_script(|dealer, push| { + boxed_test_future(async move { + let utility = recv_engine_message(dealer).await; + assert_eq!(utility[0].as_ref(), &[0x03]); + + let payload = decode_value(&utility[1]).expect("decode utility payload"); + let array = payload.as_array().expect("utility payload array"); + let call_id = array[1].as_u64().expect("call id"); + + assert_eq!(array[2], Value::from("is_scheduler_paused")); + assert_eq!(array[3], Value::Array(Vec::new())); + + send_outputs(push, utility_outputs(call_id, utility_result_value(true))).await; + }) + }) + .await; + + let response = app + .clone() + .call( + Request::builder() + .method("GET") + .uri("/is_paused") + .body(Body::empty()) + .expect("build request"), + ) + .await + .expect("call app"); + + assert_eq!(response.status(), StatusCode::OK); + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body"); + engine_task.await.expect("mock engine task"); + + assert_eq!( + serde_json::from_slice::(&body).expect("decode json"), + json!({ "is_paused": true }) + ); +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial] async fn admin_routes_are_hidden_when_dev_mode_is_disabled() { @@ -4244,6 +4399,9 @@ async fn admin_routes_are_hidden_when_dev_mode_is_disabled() { ("GET", "/is_sleeping"), ("POST", "/sleep"), ("POST", "/wake_up"), + ("GET", "/is_paused"), + ("POST", "/pause"), + ("POST", "/resume"), ("POST", "/collective_rpc"), ("POST", "/reset_prefix_cache"), ("POST", "/reset_mm_cache"),