chore: import upstream snapshot with attribution

python/sglang/srt/ray/scheduler_actor.py

@@ -0,0 +1,133 @@
+# 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 actor wrapper for SGLang Scheduler."""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Any, Dict, Optional
+
+import ray
+
+if TYPE_CHECKING:
+    from sglang.srt.server_args import PortArgs, ServerArgs
+
+
+logger = logging.getLogger(__name__)
+
+
+@ray.remote
+class SchedulerActor:
+    """Ray actor wrapper for SGLang Scheduler.
+
+    Each actor manages one GPU and runs the Scheduler + TpModelWorker stack.
+    Ray is used for process lifecycle; ZMQ handles request/response communication.
+    """
+
+    def __init__(
+        self,
+        server_args: ServerArgs,
+        port_args: PortArgs,
+        gpu_id: int,
+        tp_rank: int,
+        attn_cp_rank: int,
+        moe_dp_rank: int,
+        moe_ep_rank: int,
+        pp_rank: int,
+        dp_rank: Optional[int],
+        dist_init_addr: Optional[str] = None,
+    ):
+        import dataclasses
+
+        from sglang.srt.environ import envs
+        from sglang.srt.managers.scheduler import Scheduler, configure_scheduler_process
+        from sglang.srt.utils.numa_utils import (
+            get_numa_node_if_available,
+            numa_bind_to_node,
+        )
+
+        # Override dist_init_addr if provided (for multi-node)
+        if dist_init_addr:
+            server_args = dataclasses.replace(
+                server_args, dist_init_addr=dist_init_addr
+            )
+
+        # Get actual GPU IDs from Ray runtime context
+        accelerator_ids = ray.get_runtime_context().get_accelerator_ids()
+        assigned_gpus = accelerator_ids.get("GPU", [])
+
+        if assigned_gpus:
+            # Ray assigned specific GPU(s), use the first one
+            actual_gpu_id = int(assigned_gpus[0])
+            logger.info(f"[TP{tp_rank}] Ray assigned GPU: {actual_gpu_id}")
+        else:
+            # Fallback to passed gpu_id
+            actual_gpu_id = gpu_id
+            logger.info(f"[TP{tp_rank}] Using passed gpu_id: {gpu_id}")
+
+        # Configure worker (logging, process title, etc.)
+        dp_rank = configure_scheduler_process(
+            server_args,
+            actual_gpu_id,
+            tp_rank,
+            attn_cp_rank,
+            moe_dp_rank,
+            moe_ep_rank,
+            pp_rank,
+            dp_rank,
+        )
+
+        # Ray actors can't use the numactl subprocess-wrapping approach
+        # (SGLANG_NUMA_BIND_V2's normal path), so bind in-process via libnuma.
+        # The V1 path inside configure_scheduler_process already handles
+        # SGLANG_NUMA_BIND_V2=False.
+        if envs.SGLANG_NUMA_BIND_V2.get():
+            numa_node = get_numa_node_if_available(server_args, actual_gpu_id)
+            if numa_node is not None:
+                numa_bind_to_node(numa_node)
+                logger.info(
+                    f"[TP{tp_rank}] Bound to NUMA node {numa_node} for GPU {actual_gpu_id}"
+                )
+
+        # Create scheduler (loads model into GPU, initializes NCCL)
+        self.scheduler = Scheduler(
+            server_args=server_args,
+            port_args=port_args,
+            gpu_id=actual_gpu_id,
+            tp_rank=tp_rank,
+            moe_ep_rank=moe_ep_rank,
+            pp_rank=pp_rank,
+            attn_cp_rank=attn_cp_rank,
+            moe_dp_rank=moe_dp_rank,
+            dp_rank=dp_rank,
+        )
+
+        self._tp_rank = tp_rank
+        self._pp_rank = pp_rank
+
+    def get_info(self) -> Dict[str, Any]:
+        """Return scheduler initialization info for handshake."""
+        return self.scheduler.get_init_info()
+
+    def run_event_loop(self) -> None:
+        """Run the scheduler's event loop. Blocks until shutdown."""
+        try:
+            import torch
+
+            # Need to set the GPU id for the event loop for nccl to work
+            torch.cuda.set_device(self.scheduler.ps.gpu_id)
+            self.scheduler.run_event_loop()
+        except Exception as e:
+            logger.error(f"Scheduler PP{self._pp_rank} TP{self._tp_rank} crashed: {e}")
+            raise