launch_router.py 17.5 KB
Newer Older
1
2
import argparse
import dataclasses
3
import logging
4
import sys
5
from typing import Dict, List, Optional
6
7
8
9
10

from sglang_router import Router
from sglang_router_rs import PolicyType


11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
def setup_logger():
    logger = logging.getLogger("router")
    logger.setLevel(logging.INFO)

    formatter = logging.Formatter(
        "[Router (Python)] %(asctime)s - %(levelname)s - %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S",
    )

    handler = logging.StreamHandler()
    handler.setFormatter(formatter)
    logger.addHandler(handler)

    return logger


27
28
29
@dataclasses.dataclass
class RouterArgs:
    # Worker configuration
30
    worker_urls: List[str] = dataclasses.field(default_factory=list)
31
32
33
    host: str = "127.0.0.1"
    port: int = 30000

34
    # PD-specific configuration
35
    pd_disaggregation: bool = False  # Enable PD disaggregated mode
36
37
38
39
40
    prefill_urls: List[tuple] = dataclasses.field(
        default_factory=list
    )  # List of (url, bootstrap_port)
    decode_urls: List[str] = dataclasses.field(default_factory=list)

41
42
    # Routing policy
    policy: str = "cache_aware"
43
    worker_startup_timeout_secs: int = 300
44
    worker_startup_check_interval: int = 10
45
    cache_threshold: float = 0.5
46
47
    balance_abs_threshold: int = 32
    balance_rel_threshold: float = 1.0001
48
49
    eviction_interval: int = 60
    max_tree_size: int = 2**24
50
    max_payload_size: int = 256 * 1024 * 1024  # 256MB default for large batches
51
    log_dir: Optional[str] = None
52
    log_level: Optional[str] = None
53
54
55
56
57
    # Service discovery configuration
    service_discovery: bool = False
    selector: Dict[str, str] = dataclasses.field(default_factory=dict)
    service_discovery_port: int = 80
    service_discovery_namespace: Optional[str] = None
58
59
60
61
    # PD service discovery configuration
    prefill_selector: Dict[str, str] = dataclasses.field(default_factory=dict)
    decode_selector: Dict[str, str] = dataclasses.field(default_factory=dict)
    bootstrap_port_annotation: str = "sglang.ai/bootstrap-port"
62
63
64
    # Prometheus configuration
    prometheus_port: Optional[int] = None
    prometheus_host: Optional[str] = None
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108

    @staticmethod
    def add_cli_args(
        parser: argparse.ArgumentParser,
        use_router_prefix: bool = False,
        exclude_host_port: bool = False,
    ):
        """
        Add router-specific arguments to an argument parser.

        Args:
            parser: The argument parser to add arguments to
            use_router_prefix: If True, prefix all arguments with 'router-' to avoid conflicts
            exclude_host_port: If True, don't add host and port arguments (used when inheriting from server)
        """
        prefix = "router-" if use_router_prefix else ""

        # Worker configuration
        if not exclude_host_port:
            parser.add_argument(
                "--host",
                type=str,
                default=RouterArgs.host,
                help="Host address to bind the router server",
            )
            parser.add_argument(
                "--port",
                type=int,
                default=RouterArgs.port,
                help="Port number to bind the router server",
            )

        parser.add_argument(
            "--worker-urls",
            type=str,
            nargs="+",
            help="List of worker URLs (e.g., http://worker1:8000 http://worker2:8000)",
        )

        # Routing policy configuration
        parser.add_argument(
            f"--{prefix}policy",
            type=str,
            default=RouterArgs.policy,
109
110
111
112
113
114
            choices=["random", "round_robin", "cache_aware", "power_of_two"],
            help="Load balancing policy to use. Note: power_of_two is only available in PD disaggregated mode",
        )

        # PD-specific arguments
        parser.add_argument(
115
            f"--{prefix}pd-disaggregation",
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
            action="store_true",
            help="Enable PD (Prefill-Decode) disaggregated mode",
        )
        parser.add_argument(
            f"--{prefix}prefill",
            nargs=2,
            action="append",
            metavar=("URL", "BOOTSTRAP_PORT"),
            help="Prefill server URL and bootstrap port. Can be specified multiple times. BOOTSTRAP_PORT can be 'none' for no bootstrap port.",
        )
        parser.add_argument(
            f"--{prefix}decode",
            nargs=1,
            action="append",
            metavar=("URL",),
            help="Decode server URL. Can be specified multiple times.",
132
        )
