59 lines
2.4 KiB
Python
59 lines
2.4 KiB
Python
from typing import List, Optional
|
|
|
|
from sglang_router_rs import PolicyType
|
|
from sglang_router_rs import Router as _Router
|
|
|
|
|
|
class Router:
|
|
"""
|
|
A high-performance router for distributing requests across worker nodes.
|
|
|
|
Args:
|
|
worker_urls: List of URLs for worker nodes that will handle requests. Each URL should include
|
|
the protocol, host, and port (e.g., ['http://worker1:8000', 'http://worker2:8000'])
|
|
policy: Load balancing policy to use. Options:
|
|
- PolicyType.Random: Randomly select workers
|
|
- PolicyType.RoundRobin: Distribute requests in round-robin fashion
|
|
- PolicyType.CacheAware: Distribute requests in cache-aware fashion
|
|
host: Host address to bind the router server. Default: '127.0.0.1'
|
|
port: Port number to bind the router server. Default: 3001
|
|
cache_threshold: Cache threshold (0.0-1.0) for cache-aware routing. Routes to cached worker
|
|
if the match rate exceeds threshold, otherwise routes to the worker with the smallest
|
|
tree. Default: 0.5
|
|
cache_routing_prob: Probability of using cache-aware routing (0.0-1.0). Default 1.0 for
|
|
full cache-aware routing, suitable for perfectly divided prefix workloads. For uneven
|
|
workloads, use a lower value to better distribute requests
|
|
eviction_interval_secs: Interval in seconds between cache eviction operations in cache-aware
|
|
routing. Default: 60
|
|
max_tree_size: Maximum size of the approximation tree for cache-aware routing. Default: 2^24
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
worker_urls: List[str],
|
|
policy: PolicyType = PolicyType.RoundRobin,
|
|
host: str = "127.0.0.1",
|
|
port: int = 3001,
|
|
cache_threshold: float = 0.50,
|
|
cache_routing_prob: float = 1.0,
|
|
eviction_interval_secs: int = 60,
|
|
max_tree_size: int = 2**24,
|
|
):
|
|
self._router = _Router(
|
|
worker_urls=worker_urls,
|
|
policy=policy,
|
|
host=host,
|
|
port=port,
|
|
cache_threshold=cache_threshold,
|
|
cache_routing_prob=cache_routing_prob,
|
|
eviction_interval_secs=eviction_interval_secs,
|
|
max_tree_size=max_tree_size,
|
|
)
|
|
|
|
def start(self) -> None:
|
|
"""Start the router server.
|
|
|
|
This method blocks until the server is shut down.
|
|
"""
|
|
self._router.start()
|