feat: add SGLANG_RAY_BUNDLE_INDICES for fine-grained Ray bundle index control (#24667)
Signed-off-by: Haichuan Hu <kaisennhu@gmail.com>
This commit is contained in:
Haichuan Hu
@@ -323,6 +323,8 @@ class Envs:
|
||||
# Model Parallel
|
||||
SGLANG_USE_MESSAGE_QUEUE_BROADCASTER = EnvBool(True)
|
||||
SGLANG_ONE_VISIBLE_DEVICE_PER_PROCESS = EnvBool(False)
|
||||
# Comma-separated bundle indices for Ray Custom PG mode (e.g., "0,1,2,7").
|
||||
SGLANG_RAY_BUNDLE_INDICES = EnvStr("")
|
||||
# Override the distributed init method used by torch.distributed.init_process_group.
|
||||
# Set to "env://" to use an externally-created TCPStore via MASTER_ADDR/MASTER_PORT.
|
||||
SGLANG_DISTRIBUTED_INIT_METHOD_OVERRIDE = EnvStr(None)
|
||||
|
||||
@@ -20,15 +20,16 @@ from typing import List, Optional
|
||||
|
||||
import ray
|
||||
import zmq
|
||||
from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
|
||||
|
||||
from sglang.srt.entrypoints.engine import (
|
||||
_calculate_rank_ranges,
|
||||
_compute_parallelism_ranks,
|
||||
)
|
||||
from sglang.srt.entrypoints.engine import _calculate_rank_ranges
|
||||
from sglang.srt.layers.dp_attention import compute_dp_attention_world_info
|
||||
from sglang.srt.managers.data_parallel_controller import DataParallelController
|
||||
from sglang.srt.ray.scheduler_actor import SchedulerActor
|
||||
from sglang.srt.ray.engine import (
|
||||
_compute_world_size,
|
||||
_create_scheduler_actor,
|
||||
_get_bundle_node_ip,
|
||||
_resolve_bundle_indices,
|
||||
)
|
||||
from sglang.srt.server_args import PortArgs, ServerArgs
|
||||
from sglang.srt.utils.network import bind_port, get_zmq_socket, get_zmq_socket_on_host
|
||||
|
||||
@@ -48,7 +49,7 @@ class RayDataParallelController(DataParallelController):
|
||||
server_args: ServerArgs,
|
||||
port_args: PortArgs,
|
||||
placement_group,
|
||||
bundle_for_node: List[int],
|
||||
bundle_for_node: Optional[List[int]],
|
||||
rank0_node_ip: str,
|
||||
):
|
||||
# Set Ray-specific attributes BEFORE super().__init__() because the
|
||||
@@ -127,87 +128,129 @@ class RayDataParallelController(DataParallelController):
|
||||
):
|
||||
"""Create SchedulerActor Ray actors for one TP group (one DP rank).
|
||||
|
||||
For DP attention, dp_rank=None and worker_ports is provided; the dp_rank
|
||||
is derived from tp_rank via compute_dp_attention_world_info.
|
||||
|
||||
For regular DP, dp_rank is an integer and worker_ports is None.
|
||||
Args:
|
||||
dp_rank: DP rank for regular DP; None for DP attention (derived from tp_rank).
|
||||
worker_ports: Pre-allocated ports for DP attention; None for regular DP.
|
||||
"""
|
||||
nnodes = server_args.nnodes
|
||||
batch_start_idx = len(self.scheduler_actors)
|
||||
|
||||
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(
|
||||
nnodes, server_args.pp_size, server_args.tp_size, node_rank=node_idx
|
||||
)
|
||||
if self.server_args.placement_group is None:
|
||||
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(
|
||||
nnodes, server_args.pp_size, server_args.tp_size, node_rank=node_idx
|
||||
)
|
||||
for pp_rank in pp_range:
|
||||
for tp_rank in tp_range:
|
||||
rank_port_args = port_args
|
||||
actual_dp_rank = dp_rank
|
||||
|
||||
for pp_rank in pp_range:
|
||||
for tp_rank in tp_range:
|
||||
rank_port_args = port_args
|
||||
actual_dp_rank = dp_rank
|
||||
|
||||
if server_args.enable_dp_attention:
|
||||
# DP attention: derive dp_rank from tp_rank
|
||||
_, _, actual_dp_rank, _ = compute_dp_attention_world_info(
|
||||
server_args.enable_dp_attention,
|
||||
tp_rank,
|
||||
server_args.tp_size,
|
||||
server_args.dp_size,
|
||||
server_args.attn_cp_size,
|
||||
local_gpu_idx = (pp_rank % pp_per_node) * tp_per_node + (
|
||||
tp_rank % tp_per_node
|
||||
)
|
||||
rank_port_args = PortArgs.init_new(
|
||||
server_args, actual_dp_rank, worker_ports
|
||||
|
||||
if server_args.enable_dp_attention:
|
||||
_, _, actual_dp_rank, _ = compute_dp_attention_world_info(
|
||||
server_args.enable_dp_attention,
|
||||
tp_rank,
|
||||
server_args.tp_size,
|
||||
server_args.dp_size,
|
||||
server_args.attn_cp_size,
|
||||
)
|
||||
rank_port_args = PortArgs.init_new(
|
||||
server_args, actual_dp_rank, worker_ports
|
||||
)
|
||||
# All DP ranks share the same NCCL port (reuse TP group)
|
||||
rank_port_args.nccl_port = port_args.nccl_port
|
||||
rank_port_args.instance_id = port_args.instance_id
|
||||
# The detokenizer and tokenizer bind using the
|
||||
# original port_args addresses (127.0.0.1 when
|
||||
# dist_init_addr is unset). Scheduler actors must
|
||||
# connect to the same addresses.
|
||||
rank_port_args.detokenizer_ipc_name = (
|
||||
port_args.detokenizer_ipc_name
|
||||
)
|
||||
rank_port_args.tokenizer_ipc_name = (
|
||||
port_args.tokenizer_ipc_name
|
||||
)
|
||||
|
||||
dist_init_addr = (
|
||||
f"{self.rank0_node_ip}:{rank_port_args.nccl_port}"
|
||||
)
|
||||
# All DP ranks share the same NCCL port (reuse TP group)
|
||||
rank_port_args.nccl_port = port_args.nccl_port
|
||||
rank_port_args.instance_id = port_args.instance_id
|
||||
# The detokenizer and tokenizer bind using the
|
||||
# original port_args addresses (127.0.0.1 when
|
||||
# dist_init_addr is unset). Scheduler actors must
|
||||
# connect to the same addresses.
|
||||
rank_port_args.detokenizer_ipc_name = (
|
||||
port_args.detokenizer_ipc_name
|
||||
|
||||
actor = _create_scheduler_actor(
|
||||
pg=self.pg,
|
||||
bundle_idx=bundle_idx,
|
||||
gpu_id=local_gpu_idx,
|
||||
server_args=server_args,
|
||||
port_args=rank_port_args,
|
||||
tp_rank=tp_rank,
|
||||
pp_rank=pp_rank,
|
||||
dp_rank=actual_dp_rank,
|
||||
dist_init_addr=dist_init_addr,
|
||||
rank0_node_ip=self.rank0_node_ip,
|
||||
)
|
||||
rank_port_args.tokenizer_ipc_name = port_args.tokenizer_ipc_name
|
||||
self.scheduler_actors.append(actor)
|
||||
|
||||
local_gpu_idx = (pp_rank % pp_per_node) * tp_per_node + (
|
||||
tp_rank % tp_per_node
|
||||
else:
|
||||
world_size = _compute_world_size(server_args)
|
||||
bundle_indices = _resolve_bundle_indices(self.pg, world_size)
|
||||
|
||||
ranks_per_tp_group = server_args.tp_size * server_args.pp_size
|
||||
if dp_rank is not None:
|
||||
start_rank = dp_rank * ranks_per_tp_group
|
||||
end_rank = start_rank + ranks_per_tp_group
|
||||
# Each DP group must use its own local rank-0's node IP for
|
||||
# NCCL rendezvous, not the world rank-0's node IP.
|
||||
local_rank0_bundle_idx = bundle_indices[start_rank]
|
||||
local_rank0_node_ip = _get_bundle_node_ip(
|
||||
self.pg, local_rank0_bundle_idx
|
||||
)
|
||||
else:
|
||||
start_rank = 0
|
||||
end_rank = world_size
|
||||
local_rank0_node_ip = self.rank0_node_ip
|
||||
|
||||
for global_rank in range(start_rank, end_rank):
|
||||
local_rank = global_rank % ranks_per_tp_group
|
||||
pp_rank = local_rank // server_args.tp_size
|
||||
tp_rank = local_rank % server_args.tp_size
|
||||
rank_port_args = port_args
|
||||
actual_dp_rank = dp_rank
|
||||
|
||||
bundle_idx = bundle_indices[global_rank]
|
||||
|
||||
if server_args.enable_dp_attention:
|
||||
_, _, actual_dp_rank, _ = compute_dp_attention_world_info(
|
||||
server_args.enable_dp_attention,
|
||||
tp_rank,
|
||||
server_args.tp_size,
|
||||
server_args.dp_size,
|
||||
server_args.attn_cp_size,
|
||||
)
|
||||
|
||||
attn_cp_rank, moe_dp_rank, moe_ep_rank = _compute_parallelism_ranks(
|
||||
server_args, tp_rank
|
||||
rank_port_args = PortArgs.init_new(
|
||||
server_args, actual_dp_rank, worker_ports
|
||||
)
|
||||
rank_port_args.nccl_port = port_args.nccl_port
|
||||
rank_port_args.detokenizer_ipc_name = port_args.detokenizer_ipc_name
|
||||
rank_port_args.tokenizer_ipc_name = port_args.tokenizer_ipc_name
|
||||
|
||||
# Each DP group needs a unique dist_init_addr for its own
|
||||
# torch.distributed process group. Use nccl_port which is
|
||||
# unique per DP group (regular DP) or shared (DP attention).
|
||||
dist_init_addr = f"{self.rank0_node_ip}:{rank_port_args.nccl_port}"
|
||||
dist_init_addr = f"{local_rank0_node_ip}:{rank_port_args.nccl_port}"
|
||||
|
||||
actor = SchedulerActor.options(
|
||||
num_cpus=0,
|
||||
num_gpus=1,
|
||||
name=(
|
||||
f"sglang_scheduler_node{self.rank0_node_ip}"
|
||||
f"_dp{actual_dp_rank}_pp{pp_rank}_tp{tp_rank}"
|
||||
f"_pg{self.pg.id.hex()[:8]}_bundle{bundle_idx}"
|
||||
),
|
||||
scheduling_strategy=PlacementGroupSchedulingStrategy(
|
||||
placement_group=self.pg,
|
||||
placement_group_bundle_index=bundle_idx,
|
||||
),
|
||||
).remote(
|
||||
server_args=server_args,
|
||||
port_args=rank_port_args,
|
||||
gpu_id=local_gpu_idx,
|
||||
tp_rank=tp_rank,
|
||||
attn_cp_rank=attn_cp_rank,
|
||||
moe_dp_rank=moe_dp_rank,
|
||||
moe_ep_rank=moe_ep_rank,
|
||||
pp_rank=pp_rank,
|
||||
dp_rank=actual_dp_rank,
|
||||
dist_init_addr=dist_init_addr,
|
||||
)
|
||||
self.scheduler_actors.append(actor)
|
||||
actor = _create_scheduler_actor(
|
||||
pg=self.pg,
|
||||
bundle_idx=bundle_idx,
|
||||
gpu_id=0, # Each bundle has exactly 1 GPU
|
||||
server_args=server_args,
|
||||
port_args=rank_port_args,
|
||||
tp_rank=tp_rank,
|
||||
pp_rank=pp_rank,
|
||||
dp_rank=actual_dp_rank,
|
||||
dist_init_addr=dist_init_addr,
|
||||
rank0_node_ip=local_rank0_node_ip,
|
||||
)
|
||||
self.scheduler_actors.append(actor)
|
||||
|
||||
# Wait for all actors created in this call to initialize
|
||||
batch_actors = self.scheduler_actors[batch_start_idx:]
|
||||
|
||||
+252
-54
@@ -18,7 +18,7 @@ from __future__ import annotations
|
||||
import dataclasses
|
||||
import logging
|
||||
import threading
|
||||
from typing import Callable
|
||||
from typing import Callable, List, Optional
|
||||
|
||||
import ray
|
||||
from ray.util.placement_group import PlacementGroup
|
||||
@@ -30,6 +30,7 @@ from sglang.srt.entrypoints.engine import (
|
||||
_calculate_rank_ranges,
|
||||
_compute_parallelism_ranks,
|
||||
)
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.ray.scheduler_actor import SchedulerActor
|
||||
from sglang.srt.server_args import PortArgs, ServerArgs
|
||||
|
||||
@@ -76,9 +77,166 @@ def _find_engine_bundle(
|
||||
)
|
||||
|
||||
|
||||
def _get_bundle_node_ip(placement_group: PlacementGroup, bundle_idx: int) -> str:
|
||||
"""Get the IP address of the node where a specific bundle is located.
|
||||
|
||||
Args:
|
||||
placement_group: The placement group
|
||||
bundle_idx: Bundle index to query
|
||||
|
||||
Returns:
|
||||
IP address of the node where the bundle is located.
|
||||
"""
|
||||
|
||||
@ray.remote(num_cpus=0, num_gpus=0)
|
||||
def get_node_ip():
|
||||
return ray.util.get_node_ip_address()
|
||||
|
||||
return ray.get(
|
||||
get_node_ip.options(
|
||||
scheduling_strategy=PlacementGroupSchedulingStrategy(
|
||||
placement_group=placement_group,
|
||||
placement_group_bundle_index=bundle_idx,
|
||||
),
|
||||
).remote()
|
||||
)
|
||||
|
||||
|
||||
def _compute_world_size(server_args: ServerArgs) -> int:
|
||||
"""Compute world_size (total number of scheduler actors/GPUs needed).
|
||||
|
||||
Normal: dp_size * tp_size * pp_size; DP attention: tp_size * pp_size.
|
||||
"""
|
||||
if server_args.enable_dp_attention:
|
||||
return server_args.tp_size * server_args.pp_size
|
||||
return server_args.dp_size * server_args.tp_size * server_args.pp_size
|
||||
|
||||
|
||||
def _resolve_bundle_indices(pg: PlacementGroup, world_size: int) -> List[int]:
|
||||
"""Resolve bundle indices for Custom PG mode.
|
||||
|
||||
Parses SGLANG_RAY_BUNDLE_INDICES env var if set; otherwise returns
|
||||
sequential indices [0, 1, ..., world_size-1].
|
||||
|
||||
Args:
|
||||
pg: Placement group (used to get total_bundles count).
|
||||
world_size: Number of bundle indices expected (pre-computed via _compute_world_size).
|
||||
|
||||
Returns:
|
||||
List of bundle indices of length world_size.
|
||||
"""
|
||||
total_bundles = len(pg.bundle_specs)
|
||||
indices_str = envs.SGLANG_RAY_BUNDLE_INDICES.get()
|
||||
if not indices_str:
|
||||
return list(range(world_size))
|
||||
|
||||
indices = list(map(int, indices_str.split(",")))
|
||||
|
||||
if len(indices) != world_size:
|
||||
raise ValueError(
|
||||
f"SGLANG_RAY_BUNDLE_INDICES has {len(indices)} values, "
|
||||
f"expected {world_size}"
|
||||
)
|
||||
|
||||
if len(set(indices)) != len(indices):
|
||||
raise ValueError(f"SGLANG_RAY_BUNDLE_INDICES has duplicates: {indices}")
|
||||
|
||||
for idx in indices:
|
||||
if idx < 0 or idx >= total_bundles:
|
||||
raise ValueError(f"Bundle index {idx} out of range [0, {total_bundles})")
|
||||
|
||||
return indices
|
||||
|
||||
|
||||
def _validate_custom_placement_group(pg: PlacementGroup, world_size: int) -> None:
|
||||
"""Validate custom placement group: 1 GPU per bundle, enough GPU bundles for world_size.
|
||||
|
||||
Args:
|
||||
pg: User-provided placement group.
|
||||
world_size: Number of GPU bundles required.
|
||||
"""
|
||||
bundles = pg.bundle_specs
|
||||
gpu_bundle_count = 0
|
||||
for bundle in bundles:
|
||||
gpu_count = bundle.get("GPU", 0)
|
||||
if gpu_count > 1:
|
||||
raise ValueError(
|
||||
"Custom placement group must have exactly 1 GPU per bundle. "
|
||||
f"Found bundle with {gpu_count} GPUs."
|
||||
)
|
||||
if gpu_count > 0:
|
||||
gpu_bundle_count += 1
|
||||
|
||||
if gpu_bundle_count < world_size:
|
||||
raise ValueError(
|
||||
f"Custom placement group has {gpu_bundle_count} GPU bundles, "
|
||||
f"but needs {world_size} for world_size. "
|
||||
"Provide more bundles or reduce parallelism."
|
||||
)
|
||||
|
||||
|
||||
def _create_scheduler_actor(
|
||||
pg: PlacementGroup,
|
||||
bundle_idx: int,
|
||||
gpu_id: int,
|
||||
server_args: ServerArgs,
|
||||
port_args: PortArgs,
|
||||
tp_rank: int,
|
||||
pp_rank: int,
|
||||
dp_rank: int,
|
||||
dist_init_addr: str,
|
||||
rank0_node_ip: str,
|
||||
) -> SchedulerActor:
|
||||
"""Create a SchedulerActor on the given placement group bundle.
|
||||
|
||||
Args:
|
||||
pg: Placement group to schedule actor onto.
|
||||
bundle_idx: Bundle index within the placement group.
|
||||
gpu_id: GPU ID within the bundle (0 for custom PG, computed for auto PG).
|
||||
rank0_node_ip: IP of rank-0's node, used for NCCL rendezvous.
|
||||
dist_init_addr: Distributed init address (tcp://rank0_node_ip:nccl_port).
|
||||
"""
|
||||
attn_cp_rank, moe_dp_rank, moe_ep_rank = _compute_parallelism_ranks(
|
||||
server_args, tp_rank
|
||||
)
|
||||
|
||||
return SchedulerActor.options(
|
||||
num_cpus=0,
|
||||
num_gpus=1,
|
||||
name=(
|
||||
f"sglang_scheduler_node{rank0_node_ip}"
|
||||
f"_dp{dp_rank}_pp{pp_rank}_tp{tp_rank}"
|
||||
f"_pg{pg.id.hex()[:8]}_bundle{bundle_idx}"
|
||||
),
|
||||
scheduling_strategy=PlacementGroupSchedulingStrategy(
|
||||
placement_group=pg,
|
||||
placement_group_bundle_index=bundle_idx,
|
||||
),
|
||||
).remote(
|
||||
server_args=server_args,
|
||||
port_args=port_args,
|
||||
gpu_id=gpu_id,
|
||||
tp_rank=tp_rank,
|
||||
attn_cp_rank=attn_cp_rank,
|
||||
moe_dp_rank=moe_dp_rank,
|
||||
moe_ep_rank=moe_ep_rank,
|
||||
pp_rank=pp_rank,
|
||||
dp_rank=dp_rank,
|
||||
dist_init_addr=dist_init_addr,
|
||||
)
|
||||
|
||||
|
||||
class RayEngine(Engine):
|
||||
"""Engine using Ray actors for scheduler processes."""
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
placement_group = kwargs.pop("placement_group", None)
|
||||
if "log_level" not in kwargs:
|
||||
kwargs["log_level"] = "error"
|
||||
server_args = ServerArgs(**kwargs)
|
||||
server_args.placement_group = placement_group
|
||||
super().__init__(server_args=server_args)
|
||||
|
||||
def shutdown(self):
|
||||
"""Shutdown the engine — kill Ray scheduler actors then local processes."""
|
||||
for actor in self._scheduler_init_result.scheduler_actors:
|
||||
@@ -101,7 +259,7 @@ class RayEngine(Engine):
|
||||
Tuple of (RaySchedulerInitResult, None).
|
||||
scheduler_procs is None since Ray uses actors instead of mp.Process.
|
||||
"""
|
||||
pg = ray.util.get_current_placement_group()
|
||||
pg = server_args.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,
|
||||
@@ -130,69 +288,102 @@ class RayEngine(Engine):
|
||||
)
|
||||
ray.get(pg.ready())
|
||||
|
||||
is_custom_pg = server_args.placement_group is not None
|
||||
nnodes = server_args.nnodes
|
||||
world_size = _compute_world_size(server_args)
|
||||
|
||||
# co-located with the Engine and rank0 scheduler at the same node
|
||||
engine_bundle, engine_ip = _find_engine_bundle(pg, nnodes)
|
||||
bundle_for_node = [engine_bundle] + [
|
||||
i for i in range(nnodes) if i != engine_bundle
|
||||
]
|
||||
rank0_node_ip = engine_ip
|
||||
if not is_custom_pg:
|
||||
engine_bundle, engine_ip = _find_engine_bundle(pg, nnodes)
|
||||
bundle_for_node = [engine_bundle] + [
|
||||
i for i in range(nnodes) if i != engine_bundle
|
||||
]
|
||||
rank0_node_ip = engine_ip
|
||||
else:
|
||||
try:
|
||||
_validate_custom_placement_group(pg, world_size)
|
||||
except ValueError as e:
|
||||
logger.error(f"Custom placement group validation failed: {e}")
|
||||
raise RuntimeError(
|
||||
f"Custom placement group validation failed: {e}"
|
||||
) from e
|
||||
bundle_for_node = None
|
||||
indices_str = envs.SGLANG_RAY_BUNDLE_INDICES.get()
|
||||
rank0_bundle_idx = int(indices_str.split(",")[0]) if indices_str else 0
|
||||
rank0_node_ip = _get_bundle_node_ip(pg, rank0_bundle_idx)
|
||||
|
||||
if server_args.dp_size == 1:
|
||||
# Launch tensor parallel scheduler actors
|
||||
world_size = server_args.tp_size * server_args.pp_size
|
||||
gpus_per_node = world_size // nnodes
|
||||
|
||||
logger.info(
|
||||
f"Ray cluster: {nnodes} nodes, "
|
||||
f"Use {gpus_per_node} GPUs/node, world_size={world_size}"
|
||||
)
|
||||
|
||||
dist_init_addr = f"{rank0_node_ip}:{port_args.nccl_port}"
|
||||
logger.info(f"dist_init_addr: {dist_init_addr}")
|
||||
|
||||
scheduler_actors = []
|
||||
|
||||
for node_idx in range(nnodes):
|
||||
bundle_idx = bundle_for_node[node_idx]
|
||||
pp_range, tp_range, pp_per_node, tp_per_node = _calculate_rank_ranges(
|
||||
nnodes,
|
||||
server_args.pp_size,
|
||||
server_args.tp_size,
|
||||
node_rank=node_idx,
|
||||
if not is_custom_pg:
|
||||
gpus_per_node = world_size // nnodes
|
||||
logger.info(
|
||||
f"Ray cluster (auto PG): {nnodes} nodes, "
|
||||
f"{gpus_per_node} GPUs/node, world_size={world_size}"
|
||||
)
|
||||
for pp_rank in pp_range:
|
||||
for tp_rank in tp_range:
|
||||
local_gpu_idx = (pp_rank % pp_per_node) * tp_per_node + (
|
||||
tp_rank % tp_per_node
|
||||
)
|
||||
|
||||
attn_cp_rank, moe_dp_rank, moe_ep_rank = (
|
||||
_compute_parallelism_ranks(server_args, tp_rank)
|
||||
for node_idx in range(nnodes):
|
||||
bundle_idx = bundle_for_node[node_idx]
|
||||
pp_range, tp_range, pp_per_node, tp_per_node = (
|
||||
_calculate_rank_ranges(
|
||||
nnodes,
|
||||
server_args.pp_size,
|
||||
server_args.tp_size,
|
||||
node_rank=node_idx,
|
||||
)
|
||||
)
|
||||
for pp_rank in pp_range:
|
||||
for tp_rank in tp_range:
|
||||
local_gpu_idx = (pp_rank % pp_per_node) * tp_per_node + (
|
||||
tp_rank % tp_per_node
|
||||
)
|
||||
|
||||
actor = SchedulerActor.options(
|
||||
num_cpus=0,
|
||||
num_gpus=1,
|
||||
name=f"sglang_scheduler_node{rank0_node_ip}_pp{pp_rank}_tp{tp_rank}_pg{pg.id.hex()[:8]}_bundle{bundle_idx}",
|
||||
scheduling_strategy=PlacementGroupSchedulingStrategy(
|
||||
placement_group=pg,
|
||||
placement_group_bundle_index=bundle_idx,
|
||||
),
|
||||
).remote(
|
||||
server_args=server_args,
|
||||
port_args=port_args,
|
||||
gpu_id=local_gpu_idx,
|
||||
tp_rank=tp_rank,
|
||||
attn_cp_rank=attn_cp_rank,
|
||||
moe_dp_rank=moe_dp_rank,
|
||||
moe_ep_rank=moe_ep_rank,
|
||||
pp_rank=pp_rank,
|
||||
dp_rank=0,
|
||||
dist_init_addr=dist_init_addr,
|
||||
)
|
||||
scheduler_actors.append(actor)
|
||||
actor = _create_scheduler_actor(
|
||||
pg=pg,
|
||||
bundle_idx=bundle_idx,
|
||||
gpu_id=local_gpu_idx,
|
||||
server_args=server_args,
|
||||
port_args=port_args,
|
||||
tp_rank=tp_rank,
|
||||
pp_rank=pp_rank,
|
||||
dp_rank=0,
|
||||
dist_init_addr=dist_init_addr,
|
||||
rank0_node_ip=rank0_node_ip,
|
||||
)
|
||||
scheduler_actors.append(actor)
|
||||
|
||||
else:
|
||||
try:
|
||||
bundle_indices = _resolve_bundle_indices(pg, world_size)
|
||||
except ValueError as e:
|
||||
logger.error(f"Failed to resolve bundle indices: {e}")
|
||||
raise RuntimeError(f"Failed to resolve bundle indices: {e}") from e
|
||||
|
||||
logger.info(
|
||||
f"Ray cluster (custom PG): world_size={world_size}, "
|
||||
f"bundle_indices={bundle_indices}"
|
||||
)
|
||||
|
||||
for rank in range(world_size):
|
||||
pp_rank = rank // server_args.tp_size
|
||||
tp_rank = rank % server_args.tp_size
|
||||
bundle_idx = bundle_indices[rank]
|
||||
|
||||
actor = _create_scheduler_actor(
|
||||
pg=pg,
|
||||
bundle_idx=bundle_idx,
|
||||
gpu_id=0, # Each bundle has exactly 1 GPU
|
||||
server_args=server_args,
|
||||
port_args=port_args,
|
||||
tp_rank=tp_rank,
|
||||
pp_rank=pp_rank,
|
||||
dp_rank=0,
|
||||
dist_init_addr=dist_init_addr,
|
||||
rank0_node_ip=rank0_node_ip,
|
||||
)
|
||||
scheduler_actors.append(actor)
|
||||
|
||||
try:
|
||||
scheduler_infos = ray.get(
|
||||
@@ -228,7 +419,11 @@ class RayEngine(Engine):
|
||||
# Launch the data parallel controller
|
||||
return (
|
||||
cls._launch_dp_scheduler_processes(
|
||||
server_args, port_args, pg, bundle_for_node, rank0_node_ip
|
||||
server_args,
|
||||
port_args,
|
||||
pg,
|
||||
bundle_for_node,
|
||||
rank0_node_ip,
|
||||
),
|
||||
None,
|
||||
)
|
||||
@@ -239,7 +434,7 @@ class RayEngine(Engine):
|
||||
server_args: ServerArgs,
|
||||
port_args: PortArgs,
|
||||
pg,
|
||||
bundle_for_node: list,
|
||||
bundle_for_node: Optional[List[int]],
|
||||
rank0_node_ip: str,
|
||||
) -> RaySchedulerInitResult:
|
||||
"""Launch DP schedulers via RayDataParallelController."""
|
||||
@@ -266,6 +461,9 @@ 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.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.
|
||||
|
||||
@@ -44,6 +44,8 @@ def launch_server(
|
||||
if execute_warmup_func is None:
|
||||
execute_warmup_func = _execute_server_warmup
|
||||
|
||||
server_args.placement_group = None
|
||||
|
||||
(
|
||||
tokenizer_manager,
|
||||
template_manager,
|
||||
|
||||
Reference in New Issue
Block a user