[Rust Frontend] Add /pause, /resume, /is_paused endpoints (#44499)
Signed-off-by: Sahil Singh <sahiilsiingh37@gmail.com> Co-authored-by: Bugen Zhao <i@bugenzhao.com>
This commit is contained in:
Sahil Singh
co-authored by
Bugen Zhao
Bugen Zhao
parent
469f3dcf1d
commit
3c0b4432be
@@ -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<bool> {
|
||||
let results: Vec<bool> = 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 {
|
||||
|
||||
@@ -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))
|
||||
}
|
||||
|
||||
|
||||
@@ -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<Arc<AppState>>,
|
||||
Query(params): Query<PauseParams>,
|
||||
) -> Result<Json<StatusResponse>, 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<Arc<AppState>>) -> Result<Json<StatusResponse>, 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<Arc<AppState>>,
|
||||
) -> Result<Json<IsPausedResponse>, 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 }))
|
||||
}
|
||||
@@ -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::<serde_json::Value>(&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::<serde_json::Value>(&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::<serde_json::Value>(&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"),
|
||||
|
||||
Reference in New Issue
Block a user