Change ZMQ communication to use async ZMQ

Author: tdene

examples/inference/gpt/gpt_dynamic_inference_with_coordinator.py

@@ -39,7 +39,7 @@ async def main(
     # once you call engine.start_listening_to_data_parallel_coordinator,
     # the engine will start accepting requests from the data parallel coordinator.
     # and processing them in an asyncio coroutine. 
-    await engine.start_listening_to_data_parallel_coordinator( 
+    engine.start_listening_to_data_parallel_coordinator(
         inference_coordinator_port=port, launch_inference_coordinator=True
     )
     # if you want to use your own inference coordinator - 
@@ -51,7 +51,7 @@ async def main(
     # 5. look at InferenceClient to see how we create requests with headers. 
     if dist.get_rank() == 0: 
         client = InferenceClient(port) # submits requests to the inference coordinator
-        await client.start()
+        client.start()
         base_arrival_time = time.time_ns() / 10**9
         for request in requests:
             request.time_arrival = request.time_offset + base_arrival_time

megatron/core/inference/data_parallel_inference_coordinator.py

@@ -1,18 +1,22 @@
 # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
+import asyncio
 import faulthandler
 import logging
 import signal
 from collections import deque
 from itertools import cycle
 from multiprocessing import Event
+from typing import List, Optional, Tuple
 
 import torch
 
 from megatron.core.inference.headers import Headers
+from megatron.core.utils import get_asyncio_loop, trace_async_exceptions
 
 try:
     import zmq
+    import zmq.asyncio
 
     HAVE_ZMQ = True
 except:
@@ -25,11 +29,6 @@
 except:
     HAVE_MSGPACK = False
 
-# Register faulthandler to emit stack traces upon process kill.
-faulthandler.enable()
-faulthandler.register(signal.SIGTERM, all_threads=False, chain=True)
-faulthandler.register(signal.SIGINT, all_threads=False, chain=True)
-
 
 class DataParallelInferenceCoordinator:
     """
@@ -65,7 +64,9 @@ class DataParallelInferenceCoordinator:
         next_request_id (int): A counter for generating unique server-side request IDs.
     """
 
-    def __init__(self, inference_coordinator_port: int, data_parallel_size: int):
+    def __init__(
+        self, ready_event: Event, inference_coordinator_port: int, data_parallel_size: int
+    ):
         """
         Initializes the inference coordinator.
 
@@ -74,6 +75,8 @@ def __init__(self, inference_coordinator_port: int, data_parallel_size: int):
         ranks to connect before proceeding.
 
         Args:
+            ready_event (Event): A threading or multiprocessing event object that is set()
+                once the coordinator is ready to accept connections.
             inference_coordinator_port (int): The TCP port number to bind the server to.
             data_parallel_size (int): The number of TP-coordinator workers that are
                 expected to connect.
@@ -86,7 +89,10 @@ def __init__(self, inference_coordinator_port: int, data_parallel_size: int):
             "please install the messagepack library to use DataParallelInferenceCoordinator\n"
             "pip install msgpack"
         )
-        self.context = zmq.Context()
+        self.ready_event = ready_event
+        self.data_parallel_size = data_parallel_size
+
+        self.context = zmq.asyncio.Context.instance()
 
         # This is the central router socket
         # 1. data parallel ranks connect to this socket to register themselves
@@ -96,24 +102,8 @@ def __init__(self, inference_coordinator_port: int, data_parallel_size: int):
         #    the user that had submitted the request originally.
 
         self.router_socket = self.context.socket(zmq.ROUTER)
+        self.socket_uses_identity = True
         self.router_socket.bind(f"tcp://0.0.0.0:{inference_coordinator_port}")
-        self.data_parallel_size = data_parallel_size
-
-        logging.info("Inference Coordinator: waiting for connections from data parallel ranks...")
-        # First wait for all data parallel ranks to establish connections.
-        self.identities_of_data_parallel_ranks = deque([])
-        # time.sleep(5)  # Give data parallel ranks time to spawn and connect.
-        for _ in range(data_parallel_size):
-            identity, _ = self.router_socket.recv_multipart()
-            assert identity not in self.identities_of_data_parallel_ranks
-            self.identities_of_data_parallel_ranks.append(identity)
-        logging.info("Inference Coordinator: Connected with data parallel ranks...")
-        self.data_parallel_rank_iterator = cycle(self.identities_of_data_parallel_ranks)
-
-        self.request_id_to_client_id = {}
-        self.request_id_to_client_request_id = {}
-
-        self.next_request_id = 0
 
     def get_next_data_parallel_rank(self):
         """
@@ -124,7 +114,7 @@ def get_next_data_parallel_rank(self):
         """
         return next(self.data_parallel_rank_iterator)
 
-    def start(self):
+    def start(self, loop: Optional[asyncio.AbstractEventLoop] = None):
         """
         Starts the main event loop for the coordinator.
 
@@ -134,45 +124,70 @@ def start(self):
         handling new client connections, forwarding requests, broadcasting
         control signals, or processing replies from the engines.
         """
+        logging.info("Inference Coordinator: waiting for connections from data parallel ranks...")
+        loop = get_asyncio_loop(loop)
+
+        self.ready_event.clear()
+        self.identities_of_data_parallel_ranks = deque([])
+        self.data_parallel_rank_iterator = cycle([])
+        self.known_clients = set()
+        self.request_id_to_client_id = {}
+        self.request_id_to_client_request_id = {}
+        self.next_request_id = 0
+
+        self._send_awaitables = asyncio.Queue()
+        self.send_task = loop.create_task(self._send_task())
+        self.recv_task = loop.create_task(self._recv_task())
+
+    @trace_async_exceptions
+    async def _recv_task(self):
+        """Main loop of the inference coordinator."""
+
+        print("Inference Coordinator: waiting for connections from data parallel ranks...")
+        # First wait for all data parallel ranks to establish connections.
+        for _ in range(self.data_parallel_size):
+            identity, header, _ = await self._irecv()
+            assert header == Headers.CONNECT
+            assert identity not in self.identities_of_data_parallel_ranks
+            self.identities_of_data_parallel_ranks.append(identity)
+            print(f"Inference Coordinator: Data parallel rank connected: {identity}")
+        print("All data parallel ranks connected.")
+        logging.info("Inference Coordinator: Connected with data parallel ranks...")
+        self.data_parallel_rank_iterator = cycle(self.identities_of_data_parallel_ranks)
+        self.ready_event.set()
+        print("Inference Coordinator: Ready to accept client connections.")
+
         # Todo [Siddharth]: Make this more robust to handle invalid messages.
-        known_clients = set()
         while True:
-            sender_identity, serialized_payload = self.router_socket.recv_multipart()
-            deserialized_payload = msgpack.unpackb(serialized_payload, raw=False)
-            header = Headers(deserialized_payload[0])
+            identity, header, data = await self._irecv()
 
             if header == Headers.CONNECT:
-                if sender_identity in known_clients:
+                if identity in self.known_clients:
                     logging.info(
-                        f"Client {sender_identity} sent a duplicate connect request. Ignoring .."
+                        f"Client {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.ACK.value], use_bin_type=True)]
-                )
+                self.known_clients.add(identity)
+                self._isend(identity, Headers.ACK)
 
             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."
-                    )
+                if identity not in self.known_clients:
+                    logging.info(f"Received message from unknown client {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:]
+                client_request_id, prompt, sampling_params = data
                 # 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_id[request_id] = identity
                 self.request_id_to_client_request_id[request_id] = client_request_id
 
                 # Serialize prompt.
@@ -184,28 +199,22 @@ def start(self):
                     raise Exception("specialize for <%s> prompt." % type(prompt).__name__)
 
                 next_data_parallel_rank_identity = self.get_next_data_parallel_rank()
-                self.router_socket.send_multipart(
-                    [
-                        next_data_parallel_rank_identity,
-                        msgpack.packb(
-                            [Headers.SUBMIT_REQUEST.value, request_id, prompt, sampling_params],
-                            use_bin_type=True,
-                        ),
-                    ]
+                self._isend(
+                    next_data_parallel_rank_identity,
+                    Headers.SUBMIT_REQUEST,
+                    [request_id, prompt, sampling_params],
                 )
             elif header in [Headers.PAUSE, Headers.UNPAUSE, Headers.STOP]:
                 # control signals for the engine
                 # broadcast to all data parallel ranks
-                if sender_identity not in known_clients:
+                if identity not in self.known_clients:
                     continue
                 for data_parallel_rank_id in self.identities_of_data_parallel_ranks:
-                    self.router_socket.send_multipart(
-                        [data_parallel_rank_id, msgpack.packb([header.value], use_bin_type=True)]
-                    )
+                    self._isend(data_parallel_rank_id, header)
             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]
+                assert identity in self.identities_of_data_parallel_ranks
+                finished_requests = data
 
                 for finished_request in finished_requests:
                     fid = finished_request["request_id"]
@@ -214,15 +223,68 @@ def start(self):
                     del self.request_id_to_client_id[fid]
                     del self.request_id_to_client_request_id[fid]
 
-                    self.router_socket.send_multipart(
-                        [
-                            client_identity,
-                            msgpack.packb(
-                                [client_request_identity, finished_request], use_bin_type=True
-                            ),
-                        ]
+                    self._isend(
+                        client_identity,
+                        Headers.ENGINE_REPLY,
+                        [client_request_identity, finished_request],
                     )
 
+    @trace_async_exceptions
+    async def _send_task(self):
+        """Pop futures of sends out of a queue and await them.
+
+        For explanation why this works, refer to the documentation for zmq.asyncio:
+            'Returns a Future that resolves when sending is complete.'
+        """
+        while True:
+            await (await self._send_awaitables.get())
+            self._send_awaitables.task_done()
+
+    def _isend(
+        self, identity: bytes, header: Headers, data: Optional[List] = None
+    ) -> asyncio.Future:
+        """
+        Asynchronously send a signal to the inference coordinator.
+
+        Args:
+            identity (bytes): The ZMQ identity of the recipient.
+            header (Headers): The signal header to send.
+            data (Optional[List]): The data payload to send.
+        """
+        to_send = [identity, header.value.to_bytes()]
+        if data is not None:
+            to_send.append(msgpack.packb(data, use_bin_type=True))
+        send_awaitable = self.router_socket.send_multipart(to_send)
+        self._send_awaitables.put_nowait(send_awaitable)
+
+    async def _irecv(
+        self, deserialize: bool = True
+    ) -> Tuple[Optional[bytes], Headers, List | bytes | None]:
+        """
+        Asynchronously receive a signal from the inference coordinator.
+
+        Returns:
+            identity (Optional[bytes]): The source of the signal.
+            header (Headers): The signal header received.
+            data (List | bytes | None): The data payload received.
+        """
+        raw = await self.router_socket.recv_multipart()
+        if self.socket_uses_identity:
+            identity, header, *rest = raw
+        else:
+            header, *rest = raw
+            identity = None
+
+        header = Headers(int.from_bytes(header))
+        data = rest[0] if rest else None
+
+        if deserialize:
+            message = msgpack.unpackb(data, raw=False) if data is not None else None
+        else:
+            message = data
+
+        return identity, header, message
+
     @classmethod
     def entrypoint(
         cls, ready_event: Event, inference_coordinator_port: int, data_parallel_size: int
@@ -239,17 +301,29 @@ def entrypoint(
             inference_coordinator_port (int): The port to bind to.
             data_parallel_size (int): The number of expected TP-coordinators.
         """
-        coordinator = cls(inference_coordinator_port, data_parallel_size)
-        ready_event.set()
+        # Register faulthandler to emit stack traces upon process kill.
+        faulthandler.enable()
+        faulthandler.register(signal.SIGTERM, all_threads=False, chain=True)
+        faulthandler.register(signal.SIGINT, all_threads=False, chain=True)
+
+        print("Inference Coordinator: Initializing coordinator...")
+        coordinator = cls(ready_event, inference_coordinator_port, data_parallel_size)
+        print("Inference Coordinator: Starting coordinator...")
+        loop = get_asyncio_loop()
+        coordinator.start(loop=loop)
+        print("Inference Coordinator: Coordinator started.")
         try:
-            coordinator.start()
+            loop.run_forever()
         except KeyboardInterrupt:
             logging.info("Coordinator process interrupted. Exiting...")
+        finally:
             coordinator.stop()
 
     def stop(self):
         """
         Stops the inference coordinator, performing any necessary cleanup operations.
         """
+        self.send_task.cancel()
+        self.recv_task.cancel()
         self.router_socket.close()
         self.context.term()