chore: vendor sglang v0.5.10 snapshot
This commit is contained in:
3
third_party/sglang/sgl-model-gateway/bindings/python/src/sglang_router/__init__.py
vendored
Normal file
3
third_party/sglang/sgl-model-gateway/bindings/python/src/sglang_router/__init__.py
vendored
Normal file
@@ -0,0 +1,3 @@
|
||||
from sglang_router.version import __version__
|
||||
|
||||
__all__ = ["__version__"]
|
||||
8
third_party/sglang/sgl-model-gateway/bindings/python/src/sglang_router/__main__.py
vendored
Normal file
8
third_party/sglang/sgl-model-gateway/bindings/python/src/sglang_router/__main__.py
vendored
Normal file
@@ -0,0 +1,8 @@
|
||||
"""
|
||||
Allow running the CLI via: python -m sglang_router
|
||||
"""
|
||||
|
||||
from sglang_router.cli import main
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
107
third_party/sglang/sgl-model-gateway/bindings/python/src/sglang_router/cli.py
vendored
Executable file
107
third_party/sglang/sgl-model-gateway/bindings/python/src/sglang_router/cli.py
vendored
Executable file
@@ -0,0 +1,107 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
SGLang Model Gateway CLI
|
||||
|
||||
Provides convenient command-line interface for launching the router and server.
|
||||
|
||||
Usage:
|
||||
smg launch [args] # Launch router only
|
||||
smg server [args] # Launch router + server
|
||||
smg --help # Show help
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import sys
|
||||
from typing import List, Optional
|
||||
|
||||
from sglang_router.sglang_router_rs import (
|
||||
get_verbose_version_string,
|
||||
get_version_string,
|
||||
)
|
||||
|
||||
|
||||
def create_parser() -> argparse.ArgumentParser:
|
||||
"""Create the main CLI parser with subcommands."""
|
||||
prog_name = os.path.basename(sys.argv[0]) if sys.argv else "smg"
|
||||
parser = argparse.ArgumentParser(
|
||||
prog=prog_name,
|
||||
description="SGLang Model Gateway - High-performance inference router",
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
)
|
||||
|
||||
subparsers = parser.add_subparsers(dest="command", help="Available commands")
|
||||
|
||||
# Launch router subcommand
|
||||
launch_parser = subparsers.add_parser(
|
||||
"launch",
|
||||
help="Launch router only (requires existing worker URLs)",
|
||||
description="Launch the SGLang router with existing worker instances",
|
||||
add_help=False, # Let router handle --help
|
||||
)
|
||||
|
||||
# Launch server + router subcommand
|
||||
server_parser = subparsers.add_parser(
|
||||
"server",
|
||||
help="Launch router and server processes together",
|
||||
description="Launch both SGLang router and server processes",
|
||||
add_help=False, # Let server handle --help
|
||||
)
|
||||
|
||||
return parser
|
||||
|
||||
|
||||
def main(argv: Optional[List[str]] = None) -> None:
|
||||
"""Main CLI entry point."""
|
||||
if argv is None:
|
||||
argv = sys.argv[1:]
|
||||
|
||||
# Handle version flags before parsing
|
||||
if argv and argv[0] in ["--version", "-V", "--version-verbose"]:
|
||||
if argv[0] == "--version-verbose":
|
||||
print(get_verbose_version_string())
|
||||
else:
|
||||
print(get_version_string())
|
||||
sys.exit(0)
|
||||
|
||||
# Handle empty command - show help
|
||||
if not argv or argv[0] not in ["launch", "server", "-h", "--help"]:
|
||||
parser = create_parser()
|
||||
parser.print_help()
|
||||
sys.exit(1)
|
||||
|
||||
parser = create_parser()
|
||||
args, unknown = parser.parse_known_args(argv)
|
||||
|
||||
if args.command == "launch":
|
||||
# Import and call launch_router functions directly
|
||||
from sglang_router.launch_router import launch_router, parse_router_args
|
||||
|
||||
# All router args are in unknown
|
||||
router_args = parse_router_args(unknown)
|
||||
launch_router(router_args)
|
||||
|
||||
elif args.command == "server":
|
||||
# Import and call launch_server main with proper argv
|
||||
# Note: launch_server.main() uses argparse internally which reads sys.argv
|
||||
# We need to temporarily set sys.argv for compatibility
|
||||
import sglang_router.launch_server as launch_server_module
|
||||
|
||||
# Preserve original sys.argv
|
||||
original_argv = sys.argv
|
||||
try:
|
||||
# All server args are in unknown
|
||||
prog_name = os.path.basename(sys.argv[0]) if sys.argv else "smg"
|
||||
sys.argv = [f"{prog_name} server"] + unknown
|
||||
launch_server_module.main()
|
||||
finally:
|
||||
# Restore original sys.argv
|
||||
sys.argv = original_argv
|
||||
|
||||
else:
|
||||
parser.print_help()
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
109
third_party/sglang/sgl-model-gateway/bindings/python/src/sglang_router/launch_router.py
vendored
Normal file
109
third_party/sglang/sgl-model-gateway/bindings/python/src/sglang_router/launch_router.py
vendored
Normal file
@@ -0,0 +1,109 @@
|
||||
import argparse
|
||||
import logging
|
||||
import sys
|
||||
from typing import List, Optional
|
||||
|
||||
import setproctitle
|
||||
from sglang_router.mini_lb import MiniLoadBalancer
|
||||
from sglang_router.router_args import RouterArgs
|
||||
|
||||
logger = logging.getLogger("router")
|
||||
|
||||
try:
|
||||
from sglang_router.router import Router
|
||||
except ImportError:
|
||||
Router = None
|
||||
logger.warning(
|
||||
"Rust Router is not installed, only python MiniLB (debugging only) is available"
|
||||
)
|
||||
|
||||
|
||||
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
|
||||
"""
|
||||
setproctitle.setproctitle("sglang::router")
|
||||
try:
|
||||
# Convert to RouterArgs if needed
|
||||
if not isinstance(args, RouterArgs):
|
||||
router_args = RouterArgs.from_cli_args(args)
|
||||
else:
|
||||
router_args = args
|
||||
|
||||
if router_args.mini_lb:
|
||||
mini_lb = MiniLoadBalancer(router_args)
|
||||
mini_lb.start()
|
||||
else:
|
||||
if Router is None:
|
||||
raise RuntimeError("Rust Router is not installed")
|
||||
router_args._validate_router_args()
|
||||
router = Router.from_args(router_args)
|
||||
router.start()
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error starting router: {e}")
|
||||
raise e
|
||||
|
||||
|
||||
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:
|
||||
# Regular mode
|
||||
python -m sglang_router.launch_router --worker-urls http://worker1:8000 http://worker2:8000
|
||||
|
||||
# PD disaggregated mode with same policy for both
|
||||
python -m sglang_router.launch_router --pd-disaggregation \\
|
||||
--prefill http://prefill1:8000 9000 --prefill http://prefill2:8000 \\
|
||||
--decode http://decode1:8001 --decode http://decode2:8001 \\
|
||||
--policy cache_aware
|
||||
|
||||
# PD mode with optional bootstrap ports
|
||||
python -m sglang_router.launch_router --pd-disaggregation \\
|
||||
--prefill http://prefill1:8000 9000 \\ # With bootstrap port
|
||||
--prefill http://prefill2:8000 none \\ # Explicitly no bootstrap port
|
||||
--prefill http://prefill3:8000 \\ # Defaults to no bootstrap port
|
||||
--decode http://decode1:8001 --decode http://decode2:8001
|
||||
|
||||
# PD mode with different policies for prefill and decode
|
||||
python -m sglang_router.launch_router --pd-disaggregation \\
|
||||
--prefill http://prefill1:8000 --prefill http://prefill2:8000 \\
|
||||
--decode http://decode1:8001 --decode http://decode2:8001 \\
|
||||
--prefill-policy cache_aware --decode-policy power_of_two
|
||||
|
||||
""",
|
||||
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:])
|
||||
launch_router(router_args)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
213
third_party/sglang/sgl-model-gateway/bindings/python/src/sglang_router/launch_server.py
vendored
Normal file
213
third_party/sglang/sgl-model-gateway/bindings/python/src/sglang_router/launch_server.py
vendored
Normal file
@@ -0,0 +1,213 @@
|
||||
import argparse
|
||||
import asyncio
|
||||
import copy
|
||||
import logging
|
||||
import multiprocessing as mp
|
||||
import os
|
||||
import random
|
||||
import signal
|
||||
import sys
|
||||
import time
|
||||
from typing import List
|
||||
|
||||
import requests
|
||||
from setproctitle import setproctitle
|
||||
from sglang_router.launch_router import RouterArgs, launch_router
|
||||
|
||||
from sglang.srt.server_args import ServerArgs
|
||||
from sglang.srt.utils.network import is_port_available
|
||||
|
||||
|
||||
def setup_logger():
|
||||
logger = logging.getLogger("router")
|
||||
logger.setLevel(logging.INFO)
|
||||
|
||||
formatter = logging.Formatter(
|
||||
"[Router (Python)] %(asctime)s - %(levelname)s - %(message)s - %(filename)s:%(lineno)d",
|
||||
datefmt="%Y-%m-%d %H:%M:%S",
|
||||
)
|
||||
|
||||
handler = logging.StreamHandler()
|
||||
handler.setFormatter(formatter)
|
||||
logger.addHandler(handler)
|
||||
|
||||
return logger
|
||||
|
||||
|
||||
logger = setup_logger()
|
||||
|
||||
|
||||
# Create new process group
|
||||
def run_server(server_args, dp_rank):
|
||||
"""
|
||||
Note:
|
||||
|
||||
1. Without os.setpgrp(), all processes share the same PGID. When you press Ctrl+C, the terminal sends SIGINT to all processes in the group simultaneously.
|
||||
This can cause leaf processes to terminate first, which messes up the cleaning order and produces orphaned processes.
|
||||
|
||||
Terminal (PGID=100)
|
||||
└── Main Python Process (PGID=100)
|
||||
└── Server Process 1 (PGID=100)
|
||||
└── Scheduler 1
|
||||
└── Detokenizer 1
|
||||
└── Server Process 2 (PGID=100)
|
||||
└── Scheduler 2
|
||||
└── Detokenizer 2
|
||||
|
||||
2. With os.setpgrp(), the main Python process and its children are in a separate group. Now:
|
||||
|
||||
Terminal (PGID=100)
|
||||
└── Main Python Process (PGID=200)
|
||||
└── Server Process 1 (PGID=300)
|
||||
└── Scheduler 1
|
||||
└── Detokenizer 1
|
||||
└── Server Process 2 (PGID=400)
|
||||
└── Scheduler 2
|
||||
└── Detokenizer 2
|
||||
"""
|
||||
# create new process group
|
||||
os.setpgrp()
|
||||
|
||||
setproctitle("sglang::server")
|
||||
# Set SGLANG_DP_RANK environment variable
|
||||
os.environ["SGLANG_DP_RANK"] = str(dp_rank)
|
||||
|
||||
# Launch server in appropriate mode (HTTP or gRPC)
|
||||
if server_args.grpc_mode:
|
||||
from sglang.srt.entrypoints.grpc_server import serve_grpc
|
||||
|
||||
asyncio.run(serve_grpc(server_args))
|
||||
else:
|
||||
from sglang.srt.entrypoints.http_server import launch_server
|
||||
|
||||
launch_server(server_args)
|
||||
|
||||
|
||||
def launch_server_process(
|
||||
server_args: ServerArgs, worker_port: int, dp_id: int
|
||||
) -> mp.Process:
|
||||
"""Launch a single server process with the given args and port."""
|
||||
server_args = copy.deepcopy(server_args)
|
||||
server_args.port = worker_port
|
||||
server_args.base_gpu_id = dp_id * server_args.tp_size
|
||||
server_args.dp_size = 1
|
||||
|
||||
proc = mp.Process(target=run_server, args=(server_args, dp_id))
|
||||
proc.start()
|
||||
return proc
|
||||
|
||||
|
||||
def wait_for_server_health(host: str, port: int, timeout: int = 300) -> bool:
|
||||
"""Wait for server to be healthy by checking /health endpoint."""
|
||||
start_time = time.perf_counter()
|
||||
url = f"http://{host}:{port}/health"
|
||||
|
||||
while time.perf_counter() - start_time < timeout:
|
||||
try:
|
||||
response = requests.get(url, timeout=5)
|
||||
if response.status_code == 200:
|
||||
return True
|
||||
except requests.exceptions.RequestException:
|
||||
pass
|
||||
time.sleep(1)
|
||||
return False
|
||||
|
||||
|
||||
def find_available_ports(base_port: int, count: int) -> List[int]:
|
||||
"""Find consecutive available ports starting from base_port."""
|
||||
available_ports = []
|
||||
current_port = base_port
|
||||
|
||||
while len(available_ports) < count:
|
||||
if is_port_available(current_port):
|
||||
available_ports.append(current_port)
|
||||
current_port += random.randint(100, 1000)
|
||||
|
||||
return available_ports
|
||||
|
||||
|
||||
def cleanup_processes(processes: List[mp.Process]):
|
||||
for process in processes:
|
||||
logger.info(f"Terminating process group {process.pid}")
|
||||
try:
|
||||
os.killpg(process.pid, signal.SIGTERM)
|
||||
except ProcessLookupError:
|
||||
# Process group may already be terminated
|
||||
pass
|
||||
|
||||
# Wait for processes to terminate
|
||||
for process in processes:
|
||||
process.join(timeout=5)
|
||||
if process.is_alive():
|
||||
logger.warning(
|
||||
f"Process {process.pid} did not terminate gracefully, forcing kill"
|
||||
)
|
||||
try:
|
||||
os.killpg(process.pid, signal.SIGKILL)
|
||||
except ProcessLookupError:
|
||||
pass
|
||||
|
||||
logger.info("All process groups terminated")
|
||||
|
||||
|
||||
def main():
|
||||
# CUDA runtime isn't fork-safe, which can lead to subtle bugs or crashes
|
||||
mp.set_start_method("spawn")
|
||||
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Launch SGLang router and server processes"
|
||||
)
|
||||
|
||||
ServerArgs.add_cli_args(parser)
|
||||
RouterArgs.add_cli_args(parser, use_router_prefix=True, exclude_host_port=True)
|
||||
parser.add_argument(
|
||||
"--router-dp-worker-base-port",
|
||||
type=int,
|
||||
default=31000,
|
||||
help="Base port number for data parallel workers",
|
||||
)
|
||||
# No extra retry/CB flags here; RouterArgs.add_cli_args already defines them with router- prefix
|
||||
|
||||
args = parser.parse_args()
|
||||
server_args = ServerArgs.from_cli_args(args)
|
||||
router_args = RouterArgs.from_cli_args(args, use_router_prefix=True)
|
||||
|
||||
# Find available ports for workers
|
||||
worker_ports = find_available_ports(
|
||||
args.router_dp_worker_base_port, server_args.dp_size
|
||||
)
|
||||
|
||||
# Start server processes
|
||||
server_processes = []
|
||||
|
||||
for i, worker_port in enumerate(worker_ports):
|
||||
logger.info(f"Launching DP server process {i} on port {worker_port}")
|
||||
proc = launch_server_process(server_args, worker_port, i)
|
||||
server_processes.append(proc)
|
||||
|
||||
signal.signal(signal.SIGINT, lambda sig, frame: cleanup_processes(server_processes))
|
||||
signal.signal(
|
||||
signal.SIGTERM, lambda sig, frame: cleanup_processes(server_processes)
|
||||
)
|
||||
signal.signal(
|
||||
signal.SIGQUIT, lambda sig, frame: cleanup_processes(server_processes)
|
||||
)
|
||||
|
||||
# Update router args with worker URLs
|
||||
# Use grpc:// protocol if server is in gRPC mode, otherwise http://
|
||||
protocol = "grpc" if server_args.grpc_mode else "http"
|
||||
router_args.worker_urls = [
|
||||
f"{protocol}://{server_args.host}:{port}" for port in worker_ports
|
||||
]
|
||||
|
||||
# Start the router
|
||||
try:
|
||||
launch_router(router_args)
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to start router: {e}")
|
||||
cleanup_processes(server_processes)
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
462
third_party/sglang/sgl-model-gateway/bindings/python/src/sglang_router/mini_lb.py
vendored
Normal file
462
third_party/sglang/sgl-model-gateway/bindings/python/src/sglang_router/mini_lb.py
vendored
Normal file
@@ -0,0 +1,462 @@
|
||||
"""
|
||||
Minimal HTTP load balancer for prefill and decode servers for testing.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import ipaddress
|
||||
import logging
|
||||
import random
|
||||
import urllib
|
||||
import warnings
|
||||
from http import HTTPStatus
|
||||
from itertools import chain
|
||||
from typing import Optional
|
||||
|
||||
import aiohttp
|
||||
import orjson
|
||||
import uvicorn
|
||||
from fastapi import FastAPI, HTTPException
|
||||
from fastapi.responses import ORJSONResponse, Response, StreamingResponse
|
||||
from sglang_router.router_args import RouterArgs
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
AIOHTTP_STREAM_READ_CHUNK_SIZE = (
|
||||
1024 * 64
|
||||
) # 64KB, to prevent aiohttp's "Chunk too big" error
|
||||
|
||||
|
||||
def maybe_wrap_ipv6_address(address: str) -> str:
|
||||
try:
|
||||
ipaddress.IPv6Address(address)
|
||||
return f"[{address}]"
|
||||
except ValueError:
|
||||
return address
|
||||
|
||||
|
||||
class MiniLoadBalancer:
|
||||
def __init__(
|
||||
self,
|
||||
router_args: RouterArgs,
|
||||
):
|
||||
self._validate_router_args(router_args)
|
||||
|
||||
self.host = router_args.host
|
||||
self.port = router_args.port
|
||||
self.timeout = router_args.request_timeout_secs
|
||||
self.prefill_urls = [url[0] for url in router_args.prefill_urls]
|
||||
self.prefill_bootstrap_ports = [url[1] for url in router_args.prefill_urls]
|
||||
self.decode_urls = router_args.decode_urls
|
||||
self.test_external_dp_routing = router_args.test_external_dp_routing
|
||||
self.prefill_dp_size = None
|
||||
self.decode_dp_size = None
|
||||
|
||||
def _validate_router_args(self, router_args: RouterArgs):
|
||||
logger.warning(
|
||||
"\x1b[33mMiniLB is only for debugging purposes, it only supports random policy!\033[0m"
|
||||
)
|
||||
|
||||
# NOTE: too many arguments unsupported, just validate some important ones
|
||||
if router_args.policy != "random":
|
||||
logger.warning("[MiniLB] Overriding policy to random")
|
||||
router_args.policy = "random"
|
||||
|
||||
if not router_args.pd_disaggregation:
|
||||
raise ValueError("MiniLB only supports PD disaggregation mode")
|
||||
|
||||
if len(router_args.prefill_urls) == 0 or len(router_args.decode_urls) == 0:
|
||||
raise ValueError(
|
||||
"MiniLB requires at least one prefill and one decode server"
|
||||
)
|
||||
|
||||
def start(self):
|
||||
global lb
|
||||
lb = self
|
||||
uvicorn.run(app, host=self.host, port=self.port)
|
||||
|
||||
async def _ensure_dp_sizes(self):
|
||||
if self.prefill_dp_size is not None:
|
||||
return
|
||||
async with aiohttp.ClientSession() as session:
|
||||
async with session.get(f"{self.prefill_urls[0]}/server_info") as resp:
|
||||
info = await resp.json()
|
||||
self.prefill_dp_size = len(info.get("internal_states", [1]))
|
||||
async with session.get(f"{self.decode_urls[0]}/server_info") as resp:
|
||||
info = await resp.json()
|
||||
self.decode_dp_size = len(info.get("internal_states", [1]))
|
||||
logger.info(
|
||||
f"[MiniLB] DP sizes: prefill={self.prefill_dp_size}, decode={self.decode_dp_size}"
|
||||
)
|
||||
|
||||
def _fork_dp_requests(self, request):
|
||||
p_rank = random.randint(0, self.prefill_dp_size - 1)
|
||||
d_rank = random.randint(0, self.decode_dp_size - 1)
|
||||
|
||||
prefill_req = request.copy()
|
||||
decode_req = request.copy()
|
||||
prefill_req["routed_dp_rank"] = p_rank
|
||||
decode_req["routed_dp_rank"] = d_rank
|
||||
decode_req["disagg_prefill_dp_rank"] = p_rank
|
||||
|
||||
return prefill_req, decode_req, d_rank
|
||||
|
||||
def select_pair(self):
|
||||
assert len(self.prefill_urls) > 0, "No prefill servers available"
|
||||
assert len(self.decode_urls) > 0, "No decode servers available"
|
||||
pidx = random.randint(0, len(self.prefill_urls) - 1)
|
||||
didx = random.randint(0, len(self.decode_urls) - 1)
|
||||
return (
|
||||
self.prefill_urls[pidx],
|
||||
self.prefill_bootstrap_ports[pidx],
|
||||
self.decode_urls[didx],
|
||||
)
|
||||
|
||||
async def generate(
|
||||
self, modified_request, prefill_server, decode_server, endpoint
|
||||
) -> ORJSONResponse:
|
||||
assert endpoint[0] != "/", f"Endpoint should not start with '/': {endpoint}"
|
||||
|
||||
expected_decode_dp_rank = None
|
||||
if self.test_external_dp_routing:
|
||||
await self._ensure_dp_sizes()
|
||||
prefill_req, decode_req, expected_decode_dp_rank = self._fork_dp_requests(
|
||||
modified_request
|
||||
)
|
||||
else:
|
||||
prefill_req = modified_request
|
||||
decode_req = modified_request
|
||||
|
||||
async with aiohttp.ClientSession(
|
||||
timeout=aiohttp.ClientTimeout(
|
||||
total=self.timeout
|
||||
) # Add timeout for request reliability
|
||||
) as session:
|
||||
|
||||
tasks = [
|
||||
session.post(f"{prefill_server}/{endpoint}", json=prefill_req),
|
||||
session.post(f"{decode_server}/{endpoint}", json=decode_req),
|
||||
]
|
||||
|
||||
# Wait for both responses to complete. Prefill should end first.
|
||||
prefill_response, decode_response = await asyncio.gather(*tasks)
|
||||
|
||||
if "return_logprob" in modified_request:
|
||||
|
||||
prefill_json = await prefill_response.json()
|
||||
ret_json = await decode_response.json()
|
||||
|
||||
# merge `meta_info.input_token_logprobs` from prefill to decode
|
||||
if "meta_info" in ret_json:
|
||||
if "input_token_logprobs" in ret_json["meta_info"]:
|
||||
ret_json["meta_info"]["input_token_logprobs"] = (
|
||||
prefill_json["meta_info"]["input_token_logprobs"]
|
||||
+ ret_json["meta_info"]["input_token_logprobs"]
|
||||
)
|
||||
else:
|
||||
ret_json = await decode_response.json()
|
||||
|
||||
if expected_decode_dp_rank is not None:
|
||||
actual = ret_json.get("meta_info", {}).get("dp_rank")
|
||||
if actual != expected_decode_dp_rank:
|
||||
return ORJSONResponse(
|
||||
content={
|
||||
"error": f"DP rank mismatch: expected {expected_decode_dp_rank}, got {actual}"
|
||||
},
|
||||
status_code=500,
|
||||
)
|
||||
|
||||
return ORJSONResponse(
|
||||
content=ret_json,
|
||||
status_code=decode_response.status,
|
||||
)
|
||||
|
||||
async def generate_stream(
|
||||
self, modified_request, prefill_server, decode_server, endpoint="generate"
|
||||
):
|
||||
|
||||
if self.test_external_dp_routing:
|
||||
warnings.warn("--test-external-dp-routing is not supported with streaming")
|
||||
|
||||
assert endpoint[0] != "/", f"Endpoint should not start with '/': {endpoint}"
|
||||
|
||||
async def stream_results():
|
||||
async with aiohttp.ClientSession(
|
||||
timeout=aiohttp.ClientTimeout(
|
||||
total=self.timeout
|
||||
) # Add timeout for request reliability
|
||||
) as session:
|
||||
# Create the tasks for both prefill and decode requests
|
||||
tasks = [
|
||||
session.post(f"{prefill_server}/{endpoint}", json=modified_request),
|
||||
session.post(f"{decode_server}/{endpoint}", json=modified_request),
|
||||
]
|
||||
|
||||
# Wait for both responses to complete. Since this is streaming, they return immediately.
|
||||
prefill_response, decode_response = await asyncio.gather(*tasks)
|
||||
|
||||
if modified_request.get("return_logprob", False):
|
||||
prefill_chunks = []
|
||||
async for chunk in prefill_response.content:
|
||||
prefill_chunks.append(chunk)
|
||||
|
||||
first_prefill_chunk = (
|
||||
prefill_chunks[0].decode("utf-8")[5:].strip("\n")
|
||||
)
|
||||
first_prefill_chunk_json = orjson.loads(first_prefill_chunk)
|
||||
|
||||
async for chunk in decode_response.content:
|
||||
# Note: This is inefficient
|
||||
# merge prefill input_token_logprobs, output_token_logprobs to decode
|
||||
decoded_chunk = chunk.decode("utf-8")
|
||||
if (
|
||||
decoded_chunk
|
||||
and decoded_chunk.startswith("data:")
|
||||
and "[DONE]" not in decoded_chunk
|
||||
):
|
||||
ret_json = orjson.loads(decoded_chunk[5:].strip("\n"))
|
||||
ret_json["meta_info"]["input_token_logprobs"] = (
|
||||
first_prefill_chunk_json["meta_info"][
|
||||
"input_token_logprobs"
|
||||
]
|
||||
+ ret_json["meta_info"]["input_token_logprobs"]
|
||||
)
|
||||
|
||||
yield b"data: " + orjson.dumps(ret_json) + b"\n\n"
|
||||
else:
|
||||
yield chunk
|
||||
else:
|
||||
async for chunk in decode_response.content.iter_chunked(
|
||||
AIOHTTP_STREAM_READ_CHUNK_SIZE
|
||||
):
|
||||
yield chunk
|
||||
|
||||
return StreamingResponse(
|
||||
stream_results(),
|
||||
media_type="text/event-stream",
|
||||
)
|
||||
|
||||
|
||||
app = FastAPI()
|
||||
lb: Optional[MiniLoadBalancer] = None
|
||||
|
||||
|
||||
@app.get("/health")
|
||||
async def health_check():
|
||||
return Response(status_code=200)
|
||||
|
||||
|
||||
@app.get("/health_generate")
|
||||
async def health_generate():
|
||||
async with aiohttp.ClientSession() as session:
|
||||
# Create the tasks
|
||||
tasks = []
|
||||
for server in chain(lb.prefill_urls, lb.decode_urls):
|
||||
tasks.append(session.get(f"{server}/health_generate"))
|
||||
for i, response in enumerate(asyncio.as_completed(tasks)):
|
||||
await response
|
||||
return Response(status_code=200)
|
||||
|
||||
|
||||
@app.post("/flush_cache")
|
||||
async def flush_cache():
|
||||
async with aiohttp.ClientSession() as session:
|
||||
# Create the tasks
|
||||
tasks = []
|
||||
for server in chain(lb.prefill_urls, lb.decode_urls):
|
||||
tasks.append(session.post(f"{server}/flush_cache"))
|
||||
for i, response in enumerate(asyncio.as_completed(tasks)):
|
||||
await response
|
||||
return Response(status_code=200)
|
||||
|
||||
|
||||
# TODO: Remove `/get_server_info` alias after one release-cycle deprecation window.
|
||||
@app.get("/server_info")
|
||||
@app.get("/get_server_info")
|
||||
async def get_server_info():
|
||||
prefill_infos = []
|
||||
decode_infos = []
|
||||
all_internal_states = []
|
||||
|
||||
async with aiohttp.ClientSession() as session:
|
||||
for server in lb.prefill_urls:
|
||||
server_info = await session.get(f"{server}/server_info")
|
||||
prefill_infos.append(await server_info.json())
|
||||
for server in lb.decode_urls:
|
||||
server_info = await session.get(f"{server}/server_info")
|
||||
info_json = await server_info.json()
|
||||
decode_infos.append(info_json)
|
||||
# Extract internal_states from decode servers
|
||||
if "internal_states" in info_json:
|
||||
all_internal_states.extend(info_json["internal_states"])
|
||||
|
||||
# Return format expected by bench_one_batch_server.py
|
||||
if all_internal_states:
|
||||
return {
|
||||
"internal_states": all_internal_states,
|
||||
"prefill": prefill_infos,
|
||||
"decode": decode_infos,
|
||||
}
|
||||
else:
|
||||
# Fallback with dummy data if no internal states found
|
||||
return {
|
||||
"internal_states": [
|
||||
{
|
||||
"last_gen_throughput": 0.0,
|
||||
"avg_spec_accept_length": None,
|
||||
}
|
||||
],
|
||||
"prefill": prefill_infos,
|
||||
"decode": decode_infos,
|
||||
}
|
||||
|
||||
|
||||
async def _get_model_info_impl():
|
||||
if not lb or not lb.prefill_urls:
|
||||
raise HTTPException(
|
||||
status_code=HTTPStatus.SERVICE_UNAVAILABLE,
|
||||
detail="There is no server registered",
|
||||
)
|
||||
|
||||
target_server_url = lb.prefill_urls[0]
|
||||
endpoint_url = f"{target_server_url}/model_info"
|
||||
|
||||
async with aiohttp.ClientSession() as session:
|
||||
try:
|
||||
async with session.get(endpoint_url) as response:
|
||||
if response.status != 200:
|
||||
error_text = await response.text()
|
||||
raise HTTPException(
|
||||
status_code=HTTPStatus.BAD_GATEWAY,
|
||||
detail=(
|
||||
f"Failed to get model info from {target_server_url}"
|
||||
f"Status: {response.status}, Response: {error_text}"
|
||||
),
|
||||
)
|
||||
|
||||
model_info_json = await response.json()
|
||||
return ORJSONResponse(content=model_info_json)
|
||||
|
||||
except aiohttp.ClientError as e:
|
||||
raise HTTPException(
|
||||
status_code=HTTPStatus.SERVICE_UNAVAILABLE,
|
||||
detail=f"Failed to get model info from backend",
|
||||
)
|
||||
|
||||
|
||||
@app.get("/model_info")
|
||||
async def model_info():
|
||||
return await _get_model_info_impl()
|
||||
|
||||
|
||||
@app.get("/get_model_info")
|
||||
async def get_model_info():
|
||||
return await _get_model_info_impl()
|
||||
|
||||
|
||||
@app.post("/generate")
|
||||
async def handle_generate_request(request_data: dict):
|
||||
prefill_server, bootstrap_port, decode_server = lb.select_pair()
|
||||
|
||||
# Parse and transform prefill_server for bootstrap data
|
||||
parsed_url = urllib.parse.urlparse(prefill_server)
|
||||
hostname = maybe_wrap_ipv6_address(parsed_url.hostname)
|
||||
modified_request = request_data.copy()
|
||||
|
||||
batch_size = _get_request_batch_size(modified_request)
|
||||
if batch_size is not None:
|
||||
modified_request.update(
|
||||
{
|
||||
"bootstrap_host": [hostname] * batch_size,
|
||||
"bootstrap_port": [bootstrap_port] * batch_size,
|
||||
"bootstrap_room": [
|
||||
_generate_bootstrap_room() for _ in range(batch_size)
|
||||
],
|
||||
}
|
||||
)
|
||||
else:
|
||||
modified_request.update(
|
||||
{
|
||||
"bootstrap_host": hostname,
|
||||
"bootstrap_port": bootstrap_port,
|
||||
"bootstrap_room": _generate_bootstrap_room(),
|
||||
}
|
||||
)
|
||||
|
||||
if request_data.get("stream", False):
|
||||
return await lb.generate_stream(
|
||||
modified_request, prefill_server, decode_server, "generate"
|
||||
)
|
||||
else:
|
||||
return await lb.generate(
|
||||
modified_request, prefill_server, decode_server, "generate"
|
||||
)
|
||||
|
||||
|
||||
async def _forward_to_backend(request_data: dict, endpoint_name: str):
|
||||
prefill_server, bootstrap_port, decode_server = lb.select_pair()
|
||||
|
||||
# Parse and transform prefill_server for bootstrap data
|
||||
parsed_url = urllib.parse.urlparse(prefill_server)
|
||||
hostname = maybe_wrap_ipv6_address(parsed_url.hostname)
|
||||
modified_request = request_data.copy()
|
||||
modified_request.update(
|
||||
{
|
||||
"bootstrap_host": hostname,
|
||||
"bootstrap_port": bootstrap_port,
|
||||
"bootstrap_room": _generate_bootstrap_room(),
|
||||
}
|
||||
)
|
||||
|
||||
if request_data.get("stream", False):
|
||||
return await lb.generate_stream(
|
||||
modified_request,
|
||||
prefill_server,
|
||||
decode_server,
|
||||
endpoint=endpoint_name,
|
||||
)
|
||||
else:
|
||||
return await lb.generate(
|
||||
modified_request,
|
||||
prefill_server,
|
||||
decode_server,
|
||||
endpoint=endpoint_name,
|
||||
)
|
||||
|
||||
|
||||
@app.post("/v1/chat/completions")
|
||||
async def handle_chat_completion_request(request_data: dict):
|
||||
return await _forward_to_backend(request_data, "v1/chat/completions")
|
||||
|
||||
|
||||
@app.post("/v1/completions")
|
||||
async def handle_completion_request(request_data: dict):
|
||||
return await _forward_to_backend(request_data, "v1/completions")
|
||||
|
||||
|
||||
def _generate_bootstrap_room():
|
||||
return random.randint(0, 2**63 - 1)
|
||||
|
||||
|
||||
# We may utilize `GenerateReqInput`'s logic later
|
||||
def _get_request_batch_size(request):
|
||||
if (text := request.get("text")) is not None:
|
||||
return None if isinstance(text, str) else len(text)
|
||||
if (input_ids := request.get("input_ids")) is not None:
|
||||
return None if isinstance(input_ids[0], int) else len(input_ids)
|
||||
return None
|
||||
|
||||
|
||||
@app.get("/v1/models")
|
||||
async def get_models():
|
||||
prefill_server = lb.prefill_urls[0] # Get the first prefill server
|
||||
async with aiohttp.ClientSession() as session:
|
||||
try:
|
||||
response = await session.get(f"{prefill_server}/v1/models")
|
||||
if response.status != 200:
|
||||
raise HTTPException(
|
||||
status_code=response.status,
|
||||
detail=f"Prefill server error: Status {response.status}",
|
||||
)
|
||||
return ORJSONResponse(content=await response.json())
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
320
third_party/sglang/sgl-model-gateway/bindings/python/src/sglang_router/router.py
vendored
Normal file
320
third_party/sglang/sgl-model-gateway/bindings/python/src/sglang_router/router.py
vendored
Normal file
@@ -0,0 +1,320 @@
|
||||
from typing import Optional
|
||||
|
||||
from sglang_router.router_args import RouterArgs
|
||||
from sglang_router.sglang_router_rs import (
|
||||
BackendType,
|
||||
HistoryBackendType,
|
||||
PolicyType,
|
||||
PyApiKeyEntry,
|
||||
PyControlPlaneAuthConfig,
|
||||
PyJwtConfig,
|
||||
PyOracleConfig,
|
||||
PyPostgresConfig,
|
||||
PyRedisConfig,
|
||||
PyRole,
|
||||
)
|
||||
from sglang_router.sglang_router_rs import Router as _Router
|
||||
|
||||
|
||||
def policy_from_str(policy_str: Optional[str]) -> PolicyType:
|
||||
"""Convert policy string to PolicyType enum."""
|
||||
if policy_str is None:
|
||||
return None
|
||||
policy_map = {
|
||||
"random": PolicyType.Random,
|
||||
"round_robin": PolicyType.RoundRobin,
|
||||
"cache_aware": PolicyType.CacheAware,
|
||||
"power_of_two": PolicyType.PowerOfTwo,
|
||||
"bucket": PolicyType.Bucket,
|
||||
"manual": PolicyType.Manual,
|
||||
"consistent_hashing": PolicyType.ConsistentHashing,
|
||||
"prefix_hash": PolicyType.PrefixHash,
|
||||
}
|
||||
return policy_map[policy_str]
|
||||
|
||||
|
||||
def backend_from_str(backend_str: Optional[str]) -> BackendType:
|
||||
"""Convert backend string to BackendType enum."""
|
||||
if isinstance(backend_str, BackendType):
|
||||
return backend_str
|
||||
if backend_str is None:
|
||||
return BackendType.Sglang
|
||||
backend_map = {"sglang": BackendType.Sglang, "openai": BackendType.Openai}
|
||||
backend_lower = backend_str.lower()
|
||||
if backend_lower not in backend_map:
|
||||
raise ValueError(
|
||||
f"Unknown backend: {backend_str}. Valid options: {', '.join(backend_map.keys())}"
|
||||
)
|
||||
return backend_map[backend_lower]
|
||||
|
||||
|
||||
def history_backend_from_str(backend_str: Optional[str]) -> HistoryBackendType:
|
||||
"""Convert history backend string to HistoryBackendType enum."""
|
||||
if isinstance(backend_str, HistoryBackendType):
|
||||
return backend_str
|
||||
if backend_str is None:
|
||||
return HistoryBackendType.Memory
|
||||
backend_lower = backend_str.lower()
|
||||
if backend_lower == "memory":
|
||||
return HistoryBackendType.Memory
|
||||
elif backend_lower == "none":
|
||||
# Use getattr to access 'None' which is a Python keyword
|
||||
return getattr(HistoryBackendType, "None")
|
||||
elif backend_lower == "oracle":
|
||||
return HistoryBackendType.Oracle
|
||||
elif backend_lower == "postgres":
|
||||
return HistoryBackendType.Postgres
|
||||
elif backend_lower == "redis":
|
||||
return HistoryBackendType.Redis
|
||||
else:
|
||||
raise ValueError(f"Unknown history backend: {backend_str}")
|
||||
|
||||
|
||||
def role_from_str(role_str: str) -> PyRole:
|
||||
"""Convert role string to PyRole enum."""
|
||||
if role_str.lower() == "admin":
|
||||
return PyRole.Admin
|
||||
return PyRole.User
|
||||
|
||||
|
||||
def build_control_plane_auth_config(
|
||||
args_dict: dict,
|
||||
) -> Optional[PyControlPlaneAuthConfig]:
|
||||
"""Build control plane auth config from args dict."""
|
||||
api_keys = args_dict.get("control_plane_api_keys", [])
|
||||
jwt_issuer = args_dict.get("jwt_issuer")
|
||||
jwt_audience = args_dict.get("jwt_audience")
|
||||
audit_enabled = args_dict.get("control_plane_audit_enabled", False)
|
||||
|
||||
# Check if any auth is configured
|
||||
has_api_keys = bool(api_keys)
|
||||
has_jwt = jwt_issuer is not None and jwt_audience is not None
|
||||
|
||||
if not has_api_keys and not has_jwt:
|
||||
return None
|
||||
|
||||
# Build API key entries
|
||||
py_api_keys = []
|
||||
for key_tuple in api_keys:
|
||||
# Tuple format: (id, name, key, role)
|
||||
key_id, name, key, role = key_tuple
|
||||
py_api_keys.append(
|
||||
PyApiKeyEntry(
|
||||
id=key_id,
|
||||
name=name,
|
||||
key=key,
|
||||
role=role_from_str(role),
|
||||
)
|
||||
)
|
||||
|
||||
# Build JWT config if present
|
||||
jwt_config = None
|
||||
if has_jwt:
|
||||
jwt_config = PyJwtConfig(
|
||||
issuer=jwt_issuer,
|
||||
audience=jwt_audience,
|
||||
jwks_uri=args_dict.get("jwt_jwks_uri"),
|
||||
role_mapping=args_dict.get("jwt_role_mapping", {}),
|
||||
)
|
||||
|
||||
return PyControlPlaneAuthConfig(
|
||||
jwt=jwt_config,
|
||||
api_keys=py_api_keys,
|
||||
audit_enabled=audit_enabled,
|
||||
)
|
||||
|
||||
|
||||
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 based on cache state and load balance
|
||||
- PolicyType.PowerOfTwo: Select best of two random workers based on load (PD mode only)
|
||||
host: Host address to bind the router server. Supports IPv4, IPv6 (e.g., ::, ::1), or 0.0.0.0 for all interfaces. Default: '0.0.0.0'
|
||||
port: Port number to bind the router server. Default: 3001
|
||||
worker_startup_timeout_secs: Timeout in seconds for worker startup and registration. Large models can take significant time to load into GPU memory. Default: 1800 (30 minutes)
|
||||
worker_startup_check_interval: Interval in seconds between checks for worker initialization. Default: 10
|
||||
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
|
||||
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
|
||||
eviction_interval_secs: Interval in seconds between cache eviction operations in cache-aware
|
||||
routing. Default: 60
|
||||
max_payload_size: Maximum payload size in bytes. Default: 256MB
|
||||
max_tree_size: Maximum size of the approximation tree for cache-aware routing. Default: 2^24
|
||||
dp_aware: Enable data parallelism aware schedule. Default: False
|
||||
enable_igw: Enable IGW (Inference-Gateway) mode for multi-model support. When enabled,
|
||||
the router can manage multiple models simultaneously with per-model load balancing
|
||||
policies. Default: False
|
||||
api_key: The api key used for the authorization with the worker.
|
||||
Useful when the dp aware scheduling strategy is enabled.
|
||||
Default: None
|
||||
log_dir: Directory to store log files. If None, logs are only output to console. Default: None
|
||||
log_level: Logging level. Options: 'debug', 'info', 'warn', 'error'.
|
||||
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
|
||||
prefill_selector: Dictionary mapping of label keys to values for Kubernetes pod selection
|
||||
for prefill servers (PD mode only). Default: {}
|
||||
decode_selector: Dictionary mapping of label keys to values for Kubernetes pod selection
|
||||
for decode servers (PD mode only). Default: {}
|
||||
prometheus_port: Port to expose Prometheus metrics. Default: None
|
||||
prometheus_host: Host address to bind the Prometheus metrics server. Default: None
|
||||
pd_disaggregation: 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)
|
||||
prefill_policy: Specific load balancing policy for prefill nodes (PD mode only).
|
||||
If not specified, uses the main policy. Default: None
|
||||
decode_policy: Specific load balancing policy for decode nodes (PD mode only).
|
||||
If not specified, uses the main policy. Default: None
|
||||
request_id_headers: List of HTTP headers to check for request IDs. If not specified,
|
||||
uses common defaults: ['x-request-id', 'x-correlation-id', 'x-trace-id', 'request-id'].
|
||||
Example: ['x-my-request-id', 'x-custom-trace-id']. Default: None
|
||||
bootstrap_port_annotation: Kubernetes annotation name for bootstrap port (PD mode).
|
||||
Default: 'sglang.ai/bootstrap-port'
|
||||
request_timeout_secs: Request timeout in seconds. Default: 600
|
||||
max_concurrent_requests: Maximum number of concurrent requests allowed for rate limiting. Default: 256
|
||||
queue_size: Queue size for pending requests when max concurrent limit reached (0 = no queue, return 429 immediately). Default: 100
|
||||
queue_timeout_secs: Maximum time (in seconds) a request can wait in queue before timing out. Default: 60
|
||||
rate_limit_tokens_per_second: Token bucket refill rate (tokens per second). If not set, defaults to max_concurrent_requests. Default: None
|
||||
cors_allowed_origins: List of allowed origins for CORS. Empty list allows all origins. Default: []
|
||||
health_failure_threshold: Number of consecutive health check failures before marking worker unhealthy. Default: 3
|
||||
health_success_threshold: Number of consecutive health check successes before marking worker healthy. Default: 2
|
||||
health_check_timeout_secs: Timeout in seconds for health check requests. Default: 5
|
||||
health_check_interval_secs: Interval in seconds between runtime health checks. Default: 60
|
||||
health_check_endpoint: Health check endpoint path. Default: '/health'
|
||||
model_path: Model path for loading tokenizer (HuggingFace model ID or local path). Default: None
|
||||
tokenizer_path: Explicit tokenizer path (overrides model_path tokenizer if provided). Default: None
|
||||
server_cert_path: Path to server TLS certificate (PEM format). Default: None
|
||||
server_key_path: Path to server TLS private key (PEM format). Default: None
|
||||
"""
|
||||
|
||||
def __init__(self, router: _Router):
|
||||
self._router = router
|
||||
|
||||
@staticmethod
|
||||
def from_args(args: RouterArgs) -> "Router":
|
||||
"""Create a router from a RouterArgs instance."""
|
||||
|
||||
args_dict = vars(args)
|
||||
# Convert RouterArgs to _Router parameters
|
||||
args_dict["worker_urls"] = (
|
||||
[]
|
||||
if args_dict["service_discovery"] or args_dict["pd_disaggregation"]
|
||||
else args_dict["worker_urls"]
|
||||
)
|
||||
args_dict["policy"] = policy_from_str(args_dict["policy"])
|
||||
args_dict["prefill_urls"] = (
|
||||
args_dict["prefill_urls"] if args_dict["pd_disaggregation"] else None
|
||||
)
|
||||
args_dict["decode_urls"] = (
|
||||
args_dict["decode_urls"] if args_dict["pd_disaggregation"] else None
|
||||
)
|
||||
args_dict["prefill_policy"] = policy_from_str(args_dict["prefill_policy"])
|
||||
args_dict["decode_policy"] = policy_from_str(args_dict["decode_policy"])
|
||||
|
||||
# Convert backend
|
||||
args_dict["backend"] = backend_from_str(args_dict.get("backend"))
|
||||
|
||||
# Convert history_backend to enum first
|
||||
history_backend_raw = args_dict.get("history_backend", "memory")
|
||||
history_backend = history_backend_from_str(history_backend_raw)
|
||||
|
||||
# Convert Oracle config if needed
|
||||
oracle_config = None
|
||||
if history_backend == HistoryBackendType.Oracle:
|
||||
# Prioritize TNS alias over connect descriptor
|
||||
tns_alias = args_dict.get("oracle_tns_alias")
|
||||
connect_descriptor = args_dict.get("oracle_connect_descriptor")
|
||||
|
||||
# Use TNS alias if provided, otherwise use connect descriptor
|
||||
final_descriptor = tns_alias if tns_alias else connect_descriptor
|
||||
|
||||
oracle_config = PyOracleConfig(
|
||||
password=args_dict.get("oracle_password"),
|
||||
username=args_dict.get("oracle_username"),
|
||||
connect_descriptor=final_descriptor,
|
||||
wallet_path=args_dict.get("oracle_wallet_path"),
|
||||
pool_min=args_dict.get("oracle_pool_min", 1),
|
||||
pool_max=args_dict.get("oracle_pool_max", 16),
|
||||
pool_timeout_secs=args_dict.get("oracle_pool_timeout_secs", 30),
|
||||
)
|
||||
args_dict["oracle_config"] = oracle_config
|
||||
args_dict["history_backend"] = history_backend
|
||||
|
||||
# Convert Postgres config if needed
|
||||
postgres_config = None
|
||||
if history_backend == HistoryBackendType.Postgres:
|
||||
postgres_config = PyPostgresConfig(
|
||||
db_url=args_dict.get("postgres_db_url"),
|
||||
pool_max=args_dict.get("postgres_pool_max", 16),
|
||||
)
|
||||
args_dict["postgres_config"] = postgres_config
|
||||
|
||||
# Convert Redis config if needed
|
||||
redis_config = None
|
||||
if history_backend == HistoryBackendType.Redis:
|
||||
retention_days = args_dict.get("redis_retention_days", 30)
|
||||
# If retention_days is negative, it means persistent storage (None in Rust)
|
||||
retention_arg = None if retention_days < 0 else retention_days
|
||||
|
||||
redis_config = PyRedisConfig(
|
||||
url=args_dict.get("redis_url"),
|
||||
pool_max=args_dict.get("redis_pool_max", 16),
|
||||
retention_days=retention_arg,
|
||||
)
|
||||
args_dict["redis_config"] = redis_config
|
||||
|
||||
# Build control plane auth config
|
||||
args_dict["control_plane_auth"] = build_control_plane_auth_config(args_dict)
|
||||
|
||||
# Remove fields that shouldn't be passed to Rust Router constructor
|
||||
fields_to_remove = [
|
||||
"mini_lb",
|
||||
"test_external_dp_routing",
|
||||
"oracle_wallet_path",
|
||||
"oracle_tns_alias",
|
||||
"oracle_connect_descriptor",
|
||||
"oracle_username",
|
||||
"oracle_password",
|
||||
"oracle_pool_min",
|
||||
"oracle_pool_max",
|
||||
"oracle_pool_timeout_secs",
|
||||
"postgres_db_url",
|
||||
"postgres_pool_max",
|
||||
"redis_url",
|
||||
"redis_pool_max",
|
||||
"redis_retention_days",
|
||||
# Control plane auth fields (converted to control_plane_auth)
|
||||
"control_plane_api_keys",
|
||||
"control_plane_audit_enabled",
|
||||
"jwt_issuer",
|
||||
"jwt_audience",
|
||||
"jwt_jwks_uri",
|
||||
"jwt_role_mapping",
|
||||
]
|
||||
for field in fields_to_remove:
|
||||
args_dict.pop(field, None)
|
||||
|
||||
return Router(_Router(**args_dict))
|
||||
|
||||
def start(self) -> None:
|
||||
"""Start the router server.
|
||||
|
||||
This method blocks until the server is shut down.
|
||||
"""
|
||||
self._router.start()
|
||||
1104
third_party/sglang/sgl-model-gateway/bindings/python/src/sglang_router/router_args.py
vendored
Normal file
1104
third_party/sglang/sgl-model-gateway/bindings/python/src/sglang_router/router_args.py
vendored
Normal file
File diff suppressed because it is too large
Load Diff
1
third_party/sglang/sgl-model-gateway/bindings/python/src/sglang_router/version.py
vendored
Normal file
1
third_party/sglang/sgl-model-gateway/bindings/python/src/sglang_router/version.py
vendored
Normal file
@@ -0,0 +1 @@
|
||||
__version__ = "0.3.2"
|
||||
Reference in New Issue
Block a user