From b3cfad0a80069a0952b4eadb275316dc16b59abe Mon Sep 17 00:00:00 2001 From: Xinyu Zhang <60529799+xyuzh@users.noreply.github.com> Date: Thu, 5 Mar 2026 13:21:23 -0800 Subject: [PATCH] Add Ray actor support for scheduler process management (DP=1) (#17684) Co-authored-by: Cursor --- python/pyproject.toml | 4 + python/sglang/bench_offline_throughput.py | 73 ++- python/sglang/launch_server.py | 10 + python/sglang/srt/entrypoints/engine.py | 457 ++++++++++-------- python/sglang/srt/entrypoints/http_server.py | 80 ++-- python/sglang/srt/managers/scheduler.py | 107 +++-- python/sglang/srt/ray/__init__.py | 3 + python/sglang/srt/ray/engine.py | 174 +++++++ python/sglang/srt/ray/http_server.py | 64 +++ python/sglang/srt/ray/scheduler_actor.py | 111 +++++ python/sglang/srt/server_args.py | 6 + test/manual/test_ray_engine.py | 461 +++++++++++++++++++ 12 files changed, 1303 insertions(+), 247 deletions(-) create mode 100644 python/sglang/srt/ray/__init__.py create mode 100644 python/sglang/srt/ray/engine.py create mode 100644 python/sglang/srt/ray/http_server.py create mode 100644 python/sglang/srt/ray/scheduler_actor.py create mode 100644 test/manual/test_ray_engine.py diff --git a/python/pyproject.toml b/python/pyproject.toml index fa5c705fa..d0feacb86 100755 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -121,6 +121,10 @@ diffusion = [ "xatlas", ] +ray = [ + "ray[default]>=2.54.0", +] + tracing = [ "opentelemetry-api", "opentelemetry-exporter-otlp", diff --git a/python/sglang/bench_offline_throughput.py b/python/sglang/bench_offline_throughput.py index b334a155d..0943acd8d 100644 --- a/python/sglang/bench_offline_throughput.py +++ b/python/sglang/bench_offline_throughput.py @@ -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": diff --git a/python/sglang/launch_server.py b/python/sglang/launch_server.py index af4a41e62..fe06a9289 100644 --- a/python/sglang/launch_server.py +++ b/python/sglang/launch_server.py @@ -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 diff --git a/python/sglang/srt/entrypoints/engine.py b/python/sglang/srt/entrypoints/engine.py index 302cc9f77..87c8c0a19 100644 --- a/python/sglang/srt/entrypoints/engine.py +++ b/python/sglang/srt/entrypoints/engine.py @@ -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 diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/entrypoints/http_server.py index 4eb546f68..d57108c41 100644 --- a/python/sglang/srt/entrypoints/http_server.py +++ b/python/sglang/srt/entrypoints/http_server.py @@ -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, + ) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 023efb9f3..c18539f19 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -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() diff --git a/python/sglang/srt/ray/__init__.py b/python/sglang/srt/ray/__init__.py new file mode 100644 index 000000000..5927c789f --- /dev/null +++ b/python/sglang/srt/ray/__init__.py @@ -0,0 +1,3 @@ +from sglang.srt.ray.engine import RayEngine + +__all__ = ["RayEngine"] diff --git a/python/sglang/srt/ray/engine.py b/python/sglang/srt/ray/engine.py new file mode 100644 index 000000000..9c1fe75cb --- /dev/null +++ b/python/sglang/srt/ray/engine.py @@ -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, + ) diff --git a/python/sglang/srt/ray/http_server.py b/python/sglang/srt/ray/http_server.py new file mode 100644 index 000000000..c0580838e --- /dev/null +++ b/python/sglang/srt/ray/http_server.py @@ -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, + ) diff --git a/python/sglang/srt/ray/scheduler_actor.py b/python/sglang/srt/ray/scheduler_actor.py new file mode 100644 index 000000000..e68b78e94 --- /dev/null +++ b/python/sglang/srt/ray/scheduler_actor.py @@ -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 diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index f6f411155..03fe80e03 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -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.", diff --git a/test/manual/test_ray_engine.py b/test/manual/test_ray_engine.py new file mode 100644 index 000000000..6aa767d6b --- /dev/null +++ b/test/manual/test_ray_engine.py @@ -0,0 +1,461 @@ +"""Integration tests for RayEngine and Ray HTTP server (requires GPU + Ray). + +Tests the Ray actor scheduler backend: + - Offline inference via Engine(use_ray=True) inside a Ray actor on a placement group + - Error paths in RayEngine._launch_scheduler_processes() + - HTTP server launched via --use-ray flag + +Usage: + # 1-GPU tests + 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 + + # 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 +""" + +from __future__ import annotations + +import os +import time +import unittest + +import torch + +from sglang.test.test_utils import DEFAULT_SMALL_MODEL_NAME_FOR_TEST + +# Allow overriding the model via env var for environments without gated access +_MODEL = os.environ.get("SGLANG_TEST_MODEL", DEFAULT_SMALL_MODEL_NAME_FOR_TEST) + +try: + import ray + from ray.runtime_env import RuntimeEnv + from ray.util.placement_group import placement_group + from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy + + # Prevent Ray from overriding CUDA_VISIBLE_DEVICES so that all GPUs + # remain visible inside actors regardless of num_gpus allocation. + _env_vars = {"RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES": "1"} + if os.environ.get("HF_TOKEN"): + _env_vars["HF_TOKEN"] = os.environ["HF_TOKEN"] + _RAY_RUNTIME_ENV = RuntimeEnv(env_vars=_env_vars) + _has_ray = True +except ImportError: + _has_ray = False + _RAY_RUNTIME_ENV = None + + +_NUM_GPUS = torch.cuda.device_count() + +_SAMPLING_PARAMS = {"max_new_tokens": 32, "temperature": 0.0} + +_PROMPTS = [ + "The capital of France is", + "Explain quantum computing in simple terms:", + "Write a haiku about programming:", + "What is 2 + 2?", +] + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _create_engine_on_pg(tp_size, pp_size=1, model=_MODEL, extra_kwargs=None): + """Create an EngineActor on a placement group and wait for it to be ready. + + Returns (engine_actor, placement_group). + """ + + @ray.remote + class EngineActor: + def __init__(self, **kwargs): + from sglang.srt.ray.engine import RayEngine + + self.engine = RayEngine(**kwargs) + + def is_ready(self): + return True + + def generate(self, prompt, sampling_params): + return self.engine.generate(prompt=prompt, sampling_params=sampling_params) + + def shutdown(self): + if self.engine: + self.engine.shutdown() + self.engine = None + + total_gpus = tp_size * pp_size + pg = placement_group( + [{"CPU": 1, "GPU": total_gpus}], + strategy="STRICT_PACK", + ) + ray.get(pg.ready()) + + kwargs = dict( + model_path=model, + tp_size=tp_size, + pp_size=pp_size, + ) + if extra_kwargs: + kwargs.update(extra_kwargs) + + actor = EngineActor.options( + num_cpus=1, + num_gpus=0, + scheduling_strategy=PlacementGroupSchedulingStrategy( + placement_group=pg, + placement_group_bundle_index=0, + ), + ).remote(**kwargs) + + ray.get(actor.is_ready.remote(), timeout=600) + return actor, pg + + +def _cleanup(actor, pg): + """Shutdown engine actor and remove placement group.""" + try: + ray.get(actor.shutdown.remote(), timeout=60) + except Exception: + pass + try: + ray.util.remove_placement_group(pg) + except Exception: + pass + + +# --------------------------------------------------------------------------- +# Tests: Offline TP=1 +# --------------------------------------------------------------------------- + + +@unittest.skipUnless(_has_ray, "ray is not installed") +@unittest.skipUnless(_NUM_GPUS >= 1, "requires at least 1 GPU") +class TestRayEngineOfflineTP1(unittest.TestCase): + + @classmethod + def setUpClass(cls): + if not ray.is_initialized(): + ray.init(log_to_driver=True, runtime_env=_RAY_RUNTIME_ENV) + cls.actor, cls.pg = _create_engine_on_pg(tp_size=1) + + @classmethod + def tearDownClass(cls): + _cleanup(cls.actor, cls.pg) + ray.shutdown() + + def test_offline_generate(self): + result = ray.get( + self.actor.generate.remote("The capital of France is", _SAMPLING_PARAMS) + ) + self.assertIn("text", result) + self.assertGreater(len(result["text"]), 0) + print(f"Generated: {result['text'][:200]}") + + def test_batch_generate(self): + for prompt in _PROMPTS: + result = ray.get(self.actor.generate.remote(prompt, _SAMPLING_PARAMS)) + self.assertIn("text", result) + self.assertGreater(len(result["text"]), 0, f"Empty output for: {prompt}") + + def test_deterministic(self): + prompt = "The meaning of life is" + r1 = ray.get(self.actor.generate.remote(prompt, _SAMPLING_PARAMS)) + r2 = ray.get(self.actor.generate.remote(prompt, _SAMPLING_PARAMS)) + self.assertEqual(r1["text"], r2["text"]) + + +# --------------------------------------------------------------------------- +# Tests: Offline TP=2 +# --------------------------------------------------------------------------- + + +@unittest.skipUnless(_has_ray, "ray is not installed") +@unittest.skipUnless(_NUM_GPUS >= 2, "requires at least 2 GPUs") +class TestRayEngineOfflineTP2(unittest.TestCase): + + @classmethod + def setUpClass(cls): + if not ray.is_initialized(): + ray.init(log_to_driver=True, runtime_env=_RAY_RUNTIME_ENV) + cls.actor, cls.pg = _create_engine_on_pg(tp_size=2) + + @classmethod + def tearDownClass(cls): + _cleanup(cls.actor, cls.pg) + ray.shutdown() + + def test_offline_generate_tp2(self): + result = ray.get( + self.actor.generate.remote("The capital of France is", _SAMPLING_PARAMS) + ) + self.assertIn("text", result) + self.assertGreater(len(result["text"]), 0) + print(f"Generated (TP=2): {result['text'][:200]}") + + def test_batch_generate_tp2(self): + for prompt in _PROMPTS: + result = ray.get(self.actor.generate.remote(prompt, _SAMPLING_PARAMS)) + self.assertIn("text", result) + self.assertGreater(len(result["text"]), 0, f"Empty output for: {prompt}") + + +# --------------------------------------------------------------------------- +# Tests: Offline PP=2 +# --------------------------------------------------------------------------- + + +@unittest.skipUnless(_has_ray, "ray is not installed") +@unittest.skipUnless(_NUM_GPUS >= 2, "requires at least 2 GPUs") +class TestRayEngineOfflinePP2(unittest.TestCase): + + @classmethod + def setUpClass(cls): + if not ray.is_initialized(): + ray.init(log_to_driver=True, runtime_env=_RAY_RUNTIME_ENV) + cls.actor, cls.pg = _create_engine_on_pg(tp_size=1, pp_size=2) + + @classmethod + def tearDownClass(cls): + _cleanup(cls.actor, cls.pg) + ray.shutdown() + + def test_offline_generate_pp2(self): + result = ray.get( + self.actor.generate.remote("The capital of France is", _SAMPLING_PARAMS) + ) + self.assertIn("text", result) + self.assertGreater(len(result["text"]), 0) + print(f"Generated (PP=2): {result['text'][:200]}") + + def test_batch_generate_pp2(self): + for prompt in _PROMPTS: + result = ray.get(self.actor.generate.remote(prompt, _SAMPLING_PARAMS)) + self.assertIn("text", result) + self.assertGreater(len(result["text"]), 0, f"Empty output for: {prompt}") + + +# --------------------------------------------------------------------------- +# Tests: Error paths +# --------------------------------------------------------------------------- + + +@unittest.skipUnless(_has_ray, "ray is not installed") +@unittest.skipUnless(_NUM_GPUS >= 1, "requires at least 1 GPU") +class TestRayEngineErrors(unittest.TestCase): + + @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_dp_greater_than_1_raises(self): + """RayEngine with dp_size > 1 should raise NotImplementedError.""" + + @ray.remote + class _BadActor: + def try_create(self): + from sglang.srt.ray.engine import RayEngine + + try: + RayEngine( + model_path=_MODEL, + tp_size=1, + dp_size=2, + use_ray=True, + ) + return None + except (NotImplementedError, RuntimeError) as e: + return str(e) + + pg = placement_group([{"CPU": 1, "GPU": 1}], strategy="STRICT_PACK") + ray.get(pg.ready()) + + actor = _BadActor.options( + num_cpus=1, + num_gpus=0, + scheduling_strategy=PlacementGroupSchedulingStrategy( + placement_group=pg, + placement_group_bundle_index=0, + ), + ).remote() + + try: + error_msg = ray.get(actor.try_create.remote(), timeout=120) + self.assertIsNotNone(error_msg, "Expected error but RayEngine created OK") + self.assertIn("dp_size", error_msg.lower()) + finally: + ray.util.remove_placement_group(pg) + + def test_missing_placement_group_raises(self): + """RayEngine without a placement group should raise RuntimeError.""" + + @ray.remote(num_gpus=1) + def _try_create_without_pg(): + from sglang.srt.ray.engine import RayEngine + + try: + RayEngine( + model_path=_MODEL, + tp_size=1, + use_ray=True, + ) + return None + except RuntimeError as e: + return str(e) + + error_msg = ray.get(_try_create_without_pg.remote(), timeout=120) + self.assertIsNotNone( + error_msg, "Expected RuntimeError but RayEngine created OK" + ) + self.assertIn("placement group", error_msg.lower()) + + +# --------------------------------------------------------------------------- +# Tests: HTTP server +# --------------------------------------------------------------------------- + + +@unittest.skipUnless(_has_ray, "ray is not installed") +@unittest.skipUnless(_NUM_GPUS >= 1, "requires at least 1 GPU") +class TestRayHTTPServerTP1(unittest.TestCase): + """Test the Ray HTTP server path (launch_server.py --use-ray). + + Launches the server inside a Ray task on a placement group (mirrors + examples/anyscale/driver_online.py) and sends HTTP requests to it. + """ + + @classmethod + def setUpClass(cls): + import requests as req_lib + + if not ray.is_initialized(): + ray.init(log_to_driver=True, runtime_env=_RAY_RUNTIME_ENV) + + cls.port = 30100 + cls.pg = placement_group( + [{"CPU": 1, "GPU": 1}], + strategy="STRICT_PACK", + ) + ray.get(cls.pg.ready()) + + pg_strategy = PlacementGroupSchedulingStrategy( + placement_group=cls.pg, + placement_group_bundle_index=0, + ) + + # Resolve the node IP where the server will run + @ray.remote(num_cpus=0, num_gpus=0) + def _get_ip(): + return ray.util.get_node_ip_address() + + cls.node_ip = ray.get(_get_ip.options(scheduling_strategy=pg_strategy).remote()) + cls.base_url = f"http://{cls.node_ip}:{cls.port}" + + # Launch server as a Ray task (blocks until server exits) + @ray.remote + def _launch(**kwargs): + from sglang.srt.ray.http_server import launch_server + from sglang.srt.server_args import ServerArgs + + launch_server(ServerArgs(**kwargs)) + + cls.server_ref = _launch.options( + num_cpus=1, + num_gpus=0, + scheduling_strategy=pg_strategy, + ).remote( + model_path=_MODEL, + tp_size=1, + port=cls.port, + host="0.0.0.0", + use_ray=True, + ) + + # Wait for health check + t0 = time.time() + timeout = 600 + healthy = False + while time.time() - t0 < timeout: + ready, _ = ray.wait([cls.server_ref], timeout=0) + if ready: + try: + ray.get(cls.server_ref) + except Exception as e: + raise RuntimeError(f"Server task crashed: {e}") from e + raise RuntimeError("Server task exited before becoming healthy") + try: + if req_lib.get(f"{cls.base_url}/health", timeout=5).status_code == 200: + healthy = True + break + except req_lib.exceptions.RequestException: + pass + time.sleep(3) + + if not healthy: + ray.cancel(cls.server_ref, force=True) + raise RuntimeError(f"Server did not become healthy within {timeout}s") + + @classmethod + def tearDownClass(cls): + try: + ray.cancel(cls.server_ref, force=True) + except Exception: + pass + try: + ray.util.remove_placement_group(cls.pg) + except Exception: + pass + ray.shutdown() + + def test_health_endpoint(self): + import requests + + resp = requests.get(f"{self.base_url}/health", timeout=10) + self.assertEqual(resp.status_code, 200) + + def test_generate_endpoint(self): + import requests + + resp = requests.post( + f"{self.base_url}/generate", + json={ + "text": "The capital of France is", + "sampling_params": _SAMPLING_PARAMS, + }, + timeout=60, + ) + resp.raise_for_status() + data = resp.json() + self.assertIn("text", data) + self.assertGreater(len(data["text"]), 0) + print(f"HTTP response: {data['text'][:200]}") + + def test_generate_multiple(self): + import requests + + for prompt in _PROMPTS: + resp = requests.post( + f"{self.base_url}/generate", + json={ + "text": prompt, + "sampling_params": _SAMPLING_PARAMS, + }, + timeout=60, + ) + resp.raise_for_status() + data = resp.json() + self.assertIn("text", data) + self.assertGreater(len(data["text"]), 0, f"Empty output for: {prompt}") + + +if __name__ == "__main__": + unittest.main()