133
134
135
136
137
138
        parser.add_argument(
            f"--{prefix}worker-startup-timeout-secs",
            type=int,
            default=RouterArgs.worker_startup_timeout_secs,
            help="Timeout in seconds for worker startup",
        )
139
140
141
142
143
144
        parser.add_argument(
            f"--{prefix}worker-startup-check-interval",
            type=int,
            default=RouterArgs.worker_startup_check_interval,
            help="Interval in seconds between checks for worker startup",
        )
145
146
147
148
149
150
151
        parser.add_argument(
            f"--{prefix}cache-threshold",
            type=float,
            default=RouterArgs.cache_threshold,
            help="Cache threshold (0.0-1.0) for cache-aware routing",
        )
        parser.add_argument(
152
153
154
155
156
157
158
            f"--{prefix}balance-abs-threshold",
            type=int,
            default=RouterArgs.balance_abs_threshold,
            help="Load balancing is triggered when (max_load - min_load) > abs_threshold AND max_load > min_load * rel_threshold. Otherwise, use cache aware",
        )
        parser.add_argument(
            f"--{prefix}balance-rel-threshold",
159
            type=float,
160
161
            default=RouterArgs.balance_rel_threshold,
            help="Load balancing is triggered when (max_load - min_load) > abs_threshold AND max_load > min_load * rel_threshold. Otherwise, use cache aware",
162
163
164
165
166
167
168
169
170
171
172
173
174
        )
        parser.add_argument(
            f"--{prefix}eviction-interval",
            type=int,
            default=RouterArgs.eviction_interval,
            help="Interval in seconds between cache eviction operations",
        )
        parser.add_argument(
            f"--{prefix}max-tree-size",
            type=int,
            default=RouterArgs.max_tree_size,
            help="Maximum size of the approximation tree for cache-aware routing",
        )
175
176
177
178
179
180
        parser.add_argument(
            f"--{prefix}max-payload-size",
            type=int,
            default=RouterArgs.max_payload_size,
            help="Maximum payload size in bytes",
        )
181
182
183
184
185
186
        parser.add_argument(
            f"--{prefix}log-dir",
            type=str,
            default=None,
            help="Directory to store log files. If not specified, logs are only output to console.",
        )
187
188
189
190
191
192
193
        parser.add_argument(
            f"--{prefix}log-level",
            type=str,
            default="info",
            choices=["debug", "info", "warning", "error", "critical"],
            help="Set the logging level. If not specified, defaults to INFO.",
        )
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
        parser.add_argument(
            f"--{prefix}service-discovery",
            action="store_true",
            help="Enable Kubernetes service discovery",
        )
        parser.add_argument(
            f"--{prefix}selector",
            type=str,
            nargs="+",
            help="Label selector for Kubernetes service discovery (format: key1=value1 key2=value2)",
        )
        parser.add_argument(
            f"--{prefix}service-discovery-port",
            type=int,
            default=RouterArgs.service_discovery_port,
            help="Port to use for discovered worker pods",
        )
        parser.add_argument(
            f"--{prefix}service-discovery-namespace",
            type=str,
            help="Kubernetes namespace to watch for pods. If not provided, watches all namespaces (requires cluster-wide permissions)",
        )
216
217
218
219
220
221
222
223
224
225
226
227
        parser.add_argument(
            f"--{prefix}prefill-selector",
            type=str,
            nargs="+",
            help="Label selector for prefill server pods in PD mode (format: key1=value1 key2=value2)",
        )
        parser.add_argument(
            f"--{prefix}decode-selector",
            type=str,
            nargs="+",
            help="Label selector for decode server pods in PD mode (format: key1=value1 key2=value2)",
        )
228
229
230
231
232
233
234
235
236
237
238
239
240
        # Prometheus configuration
        parser.add_argument(
            f"--{prefix}prometheus-port",
            type=int,
            default=29000,
            help="Port to expose Prometheus metrics. If not specified, Prometheus metrics are disabled",
        )
        parser.add_argument(
            f"--{prefix}prometheus-host",
            type=str,
            default="127.0.0.1",
            help="Host address to bind the Prometheus metrics server",
        )
241
242
243
244
245
246
247
248
249
250
251
252
253

    @classmethod
    def from_cli_args(
        cls, args: argparse.Namespace, use_router_prefix: bool = False
    ) -> "RouterArgs":
        """
        Create RouterArgs instance from parsed command line arguments.

        Args:
            args: Parsed command line arguments
            use_router_prefix: If True, look for arguments with 'router-' prefix
        """
        prefix = "router_" if use_router_prefix else ""
