A high-throughput and memory-efficient inference and serving engine for LLMs
[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>
Sahil Singh committed
Jun 8, 2026 at 14:58 UTC
3c0b4432be91d4a5fda590f3fb3236cd087882c5
4 files changed
+284
rust/src/engine-core-client/src/client.rs
+32
@@ -692,6 +692,38 @@ impl EngineCoreClient {
692
Ok(())
693
}
694
695
+ /// Pause the scheduler so generation can be halted
696
+ pub async fn pause_scheduler(&self, mode: &str, clear_cache: bool) -> Result<()> {
697
+ self.call_utility::<(), _>("pause_scheduler", (mode, clear_cache)).await?;
698
+ Ok(())
699
+ }
700
+
701
+ /// Resume the scheduler after a pause
702
+ pub async fn resume_scheduler(&self) -> Result<()> {
703
+ self.call_utility::<(), _>("resume_scheduler", ()).await?;
704
+ Ok(())
705
+ }
706
+
707
+ /// Return whether the scheduler is currently in any pause state.
708
+ pub async fn is_scheduler_paused(&self) -> Result<bool> {
709
+ let results: Vec<bool> = self.call_utility("is_scheduler_paused", ()).await?;
710
+ // `engine_count >= 1` is enforced during startup handshake, so `results`
711
+ // is normally non-empty; fall back to a fail-loud error rather than
712
+ // indexing in case that invariant is ever bypassed.
713
+ let first = *results.first().ok_or_else(|| Error::InconsistentUtilityResults {
714
+ method: "is_scheduler_paused".to_string(),
715
+ values: "[]".to_string(),
716
+ })?;
717
+ if results.iter().all(|&v| v == first) {
718
+ Ok(first)
719
+ } else {
720
+ Err(Error::InconsistentUtilityResults {
721
+ method: "is_scheduler_paused".to_string(),
722
+ values: format!("{results:?}"),
723
+ })
724
+ }
725
+ }
726
+
727
/// Shut down local client tasks and close transport state.
728
pub async fn shutdown(self) -> Result<()> {
729
let Self {
rust/src/server/src/routes.rs
+4
@@ -6,6 +6,7 @@ mod load;
6
mod lora;
7
mod metrics;
8
pub(crate) mod openai;
9
+mod pause;
10
mod server_info;
11
mod sleep;
12
mod version;
@@ -90,6 +91,9 @@ fn build_router_with_options(
91
.route("/sleep", post(sleep::sleep))
92
.route("/wake_up", post(sleep::wake_up))
93
.route("/is_sleeping", get(sleep::is_sleeping))
94
+ .route("/pause", post(pause::pause))
95
+ .route("/resume", post(pause::resume))
96
+ .route("/is_paused", get(pause::is_paused))
97
.route("/server_info", get(server_info::server_info))
98
}
99
rust/src/server/src/routes/pause.rs
new
+90
@@ -0,0 +1,90 @@
1
+use std::sync::Arc;
2
+
3
+use axum::Json;
4
+use axum::extract::{Query, State};
5
+use serde::{Deserialize, Serialize};
6
+
7
+use crate::error::ApiError;
8
+use crate::state::AppState;
9
+use crate::utils::utility_call_error;
10
+
11
+#[derive(Debug, Deserialize)]
12
+pub(crate) struct PauseParams {
13
+ #[serde(default = "default_pause_mode")]
14
+ mode: String,
15
+ #[serde(default = "default_clear_cache")]
16
+ clear_cache: bool,
17
+}
18
+
19
+#[derive(Serialize)]
20
+pub(crate) struct StatusResponse {
21
+ status: &'static str,
22
+}
23
+
24
+#[derive(Serialize)]
25
+pub(crate) struct IsPausedResponse {
26
+ is_paused: bool,
27
+}
28
+
29
+/// Pause modes accepted by the engine (`PauseMode` in the Python frontend).
30
+const VALID_PAUSE_MODES: [&str; 3] = ["abort", "wait", "keep"];
31
+
32
+fn default_pause_mode() -> String {
33
+ "abort".to_string()
34
+}
35
+
36
+const fn default_clear_cache() -> bool {
37
+ true
38
+}
39
+
40
+// TODO: the Python frontend also accepts the deprecated
41
+// `wait_for_inflight_requests` flag (equivalent to `mode="wait"`); it is
42
+// intentionally omitted here in favor of the `mode` parameter.
43
+
44
+/// Pause the scheduler so generation can be halted (e.g. for weight updates).
45
+pub async fn pause(
46
+ State(state): State<Arc<AppState>>,
47
+ Query(params): Query<PauseParams>,
48
+) -> Result<Json<StatusResponse>, ApiError> {
49
+ if !VALID_PAUSE_MODES.contains(¶ms.mode.as_str()) {
50
+ return Err(ApiError::invalid_request(
51
+ format!(
52
+ "Invalid pause mode '{}'; expected one of: abort, wait, keep",
53
+ params.mode
54
+ ),
55
+ Some("mode"),
56
+ ));
57
+ }
58
+
59
+ state
60
+ .engine_core_client()
61
+ .pause_scheduler(¶ms.mode, params.clear_cache)
62
+ .await
63
+ .map_err(|error| utility_call_error("pause", error))?;
64
+
65
+ Ok(Json(StatusResponse { status: "paused" }))
66
+}
67
+
68
+/// Resume the scheduler after a pause.
69
+pub async fn resume(State(state): State<Arc<AppState>>) -> Result<Json<StatusResponse>, ApiError> {
70
+ state
71
+ .engine_core_client()
72
+ .resume_scheduler()
73
+ .await
74
+ .map_err(|error| utility_call_error("resume", error))?;
75
+
76
+ Ok(Json(StatusResponse { status: "resumed" }))
77
+}
78
+
79
+/// Return whether the scheduler is currently paused.
80
+pub async fn is_paused(
81
+ State(state): State<Arc<AppState>>,
82
+) -> Result<Json<IsPausedResponse>, ApiError> {
83
+ let is_paused = state
84
+ .engine_core_client()
85
+ .is_scheduler_paused()
86
+ .await
87
+ .map_err(|error| utility_call_error("is_paused", error))?;
88
+
89
+ Ok(Json(IsPausedResponse { is_paused }))
90
+}
rust/src/server/src/routes/tests.rs
+158
@@ -4228,6 +4228,161 @@ async fn is_sleeping_route_returns_json_payload() {
4228
);
4229
}
4230
4231
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
4232
+#[serial]
4233
+async fn pause_route_uses_python_compatible_default_query_values() {
4234
+ let (app, engine_task) = test_admin_app_with_engine_script(|dealer, push| {
4235
+ boxed_test_future(async move {
4236
+ let utility = recv_engine_message(dealer).await;
4237
+ assert_eq!(utility[0].as_ref(), &[0x03]);
4238
+
4239
+ let payload = decode_value(&utility[1]).expect("decode utility payload");
4240
+ let array = payload.as_array().expect("utility payload array");
4241
+ let call_id = array[1].as_u64().expect("call id");
4242
+
4243
+ assert_eq!(array[2], Value::from("pause_scheduler"));
4244
+ assert_eq!(
4245
+ array[3],
4246
+ Value::Array(vec![Value::from("abort"), Value::from(true)])
4247
+ );
4248
+
4249
+ send_outputs(push, utility_outputs(call_id, utility_none_result())).await;
4250
+ })
4251
+ })
4252
+ .await;
4253
+
4254
+ let response = app
4255
+ .clone()
4256
+ .call(
4257
+ Request::builder()
4258
+ .method("POST")
4259
+ .uri("/pause")
4260
+ .body(Body::empty())
4261
+ .expect("build request"),
4262
+ )
4263
+ .await
4264
+ .expect("call app");
4265
+
4266
+ assert_eq!(response.status(), StatusCode::OK);
4267
+ let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
4268
+ engine_task.await.expect("mock engine task");
4269
+
4270
+ assert_eq!(
4271
+ serde_json::from_slice::<serde_json::Value>(&body).expect("decode json"),
4272
+ json!({ "status": "paused" })
4273
+ );
4274
+}
4275
+
4276
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
4277
+#[serial]
4278
+async fn pause_route_rejects_invalid_mode() {
4279
+ let (app, engine_task) =
4280
+ test_admin_app_with_engine_script(|_dealer, _push| boxed_test_future(async move {})).await;
4281
+
4282
+ let response = app
4283
+ .clone()
4284
+ .call(
4285
+ Request::builder()
4286
+ .method("POST")
4287
+ .uri("/pause?mode=banana")
4288
+ .body(Body::empty())
4289
+ .expect("build request"),
4290
+ )
4291
+ .await
4292
+ .expect("call app");
4293
+
4294
+ assert_eq!(response.status(), StatusCode::BAD_REQUEST);
4295
+ let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
4296
+ let json: serde_json::Value = serde_json::from_slice(&body).expect("decode json");
4297
+ assert_eq!(json["error"]["type"], "invalid_request_error");
4298
+ assert_eq!(json["error"]["param"], "mode");
4299
+ engine_task.abort_and_join().await;
4300
+}
4301
+
4302
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
4303
+#[serial]
4304
+async fn resume_route_sends_no_args() {
4305
+ let (app, engine_task) = test_admin_app_with_engine_script(|dealer, push| {
4306
+ boxed_test_future(async move {
4307
+ let utility = recv_engine_message(dealer).await;
4308
+ assert_eq!(utility[0].as_ref(), &[0x03]);
4309
+
4310
+ let payload = decode_value(&utility[1]).expect("decode utility payload");
4311
+ let array = payload.as_array().expect("utility payload array");
4312
+ let call_id = array[1].as_u64().expect("call id");
4313
+
4314
+ assert_eq!(array[2], Value::from("resume_scheduler"));
4315
+ assert_eq!(array[3], Value::Array(Vec::new()));
4316
+
4317
+ send_outputs(push, utility_outputs(call_id, utility_none_result())).await;
4318
+ })
4319
+ })
4320
+ .await;
4321
+
4322
+ let response = app
4323
+ .clone()
4324
+ .call(
4325
+ Request::builder()
4326
+ .method("POST")
4327
+ .uri("/resume")
4328
+ .body(Body::empty())
4329
+ .expect("build request"),
4330
+ )
4331
+ .await
4332
+ .expect("call app");
4333
+
4334
+ assert_eq!(response.status(), StatusCode::OK);
4335
+ let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
4336
+ engine_task.await.expect("mock engine task");
4337
+
4338
+ assert_eq!(
4339
+ serde_json::from_slice::<serde_json::Value>(&body).expect("decode json"),
4340
+ json!({ "status": "resumed" })
4341
+ );
4342
+}
4343
+
4344
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
4345
+#[serial]
4346
+async fn is_paused_route_returns_json_payload() {
4347
+ let (app, engine_task) = test_admin_app_with_engine_script(|dealer, push| {
4348
+ boxed_test_future(async move {
4349
+ let utility = recv_engine_message(dealer).await;
4350
+ assert_eq!(utility[0].as_ref(), &[0x03]);
4351
+
4352
+ let payload = decode_value(&utility[1]).expect("decode utility payload");
4353
+ let array = payload.as_array().expect("utility payload array");
4354
+ let call_id = array[1].as_u64().expect("call id");
4355
+
4356
+ assert_eq!(array[2], Value::from("is_scheduler_paused"));
4357
+ assert_eq!(array[3], Value::Array(Vec::new()));
4358
+
4359
+ send_outputs(push, utility_outputs(call_id, utility_result_value(true))).await;
4360
+ })
4361
+ })
4362
+ .await;
4363
+
4364
+ let response = app
4365
+ .clone()
4366
+ .call(
4367
+ Request::builder()
4368
+ .method("GET")
4369
+ .uri("/is_paused")
4370
+ .body(Body::empty())
4371
+ .expect("build request"),
4372
+ )
4373
+ .await
4374
+ .expect("call app");
4375
+
4376
+ assert_eq!(response.status(), StatusCode::OK);
4377
+ let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
4378
+ engine_task.await.expect("mock engine task");
4379
+
4380
+ assert_eq!(
4381
+ serde_json::from_slice::<serde_json::Value>(&body).expect("decode json"),
4382
+ json!({ "is_paused": true })
4383
+ );
4384
+}
4385
+
4386
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
4387
#[serial]
4388
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() {
4399
("GET", "/is_sleeping"),
4400
("POST", "/sleep"),
4401
("POST", "/wake_up"),
4402
+ ("GET", "/is_paused"),
4403
+ ("POST", "/pause"),
4404
+ ("POST", "/resume"),
4405
("POST", "/collective_rpc"),
4406
("POST", "/reset_prefix_cache"),
4407
("POST", "/reset_mm_cache"),