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

Clean up the __init__ of TokenizerManager and DetokenizerManager (#15796)

Browse Source
This commit is contained in:
Lianmin Zheng
2025-12-25 10:30:27 -08:00
committed by GitHub
parent 6ef543f99d
commit 8087ef126f
4 changed files with 130 additions and 54 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -1,3 +1,9 @@
"""
The definition of CustomOps for multi hardware dispatching.
TODO: Move this to python/sglang/srt/layers/custom_op.py
"""
from torch import nn
from sglang.srt.utils import (

31
python/sglang/srt/managers/detokenizer_manager.py
View File

@@ -79,6 +79,18 @@ class DetokenizerManager(MultiHttpWorkerDetokenizerMixin):
port_args: PortArgs,
):
# Init inter-process communication
self.init_ipc_channels(port_args)
# Init tokenizer
self.init_tokenizer(server_args)
# Init running status
self.init_running_status(server_args)
# Init dispatcher
self.init_request_dispatcher()
def init_ipc_channels(self, port_args: PortArgs):
context = zmq.Context(2)
self.recv_from_scheduler = get_zmq_socket(
context, zmq.PULL, port_args.detokenizer_ipc_name, True
@@ -87,7 +99,7 @@ class DetokenizerManager(MultiHttpWorkerDetokenizerMixin):
context, zmq.PUSH, port_args.tokenizer_ipc_name, False
)
# Init tokenizer
def init_tokenizer(self, server_args: ServerArgs):
if server_args.skip_tokenizer_init:
self.tokenizer = None
else:
@@ -98,12 +110,20 @@ class DetokenizerManager(MultiHttpWorkerDetokenizerMixin):
revision=server_args.revision,
)
def init_running_status(self, server_args: ServerArgs):
self.decode_status = LimitedCapacityDict(capacity=DETOKENIZER_MAX_STATES)
self.is_dummy = False
self.is_tool_call_parser_gpt_oss = server_args.tool_call_parser == "gpt-oss"
self.disable_tokenizer_batch_decode = server_args.disable_tokenizer_batch_decode
# Init dispatcher
self.watchdog = Watchdog.create(
debug_name="DetokenizerManager",
watchdog_timeout=server_args.soft_watchdog_timeout,
soft=True,
test_stuck_time=envs.SGLANG_TEST_STUCK_DETOKENIZER.get(),
)
def init_request_dispatcher(self):
self._request_dispatcher = TypeBasedDispatcher(
[
(BatchEmbeddingOutput, self.handle_batch_embedding_out),
@@ -113,13 +133,6 @@ class DetokenizerManager(MultiHttpWorkerDetokenizerMixin):
]
)
self.watchdog = Watchdog.create(
debug_name="DetokenizerManager",
watchdog_timeout=server_args.soft_watchdog_timeout,
soft=True,
test_stuck_time=envs.SGLANG_TEST_STUCK_DETOKENIZER.get(),
)
def event_loop(self):
"""The event loop that handles requests"""
while True:

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