254
255
256
257
258
259
        worker_urls = getattr(args, "worker_urls", [])

        # Parse PD URLs
        prefill_urls = cls._parse_prefill_urls(getattr(args, f"{prefix}prefill", None))
        decode_urls = cls._parse_decode_urls(getattr(args, f"{prefix}decode", None))

260
        return cls(
261
            worker_urls=worker_urls,
262
263
            host=args.host,
            port=args.port,
264
            pd_disaggregation=getattr(args, f"{prefix}pd_disaggregation", False),
265
266
            prefill_urls=prefill_urls,
            decode_urls=decode_urls,
267
            policy=getattr(args, f"{prefix}policy"),
268
269
270
            worker_startup_timeout_secs=getattr(
                args, f"{prefix}worker_startup_timeout_secs"
            ),
271
272
273
            worker_startup_check_interval=getattr(
                args, f"{prefix}worker_startup_check_interval"
            ),
274
            cache_threshold=getattr(args, f"{prefix}cache_threshold"),
275
276
            balance_abs_threshold=getattr(args, f"{prefix}balance_abs_threshold"),
            balance_rel_threshold=getattr(args, f"{prefix}balance_rel_threshold"),
277
278
            eviction_interval=getattr(args, f"{prefix}eviction_interval"),
            max_tree_size=getattr(args, f"{prefix}max_tree_size"),
279
            max_payload_size=getattr(args, f"{prefix}max_payload_size"),
280
            log_dir=getattr(args, f"{prefix}log_dir", None),
281
            log_level=getattr(args, f"{prefix}log_level", None),
282
283
284
285
286
287
            service_discovery=getattr(args, f"{prefix}service_discovery", False),
            selector=cls._parse_selector(getattr(args, f"{prefix}selector", None)),
            service_discovery_port=getattr(args, f"{prefix}service_discovery_port"),
            service_discovery_namespace=getattr(
                args, f"{prefix}service_discovery_namespace", None
            ),
288
289
290
291
292
293
294
            prefill_selector=cls._parse_selector(
                getattr(args, f"{prefix}prefill_selector", None)
            ),
            decode_selector=cls._parse_selector(
                getattr(args, f"{prefix}decode_selector", None)
            ),
            bootstrap_port_annotation="sglang.ai/bootstrap-port",  # Mooncake-specific annotation
295
296
            prometheus_port=getattr(args, f"{prefix}prometheus_port", None),
            prometheus_host=getattr(args, f"{prefix}prometheus_host", None),
297
298
        )

299
300
301
302
303
304
305
306
307
308
309
310
    @staticmethod
    def _parse_selector(selector_list):
        if not selector_list:
            return {}

        selector = {}
        for item in selector_list:
            if "=" in item:
                key, value = item.split("=", 1)
                selector[key] = value
        return selector

311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
    @staticmethod
    def _parse_prefill_urls(prefill_list):
        """Parse prefill URLs from --prefill arguments.

        Format: --prefill URL BOOTSTRAP_PORT
        Example: --prefill http://prefill1:8080 9000 --prefill http://prefill2:8080 none
        """
        if not prefill_list:
            return []

        prefill_urls = []
        for url, bootstrap_port_str in prefill_list:
            # Handle 'none' as None
            if bootstrap_port_str.lower() == "none":
                bootstrap_port = None
            else:
                try:
                    bootstrap_port = int(bootstrap_port_str)
                except ValueError:
                    raise ValueError(
                        f"Invalid bootstrap port: {bootstrap_port_str}. Must be a number or 'none'"
                    )

            prefill_urls.append((url, bootstrap_port))

        return prefill_urls

    @staticmethod
    def _parse_decode_urls(decode_list):
        """Parse decode URLs from --decode arguments.

        Format: --decode URL
        Example: --decode http://decode1:8081 --decode http://decode2:8081
        """
        if not decode_list:
            return []

        # decode_list is a list of single-element lists due to nargs=1
        return [url[0] for url in decode_list]

351
352
353
354
355
356
357

def policy_from_str(policy_str: str) -> PolicyType:
    """Convert policy string to PolicyType enum."""
    policy_map = {
        "random": PolicyType.Random,
        "round_robin": PolicyType.RoundRobin,
        "cache_aware": PolicyType.CacheAware,
358
        "power_of_two": PolicyType.PowerOfTwo,
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
    }
    return policy_map[policy_str]


