chore: import upstream snapshot with attribution

python/sglang/srt/managers/data_parallel_controller.py

@@ -0,0 +1,715 @@
+# Copyright 2023-2024 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""A controller that dispatches requests to multiple data parallel workers."""
+
+import faulthandler
+import logging
+import multiprocessing as mp
+import signal
+import threading
+import time
+from enum import Enum, auto
+from typing import Callable, List, Optional
+
+import psutil
+import setproctitle
+import zmq
+
+from sglang.srt.environ import envs
+from sglang.srt.layers.dp_attention import compute_dp_attention_world_info
+from sglang.srt.managers.io_struct import (
+    ActiveRanksOutput,
+    BatchTokenizedEmbeddingReqInput,
+    BatchTokenizedGenerateReqInput,
+    BlockReqInput,
+    ProfileReq,
+    TokenizedEmbeddingReqInput,
+    TokenizedGenerateReqInput,
+    sock_recv,
+    sock_send,
+    unwrap_from_pickle,
+    wrap_as_pickle,
+)
+from sglang.srt.managers.load_snapshot import create_load_snapshot_reader
+from sglang.srt.managers.schedule_batch import Req
+from sglang.srt.managers.scheduler import run_scheduler_process
+from sglang.srt.observability.cpu_monitor import start_cpu_monitor_thread
+from sglang.srt.observability.req_time_stats import DPControllerReqTimeStats
+from sglang.srt.observability.trace import process_tracing_init, trace_set_thread_info
+from sglang.srt.server_args import (
+    DP_ATTENTION_HANDSHAKE_PORT_DELTA,
+    PortArgs,
+    ServerArgs,
+)
+from sglang.srt.utils import numa_utils
+from sglang.srt.utils.common import (
+    configure_logger,
+    kill_itself_when_parent_died,
+    maybe_reindex_device_id,
+)
+from sglang.srt.utils.network import (
+    NetworkAddress,
+    bind_port,
+    get_zmq_socket,
+    get_zmq_socket_on_host,
+)
+from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter
+from sglang.srt.utils.watchdog import Watchdog
+from sglang.utils import TypeBasedDispatcher, get_exception_traceback
+
+logger = logging.getLogger(__name__)
+
+SCHEDULER_PIDS_ARG = "scheduler_pids"
+
+
+class LoadBalanceMethod(Enum):
+    """Load balance method."""
+
+    ROUND_ROBIN = auto()
+    FOLLOW_BOOTSTRAP_ROOM = auto()
+    TOTAL_REQUESTS = auto()
+    TOTAL_TOKENS = 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, dp_size: int):
+        self.dp_size = dp_size
+        self.total_requests = [0] * dp_size
+        self.total_tokens = [0] * dp_size
+        self.last_timestamp = [0.0] * dp_size
+
+    def update_budget(self, loads):
+        """Update budget from shm snapshots, skipping stale reads."""
+        for load in loads:
+            if load.timestamp == self.last_timestamp[load.dp_rank]:
+                continue
+            self.last_timestamp[load.dp_rank] = load.timestamp
+            self.total_requests[load.dp_rank] = (
+                load.num_running_reqs + load.num_waiting_reqs
+            )
+            self.total_tokens[load.dp_rank] = load.num_total_tokens
+
+    def dispatch(self, method: LoadBalanceMethod, estimated_tokens: int = 0):
+        if method == LoadBalanceMethod.TOTAL_REQUESTS:
+            target_rank = self.total_requests.index(min(self.total_requests))
+        elif method == LoadBalanceMethod.TOTAL_TOKENS:
+            # Use total_requests as a tie-breaker when total_tokens are equal
+            target_rank = min(
+                range(self.dp_size),
+                key=lambda i: (self.total_tokens[i], self.total_requests[i]),
+            )
+        else:
+            return None
+
+        # Increment the load of that worker by one as a heuristic
+        self.total_requests[target_rank] += 1
+        self.total_tokens[target_rank] += estimated_tokens
+        return target_rank
+
+
+class DataParallelController:
+    """A controller that dispatches requests to multiple data parallel workers."""
+
+    def __init__(
+        self,
+        server_args: ServerArgs,
+        port_args: PortArgs,
+        run_scheduler_process_func: Callable,
+    ) -> None:
+        # Parse args
+        self.server_args = server_args
+        self.port_args = port_args
+        self.load_balance_method = LoadBalanceMethod.from_str(
+            server_args.load_balance_method
+        )
+        self.run_scheduler_process_func = run_scheduler_process_func
+
+        # Init inter-process communication
+        self.context = zmq.Context(1 + server_args.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
+            )
+
+        # Dispatch method
+        self.round_robin_counter = 0
+        dispatch_lookup = {
+            LoadBalanceMethod.ROUND_ROBIN: self.round_robin_scheduler,
+            LoadBalanceMethod.FOLLOW_BOOTSTRAP_ROOM: self.follow_bootstrap_room_scheduler,
+            LoadBalanceMethod.TOTAL_REQUESTS: self.total_requests_scheduler,
+            LoadBalanceMethod.TOTAL_TOKENS: self.total_tokens_scheduler,
+        }
+        self.dispatching = dispatch_lookup[self.load_balance_method]
+        self.refresh_load_budget_on_dispatch = self.load_balance_method in (
+            LoadBalanceMethod.TOTAL_REQUESTS,
+            LoadBalanceMethod.TOTAL_TOKENS,
+        )
+
+        # Load balance budget
+        self.dp_budget = DPBudget(server_args.dp_size)
+        self.load_snapshot_reader = create_load_snapshot_reader(
+            server_args,
+            port_args,
+            caller="DataParallelController",
+        )
+        self._last_refresh_time = 0.0
+
+        # To protect changing env vars to set CUDA_VISIBLE_DEVICES.
+        self.env_lock = threading.Lock()
+
+        # Launch data parallel workers
+        self.scheduler_procs = []
+        self.workers: List[zmq.Socket] = [None] * server_args.dp_size
+        self.status: List[bool] = [True] * server_args.dp_size
+
+        if server_args.enable_dp_attention:
+            self.launch_dp_attention_schedulers(server_args, port_args)
+            # When local control broadcast is enabled, send control messages to
+            # every DP group leader (attn_tp_rank=0) so each leader broadcasts
+            # within its own attn_tp_group instead of the full tp_group.
+            # Otherwise fall back to the original behaviour: send to only the
+            # first leader, which then broadcasts over the full tp_group.
+            local_ctrl = server_args.enable_dp_attention_local_control_broadcast
+            self.control_message_step = 1 if local_ctrl else server_args.tp_size
+        else:
+            self.launch_dp_schedulers(server_args, port_args)
+            self.control_message_step = 1
+
+        self.init_dispatcher()
+
+        self.soft_watchdog = Watchdog.create(
+            debug_name="DataParallelController",
+            watchdog_timeout=server_args.soft_watchdog_timeout,
+            soft=True,
+            test_stuck_time=envs.SGLANG_TEST_STUCK_DP_CONTROLLER.get(),
+        )
+
+        if server_args.enable_metrics:
+            start_cpu_monitor_thread("data_parallel_controller")
+
+    def send_to_all_workers(self, obj):
+        for i, worker in enumerate(self.workers):
+            if self.status[i]:
+                sock_send(worker, 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]:
+            sock_send(worker, obj)
+
+    def update_active_ranks(self, ranks: ActiveRanksOutput):
+        self.status = ranks.status
+
+    def refresh_load_budget(self):
+        # Throttle to at most once per 20ms.  When a burst of requests
+        # arrives, dispatching_with_trace() calls this before every
+        # dispatch.  Each call reads the latest scheduler snapshot and
+        # overwrites the speculative +1 increments that DPBudget.dispatch()
+        # added for previously dispatched requests in this burst.  Without
+        # throttling, the budget resets to the (stale) scheduler-reported
+        # value on every request, causing the entire burst to land on a
+        # single DP rank.  The 20ms interval lets the burst complete
+        # using speculative counters, then refreshes from the real
+        # scheduler load for the next batch.
+        now = time.perf_counter()
+        if now - self._last_refresh_time < 0.02:
+            return
+        self._last_refresh_time = now
+        self.dp_budget.update_budget(self.load_snapshot_reader.read_all())
+
+    def dispatching_with_trace(self, req: Req, refresh_load_budget: bool = True):
+        if refresh_load_budget and self.refresh_load_budget_on_dispatch:
+            self.refresh_load_budget()
+
+        time_stats = DPControllerReqTimeStats.new_from_obj(
+            unwrap_from_pickle(req.time_stats)
+        )
+
+        time_stats.set_dp_dispatch_time()
+        req.time_stats = wrap_as_pickle(time_stats)
+        self.dispatching(req)
+        req.time_stats = time_stats
+        req.time_stats.set_dp_dispatch_finish_time()
+
+    def dispatch_batch_generate(self, batch_req: BatchTokenizedGenerateReqInput):
+        if self.refresh_load_budget_on_dispatch:
+            self.refresh_load_budget()
+        for req in batch_req:
+            self.dispatching_with_trace(req, refresh_load_budget=False)
+
+    def dispatch_batch_embedding(self, batch_req: BatchTokenizedEmbeddingReqInput):
+        if self.refresh_load_budget_on_dispatch:
+            self.refresh_load_budget()
+        for req in batch_req:
+            self.dispatching_with_trace(req, refresh_load_budget=False)
+
+    def init_dispatcher(self):
+        self._request_dispatcher = TypeBasedDispatcher(
+            [
+                (TokenizedGenerateReqInput, self.dispatching_with_trace),
+                (TokenizedEmbeddingReqInput, self.dispatching_with_trace),
+                (BatchTokenizedGenerateReqInput, self.dispatch_batch_generate),
+                (BatchTokenizedEmbeddingReqInput, self.dispatch_batch_embedding),
+                (BlockReqInput, self.send_to_all_workers),
+                (ProfileReq, self.send_to_all_workers),
+                (ActiveRanksOutput, self.update_active_ranks),
+            ]
+        )
+        self._request_dispatcher.add_fallback_fn(self.send_control_message)
+
+    def launch_dp_schedulers(self, server_args, port_args):
+        base_gpu_id = 0
+
+        threads = []
+        sockets = []
+        ready_events = []
+        for dp_rank in range(server_args.dp_size):
+            tmp_port_args = PortArgs.init_new(server_args)
+            tmp_port_args.tokenizer_ipc_name = port_args.tokenizer_ipc_name
+            tmp_port_args.detokenizer_ipc_name = port_args.detokenizer_ipc_name
+            tmp_port_args.instance_id = port_args.instance_id
+
+            # This port is checked free in PortArgs.init_new.
+            # We hold it first so that the next dp worker gets a different port
+            sockets.append(bind_port(tmp_port_args.nccl_port))
+
+            ready_event = threading.Event()
+            ready_events.append(ready_event)
+
+            # Create a thread for each worker
+            thread = threading.Thread(
+                target=self.launch_tensor_parallel_group_thread,
+                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.pp_size * server_args.gpu_id_step
+            )
+
+            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,
+                )
+
+        # Free all sockets before starting the threads to launch TP workers
+        for sock in sockets:
+            sock.close()
+
+        # Start all threads
+        for thread in threads:
+            thread.start()
+        for event in ready_events:
+            event.wait()
+
+    def launch_tensor_parallel_group_thread(
+        self,
+        server_args: ServerArgs,
+        port_args: PortArgs,
+        base_gpu_id: int,
+        dp_rank: int,
+        ready_event: threading.Event,
+    ):
+        self.launch_tensor_parallel_group(server_args, port_args, base_gpu_id, dp_rank)
+        ready_event.set()
+
+        # This thread cannot be closed because otherwise the `kill_itself_when_parent_died`
+        # function in scheduler.py will kill the scheduler.
+        while True:
+            time.sleep(30 * 24 * 3600)
+
+    def _broadcast_worker_ports(
+        self, server_args: ServerArgs, worker_ports: Optional[List[int]] = None
+    ) -> List[int]:
+        """Broadcast worker ports from node 0 to all other nodes.
+
+        Node 0 acts as the server, waiting for all other nodes to connect and
+        sending them the pre-allocated worker ports. Other nodes act as clients,
+        connecting to node 0 to receive their copy of the worker ports.
+
+        Args:
+            server_args: Server arguments containing node configuration.
+            worker_ports: Pre-allocated worker ports to broadcast.
+
+        Returns:
+            List of worker ports (same on all nodes after broadcast).
+        """
+        # Determine the endpoint for inter-node communication
+        if server_args.dist_init_addr is None:
+            na = NetworkAddress(
+                server_args.host or "127.0.0.1",
+                server_args.port + DP_ATTENTION_HANDSHAKE_PORT_DELTA,
+            )
+        else:
+            na = NetworkAddress.parse(server_args.dist_init_addr)
+            na = NetworkAddress(na.host, na.port + DP_ATTENTION_HANDSHAKE_PORT_DELTA)
+        endpoint = na.to_tcp()
+
+        if server_args.node_rank == 0:
+            # Node 0: Broadcast worker ports to all other nodes
+            return self._broadcast_ports_as_server(
+                endpoint, server_args.nnodes - 1, worker_ports
+            )
+        else:
+            # Other nodes: Receive worker ports from node 0
+            return self._receive_ports_as_client(endpoint, server_args.node_rank)
+
+    def _broadcast_ports_as_server(
+        self, endpoint: str, expected_clients: int, worker_ports: List[int]
+    ) -> List[int]:
+        """Broadcast worker ports to all client nodes."""
+        logger.debug(f"Broadcasting worker ports to {expected_clients} client nodes")
+        logger.debug(f"Worker ports: {worker_ports}")
+
+        rep_socket = get_zmq_socket(self.context, zmq.REP, endpoint, True)
+
+        try:
+            connected_clients = 0
+            while connected_clients < expected_clients:
+                # Wait for client handshake
+                client_rank = sock_recv(rep_socket)
+                logger.debug(f"Received handshake from node {client_rank}")
+
+                # Send worker ports to client
+                sock_send(rep_socket, wrap_as_pickle(worker_ports))
+                connected_clients += 1
+                logger.debug(
+                    f"Sent worker ports to {connected_clients}/{expected_clients} nodes"
+                )
+
+            logger.debug("Worker port broadcast completed")
+            return worker_ports
+        finally:
+            if self.server_args.elastic_ep_backend is None:
+                rep_socket.close()
+            else:
+                threading.Thread(
+                    target=self._reply_ports_as_server,
+                    args=(rep_socket, worker_ports),
+                    daemon=True,
+                ).start()
+
+    def _reply_ports_as_server(self, rep_socket: zmq.Socket, worker_ports: List[int]):
+        """
+        Runs as a background thread to broadcast worker ports for recovered EP ranks
+        """
+        while True:
+            # Wait for client handshake
+            try:
+                client_rank = sock_recv(rep_socket)
+            except Exception:
+                logger.exception(
+                    "Failed to recv/decode handshake in reply thread; continue"
+                )
+                continue
+            logger.debug(f"Received handshake from node {client_rank}")
+
+            # Send worker ports to client
+            sock_send(rep_socket, wrap_as_pickle(worker_ports))
+            logger.debug(f"Sent worker ports to node {client_rank}")
+
+    def _receive_ports_as_client(self, endpoint: str, node_rank: int) -> List[int]:
+        """Receive worker ports from the server node."""
+        logger.debug(f"Connecting to node 0 to receive worker ports")
+
+        req_socket = get_zmq_socket(self.context, zmq.REQ, endpoint, False)
+        req_socket.setsockopt(zmq.RCVTIMEO, 600 * 1000)  # 10 minute timeout
+        req_socket.setsockopt(zmq.SNDTIMEO, 600 * 1000)
+
+        try:
+            # Send handshake with our node rank
+            sock_send(req_socket, wrap_as_pickle(str(node_rank)))
+
+            # Receive worker ports
+            worker_ports = sock_recv(req_socket)
+            logger.debug(f"Received {len(worker_ports)} worker ports from node 0")
+            return worker_ports
+        except zmq.Again:
+            logger.error("Timeout waiting for worker ports from node 0")
+            raise RuntimeError(
+                "Failed to receive worker ports from node 0 within timeout"
+            )
+        finally:
+            req_socket.close()
+
+    def launch_dp_attention_schedulers(
+        self, server_args: ServerArgs, port_args: PortArgs
+    ):
+        if server_args.dist_init_addr is None:
+            bind_host = "127.0.0.1"
+        else:
+            bind_host = NetworkAddress.parse(server_args.dist_init_addr).host
+
+        # Pre-allocate worker ports on node 0 to avoid conflicts
+        worker_ports = []
+        if server_args.node_rank == 0:
+            for dp_rank in range(server_args.dp_size):
+                worker_port, worker_socket = get_zmq_socket_on_host(
+                    self.context, zmq.PUSH, host=bind_host
+                )
+                worker_ports.append(worker_port)
+                self.workers[dp_rank] = worker_socket
+                logger.debug(
+                    "Assigned port %s to worker %s on host %s",
+                    worker_port,
+                    dp_rank,
+                    bind_host,
+                )
+
+        broadcasted_ports = self._broadcast_worker_ports(
+            server_args, worker_ports if worker_ports else None
+        )
+        self.launch_tensor_parallel_group(
+            server_args, port_args, 0, None, broadcasted_ports
+        )
+
+    def launch_tensor_parallel_group(
+        self,
+        server_args: ServerArgs,
+        port_args: PortArgs,
+        base_gpu_id: int,
+        dp_rank: Optional[int],
+        worker_ports: Optional[List[int]] = None,
+    ):
+        if not server_args.enable_dp_attention:
+            logger.info(f"Launch DP{dp_rank} starting at GPU #{base_gpu_id}.")
+
+        memory_saver_adapter = TorchMemorySaverAdapter.create(
+            enable=server_args.enable_memory_saver
+        )
+
+        scheduler_pipe_readers = []
+
+        pp_size_per_node = max(server_args.pp_size // server_args.nnodes, 1)
+        nnodes_per_pp_rank = max(server_args.nnodes // server_args.pp_size, 1)
+        pp_rank_range = range(
+            pp_size_per_node * (server_args.node_rank // nnodes_per_pp_rank),
+            pp_size_per_node * (server_args.node_rank // nnodes_per_pp_rank + 1),
+        )
+
+        nnodes_per_tp_group = nnodes_per_pp_rank
+        tp_size_per_node = server_args.tp_size // nnodes_per_tp_group
+        tp_rank_range = range(
+            tp_size_per_node * (server_args.node_rank % nnodes_per_tp_group),
+            tp_size_per_node * (server_args.node_rank % nnodes_per_tp_group + 1),
+        )
+
+        attn_cp_rank = 0
+        moe_dp_rank = 0
+        for pp_rank in pp_rank_range:
+            for tp_rank in tp_rank_range:
+                rank_port_args = port_args
+
+                if server_args.enable_dp_attention:
+                    # dp attention has different sharding logic
+                    _, _, dp_rank, _ = compute_dp_attention_world_info(
+                        server_args.enable_dp_attention,
+                        tp_rank,
+                        server_args.tp_size,
+                        server_args.dp_size,
+                        server_args.attn_cp_size,
+                    )
+                    # compute zmq ports for this dp rank
+                    rank_port_args = PortArgs.init_new(
+                        server_args, dp_rank, worker_ports
+                    )
+                    # Data parallelism reuses the tensor parallelism group,
+                    # so all dp ranks should use the same nccl port.
+                    rank_port_args.nccl_port = port_args.nccl_port
+                    rank_port_args.instance_id = port_args.instance_id
+
+                reader, writer = mp.Pipe(duplex=False)
+                gpu_id = (
+                    server_args.base_gpu_id
+                    + base_gpu_id
+                    + ((pp_rank % pp_size_per_node) * tp_size_per_node)
+                    + (tp_rank % tp_size_per_node) * server_args.gpu_id_step
+                )
+                attn_dp_size = (
+                    server_args.dp_size if server_args.enable_dp_attention else 1
+                )
+
+                # Parallelism hierarchy (outermost to innermost):
+                # - Attention: Global(TP) -> DP -> ATTN_CP -> ATTN_TP (innermost)
+                # - MoE: Global(TP) -> MOE_DP -> EP -> MOE_TP (innermost)
+                attn_tp_size = (
+                    server_args.tp_size // attn_dp_size // server_args.attn_cp_size
+                )
+                attn_cp_rank = (tp_rank // attn_tp_size) % server_args.attn_cp_size
+                moe_dp_rank = tp_rank // (
+                    server_args.tp_size // server_args.moe_dp_size
+                )
+                moe_ep_rank = (
+                    tp_rank
+                    % (server_args.tp_size // server_args.moe_dp_size)
+                    // (
+                        server_args.tp_size
+                        // server_args.moe_dp_size
+                        // server_args.ep_size
+                    )
+                )
+
+                with self.env_lock, maybe_reindex_device_id(gpu_id) as gpu_id:
+                    proc = mp.Process(
+                        target=self.run_scheduler_process_func,
+                        args=(
+                            server_args,
+                            rank_port_args,
+                            gpu_id,
+                            tp_rank,
+                            attn_cp_rank,
+                            moe_dp_rank,
+                            moe_ep_rank,
+                            pp_rank,
+                            dp_rank,
+                            writer,
+                        ),
+                    )
+                    with (
+                        memory_saver_adapter.configure_subprocess(),
+                        numa_utils.configure_subprocess(server_args, gpu_id),
+                    ):
+                        proc.start()
+                self.scheduler_procs.append(proc)
+                scheduler_pipe_readers.append(reader)
+
+        # Wait for model to finish loading
+        scheduler_info = []
+        for i in range(len(scheduler_pipe_readers)):
+            scheduler_info.append(scheduler_pipe_readers[i].recv())
+
+        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.routed_dp_rank is not None:
+            logger.debug(f"Direct routing to DP rank {req.routed_dp_rank}")
+            sock_send(self.workers[req.routed_dp_rank], req)
+            return True
+        return False
+
+    def round_robin_scheduler(self, req: Req):
+        if self.maybe_external_dp_rank_routing(req):
+            return
+
+        while True:
+            if self.status[self.round_robin_counter]:
+                logger.debug(f"Choose worker {self.round_robin_counter}")
+                sock_send(self.workers[self.round_robin_counter], req)
+                self.round_robin_counter = (self.round_robin_counter + 1) % len(
+                    self.workers
+                )
+                break
+            self.round_robin_counter = (self.round_robin_counter + 1) % len(
+                self.workers
+            )
+
+    def follow_bootstrap_room_scheduler(self, req: Req):
+        if self.maybe_external_dp_rank_routing(req):
+            return
+
+        assert req.bootstrap_room is not None, (
+            "req.bootstrap_room should not be None. Do not send requests directly to "
+            "prefill or decode instances; send to the router instead."
+        )
+        target_rank = req.bootstrap_room % len(self.workers)
+        sock_send(self.workers[target_rank], req)
+
+    def total_requests_scheduler(self, req: Req):
+        if self.maybe_external_dp_rank_routing(req):
+            return
+        target_worker = self.dp_budget.dispatch(LoadBalanceMethod.TOTAL_REQUESTS)
+        sock_send(self.workers[target_worker], req)
+
+    def total_tokens_scheduler(self, req: Req):
+        if self.maybe_external_dp_rank_routing(req):
+            return
+        estimated_tokens = len(req.input_ids)
+        target_worker = self.dp_budget.dispatch(
+            LoadBalanceMethod.TOTAL_TOKENS, estimated_tokens=estimated_tokens
+        )
+        sock_send(self.workers[target_worker], req)
+
+    def event_loop(self):
+        while True:
+            while True:
+                self.soft_watchdog.feed()
+                try:
+                    recv_req = sock_recv(self.recv_from_tokenizer, flags=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,
+    run_scheduler_process_func: Callable = run_scheduler_process,
+):
+    setproctitle.setproctitle("sglang::data_parallel_controller")
+    faulthandler.enable()
+    kill_itself_when_parent_died()
+    parent_process = psutil.Process().parent()
+
+    configure_logger(server_args)
+    if server_args.enable_trace:
+        process_tracing_init(
+            server_args.otlp_traces_endpoint,
+            "sglang",
+            trace_modules=server_args.trace_modules,
+        )
+        thread_label = "DP Controller"
+        if server_args.disaggregation_mode == "prefill":
+            thread_label = "Prefill DP Controller"
+        elif server_args.disaggregation_mode == "decode":
+            thread_label = "Decode DP Controller"
+        trace_set_thread_info(thread_label)
+
+    try:
+        controller = DataParallelController(
+            server_args, port_args, run_scheduler_process_func
+        )
+        scheduler_pids = [
+            proc.pid for proc in controller.scheduler_procs if proc is not None
+        ]
+        pipe_writer.send(
+            {
+                "status": "ready",
+                "max_total_num_tokens": controller.max_total_num_tokens,
+                "max_req_input_len": controller.max_req_input_len,
+                SCHEDULER_PIDS_ARG: scheduler_pids,
+            }
+        )
+        if server_args.node_rank == 0:
+            controller.event_loop()
+        for proc in controller.scheduler_procs:
+            proc.join()
+            logger.error(
+                f"Scheduler or DataParallelController {proc.pid} terminated with {proc.exitcode}"
+            )
+    except Exception:
+        traceback = get_exception_traceback()
+        logger.error(f"DataParallelController hit an exception: {traceback}")
+        parent_process.send_signal(signal.SIGQUIT)