clean up

Author: ryanolson

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

@@ -1,8 +1,8 @@
 // SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 // SPDX-License-Identifier: Apache-2.0
 
-pub mod slot;
 pub mod recorder;
+pub mod slot;
 
 use super::*;
 use dynamo_runtime::DistributedRuntime;
@@ -28,8 +28,8 @@ use std::{
     collections::HashSet,
     sync::{Arc, Mutex},
 };
-use tokio::sync::mpsc;
 use tokio;
+use tokio::sync::mpsc;
 
 type VllmLocality = Logical<DistributedLeaderWorkerResources>;
 
@@ -38,10 +38,9 @@ impl From<SlotError> for PyErr {
         to_pyerr(err)
     }
 }
+use anyhow;
 use dynamo_llm::recorder::Recorder;
 use tokio_util::sync::CancellationToken;
-use anyhow;
-
 
 pub trait Leader: Send + Sync + std::fmt::Debug {
     fn get_num_new_matched_tokens(
@@ -83,7 +82,6 @@ pub struct KvConnectorLeader {
     iteration_counter: u64,
 }
 
-
 impl KvConnectorLeader {
     fn new(
         worker_id: String,
@@ -137,7 +135,9 @@ impl Leader for KvConnectorLeader {
         debug_assert!(num_computed_tokens % self.block_size == 0);
 
         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 mut slot = shared_slot
+            .lock()
+            .map_err(|e| anyhow::anyhow!("Failed to lock slot: {}", e))?;
 
         // early exit if we cannot match full block
         if (slot.sequence().total_tokens() - num_computed_tokens) < self.block_size {
@@ -180,7 +180,9 @@ impl Leader for KvConnectorLeader {
         );
 
         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 mut slot = shared_slot
+            .lock()
+            .map_err(|e| anyhow::anyhow!("Failed to lock slot: {}", e))?;
 
         // we have not yet advanced the computed position, but now we can, since we have an indication that we have
         // necessary gpu blocks into which we will load the external tokens.
@@ -207,7 +209,7 @@ impl Leader for KvConnectorLeader {
     }
 
     #[tracing::instrument(level = "debug", skip_all, fields(iteration = self.iteration_counter + 1))]
-    pub fn build_connector_metadata(
+    fn build_connector_metadata(
         &mut self,
         scheduler_output: SchedulerOutput,
     ) -> anyhow::Result<Vec<u8>> {
@@ -234,7 +236,9 @@ impl Leader for KvConnectorLeader {
         // 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)?;
-            let mut slot = shared_slot.lock().map_err(|e| anyhow::anyhow!("Failed to lock slot: {}", e))?;
+            let mut slot = shared_slot
+                .lock()
+                .map_err(|e| anyhow::anyhow!("Failed to lock slot: {}", e))?;
 
             md.create_slot(request_id.clone());
 
@@ -256,7 +260,9 @@ impl Leader for KvConnectorLeader {
             );
 
             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 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,7 +303,9 @@ impl Leader for KvConnectorLeader {
             );
 
             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 mut slot = shared_slot
+                .lock()
+                .map_err(|e| anyhow::anyhow!("Failed to lock slot: {}", e))?;
 
             let scheduled_tokens = *scheduler_output
                 .num_scheduled_tokens
@@ -322,16 +330,23 @@ impl Leader for KvConnectorLeader {
         }
 
         tracing::debug!("metadata: {md:#?}");
-        serde_json::to_vec(&md).map_err(|e| anyhow::anyhow!("Failed to serialize connector metadata: {}", e))
+        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>) -> anyhow::Result<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)?;
 
         // mark the slot as finished
-        let mut slot = shared_slot.lock().map_err(|e| anyhow::anyhow!("Failed to lock slot: {}", e))?;
+        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
@@ -413,7 +428,9 @@ 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).map_err(to_pyerr)
+        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(
@@ -422,25 +439,30 @@ 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).map_err(to_pyerr)
+        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).map_err(to_pyerr)
+    fn build_connector_metadata(&mut self, scheduler_output: SchedulerOutput) -> PyResult<Vec<u8>> {
+        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).map_err(to_pyerr)
+        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).map_err(to_pyerr)
+        self.connector_leader
+            .create_slot(request, tokens)
+            .map_err(to_pyerr)
     }
 }

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

@@ -1,7 +1,6 @@
 use super::*;
 use anyhow;
 
-
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub enum Action {
     GetNumNewMatchedTokens(GetNumNewMatchedTokensInput, GetNumNewMatchedTokensOutput),
@@ -33,8 +32,7 @@ pub struct UpdateStateAfterAllocInput {
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
-pub struct UpdateStateAfterAllocOutput {
-}
+pub struct UpdateStateAfterAllocOutput {}
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct BuildConnectorMetaInput {
@@ -43,7 +41,7 @@ pub struct BuildConnectorMetaInput {
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct BuildConnectorMetaOutput {
-    metadata: Vec<u8>,
+    metadata: ConnectorMetadata,
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
@@ -74,8 +72,7 @@ pub struct CreateSlotInput {
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
-pub struct CreateSlotOutput {
-}
+pub struct CreateSlotOutput {}
 
 #[derive(Debug)]
 pub struct KvConnectorLeaderRecorder {
@@ -109,9 +106,11 @@ impl KvConnectorLeaderRecorder {
         let output_path = "/tmp/records.jsonl";
         tracing::info!("recording events to {}", output_path);
 
-        let recorder = drt.runtime().primary().block_on(async {
-            Recorder::new(token, &output_path, None, None, None).await
-        }).unwrap();
+        let recorder = drt
+            .runtime()
+            .primary()
+            .block_on(async { Recorder::new(token, &output_path, None, None, None).await })
+            .unwrap();
 
         let connector_leader = KvConnectorLeader {
             slot_manager: ConnectorSlotManager::new(block_manager.clone(), leader, drt.clone()),
@@ -124,7 +123,10 @@ impl KvConnectorLeaderRecorder {
         let (unbounded_tx, unbounded_rx) = mpsc::unbounded_channel();
         let recorder_tx = recorder.event_sender();
 
-        let _ = drt.runtime().primary().spawn(Self::forward_unbounded_to_sender(unbounded_rx, recorder_tx));
+        // todo(kvbm): make this a critical task
+        drt.runtime()
+            .primary()
+            .spawn(Self::forward_unbounded_to_sender(unbounded_rx, recorder_tx));
 
         Self {
             _recorder: recorder,
@@ -160,16 +162,22 @@ impl Leader for KvConnectorLeaderRecorder {
     ) -> anyhow::Result<(usize, bool)> {
         let input_copy = GetNumNewMatchedTokensInput {
             request_id: request_id.clone(),
-            request_num_tokens: request_num_tokens.clone(),
-            num_computed_tokens: num_computed_tokens.clone(),
+            request_num_tokens,
+            num_computed_tokens,
         };
-        let output = self.connector_leader.get_num_new_matched_tokens(request_id, request_num_tokens, num_computed_tokens);
-        let output_copy = output.as_ref().unwrap().clone();
-        let _ = self.unbounded_tx.send(Action::GetNumNewMatchedTokens(input_copy, GetNumNewMatchedTokensOutput {
-                num_new_matched_tokens: output_copy.0,
-                has_matched: output_copy.1,
-            }));
-        output
+        let output = self.connector_leader.get_num_new_matched_tokens(
+            request_id,
+            request_num_tokens,
+            num_computed_tokens,
+        )?;
+        let _ = self.unbounded_tx.send(Action::GetNumNewMatchedTokens(
+            input_copy,
+            GetNumNewMatchedTokensOutput {
+                num_new_matched_tokens: output.0,
+                has_matched: output.1,
+            },
+        ));
+        Ok(output)
     }
 
     /// We drop the need to pass in the KvCacheBlocks and the num_external_tokens as they are captured
@@ -186,10 +194,17 @@ impl Leader for KvConnectorLeaderRecorder {
         let input_copy = UpdateStateAfterAllocInput {
             request_id: request_id.clone(),
             block_ids: block_ids.clone(),
-            num_external_tokens: num_external_tokens.clone(),
+            num_external_tokens,
         };
-        let _ = self.connector_leader.update_state_after_alloc(request_id, block_ids, num_external_tokens).unwrap();
-        let _ = self.unbounded_tx.send(Action::UpdateStateAfterAlloc(input_copy, UpdateStateAfterAllocOutput {}));
+        self.connector_leader.update_state_after_alloc(
+            request_id,
+            block_ids,
+            num_external_tokens,
+        )?;
+        let _ = self.unbounded_tx.send(Action::UpdateStateAfterAlloc(
+            input_copy,
+            UpdateStateAfterAllocOutput {},
+        ));
         Ok(())
     }
 
@@ -200,39 +215,48 @@ impl Leader for KvConnectorLeaderRecorder {
         let input_copy = BuildConnectorMetaInput {
             scheduler_output: scheduler_output.clone(),
         };
-        let output = self.connector_leader.build_connector_metadata(scheduler_output);
-        let output_copy = output.as_ref().unwrap().clone();
-        let _ = self.unbounded_tx
-                .send(Action::BuildConnectorMeta(input_copy, BuildConnectorMetaOutput {
-                    metadata: output_copy,
-                }));
-        output
+        let output = self
+            .connector_leader
+            .build_connector_metadata(scheduler_output)?;
+        let _ = self.unbounded_tx.send(Action::BuildConnectorMeta(
+            input_copy,
+            BuildConnectorMetaOutput {
+                metadata: serde_json::from_slice(&output)?,
+            },
+        ));
+        Ok(output)
     }
 
-    fn request_finished(&mut self, request_id: String, block_ids: Vec<BlockId>) -> anyhow::Result<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(),
         };
-        let output = self.connector_leader.request_finished(request_id, block_ids);
-        let output_copy = output.as_ref().unwrap().clone();
-        let _ = self.unbounded_tx
-                .send(Action::RequestFinished(input_copy, RequestFinishedOutput {
-                    is_finished: output_copy,
-                }));
-        output
+        let output = self
+            .connector_leader
+            .request_finished(request_id, block_ids)?;
+        let _ = self.unbounded_tx.send(Action::RequestFinished(
+            input_copy,
+            RequestFinishedOutput {
+                is_finished: output,
+            },
+        ));
+        Ok(output)
     }
 
     fn has_slot(&self, request_id: String) -> bool {
         let input_copy = HasSlotInput {
             request_id: request_id.clone(),
         };
         let output = self.connector_leader.has_slot(request_id);
-        let output_copy = output.clone();
-        let _ = self.unbounded_tx
-                .send(Action::HasSlot(input_copy, HasSlotOutput {
-                    result: output_copy,
-                }));
+        let _ = self.unbounded_tx.send(Action::HasSlot(
+            input_copy,
+            HasSlotOutput { result: output },
+        ));
         output
     }
 
@@ -244,7 +268,9 @@ impl Leader for KvConnectorLeaderRecorder {
             tokens: tokens.clone(),
         };
         let _ = self.connector_leader.create_slot(request, tokens);
-        let _ = self.unbounded_tx.send(Action::CreateSlot(input_copy, CreateSlotOutput {}));
+        let _ = self
+            .unbounded_tx
+            .send(Action::CreateSlot(input_copy, CreateSlotOutput {}));
         Ok(())
     }
-}
+}