fix(connector): adapt handle_preemptions to vLLM #34805 interface change

python/pegaflow/connector/__init__.py

@@ -186,9 +186,15 @@ def unregister_context(self) -> None:
         if self._worker:
             self._worker.unregister_context()
 
-    def handle_preemptions(self, preempted_req_ids: set[str]) -> None:
-        if self._worker:
-            self._worker.handle_preemptions(preempted_req_ids)
+    def handle_preemptions(self, preempted) -> None:
+        if not self._worker:
+            return
+        # Compat: old vLLM passes set[str], new vLLM passes KVConnectorMetadata
+        if isinstance(preempted, set):
+            preempted_req_ids = preempted
+        else:
+            preempted_req_ids = getattr(preempted, "preempted_req_ids", None) or set()
+        self._worker.handle_preemptions(preempted_req_ids)
 
     # ==============================
     # Scheduler-side methods

python/pegaflow/connector/common.py

@@ -98,11 +98,13 @@ def __init__(
         self,
         load_intents: dict[str, LoadIntent] | None = None,
         save_intents: dict[str, SaveIntent] | None = None,
+        preempted_req_ids: set[str] | None = None,
     ):
         super().__init__()
         # Maps request_id -> intent
         self.load_intents: dict[str, LoadIntent] = load_intents or {}
         self.save_intents: dict[str, SaveIntent] = save_intents or {}
+        self.preempted_req_ids: set[str] = preempted_req_ids or set()
 
     def __repr__(self) -> str:
         return (

python/pegaflow/connector/scheduler.py

@@ -292,6 +292,7 @@ def build_connector_meta(self, scheduler_output: "SchedulerOutput") -> PegaConne
         return PegaConnectorMetadata(
             load_intents=load_intents,
             save_intents=save_intents,
+            preempted_req_ids=scheduler_output.preempted_req_ids or None,
         )
 
     def _consume_save_intent(self, req_id: str) -> SaveIntent | None:

python/pegaflow/connector/worker.py

@@ -571,7 +571,7 @@ def _handle_save_completion(
                 suffix,
             )
 
-    def handle_preemptions(self, preempted_req_ids: set[str]) -> None:
+    def handle_preemptions(self, preempted_req_ids: set[str] | None) -> None:
         """Wait for preempted requests' saves to complete before blocks are reused.
 
         Called by vLLM BEFORE preempted blocks are overwritten. This prevents