276 lines
8.6 KiB
Python
276 lines
8.6 KiB
Python
"""Focused runtime tests for the Python router and server-launch helpers."""
|
|
|
|
|
|
def _install_sglang_stubs(monkeypatch):
|
|
"""Install lightweight stubs for sglang.srt to avoid heavy deps during unit tests."""
|
|
import sys
|
|
import types
|
|
|
|
sglang_mod = types.ModuleType("sglang")
|
|
srt_mod = types.ModuleType("sglang.srt")
|
|
entry_mod = types.ModuleType("sglang.srt.entrypoints")
|
|
http_server_mod = types.ModuleType("sglang.srt.entrypoints.http_server")
|
|
server_args_mod = types.ModuleType("sglang.srt.server_args")
|
|
# sglang.srt.utils was refactored from a module into a package; launch_server.py
|
|
# imports from sglang.srt.utils.network, so the stub must model the submodule.
|
|
utils_mod = types.ModuleType("sglang.srt.utils")
|
|
network_mod = types.ModuleType("sglang.srt.utils.network")
|
|
|
|
def launch_server(_args):
|
|
return None
|
|
|
|
class ServerArgs:
|
|
# Minimal fields used by launch_server_process
|
|
def __init__(self):
|
|
self.port = 0
|
|
self.base_gpu_id = 0
|
|
self.dp_size = 1
|
|
self.tp_size = 1
|
|
|
|
@staticmethod
|
|
def add_cli_args(_parser):
|
|
return None
|
|
|
|
@staticmethod
|
|
def from_cli_args(_args):
|
|
sa = ServerArgs()
|
|
if hasattr(_args, "dp_size"):
|
|
sa.dp_size = _args.dp_size
|
|
if hasattr(_args, "tp_size"):
|
|
sa.tp_size = _args.tp_size
|
|
if hasattr(_args, "host"):
|
|
sa.host = _args.host
|
|
else:
|
|
sa.host = "127.0.0.1"
|
|
return sa
|
|
|
|
def is_port_available(_port: int) -> bool:
|
|
return True
|
|
|
|
http_server_mod.launch_server = launch_server
|
|
server_args_mod.ServerArgs = ServerArgs
|
|
network_mod.is_port_available = is_port_available
|
|
utils_mod.network = network_mod
|
|
|
|
# Also stub external deps imported at module top-level
|
|
def _dummy_get(*_a, **_k):
|
|
raise NotImplementedError
|
|
|
|
requests_stub = types.SimpleNamespace(
|
|
exceptions=types.SimpleNamespace(RequestException=Exception), get=_dummy_get
|
|
)
|
|
setproctitle_stub = types.SimpleNamespace(setproctitle=lambda *_a, **_k: None)
|
|
|
|
monkeypatch.setitem(sys.modules, "requests", requests_stub)
|
|
monkeypatch.setitem(sys.modules, "setproctitle", setproctitle_stub)
|
|
|
|
monkeypatch.setitem(sys.modules, "sglang", sglang_mod)
|
|
monkeypatch.setitem(sys.modules, "sglang.srt", srt_mod)
|
|
monkeypatch.setitem(sys.modules, "sglang.srt.entrypoints", entry_mod)
|
|
monkeypatch.setitem(
|
|
sys.modules, "sglang.srt.entrypoints.http_server", http_server_mod
|
|
)
|
|
monkeypatch.setitem(sys.modules, "sglang.srt.server_args", server_args_mod)
|
|
monkeypatch.setitem(sys.modules, "sglang.srt.utils", utils_mod)
|
|
monkeypatch.setitem(sys.modules, "sglang.srt.utils.network", network_mod)
|
|
|
|
|
|
def test_router_defaults_and_start(monkeypatch):
|
|
"""Router wrapper: defaults normalization and start() call.
|
|
|
|
Mocks the Rust-backed _Router to avoid native deps.
|
|
"""
|
|
from sglang_router import router as router_mod
|
|
|
|
captured = {}
|
|
|
|
class FakeRouter:
|
|
def __init__(self, **kwargs):
|
|
captured.update(kwargs)
|
|
|
|
def start(self):
|
|
captured["started"] = True
|
|
|
|
monkeypatch.setattr(router_mod, "_Router", FakeRouter, raising=True)
|
|
|
|
from sglang_router.router_args import RouterArgs as _RouterArgs
|
|
|
|
Router = router_mod.Router
|
|
args = _RouterArgs(
|
|
worker_urls=["http://w1:8000"],
|
|
policy="round_robin",
|
|
selector=None,
|
|
prefill_selector=None,
|
|
decode_selector=None,
|
|
cors_allowed_origins=None,
|
|
)
|
|
|
|
r = Router.from_args(args)
|
|
|
|
# Defaults preserved/normalized by Router.from_args
|
|
assert captured["selector"] is None
|
|
assert captured["prefill_selector"] is None
|
|
assert captured["decode_selector"] is None
|
|
assert captured["cors_allowed_origins"] is None
|
|
assert captured["worker_urls"] == ["http://w1:8000"]
|
|
from sglang_router.sglang_router_rs import PolicyType
|
|
|
|
assert captured["policy"] == PolicyType.RoundRobin
|
|
|
|
r.start()
|
|
assert captured.get("started") is True
|
|
|
|
|
|
def test_find_available_ports_and_wait_health(monkeypatch):
|
|
"""launch_server helpers: port finding and health waiting with transient error."""
|
|
_install_sglang_stubs(monkeypatch)
|
|
import importlib
|
|
|
|
ls = importlib.import_module("sglang_router.launch_server")
|
|
|
|
# Deterministic increments
|
|
monkeypatch.setattr(ls.random, "randint", lambda a, b: 100)
|
|
ports = ls.find_available_ports(30000, 3)
|
|
assert ports == [30000, 30100, 30200]
|
|
|
|
calls = {"n": 0}
|
|
|
|
class Ok:
|
|
status_code = 200
|
|
|
|
def fake_get(_url, timeout=5):
|
|
calls["n"] += 1
|
|
if calls["n"] == 1:
|
|
raise ls.requests.exceptions.RequestException("boom")
|
|
return Ok()
|
|
|
|
monkeypatch.setattr(ls.requests, "get", fake_get)
|
|
monkeypatch.setattr(ls.time, "sleep", lambda _s: None)
|
|
base = {"t": 0.0}
|
|
monkeypatch.setattr(
|
|
ls.time,
|
|
"perf_counter",
|
|
lambda: base.__setitem__("t", base["t"] + 0.1) or base["t"],
|
|
)
|
|
|
|
assert ls.wait_for_server_health("127.0.0.1", 12345, timeout=1)
|
|
|
|
|
|
def test_launch_server_process_and_cleanup(monkeypatch):
|
|
"""launch_server: process creation args and cleanup SIGTERM/SIGKILL logic."""
|
|
_install_sglang_stubs(monkeypatch)
|
|
import importlib
|
|
|
|
ls = importlib.import_module("sglang_router.launch_server")
|
|
|
|
created = {}
|
|
|
|
class FakeProcess:
|
|
def __init__(self, target, args):
|
|
created["target"] = target
|
|
created["args"] = args
|
|
self.pid = 4242
|
|
self._alive = True
|
|
|
|
def start(self):
|
|
created["started"] = True
|
|
|
|
def join(self, timeout=None):
|
|
return None
|
|
|
|
def is_alive(self):
|
|
return self._alive
|
|
|
|
monkeypatch.setattr(ls.mp, "Process", FakeProcess)
|
|
|
|
import sys as _sys
|
|
|
|
SA = _sys.modules["sglang.srt.server_args"].ServerArgs
|
|
sa = SA()
|
|
sa.tp_size = 2
|
|
|
|
ls.launch_server_process(sa, worker_port=31001, dp_id=3)
|
|
assert created.get("started") is True
|
|
targ, targ_args = created["target"], created["args"]
|
|
assert targ is ls.run_server
|
|
passed_sa = targ_args[0]
|
|
assert passed_sa.port == 31001
|
|
assert passed_sa.base_gpu_id == 3 * 2
|
|
assert passed_sa.dp_size == 1
|
|
|
|
# cleanup_processes
|
|
p1 = FakeProcess(target=None, args=())
|
|
p1._alive = False
|
|
p2 = FakeProcess(target=None, args=())
|
|
p2._alive = True
|
|
|
|
calls = []
|
|
|
|
def fake_killpg(pid, sig):
|
|
calls.append((pid, sig))
|
|
|
|
monkeypatch.setattr(ls.os, "killpg", fake_killpg)
|
|
|
|
ls.cleanup_processes([p1, p2])
|
|
|
|
import signal as _sig
|
|
|
|
assert (p1.pid, _sig.SIGTERM) in calls and (p2.pid, _sig.SIGTERM) in calls
|
|
assert (p2.pid, _sig.SIGKILL) in calls
|
|
|
|
|
|
def test_launch_server_process_declares_on_a_resolved_record(monkeypatch):
|
|
"""A record that carries its resolution takes the declaration channel.
|
|
|
|
The stub above has no `replace_resolved`, so it exercises the older wheels'
|
|
path. A current `ServerArgs` refuses plain assignment once resolution has
|
|
finished; the per-worker values reach the child as a declaration on a copy,
|
|
and the parent keeps what the operator passed.
|
|
"""
|
|
_install_sglang_stubs(monkeypatch)
|
|
import importlib
|
|
|
|
ls = importlib.import_module("sglang_router.launch_server")
|
|
|
|
created = {}
|
|
|
|
class FakeProcess:
|
|
def __init__(self, target, args):
|
|
created["target"] = target
|
|
created["args"] = args
|
|
self.pid = 4243
|
|
|
|
def start(self):
|
|
created["started"] = True
|
|
|
|
monkeypatch.setattr(ls.mp, "Process", FakeProcess)
|
|
|
|
calls = []
|
|
|
|
class ResolvedServerArgs:
|
|
def __init__(self, **fields):
|
|
self.port = fields.get("port", 30000)
|
|
self.base_gpu_id = fields.get("base_gpu_id", 0)
|
|
self.dp_size = fields.get("dp_size", 4)
|
|
self.tp_size = fields.get("tp_size", 2)
|
|
|
|
def replace_resolved(self, source, **changes):
|
|
calls.append((source, dict(changes)))
|
|
fields = dict(vars(self))
|
|
fields.update(changes)
|
|
return ResolvedServerArgs(**fields)
|
|
|
|
parent = ResolvedServerArgs()
|
|
proc = ls.launch_server_process(parent, worker_port=31002, dp_id=3)
|
|
|
|
assert created.get("started") is True
|
|
assert proc.pid == 4243
|
|
assert len(calls) == 1
|
|
source, changes = calls[0]
|
|
assert source == "sglang_router.launch_server_process"
|
|
assert changes == {"port": 31002, "base_gpu_id": 6, "dp_size": 1}
|
|
|
|
worker = created["args"][0]
|
|
assert (worker.port, worker.base_gpu_id, worker.dp_size) == (31002, 6, 1)
|
|
assert (parent.port, parent.base_gpu_id, parent.dp_size) == (30000, 0, 4)
|