refactor: separate pyo3 from rust in vllm connector worker 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/leader.rs

@@ -30,7 +30,6 @@ use std::{
 };
 use tokio::sync::mpsc;
 use tokio;
-use pyo3_async_runtimes;
 
 type VllmLocality = Logical<DistributedLeaderWorkerResources>;
 
@@ -41,6 +40,7 @@ impl From<SlotError> for PyErr {
 }
 use dynamo_llm::recorder::Recorder;
 use tokio_util::sync::CancellationToken;
+use anyhow;
 
 
 pub trait Leader: Send + Sync + std::fmt::Debug {
@@ -49,29 +49,29 @@ pub trait Leader: Send + Sync + std::fmt::Debug {
         request_id: String,
         request_num_tokens: usize,
         num_computed_tokens: usize,
-    ) -> PyResult<(usize, bool)>;
+    ) -> anyhow::Result<(usize, bool)>;
 
     fn update_state_after_alloc(
         &mut self,
         request_id: String,
         block_ids: Vec<BlockId>,
         num_external_tokens: usize,
-    ) -> PyResult<()>;
+    ) -> anyhow::Result<()>;
 
     fn build_connector_metadata(
         &mut self,
         scheduler_output: SchedulerOutput,
-    ) -> PyResult<Vec<u8>>;
+    ) -> anyhow::Result<Vec<u8>>;
 
     fn request_finished(
         &mut self,
         request_id: String,
         block_ids: Vec<BlockId>,
-    ) -> PyResult<bool>;
+    ) -> anyhow::Result<bool>;
 
     fn has_slot(&self, request_id: String) -> bool;
 
-    fn create_slot(&mut self, request: KvbmRequest, tokens: Vec<u32>) -> PyResult<()>;
+    fn create_slot(&mut self, request: KvbmRequest, tokens: Vec<u32>) -> anyhow::Result<()>;
 }
 
 #[derive(Debug)]
