diff --git a/rust/src/engine-core-client/src/mock_engine.rs b/rust/src/engine-core-client/src/mock_engine.rs index 32cd48c396f..11c012b1f16 100644 --- a/rust/src/engine-core-client/src/mock_engine.rs +++ b/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, } } diff --git a/rust/src/engine-core-client/src/protocol/handshake.rs b/rust/src/engine-core-client/src/protocol/handshake.rs index d659dc8a244..3ca8774b2d6 100644 --- a/rust/src/engine-core-client/src/protocol/handshake.rs +++ b/rust/src/engine-core-client/src/protocol/handshake.rs @@ -28,7 +28,7 @@ pub struct ReadyMessage { /// profiling). /// /// Original Python definition: -/// +/// #[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, /// 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, + /// Maximum achievable request concurrency given the KV cache, if reported. + pub kv_cache_max_concurrency: Option, } /// Frontend-owned ZMQ addresses that are sent to the engine during startup diff --git a/rust/src/engine-core-client/src/tests/client.rs b/rust/src/engine-core-client/src/tests/client.rs index 32530d6d385..322eebfd83d 100644 --- a/rust/src/engine-core-client/src/tests/client.rs +++ b/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 { + 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)] diff --git a/rust/src/engine-core-client/src/tests/python_compat.py b/rust/src/engine-core-client/src/tests/python_compat.py index bb81a6df1ad..89179b3fbfe 100755 --- a/rust/src/engine-core-client/src/tests/python_compat.py +++ b/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())