refactor: kvbm modularity DIS-657 Eliminate ETCD from the leader-worker initialization (#3202)

Signed-off-by: richardhuo-nv <rihuo@nvidia.com>

Author: richardhuo-nv

Cargo.lock

@@ -54,10 +54,6 @@ spec:
       envs:
         - name: DYN_KVBM_CPU_CACHE_GB
           value: "100"
-        - name: DYN_KVBM_BARRIER_ID_PREFIX
-          valueFrom:
-            fieldRef:
-              fieldPath: metadata.name
       extraPodSpec:
         mainContainer:
           image: nvcr.io/nvidia/ai-dynamo/vllm-runtime:my-tag

components/backends/vllm/deploy/disagg_kvbm_2p2d.yaml

@@ -58,10 +58,6 @@ spec:
       envs:
         - name: DYN_KVBM_CPU_CACHE_GB
           value: "100"
-        - name: DYN_KVBM_BARRIER_ID_PREFIX
-          valueFrom:
-            fieldRef:
-              fieldPath: metadata.name
       extraPodSpec:
         mainContainer:
           image: nvcr.io/nvidia/ai-dynamo/vllm-runtime:my-tag

components/backends/vllm/deploy/disagg_kvbm_tp2.yaml

@@ -15,7 +15,6 @@ CUDA_VISIBLE_DEVICES=1 python3 -m dynamo.vllm --model Qwen/Qwen3-0.6B --connecto
 # run prefill workers on GPU 2 and 3 with KVBM enabled using 20GB of CPU cache
 # NOTE: use different barrier id prefixes for each prefill worker to avoid conflicts
 # NOTE: remove --enforce-eager for production use
-DYN_KVBM_BARRIER_ID_PREFIX=kvbm_0 \
 DYN_KVBM_CPU_CACHE_GB=20 \
 CUDA_VISIBLE_DEVICES=2 \
   python3 -m dynamo.vllm \
@@ -24,7 +23,8 @@ CUDA_VISIBLE_DEVICES=2 \
     --connector kvbm nixl \
     --enforce-eager &
 
-DYN_KVBM_BARRIER_ID_PREFIX=kvbm_1 \
+DYN_KVBM_LEADER_ZMQ_PUB_PORT=56003 \
+DYN_KVBM_LEADER_ZMQ_ACK_PORT=56004 \
 DYN_KVBM_CPU_CACHE_GB=20 \
 CUDA_VISIBLE_DEVICES=3 \
   python3 -m dynamo.vllm \

components/backends/vllm/launch/disagg_kvbm_2p2d.sh

@@ -8,5 +8,5 @@ mod utils;
 mod worker;
 
 pub use leader::KvbmLeader;
-pub use utils::get_barrier_id_prefix;
+pub use utils::{get_leader_zmq_ack_url, get_leader_zmq_pub_url};
 pub use worker::{KvbmWorker, PyLayoutType, VllmTensor};

lib/bindings/python/Cargo.lock

@@ -2,12 +2,11 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use super::*;
-use utils::get_barrier_id_prefix;
-
 use derive_getters::Dissolve;
 use llm_rs::block_manager::distributed::{
     KvbmLeader as KvbmLeaderImpl, KvbmLeaderConfig, KvbmLeaderNumBlocksConfig,
 };
+use utils::{get_leader_zmq_ack_url, get_leader_zmq_pub_url};
 
 const CPU_CACHE: &str = "DYN_KVBM_CPU_CACHE_GB";
 const CPU_CACHE_OVERRIDE: &str = "DYN_KVBM_CPU_CACHE_OVERRIDE_NUM_BLOCKS";
@@ -72,17 +71,16 @@ impl KvbmLeader {
     #[new]
     #[pyo3(signature = (world_size, drt))]
     fn new(world_size: usize, drt: DistributedRuntime) -> PyResult<Self> {
-        let barrier_id_prefix = get_barrier_id_prefix();
         let leader_init_timeout_sec: u64 =
             get_leader_init_timeout_secs(LEADER_WORKER_INIT_TIMEOUT_SECS);
 
         let config = KvbmLeaderConfig::builder()
-            .barrier_id_prefix(barrier_id_prefix)
             .world_size(world_size)
             .leader_init_timeout_secs(leader_init_timeout_sec)
-            .drt(drt.inner().clone())
             .host_blocks_config(get_blocks_config(CPU_CACHE, CPU_CACHE_OVERRIDE))
             .disk_blocks_config(get_blocks_config(DISK_CACHE, DISK_CACHE_OVERRIDE))
+            .leader_pub_url(get_leader_zmq_pub_url())
+            .leader_ack_url(get_leader_zmq_ack_url())
             .build()
             .map_err(to_pyerr)?;
 

lib/bindings/python/rust/llm/block_manager/distributed.rs

@@ -1,9 +1,64 @@
 // SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 // SPDX-License-Identifier: Apache-2.0
+use std::env;
 
-pub fn get_barrier_id_prefix() -> String {
-    std::env::var("DYN_KVBM_BARRIER_ID_PREFIX")
+const DEFAULT_LEADER_ZMQ_HOST: &str = "127.0.0.1";
+const DEFAULT_LEADER_ZMQ_PUB_PORT: u16 = 56001;
+const DEFAULT_LEADER_ZMQ_ACK_PORT: u16 = 56002;
+
+fn read_env_trimmed(key: &str) -> Option<String> {
+    env::var(key)
         .ok()
-        .filter(|s| !s.trim().is_empty())
-        .unwrap_or_else(|| "kvbm".to_string())
+        .map(|s| s.trim().to_string())
+        .filter(|s| !s.is_empty())
+}
+
+fn parse_port_u16(s: &str) -> Option<u16> {
+    match s.parse::<u32>() {
+        Ok(v) if (1..=65535).contains(&v) => Some(v as u16),
+        _ => None,
+    }
+}
+
+fn validated_port_from_env(key: &str, default_port: u16) -> u16 {
+    if let Some(val) = read_env_trimmed(key) {
+        if let Some(p) = parse_port_u16(&val) {
+            if p < 1024 {
+                tracing::warn!("{key} is a privileged port ({p}); binding may require extra caps");
+            }
+            return p;
+        } else {
+            tracing::warn!("{key} invalid value '{val}', falling back to default {default_port}");
+        }
+    }
+    default_port
+}
+
+fn get_leader_zmq_host() -> String {
+    read_env_trimmed("DYN_KVBM_LEADER_ZMQ_HOST")
+        .unwrap_or_else(|| DEFAULT_LEADER_ZMQ_HOST.to_string())
+}
+
+fn get_leader_zmq_pub_port() -> String {
+    validated_port_from_env("DYN_KVBM_LEADER_ZMQ_PUB_PORT", DEFAULT_LEADER_ZMQ_PUB_PORT).to_string()
+}
+
+fn get_leader_zmq_ack_port() -> String {
+    validated_port_from_env("DYN_KVBM_LEADER_ZMQ_ACK_PORT", DEFAULT_LEADER_ZMQ_ACK_PORT).to_string()
+}
+
+pub fn get_leader_zmq_pub_url() -> String {
+    format!(
+        "tcp://{}:{}",
+        get_leader_zmq_host(),
+        get_leader_zmq_pub_port()
+    )
+}
+
+pub fn get_leader_zmq_ack_url() -> String {
+    format!(
+        "tcp://{}:{}",
+        get_leader_zmq_host(),
+        get_leader_zmq_ack_port()
+    )
 }

lib/bindings/python/rust/llm/block_manager/distributed/leader.rs

@@ -1,10 +1,11 @@
 // SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 // SPDX-License-Identifier: Apache-2.0
 
+use utils::{get_leader_zmq_ack_url, get_leader_zmq_pub_url};
+
 use super::*;
 
 use std::sync::Arc;
-use utils::get_barrier_id_prefix;
 
 use llm_rs::block_manager::distributed::{
     BlockTransferHandler as RustBlockTransferHandler, KvbmWorker as KvbmWorkerImpl,
@@ -171,16 +172,13 @@ impl KvbmWorker {
             vllm_tensors.push(Arc::new(vllm_tensor));
         }
 
-        let barrier_id_prefix = get_barrier_id_prefix();
-
         let config = KvbmWorkerConfig::builder()
             .drt(drt)
             .num_device_blocks(num_device_blocks)
             .page_size(page_size)
             .tensors(vllm_tensors)
             .device_id(device_id)
             .dtype_width_bytes(dtype_width_bytes)
-            .barrier_id_prefix(barrier_id_prefix)
             .device_layout_type(
                 device_layout_type
                     .map(|py_layout| py_layout.into())
@@ -196,6 +194,8 @@ impl KvbmWorker {
                     .map(|py_layout| py_layout.into())
                     .unwrap_or(LayoutType::FullyContiguous),
             )
+            .leader_pub_url(get_leader_zmq_pub_url())
+            .leader_ack_url(get_leader_zmq_ack_url())
             .build()
             .map_err(to_pyerr)?;
 

lib/bindings/python/rust/llm/block_manager/distributed/utils.rs

@@ -150,9 +150,6 @@ impl KvConnectorLeader {
 
                 let _ = slot_manager_cell.set(sm);
 
-                // another barrier sync to make sure worker init won't return before leader is ready
-                let _ = leader.run_leader_readiness_barrier_blocking(drt);
-
                 if leader_ready_tx.send("finished".to_string()).is_err() {
                     tracing::error!("main routine receiver dropped before result was sent");
                 }