refactor: separate python and rust in vllm connector worker, similar to leader

Author: ziqifan617

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

@@ -47,7 +47,7 @@ fn _vllm_integration(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<BlockStates>()?;
     m.add_class::<SlotUpdate>()?;
 
-    m.add_class::<connector::worker::KvConnectorWorker>()?;
+    m.add_class::<connector::worker::PyKvConnectorWorker>()?;
     m.add_class::<connector::leader::PyKvConnectorLeader>()?;
     m.add_class::<connector::SchedulerOutput>()?;
     Ok(())

lib/bindings/python/rust/llm/block_manager/vllm/connector/worker.rs

@@ -21,7 +21,28 @@ use dynamo_llm::block_manager::storage::torch::TorchTensor;
 use dynamo_runtime::utils::task::CriticalTaskExecutionHandle;
 use dynamo_runtime::DistributedRuntime;
 
-#[pyclass]
+pub trait Worker: Send + Sync {
+    fn register_kv_caches(
+        &mut self,
+        num_device_blocks: usize,
+        page_size: usize,
+        device_id: usize,
+        dtype_width_bytes: usize,
+        kv_caches: Vec<(String, Py<PyAny>)>,
+    ) -> PyResult<()>;
+
+    fn bind_connector_metadata(&mut self, metadata: Vec<u8>) -> PyResult<()>;
+
+    fn clear_connector_metadata(&mut self);
+
+    fn save_kv_layer(&mut self, layer_name: String, kv_layer: Py<PyAny>) -> PyResult<()>;
+
+    fn get_finished(
+        &mut self,
+        finished_requests: HashSet<String>,
+    ) -> (HashSet<String>, HashSet<String>);
+}
+
 pub struct KvConnectorWorker {
     drt: DistributedRuntime,
     kvbm_worker: OnceLock<KvbmWorker>,
@@ -45,9 +66,7 @@ pub struct KvConnectorWorker {
     layers_complete: usize,
 }
 
-#[pymethods]
 impl KvConnectorWorker {
-    #[new]
     fn new(py_drt: PyDistributedRuntime, vllm_worker_id: String) -> PyResult<Self> {
         let drt = py_drt.inner.clone();
         let runtime = drt.runtime().primary();
@@ -85,12 +104,14 @@ impl KvConnectorWorker {
             layers_complete: 0,
         })
     }
+}
 
+impl Worker for KvConnectorWorker {
     /// Registers the KV caches with the KVBM worker.
     ///
     /// The Dynamo KVBM worker is lazily initialized when the first KV cache is registered.
     /// This process establishes a connection between all KVBM workers and the leader.
-    pub fn register_kv_caches(
+    fn register_kv_caches(
         &mut self,
         num_device_blocks: usize,
         page_size: usize,
@@ -148,7 +169,7 @@ impl KvConnectorWorker {
     /// Loads the metadata from the leader.
     /// This action translates the metadata into a set of actions that the worker will perform.
     /// All actions much be assigned to a slot before [`KvConnectorWorker::clear_metadata`] is called.
-    pub fn bind_connector_metadata(&mut self, metadata: Vec<u8>) -> PyResult<()> {
+    fn bind_connector_metadata(&mut self, metadata: Vec<u8>) -> PyResult<()> {
         debug_assert!(!self.bound, "connector metadata already bound");
         let metadata: ConnectorMetadata = serde_json::from_slice(&metadata).map_err(to_pyerr)?;
         self.bound = true;
@@ -214,7 +235,7 @@ impl KvConnectorWorker {
     }
 
     /// Clears the connector metadata and marks the iteration as complete.
-    pub fn clear_connector_metadata(&mut self) {
+    fn clear_connector_metadata(&mut self) {
         tracing::debug!(iteration = self.iteration, "clearing connector metadata");
         debug_assert!(self.bound, "connector metadata not bound");
         self.bound = false;
@@ -227,7 +248,7 @@ impl KvConnectorWorker {
 
     /// Trigger layer-wise completion signals.
     /// Trigger block-wise completion signals afer last layer.
-    pub fn save_kv_layer(&mut self, _layer_name: String, _kv_layer: Py<PyAny>) -> PyResult<()> {
+    fn save_kv_layer(&mut self, layer_name: String, kv_layer: Py<PyAny>) -> PyResult<()> {
         self.layers_complete += 1;
         if self.layers_complete == self.kv_caches.len() {
             let offloading_operations = std::mem::take(&mut self.offloading_operations);
@@ -238,7 +259,7 @@ impl KvConnectorWorker {
         Ok(())
     }
 
-    pub fn get_finished(
+    fn get_finished(
         &mut self,
         finished_requests: HashSet<String>,
     ) -> (HashSet<String>, HashSet<String>) {
@@ -321,3 +342,48 @@ impl KvConnectorWorker {
         (is_finished_offloading, is_finished_onboarding)
     }
 }
+
+#[pyclass]
+pub struct PyKvConnectorWorker {
+    connector_worker: Box<dyn Worker>,
+}
+
+#[pymethods]
+impl PyKvConnectorWorker {
+    #[new]
+    #[pyo3(signature = (py_drt, vllm_worker_id))]
+    pub fn new(py_drt: PyDistributedRuntime, vllm_worker_id: String) -> PyResult<Self> {
+        let connector_worker: Box<dyn Worker> = Box::new(KvConnectorWorker::new(py_drt, vllm_worker_id)?);
+        Ok(Self { connector_worker })
+    }
+
+    pub fn register_kv_caches(
+        &mut self,
+        num_device_blocks: usize,
+        page_size: usize,
+        device_id: usize,
+        dtype_width_bytes: usize,
+        kv_caches: Vec<(String, Py<PyAny>)>,
+    ) -> PyResult<()> {
+        self.connector_worker.register_kv_caches(num_device_blocks, page_size, device_id, dtype_width_bytes, kv_caches)
+    }
+
+    pub fn bind_connector_metadata(&mut self, metadata: Vec<u8>) -> PyResult<()> {
+        self.connector_worker.bind_connector_metadata(metadata)
+    }
+
+    pub fn clear_connector_metadata(&mut self) {
+        self.connector_worker.clear_connector_metadata()
+    }
+
+    pub fn save_kv_layer(&mut self, layer_name: String, kv_layer: Py<PyAny>) -> PyResult<()> {
+        self.connector_worker.save_kv_layer(layer_name, kv_layer)
+    }
+
+    pub fn get_finished(
+        &mut self,
+        finished_requests: HashSet<String>,
+    ) -> (HashSet<String>, HashSet<String>) {
+        self.connector_worker.get_finished(finished_requests)
+    }
+}

lib/bindings/python/src/dynamo/llm/vllm_integration/rust.py

@@ -16,7 +16,7 @@
     BlockStates = getattr(_vllm_integration, "BlockStates")
     SlotUpdate = getattr(_vllm_integration, "SlotUpdate")
 
-    KvConnectorWorker = getattr(_vllm_integration, "KvConnectorWorker")
+    KvConnectorWorker = getattr(_vllm_integration, "PyKvConnectorWorker")
     KvConnectorLeader = getattr(_vllm_integration, "PyKvConnectorLeader")
     SchedulerOutput = getattr(_vllm_integration, "SchedulerOutput")