@tduncan / Nvidia-Optimization-Stack-VLLM / commits / e5f3d08d72

[Frontend][RL] Align sleep-mode API responses and operation metrics (#52864)

Signed-off-by: Ronald1995 <ronaldautomobile@163.com> Signed-off-by: aoshen02 <aoshen02@users.noreply.github.com> Co-authored-by: aoshen02 <aoshen02@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

Ronald committed Sep 30, 2026 at 14:16 UTC e5f3d08d722870a5301560fad0cfa65b7e333900
14 files changed +834 -48
rust/src/engine-core-client/src/tests/client.rs
+83
@@ -2290,6 +2290,89 @@ async fn multi_engine_abort_is_grouped_and_utility_fans_out_to_all_engines() {
2290 client.shutdown().await.unwrap();
2291 }
2292
2293 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
2294 +async fn wake_up_requires_every_engine_to_be_fully_awake() {
2295 + init_tracing();
2296 + let ipc = IpcNamespace::new().unwrap();
2297 + let handshake_address = ipc.handshake_endpoint();
2298 + let mut engines = Vec::new();
2299 + for engine_index in 0..2 {
2300 + let engine = spawn_mock_engine_task(
2301 + handshake_address.clone(),
2302 + EngineId::from_engine_index(engine_index).into_frame().to_vec(),
2303 + move |dealer, push| {
2304 + Box::pin(async move {
2305 + let results = if engine_index == 0 {
2306 + [true, false, true]
2307 + } else {
2308 + [false, true, true]
2309 + };
2310 + for (attempt, fully_awake) in results.into_iter().enumerate() {
2311 + let utility = recv_engine_message(dealer).await;
2312 + assert_eq!(utility[0].as_ref(), &[0x03]);
2313 + let payload = decode_value(&utility[1]);
2314 + let array = payload.as_array().expect("utility payload array");
2315 + let call_id = array[1].as_u64().expect("call_id");
2316 + assert_eq!(array[2], Value::from("wake_up"));
2317 + let tags = if attempt == 1 {
2318 + Value::Array(vec![Value::from("weights")])
2319 + } else {
2320 + Value::Nil
2321 + };
2322 + assert_eq!(array[3], Value::Array(vec![tags]));
2323 + send_outputs(
2324 + push,
2325 + UtilityCallOutput {
2326 + engine_index: u32::from(engine_index),
2327 + timestamp: 0.0,
2328 + output: UtilityOutput {
2329 + call_id: call_id.into(),
2330 + failure_message: None,
2331 + result: Some(utility_result_value(fully_awake)),
2332 + },
2333 + }
2334 + .into(),
2335 + )
2336 + .await;
2337 + }
2338 + })
2339 + },
2340 + );
2341 + engines.push(engine);
2342 + }
2343 + let client = connect_client_with_ipc(
2344 + handshake_test_config(
2345 + handshake_address,
2346 + 2,
2347 + "test-model",
2348 + Duration::from_secs(2),
2349 + 5,
2350 + None,
2351 + ),
2352 + &ipc,
2353 + )
2354 + .await;
2355 +
2356 + for (tags, expected) in [
2357 + (None, false),
2358 + (Some(vec!["weights".to_owned()]), false),
2359 + (None, true),
2360 + ] {
2361 + assert_eq!(
2362 + timeout(Duration::from_secs(2), client.wake_up(tags))
2363 + .await
2364 + .expect("wake timeout")
2365 + .expect("wake result"),
2366 + expected
2367 + );
2368 + }
2369 + for (shutdown_tx, engine_task) in engines {
2370 + let _ = shutdown_tx.send(());
2371 + engine_task.await.unwrap();
2372 + }
2373 + client.shutdown().await.unwrap();
2374 +}
2375 +
2376 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
2377 async fn collective_rpc_flattens_results_from_all_engines() {
2378 init_tracing();
rust/src/metrics/src/api_server.rs
+118 -1
@@ -1,11 +1,13 @@
1 // SPDX-License-Identifier: Apache-2.0
2 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 +use std::time::Instant;
5 +
6 use prometheus_client::encoding::EncodeLabelSet;
7 use prometheus_client::metrics::family::Family;
8 use prometheus_client::registry::Registry;
9
8 -use crate::{Histogram, U64Counter};
10 +use crate::{Histogram, U64Counter, U64Gauge};
11
12 const HTTP_REQUEST_DURATION_BUCKETS: [f64; 3] = [0.1, 0.5, 1.0];
13 const HTTP_REQUEST_DURATION_HIGHR_BUCKETS: [f64; 21] = [
@@ -21,6 +23,10 @@ fn http_request_duration_highr_histogram() -> Histogram {
23 Histogram::new(HTTP_REQUEST_DURATION_HIGHR_BUCKETS.iter().copied())
24 }
25
26 +fn sleep_mode_operation_duration_histogram() -> Histogram {
27 + Histogram::new([0.01, 0.1, 1.0, 10.0, 30.0, 60.0, 120.0, 300.0, 600.0])
28 +}
29 +
30 #[derive(Clone, Debug, Hash, PartialEq, Eq, EncodeLabelSet)]
31 pub struct HttpRequestLabels {
32 pub method: String,
@@ -34,15 +40,25 @@ pub struct HttpHandlerLabels {
40 pub handler: String,
41 }
42
43 +#[derive(Clone, Debug, Hash, PartialEq, Eq, EncodeLabelSet)]
44 +pub struct SleepModeOperationLabels {
45 + pub operation: &'static str,
46 +}
47 +
48 pub(crate) type HttpRequestCounterFamily = Family<HttpRequestLabels, U64Counter>;
49 pub(crate) type HttpHandlerHistogramFamily =
50 Family<HttpHandlerLabels, Histogram, fn() -> Histogram>;
51 +pub(crate) type SleepModeOperationHistogramFamily =
52 + Family<SleepModeOperationLabels, Histogram, fn() -> Histogram>;
53 +pub(crate) type SleepModeOperationGaugeFamily = Family<SleepModeOperationLabels, U64Gauge>;
54
55 /// API-server Prometheus families exported from the HTTP middleware layer.
56 pub struct ApiServerMetrics {
57 pub http_requests: HttpRequestCounterFamily,
58 pub http_request_duration_seconds: HttpHandlerHistogramFamily,
59 pub http_request_duration_highr_seconds: Histogram,
60 + pub sleep_mode_operation_duration_seconds: SleepModeOperationHistogramFamily,
61 + pub sleep_mode_operations_in_flight: SleepModeOperationGaugeFamily,
62 }
63
64 impl ApiServerMetrics {
@@ -71,10 +87,111 @@ impl ApiServerMetrics {
87 http_request_duration_highr_seconds.clone(),
88 );
89
90 + let sleep_mode_operation_duration_seconds = Family::new_with_constructor(
91 + sleep_mode_operation_duration_histogram as fn() -> Histogram,
92 + );
93 + registry.register(
94 + "vllm:sleep_mode_operation_duration_seconds",
95 + "Duration of one sleep-mode operation.",
96 + sleep_mode_operation_duration_seconds.clone(),
97 + );
98 +
99 + let sleep_mode_operations_in_flight = SleepModeOperationGaugeFamily::default();
100 + registry.register(
101 + "vllm:sleep_mode_operations_in_flight",
102 + "Sleep-mode operations currently awaited.",
103 + sleep_mode_operations_in_flight.clone(),
104 + );
105 +
106 Self {
107 http_requests,
108 http_request_duration_seconds,
109 http_request_duration_highr_seconds,
110 + sleep_mode_operation_duration_seconds,
111 + sleep_mode_operations_in_flight,
112 + }
113 + }
114 +
115 + /// Record one dispatched sleep-mode engine operation.
116 + pub fn record_sleep_mode_operation(
117 + &self,
118 + operation: &'static str,
119 + ) -> SleepModeOperationRecorder {
120 + let labels = SleepModeOperationLabels { operation };
121 + let in_flight = self.sleep_mode_operations_in_flight.get_or_create_owned(&labels);
122 + in_flight.inc();
123 + SleepModeOperationRecorder {
124 + started_at: Instant::now(),
125 + duration: self.sleep_mode_operation_duration_seconds.get_or_create_owned(&labels),
126 + in_flight,
127 + }
128 + }
129 +}
130 +
131 +/// Records duration and clears in-flight tracking on completion or cancellation.
132 +pub struct SleepModeOperationRecorder {
133 + started_at: Instant,
134 + duration: Histogram,
135 + in_flight: U64Gauge,
136 +}
137 +
138 +impl Drop for SleepModeOperationRecorder {
139 + fn drop(&mut self) {
140 + self.in_flight.dec();
141 + self.duration.observe(self.started_at.elapsed().as_secs_f64());
142 + }
143 +}
144 +
145 +#[cfg(test)]
146 +mod tests {
147 + use prometheus_client::encoding::text::encode;
148 + use prometheus_client::registry::Registry;
149 +
150 + use super::ApiServerMetrics;
151 +
152 + fn rendered_metrics(registry: &Registry) -> String {
153 + let mut rendered = String::new();
154 + encode(&mut rendered, registry).expect("encode metrics");
155 + rendered
156 + }
157 +
158 + #[test]
159 + fn sleep_mode_recorder_tracks_duration_and_concurrency() {
160 + let mut registry = Registry::default();
161 + let metrics = ApiServerMetrics::register(&mut registry);
162 + for operation in ["sleep", "release_kv_cache_memory", "wake"] {
163 + let recorder = metrics.record_sleep_mode_operation(operation);
164 + let active = rendered_metrics(&registry);
165 + assert!(active.contains(&format!(
166 + "vllm:sleep_mode_operations_in_flight{{operation=\"{operation}\"}} 1"
167 + )));
168 +
169 + drop(recorder);
170 + let succeeded = rendered_metrics(&registry);
171 + assert!(succeeded.contains(&format!(
172 + "vllm:sleep_mode_operations_in_flight{{operation=\"{operation}\"}} 0"
173 + )));
174 + assert!(succeeded.contains(&format!(
175 + "vllm:sleep_mode_operation_duration_seconds_count{{operation=\"{operation}\"}} 1"
176 + )));
177 +
178 + let first = metrics.record_sleep_mode_operation(operation);
179 + let second = metrics.record_sleep_mode_operation(operation);
180 + assert!(rendered_metrics(&registry).contains(&format!(
181 + "vllm:sleep_mode_operations_in_flight{{operation=\"{operation}\"}} 2"
182 + )));
183 + drop(first);
184 + assert!(rendered_metrics(&registry).contains(&format!(
185 + "vllm:sleep_mode_operations_in_flight{{operation=\"{operation}\"}} 1"
186 + )));
187 + drop(second);
188 + let rendered = rendered_metrics(&registry);
189 + assert!(rendered.contains(&format!(
190 + "vllm:sleep_mode_operations_in_flight{{operation=\"{operation}\"}} 0"
191 + )));
192 + assert!(rendered.contains(&format!(
193 + "vllm:sleep_mode_operation_duration_seconds_count{{operation=\"{operation}\"}} 3"
194 + )));
195 }
196 }
197 }
rust/src/server/src/middleware/metrics.rs
-4
@@ -25,10 +25,6 @@ const EXCLUDED_HANDLERS: &[&str] = &[
25 "/reset_mm_cache",
26 "/reset_encoder_cache",
27 "/collective_rpc",
28 - "/sleep",
29 - "/release_kv_cache_memory",
30 - "/wake_up",
31 - "/is_sleeping",
28 ];
29
30 /// Map the request method to a bounded Prometheus label.
rust/src/server/src/routes/sleep.rs
+47 -12
@@ -6,9 +6,9 @@ use std::sync::Arc;
6 use axum::Json;
7 use axum::extract::rejection::QueryRejection;
8 use axum::extract::{Query, RawQuery, State};
9 -use axum::http::StatusCode;
9 use serde::{Deserialize, Serialize};
10 use vllm_engine_core_client::protocol::utility::PauseMode;
11 +use vllm_metrics::METRICS;
12
13 use crate::error::ApiError;
14 use crate::state::AppState;
@@ -19,6 +19,24 @@ pub(crate) struct IsSleepingResponse {
19 is_sleeping: bool,
20 }
21
22 +#[derive(Serialize)]
23 +pub struct SleepResponse {
24 + status: &'static str,
25 + level: u32,
26 +}
27 +
28 +#[derive(Serialize)]
29 +pub struct ReleaseKvCacheMemoryResponse {
30 + status: &'static str,
31 +}
32 +
33 +#[derive(Serialize)]
34 +pub struct WakeUpResponse {
35 + status: &'static str,
36 + /// Tags requested by the caller; null means wake all resources.
37 + tags: Option<Vec<String>>,
38 +}
39 +
40 #[derive(Debug, Deserialize)]
41 pub(crate) struct SleepParams {
42 #[serde(default = "default_sleep_level")]
@@ -32,43 +50,56 @@ const fn default_sleep_level() -> u32 {
50 }
51
52 fn invalid_query(error: QueryRejection) -> ApiError {
35 - ApiError::invalid_request(error.body_text(), Some("mode"))
53 + ApiError::invalid_request(error.body_text(), None)
54 }
55
56 /// Put the engine to sleep.
57 pub async fn sleep(
58 State(state): State<Arc<AppState>>,
59 params: Result<Query<SleepParams>, QueryRejection>,
42 -) -> Result<StatusCode, ApiError> {
43 - let Query(params) = params.map_err(invalid_query)?;
60 +) -> Result<Json<SleepResponse>, ApiError> {
61 + let Query(SleepParams { level, mode }) = params.map_err(invalid_query)?;
62 + if level > 2 {
63 + return Err(ApiError::invalid_request(
64 + "level must be between 0 and 2",
65 + Some("level"),
66 + ));
67 + }
68 + let _recorder = METRICS.api_server.record_sleep_mode_operation("sleep");
69
70 state
71 .engine_core_client()
47 - .sleep(params.level, params.mode)
72 + .sleep(level, mode)
73 .await
74 .map_err(|error| utility_call_error("sleep", error))?;
75
51 - Ok(StatusCode::OK)
76 + Ok(Json(SleepResponse {
77 + status: "sleeping",
78 + level,
79 + }))
80 }
81
82 /// Release KV cache memory while keeping model weights resident.
83 pub async fn release_kv_cache_memory(
84 State(state): State<Arc<AppState>>,
57 -) -> Result<StatusCode, ApiError> {
85 +) -> Result<Json<ReleaseKvCacheMemoryResponse>, ApiError> {
86 + let _recorder = METRICS.api_server.record_sleep_mode_operation("release_kv_cache_memory");
87 state
88 .engine_core_client()
89 .release_kv_cache_memory()
90 .await
91 .map_err(|error| utility_call_error("release_kv_cache_memory", error))?;
92
64 - Ok(StatusCode::OK)
93 + Ok(Json(ReleaseKvCacheMemoryResponse {
94 + status: "kv_cache_released",
95 + }))
96 }
97
98 /// Wake the engine from sleep mode.
99 pub async fn wake_up(
100 State(state): State<Arc<AppState>>,
101 RawQuery(query): RawQuery,
71 -) -> Result<StatusCode, ApiError> {
102 +) -> Result<Json<WakeUpResponse>, ApiError> {
103 let tags = query.and_then(|query| {
104 let tags = url::form_urlencoded::parse(query.as_bytes())
105 .filter(|(key, _)| key == "tags")
@@ -77,13 +108,17 @@ pub async fn wake_up(
108 (!tags.is_empty()).then_some(tags)
109 });
110
80 - state
111 + let _recorder = METRICS.api_server.record_sleep_mode_operation("wake");
112 + let fully_awake = state
113 .engine_core_client()
82 - .wake_up(tags)
114 + .wake_up(tags.clone())
115 .await
116 .map_err(|error| utility_call_error("wake_up", error))?;
117
86 - Ok(StatusCode::OK)
118 + Ok(Json(WakeUpResponse {
119 + status: if fully_awake { "awake" } else { "sleeping" },
120 + tags,
121 + }))
122 }
123
124 /// Return whether the engine is currently sleeping at any level.
rust/src/server/src/routes/tests.rs
+279 -5
@@ -5892,6 +5892,7 @@ async fn weight_checker_route_rejects_invalid_requests_before_engine_calls() {
5892 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
5893 #[serial]
5894 async fn sleep_route_uses_python_compatible_default_query_values() {
5895 + let before = METRICS.render().expect("render metrics");
5896 let (app, engine_task) = test_admin_app_with_engine_script(|dealer, push| {
5897 boxed_test_future(async move {
5898 let utility = recv_engine_message(dealer).await;
@@ -5927,13 +5928,95 @@ async fn sleep_route_uses_python_compatible_default_query_values() {
5928 let status = response.status();
5929 let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
5930 assert_eq!(status, StatusCode::OK, "{}", String::from_utf8_lossy(&body));
5930 - assert!(body.is_empty());
5931 + assert_eq!(
5932 + serde_json::from_slice::<serde_json::Value>(&body).expect("decode json"),
5933 + json!({ "status": "sleeping", "level": 1 })
5934 + );
5935 engine_task.await.expect("mock engine task");
5936 + let after = METRICS.render().expect("render metrics");
5937 + assert_eq!(
5938 + metric_delta(
5939 + &before,
5940 + &after,
5941 + "http_requests_total",
5942 + Some("method=\"POST\",status=\"2xx\",handler=\"/sleep\""),
5943 + ),
5944 + 1.0
5945 + );
5946 + assert_eq!(
5947 + metric_delta(
5948 + &before,
5949 + &after,
5950 + "vllm:sleep_mode_operation_duration_seconds_count",
5951 + Some("operation=\"sleep\""),
5952 + ),
5953 + 1.0
5954 + );
5955 + assert_eq!(
5956 + metric_value(
5957 + &after,
5958 + "vllm:sleep_mode_operations_in_flight",
5959 + Some("operation=\"sleep\""),
5960 + ),
5961 + Some(0.0)
5962 + );
5963 + assert_eq!(
5964 + metric_delta(
5965 + &before,
5966 + &after,
5967 + "http_request_duration_seconds_count",
5968 + Some("method=\"POST\",handler=\"/sleep\""),
5969 + ),
5970 + 1.0
5971 + );
5972 +}
5973 +
5974 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
5975 +#[serial]
5976 +async fn sleep_route_rejects_invalid_query_before_engine_dispatch() {
5977 + let before = METRICS.render().expect("render metrics");
5978 + let (app, engine_task) =
5979 + test_admin_app_with_engine_script(|_dealer, _push| boxed_test_future(async move {})).await;
5980 + for query in ["level=invalid", "level=-1", "level=3", "mode=invalid"] {
5981 + let response = app
5982 + .clone()
5983 + .call(
5984 + Request::builder()
5985 + .method("POST")
5986 + .uri(format!("/sleep?{query}"))
5987 + .body(Body::empty())
5988 + .expect("build request"),
5989 + )
5990 + .await
5991 + .expect("call app");
5992 + assert_eq!(response.status(), StatusCode::BAD_REQUEST, "{query}");
5993 + }
5994 + engine_task.abort_and_join().await;
5995 + let after = METRICS.render().expect("render metrics");
5996 + assert_eq!(
5997 + metric_delta(
5998 + &before,
5999 + &after,
6000 + "http_requests_total",
6001 + Some("method=\"POST\",status=\"4xx\",handler=\"/sleep\""),
6002 + ),
6003 + 4.0
6004 + );
6005 + for metric in [
6006 + "vllm:sleep_mode_operation_duration_seconds_count",
6007 + "vllm:sleep_mode_operations_in_flight",
6008 + ] {
6009 + assert_eq!(
6010 + metric_delta(&before, &after, metric, Some("operation=\"sleep\"")),
6011 + 0.0
6012 + );
6013 + }
6014 }
6015
6016 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
6017 #[serial]
6018 async fn release_kv_cache_memory_route_sends_expected_utility_call() {
6019 + let before = METRICS.render().expect("render metrics");
6020 let (app, engine_task) = test_admin_app_with_engine_script(|dealer, push| {
6021 boxed_test_future(async move {
6022 let utility = recv_engine_message(dealer).await;
@@ -5966,13 +6049,177 @@ async fn release_kv_cache_memory_route_sends_expected_utility_call() {
6049 let status = response.status();
6050 let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
6051 assert_eq!(status, StatusCode::OK, "{}", String::from_utf8_lossy(&body));
5969 - assert!(body.is_empty());
6052 + assert_eq!(
6053 + serde_json::from_slice::<serde_json::Value>(&body).expect("decode json"),
6054 + json!({ "status": "kv_cache_released" })
6055 + );
6056 engine_task.await.expect("mock engine task");
6057 + let after = METRICS.render().expect("render metrics");
6058 + assert_eq!(
6059 + metric_delta(
6060 + &before,
6061 + &after,
6062 + "http_requests_total",
6063 + Some("method=\"POST\",status=\"2xx\",handler=\"/release_kv_cache_memory\""),
6064 + ),
6065 + 1.0
6066 + );
6067 +}
6068 +
6069 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
6070 +#[serial]
6071 +async fn release_kv_cache_memory_error_records_http_status_and_operation_duration() {
6072 + let before = METRICS.render().expect("render metrics");
6073 + let (app, engine_task) = test_admin_app_with_engine_script(|dealer, push| {
6074 + boxed_test_future(async move {
6075 + let utility = recv_engine_message(dealer).await;
6076 + let payload = decode_value(&utility[1]).expect("decode utility payload");
6077 + let call_id =
6078 + payload.as_array().expect("utility payload array")[1].as_u64().expect("call id");
6079 + send_outputs(
6080 + push,
6081 + UtilityCallOutput {
6082 + output: UtilityOutput {
6083 + call_id: call_id.into(),
6084 + failure_message: Some("requires a completed pause first".to_owned()),
6085 + result: None,
6086 + },
6087 + ..Default::default()
6088 + }
6089 + .into(),
6090 + )
6091 + .await;
6092 + })
6093 + })
6094 + .await;
6095 + let response = app
6096 + .clone()
6097 + .call(
6098 + Request::builder()
6099 + .method("POST")
6100 + .uri("/release_kv_cache_memory")
6101 + .body(Body::empty())
6102 + .expect("build request"),
6103 + )
6104 + .await
6105 + .expect("call app");
6106 + assert_eq!(response.status(), StatusCode::INTERNAL_SERVER_ERROR);
6107 + engine_task.await.expect("mock engine task");
6108 + let after = METRICS.render().expect("render metrics");
6109 + assert_eq!(
6110 + metric_delta(
6111 + &before,
6112 + &after,
6113 + "http_requests_total",
6114 + Some("method=\"POST\",status=\"5xx\",handler=\"/release_kv_cache_memory\""),
6115 + ),
6116 + 1.0
6117 + );
6118 + assert_eq!(
6119 + metric_delta(
6120 + &before,
6121 + &after,
6122 + "vllm:sleep_mode_operation_duration_seconds_count",
6123 + Some("operation=\"release_kv_cache_memory\""),
6124 + ),
6125 + 1.0
6126 + );
6127 + assert_eq!(
6128 + metric_value(
6129 + &after,
6130 + "vllm:sleep_mode_operations_in_flight",
6131 + Some("operation=\"release_kv_cache_memory\""),
6132 + ),
6133 + Some(0.0)
6134 + );
6135 +}
6136 +
6137 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
6138 +#[serial]
6139 +async fn cancelled_wake_operation_clears_in_flight_and_is_scraped() {
6140 + let before = METRICS.render().expect("render metrics");
6141 + let (dispatched_tx, dispatched_rx) = tokio::sync::oneshot::channel();
6142 + let (app, engine_task) = test_admin_app_with_engine_script(move |dealer, _push| {
6143 + boxed_test_future(async move {
6144 + let utility = recv_engine_message(dealer).await;
6145 + let payload = decode_value(&utility[1]).expect("decode utility payload");
6146 + assert_eq!(
6147 + payload.as_array().expect("utility payload array")[2],
6148 + Value::from("wake_up")
6149 + );
6150 + dispatched_tx.send(()).expect("notify dispatch");
6151 + std::future::pending::<()>().await;
6152 + })
6153 + })
6154 + .await;
6155 + let mut request_app = app.clone();
6156 + let request_task = tokio::spawn(async move {
6157 + request_app
6158 + .call(
6159 + Request::builder()
6160 + .method("POST")
6161 + .uri("/wake_up")
6162 + .body(Body::empty())
6163 + .expect("build request"),
6164 + )
6165 + .await
6166 + });
6167 + tokio::time::timeout(Duration::from_secs(2), dispatched_rx)
6168 + .await
6169 + .expect("wake dispatch timeout")
6170 + .expect("wake dispatched");
6171 + let active = METRICS.render().expect("render metrics");
6172 + assert_eq!(
6173 + metric_delta(
6174 + &before,
6175 + &active,
6176 + "vllm:sleep_mode_operations_in_flight",
6177 + Some("operation=\"wake\""),
6178 + ),
6179 + 1.0
6180 + );
6181 + request_task.abort();
6182 + assert!(request_task.await.expect_err("request cancelled").is_cancelled());
6183 + engine_task.abort_and_join().await;
6184 +
6185 + let response = app
6186 + .clone()
6187 + .call(
6188 + Request::builder()
6189 + .uri("/metrics")
6190 + .body(Body::empty())
6191 + .expect("build scrape request"),
6192 + )
6193 + .await
6194 + .expect("scrape metrics");
6195 + assert_eq!(response.status(), StatusCode::OK);
6196 + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read metrics");
6197 + let after = String::from_utf8(body.to_vec()).expect("utf8 metrics");
6198 + assert_eq!(
6199 + metric_value(
6200 + &after,
6201 + "vllm:sleep_mode_operations_in_flight",
6202 + Some("operation=\"wake\""),
6203 + ),
6204 + Some(0.0)
6205 + );
6206 + assert_eq!(
6207 + metric_delta(
6208 + &before,
6209 + &after,
6210 + "vllm:sleep_mode_operation_duration_seconds_count",
6211 + Some("operation=\"wake\""),
6212 + ),
6213 + 1.0
6214 + );
6215 + assert!(!after.contains("rl_sleep_mode"));
6216 + assert!(!after.contains("sleep_mode_operations_total"));
6217 }
6218
6219 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
6220 #[serial]
6221 async fn wake_up_route_without_tags_accepts_bool_result() {
6222 + let before = METRICS.render().expect("render metrics");
6223 let (app, engine_task) = test_admin_app_with_engine_script(|dealer, push| {
6224 boxed_test_future(async move {
6225 let utility = recv_engine_message(dealer).await;
@@ -6005,8 +6252,21 @@ async fn wake_up_route_without_tags_accepts_bool_result() {
6252 let status = response.status();
6253 let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
6254 assert_eq!(status, StatusCode::OK, "{}", String::from_utf8_lossy(&body));
6008 - assert!(body.is_empty());
6255 + assert_eq!(
6256 + serde_json::from_slice::<serde_json::Value>(&body).expect("decode json"),
6257 + json!({ "status": "awake", "tags": null })
6258 + );
6259 engine_task.await.expect("mock engine task");
6260 + let after = METRICS.render().expect("render metrics");
6261 + assert_eq!(
6262 + metric_delta(
6263 + &before,
6264 + &after,
6265 + "http_requests_total",
6266 + Some("method=\"POST\",status=\"2xx\",handler=\"/wake_up\""),
6267 + ),
6268 + 1.0
6269 + );
6270 }
6271
6272 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
@@ -6030,7 +6290,7 @@ async fn wake_up_route_accepts_repeated_tags() {
6290 ])])
6291 );
6292
6033 - send_outputs(push, utility_outputs(call_id, utility_result_value(true))).await;
6293 + send_outputs(push, utility_outputs(call_id, utility_result_value(false))).await;
6294 })
6295 })
6296 .await;
@@ -6050,13 +6310,17 @@ async fn wake_up_route_accepts_repeated_tags() {
6310 let status = response.status();
6311 let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
6312 assert_eq!(status, StatusCode::OK, "{}", String::from_utf8_lossy(&body));
6053 - assert!(body.is_empty());
6313 + assert_eq!(
6314 + serde_json::from_slice::<serde_json::Value>(&body).expect("decode json"),
6315 + json!({ "status": "sleeping", "tags": ["weights", "kv_cache"] })
6316 + );
6317 engine_task.await.expect("mock engine task");
6318 }
6319
6320 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
6321 #[serial]
6322 async fn is_sleeping_route_returns_json_payload() {
6323 + let before = METRICS.render().expect("render metrics");
6324 let (app, engine_task) = test_admin_app_with_engine_script(|dealer, push| {
6325 boxed_test_future(async move {
6326 let utility = recv_engine_message(dealer).await;
@@ -6094,6 +6358,16 @@ async fn is_sleeping_route_returns_json_payload() {
6358 serde_json::from_slice::<serde_json::Value>(&body).expect("decode json"),
6359 json!({ "is_sleeping": true })
6360 );
6361 + let after = METRICS.render().expect("render metrics");
6362 + assert_eq!(
6363 + metric_delta(
6364 + &before,
6365 + &after,
6366 + "http_requests_total",
6367 + Some("method=\"GET\",status=\"2xx\",handler=\"/is_sleeping\""),
6368 + ),
6369 + 1.0
6370 + );
6371 }
6372
6373 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
tests/entrypoints/serve/dev/test_sleep.py
+179 -5
@@ -1,35 +1,209 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 +import asyncio
5 +from types import SimpleNamespace
6 from unittest.mock import AsyncMock
7
8 import pytest
9 import requests
10 from fastapi import FastAPI
11 from fastapi.testclient import TestClient
12 +from prometheus_client import CollectorRegistry
13 from prometheus_client.parser import text_string_to_metric_families
14
15 from tests.utils import RemoteOpenAIServer
16 +from vllm.entrypoints.serve.dev.sleep import metrics as sleep_metrics
17 from vllm.entrypoints.serve.dev.sleep.api_router import attach_router
18 +from vllm.entrypoints.serve.dev.sleep.metrics import SleepModeOperationMetrics
19 +from vllm.entrypoints.serve.exception_handling.register import init_exception_handler
20 +from vllm.entrypoints.serve.instrumentator.metrics import (
21 + attach_router as attach_metrics_router,
22 +)
23 +from vllm.v1.metrics import prometheus as prometheus_metrics
24
25 MODEL_NAME = "meta-llama/Llama-3.2-1B"
26
27
18 -@pytest.mark.cpu_test
19 -@pytest.mark.parametrize("fails", [False, True])
20 -def test_release_kv_cache_memory_route(fails):
28 +@pytest.fixture
29 +def sleep_route_app(monkeypatch):
30 app = FastAPI()
31 + app.state.args = SimpleNamespace(log_error_stack=False)
32 app.state.engine_client = AsyncMock()
33 + metrics = SleepModeOperationMetrics(CollectorRegistry())
34 + monkeypatch.setattr(
35 + "vllm.entrypoints.serve.dev.sleep.api_router.sleep_mode_operation_metrics",
36 + lambda: metrics,
37 + )
38 +
39 + attach_router(app)
40 + init_exception_handler(app)
41 + return app, metrics
42 +
43 +
44 +@pytest.mark.cpu_test
45 +@pytest.mark.parametrize("level", [0, 1, 2])
46 +def test_sleep_route_response_and_engine_arguments(sleep_route_app, level):
47 + app, _ = sleep_route_app
48 + with TestClient(app) as client:
49 + response = client.post("/sleep", params={"level": level})
50 + assert response.status_code == 200
51 + assert response.json() == {"status": "sleeping", "level": level}
52 + app.state.engine_client.sleep.assert_awaited_once_with(level, "abort")
53 +
54 +
55 +@pytest.mark.cpu_test
56 +@pytest.mark.parametrize(
57 + ("query", "expected_param"),
58 + [
59 + ("level=invalid", "query.level"),
60 + ("level=-1", "query.level"),
61 + ("level=3", "query.level"),
62 + ("mode=invalid", "query.mode"),
63 + ],
64 +)
65 +def test_sleep_route_rejects_invalid_query_before_dispatch(
66 + sleep_route_app, query, expected_param
67 +):
68 + app, metrics = sleep_route_app
69 + with TestClient(app) as client:
70 + response = client.post(f"/sleep?{query}")
71 + assert response.status_code == 400
72 + assert response.json()["error"]["param"] == expected_param
73 + app.state.engine_client.sleep.assert_not_awaited()
74 + assert list(metrics.duration.collect()[0].samples) == []
75 + assert list(metrics.in_flight.collect()[0].samples) == []
76 +
77 +
78 +@pytest.mark.cpu_test
79 +@pytest.mark.parametrize(
80 + ("fully_awake", "tags", "expected"),
81 + [
82 + (True, "", {"status": "awake", "tags": None}),
83 + (False, "?tags=weights", {"status": "sleeping", "tags": ["weights"]}),
84 + ],
85 +)
86 +def test_wake_route_maps_engine_result_without_state_query(
87 + sleep_route_app, fully_awake, tags, expected
88 +):
89 + app, _ = sleep_route_app
90 + app.state.engine_client.wake_up.return_value = fully_awake
91 + with TestClient(app) as client:
92 + response = client.post(f"/wake_up{tags}")
93 + assert response.status_code == 200
94 + assert response.json() == expected
95 + app.state.engine_client.wake_up.assert_awaited_once_with(expected["tags"])
96 + app.state.engine_client.is_sleeping.assert_not_awaited()
97 +
98 +
99 +@pytest.mark.cpu_test
100 +@pytest.mark.parametrize("fails", [False, True])
101 +def test_release_kv_cache_memory_route(sleep_route_app, fails):
102 + app, metrics = sleep_route_app
103 release = app.state.engine_client.release_kv_cache_memory
104 if fails:
105 release.side_effect = RuntimeError("requires a completed pause first")
26 - attach_router(app)
27 -
106 with TestClient(app, raise_server_exceptions=False) as client:
107 response = client.post("/release_kv_cache_memory")
108
109 assert response.status_code == (500 if fails else 200)
110 + if not fails:
111 + assert response.json() == {"status": "kv_cache_released"}
112 release.assert_awaited_once_with()
113 + assert metrics.in_flight.labels("release_kv_cache_memory")._value.get() == 0
114 + assert _duration_count(metrics, "release_kv_cache_memory") == 1
115 +
116 +
117 +@pytest.mark.cpu_test
118 +@pytest.mark.parametrize("operation", ["sleep", "release_kv_cache_memory", "wake"])
119 +@pytest.mark.parametrize("exception", [RuntimeError, asyncio.CancelledError])
120 +def test_sleep_mode_recorder_tracks_duration_and_cleans_up(operation, exception):
121 + metrics = SleepModeOperationMetrics(CollectorRegistry())
122 + with metrics.record(operation):
123 + assert metrics.in_flight.labels(operation)._value.get() == 1
124 + with pytest.raises(exception), metrics.record(operation):
125 + raise exception()
126 + assert metrics.in_flight.labels(operation)._value.get() == 0
127 + assert metrics.duration.labels(operation)._sum.get() >= 0
128 + assert _duration_count(metrics, operation) == 2
129 +
130 +
131 +def _duration_count(metrics, operation):
132 + return next(
133 + sample.value
134 + for sample in metrics.duration.collect()[0].samples
135 + if sample.name.endswith("_count") and sample.labels == {"operation": operation}
136 + )
137 +
138 +
139 +@pytest.mark.cpu_test
140 +@pytest.mark.parametrize(
141 + ("path", "method", "operation", "status", "fails"),
142 + [
143 + ("/sleep", "POST", "sleep", 200, False),
144 + ("/release_kv_cache_memory", "POST", "release_kv_cache_memory", 200, False),
145 + ("/wake_up", "POST", "wake", 200, False),
146 + ("/is_sleeping", "GET", None, 200, False),
147 + ("/sleep?level=3", "POST", None, 400, False),
148 + ("/release_kv_cache_memory", "POST", "release_kv_cache_memory", 500, True),
149 + ],
150 +)
151 +def test_sleep_routes_visible_on_production_metrics_endpoint(
152 + monkeypatch, path, method, operation, status, fails
153 +):
154 + registry = CollectorRegistry()
155 + monkeypatch.delenv("PROMETHEUS_MULTIPROC_DIR", raising=False)
156 + monkeypatch.setattr(sleep_metrics, "REGISTRY", registry)
157 + monkeypatch.setattr(sleep_metrics, "_metrics", None)
158 + monkeypatch.setattr(prometheus_metrics, "REGISTRY", registry)
159 +
160 + app = FastAPI()
161 + app.state.args = SimpleNamespace(log_error_stack=False)
162 + app.state.engine_client = AsyncMock()
163 + app.state.engine_client.wake_up.return_value = True
164 + app.state.engine_client.is_sleeping.return_value = False
165 + if fails:
166 + app.state.engine_client.release_kv_cache_memory.side_effect = RuntimeError(
167 + "engine failed"
168 + )
169 + attach_router(app)
170 + init_exception_handler(app)
171 + attach_metrics_router(app)
172 +
173 + with TestClient(app, raise_server_exceptions=False) as client:
174 + assert client.request(method, path).status_code == status
175 + response = client.get("/metrics")
176 +
177 + assert response.status_code == 200
178 + samples = [
179 + sample
180 + for family in text_string_to_metric_families(response.text)
181 + for sample in family.samples
182 + ]
183 + assert any(
184 + sample.name == "http_requests_total"
185 + and sample.labels
186 + == {
187 + "handler": path.split("?")[0],
188 + "method": method,
189 + "status": f"{status // 100}xx",
190 + }
191 + and sample.value == 1
192 + for sample in samples
193 + )
194 + assert not any("sleep_mode_operations_total" in sample.name for sample in samples)
195 + assert any(
196 + sample.name == "vllm:sleep_mode_operations_in_flight"
197 + and sample.labels == {"operation": operation}
198 + and sample.value == 0
199 + for sample in samples
200 + ) == (operation is not None)
201 + assert any(
202 + sample.name == "vllm:sleep_mode_operation_duration_seconds_count"
203 + and sample.labels == {"operation": operation}
204 + and sample.value == 1
205 + for sample in samples
206 + ) == (operation is not None)
207
208
209 def test_sleep_mode():
tests/v1/engine/test_wake_up_contract.py new
+32
@@ -0,0 +1,32 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from unittest.mock import AsyncMock, call
5 +
6 +import pytest
7 +
8 +from vllm.v1.engine.core_client import DPLBAsyncMPClient
9 +
10 +pytestmark = pytest.mark.cpu_test
11 +
12 +
13 +@pytest.mark.asyncio
14 +@pytest.mark.parametrize(
15 + "results, expected",
16 + [
17 + ([True, True], True),
18 + ([True, False], False),
19 + ([False, True], False),
20 + ([False, False], False),
21 + ],
22 +)
23 +async def test_dp_wake_up_requires_all_engines_awake(results, expected):
24 + client = object.__new__(DPLBAsyncMPClient)
25 + client.core_engines = [b"engine-0", b"engine-1"]
26 + client._call_utility_async = AsyncMock(side_effect=results)
27 +
28 + assert await client.wake_up_async(["weights"]) is expected
29 + assert client._call_utility_async.await_args_list == [
30 + call("wake_up", ["weights"], engine=b"engine-0"),
31 + call("wake_up", ["weights"], engine=b"engine-1"),
32 + ]
vllm/engine/protocol.py
+2 -2
@@ -199,8 +199,8 @@ class EngineClient(ABC):
199 ...
200
201 @abstractmethod
202 - async def wake_up(self, tags: list[str] | None = None) -> None:
203 - """Wake up the engine."""
202 + async def wake_up(self, tags: list[str] | None = None) -> bool:
203 + """Wake the tagged resources; return whether every engine is fully awake."""
204 ...
205
206 @abstractmethod
vllm/entrypoints/llm.py
+5 -2
@@ -839,7 +839,7 @@ class LLM(BeamSearchOfflineMixin, PoolingOfflineMixin, OfflineInferenceMixin):
839 """
840 self.llm_engine.release_kv_cache_memory()
841
842 - def wake_up(self, tags: list[str] | None = None):
842 + def wake_up(self, tags: list[str] | None = None) -> bool:
843 """Wake up the engine from sleep mode. See the [sleep][vllm.LLM.sleep]
844 method for more details.
845
@@ -851,8 +851,11 @@ class LLM(BeamSearchOfflineMixin, PoolingOfflineMixin, OfflineInferenceMixin):
851 (or None) before the engine is used again.
852 Use tags=["scheduling"] to resume from level 0 sleep.
853
854 + Returns:
855 + Whether the engine is fully awake.
856 +
857 """
855 - self.llm_engine.wake_up(tags)
858 + return self.llm_engine.wake_up(tags)
859
860 def get_metrics(self) -> list["Metric"]:
861 """Return a snapshot of aggregated metrics from Prometheus.
vllm/entrypoints/serve/dev/sleep/api_router.py
+25 -14
@@ -2,11 +2,16 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
5 -from fastapi import APIRouter, FastAPI, Request
6 -from fastapi.responses import JSONResponse, Response
5 +from typing import Annotated
6 +
7 +from fastapi import APIRouter, FastAPI, Query, Request
8 +from fastapi.responses import JSONResponse
9
10 from vllm.engine.protocol import EngineClient
11 from vllm.logger import init_logger
12 +from vllm.v1.engine import PauseMode
13 +
14 +from .metrics import sleep_mode_operation_metrics
15
16 logger = init_logger(__name__)
17
@@ -19,29 +24,35 @@ router = APIRouter()
24
25
26 @router.post("/sleep")
22 -async def sleep(raw_request: Request):
23 - # get POST params
24 - level = raw_request.query_params.get("level", "1")
25 - mode = raw_request.query_params.get("mode", "abort")
26 - await engine_client(raw_request).sleep(int(level), mode)
27 - return Response(status_code=200)
27 +async def sleep(
28 + raw_request: Request,
29 + level: Annotated[int, Query(ge=0, le=2)] = 1,
30 + mode: Annotated[PauseMode, Query()] = "abort",
31 +) -> JSONResponse:
32 + with sleep_mode_operation_metrics().record("sleep"):
33 + await engine_client(raw_request).sleep(level, mode)
34 + return JSONResponse(content={"status": "sleeping", "level": level})
35
36
37 @router.post("/release_kv_cache_memory")
31 -async def release_kv_cache_memory(raw_request: Request):
32 - await engine_client(raw_request).release_kv_cache_memory()
33 - return Response(status_code=200)
38 +async def release_kv_cache_memory(raw_request: Request) -> JSONResponse:
39 + with sleep_mode_operation_metrics().record("release_kv_cache_memory"):
40 + await engine_client(raw_request).release_kv_cache_memory()
41 + return JSONResponse(content={"status": "kv_cache_released"})
42
43
44 @router.post("/wake_up")
37 -async def wake_up(raw_request: Request):
45 +async def wake_up(raw_request: Request) -> JSONResponse:
46 tags = raw_request.query_params.getlist("tags")
47 if tags == []:
48 # set to None to wake up all tags if no tags are provided
49 tags = None
50 logger.info("wake up the engine with tags: %s", tags)
43 - await engine_client(raw_request).wake_up(tags)
44 - return Response(status_code=200)
51 + with sleep_mode_operation_metrics().record("wake"):
52 + fully_awake = await engine_client(raw_request).wake_up(tags)
53 + return JSONResponse(
54 + content={"status": "awake" if fully_awake else "sleeping", "tags": tags}
55 + )
56
57
58 @router.get("/is_sleeping")
vllm/entrypoints/serve/dev/sleep/metrics.py new
+57
@@ -0,0 +1,57 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Metrics for dispatched sleep-mode frontend operations."""
4 +
5 +from collections.abc import Iterator
6 +from contextlib import contextmanager
7 +from threading import Lock
8 +from time import perf_counter
9 +from typing import Literal
10 +
11 +from prometheus_client import REGISTRY, CollectorRegistry, Gauge, Histogram
12 +
13 +Operation = Literal["sleep", "release_kv_cache_memory", "wake"]
14 +BUCKETS = (0.01, 0.1, 1, 10, 30, 60, 120, 300, 600)
15 +
16 +
17 +class SleepModeOperationMetrics:
18 + def __init__(self, registry: CollectorRegistry | None = None):
19 + registry = REGISTRY if registry is None else registry
20 + self.duration = Histogram(
21 + "vllm:sleep_mode_operation_duration_seconds",
22 + "Duration of one sleep-mode operation.",
23 + ["operation"],
24 + buckets=BUCKETS,
25 + registry=registry,
26 + )
27 + self.in_flight = Gauge(
28 + "vllm:sleep_mode_operations_in_flight",
29 + "Sleep-mode operations currently awaited.",
30 + ["operation"],
31 + multiprocess_mode="livesum",
32 + registry=registry,
33 + )
34 +
35 + @contextmanager
36 + def record(self, operation: Operation) -> Iterator[None]:
37 + active = self.in_flight.labels(operation)
38 + started = perf_counter()
39 + active.inc()
40 + try:
41 + yield
42 + finally:
43 + active.dec()
44 + self.duration.labels(operation).observe(perf_counter() - started)
45 +
46 +
47 +_metrics: SleepModeOperationMetrics | None = None
48 +_metrics_lock = Lock()
49 +
50 +
51 +def sleep_mode_operation_metrics() -> SleepModeOperationMetrics:
52 + global _metrics
53 + if _metrics is None:
54 + with _metrics_lock:
55 + if _metrics is None:
56 + _metrics = SleepModeOperationMetrics()
57 + return _metrics
vllm/v1/engine/async_llm.py
+3 -1
@@ -1110,12 +1110,14 @@ class AsyncLLM(EngineClient):
1110 if self.logger_manager is not None:
1111 self.logger_manager.record_sleep_state(1, 0)
1112
1113 - async def wake_up(self, tags: list[str] | None = None) -> None:
1113 + async def wake_up(self, tags: list[str] | None = None) -> bool:
1114 fully_awake = await self.engine_core.wake_up_async(tags)
1115
1116 if self.logger_manager is not None and fully_awake:
1117 self.logger_manager.record_sleep_state(0, 0)
1118
1119 + return fully_awake
1120 +
1121 async def checkpoint_prepare(self) -> None:
1122 await self.collective_rpc("checkpoint_prepare")
1123
vllm/v1/engine/core_client.py
+1 -1
@@ -1320,7 +1320,7 @@ class AsyncMPClient(MPClient):
1320 await self.call_utility_async("release_kv_cache_memory")
1321
1322 async def wake_up_async(self, tags: list[str] | None = None) -> bool:
1323 - return await self.call_utility_async("wake_up", tags)
1323 + return all(await self.call_utility_all_async("wake_up", tags))
1324
1325 async def is_sleeping_async(self) -> bool:
1326 return await self.call_utility_consensus_async("is_sleeping")
vllm/v1/engine/llm_engine.py
+3 -1
@@ -386,12 +386,14 @@ class LLMEngine:
386 if self.logger_manager is not None:
387 self.logger_manager.record_sleep_state(1, 0)
388
389 - def wake_up(self, tags: list[str] | None = None):
389 + def wake_up(self, tags: list[str] | None = None) -> bool:
390 fully_awake = self.engine_core.wake_up(tags)
391
392 if self.logger_manager is not None and fully_awake:
393 self.logger_manager.record_sleep_state(0, 0)
394
395 + return fully_awake
396 +
397 def is_sleeping(self) -> bool:
398 return self.engine_core.is_sleeping()
399