Slime Router

class StringRadixTrie:
    """
    String-based Radix Trie for efficient prefix matching and token caching.
    Features:
    - Efficient string prefix matching
    - Token ID caching for matched prefixes
    - Thread-safe operations
    - Weight version tracking
    - Automatic garbage collection based on weight version thresholds
    """
    
    def __init__(self, max_cache_size: int = 10000, gc_threshold_k: int = 5, 
                 tokenizer=None, verbose: bool = False):
        self.max_cache_size = max_cache_size
        self.gc_threshold_k = gc_threshold_k
        self.tokenizer = tokenizer
        self.root = StringTreeNode()

def find_longest_prefix(self, text: str) -> MatchResult:
    """Find the longest cached prefix for the given text."""
    matched_tokens = []
    matched_logp = []
    matched_loss_mask = []
    current_node = self.root
    remaining_text = text
    
    while remaining_text:
        # Find best matching child
        for child_node in current_node.children:
            if remaining_text.startswith(child_node.string_key):
                current_node = child_node
                matched_tokens.extend(child_node.token_ids)
                matched_logp.extend(child_node.logp)
                matched_loss_mask.extend(child_node.loss_mask)
                break
    
    return MatchResult(matched_prefix, matched_tokens, matched_logp, 
                      matched_loss_mask, remaining_text, current_node)

class RadixTreeMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next):
        if path != "/generate":
            return await call_next(request)
        
        # Retrieve cached tokens for input text
        input_tokens, input_logprobs, input_loss_mask = \
            self.radix_tree.retrieve_from_text(input_text, return_logprob=True)
        
        # Get response from SGLang
        response = await call_next(request)
        
        # Cache the full trajectory (input + output)
        if "output_token_logprobs" in response_data.get("meta_info", {}):
            full_token_ids = input_tokens + generated_token_ids
            full_logprobs = input_logprobs + generated_token_logprobs
            self.radix_tree.insert(full_text, full_token_ids, full_logprobs, 
                                  full_loss_mask, weight_version)

class SlimeRouter:
    def __init__(self, args, verbose=False):
        self.app = FastAPI()
        
        # URL -> Active Request Count (load state)
        self.worker_request_counts: dict[str, int] = {}
        # URL -> Consecutive Failures
        self.worker_failure_counts: dict[str, int] = {}
        # Quarantined workers excluded from routing pool
        self.dead_workers: set[str] = set()
        
        # Setup HTTP client with connection pooling
        self.client = httpx.AsyncClient(
            limits=httpx.Limits(max_connections=max_connections),
            timeout=httpx.Timeout(timeout),
        )
        
        # Load middleware plugins
        for middleware_path in args.slime_router_middleware_paths or []:
            middleware = load_function(middleware_path)
            self.app.add_middleware(middleware, router=self)

def _use_url(self):
    """Select worker URL with minimal active requests."""
    if not self.dead_workers:
        url = min(self.worker_request_counts, key=self.worker_request_counts.get)
    else:
        valid_workers = (w for w in self.worker_request_counts 
                        if w not in self.dead_workers)
        url = min(valid_workers, key=self.worker_request_counts.get)
    
    self.worker_request_counts[url] += 1
    return url

async def _health_check_loop(self):
    while True:
        await asyncio.sleep(interval)
        
        urls = [u for u in self.worker_request_counts if u not in self.dead_workers]
        results = await asyncio.gather(*(
            self._check_worker_health(url) for url in urls
        ))
        
        for url, is_healthy in results:
            if not is_healthy:
                failures = self.worker_failure_counts.get(url, 0) + 1
                self.worker_failure_counts[url] = failures
                
                if failures >= threshold:
                    logger.warning(f"Worker {url} failed {threshold} checks. Marking DEAD.")
                    self.dead_workers.add(url)
            else:
                self.worker_failure_counts[url] = 0