fix: potential out-of-bound (#420)

Signed-off-by: Tianer Zhou <[email protected]>

Author: ezhoureal

lib/llm/src/kv_router/scheduler.rs

@@ -245,8 +245,7 @@ pub fn select_worker(
         let kv_load_ratio = w.data.kv_active_blocks as f64 / w.data.kv_total_blocks as f64;
         let load_deviation = kv_load_ratio - workers.load_avg;
 
-        // [FIXME] multiple endpoints of the same worker cause out of bound error
-        let worker_id = workers.worker_ids[i];
+        let worker_id = w.worker_id();
         let overlap_score = request.overlap.scores.get(&worker_id).map_or(0, |x| *x);
         let overlap_score = overlap_score as usize * kv_block_size;
 

lib/llm/src/kv_router/scoring.rs

@@ -16,14 +16,12 @@
 //! Scoring functions for the KV router.
 
 use serde::{Deserialize, Serialize};
-use std::collections::HashSet;
 
 use crate::kv_router::scheduler::Endpoint;
 
 #[derive(Debug, Default, Serialize, Deserialize, Clone)]
 pub struct ProcessedEndpoints {
     pub endpoints: Vec<Endpoint>,
-    pub worker_ids: Vec<i64>,
     pub load_avg: f64,
     pub load_std: f64,
 }
@@ -43,12 +41,8 @@ impl ProcessedEndpoints {
             / load_values.len() as f64;
         let load_std = variance.sqrt();
 
-        let worker_ids: HashSet<i64> = endpoints.iter().map(|x| x.worker_id()).collect();
-        let worker_ids: Vec<i64> = worker_ids.into_iter().collect();
-
         ProcessedEndpoints {
             endpoints,
-            worker_ids,
             load_avg,
             load_std,
         }