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
@@ -47,6 +47,62 @@
|
||||
"Please see [the examples](https://github.com/sgl-project/sglang/tree/main/examples/runtime/engine) for further use cases."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Ray Integration\n",
|
||||
"\n",
|
||||
"When running in a Ray cluster, you can use `RayEngine` with a custom placement group for fine-grained GPU placement control.\n",
|
||||
"\n",
|
||||
"### Custom Placement Groups\n",
|
||||
"\n",
|
||||
"Pass a `placement_group` with 1-GPU-per-bundle bundles to control exactly which GPUs are used. Each bundle should have exactly 1 GPU for deterministic mapping.\n",
|
||||
"\n",
|
||||
"```python\n",
|
||||
"import ray\n",
|
||||
"from ray.util.placement_group import placement_group\n",
|
||||
"from sglang.srt.ray.engine import RayEngine\n",
|
||||
"\n",
|
||||
"ray.init()\n",
|
||||
"\n",
|
||||
"# Create placement group with specific GPU bundles\n",
|
||||
"pg = placement_group(\n",
|
||||
" [{\"GPU\": 1} for _ in range(4)], # 4 bundles, each with 1 GPU\n",
|
||||
" strategy=\"STRICT_PACK\",\n",
|
||||
")\n",
|
||||
"ray.get(pg.ready())\n",
|
||||
"\n",
|
||||
"# Launch RayEngine on custom placement group\n",
|
||||
"engine = RayEngine(\n",
|
||||
" model_path=\"meta-llama/Meta-Llama-3-8B-Instruct\",\n",
|
||||
" tp_size=4,\n",
|
||||
" use_ray=True,\n",
|
||||
" placement_group=pg,\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"# Optional: specify exact bundle indices via environment variable\n",
|
||||
"# export SGLANG_RAY_BUNDLE_INDICES=\"0,1,2,3\"\n",
|
||||
"```\n",
|
||||
"\n",
|
||||
"### Bundle Index Control\n",
|
||||
"\n",
|
||||
"Use `SGLANG_RAY_BUNDLE_INDICES` environment variable to specify which placement group bundles to use for each worker rank. This enables:\n",
|
||||
"- Skipping unhealthy GPUs\n",
|
||||
"- Topology-aware placement (e.g., NVLink-connected GPUs)\n",
|
||||
"- Non-sequential bundle assignment\n",
|
||||
"\n",
|
||||
"```bash\n",
|
||||
"# Use bundles 0,1,2,7 (skip bundles 3-6) for tp_size=4\n",
|
||||
"export SGLANG_RAY_BUNDLE_INDICES=\"0,1,2,7\"\n",
|
||||
"\n",
|
||||
"# Place workers on NVLink-connected GPUs\n",
|
||||
"export SGLANG_RAY_BUNDLE_INDICES=\"0,1,2,3\"\n",
|
||||
"```\n",
|
||||
"\n",
|
||||
"The number of indices must match `world_size` (`tp_size * pp_size * dp_size`, or `tp_size * pp_size` when `enable_dp_attention=True`)."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
|
||||
@@ -629,6 +629,11 @@ SGLang supports various environment variables that can be used to configure its
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Set one visible device per process for distributed computing</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>`false`</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>SGLANG_RAY_BUNDLE_INDICES</code></td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Comma-separated bundle indices for Ray actor placement (e.g., <code>"0,1,2,3"</code>). Must match world_size. Enables fine-grained GPU assignment in custom placement groups.</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>Not set</td>
|
||||
</tr>
|
||||
</tbody>
|
||||
</table>
|
||||
|
||||
|
||||
@@ -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,27 +128,29 @@ 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)
|
||||
|
||||
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
|
||||
|
||||
local_gpu_idx = (pp_rank % pp_per_node) * tp_per_node + (
|
||||
tp_rank % tp_per_node
|
||||
)
|
||||
|
||||
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,
|
||||
@@ -168,44 +171,84 @@ class RayDataParallelController(DataParallelController):
|
||||
rank_port_args.detokenizer_ipc_name = (
|
||||
port_args.detokenizer_ipc_name
|
||||
)
|
||||
rank_port_args.tokenizer_ipc_name = port_args.tokenizer_ipc_name
|
||||
|
||||
local_gpu_idx = (pp_rank % pp_per_node) * tp_per_node + (
|
||||
tp_rank % tp_per_node
|
||||
rank_port_args.tokenizer_ipc_name = (
|
||||
port_args.tokenizer_ipc_name
|
||||
)
|
||||
|
||||
attn_cp_rank, moe_dp_rank, moe_ep_rank = _compute_parallelism_ranks(
|
||||
server_args, tp_rank
|
||||
dist_init_addr = (
|
||||
f"{self.rank0_node_ip}:{rank_port_args.nccl_port}"
|
||||
)
|
||||
|
||||
# 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}"
|
||||
|
||||
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(
|
||||
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,
|
||||
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,
|
||||
rank0_node_ip=self.rank0_node_ip,
|
||||
)
|
||||
self.scheduler_actors.append(actor)
|
||||
|
||||
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,
|
||||
)
|
||||
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
|
||||
|
||||
dist_init_addr = f"{local_rank0_node_ip}:{rank_port_args.nccl_port}"
|
||||
|
||||
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)
|
||||
|
||||
|
||||
+230
-32
@@ -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,67 +288,100 @@ 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
|
||||
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 = []
|
||||
|
||||
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 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(
|
||||
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
|
||||
)
|
||||
|
||||
attn_cp_rank, moe_dp_rank, moe_ep_rank = (
|
||||
_compute_parallelism_ranks(server_args, tp_rank)
|
||||
)
|
||||
|
||||
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(
|
||||
actor = _create_scheduler_actor(
|
||||
pg=pg,
|
||||
bundle_idx=bundle_idx,
|
||||
gpu_id=local_gpu_idx,
|
||||
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,
|
||||
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)
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -3,6 +3,7 @@
|
||||
Tests the Ray actor scheduler backend:
|
||||
- Offline inference via Engine(use_ray=True) inside a Ray actor on a placement group
|
||||
- Data parallel (DP) and DP attention support
|
||||
- Custom placement_group and SGLANG_RAY_BUNDLE_INDICES for fine-grained bundle control
|
||||
- Error paths in RayEngine._launch_scheduler_processes()
|
||||
- HTTP server launched via --use-ray flag
|
||||
|
||||
@@ -11,12 +12,14 @@ Usage:
|
||||
python -m pytest test/manual/test_ray_engine.py::TestRayEngineOfflineTP1 -v -s
|
||||
python -m pytest test/manual/test_ray_engine.py::TestRayEngineErrors -v -s
|
||||
python -m pytest test/manual/test_ray_engine.py::TestRayHTTPServerTP1 -v -s
|
||||
python -m pytest test/manual/test_ray_engine.py::TestRayEnginePlacementGroupErrors -v -s
|
||||
|
||||
# 2-GPU tests
|
||||
python -m pytest test/manual/test_ray_engine.py::TestRayEngineOfflineTP2 -v -s
|
||||
python -m pytest test/manual/test_ray_engine.py::TestRayEngineOfflinePP2 -v -s
|
||||
python -m pytest test/manual/test_ray_engine.py::TestRayEngineOfflineDP2 -v -s
|
||||
python -m pytest test/manual/test_ray_engine.py::TestRayEngineOfflineDPAttention -v -s
|
||||
python -m pytest test/manual/test_ray_engine.py::TestRayEnginePlacementGroup -v -s
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -516,5 +519,217 @@ class TestRayHTTPServerTP1(unittest.TestCase):
|
||||
self.assertGreater(len(data["text"]), 0, f"Empty output for: {prompt}")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests: Custom placement_group and SGLANG_RAY_BUNDLE_INDICES
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@unittest.skipUnless(_has_ray, "ray is not installed")
|
||||
@unittest.skipUnless(_NUM_GPUS >= 2, "requires at least 2 GPUs")
|
||||
class TestRayEnginePlacementGroup(unittest.TestCase):
|
||||
"""Test RayEngine with custom placement_group and SGLANG_RAY_BUNDLE_INDICES."""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
if not ray.is_initialized():
|
||||
ray.init(log_to_driver=True, runtime_env=_RAY_RUNTIME_ENV)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
ray.shutdown()
|
||||
|
||||
def test_custom_pg_dp1_tp2(self):
|
||||
"""Test custom placement_group with dp_size=1, tp_size=2."""
|
||||
from sglang.srt.ray.engine import RayEngine
|
||||
|
||||
pg = placement_group([{"GPU": 1}] * 2, strategy="STRICT_PACK")
|
||||
ray.get(pg.ready())
|
||||
|
||||
engine = RayEngine(
|
||||
model_path=_MODEL,
|
||||
tp_size=2,
|
||||
placement_group=pg,
|
||||
use_ray=True,
|
||||
)
|
||||
|
||||
result = engine.generate("The capital of France is", _SAMPLING_PARAMS)
|
||||
self.assertIn("text", result)
|
||||
self.assertGreater(len(result["text"]), 0)
|
||||
print(f"Generated (dp=1, tp=2, custom PG): {result['text'][:200]}")
|
||||
|
||||
engine.shutdown()
|
||||
ray.util.remove_placement_group(pg)
|
||||
|
||||
def test_bundle_indices_dp1_tp2(self):
|
||||
"""Test SGLANG_RAY_BUNDLE_INDICES with dp_size=1, tp_size=2."""
|
||||
from sglang.srt.ray.engine import RayEngine
|
||||
|
||||
os.environ["SGLANG_RAY_BUNDLE_INDICES"] = "0,1"
|
||||
|
||||
pg = placement_group([{"GPU": 1}] * 2, strategy="STRICT_PACK")
|
||||
ray.get(pg.ready())
|
||||
|
||||
engine = RayEngine(
|
||||
model_path=_MODEL,
|
||||
tp_size=2,
|
||||
placement_group=pg,
|
||||
use_ray=True,
|
||||
)
|
||||
|
||||
result = engine.generate("The capital of France is", _SAMPLING_PARAMS)
|
||||
self.assertIn("text", result)
|
||||
self.assertGreater(len(result["text"]), 0)
|
||||
print(f"Generated (dp=1, tp=2, indices=0,1): {result['text'][:200]}")
|
||||
|
||||
engine.shutdown()
|
||||
ray.util.remove_placement_group(pg)
|
||||
del os.environ["SGLANG_RAY_BUNDLE_INDICES"]
|
||||
|
||||
def test_custom_pg_dp2_tp1(self):
|
||||
"""Test custom placement_group with dp_size=2, tp_size=1."""
|
||||
from sglang.srt.ray.engine import RayEngine
|
||||
|
||||
pg = placement_group([{"GPU": 1}] * 2, strategy="STRICT_PACK")
|
||||
ray.get(pg.ready())
|
||||
|
||||
engine = RayEngine(
|
||||
model_path=_MODEL,
|
||||
tp_size=1,
|
||||
dp_size=2,
|
||||
placement_group=pg,
|
||||
use_ray=True,
|
||||
)
|
||||
|
||||
result = engine.generate("The capital of France is", _SAMPLING_PARAMS)
|
||||
self.assertIn("text", result)
|
||||
self.assertGreater(len(result["text"]), 0)
|
||||
print(f"Generated (dp=2, tp=1, custom PG): {result['text'][:200]}")
|
||||
|
||||
engine.shutdown()
|
||||
ray.util.remove_placement_group(pg)
|
||||
|
||||
def test_bundle_indices_skip_bundle(self):
|
||||
"""Test skipping unhealthy GPU by using bundle_indices."""
|
||||
from sglang.srt.ray.engine import RayEngine
|
||||
|
||||
os.environ["SGLANG_RAY_BUNDLE_INDICES"] = "1" # Skip bundle 0
|
||||
|
||||
pg = placement_group([{"GPU": 1}] * 2, strategy="STRICT_PACK")
|
||||
ray.get(pg.ready())
|
||||
|
||||
engine = RayEngine(
|
||||
model_path=_MODEL,
|
||||
tp_size=1,
|
||||
placement_group=pg,
|
||||
use_ray=True,
|
||||
)
|
||||
|
||||
result = engine.generate("The capital of France is", _SAMPLING_PARAMS)
|
||||
self.assertIn("text", result)
|
||||
self.assertGreater(len(result["text"]), 0)
|
||||
print(f"Generated (tp=1, skip bundle 0): {result['text'][:200]}")
|
||||
|
||||
engine.shutdown()
|
||||
ray.util.remove_placement_group(pg)
|
||||
del os.environ["SGLANG_RAY_BUNDLE_INDICES"]
|
||||
|
||||
def test_custom_pg_dp_attention(self):
|
||||
"""Test custom placement_group with enable_dp_attention=True."""
|
||||
from sglang.srt.ray.engine import RayEngine
|
||||
|
||||
pg = placement_group([{"GPU": 1}] * 2, strategy="STRICT_PACK")
|
||||
ray.get(pg.ready())
|
||||
|
||||
engine = RayEngine(
|
||||
model_path=_DP_ATTN_MODEL,
|
||||
tp_size=2,
|
||||
placement_group=pg,
|
||||
use_ray=True,
|
||||
enable_dp_attention=True,
|
||||
)
|
||||
|
||||
result = engine.generate("The capital of France is", _SAMPLING_PARAMS)
|
||||
self.assertIn("text", result)
|
||||
self.assertGreater(len(result["text"]), 0)
|
||||
print(f"Generated (dp attention, tp=2, custom PG): {result['text'][:200]}")
|
||||
|
||||
engine.shutdown()
|
||||
ray.util.remove_placement_group(pg)
|
||||
|
||||
|
||||
@unittest.skipUnless(_has_ray, "ray is not installed")
|
||||
@unittest.skipUnless(_NUM_GPUS >= 1, "requires at least 1 GPU")
|
||||
class TestRayEnginePlacementGroupErrors(unittest.TestCase):
|
||||
"""Test error handling for placement_group and bundle indices."""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
if not ray.is_initialized():
|
||||
ray.init(log_to_driver=True, runtime_env=_RAY_RUNTIME_ENV)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
ray.shutdown()
|
||||
|
||||
def test_multi_gpu_bundle_raises_error(self):
|
||||
"""Custom PG with multi-GPU bundles should raise an error."""
|
||||
|
||||
@ray.remote(num_gpus=0)
|
||||
def _try_multi_gpu_bundle():
|
||||
pg = placement_group([{"GPU": 2}], strategy="STRICT_PACK")
|
||||
ray.get(pg.ready())
|
||||
|
||||
from sglang.srt.ray.engine import RayEngine
|
||||
|
||||
try:
|
||||
RayEngine(
|
||||
model_path=_MODEL,
|
||||
tp_size=2,
|
||||
placement_group=pg,
|
||||
use_ray=True,
|
||||
)
|
||||
return None
|
||||
except Exception as e:
|
||||
return str(e)
|
||||
finally:
|
||||
ray.util.remove_placement_group(pg)
|
||||
|
||||
error_msg = ray.get(_try_multi_gpu_bundle.remote(), timeout=120)
|
||||
self.assertIsNotNone(error_msg)
|
||||
self.assertIn("exactly 1 GPU per bundle", error_msg)
|
||||
|
||||
def test_invalid_bundle_index_raises_error(self):
|
||||
"""SGLANG_RAY_BUNDLE_INDICES with invalid index should raise an error."""
|
||||
|
||||
@ray.remote(num_gpus=0)
|
||||
def _try_invalid_bundle_index():
|
||||
import os
|
||||
|
||||
os.environ["SGLANG_RAY_BUNDLE_INDICES"] = "0,10"
|
||||
|
||||
pg = placement_group([{"GPU": 1}] * 2, strategy="STRICT_PACK")
|
||||
ray.get(pg.ready())
|
||||
|
||||
from sglang.srt.ray.engine import RayEngine
|
||||
|
||||
try:
|
||||
RayEngine(
|
||||
model_path=_MODEL,
|
||||
tp_size=2,
|
||||
placement_group=pg,
|
||||
use_ray=True,
|
||||
)
|
||||
return None
|
||||
except Exception as e:
|
||||
return str(e)
|
||||
finally:
|
||||
os.environ.pop("SGLANG_RAY_BUNDLE_INDICES", None)
|
||||
ray.util.remove_placement_group(pg)
|
||||
|
||||
error_msg = ray.get(_try_invalid_bundle_index.remote(), timeout=120)
|
||||
self.assertIsNotNone(error_msg)
|
||||
self.assertIn("out of range", error_msg)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
Reference in New Issue
Block a user