[Rust Frontend][gRPC] Add RL lifecycle control (#51316)
Signed-off-by: Connor Carpenter <connorc@nvidia.com> Co-authored-by: Nick Hill <nickhill123@gmail.com>
Connor Carpenter committed
Aug 14, 2026 at 11:32 UTC
f473870ecf7a5a83958dbc116468663224f45b05
17 files changed
+614
-16
docs/training/weight_transfer/README.md
+2
@@ -158,6 +158,8 @@ When running vLLM as an HTTP server, the following endpoints are available for w
158
!!! note
159
The HTTP weight transfer endpoints require `VLLM_SERVER_DEV_MODE=1` to be set.
160
161
+The Rust frontend's optional gRPC `Control` service exposes the same pause, sleep, weight-transfer, and weight-version lifecycle for trusted sidecars. The `ServerInfo.rl_capabilities` response reports whether weight transfer and sleep mode were configured. Backend-specific `init_info` and `update_info` remain JSON metadata; model tensors continue to move over the configured NCCL, IPC, or sparse-NCCL transport.
162
+
163
## Extending the System
164
165
Every piece of the system is replaceable: the weights you send (`WeightSource`),
docs/usage/security.md
+4
-3
@@ -345,7 +345,7 @@ vLLM supports loading out-of-tree HTTP routes via the `vllm.endpoint_plugins` en
345
346
## gRPC Interface
347
348
-vLLM provides an optional gRPC Generate service on a separate TCP port, enabled via the `--grpc-port` flag. When not specified, no gRPC server is started. The gRPC listener binds to the same host address as the HTTP server.
348
+vLLM provides optional gRPC `Inference` and `Control` services on a separate TCP port, enabled via the `--grpc-port` flag. When not specified, no gRPC server is started. The gRPC listener binds to the same host address as the HTTP server.
349
350
**Warning:** The gRPC interface is **insecure by default** — it does not implement authentication, authorization, or encryption. It should be considered a private, internal interface intended for use only between co-located services within a trusted network. Do not expose the gRPC port to the public internet or untrusted clients. If you enable the gRPC interface, protect it via network-level access controls such as firewall rules, network segmentation, or deployment on an isolated private network.
351
@@ -354,8 +354,9 @@ vLLM provides an optional gRPC Generate service on a separate TCP port, enabled
354
An attacker who can reach the gRPC port can:
355
356
1. **Run arbitrary inference** via the `Generate` and `GenerateStream` RPCs without any credentials
357
-2. **Consume GPU and compute resources** by submitting unbounded generation requests
358
-3. **Cause Denial of Service** by exploiting bugs in the gRPC interface that can crash vLLM.
357
+2. **Mutate engine state** by pausing generation, sleeping the engine, or initiating configured RL weight updates through the `Control` service
358
+3. **Consume GPU and compute resources** by submitting unbounded generation requests
359
+4. **Cause Denial of Service** by exploiting bugs in the gRPC interface that can crash vLLM.
360
361
### Recommendations
362
rust/proto/control.proto
+81
@@ -9,6 +9,21 @@ service Control {
9
rpc GetModelInfo (GetModelInfoRequest) returns (ModelInfo) {}
10
rpc Abort (AbortRequest) returns (AbortResponse) {}
11
rpc GetKvEventSources (GetKvEventSourcesRequest) returns (GetKvEventSourcesResponse) {}
12
+
13
+ // Reinforcement-learning lifecycle and weight updates.
14
+ rpc PauseGeneration (PauseGenerationRequest) returns (PauseGenerationResponse) {}
15
+ rpc ResumeGeneration (ResumeGenerationRequest) returns (ResumeGenerationResponse) {}
16
+ rpc IsPaused (IsPausedRequest) returns (IsPausedResponse) {}
17
+ rpc Sleep (SleepRequest) returns (SleepResponse) {}
18
+ rpc WakeUp (WakeUpRequest) returns (WakeUpResponse) {}
19
+ rpc IsSleeping (IsSleepingRequest) returns (IsSleepingResponse) {}
20
+ rpc InitWeightTransferEngine (InitWeightTransferEngineRequest) returns (InitWeightTransferEngineResponse) {}
21
+ rpc StartWeightUpdate (StartWeightUpdateRequest) returns (StartWeightUpdateResponse) {}
22
+ rpc StartDraftWeightUpdate (StartDraftWeightUpdateRequest) returns (StartDraftWeightUpdateResponse) {}
23
+ rpc UpdateWeights (UpdateWeightsRequest) returns (UpdateWeightsResponse) {}
24
+ rpc FinishWeightUpdate (FinishWeightUpdateRequest) returns (FinishWeightUpdateResponse) {}
25
+ rpc UpdateWeightVersion (UpdateWeightVersionRequest) returns (UpdateWeightVersionResponse) {}
26
+ rpc GetWeightVersion (GetWeightVersionRequest) returns (GetWeightVersionResponse) {}
27
}
28
29
message GetServerInfoRequest {}
@@ -23,6 +38,14 @@ message ServerInfo {
38
uint64 total_kv_blocks = 7;
39
uint64 max_running_requests = 8;
40
uint64 max_batched_tokens = 9;
41
+ RlCapabilities rl_capabilities = 11;
42
+}
43
+
44
+message RlCapabilities {
45
+ bool weight_transfer_enabled = 1;
46
+ string weight_transfer_backend = 2;
47
+ bool sleep_mode_enabled = 3;
48
+ bool draft_weight_updates_enabled = 4;
49
}
50
51
message ParallelismInfo {
@@ -53,6 +76,64 @@ message AbortRequest {
76
77
message AbortResponse {}
78
79
+// ======================================================================================
80
+// Reinforcement-learning control
81
+// ======================================================================================
82
+
83
+enum PauseMode {
84
+ PAUSE_MODE_UNSPECIFIED = 0;
85
+ PAUSE_MODE_ABORT = 1;
86
+ PAUSE_MODE_WAIT = 2;
87
+ PAUSE_MODE_KEEP = 3;
88
+}
89
+
90
+message PauseGenerationRequest {
91
+ PauseMode mode = 1;
92
+ optional bool clear_cache = 2;
93
+}
94
+message PauseGenerationResponse {}
95
+
96
+message ResumeGenerationRequest {}
97
+message ResumeGenerationResponse {}
98
+
99
+message IsPausedRequest {}
100
+message IsPausedResponse { bool paused = 1; }
101
+
102
+message SleepRequest {
103
+ optional uint32 level = 1;
104
+ PauseMode mode = 2;
105
+}
106
+message SleepResponse {}
107
+
108
+message WakeUpRequest { repeated string tags = 1; }
109
+message WakeUpResponse {}
110
+
111
+message IsSleepingRequest {}
112
+message IsSleepingResponse { bool sleeping = 1; }
113
+
114
+// The payloads are backend-specific JSON objects. Tensor data remains on the
115
+// configured NCCL, IPC, or sparse-NCCL transport rather than crossing gRPC.
116
+message InitWeightTransferEngineRequest { bytes init_info_json = 1; }
117
+message InitWeightTransferEngineResponse {}
118
+
119
+message StartWeightUpdateRequest {}
120
+message StartWeightUpdateResponse {}
121
+
122
+message StartDraftWeightUpdateRequest {}
123
+message StartDraftWeightUpdateResponse {}
124
+
125
+message UpdateWeightsRequest { bytes update_info_json = 1; }
126
+message UpdateWeightsResponse {}
127
+
128
+message FinishWeightUpdateRequest { optional string weight_version = 1; }
129
+message FinishWeightUpdateResponse {}
130
+
131
+message UpdateWeightVersionRequest { string weight_version = 1; }
132
+message UpdateWeightVersionResponse {}
133
+
134
+message GetWeightVersionRequest {}
135
+message GetWeightVersionResponse { string weight_version = 1; }
136
+
137
// ======================================================================================
138
// KV discovery
139
// ======================================================================================
rust/src/engine-core-client/src/client.rs
+73
@@ -1,12 +1,14 @@
1
// SPDX-License-Identifier: Apache-2.0
2
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+use std::collections::BTreeMap;
5
use std::sync::Arc;
6
use std::time::Duration;
7
8
use futures::future::{join_all, try_join_all};
9
use itertools::Itertools;
10
use serde::Serialize;
11
+use serde_json::Value as JsonValue;
12
use tokio::sync::mpsc;
13
use tokio_util::task::AbortOnDropHandle;
14
use tracing::{debug, info, trace};
@@ -674,6 +676,77 @@ impl EngineCoreClient {
676
.collect())
677
}
678
679
+ /// Initialize the configured RL weight-transfer backend.
680
+ pub async fn init_weight_transfer_engine(&self, init_info: JsonValue) -> Result<()> {
681
+ self.collective_rpc(
682
+ "init_weight_transfer_engine",
683
+ None,
684
+ Vec::<JsonValue>::new(),
685
+ BTreeMap::from([("init_info".to_string(), init_info)]),
686
+ )
687
+ .await?;
688
+ Ok(())
689
+ }
690
+
691
+ /// Start a weight update for the base model.
692
+ pub async fn start_weight_update(&self) -> Result<()> {
693
+ self.collective_rpc(
694
+ "start_weight_update",
695
+ None,
696
+ Vec::<JsonValue>::new(),
697
+ BTreeMap::<String, JsonValue>::new(),
698
+ )
699
+ .await?;
700
+ Ok(())
701
+ }
702
+
703
+ /// Start a weight update for the speculative draft model.
704
+ pub async fn start_draft_weight_update(&self) -> Result<()> {
705
+ self.collective_rpc(
706
+ "start_draft_weight_update",
707
+ None,
708
+ Vec::<JsonValue>::new(),
709
+ BTreeMap::<String, JsonValue>::new(),
710
+ )
711
+ .await?;
712
+ Ok(())
713
+ }
714
+
715
+ /// Apply one backend-specific weight metadata chunk.
716
+ pub async fn update_weights(&self, update_info: JsonValue) -> Result<()> {
717
+ self.collective_rpc(
718
+ "update_weights",
719
+ None,
720
+ Vec::<JsonValue>::new(),
721
+ BTreeMap::from([("update_info".to_string(), update_info)]),
722
+ )
723
+ .await?;
724
+ Ok(())
725
+ }
726
+
727
+ /// Finish the current weight update.
728
+ pub async fn finish_weight_update(&self) -> Result<()> {
729
+ self.collective_rpc(
730
+ "finish_weight_update",
731
+ None,
732
+ Vec::<JsonValue>::new(),
733
+ BTreeMap::<String, JsonValue>::new(),
734
+ )
735
+ .await?;
736
+ Ok(())
737
+ }
738
+
739
+ /// Set the committed weight version on every connected engine.
740
+ pub async fn set_weight_version(&self, weight_version: &str) -> Result<()> {
741
+ self.call_utility::<(), _>("set_weight_version", (weight_version,)).await?;
742
+ Ok(())
743
+ }
744
+
745
+ /// Return the committed weight version agreed on by every connected engine.
746
+ pub async fn get_weight_version(&self) -> Result<String> {
747
+ self.call_utility_consensus("get_weight_version", ()).await
748
+ }
749
+
750
/// Return whether the engine is currently sleeping at any level.
751
pub async fn is_sleeping(&self) -> Result<bool> {
752
self.call_utility_consensus("is_sleeping", ()).await
rust/src/engine-core-client/src/mock_engine.rs
+3
@@ -67,6 +67,9 @@ pub fn default_ready_response() -> EngineCoreReadyResponse {
67
kv_cache_size_tokens: None,
68
kv_cache_max_concurrency: None,
69
kv_events_config: None,
70
+ weight_transfer_backend: None,
71
+ enable_sleep_mode: false,
72
+ supports_draft_weight_updates: false,
73
}
74
}
75
rust/src/engine-core-client/src/protocol/handshake.rs
+9
@@ -87,6 +87,15 @@ pub struct EngineCoreReadyResponse {
87
/// KV-event publisher configuration, if configured.
88
#[serde(default)]
89
pub kv_events_config: Option<KvEventsConfig>,
90
+ /// Configured RL weight-transfer backend, if weight transfer is enabled.
91
+ #[serde(default)]
92
+ pub weight_transfer_backend: Option<String>,
93
+ /// Whether the engine was started with sleep mode enabled.
94
+ #[serde(default)]
95
+ pub enable_sleep_mode: bool,
96
+ /// Whether the engine has a speculative draft model that can be updated.
97
+ #[serde(default)]
98
+ pub supports_draft_weight_updates: bool,
99
}
100
101
/// Frontend-owned ZMQ addresses that are sent to the engine during startup
rust/src/engine-core-client/src/tests/client.rs
+6
@@ -2686,6 +2686,12 @@ fn python_msgpack_fixtures_match_rust_encoding() {
2686
2687
let ready_response: EngineCoreReadyResponse =
2688
rmp_serde::from_slice(&hex::decode(ready_response_hex).unwrap()).unwrap();
2689
+ assert_eq!(
2690
+ ready_response.weight_transfer_backend.as_deref(),
2691
+ Some("nccl")
2692
+ );
2693
+ assert!(ready_response.enable_sleep_mode);
2694
+ assert!(ready_response.supports_draft_weight_updates);
2695
let kv_events_config = ready_response.kv_events_config.expect("KV events config should decode");
2696
assert!(kv_events_config.enable_kv_cache_events);
2697
assert_eq!(kv_events_config.publisher, "zmq");
rust/src/engine-core-client/src/tests/python_compat.py
+6
@@ -412,6 +412,9 @@ class EngineCoreReadyResponse:
412
kv_cache_size_tokens: int | None = None
413
kv_cache_max_concurrency: float | None = None
414
kv_events_config: KVEventsConfig | None = None
415
+ weight_transfer_backend: str | None = None
416
+ enable_sleep_mode: bool = False
417
+ supports_draft_weight_updates: bool = False
418
419
420
ready_response = EngineCoreReadyResponse(
@@ -430,6 +433,9 @@ ready_response = EngineCoreReadyResponse(
433
max_num_seqs=256,
434
max_num_batched_tokens=8192,
435
instance_id="test-instance",
436
+ weight_transfer_backend="nccl",
437
+ enable_sleep_mode=True,
438
+ supports_draft_weight_updates=True,
439
kv_events_config=KVEventsConfig(
440
enable_kv_cache_events=True,
441
publisher="zmq",
rust/src/server/src/grpc/control.rs
+283
-1
@@ -3,9 +3,13 @@
3
4
use std::sync::Arc;
5
6
+use serde_json::Value as JsonValue;
7
use thiserror_ext::AsReport as _;
8
+use tokio::sync::Mutex;
9
use tonic::{Request, Response, Status};
10
+use vllm_engine_core_client::EngineCoreClient;
11
use vllm_engine_core_client::protocol::handshake::EngineCoreReadyResponse;
12
+use vllm_engine_core_client::protocol::utility::PauseMode as EnginePauseMode;
13
14
use super::{ControlServer, pb};
15
use crate::state::AppState;
@@ -15,17 +19,25 @@ pub(crate) type ControlGrpcService = ControlServer<ControlServiceImpl>;
19
/// gRPC control service backed by the shared application state.
20
pub struct ControlServiceImpl {
21
state: Arc<AppState>,
22
+ rl_lock: Mutex<()>,
23
}
24
25
impl ControlServiceImpl {
26
pub fn new(state: Arc<AppState>) -> Self {
22
- Self { state }
27
+ Self {
28
+ state,
29
+ rl_lock: Mutex::new(()),
30
+ }
31
}
32
33
fn ready(&self) -> &EngineCoreReadyResponse {
34
self.state.engine_core_client().ready_response()
35
}
36
37
+ fn client(&self) -> &EngineCoreClient {
38
+ self.state.engine_core_client()
39
+ }
40
+
41
fn parallelism_info(&self) -> pb::ParallelismInfo {
42
let ready = self.ready();
43
pb::ParallelismInfo {
@@ -36,10 +48,91 @@ impl ControlServiceImpl {
48
decode_context_parallel_size: ready.decode_context_parallel_size,
49
}
50
}
51
+
52
+ fn weight_transfer_backend(&self) -> Option<&str> {
53
+ let responses = self.client().ready_responses();
54
+ let backend = responses.first()?.weight_transfer_backend.as_deref()?;
55
+ responses
56
+ .iter()
57
+ .all(|ready| ready.weight_transfer_backend.as_deref() == Some(backend))
58
+ .then_some(backend)
59
+ }
60
+
61
+ fn sleep_mode_enabled(&self) -> bool {
62
+ self.client().ready_responses().iter().all(|ready| ready.enable_sleep_mode)
63
+ }
64
+
65
+ fn draft_weight_updates_enabled(&self) -> bool {
66
+ self.client()
67
+ .ready_responses()
68
+ .iter()
69
+ .all(|ready| ready.supports_draft_weight_updates)
70
+ }
71
+
72
+ fn rl_capabilities(&self) -> pb::RlCapabilities {
73
+ let backend = self.weight_transfer_backend();
74
+ pb::RlCapabilities {
75
+ weight_transfer_enabled: backend.is_some(),
76
+ weight_transfer_backend: backend.unwrap_or_default().to_string(),
77
+ sleep_mode_enabled: self.sleep_mode_enabled(),
78
+ draft_weight_updates_enabled: self.draft_weight_updates_enabled(),
79
+ }
80
+ }
81
+
82
+ fn require_weight_transfer(&self) -> Result<(), Status> {
83
+ self.weight_transfer_backend().map(|_| ()).ok_or_else(|| {
84
+ Status::failed_precondition(
85
+ "weight transfer is not configured; start vLLM with --weight-transfer-config",
86
+ )
87
+ })
88
+ }
89
+
90
+ fn require_sleep_mode(&self) -> Result<(), Status> {
91
+ self.sleep_mode_enabled().then_some(()).ok_or_else(|| {
92
+ Status::failed_precondition(
93
+ "sleep mode is not configured; start vLLM with --enable-sleep-mode",
94
+ )
95
+ })
96
+ }
97
}
98
99
const GRPC_API_VERSION: &str = "vllm";
100
101
+fn utility_status(method: &'static str, error: vllm_engine_core_client::Error) -> Status {
102
+ Status::internal(format!("{method} failed: {}", error.to_report_string()))
103
+}
104
+
105
+fn pause_mode(mode: i32) -> Result<EnginePauseMode, Status> {
106
+ match pb::PauseMode::try_from(mode) {
107
+ Ok(pb::PauseMode::Unspecified | pb::PauseMode::Abort) => Ok(EnginePauseMode::Abort),
108
+ Ok(pb::PauseMode::Wait) => Ok(EnginePauseMode::Wait),
109
+ Ok(pb::PauseMode::Keep) => Ok(EnginePauseMode::Keep),
110
+ Err(_) => Err(Status::invalid_argument("invalid pause mode")),
111
+ }
112
+}
113
+
114
+fn json_object(bytes: &[u8], field: &'static str) -> Result<JsonValue, Status> {
115
+ let value = serde_json::from_slice::<JsonValue>(bytes).map_err(|error| {
116
+ Status::invalid_argument(format!(
117
+ "{field} must contain valid JSON: {}",
118
+ error.to_report_string()
119
+ ))
120
+ })?;
121
+ if !value.is_object() {
122
+ return Err(Status::invalid_argument(format!(
123
+ "{field} must contain a JSON object"
124
+ )));
125
+ }
126
+ Ok(value)
127
+}
128
+
129
+fn weight_version(value: String) -> Result<String, Status> {
130
+ if value.trim().is_empty() {
131
+ return Err(Status::invalid_argument("weight_version must not be empty"));
132
+ }
133
+ Ok(value)
134
+}
135
+
136
#[tonic::async_trait]
137
impl pb::control_server::Control for ControlServiceImpl {
138
async fn get_server_info(
@@ -57,6 +150,7 @@ impl pb::control_server::Control for ControlServiceImpl {
150
total_kv_blocks: self.state.engine_core_client().total_num_gpu_blocks(),
151
max_running_requests: ready.max_num_seqs,
152
max_batched_tokens: ready.max_num_batched_tokens,
153
+ rl_capabilities: Some(self.rl_capabilities()),
154
}))
155
}
156
@@ -112,6 +206,194 @@ impl pb::control_server::Control for ControlServiceImpl {
206
let sources = client.ready_responses().into_iter().filter_map(kv_event_source).collect();
207
Ok(Response::new(pb::GetKvEventSourcesResponse { sources }))
208
}
209
+
210
+ async fn pause_generation(
211
+ &self,
212
+ request: Request<pb::PauseGenerationRequest>,
213
+ ) -> Result<Response<pb::PauseGenerationResponse>, Status> {
214
+ let request = request.into_inner();
215
+ let mode = pause_mode(request.mode)?;
216
+ let clear_cache = request.clear_cache.unwrap_or(true);
217
+ let _guard = self.rl_lock.lock().await;
218
+ self.client()
219
+ .pause_scheduler(mode, clear_cache)
220
+ .await
221
+ .map_err(|error| utility_status("pause_generation", error))?;
222
+ Ok(Response::new(pb::PauseGenerationResponse {}))
223
+ }
224
+
225
+ async fn resume_generation(
226
+ &self,
227
+ _request: Request<pb::ResumeGenerationRequest>,
228
+ ) -> Result<Response<pb::ResumeGenerationResponse>, Status> {
229
+ let _guard = self.rl_lock.lock().await;
230
+ self.client()
231
+ .resume_scheduler()
232
+ .await
233
+ .map_err(|error| utility_status("resume_generation", error))?;
234
+ Ok(Response::new(pb::ResumeGenerationResponse {}))
235
+ }
236
+
237
+ async fn is_paused(
238
+ &self,
239
+ _request: Request<pb::IsPausedRequest>,
240
+ ) -> Result<Response<pb::IsPausedResponse>, Status> {
241
+ let paused = self
242
+ .client()
243
+ .is_scheduler_paused()
244
+ .await
245
+ .map_err(|error| utility_status("is_paused", error))?;
246
+ Ok(Response::new(pb::IsPausedResponse { paused }))
247
+ }
248
+
249
+ async fn sleep(
250
+ &self,
251
+ request: Request<pb::SleepRequest>,
252
+ ) -> Result<Response<pb::SleepResponse>, Status> {
253
+ self.require_sleep_mode()?;
254
+ let request = request.into_inner();
255
+ let mode = pause_mode(request.mode)?;
256
+ let level = request.level.unwrap_or(1);
257
+ let _guard = self.rl_lock.lock().await;
258
+ self.client()
259
+ .sleep(level, mode)
260
+ .await
261
+ .map_err(|error| utility_status("sleep", error))?;
262
+ Ok(Response::new(pb::SleepResponse {}))
263
+ }
264
+
265
+ async fn wake_up(
266
+ &self,
267
+ request: Request<pb::WakeUpRequest>,
268
+ ) -> Result<Response<pb::WakeUpResponse>, Status> {
269
+ self.require_sleep_mode()?;
270
+ let tags = request.into_inner().tags;
271
+ let tags = (!tags.is_empty()).then_some(tags);
272
+ let _guard = self.rl_lock.lock().await;
273
+ self.client()
274
+ .wake_up(tags)
275
+ .await
276
+ .map_err(|error| utility_status("wake_up", error))?;
277
+ Ok(Response::new(pb::WakeUpResponse {}))
278
+ }
279
+
280
+ async fn is_sleeping(
281
+ &self,
282
+ _request: Request<pb::IsSleepingRequest>,
283
+ ) -> Result<Response<pb::IsSleepingResponse>, Status> {
284
+ let sleeping = self
285
+ .client()
286
+ .is_sleeping()
287
+ .await
288
+ .map_err(|error| utility_status("is_sleeping", error))?;
289
+ Ok(Response::new(pb::IsSleepingResponse { sleeping }))
290
+ }
291
+
292
+ async fn init_weight_transfer_engine(
293
+ &self,
294
+ request: Request<pb::InitWeightTransferEngineRequest>,
295
+ ) -> Result<Response<pb::InitWeightTransferEngineResponse>, Status> {
296
+ self.require_weight_transfer()?;
297
+ let init_info = json_object(&request.into_inner().init_info_json, "init_info_json")?;
298
+ let _guard = self.rl_lock.lock().await;
299
+ self.client()
300
+ .init_weight_transfer_engine(init_info)
301
+ .await
302
+ .map_err(|error| utility_status("init_weight_transfer_engine", error))?;
303
+ Ok(Response::new(pb::InitWeightTransferEngineResponse {}))
304
+ }
305
+
306
+ async fn start_weight_update(
307
+ &self,
308
+ _request: Request<pb::StartWeightUpdateRequest>,
309
+ ) -> Result<Response<pb::StartWeightUpdateResponse>, Status> {
310
+ self.require_weight_transfer()?;
311
+ let _guard = self.rl_lock.lock().await;
312
+ self.client()
313
+ .start_weight_update()
314
+ .await
315
+ .map_err(|error| utility_status("start_weight_update", error))?;
316
+ Ok(Response::new(pb::StartWeightUpdateResponse {}))
317
+ }
318
+
319
+ async fn start_draft_weight_update(
320
+ &self,
321
+ _request: Request<pb::StartDraftWeightUpdateRequest>,
322
+ ) -> Result<Response<pb::StartDraftWeightUpdateResponse>, Status> {
323
+ self.require_weight_transfer()?;
324
+ if !self.draft_weight_updates_enabled() {
325
+ return Err(Status::failed_precondition(
326
+ "draft weight updates require a configured speculative draft model",
327
+ ));
328
+ }
329
+ let _guard = self.rl_lock.lock().await;
330
+ self.client()
331
+ .start_draft_weight_update()
332
+ .await
333
+ .map_err(|error| utility_status("start_draft_weight_update", error))?;
334
+ Ok(Response::new(pb::StartDraftWeightUpdateResponse {}))
335
+ }
336
+
337
+ async fn update_weights(
338
+ &self,
339
+ request: Request<pb::UpdateWeightsRequest>,
340
+ ) -> Result<Response<pb::UpdateWeightsResponse>, Status> {
341
+ self.require_weight_transfer()?;
342
+ let update_info = json_object(&request.into_inner().update_info_json, "update_info_json")?;
343
+ let _guard = self.rl_lock.lock().await;
344
+ self.client()
345
+ .update_weights(update_info)
346
+ .await
347
+ .map_err(|error| utility_status("update_weights", error))?;
348
+ Ok(Response::new(pb::UpdateWeightsResponse {}))
349
+ }
350
+
351
+ async fn finish_weight_update(
352
+ &self,
353
+ request: Request<pb::FinishWeightUpdateRequest>,
354
+ ) -> Result<Response<pb::FinishWeightUpdateResponse>, Status> {
355
+ self.require_weight_transfer()?;
356
+ let version = request.into_inner().weight_version.map(weight_version).transpose()?;
357
+ let _guard = self.rl_lock.lock().await;
358
+ self.client()
359
+ .finish_weight_update()
360
+ .await
361
+ .map_err(|error| utility_status("finish_weight_update", error))?;
362
+ if let Some(version) = version {
363
+ self.client()
364
+ .set_weight_version(&version)
365
+ .await
366
+ .map_err(|error| utility_status("update_weight_version", error))?;
367
+ }
368
+ Ok(Response::new(pb::FinishWeightUpdateResponse {}))
369
+ }
370
+
371
+ async fn update_weight_version(
372
+ &self,
373
+ request: Request<pb::UpdateWeightVersionRequest>,
374
+ ) -> Result<Response<pb::UpdateWeightVersionResponse>, Status> {
375
+ let version = weight_version(request.into_inner().weight_version)?;
376
+ let _guard = self.rl_lock.lock().await;
377
+ self.client()
378
+ .set_weight_version(&version)
379
+ .await
380
+ .map_err(|error| utility_status("update_weight_version", error))?;
381
+ Ok(Response::new(pb::UpdateWeightVersionResponse {}))
382
+ }
383
+
384
+ async fn get_weight_version(
385
+ &self,
386
+ _request: Request<pb::GetWeightVersionRequest>,
387
+ ) -> Result<Response<pb::GetWeightVersionResponse>, Status> {
388
+ let weight_version = self
389
+ .client()
390
+ .get_weight_version()
391
+ .await
392
+ .map_err(|error| utility_status("get_weight_version", error))?;
393
+ Ok(Response::new(pb::GetWeightVersionResponse {
394
+ weight_version,
395
+ }))
396
+ }
397
}
398
399
pub(super) fn kv_event_source(response: &EngineCoreReadyResponse) -> Option<pb::KvEventSource> {
rust/src/server/src/grpc/tests.rs
+109
-11
@@ -12,6 +12,7 @@ use std::time::Duration;
12
use futures::StreamExt as _;
13
use hyper_util::rt::TokioIo;
14
use openssl::ssl::{SslConnector, SslFiletype, SslMethod};
15
+use rmpv::Value;
16
use serial_test::serial;
17
use tokio::io::{AsyncReadExt as _, AsyncWriteExt as _};
18
use tokio::net::TcpStream;
@@ -30,14 +31,15 @@ use vllm_engine_core_client::mock_engine::{
31
DEFAULT_MOCK_BLOCK_SIZE, DEFAULT_MOCK_MAX_MODEL_LEN, DEFAULT_MOCK_NUM_GPU_BLOCKS,
32
default_ready_response,
33
};
33
-use vllm_engine_core_client::protocol::handshake::KvEventsConfig;
34
+use vllm_engine_core_client::protocol::decode_value;
35
+use vllm_engine_core_client::protocol::handshake::{EngineCoreReadyResponse, KvEventsConfig};
36
use vllm_engine_core_client::protocol::output::{
37
EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs, RequestBatchOutputs,
38
+ UtilityCallOutput,
39
};
40
use vllm_engine_core_client::protocol::request::EngineCoreRequest;
38
-use vllm_engine_core_client::test_utils::{
39
- IpcNamespace, spawn_mock_engine_task, spawn_mock_engine_task_with_ready,
40
-};
41
+use vllm_engine_core_client::protocol::utility::{UtilityOutput, UtilityResultEnvelope};
42
+use vllm_engine_core_client::test_utils::{IpcNamespace, spawn_mock_engine_task_with_ready};
43
use vllm_engine_core_client::{EngineCoreClient, EngineCoreClientConfig, EngineId, TransportMode};
44
use vllm_llm::Llm;
45
use vllm_text::tokenizer::DynTokenizer;
@@ -291,13 +293,10 @@ async fn setup_grpc_service_with_backend<F>(
293
where
294
F: FnOnce(&EngineCoreRequest) + Send + 'static,
295
{
294
- let ipc = IpcNamespace::new().expect("create ipc namespace");
295
- let handshake_address = ipc.handshake_endpoint();
296
- let engine_id = engine_id.into();
297
-
298
- let engine_task = MockEngineTask::new(spawn_mock_engine_task(
299
- handshake_address.clone(),
300
- engine_id.clone(),
296
+ setup_grpc_service_with_engine_script(
297
+ engine_id,
298
+ default_ready_response(),
299
+ backend,
300
move |dealer, push| {
301
boxed_test_future(async move {
302
let add = recv_engine_message(dealer).await;
@@ -311,6 +310,33 @@ where
310
.await;
311
})
312
},
313
+ )
314
+ .await
315
+}
316
+
317
+async fn setup_grpc_service_with_engine_script<F>(
318
+ engine_id: impl Into<EngineId>,
319
+ ready: EngineCoreReadyResponse,
320
+ backend: Arc<dyn ChatTextBackend>,
321
+ script: F,
322
+) -> (
323
+ InferenceServer<InferenceServiceImpl>,
324
+ ControlServer<ControlServiceImpl>,
325
+ tokio::sync::watch::Receiver<bool>,
326
+ MockEngineTask,
327
+)
328
+where
329
+ F: for<'a> FnOnce(&'a mut DealerSocket, &'a mut PushSocket) -> TestFuture<'a> + Send + 'static,
330
+{
331
+ let ipc = IpcNamespace::new().expect("create ipc namespace");
332
+ let handshake_address = ipc.handshake_endpoint();
333
+ let engine_id = engine_id.into();
334
+
335
+ let engine_task = MockEngineTask::new(spawn_mock_engine_task_with_ready(
336
+ handshake_address.clone(),
337
+ engine_id.clone(),
338
+ ready,
339
+ script,
340
));
341
342
let client = EngineCoreClient::connect(
@@ -1509,6 +1535,11 @@ async fn control_reports_server_and_model_info() {
1535
assert_eq!(parallelism.data_parallel_size, 1);
1536
assert_eq!(parallelism.data_parallel_rank, 0);
1537
assert_eq!(parallelism.decode_context_parallel_size, 1);
1538
+ let rl = server.rl_capabilities.expect("RL capabilities");
1539
+ assert!(!rl.weight_transfer_enabled);
1540
+ assert!(rl.weight_transfer_backend.is_empty());
1541
+ assert!(!rl.sleep_mode_enabled);
1542
+ assert!(!rl.draft_weight_updates_enabled);
1543
1544
let model = client
1545
.get_model_info(pb::GetModelInfoRequest {})
@@ -1527,6 +1558,65 @@ async fn control_reports_server_and_model_info() {
1558
server_task.abort();
1559
}
1560
1561
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1562
+#[serial]
1563
+async fn control_forwards_weight_update_without_pause_guard() {
1564
+ let mut ready = default_ready_response();
1565
+ ready.weight_transfer_backend = Some("nccl".to_string());
1566
+ let (inference_service, control_service, engine_health, engine_task) =
1567
+ setup_grpc_service_with_engine_script(
1568
+ b"engine-grpc-rl".to_vec(),
1569
+ ready,
1570
+ Arc::new(FakeTextBackend),
1571
+ |dealer, push| {
1572
+ boxed_test_future(async move {
1573
+ let frames = recv_engine_message(dealer).await;
1574
+ assert_eq!(frames[0].as_ref(), &[0x03]);
1575
+ let payload = decode_value(&frames[1]).expect("decode utility payload");
1576
+ let fields = payload.as_array().expect("utility payload array");
1577
+ let call_id = fields[1].as_u64().expect("utility call id");
1578
+ assert_eq!(fields[2].as_str(), Some("collective_rpc"));
1579
+ let args = fields[3].as_array().expect("collective_rpc arguments");
1580
+ assert_eq!(args[0].as_str(), Some("update_weights"));
1581
+ send_outputs(
1582
+ push,
1583
+ UtilityCallOutput {
1584
+ output: UtilityOutput {
1585
+ call_id: call_id.into(),
1586
+ failure_message: None,
1587
+ result: Some(UtilityResultEnvelope::without_type_info(
1588
+ Value::Array(vec![Value::Nil]),
1589
+ )),
1590
+ },
1591
+ ..Default::default()
1592
+ }
1593
+ .into(),
1594
+ )
1595
+ .await;
1596
+ })
1597
+ },
1598
+ )
1599
+ .await;
1600
+ let (channel, server_task) = start_grpc_test_server(
1601
+ inference_service,
1602
+ control_service,
1603
+ engine_health,
1604
+ tokio_util::sync::CancellationToken::new(),
1605
+ )
1606
+ .await;
1607
+ let mut client = ControlClient::new(channel);
1608
+
1609
+ client
1610
+ .update_weights(pb::UpdateWeightsRequest {
1611
+ update_info_json: br#"{"names":["model.weight"]}"#.to_vec(),
1612
+ })
1613
+ .await
1614
+ .expect("forward weight update without a pause probe");
1615
+
1616
+ engine_task.await.expect("mock engine task");
1617
+ server_task.abort();
1618
+}
1619
+
1620
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1621
async fn control_aggregates_multi_engine_capacity() {
1622
let ipc = IpcNamespace::new().expect("create ipc namespace");
@@ -1536,6 +1626,9 @@ async fn control_aggregates_multi_engine_capacity() {
1626
ready_0.max_model_len = 8_192;
1627
ready_0.num_gpu_blocks = 10;
1628
ready_0.data_parallel_size = 2;
1629
+ ready_0.weight_transfer_backend = Some("nccl".to_string());
1630
+ ready_0.enable_sleep_mode = true;
1631
+ ready_0.supports_draft_weight_updates = true;
1632
1633
let mut ready_1 = default_ready_response();
1634
ready_1.max_model_len = 4_096;
@@ -1586,6 +1679,11 @@ async fn control_aggregates_multi_engine_capacity() {
1679
.into_inner();
1680
assert_eq!(server.max_model_len, 4_096);
1681
assert_eq!(server.total_kv_blocks, 30);
1682
+ let rl = server.rl_capabilities.expect("RL capabilities");
1683
+ assert!(!rl.weight_transfer_enabled);
1684
+ assert!(rl.weight_transfer_backend.is_empty());
1685
+ assert!(!rl.sleep_mode_enabled);
1686
+ assert!(!rl.draft_weight_updates_enabled);
1687
assert_eq!(server.parallelism.unwrap().data_parallel_size, 4);
1688
1689
drop(engine_tasks);
rust/src/server/src/lib.rs
+2
-1
@@ -217,7 +217,8 @@ where
217
health_reporter.set_serving::<grpc::InferenceGrpcService>().await;
218
health_reporter.set_serving::<grpc::ControlGrpcService>().await;
219
let control_service =
220
- grpc::ControlGrpcService::new(grpc::ControlServiceImpl::new(state.clone()));
220
+ grpc::ControlGrpcService::new(grpc::ControlServiceImpl::new(state.clone()))
221
+ .max_decoding_message_size(DEFAULT_REQUEST_BODY_LIMIT_BYTES);
222
let inference_service =
223
grpc::InferenceGrpcService::new(grpc::InferenceServiceImpl::new(state.clone()))
224
.max_decoding_message_size(DEFAULT_REQUEST_BODY_LIMIT_BYTES);
tests/v1/engine/test_engine_core_client.py
+1
@@ -1292,6 +1292,7 @@ def test_engine_core_proc_instantiation_cuda_empty(monkeypatch: pytest.MonkeyPat
1292
mock_executor.get_kv_cache_specs.return_value = [{"default": mock_spec}]
1293
mock_executor.determine_available_memory.return_value = [1024 * 1024 * 1024]
1294
mock_executor.initialize_from_config.return_value = None
1295
+ mock_executor.supports_draft_weight_updates.return_value = False
1296
1297
return mock_executor
1298
vllm/v1/engine/__init__.py
+3
@@ -92,6 +92,9 @@ class EngineCoreReadyResponse:
92
kv_cache_size_tokens: int | None = None
93
kv_cache_max_concurrency: float | None = None
94
kv_events_config: KVEventsConfig | None = None
95
+ weight_transfer_backend: str | None = None
96
+ enable_sleep_mode: bool = False
97
+ supports_draft_weight_updates: bool = False
98
99
100
class EngineCoreRequest(
vllm/v1/engine/core.py
+9
@@ -1635,6 +1635,15 @@ class EngineCoreProc(EngineCore):
1635
max_num_batched_tokens=scheduler_config.max_num_batched_tokens,
1636
instance_id=self.vllm_config.instance_id,
1637
kv_events_config=self.scheduler.get_kv_event_publisher_config(),
1638
+ weight_transfer_backend=(
1639
+ self.vllm_config.weight_transfer_config.backend
1640
+ if self.vllm_config.weight_transfer_config is not None
1641
+ else None
1642
+ ),
1643
+ enable_sleep_mode=self.vllm_config.model_config.enable_sleep_mode,
1644
+ supports_draft_weight_updates=(
1645
+ self.model_executor.supports_draft_weight_updates()
1646
+ ),
1647
)
1648
1649
def process_input_sockets(
vllm/v1/executor/abstract.py
+6
@@ -289,6 +289,12 @@ class Executor(ABC):
289
output = self.collective_rpc("get_supported_tasks")
290
return output[0]
291
292
+ def supports_draft_weight_updates(self) -> bool:
293
+ worker_support: list[bool] = self.collective_rpc(
294
+ "supports_draft_weight_updates"
295
+ )
296
+ return all(worker_support)
297
+
298
def add_lora(self, lora_request: LoRARequest) -> bool:
299
assert lora_request.lora_int_id > 0, "lora_id must be greater than 0."
300
return all(self.collective_rpc("add_lora", args=(lora_request,)))
vllm/v1/worker/gpu_worker.py
+13
@@ -869,6 +869,19 @@ class Worker(WorkerBase):
869
def get_draft_model(self) -> nn.Module | None:
870
return self.model_runner.get_draft_model()
871
872
+ def supports_draft_weight_updates(self) -> bool:
873
+ engine = self.weight_transfer_engine
874
+ speculative_config = self.speculative_config
875
+ get_draft_model = getattr(self.model_runner, "get_draft_model", None)
876
+ return (
877
+ engine is not None
878
+ and engine.supports_draft_weight_update
879
+ and callable(get_draft_model)
880
+ and get_draft_model() is not None
881
+ and speculative_config is not None
882
+ and speculative_config.draft_model_config is not None
883
+ )
884
+
885
def _set_draft_weight_update_target(self) -> None:
886
assert self.weight_transfer_engine is not None
887
vllm/v1/worker/worker_base.py
+4
@@ -125,6 +125,10 @@ class WorkerBase:
125
def get_model(self) -> nn.Module:
126
raise NotImplementedError
127
128
+ def supports_draft_weight_updates(self) -> bool:
129
+ """Whether this worker can update its configured speculative model."""
130
+ return False
131
+
132
def apply_model(self, fn: Callable[[nn.Module], _R]) -> _R:
133
"""Apply a function on the model inside this worker."""
134
return fn(self.get_model())