sglang 0.5.2rc2__py3-none-any.whl → 0.5.3__py3-none-any.whl

sglang/srt/managers/data_parallel_controller.py

@@ -13,16 +13,15 @@
 # ==============================================================================
 """A controller that dispatches requests to multiple data parallel workers."""
 
+import faulthandler
 import logging
 import multiprocessing as mp
 import signal
-import struct
-import sys
 import threading
 import time
+from collections import deque
 from enum import Enum, auto
-from multiprocessing import shared_memory
-from typing import Dict, List
+from typing import List
 
 import psutil
 import setproctitle
@@ -33,14 +32,19 @@ from sglang.srt.managers.io_struct import (
     BlockReqInput,
     TokenizedEmbeddingReqInput,
     TokenizedGenerateReqInput,
+    WatchLoadUpdateReq,
 )
 from sglang.srt.managers.schedule_batch import Req
 from sglang.srt.managers.scheduler import run_scheduler_process
-from sglang.srt.managers.utils import DPBalanceMeta
 from sglang.srt.server_args import PortArgs, ServerArgs
 from sglang.srt.torch_memory_saver_adapter import TorchMemorySaverAdapter
-from sglang.srt.utils import bind_port, configure_logger, get_zmq_socket
-from sglang.utils import get_exception_traceback
+from sglang.srt.utils import (
+    bind_port,
+    configure_logger,
+    get_zmq_socket,
+    kill_itself_when_parent_died,
+)
+from sglang.utils import TypeBasedDispatcher, get_exception_traceback
 
 logger = logging.getLogger(__name__)
 
@@ -61,18 +65,48 @@ class LoadBalanceMethod(Enum):
             raise ValueError(f"Invalid load balance method: {method}") from exc
 
 
+class DPBudget:
+    def __init__(self):
+        # TODO: support minimum tokens method
+        self.budget_queue = deque()
+
+    def update_budget(self, load_update: WatchLoadUpdateReq):
+        """Update the budget queue.
+        Use num_reqs instead of num_waiting_reqs to balance decode running batch.
+        """
+        loads = load_update.loads
+        self.budget_queue.clear()
+
+        num_reqs = [load.num_reqs for load in loads]
+        if not num_reqs:
+            return
+
+        max_num_reqs = max(num_reqs)
+        if all(x == max_num_reqs for x in num_reqs):
+            return
+
+        while any(x != num_reqs[0] for x in num_reqs):
+            min_load = min(num_reqs)
+            min_indices = [i for i, x in enumerate(num_reqs) if x == min_load]
+            second_min_load = min(x for x in num_reqs if x > min_load)
+            self.budget_queue.extend(
+                [loads[i].dp_rank for i in min_indices] * (second_min_load - min_load)
+            )
+            for idx in min_indices:
+                num_reqs[idx] = second_min_load
+
+    def dispatch(self):
+        if self.budget_queue:
+            return self.budget_queue.popleft()
+        return None
+
+
 class DataParallelController:
     """A controller that dispatches requests to multiple data parallel workers."""
 
-    def __init__(
-        self,
-        server_args: ServerArgs,
-        port_args: PortArgs,
-        dp_balance_meta: DPBalanceMeta,
-    ) -> None:
+    def __init__(self, server_args: ServerArgs, port_args: PortArgs) -> None:
         # for dp balance
         self.global_balance_id = 0
-        self.balance_meta = dp_balance_meta
 
         # Parse args
         self.max_total_num_tokens = None
@@ -98,9 +132,12 @@ class DataParallelController:
         }
         self.dispatching = dispatch_lookup[self.load_balance_method]
 
+        # Load balance budget
+        self.dp_budget = DPBudget()
+
         # Launch data parallel workers
         self.scheduler_procs = []
-        self.workers = [None] * server_args.dp_size
+        self.workers: List[zmq.Socket] = [None] * server_args.dp_size
 
         if server_args.enable_dp_attention:
             dp_port_args = self.launch_dp_attention_schedulers(server_args, port_args)
@@ -121,6 +158,31 @@ class DataParallelController:
 
         self.max_req_input_len = None
 
