[rust] cache-aware DP - approx tree (#1934)

rust/src/router.rs

@@ -1,38 +1,90 @@
+use crate::tree::RadixTree;
 use actix_web::http::header::{HeaderValue, CONTENT_TYPE};
 use actix_web::{HttpRequest, HttpResponse};
 use bytes::Bytes;
 use futures_util::TryStreamExt;
+use std::collections::HashMap;
 use std::fmt::Debug;
+use std::sync::atomic::AtomicUsize;
+use std::sync::{Arc, Mutex};
+use tokenizers::tokenizer::Tokenizer;
 
 #[derive(Debug)]
 pub enum Router {
     RoundRobin {
         worker_urls: Vec<String>,
-        current_index: std::sync::atomic::AtomicUsize,
+        current_index: AtomicUsize,
     },
     Random {
         worker_urls: Vec<String>,
     },
+    ApproxTree {
+        worker_urls: Vec<String>,
+        // TODO: don't lock the whole tree
+        url_to_tree: Arc<Mutex<HashMap<String, RadixTree>>>,
+        tokenizer: Tokenizer,
+        url_to_count: Arc<Mutex<HashMap<String, usize>>>,
+        cache_threshold: f32,
+    },
+}
+
+pub enum PolicyConfig {
+    RandomConfig,
+    RoundRobinConfig,
+    ApproxTreeConfig {
+        tokenizer_path: String,
+        cache_threshold: f32,
+    },
+}
+
+fn get_token_ids_from_request(body: &Bytes, tokenizer: &Tokenizer) -> Vec<u32> {
+    // 1. convert body to json
+    let json = serde_json::from_slice::<serde_json::Value>(body).unwrap();
+    // 2. get the text field
+    let text = json.get("text").and_then(|t| t.as_str()).unwrap_or("");
+    // 3. tokenize the text field
+    let tokens = tokenizer.encode(text, false).unwrap();
+
+    tokens.get_ids().to_vec()
 }
 
 impl Router {
-    pub fn new(worker_urls: Vec<String>, policy: String) -> Self {
-        match policy.to_lowercase().as_str() {
-            "random" => Router::Random { worker_urls },
-            "round_robin" => Router::RoundRobin {
+    pub fn new(worker_urls: Vec<String>, policy_config: PolicyConfig) -> Self {
+        match policy_config {
+            PolicyConfig::RandomConfig => Router::Random { worker_urls },
+            PolicyConfig::RoundRobinConfig => Router::RoundRobin {
                 worker_urls,
                 current_index: std::sync::atomic::AtomicUsize::new(0),
             },
-            _ => panic!(
-                "Unknown routing policy: {}. The available policies are 'random' and 'round_robin'",
-                policy
-            ),
+            PolicyConfig::ApproxTreeConfig {
+                tokenizer_path,
+                cache_threshold,
+            } => {
+                let mut url_to_tree = HashMap::new();
+                let mut url_to_count = HashMap::new();
+
+                for url in &worker_urls {
+                    url_to_tree.insert(url.clone(), RadixTree::new());
+                    url_to_count.insert(url.clone(), 0);
+                }
+
+                Router::ApproxTree {
+                    worker_urls,
+                    url_to_tree: Arc::new(Mutex::new(url_to_tree)),
+                    // TODO: rust ::from_pretrained cannot load from local file, so use ::from_file to load local file
+                    tokenizer: Tokenizer::from_file(tokenizer_path).unwrap(),
+                    url_to_count: Arc::new(Mutex::new(url_to_count)),
+                    cache_threshold,
+                }
+            }
         }
     }
 
     pub fn get_first(&self) -> Option<String> {
         match self {
-            Router::RoundRobin { worker_urls, .. } | Router::Random { worker_urls } => {
+            Router::RoundRobin { worker_urls, .. }
+            | Router::Random { worker_urls }
+            | Router::ApproxTree { worker_urls, .. } => {
                 if worker_urls.is_empty() {
                     None
                 } else {
@@ -48,26 +100,96 @@ impl Router {
         req: HttpRequest,
         body: Bytes,
     ) -> HttpResponse {
+        let mut input_ids: Vec<u32> = Vec::new();
+        if let Router::ApproxTree { tokenizer, .. } = self {
+            input_ids = get_token_ids_from_request(&body, tokenizer);
+        }
+
         let worker_url = match self {
             Router::RoundRobin {
                 worker_urls,
                 current_index,
             } => {
-                current_index
+                let idx = current_index
                     .fetch_update(
                         std::sync::atomic::Ordering::SeqCst,
                         std::sync::atomic::Ordering::SeqCst,
                         |x| Some((x + 1) % worker_urls.len()),
                     )
-                    .expect_err("Error updating index in round robin");
+                    .unwrap();
 
-                &worker_urls[current_index.load(std::sync::atomic::Ordering::SeqCst)]
+                worker_urls[idx].clone()
             }
+
             Router::Random { worker_urls } => {
-                &worker_urls[rand::random::<usize>() % worker_urls.len()]
+                worker_urls[rand::random::<usize>() % worker_urls.len()].clone()
+            }
+
+            Router::ApproxTree {
+                worker_urls,
+                url_to_tree,
+                url_to_count,
+                cache_threshold,
+                ..
+            } => {
+                // TODO: pipeline the locks. Release one earlier.
+
+                let mut max_matched_rate = 0.0;
+                let mut max_matched_idx = 0;
+
+                let locked_url_to_tree = url_to_tree.lock().unwrap();
+
+                // 1. Find the highest matched worker
+                for (i, url) in worker_urls.iter().enumerate() {
+                    let tree = locked_url_to_tree.get(url).unwrap();
+                    let matched = tree.prefix_match(&input_ids[..]).len();
+                    let matched_rate = matched as f32 / input_ids.len() as f32;
+
+                    if matched_rate > max_matched_rate {
+                        max_matched_rate = matched_rate;
+                        max_matched_idx = i;
+                    }
+                }
+
+                // 2. If the rate is higher than the threshold, select the worker. If not, select the worker with the shortest queue
+                if max_matched_rate > *cache_threshold {
+                    worker_urls[max_matched_idx].clone()
+                } else {
+                    // pick the shortest queue from url_to_count
+                    let locked_url_to_count = url_to_count.lock().unwrap();
+
+                    let mut min_count = std::usize::MAX;
+                    let mut min_count_id = 0;
+
+                    for (i, url) in worker_urls.iter().enumerate() {
+                        let count = locked_url_to_count.get(url).unwrap();
+                        if *count < min_count {
+                            min_count = *count;
+                            min_count_id = i;
+                        }
+                    }
+
+                    worker_urls[min_count_id].clone()
+                }
             }
         };
 
+        if let Router::ApproxTree {
+            url_to_tree,
+            url_to_count,
+            ..
+        } = self
+        {
+            // Insert input_ids to the tree
+            let mut locked_url_to_tree = url_to_tree.lock().unwrap();
+            let selected_tree = locked_url_to_tree.get_mut(&worker_url).unwrap();
+            selected_tree.insert(&input_ids[..]);
+
+            let mut locked_url_to_count = url_to_count.lock().unwrap();
+            let count = locked_url_to_count.get_mut(&worker_url).unwrap();
+            *count += 1;
+        }
+
         // Check if client requested streaming
         let is_stream = serde_json::from_slice::<serde_json::Value>(&body)
             .map(|v| v.get("stream").and_then(|s| s.as_bool()).unwrap_or(false))
@@ -94,11 +216,19 @@ impl Router {
             .unwrap_or(actix_web::http::StatusCode::INTERNAL_SERVER_ERROR);
 
         if !is_stream {
+            // TODO: do the correction on the tree based on the cached input_ids
+            if let Router::ApproxTree { url_to_count, .. } = self {
+                let mut locked_url_to_count = url_to_count.lock().unwrap();
+                let count = locked_url_to_count.get_mut(&worker_url).unwrap();
+                *count -= 1;
+            }
+
             match res.bytes().await {
                 Ok(body) => HttpResponse::build(status).body(body.to_vec()),
                 Err(_) => HttpResponse::InternalServerError().finish(),
             }
         } else {
+            // TODO: do the correction on the tree based on the cached input_ids. The streaming might be tricker to handle
             HttpResponse::build(status)
                 .insert_header((CONTENT_TYPE, HeaderValue::from_static("text/event-stream")))
                 .streaming(res.bytes_stream().map_err(|_| {