[Bugfix][Rust] Sync EngineCoreReadyResponse with the Python dataclass (#45557)

Co-authored-by: Bugen Zhao <i@bugenzhao.com>
Signed-off-by: Will Eaton <weaton@redhat.com>
Signed-off-by: Bugen Zhao <i@bugenzhao.com>

rust/src/engine-core-client/src/mock_engine.rs

@@ -15,6 +15,8 @@ use crate::protocol::{ModelDtype, decode_msgpack, encode_msgpack};
 pub const DEFAULT_MOCK_MAX_MODEL_LEN: u64 = 1024 * 1024;
 /// Default KV block count advertised by reusable mock engine helpers.
 pub const DEFAULT_MOCK_NUM_GPU_BLOCKS: u64 = 0;
+/// Default KV block size (tokens per block)
+pub const DEFAULT_MOCK_BLOCK_SIZE: u64 = 16;
 
 /// Startup behavior for one mock engine joining a frontend.
 #[derive(Debug, Clone)]
@@ -46,9 +48,12 @@ pub fn default_ready_response() -> EngineCoreReadyResponse {
     EngineCoreReadyResponse {
         max_model_len: DEFAULT_MOCK_MAX_MODEL_LEN,
         num_gpu_blocks: DEFAULT_MOCK_NUM_GPU_BLOCKS,
+        block_size: DEFAULT_MOCK_BLOCK_SIZE,
         dp_stats_address: None,
         dtype: ModelDtype::Float32,
         vllm_version: "test-vllm-version".to_string(),
+        kv_cache_size_tokens: None,
+        kv_cache_max_concurrency: None,
     }
 }
 

rust/src/engine-core-client/src/protocol/handshake.rs

@@ -28,7 +28,7 @@ pub struct ReadyMessage {
 /// profiling).
 ///
 /// Original Python definition:
-/// <https://github.com/vllm-project/vllm/blob/c8d98f81f6/vllm/v1/engine/__init__.py#L67-L77>
+/// <https://github.com/vllm-project/vllm/blob/c9340e6f35/vllm/v1/engine/__init__.py#L68-L80>
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct EngineCoreReadyResponse {
     /// Engine-reported maximum model context length (auto-fitted after
@@ -36,12 +36,18 @@ pub struct EngineCoreReadyResponse {
     pub max_model_len: u64,
     /// Number of GPU blocks available for KV cache on this engine.
     pub num_gpu_blocks: u64,
+    /// KV cache block size (tokens per block).
+    pub block_size: u64,
     /// DP coordinator stats publish address, if applicable.
     pub dp_stats_address: Option<String>,
     /// Effective model dtype after Python vLLM resolves `--dtype`.
     pub dtype: ModelDtype,
     /// Python vLLM version reported by the engine process.
     pub vllm_version: String,
+    /// Total KV cache capacity in tokens, if reported.
+    pub kv_cache_size_tokens: Option<u64>,
+    /// Maximum achievable request concurrency given the KV cache, if reported.
+    pub kv_cache_max_concurrency: Option<f64>,
 }
 
 /// Frontend-owned ZMQ addresses that are sent to the engine during startup

rust/src/engine-core-client/src/tests/client.rs

@@ -2445,6 +2445,7 @@ fn python_msgpack_fixtures_match_rust_encoding() {
     let inline_prompt_frames = lines.next().expect("missing inline prompt logprobs fixture line");
     let multipart_prompt_frames =
         lines.next().expect("missing multipart prompt logprobs fixture line");
+    let ready_response_hex = lines.next().expect("missing ready response fixture line");
 
     let request_bytes = hex::decode(request_hex).unwrap();
     let multimodal_request_bytes = hex::decode(multimodal_request_hex).unwrap();
@@ -2554,6 +2555,23 @@ fn python_msgpack_fixtures_match_rust_encoding() {
             .as_ref()
             .expect("multipart prompt logprobs decoded"),
     );
+
+    let map_keys = |bytes: &[u8]| -> BTreeSet<String> {
+        match decode_value(bytes) {
+            Value::Map(entries) => entries
+                .into_iter()
+                .filter_map(|(key, _)| key.as_str().map(str::to_owned))
+                .collect(),
+            other => panic!("ready response should encode as a map, got {other:?}"),
+        }
+    };
+    let python_ready_keys = map_keys(&hex::decode(ready_response_hex).unwrap());
+    let rust_ready_keys =
+        map_keys(&rmp_serde::to_vec_named(&crate::mock_engine::default_ready_response()).unwrap());
+    assert_eq!(
+        rust_ready_keys, python_ready_keys,
+        "EngineCoreReadyResponse drifted from the Python dataclass",
+    );
 }
 
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]

rust/src/engine-core-client/src/tests/python_compat.py

@@ -10,6 +10,7 @@
 # ]
 # ///
 
+from dataclasses import dataclass
 from enum import Enum, IntEnum
 
 import msgpack
@@ -337,6 +338,28 @@ multipart_prompt_logprobs = engine_outputs_wire(
     )
 )
 
+
+@dataclass
+class EngineCoreReadyResponse:
+    max_model_len: int
+    num_gpu_blocks: int
+    block_size: int
+    dp_stats_address: str | None
+    dtype: str
+    vllm_version: str
+    kv_cache_size_tokens: int | None = None
+    kv_cache_max_concurrency: float | None = None
+
+
+ready_response = EngineCoreReadyResponse(
+    max_model_len=32768,
+    num_gpu_blocks=1000,
+    block_size=16,
+    dp_stats_address=None,
+    dtype="float32",
+    vllm_version="0.0.0",
+)
+
 print(msgspec.msgpack.encode(request).hex())
 print(msgpack.packb(multimodal_request_wire, use_bin_type=True).hex())
 print(msgspec.msgpack.encode(outputs).hex())
@@ -354,3 +377,4 @@ print(
         for frame in encode_output_frames(multipart_prompt_logprobs, size_threshold=1)
     )
 )
+print(msgspec.msgpack.encode(ready_response).hex())