support MPServer and embedded server for granian to enable muti tokenizer worker (#28573)
This commit is contained in:
Rain Jiang
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user