NIXL: re-work send timeout tracking on prefill side

In a prefill instance, we need to free KV blocks that have
not been fetched after a timeout. See #20139.

In #26012, we're trying to deal with corner cases involved
with doing this request timeout tracking on the worker side.
This PR proposes moving all of this to the scheduler side,
hopefully making the logic simpler.

Note the expiry timer is switched back to monotonic time
because the timestamp is no longer sent across process
boundaries.

Signed-off-by: Mark McLoughlin <markmc@redhat.com>

Author: markmc

tests/v1/kv_connector/unit/test_nixl_connector.py

@@ -40,7 +40,6 @@
 from vllm.sampling_params import SamplingParams
 from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
 from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput
-from vllm.v1.request import RequestStatus
 
 from .utils import create_request, create_scheduler, create_vllm_config
 
@@ -1078,69 +1077,3 @@ def test_shutdown_cleans_up_resources(dist_init):
         assert mock_dereg.call_count == 2
         mock_dereg.assert_any_call("desc1")
         mock_dereg.assert_any_call("desc2")
-
-
-@patch(
-    "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.NixlWrapper",
-    FakeNixlWrapper,
-)
-def test_aborted_request_removed_from_worker_in_batch(dist_init):
-    """
-    Create and schedule a request so that P adds it to in-batch tracking via
-    the real scheduler, then simulate an abort (request not in next scheduler
-    iteration) and verify the worker no longer tracks it as in-batch.
-    """
-    vllm_config = create_vllm_config()
-
-    scheduler = create_scheduler(vllm_config)
-    # KVConnector Worker in P
-    connector = NixlConnector(vllm_config, KVConnectorRole.WORKER)
-    connector.connector_worker = FakeNixlConnectorWorker(
-        vllm_config, connector.engine_id, hand_shake_latency=0
-    )
-
-    # Create a request that triggers do_remote_decode so that
-    # the scheduler adds it to reqs_in_batch
-    req = create_request(request_id=1, do_remote_decode=True, max_tokens=1)
-    scheduler.add_request(req)
-
-    # First scheduling pass - examinate build_connector_meta output
-    sched_out = scheduler.schedule()
-    kv_meta = sched_out.kv_connector_metadata
-    assert kv_meta is not None
-    assert isinstance(kv_meta, NixlConnectorMetadata)
-    assert req.request_id in kv_meta.reqs_in_batch
-
-    #### Model Runner start ####
-    # Bind scheduler-produced metadata and start worker processing.
-    connector.bind_connector_metadata(kv_meta)
-
-    dummy_ctx = ForwardContext(
-        no_compile_layers={},
-        attn_metadata={},
-        virtual_engine=0,
-    )
-    connector.start_load_kv(dummy_ctx)
-
-    # Ensure it was tracked by the worker
-    assert req.request_id in connector.connector_worker._reqs_to_process
-
-    #### Model Runner end ####
-
-    # Abort request - request_finished call in connector scheduler
-    scheduler.finish_requests(req.request_id, RequestStatus.FINISHED_ABORTED)
-    # Second scheduling pass - build metadata with aborted request
-    sched_out2 = scheduler.schedule()
-    kv_meta2 = sched_out2.kv_connector_metadata
-    assert kv_meta2 is not None
-    assert isinstance(kv_meta2, NixlConnectorMetadata)
-    assert req.request_id not in kv_meta2.reqs_in_batch
-
-    # Bind empty/abort metadata and run worker step
-    #### Model Runner start ####
-    connector.bind_connector_metadata(kv_meta2)
-    connector.start_load_kv(dummy_ctx)
-
-    # After abort, the worker should not keep tracking it as "in-batch"
-    assert req.request_id not in connector.connector_worker._reqs_to_process
-    #### Model Runner end ####

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

@@ -47,6 +47,7 @@
 if TYPE_CHECKING:
     from vllm.attention.backends.abstract import AttentionMetadata
     from vllm.v1.core.kv_cache_manager import KVCacheBlocks
+    from vllm.v1.outputs import KVConnectorOutput
     from vllm.v1.request import Request
 
 Transfer = tuple[int, float]  # (xfer_handle, start_time)