+        self.init_dispatcher()
+
+    def send_to_all_workers(self, obj):
+        for worker in self.workers:
+            worker.send_pyobj(obj)
+
+    def send_control_message(self, obj):
+        # Send control messages to first worker of tp group
+        for worker in self.workers[:: self.control_message_step]:
+            worker.send_pyobj(obj)
+
+    def handle_load_update_req(self, obj):
+        self.dp_budget.update_budget(obj)
+
+    def init_dispatcher(self):
+        self._request_dispatcher = TypeBasedDispatcher(
+            [
+                (TokenizedGenerateReqInput, self.dispatching),
+                (TokenizedEmbeddingReqInput, self.dispatching),
+                (BlockReqInput, self.send_to_all_workers),
+                (WatchLoadUpdateReq, self.handle_load_update_req),
+            ]
+        )
+        self._request_dispatcher.add_fallback_fn(self.send_control_message)
+
     def launch_dp_schedulers(self, server_args, port_args):
         base_gpu_id = 0
 
@@ -147,7 +209,9 @@ class DataParallelController:
                 args=(server_args, tmp_port_args, base_gpu_id, dp_rank, ready_event),
             )
             threads.append(thread)
-            base_gpu_id += server_args.tp_size * server_args.gpu_id_step
+            base_gpu_id += (
+                server_args.tp_size * server_args.pp_size * server_args.gpu_id_step
+            )
 
         # Free all sockets before starting the threads to launch TP workers
         for sock in sockets:
@@ -250,7 +314,6 @@ class DataParallelController:
                         pp_rank,
                         dp_rank,
                         writer,
-                        self.balance_meta,
                     ),
                 )
                 with memory_saver_adapter.configure_subprocess():
@@ -266,52 +329,43 @@ class DataParallelController:
         self.max_total_num_tokens = scheduler_info[0]["max_total_num_tokens"]
         self.max_req_input_len = scheduler_info[0]["max_req_input_len"]
 
+    def maybe_external_dp_rank_routing(self, req: Req):
+        if req.data_parallel_rank is not None:
+            logger.debug(f"Direct routing to DP rank {req.data_parallel_rank}")
+            self.workers[req.data_parallel_rank].send_pyobj(req)
+            return True
+        return False
+
     def round_robin_scheduler(self, req: Req):
+        if self.maybe_external_dp_rank_routing(req):
+            return
+
         if self.server_args.disaggregation_mode == "null":
-            if req.data_parallel_rank is not None:
-                logger.debug(f"Direct routing to DP rank {req.data_parallel_rank}")
-                self.workers[req.data_parallel_rank].send_pyobj(req)
-            else:
-                self.workers[self.round_robin_counter].send_pyobj(req)
-                self.round_robin_counter = (self.round_robin_counter + 1) % len(
-                    self.workers
-                )
+            self.workers[self.round_robin_counter].send_pyobj(req)
+            self.round_robin_counter = (self.round_robin_counter + 1) % len(
+                self.workers
+            )
         else:
-            if req.data_parallel_rank is not None:
-                logger.debug(f"Direct routing to DP rank {req.data_parallel_rank}")
-                self.workers[req.data_parallel_rank].send_pyobj(req)
-            else:
-                self.workers[req.bootstrap_room % len(self.workers)].send_pyobj(req)
-
-    def shortest_queue_scheduler(self, input_requests):
-        raise NotImplementedError()
+            self.workers[req.bootstrap_room % len(self.workers)].send_pyobj(req)
+
+    def shortest_queue_scheduler(self, req):
+        if self.maybe_external_dp_rank_routing(req):
+            return
+        target_worker = self.dp_budget.dispatch()
+        if target_worker is None:
+            self.round_robin_scheduler(req)
+        else:
+            self.workers[target_worker].send_pyobj(req)
 
     def minimum_tokens_scheduler(self, req):
-        # This variable corresponds to the balance_id in TokenizedGenerateReqInput.
-        # We use it to to control the number of onfly tokens (requests dispatched to workers but not yet received).
-        def get_next_global_balance_id() -> int:
-            INT32_MAX = 2147483647
-            current_id = self.global_balance_id
-            self.global_balance_id = (self.global_balance_id + 1) % INT32_MAX
-            return current_id
-
-        req.dp_balance_id = get_next_global_balance_id()
-        with self.balance_meta.mutex:
-            # 1. local_tokens represents the tokens currently inferring on the worker,
-            #  while onfly refers to the requests dispatched by the dispatcher but not yet received by the scheduler.
-            onfly_info = self.balance_meta.get_shared_onfly()
-            local_tokens = self.balance_meta.get_shared_local_tokens()
-            total_tokens = [
-                local_token + sum(onfly_dict.values())
-                for local_token, onfly_dict in zip(local_tokens, onfly_info)
-            ]
-            target_worker = total_tokens.index(min(total_tokens))
-            onfly_info[target_worker][req.dp_balance_id] = len(req.input_ids)
-            # 2. write the new onfly info to the shm
-            self.balance_meta.set_shared_onfly_info(onfly_info)
+        if self.maybe_external_dp_rank_routing(req):
+            return
 
