From c691e8b39618651bd9da32851843dd026a93ab68 Mon Sep 17 00:00:00 2001
From: Justin Ji <justinji@college.harvard.edu>
Date: Fri, 2 May 2025 10:35:17 -0700
Subject: [PATCH 1/6] Add prefix tree class

Signed-off-by: Justin Ji <justinji@college.harvard.edu>
---
 .../serve/deployments/routers/prefix_tree.py  | 232 ++++++++++++++++++
 1 file changed, 232 insertions(+)
 create mode 100644 python/ray/llm/_internal/serve/deployments/routers/prefix_tree.py

diff --git a/python/ray/llm/_internal/serve/deployments/routers/prefix_tree.py b/python/ray/llm/_internal/serve/deployments/routers/prefix_tree.py
new file mode 100644
index 000000000000..a3dc8cce0d0a
--- /dev/null
+++ b/python/ray/llm/_internal/serve/deployments/routers/prefix_tree.py
@@ -0,0 +1,232 @@
+from ray import serve
+import time
+from collections import defaultdict
+from threading import Lock, RLock
+import json
+from typing import Optional, List, Tuple, Dict, Any, TypeVar, Union, cast
+
+class Node:
+    """
+    Node in a prefix tree that tracks tenant access time.
+    
+    Each node represents a segment of text and can belong to multiple tenants.
+    """
+    def __init__(self, text: str = "", parent: Optional["Node"] = None):
+        self.text: str = text
+        self.parent: Optional["Node"] = parent
+        self.children: Dict[str, "Node"] = {}  # Maps char -> Node
+        self.tenant_last_access_time: Dict[str, int] = {}  # Maps tenant -> timestamp in ms (int)
+        
+    def __str__(self) -> str:
+        return f"Node(text='{self.text}', tenants={list(self.tenant_last_access_time.keys())})"
+
+@serve.deployment(name="TreeDeployment")
+class PrefixTree:
+    """
+    Thread-safe multi-tenant prefix tree (approximate radix tree).
+    
+    Features:
+    1. Stores data for multiple tenants in the same tree structure
+    2. Node-level locking for concurrent access
+    3. Leaf LRU eviction based on tenant access time
+    """
+    def __init__(self) -> None:
+        self.root: Node = Node()
+        self.tenant_char_count: Dict[str, int] = defaultdict(int)  # Maps tenant -> character count
+        self.lock: Lock = Lock()  # For operations that need to lock the entire tree
+ 
+    @staticmethod
+    def shared_prefix_count(a: str, b: str) -> int:
+        """Count the number of shared characters at the beginning of two strings."""
+        i: int = 0
+        for char_a, char_b in zip(a, b):
+            if char_a == char_b:
+                i += 1
+            else:
+                break
+        return i
+
+    def insert(self, text: str, tenant: str) -> None:
+        """Insert text into tree with given tenant."""
+        with self.lock:
+            curr_node: Node = self.root
+            timestamp_ms: int = int(time.time() * 1000)
+            i: int = 0
+            while i < len(text):
+                curr_node.tenant_last_access_time[tenant] = timestamp_ms
+                first_char: str = text[i]
+                curr_text: str = text[i:]
+                if first_char not in curr_node.children:
+                    # No match, create new node
+                    # e.g. curr_node.children = {}, curr_text = "hello" -> curr_node.children = {"h": Node("hello")}
+                    new_node: Node = Node(text=curr_text, parent=curr_node)
+                    new_node.tenant_last_access_time[tenant] = timestamp_ms
+                    
+                    # Increment char count for tenant
+                    self.tenant_char_count[tenant] += len(curr_text)
+                    
+                    curr_node.children[first_char] = new_node
+                else:
+                    # Match found, check if need to split
+                    matched_node: Node = curr_node.children[first_char]
+                    shared_count: int = self.shared_prefix_count(matched_node.text, curr_text)
+                    
+                    if shared_count < len(matched_node.text):
+                        # Partial match, split at matched point
+                        # Example:
+                        ## Before update:
+                        ### curr_node.children = {"h": Node("helloworld")}, curr_text = "hellothere" -> shared_count = 5
+                        ### matched_node = Node("helloworld")
+
+                        ## During update:
+                        ### Increment tenant_char_count[tenant] by shared_count if matched_node has not seen this tenant before
+
+                        ## After update:
+                        ### curr_node.children = {"h": Node("hello", children = {"w": Node("world")})}
+                        ### parent_node = Node("hello"), matched_node = Node("world")
+                        ### Update tenant_last_access_time for parent_node, NOT matched_node
+                        ### (new) curr_text = "there", (new) curr_node = parent_node
+                        ### Continue adding "there" to tree in next iteration
+
+                        matched_text: str = matched_node.text[:shared_count]
+                        remaining_text: str = matched_node.text[shared_count:]
+
+                        # Update tenant char count for the new split node
+                        if tenant not in matched_node.tenant_last_access_time:
+                            self.tenant_char_count[tenant] += shared_count
+                        
+                        # Create new parent node
+                        new_parent: Node = Node(text=matched_text, parent=curr_node)
+                        new_parent.tenant_last_access_time = matched_node.tenant_last_access_time.copy()
+                        
+                        # Update matched_node
+                        matched_node.text = remaining_text
+                        matched_node.parent = new_parent
+
+                        # Connect new parent node to matched_node
+                        new_parent.children[remaining_text[0]] = matched_node
+
+                        # Connect current node to new parent
+                        curr_node.children[first_char] = new_parent
+                        
+                        # Move down the tree
+                        curr_node = new_parent
+                        i += shared_count
+                    else:
+                        # Full match
+
+                        # Update tenant char count if this is a new tenant for this node
+                        if tenant not in matched_node.tenant_last_access_time:
+                            self.tenant_char_count[tenant] += shared_count
+                        
+                        # # Update tenant last access time
+                        # matched_node.tenant_last_access_time[tenant] = timestamp_ms
+
+                        # Move down the tree
+                        curr_node = matched_node
+                        i += shared_count
+
+    def prefix_match(self, text: str, available_tenants: Optional[List[str]] = None) -> Tuple[str, Optional[List[str]]]:
+        """
+        Match text against tree and return (matched_text, matched_tenants).
+        Does not update access time for the matched tenants (only updates when insert() is called).
+        """
+        with self.lock:
+            curr_node: Node = self.root
+            i: int = 0
+            text_len: int = len(text)
+            
+            while i < text_len:
+                first_char: str = text[i]
+                curr_text: str = text[i:]
+                
+                if first_char in curr_node.children:
+                    matched_node: Node = curr_node.children[first_char]
+
+                    # Check if any of the available tenants match this node
+                    if available_tenants:
+                        if not any(tenant in matched_node.tenant_last_access_time for tenant in available_tenants):
+                            break
+
+                    shared_count: int = self.shared_prefix_count(matched_node.text, curr_text)
+                    i += shared_count
+                    curr_node = matched_node
+
+                    if shared_count < len(matched_node.text):
+                        # Partial match, stop here
+                        break
+                else:
+                    # No match found, stop here
+                    break
+            
+            # Select the tenants in available_tenants that are in the current node
+            selected_tenants: Optional[List[str]] = None
+            if available_tenants:
+                matching_tenants = [tenant for tenant in available_tenants if tenant in curr_node.tenant_last_access_time]
+                if matching_tenants:
+                    selected_tenants = matching_tenants
+            else:
+                if curr_node.tenant_last_access_time:
+                    selected_tenants = list(curr_node.tenant_last_access_time)
+            
+            ret_text: str = text[:i]
+            return ret_text, selected_tenants
+    
+    def get_smallest_tenant(self) -> Optional[str]:
+        """Get the tenant with the smallest total character count."""
+        with self.lock:
+            if not self.tenant_char_count:
+                # Return first worker if no data yet
+                return None
+            
+            return min(self.tenant_char_count.items(), key=lambda x: x[1])[0]
+    
+    def evict_tenant_by_size(self, max_size: int) -> None:
+        """Evict nodes for tenants that exceed the maximum tree size."""
+        with self.lock:
+            # Get total tree size
+            total_size: int = sum(self.tenant_char_count.values())
+            
+            # If tree is smaller than max size, no need to evict
+            if total_size <= max_size:
+                return
+            
+            # Calculate how much we need to evict
+            excess: int = total_size - max_size
+            
+            # Sort tenants by size (largest first)
+            sorted_tenants: List[Tuple[str, int]] = sorted(
+                self.tenant_char_count.items(), 
+                key=lambda x: x[1], 
+                reverse=True
+            )
+            
+            # Evict from largest tenants first
+            for tenant, size in sorted_tenants:
+                # If we've evicted enough, stop
+                if excess <= 0:
+                    break
+                
+                # Calculate how much to evict from this tenant
+                # Evict at most half of the tenant's size
+                evict_amount: int = min(excess, size // 2)
+                
+                if evict_amount > 0:
+                    # print(f"Evicting {evict_amount} chars from tenant {tenant}")
+                    self.tenant_char_count[tenant] -= evict_amount
+                    excess -= evict_amount
+            
+            # print(f"Tree eviction complete. New size: {sum(self.tenant_char_count.values())}")
+    
+    def get_tenant_char_count(self) -> Dict[str, int]:
+        """Get character count for each tenant."""
+        with self.lock:
+            return dict(self.tenant_char_count)
+            
+    def remove_tenant(self, tenant: str) -> None:
+        """Remove all nodes belonging to a tenant."""
+        # Would require a traversal of the tree and removing the tenant
+        # from tenant_last_access_time. Simplifying for now.
+        with self.lock:
+            if tenant in self.tenant_char_count:
+                del self.tenant_char_count[tenant] 
\ No newline at end of file

From 939aa7e84d736b2749d167fd831c20f23f1ccd18 Mon Sep 17 00:00:00 2001
From: Justin Ji <justinji@college.harvard.edu>
Date: Fri, 2 May 2025 10:51:09 -0700
Subject: [PATCH 2/6] Linting

Signed-off-by: Justin Ji <justinji@college.harvard.edu>
---
 .../serve/deployments/routers/prefix_tree.py  | 170 ++++++++++--------
 1 file changed, 100 insertions(+), 70 deletions(-)

diff --git a/python/ray/llm/_internal/serve/deployments/routers/prefix_tree.py b/python/ray/llm/_internal/serve/deployments/routers/prefix_tree.py
index a3dc8cce0d0a..fd79d1be41d7 100644
--- a/python/ray/llm/_internal/serve/deployments/routers/prefix_tree.py
+++ b/python/ray/llm/_internal/serve/deployments/routers/prefix_tree.py
@@ -1,40 +1,50 @@
 from ray import serve
 import time
 from collections import defaultdict
-from threading import Lock, RLock
-import json
-from typing import Optional, List, Tuple, Dict, Any, TypeVar, Union, cast
+from threading import Lock
+from typing import Optional, List, Tuple, Dict
+
 
 class Node:
     """
     Node in a prefix tree that tracks tenant access time.
-    
+
     Each node represents a segment of text and can belong to multiple tenants.
     """
+
     def __init__(self, text: str = "", parent: Optional["Node"] = None):
         self.text: str = text
         self.parent: Optional["Node"] = parent
         self.children: Dict[str, "Node"] = {}  # Maps char -> Node
-        self.tenant_last_access_time: Dict[str, int] = {}  # Maps tenant -> timestamp in ms (int)
-        
+        self.tenant_last_access_time: Dict[
+            str, int
+        ] = {}  # Maps tenant -> timestamp in ms (int)
+
     def __str__(self) -> str:
         return f"Node(text='{self.text}', tenants={list(self.tenant_last_access_time.keys())})"
 
+
 @serve.deployment(name="TreeDeployment")
 class PrefixTree:
     """
     Thread-safe multi-tenant prefix tree (approximate radix tree).
-    
+
     Features:
     1. Stores data for multiple tenants in the same tree structure
     2. Node-level locking for concurrent access
     3. Leaf LRU eviction based on tenant access time
     """
+
     def __init__(self) -> None:
         self.root: Node = Node()
-        self.tenant_char_count: Dict[str, int] = defaultdict(int)  # Maps tenant -> character count
+        self.tenant_char_count: Dict[str, int] = defaultdict(
+            int
+        )  # Maps tenant -> character count
         self.lock: Lock = Lock()  # For operations that need to lock the entire tree
- 
+        self.tenant_nodes: Dict[str, List[Node]] = defaultdict(
+            list
+        )  # Maps tenant -> list of nodes it belongs to
+
     @staticmethod
     def shared_prefix_count(a: str, b: str) -> int:
         """Count the number of shared characters at the beginning of two strings."""
@@ -61,16 +71,19 @@ def insert(self, text: str, tenant: str) -> None:
                     # e.g. curr_node.children = {}, curr_text = "hello" -> curr_node.children = {"h": Node("hello")}
                     new_node: Node = Node(text=curr_text, parent=curr_node)
                     new_node.tenant_last_access_time[tenant] = timestamp_ms
-                    
-                    # Increment char count for tenant
+
+                    # Increment char count for tenant and add node to tenant_nodes
                     self.tenant_char_count[tenant] += len(curr_text)
-                    
+                    self.tenant_nodes[tenant].append(new_node)
+
                     curr_node.children[first_char] = new_node
                 else:
                     # Match found, check if need to split
                     matched_node: Node = curr_node.children[first_char]
-                    shared_count: int = self.shared_prefix_count(matched_node.text, curr_text)
-                    
+                    shared_count: int = self.shared_prefix_count(
+                        matched_node.text, curr_text
+                    )
+
                     if shared_count < len(matched_node.text):
                         # Partial match, split at matched point
                         # Example:
@@ -94,11 +107,13 @@ def insert(self, text: str, tenant: str) -> None:
                         # Update tenant char count for the new split node
                         if tenant not in matched_node.tenant_last_access_time:
                             self.tenant_char_count[tenant] += shared_count
-                        
+
                         # Create new parent node
                         new_parent: Node = Node(text=matched_text, parent=curr_node)
-                        new_parent.tenant_last_access_time = matched_node.tenant_last_access_time.copy()
-                        
+                        new_parent.tenant_last_access_time = (
+                            matched_node.tenant_last_access_time.copy()
+                        )
+                        self.tenant_nodes[tenant].append(new_parent)
                         # Update matched_node
                         matched_node.text = remaining_text
                         matched_node.parent = new_parent
@@ -108,7 +123,7 @@ def insert(self, text: str, tenant: str) -> None:
 
                         # Connect current node to new parent
                         curr_node.children[first_char] = new_parent
-                        
+
                         # Move down the tree
                         curr_node = new_parent
                         i += shared_count
@@ -118,7 +133,7 @@ def insert(self, text: str, tenant: str) -> None:
                         # Update tenant char count if this is a new tenant for this node
                         if tenant not in matched_node.tenant_last_access_time:
                             self.tenant_char_count[tenant] += shared_count
-                        
+
                         # # Update tenant last access time
                         # matched_node.tenant_last_access_time[tenant] = timestamp_ms
 
@@ -126,7 +141,9 @@ def insert(self, text: str, tenant: str) -> None:
                         curr_node = matched_node
                         i += shared_count
 
-    def prefix_match(self, text: str, available_tenants: Optional[List[str]] = None) -> Tuple[str, Optional[List[str]]]:
+    def prefix_match(
+        self, text: str, available_tenants: Optional[List[str]] = None
+    ) -> Tuple[str, Optional[List[str]]]:
         """
         Match text against tree and return (matched_text, matched_tenants).
         Does not update access time for the matched tenants (only updates when insert() is called).
@@ -135,20 +152,25 @@ def prefix_match(self, text: str, available_tenants: Optional[List[str]] = None)
             curr_node: Node = self.root
             i: int = 0
             text_len: int = len(text)
-            
+
             while i < text_len:
                 first_char: str = text[i]
                 curr_text: str = text[i:]
-                
+
                 if first_char in curr_node.children:
                     matched_node: Node = curr_node.children[first_char]
 
                     # Check if any of the available tenants match this node
                     if available_tenants:
-                        if not any(tenant in matched_node.tenant_last_access_time for tenant in available_tenants):
+                        if not any(
+                            tenant in matched_node.tenant_last_access_time
+                            for tenant in available_tenants
+                        ):
                             break
 
-                    shared_count: int = self.shared_prefix_count(matched_node.text, curr_text)
+                    shared_count: int = self.shared_prefix_count(
+                        matched_node.text, curr_text
+                    )
                     i += shared_count
                     curr_node = matched_node
 
@@ -158,75 +180,83 @@ def prefix_match(self, text: str, available_tenants: Optional[List[str]] = None)
                 else:
                     # No match found, stop here
                     break
-            
+
             # Select the tenants in available_tenants that are in the current node
             selected_tenants: Optional[List[str]] = None
             if available_tenants:
-                matching_tenants = [tenant for tenant in available_tenants if tenant in curr_node.tenant_last_access_time]
+                matching_tenants = [
+                    tenant
+                    for tenant in available_tenants
+                    if tenant in curr_node.tenant_last_access_time
+                ]
                 if matching_tenants:
                     selected_tenants = matching_tenants
             else:
                 if curr_node.tenant_last_access_time:
                     selected_tenants = list(curr_node.tenant_last_access_time)
-            
+
             ret_text: str = text[:i]
             return ret_text, selected_tenants
-    
+
     def get_smallest_tenant(self) -> Optional[str]:
         """Get the tenant with the smallest total character count."""
         with self.lock:
             if not self.tenant_char_count:
                 # Return first worker if no data yet
                 return None
-            
+
             return min(self.tenant_char_count.items(), key=lambda x: x[1])[0]
-    
-    def evict_tenant_by_size(self, max_size: int) -> None:
-        """Evict nodes for tenants that exceed the maximum tree size."""
-        with self.lock:
-            # Get total tree size
-            total_size: int = sum(self.tenant_char_count.values())
-            
-            # If tree is smaller than max size, no need to evict
-            if total_size <= max_size:
-                return
-            
-            # Calculate how much we need to evict
-            excess: int = total_size - max_size
-            
-            # Sort tenants by size (largest first)
-            sorted_tenants: List[Tuple[str, int]] = sorted(
-                self.tenant_char_count.items(), 
-                key=lambda x: x[1], 
-                reverse=True
-            )
-            
-            # Evict from largest tenants first
-            for tenant, size in sorted_tenants:
-                # If we've evicted enough, stop
-                if excess <= 0:
-                    break
-                
-                # Calculate how much to evict from this tenant
-                # Evict at most half of the tenant's size
-                evict_amount: int = min(excess, size // 2)
-                
-                if evict_amount > 0:
-                    # print(f"Evicting {evict_amount} chars from tenant {tenant}")
-                    self.tenant_char_count[tenant] -= evict_amount
-                    excess -= evict_amount
-            
-            # print(f"Tree eviction complete. New size: {sum(self.tenant_char_count.values())}")
-    
+
+    # def evict_tenant_by_size(self, max_size: int) -> None:
+    #     """Evict nodes for tenants that exceed the maximum tree size."""
+    #     with self.lock:
+    #         # Get total tree size
+    #         total_size: int = sum(self.tenant_char_count.values())
+
+    #         # If tree is smaller than max size, no need to evict
+    #         if total_size <= max_size:
+    #             return
+
+    #         # Calculate how much we need to evict
+    #         excess: int = total_size - max_size
+
+    #         # Sort tenants by size (largest first)
+    #         sorted_tenants: List[Tuple[str, int]] = sorted(
+    #             self.tenant_char_count.items(),
+    #             key=lambda x: x[1],
+    #             reverse=True
+    #         )
+
+    #         # Evict from largest tenants first
+    #         for tenant, size in sorted_tenants:
+    #             # If we've evicted enough, stop
+    #             if excess <= 0:
+    #                 break
+
+    #             # Calculate how much to evict from this tenant
+    #             # Evict at most half of the tenant's size
+    #             evict_amount: int = min(excess, size // 2)
+
+    #             if evict_amount > 0:
+    #                 # print(f"Evicting {evict_amount} chars from tenant {tenant}")
+    #                 self.tenant_char_count[tenant] -= evict_amount
+    #                 excess -= evict_amount
+
+    #         # print(f"Tree eviction complete. New size: {sum(self.tenant_char_count.values())}")
+
     def get_tenant_char_count(self) -> Dict[str, int]:
         """Get character count for each tenant."""
         with self.lock:
             return dict(self.tenant_char_count)
-            
+
     def remove_tenant(self, tenant: str) -> None:
         """Remove all nodes belonging to a tenant."""
         # Would require a traversal of the tree and removing the tenant
         # from tenant_last_access_time. Simplifying for now.
         with self.lock:
-            if tenant in self.tenant_char_count:
-                del self.tenant_char_count[tenant] 
\ No newline at end of file
+            for node in self.tenant_nodes[tenant]:
+                node.tenant_last_access_time.pop(tenant, None)
+                if not node.tenant_last_access_time:
+                    node.parent.children.pop(node.text[0], None)
+            self.tenant_char_count.pop(tenant, None)
+            self.tenant_nodes.pop(tenant, None)

From 8eb8a9c711abef1bf8616587cc6f9bf76d51513c Mon Sep 17 00:00:00 2001
From: Justin Ji <justinji@college.harvard.edu>
Date: Fri, 2 May 2025 12:54:09 -0700
Subject: [PATCH 3/6] Add test cases

Signed-off-by: Justin Ji <justinji@college.harvard.edu>
---
 .../serve/deployments/routers/prefix_tree.py  | 113 ++++++++------
 .../serve/cpu/deployments/test_prefix_tree.py | 147 ++++++++++++++++++
 2 files changed, 213 insertions(+), 47 deletions(-)
 create mode 100644 python/ray/llm/tests/serve/cpu/deployments/test_prefix_tree.py

diff --git a/python/ray/llm/_internal/serve/deployments/routers/prefix_tree.py b/python/ray/llm/_internal/serve/deployments/routers/prefix_tree.py
index fd79d1be41d7..9d208edb14fd 100644
--- a/python/ray/llm/_internal/serve/deployments/routers/prefix_tree.py
+++ b/python/ray/llm/_internal/serve/deployments/routers/prefix_tree.py
@@ -1,7 +1,7 @@
 from ray import serve
 import time
 from collections import defaultdict
-from threading import Lock
+from threading import RLock
 from typing import Optional, List, Tuple, Dict
 
 
@@ -40,11 +40,23 @@ def __init__(self) -> None:
         self.tenant_char_count: Dict[str, int] = defaultdict(
             int
         )  # Maps tenant -> character count
-        self.lock: Lock = Lock()  # For operations that need to lock the entire tree
+        self.lock: RLock = RLock()  # For operations that need to lock the entire tree
         self.tenant_nodes: Dict[str, List[Node]] = defaultdict(
             list
         )  # Maps tenant -> list of nodes it belongs to
 
+    def get_root(self) -> Node:
+        return self.root
+
+    def get_tenant_char_count(self) -> Dict[str, int]:
+        return self.tenant_char_count
+
+    def get_tenant_nodes(self) -> Dict[str, List[Node]]:
+        return self.tenant_nodes
+
+    def to_string(self) -> str:
+        return f"PrefixTree(root={self.root.__str__()}, tenant_char_count={self.tenant_char_count}, tenant_nodes={self.tenant_nodes})"
+
     @staticmethod
     def shared_prefix_count(a: str, b: str) -> int:
         """Count the number of shared characters at the beginning of two strings."""
@@ -202,61 +214,68 @@ def get_smallest_tenant(self) -> Optional[str]:
         """Get the tenant with the smallest total character count."""
         with self.lock:
             if not self.tenant_char_count:
-                # Return first worker if no data yet
                 return None
 
             return min(self.tenant_char_count.items(), key=lambda x: x[1])[0]
 
-    # def evict_tenant_by_size(self, max_size: int) -> None:
-    #     """Evict nodes for tenants that exceed the maximum tree size."""
-    #     with self.lock:
-    #         # Get total tree size
-    #         total_size: int = sum(self.tenant_char_count.values())
-
-    #         # If tree is smaller than max size, no need to evict
-    #         if total_size <= max_size:
-    #             return
+    def get_tenant_char_count(self) -> Dict[str, int]:
+        """Get character count for each tenant."""
+        with self.lock:
+            return dict(self.tenant_char_count)
 
-    #         # Calculate how much we need to evict
-    #         excess: int = total_size - max_size
+    def remove_tenant_entirely(self, tenant: str) -> int:
+        """Remove all nodes belonging to a tenant, returns the number of characters removed. Also removes the tenant from tenant_char_count and tenant_nodes."""
+        with self.lock:
+            total_chars_removed: int = 0
+            for node in self.tenant_nodes[tenant].copy():
+                total_chars_removed += self.remove_tenant_single_node(tenant, node)
+            self.tenant_nodes.pop(tenant, None)
+            self.tenant_char_count.pop(tenant, None)
+            return total_chars_removed
 
-    #         # Sort tenants by size (largest first)
-    #         sorted_tenants: List[Tuple[str, int]] = sorted(
-    #             self.tenant_char_count.items(),
-    #             key=lambda x: x[1],
-    #             reverse=True
-    #         )
+    def remove_tenant_single_node(self, tenant: str, node: Node) -> int:
+        """Remove a single node belonging to a tenant, returns the number of characters removed."""
+        with self.lock:
+            removed_chars_len: int = len(node.text)
+            self.tenant_char_count[tenant] -= removed_chars_len
+            self.tenant_nodes[tenant].remove(node)
+            node.tenant_last_access_time.pop(tenant, None)
 
-    #         # Evict from largest tenants first
-    #         for tenant, size in sorted_tenants:
-    #             # If we've evicted enough, stop
-    #             if excess <= 0:
-    #                 break
+            # If this node has no more tenants, remove it from the parent
+            if not node.tenant_last_access_time:
+                node.parent.children.pop(node.text[0], None)
 
-    #             # Calculate how much to evict from this tenant
-    #             # Evict at most half of the tenant's size
-    #             evict_amount: int = min(excess, size // 2)
+            return removed_chars_len
 
-    #             if evict_amount > 0:
-    #                 # print(f"Evicting {evict_amount} chars from tenant {tenant}")
-    #                 self.tenant_char_count[tenant] -= evict_amount
-    #                 excess -= evict_amount
+    def evict_tenant(self, tenant: str, min_remove_size: int) -> int:
+        """Evict nodes from a tenant until the removed character count is at least min_remove_size.
 
-    #         # print(f"Tree eviction complete. New size: {sum(self.tenant_char_count.values())}")
+        Args:
+            tenant: The tenant to evict nodes from
+            min_remove_size: Minimum number of characters to remove
 
-    def get_tenant_char_count(self) -> Dict[str, int]:
-        """Get character count for each tenant."""
+        Returns:
+            int: The actual number of characters removed
+        """
         with self.lock:
-            return dict(self.tenant_char_count)
+            if tenant not in self.tenant_nodes or not self.tenant_nodes[tenant]:
+                return 0
 
-    def remove_tenant(self, tenant: str) -> None:
-        """Remove all nodes belonging to a tenant."""
-        # Would require a traversal of the tree and removing the tenant
-        # from tenant_last_access_time. Simplifying for now.
-        with self.lock:
-            for node in self.tenant_nodes[tenant]:
-                node.tenant_last_access_time.pop(tenant, None)
-                if not node.tenant_last_access_time:
-                    node.parent.children.pop(node.text[0], None)
-            self.tenant_char_count.pop(tenant, None)
-            self.tenant_nodes.pop(tenant, None)
+            # Sort nodes by last access time (oldest first)
+            nodes_to_evict = sorted(
+                self.tenant_nodes[tenant],
+                key=lambda node: node.tenant_last_access_time.get(tenant, 0),
+            )
+
+            total_chars_removed: int = 0
+
+            # Remove nodes until we've reached the minimum removal size
+            for node in nodes_to_evict.copy():
+                # Use existing function to remove tenant from node
+                total_chars_removed += self.remove_tenant_single_node(tenant, node)
+
+                # Check if we've removed enough characters
+                if total_chars_removed >= min_remove_size:
+                    break
+
+            return total_chars_removed
diff --git a/python/ray/llm/tests/serve/cpu/deployments/test_prefix_tree.py b/python/ray/llm/tests/serve/cpu/deployments/test_prefix_tree.py
new file mode 100644
index 000000000000..b8d6eda04a8d
--- /dev/null
+++ b/python/ray/llm/tests/serve/cpu/deployments/test_prefix_tree.py
@@ -0,0 +1,147 @@
+import pytest
+import time
+import ray
+from ray import serve
+
+from ray.llm._internal.serve.deployments.routers.prefix_tree import PrefixTree
+
+
+@pytest.fixture(scope="module", autouse=True)
+def serve_instance():
+    # Start Ray and Serve once per test module
+    ray.init(ignore_reinit_error=True)
+    serve.start(detached=True)
+    yield
+    serve.shutdown()
+    ray.shutdown()
+
+
+@pytest.mark.asyncio
+async def test_insert_and_basic_match():
+    # Deploy a clean PrefixTree
+    tree = serve.run(PrefixTree.bind())
+
+    # Insert and match exact string
+    await tree.insert.remote("hello", "tenant-A")
+    matched_text, tenants = await tree.prefix_match.remote("hello")
+    assert matched_text == "hello"
+    assert tenants == ["tenant-A"]
+
+
+@pytest.mark.asyncio
+async def test_tree_splits_nodes_on_partial_match():
+    tree = serve.run(PrefixTree.bind())
+    await tree.insert.remote("helloworld", "A")
+    await tree.insert.remote("hellothere", "B")
+
+    # After inserting both, the root should have one child "h"
+    root = await tree.get_root.remote()
+    h_node = root.children.get("h")
+    assert h_node is not None
+
+
+@pytest.mark.asyncio
+async def test_no_match():
+    tree = serve.run(PrefixTree.bind())
+    matched_text, tenants = await tree.prefix_match.remote("hello")
+    assert matched_text == ""
+    assert tenants is None
+
+
+@pytest.mark.asyncio
+async def test_duplicate_insertion_no_double_count():
+    tree = serve.run(PrefixTree.bind())
+
+    await tree.insert.remote("foo", "T1")
+    await tree.insert.remote("foo", "T1")  # duplicate
+
+    counts = await tree.get_tenant_char_count.remote()
+    # Should count 'foo' only once
+    assert counts.get("T1", 0) == 3
+
+
+@pytest.mark.asyncio
+async def test_shared_prefix_splitting_and_branching():
+    tree = serve.run(PrefixTree.bind())
+
+    await tree.insert.remote("helloworld", "A")
+    await tree.insert.remote("hellothere", "B")
+
+    text_a, tenants_a = await tree.prefix_match.remote("helloworld")
+    text_b, tenants_b = await tree.prefix_match.remote("hellothere")
+
+    assert text_a == "helloworld"
+    assert tenants_a == ["A"]
+    assert text_b == "hellothere"
+    assert tenants_b == ["B"]
+
+
+@pytest.mark.asyncio
+async def test_prefix_match_partial_and_filter():
+    tree = serve.run(PrefixTree.bind())
+
+    await tree.insert.remote("apple", "X")
+    await tree.insert.remote("apricot", "Y")
+
+    # Partial match for 'application' -> 'appl'
+    text, tenants = await tree.prefix_match.remote("application")
+    assert text == "appl"
+    assert tenants == ["X"]
+
+    # Filter by available_tenants=['X'] on 'apricot'
+    text_fx, tenants_fx = await tree.prefix_match.remote("apricot", ["X"])
+    assert text_fx == "ap"
+    assert tenants_fx == ["X"]
+
+    # Filter by non-existent tenant yields no tenants
+    text_fz, tenants_fz = await tree.prefix_match.remote("apricot", ["Z"])
+    assert text_fz == ""
+    assert tenants_fz is None
+
+
+@pytest.mark.asyncio
+async def test_remove_and_get_smallest_and_evict():
+    tree = serve.run(PrefixTree.bind())
+
+    # Test removal and char count
+    await tree.insert.remote("cat", "T1")
+    await tree.insert.remote("dog", "T1")
+
+    counts = await tree.get_tenant_char_count.remote()
+    assert counts.get("T1") == len("cat") + len("dog")
+
+    # Remove entire tenant
+    removed = await tree.remove_tenant_entirely.remote("T1")
+    assert removed == len("cat") + len("dog")
+
+    counts_after = await tree.get_tenant_char_count.remote()
+    assert "T1" not in counts_after
+
+    # Test eviction LRU behavior
+    await tree.insert.remote("a", "T2")
+    time.sleep(0.001)
+    await tree.insert.remote("bb", "T2")
+    time.sleep(0.001)
+    await tree.insert.remote("ccc", "T2")
+
+    before = (await tree.get_tenant_char_count.remote())["T2"]
+    evicted = await tree.evict_tenant.remote("T2", 2)
+    after = (await tree.get_tenant_char_count.remote())["T2"]
+
+    assert evicted >= 2
+    assert before - after == evicted
+
+
+@pytest.mark.asyncio
+async def test_get_smallest_tenant():
+    tree = serve.run(PrefixTree.bind())
+    await tree.insert.remote("aaaa", "A")
+    await tree.insert.remote("bb", "B")
+    smallest = await tree.get_smallest_tenant.remote()
+    assert smallest == "B"
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))

From 3e55d54d26b7343c1fa2660763b28dfb3df16b1a Mon Sep 17 00:00:00 2001
From: Justin Ji <justinji@college.harvard.edu>
Date: Fri, 2 May 2025 18:08:17 -0700
Subject: [PATCH 4/6] Implement eviction, tests passing

Signed-off-by: Justin Ji <justinji@college.harvard.edu>
---
 .../serve/deployments/routers/prefix_tree.py  | 202 ++++++---
 .../serve/cpu/deployments/test_prefix_tree.py | 397 ++++++++++++++----
 2 files changed, 455 insertions(+), 144 deletions(-)

diff --git a/python/ray/llm/_internal/serve/deployments/routers/prefix_tree.py b/python/ray/llm/_internal/serve/deployments/routers/prefix_tree.py
index 9d208edb14fd..a5de6674aacf 100644
--- a/python/ray/llm/_internal/serve/deployments/routers/prefix_tree.py
+++ b/python/ray/llm/_internal/serve/deployments/routers/prefix_tree.py
@@ -2,7 +2,7 @@
 import time
 from collections import defaultdict
 from threading import RLock
-from typing import Optional, List, Tuple, Dict
+from typing import Optional, List, Tuple, Dict, Set
 
 
 class Node:
@@ -20,8 +20,8 @@ def __init__(self, text: str = "", parent: Optional["Node"] = None):
             str, int
         ] = {}  # Maps tenant -> timestamp in ms (int)
 
-    def __str__(self) -> str:
-        return f"Node(text='{self.text}', tenants={list(self.tenant_last_access_time.keys())})"
+    def to_string(self) -> str:
+        return f"Node(text='{self.text}', parent={self.parent}, children={self.children}, tenant_last_access_time={self.tenant_last_access_time})"
 
 
 @serve.deployment(name="TreeDeployment")
@@ -36,26 +36,30 @@ class PrefixTree:
     """
 
     def __init__(self) -> None:
+        self.lock: RLock = RLock()
         self.root: Node = Node()
-        self.tenant_char_count: Dict[str, int] = defaultdict(
-            int
-        )  # Maps tenant -> character count
-        self.lock: RLock = RLock()  # For operations that need to lock the entire tree
-        self.tenant_nodes: Dict[str, List[Node]] = defaultdict(
-            list
-        )  # Maps tenant -> list of nodes it belongs to
-
-    def get_root(self) -> Node:
-        return self.root
+        self.tenants: Set[str] = set()
+        self.tenant_char_count: Dict[str, int] = {}
+        self.tenant_nodes: Dict[str, Set[Node]] = {}
 
-    def get_tenant_char_count(self) -> Dict[str, int]:
-        return self.tenant_char_count
-
-    def get_tenant_nodes(self) -> Dict[str, List[Node]]:
-        return self.tenant_nodes
+    def reset(self) -> None:
+        """Reset the tree to an empty state."""
+        with self.lock:
+            self.root = Node()
+            self.tenants = set()
+            self.tenant_char_count = {}
+            self.tenant_nodes = {}
+
+    def to_dict(self) -> Dict:
+        return {
+            "root": self.root,
+            "tenants": self.tenants,
+            "tenant_char_count": self.tenant_char_count,
+            "tenant_nodes": self.tenant_nodes
+        }
 
     def to_string(self) -> str:
-        return f"PrefixTree(root={self.root.__str__()}, tenant_char_count={self.tenant_char_count}, tenant_nodes={self.tenant_nodes})"
+        return f"PrefixTree(root={self.root.__str__()}, tenants={self.tenants}, tenant_char_count={self.tenant_char_count}, tenant_nodes={self.tenant_nodes})"
 
     @staticmethod
     def shared_prefix_count(a: str, b: str) -> int:
@@ -68,14 +72,69 @@ def shared_prefix_count(a: str, b: str) -> int:
                 break
         return i
 
-    def insert(self, text: str, tenant: str) -> None:
-        """Insert text into tree with given tenant."""
+    # def insert(self, text: str, tenant: str) -> Node:
+    #     """Insert text into tree with given tenant. Returns the node that was inserted (or the existing node if it was updated)."""
+    #     with self.lock:
+    #         if tenant not in self.tenants:
+    #             raise ValueError(f"Cannot insert text for tenant '{tenant}': tenant does not exist")
+
+    #         curr_node: Node = self.root
+    #         timestamp_ms: int = int(time.time() * 1000)
+    #         i: int = 0
+    #         while i < len(text):
+    #             self.tenant_nodes[tenant].add(curr_node)
+
+    #             first_char: str = text[i]
+    #             curr_text: str = text[i:]
+    #             if first_char not in curr_node.children:
+    #                 # No match, create new node
+    #                 # e.g. curr_node.children = {}, curr_text = "hello" -> curr_node.children = {"h": Node("hello")}
+    #                 new_node: Node = Node(text=curr_text, parent=curr_node)
+    #                 new_node.tenant_last_access_time[tenant] = timestamp_ms
+    #                 curr_node.children[first_char] = new_node
+
+    #             # Match found, check if need to split
+    #             matched_node: Node = curr_node.children[first_char]
+    #             shared_count: int = self.shared_prefix_count(
+    #                 matched_node.text, curr_text
+    #             )
+    #             if shared_count == len(matched_node.text):
+    #                 # Full match, move down the tree
+    #                 if tenant not in matched_node.tenant_last_access_time:
+    #                     self.tenant_char_count[tenant] += shared_count
+    #                 matched_node.tenant_last_access_time[tenant] = timestamp_ms
+    #                 curr_node = matched_node
+    #             else:
+    #                 # Partial match, split at matched point
+    #                 matched_text: str = matched_node.text[:shared_count]
+    #                 remaining_text: str = matched_node.text[shared_count:]
+    #                 new_parent: Node = Node(text=matched_text, parent=curr_node)
+    #                 matched_node.text = remaining_text
+    #                 matched_node.parent = new_parent
+    #                 new_parent.children[remaining_text[0]] = matched_node
+    #                 if tenant not in new_parent.tenant_last_access_time:
+    #                     self.tenant_char_count[tenant] += shared_count
+    #                 new_parent.tenant_last_access_time[tenant] = timestamp_ms
+    #                 curr_node = new_parent
+
+    #             i += shared_count
+
+    #         self.tenant_nodes[tenant].add(curr_node)
+    #         return curr_node
+            
+    def insert(self, text: str, tenant: str) -> Node:
+        """Insert text into tree with given tenant. Returns the node that was inserted (or the existing node if it was updated)."""
         with self.lock:
+            if tenant not in self.tenants:
+                raise ValueError(f"Cannot insert text for tenant '{tenant}': tenant does not exist")
+
             curr_node: Node = self.root
             timestamp_ms: int = int(time.time() * 1000)
             i: int = 0
             while i < len(text):
                 curr_node.tenant_last_access_time[tenant] = timestamp_ms
+                self.tenant_nodes[tenant].add(curr_node)
+
                 first_char: str = text[i]
                 curr_text: str = text[i:]
                 if first_char not in curr_node.children:
@@ -86,7 +145,7 @@ def insert(self, text: str, tenant: str) -> None:
 
                     # Increment char count for tenant and add node to tenant_nodes
                     self.tenant_char_count[tenant] += len(curr_text)
-                    self.tenant_nodes[tenant].append(new_node)
+                    self.tenant_nodes[tenant].add(new_node)
 
                     curr_node.children[first_char] = new_node
                 else:
@@ -125,7 +184,6 @@ def insert(self, text: str, tenant: str) -> None:
                         new_parent.tenant_last_access_time = (
                             matched_node.tenant_last_access_time.copy()
                         )
-                        self.tenant_nodes[tenant].append(new_parent)
                         # Update matched_node
                         matched_node.text = remaining_text
                         matched_node.parent = new_parent
@@ -152,14 +210,28 @@ def insert(self, text: str, tenant: str) -> None:
                         # Move down the tree
                         curr_node = matched_node
                         i += shared_count
-
+            curr_node.tenant_last_access_time[tenant] = timestamp_ms
+            self.tenant_nodes[tenant].add(curr_node)
+            return curr_node
     def prefix_match(
         self, text: str, available_tenants: Optional[List[str]] = None
     ) -> Tuple[str, Optional[List[str]]]:
         """
         Match text against tree and return (matched_text, matched_tenants).
         Does not update access time for the matched tenants (only updates when insert() is called).
+        If available_tenants is not provided, all tenants are considered.
         """
+        if available_tenants:
+            # Filter available_tenants to only include those that exist in the tree
+            available_tenants = [
+                tenant for tenant in available_tenants 
+                if tenant in self.tenants
+            ]
+            if not available_tenants:
+                return "", None
+        else:
+            available_tenants = list(self.tenants)
+
         with self.lock:
             curr_node: Node = self.root
             i: int = 0
@@ -173,12 +245,11 @@ def prefix_match(
                     matched_node: Node = curr_node.children[first_char]
 
                     # Check if any of the available tenants match this node
-                    if available_tenants:
-                        if not any(
-                            tenant in matched_node.tenant_last_access_time
-                            for tenant in available_tenants
-                        ):
-                            break
+                    if not any(
+                        tenant in matched_node.tenant_last_access_time
+                        for tenant in available_tenants
+                    ):
+                        break
 
                     shared_count: int = self.shared_prefix_count(
                         matched_node.text, curr_text
@@ -195,59 +266,65 @@ def prefix_match(
 
             # Select the tenants in available_tenants that are in the current node
             selected_tenants: Optional[List[str]] = None
-            if available_tenants:
-                matching_tenants = [
-                    tenant
-                    for tenant in available_tenants
-                    if tenant in curr_node.tenant_last_access_time
-                ]
-                if matching_tenants:
-                    selected_tenants = matching_tenants
-            else:
-                if curr_node.tenant_last_access_time:
-                    selected_tenants = list(curr_node.tenant_last_access_time)
+            matching_tenants = [
+                tenant
+                for tenant in available_tenants
+                if tenant in curr_node.tenant_last_access_time
+            ]
+            if matching_tenants:
+                selected_tenants = matching_tenants
 
             ret_text: str = text[:i]
             return ret_text, selected_tenants
 
-    def get_smallest_tenant(self) -> Optional[str]:
-        """Get the tenant with the smallest total character count."""
+
+    def add_tenant(self, tenant: str) -> None:
+        """Add a tenant to the tree."""
         with self.lock:
-            if not self.tenant_char_count:
-                return None
+            if tenant in self.tenants:
+                raise ValueError(f"Cannot add tenant '{tenant}': tenant already exists")
 
-            return min(self.tenant_char_count.items(), key=lambda x: x[1])[0]
+            self.tenants.add(tenant)
+            self.tenant_char_count[tenant] = 0
+            self.tenant_nodes[tenant] = set()
 
-    def get_tenant_char_count(self) -> Dict[str, int]:
-        """Get character count for each tenant."""
-        with self.lock:
-            return dict(self.tenant_char_count)
 
-    def remove_tenant_entirely(self, tenant: str) -> int:
-        """Remove all nodes belonging to a tenant, returns the number of characters removed. Also removes the tenant from tenant_char_count and tenant_nodes."""
+    def remove_tenant(self, tenant: str) -> int:
+        """Remove a tenant's nodes from the tree, returns the number of characters removed. Also removes the tenant from tenants, tenant_char_count, and tenant_nodes."""
         with self.lock:
+            if tenant not in self.tenants:
+                raise ValueError(f"Cannot remove tenant '{tenant}': tenant does not exist")
+
             total_chars_removed: int = 0
             for node in self.tenant_nodes[tenant].copy():
                 total_chars_removed += self.remove_tenant_single_node(tenant, node)
+            
+            self.tenants.remove(tenant)
             self.tenant_nodes.pop(tenant, None)
             self.tenant_char_count.pop(tenant, None)
             return total_chars_removed
 
+
     def remove_tenant_single_node(self, tenant: str, node: Node) -> int:
         """Remove a single node belonging to a tenant, returns the number of characters removed."""
         with self.lock:
+            if tenant not in self.tenants:
+                raise ValueError(f"Cannot remove tenant '{tenant}': tenant does not exist")
+            if node not in self.tenant_nodes[tenant] or tenant not in node.tenant_last_access_time:
+                raise ValueError(f"Cannot remove node '{node.text}' from tenant '{tenant}': tenant does not have this node")
+
             removed_chars_len: int = len(node.text)
             self.tenant_char_count[tenant] -= removed_chars_len
             self.tenant_nodes[tenant].remove(node)
             node.tenant_last_access_time.pop(tenant, None)
-
             # If this node has no more tenants, remove it from the parent
-            if not node.tenant_last_access_time:
+            if not node.tenant_last_access_time and node.parent:
                 node.parent.children.pop(node.text[0], None)
 
             return removed_chars_len
 
-    def evict_tenant(self, tenant: str, min_remove_size: int) -> int:
+
+    def evict_tenant_by_LRU(self, tenant: str, min_remove_size: int) -> int:
         """Evict nodes from a tenant until the removed character count is at least min_remove_size.
 
         Args:
@@ -259,7 +336,10 @@ def evict_tenant(self, tenant: str, min_remove_size: int) -> int:
         """
         with self.lock:
             if tenant not in self.tenant_nodes or not self.tenant_nodes[tenant]:
-                return 0
+                raise ValueError(f"Cannot evict tenant '{tenant}': tenant does not exist or has no nodes")
+
+            if self.tenant_char_count[tenant] < min_remove_size:
+                raise ValueError(f"Cannot evict tenant '{tenant}': total character count is less than min_remove_size")
 
             # Sort nodes by last access time (oldest first)
             nodes_to_evict = sorted(
@@ -267,6 +347,7 @@ def evict_tenant(self, tenant: str, min_remove_size: int) -> int:
                 key=lambda node: node.tenant_last_access_time.get(tenant, 0),
             )
 
+
             total_chars_removed: int = 0
 
             # Remove nodes until we've reached the minimum removal size
@@ -279,3 +360,12 @@ def evict_tenant(self, tenant: str, min_remove_size: int) -> int:
                     break
 
             return total_chars_removed
+
+
+    def get_smallest_tenant(self) -> Optional[str]:
+        """Get the tenant with the smallest total character count."""
+        with self.lock:
+            if not self.tenant_char_count:
+                return None
+
+            return min(self.tenant_char_count.items(), key=lambda x: x[1])[0]
\ No newline at end of file
diff --git a/python/ray/llm/tests/serve/cpu/deployments/test_prefix_tree.py b/python/ray/llm/tests/serve/cpu/deployments/test_prefix_tree.py
index b8d6eda04a8d..e5777b73ca41 100644
--- a/python/ray/llm/tests/serve/cpu/deployments/test_prefix_tree.py
+++ b/python/ray/llm/tests/serve/cpu/deployments/test_prefix_tree.py
@@ -17,131 +17,352 @@ def serve_instance():
 
 
 @pytest.mark.asyncio
-async def test_insert_and_basic_match():
-    # Deploy a clean PrefixTree
+async def test_add_tenant():
+    """Test adding tenants to the tree."""
     tree = serve.run(PrefixTree.bind())
-
-    # Insert and match exact string
-    await tree.insert.remote("hello", "tenant-A")
-    matched_text, tenants = await tree.prefix_match.remote("hello")
-    assert matched_text == "hello"
-    assert tenants == ["tenant-A"]
+    
+    # 1. Test basic tenant addition
+    await tree.reset.remote()
+    await tree.add_tenant.remote("tenant_1")
+    tree_rep = await tree.to_dict.remote()
+    assert "tenant_1" in tree_rep["tenants"]
+    assert tree_rep["tenant_char_count"]["tenant_1"] == 0
+    assert tree_rep["tenant_nodes"]["tenant_1"] == set()
+    
+    # 2. Test adding duplicate tenant raises ValueError
+    await tree.reset.remote()
+    await tree.add_tenant.remote("tenant_1")
+    with pytest.raises(ValueError):
+        await tree.add_tenant.remote("tenant_1")
 
 
 @pytest.mark.asyncio
-async def test_tree_splits_nodes_on_partial_match():
+async def test_insert():
+    """Test the insert functionality of PrefixTree."""
     tree = serve.run(PrefixTree.bind())
-    await tree.insert.remote("helloworld", "A")
-    await tree.insert.remote("hellothere", "B")
-
-    # After inserting both, the root should have one child "h"
-    root = await tree.get_root.remote()
+    
+    # 1. Test basic insertion
+    await tree.reset.remote()
+    await tree.add_tenant.remote("tenant_1")
+    await tree.insert.remote("hello", "tenant_1")
+    matched_text, tenants = await tree.prefix_match.remote("hello")
+    assert matched_text == "hello"
+    assert tenants == ["tenant_1"]
+    
+    tree_rep = await tree.to_dict.remote()
+    assert tree_rep["tenant_char_count"]["tenant_1"] == 5
+    assert len(tree_rep["tenant_nodes"]["tenant_1"]) == 2
+    
+    # 2. Test duplicate insertion doesn't double count
+    await tree.reset.remote()
+    await tree.add_tenant.remote("tenant_1")
+    await tree.add_tenant.remote("tenant_2")
+    await tree.insert.remote("foo", "tenant_1")
+    await tree.insert.remote("foo", "tenant_1")  # duplicate
+    await tree.insert.remote("bar", "tenant_2")
+    
+    tree_rep = await tree.to_dict.remote()
+    assert tree_rep["tenant_char_count"]["tenant_1"] == 3
+    assert tree_rep["tenant_char_count"]["tenant_2"] == 3
+
+    # 3. Test node splitting on partial match
+    await tree.reset.remote()
+    await tree.add_tenant.remote("tenant_1")
+    await tree.add_tenant.remote("tenant_2")
+    await tree.insert.remote("helloworld", "tenant_1")
+    await tree.insert.remote("hellothere", "tenant_2")
+    
+    tree_rep = await tree.to_dict.remote()
+    root = tree_rep["root"]
     h_node = root.children.get("h")
     assert h_node is not None
+    assert h_node.text == "hello"
+    assert h_node.children.get("w").text == "world"
+    assert h_node.children.get("t").text == "there"
+    
+    # 4. Test inserting for non-existent tenant raises ValueError
+    await tree.reset.remote()
+    with pytest.raises(ValueError):
+        await tree.insert.remote("hello", "nonexistent_tenant")
 
 
 @pytest.mark.asyncio
-async def test_no_match():
+async def test_prefix_match():
+    """Test the prefix_match functionality of PrefixTree."""
     tree = serve.run(PrefixTree.bind())
+    
+    # 1. Test no match
+    await tree.reset.remote()
     matched_text, tenants = await tree.prefix_match.remote("hello")
     assert matched_text == ""
     assert tenants is None
-
-
-@pytest.mark.asyncio
-async def test_duplicate_insertion_no_double_count():
-    tree = serve.run(PrefixTree.bind())
-
-    await tree.insert.remote("foo", "T1")
-    await tree.insert.remote("foo", "T1")  # duplicate
-
-    counts = await tree.get_tenant_char_count.remote()
-    # Should count 'foo' only once
-    assert counts.get("T1", 0) == 3
-
-
-@pytest.mark.asyncio
-async def test_shared_prefix_splitting_and_branching():
-    tree = serve.run(PrefixTree.bind())
-
-    await tree.insert.remote("helloworld", "A")
-    await tree.insert.remote("hellothere", "B")
-
+    
+    # 2. Test match with non-existing prefix returns empty string and all tenants
+    await tree.reset.remote()
+    await tree.add_tenant.remote("tenant_1")
+    await tree.add_tenant.remote("tenant_2")
+    await tree.insert.remote("hello", "tenant_1")
+    await tree.insert.remote("hellothere", "tenant_2")
+    matched_text, tenants = await tree.prefix_match.remote("foobar")
+    assert matched_text == ""
+    assert len(tenants) == 2
+    assert "tenant_1" in tenants
+    assert "tenant_2" in tenants
+    
+    # 3. Test exact match
+    await tree.reset.remote()
+    await tree.add_tenant.remote("tenant_1")
+    await tree.insert.remote("hello", "tenant_1")
+    matched_text, tenants = await tree.prefix_match.remote("hello")
+    assert matched_text == "hello"
+    assert tenants == ["tenant_1"]
+    
+    
+    # 4. Test partial match
+    await tree.reset.remote()
+    await tree.add_tenant.remote("tenant_1")
+    await tree.add_tenant.remote("tenant_2")
+    await tree.insert.remote("apple", "tenant_1")
+    await tree.insert.remote("apricot", "tenant_2")
+    text, tenants = await tree.prefix_match.remote("application")
+    assert text == "appl"
+    assert tenants == ["tenant_1"]
+    
+    # 5. Test match by tenant
+    await tree.reset.remote()
+    await tree.add_tenant.remote("tenant_1")
+    await tree.add_tenant.remote("tenant_2")
+    await tree.insert.remote("apple", "tenant_1")
+    await tree.insert.remote("apricot", "tenant_2")
+    text, tenants = await tree.prefix_match.remote("application", ["tenant_2"])
+    assert text == "ap"
+    assert tenants == ["tenant_2"]
+    
+    # 6. Test match by non-existent tenant
+    await tree.reset.remote()
+    await tree.add_tenant.remote("tenant_1")
+    await tree.add_tenant.remote("tenant_2")
+    await tree.insert.remote("apple", "tenant_1")
+    await tree.insert.remote("apricot", "tenant_2")
+    text, tenants = await tree.prefix_match.remote("application", ["tenant_3"])
+    assert text == ""
+    assert tenants is None
+    
+    # 7. Test shared prefix matching with branches
+    await tree.reset.remote()
+    await tree.add_tenant.remote("tenant_1")
+    await tree.add_tenant.remote("tenant_2")
+    await tree.insert.remote("helloworld", "tenant_1")
+    await tree.insert.remote("hellothere", "tenant_2")
     text_a, tenants_a = await tree.prefix_match.remote("helloworld")
-    text_b, tenants_b = await tree.prefix_match.remote("hellothere")
-
+    text_b, tenants_b = await tree.prefix_match.remote("hellothereworld")
     assert text_a == "helloworld"
-    assert tenants_a == ["A"]
+    assert tenants_a == ["tenant_1"]
     assert text_b == "hellothere"
-    assert tenants_b == ["B"]
+    assert tenants_b == ["tenant_2"]
 
 
 @pytest.mark.asyncio
-async def test_prefix_match_partial_and_filter():
+async def test_remove_tenant():
+    """Test removing a tenant from the tree."""
     tree = serve.run(PrefixTree.bind())
-
-    await tree.insert.remote("apple", "X")
-    await tree.insert.remote("apricot", "Y")
-
-    # Partial match for 'application' -> 'appl'
-    text, tenants = await tree.prefix_match.remote("application")
-    assert text == "appl"
-    assert tenants == ["X"]
-
-    # Filter by available_tenants=['X'] on 'apricot'
-    text_fx, tenants_fx = await tree.prefix_match.remote("apricot", ["X"])
-    assert text_fx == "ap"
-    assert tenants_fx == ["X"]
-
-    # Filter by non-existent tenant yields no tenants
-    text_fz, tenants_fz = await tree.prefix_match.remote("apricot", ["Z"])
-    assert text_fz == ""
-    assert tenants_fz is None
-
+    
+    # 1. Test basic tenant removal
+    await tree.reset.remote()
+    await tree.add_tenant.remote("tenant_1")
+    await tree.insert.remote("hello", "tenant_1")
+    removed = await tree.remove_tenant.remote("tenant_1")
+    assert removed == 5
+    
+    tree_rep = await tree.to_dict.remote()
+    assert "tenant_1" not in tree_rep["tenants"]
+    assert "tenant_1" not in tree_rep["tenant_char_count"]
+    assert "tenant_1" not in tree_rep["tenant_nodes"]
+    
+    # 2. Test removing tenant with multiple nodes
+    await tree.reset.remote()
+    await tree.add_tenant.remote("tenant_1")
+    await tree.insert.remote("cat", "tenant_1")
+    await tree.insert.remote("dog", "tenant_1")
+    removed = await tree.remove_tenant.remote("tenant_1")
+    assert removed == len("cat") + len("dog")
+    
+    # 3. Test removing non-existent tenant raises ValueError
+    await tree.reset.remote()
+    with pytest.raises(ValueError):
+        await tree.remove_tenant.remote("nonexistent_tenant")
+    
+    # 4. Test tree structure after removing tenant
+    await tree.reset.remote()
+    await tree.add_tenant.remote("tenant_1")
+    await tree.add_tenant.remote("tenant_2")
+    await tree.insert.remote("hello", "tenant_1")
+    await tree.insert.remote("hello", "tenant_2")
+
+    # Remove tenant_1, verify tenant_2 still works
+    await tree.remove_tenant.remote("tenant_1")
+    
+    tree_rep = await tree.to_dict.remote()
+    assert "tenant_1" not in tree_rep["tenants"]
+    assert "tenant_2" in tree_rep["tenants"]
+    
+    matched_text, tenants = await tree.prefix_match.remote("hello")
+    assert matched_text == "hello"
+    assert tenants == ["tenant_2"]
+    
+    # 5. Test removing the last tenant from a node removes the node
+    await tree.reset.remote()
+    await tree.add_tenant.remote("tenant_1")
+    await tree.add_tenant.remote("tenant_2")
+    await tree.insert.remote("unique1", "tenant_1")
+    await tree.insert.remote("unique2", "tenant_2")
+    
+    # Remove tenant_1
+    await tree.remove_tenant.remote("tenant_1")
+    
+    tree_rep = await tree.to_dict.remote()
+    root = tree_rep["root"]
+    # 'u' node should only have one child now ('2' from unique2)
+    assert 'u' in root.children
+    assert '2' in root.children['u'].children  # '2' from unique2
+    assert len(root.children['u'].children) == 1
+    
 
 @pytest.mark.asyncio
-async def test_remove_and_get_smallest_and_evict():
+async def test_remove_tenant_single_node():
+    """Test removing a single node for a tenant."""
     tree = serve.run(PrefixTree.bind())
+    
+
+    # # 1. Test removing a single node
+    # TEST FAILS: Ray creates new node instances when making remote calls?
+    # The node from insert.remote() is not identity-equal to the one in tenant_nodes
+
+    # await tree.reset.remote()
+    # await tree.add_tenant.remote("tenant_1")
+    # h_node = await tree.insert.remote("hello", "tenant_1")
+    
+    # removed = await tree.remove_tenant_single_node.remote("tenant_1", h_node)
+    # assert removed == 5
+    
+    # tree_rep = await tree.to_dict.remote()
+    # assert tree_rep["tenant_char_count"]["tenant_1"] == 0
+    # assert tree_rep["tenant_nodes"]["tenant_1"] == set()
+    
+    # 2. Test removing node for non-existent tenant raises ValueError
+    await tree.reset.remote()
+    await tree.add_tenant.remote("tenant_1")
+    await tree.insert.remote("hello", "tenant_1")
+    
+    tree_rep = await tree.to_dict.remote()
+    root = tree_rep["root"]
+    h_node = root.children.get("h")
+    
+    with pytest.raises(ValueError):
+        await tree.remove_tenant_single_node.remote("nonexistent_tenant", h_node)
+    
+    # 3. Test removing node that doesn't belong to tenant raises ValueError
+    await tree.reset.remote()
+    await tree.add_tenant.remote("tenant_1")
+    await tree.add_tenant.remote("tenant_2")
+    await tree.insert.remote("hello", "tenant_1")
+    
+    tree_rep = await tree.to_dict.remote()
+    root = tree_rep["root"]
+    h_node = root.children.get("h")
+    
+    with pytest.raises(ValueError):
+        await tree.remove_tenant_single_node.remote("tenant_2", h_node)
 
-    # Test removal and char count
-    await tree.insert.remote("cat", "T1")
-    await tree.insert.remote("dog", "T1")
-
-    counts = await tree.get_tenant_char_count.remote()
-    assert counts.get("T1") == len("cat") + len("dog")
-
-    # Remove entire tenant
-    removed = await tree.remove_tenant_entirely.remote("T1")
-    assert removed == len("cat") + len("dog")
-
-    counts_after = await tree.get_tenant_char_count.remote()
-    assert "T1" not in counts_after
 
-    # Test eviction LRU behavior
-    await tree.insert.remote("a", "T2")
+@pytest.mark.asyncio
+async def test_evict_tenant_by_LRU():
+    """Test the evict_tenant_by_LRU functionality of PrefixTree."""
+    tree = serve.run(PrefixTree.bind())
+    
+    # 1. Test eviction with LRU ordering
+    await tree.reset.remote()
+    await tree.add_tenant.remote("tenant_1")
+    await tree.insert.remote("a", "tenant_1")
     time.sleep(0.001)
-    await tree.insert.remote("bb", "T2")
+    await tree.insert.remote("bb", "tenant_1")
     time.sleep(0.001)
-    await tree.insert.remote("ccc", "T2")
-
-    before = (await tree.get_tenant_char_count.remote())["T2"]
-    evicted = await tree.evict_tenant.remote("T2", 2)
-    after = (await tree.get_tenant_char_count.remote())["T2"]
-
-    assert evicted >= 2
+    await tree.insert.remote("ccc", "tenant_1")
+    
+    tree_rep = await tree.to_dict.remote()
+    before = tree_rep["tenant_char_count"]["tenant_1"]
+    
+    evicted = await tree.evict_tenant_by_LRU.remote("tenant_1", 2)
+    
+    tree_rep = await tree.to_dict.remote()
+    after = tree_rep["tenant_char_count"]["tenant_1"]
+    
+    assert evicted == 3
     assert before - after == evicted
+    assert "tenant_1" in tree_rep["tenants"]
+    
+    # 2. Test eviction of non-existent tenant raises ValueError
+    await tree.reset.remote()
+    with pytest.raises(ValueError):
+        await tree.evict_tenant_by_LRU.remote("nonexistent_tenant", 5)
+
+    # 3. Test eviction of tenant with insufficient characters raises ValueError
+    await tree.reset.remote()
+    await tree.add_tenant.remote("tenant_2")
+    await tree.insert.remote("xyz", "tenant_2")
+    with pytest.raises(ValueError):
+        await tree.evict_tenant_by_LRU.remote("tenant_2", 4)
+
+    # 4. Test eviction of all tenant data
+    await tree.reset.remote()
+    await tree.add_tenant.remote("tenant_2")
+    await tree.insert.remote("xyz", "tenant_2")
+    
+    tree_rep = await tree.to_dict.remote()
+    total_size = tree_rep["tenant_char_count"]["tenant_2"]
+    
+    evicted = await tree.evict_tenant_by_LRU.remote("tenant_2", total_size)
+    assert evicted == total_size
+    
+    tree_rep = await tree.to_dict.remote()
+    assert "tenant_2" in tree_rep["tenants"]
 
 
 @pytest.mark.asyncio
 async def test_get_smallest_tenant():
+    """Test the get_smallest_tenant functionality of PrefixTree."""
     tree = serve.run(PrefixTree.bind())
-    await tree.insert.remote("aaaa", "A")
-    await tree.insert.remote("bb", "B")
+    
+    # 1. Test with empty tree
+    await tree.reset.remote()
     smallest = await tree.get_smallest_tenant.remote()
-    assert smallest == "B"
+    assert smallest is None
+    
+    # 2. Test with multiple tenants of different sizes
+    await tree.reset.remote()
+    await tree.add_tenant.remote("tenant_1")
+    await tree.add_tenant.remote("tenant_2")
+    await tree.add_tenant.remote("tenant_3")
+    await tree.insert.remote("aaaa", "tenant_1")
+    await tree.insert.remote("bb", "tenant_2")
+    await tree.insert.remote("c", "tenant_3")
+    
+    smallest = await tree.get_smallest_tenant.remote()
+    assert smallest == "tenant_3"
+    
+    # 3. Test after removing the smallest tenant
+    await tree.reset.remote()
+    await tree.add_tenant.remote("tenant_1")
+    await tree.add_tenant.remote("tenant_2")
+    await tree.add_tenant.remote("tenant_3")
+    await tree.insert.remote("aaaa", "tenant_1")
+    await tree.insert.remote("bb", "tenant_2")
+    await tree.insert.remote("c", "tenant_3")
+    await tree.remove_tenant.remote("tenant_3")
+    smallest = await tree.get_smallest_tenant.remote()
+    assert smallest == "tenant_2"
 
 
 if __name__ == "__main__":
     import sys
-
     sys.exit(pytest.main(["-v", __file__]))

From f1f6e95fcdcf14e8f92fe9c5e9ba79fe62b37bbd Mon Sep 17 00:00:00 2001
From: Justin Ji <justinji@college.harvard.edu>
Date: Fri, 2 May 2025 18:09:55 -0700
Subject: [PATCH 5/6] Linting

Signed-off-by: Justin Ji <justinji@college.harvard.edu>
---
 .../serve/deployments/routers/prefix_tree.py  |  48 ++++----
 .../serve/cpu/deployments/test_prefix_tree.py | 103 +++++++++---------
 2 files changed, 79 insertions(+), 72 deletions(-)

diff --git a/python/ray/llm/_internal/serve/deployments/routers/prefix_tree.py b/python/ray/llm/_internal/serve/deployments/routers/prefix_tree.py
index a5de6674aacf..60c76b2f00ab 100644
--- a/python/ray/llm/_internal/serve/deployments/routers/prefix_tree.py
+++ b/python/ray/llm/_internal/serve/deployments/routers/prefix_tree.py
@@ -1,6 +1,5 @@
 from ray import serve
 import time
-from collections import defaultdict
 from threading import RLock
 from typing import Optional, List, Tuple, Dict, Set
 
@@ -55,7 +54,7 @@ def to_dict(self) -> Dict:
             "root": self.root,
             "tenants": self.tenants,
             "tenant_char_count": self.tenant_char_count,
-            "tenant_nodes": self.tenant_nodes
+            "tenant_nodes": self.tenant_nodes,
         }
 
     def to_string(self) -> str:
@@ -121,12 +120,14 @@ def shared_prefix_count(a: str, b: str) -> int:
 
     #         self.tenant_nodes[tenant].add(curr_node)
     #         return curr_node
-            
+
     def insert(self, text: str, tenant: str) -> Node:
         """Insert text into tree with given tenant. Returns the node that was inserted (or the existing node if it was updated)."""
         with self.lock:
             if tenant not in self.tenants:
-                raise ValueError(f"Cannot insert text for tenant '{tenant}': tenant does not exist")
+                raise ValueError(
+                    f"Cannot insert text for tenant '{tenant}': tenant does not exist"
+                )
 
             curr_node: Node = self.root
             timestamp_ms: int = int(time.time() * 1000)
@@ -213,6 +214,7 @@ def insert(self, text: str, tenant: str) -> Node:
             curr_node.tenant_last_access_time[tenant] = timestamp_ms
             self.tenant_nodes[tenant].add(curr_node)
             return curr_node
+
     def prefix_match(
         self, text: str, available_tenants: Optional[List[str]] = None
     ) -> Tuple[str, Optional[List[str]]]:
@@ -224,8 +226,7 @@ def prefix_match(
         if available_tenants:
             # Filter available_tenants to only include those that exist in the tree
             available_tenants = [
-                tenant for tenant in available_tenants 
-                if tenant in self.tenants
+                tenant for tenant in available_tenants if tenant in self.tenants
             ]
             if not available_tenants:
                 return "", None
@@ -277,7 +278,6 @@ def prefix_match(
             ret_text: str = text[:i]
             return ret_text, selected_tenants
 
-
     def add_tenant(self, tenant: str) -> None:
         """Add a tenant to the tree."""
         with self.lock:
@@ -288,30 +288,37 @@ def add_tenant(self, tenant: str) -> None:
             self.tenant_char_count[tenant] = 0
             self.tenant_nodes[tenant] = set()
 
-
     def remove_tenant(self, tenant: str) -> int:
         """Remove a tenant's nodes from the tree, returns the number of characters removed. Also removes the tenant from tenants, tenant_char_count, and tenant_nodes."""
         with self.lock:
             if tenant not in self.tenants:
-                raise ValueError(f"Cannot remove tenant '{tenant}': tenant does not exist")
+                raise ValueError(
+                    f"Cannot remove tenant '{tenant}': tenant does not exist"
+                )
 
             total_chars_removed: int = 0
             for node in self.tenant_nodes[tenant].copy():
                 total_chars_removed += self.remove_tenant_single_node(tenant, node)
-            
+
             self.tenants.remove(tenant)
             self.tenant_nodes.pop(tenant, None)
             self.tenant_char_count.pop(tenant, None)
             return total_chars_removed
 
-
     def remove_tenant_single_node(self, tenant: str, node: Node) -> int:
         """Remove a single node belonging to a tenant, returns the number of characters removed."""
         with self.lock:
             if tenant not in self.tenants:
-                raise ValueError(f"Cannot remove tenant '{tenant}': tenant does not exist")
-            if node not in self.tenant_nodes[tenant] or tenant not in node.tenant_last_access_time:
-                raise ValueError(f"Cannot remove node '{node.text}' from tenant '{tenant}': tenant does not have this node")
+                raise ValueError(
+                    f"Cannot remove tenant '{tenant}': tenant does not exist"
+                )
+            if (
+                node not in self.tenant_nodes[tenant]
+                or tenant not in node.tenant_last_access_time
+            ):
+                raise ValueError(
+                    f"Cannot remove node '{node.text}' from tenant '{tenant}': tenant does not have this node"
+                )
 
             removed_chars_len: int = len(node.text)
             self.tenant_char_count[tenant] -= removed_chars_len
@@ -323,7 +330,6 @@ def remove_tenant_single_node(self, tenant: str, node: Node) -> int:
 
             return removed_chars_len
 
-
     def evict_tenant_by_LRU(self, tenant: str, min_remove_size: int) -> int:
         """Evict nodes from a tenant until the removed character count is at least min_remove_size.
 
@@ -336,10 +342,14 @@ def evict_tenant_by_LRU(self, tenant: str, min_remove_size: int) -> int:
         """
         with self.lock:
             if tenant not in self.tenant_nodes or not self.tenant_nodes[tenant]:
-                raise ValueError(f"Cannot evict tenant '{tenant}': tenant does not exist or has no nodes")
+                raise ValueError(
+                    f"Cannot evict tenant '{tenant}': tenant does not exist or has no nodes"
+                )
 
             if self.tenant_char_count[tenant] < min_remove_size:
-                raise ValueError(f"Cannot evict tenant '{tenant}': total character count is less than min_remove_size")
+                raise ValueError(
+                    f"Cannot evict tenant '{tenant}': total character count is less than min_remove_size"
+                )
 
             # Sort nodes by last access time (oldest first)
             nodes_to_evict = sorted(
@@ -347,7 +357,6 @@ def evict_tenant_by_LRU(self, tenant: str, min_remove_size: int) -> int:
                 key=lambda node: node.tenant_last_access_time.get(tenant, 0),
             )
 
-
             total_chars_removed: int = 0
 
             # Remove nodes until we've reached the minimum removal size
@@ -361,11 +370,10 @@ def evict_tenant_by_LRU(self, tenant: str, min_remove_size: int) -> int:
 
             return total_chars_removed
 
-
     def get_smallest_tenant(self) -> Optional[str]:
         """Get the tenant with the smallest total character count."""
         with self.lock:
             if not self.tenant_char_count:
                 return None
 
-            return min(self.tenant_char_count.items(), key=lambda x: x[1])[0]
\ No newline at end of file
+            return min(self.tenant_char_count.items(), key=lambda x: x[1])[0]
diff --git a/python/ray/llm/tests/serve/cpu/deployments/test_prefix_tree.py b/python/ray/llm/tests/serve/cpu/deployments/test_prefix_tree.py
index e5777b73ca41..f66cd4641862 100644
--- a/python/ray/llm/tests/serve/cpu/deployments/test_prefix_tree.py
+++ b/python/ray/llm/tests/serve/cpu/deployments/test_prefix_tree.py
@@ -20,7 +20,7 @@ def serve_instance():
 async def test_add_tenant():
     """Test adding tenants to the tree."""
     tree = serve.run(PrefixTree.bind())
-    
+
     # 1. Test basic tenant addition
     await tree.reset.remote()
     await tree.add_tenant.remote("tenant_1")
@@ -28,7 +28,7 @@ async def test_add_tenant():
     assert "tenant_1" in tree_rep["tenants"]
     assert tree_rep["tenant_char_count"]["tenant_1"] == 0
     assert tree_rep["tenant_nodes"]["tenant_1"] == set()
-    
+
     # 2. Test adding duplicate tenant raises ValueError
     await tree.reset.remote()
     await tree.add_tenant.remote("tenant_1")
@@ -40,7 +40,7 @@ async def test_add_tenant():
 async def test_insert():
     """Test the insert functionality of PrefixTree."""
     tree = serve.run(PrefixTree.bind())
-    
+
     # 1. Test basic insertion
     await tree.reset.remote()
     await tree.add_tenant.remote("tenant_1")
@@ -48,11 +48,11 @@ async def test_insert():
     matched_text, tenants = await tree.prefix_match.remote("hello")
     assert matched_text == "hello"
     assert tenants == ["tenant_1"]
-    
+
     tree_rep = await tree.to_dict.remote()
     assert tree_rep["tenant_char_count"]["tenant_1"] == 5
     assert len(tree_rep["tenant_nodes"]["tenant_1"]) == 2
-    
+
     # 2. Test duplicate insertion doesn't double count
     await tree.reset.remote()
     await tree.add_tenant.remote("tenant_1")
@@ -60,7 +60,7 @@ async def test_insert():
     await tree.insert.remote("foo", "tenant_1")
     await tree.insert.remote("foo", "tenant_1")  # duplicate
     await tree.insert.remote("bar", "tenant_2")
-    
+
     tree_rep = await tree.to_dict.remote()
     assert tree_rep["tenant_char_count"]["tenant_1"] == 3
     assert tree_rep["tenant_char_count"]["tenant_2"] == 3
@@ -71,7 +71,7 @@ async def test_insert():
     await tree.add_tenant.remote("tenant_2")
     await tree.insert.remote("helloworld", "tenant_1")
     await tree.insert.remote("hellothere", "tenant_2")
-    
+
     tree_rep = await tree.to_dict.remote()
     root = tree_rep["root"]
     h_node = root.children.get("h")
@@ -79,7 +79,7 @@ async def test_insert():
     assert h_node.text == "hello"
     assert h_node.children.get("w").text == "world"
     assert h_node.children.get("t").text == "there"
-    
+
     # 4. Test inserting for non-existent tenant raises ValueError
     await tree.reset.remote()
     with pytest.raises(ValueError):
@@ -90,13 +90,13 @@ async def test_insert():
 async def test_prefix_match():
     """Test the prefix_match functionality of PrefixTree."""
     tree = serve.run(PrefixTree.bind())
-    
+
     # 1. Test no match
     await tree.reset.remote()
     matched_text, tenants = await tree.prefix_match.remote("hello")
     assert matched_text == ""
     assert tenants is None
-    
+
     # 2. Test match with non-existing prefix returns empty string and all tenants
     await tree.reset.remote()
     await tree.add_tenant.remote("tenant_1")
@@ -108,7 +108,7 @@ async def test_prefix_match():
     assert len(tenants) == 2
     assert "tenant_1" in tenants
     assert "tenant_2" in tenants
-    
+
     # 3. Test exact match
     await tree.reset.remote()
     await tree.add_tenant.remote("tenant_1")
@@ -116,8 +116,7 @@ async def test_prefix_match():
     matched_text, tenants = await tree.prefix_match.remote("hello")
     assert matched_text == "hello"
     assert tenants == ["tenant_1"]
-    
-    
+
     # 4. Test partial match
     await tree.reset.remote()
     await tree.add_tenant.remote("tenant_1")
@@ -127,7 +126,7 @@ async def test_prefix_match():
     text, tenants = await tree.prefix_match.remote("application")
     assert text == "appl"
     assert tenants == ["tenant_1"]
-    
+
     # 5. Test match by tenant
     await tree.reset.remote()
     await tree.add_tenant.remote("tenant_1")
@@ -137,7 +136,7 @@ async def test_prefix_match():
     text, tenants = await tree.prefix_match.remote("application", ["tenant_2"])
     assert text == "ap"
     assert tenants == ["tenant_2"]
-    
+
     # 6. Test match by non-existent tenant
     await tree.reset.remote()
     await tree.add_tenant.remote("tenant_1")
@@ -147,7 +146,7 @@ async def test_prefix_match():
     text, tenants = await tree.prefix_match.remote("application", ["tenant_3"])
     assert text == ""
     assert tenants is None
-    
+
     # 7. Test shared prefix matching with branches
     await tree.reset.remote()
     await tree.add_tenant.remote("tenant_1")
@@ -166,19 +165,19 @@ async def test_prefix_match():
 async def test_remove_tenant():
     """Test removing a tenant from the tree."""
     tree = serve.run(PrefixTree.bind())
-    
+
     # 1. Test basic tenant removal
     await tree.reset.remote()
     await tree.add_tenant.remote("tenant_1")
     await tree.insert.remote("hello", "tenant_1")
     removed = await tree.remove_tenant.remote("tenant_1")
     assert removed == 5
-    
+
     tree_rep = await tree.to_dict.remote()
     assert "tenant_1" not in tree_rep["tenants"]
     assert "tenant_1" not in tree_rep["tenant_char_count"]
     assert "tenant_1" not in tree_rep["tenant_nodes"]
-    
+
     # 2. Test removing tenant with multiple nodes
     await tree.reset.remote()
     await tree.add_tenant.remote("tenant_1")
@@ -186,12 +185,12 @@ async def test_remove_tenant():
     await tree.insert.remote("dog", "tenant_1")
     removed = await tree.remove_tenant.remote("tenant_1")
     assert removed == len("cat") + len("dog")
-    
+
     # 3. Test removing non-existent tenant raises ValueError
     await tree.reset.remote()
     with pytest.raises(ValueError):
         await tree.remove_tenant.remote("nonexistent_tenant")
-    
+
     # 4. Test tree structure after removing tenant
     await tree.reset.remote()
     await tree.add_tenant.remote("tenant_1")
@@ -201,38 +200,37 @@ async def test_remove_tenant():
 
     # Remove tenant_1, verify tenant_2 still works
     await tree.remove_tenant.remote("tenant_1")
-    
+
     tree_rep = await tree.to_dict.remote()
     assert "tenant_1" not in tree_rep["tenants"]
     assert "tenant_2" in tree_rep["tenants"]
-    
+
     matched_text, tenants = await tree.prefix_match.remote("hello")
     assert matched_text == "hello"
     assert tenants == ["tenant_2"]
-    
+
     # 5. Test removing the last tenant from a node removes the node
     await tree.reset.remote()
     await tree.add_tenant.remote("tenant_1")
     await tree.add_tenant.remote("tenant_2")
     await tree.insert.remote("unique1", "tenant_1")
     await tree.insert.remote("unique2", "tenant_2")
-    
+
     # Remove tenant_1
     await tree.remove_tenant.remote("tenant_1")
-    
+
     tree_rep = await tree.to_dict.remote()
     root = tree_rep["root"]
     # 'u' node should only have one child now ('2' from unique2)
-    assert 'u' in root.children
-    assert '2' in root.children['u'].children  # '2' from unique2
-    assert len(root.children['u'].children) == 1
-    
+    assert "u" in root.children
+    assert "2" in root.children["u"].children  # '2' from unique2
+    assert len(root.children["u"].children) == 1
+
 
 @pytest.mark.asyncio
 async def test_remove_tenant_single_node():
     """Test removing a single node for a tenant."""
     tree = serve.run(PrefixTree.bind())
-    
 
     # # 1. Test removing a single node
     # TEST FAILS: Ray creates new node instances when making remote calls?
@@ -241,36 +239,36 @@ async def test_remove_tenant_single_node():
     # await tree.reset.remote()
     # await tree.add_tenant.remote("tenant_1")
     # h_node = await tree.insert.remote("hello", "tenant_1")
-    
+
     # removed = await tree.remove_tenant_single_node.remote("tenant_1", h_node)
     # assert removed == 5
-    
+
     # tree_rep = await tree.to_dict.remote()
     # assert tree_rep["tenant_char_count"]["tenant_1"] == 0
     # assert tree_rep["tenant_nodes"]["tenant_1"] == set()
-    
+
     # 2. Test removing node for non-existent tenant raises ValueError
     await tree.reset.remote()
     await tree.add_tenant.remote("tenant_1")
     await tree.insert.remote("hello", "tenant_1")
-    
+
     tree_rep = await tree.to_dict.remote()
     root = tree_rep["root"]
     h_node = root.children.get("h")
-    
+
     with pytest.raises(ValueError):
         await tree.remove_tenant_single_node.remote("nonexistent_tenant", h_node)
-    
+
     # 3. Test removing node that doesn't belong to tenant raises ValueError
     await tree.reset.remote()
     await tree.add_tenant.remote("tenant_1")
     await tree.add_tenant.remote("tenant_2")
     await tree.insert.remote("hello", "tenant_1")
-    
+
     tree_rep = await tree.to_dict.remote()
     root = tree_rep["root"]
     h_node = root.children.get("h")
-    
+
     with pytest.raises(ValueError):
         await tree.remove_tenant_single_node.remote("tenant_2", h_node)
 
@@ -279,7 +277,7 @@ async def test_remove_tenant_single_node():
 async def test_evict_tenant_by_LRU():
     """Test the evict_tenant_by_LRU functionality of PrefixTree."""
     tree = serve.run(PrefixTree.bind())
-    
+
     # 1. Test eviction with LRU ordering
     await tree.reset.remote()
     await tree.add_tenant.remote("tenant_1")
@@ -288,19 +286,19 @@ async def test_evict_tenant_by_LRU():
     await tree.insert.remote("bb", "tenant_1")
     time.sleep(0.001)
     await tree.insert.remote("ccc", "tenant_1")
-    
+
     tree_rep = await tree.to_dict.remote()
     before = tree_rep["tenant_char_count"]["tenant_1"]
-    
+
     evicted = await tree.evict_tenant_by_LRU.remote("tenant_1", 2)
-    
+
     tree_rep = await tree.to_dict.remote()
     after = tree_rep["tenant_char_count"]["tenant_1"]
-    
+
     assert evicted == 3
     assert before - after == evicted
     assert "tenant_1" in tree_rep["tenants"]
-    
+
     # 2. Test eviction of non-existent tenant raises ValueError
     await tree.reset.remote()
     with pytest.raises(ValueError):
@@ -317,13 +315,13 @@ async def test_evict_tenant_by_LRU():
     await tree.reset.remote()
     await tree.add_tenant.remote("tenant_2")
     await tree.insert.remote("xyz", "tenant_2")
-    
+
     tree_rep = await tree.to_dict.remote()
     total_size = tree_rep["tenant_char_count"]["tenant_2"]
-    
+
     evicted = await tree.evict_tenant_by_LRU.remote("tenant_2", total_size)
     assert evicted == total_size
-    
+
     tree_rep = await tree.to_dict.remote()
     assert "tenant_2" in tree_rep["tenants"]
 
@@ -332,12 +330,12 @@ async def test_evict_tenant_by_LRU():
 async def test_get_smallest_tenant():
     """Test the get_smallest_tenant functionality of PrefixTree."""
     tree = serve.run(PrefixTree.bind())
-    
+
     # 1. Test with empty tree
     await tree.reset.remote()
     smallest = await tree.get_smallest_tenant.remote()
     assert smallest is None
-    
+
     # 2. Test with multiple tenants of different sizes
     await tree.reset.remote()
     await tree.add_tenant.remote("tenant_1")
@@ -346,10 +344,10 @@ async def test_get_smallest_tenant():
     await tree.insert.remote("aaaa", "tenant_1")
     await tree.insert.remote("bb", "tenant_2")
     await tree.insert.remote("c", "tenant_3")
-    
+
     smallest = await tree.get_smallest_tenant.remote()
     assert smallest == "tenant_3"
-    
+
     # 3. Test after removing the smallest tenant
     await tree.reset.remote()
     await tree.add_tenant.remote("tenant_1")
@@ -365,4 +363,5 @@ async def test_get_smallest_tenant():
 
 if __name__ == "__main__":
     import sys
+
     sys.exit(pytest.main(["-v", __file__]))

From ee9568df4a8942ab505beebab19a81a52de38e55 Mon Sep 17 00:00:00 2001
From: Justin Ji <justinji@college.harvard.edu>
Date: Mon, 5 May 2025 18:10:21 -0700
Subject: [PATCH 6/6] Address comments

Signed-off-by: Justin Ji <justinji@college.harvard.edu>
---
 .../serve/deployments/routers/prefix_tree.py  | 379 --------------
 .../prefix_aware/prefix_tree.py               | 486 ++++++++++++++++++
 .../serve/cpu/deployments/test_prefix_tree.py | 168 +++---
 3 files changed, 562 insertions(+), 471 deletions(-)
 delete mode 100644 python/ray/llm/_internal/serve/deployments/routers/prefix_tree.py
 create mode 100644 python/ray/llm/_internal/serve/replica_scheduler/prefix_aware/prefix_tree.py

diff --git a/python/ray/llm/_internal/serve/deployments/routers/prefix_tree.py b/python/ray/llm/_internal/serve/deployments/routers/prefix_tree.py
deleted file mode 100644
index 60c76b2f00ab..000000000000
--- a/python/ray/llm/_internal/serve/deployments/routers/prefix_tree.py
+++ /dev/null
@@ -1,379 +0,0 @@
-from ray import serve
-import time
-from threading import RLock
-from typing import Optional, List, Tuple, Dict, Set
-
-
-class Node:
-    """
-    Node in a prefix tree that tracks tenant access time.
-
-    Each node represents a segment of text and can belong to multiple tenants.
-    """
-
-    def __init__(self, text: str = "", parent: Optional["Node"] = None):
-        self.text: str = text
-        self.parent: Optional["Node"] = parent
-        self.children: Dict[str, "Node"] = {}  # Maps char -> Node
-        self.tenant_last_access_time: Dict[
-            str, int
-        ] = {}  # Maps tenant -> timestamp in ms (int)
-
-    def to_string(self) -> str:
-        return f"Node(text='{self.text}', parent={self.parent}, children={self.children}, tenant_last_access_time={self.tenant_last_access_time})"
-
-
-@serve.deployment(name="TreeDeployment")
-class PrefixTree:
-    """
-    Thread-safe multi-tenant prefix tree (approximate radix tree).
-
-    Features:
-    1. Stores data for multiple tenants in the same tree structure
-    2. Node-level locking for concurrent access
-    3. Leaf LRU eviction based on tenant access time
-    """
-
-    def __init__(self) -> None:
-        self.lock: RLock = RLock()
-        self.root: Node = Node()
-        self.tenants: Set[str] = set()
-        self.tenant_char_count: Dict[str, int] = {}
-        self.tenant_nodes: Dict[str, Set[Node]] = {}
-
-    def reset(self) -> None:
-        """Reset the tree to an empty state."""
-        with self.lock:
-            self.root = Node()
-            self.tenants = set()
-            self.tenant_char_count = {}
-            self.tenant_nodes = {}
-
-    def to_dict(self) -> Dict:
-        return {
-            "root": self.root,
-            "tenants": self.tenants,
-            "tenant_char_count": self.tenant_char_count,
-            "tenant_nodes": self.tenant_nodes,
-        }
-
-    def to_string(self) -> str:
-        return f"PrefixTree(root={self.root.__str__()}, tenants={self.tenants}, tenant_char_count={self.tenant_char_count}, tenant_nodes={self.tenant_nodes})"
-
-    @staticmethod
-    def shared_prefix_count(a: str, b: str) -> int:
-        """Count the number of shared characters at the beginning of two strings."""
-        i: int = 0
-        for char_a, char_b in zip(a, b):
-            if char_a == char_b:
-                i += 1
-            else:
-                break
-        return i
-
-    # def insert(self, text: str, tenant: str) -> Node:
-    #     """Insert text into tree with given tenant. Returns the node that was inserted (or the existing node if it was updated)."""
-    #     with self.lock:
-    #         if tenant not in self.tenants:
-    #             raise ValueError(f"Cannot insert text for tenant '{tenant}': tenant does not exist")
-
-    #         curr_node: Node = self.root
-    #         timestamp_ms: int = int(time.time() * 1000)
-    #         i: int = 0
-    #         while i < len(text):
-    #             self.tenant_nodes[tenant].add(curr_node)
-
-    #             first_char: str = text[i]
-    #             curr_text: str = text[i:]
-    #             if first_char not in curr_node.children:
-    #                 # No match, create new node
-    #                 # e.g. curr_node.children = {}, curr_text = "hello" -> curr_node.children = {"h": Node("hello")}
-    #                 new_node: Node = Node(text=curr_text, parent=curr_node)
-    #                 new_node.tenant_last_access_time[tenant] = timestamp_ms
-    #                 curr_node.children[first_char] = new_node
-
-    #             # Match found, check if need to split
-    #             matched_node: Node = curr_node.children[first_char]
-    #             shared_count: int = self.shared_prefix_count(
-    #                 matched_node.text, curr_text
-    #             )
-    #             if shared_count == len(matched_node.text):
-    #                 # Full match, move down the tree
-    #                 if tenant not in matched_node.tenant_last_access_time:
-    #                     self.tenant_char_count[tenant] += shared_count
-    #                 matched_node.tenant_last_access_time[tenant] = timestamp_ms
-    #                 curr_node = matched_node
-    #             else:
-    #                 # Partial match, split at matched point
-    #                 matched_text: str = matched_node.text[:shared_count]
-    #                 remaining_text: str = matched_node.text[shared_count:]
-    #                 new_parent: Node = Node(text=matched_text, parent=curr_node)
-    #                 matched_node.text = remaining_text
-    #                 matched_node.parent = new_parent
-    #                 new_parent.children[remaining_text[0]] = matched_node
-    #                 if tenant not in new_parent.tenant_last_access_time:
-    #                     self.tenant_char_count[tenant] += shared_count
-    #                 new_parent.tenant_last_access_time[tenant] = timestamp_ms
-    #                 curr_node = new_parent
-
-    #             i += shared_count
-
-    #         self.tenant_nodes[tenant].add(curr_node)
-    #         return curr_node
-
-    def insert(self, text: str, tenant: str) -> Node:
-        """Insert text into tree with given tenant. Returns the node that was inserted (or the existing node if it was updated)."""
-        with self.lock:
-            if tenant not in self.tenants:
-                raise ValueError(
-                    f"Cannot insert text for tenant '{tenant}': tenant does not exist"
-                )
-
-            curr_node: Node = self.root
-            timestamp_ms: int = int(time.time() * 1000)
-            i: int = 0
-            while i < len(text):
-                curr_node.tenant_last_access_time[tenant] = timestamp_ms
-                self.tenant_nodes[tenant].add(curr_node)
-
-                first_char: str = text[i]
-                curr_text: str = text[i:]
-                if first_char not in curr_node.children:
-                    # No match, create new node
-                    # e.g. curr_node.children = {}, curr_text = "hello" -> curr_node.children = {"h": Node("hello")}
-                    new_node: Node = Node(text=curr_text, parent=curr_node)
-                    new_node.tenant_last_access_time[tenant] = timestamp_ms
-
-                    # Increment char count for tenant and add node to tenant_nodes
-                    self.tenant_char_count[tenant] += len(curr_text)
-                    self.tenant_nodes[tenant].add(new_node)
-
-                    curr_node.children[first_char] = new_node
-                else:
-                    # Match found, check if need to split
-                    matched_node: Node = curr_node.children[first_char]
-                    shared_count: int = self.shared_prefix_count(
-                        matched_node.text, curr_text
-                    )
-
-                    if shared_count < len(matched_node.text):
-                        # Partial match, split at matched point
-                        # Example:
-                        ## Before update:
-                        ### curr_node.children = {"h": Node("helloworld")}, curr_text = "hellothere" -> shared_count = 5
-                        ### matched_node = Node("helloworld")
-
-                        ## During update:
-                        ### Increment tenant_char_count[tenant] by shared_count if matched_node has not seen this tenant before
-
-                        ## After update:
-                        ### curr_node.children = {"h": Node("hello", children = {"w": Node("world")})}
-                        ### parent_node = Node("hello"), matched_node = Node("world")
-                        ### Update tenant_last_access_time for parent_node, NOT matched_node
-                        ### (new) curr_text = "there", (new) curr_node = parent_node
-                        ### Continue adding "there" to tree in next iteration
-
-                        matched_text: str = matched_node.text[:shared_count]
-                        remaining_text: str = matched_node.text[shared_count:]
-
-                        # Update tenant char count for the new split node
-                        if tenant not in matched_node.tenant_last_access_time:
-                            self.tenant_char_count[tenant] += shared_count
-
-                        # Create new parent node
-                        new_parent: Node = Node(text=matched_text, parent=curr_node)
-                        new_parent.tenant_last_access_time = (
-                            matched_node.tenant_last_access_time.copy()
-                        )
-                        # Update matched_node
-                        matched_node.text = remaining_text
-                        matched_node.parent = new_parent
-
-                        # Connect new parent node to matched_node
-                        new_parent.children[remaining_text[0]] = matched_node
-
-                        # Connect current node to new parent
-                        curr_node.children[first_char] = new_parent
-
-                        # Move down the tree
-                        curr_node = new_parent
-                        i += shared_count
-                    else:
-                        # Full match
-
-                        # Update tenant char count if this is a new tenant for this node
-                        if tenant not in matched_node.tenant_last_access_time:
-                            self.tenant_char_count[tenant] += shared_count
-
-                        # # Update tenant last access time
-                        # matched_node.tenant_last_access_time[tenant] = timestamp_ms
-
-                        # Move down the tree
-                        curr_node = matched_node
-                        i += shared_count
-            curr_node.tenant_last_access_time[tenant] = timestamp_ms
-            self.tenant_nodes[tenant].add(curr_node)
-            return curr_node
-
-    def prefix_match(
-        self, text: str, available_tenants: Optional[List[str]] = None
-    ) -> Tuple[str, Optional[List[str]]]:
-        """
-        Match text against tree and return (matched_text, matched_tenants).
-        Does not update access time for the matched tenants (only updates when insert() is called).
-        If available_tenants is not provided, all tenants are considered.
-        """
-        if available_tenants:
-            # Filter available_tenants to only include those that exist in the tree
-            available_tenants = [
-                tenant for tenant in available_tenants if tenant in self.tenants
-            ]
-            if not available_tenants:
-                return "", None
-        else:
-            available_tenants = list(self.tenants)
-
-        with self.lock:
-            curr_node: Node = self.root
-            i: int = 0
-            text_len: int = len(text)
-
-            while i < text_len:
-                first_char: str = text[i]
-                curr_text: str = text[i:]
-
-                if first_char in curr_node.children:
-                    matched_node: Node = curr_node.children[first_char]
-
-                    # Check if any of the available tenants match this node
-                    if not any(
-                        tenant in matched_node.tenant_last_access_time
-                        for tenant in available_tenants
-                    ):
-                        break
-
-                    shared_count: int = self.shared_prefix_count(
-                        matched_node.text, curr_text
-                    )
-                    i += shared_count
-                    curr_node = matched_node
-
-                    if shared_count < len(matched_node.text):
-                        # Partial match, stop here
-                        break
-                else:
-                    # No match found, stop here
-                    break
-
-            # Select the tenants in available_tenants that are in the current node
-            selected_tenants: Optional[List[str]] = None
-            matching_tenants = [
-                tenant
-                for tenant in available_tenants
-                if tenant in curr_node.tenant_last_access_time
-            ]
-            if matching_tenants:
-                selected_tenants = matching_tenants
-
-            ret_text: str = text[:i]
-            return ret_text, selected_tenants
-
-    def add_tenant(self, tenant: str) -> None:
-        """Add a tenant to the tree."""
-        with self.lock:
-            if tenant in self.tenants:
-                raise ValueError(f"Cannot add tenant '{tenant}': tenant already exists")
-
-            self.tenants.add(tenant)
-            self.tenant_char_count[tenant] = 0
-            self.tenant_nodes[tenant] = set()
-
-    def remove_tenant(self, tenant: str) -> int:
-        """Remove a tenant's nodes from the tree, returns the number of characters removed. Also removes the tenant from tenants, tenant_char_count, and tenant_nodes."""
-        with self.lock:
-            if tenant not in self.tenants:
-                raise ValueError(
-                    f"Cannot remove tenant '{tenant}': tenant does not exist"
-                )
-
-            total_chars_removed: int = 0
-            for node in self.tenant_nodes[tenant].copy():
-                total_chars_removed += self.remove_tenant_single_node(tenant, node)
-
-            self.tenants.remove(tenant)
-            self.tenant_nodes.pop(tenant, None)
-            self.tenant_char_count.pop(tenant, None)
-            return total_chars_removed
-
-    def remove_tenant_single_node(self, tenant: str, node: Node) -> int:
-        """Remove a single node belonging to a tenant, returns the number of characters removed."""
-        with self.lock:
-            if tenant not in self.tenants:
-                raise ValueError(
-                    f"Cannot remove tenant '{tenant}': tenant does not exist"
-                )
-            if (
-                node not in self.tenant_nodes[tenant]
-                or tenant not in node.tenant_last_access_time
-            ):
-                raise ValueError(
-                    f"Cannot remove node '{node.text}' from tenant '{tenant}': tenant does not have this node"
-                )
-
-            removed_chars_len: int = len(node.text)
-            self.tenant_char_count[tenant] -= removed_chars_len
-            self.tenant_nodes[tenant].remove(node)
-            node.tenant_last_access_time.pop(tenant, None)
-            # If this node has no more tenants, remove it from the parent
-            if not node.tenant_last_access_time and node.parent:
-                node.parent.children.pop(node.text[0], None)
-
-            return removed_chars_len
-
-    def evict_tenant_by_LRU(self, tenant: str, min_remove_size: int) -> int:
-        """Evict nodes from a tenant until the removed character count is at least min_remove_size.
-
-        Args:
-            tenant: The tenant to evict nodes from
-            min_remove_size: Minimum number of characters to remove
-
-        Returns:
-            int: The actual number of characters removed
-        """
-        with self.lock:
-            if tenant not in self.tenant_nodes or not self.tenant_nodes[tenant]:
-                raise ValueError(
-                    f"Cannot evict tenant '{tenant}': tenant does not exist or has no nodes"
-                )
-
-            if self.tenant_char_count[tenant] < min_remove_size:
-                raise ValueError(
-                    f"Cannot evict tenant '{tenant}': total character count is less than min_remove_size"
-                )
-
-            # Sort nodes by last access time (oldest first)
-            nodes_to_evict = sorted(
-                self.tenant_nodes[tenant],
-                key=lambda node: node.tenant_last_access_time.get(tenant, 0),
-            )
-
-            total_chars_removed: int = 0
-
-            # Remove nodes until we've reached the minimum removal size
-            for node in nodes_to_evict.copy():
-                # Use existing function to remove tenant from node
-                total_chars_removed += self.remove_tenant_single_node(tenant, node)
-
-                # Check if we've removed enough characters
-                if total_chars_removed >= min_remove_size:
-                    break
-
-            return total_chars_removed
-
-    def get_smallest_tenant(self) -> Optional[str]:
-        """Get the tenant with the smallest total character count."""
-        with self.lock:
-            if not self.tenant_char_count:
-                return None
-
-            return min(self.tenant_char_count.items(), key=lambda x: x[1])[0]
diff --git a/python/ray/llm/_internal/serve/replica_scheduler/prefix_aware/prefix_tree.py b/python/ray/llm/_internal/serve/replica_scheduler/prefix_aware/prefix_tree.py
new file mode 100644
index 000000000000..90fcb765571e
--- /dev/null
+++ b/python/ray/llm/_internal/serve/replica_scheduler/prefix_aware/prefix_tree.py
@@ -0,0 +1,486 @@
+from __future__ import annotations
+
+import heapq
+import logging
+import os
+from threading import RLock
+from typing import Dict, List, Optional, Set, Tuple
+
+from ray import serve
+
+# Logger for this module
+logger = logging.getLogger(__name__)
+
+
+class Node:
+    """
+    Node in a prefix tree that tracks tenant access time.
+
+    Each node represents a segment of text and can belong to multiple tenants.
+
+    Example tree structure:
+        Representing the strings inserted in order:
+            - "helloworld"  at time 1 by tenant_1
+            - "hellothere"  at time 2 by tenant_2
+            - "hellothomas" at time 3 by tenant_2
+
+               root: []
+         {tenant_1: 1, tenant_2: 3}
+                      |
+                   (h)|
+                      |
+                   [hello]
+          {tenant_1: 1, tenant_2: 3}
+                   /     \
+               (w)/       \\(t)
+                 /         \
+             [world]      [th]
+         {tenant_1: 1} {tenant_2: 3}
+                          /  \
+                      (e)/    \\(o)
+                        /      \
+                     [ere]   [omas]
+              {tenant_2: 2} {tenant_2: 3}
+          
+        Legend for each node:
+        - [text] = Node.text
+        - {tenant, timestamp} = Node.tenant_last_access_time
+        - (x) = edge label (first character used as key for parent's children)
+    """
+
+    def __init__(self, text: str = "", parent: Optional[Node] = None) -> None:
+        """
+        Initialize a node in the prefix tree.
+
+        Args:
+            text: The text segment this node represents
+            parent: The parent node of this node
+        """
+        self.text: str = text
+        self.parent: Optional[Node] = parent
+        self.children: Dict[str, Node] = {}  # Maps first character to child node
+        self.tenant_last_access_time: Dict[str, int] = (
+            {}
+        )  # Maps tenant ID to last access timestamp (in milliseconds)
+
+    def __repr__(self) -> str:
+        return f"Node(text='{self.text}', children={list(self.children.keys())}, tenants={list(self.tenant_last_access_time.keys())})"
+
+
+class TenantHeapNode:
+    """
+    Wrapper class for storing nodes in a min-heap, ordered by tenant access time.
+    Used for efficient LRU eviction of tenant nodes.
+    """
+
+    def __init__(self, node: Node, tenant: str) -> None:
+        """
+        Initialize a heap node for efficient LRU tenant management.
+
+        Args:
+            node: The prefix tree node this heap node refers to
+            tenant: The tenant ID this heap node is associated with
+        """
+        self.node = node
+        self.tenant_ordering_key = tenant
+
+    def __lt__(self, other: TenantHeapNode) -> bool:
+        """
+        Compare heap nodes based on tenant's last access time.
+
+        Args:
+            other: Another TenantHeapNode to compare with
+
+        Returns:
+            True if this node's tenant access time is earlier than the other's
+        """
+        return (
+            self.node.tenant_last_access_time[self.tenant_ordering_key]
+            < other.node.tenant_last_access_time[other.tenant_ordering_key]
+        )
+
+    def __repr__(self) -> str:
+        return f"TenantHeapNode(node={self.node}, tenant_ordering_key={self.tenant_ordering_key})"
+
+
+@serve.deployment(name="TreeDeployment")
+class PrefixTree:
+    """
+    Thread-safe multi-tenant prefix tree (approximate radix tree).
+
+    Features:
+    1. Stores data for multiple tenants in the same tree structure
+    2. Thread-safe with node-level locking for concurrent access
+    3. LRU eviction based on tenant access time
+    4. Efficient prefix matching across multiple tenants
+    """
+
+    def __init__(self) -> None:
+        """Initialize an empty prefix tree."""
+        self.lock: RLock = RLock()
+        self.root: Node = Node()
+        self.tenants: Set[str] = set()  # Set of tenant IDs
+        self.tenant_char_count: Dict[str, int] = (
+            {}
+        )  # Tracks total character count per tenant
+        self.tenant_nodes: Dict[str, Set[Node]] = (
+            {}
+        )  # Maps tenant ID to set of nodes belonging to that tenant
+        self.tenant_nodes_sorted: Dict[str, List[TenantHeapNode]] = (
+            {}
+        )  # Maps tenant ID to heap of nodes for LRU eviction
+
+    def reset(self) -> None:
+        """Reset the tree to an empty state."""
+        with self.lock:
+            self.root = Node()
+            self.tenants = set()
+            self.tenant_char_count = {}
+            self.tenant_nodes = {}
+            self.tenant_nodes_sorted = {}
+
+    def to_dict(self) -> Dict:
+        """
+        Convert tree to dictionary for serialization.
+
+        Returns:
+            Dictionary representation of the tree
+        """
+        return {
+            "root": self.root,
+            "tenants": self.tenants,
+            "tenant_char_count": self.tenant_char_count,
+            "tenant_nodes": self.tenant_nodes,
+            "tenant_nodes_sorted": self.tenant_nodes_sorted,
+        }
+
+    def to_string(self) -> str:
+        """String representation of the tree."""
+        return f"PrefixTree(tenants={self.tenants}, tenant_char_count={self.tenant_char_count}, tenant_nodes={self.tenant_nodes}, tenant_nodes_sorted={self.tenant_nodes_sorted})"
+
+    @staticmethod
+    def _shared_prefix_count(a: str, b: str) -> int:
+        """
+        Count the number of shared characters at the beginning of two strings.
+
+        Args:
+            a: First string
+            b: Second string
+
+        Returns:
+            Number of matching characters at the beginning
+        """
+        return len(os.path.commonprefix([a, b]))
+
+    def insert(self, text: str, tenant: str, timestamp_ms: int) -> Node:
+        """
+        Insert text into tree for a specific tenant.
+
+        If the tenant doesn't exist, it will be automatically added.
+
+        Args:
+            text: Text to insert
+            tenant: Tenant ID
+            timestamp_ms: Current timestamp in milliseconds
+
+        Returns:
+            The node that was inserted or updated
+        """
+        with self.lock:
+            if tenant not in self.tenants:
+                self._add_tenant(tenant)
+
+            curr_node: Node = self.root
+            i: int = 0
+
+            while i <= len(text):
+                # Invariant: assume curr_node has not been visited by tenant yet
+                # Update tenant info for current node
+                if tenant not in curr_node.tenant_last_access_time:
+                    self.tenant_nodes[tenant].add(curr_node)
+                    self.tenant_char_count[tenant] += len(curr_node.text)
+                    self.tenant_nodes_sorted[tenant].append(
+                        TenantHeapNode(curr_node, tenant)
+                    )
+
+                curr_node.tenant_last_access_time[tenant] = timestamp_ms
+                heapq.heapify(self.tenant_nodes_sorted[tenant])
+
+                if i == len(text):
+                    break
+
+                first_char: str = text[i]
+                curr_text: str = text[i:]
+
+                if first_char not in curr_node.children:
+                    # No match, create new node. Don't update new node as "visited" by tenant yet; it will be done in the code below.
+                    # e.g. curr_node.children = {}, curr_text = "hello" -> curr_node.children = {"h": Node("hello")}
+                    new_node: Node = Node(text=curr_text, parent=curr_node)
+                    curr_node.children[first_char] = new_node
+
+                # Match found, check if we need to split
+                matched_node: Node = curr_node.children[first_char]
+                shared_count: int = self._shared_prefix_count(
+                    matched_node.text, curr_text
+                )
+
+                if shared_count < len(matched_node.text):
+                    # Partial match, split node at matched point
+                    # Example:
+                    ## Before update:
+                    ### curr_node.children = {"h": Node("helloworld")}, curr_text = "hellothere" -> shared_count = 5
+                    ### matched_node = Node("helloworld")
+
+                    ## During update:
+                    ### Increment tenant_char_count[tenant] by shared_count if matched_node has not seen this tenant before
+
+                    ## After update:
+                    ### curr_node.children = {"h": Node("hello", children = {"w": Node("world")})}
+                    ### parent_node = Node("hello"), matched_node = Node("world")
+                    ### Update tenant_last_access_time for parent_node, NOT matched_node
+                    ### (new) curr_text = "there", (new) curr_node = parent_node
+                    ### Continue adding "there" to tree in next iteration
+
+                    matched_text: str = matched_node.text[:shared_count]
+                    remaining_text: str = matched_node.text[shared_count:]
+
+                    # Create new intermediate node
+                    new_parent: Node = Node(text=matched_text, parent=curr_node)
+                    new_parent.tenant_last_access_time = (
+                        matched_node.tenant_last_access_time.copy()
+                    )
+
+                    # Update existing matched node
+                    matched_node.text = remaining_text
+                    matched_node.parent = new_parent
+
+                    # Connect nodes
+                    new_parent.children[remaining_text[0]] = matched_node
+                    curr_node.children[first_char] = new_parent
+
+                    # Continue traversal
+                    curr_node = new_parent
+                    i += shared_count
+                else:
+                    # Full match, continue down the tree
+                    curr_node = matched_node
+                    i += shared_count
+
+            return curr_node
+
+    def prefix_match(
+        self, text: str, available_tenants: Optional[List[str]] = None
+    ) -> Tuple[str, Optional[List[str]]]:
+        """
+        Match text against tree and return matched text and matching tenants.
+
+        Args:
+            text: Text to match
+            available_tenants: List of tenants to match against (or None for all)
+
+        Returns:
+            Tuple of (matched_text, matched_tenant_ids)
+        """
+        if available_tenants:
+            # Filter available_tenants to only include those in the tree
+            available_tenants = [
+                tenant for tenant in available_tenants if tenant in self.tenants
+            ]
+            if not available_tenants:
+                return "", None
+        else:
+            available_tenants = list(self.tenants)
+
+        with self.lock:
+            curr_node: Node = self.root
+            i: int = 0
+            text_len: int = len(text)
+
+            while i < text_len:
+                first_char: str = text[i]
+                curr_text: str = text[i:]
+
+                if first_char in curr_node.children:
+                    matched_node: Node = curr_node.children[first_char]
+
+                    # Check if any available tenants match this node
+                    if not any(
+                        tenant in matched_node.tenant_last_access_time
+                        for tenant in available_tenants
+                    ):
+                        break
+
+                    shared_count: int = self._shared_prefix_count(
+                        matched_node.text, curr_text
+                    )
+                    i += shared_count
+                    curr_node = matched_node
+
+                    if shared_count < len(matched_node.text):
+                        # Partial match, stop here
+                        break
+                else:
+                    # No match found, stop here
+                    break
+
+            # Find tenants in current node that match available tenants
+            matching_tenants = [
+                tenant
+                for tenant in available_tenants
+                if tenant in curr_node.tenant_last_access_time
+            ]
+
+            selected_tenants: Optional[List[str]] = (
+                matching_tenants if matching_tenants else None
+            )
+            matched_text: str = text[:i]
+
+            return matched_text, selected_tenants
+
+    def remove_tenant(self, tenant: str) -> int:
+        """
+        Remove a tenant and all its nodes from the tree.
+
+        Args:
+            tenant: Tenant ID to remove
+
+        Returns:
+            Number of characters removed
+
+        Raises:
+            ValueError: If tenant does not exist
+        """
+        with self.lock:
+            if tenant not in self.tenants:
+                raise ValueError(
+                    f"Cannot remove tenant '{tenant}': tenant does not exist"
+                )
+
+            total_chars_removed: int = 0
+            for node in self.tenant_nodes[tenant].copy():
+                total_chars_removed += self._remove_tenant_single_node(tenant, node)
+
+            self.tenants.remove(tenant)
+            self.tenant_nodes.pop(tenant, None)
+            self.tenant_char_count.pop(tenant, None)
+            self.tenant_nodes_sorted.pop(tenant, None)
+
+            return total_chars_removed
+
+    def _remove_tenant_single_node(self, tenant: str, node: Node) -> int:
+        """
+        Remove a tenant from a single node.
+
+        Args:
+            tenant: Tenant ID to remove
+            node: Node to remove tenant from
+
+        Returns:
+            Number of characters removed
+
+        Raises:
+            ValueError: If tenant does not exist or node doesn't belong to tenant
+        """
+        with self.lock:
+            if tenant not in self.tenants:
+                raise ValueError(
+                    f"Cannot remove tenant '{tenant}': tenant does not exist"
+                )
+
+            if (
+                node not in self.tenant_nodes[tenant]
+                or tenant not in node.tenant_last_access_time
+            ):
+                raise ValueError(
+                    f"Cannot remove node '{node.text}' from tenant '{tenant}': "
+                    f"tenant does not have this node"
+                )
+
+            removed_chars_len: int = len(node.text)
+            self.tenant_char_count[tenant] -= removed_chars_len
+            self.tenant_nodes[tenant].remove(node)
+            node.tenant_last_access_time.pop(tenant, None)
+
+            # Clean up empty nodes
+            if not node.tenant_last_access_time and node.parent:
+                if (
+                    node.text and node.text[0] in node.parent.children
+                ):  # Defensive check
+                    node.parent.children.pop(node.text[0], None)
+
+            return removed_chars_len
+
+    def evict_tenant_by_LRU(self, tenant: str, min_remove_size: int) -> int:
+        """
+        Evict least recently used nodes for a tenant until minimum size is freed.
+
+        Args:
+            tenant: The tenant to evict nodes from
+            min_remove_size: Minimum number of characters to remove
+
+        Returns:
+            Actual number of characters removed
+
+        Raises:
+            ValueError: If tenant doesn't exist or has insufficient nodes
+        """
+        with self.lock:
+            if tenant not in self.tenant_nodes or not self.tenant_nodes[tenant]:
+                raise ValueError(
+                    f"Cannot evict tenant '{tenant}': tenant does not exist or has no nodes"
+                )
+
+            if self.tenant_char_count[tenant] < min_remove_size:
+                raise ValueError(
+                    f"Cannot evict tenant '{tenant}': total character count "
+                    f"({self.tenant_char_count[tenant]}) is less than min_remove_size "
+                    f"({min_remove_size})"
+                )
+
+            total_chars_removed: int = 0
+
+            # Directly use the tenant's priority queue
+            while (
+                total_chars_removed < min_remove_size
+                and self.tenant_nodes_sorted[tenant]
+            ):
+                heap_node: TenantHeapNode = heapq.heappop(
+                    self.tenant_nodes_sorted[tenant]
+                )
+                total_chars_removed += self._remove_tenant_single_node(
+                    tenant, heap_node.node
+                )
+
+            return total_chars_removed
+
+    def get_smallest_tenant(self) -> Optional[str]:
+        """
+        Get the tenant with the smallest total character count.
+
+        Returns:
+            Tenant ID with smallest character count, or None if no tenants
+        """
+        with self.lock:
+            if not self.tenant_char_count:
+                return None
+
+        return min(self.tenant_char_count, key=self.tenant_char_count.get, default=None)
+
+    def _add_tenant(self, tenant: str) -> None:
+        """
+        Add a new tenant to the tree.
+
+        If the tenant already exists, this is a no-op with a warning log.
+
+        Args:
+            tenant: Tenant ID to add
+        """
+        with self.lock:
+            if tenant in self.tenants:
+                logger.warning(f"Tenant '{tenant}' already exists. No action taken.")
+                return
+
+            self.tenants.add(tenant)
+            self.tenant_char_count[tenant] = 0
+            self.tenant_nodes[tenant] = set()
+            self.tenant_nodes_sorted[tenant] = []
diff --git a/python/ray/llm/tests/serve/cpu/deployments/test_prefix_tree.py b/python/ray/llm/tests/serve/cpu/deployments/test_prefix_tree.py
index f66cd4641862..f9214546af86 100644
--- a/python/ray/llm/tests/serve/cpu/deployments/test_prefix_tree.py
+++ b/python/ray/llm/tests/serve/cpu/deployments/test_prefix_tree.py
@@ -3,7 +3,9 @@
 import ray
 from ray import serve
 
-from ray.llm._internal.serve.deployments.routers.prefix_tree import PrefixTree
+from ray.llm._internal.serve.replica_scheduler.prefix_aware.prefix_tree import (
+    PrefixTree,
+)
 
 
 @pytest.fixture(scope="module", autouse=True)
@@ -18,22 +20,25 @@ def serve_instance():
 
 @pytest.mark.asyncio
 async def test_add_tenant():
-    """Test adding tenants to the tree."""
+    """Test adding tenants to the tree via the private _add_tenant method."""
     tree = serve.run(PrefixTree.bind())
 
     # 1. Test basic tenant addition
     await tree.reset.remote()
-    await tree.add_tenant.remote("tenant_1")
+    await tree._add_tenant.remote("tenant_1")
     tree_rep = await tree.to_dict.remote()
     assert "tenant_1" in tree_rep["tenants"]
     assert tree_rep["tenant_char_count"]["tenant_1"] == 0
     assert tree_rep["tenant_nodes"]["tenant_1"] == set()
 
-    # 2. Test adding duplicate tenant raises ValueError
+    # 2. Test adding duplicate tenant logs warning but doesn't raise error
     await tree.reset.remote()
-    await tree.add_tenant.remote("tenant_1")
-    with pytest.raises(ValueError):
-        await tree.add_tenant.remote("tenant_1")
+    await tree._add_tenant.remote("tenant_1")
+    # This should not raise an error now
+    await tree._add_tenant.remote("tenant_1")
+    # Verify the tenant still exists
+    tree_rep = await tree.to_dict.remote()
+    assert "tenant_1" in tree_rep["tenants"]
 
 
 @pytest.mark.asyncio
@@ -43,8 +48,8 @@ async def test_insert():
 
     # 1. Test basic insertion
     await tree.reset.remote()
-    await tree.add_tenant.remote("tenant_1")
-    await tree.insert.remote("hello", "tenant_1")
+    # No need to call add_tenant first - insert will do it automatically
+    await tree.insert.remote("hello", "tenant_1", int(time.time() * 1000))
     matched_text, tenants = await tree.prefix_match.remote("hello")
     assert matched_text == "hello"
     assert tenants == ["tenant_1"]
@@ -55,11 +60,10 @@ async def test_insert():
 
     # 2. Test duplicate insertion doesn't double count
     await tree.reset.remote()
-    await tree.add_tenant.remote("tenant_1")
-    await tree.add_tenant.remote("tenant_2")
-    await tree.insert.remote("foo", "tenant_1")
-    await tree.insert.remote("foo", "tenant_1")  # duplicate
-    await tree.insert.remote("bar", "tenant_2")
+    # Insert automatically adds tenants
+    await tree.insert.remote("foo", "tenant_1", int(time.time() * 1000))
+    await tree.insert.remote("foo", "tenant_1", int(time.time() * 1000))  # duplicate
+    await tree.insert.remote("bar", "tenant_2", int(time.time() * 1000))
 
     tree_rep = await tree.to_dict.remote()
     assert tree_rep["tenant_char_count"]["tenant_1"] == 3
@@ -67,10 +71,8 @@ async def test_insert():
 
     # 3. Test node splitting on partial match
     await tree.reset.remote()
-    await tree.add_tenant.remote("tenant_1")
-    await tree.add_tenant.remote("tenant_2")
-    await tree.insert.remote("helloworld", "tenant_1")
-    await tree.insert.remote("hellothere", "tenant_2")
+    await tree.insert.remote("helloworld", "tenant_1", int(time.time() * 1000))
+    await tree.insert.remote("hellothere", "tenant_2", int(time.time() * 1000))
 
     tree_rep = await tree.to_dict.remote()
     root = tree_rep["root"]
@@ -80,10 +82,15 @@ async def test_insert():
     assert h_node.children.get("w").text == "world"
     assert h_node.children.get("t").text == "there"
 
-    # 4. Test inserting for non-existent tenant raises ValueError
+    # 4. Test inserting for non-existent tenant automatically adds the tenant
     await tree.reset.remote()
-    with pytest.raises(ValueError):
-        await tree.insert.remote("hello", "nonexistent_tenant")
+    # This should not raise an error now
+    await tree.insert.remote("hello", "nonexistent_tenant", int(time.time() * 1000))
+
+    # Verify the tenant was added
+    tree_rep = await tree.to_dict.remote()
+    assert "nonexistent_tenant" in tree_rep["tenants"]
+    assert tree_rep["tenant_char_count"]["nonexistent_tenant"] == 5
 
 
 @pytest.mark.asyncio
@@ -91,18 +98,16 @@ async def test_prefix_match():
     """Test the prefix_match functionality of PrefixTree."""
     tree = serve.run(PrefixTree.bind())
 
-    # 1. Test no match
-    await tree.reset.remote()
-    matched_text, tenants = await tree.prefix_match.remote("hello")
-    assert matched_text == ""
-    assert tenants is None
+    # # 1. Test no match
+    # await tree.reset.remote()
+    # matched_text, tenants = await tree.prefix_match.remote("hello")
+    # assert matched_text == ""
+    # assert tenants is None
 
     # 2. Test match with non-existing prefix returns empty string and all tenants
     await tree.reset.remote()
-    await tree.add_tenant.remote("tenant_1")
-    await tree.add_tenant.remote("tenant_2")
-    await tree.insert.remote("hello", "tenant_1")
-    await tree.insert.remote("hellothere", "tenant_2")
+    await tree.insert.remote("hello", "tenant_1", int(time.time() * 1000))
+    await tree.insert.remote("hellothere", "tenant_2", int(time.time() * 1000))
     matched_text, tenants = await tree.prefix_match.remote("foobar")
     assert matched_text == ""
     assert len(tenants) == 2
@@ -111,48 +116,39 @@ async def test_prefix_match():
 
     # 3. Test exact match
     await tree.reset.remote()
-    await tree.add_tenant.remote("tenant_1")
-    await tree.insert.remote("hello", "tenant_1")
+    await tree.insert.remote("hello", "tenant_1", int(time.time() * 1000))
     matched_text, tenants = await tree.prefix_match.remote("hello")
     assert matched_text == "hello"
     assert tenants == ["tenant_1"]
 
     # 4. Test partial match
     await tree.reset.remote()
-    await tree.add_tenant.remote("tenant_1")
-    await tree.add_tenant.remote("tenant_2")
-    await tree.insert.remote("apple", "tenant_1")
-    await tree.insert.remote("apricot", "tenant_2")
+    await tree.insert.remote("apple", "tenant_1", int(time.time() * 1000))
+    await tree.insert.remote("apricot", "tenant_2", int(time.time() * 1000))
     text, tenants = await tree.prefix_match.remote("application")
     assert text == "appl"
     assert tenants == ["tenant_1"]
 
     # 5. Test match by tenant
     await tree.reset.remote()
-    await tree.add_tenant.remote("tenant_1")
-    await tree.add_tenant.remote("tenant_2")
-    await tree.insert.remote("apple", "tenant_1")
-    await tree.insert.remote("apricot", "tenant_2")
+    await tree.insert.remote("apple", "tenant_1", int(time.time() * 1000))
+    await tree.insert.remote("apricot", "tenant_2", int(time.time() * 1000))
     text, tenants = await tree.prefix_match.remote("application", ["tenant_2"])
     assert text == "ap"
     assert tenants == ["tenant_2"]
 
     # 6. Test match by non-existent tenant
     await tree.reset.remote()
-    await tree.add_tenant.remote("tenant_1")
-    await tree.add_tenant.remote("tenant_2")
-    await tree.insert.remote("apple", "tenant_1")
-    await tree.insert.remote("apricot", "tenant_2")
+    await tree.insert.remote("apple", "tenant_1", int(time.time() * 1000))
+    await tree.insert.remote("apricot", "tenant_2", int(time.time() * 1000))
     text, tenants = await tree.prefix_match.remote("application", ["tenant_3"])
     assert text == ""
     assert tenants is None
 
     # 7. Test shared prefix matching with branches
     await tree.reset.remote()
-    await tree.add_tenant.remote("tenant_1")
-    await tree.add_tenant.remote("tenant_2")
-    await tree.insert.remote("helloworld", "tenant_1")
-    await tree.insert.remote("hellothere", "tenant_2")
+    await tree.insert.remote("helloworld", "tenant_1", int(time.time() * 1000))
+    await tree.insert.remote("hellothere", "tenant_2", int(time.time() * 1000))
     text_a, tenants_a = await tree.prefix_match.remote("helloworld")
     text_b, tenants_b = await tree.prefix_match.remote("hellothereworld")
     assert text_a == "helloworld"
@@ -168,8 +164,7 @@ async def test_remove_tenant():
 
     # 1. Test basic tenant removal
     await tree.reset.remote()
-    await tree.add_tenant.remote("tenant_1")
-    await tree.insert.remote("hello", "tenant_1")
+    await tree.insert.remote("hello", "tenant_1", int(time.time() * 1000))
     removed = await tree.remove_tenant.remote("tenant_1")
     assert removed == 5
 
@@ -180,9 +175,8 @@ async def test_remove_tenant():
 
     # 2. Test removing tenant with multiple nodes
     await tree.reset.remote()
-    await tree.add_tenant.remote("tenant_1")
-    await tree.insert.remote("cat", "tenant_1")
-    await tree.insert.remote("dog", "tenant_1")
+    await tree.insert.remote("cat", "tenant_1", int(time.time() * 1000))
+    await tree.insert.remote("dog", "tenant_1", int(time.time() * 1000))
     removed = await tree.remove_tenant.remote("tenant_1")
     assert removed == len("cat") + len("dog")
 
@@ -193,10 +187,8 @@ async def test_remove_tenant():
 
     # 4. Test tree structure after removing tenant
     await tree.reset.remote()
-    await tree.add_tenant.remote("tenant_1")
-    await tree.add_tenant.remote("tenant_2")
-    await tree.insert.remote("hello", "tenant_1")
-    await tree.insert.remote("hello", "tenant_2")
+    await tree.insert.remote("hello", "tenant_1", int(time.time() * 1000))
+    await tree.insert.remote("hello", "tenant_2", int(time.time() * 1000))
 
     # Remove tenant_1, verify tenant_2 still works
     await tree.remove_tenant.remote("tenant_1")
@@ -211,10 +203,8 @@ async def test_remove_tenant():
 
     # 5. Test removing the last tenant from a node removes the node
     await tree.reset.remote()
-    await tree.add_tenant.remote("tenant_1")
-    await tree.add_tenant.remote("tenant_2")
-    await tree.insert.remote("unique1", "tenant_1")
-    await tree.insert.remote("unique2", "tenant_2")
+    await tree.insert.remote("unique1", "tenant_1", int(time.time() * 1000))
+    await tree.insert.remote("unique2", "tenant_2", int(time.time() * 1000))
 
     # Remove tenant_1
     await tree.remove_tenant.remote("tenant_1")
@@ -228,7 +218,7 @@ async def test_remove_tenant():
 
 
 @pytest.mark.asyncio
-async def test_remove_tenant_single_node():
+async def test__remove_tenant_single_node():
     """Test removing a single node for a tenant."""
     tree = serve.run(PrefixTree.bind())
 
@@ -237,10 +227,10 @@ async def test_remove_tenant_single_node():
     # The node from insert.remote() is not identity-equal to the one in tenant_nodes
 
     # await tree.reset.remote()
-    # await tree.add_tenant.remote("tenant_1")
-    # h_node = await tree.insert.remote("hello", "tenant_1")
+    # await tree.insert.remote("hello", "tenant_1", int(time.time() * 1000))
+    # h_node = await tree.insert.remote("hello", "tenant_1", int(time.time() * 1000))
 
-    # removed = await tree.remove_tenant_single_node.remote("tenant_1", h_node)
+    # removed = await tree._remove_tenant_single_node.remote("tenant_1", h_node)
     # assert removed == 5
 
     # tree_rep = await tree.to_dict.remote()
@@ -249,28 +239,26 @@ async def test_remove_tenant_single_node():
 
     # 2. Test removing node for non-existent tenant raises ValueError
     await tree.reset.remote()
-    await tree.add_tenant.remote("tenant_1")
-    await tree.insert.remote("hello", "tenant_1")
+    await tree.insert.remote("hello", "tenant_1", int(time.time() * 1000))
 
     tree_rep = await tree.to_dict.remote()
     root = tree_rep["root"]
     h_node = root.children.get("h")
 
     with pytest.raises(ValueError):
-        await tree.remove_tenant_single_node.remote("nonexistent_tenant", h_node)
+        await tree._remove_tenant_single_node.remote("nonexistent_tenant", h_node)
 
     # 3. Test removing node that doesn't belong to tenant raises ValueError
     await tree.reset.remote()
-    await tree.add_tenant.remote("tenant_1")
-    await tree.add_tenant.remote("tenant_2")
-    await tree.insert.remote("hello", "tenant_1")
+    await tree.insert.remote("hello", "tenant_1", int(time.time() * 1000))
+    await tree.insert.remote("world", "tenant_2", int(time.time() * 1000))
 
     tree_rep = await tree.to_dict.remote()
     root = tree_rep["root"]
     h_node = root.children.get("h")
 
     with pytest.raises(ValueError):
-        await tree.remove_tenant_single_node.remote("tenant_2", h_node)
+        await tree._remove_tenant_single_node.remote("tenant_2", h_node)
 
 
 @pytest.mark.asyncio
@@ -280,12 +268,14 @@ async def test_evict_tenant_by_LRU():
 
     # 1. Test eviction with LRU ordering
     await tree.reset.remote()
-    await tree.add_tenant.remote("tenant_1")
-    await tree.insert.remote("a", "tenant_1")
+    current_time = int(time.time() * 1000)
+    await tree.insert.remote("a", "tenant_1", current_time)
     time.sleep(0.001)
-    await tree.insert.remote("bb", "tenant_1")
+    current_time = int(time.time() * 1000)
+    await tree.insert.remote("bb", "tenant_1", current_time)
     time.sleep(0.001)
-    await tree.insert.remote("ccc", "tenant_1")
+    current_time = int(time.time() * 1000)
+    await tree.insert.remote("ccc", "tenant_1", current_time)
 
     tree_rep = await tree.to_dict.remote()
     before = tree_rep["tenant_char_count"]["tenant_1"]
@@ -306,15 +296,13 @@ async def test_evict_tenant_by_LRU():
 
     # 3. Test eviction of tenant with insufficient characters raises ValueError
     await tree.reset.remote()
-    await tree.add_tenant.remote("tenant_2")
-    await tree.insert.remote("xyz", "tenant_2")
+    await tree.insert.remote("xyz", "tenant_2", int(time.time() * 1000))
     with pytest.raises(ValueError):
         await tree.evict_tenant_by_LRU.remote("tenant_2", 4)
 
     # 4. Test eviction of all tenant data
     await tree.reset.remote()
-    await tree.add_tenant.remote("tenant_2")
-    await tree.insert.remote("xyz", "tenant_2")
+    await tree.insert.remote("xyz", "tenant_2", int(time.time() * 1000))
 
     tree_rep = await tree.to_dict.remote()
     total_size = tree_rep["tenant_char_count"]["tenant_2"]
@@ -338,24 +326,20 @@ async def test_get_smallest_tenant():
 
     # 2. Test with multiple tenants of different sizes
     await tree.reset.remote()
-    await tree.add_tenant.remote("tenant_1")
-    await tree.add_tenant.remote("tenant_2")
-    await tree.add_tenant.remote("tenant_3")
-    await tree.insert.remote("aaaa", "tenant_1")
-    await tree.insert.remote("bb", "tenant_2")
-    await tree.insert.remote("c", "tenant_3")
+    current_time = int(time.time() * 1000)
+    await tree.insert.remote("aaaa", "tenant_1", current_time)
+    await tree.insert.remote("bb", "tenant_2", current_time)
+    await tree.insert.remote("c", "tenant_3", current_time)
 
     smallest = await tree.get_smallest_tenant.remote()
     assert smallest == "tenant_3"
 
     # 3. Test after removing the smallest tenant
     await tree.reset.remote()
-    await tree.add_tenant.remote("tenant_1")
-    await tree.add_tenant.remote("tenant_2")
-    await tree.add_tenant.remote("tenant_3")
-    await tree.insert.remote("aaaa", "tenant_1")
-    await tree.insert.remote("bb", "tenant_2")
-    await tree.insert.remote("c", "tenant_3")
+    current_time = int(time.time() * 1000)
+    await tree.insert.remote("aaaa", "tenant_1", current_time)
+    await tree.insert.remote("bb", "tenant_2", current_time)
+    await tree.insert.remote("c", "tenant_3", current_time)
     await tree.remove_tenant.remote("tenant_3")
     smallest = await tree.get_smallest_tenant.remote()
     assert smallest == "tenant_2"