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:
@@ -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(
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user