diff --git a/python/sglang/srt/entrypoints/engine.py b/python/sglang/srt/entrypoints/engine.py index 912d1b9ec..921c31dc2 100644 --- a/python/sglang/srt/entrypoints/engine.py +++ b/python/sglang/srt/entrypoints/engine.py @@ -27,15 +27,14 @@ import random import signal import threading import time -from typing import AsyncIterator, Dict, Iterator, List, Optional, Tuple, Union - -import zmq +from typing import AsyncIterator, Callable, Dict, Iterator, List, Optional, Tuple, Union # Fix a bug of Python threading setattr(threading, "_register_atexit", lambda *args, **kwargs: None) import torch import uvloop +import zmq from sglang.srt.entrypoints.EngineBase import EngineBase from sglang.srt.managers.data_parallel_controller import ( @@ -90,6 +89,102 @@ asyncio.set_event_loop_policy(uvloop.EventLoopPolicy()) _is_cuda = is_cuda() +def _launch_subprocesses( + server_args: ServerArgs, port_args: Optional[PortArgs] = None +) -> Tuple[TokenizerManager, TemplateManager, Dict, PortArgs]: + """ + 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=}") + + # If using model from www.modelscope.cn, first download the model + server_args.model_path, server_args.tokenizer_path = prepare_model_and_tokenizer( + server_args.model_path, server_args.tokenizer_path + ) + + # Launch scheduler processes + scheduler_procs, scheduler_pipe_readers = _launch_scheduler_processes( + server_args=server_args, + port_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. + + for reader in scheduler_pipe_readers: + data = reader.recv() + assert data["status"] == "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, None, 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, None, port_args + + # Launch detokenizer process + detoken_proc = mp.Process( + target=run_detokenizer_process, + 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( + 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_infos = [] + for i in range(len(scheduler_pipe_readers)): + try: + data = scheduler_pipe_readers[i].recv() + except EOFError: + logger.error( + f"Rank {i} scheduler is dead. Please check if there are relevant logs." + ) + scheduler_procs[i].join() + logger.error(f"Exit code: {scheduler_procs[i].exitcode}") + raise + + if data["status"] != "ready": + raise RuntimeError( + "Initialization failed. Please see the error messages above." + ) + scheduler_infos.append(data) + + # Get back some info from scheduler to tokenizer_manager + scheduler_info = scheduler_infos[0] + tokenizer_manager.max_req_input_len = scheduler_info["max_req_input_len"] + + return tokenizer_manager, template_manager, scheduler_info, port_args + + class Engine(EngineBase): """ The entry point to the inference engine. @@ -104,6 +199,11 @@ class Engine(EngineBase): 2. Inter-process communication is done through IPC (each process uses a different port) via the ZMQ library. """ + # Some fields to allow people to override the server args + # and launch processes for their private forks. + launch_subprocesses_func: Callable = staticmethod(_launch_subprocesses) + server_args_class: ServerArgs = ServerArgs + def __init__(self, **kwargs): """ The arguments of this function is the same as `sglang/srt/server_args.py::ServerArgs`. @@ -119,7 +219,7 @@ class Engine(EngineBase): if "log_level" not in kwargs: # Do not print logs by default kwargs["log_level"] = "error" - server_args = ServerArgs(**kwargs) + server_args = self.server_args_class(**kwargs) self.server_args = server_args logger.info(f"{server_args=}") @@ -128,7 +228,7 @@ class Engine(EngineBase): # Launch subprocesses tokenizer_manager, template_manager, scheduler_info, port_args = ( - _launch_subprocesses(server_args=server_args) + self.launch_subprocesses_func(server_args=server_args) ) self.tokenizer_manager = tokenizer_manager self.template_manager = template_manager @@ -780,28 +880,14 @@ def _init_tokenizer_manager( return tokenizer_manager, template_manager -def _launch_subprocesses( - server_args: ServerArgs, port_args: Optional[PortArgs] = None -) -> Tuple[TokenizerManager, TemplateManager, Dict, PortArgs]: - """ - 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=}") - - # If using model from www.modelscope.cn, first download the model. - server_args.model_path, server_args.tokenizer_path = prepare_model_and_tokenizer( - server_args.model_path, server_args.tokenizer_path - ) - +def _launch_scheduler_processes( + server_args: ServerArgs, + port_args: PortArgs, + run_scheduler_process_func: Callable = run_scheduler_process, + run_data_parallel_controller_process_func: Callable = run_data_parallel_controller_process, +): scheduler_procs = [] + if server_args.dp_size == 1: # Launch tensor parallel scheduler processes memory_saver_adapter = TorchMemorySaverAdapter.create( @@ -835,7 +921,7 @@ def _launch_subprocesses( with maybe_reindex_device_id(gpu_id) as gpu_id: proc = mp.Process( - target=run_scheduler_process, + target=run_scheduler_process_func, args=( server_args, port_args, @@ -859,76 +945,10 @@ def _launch_subprocesses( reader, writer = mp.Pipe(duplex=False) scheduler_pipe_readers = [reader] proc = mp.Process( - target=run_data_parallel_controller_process, + target=run_data_parallel_controller_process_func, args=(server_args, port_args, writer), ) proc.start() scheduler_procs.append(proc) - 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. - - for reader in scheduler_pipe_readers: - data = reader.recv() - assert data["status"] == "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, None, 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, None, port_args - - # Launch detokenizer process - detoken_proc = mp.Process( - target=run_detokenizer_process, - 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( - 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_infos = [] - for i in range(len(scheduler_pipe_readers)): - try: - data = scheduler_pipe_readers[i].recv() - except EOFError: - logger.error( - f"Rank {i} scheduler is dead. Please check if there are relevant logs." - ) - scheduler_procs[i].join() - logger.error(f"Exit code: {scheduler_procs[i].exitcode}") - raise - - if data["status"] != "ready": - raise RuntimeError( - "Initialization failed. Please see the error messages above." - ) - scheduler_infos.append(data) - - # Assume all schedulers have the same scheduler_info - scheduler_info = scheduler_infos[0] - tokenizer_manager.max_req_input_len = scheduler_info["max_req_input_len"] - - return tokenizer_manager, template_manager, scheduler_info, port_args + return scheduler_procs, scheduler_pipe_readers diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/entrypoints/http_server.py index 5eb908d62..fa14836dc 100644 --- a/python/sglang/srt/entrypoints/http_server.py +++ b/python/sglang/srt/entrypoints/http_server.py @@ -25,16 +25,22 @@ import os import tempfile import threading import time +from contextlib import asynccontextmanager from http import HTTPStatus -from typing import Any, AsyncIterator, Callable, Dict, List, Optional, Union - -from sglang.srt.tracing.trace import process_tracing_init, trace_set_thread_info +from typing import ( + Any, + AsyncGenerator, + AsyncIterator, + Callable, + Dict, + List, + Optional, + Union, +) # Fix a bug of Python threading setattr(threading, "_register_atexit", lambda *args, **kwargs: None) -from contextlib import asynccontextmanager -from typing import AsyncGenerator import numpy as np import orjson @@ -119,6 +125,7 @@ from sglang.srt.managers.tokenizer_manager import ServerStatus, TokenizerManager from sglang.srt.metrics.func_timer import enable_func_timer from sglang.srt.parser.reasoning_parser import ReasoningParser from sglang.srt.server_args import PortArgs, ServerArgs +from sglang.srt.tracing.trace import process_tracing_init, trace_set_thread_info from sglang.srt.utils import ( add_api_key_middleware, add_prometheus_middleware, @@ -134,6 +141,7 @@ from sglang.version import __version__ logger = logging.getLogger(__name__) asyncio.set_event_loop_policy(uvloop.EventLoopPolicy()) +# Global constants HEALTH_CHECK_TIMEOUT = int(os.getenv("SGLANG_HEALTH_CHECK_TIMEOUT", 20)) WAIT_WEIGHTS_READY_TIMEOUT = int(os.getenv("SGLANG_WAIT_WEIGHTS_READY_TIMEOUT", 120)) @@ -154,8 +162,15 @@ def set_global_state(global_state: _GlobalState): _global_state = global_state +def get_global_state() -> _GlobalState: + return _global_state + + async def init_multi_tokenizer() -> ServerArgs: - """Read args information from shm and init tokenizer manager for current process""" + """ + Initialization function for multi-process tokenizer mode. + It read args information from shm and inits tokenizer manager for current process. + """ # Read configuration from shared memory main_pid = get_main_process_id() @@ -206,16 +221,12 @@ async def init_multi_tokenizer() -> ServerArgs: async def lifespan(fast_api_app: FastAPI): if getattr(fast_api_app, "is_single_tokenizer_mode", False): server_args = fast_api_app.server_args - warmup_thread_args = fast_api_app.warmup_thread_args + warmup_thread_kwargs = fast_api_app.warmup_thread_kwargs thread_label = "Tokenizer" else: # Initialize multi-tokenizer support for worker processes server_args = await init_multi_tokenizer() - warmup_thread_args = ( - server_args, - None, - None, - ) + warmup_thread_kwargs = dict(server_args=server_args) thread_label = f"MultiTokenizer-{_global_state.tokenizer_manager.worker_id}" # Add prometheus middleware @@ -298,7 +309,7 @@ async def lifespan(fast_api_app: FastAPI): # Execute the general warmup warmup_thread = threading.Thread( target=_wait_and_warmup, - args=warmup_thread_args, + kwargs=warmup_thread_kwargs, ) warmup_thread.start() @@ -434,8 +445,9 @@ async def health_generate(request: Request) -> Response: rid = f"HEALTH_CHECK_{time.time()}" if _global_state.tokenizer_manager.is_image_gen: - # Keep this branch for some internal use cases. - raise NotImplementedError("Image generation is not supported yet.") + gri = _global_state.tokenizer_manager.get_image_gen_health_check_request( + rid, sampling_params + ) elif _global_state.tokenizer_manager.is_generation: gri = GenerateReqInput( rid=rid, @@ -1366,106 +1378,6 @@ def _create_error_response(e): ) -def launch_server( - server_args: ServerArgs, - pipe_finish_writer: Optional[multiprocessing.connection.Connection] = None, - 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. - """ - tokenizer_manager, template_manager, scheduler_info, port_args = ( - _launch_subprocesses(server_args=server_args) - ) - - set_global_state( - _GlobalState( - tokenizer_manager=tokenizer_manager, - template_manager=template_manager, - scheduler_info=scheduler_info, - ) - ) - - if server_args.enable_metrics: - add_prometheus_track_response_middleware(app) - - # Pass additional arguments to the lifespan function. - # They will be used for additional initialization setups. - if server_args.tokenizer_worker_num == 1: - # If it is single tokenizer mode, we can pass the arguments by attributes of the app object. - app.is_single_tokenizer_mode = True - app.server_args = server_args - app.warmup_thread_args = ( - server_args, - pipe_finish_writer, - launch_callback, - ) - - # Add api key authorization - # This is only supported in single tokenizer mode. - if server_args.api_key: - add_api_key_middleware(app, server_args.api_key) - else: - # If it is multi-tokenizer mode, we need to write the arguments to shared memory - # for other worker processes to read. - app.is_single_tokenizer_mode = False - multi_tokenizer_args_shm = write_data_for_multi_tokenizer( - port_args, server_args, scheduler_info - ) - - try: - # Update logging configs - set_uvicorn_logging_configs() - - # Listen for HTTP requests - if server_args.tokenizer_worker_num == 1: - uvicorn.run( - app, - host=server_args.host, - port=server_args.port, - root_path=server_args.fastapi_root_path, - log_level=server_args.log_level_http or server_args.log_level, - timeout_keep_alive=5, - loop="uvloop", - ) - else: - from uvicorn.config import LOGGING_CONFIG - - LOGGING_CONFIG["loggers"]["sglang.srt.entrypoints.http_server"] = { - "handlers": ["default"], - "level": "INFO", - "propagate": False, - } - monkey_patch_uvicorn_multiprocessing() - - uvicorn.run( - "sglang.srt.entrypoints.http_server:app", - host=server_args.host, - port=server_args.port, - root_path=server_args.fastapi_root_path, - log_level=server_args.log_level_http or server_args.log_level, - timeout_keep_alive=5, - loop="uvloop", - workers=server_args.tokenizer_worker_num, - ) - finally: - if server_args.tokenizer_worker_num > 1: - multi_tokenizer_args_shm.unlink() - _global_state.tokenizer_manager.socket_mapping.clear_all_sockets() - - # Minimal 32x32 black PNG (base64, GLM4v requires at least 32x32 sized image) MINIMUM_PNG_PICTURE_BASE64 = "iVBORw0KGgoAAAANSUhEUgAAACAAAAAgCAYAAABzenr0AAAACXBIWXMAAA7EAAAOxAGVKw4bAAAAbUlEQVRYhe3VsQ2AMAxE0Y/lIgNQULD/OqyCMgCihCKSG4yRuKuiNH6JLsoEbMACOGBcua9HOR7Y6w6swBwMy0qLTpkeI77qdEBpBFAHBBDAGH8WrwJKI4AAegUCfAKgEgpQDvh3CR3oQCuav58qlAw73kKCSgAAAABJRU5ErkJggg==" @@ -1501,9 +1413,8 @@ def _execute_server_warmup( model_info = res.json() + # Construct a warmup request is_vlm = bool(model_info.get("has_image_understanding", False)) - - # Send a warmup request if model_info["is_generation"]: if is_vlm and not server_args.skip_tokenizer_init: request_name = "/v1/chat/completions" @@ -1554,7 +1465,7 @@ def _execute_server_warmup( if server_args.dp_size == 1: json_data["text"] = json_data["text"][0] - # Debug dumping + # Config debug dumping if server_args.debug_tensor_dump_input_file: json_data.pop("text", None) json_data["input_ids"] = np.load( @@ -1562,8 +1473,9 @@ def _execute_server_warmup( ).tolist() json_data["sampling_params"]["max_new_tokens"] = 0 + # Send a warmup request + warmup_timeout = envs.SGLANG_WARMUP_TIMEOUT.get() try: - warmup_timeout = envs.SGLANG_WARMUP_TIMEOUT.get() if server_args.disaggregation_mode == "null": res = requests.post( url + request_name, @@ -1627,13 +1539,16 @@ def _execute_server_warmup( def _wait_and_warmup( server_args: ServerArgs, - pipe_finish_writer: Optional[multiprocessing.connection.Connection], + pipe_finish_writer: Optional[multiprocessing.connection.Connection] = None, launch_callback: Optional[Callable[[], None]] = None, + execute_warmup_func: Callable = _execute_server_warmup, ): if server_args.checkpoint_engine_wait_weights_before_ready: _wait_weights_ready() + + # Send a warmup request if not server_args.skip_server_warmup: - if not _execute_server_warmup( + if not execute_warmup_func( server_args, pipe_finish_writer, ): @@ -1641,6 +1556,7 @@ def _wait_and_warmup( else: _global_state.tokenizer_manager.server_status = ServerStatus.Up + # The server is ready for requests logger.info("The server is fired up and ready to roll!") if pipe_finish_writer is not None: @@ -1675,3 +1591,106 @@ def _wait_weights_ready(): f"Consider increasing SGLANG_WAIT_WEIGHTS_READY_TIMEOUT environment variable. " f"Current status: initial_weights_loaded={_global_state.tokenizer_manager.initial_weights_loaded}" ) + + +def launch_server( + server_args: ServerArgs, + pipe_finish_writer: Optional[multiprocessing.connection.Connection] = None, + launch_subprocesses_func: Callable = _launch_subprocesses, + 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. + """ + tokenizer_manager, template_manager, scheduler_info, port_args = ( + launch_subprocesses_func(server_args=server_args) + ) + + set_global_state( + _GlobalState( + tokenizer_manager=tokenizer_manager, + template_manager=template_manager, + scheduler_info=scheduler_info, + ) + ) + + if server_args.enable_metrics: + add_prometheus_track_response_middleware(app) + + # Pass additional arguments to the lifespan function. + # They will be used for additional initialization setups. + if server_args.tokenizer_worker_num == 1: + # If it is single tokenizer mode, we can pass the arguments by attributes of the app object. + app.is_single_tokenizer_mode = True + app.server_args = server_args + app.warmup_thread_kwargs = dict( + server_args=server_args, + pipe_finish_writer=pipe_finish_writer, + launch_callback=launch_callback, + execute_warmup_func=execute_warmup_func, + ) + + # Add api key authorization + # This is only supported in single tokenizer mode. + if server_args.api_key: + add_api_key_middleware(app, server_args.api_key) + else: + # If it is multi-tokenizer mode, we need to write the arguments to shared memory + # for other worker processes to read. + app.is_single_tokenizer_mode = False + multi_tokenizer_args_shm = write_data_for_multi_tokenizer( + port_args, server_args, scheduler_info + ) + + try: + # Update logging configs + set_uvicorn_logging_configs() + + # Listen for HTTP requests + if server_args.tokenizer_worker_num == 1: + uvicorn.run( + app, + host=server_args.host, + port=server_args.port, + root_path=server_args.fastapi_root_path, + log_level=server_args.log_level_http or server_args.log_level, + timeout_keep_alive=5, + loop="uvloop", + ) + else: + from uvicorn.config import LOGGING_CONFIG + + LOGGING_CONFIG["loggers"]["sglang.srt.entrypoints.http_server"] = { + "handlers": ["default"], + "level": "INFO", + "propagate": False, + } + monkey_patch_uvicorn_multiprocessing() + + uvicorn.run( + "sglang.srt.entrypoints.http_server:app", + host=server_args.host, + port=server_args.port, + root_path=server_args.fastapi_root_path, + log_level=server_args.log_level_http or server_args.log_level, + timeout_keep_alive=5, + loop="uvloop", + workers=server_args.tokenizer_worker_num, + ) + finally: + if server_args.tokenizer_worker_num > 1: + multi_tokenizer_args_shm.unlink() + _global_state.tokenizer_manager.socket_mapping.clear_all_sockets()