-        # logger.info(f"dp workers {local_tokens=}, {onfly_info=}, {target_worker=}")
-        self.workers[target_worker].send_pyobj(req)
+        logger.warning(
+            "The 'minimum_tokens' load balancing method is deprecated for now and will introduced later."
+            "Fall back to 'round_robin_scheduler'"
+        )
+        self.round_robin_scheduler(req)
 
     def event_loop(self):
         while True:
@@ -320,22 +374,7 @@ class DataParallelController:
                     recv_req = self.recv_from_tokenizer.recv_pyobj(zmq.NOBLOCK)
                 except zmq.ZMQError:
                     break
-
-                if isinstance(
-                    recv_req,
-                    (
-                        TokenizedGenerateReqInput,
-                        TokenizedEmbeddingReqInput,
-                    ),
-                ):
-                    self.dispatching(recv_req)
-                elif isinstance(recv_req, BlockReqInput):
-                    for worker in self.workers:
-                        worker.send_pyobj(recv_req)
-                else:
-                    # Send other control messages to first worker of tp group
-                    for worker in self.workers[:: self.control_message_step]:
-                        worker.send_pyobj(recv_req)
+                self._request_dispatcher(recv_req)
 
 
 def run_data_parallel_controller_process(
@@ -343,15 +382,14 @@ def run_data_parallel_controller_process(
     port_args: PortArgs,
     pipe_writer,
 ):
+    kill_itself_when_parent_died()
     setproctitle.setproctitle("sglang::data_parallel_controller")
+    faulthandler.enable()
     configure_logger(server_args)
     parent_process = psutil.Process().parent()
-    balance_meta = DPBalanceMeta(server_args.dp_size)
 
     try:
-        controller = DataParallelController(
-            server_args, port_args, dp_balance_meta=balance_meta
-        )
+        controller = DataParallelController(server_args, port_args)
         pipe_writer.send(
             {
                 "status": "ready",
@@ -370,6 +408,3 @@ def run_data_parallel_controller_process(
         traceback = get_exception_traceback()
         logger.error(f"DataParallelController hit an exception: {traceback}")
         parent_process.send_signal(signal.SIGQUIT)
-    finally:
-        # we need to destruct mp.Manager() in balance_meta
-        balance_meta.destructor()

sglang/srt/managers/detokenizer_manager.py

@@ -24,17 +24,16 @@ import psutil
 import setproctitle
 import zmq
 
-from sglang.srt.hf_transformers_utils import get_tokenizer
 from sglang.srt.managers.io_struct import (
-    BatchEmbeddingOut,
+    BatchEmbeddingOutput,
     BatchMultimodalDecodeReq,
-    BatchMultimodalOut,
-    BatchStrOut,
-    BatchTokenIDOut,
+    BatchMultimodalOutput,
+    BatchStrOutput,
+    BatchTokenIDOutput,
     FreezeGCReq,
     MultiTokenizerRegisterReq,
 )
-from sglang.srt.managers.multi_tokenizer_mixin import MultiTokenizerMixin
+from sglang.srt.managers.multi_tokenizer_mixin import MultiHttpWorkerDetokenizerMixin
 from sglang.srt.server_args import PortArgs, ServerArgs
 from sglang.srt.utils import (
     configure_logger,
@@ -42,6 +41,7 @@ from sglang.srt.utils import (
     get_zmq_socket,
     kill_itself_when_parent_died,
 )
+from sglang.srt.utils.hf_transformers_utils import get_tokenizer
 from sglang.utils import (
     TypeBasedDispatcher,
     find_printable_text,
@@ -69,7 +69,7 @@ class DecodeStatus:
     sent_offset: int = 0
 
 
-class DetokenizerManager(MultiTokenizerMixin):
+class DetokenizerManager(MultiHttpWorkerDetokenizerMixin):
     """DetokenizerManager is a process that detokenizes the token ids."""
 
     def __init__(
@@ -101,8 +101,8 @@ class DetokenizerManager(MultiTokenizerMixin):
 
         self._request_dispatcher = TypeBasedDispatcher(
             [
-                (BatchEmbeddingOut, self.handle_batch_embedding_out),
-                (BatchTokenIDOut, self.handle_batch_token_id_out),
+                (BatchEmbeddingOutput, self.handle_batch_embedding_out),
+                (BatchTokenIDOutput, self.handle_batch_token_id_out),
                 (BatchMultimodalDecodeReq, self.handle_multimodal_decode_req),
                 (MultiTokenizerRegisterReq, lambda x: x),
                 (FreezeGCReq, self.handle_freeze_gc_req),
@@ -145,11 +145,11 @@ class DetokenizerManager(MultiTokenizerMixin):
             return output[:-1]
         return output
 
-    def handle_batch_embedding_out(self, recv_obj: BatchEmbeddingOut):
+    def handle_batch_embedding_out(self, recv_obj: BatchEmbeddingOutput):
         # If it is embedding model, no detokenization is needed.
         return recv_obj
 
-    def handle_batch_token_id_out(self, recv_obj: BatchTokenIDOut):
+    def handle_batch_token_id_out(self, recv_obj: BatchTokenIDOutput):
         bs = len(recv_obj.rids)
 
         # Initialize decode status
@@ -224,7 +224,7 @@ class DetokenizerManager(MultiTokenizerMixin):
             s.sent_offset = len(output_str)
             output_strs.append(incremental_output)
 
-        return BatchStrOut(
+        return BatchStrOutput(
             rids=recv_obj.rids,
             finished_reasons=recv_obj.finished_reasons,
             output_strs=output_strs,
@@ -246,17 +246,21 @@ class DetokenizerManager(MultiTokenizerMixin):
             output_token_ids_logprobs_val=recv_obj.output_token_ids_logprobs_val,
             output_token_ids_logprobs_idx=recv_obj.output_token_ids_logprobs_idx,
             output_hidden_states=recv_obj.output_hidden_states,
+            placeholder_tokens_idx=None,
+            placeholder_tokens_val=None,
         )
 
     def handle_multimodal_decode_req(self, recv_obj: BatchMultimodalDecodeReq):
         outputs = self.tokenizer.detokenize(recv_obj)
-        return BatchMultimodalOut(
+        return BatchMultimodalOutput(
             rids=recv_obj.rids,
             finished_reasons=recv_obj.finished_reasons,
             outputs=outputs,
             prompt_tokens=recv_obj.prompt_tokens,
             completion_tokens=recv_obj.completion_tokens,
             cached_tokens=recv_obj.cached_tokens,
+            placeholder_tokens_idx=None,
+            placeholder_tokens_val=None,
         )
 
     def handle_freeze_gc_req(self, recv_req: FreezeGCReq):
@@ -289,11 +293,11 @@ def run_detokenizer_process(
     try:
         manager = DetokenizerManager(server_args, port_args)
         if server_args.tokenizer_worker_num > 1:
-            manager.multi_tokenizer_manager_event_loop()
+            manager.multi_http_worker_event_loop()
         else:
             manager.event_loop()
     except Exception:
-        manager.clear_tokenizer_mapping()
+        manager.maybe_clear_socket_mapping()
         traceback = get_exception_traceback()
         logger.error(f"DetokenizerManager hit an exception: {traceback}")
         parent_process.send_signal(signal.SIGQUIT)

sglang/srt/managers/disagg_service.py

@@ -0,0 +1,46 @@
+"""Start bootstrap/kv-store-related server"""
+
+import os
+from typing import Type
+
+from sglang.srt.disaggregation.base import BaseKVBootstrapServer
+from sglang.srt.disaggregation.utils import (
+    DisaggregationMode,
+    KVClassType,
+    TransferBackend,
+    get_kv_class,
+)
+from sglang.srt.server_args import ServerArgs
+
+
+def start_disagg_service(
+    server_args: ServerArgs,
+):
+    # Start kv boostrap server on prefill
+    disagg_mode = DisaggregationMode(server_args.disaggregation_mode)
+    transfer_backend = TransferBackend(server_args.disaggregation_transfer_backend)
+
+    if disagg_mode == DisaggregationMode.PREFILL:
+        # only start bootstrap server on prefill tm
+        kv_bootstrap_server_class: Type[BaseKVBootstrapServer] = get_kv_class(
+            transfer_backend, KVClassType.BOOTSTRAP_SERVER
+        )
+        bootstrap_server: BaseKVBootstrapServer = kv_bootstrap_server_class(
+            host=server_args.host,
+            port=server_args.disaggregation_bootstrap_port,
+        )
+        is_create_store = (
+            server_args.node_rank == 0 and transfer_backend == TransferBackend.ASCEND
+        )
+        if is_create_store:
+            try:
+                from mf_adapter import create_config_store
+
+                ascend_url = os.getenv("ASCEND_MF_STORE_URL")
+                create_config_store(ascend_url)
+            except Exception as e:
+                error_message = f"Failed create mf store, invalid ascend_url."
+                error_message += f" With exception {e}"
+                raise error_message
+
+        return bootstrap_server