chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Add Ray actor support for scheduler process management (DP=1) (#17684)

Browse Source 
Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
Xinyu Zhang
2026-03-05 13:21:23 -08:00
committed by GitHub
parent 07e7603c0c
commit b3cfad0a80
12 changed files with 1303 additions and 247 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
python/pyproject.toml
View File

@@ -121,6 +121,10 @@ diffusion = [
"xatlas",
]
ray = [
"ray[default]>=2.54.0",
]
tracing = [
"opentelemetry-api",
"opentelemetry-exporter-otlp",

73
python/sglang/bench_offline_throughput.py
View File

@@ -323,12 +323,83 @@ def monitor_trace_file(known_files, directory, interval=1):
break
def _create_ray_engine_backend(server_args: ServerArgs):
"""Create a RayEngine inside a Ray actor on a placement group.
RayEngine requires a placement group, so we launch it inside a Ray actor
and return a lightweight proxy that forwards calls via ray.get().
"""
import ray
from ray.runtime_env import RuntimeEnv
from ray.util.placement_group import placement_group
from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
env_vars = {"RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES": "1"}
if os.environ.get("HF_TOKEN"):
env_vars["HF_TOKEN"] = os.environ["HF_TOKEN"]
if not ray.is_initialized():
ray.init(runtime_env=RuntimeEnv(env_vars=env_vars))
total_gpus = server_args.tp_size * server_args.pp_size
pg = placement_group([{"CPU": 1, "GPU": total_gpus}], strategy="STRICT_PACK")
ray.get(pg.ready())
@ray.remote
class _EngineActor:
def __init__(self, **kwargs):
from sglang.srt.ray.engine import RayEngine
self.engine = RayEngine(**kwargs)
def call(self, method, **kwargs):
return getattr(self.engine, method)(**kwargs)
actor = _EngineActor.options(
num_cpus=1,
num_gpus=0,
scheduling_strategy=PlacementGroupSchedulingStrategy(
placement_group=pg,
placement_group_bundle_index=0,
),
).remote(**dataclasses.asdict(server_args))
class _Proxy:
"""Forwards method calls to the remote RayEngine actor."""
def generate(self, **kwargs):
return ray.get(actor.call.remote("generate", **kwargs))
def get_server_info(self, **kwargs):
return ray.get(actor.call.remote("get_server_info", **kwargs))
def start_profile(self, **kwargs):
return ray.get(actor.call.remote("start_profile", **kwargs))
def stop_profile(self, **kwargs):
return ray.get(actor.call.remote("stop_profile", **kwargs))
def shutdown(self):
try:
ray.get(actor.call.remote("shutdown"), timeout=60)
except Exception:
pass
try:
ray.util.remove_placement_group(pg)
except Exception:
pass
return _Proxy()
def throughput_test(
server_args: ServerArgs,
bench_args: BenchArgs,
):
if bench_args.backend == "engine":
backend = Engine(**dataclasses.asdict(server_args))
if server_args.use_ray:
backend = _create_ray_engine_backend(server_args)
else:
backend = Engine(**dataclasses.asdict(server_args))
if not backend:
raise ValueError("Please provide valid engine arguments")
elif bench_args.backend == "runtime":

10
python/sglang/launch_server.py
View File

@@ -28,6 +28,16 @@ def run_server(server_args):
from sglang.srt.entrypoints.grpc_server import serve_grpc
asyncio.run(serve_grpc(server_args))
elif server_args.use_ray:
try:
from sglang.srt.ray.http_server import launch_server
except ImportError:
raise ImportError(
"Ray is required for --use-ray mode. "
"Install it with: pip install 'sglang[ray]'"
)
launch_server(server_args)
else:
# Default mode: HTTP mode.
from sglang.srt.entrypoints.http_server import launch_server

457
python/sglang/srt/entrypoints/engine.py
View File

@@ -27,7 +27,17 @@ import random
import signal
import threading
import time
from typing import AsyncIterator, Callable, Dict, Iterator, List, Optional, Tuple, Union
from typing import (
Any,
AsyncIterator,
Callable,
Dict,
Iterator,
List,
Optional,
Tuple,
Union,
)
# Fix a bug of Python threading
setattr(threading, "_register_atexit", lambda *args, **kwargs: None)
@@ -96,6 +106,15 @@ asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
_is_cuda = is_cuda()
@dataclasses.dataclass
class SchedulerInitResult:
"""Result from launching schedulers."""
scheduler_infos: List[Dict[str, Any]]
wait_for_ready: Callable[[], None] = lambda: None
wait_for_completion: Callable[[], None] = lambda: None
def init_tokenizer_manager(
server_args: ServerArgs,
port_args: PortArgs,
@@ -161,21 +180,24 @@ class Engine(EngineBase):
atexit.register(self.shutdown)
# Launch subprocesses
tokenizer_manager, template_manager, scheduler_infos, port_args = (
_launch_subprocesses(
server_args=server_args,
init_tokenizer_manager_func=self.init_tokenizer_manager_func,
run_scheduler_process_func=self.run_scheduler_process_func,
run_detokenizer_process_func=self.run_detokenizer_process_func,
)
(
tokenizer_manager,
template_manager,
port_args,
scheduler_init_result,
) = self._launch_subprocesses(
server_args=server_args,
init_tokenizer_manager_func=self.init_tokenizer_manager_func,
run_scheduler_process_func=self.run_scheduler_process_func,
run_detokenizer_process_func=self.run_detokenizer_process_func,
)
self.tokenizer_manager = tokenizer_manager
self.template_manager = template_manager
self.scheduler_info = scheduler_infos[0]
self._scheduler_init_result = scheduler_init_result
self.port_args = port_args
self.remote_instance_transfer_engine_info = (
parse_remote_instance_transfer_engine_info_from_scheduler_infos(
scheduler_infos
scheduler_init_result.scheduler_infos
)
)
@@ -477,6 +499,201 @@ class Engine(EngineBase):
ret = self.loop.run_until_complete(generator.__anext__())
return ret
@classmethod
def _launch_scheduler_processes(
cls,
server_args: ServerArgs,
port_args: PortArgs,
run_scheduler_process_func: Callable,
) -> SchedulerInitResult:
"""Launch scheduler processes using multiprocessing.
Override in subclasses for different backends (e.g. Ray).
"""
scheduler_procs = []
if server_args.dp_size == 1:
# Launch tensor parallel scheduler processes
memory_saver_adapter = TorchMemorySaverAdapter.create(
enable=server_args.enable_memory_saver
)
scheduler_pipe_readers = []
pp_rank_range, tp_rank_range, pp_size_per_node, tp_size_per_node = (
_calculate_rank_ranges(
server_args.nnodes,
server_args.pp_size,
server_args.tp_size,
server_args.node_rank,
)
)
for pp_rank in pp_rank_range:
for tp_rank in tp_rank_range:
reader, writer = mp.Pipe(duplex=False)
gpu_id = (
server_args.base_gpu_id
+ ((pp_rank % pp_size_per_node) * tp_size_per_node)
+ (tp_rank % tp_size_per_node) * server_args.gpu_id_step
)
attn_cp_rank, moe_dp_rank, moe_ep_rank = _compute_parallelism_ranks(
server_args, tp_rank
)
with maybe_reindex_device_id(gpu_id) as gpu_id:
proc = mp.Process(
target=run_scheduler_process_func,
args=(
server_args,
port_args,
gpu_id,
tp_rank,
attn_cp_rank,
moe_dp_rank,
moe_ep_rank,
pp_rank,
None,
writer,
),
)
with memory_saver_adapter.configure_subprocess(), numa_utils.configure_subprocess(
server_args, gpu_id
):
proc.start()
scheduler_procs.append(proc)
scheduler_pipe_readers.append(reader)
else:
# Launch the data parallel controller
reader, writer = mp.Pipe(duplex=False)
scheduler_pipe_readers = [reader]
proc = mp.Process(
target=run_data_parallel_controller_process,
kwargs=dict(
server_args=server_args,
port_args=port_args,
pipe_writer=writer,
run_scheduler_process_func=run_scheduler_process_func,
),
)
proc.start()
scheduler_procs.append(proc)
scheduler_infos = []
def wait_for_ready():
infos = _wait_for_scheduler_ready(scheduler_pipe_readers, scheduler_procs)
scheduler_infos.extend(infos)
def wait_for_completion():
for proc in scheduler_procs:
proc.join()
logger.error(
f"Scheduler or DataParallelController {proc.pid} "
f"terminated with {proc.exitcode}"
)
return SchedulerInitResult(
scheduler_infos=scheduler_infos,
wait_for_ready=wait_for_ready,
wait_for_completion=wait_for_completion,
)
@classmethod
def _launch_subprocesses(
cls,
server_args: ServerArgs,
init_tokenizer_manager_func: Callable,
run_scheduler_process_func: Callable,
run_detokenizer_process_func: Callable,
port_args: Optional[PortArgs] = None,
) -> Tuple[TokenizerManager, TemplateManager, PortArgs, SchedulerInitResult]:
"""Launch the TokenizerManager in the main process, the Scheduler in a subprocess, and the DetokenizerManager in another subprocess.
Returns:
Tuple of (tokenizer_manager, template_manager, port_args, scheduler_init_result).
"""
# Configure global environment
configure_logger(server_args)
_set_envs_and_config(server_args)
server_args.check_server_args()
# Allocate ports for inter-process communications
if port_args is None:
port_args = PortArgs.init_new(server_args)
logger.info(f"{server_args=}")
# Launch scheduler processes
scheduler_init_result = cls._launch_scheduler_processes(
server_args, port_args, run_scheduler_process_func
)
if (
server_args.enable_elastic_expert_backup
and server_args.elastic_ep_backend is not None
):
run_expert_backup_manager(server_args, port_args)
if server_args.node_rank >= 1:
# In multi-node cases, non-zero rank nodes do not need to run tokenizer or detokenizer,
# so they can just wait here.
scheduler_init_result.wait_for_ready()
if os.getenv("SGLANG_BLOCK_NONZERO_RANK_CHILDREN") == "0":
# When using `Engine` as a Python API, we don't want to block here.
return (
None,
None,
port_args,
scheduler_init_result,
)
launch_dummy_health_check_server(
server_args.host, server_args.port, server_args.enable_metrics
)
scheduler_init_result.wait_for_completion()
return (
None,
None,
port_args,
scheduler_init_result,
)
# Launch detokenizer process
detoken_proc = mp.Process(
target=run_detokenizer_process_func,
args=(
server_args,
port_args,
),
)
detoken_proc.start()
# Init tokenizer manager first, as the bootstrap server is initialized here
if server_args.tokenizer_worker_num == 1:
tokenizer_manager, template_manager = init_tokenizer_manager_func(
server_args, port_args
)
else:
# Launch multi-tokenizer router
tokenizer_manager = MultiTokenizerRouter(server_args, port_args)
template_manager = None
# Wait for the model to finish loading
scheduler_init_result.wait_for_ready()
# Get back some info from scheduler to tokenizer_manager
tokenizer_manager.max_req_input_len = scheduler_init_result.scheduler_infos[0][
"max_req_input_len"
]
return (
tokenizer_manager,
template_manager,
port_args,
scheduler_init_result,
)
def shutdown(self):
"""Shutdown the engine"""
kill_process_tree(os.getpid(), include_parent=False)
@@ -557,7 +774,7 @@ class Engine(EngineBase):
)
return {
**dataclasses.asdict(self.tokenizer_manager.server_args),
**self.scheduler_info,
**self._scheduler_init_result.scheduler_infos[0],
"internal_states": internal_states,
"version": __version__,
}
@@ -959,188 +1176,56 @@ def _wait_for_scheduler_ready(
return scheduler_infos
def _launch_scheduler_processes(
server_args: ServerArgs,
port_args: PortArgs,
run_scheduler_process_func: Callable,
):
scheduler_procs = []
def _calculate_rank_ranges(
nnodes: int, pp_size: int, tp_size: int, node_rank: int
) -> Tuple[range, range, int, int]:
"""Calculate pp_rank_range and tp_rank_range for a given node.
if server_args.dp_size == 1:
# Launch tensor parallel scheduler processes
memory_saver_adapter = TorchMemorySaverAdapter.create(
enable=server_args.enable_memory_saver
)
scheduler_pipe_readers = []
Args:
nnodes: Total number of nodes.
pp_size: Pipeline parallel size.
tp_size: Tensor parallel size.
node_rank: The rank of the node to compute ranges for.
pp_size_per_node = max(server_args.pp_size // server_args.nnodes, 1)
nnodes_per_pp_rank = max(server_args.nnodes // server_args.pp_size, 1)
pp_rank_range = range(
pp_size_per_node * (server_args.node_rank // nnodes_per_pp_rank),
pp_size_per_node * (server_args.node_rank // nnodes_per_pp_rank + 1),
)
nnodes_per_tp_group = nnodes_per_pp_rank
tp_size_per_node = server_args.tp_size // nnodes_per_tp_group
tp_rank_range = range(
tp_size_per_node * (server_args.node_rank % nnodes_per_tp_group),
tp_size_per_node * (server_args.node_rank % nnodes_per_tp_group + 1),
)
for pp_rank in pp_rank_range:
for tp_rank in tp_rank_range:
reader, writer = mp.Pipe(duplex=False)
gpu_id = (
server_args.base_gpu_id
+ ((pp_rank % pp_size_per_node) * tp_size_per_node)
+ (tp_rank % tp_size_per_node) * server_args.gpu_id_step
)
attn_dp_size = (
server_args.dp_size if server_args.enable_dp_attention else 1
)
# Parallelism hierarchy (outermost to innermost):
# - Attention: Global(TP) -> DP -> ATTN_CP -> ATTN_TP (innermost)
# - MoE: Global(TP) -> MOE_DP -> EP -> MOE_TP (innermost)
attn_tp_size = (
server_args.tp_size // attn_dp_size // server_args.attn_cp_size
)
attn_cp_rank = (tp_rank // attn_tp_size) % server_args.attn_cp_size
moe_dp_rank = tp_rank // (
server_args.tp_size // server_args.moe_dp_size
)
moe_ep_rank = (
tp_rank
% (server_args.tp_size // server_args.moe_dp_size)
// (
server_args.tp_size
// server_args.moe_dp_size
// server_args.ep_size
)
)
with maybe_reindex_device_id(gpu_id) as gpu_id:
proc = mp.Process(
target=run_scheduler_process_func,
args=(
server_args,
port_args,
gpu_id,
tp_rank,
attn_cp_rank,
moe_dp_rank,
moe_ep_rank,
pp_rank,
None,
writer,
),
)
with memory_saver_adapter.configure_subprocess(), numa_utils.configure_subprocess(
server_args, gpu_id
):
proc.start()
scheduler_procs.append(proc)
scheduler_pipe_readers.append(reader)
else:
# Launch the data parallel controller
reader, writer = mp.Pipe(duplex=False)
scheduler_pipe_readers = [reader]
proc = mp.Process(
target=run_data_parallel_controller_process,
kwargs=dict(
server_args=server_args,
port_args=port_args,
pipe_writer=writer,
run_scheduler_process_func=run_scheduler_process_func,
),
)
proc.start()
scheduler_procs.append(proc)
return scheduler_procs, scheduler_pipe_readers
def _launch_subprocesses(
server_args: ServerArgs,
init_tokenizer_manager_func: Callable,
run_scheduler_process_func: Callable,
run_detokenizer_process_func: Callable,
port_args: Optional[PortArgs] = None,
) -> Tuple[TokenizerManager, TemplateManager, Tuple[Dict], PortArgs]:
Returns:
A tuple of (pp_rank_range, tp_rank_range, pp_size_per_node, tp_size_per_node):
- pp_rank_range: range of pipeline-parallel ranks assigned to this node.
- tp_rank_range: range of tensor-parallel ranks assigned to this node.
- pp_size_per_node: number of PP ranks per node.
- tp_size_per_node: number of TP ranks per node.
"""
Launch the TokenizerManager in the main process, the Scheduler in a subprocess, and the DetokenizerManager in another subprocess.
"""
# Configure global environment
configure_logger(server_args)
_set_envs_and_config(server_args)
server_args.check_server_args()
# Allocate ports for inter-process communications
if port_args is None:
port_args = PortArgs.init_new(server_args)
logger.info(f"{server_args=}")
# Launch scheduler processes
scheduler_procs, scheduler_pipe_readers = _launch_scheduler_processes(
server_args=server_args,
port_args=port_args,
run_scheduler_process_func=run_scheduler_process_func,
pp_size_per_node = max(pp_size // nnodes, 1)
nnodes_per_pp_rank = max(nnodes // pp_size, 1)
pp_rank_range = range(
pp_size_per_node * (node_rank // nnodes_per_pp_rank),
pp_size_per_node * (node_rank // nnodes_per_pp_rank + 1),
)
if (
server_args.enable_elastic_expert_backup
and server_args.elastic_ep_backend is not None
):
run_expert_backup_manager(server_args, port_args)
if server_args.node_rank >= 1:
# In multi-node cases, non-zero rank nodes do not need to run tokenizer or detokenizer,
# so they can just wait here.
scheduler_infos = _wait_for_scheduler_ready(
scheduler_pipe_readers, scheduler_procs
)
if os.getenv("SGLANG_BLOCK_NONZERO_RANK_CHILDREN") == "0":
# When using `Engine` as a Python API, we don't want to block here.
return None, None, scheduler_infos, port_args
launch_dummy_health_check_server(
server_args.host, server_args.port, server_args.enable_metrics
)
for proc in scheduler_procs:
proc.join()
logger.error(
f"Scheduler or DataParallelController {proc.pid} terminated with {proc.exitcode}"
)
return None, None, scheduler_infos, port_args
# Launch detokenizer process
detoken_proc = mp.Process(
target=run_detokenizer_process_func,
args=(
server_args,
port_args,
),
nnodes_per_tp_group = nnodes_per_pp_rank
tp_size_per_node = tp_size // nnodes_per_tp_group
tp_rank_range = range(
tp_size_per_node * (node_rank % nnodes_per_tp_group),
tp_size_per_node * (node_rank % nnodes_per_tp_group + 1),
)
detoken_proc.start()
# Init tokenizer manager first, as the bootstrap server is initialized here
if server_args.tokenizer_worker_num == 1:
tokenizer_manager, template_manager = init_tokenizer_manager_func(
server_args, port_args
)
else:
# Launch multi-tokenizer router
tokenizer_manager = MultiTokenizerRouter(server_args, port_args)
template_manager = None
return pp_rank_range, tp_rank_range, pp_size_per_node, tp_size_per_node
# Wait for the model to finish loading
scheduler_infos = _wait_for_scheduler_ready(scheduler_pipe_readers, scheduler_procs)
# Get back some info from scheduler to tokenizer_manager
tokenizer_manager.max_req_input_len = scheduler_infos[0]["max_req_input_len"]
def _compute_parallelism_ranks(
server_args: ServerArgs, tp_rank: int
) -> Tuple[int, int, int]:
"""Compute attention-CP, MoE-DP, and MoE-EP ranks for a TP rank."""
attn_dp_size = server_args.dp_size if server_args.enable_dp_attention else 1
return tokenizer_manager, template_manager, scheduler_infos, port_args
# Parallelism hierarchy (outermost to innermost):
# - Attention: Global(TP) -> DP -> ATTN_CP -> ATTN_TP (innermost)
# - MoE: Global(TP) -> MOE_DP -> EP -> MOE_TP (innermost)
attn_tp_size = server_args.tp_size // attn_dp_size // server_args.attn_cp_size
attn_cp_rank = (tp_rank // attn_tp_size) % server_args.attn_cp_size
moe_dp_rank = tp_rank // (server_args.tp_size // server_args.moe_dp_size)
moe_ep_rank = (
tp_rank
% (server_args.tp_size // server_args.moe_dp_size)
// (server_args.tp_size // server_args.moe_dp_size // server_args.ep_size)
)
return attn_cp_rank, moe_dp_rank, moe_ep_rank

80
python/sglang/srt/entrypoints/http_server.py
View File

@@ -67,7 +67,7 @@ from sglang.srt.entrypoints.anthropic.protocol import (
)
from sglang.srt.entrypoints.anthropic.serving import AnthropicServing
from sglang.srt.entrypoints.engine import (
_launch_subprocesses,
Engine,
init_tokenizer_manager,
run_detokenizer_process,
run_scheduler_process,
@@ -1958,28 +1958,18 @@ def _wait_weights_ready():
)
def launch_server(
def _setup_and_run_http_server(
server_args: ServerArgs,
init_tokenizer_manager_func: Callable = init_tokenizer_manager,
run_scheduler_process_func: Callable = run_scheduler_process,
run_detokenizer_process_func: Callable = run_detokenizer_process,
tokenizer_manager,
template_manager,
port_args: PortArgs,
scheduler_infos: List[Dict],
execute_warmup_func: Callable = _execute_server_warmup,
launch_callback: Optional[Callable[[], None]] = None,
):
"""
Launch SRT (SGLang Runtime) Server.
"""Set up global state, configure middleware, and run uvicorn.
The SRT server consists of an HTTP server and an SRT engine.
- HTTP server: A FastAPI server that routes requests to the engine.
- The engine consists of three components:
1. TokenizerManager: Tokenizes the requests and sends them to the scheduler.
2. Scheduler (subprocess): Receives requests from the Tokenizer Manager, schedules batches, forwards them, and sends the output tokens to the Detokenizer Manager.
3. DetokenizerManager (subprocess): Detokenizes the output tokens and sends the result back to the Tokenizer Manager.
Note:
1. The HTTP server, Engine, and TokenizerManager all run in the main process.
2. Inter-process communication is done through IPC (each process uses a different port) via the ZMQ library.
Called by launch_server after subprocesses have been launched.
"""
# Reserve the HTTP port before launching subprocesses to fail fast if port is unavailable.
# This prevents wasting time loading models only to discover port conflicts later.
@@ -1987,16 +1977,6 @@ def launch_server(
multi_tokenizer_args_shm = None
try:
# Launch subprocesses
tokenizer_manager, template_manager, scheduler_infos, port_args = (
_launch_subprocesses(
server_args=server_args,
init_tokenizer_manager_func=init_tokenizer_manager_func,
run_scheduler_process_func=run_scheduler_process_func,
run_detokenizer_process_func=run_detokenizer_process_func,
)
)
# Parse info got from the schedulers
remote_instance_transfer_engine_info = (
parse_remote_instance_transfer_engine_info_from_scheduler_infos(
@@ -2162,3 +2142,47 @@ def launch_server(
multi_tokenizer_args_shm.unlink()
if _global_state is not None:
_global_state.tokenizer_manager.socket_mapping.clear_all_sockets()
def launch_server(
server_args: ServerArgs,
init_tokenizer_manager_func: Callable = init_tokenizer_manager,
run_scheduler_process_func: Callable = run_scheduler_process,
run_detokenizer_process_func: Callable = run_detokenizer_process,
execute_warmup_func: Callable = _execute_server_warmup,
launch_callback: Optional[Callable[[], None]] = None,
):
"""
Launch SRT (SGLang Runtime) Server.
The SRT server consists of an HTTP server and an SRT engine.
- HTTP server: A FastAPI server that routes requests to the engine.
- The engine consists of three components:
1. TokenizerManager: Tokenizes the requests and sends them to the scheduler.
2. Scheduler (subprocess): Receives requests from the Tokenizer Manager, schedules batches, forwards them, and sends the output tokens to the Detokenizer Manager.
3. DetokenizerManager (subprocess): Detokenizes the output tokens and sends the result back to the Tokenizer Manager.
Note:
1. The HTTP server, Engine, and TokenizerManager all run in the main process.
2. Inter-process communication is done through IPC (each process uses a different port) via the ZMQ library.
"""
# Launch subprocesses
tokenizer_manager, template_manager, port_args, scheduler_init_result = (
Engine._launch_subprocesses(
server_args=server_args,
init_tokenizer_manager_func=init_tokenizer_manager_func,
run_scheduler_process_func=run_scheduler_process_func,
run_detokenizer_process_func=run_detokenizer_process_func,
)
)
_setup_and_run_http_server(
server_args,
tokenizer_manager,
template_manager,
port_args,
scheduler_init_result.scheduler_infos,
execute_warmup_func=execute_warmup_func,
launch_callback=launch_callback,
)

107
python/sglang/srt/managers/scheduler.py
View File

@@ -22,7 +22,7 @@ import time
from collections import deque
from dataclasses import dataclass
from http import HTTPStatus
from typing import Any, Deque, List, Optional, Tuple, Union
from typing import Any, Deque, Dict, List, Optional, Tuple, Union
import psutil
import setproctitle
@@ -1111,6 +1111,45 @@ class Scheduler(
"Request running timeout reached.", HTTPStatus.SERVICE_UNAVAILABLE
)
def get_init_info(self) -> Dict[str, Any]:
"""Return scheduler initialization info for handshake.
This method provides the initialization info needed by the tokenizer manager
and other components to verify the scheduler is ready.
"""
result_dict = {
"status": "ready",
"max_total_num_tokens": self.max_total_num_tokens,
"max_req_input_len": self.max_req_input_len,
}
if self.server_args.remote_instance_weight_loader_use_transfer_engine():
(
remote_instance_transfer_engine_session_id,
remote_instance_transfer_engine_weights_info_dict,
) = self.get_remote_instance_transfer_engine_info()
result_dict.update(
{
"tp_rank": self.tp_rank,
"remote_instance_transfer_engine_session_id": remote_instance_transfer_engine_session_id,
"remote_instance_transfer_engine_weights_info_dict": remote_instance_transfer_engine_weights_info_dict,
}
)
return result_dict
def run_event_loop(self) -> None:
"""Run the scheduler's event loop.
Sets up the schedule stream and dispatches to the appropriate event loop.
The event loop blocks until shutdown.
"""
self.schedule_stream = self.device_module.Stream(priority=0)
if self.device == "cpu":
self.schedule_stream.synchronize = lambda: None # No-op for CPU
with CudaStreamContext(self.schedule_stream):
dispatch_event_loop(self)
@DynamicGradMode()
def event_loop_normal(self):
"""A normal scheduler loop."""
@@ -3110,23 +3149,26 @@ def dispatch_event_loop(scheduler: Scheduler):
scheduler.event_loop_normal_disagg_decode()
def run_scheduler_process(
def configure_scheduler(
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],
pipe_writer,
):
) -> Optional[int]:
"""Configure scheduler worker: logging, process title, etc.
Returns:
dp_rank
"""
# Generate the logger prefix
prefix = ""
if dp_rank is None and "SGLANG_DP_RANK" in os.environ:
# [For Router] if env var "SGLANG_DP_RANK" exist, set dp_rank to the value of the env var
dp_rank = int(os.environ["SGLANG_DP_RANK"])
prefix = ""
if dp_rank is not None:
prefix += f" DP{dp_rank}"
if server_args.pp_size > 1:
@@ -3143,13 +3185,33 @@ def run_scheduler_process(
# Config the process
setproctitle.setproctitle(f"sglang::scheduler{prefix.replace(' ', '_')}")
faulthandler.enable()
kill_itself_when_parent_died()
parent_process = psutil.Process().parent()
# Configure the logger
configure_logger(server_args, prefix=prefix)
suppress_other_loggers()
return dp_rank
def run_scheduler_process(
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],
pipe_writer,
):
dp_rank = configure_scheduler(
server_args, tp_rank, attn_cp_rank, moe_dp_rank, moe_ep_rank, pp_rank, dp_rank
)
kill_itself_when_parent_died()
parent_process = psutil.Process().parent()
# Set cpu affinity to this gpu process
if get_bool_env_var("SGLANG_SET_CPU_AFFINITY"):
set_gpu_proc_affinity(
@@ -3183,31 +3245,12 @@ def run_scheduler_process(
moe_dp_rank,
dp_rank,
)
result_dict = {
"status": "ready",
"max_total_num_tokens": scheduler.max_total_num_tokens,
"max_req_input_len": scheduler.max_req_input_len,
}
if server_args.remote_instance_weight_loader_use_transfer_engine():
(
remote_instance_transfer_engine_session_id,
remote_instance_transfer_engine_weights_info_dict,
) = scheduler.get_remote_instance_transfer_engine_info()
result_dict.update(
{
"tp_rank": tp_rank,
"remote_instance_transfer_engine_session_id": remote_instance_transfer_engine_session_id,
"remote_instance_transfer_engine_weights_info_dict": remote_instance_transfer_engine_weights_info_dict,
}
)
pipe_writer.send(result_dict)
# Send initialization info back to the parent process
pipe_writer.send(scheduler.get_init_info())
scheduler.schedule_stream = scheduler.device_module.Stream(priority=0)
if scheduler.device == "cpu":
scheduler.schedule_stream.synchronize = lambda: None # No-op for CPU
with CudaStreamContext(scheduler.schedule_stream):
dispatch_event_loop(scheduler)
# Run the event loop (blocks until shutdown)
scheduler.run_event_loop()
except Exception:
traceback = get_exception_traceback()

3
python/sglang/srt/ray/__init__.py Normal file
View File

@@ -0,0 +1,3 @@
from sglang.srt.ray.engine import RayEngine
__all__ = ["RayEngine"]

174
python/sglang/srt/ray/engine.py Normal file
View File

@@ -0,0 +1,174 @@
# 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
from typing import Callable
import ray
from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
from sglang.srt.entrypoints.engine import (
Engine,
SchedulerInitResult,
_calculate_rank_ranges,
_compute_parallelism_ranks,
)
from sglang.srt.ray.scheduler_actor import SchedulerActor
from sglang.srt.server_args import ZMQ_TCP_PORT_DELTA, 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 _get_rank0_node_ip(placement_group) -> str:
"""Get the IP address of the node where rank 0 will run.
Uses a probe task to discover the IP of the placement group's first bundle node.
This is needed because rank 0 starts the TCPStore server for torch.distributed,
so dist_init_addr must be the IP of the node where rank 0 runs, not the driver node.
"""
@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=0,
),
).remote()
)
class RayEngine(Engine):
"""Engine using Ray actors for scheduler processes."""
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,
) -> SchedulerInitResult:
"""Launch schedulers as Ray actors."""
if server_args.dp_size > 1:
raise NotImplementedError(
"Ray support for dp_size > 1 is not yet implemented. "
"Set dp_size=1 or use_ray=False."
)
pg = ray.util.get_current_placement_group()
if pg is None:
raise RuntimeError(
"use_ray=True requires a placement group, but none was detected. "
"Schedule the Engine actor onto a placement group"
)
world_size = server_args.tp_size * server_args.pp_size
nnodes = server_args.nnodes
gpus_per_node = world_size // nnodes
logger.info(
f"Ray cluster: {nnodes} nodes, "
f"Use {gpus_per_node} GPUs/node, world_size={world_size}"
)
rank0_node_ip = _get_rank0_node_ip(pg)
dist_init_addr = f"{rank0_node_ip}:{server_args.port + ZMQ_TCP_PORT_DELTA}"
logger.info(f"dist_init_addr: {dist_init_addr}")
scheduler_actors = []
for node_idx in range(nnodes):
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_rank0node={rank0_node_ip}_pp{pp_rank}_tp{tp_rank}",
scheduling_strategy=PlacementGroupSchedulingStrategy(
placement_group=pg,
placement_group_bundle_index=node_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)
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,
)

64
python/sglang/srt/ray/http_server.py Normal file
View File

@@ -0,0 +1,64 @@
# 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 HTTP server launcher."""
from typing import Callable, Optional
from sglang.srt.entrypoints.engine import (
init_tokenizer_manager,
run_detokenizer_process,
run_scheduler_process,
)
from sglang.srt.server_args import ServerArgs
def launch_server(
server_args: ServerArgs,
init_tokenizer_manager_func: Callable = init_tokenizer_manager,
run_scheduler_process_func: Callable = run_scheduler_process,
run_detokenizer_process_func: Callable = run_detokenizer_process,
execute_warmup_func: Optional[Callable] = None,
launch_callback: Optional[Callable[[], None]] = None,
):
"""Launch HTTP server with Ray-based scheduler actors.
Mirrors http_server.launch_server() but uses RayEngine for scheduler launching.
"""
from sglang.srt.entrypoints.http_server import (
_execute_server_warmup,
_setup_and_run_http_server,
)
from sglang.srt.ray.engine import RayEngine
if execute_warmup_func is None:
execute_warmup_func = _execute_server_warmup
tokenizer_manager, template_manager, port_args, scheduler_init_result = (
RayEngine._launch_subprocesses(
server_args,
init_tokenizer_manager_func=init_tokenizer_manager_func,
run_scheduler_process_func=run_scheduler_process_func,
run_detokenizer_process_func=run_detokenizer_process_func,
)
)
_setup_and_run_http_server(
server_args,
tokenizer_manager,
template_manager,
port_args,
scheduler_init_result.scheduler_infos,
execute_warmup_func=execute_warmup_func,
launch_callback=launch_callback,
)

111
python/sglang/srt/ray/scheduler_actor.py Normal file
View File

@@ -0,0 +1,111 @@
# 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.managers.scheduler import Scheduler, configure_scheduler
# 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(
server_args,
tp_rank,
attn_cp_rank,
moe_dp_rank,
moe_ep_rank,
pp_rank,
dp_rank,
)
# 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:
self.scheduler.run_event_loop()
except Exception as e:
logger.error(f"Scheduler PP{self._pp_rank} TP{self._tp_rank} crashed: {e}")
raise

6
python/sglang/srt/server_args.py
View File

@@ -376,6 +376,7 @@ class ServerArgs:
base_gpu_id: int = 0
gpu_id_step: int = 1
sleep_on_idle: bool = False
use_ray: bool = False
custom_sigquit_handler: Optional[Callable] = None
# Logging
@@ -3726,6 +3727,11 @@ class ServerArgs:
action="store_true",
help="Reduce CPU usage when sglang is idle.",
)
parser.add_argument(
"--use-ray",
action="store_true",
help="Use Ray actors for scheduler process management.",
)
parser.add_argument(
"--custom-sigquit-handler",
help="Register a custom sigquit handler so you can do additional cleanup after the server is shutdown. This is only available for Engine, not for CLI.",