def launch_router(args: argparse.Namespace) -> Optional[Router]:
    """
    Launch the SGLang router with the configuration from parsed arguments.

    Args:
        args: Namespace object containing router configuration
            Can be either raw argparse.Namespace or converted RouterArgs

    Returns:
        Router instance if successful, None if failed
    """
374
    logger = logging.getLogger("router")
375
376
377
378
379
380
381
    try:
        # Convert to RouterArgs if needed
        if not isinstance(args, RouterArgs):
            router_args = RouterArgs.from_cli_args(args)
        else:
            router_args = args

382
        # Validate configuration based on mode
383
384
385
386
387
388
389
        if router_args.pd_disaggregation:
            # Validate PD configuration - skip URL requirements if using service discovery
            if not router_args.service_discovery:
                if not router_args.prefill_urls:
                    raise ValueError("PD disaggregation mode requires --prefill")
                if not router_args.decode_urls:
                    raise ValueError("PD disaggregation mode requires --decode")
390
391

        # Create router with unified constructor
392
        router = Router(
393
            worker_urls=(
394
395
396
                []
                if router_args.service_discovery or router_args.pd_disaggregation
                else router_args.worker_urls
397
            ),
398
399
            host=router_args.host,
            port=router_args.port,
400
401
            policy=policy_from_str(router_args.policy),
            worker_startup_timeout_secs=router_args.worker_startup_timeout_secs,
402
            worker_startup_check_interval=router_args.worker_startup_check_interval,
403
            cache_threshold=router_args.cache_threshold,
404
405
            balance_abs_threshold=router_args.balance_abs_threshold,
            balance_rel_threshold=router_args.balance_rel_threshold,
406
407
            eviction_interval_secs=router_args.eviction_interval,
            max_tree_size=router_args.max_tree_size,
408
            max_payload_size=router_args.max_payload_size,
409
            log_dir=router_args.log_dir,
410
            log_level=router_args.log_level,
411
412
413
414
            service_discovery=router_args.service_discovery,
            selector=router_args.selector,
            service_discovery_port=router_args.service_discovery_port,
            service_discovery_namespace=router_args.service_discovery_namespace,
415
416
            prefill_selector=router_args.prefill_selector,
            decode_selector=router_args.decode_selector,
417
418
            prometheus_port=router_args.prometheus_port,
            prometheus_host=router_args.prometheus_host,
419
            pd_disaggregation=router_args.pd_disaggregation,
420
            prefill_urls=(
421
                router_args.prefill_urls if router_args.pd_disaggregation else None
422
423
            ),
            decode_urls=(
424
                router_args.decode_urls if router_args.pd_disaggregation else None
425
            ),
426
427
428
429
430
431
        )

        router.start()
        return router

    except Exception as e:
432
        logger.error(f"Error starting router: {e}")
433
        raise e
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453


class CustomHelpFormatter(
    argparse.RawDescriptionHelpFormatter, argparse.ArgumentDefaultsHelpFormatter
):
    """Custom formatter that preserves both description formatting and shows defaults"""

    pass


def parse_router_args(args: List[str]) -> RouterArgs:
    """Parse command line arguments and return RouterArgs instance."""
    parser = argparse.ArgumentParser(
        description="""SGLang Router - High-performance request distribution across worker nodes

Usage:
This launcher enables starting a router with individual worker instances. It is useful for
multi-node setups or when you want to start workers and router separately.

Examples:
454
  # Regular mode
455
  python -m sglang_router.launch_router --worker-urls http://worker1:8000 http://worker2:8000
456
457

  # PD disaggregated mode
458
  python -m sglang_router.launch_router --pd-disaggregation \\
459
460
461
    --prefill http://prefill1:8000 9000 --prefill http://prefill2:8000 none \\
    --decode http://decode1:8001 --decode http://decode2:8001 \\
    --policy cache_aware
462
463
464
465
466
467
468
469
470
471
472

    """,
        formatter_class=CustomHelpFormatter,
    )

    RouterArgs.add_cli_args(parser, use_router_prefix=False)
    return RouterArgs.from_cli_args(parser.parse_args(args), use_router_prefix=False)


def main() -> None:
    router_args = parse_router_args(sys.argv[1:])
473
    launch_router(router_args)
474
475
476
477


if __name__ == "__main__":
    main()