config: pass the Ray placement group as a launch argument

`RayEngine.__init__` set the caller's `PlacementGroup` on `ServerArgs` as an
undeclared attribute, and every consumer read it back off the config object:
the non-DP launch path took `pg` and `is_custom_pg` from it, the DP path had to
re-attach the handle after `dataclasses.replace` dropped it, the DP controller
asked the config whether the group was the caller's, and the Ray HTTP launcher
had to clear the field so a stale handle could not leak in.

A live cluster object is not config. It now travels as a `placement_group`
argument on the two launch hooks — `Engine._launch_subprocesses` and
`Engine._launch_scheduler_processes` — defaulting to `None`, with `RayEngine`
holding the caller's group in `_placement_group` and the DP path passing
`is_custom_pg` down to the controller. The public API is unchanged:
`RayEngine(placement_group=pg, ...)` still works, `pg` still falls back to the
ambient group, and the HTTP launcher gets `None` from the default.

Writer ratchet 18 -> 15.
This commit is contained in:
Cheng Wan
2026-08-05 19:28:39 -07:00
committed by GitHub
parent f9b954ddb1
commit 1d47952c7c
5 changed files with 29 additions and 18 deletions
+14 -1
View File
@@ -213,6 +213,10 @@ class Engine(EngineScoreMixin, EngineBase):
run_scheduler_process_func: Callable = staticmethod(run_scheduler_process)
run_detokenizer_process_func: Callable = staticmethod(run_detokenizer_process)
# Backend-specific launch handle: the Ray engine schedules its actors onto a
# placement group. Not config — a live cluster object.
_placement_group = None
def __init__(self, **kwargs):
"""
The arguments of this function is the same as `sglang/srt/server_args.py::ServerArgs`.
@@ -264,6 +268,7 @@ class Engine(EngineScoreMixin, EngineBase):
init_tokenizer_manager_func=self.init_tokenizer_manager_func,
run_scheduler_process_func=self.run_scheduler_process_func,
run_detokenizer_process_func=self.run_detokenizer_process_func,
placement_group=self._placement_group,
)
self.tokenizer_manager = tokenizer_manager
self.template_manager = template_manager
@@ -825,6 +830,8 @@ class Engine(EngineScoreMixin, EngineBase):
server_args: ServerArgs,
port_args: PortArgs,
run_scheduler_process_func: Callable,
*,
placement_group=None,
) -> Tuple[SchedulerInitResult, Optional[List]]:
"""Launch scheduler processes using multiprocessing.
Override in subclasses for different backends (e.g. Ray).
@@ -1029,6 +1036,7 @@ class Engine(EngineScoreMixin, EngineBase):
run_scheduler_process_func: Callable,
run_detokenizer_process_func: Callable,
port_args: Optional[PortArgs] = None,
placement_group=None,
) -> Tuple[
TokenizerManager,
TemplateManager,
@@ -1089,8 +1097,13 @@ class Engine(EngineScoreMixin, EngineBase):
weight_cache_daemon_procs = cls._launch_weight_cache_daemons(server_args)
# Launch scheduler processes
# Passed only when there is one: this hook is an override point, and a
# subclass written against the three-argument signature must keep working.
launch_kwargs = (
{"placement_group": placement_group} if placement_group is not None else {}
)
scheduler_init_result, scheduler_procs = cls._launch_scheduler_processes(
server_args, port_args, run_scheduler_process_func
server_args, port_args, run_scheduler_process_func, **launch_kwargs
)
scheduler_init_result.engine_info_bootstrap_server = (
engine_info_bootstrap_server
@@ -52,11 +52,13 @@ class RayDataParallelController(DataParallelController):
placement_group,
bundle_for_node: Optional[List[int]],
rank0_node_ip: str,
is_custom_pg: bool = False,
):
# Set Ray-specific attributes BEFORE super().__init__() because the
# parent constructor calls launch_dp_schedulers / launch_dp_attention_schedulers
# which we override, and those methods need these attributes.
self.pg = placement_group
self.is_custom_pg = is_custom_pg
self.bundle_for_node = bundle_for_node
self.rank0_node_ip = rank0_node_ip
self.scheduler_actors: List = []
@@ -137,7 +139,7 @@ class RayDataParallelController(DataParallelController):
nnodes = server_args.nnodes
batch_start_idx = len(self.scheduler_actors)
if self.server_args.placement_group is None:
if not self.is_custom_pg:
for node_idx in range(nnodes):
bundle_idx = self.bundle_for_node[node_idx]
pp_range, tp_range, pp_per_node, tp_per_node = _calculate_rank_ranges(
+11 -13
View File
@@ -230,12 +230,12 @@ class RayEngine(Engine):
"""Engine using Ray actors for scheduler processes."""
def __init__(self, **kwargs):
placement_group = kwargs.pop("placement_group", None)
# Set before super().__init__(): it launches the subprocesses, which need
# the group to schedule the scheduler actors onto.
self._placement_group = kwargs.pop("placement_group", None)
if "log_level" not in kwargs:
kwargs["log_level"] = "error"
server_args = ServerArgs(**kwargs)
server_args.override("ray.placement_group", placement_group=placement_group)
super().__init__(server_args=server_args)
super().__init__(server_args=ServerArgs(**kwargs))
def shutdown(self):
"""Shutdown the engine — kill Ray scheduler actors then local processes."""
@@ -252,6 +252,8 @@ class RayEngine(Engine):
server_args: ServerArgs,
port_args: PortArgs,
run_scheduler_process_func: Callable,
*,
placement_group=None,
) -> tuple[SchedulerInitResult, None]:
"""Launch schedulers as Ray actors.
@@ -259,7 +261,7 @@ class RayEngine(Engine):
Tuple of (RaySchedulerInitResult, None).
scheduler_procs is None since Ray uses actors instead of mp.Process.
"""
pg = server_args.placement_group or ray.util.get_current_placement_group()
pg = placement_group or ray.util.get_current_placement_group()
if pg is None:
from ray.util.placement_group import (
placement_group as create_placement_group,
@@ -288,7 +290,7 @@ class RayEngine(Engine):
)
ray.get(pg.ready())
is_custom_pg = server_args.placement_group is not None
is_custom_pg = placement_group is not None
nnodes = server_args.nnodes
world_size = _compute_world_size(server_args)
@@ -424,6 +426,7 @@ class RayEngine(Engine):
pg,
bundle_for_node,
rank0_node_ip,
is_custom_pg,
),
None,
)
@@ -436,6 +439,7 @@ class RayEngine(Engine):
pg,
bundle_for_node: Optional[List[int]],
rank0_node_ip: str,
is_custom_pg: bool = False,
) -> RaySchedulerInitResult:
"""Launch DP schedulers via RayDataParallelController."""
from sglang.srt.ray.data_parallel_controller import (
@@ -461,16 +465,10 @@ class RayEngine(Engine):
server_args,
dist_init_addr=f"{rank0_node_ip}:{port_args.nccl_port}",
)
# dataclasses.replace only copies declared fields; placement_group is
# a dynamic attribute that must be manually appended after the rebuild.
dp_server_args.override(
"ray.placement_group", placement_group=server_args.placement_group
)
# Create the DP controller in-process. This blocks until all actors
# are initialized and their event loops have started.
controller = RayDataParallelController(
dp_server_args, port_args, pg, bundle_for_node, rank0_node_ip
dp_server_args, port_args, pg, bundle_for_node, rank0_node_ip, is_custom_pg
)
# Start the DP controller's event loop in a daemon thread.
-2
View File
@@ -44,8 +44,6 @@ def launch_server(
if execute_warmup_func is None:
execute_warmup_func = _execute_server_warmup
server_args.override("ray.http_server.clear_placement_group", placement_group=None)
(
tokenizer_manager,
template_manager,