diff --git a/docs_new/docs/basic_usage/offline_engine_api.ipynb b/docs_new/docs/basic_usage/offline_engine_api.ipynb index fe8a9e304..963bd3f73 100644 --- a/docs_new/docs/basic_usage/offline_engine_api.ipynb +++ b/docs_new/docs/basic_usage/offline_engine_api.ipynb @@ -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": {}, diff --git a/docs_new/docs/references/environment_variables.mdx b/docs_new/docs/references/environment_variables.mdx index aa23ecfbe..e5e6d7189 100644 --- a/docs_new/docs/references/environment_variables.mdx +++ b/docs_new/docs/references/environment_variables.mdx @@ -629,6 +629,11 @@ SGLang supports various environment variables that can be used to configure its Set one visible device per process for distributed computing `false` + + SGLANG_RAY_BUNDLE_INDICES + Comma-separated bundle indices for Ray actor placement (e.g., "0,1,2,3"). Must match world_size. Enables fine-grained GPU assignment in custom placement groups. + Not set + diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 33adc21ee..5e0b95f1d 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -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) diff --git a/python/sglang/srt/ray/data_parallel_controller.py b/python/sglang/srt/ray/data_parallel_controller.py index e658d5a06..eebef6b7b 100644 --- a/python/sglang/srt/ray/data_parallel_controller.py +++ b/python/sglang/srt/ray/data_parallel_controller.py @@ -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:] diff --git a/python/sglang/srt/ray/engine.py b/python/sglang/srt/ray/engine.py index 5cbbb335a..b76895ed0 100644 --- a/python/sglang/srt/ray/engine.py +++ b/python/sglang/srt/ray/engine.py @@ -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. diff --git a/python/sglang/srt/ray/http_server.py b/python/sglang/srt/ray/http_server.py index c2acda83e..d57f845d9 100644 --- a/python/sglang/srt/ray/http_server.py +++ b/python/sglang/srt/ray/http_server.py @@ -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, diff --git a/test/manual/test_ray_engine.py b/test/manual/test_ray_engine.py index f91e09fcf..0eeb996ba 100644 --- a/test/manual/test_ray_engine.py +++ b/test/manual/test_ray_engine.py @@ -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()