some prelim cleanups

Author: PeaBrane

lib/llm/src/kv_router.rs

@@ -25,9 +25,11 @@ pub mod scoring;
 
 use crate::{
     kv_router::{
-        indexer::{KvIndexer, KvIndexerInterface, RouterEvent},
+        indexer::{KvIndexer, KvIndexerInterface},
         metrics_aggregator::KvMetricsAggregator,
-        protocols::{LocalBlockHash, RouterRequest, RouterResponse, WorkerSelectionResult},
+        protocols::{
+            LocalBlockHash, RouterEvent, RouterRequest, RouterResponse, WorkerSelectionResult,
+        },
         scheduler::{KvScheduler, KvSchedulerError, SchedulingRequest},
         scoring::ProcessedEndpoints,
     },

lib/llm/src/kv_router/indexer.rs

@@ -77,9 +77,6 @@ pub enum KvRouterError {
     IndexerDroppedRequest,
 }
 
-/// Identifier of a LLM worker which emits events to the router.
-pub type WorkerId = i64;
-
 /// A shared reference to a [`RadixBlock`].
 type SharedRadixBlock = Rc<RefCell<RadixBlock>>;
 
@@ -133,31 +130,6 @@ pub fn compute_block_hash_for_seq(tokens: &[u32], kv_block_size: usize) -> Vec<L
         .collect()
 }
 
-/// A [`KvCacheEvent`] on a specific LLM worker denoted by [`WorkerId`].
-#[derive(Debug, Clone, Serialize, Deserialize)]
-pub struct RouterEvent {
-    /// The ID of the worker emitting the event.
-    worker_id: WorkerId,
-    /// The cache event associated with the worker.
-    event: KvCacheEvent,
-}
-
-impl RouterEvent {
-    /// Create a new `RouterEvent`.
-    ///
-    /// ### Arguments
-    ///
-    /// * `worker_id` - The ID of the worker emitting the event.
-    /// * `event` - The cache event.
-    ///
-    /// ### Returns
-    ///
-    /// A new `RouterEvent`.
-    pub fn new(worker_id: WorkerId, event: KvCacheEvent) -> Self {
-        Self { worker_id, event }
-    }
-}
-
 /// A block in the Radix Tree.
 #[derive(Debug)]
 struct RadixBlock {

lib/llm/src/kv_router/protocols.rs

@@ -21,15 +21,18 @@ pub struct RouterRequest {
     pub tokens: Vec<Token>,
 }
 
+/// Identifier of a LLM worker which emits events to the router.
+pub type WorkerId = i64;
+
 #[derive(Debug, Clone, Serialize, Deserialize, Default)]
 pub struct RouterResponse {
-    pub worker_id: i64,
+    pub worker_id: WorkerId,
 }
 
 #[derive(Debug)]
 pub struct WorkerSelectionResult {
     /// The worker id of the selected worker
-    pub worker_id: i64,
+    pub worker_id: WorkerId,
 
     /// The total number of blocks required to prefill the request
     pub required_blocks: u64,
@@ -58,14 +61,14 @@ pub struct ForwardPassMetrics {
 
 /// A [`LocalBlockHash`] is a hash computed from the tokens_ids, extra_token_ids and the optional
 /// lora_id of a block.
-#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Ord, PartialOrd)]
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
 pub struct LocalBlockHash(pub u64);
 
 /// A sequence aware hash of a block where the hash is computed from the tokens_ids, extra_token_ids
 /// and the optional lora_id of a block, PLUS the hash of the parent block.
 ///
 /// In this case, the hashing function is external and unknown.
-#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Ord, PartialOrd)]
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
 pub struct ExternalSequenceBlockHash(pub u64);
 
 // Implement From trait for convenient conversion
@@ -138,6 +141,31 @@ pub struct KvCacheRemoveData {
     pub block_hashes: Vec<ExternalSequenceBlockHash>,
 }
 
+/// A [`KvCacheEvent`] on a specific LLM worker denoted by [`WorkerId`].
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct RouterEvent {
+    /// The ID of the worker emitting the event.
+    pub worker_id: WorkerId,
+    /// The cache event associated with the worker.
+    pub event: KvCacheEvent,
+}
+
+impl RouterEvent {
+    /// Create a new `RouterEvent`.
+    ///
+    /// ### Arguments
+    ///
+    /// * `worker_id` - The ID of the worker emitting the event.
+    /// * `event` - The cache event.
+    ///
+    /// ### Returns
+    ///
+    /// A new `RouterEvent`.
+    pub fn new(worker_id: WorkerId, event: KvCacheEvent) -> Self {
+        Self { worker_id, event }
+    }
+}
+
 impl Serialize for LocalBlockHash {
     fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
     where

lib/llm/src/kv_router/publisher.rs

@@ -14,9 +14,7 @@
 // limitations under the License.
 
 use crate::kv_router::{
-    indexer::{compute_block_hash_for_seq, RouterEvent},
-    protocols::*,
-    KV_EVENT_SUBJECT, KV_METRICS_ENDPOINT,
+    indexer::compute_block_hash_for_seq, protocols::*, KV_EVENT_SUBJECT, KV_METRICS_ENDPOINT,
 };
 use async_trait::async_trait;
 use dynamo_runtime::traits::{events::EventPublisher, DistributedRuntimeProvider, RuntimeProvider};

lib/llm/src/kv_router/recorder.rs

@@ -13,7 +13,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use crate::kv_router::indexer::RouterEvent;
+use crate::kv_router::protocols::*;
 use crate::recorder::Recorder;
 
 // Type alias for backward compatibility
@@ -23,8 +23,6 @@ pub type KvRecorder = Recorder<RouterEvent>;
 mod tests {
     use super::*;
     use crate::kv_router::indexer::KvIndexer;
-    use crate::kv_router::indexer::WorkerId;
-    use crate::kv_router::protocols::*;
     use std::time::Duration;
     use tempfile::tempdir;
     use tokio::fs;