sgl-project--sglang/python/sglang/srt/ray/data_parallel_controller.py

# Copyright 2023-2024 SGLang Team
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Ray-aware DataParallelController that launches SchedulerActors instead of mp.Process."""

from __future__ import annotations

import logging
from typing import List, Optional

import ray
import zmq

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.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

logger = logging.getLogger(__name__)


class RayDataParallelController(DataParallelController):
    """DataParallelController that uses Ray actors for scheduler processes.

    Overrides the process-spawning methods to create SchedulerActor Ray actors
    instead of mp.Process. Runs in-process (not as a separate mp.Process) and
    reuses the parent's event_loop, dispatching, and ZMQ routing.
    """

    def __init__(
        self,
        server_args: ServerArgs,
        port_args: PortArgs,
        placement_group,
        bundle_for_node: Optional[List[int]],
        rank0_node_ip: str,
    ):
        # Set Ray-specific attributes BEFORE super().__init__() because the
        # parent constructor calls launch_dp_schedulers / launch_dp_attention_schedulers
        # which we override, and those methods need these attributes.
        self.pg = placement_group
        self.bundle_for_node = bundle_for_node
        self.rank0_node_ip = rank0_node_ip
        self.scheduler_actors: List = []
        self.event_loop_refs: List = []

        # super().__init__ will call our overridden launch methods via MRO.
        # Pass run_scheduler_process_func=None since we don't spawn mp.Process.
        super().__init__(server_args, port_args, run_scheduler_process_func=None)

    def launch_dp_schedulers(self, server_args: ServerArgs, port_args: PortArgs):
        """Override: launch Ray scheduler actors per DP rank."""
        sockets = []
        dp_port_args_list = []

        for dp_rank in range(server_args.dp_size):
            tmp_port_args = PortArgs.init_new(server_args)
            tmp_port_args.tokenizer_ipc_name = port_args.tokenizer_ipc_name
            tmp_port_args.detokenizer_ipc_name = port_args.detokenizer_ipc_name
            tmp_port_args.instance_id = port_args.instance_id

            # Hold NCCL port so the next DP rank gets a different one
            sockets.append(bind_port(tmp_port_args.nccl_port))
            dp_port_args_list.append(tmp_port_args)

            # Create ZMQ PUSH socket for this DP rank (controller → scheduler)
            if server_args.node_rank == 0:
                self.workers[dp_rank] = get_zmq_socket(
                    self.context,
                    zmq.PUSH,
                    tmp_port_args.scheduler_input_ipc_name,
                    True,
                )

        # Release held ports before creating actors
        for sock in sockets:
            sock.close()

        # Create actors for each DP rank sequentially
        for dp_rank in range(server_args.dp_size):
            self._launch_ray_tp_group(server_args, dp_port_args_list[dp_rank], dp_rank)

    def launch_dp_attention_schedulers(
        self, server_args: ServerArgs, port_args: PortArgs
    ):
        """Override: pre-allocate ports, skip broadcast, create Ray actors."""
        # Pre-allocate worker ports on the controller node, binding to the
        # rank-0 node IP instead of tcp://* to avoid exposing unauthenticated
        # ZMQ sockets (CVE-2026-3060).
        worker_ports = []
        for dp_rank in range(server_args.dp_size):
            worker_port, worker_socket = get_zmq_socket_on_host(
                self.context, zmq.PUSH, host=self.rank0_node_ip
            )
            worker_ports.append(worker_port)
            self.workers[dp_rank] = worker_socket
            logger.debug(f"Assigned port {worker_port} to worker {dp_rank}")

        # Skip _broadcast_worker_ports — Ray creates all actors centrally,
        # so there's no need for the inter-node handshake protocol.
        self._launch_ray_tp_group(
            server_args, port_args, dp_rank=None, worker_ports=worker_ports
        )

    def _launch_ray_tp_group(
        self,
        server_args: ServerArgs,
        port_args: PortArgs,
        dp_rank: Optional[int],
        worker_ports: Optional[List[int]] = None,
    ):
        """Create SchedulerActor Ray actors for one TP group (one DP rank).

        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:
                            _, _, 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}"
                        )

                        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,
                        )
                        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)

        # Wait for all actors created in this call to initialize
        batch_actors = self.scheduler_actors[batch_start_idx:]
        try:
            scheduler_infos = ray.get(
                [actor.get_info.remote() for actor in batch_actors]
            )
        except ray.exceptions.RayActorError as e:
            for actor in self.scheduler_actors:
                try:
                    ray.kill(actor)
                except Exception:
                    logger.error(f"Failed to kill Ray scheduler actor: {actor}")
            raise RuntimeError(f"Scheduler actor failed to initialize: {e}")

        # Store init info from the first actor (same across all actors)
        if scheduler_infos:
            self.max_total_num_tokens = scheduler_infos[0]["max_total_num_tokens"]
            self.max_req_input_len = scheduler_infos[0]["max_req_input_len"]

        # Start event loops (non-blocking — runs until actor is killed)
        self.event_loop_refs.extend(
            [actor.run_event_loop.remote() for actor in batch_actors]
        )

    # Override launch_tensor_parallel_group to be a no-op since we don't use it.
    # The parent's launch_dp_schedulers/launch_dp_attention_schedulers call this,
    # but our overrides call _launch_ray_tp_group instead.
    def launch_tensor_parallel_group(
        self,
        server_args: ServerArgs,
        port_args: PortArgs,
        base_gpu_id: int,
        dp_rank: Optional[int],
        worker_ports: Optional[List[int]] = None,
    ):
        raise RuntimeError(
            "RayDataParallelController should not call launch_tensor_parallel_group. "
            "Use _launch_ray_tp_group instead."
        )