chore: import upstream snapshot with attribution

python/tokenspeed/runtime/engine/data_parallel_controller.py

@@ -0,0 +1,381 @@
+# Copyright (c) 2026 LightSeek Foundation
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+"""A controller that dispatches requests to multiple data parallel workers."""
+
+import copy
+import multiprocessing as mp
+import os
+import signal
+import threading
+from collections import deque
+from enum import Enum, auto
+
+import psutil
+import setproctitle
+import zmq
+
+from tokenspeed.runtime.engine.event_loop import run_event_loop
+from tokenspeed.runtime.engine.io_struct import (
+    BlockReqInput,
+    TokenizedEmbeddingReqInput,
+    TokenizedGenerateReqInput,
+    WatchLoadUpdateReq,
+)
+from tokenspeed.runtime.engine.request import Req
+from tokenspeed.runtime.utils import (
+    configure_logger,
+    get_colorful_logger,
+    get_zmq_socket,
+)
+from tokenspeed.runtime.utils.dispatch import TypeBasedDispatcher
+from tokenspeed.runtime.utils.exceptions import get_exception_traceback
+from tokenspeed.runtime.utils.process import register_usr_signal
+from tokenspeed.runtime.utils.server_args import PortArgs, ServerArgs
+
+logger = get_colorful_logger(__name__)
+
+
+class LoadBalanceMethod(Enum):
+    """Load balance method."""
+
+    ROUND_ROBIN = auto()
+    SHORTEST_QUEUE = auto()
+    MINIMUM_CACHE_USAGE = auto()
+
+    @classmethod
+    def from_str(cls, method: str):
+        method = method.upper()
+        try:
+            return cls[method]
+        except KeyError as exc:
+            raise ValueError(f"Invalid load balance method: {method}") from exc
+
+
+class DPBudget:
+    def __init__(self, method: LoadBalanceMethod):
+        # Use different metrics for load balancing:
+        # - SHORTEST_QUEUE: by num_reqs (running + waiting)
+        # - MINIMUM_CACHE_USAGE: by num_pages (page usage)
+        self.method = method
+        self.budget_queue = deque()
+
+    def update_budget(self, load_update: WatchLoadUpdateReq):
+        """Update the budget queue.
+
+        For SHORTEST_QUEUE, use num_reqs instead of num_waiting_reqs to balance decode running batch.
+        For MINIMUM_CACHE_USAGE, use num_pages as cache usage metric.
+        """
+        # method update_budget and method dispatch happen in the same thread, so clearing budget_queue is safe
+        self.budget_queue.clear()
+
+        loads = load_update.loads
+        if not loads:
+            return
+
+        if self.method == LoadBalanceMethod.MINIMUM_CACHE_USAGE:
+            metrics = [load.num_pages for load in loads]
+        else:
+            metrics = [load.num_reqs for load in loads]
+
+        max_metric = max(metrics)
+        if all(x == max_metric for x in metrics):
+            return
+
+        while any(x != metrics[0] for x in metrics):
+            min_load = min(metrics)
+            min_indices = [i for i, x in enumerate(metrics) if x == min_load]
+            second_min_load = min(x for x in metrics 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:
+                metrics[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) -> None:
+        # Parse args
+        self.max_total_num_tokens = None
+        self.max_req_input_len = None
+        self.max_num_seqs = None
+        self.chunked_prefill_size = None
+        self.max_model_len = None
+        self.server_args = server_args
+        self.port_args = port_args
+        self.load_balance_method = LoadBalanceMethod.from_str(
+            server_args.load_balance_method
+        )
+
+        # Init inter-process communication
+        self.context = zmq.Context(1 + server_args.mapping.attn.dp_size)
+        if server_args.node_rank == 0:
+            self.recv_from_tokenizer = get_zmq_socket(
+                self.context, zmq.PULL, port_args.scheduler_input_ipc_name, False
+            )
+        # dp_worker for fixed data dispatch can be set by SINGLE_WORKER_ID environment variable
+        robin_scheduler = (
+            self.round_robin_scheduler
+            if os.environ.get("SINGLE_WORKER_ID", "-1") == "-1"
+            else self.single_robin_scheduler
+        )
+        # Dispatch method
+        self.round_robin_counter = 0
+        dispatch_lookup = {
+            LoadBalanceMethod.ROUND_ROBIN: robin_scheduler,
+            LoadBalanceMethod.SHORTEST_QUEUE: self.budget_scheduler,
+            LoadBalanceMethod.MINIMUM_CACHE_USAGE: self.budget_scheduler,
+        }
+        self.dispatching = dispatch_lookup[self.load_balance_method]
+
+        # Load balance budget
+        self.dp_budget = DPBudget(self.load_balance_method)
+
+        # Launch data parallel workers
+        self.scheduler_procs = []
+        self.workers = [None] * server_args.mapping.attn.dp_size
+
+        self.launch_dp_schedulers(server_args, port_args)
+
+        # Workers are already created in launch_dp_schedulers before starting scheduler threads
+
+        if server_args.mapping.has_attn_dp:
+            self.control_message_step = server_args.mapping.attn.tp_size
+        else:
+            self.control_message_step = 1
+
+        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):
+        threads = []
+        dp_port_args = []
+
+        # Parse dist_init_addr from port_args to create per-dp-rank ports
+        # Extract base info from the passed port_args
+        base_scheduler_port = int(port_args.scheduler_input_ipc_name.split(":")[-1])
+        dist_init_host = port_args.scheduler_input_ipc_name.split("//")[1].split(":")[0]
+
+        # port_args.scheduler_input_ipc_name (base_scheduler_port) is used by:
+        # TokenizerManager -> DataParallelController
+        #
+        # For DataParallelController -> Scheduler[dp_rank], we need different ports.
+        # Following the same logic as PortArgs.init_new with dp_rank parameter:
+        # scheduler_input_port = port_base + 4 + dp_rank
+        # Since base_scheduler_port = port_base + 4, we have:
+        # scheduler_input_port = base_scheduler_port + dp_rank
+        #
+        # But we need to avoid conflict with TokenizerManager's port (base_scheduler_port).
+        # So we start from base_scheduler_port + 1 for dp_rank=0.
+
+        for dp_rank in range(server_args.mapping.attn.dp_size):
+            # Create port_args for each dp_rank by adjusting scheduler_input_port
+            # This avoids calling PortArgs.init_new which might use default port
+            # Use base_scheduler_port + 1 + dp_rank to avoid conflict with TokenizerManager
+            scheduler_input_port = base_scheduler_port + 1 + dp_rank
+            tmp_port_args = PortArgs(
+                tokenizer_ipc_name=port_args.tokenizer_ipc_name,
+                scheduler_input_ipc_name=f"tcp://{dist_init_host}:{scheduler_input_port}",
+                nccl_port=port_args.nccl_port,
+                rpc_ipc_name=port_args.rpc_ipc_name,
+                metrics_ipc_name=port_args.metrics_ipc_name,
+                tokenizer_worker_ipc_name=port_args.tokenizer_worker_ipc_name,
+            )
+            dp_port_args.append(tmp_port_args)
+
+            # Bind to scheduler_input_ipc_name BEFORE starting scheduler threads
+            # This ensures the port is available when scheduler tries to connect
+            if server_args.node_rank == 0:
+                self.workers[dp_rank] = get_zmq_socket(
+                    self.context,
+                    zmq.PUSH,
+                    tmp_port_args.scheduler_input_ipc_name,
+                    True,  # bind
+                )
+
+        if not server_args.mapping.attn.has_dp:
+            dp_rank_range = range(0, 1)
+        else:
+            dp_ranks_per_node = (
+                server_args.mapping.attn.dp_size // server_args.mapping.nnodes
+            )
+            dp_rank_range = range(
+                dp_ranks_per_node * server_args.node_rank,
+                dp_ranks_per_node * (server_args.node_rank + 1),
+            )
+        for dp_rank in dp_rank_range:
+            # Create a thread for each worker
+            thread = threading.Thread(
+                target=self.launch_tensor_parallel_group,
+                args=(server_args, dp_port_args[dp_rank], dp_rank),
+            )
+            threads.append(thread)
+
+        # Start all threads
+        for thread in threads:
+            thread.start()
+        for thread in threads:
+            thread.join()
+
+        return dp_port_args
+
+    def launch_tensor_parallel_group(
+        self,
+        server_args: ServerArgs,
+        port_args: PortArgs,
+        dp_rank: int,
+    ):
+        scheduler_pipe_readers = []
+        mapping_template = server_args.mapping
+        attn_tp_size = mapping_template.attn.tp_size
+
+        if attn_tp_size > mapping_template.nprocs_per_node:
+            attn_tp_ranks_per_node = attn_tp_size // mapping_template.nnodes
+            attn_tp_rank_range = range(
+                attn_tp_ranks_per_node * server_args.node_rank,
+                attn_tp_ranks_per_node * (server_args.node_rank + 1),
+            )
+        else:
+            attn_tp_rank_range = range(0, attn_tp_size)
+        for attn_tp_rank in attn_tp_rank_range:
+            reader, writer = mp.Pipe(duplex=False)
+            global_rank = dp_rank * attn_tp_size + attn_tp_rank
+
+            # Create per-rank server_args with rank-initialized mapping
+            rank_server_args = copy.copy(server_args)
+            rank_server_args.mapping = copy.deepcopy(mapping_template)
+            rank_server_args.mapping.rank = global_rank
+
+            proc = mp.Process(
+                target=run_event_loop,
+                args=(
+                    rank_server_args,
+                    port_args,
+                    writer,
+                ),
+            )
+            proc.start()
+            self.scheduler_procs.append(proc)
+            scheduler_pipe_readers.append(reader)
+        # Wait for model to finish loading
+        scheduler_info = [reader.recv() for reader in scheduler_pipe_readers]
+
+        self.max_total_num_tokens = scheduler_info[0]["max_total_num_tokens"]
+        self.max_req_input_len = scheduler_info[0]["max_req_input_len"]
+        self.max_num_seqs = scheduler_info[0]["max_num_seqs"]
+        self.chunked_prefill_size = scheduler_info[0]["chunked_prefill_size"]
+        self.max_model_len = scheduler_info[0]["max_model_len"]
+
+    def round_robin_scheduler(self, req: Req):
+        if self.server_args.disaggregation_mode == "null":
+            self.workers[self.round_robin_counter].send_pyobj(req)
+            self.round_robin_counter = (self.round_robin_counter + 1) % len(
+                self.workers
+            )
+        else:
+            self.workers[req.bootstrap_room % len(self.workers)].send_pyobj(req)
+
+    def single_robin_scheduler(self, req):
+        worker_id = int(os.environ.get("SINGLE_WORKER_ID", "-1"))
+        if not 0 <= worker_id < self.server_args.mapping.attn.dp_size - 1:
+            raise ValueError(f"Invalid SINGLE_WORKER_ID:{worker_id}")
+        self.workers[worker_id].send_pyobj(req)
+
+    def budget_scheduler(self, req):
+        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 event_loop(self):
+        while True:
+            while True:
+                try:
+                    recv_req = self.recv_from_tokenizer.recv_pyobj(zmq.NOBLOCK)
+                except zmq.ZMQError:
+                    break
+                self._request_dispatcher(recv_req)
+
+
+def run_data_parallel_controller_process(
+    server_args: ServerArgs,
+    port_args: PortArgs,
+    pipe_writer,
+):
+    setproctitle.setproctitle("tokenspeed::data_parallel_controller")
+    configure_logger(server_args)
+    parent_process = psutil.Process().parent()
+    register_usr_signal()
+
+    try:
+        controller = DataParallelController(server_args, port_args)
+        pipe_writer.send(
+            {
+                "status": "ready",
+                "max_total_num_tokens": controller.max_total_num_tokens,
+                "max_req_input_len": controller.max_req_input_len,
+                "max_num_seqs": controller.max_num_seqs,
+                "chunked_prefill_size": controller.chunked_prefill_size,
+                "max_model_len": controller.max_model_len,
+            }
+        )
+        if server_args.node_rank == 0:
+            controller.event_loop()
+        for proc in controller.scheduler_procs:
+            proc.join()
+            logger.error(
+                "Scheduler or DataParallelController %s terminated with %s",
+                proc.pid,
+                proc.exitcode,
+            )
+    except Exception:
+        traceback = get_exception_traceback()
+        logger.error("DataParallelController hit an exception: %s", traceback)
+        parent_process.send_signal(signal.SIGUSR1)