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_scheduler_process_func: Callable = staticmethod(run_scheduler_process)
|
||||||
run_detokenizer_process_func: Callable = staticmethod(run_detokenizer_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):
|
def __init__(self, **kwargs):
|
||||||
"""
|
"""
|
||||||
The arguments of this function is the same as `sglang/srt/server_args.py::ServerArgs`.
|
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,
|
init_tokenizer_manager_func=self.init_tokenizer_manager_func,
|
||||||
run_scheduler_process_func=self.run_scheduler_process_func,
|
run_scheduler_process_func=self.run_scheduler_process_func,
|
||||||
run_detokenizer_process_func=self.run_detokenizer_process_func,
|
run_detokenizer_process_func=self.run_detokenizer_process_func,
|
||||||
|
placement_group=self._placement_group,
|
||||||
)
|
)
|
||||||
self.tokenizer_manager = tokenizer_manager
|
self.tokenizer_manager = tokenizer_manager
|
||||||
self.template_manager = template_manager
|
self.template_manager = template_manager
|
||||||
@@ -825,6 +830,8 @@ class Engine(EngineScoreMixin, EngineBase):
|
|||||||
server_args: ServerArgs,
|
server_args: ServerArgs,
|
||||||
port_args: PortArgs,
|
port_args: PortArgs,
|
||||||
run_scheduler_process_func: Callable,
|
run_scheduler_process_func: Callable,
|
||||||
|
*,
|
||||||
|
placement_group=None,
|
||||||
) -> Tuple[SchedulerInitResult, Optional[List]]:
|
) -> Tuple[SchedulerInitResult, Optional[List]]:
|
||||||
"""Launch scheduler processes using multiprocessing.
|
"""Launch scheduler processes using multiprocessing.
|
||||||
Override in subclasses for different backends (e.g. Ray).
|
Override in subclasses for different backends (e.g. Ray).
|
||||||
@@ -1029,6 +1036,7 @@ class Engine(EngineScoreMixin, EngineBase):
|
|||||||
run_scheduler_process_func: Callable,
|
run_scheduler_process_func: Callable,
|
||||||
run_detokenizer_process_func: Callable,
|
run_detokenizer_process_func: Callable,
|
||||||
port_args: Optional[PortArgs] = None,
|
port_args: Optional[PortArgs] = None,
|
||||||
|
placement_group=None,
|
||||||
) -> Tuple[
|
) -> Tuple[
|
||||||
TokenizerManager,
|
TokenizerManager,
|
||||||
TemplateManager,
|
TemplateManager,
|
||||||
@@ -1089,8 +1097,13 @@ class Engine(EngineScoreMixin, EngineBase):
|
|||||||
weight_cache_daemon_procs = cls._launch_weight_cache_daemons(server_args)
|
weight_cache_daemon_procs = cls._launch_weight_cache_daemons(server_args)
|
||||||
|
|
||||||
# Launch scheduler processes
|
# 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(
|
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 = (
|
scheduler_init_result.engine_info_bootstrap_server = (
|
||||||
engine_info_bootstrap_server
|
engine_info_bootstrap_server
|
||||||
|
|||||||
@@ -52,11 +52,13 @@ class RayDataParallelController(DataParallelController):
|
|||||||
placement_group,
|
placement_group,
|
||||||
bundle_for_node: Optional[List[int]],
|
bundle_for_node: Optional[List[int]],
|
||||||
rank0_node_ip: str,
|
rank0_node_ip: str,
|
||||||
|
is_custom_pg: bool = False,
|
||||||
):
|
):
|
||||||
# Set Ray-specific attributes BEFORE super().__init__() because the
|
# Set Ray-specific attributes BEFORE super().__init__() because the
|
||||||
# parent constructor calls launch_dp_schedulers / launch_dp_attention_schedulers
|
# parent constructor calls launch_dp_schedulers / launch_dp_attention_schedulers
|
||||||
# which we override, and those methods need these attributes.
|
# which we override, and those methods need these attributes.
|
||||||
self.pg = placement_group
|
self.pg = placement_group
|
||||||
|
self.is_custom_pg = is_custom_pg
|
||||||
self.bundle_for_node = bundle_for_node
|
self.bundle_for_node = bundle_for_node
|
||||||
self.rank0_node_ip = rank0_node_ip
|
self.rank0_node_ip = rank0_node_ip
|
||||||
self.scheduler_actors: List = []
|
self.scheduler_actors: List = []
|
||||||
@@ -137,7 +139,7 @@ class RayDataParallelController(DataParallelController):
|
|||||||
nnodes = server_args.nnodes
|
nnodes = server_args.nnodes
|
||||||
batch_start_idx = len(self.scheduler_actors)
|
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):
|
for node_idx in range(nnodes):
|
||||||
bundle_idx = self.bundle_for_node[node_idx]
|
bundle_idx = self.bundle_for_node[node_idx]
|
||||||
pp_range, tp_range, pp_per_node, tp_per_node = _calculate_rank_ranges(
|
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."""
|
"""Engine using Ray actors for scheduler processes."""
|
||||||
|
|
||||||
def __init__(self, **kwargs):
|
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:
|
if "log_level" not in kwargs:
|
||||||
kwargs["log_level"] = "error"
|
kwargs["log_level"] = "error"
|
||||||
server_args = ServerArgs(**kwargs)
|
super().__init__(server_args=ServerArgs(**kwargs))
|
||||||
server_args.override("ray.placement_group", placement_group=placement_group)
|
|
||||||
super().__init__(server_args=server_args)
|
|
||||||
|
|
||||||
def shutdown(self):
|
def shutdown(self):
|
||||||
"""Shutdown the engine — kill Ray scheduler actors then local processes."""
|
"""Shutdown the engine — kill Ray scheduler actors then local processes."""
|
||||||
@@ -252,6 +252,8 @@ class RayEngine(Engine):
|
|||||||
server_args: ServerArgs,
|
server_args: ServerArgs,
|
||||||
port_args: PortArgs,
|
port_args: PortArgs,
|
||||||
run_scheduler_process_func: Callable,
|
run_scheduler_process_func: Callable,
|
||||||
|
*,
|
||||||
|
placement_group=None,
|
||||||
) -> tuple[SchedulerInitResult, None]:
|
) -> tuple[SchedulerInitResult, None]:
|
||||||
"""Launch schedulers as Ray actors.
|
"""Launch schedulers as Ray actors.
|
||||||
|
|
||||||
@@ -259,7 +261,7 @@ class RayEngine(Engine):
|
|||||||
Tuple of (RaySchedulerInitResult, None).
|
Tuple of (RaySchedulerInitResult, None).
|
||||||
scheduler_procs is None since Ray uses actors instead of mp.Process.
|
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:
|
if pg is None:
|
||||||
from ray.util.placement_group import (
|
from ray.util.placement_group import (
|
||||||
placement_group as create_placement_group,
|
placement_group as create_placement_group,
|
||||||
@@ -288,7 +290,7 @@ class RayEngine(Engine):
|
|||||||
)
|
)
|
||||||
ray.get(pg.ready())
|
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
|
nnodes = server_args.nnodes
|
||||||
world_size = _compute_world_size(server_args)
|
world_size = _compute_world_size(server_args)
|
||||||
|
|
||||||
@@ -424,6 +426,7 @@ class RayEngine(Engine):
|
|||||||
pg,
|
pg,
|
||||||
bundle_for_node,
|
bundle_for_node,
|
||||||
rank0_node_ip,
|
rank0_node_ip,
|
||||||
|
is_custom_pg,
|
||||||
),
|
),
|
||||||
None,
|
None,
|
||||||
)
|
)
|
||||||
@@ -436,6 +439,7 @@ class RayEngine(Engine):
|
|||||||
pg,
|
pg,
|
||||||
bundle_for_node: Optional[List[int]],
|
bundle_for_node: Optional[List[int]],
|
||||||
rank0_node_ip: str,
|
rank0_node_ip: str,
|
||||||
|
is_custom_pg: bool = False,
|
||||||
) -> RaySchedulerInitResult:
|
) -> RaySchedulerInitResult:
|
||||||
"""Launch DP schedulers via RayDataParallelController."""
|
"""Launch DP schedulers via RayDataParallelController."""
|
||||||
from sglang.srt.ray.data_parallel_controller import (
|
from sglang.srt.ray.data_parallel_controller import (
|
||||||
@@ -461,16 +465,10 @@ class RayEngine(Engine):
|
|||||||
server_args,
|
server_args,
|
||||||
dist_init_addr=f"{rank0_node_ip}:{port_args.nccl_port}",
|
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
|
# Create the DP controller in-process. This blocks until all actors
|
||||||
# are initialized and their event loops have started.
|
# are initialized and their event loops have started.
|
||||||
controller = RayDataParallelController(
|
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.
|
# Start the DP controller's event loop in a daemon thread.
|
||||||
|
|||||||
@@ -44,8 +44,6 @@ def launch_server(
|
|||||||
if execute_warmup_func is None:
|
if execute_warmup_func is None:
|
||||||
execute_warmup_func = _execute_server_warmup
|
execute_warmup_func = _execute_server_warmup
|
||||||
|
|
||||||
server_args.override("ray.http_server.clear_placement_group", placement_group=None)
|
|
||||||
|
|
||||||
(
|
(
|
||||||
tokenizer_manager,
|
tokenizer_manager,
|
||||||
template_manager,
|
template_manager,
|
||||||
|
|||||||
@@ -49,7 +49,7 @@ _EXCLUDED = (
|
|||||||
"multimodal_gen",
|
"multimodal_gen",
|
||||||
)
|
)
|
||||||
|
|
||||||
_BASELINE = 18
|
_BASELINE = 15
|
||||||
|
|
||||||
|
|
||||||
class TestServerArgsWriterRatchet(CustomTestCase):
|
class TestServerArgsWriterRatchet(CustomTestCase):
|
||||||
|
|||||||
Reference in New Issue
Block a user