@@ -117,9 +118,6 @@ class NixlConnectorMetadata(KVConnectorMetadata):
     def __init__(self):
         self.reqs_to_recv: dict[ReqId, ReqMeta] = {}
         self.reqs_to_save: dict[ReqId, ReqMeta] = {}
-        self.reqs_to_send: dict[ReqId, float] = {}
-        self.reqs_in_batch: set[ReqId] = set()
-        self.reqs_not_processed: set[ReqId] = set()
 
     def add_new_req(
         self,
@@ -210,6 +208,13 @@ def build_connector_meta(
         assert self.connector_scheduler is not None
         return self.connector_scheduler.build_connector_meta(scheduler_output)
 
+    def update_connector_output(
+        self,
+        connector_output: "KVConnectorOutput",
+    ):
+        assert self.connector_scheduler is not None
+        return self.connector_scheduler.update_connector_output(connector_output)
+
     def request_finished(
         self,
         request: "Request",
@@ -278,6 +283,99 @@ def shutdown(self):
             self.connector_worker.shutdown()
 
 
+class ReqsNeedSendTracker:
+    @dataclass
+    class RequestTimer:
+        """Timer for requests that need to be sent for remote decode."""
+
+        expiry_time: float
+        """Expiry time to avoid stranded KV blocks that are never fetched."""
+        consumer_count: int
+        """Consumer notification count - with heterogeneous TP, P must wait
+        for all assigned D TP workers to finish reading before safely freeing
+        the blocks."""
+
+    def __init__(self):
+        self._reqs_need_send: dict[ReqId, ReqsNeedSendTracker.RequestTimer] = {}
+        self._timeout = envs.VLLM_NIXL_ABORT_REQUEST_TIMEOUT
+
+    def start_timer(self, req_id):
+        self._reqs_need_send[req_id] = self.RequestTimer(
+            time.monotonic() + self._timeout, 0
+        )
+
+    def delete_timer(self, req_id):
+        if req_id not in self._reqs_need_send:
+            return
+        logger.debug("Deleting KV transfer timeout for request %s", req_id)
+        del self._reqs_need_send[req_id]
+
+    def _process_finished_notifs(self, finished_notifs: set[str]) -> set[str]:
+        """Process notifications from D and track consumer completion.
+
+        The notification strings are in format "req_id:tp_ratio".
+
+        Return request IDs that have completed sending to all consumers, to be
+        used by the scheduler via KVConnectorOutput.finished_sending.
+        """
+        finished_sending: set[str] = set()
+        for notif in finished_notifs or ():
+            try:
+                req_id, tp_ratio = notif.rsplit(":", 1)
+            except (ValueError, TypeError) as e:
+                raise ValueError(f"Invalid notification: {notif}") from e
+
+            # Sent notifications received after we already timed out
+            if req_id not in self._reqs_need_send:
+                logger.debug(
+                    "Already finished or expired KV transfer for request %s", req_id
+                )
+                continue
+
+            # Wait all consumers (D) to be done reading before freeing.
+            request_timer = self._reqs_need_send[req_id]
+            request_timer.consumer_count += 1
+            if request_timer.consumer_count < int(tp_ratio):
+                continue
+
+            logger.debug(
+                "KV transfer finished for request %s after retrieval by %d "
+                "decode worker(s).",
+                req_id,
+                request_timer.consumer_count,
+            )
+            del self._reqs_need_send[req_id]
+            finished_sending.add(req_id)
+
+        return finished_sending
+
+    def _abort_expired_requests(self, finished_sending: set[str]) -> set[str]:
+        """Abort requests that have passed their expiry timeout.
+
+        Adds aborted requests to KVConnectorOutput.finished_sending.
+        """
+        now = time.monotonic()
+        while self._reqs_need_send:
+            req_id, request_timer = next(iter(self._reqs_need_send.items()))
+            # Insertion-ordered dict; oldest first so we can exit early.
+            if now < request_timer.expiry_time:
+                break
+            logger.warning(
+                "Releasing expired KV blocks for request %s which were "
+                "retrieved by %d decode worker(s) within %d seconds.",
+                req_id,
+                request_timer.consumer_count,
+                self._timeout,
+            )
+            del self._reqs_need_send[req_id]
+            finished_sending.add(req_id)
+        return finished_sending
+
+    def reqs_finished_sending(self, finished_notifs: set[str]) -> set[str]:
+        finished_sending = self._process_finished_notifs(finished_notifs)
+        return self._abort_expired_requests(finished_sending)
+
+
 class NixlConnectorScheduler:
     """Implementation of Scheduler side methods"""
 
@@ -299,12 +397,8 @@ def __init__(self, vllm_config: VllmConfig, engine_id: str):
         # the scheduler. Used to make metadata passed to Worker.
         self._reqs_need_recv: dict[ReqId, tuple[Request, list[int]]] = {}
         self._reqs_need_save: dict[ReqId, tuple[Request, list[int]]] = {}
-        # Reqs to send and their expiration time
-        self._reqs_need_send: dict[ReqId, float] = {}
-        self._reqs_in_batch: set[ReqId] = set()
-        # Reqs to remove from processed set because they're not to send after
-        # remote prefill or aborted.
-        self._reqs_not_processed: set[ReqId] = set()
+
+        self._reqs_need_send = ReqsNeedSendTracker()
 
     def get_num_new_matched_tokens(
         self, request: "Request", num_computed_tokens: int
@@ -355,8 +449,6 @@ def update_state_after_alloc(
         if not params:
             return
 
-        if params.get("do_remote_decode"):
-            self._reqs_in_batch.add(request.request_id)
         if self.use_host_buffer and params.get("do_remote_decode"):
             # NOTE: when accelerator is not directly supported by Nixl,
             # prefilled blocks need to be saved to host memory before transfer.
@@ -428,19 +520,20 @@ def build_connector_meta(
                 save_to_host=True,
             )
 
-        meta.reqs_to_send = self._reqs_need_send
-        meta.reqs_in_batch = self._reqs_in_batch
-        meta.reqs_not_processed = self._reqs_not_processed
-
         # Clear the list once workers start the transfers
         self._reqs_need_recv.clear()
         self._reqs_need_save.clear()
-        self._reqs_in_batch = set()
-        self._reqs_not_processed = set()
-        self._reqs_need_send = {}
 
         return meta
 
+    def update_connector_output(
+        self,
+        connector_output: "KVConnectorOutput",
+    ):
+        connector_output.finished_sending = self._reqs_need_send.reqs_finished_sending(
+            connector_output.finished_sending
+        )
+
     def request_finished(
         self,
         request: "Request",
@@ -474,10 +567,10 @@ def request_finished(
 
         if not params.get("do_remote_decode"):
             return False, None
+
         if request.status != RequestStatus.FINISHED_LENGTH_CAPPED:
-            # Also include the case of a P/D Prefill request with immediate
-            # block free (eg abort). Stop tracking this request.
-            self._reqs_not_processed.add(request.request_id)
+            # Request aborted after we delayed freeing the blocks?
+            self._reqs_need_send.delete_timer(request.request_id)
             return False, None
 
         # TODO: check whether block_ids actually ever be 0. If not we could
@@ -486,9 +579,7 @@ def request_finished(
 
         if delay_free_blocks:
             # Prefill request on remote. It will be read from D upon completion
-            self._reqs_need_send[request.request_id] = (
-                time.perf_counter() + envs.VLLM_NIXL_ABORT_REQUEST_TIMEOUT
-            )
+            self._reqs_need_send.start_timer(request.request_id)
 
         return delay_free_blocks, dict(
             do_remote_prefill=True,
@@ -609,10 +700,6 @@ def __init__(self, vllm_config: VllmConfig, engine_id: str):
         # [req_id -> list[handle]]
         self._recving_metadata: dict[ReqId, ReqMeta] = {}
         self._recving_transfers = defaultdict[ReqId, list[Transfer]](list)
-        # Track the expiration time of requests that are waiting to be sent.
-        self._reqs_to_send: dict[ReqId, float] = {}
-        # Set of requests that have been part of a batch, regardless of status.
-        self._reqs_to_process: set[ReqId] = set()
 
         # Background thread for handling new handshake requests.
         self._nixl_handshake_listener_t: Optional[threading.Thread] = None
@@ -654,9 +741,6 @@ def __init__(self, vllm_config: VllmConfig, engine_id: str):
         logger.debug("Detected kv cache layout %s", self.kv_cache_layout)
 
         self._tp_size: dict[EngineId, int] = {self.engine_id: self.world_size}
-        # With heterogeneous TP, P must wait for all assigned D TP workers to
-        # finish reading before safely freeing the blocks.
-        self.consumer_notification_counts_by_req = defaultdict[ReqId, int](int)
         self.xfer_stats = NixlKVConnectorStats()
 
     @staticmethod
@@ -1220,25 +1304,6 @@ def get_finished(self) -> tuple[set[str], set[str]]:
                 assert meta, f"{req_id} not found in recving_metadata list"
                 self.sync_recved_kv_to_device(req_id, meta)
 
-        # Handle timeout to avoid stranding blocks on remote.
-        now = time.perf_counter()
-        while self._reqs_to_send:
-            req_id, expires = next(iter(self._reqs_to_send.items()))
-            # Sorted dict, oldest requests are put first so we can exit early.
-            if now < expires:
-                break
-            count = self.consumer_notification_counts_by_req.pop(req_id, 0)
-            logger.warning(
-                "Releasing expired KV blocks for request %s which were "
-                "retrieved by %d decode worker(s) within %d seconds.",
-                req_id,
-                count,
-                envs.VLLM_NIXL_ABORT_REQUEST_TIMEOUT,
-            )
-            self._reqs_to_process.remove(req_id)
-            del self._reqs_to_send[req_id]
-            done_sending.add(req_id)
-
         return done_sending, done_recving
 
     def _get_new_notifs(self) -> set[str]:
@@ -1250,26 +1315,8 @@ def _get_new_notifs(self) -> set[str]:
         notified_req_ids: set[str] = set()
         for notifs in self.nixl_wrapper.get_new_notifs().values():
             for notif in notifs:
-                req_id, tp_ratio = notif.decode("utf-8").rsplit(":", 1)
-                if (
-                    req_id not in self._reqs_to_send
-                    and req_id not in self._reqs_to_process
-                ):
-                    logger.error(
-                        "Potentially invalid KV blocks for "
-                        "unrecognized request %s were retrieved by "
-                        "a decode worker. They may have expired.",
-                        req_id,
-                    )
-                    continue
-
-                self.consumer_notification_counts_by_req[req_id] += 1
-                # Wait all consumers (D) to be done reading before freeing.
-                if self.consumer_notification_counts_by_req[req_id] == int(tp_ratio):
-                    notified_req_ids.add(req_id)
-                    del self.consumer_notification_counts_by_req[req_id]
-                    self._reqs_to_process.remove(req_id)
-                    self._reqs_to_send.pop(req_id, None)
+                # Note - this is in req_id:tp_ratio format
+                notified_req_ids.add(notif.decode("utf-8"))
         return notified_req_ids
 
     def _pop_done_transfers(
@@ -1333,24 +1380,6 @@ def start_load_kv(self, metadata: NixlConnectorMetadata):
         while not self._ready_requests.empty():
             self._read_blocks_for_req(*self._ready_requests.get_nowait())
 
-        # Keep around the requests that have been part of a batch. This is
-        # needed because async scheduling pushes the misalignment between the
-        # moment in which requests expiration is set (P side) and the moment in
-        # which blocks are read from D. As P can now more easily lag behind D
-        # while processing the next batch, we make sure to only set an
-        # expiration for requests that have not been read from D yet.
-        for req_id in metadata.reqs_in_batch:
-            self._reqs_to_process.add(req_id)
-
-        # Remove all requests that are not to be processed (eg aborted).
-        for req_id in metadata.reqs_not_processed:
-            self._reqs_to_process.discard(req_id)
-
-        # Add to requests that are waiting to be read and track expiration.
-        for req_id, expiration_time in metadata.reqs_to_send.items():
-            if req_id in self._reqs_to_process:
-                self._reqs_to_send[req_id] = expiration_time
-
     def _read_blocks_for_req(self, req_id: str, meta: ReqMeta):
         logger.debug(
             "Remote agent %s available, calling _read_blocks for req %s",