support MPServer and embedded server for granian to enable muti tokenizer worker (#28573)

This commit is contained in:
Rain Jiang
2026-06-18 17:59:12 -07:00
committed by GitHub
parent ea407df4b0
commit ef01618dfb
5 changed files with 131 additions and 128 deletions
+106 -108
View File
@@ -207,32 +207,6 @@ def get_global_state() -> _GlobalState:
return _global_state
async def _init_granian_worker() -> ServerArgs:
main_pid = get_main_process_id()
port_args, server_args, scheduler_info = read_from_shared_memory(
f"multi_tokenizer_args_{main_pid}"
)
tokenizer_manager = TokenizerManager(server_args, port_args)
template_manager = TemplateManager()
template_manager.initialize_templates(
tokenizer_manager=tokenizer_manager,
model_path=server_args.model_path,
chat_template=server_args.chat_template,
completion_template=server_args.completion_template,
)
tokenizer_manager.max_req_input_len = scheduler_info["max_req_input_len"]
set_global_state(
_GlobalState(
tokenizer_manager=tokenizer_manager,
template_manager=template_manager,
scheduler_info=scheduler_info,
)
)
return server_args
async def init_multi_tokenizer() -> ServerArgs:
"""
Initialization function for multi-process tokenizer mode.
@@ -290,10 +264,6 @@ async def lifespan(fast_api_app: FastAPI):
server_args = fast_api_app.server_args
warmup_thread_kwargs = fast_api_app.warmup_thread_kwargs
thread_label = "Tokenizer"
elif envs.SGLANG_GRANIAN_PARENT_PID.get() is not None:
server_args = await _init_granian_worker()
warmup_thread_kwargs = dict(server_args=server_args)
thread_label = "Tokenizer"
else:
# Initialize multi-tokenizer support for worker processes
server_args = await init_multi_tokenizer()
@@ -2209,51 +2179,77 @@ def _wait_weights_ready():
)
def _close_main_process_sockets():
"""Close the main process's ZMQ sockets before spawning Granian workers.
def _run_granian_server(
host,
port,
log_level,
tokenizer_worker_num=1,
ssl_certfile=None,
ssl_keyfile=None,
ssl_ca_certs=None,
ssl_keyfile_password=None,
ssl_verify=False, # MTls is not supported
backlog=2048,
backpressure=2048,
):
"""Serve the in-process ASGI app with Granian (embedded mode) over HTTP/2.
Granian workers create their own TokenizerManager with fresh ZMQ sockets.
The main process must release its sockets first to avoid binding conflicts
on the same IPC addresses.
Unlike Granian's default multi-process server, the embedded server runs a
single worker as an asyncio task inside the current process. It therefore
serves the live ``app`` object directly and reuses the already-initialized
global state (tokenizer manager, templates, ...) through the normal
single-tokenizer lifespan path -- no shared memory or worker re-init needed.
The event loop is uvloop. The default backlog and backpressure values are set
exactly like uvicorn's defaults.
"""
if _global_state is None or _global_state.tokenizer_manager is None:
return
tm = _global_state.tokenizer_manager
for attr in ("recv_from_detokenizer", "send_to_scheduler"):
sock = getattr(tm, attr, None)
if sock is None:
continue
inner = getattr(sock, "socket", None)
if inner is not None:
inner.close()
elif hasattr(sock, "close"):
sock.close()
setattr(tm, attr, None)
import signal
def _run_granian_server(server_args: ServerArgs):
"""Launch Granian with HTTP/2 support"""
from granian import Granian
from granian.constants import HTTPModes, Interfaces, Loops
from granian.server.embed import Server as GranianEmbeddedServer
Server = GranianEmbeddedServer if tokenizer_worker_num == 1 else Granian
target = (
app if tokenizer_worker_num == 1 else "sglang.srt.entrypoints.http_server:app"
)
granian_kwargs = dict(
target="sglang.srt.entrypoints.http_server:app",
address=server_args.host,
port=server_args.port,
target=target,
address=host,
port=port,
interface=Interfaces.ASGI,
http=HTTPModes.auto,
loop=Loops.uvloop,
log_level=server_args.log_level_http or server_args.log_level or "info",
workers=1,
log_level=log_level,
ssl_cert=ssl_certfile,
ssl_key=ssl_keyfile,
ssl_key_password=ssl_keyfile_password,
ssl_ca=ssl_ca_certs,
ssl_client_verify=ssl_verify,
backlog=backlog,
backpressure=backpressure,
)
ssl_enabled = server_args.ssl_certfile and server_args.ssl_keyfile
if ssl_enabled:
granian_kwargs["ssl_cert"] = server_args.ssl_certfile
granian_kwargs["ssl_key"] = server_args.ssl_keyfile
if tokenizer_worker_num > 1:
granian_kwargs["workers"] = tokenizer_worker_num
granian_kwargs["loop"] = Loops.uvloop
server = Granian(**granian_kwargs)
server.serve()
server = Server(**granian_kwargs)
if tokenizer_worker_num == 1:
async def serve():
# The embedded server does not install its own signal handlers, so wire
# SIGINT/SIGTERM to a graceful stop, mirroring uvicorn's behavior.
loop = asyncio.get_running_loop()
for sig in (signal.SIGINT, signal.SIGTERM):
try:
loop.add_signal_handler(sig, server.stop)
except (NotImplementedError, ValueError):
pass
await server.serve()
uvloop.run(serve())
else:
server.serve()
def _setup_and_run_http_server(
@@ -2286,35 +2282,6 @@ def _setup_and_run_http_server(
if server_args.enable_metrics:
add_prometheus_track_response_middleware(app)
# Use Granian for HTTP/2 server
if server_args.enable_http2:
# Reuse the multi-tokenizer shared memory mechanism to pass
# init args (port_args, server_args, scheduler_info) to
# Granian workers, which are independent processes.
multi_tokenizer_args_shm = write_data_for_multi_tokenizer(
port_args, server_args, scheduler_infos[0]
)
try:
if server_args.ssl_certfile:
logger.info(
f"SSL enabled: certfile={server_args.ssl_certfile}, "
f"keyfile={server_args.ssl_keyfile}"
)
logger.info(
f"Starting Granian HTTP/2 server on "
f"{server_args.host}:{server_args.port}"
)
# Propagate the main process PID via os.environ so Granian
# workers (forked or spawned) can locate the shared memory
# segment created above.
envs.SGLANG_GRANIAN_PARENT_PID.set(os.getpid())
_close_main_process_sockets()
_run_granian_server(server_args)
finally:
if multi_tokenizer_args_shm is not None:
multi_tokenizer_args_shm.unlink()
return
# Pass additional arguments to the lifespan function.
# They will be used for additional initialization setups.
if server_args.tokenizer_worker_num == 1:
@@ -2366,7 +2333,22 @@ def _setup_and_run_http_server(
# Listen for HTTP requests
if server_args.tokenizer_worker_num == 1:
if server_args.enable_ssl_refresh:
if server_args.enable_http2:
logger.info(
f"Starting embedded Granian HTTP/2 server on "
f"{server_args.host}:{server_args.port}"
)
_run_granian_server(
host=server_args.host,
port=server_args.port,
log_level=server_args.log_level_http or server_args.log_level,
ssl_certfile=server_args.ssl_certfile,
ssl_keyfile=server_args.ssl_keyfile,
ssl_ca_certs=server_args.ssl_ca_certs,
ssl_keyfile_password=server_args.ssl_keyfile_password,
ssl_verify=False, # No MTLS supported for now.
)
elif server_args.enable_ssl_refresh:
# Use Config/Server API for access to the SSLContext.
config = uvicorn.Config(
app,
@@ -2435,21 +2417,37 @@ def _setup_and_run_http_server(
"SSL refresh will be disabled."
)
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=envs.SGLANG_TIMEOUT_KEEP_ALIVE.get(),
timeout_worker_healthcheck=envs.SGLANG_UVICORN_WORKER_HEALTHCHECK_TIMEOUT.get(),
loop="uvloop",
workers=server_args.tokenizer_worker_num,
ssl_keyfile=server_args.ssl_keyfile,
ssl_certfile=server_args.ssl_certfile,
ssl_ca_certs=server_args.ssl_ca_certs,
ssl_keyfile_password=server_args.ssl_keyfile_password,
)
if server_args.enable_http2:
logger.info(
f"Starting embedded Granian HTTP/2 server on "
f"{server_args.host}:{server_args.port}"
)
_run_granian_server(
host=server_args.host,
port=server_args.port,
log_level=server_args.log_level_http or server_args.log_level,
tokenizer_worker_num=server_args.tokenizer_worker_num,
ssl_certfile=server_args.ssl_certfile,
ssl_keyfile=server_args.ssl_keyfile,
ssl_ca_certs=server_args.ssl_ca_certs,
ssl_keyfile_password=server_args.ssl_keyfile_password,
)
else:
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=envs.SGLANG_TIMEOUT_KEEP_ALIVE.get(),
timeout_worker_healthcheck=envs.SGLANG_UVICORN_WORKER_HEALTHCHECK_TIMEOUT.get(),
loop="uvloop",
workers=server_args.tokenizer_worker_num,
ssl_keyfile=server_args.ssl_keyfile,
ssl_certfile=server_args.ssl_certfile,
ssl_ca_certs=server_args.ssl_ca_certs,
ssl_keyfile_password=server_args.ssl_keyfile_password,
)
finally:
if server_args.tokenizer_worker_num > 1:
if multi_tokenizer_args_shm is not None:
-3
View File
@@ -708,9 +708,6 @@ class Envs:
# too short when many workers cold-start and load tokenizers in parallel.
SGLANG_UVICORN_WORKER_HEALTHCHECK_TIMEOUT = EnvInt(10)
# HTTP/2 Server
SGLANG_GRANIAN_PARENT_PID = EnvInt(None)
# Health Check
SGLANG_ENABLE_HEALTH_ENDPOINT_GENERATION = EnvBool(True)
@@ -684,16 +684,9 @@ async def print_exception_wrapper(func):
def get_main_process_id() -> int:
"""Get the main process ID.
Supports override via SGLANG_GRANIAN_PARENT_PID for workers whose
multiprocessing parent PID differs from the shared-memory owner.
"""
from sglang.srt.environ import envs
override = envs.SGLANG_GRANIAN_PARENT_PID.get()
if override is not None:
return override
Get the main process ID.
"""
return multiprocessing.current_process()._parent_pid
-6
View File
@@ -1178,12 +1178,6 @@ class ServerArgs:
"Use Uvicorn (the default) or handle certificate rotation externally."
)
if self.tokenizer_worker_num > 1:
raise ValueError(
"--enable-http2 does not yet support --tokenizer-worker-num > 1. "
"Multi-worker HTTP/2 support will be added in a future release."
)
def _handle_multimodal(self):
"""Validate mm_process_config structure before model loading."""
if self.mm_process_config is not None: