# 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)