Refactor data parallel coordinator to enable modular handlers

megatron/core/inference/data_parallel_inference_coordinator/__init__.py

@@ -0,0 +1,22 @@
+# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+
+"""Data parallel inference coordinator package.
+
+The coordinator class itself lives in coordinator.py; message handlers are in
+handlers.py and the control-signal state machine in state.py. This module
+re-exports the public names so existing imports of
+``megatron.core.inference.data_parallel_inference_coordinator`` keep working.
+"""
+
+from .coordinator import DataParallelInferenceCoordinator
+from .handlers import HANDLERS, message_handler
+from .state import CONTROL_TRANSITIONS, ControlTransition, CoordinatorState
+
+__all__ = [
+    "DataParallelInferenceCoordinator",
+    "CoordinatorState",
+    "ControlTransition",
+    "CONTROL_TRANSITIONS",
+    "HANDLERS",
+    "message_handler",
+]

megatron/core/inference/data_parallel_inference_coordinator.py → megatron/core/inference/data_parallel_inference_coordinator/coordinator.py

@@ -7,7 +7,6 @@
 import signal
 import socket
 from collections import deque
-from enum import Enum, auto
 from multiprocessing import Event
 from multiprocessing.connection import Connection
 
@@ -21,6 +20,9 @@
     TextGenerationController,
 )
 
+from .handlers import HANDLERS
+from .state import CoordinatorState
+
 try:
     import zmq
 
@@ -63,6 +65,10 @@ class DataParallelInferenceCoordinator:
     5.  **Control Signal Broadcasting**: It relays control signals (e.g., PAUSE, STOP)
         from a client to all connected data parallel ranks.
 
+    Message handling is split out into handlers.py: the event loop in start()
+    dispatches each message to the handler registered for its header, so
+    supporting a new message type requires no changes here.
+
     Attributes:
         router_socket (zmq.Socket): The central ZMQ ROUTER socket for all communication.
         data_parallel_size (int): The number of data parallel workers to expect.
@@ -75,13 +81,9 @@ class DataParallelInferenceCoordinator:
         next_request_id (int): A counter for generating unique server-side request IDs.
     """
 
-    class CoordinatorState(Enum):
-        """State machine for the coordinator."""
-
-        RUNNING = auto()
-        PAUSED = auto()
-        SUSPENDED = auto()
-        STOPPING = auto()
+    # Exposed as a class attribute for backwards compatibility; the canonical
+    # definition lives in state.py.
+    CoordinatorState = CoordinatorState
 
     def __init__(
         self,
@@ -189,7 +191,7 @@ def __init__(
 
         self.next_request_id = 0
         self.tokenizer = tokenizer
-        self.state = self.CoordinatorState.RUNNING
+        self.state = CoordinatorState.RUNNING
 
         # Prefix caching state for routing.
         self.block_size_tokens = block_size_tokens
@@ -220,6 +222,12 @@ def __init__(
         self._hash_table: dict[int, dict[int, int]] = {}
         self._hash_assignment_counter = 0
 
+        # Clients that have completed the CONNECT handshake.
+        self.known_clients = set()
+
+        # Header -> handler dispatch table, sourced from the handler registry.
+        self._handlers = dict(HANDLERS)
+
     def get_least_loaded_data_parallel_rank(self):
         """
         Selects the data parallel rank with the fewest in-flight requests.
@@ -307,6 +315,12 @@ def _send_to_engine(self, identity, payload):
                 return False
             raise
 
+    def _broadcast_to_engines(self, payload):
+        """Send a deserialized payload to every connected data parallel rank."""
+        serialized = msgpack.packb(payload, use_bin_type=True)
+        for data_parallel_rank_id in list(self.identities_of_data_parallel_ranks):
+            self._send_to_engine(data_parallel_rank_id, serialized)
+
     def compute_request_hashes(self, prompt):
         """Compute block hashes for a prompt on CPU.
 
@@ -408,212 +422,33 @@ def start(self):
         Starts the main event loop for the coordinator.
 
         This method runs an infinite loop, continuously listening for incoming
-        messages on the ZMQ ROUTER socket. It parses the message header to
-        determine the message type and takes appropriate action, such as
-        handling new client connections, forwarding requests, broadcasting
-        control signals, or processing replies from the engines.
+        messages on the ZMQ ROUTER socket. It reads the message header and
+        dispatches to the handler registered for it (see handlers.py).
+        A handler that returns a truthy value stops the loop.
         """
         # Todo [Siddharth]: Make this more robust to handle invalid messages.