@@ -128,16 +128,16 @@ impl Leader for KvConnectorLeader {
         request_id: String,
         request_num_tokens: usize,
         num_computed_tokens: usize,
-    ) -> PyResult<(usize, bool)> {
+    ) -> anyhow::Result<(usize, bool)> {
         tracing::debug!(
             "request_num_tokens: {request_num_tokens}; num_computed_tokens: {num_computed_tokens}"
         );
 
         // the number of device matched tokens should be less than or equal to the number of tokens in the request
         debug_assert!(num_computed_tokens % self.block_size == 0);
 
-        let shared_slot = self.slot_manager.get_slot(&request_id).map_err(to_pyerr)?;
-        let mut slot = shared_slot.lock().map_err(to_pyerr)?;
+        let shared_slot = self.slot_manager.get_slot(&request_id)?;
+        let mut slot = shared_slot.lock().map_err(|e| anyhow::anyhow!("Failed to lock slot: {}", e))?;
 
         // vllm is telling us that the tokens have been computed, since we do not have insight into the device pool
         // we accept this and advance the computed position
@@ -179,16 +179,16 @@ impl Leader for KvConnectorLeader {
         request_id: String,
         block_ids: Vec<BlockId>,
         num_external_tokens: usize,
-    ) -> PyResult<()> {
+    ) -> anyhow::Result<()> {
         tracing::debug!(
             request_id,
             "num_device_blocks: {}; num_external_tokens: {}",
             block_ids.len(),
             num_external_tokens
         );
 
-        let shared_slot = self.slot_manager.get_slot(&request_id).map_err(to_pyerr)?;
-        let mut slot = shared_slot.lock().map_err(to_pyerr)?;
+        let shared_slot = self.slot_manager.get_slot(&request_id)?;
+        let mut slot = shared_slot.lock().map_err(|e| anyhow::anyhow!("Failed to lock slot: {}", e))?;
 
         slot.append_mutable_device_blocks(&block_ids)?;
 
@@ -211,7 +211,7 @@ impl Leader for KvConnectorLeader {
     fn build_connector_metadata(
         &mut self,
         scheduler_output: SchedulerOutput,
-    ) -> PyResult<Vec<u8>> {
+    ) -> anyhow::Result<Vec<u8>> {
         // the iteration counter is used to track the number of times we have built the connector metadata
         // all connetor operations have the iteration counter at which they were issued.
         // this allows operations to be lazily enqueued to the transfer engine
@@ -234,8 +234,8 @@ impl Leader for KvConnectorLeader {
         // This is kind of a nice abstraction as it keeps the events simplier; however, we now create the request-slot
         // once for onboarding (this loop), then again for prefill/decode (new_requests loop).
         for request_id in onboarding_slots.iter() {
-            let shared_slot = self.slot_manager.get_slot(request_id).map_err(to_pyerr)?;
-            let mut slot = shared_slot.lock().map_err(to_pyerr)?;
+            let shared_slot = self.slot_manager.get_slot(request_id)?;
+            let mut slot = shared_slot.lock().map_err(|e| anyhow::anyhow!("Failed to lock slot: {}", e))?;
 
             md.create_slot(request_id.clone());
 
@@ -256,8 +256,8 @@ impl Leader for KvConnectorLeader {
                 "request_id {request_id} not found in inflight_requests: "
             );
 
-            let shared_slot = self.slot_manager.get_slot(request_id).map_err(to_pyerr)?;
-            let mut slot = shared_slot.lock().map_err(to_pyerr)?;
+            let shared_slot = self.slot_manager.get_slot(request_id)?;
+            let mut slot = shared_slot.lock().map_err(|e| anyhow::anyhow!("Failed to lock slot: {}", e))?;
 
             // inform the worker that a new request-slot should be created
             md.create_slot(new_req.request_id.clone());
@@ -297,8 +297,8 @@ impl Leader for KvConnectorLeader {
                 "request_id {request_id} not found in inflight_requests: "
             );
 
-            let shared_slot = self.slot_manager.get_slot(request_id).map_err(to_pyerr)?;
-            let mut slot = shared_slot.lock().map_err(to_pyerr)?;
+            let shared_slot = self.slot_manager.get_slot(request_id)?;
+            let mut slot = shared_slot.lock().map_err(|e| anyhow::anyhow!("Failed to lock slot: {}", e))?;
 
             let scheduled_tokens = *scheduler_output
                 .num_scheduled_tokens
@@ -323,16 +323,16 @@ impl Leader for KvConnectorLeader {
         }
 
         tracing::debug!("scheduler_output: {scheduler_output:#?}");
-        serde_json::to_vec(&md).map_err(to_pyerr)
+        serde_json::to_vec(&md).map_err(|e| anyhow::anyhow!("Failed to serialize connector metadata: {}", e))
     }
 
-    fn request_finished(&mut self, request_id: String, block_ids: Vec<BlockId>) -> PyResult<bool> {
+    fn request_finished(&mut self, request_id: String, block_ids: Vec<BlockId>) -> anyhow::Result<bool> {
         tracing::debug!("Request finished: {request_id}; block_ids: {block_ids:?}");
         // grab the slot
-        let shared_slot = self.slot_manager.get_slot(&request_id).map_err(to_pyerr)?;
+        let shared_slot = self.slot_manager.get_slot(&request_id)?;
 
         // mark the slot as finished
-        let mut slot = shared_slot.lock().map_err(to_pyerr)?;
+        let mut slot = shared_slot.lock().map_err(|e| anyhow::anyhow!("Failed to lock slot: {}", e))?;
         slot.mark_as_finished(self.iteration_counter)?;
 
         // todo: allow the request to resolve when it should exit
@@ -361,7 +361,7 @@ impl Leader for KvConnectorLeader {
 
     /// Create a new slot for the given request ID.
     /// This is used to create a new slot for the request.
-    fn create_slot(&mut self, request: KvbmRequest, tokens: Vec<u32>) -> PyResult<()> {
+    fn create_slot(&mut self, request: KvbmRequest, tokens: Vec<u32>) -> anyhow::Result<()> {
         self.slot_manager
             .create_slot(&request.request_id, tokens, request.salt_hash)?;
 
@@ -414,7 +414,7 @@ impl PyKvConnectorLeader {
         request_num_tokens: usize,
         num_computed_tokens: usize,
     ) -> PyResult<(usize, bool)> {
-        self.connector_leader.get_num_new_matched_tokens(request_id, request_num_tokens, num_computed_tokens)
+        self.connector_leader.get_num_new_matched_tokens(request_id, request_num_tokens, num_computed_tokens).map_err(to_pyerr)
     }
 
     fn update_state_after_alloc(
@@ -423,25 +423,25 @@ impl PyKvConnectorLeader {
         block_ids: Vec<BlockId>,
         num_external_tokens: usize,
     ) -> PyResult<()> {
-        self.connector_leader.update_state_after_alloc(request_id, block_ids, num_external_tokens)
+        self.connector_leader.update_state_after_alloc(request_id, block_ids, num_external_tokens).map_err(to_pyerr)
     }
 
     fn build_connector_metadata(
         &mut self,
         scheduler_output: SchedulerOutput,
     ) -> PyResult<Vec<u8>> {
-        self.connector_leader.build_connector_metadata(scheduler_output)
+        self.connector_leader.build_connector_metadata(scheduler_output).map_err(to_pyerr)
     }
 
     fn request_finished(&mut self, request_id: &str, block_ids: Vec<BlockId>) -> PyResult<bool> {
-        self.connector_leader.request_finished(request_id.to_string(), block_ids)
+        self.connector_leader.request_finished(request_id.to_string(), block_ids).map_err(to_pyerr)
     }
 
     fn has_slot(&self, request_id: &str) -> bool {
         self.connector_leader.has_slot(request_id.to_string())
     }
 
     fn create_slot(&mut self, request: KvbmRequest, tokens: Vec<u32>) -> PyResult<()> {
-        self.connector_leader.create_slot(request, tokens)
+        self.connector_leader.create_slot(request, tokens).map_err(to_pyerr)
     }
 }

lib/bindings/python/rust/llm/block_manager/vllm/connector/leader/recorder.rs

@@ -1,4 +1,5 @@
 use super::*;
+use anyhow;
 
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
@@ -108,9 +109,7 @@ impl KvConnectorLeaderRecorder {
         let output_path = "/tmp/records.jsonl";
         tracing::info!("recording events to {}", output_path);
 
-        // Create recorder synchronously using pyo3 async runtime
-        let runtime = pyo3_async_runtimes::tokio::get_runtime();
-        let recorder = runtime.block_on(async {
+        let recorder = drt.runtime().primary().block_on(async {
             Recorder::new(token, &output_path, None, None, None).await
         }).unwrap();
 
@@ -125,7 +124,7 @@ impl KvConnectorLeaderRecorder {
         let (unbounded_tx, unbounded_rx) = mpsc::unbounded_channel();
         let recorder_tx = recorder.event_sender();
 
-        let _ = runtime.spawn(Self::forward_unbounded_to_sender(unbounded_rx, recorder_tx));
+        let _ = drt.runtime().primary().spawn(Self::forward_unbounded_to_sender(unbounded_rx, recorder_tx));
 
         Self {
             _recorder: recorder,
@@ -158,7 +157,7 @@ impl Leader for KvConnectorLeaderRecorder {
         request_id: String,
         request_num_tokens: usize,
         num_computed_tokens: usize,
-    ) -> PyResult<(usize, bool)> {
+    ) -> anyhow::Result<(usize, bool)> {
         let input_copy = GetNumNewMatchedTokensInput {
             request_id: request_id.clone(),
             request_num_tokens: request_num_tokens.clone(),
@@ -183,7 +182,7 @@ impl Leader for KvConnectorLeaderRecorder {
         request_id: String,
         block_ids: Vec<BlockId>,
         num_external_tokens: usize,
-    ) -> PyResult<()> {
+    ) -> anyhow::Result<()> {
         let input_copy = UpdateStateAfterAllocInput {
             request_id: request_id.clone(),
             block_ids: block_ids.clone(),
@@ -197,7 +196,7 @@ impl Leader for KvConnectorLeaderRecorder {
     fn build_connector_metadata(
         &mut self,
         scheduler_output: SchedulerOutput,
-    ) -> PyResult<Vec<u8>> {
+    ) -> anyhow::Result<Vec<u8>> {
         let input_copy = BuildConnectorMetaInput {
             scheduler_output: scheduler_output.clone(),
         };
@@ -210,7 +209,7 @@ impl Leader for KvConnectorLeaderRecorder {
         output
     }
 
-    fn request_finished(&mut self, request_id: String, block_ids: Vec<BlockId>) -> PyResult<bool> {
+    fn request_finished(&mut self, request_id: String, block_ids: Vec<BlockId>) -> anyhow::Result<bool> {
         let input_copy = RequestFinishedInput {
             request_id: request_id.clone(),
             block_ids: block_ids.clone(),
@@ -239,7 +238,7 @@ impl Leader for KvConnectorLeaderRecorder {
 
     /// Create a new slot for the given request ID.
     /// This is used to create a new slot for the request.
-    fn create_slot(&mut self, request: KvbmRequest, tokens: Vec<u32>) -> PyResult<()> {
+    fn create_slot(&mut self, request: KvbmRequest, tokens: Vec<u32>) -> anyhow::Result<()> {
         let input_copy = CreateSlotInput {
             request: request.clone(),
             tokens: tokens.clone(),