diff --git a/python/sglang/srt/custom_op.py b/python/sglang/srt/custom_op.py index ea3c06e6d..90cfa8e12 100644 --- a/python/sglang/srt/custom_op.py +++ b/python/sglang/srt/custom_op.py @@ -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 ( diff --git a/python/sglang/srt/managers/detokenizer_manager.py b/python/sglang/srt/managers/detokenizer_manager.py index 0e64a80b5..951439b15 100644 --- a/python/sglang/srt/managers/detokenizer_manager.py +++ b/python/sglang/srt/managers/detokenizer_manager.py @@ -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: diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 141d95d5d..bf8db9bcf 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -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 diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py index 9be7d4381..455883fcc 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -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"