router.py 5.55 KB
Newer Older
1
from typing import Dict, List, Optional
2
3
4
5
6
7
8
9
10
11

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:
12
13
        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'])
14
15
16
        policy: Load balancing policy to use. Options:
            - PolicyType.Random: Randomly select workers
            - PolicyType.RoundRobin: Distribute requests in round-robin fashion
17
            - PolicyType.CacheAware: Distribute requests based on cache state and load balance
18
            - PolicyType.PowerOfTwo: Select best of two random workers based on load (PD mode only)
19
20
        host: Host address to bind the router server. Default: '127.0.0.1'
        port: Port number to bind the router server. Default: 3001
21
        worker_startup_timeout_secs: Timeout in seconds for worker startup. Default: 300
22
        worker_startup_check_interval: Interval in seconds between checks for worker initialization. Default: 10
23
24
25
        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
26
27
28
29
        balance_abs_threshold: Load balancing is triggered when (max_load - min_load) > abs_threshold
            AND max_load > min_load * rel_threshold. Otherwise, use cache aware. Default: 32
        balance_rel_threshold: Load balancing is triggered when (max_load - min_load) > abs_threshold
            AND max_load > min_load * rel_threshold. Otherwise, use cache aware. Default: 1.0001
30
31
        eviction_interval_secs: Interval in seconds between cache eviction operations in cache-aware
            routing. Default: 60
32
        max_payload_size: Maximum payload size in bytes. Default: 256MB
33
        max_tree_size: Maximum size of the approximation tree for cache-aware routing. Default: 2^24
34
        verbose: Enable verbose logging. Default: False
35
        log_dir: Directory to store log files. If None, logs are only output to console. Default: None
36
37
38
39
40
41
42
43
        service_discovery: Enable Kubernetes service discovery. When enabled, the router will
            automatically discover worker pods based on the selector. Default: False
        selector: Dictionary mapping of label keys to values for Kubernetes pod selection.
            Example: {"app": "sglang-worker"}. Default: {}
        service_discovery_port: Port to use for service discovery. The router will generate
            worker URLs using this port. Default: 80
        service_discovery_namespace: Kubernetes namespace to watch for pods. If not provided,
            watches pods across all namespaces (requires cluster-wide permissions). Default: None
44
45
        prometheus_port: Port to expose Prometheus metrics. Default: None
        prometheus_host: Host address to bind the Prometheus metrics server. Default: None
46
47
48
        pd_disaggregated: Enable PD (Prefill-Decode) disaggregated mode. Default: False
        prefill_urls: List of (url, bootstrap_port) tuples for prefill servers (PD mode only)
        decode_urls: List of URLs for decode servers (PD mode only)
49
50
51
52
53
54
55
56
    """

    def __init__(
        self,
        worker_urls: List[str],
        policy: PolicyType = PolicyType.RoundRobin,
        host: str = "127.0.0.1",
        port: int = 3001,
57
        worker_startup_timeout_secs: int = 300,
58
        worker_startup_check_interval: int = 10,
59
        cache_threshold: float = 0.50,
60
61
        balance_abs_threshold: int = 32,
        balance_rel_threshold: float = 1.0001,
62
63
        eviction_interval_secs: int = 60,
        max_tree_size: int = 2**24,
64
        max_payload_size: int = 256 * 1024 * 1024,  # 256MB
65
        verbose: bool = False,
66
        log_dir: Optional[str] = None,
67
68
69
70
        service_discovery: bool = False,
        selector: Dict[str, str] = None,
        service_discovery_port: int = 80,
        service_discovery_namespace: Optional[str] = None,
71
72
        prometheus_port: Optional[int] = None,
        prometheus_host: Optional[str] = None,
73
74
75
        pd_disaggregated: bool = False,
        prefill_urls: Optional[List[tuple]] = None,
        decode_urls: Optional[List[str]] = None,
76
    ):
77
78
79
        if selector is None:
            selector = {}

80
81
82
83
84
        self._router = _Router(
            worker_urls=worker_urls,
            policy=policy,
            host=host,
            port=port,
85
            worker_startup_timeout_secs=worker_startup_timeout_secs,
86
            worker_startup_check_interval=worker_startup_check_interval,
87
            cache_threshold=cache_threshold,
88
89
            balance_abs_threshold=balance_abs_threshold,
            balance_rel_threshold=balance_rel_threshold,
90
91
            eviction_interval_secs=eviction_interval_secs,
            max_tree_size=max_tree_size,
92
            max_payload_size=max_payload_size,
93
            verbose=verbose,
94
            log_dir=log_dir,
95
96
97
98
            service_discovery=service_discovery,
            selector=selector,
            service_discovery_port=service_discovery_port,
            service_discovery_namespace=service_discovery_namespace,
99
100
            prometheus_port=prometheus_port,
            prometheus_host=prometheus_host,
101
102
103
            pd_disaggregated=pd_disaggregated,
            prefill_urls=prefill_urls,
            decode_urls=decode_urls,
104
105
106
107
108
109
110
111
        )

    def start(self) -> None:
        """Start the router server.

        This method blocks until the server is shut down.
        """
        self._router.start()