diff --git a/lib/sidecar/vllm/README.md b/lib/sidecar/vllm/README.md index e265603d4b37..205d77391272 100644 --- a/lib/sidecar/vllm/README.md +++ b/lib/sidecar/vllm/README.md @@ -26,6 +26,7 @@ It is a standalone Rust executable. - Sampling, stop conditions, structured output, logprobs, cache options, and priority - Opaque `kv_transfer_params` handoff - Data-parallel rank routing and KV-event source discovery +- Capability-gated RL pause/resume, sleep/wake, weight-transfer, and weight-version controls through native gRPC - Image URL and data-URI inputs, including media UUIDs The protocol does not support LoRA, encode workers, beam search, `n > 1`, @@ -56,7 +57,29 @@ dynamo-vllm-sidecar \ Use `VLLM_GRPC_ENDPOINT` instead of `--vllm-endpoint` when the endpoint is provided through the environment. -The sidecar discovers `model_id`, the served name, context length, KV capacity, scheduler limits, data-parallel topology, and KV-event sources through `vllm.Control`. `model_id` must be readable locally or fetchable by Dynamo for tokenization and chat templates. +### RL workflows + +Start vLLM with the capabilities required by the workflow, then opt the sidecar into RL discovery: + +```bash +vllm-rs serve Qwen/Qwen3-0.6B \ + --host 127.0.0.1 \ + --grpc-port 50051 \ + --enable-sleep-mode \ + --weight-transfer-config '{"backend":"nccl"}' + +DYN_SYSTEM_PORT=8081 dynamo-vllm-sidecar \ + --vllm-endpoint 127.0.0.1:50051 \ + --enable-rl +``` + +`--enable-rl` (or `DYN_ENABLE_RL=true`) requires the Dynamo system server (`DYN_SYSTEM_PORT=0` or a positive port) and registers `dyn://..rl`, which lets the Dynamo frontend discover this worker and its `/engine/control/*` and `/engine/update/*` routes through `/v1/rl/workers`. The sidecar advertises pause/resume, sleep-status, and weight-version controls when the vLLM server reports the RL gRPC API; mutating sleep/wake routes require `--enable-sleep-mode`, weight-transfer routes require `--weight-transfer-config`, and draft updates require speculative decoding support. + +The update request bodies match vLLM's RL HTTP schemas: `init_weight_transfer_engine` requires `{"init_info": {...}}`, `update_weights` requires `{"update_info": {...}}`, `finish_weight_update` accepts `{"weight_version": "..."}`, and `update_weight_version` requires `{"new_version": "..."}`. Weight tensors remain on the configured NCCL, IPC, or sparse-NCCL transport; only backend metadata crosses gRPC. + +The RL endpoint and engine routes are unauthenticated administrative surfaces that can pause serving, release GPU memory, and replace model weights. Enable them only on trusted request and system networks. + +The sidecar discovers `model_id`, the served name, context length, KV capacity, scheduler limits, data-parallel topology, and KV-event sources through `vllm.Control`. `model_id` must be readable locally or fetchable by Dynamo for tokenization and chat templates. Parser defaults are not advertised because the current inference protocol cannot preserve all parser-related request semantics. The sidecar currently supports one vLLM frontend hosting the complete data-parallel group starting at rank 0. Control reports the global size; Dynamo forwards the selected rank as `x-data-parallel-rank` gRPC metadata on each generation request. Partial and hybrid rank ownership are unsupported because the protocol does not report the locally hosted rank count, and a nonzero starting rank is rejected. When KV routing is enabled, Control must return one unique ZMQ event source for every rank in the group. @@ -88,6 +111,8 @@ cargo run -p dynamo-vllm-sidecar --bin dynamo-vllm-sidecar -- \ --vllm-endpoint 127.0.0.1:50051 ``` +The mocker does not advertise RL capabilities; use a compatible vLLM server for RL route testing. + See [`../../mocker/servers/vllm/README.md`](../../mocker/servers/vllm/README.md) for aggregated and prefill/decode examples, supported Mocker configuration, and fidelity limits. diff --git a/lib/sidecar/vllm/src/client.rs b/lib/sidecar/vllm/src/client.rs index 405045cb27fb..7913080954f0 100644 --- a/lib/sidecar/vllm/src/client.rs +++ b/lib/sidecar/vllm/src/client.rs @@ -9,6 +9,7 @@ use dynamo_sidecar_common::{ }; use tokio::time::{Instant, sleep_until, timeout_at}; use tonic::metadata::MetadataValue; +use tonic::transport::Channel; use tonic_health::pb::health_check_response::ServingStatus; use tonic_health::pb::{HealthCheckRequest, health_client::HealthClient}; @@ -47,6 +48,12 @@ impl VllmClient { self.pool.len() } + pub(crate) fn control_client(&self) -> pb::control_client::ControlClient { + pb::control_client::ControlClient::new(self.pool.next_channel()) + .max_encoding_message_size(DEFAULT_MAX_GRPC_MESSAGE_SIZE) + .max_decoding_message_size(DEFAULT_MAX_GRPC_MESSAGE_SIZE) + } + pub(crate) async fn wait_for_services( &self, services: &[&str], @@ -115,10 +122,7 @@ impl VllmClient { &self, startup_deadline: Instant, ) -> Result<(pb::ModelInfo, pb::ServerInfo), DynamoError> { - let channel = self.pool.next_channel(); - let mut client = pb::control_client::ControlClient::new(channel) - .max_encoding_message_size(DEFAULT_MAX_GRPC_MESSAGE_SIZE) - .max_decoding_message_size(DEFAULT_MAX_GRPC_MESSAGE_SIZE); + let mut client = self.control_client(); let model = timeout_at( startup_deadline, client.get_model_info(pb::GetModelInfoRequest {}), diff --git a/lib/sidecar/vllm/src/engine.rs b/lib/sidecar/vllm/src/engine.rs index 0917af391d31..5c7936f50f0e 100644 --- a/lib/sidecar/vllm/src/engine.rs +++ b/lib/sidecar/vllm/src/engine.rs @@ -10,6 +10,7 @@ use dynamo_backend_common::{ }; use dynamo_sidecar_common::{GrpcEndpoint, GrpcTransportConfig}; use futures::stream::BoxStream; +use serde_json::{Map, Value, json}; use tokio::sync::OnceCell; use tokio::time::Instant; use tokio_util::sync::CancellationToken; @@ -138,6 +139,12 @@ impl VllmSidecarEngine { }; Ok((engine, config)) } + + fn started_client(&self) -> Result<&VllmClient, DynamoError> { + self.client + .get() + .ok_or_else(|| client::engine_shutdown("vLLM sidecar is not started")) + } } #[async_trait] @@ -332,6 +339,198 @@ impl LLMEngine for VllmSidecarEngine { })) } + async fn supported_controls(&self) -> Result, DynamoError> { + let Some(capabilities) = self.model.rl_capabilities() else { + return Ok(Vec::new()); + }; + let mut controls = vec![ + "pause_generation".to_string(), + "resume_generation".to_string(), + "is_paused".to_string(), + "is_sleeping".to_string(), + "get_weight_version".to_string(), + ]; + if capabilities.sleep_mode_enabled { + controls.extend(["sleep".to_string(), "wake_up".to_string()]); + } + Ok(controls) + } + + fn validate_engine_control(&self, control: &str, body: &Value) -> Result<(), DynamoError> { + let body = request_object(body)?; + match control { + "pause_generation" => { + pause_mode(body, "pause_generation")?; + optional_bool(body, "clear_cache")?; + } + "sleep" => { + sleep_level(body)?; + pause_mode(body, "sleep")?; + } + "wake_up" => { + wake_tags(body)?; + } + _ => {} + } + Ok(()) + } + + async fn engine_control(&self, control: String, body: Value) -> Result { + if !self.supported_controls().await?.contains(&control) { + return Ok(unsupported("control", &control)); + } + let body = request_object(&body)?; + let mut grpc = self.started_client()?.control_client(); + match control.as_str() { + "pause_generation" => { + let mode = pause_mode(body, "pause_generation")?; + let clear_cache = optional_bool(body, "clear_cache")?; + grpc.pause_generation(crate::proto::PauseGenerationRequest { + mode: mode as i32, + clear_cache, + }) + .await + .map_err(|status| client::status_to_dynamo("PauseGeneration", status))?; + Ok(json!({"status": "paused"})) + } + "resume_generation" => { + grpc.resume_generation(crate::proto::ResumeGenerationRequest {}) + .await + .map_err(|status| client::status_to_dynamo("ResumeGeneration", status))?; + let sleeping = if self + .model + .rl_capabilities() + .is_some_and(|capabilities| capabilities.sleep_mode_enabled) + { + grpc_is_sleeping(&mut grpc).await? + } else { + false + }; + if sleeping { + Ok(json!({"status": "resumed", "is_sleeping": true})) + } else { + Ok(json!({"status": "resumed"})) + } + } + "is_paused" => { + let response = grpc + .is_paused(crate::proto::IsPausedRequest {}) + .await + .map_err(|status| client::status_to_dynamo("IsPaused", status))? + .into_inner(); + Ok(json!({"is_paused": response.paused})) + } + "sleep" => { + let level = sleep_level(body)?; + let mode = pause_mode(body, "sleep")?; + grpc.sleep(crate::proto::SleepRequest { + level, + mode: mode as i32, + }) + .await + .map_err(|status| client::status_to_dynamo("Sleep", status))?; + Ok(json!({"status": "sleeping"})) + } + "wake_up" => { + let tags = wake_tags(body)?; + grpc.wake_up(crate::proto::WakeUpRequest { tags }) + .await + .map_err(|status| client::status_to_dynamo("WakeUp", status))?; + if grpc_is_sleeping(&mut grpc).await? { + Ok(json!({"status": "partially_awake", "is_sleeping": true})) + } else { + Ok(json!({"status": "awake"})) + } + } + "is_sleeping" => Ok(json!({"is_sleeping": grpc_is_sleeping(&mut grpc).await?})), + "get_weight_version" => { + let response = grpc + .get_weight_version(crate::proto::GetWeightVersionRequest {}) + .await + .map_err(|status| client::status_to_dynamo("GetWeightVersion", status))? + .into_inner(); + Ok(json!({"weight_version": response.weight_version})) + } + _ => Ok(unsupported("control", &control)), + } + } + + async fn supported_updates(&self) -> Result, DynamoError> { + let Some(capabilities) = self.model.rl_capabilities() else { + return Ok(Vec::new()); + }; + let mut updates = vec!["update_weight_version".to_string()]; + if capabilities.weight_transfer_enabled { + updates.extend([ + "init_weight_transfer_engine".to_string(), + "start_weight_update".to_string(), + "update_weights".to_string(), + "finish_weight_update".to_string(), + ]); + if capabilities.draft_weight_updates_enabled { + updates.push("start_draft_weight_update".to_string()); + } + } + Ok(updates) + } + + async fn engine_update(&self, update: String, body: Value) -> Result { + if !self.supported_updates().await?.contains(&update) { + return Ok(unsupported("update", &update)); + } + let body = request_object(&body)?; + let mut grpc = self.started_client()?.control_client(); + match update.as_str() { + "init_weight_transfer_engine" => { + let init_info_json = required_object_json(body, "init_info")?; + grpc.init_weight_transfer_engine(crate::proto::InitWeightTransferEngineRequest { + init_info_json, + }) + .await + .map_err(|status| client::status_to_dynamo("InitWeightTransferEngine", status))?; + Ok(json!({"message": "Weight transfer initialized"})) + } + "start_weight_update" => { + grpc.start_weight_update(crate::proto::StartWeightUpdateRequest {}) + .await + .map_err(|status| client::status_to_dynamo("StartWeightUpdate", status))?; + Ok(json!({"message": "Weight update started"})) + } + "start_draft_weight_update" => { + grpc.start_draft_weight_update(crate::proto::StartDraftWeightUpdateRequest {}) + .await + .map_err(|status| client::status_to_dynamo("StartDraftWeightUpdate", status))?; + Ok(json!({"message": "Draft weight update started"})) + } + "update_weights" => { + let update_info_json = required_object_json(body, "update_info")?; + grpc.update_weights(crate::proto::UpdateWeightsRequest { update_info_json }) + .await + .map_err(|status| client::status_to_dynamo("UpdateWeights", status))?; + Ok(json!({"message": "Weights updated"})) + } + "finish_weight_update" => { + let weight_version = optional_string(body, "weight_version")?; + grpc.finish_weight_update(crate::proto::FinishWeightUpdateRequest { + weight_version, + }) + .await + .map_err(|status| client::status_to_dynamo("FinishWeightUpdate", status))?; + Ok(json!({"message": "Weight update finished"})) + } + "update_weight_version" => { + let weight_version = required_string(body, "new_version")?; + grpc.update_weight_version(crate::proto::UpdateWeightVersionRequest { + weight_version: weight_version.clone(), + }) + .await + .map_err(|status| client::status_to_dynamo("UpdateWeightVersion", status))?; + Ok(json!({"success": true, "new_version": weight_version})) + } + _ => Ok(unsupported("update", &update)), + } + } + async fn cleanup(&self) -> Result<(), DynamoError> { self.cancel.cancel(); Ok(()) @@ -406,6 +605,131 @@ fn zmq_connect_endpoint(endpoint: &str, grpc_endpoint: &GrpcEndpoint) -> String format!("tcp://{}:{port}", grpc_endpoint.authority_host()) } +fn unsupported(kind: &str, name: &str) -> Value { + json!({ + "status": "error", + "message": format!("unsupported engine {kind}: {name}"), + }) +} + +async fn grpc_is_sleeping( + grpc: &mut crate::proto::control_client::ControlClient, +) -> Result { + grpc.is_sleeping(crate::proto::IsSleepingRequest {}) + .await + .map_err(|status| client::status_to_dynamo("IsSleeping", status)) + .map(|response| response.into_inner().sleeping) +} + +fn request_object(body: &Value) -> Result<&Map, DynamoError> { + body.as_object() + .ok_or_else(|| client::invalid_argument("engine request body must be a JSON object")) +} + +fn optional_bool(body: &Map, field: &str) -> Result, DynamoError> { + match body.get(field) { + None | Some(Value::Null) => Ok(None), + Some(Value::Bool(value)) => Ok(Some(*value)), + Some(_) => Err(client::invalid_argument(format!( + "`{field}` must be a boolean" + ))), + } +} + +fn optional_u32(body: &Map, field: &str) -> Result, DynamoError> { + match body.get(field) { + None | Some(Value::Null) => Ok(None), + Some(value) => value + .as_u64() + .and_then(|value| u32::try_from(value).ok()) + .map(Some) + .ok_or_else(|| client::invalid_argument(format!("`{field}` must be a uint32"))), + } +} + +fn sleep_level(body: &Map) -> Result, DynamoError> { + let level = optional_u32(body, "level")?; + if level.is_some_and(|level| level > 2) { + return Err(client::invalid_argument( + "`level` must be one of 0, 1, or 2", + )); + } + Ok(level) +} + +fn optional_string(body: &Map, field: &str) -> Result, DynamoError> { + match body.get(field) { + None | Some(Value::Null) => Ok(None), + Some(Value::String(value)) => Ok(Some(value.clone())), + Some(_) => Err(client::invalid_argument(format!( + "`{field}` must be a string" + ))), + } +} + +fn required_string(body: &Map, field: &str) -> Result { + optional_string(body, field)? + .filter(|value| !value.is_empty()) + .ok_or_else(|| client::invalid_argument(format!("missing non-empty `{field}` string"))) +} + +fn pause_mode( + body: &Map, + operation: &str, +) -> Result { + match optional_string(body, "mode")?.as_deref().unwrap_or("abort") { + "abort" => Ok(crate::proto::PauseMode::Abort), + "wait" => Ok(crate::proto::PauseMode::Wait), + "keep" => Ok(crate::proto::PauseMode::Keep), + value => Err(client::invalid_argument(format!( + "{operation} mode must be abort, wait, or keep; got `{value}`" + ))), + } +} + +fn optional_strings( + body: &Map, + field: &str, +) -> Result>, DynamoError> { + match body.get(field) { + None | Some(Value::Null) => Ok(None), + Some(Value::Array(values)) => values + .iter() + .map(|value| { + value.as_str().map(ToString::to_string).ok_or_else(|| { + client::invalid_argument(format!("`{field}` must contain only strings")) + }) + }) + .collect::, _>>() + .map(Some), + Some(_) => Err(client::invalid_argument(format!( + "`{field}` must be an array of strings" + ))), + } +} + +fn wake_tags(body: &Map) -> Result, DynamoError> { + let tags = optional_strings(body, "tags")?.unwrap_or_default(); + if let Some(tag) = tags + .iter() + .find(|tag| !matches!(tag.as_str(), "weights" | "kv_cache" | "scheduling")) + { + return Err(client::invalid_argument(format!( + "wake_up tag must be weights, kv_cache, or scheduling; got `{tag}`" + ))); + } + Ok(tags) +} + +fn required_object_json(body: &Map, field: &str) -> Result, DynamoError> { + let value = body + .get(field) + .and_then(Value::as_object) + .ok_or_else(|| client::invalid_argument(format!("missing `{field}` JSON object")))?; + serde_json::to_vec(value) + .map_err(|error| client::invalid_argument(format!("invalid `{field}`: {error}"))) +} + fn bootstrap_discover( endpoint: &GrpcEndpoint, transport: GrpcTransportConfig, diff --git a/lib/sidecar/vllm/src/model.rs b/lib/sidecar/vllm/src/model.rs index d7d8540b64e3..313ee3da146e 100644 --- a/lib/sidecar/vllm/src/model.rs +++ b/lib/sidecar/vllm/src/model.rs @@ -89,9 +89,19 @@ impl DiscoveredModel { "data-parallel size changed between bootstrap and startup: expected {expected_dp_size}, observed {observed_dp_size}" ))); } + if self.server.rl_capabilities != observed.server.rl_capabilities { + return Err(client::protocol_error(format!( + "RL capabilities changed between bootstrap and startup: expected {:?}, observed {:?}", + self.server.rl_capabilities, observed.server.rl_capabilities + ))); + } Ok(()) } + pub(crate) fn rl_capabilities(&self) -> Option<&pb::RlCapabilities> { + self.server.rl_capabilities.as_ref() + } + pub(crate) fn engine_config(&self) -> EngineConfig { let parallelism = self.server.parallelism.as_ref(); EngineConfig { diff --git a/lib/sidecar/vllm/src/tests.rs b/lib/sidecar/vllm/src/tests.rs index 8980daafbbe0..e29576249fcb 100644 --- a/lib/sidecar/vllm/src/tests.rs +++ b/lib/sidecar/vllm/src/tests.rs @@ -1,7 +1,7 @@ // SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -use std::collections::BTreeSet; +use std::collections::{BTreeMap, BTreeSet}; use std::net::SocketAddr; use std::num::NonZeroUsize; use std::pin::Pin; @@ -45,6 +45,19 @@ struct FakeVllm { first_token_pending: Arc, release_first_token: Arc, server_stream_dropped: Arc, + control_calls: Arc>>, + paused: Arc, + sleeping_tags: Arc>>, + weight_version: Arc>, +} + +impl FakeVllm { + async fn record_control(&self, name: &str, body: serde_json::Value) { + self.control_calls + .lock() + .await + .push((name.to_string(), body)); + } } struct DropSignal(Arc); @@ -259,100 +272,154 @@ impl pb::control_server::Control for FakeVllm { async fn pause_generation( &self, - _request: Request, + request: Request, ) -> Result, Status> { - Err(rl_control_unavailable()) + let request = request.into_inner(); + self.record_control( + "pause_generation", + json!({"mode": request.mode, "clear_cache": request.clear_cache}), + ) + .await; + self.paused.store(true, Ordering::SeqCst); + Ok(Response::new(pb::PauseGenerationResponse {})) } async fn resume_generation( &self, _request: Request, ) -> Result, Status> { - Err(rl_control_unavailable()) + self.record_control("resume_generation", json!({})).await; + self.paused.store(false, Ordering::SeqCst); + Ok(Response::new(pb::ResumeGenerationResponse {})) } async fn is_paused( &self, _request: Request, ) -> Result, Status> { - Err(rl_control_unavailable()) + Ok(Response::new(pb::IsPausedResponse { + paused: self.paused.load(Ordering::SeqCst), + })) } async fn sleep( &self, - _request: Request, + request: Request, ) -> Result, Status> { - Err(rl_control_unavailable()) + let request = request.into_inner(); + self.record_control( + "sleep", + json!({"level": request.level, "mode": request.mode}), + ) + .await; + let mut sleeping_tags = self.sleeping_tags.lock().await; + *sleeping_tags = if request.level == Some(0) { + BTreeSet::from(["scheduling".to_string()]) + } else { + BTreeSet::from(["kv_cache".to_string(), "weights".to_string()]) + }; + Ok(Response::new(pb::SleepResponse {})) } async fn wake_up( &self, - _request: Request, + request: Request, ) -> Result, Status> { - Err(rl_control_unavailable()) + let tags = request.into_inner().tags; + self.record_control("wake_up", json!({"tags": tags.clone()})) + .await; + let mut sleeping_tags = self.sleeping_tags.lock().await; + if tags.is_empty() { + sleeping_tags.clear(); + } else { + for tag in tags { + sleeping_tags.remove(&tag); + } + } + Ok(Response::new(pb::WakeUpResponse {})) } async fn is_sleeping( &self, _request: Request, ) -> Result, Status> { - Err(rl_control_unavailable()) + Ok(Response::new(pb::IsSleepingResponse { + sleeping: !self.sleeping_tags.lock().await.is_empty(), + })) } async fn init_weight_transfer_engine( &self, - _request: Request, + request: Request, ) -> Result, Status> { - Err(rl_control_unavailable()) + let body = serde_json::from_slice(&request.into_inner().init_info_json) + .map_err(|error| Status::invalid_argument(error.to_string()))?; + self.record_control("init_weight_transfer_engine", body) + .await; + Ok(Response::new(pb::InitWeightTransferEngineResponse {})) } async fn start_weight_update( &self, _request: Request, ) -> Result, Status> { - Err(rl_control_unavailable()) + self.record_control("start_weight_update", json!({})).await; + Ok(Response::new(pb::StartWeightUpdateResponse {})) } async fn start_draft_weight_update( &self, _request: Request, ) -> Result, Status> { - Err(rl_control_unavailable()) + self.record_control("start_draft_weight_update", json!({})) + .await; + Ok(Response::new(pb::StartDraftWeightUpdateResponse {})) } async fn update_weights( &self, - _request: Request, + request: Request, ) -> Result, Status> { - Err(rl_control_unavailable()) + let body = serde_json::from_slice(&request.into_inner().update_info_json) + .map_err(|error| Status::invalid_argument(error.to_string()))?; + self.record_control("update_weights", body).await; + Ok(Response::new(pb::UpdateWeightsResponse {})) } async fn finish_weight_update( &self, - _request: Request, + request: Request, ) -> Result, Status> { - Err(rl_control_unavailable()) + let version = request.into_inner().weight_version; + if let Some(version) = &version { + self.weight_version.lock().await.clone_from(version); + } + self.record_control("finish_weight_update", json!({"weight_version": version})) + .await; + Ok(Response::new(pb::FinishWeightUpdateResponse {})) } async fn update_weight_version( &self, - _request: Request, + request: Request, ) -> Result, Status> { - Err(rl_control_unavailable()) + let version = request.into_inner().weight_version; + self.weight_version.lock().await.clone_from(&version); + self.record_control("update_weight_version", json!({"weight_version": version})) + .await; + Ok(Response::new(pb::UpdateWeightVersionResponse {})) } async fn get_weight_version( &self, _request: Request, ) -> Result, Status> { - Err(rl_control_unavailable()) + Ok(Response::new(pb::GetWeightVersionResponse { + weight_version: self.weight_version.lock().await.clone(), + })) } } -fn rl_control_unavailable() -> Status { - Status::unimplemented("RL control RPCs are not implemented in this protocol-only test server") -} - fn model_info() -> pb::ModelInfo { pb::ModelInfo { model_id: "model-source".to_string(), @@ -383,7 +450,12 @@ fn server_info() -> pb::ServerInfo { total_kv_blocks: 4096, max_running_requests: 128, max_batched_tokens: 2048, - rl_capabilities: None, + rl_capabilities: Some(pb::RlCapabilities { + weight_transfer_enabled: true, + weight_transfer_backend: "nccl".to_string(), + sleep_mode_enabled: true, + draft_weight_updates_enabled: true, + }), } } @@ -659,6 +731,16 @@ fn engine( mode: DisaggregationMode, connections: usize, model: pb::ModelInfo, +) -> VllmSidecarEngine { + engine_with_server_info(endpoint, mode, connections, model, server_info()) +} + +fn engine_with_server_info( + endpoint: &str, + mode: DisaggregationMode, + connections: usize, + model: pb::ModelInfo, + server: pb::ServerInfo, ) -> VllmSidecarEngine { let transport = GrpcTransportConfig { connections: NonZeroUsize::new(connections).expect("non-zero connection count"), @@ -666,7 +748,7 @@ fn engine( }; VllmSidecarEngine::new( GrpcEndpoint::parse(endpoint, "--vllm-endpoint").expect("valid test endpoint"), - DiscoveredModel::from_proto(model, server_info()).expect("valid discovery"), + DiscoveredModel::from_proto(model, server).expect("valid discovery"), mode, transport, ) @@ -849,6 +931,235 @@ async fn aggregated_generation_converts_request_stream_and_usage() { ); } +#[tokio::test] +async fn rl_engine_routes_preserve_lifecycle_payloads_and_version() { + let server = FakeServer::start(FakeVllm::default()).await; + let engine = engine( + &server.endpoint, + DisaggregationMode::Aggregated, + 1, + model_info(), + ); + engine.start(0).await.expect("start"); + + assert_eq!( + engine + .supported_controls() + .await + .unwrap() + .into_iter() + .collect::>(), + [ + "get_weight_version", + "is_paused", + "is_sleeping", + "pause_generation", + "resume_generation", + "sleep", + "wake_up", + ] + .into_iter() + .map(String::from) + .collect() + ); + assert_eq!( + engine + .supported_updates() + .await + .unwrap() + .into_iter() + .collect::>(), + [ + "finish_weight_update", + "init_weight_transfer_engine", + "start_draft_weight_update", + "start_weight_update", + "update_weight_version", + "update_weights", + ] + .into_iter() + .map(String::from) + .collect() + ); + + // Regression: unsupported sleep levels or wake tags can enter vLLM's + // destructive/partial sleep paths while still returning gRPC success. + for (control, body, expected) in [ + ("sleep", json!({"level": 3}), "one of 0, 1, or 2"), + ( + "wake_up", + json!({"tags": ["unknown"]}), + "weights, kv_cache, or scheduling", + ), + ] { + let error = engine + .engine_control(control.to_string(), body) + .await + .expect_err("unsupported lifecycle value must fail before gRPC"); + assert!(error.to_string().contains(expected), "unexpected {error}"); + } + + for (control, body, expected) in [ + ("is_paused", json!({}), json!({"is_paused": false})), + ( + "pause_generation", + json!({"mode": "keep", "clear_cache": false}), + json!({"status": "paused"}), + ), + ("is_paused", json!({}), json!({"is_paused": true})), + ("resume_generation", json!({}), json!({"status": "resumed"})), + ("is_paused", json!({}), json!({"is_paused": false})), + ("is_sleeping", json!({}), json!({"is_sleeping": false})), + ( + "sleep", + json!({"level": 2, "mode": "wait"}), + json!({"status": "sleeping"}), + ), + ("is_sleeping", json!({}), json!({"is_sleeping": true})), + ( + "wake_up", + json!({"tags": ["weights"]}), + json!({"status": "partially_awake", "is_sleeping": true}), + ), + ("is_sleeping", json!({}), json!({"is_sleeping": true})), + ] { + assert_eq!( + engine + .engine_control(control.to_string(), body) + .await + .unwrap(), + expected, + "unexpected {control} response" + ); + } + + for (update, body, expected) in [ + ( + "init_weight_transfer_engine", + json!({"init_info": {"master_addr": "trainer", "master_port": 1234}}), + json!({"message": "Weight transfer initialized"}), + ), + ( + "start_weight_update", + json!({}), + json!({"message": "Weight update started"}), + ), + ( + "start_draft_weight_update", + json!({}), + json!({"message": "Draft weight update started"}), + ), + ( + "update_weights", + json!({"update_info": {"names": ["layer.weight"], "shape": [4, 8]}}), + json!({"message": "Weights updated"}), + ), + ( + "finish_weight_update", + json!({"weight_version": "step-42"}), + json!({"message": "Weight update finished"}), + ), + ] { + assert_eq!( + engine + .engine_update(update.to_string(), body) + .await + .unwrap(), + expected, + "unexpected {update} response" + ); + } + assert_eq!( + engine + .engine_control("get_weight_version".to_string(), json!({})) + .await + .unwrap(), + json!({"weight_version": "step-42"}) + ); + assert_eq!( + engine + .engine_update( + "update_weight_version".to_string(), + json!({"new_version": "step-43"}), + ) + .await + .unwrap(), + json!({"success": true, "new_version": "step-43"}) + ); + assert_eq!( + engine + .engine_control("get_weight_version".to_string(), json!({})) + .await + .unwrap(), + json!({"weight_version": "step-43"}) + ); + + let calls = server.service.control_calls.lock().await; + let actual = calls.iter().cloned().collect::>(); + let expected = BTreeMap::from([ + ( + "pause_generation".to_string(), + json!({"mode": pb::PauseMode::Keep as i32, "clear_cache": false}), + ), + ("resume_generation".to_string(), json!({})), + ( + "sleep".to_string(), + json!({"level": 2, "mode": pb::PauseMode::Wait as i32}), + ), + ("wake_up".to_string(), json!({"tags": ["weights"]})), + ( + "init_weight_transfer_engine".to_string(), + json!({"master_addr": "trainer", "master_port": 1234}), + ), + ("start_weight_update".to_string(), json!({})), + ("start_draft_weight_update".to_string(), json!({})), + ( + "update_weights".to_string(), + json!({"names": ["layer.weight"], "shape": [4, 8]}), + ), + ( + "finish_weight_update".to_string(), + json!({"weight_version": "step-42"}), + ), + ( + "update_weight_version".to_string(), + json!({"weight_version": "step-43"}), + ), + ]); + assert_eq!(calls.len(), expected.len(), "each mutating RPC runs once"); + assert_eq!(actual, expected); +} + +/// Regression: vLLM exposes sleep status independently of CUDA sleep-mode +/// allocation support, so capability discovery must not hide the status RPC +/// when only the mutating sleep/wake operations are disabled. +#[tokio::test] +async fn sleep_status_remains_advertised_without_sleep_mode() { + let mut server = server_info(); + server + .rl_capabilities + .as_mut() + .expect("RL capabilities") + .sleep_mode_enabled = false; + let engine = engine_with_server_info( + "http://127.0.0.1:1", + DisaggregationMode::Aggregated, + 1, + model_info(), + server, + ); + + let controls = engine + .supported_controls() + .await + .unwrap() + .into_iter() + .collect::>(); + assert!(controls.contains("is_sleeping")); + assert!(!controls.contains("sleep")); + assert!(!controls.contains("wake_up")); +} + #[tokio::test] async fn multimodal_image_is_forwarded_with_uuid() { let service = FakeVllm::default();