Restore blocking CONNECT functionality

Author: tdene

megatron/core/inference/data_parallel_inference_coordinator.py

@@ -134,8 +134,12 @@ def start(self, loop: Optional[asyncio.AbstractEventLoop] = None):
         self.request_id_to_client_id = {}
         self.request_id_to_client_request_id = {}
         self.next_request_id = 0
-
         self._send_awaitables = asyncio.Queue()
+
+        # Attempt to connect, and do not allow any sends until we are connected.
+        self.is_running = asyncio.Event()
+
+        self.startup_sends_task = loop.create_task(self._startup_sends_task())
         self.send_task = loop.create_task(self._send_task())
         self.recv_task = loop.create_task(self._recv_task())
 
@@ -144,31 +148,30 @@ 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.
         while True:
             identity, header, data = await self._irecv()
 
-            if header == Headers.CONNECT:
+            if header == Headers.ENGINE_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}")
+                if len(self.identities_of_data_parallel_ranks) == self.data_parallel_size:
+                    self.data_parallel_rank_iterator = cycle(self.identities_of_data_parallel_ranks)
+                    self.ready_event.set()
+                    self.is_running.set()
+                    print("All data parallel ranks connected.")
+                    logging.info("Inference Coordinator: Connected with data parallel ranks...")
+                    print("Inference Coordinator: Ready to accept client connections.")
+
+            elif header == Headers.CLIENT_CONNECT:
                 if identity in self.known_clients:
                     logging.info(
                         f"Client {identity} sent a duplicate connect request. Ignoring .."
                     )
                     continue
-
                 self.known_clients.add(identity)
+                # Due to the `startup_sends` logic, this will not be sent until we are connected.
                 self._isend(identity, Headers.ACK)
 
             elif header == Headers.SUBMIT_REQUEST:
@@ -240,6 +243,13 @@ async def _send_task(self):
             await (await self._send_awaitables.get())
             self._send_awaitables.task_done()
 
+    @trace_async_exceptions
+    async def _startup_sends_task(self):
+        """Before a connection is established, we queue up sends for later."""
+        await self.is_running()
+        for (header, data) in self._startup_sends:
+            self._isend(header, data)
+
     def _isend(
         self, identity: bytes, header: Headers, data: Optional[List] = None
     ) -> asyncio.Future:
@@ -251,6 +261,12 @@ def _isend(
             header (Headers): The signal header to send.
             data (Optional[List]): The data payload to send.
         """
+        # If we have not connected yet, wait on sends.
+        if not self.is_running.is_set():
+            self._startup_sends.append((identity, header, data))
+            return
+
+        # Once we are connected, we do an atomic send and await its completion later.
         to_send = [identity, header.value.to_bytes()]
         if data is not None:
             to_send.append(msgpack.packb(data, use_bin_type=True))

megatron/core/inference/engines/dynamic_engine.py

@@ -416,8 +416,8 @@ def start_listening_to_data_parallel_coordinator(
             self.socket_for_receiving_requests.setsockopt(zmq.IDENTITY, identity.encode('utf-8'))
             self.socket_for_receiving_requests.connect(dp_addr[0])
 
-            # send empty string. this is used to register with the coordinator.
-            self._isend(self.socket_for_receiving_requests, Headers.CONNECT, b"")
+            # Register with the coordinator.
+            self._isend(self.socket_for_receiving_requests, Headers.ENGINE_CONNECT)
 
             # 2. Create a publisher socket. This is used to publish or broadcast
             #    requests within the model parallel group

megatron/core/inference/headers.py

@@ -8,11 +8,12 @@ class Headers(Enum):
     Enum representing headers used for communication with the inference-coordinator.
     """
 
-    CONNECT = 0
-    ACK = 1
-    MICROBATCH_SYNC = 2
-    SUBMIT_REQUEST = 3
-    ENGINE_REPLY = 4
-    PAUSE = 5
-    UNPAUSE = 6
-    STOP = 7
+    ENGINE_CONNECT = 0
+    CLIENT_CONNECT = 1
+    ACK = 2
+    MICROBATCH_SYNC = 3
+    SUBMIT_REQUEST = 4
+    ENGINE_REPLY = 5
+    PAUSE = 6
+    UNPAUSE = 7
+    STOP = 8

megatron/core/inference/inference_client.py

@@ -121,7 +121,7 @@ async def _recv_task(self):
             try:
                 _, header, data = await self._irecv()
 
-                assert header == Headers.ACK or self.initial_reply
+                assert header == Headers.ACK or self.is_running.is_set()
                 if header == Headers.ENGINE_REPLY:
                     request_id, reply = data
                     reply['latency'] = time.perf_counter() - self.request_submission_times.pop(
@@ -130,7 +130,7 @@ async def _recv_task(self):
                     completion_future = self.completion_futures.pop(request_id)
                     completion_future.set_result(DynamicInferenceRequest.deserialize(reply))
                 elif header == Headers.ACK:
-                    self.initial_reply = True
+                    self.is_running.set()
             except asyncio.CancelledError:
                 break
 
@@ -149,9 +149,12 @@ def start(self, loop: Optional[asyncio.AbstractEventLoop] = None):
         self.next_request_id = 0
         self._send_awaitables = asyncio.Queue()
 
-        self.initial_reply = False
-        self._isend(Headers.CONNECT)
+        # Attempt to connect, and do not allow any sends until we are connected.
+        self.is_running = asyncio.Event()
+        self._startup_sends = []
+        self._isend(Headers.CLIENT_CONNECT)
 
+        self.startup_sends_task = loop.create_task(self._startup_sends_task())
         self.send_task = loop.create_task(self._send_task())
         self.recv_task = loop.create_task(self._recv_task())
 
@@ -166,6 +169,13 @@ async def _send_task(self):
             await (await self._send_awaitables.get())
             self._send_awaitables.task_done()
 
+    @trace_async_exceptions
+    async def _startup_sends_task(self):
+        """Before a connection is established, we queue up sends for later."""
+        await self.is_running()
+        for (header, data) in self._startup_sends:
+            self._isend(header, data)
+
     def _isend(self, header: Headers, data: Optional[List] = None) -> asyncio.Future:
         """
         Asynchronously send a signal to the inference coordinator.
@@ -174,6 +184,12 @@ def _isend(self, header: Headers, data: Optional[List] = None) -> asyncio.Future
             header (Headers): The signal header to send.
             data (Optional[List]): The data payload to send.
         """
+        # If we have not connected yet, wait on sends.
+        if not self.is_running.is_set():
+            self._startup_sends.append((header, data))
+            return
+
+        # Once we are connected, we do an atomic send and await its completion later.
         to_send = [header.value.to_bytes()]
         if data is not None:
             to_send.append(msgpack.packb(data, use_bin_type=True))