diff --git a/miles/backends/sglang_utils/sglang_engine.py b/miles/backends/sglang_utils/sglang_engine.py index 0fa47831e7..4b596bc082 100644 --- a/miles/backends/sglang_utils/sglang_engine.py +++ b/miles/backends/sglang_utils/sglang_engine.py @@ -48,13 +48,7 @@ def build_server_url(host: str, port: int) -> str: return f"http://{format_v6_uri(host)}:{port}" -@dataclasses.dataclass(frozen=True) -class EngineLaunchPlan: - cmd: str - api_key: str | None - - -def compute_engine_launch_plan( +def compute_engine_launch_cmd( args, *, node_rank: int, @@ -63,7 +57,7 @@ def compute_engine_launch_plan( sglang_overrides: dict, num_gpus_per_engine: int, addr_and_ports: dict, -) -> EngineLaunchPlan: +) -> str: server_args_dict = _compute_server_args( args, node_rank=node_rank, @@ -80,8 +74,11 @@ def compute_engine_launch_plan( ) launch_args = {**server_args_dict, "host": server_args_dict["host"].strip("[]")} - cmd = shlex.join([sys.executable, "-m", "sglang.launch_server", *server_args_to_argv(launch_args)]) - return EngineLaunchPlan(cmd=cmd, api_key=server_args_dict.get("api_key")) + return shlex.join([sys.executable, "-m", "sglang.launch_server", *server_args_to_argv(launch_args)]) + + +def compute_api_key(args, *, sglang_overrides: dict) -> str | None: + return sglang_overrides.get("api_key", args.sglang_api_key) def _compute_server_args( diff --git a/miles/ray/rollout/server_cell.py b/miles/ray/rollout/server_cell.py index 81f86d83d9..0dd4b7435b 100644 --- a/miles/ray/rollout/server_cell.py +++ b/miles/ray/rollout/server_cell.py @@ -10,7 +10,12 @@ from sglang.srt.constants import GPU_MEMORY_TYPE_WEIGHTS from miles.backends.sglang_utils.sglang_api_client import SGLangApiClient, wait_server_healthy -from miles.backends.sglang_utils.sglang_engine import build_server_url, compute_engine_launch_plan, format_v6_uri +from miles.backends.sglang_utils.sglang_engine import ( + build_server_url, + compute_api_key, + compute_engine_launch_cmd, + format_v6_uri, +) from miles.backends.sglang_utils.sglang_router_api_client import SGLangRouterApiClient, use_legacy_router_api from miles.ray.rollout.cell_state import ( AddrInfo, @@ -154,8 +159,8 @@ async def start_engines(self, port_allocator: PortAllocator) -> None: ] ) - plans = { - rank: compute_engine_launch_plan( + launch_cmds = { + rank: compute_engine_launch_cmd( self.args, node_rank=local_index, worker_type=self.worker_type, @@ -169,14 +174,14 @@ async def start_engines(self, port_allocator: PortAllocator) -> None: await asyncio.gather( *[ - actor.run.remote(cmd=plans[global_rank].cmd, envs={}) + actor.run.remote(cmd=launch_cmds[global_rank], envs={}) for global_rank, actor in zip(global_ranks, actor_handles, strict=True) ] ) await wait_server_healthy( server_url=self.addr_info.server_url, - api_key=plans[global_ranks[0]].api_key, + api_key=compute_api_key(self.args, sglang_overrides=self.sglang_overrides), is_process_alive=functools.partial(_engine_actor_is_alive, self.primary_actor_handle), ) diff --git a/tests/fast/backends/sglang_utils/conftest.py b/tests/fast/backends/sglang_utils/conftest.py index 11e250fab2..8b220b2671 100644 --- a/tests/fast/backends/sglang_utils/conftest.py +++ b/tests/fast/backends/sglang_utils/conftest.py @@ -49,6 +49,7 @@ def make_engine_args(**overrides: Any) -> Namespace: use_rollout_indexer_replay=False, fp16=False, lora_rank=0, + sglang_api_key=None, lora_adapter_path=None, multi_lora=False, colocate=False, diff --git a/tests/fast/backends/sglang_utils/test_sglang_engine.py b/tests/fast/backends/sglang_utils/test_sglang_engine.py index 5d91ca73e8..85a4d136ee 100644 --- a/tests/fast/backends/sglang_utils/test_sglang_engine.py +++ b/tests/fast/backends/sglang_utils/test_sglang_engine.py @@ -9,10 +9,10 @@ pytest.importorskip("sglang") from miles.backends.sglang_utils.server_args_utils import parse_server_args_argv -from miles.backends.sglang_utils.sglang_engine import compute_engine_launch_plan +from miles.backends.sglang_utils.sglang_engine import compute_api_key, compute_engine_launch_cmd -def _plan(*, worker_type: str = "regular", args=None, addr_overrides: dict | None = None, **kwargs): +def _cmd(*, worker_type: str = "regular", args=None, addr_overrides: dict | None = None, **kwargs) -> str: addr_and_ports = dict( host="10.0.0.1", port=30000, @@ -21,7 +21,7 @@ def _plan(*, worker_type: str = "regular", args=None, addr_overrides: dict | Non dist_init_addr="10.0.0.1:20000", ) addr_and_ports.update(addr_overrides or {}) - return compute_engine_launch_plan( + return compute_engine_launch_cmd( args or make_engine_args(), node_rank=0, worker_type=worker_type, @@ -33,11 +33,10 @@ def _plan(*, worker_type: str = "regular", args=None, addr_overrides: dict | Non ) -class TestComputeEngineLaunchPlan: +class TestComputeEngineLaunchCmd: def test_the_command_launches_sglang_with_the_allocated_addressing(self): - """The plan renders one launch_server command carrying the addr map.""" - plan = _plan() - tokens = shlex.split(plan.cmd) + """The rendered launch_server command carries the addr map.""" + tokens = shlex.split(_cmd()) assert tokens[:3] == [sys.executable, "-m", "sglang.launch_server"] parsed = parse_server_args_argv(tokens[3:]) assert parsed.host == "10.0.0.1" and parsed.port == 30000 @@ -53,24 +52,36 @@ def test_every_plan_picks_a_fresh_random_seed(self): def test_a_bracketed_v6_host_is_stripped_for_the_server_but_kept_in_dist_addr(self): """sglang binds a bare v6 host while the rendezvous addr stays bracketed.""" - plan = _plan(addr_overrides=dict(host="[fd00::2]", port=31007, dist_init_addr="[fd00::1]:15003")) - parsed = parse_server_args_argv(shlex.split(plan.cmd)[3:]) + cmd = _cmd(addr_overrides=dict(host="[fd00::2]", port=31007, dist_init_addr="[fd00::1]:15003")) + parsed = parse_server_args_argv(shlex.split(cmd)[3:]) assert parsed.host == "fd00::2" assert parsed.dist_init_addr == "[fd00::1]:15003" - def test_a_prefill_plan_carries_the_bootstrap_port(self): + def test_a_prefill_command_carries_the_bootstrap_port(self): """PD-disaggregation prefill flags survive into the command.""" - plan = _plan(worker_type="prefill", addr_overrides=dict(disaggregation_bootstrap_port=20090)) - parsed = parse_server_args_argv(shlex.split(plan.cmd)[3:]) + cmd = _cmd(worker_type="prefill", addr_overrides=dict(disaggregation_bootstrap_port=20090)) + parsed = parse_server_args_argv(shlex.split(cmd)[3:]) assert parsed.disaggregation_mode == "prefill" assert parsed.disaggregation_bootstrap_port == 20090 - def test_the_plan_exposes_the_api_key_for_the_health_wait(self): - """The driver-side health wait needs the same api key the server got.""" - args = make_engine_args() - args.sglang_api_key = "secret" - assert _plan(args=args).api_key == "secret" + def test_the_command_carries_the_api_key_from_args(self): + """--sglang-api-key reaches the server through the generic passthrough.""" + cmd = _cmd(args=make_engine_args(sglang_api_key="secret")) + parsed = parse_server_args_argv(shlex.split(cmd)[3:]) + assert parsed.api_key == "secret" + + +class TestComputeApiKey: + def test_the_args_key_is_used_when_no_override_exists(self): + """The health wait needs the same key the generic passthrough gave the server.""" + args = make_engine_args(sglang_api_key="secret") + assert compute_api_key(args, sglang_overrides={}) == "secret" + + def test_an_override_key_wins_over_the_args_key(self): + """Overrides beat args exactly like they do in the rendered command.""" + args = make_engine_args(sglang_api_key="from-args") + assert compute_api_key(args, sglang_overrides={"api_key": "from-override"}) == "from-override" - def test_the_plan_has_no_api_key_when_the_server_has_none(self): + def test_no_key_anywhere_means_the_health_wait_sends_none(self): """No key configured means the health wait sends none.""" - assert _plan().api_key is None + assert compute_api_key(make_engine_args(), sglang_overrides={}) is None