sgl-project--sglang/python/sglang/srt/ray/engine.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.
# ==============================================================================
"""RayEngine - Engine subclass that launches schedulers as Ray actors."""

from __future__ import annotations

import dataclasses
import logging
import threading
from typing import Callable, List, Optional

import ray
from ray.util.placement_group import PlacementGroup
from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy

from sglang.srt.entrypoints.engine import (
    Engine,
    SchedulerInitResult,
    _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

logger = logging.getLogger(__name__)


@dataclasses.dataclass
class RaySchedulerInitResult(SchedulerInitResult):
    """SchedulerInitResult that also holds Ray actor handles for cleanup."""

    scheduler_actors: list = dataclasses.field(default_factory=list)


def _find_engine_bundle(
    placement_group: PlacementGroup, nnodes: int
) -> tuple[int, str]:
    """Find which placement group bundle is on the same node as the Engine.
    Rank0 scheduler must be co-located with the Engine. Returns (bundle_index, engine_ip).
    """
    engine_ip = ray.util.get_node_ip_address()

    @ray.remote(num_cpus=0, num_gpus=0)
    def get_node_ip():
        return ray.util.get_node_ip_address()

    bundle_ips = ray.get(
        [
            get_node_ip.options(
                scheduling_strategy=PlacementGroupSchedulingStrategy(
                    placement_group=placement_group,
                    placement_group_bundle_index=i,
                ),
            ).remote()
            for i in range(nnodes)
        ]
    )

    try:
        return bundle_ips.index(engine_ip), engine_ip
    except ValueError:
        raise RuntimeError(
            f"Engine node {engine_ip} not found in any placement group bundle {bundle_ips}. "
            f"Rank-0 scheduler must be co-located with the Engine."
        )


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.override("ray.placement_group", 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:
            try:
                ray.kill(actor)
            except Exception:
                logger.error(f"Failed to kill Ray scheduler actor: {actor}")
        super().shutdown()

    @classmethod
    def _launch_scheduler_processes(
        cls,
        server_args: ServerArgs,
        port_args: PortArgs,
        run_scheduler_process_func: Callable,
    ) -> tuple[SchedulerInitResult, None]:
        """Launch schedulers as Ray actors.

        Returns:
            Tuple of (RaySchedulerInitResult, None).
            scheduler_procs is None since Ray uses actors instead of mp.Process.
        """
        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,
            )

            if server_args.enable_dp_attention:
                total_gpus = server_args.tp_size * server_args.pp_size
            else:
                total_gpus = (
                    server_args.dp_size * server_args.tp_size * server_args.pp_size
                )

            nnodes = server_args.nnodes
            gpus_per_node = total_gpus // nnodes
            strategy = "STRICT_PACK" if nnodes == 1 else "SPREAD"

            logger.info(
                "No placement group detected. Auto-creating one with "
                f"{nnodes} bundle(s), {gpus_per_node} GPU(s)/bundle, "
                "placement group explicitly and schedule the Engine onto it."
            )

            pg = create_placement_group(
                [{"CPU": 1, "GPU": gpus_per_node}] * nnodes,
                strategy=strategy,
            )
            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)

        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:
            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(
                            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 = _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(
                    [actor.get_info.remote() for actor in scheduler_actors]
                )
            except ray.exceptions.RayActorError as e:
                for actor in 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}")

            event_loop_refs = [
                actor.run_event_loop.remote() for actor in scheduler_actors
            ]

            def wait_for_completion():
                try:
                    ray.get(event_loop_refs)
                except Exception as e:
                    logger.error(f"Ray scheduler actor terminated with error: {e}")

            return (
                RaySchedulerInitResult(
                    scheduler_infos=scheduler_infos,
                    wait_for_completion=wait_for_completion,
                    scheduler_actors=scheduler_actors,
                ),
                None,
            )
        else:
            # Launch the data parallel controller
            return (
                cls._launch_dp_scheduler_processes(
                    server_args,
                    port_args,
                    pg,
                    bundle_for_node,
                    rank0_node_ip,
                ),
                None,
            )

    @classmethod
    def _launch_dp_scheduler_processes(
        cls,
        server_args: ServerArgs,
        port_args: PortArgs,
        pg,
        bundle_for_node: Optional[List[int]],
        rank0_node_ip: str,
    ) -> RaySchedulerInitResult:
        """Launch DP schedulers via RayDataParallelController."""
        from sglang.srt.ray.data_parallel_controller import (
            RayDataParallelController,
        )

        if server_args.enable_dp_attention:
            # DP attention folds DP into TP — total GPUs = tp_size * pp_size
            total_gpus = server_args.tp_size * server_args.pp_size
        else:
            total_gpus = server_args.dp_size * server_args.tp_size * server_args.pp_size
        gpus_per_node = total_gpus // server_args.nnodes
        logger.info(
            f"Ray DP cluster: {server_args.nnodes} nodes, "
            f"{gpus_per_node} GPUs/node, dp_size={server_args.dp_size}, "
            f"tp_size={server_args.tp_size}, pp_size={server_args.pp_size}, "
            f"enable_dp_attention={server_args.enable_dp_attention}"
        )

        # Set dist_init_addr on server_args so PortArgs.init_new() can compute
        # TCP addresses correctly (required for DP attention path).
        dp_server_args = dataclasses.replace(
            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.override(
            "ray.placement_group", 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.
        controller = RayDataParallelController(
            dp_server_args, port_args, pg, bundle_for_node, rank0_node_ip
        )

        # Start the DP controller's event loop in a daemon thread.
        # It routes requests from the tokenizer to per-DP-rank schedulers.
        dp_thread = threading.Thread(
            target=controller.event_loop, daemon=True, name="dp_controller"
        )
        dp_thread.start()

        scheduler_infos = [
            {
                "max_total_num_tokens": controller.max_total_num_tokens,
                "max_req_input_len": controller.max_req_input_len,
            }
        ]

        event_loop_refs = controller.event_loop_refs

        def wait_for_completion():
            try:
                ray.get(event_loop_refs)
            except Exception as e:
                logger.error(f"Ray scheduler actor terminated with error: {e}")

        return RaySchedulerInitResult(
            scheduler_infos=scheduler_infos,
            wait_for_completion=wait_for_completion,
            scheduler_actors=controller.scheduler_actors,
        )