Merge pull request #89 from njhill/pdl-fix-send-rcv-order

[BugFix] Fix ordering of KVConnector finished send/rcv sets

Author: njhill

vllm/distributed/kv_transfer/kv_connector/v1/base.py

@@ -183,7 +183,8 @@ def get_finished(
         finished generating tokens.
 
         Returns:
-            ids of requests that have finished asynchronous (recving, sending).
+            ids of requests that have finished asynchronous transfer,
+            tuple of (sending/saving ids, recving/loading ids).
             The finished saves/sends req ids must belong to a set provided in a
             call to this method (this call or a prior one).
         """

vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py

@@ -89,10 +89,10 @@ def wait_for_save(self):
     def get_finished(
         self, finished_req_ids: set[str]
     ) -> tuple[Optional[set[str]], Optional[set[str]]]:
-        finished_recving: set[str] = set()
         finished_sending: set[str] = set()
+        finished_recving: set[str] = set()
         for c in self._connectors:
-            recving, sending = c.get_finished(finished_req_ids)
+            sending, recving = c.get_finished(finished_req_ids)
             if not recving and not sending:
                 continue
             # Aggregate finished recving request ids.
@@ -111,7 +111,7 @@ def get_finished(
                 else:
                     self._extra_async_saves[req_id] = extra_pending - 1
 
-        return finished_recving or None, finished_sending or None
+        return finished_sending or None, finished_recving or None
 
     # ==============================
     # Scheduler-side methods