@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 7632f767f0

[Rust Frontend] Support HTTP RL weight synchronization (#56567)

Co-authored-by: Bugen Zhao <i@bugenzhao.com> Co-authored-by: Codex <noreply@openai.com> Signed-off-by: Wang Zupeng <zupenwang@gmail.com> Signed-off-by: Bugen Zhao <i@bugenzhao.com>

Zupeng Wang committed Sep 14, 2026 at 11:37 UTC 7632f767f0d896a902566acab29775816b39a92d
8 files changed +483 -6
rust/README.md
+11
@@ -43,6 +43,17 @@ For example:
43 VLLM_USE_RUST_FRONTEND=1 vllm serve Qwen/Qwen3-0.6B
44 ```
45
46 +### RL weight synchronization
47 +
48 +With `VLLM_SERVER_DEV_MODE=1`, the Rust frontend supports the HTTP weight-transfer
49 +lifecycle used by `HTTPVLLMWeightSyncClient`: initialization, starting an update,
50 +transferring weights, and finishing the update. It also supports draft-model
51 +updates and `/update_weight_version` and `/weight_info` for version tracking.
52 +
53 +For trainer-side usage, see the [IPC](../examples/rl/rlhf_http_ipc.py) and
54 +[NCCL](../examples/rl/rlhf_http_nccl.py) examples. Configure the weight-transfer
55 +backend on the server and pause generation during updates, as shown there.
56 +
57 ### External Engine
58
59 `vllm-rs serve` can be run standalone with `--data-parallel-size-local 0` when the Python engines
rust/src/server/src/error.rs
+7
@@ -2,6 +2,7 @@
2 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 use axum::Json;
5 +use axum::extract::rejection::JsonRejection;
6 use axum::http::StatusCode;
7 use axum::response::{IntoResponse, Response};
8 use thiserror_ext::{AsReport as _, Construct, Macro};
@@ -70,6 +71,12 @@ impl ApiError {
71 }
72 }
73
74 +impl From<JsonRejection> for ApiError {
75 + fn from(error: JsonRejection) -> Self {
76 + Self::json_parse_error(error.body_text())
77 + }
78 +}
79 +
80 impl IntoResponse for ApiError {
81 fn into_response(self) -> Response {
82 (self.status_code(), Json(self.to_error_response())).into_response()
rust/src/server/src/routes.rs
+23
@@ -17,6 +17,7 @@ mod server_info;
17 mod sleep;
18 mod tokenize;
19 mod version;
20 +mod weight_transfer;
21 mod world_size;
22
23 use std::sync::Arc;
@@ -139,6 +140,28 @@ fn build_router_with_options(
140 .route("/reset_mm_cache", post(cache::reset_mm_cache))
141 .route("/reset_encoder_cache", post(cache::reset_encoder_cache))
142 .route("/collective_rpc", post(collective_rpc::collective_rpc))
143 + .route(
144 + "/init_weight_transfer_engine",
145 + post(weight_transfer::init_weight_transfer_engine),
146 + )
147 + .route(
148 + "/start_weight_update",
149 + post(weight_transfer::start_weight_update),
150 + )
151 + .route(
152 + "/start_draft_weight_update",
153 + post(weight_transfer::start_draft_weight_update),
154 + )
155 + .route("/update_weights", post(weight_transfer::update_weights))
156 + .route(
157 + "/finish_weight_update",
158 + post(weight_transfer::finish_weight_update),
159 + )
160 + .route(
161 + "/update_weight_version",
162 + post(weight_transfer::update_weight_version),
163 + )
164 + .route("/weight_info", get(weight_transfer::weight_info))
165 .route("/abort_requests", post(abort_requests::abort_requests))
166 .route("/sleep", post(sleep::sleep))
167 .route("/wake_up", post(sleep::wake_up))
rust/src/server/src/routes/abort_requests.rs
+1 -1
@@ -22,7 +22,7 @@ pub async fn abort_requests(
22 State(state): State<Arc<AppState>>,
23 body: Result<Json<AbortRequestsRequest>, JsonRejection>,
24 ) -> Result<StatusCode, ApiError> {
25 - let Json(body) = body.map_err(|error| ApiError::json_parse_error(error.body_text()))?;
25 + let Json(body) = body?;
26 // Empty/missing `request_ids` aborts all in-flight requests.
27 let request_ids = body.request_ids.unwrap_or_default();
28
rust/src/server/src/routes/collective_rpc.rs
+1 -1
@@ -36,7 +36,7 @@ pub async fn collective_rpc(
36 State(state): State<Arc<AppState>>,
37 body: Result<Json<CollectiveRpcRequest>, JsonRejection>,
38 ) -> Result<Json<CollectiveRpcResponse>, ApiError> {
39 - let Json(body) = body.map_err(|error| ApiError::json_parse_error(error.body_text()))?;
39 + let Json(body) = body?;
40 let method = body.method.ok_or_else(|| {
41 ApiError::invalid_request(
42 "Missing 'method' in request body".to_string(),
rust/src/server/src/routes/openai/utils/validated_json.rs
+1 -4
@@ -5,7 +5,6 @@
5 //! Variation of https://github.com/lightseekorg/smg/blob/main/crates/protocols/src/validated.rs
6
7 use axum::Json;
8 -use axum::extract::rejection::JsonRejection;
8 use axum::extract::{FromRequest, Request};
9 use serde::de::DeserializeOwned;
10 use validator::Validate;
@@ -30,9 +29,7 @@ where
29 type Rejection = ApiError;
30
31 async fn from_request(req: Request, state: &S) -> Result<Self, Self::Rejection> {
33 - let Json(mut data) = Json::<T>::from_request(req, state)
34 - .await
35 - .map_err(|err: JsonRejection| ApiError::json_parse_error(err.body_text()))?;
32 + let Json(mut data) = Json::<T>::from_request(req, state).await?;
33
34 data.normalize();
35
rust/src/server/src/routes/tests.rs
+249
@@ -6113,6 +6113,13 @@ async fn admin_routes_are_hidden_when_dev_mode_is_disabled() {
6113 ("POST", "/reset_prefix_cache"),
6114 ("POST", "/reset_mm_cache"),
6115 ("POST", "/reset_encoder_cache"),
6116 + ("POST", "/init_weight_transfer_engine"),
6117 + ("POST", "/start_weight_update"),
6118 + ("POST", "/start_draft_weight_update"),
6119 + ("POST", "/update_weights"),
6120 + ("POST", "/finish_weight_update"),
6121 + ("POST", "/update_weight_version"),
6122 + ("GET", "/weight_info"),
6123 ] {
6124 let response = app
6125 .clone()
@@ -6132,6 +6139,248 @@ async fn admin_routes_are_hidden_when_dev_mode_is_disabled() {
6139 engine_task.abort_and_join().await;
6140 }
6141
6142 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
6143 +#[serial]
6144 +async fn weight_transfer_routes_support_the_http_training_lifecycle() {
6145 + let (mut app, engine_task) = test_admin_app_with_engine_script(|dealer, push| {
6146 + boxed_test_future(async move {
6147 + let mut calls = Vec::new();
6148 + for _ in 0..10 {
6149 + let utility = recv_engine_message(dealer).await;
6150 + assert_eq!(utility[0].as_ref(), &[0x03]);
6151 + let payload = decode_value(&utility[1]).expect("decode utility");
6152 + let array = payload.as_array().expect("utility array");
6153 + let call_id = array[1].as_u64().expect("call id");
6154 + let args: serde_json::Value =
6155 + rmpv::ext::from_value(array[3].clone()).expect("decode args");
6156 + calls.push(format!("{} {args}", array[2].as_str().expect("method")));
6157 + let result = if array[2] == Value::from("get_weight_version") {
6158 + utility_result_value("step-2")
6159 + } else if array[2] == Value::from("collective_rpc") {
6160 + utility_result_value(vec![(), ()])
6161 + } else {
6162 + utility_none_result()
6163 + };
6164 + send_outputs(push, utility_outputs(call_id, result)).await;
6165 + }
6166 + expect_test::expect![[r#"
6167 + [
6168 + "collective_rpc [\"init_weight_transfer_engine\",null,[],{\"init_info\":{\"master_address\":\"127.0.0.1\",\"master_port\":29500,\"rank_offset\":1,\"world_size\":3}}]",
6169 + "collective_rpc [\"start_weight_update\",null,[],{}]",
6170 + "collective_rpc [\"update_weights\",null,[],{\"update_info\":{\"names\":[\"model.embed_tokens.weight\"],\"dtype_names\":[\"bfloat16\"],\"shapes\":[[32,16]],\"is_checkpoint_format\":true}}]",
6171 + "collective_rpc [\"update_weights\",null,[],{\"update_info\":[{\"rank\":0,\"ipc_handles_pickled\":\"opaque-handle-0\"},{\"rank\":1,\"ipc_handles_pickled\":\"opaque-handle-1\"}]}]",
6172 + "collective_rpc [\"finish_weight_update\",null,[],{}]",
6173 + "set_weight_version [\"step-1\"]",
6174 + "collective_rpc [\"start_draft_weight_update\",null,[],{}]",
6175 + "collective_rpc [\"finish_weight_update\",null,[],{}]",
6176 + "set_weight_version [\"step-2\"]",
6177 + "get_weight_version []",
6178 + ]
6179 + "#]].assert_debug_eq(&calls);
6180 + })
6181 + })
6182 + .await;
6183 +
6184 + let mut responses = Vec::new();
6185 + for (method, path, body) in [
6186 + (
6187 + "POST",
6188 + "/init_weight_transfer_engine",
6189 + Some(json!({
6190 + "init_info": {"master_address": "127.0.0.1", "master_port": 29500,
6191 + "rank_offset": 1, "world_size": 3}
6192 + })),
6193 + ),
6194 + ("POST", "/start_weight_update", None),
6195 + (
6196 + "POST",
6197 + "/update_weights",
6198 + Some(json!({
6199 + "update_info": {"names": ["model.embed_tokens.weight"],
6200 + "dtype_names": ["bfloat16"], "shapes": [[32, 16]],
6201 + "is_checkpoint_format": true}
6202 + })),
6203 + ),
6204 + (
6205 + "POST",
6206 + "/update_weights",
6207 + Some(json!({
6208 + "update_info": [{"rank": 0, "ipc_handles_pickled": "opaque-handle-0"},
6209 + {"rank": 1, "ipc_handles_pickled": "opaque-handle-1"}]
6210 + })),
6211 + ),
6212 + (
6213 + "POST",
6214 + "/finish_weight_update",
6215 + Some(json!({"weight_version": "step-1"})),
6216 + ),
6217 + ("POST", "/start_draft_weight_update", None),
6218 + ("POST", "/finish_weight_update", None),
6219 + (
6220 + "POST",
6221 + "/update_weight_version",
6222 + Some(json!({"new_version": "step-2"})),
6223 + ),
6224 + ("GET", "/weight_info", None),
6225 + ] {
6226 + let mut request = Request::builder().method(method).uri(path);
6227 + let body = match body {
6228 + Some(body) => {
6229 + request = request.header("content-type", "application/json");
6230 + Body::from(body.to_string())
6231 + }
6232 + None => Body::empty(),
6233 + };
6234 + let response =
6235 + app.call(request.body(body).expect("build request")).await.expect("call app");
6236 + assert_eq!(response.status(), StatusCode::OK, "{path}");
6237 + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
6238 + responses.push(
6239 + serde_json::from_slice::<serde_json::Value>(&body)
6240 + .expect("decode json")
6241 + .to_string(),
6242 + );
6243 + }
6244 + engine_task.await.expect("mock engine task");
6245 + expect_test::expect![[r#"
6246 + [
6247 + "{\"message\":\"Weight transfer initialized\"}",
6248 + "{\"message\":\"Weight update started\"}",
6249 + "{\"message\":\"Weights updated\"}",
6250 + "{\"message\":\"Weights updated\"}",
6251 + "{\"message\":\"Weight update finished\"}",
6252 + "{\"message\":\"Draft weight update started\"}",
6253 + "{\"message\":\"Weight update finished\"}",
6254 + "{\"success\":true,\"new_version\":\"step-2\"}",
6255 + "{\"weight_version\":\"step-2\"}",
6256 + ]
6257 + "#]]
6258 + .assert_debug_eq(&responses);
6259 +}
6260 +
6261 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
6262 +#[serial]
6263 +async fn weight_transfer_routes_reject_invalid_payloads_before_engine_calls() {
6264 + let (mut app, engine_task) = test_admin_app_with_engine_script(|dealer, _push| {
6265 + boxed_test_future(async move {
6266 + let message = recv_engine_message(dealer).await;
6267 + panic!("invalid payload reached engine: {message:?}");
6268 + })
6269 + })
6270 + .await;
6271 +
6272 + for (path, body) in [
6273 + ("/init_weight_transfer_engine", "{"),
6274 + ("/init_weight_transfer_engine", "{}"),
6275 + ("/init_weight_transfer_engine", r#"{"init_info":null}"#),
6276 + ("/init_weight_transfer_engine", r#"{"init_info":[]}"#),
6277 + ("/update_weights", "{"),
6278 + ("/update_weights", "{}"),
6279 + ("/update_weights", r#"{"update_info":null}"#),
6280 + ("/update_weights", r#"{"update_info":1}"#),
6281 + ("/update_weights", r#"{"update_info":[{},null]}"#),
6282 + ("/finish_weight_update", "{"),
6283 + ("/finish_weight_update", r#"{"weight_version":1}"#),
6284 + ("/update_weight_version", "{}"),
6285 + ("/update_weight_version", r#"{"new_version":null}"#),
6286 + ] {
6287 + let response = app
6288 + .call(
6289 + Request::builder()
6290 + .method("POST")
6291 + .uri(path)
6292 + .header("content-type", "application/json")
6293 + .body(Body::from(body))
6294 + .expect("build request"),
6295 + )
6296 + .await
6297 + .expect("call app");
6298 + assert_eq!(response.status(), StatusCode::BAD_REQUEST, "{path}: {body}");
6299 + }
6300 + engine_task.abort_and_join().await;
6301 +}
6302 +
6303 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
6304 +#[serial]
6305 +async fn weight_transfer_routes_do_not_commit_version_when_finish_fails() {
6306 + let (mut app, engine_task) = test_admin_app_with_engine_script(|dealer, push| {
6307 + boxed_test_future(async move {
6308 + let utility = recv_engine_message(dealer).await;
6309 + let payload = decode_value(&utility[1]).expect("decode utility");
6310 + let array = payload.as_array().expect("utility array");
6311 + assert_eq!(array[2], Value::from("collective_rpc"));
6312 + assert_eq!(
6313 + array[3].as_array().expect("args")[0],
6314 + Value::from("finish_weight_update")
6315 + );
6316 + send_outputs(
6317 + push,
6318 + UtilityCallOutput {
6319 + output: UtilityOutput {
6320 + call_id: array[1].as_u64().expect("call id").into(),
6321 + failure_message: Some("weight reload failed".to_string()),
6322 + result: None,
6323 + },
6324 + ..Default::default()
6325 + }
6326 + .into(),
6327 + )
6328 + .await;
6329 + let utility = recv_engine_message(dealer).await;
6330 + let payload = decode_value(&utility[1]).expect("decode utility");
6331 + let array = payload.as_array().expect("utility array");
6332 + assert_eq!(array[2], Value::from("get_weight_version"));
6333 + send_outputs(
6334 + push,
6335 + utility_outputs(
6336 + array[1].as_u64().expect("call id"),
6337 + utility_result_value("step-0"),
6338 + ),
6339 + )
6340 + .await;
6341 + })
6342 + })
6343 + .await;
6344 +
6345 + let response = app
6346 + .call(
6347 + Request::builder()
6348 + .method("POST")
6349 + .uri("/finish_weight_update")
6350 + .header("content-type", "application/json")
6351 + .body(Body::from(r#"{"weight_version":"step-1"}"#))
6352 + .expect("build request"),
6353 + )
6354 + .await
6355 + .expect("call app");
6356 + assert_eq!(response.status(), StatusCode::INTERNAL_SERVER_ERROR);
6357 + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
6358 + let error: serde_json::Value = serde_json::from_slice(&body).expect("decode json");
6359 + assert!(
6360 + error["error"]["message"]
6361 + .as_str()
6362 + .expect("error message")
6363 + .contains("weight reload failed")
6364 + );
6365 +
6366 + let response = app
6367 + .call(
6368 + Request::builder()
6369 + .uri("/weight_info")
6370 + .body(Body::empty())
6371 + .expect("build request"),
6372 + )
6373 + .await
6374 + .expect("call app");
6375 + assert_eq!(response.status(), StatusCode::OK);
6376 + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
6377 + assert_eq!(
6378 + serde_json::from_slice::<serde_json::Value>(&body).expect("decode json"),
6379 + json!({"weight_version": "step-0"})
6380 + );
6381 + engine_task.await.expect("mock engine task");
6382 +}
6383 +
6384 // ========================= Stop string tests =========================
6385
6386 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
rust/src/server/src/routes/weight_transfer.rs new
+190
@@ -0,0 +1,190 @@
1 +// SPDX-License-Identifier: Apache-2.0
2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +use std::sync::Arc;
5 +
6 +use axum::Json;
7 +use axum::extract::State;
8 +use axum::extract::rejection::JsonRejection;
9 +use serde::{Deserialize, Serialize};
10 +use serde_json::Value as JsonValue;
11 +
12 +use crate::error::{ApiError, invalid_request};
13 +use crate::state::AppState;
14 +use crate::utils::utility_call_error;
15 +
16 +#[derive(Deserialize)]
17 +pub(crate) struct InitWeightTransferRequest {
18 + init_info: Option<JsonValue>,
19 +}
20 +
21 +#[derive(Deserialize)]
22 +pub(crate) struct UpdateWeightsRequest {
23 + update_info: Option<JsonValue>,
24 +}
25 +
26 +#[derive(Default, Deserialize)]
27 +pub(crate) struct FinishWeightUpdateRequest {
28 + weight_version: Option<String>,
29 +}
30 +
31 +#[derive(Deserialize)]
32 +pub(crate) struct UpdateWeightVersionRequest {
33 + new_version: String,
34 +}
35 +
36 +#[derive(Serialize)]
37 +pub(crate) struct MessageResponse {
38 + message: &'static str,
39 +}
40 +
41 +#[derive(Serialize)]
42 +pub(crate) struct UpdateWeightVersionResponse {
43 + success: bool,
44 + new_version: String,
45 +}
46 +
47 +#[derive(Serialize)]
48 +pub(crate) struct WeightInfoResponse {
49 + weight_version: String,
50 +}
51 +
52 +/// Initialize weight transfer for RL training.
53 +pub async fn init_weight_transfer_engine(
54 + State(state): State<Arc<AppState>>,
55 + body: Result<Json<InitWeightTransferRequest>, JsonRejection>,
56 +) -> Result<Json<MessageResponse>, ApiError> {
57 + let Json(body) = body?;
58 + let init_info = body.init_info.filter(JsonValue::is_object).ok_or_else(|| {
59 + invalid_request!(
60 + param = Some("init_info"),
61 + "'init_info' must be a JSON object"
62 + )
63 + })?;
64 +
65 + state
66 + .engine_core_client()
67 + .init_weight_transfer_engine(init_info)
68 + .await
69 + .map_err(|error| utility_call_error("init_weight_transfer_engine", error))?;
70 +
71 + Ok(Json(MessageResponse {
72 + message: "Weight transfer initialized",
73 + }))
74 +}
75 +
76 +/// Start a weight update transaction.
77 +pub async fn start_weight_update(
78 + State(state): State<Arc<AppState>>,
79 +) -> Result<Json<MessageResponse>, ApiError> {
80 + state
81 + .engine_core_client()
82 + .start_weight_update()
83 + .await
84 + .map_err(|error| utility_call_error("start_weight_update", error))?;
85 +
86 + Ok(Json(MessageResponse {
87 + message: "Weight update started",
88 + }))
89 +}
90 +
91 +/// Start a draft-model weight update transaction.
92 +pub async fn start_draft_weight_update(
93 + State(state): State<Arc<AppState>>,
94 +) -> Result<Json<MessageResponse>, ApiError> {
95 + state
96 + .engine_core_client()
97 + .start_draft_weight_update()
98 + .await
99 + .map_err(|error| utility_call_error("start_draft_weight_update", error))?;
100 +
101 + Ok(Json(MessageResponse {
102 + message: "Draft weight update started",
103 + }))
104 +}
105 +
106 +/// Update model weights with backend-specific transfer metadata.
107 +pub async fn update_weights(
108 + State(state): State<Arc<AppState>>,
109 + body: Result<Json<UpdateWeightsRequest>, JsonRejection>,
110 +) -> Result<Json<MessageResponse>, ApiError> {
111 + let Json(body) = body?;
112 + let update_info = body
113 + .update_info
114 + .filter(|info| {
115 + info.is_object()
116 + || info.as_array().is_some_and(|items| items.iter().all(JsonValue::is_object))
117 + })
118 + .ok_or_else(|| {
119 + invalid_request!(
120 + param = Some("update_info"),
121 + "'update_info' must be a JSON object or a list of per-worker JSON objects",
122 + )
123 + })?;
124 +
125 + state
126 + .engine_core_client()
127 + .update_weights(update_info)
128 + .await
129 + .map_err(|error| utility_call_error("update_weights", error))?;
130 +
131 + Ok(Json(MessageResponse {
132 + message: "Weights updated",
133 + }))
134 +}
135 +
136 +/// Finish a weight update transaction and optionally set the weight version.
137 +pub async fn finish_weight_update(
138 + State(state): State<Arc<AppState>>,
139 + body: Result<Option<Json<FinishWeightUpdateRequest>>, JsonRejection>,
140 +) -> Result<Json<MessageResponse>, ApiError> {
141 + // HTTPVLLMWeightSyncClient omits the body when no version is supplied.
142 + let request = body?.map(|Json(request)| request).unwrap_or_default();
143 +
144 + let client = state.engine_core_client();
145 + client
146 + .finish_weight_update()
147 + .await
148 + .map_err(|error| utility_call_error("finish_weight_update", error))?;
149 + if let Some(version) = request.weight_version {
150 + client
151 + .set_weight_version(&version)
152 + .await
153 + .map_err(|error| utility_call_error("set_weight_version", error))?;
154 + }
155 +
156 + Ok(Json(MessageResponse {
157 + message: "Weight update finished",
158 + }))
159 +}
160 +
161 +/// Update the weight version identifier.
162 +pub async fn update_weight_version(
163 + State(state): State<Arc<AppState>>,
164 + body: Result<Json<UpdateWeightVersionRequest>, JsonRejection>,
165 +) -> Result<Json<UpdateWeightVersionResponse>, ApiError> {
166 + let Json(body) = body?;
167 + state
168 + .engine_core_client()
169 + .set_weight_version(&body.new_version)
170 + .await
171 + .map_err(|error| utility_call_error("set_weight_version", error))?;
172 +
173 + Ok(Json(UpdateWeightVersionResponse {
174 + success: true,
175 + new_version: body.new_version,
176 + }))
177 +}
178 +
179 +/// Get current weight metadata.
180 +pub async fn weight_info(
181 + State(state): State<Arc<AppState>>,
182 +) -> Result<Json<WeightInfoResponse>, ApiError> {
183 + let weight_version = state
184 + .engine_core_client()
185 + .get_weight_version()
186 + .await
187 + .map_err(|error| utility_call_error("get_weight_version", error))?;
188 +
189 + Ok(Json(WeightInfoResponse { weight_version }))
190 +}