-        known_clients = set()
         while True:
             sender_identity, serialized_payload = self.router_socket.recv_multipart()
 
-            # Allow for re-registration if connecting to a running coordinator.
+            # An empty payload is a data parallel rank (re-)registering itself.
             if serialized_payload == b"":
-                if sender_identity not in self.identities_of_data_parallel_ranks:
-                    self.identities_of_data_parallel_ranks.append(sender_identity)
-                    self._register_rank_identity(sender_identity)
+                self._handle_rank_registration(sender_identity)
                 continue
 
             deserialized_payload = msgpack.unpackb(serialized_payload, raw=False)
             header = Headers(deserialized_payload[0])
 
-            if header == Headers.CONNECT:
-                if sender_identity in known_clients:
-                    logging.info(
-                        f"Client {sender_identity} sent a duplicate connect request. Ignoring .."
-                    )
-                    continue
-
-                # print(f"New client connected: {sender_identity}")
-                known_clients.add(sender_identity)
-                self.router_socket.send_multipart(
-                    [sender_identity, msgpack.packb([Headers.CONNECT_ACK.value], use_bin_type=True)]
-                )
-
-            elif header == Headers.SUBMIT_REQUEST:
-                # ToDo [Siddharth]: We might want to tokenize the prompt on the
-                # assigned data parallel rank for this process instead
-                # of the coordinator.
-
-                # Message from a known client
-                if sender_identity not in known_clients:
-                    logging.info(
-                        f"Received message from unknown client {sender_identity}. Ignoring."
-                    )
-                    continue
-                # this is a message from a client.
-                # route it to a data parallel rank
-                client_request_id, prompt, sampling_params = deserialized_payload[1:]
-                # map client request_id to server request_id
-                # necessary because multiple clients might have the same request_id.
-                request_id = self.next_request_id
-                self.next_request_id += 1
-                self.request_id_to_client_id[request_id] = sender_identity
-                self.request_id_to_client_request_id[request_id] = client_request_id
-
-                # Serialize prompt.
-                if isinstance(prompt, (str, list)):
-                    pass
-                elif isinstance(prompt, torch.Tensor):
-                    prompt = prompt.tolist()
-                else:
-                    raise Exception("specialize for <%s> prompt." % type(prompt).__name__)
-
-                payload = msgpack.packb(
-                    [Headers.SUBMIT_REQUEST.value, request_id, prompt, sampling_params],
-                    use_bin_type=True,
-                )
-
-                request_hashes = self.compute_request_hashes(prompt)
-                if (
-                    self.prefix_caching_coordinator_policy
-                    == PrefixCachingCoordinatorPolicy.FIRST_PREFIX_BLOCK
-                ):
-                    request_hashes = request_hashes[:1]
-
-                # Account for the fact that some engines may have died.
-                for _ in range(len(self.identities_of_data_parallel_ranks)):
-                    next_identity = self.get_best_data_parallel_rank(request_hashes)
-                    if self._send_to_engine(next_identity, payload):
-                        break
-                else:
-                    # If all engines have died, we are in an abnormal state, and must exit cleanly.
-                    logging.error("Coordinator: no reachable engines for request %d", request_id)
-                    del self.request_id_to_client_id[request_id]
-                    del self.request_id_to_client_request_id[request_id]
-                    return
-
-                self.request_id_to_rank[request_id] = next_identity
-                self._pending_counts[self.identity_to_rank_index[next_identity]] += 1
-                if request_hashes:
-                    self._update_rank_hashes(next_identity, request_hashes)
-                if self.schedule_records is not None:
-                    self.schedule_records.append(
-                        {
-                            "request_id": request_id,
-                            "rank_index": self.identity_to_rank_index[next_identity],
-                            "num_hashes": len(request_hashes),
-                        }
-                    )
-
-            elif header in (
-                Headers.PAUSE,
-                Headers.UNPAUSE,
-                Headers.SUSPEND,
-                Headers.RESUME,
-                Headers.SET_GENERATION_EPOCH,
-                Headers.STOP,
-            ):
-                # Start by checking the current state against the control signal.
-                if sender_identity not in known_clients:
-                    logging.warning("Coordinator: ignoring signal from unknown client.")
-                    continue
-
-                if header == Headers.PAUSE:
-                    idem_states = (self.CoordinatorState.PAUSED, self.CoordinatorState.SUSPENDED)
-                    if self.state == self.CoordinatorState.RUNNING:
-                        self.state = self.CoordinatorState.PAUSED
-                    elif self.state in idem_states:
-                        # Already paused/suspended, ignore redundant PAUSE.
-                        continue
-                    else:
-                        logging.warning("Coordinator: ignoring PAUSE in state %s", self.state)
-                        continue
-                elif header == Headers.UNPAUSE:
-                    if self.state != self.CoordinatorState.PAUSED:
-                        logging.warning("Coordinator: ignoring UNPAUSE in state %s", self.state)
-                        continue
-                    self.state = self.CoordinatorState.RUNNING
-                elif header == Headers.SUSPEND:
-                    if self.state != self.CoordinatorState.PAUSED:
-                        logging.warning("Coordinator: ignoring SUSPEND in state %s", self.state)
-                        continue
-                    self.state = self.CoordinatorState.SUSPENDED
-                elif header == Headers.RESUME:
-                    if self.state != self.CoordinatorState.SUSPENDED:
-                        logging.warning("Coordinator: ignoring RESUME in state %s", self.state)
-                        continue
-                    self.state = self.CoordinatorState.PAUSED
-                elif header == Headers.STOP:
-                    good_states = (self.CoordinatorState.PAUSED, self.CoordinatorState.SUSPENDED)
-                    if self.state not in good_states:
-                        logging.warning("Coordinator: ignoring STOP in state %s", self.state)
-                        continue
-                    self.state = self.CoordinatorState.STOPPING
-
-                # Broadcast the control signal if we're in a good state.
-                # Forward the full deserialized payload so that data-bearing
-                # signals (e.g. SET_GENERATION_EPOCH) retain their arguments.
-                broadcast_payload = msgpack.packb(deserialized_payload, use_bin_type=True)
-                for data_parallel_rank_id in list(self.identities_of_data_parallel_ranks):
-                    self._send_to_engine(data_parallel_rank_id, broadcast_payload)
-
-                # STOP affects engines; reset coordinator to RUNNING to allow future engines.
-                if header == Headers.STOP:
-                    self.state = self.CoordinatorState.RUNNING
-
-            elif header in (Headers.START_CUDA_PROFILER, Headers.STOP_CUDA_PROFILER):
-                # Profiler control: broadcast to every connected DP engine. Not a
-                # state transition, so no CoordinatorState checks — just forward.
-                if sender_identity not in known_clients:
-                    logging.warning("Coordinator: ignoring profiler signal from unknown client.")
-                    continue
-                broadcast_payload = msgpack.packb(deserialized_payload, use_bin_type=True)
-                for data_parallel_rank_id in list(self.identities_of_data_parallel_ranks):
-                    self._send_to_engine(data_parallel_rank_id, broadcast_payload)
-
-            elif header == Headers.ENGINE_REPLY:
-                # This is the output of a single engine step on some data parallel rank.
-                assert sender_identity in self.identities_of_data_parallel_ranks
-                finished_requests = deserialized_payload[1]
-
-                for finished_request in finished_requests:
-                    self.detokenize(finished_request)
-                    fid = finished_request["request_id"]
-                    client_identity = self.request_id_to_client_id[fid]
-                    client_request_identity = self.request_id_to_client_request_id[fid]
-                    del self.request_id_to_client_id[fid]
-                    del self.request_id_to_client_request_id[fid]
-                    assigned_rank = self.request_id_to_rank.pop(fid, None)
-                    if assigned_rank is not None:
-                        idx = self.identity_to_rank_index.get(assigned_rank)
-                        if idx is not None:
-                            assert self._pending_counts[idx] >= 1
-                            self._pending_counts[idx] -= 1
-
-                    self.router_socket.send_multipart(
-                        [
-                            client_identity,
-                            msgpack.packb(
-                                [header.value, client_request_identity, finished_request],
-                                use_bin_type=True,
-                            ),
-                        ]
-                    )
-
-            elif header == Headers.SHUTDOWN:
-                if sender_identity not in known_clients:
-                    logging.warning("Coordinator: ignoring signal from unknown client.")
-                    continue
+            handler = self._handlers.get(header)
+            if handler is None:
+                raise UnknownHeaderError(header)
+            if handler(self, sender_identity, deserialized_payload):
                 break
 
-            elif header == Headers.DISCONNECT:
-                if sender_identity in self.identities_of_data_parallel_ranks:
-                    self._remove_engine(sender_identity)
-
-            else:
-                raise UnknownHeaderError(header)
+    def _handle_rank_registration(self, sender_identity):
+        """Register a data parallel rank that connected to a running coordinator."""
+        if sender_identity not in self.identities_of_data_parallel_ranks:
+            self.identities_of_data_parallel_ranks.append(sender_identity)
+            self._register_rank_identity(sender_identity)
 
     def detokenize(self, finished_request):
         """