[Rust Frontend][gRPC] Advertise LoRA capabilities (#52031)

Signed-off-by: Connor Carpenter <connorc@nvidia.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

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

@@ -286,6 +286,7 @@ impl EngineCoreClient {
         config: EngineCoreClientConfig,
         connected: transport::ConnectedTransport,
     ) -> Result<Self> {
+        validate_lora_capabilities(&connected.engines)?;
         let (output_tx, output_rx) = mpsc::channel(64);
         let (abort_tx, abort_rx) = mpsc::unbounded_channel();
         let engines = connected.engines;
@@ -476,6 +477,32 @@ impl EngineCoreClient {
     }
 }
 
+fn validate_lora_capabilities(engines: &[ConnectedEngine]) -> Result<()> {
+    let first = engines.first().expect("engine core client requires at least one engine");
+    for engine in engines {
+        let ready = &engine.ready_response;
+        if ready.supports_lora != (ready.max_loras > 0) {
+            return Err(Error::UnexpectedHandshakeMessage {
+                message: format!(
+                    "engine {:?} reported inconsistent LoRA capability (supports_lora={}, max_loras={})",
+                    engine.engine_id, ready.supports_lora, ready.max_loras
+                ),
+            });
+        }
+        if ready.supports_lora != first.ready_response.supports_lora
+            || ready.max_loras != first.ready_response.max_loras
+        {
+            return Err(Error::UnexpectedHandshakeMessage {
+                message: format!(
+                    "engine {:?} reported LoRA capability inconsistent with engine {:?}",
+                    engine.engine_id, first.engine_id
+                ),
+            });
+        }
+    }
+    Ok(())
+}
+
 // Client API implementation.
 impl EngineCoreClient {
     /// Add a new request to the engine and return a per-request raw output

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

@@ -64,6 +64,8 @@ pub fn default_ready_response() -> EngineCoreReadyResponse {
         max_num_seqs: 256,
         max_num_batched_tokens: 8192,
         instance_id: "test-instance".to_string(),
+        supports_lora: false,
+        max_loras: 0,
         kv_cache_size_tokens: None,
         kv_cache_max_concurrency: None,
         kv_events_config: None,

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

@@ -80,6 +80,10 @@ pub struct EngineCoreReadyResponse {
     pub max_num_batched_tokens: u64,
     /// Unique identifier for this server instance.
     pub instance_id: String,
+    /// Whether the engine was started with LoRA support enabled.
+    pub supports_lora: bool,
+    /// Maximum number of LoRA adapters the engine may keep active.
+    pub max_loras: u32,
     /// 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.

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

@@ -2686,6 +2686,8 @@ fn python_msgpack_fixtures_match_rust_encoding() {
 
     let ready_response: EngineCoreReadyResponse =
         rmp_serde::from_slice(&hex::decode(ready_response_hex).unwrap()).unwrap();
+    assert!(ready_response.supports_lora);
+    assert_eq!(ready_response.max_loras, 8);
     assert_eq!(
         ready_response.weight_transfer_backend.as_deref(),
         Some("nccl")

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

@@ -409,6 +409,8 @@ class EngineCoreReadyResponse:
     max_num_seqs: int
     max_num_batched_tokens: int
     instance_id: str
+    supports_lora: bool
+    max_loras: int
     kv_cache_size_tokens: int | None = None
     kv_cache_max_concurrency: float | None = None
     kv_events_config: KVEventsConfig | None = None
@@ -433,6 +435,8 @@ ready_response = EngineCoreReadyResponse(
     max_num_seqs=256,
     max_num_batched_tokens=8192,
     instance_id="test-instance",
+    supports_lora=True,
+    max_loras=8,
     weight_transfer_backend="nccl",
     enable_sleep_mode=True,
     supports_draft_weight_updates=True,