@@ -317,7 +317,7 @@ class Scheduler(
self.init_metrics(tp_rank, pp_rank, dp_rank)
# Init inter-process communication
self.init_sockets(server_args, port_args)
self.init_ipc_channels(port_args)
# Init PD-multiplexing context
if self.enable_pdmux:
@@ -373,7 +373,7 @@ class Scheduler(
else None
)
def init_sockets(self, server_args: ServerArgs, port_args: PortArgs):
def init_ipc_channels(self, port_args: PortArgs):
context = zmq.Context(2)
self.idle_sleeper = None
@@ -388,7 +388,7 @@ class Scheduler(
send_to_tokenizer = get_zmq_socket(
context, zmq.PUSH, port_args.tokenizer_ipc_name, False
)
if server_args.skip_tokenizer_init:
if self.server_args.skip_tokenizer_init:
# Directly send to the TokenizerManager
send_to_detokenizer = get_zmq_socket(
context, zmq.PUSH, port_args.tokenizer_ipc_name, False

141
python/sglang/srt/managers/tokenizer_manager.py
View File

@@ -181,11 +181,6 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerMultiItemMixi
# Parse args
self.server_args = server_args
self.enable_metrics = server_args.enable_metrics
self.request_logger = RequestLogger(
log_requests=server_args.log_requests,
log_requests_level=server_args.log_requests_level,
log_requests_format=server_args.log_requests_format,
)
self.preferred_sampling_params = server_args.preferred_sampling_params
self.crash_dump_folder = server_args.crash_dump_folder
self.enable_trace = server_args.enable_trace
@@ -210,6 +205,36 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerMultiItemMixi
set_global_server_args_for_tokenizer(server_args)
# Initialize tokenizer and multimodalprocessor
self.init_tokenizer_and_processor()
# Init inter-process communication
self.init_ipc_channels(port_args)
# Init running status
self.init_running_status()
# Init logging and dumping
self.init_request_logging_and_dumping()
# Init weight update
self.init_weight_update()
# Init LoRA status
self.init_lora()
# Init PD disaggregation and encoder disaggregation
self.init_disaggregation()
# Init metric collector and watchdog
self.init_metric_collector_watchdog()
# Init request dispatcher
self.init_request_dispatcher()
def init_tokenizer_and_processor(self):
server_args = self.server_args
# Initialize tokenizer and processor
if self.model_config.is_multimodal:
import_processors("sglang.srt.multimodal.processors")
@@ -264,7 +289,7 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerMultiItemMixi
else:
self.async_dynamic_batch_tokenizer = None
# Init inter-process communication
def init_ipc_channels(self, port_args: PortArgs):
context = zmq.asyncio.Context(2)
self.recv_from_detokenizer = get_zmq_socket(
context, zmq.PULL, port_args.tokenizer_ipc_name, True
@@ -284,13 +309,7 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerMultiItemMixi
# Make sure that each request carries the tokenizer_ipc_name for response routing
self.send_to_scheduler = SenderWrapper(port_args, send_to_scheduler)
# E Disaggregation
if self.server_args.language_only:
self.mm_receiver = MMReceiver(
server_args,
dtype=self.model_config.dtype,
)
def init_running_status(self):
# Request states
self._chosen_loop = None
self.rid_to_state: Dict[str, ReqState] = {}
@@ -301,10 +320,17 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerMultiItemMixi
self.gracefully_exit = False
self.last_receive_tstamp = 0
# Initial weights status
self.initial_weights_loaded = True
if server_args.checkpoint_engine_wait_weights_before_ready:
self.initial_weights_loaded = False
# For load balancing
self.current_load = 0
self.current_load_lock = asyncio.Lock()
def init_request_logging_and_dumping(self):
# Request logging
self.request_logger = RequestLogger(
log_requests=self.server_args.log_requests,
log_requests_level=self.server_args.log_requests_level,
log_requests_format=self.server_args.log_requests_format,
)
# Dumping
self.dump_requests_folder = "" # By default do not dump
@@ -319,6 +345,12 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerMultiItemMixi
self.server_args, obj_skip_names, out_skip_names
)
def init_weight_update(self):
# Initial weights status
self.initial_weights_loaded = True
if self.server_args.checkpoint_engine_wait_weights_before_ready:
self.initial_weights_loaded = False
# Session
self.session_futures = {} # session_id -> asyncio event
@@ -331,6 +363,7 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerMultiItemMixi
self.is_pause = False
self.is_pause_cond = asyncio.Condition()
def init_lora(self):
# LoRA
# Initialize the `LoRARegistry` with initial LoRA adapter paths provided in `server_args`.
# The registry dynamically updates as adapters are loaded / unloaded during runtime. It
@@ -349,27 +382,32 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerMultiItemMixi
for lora_ref in self.server_args.lora_paths:
self.lora_ref_cache[lora_ref.lora_name] = lora_ref
# Disaggregation
def init_disaggregation(self):
# PD Disaggregation
self.disaggregation_mode = DisaggregationMode(
self.server_args.disaggregation_mode
)
self.bootstrap_server = start_disagg_service(self.server_args)
# For load balancing
self.current_load = 0
self.current_load_lock = asyncio.Lock()
# Encoder Disaggregation
if self.server_args.language_only:
self.mm_receiver = MMReceiver(
self.server_args,
dtype=self.model_config.dtype,
)
def init_metric_collector_watchdog(self):
# Metrics
if self.enable_metrics:
labels = {
"model_name": self.server_args.served_model_name,
# TODO: Add lora name/path in the future,
}
if server_args.tokenizer_metrics_allowed_custom_labels:
for label in server_args.tokenizer_metrics_allowed_custom_labels:
if self.server_args.tokenizer_metrics_allowed_custom_labels:
for label in self.server_args.tokenizer_metrics_allowed_custom_labels:
labels[label] = ""
self.metrics_collector = TokenizerMetricsCollector(
server_args=server_args,
server_args=self.server_args,
labels=labels,
bucket_time_to_first_token=self.server_args.bucket_time_to_first_token,
bucket_e2e_request_latency=self.server_args.bucket_e2e_request_latency,
@@ -377,11 +415,16 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerMultiItemMixi
collect_tokens_histogram=self.server_args.collect_tokens_histogram,
)
# Configure GC warning
if self.server_args.gc_warning_threshold_secs > 0.0:
configure_gc_warning(self.server_args.gc_warning_threshold_secs)
self.watchdog = Watchdog.create(
debug_name="TokenizerManager",
watchdog_timeout=self.server_args.soft_watchdog_timeout,
soft=True,
test_stuck_time=envs.SGLANG_TEST_STUCK_TOKENIZER.get(),
)
# Dispatcher and communicators
def init_request_dispatcher(self):
self._result_dispatcher = TypeBasedDispatcher(
[
(
@@ -404,35 +447,26 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerMultiItemMixi
(HealthCheckOutput, lambda x: None),
]
)
self.init_communicators(server_args)
self.watchdog = Watchdog.create(
debug_name="TokenizerManager",
watchdog_timeout=server_args.soft_watchdog_timeout,
soft=True,
test_stuck_time=envs.SGLANG_TEST_STUCK_TOKENIZER.get(),
)
self.init_communicators(self.server_args)
async def generate_request(
self,
obj: Union[GenerateReqInput, EmbeddingReqInput],
request: Optional[fastapi.Request] = None,
trace_parent: Optional[str] = None,
):
created_time = obj.received_time if obj.received_time else time.time()
self.auto_create_handle_loop()
obj.normalize_batch_and_arguments()
if (
self.server_args.language_only
and isinstance(obj, GenerateReqInput)
and self.server_args.encoder_transfer_backend == "zmq_to_scheduler"
and obj.contains_mm_input()
):
self.mm_receiver.send_encode_request(obj)
if self.enable_trace:
self._trace_request_start(obj, created_time, request)
self._trace_request_start(obj, created_time, request, trace_parent)
if self.server_args.language_only:
self._handle_epd_disaggregation_encode_request(obj)
if self.server_args.tokenizer_worker_num > 1:
self._attach_multi_http_worker_info(obj)
# Log the request
self.request_logger.log_received_request(obj, self.tokenizer)
async with self.is_pause_cond:
@@ -1477,6 +1511,7 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerMultiItemMixi
"Signal SIGTERM received while health check failed. Force exiting."
)
self.dump_requests_before_crash()
self.force_exit_handler()
break
elif get_bool_env_var("SGL_FORCE_SHUTDOWN"):
@@ -1484,6 +1519,7 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerMultiItemMixi
logger.error(
"Signal SIGTERM received while force shutdown flag set. Force exiting."
)
self.force_exit_handler()
break
logger.info(
@@ -1498,6 +1534,10 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerMultiItemMixi
kill_process_tree(os.getpid(), include_parent=True)
sys.exit(0)
def force_exit_handler(self):
"""Put some custom force exit logic here."""
pass
async def handle_loop(self):
"""The event loop that handles requests"""
while True:
@@ -2137,6 +2177,7 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerMultiItemMixi
obj: Union[GenerateReqInput, EmbeddingReqInput],
created_time: Optional[float] = None,
request: Optional[fastapi.Request] = None,
trace_parent: Optional[str] = None,
):
external_trace_header = None
if request:
@@ -2144,6 +2185,11 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerMultiItemMixi
trace_set_remote_propagate_context(request.headers["trace_context"])
else:
external_trace_header = extract_trace_headers(request.headers)
elif trace_parent:
# When the request comes form the rust grpc server there isn't a
# real request object but we still need to propagate the traceparent from
# the traceparent that is explicitly passed in
external_trace_header = {"trace_parent": trace_parent}
if obj.is_single:
bootstrap_room = (
@@ -2175,6 +2221,17 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerMultiItemMixi
"", obj.rid[i], ts=int(created_time * 1e9), anonymous=True
)
def _handle_epd_disaggregation_encode_request(
self, obj: Union[GenerateReqInput, EmbeddingReqInput]
):
"""Handle EPD-disaggregation mode encoding request."""
if (
isinstance(obj, GenerateReqInput)
and self.server_args.encoder_transfer_backend == "zmq_to_scheduler"
and obj.contains_mm_input()
):
self.mm_receiver.send_encode_request(obj)
class ServerStatus(Enum):
Up = "Up"