[Bugfix][KV Connector] Reap expired NIXL leases behind a heartbeated head (#58292)
Signed-off-by: GokayAI <60583610+gokay-ai@users.noreply.github.com> Co-authored-by: GokayAI <gokay-ai@users.noreply.github.com>
GokayAI committed
Sep 25, 2026 at 20:23 UTC
c89b5b4176066d689a5fe40b0bea8b425ac0c65b
2 files changed
+117
-18
tests/v1/kv_connector/unit/test_nixl_heartbeat.py
+88
@@ -30,6 +30,10 @@ def _worker_stub():
30
w = object.__new__(NixlConnectorWorker)
31
w._reqs_to_send = {}
32
w._lease_extension = 20
33
+ w._reqs_to_process = set()
34
+ w.consumer_notification_counts_by_req = {}
35
+ w.expected_consumer_notifications_by_req = {}
36
+ w.xfer_stats = MagicMock()
37
return w
38
39
@@ -163,3 +167,87 @@ def test_handle_heartbeat():
167
assert w._reqs_to_send["req-b"] >= far_future
168
# req-unknown: not added.
169
assert "req-unknown" not in w._reqs_to_send
170
+
171
+
172
+def test_reaper_reclaims_expired_leases_behind_heartbeated_head(monkeypatch):
173
+ """A heartbeated head entry must not strand expired leases behind it.
174
+
175
+ Mirrors the stdlib repro in https://github.com/vllm-project/vllm/issues/58222:
176
+ A is inserted first and kept alive via heartbeats; B and C expire and
177
+ must still be reaped. The old early-exit scan returned nothing here.
178
+ """
179
+ w = _worker_stub()
180
+ w._reqs_to_process.update(("A", "B", "C"))
181
+ # Insertion order is A, B, C. Heartbeats will keep A at the head.
182
+ w._reqs_to_send = {"A": 30.0, "B": 32.0, "C": 34.0}
183
+
184
+ from vllm.distributed.kv_transfer.kv_connector.v1.nixl import base_worker
185
+
186
+ now = {"t": 0.0}
187
+ monkeypatch.setattr(base_worker.time, "perf_counter", lambda: now["t"])
188
+ for t in (20.0, 40.0, 60.0, 80.0):
189
+ now["t"] = t
190
+ w._handle_heartbeat("A")
191
+
192
+ # A was renewed in place: still the dict head, expiry 100.
193
+ assert next(iter(w._reqs_to_send)) == "A"
194
+ assert w._reqs_to_send["A"] == pytest.approx(100.0)
195
+
196
+ done: set[str] = set()
197
+ now["t"] = 50.0
198
+ w._reap_expired_send_leases(done)
199
+
200
+ assert done == {"B", "C"}
201
+ assert list(w._reqs_to_send) == ["A"]
202
+ assert w._reqs_to_process == {"A"}
203
+ assert w.xfer_stats.record_kv_expired_req.call_count == 2
204
+
205
+ # now == expires is expired (now >= expires): A is due at 100.
206
+ done_eq: set[str] = set()
207
+ now["t"] = 100.0
208
+ w._reap_expired_send_leases(done_eq)
209
+ assert done_eq == {"A"}
210
+ assert not w._reqs_to_send
211
+ assert not w._reqs_to_process
212
+ assert w.xfer_stats.record_kv_expired_req.call_count == 3
213
+
214
+
215
+def test_reaper_reclaims_when_now_equals_expiry(monkeypatch):
216
+ """The reaper must reclaim on the now == expires boundary."""
217
+ w = _worker_stub()
218
+ w._reqs_to_process.add("eq")
219
+ w._reqs_to_send = {"eq": 50.0}
220
+
221
+ from vllm.distributed.kv_transfer.kv_connector.v1.nixl import base_worker
222
+
223
+ monkeypatch.setattr(base_worker.time, "perf_counter", lambda: 50.0)
224
+
225
+ done: set[str] = set()
226
+ w._reap_expired_send_leases(done)
227
+
228
+ assert done == {"eq"}
229
+ assert not w._reqs_to_send
230
+ assert not w._reqs_to_process
231
+ w.xfer_stats.record_kv_expired_req.assert_called_once()
232
+
233
+
234
+def test_reaper_reclaims_shorter_lease_behind_later_deadline(monkeypatch):
235
+ """A later-inserted longer TTL must not block an earlier-expiring entry.
236
+
237
+ Delete+reinsert on heartbeat cannot restore expiry order in this case
238
+ (issue #58222), so the reaper must scan the whole map.
239
+ """
240
+ w = _worker_stub()
241
+ w._reqs_to_process.update(("X", "Y"))
242
+ w._reqs_to_send = {"X": 55.0, "Y": 46.0}
243
+
244
+ from vllm.distributed.kv_transfer.kv_connector.v1.nixl import base_worker
245
+
246
+ monkeypatch.setattr(base_worker.time, "perf_counter", lambda: 50.0)
247
+
248
+ done: set[str] = set()
249
+ w._reap_expired_send_leases(done)
250
+
251
+ assert done == {"Y"}
252
+ assert list(w._reqs_to_send) == ["X"]
253
+ assert w._reqs_to_process == {"X"}
vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py
+29
-18
@@ -2951,24 +2951,7 @@ class NixlBaseConnectorWorker:
2951
self._sync_device_after_direct_recv(direct_device_recving)
2952
2953
# Handle timeout to avoid stranding blocks on remote.
2954
- now = time.perf_counter()
2955
- while self._reqs_to_send:
2956
- req_id, expires = next(iter(self._reqs_to_send.items()))
2957
- # Sorted dict, oldest requests are put first so we can exit early.
2958
- if now < expires:
2959
- break
2960
- count = self.consumer_notification_counts_by_req.pop(req_id, 0)
2961
- self.expected_consumer_notifications_by_req.pop(req_id, None)
2962
- self.xfer_stats.record_kv_expired_req()
2963
- logger.warning(
2964
- "Releasing expired KV blocks for request %s which were "
2965
- "retrieved by %d remote worker(s) before lease expired.",
2966
- req_id,
2967
- count,
2968
- )
2969
- self._reqs_to_process.remove(req_id)
2970
- del self._reqs_to_send[req_id]
2971
- done_sending.add(req_id)
2954
+ self._reap_expired_send_leases(done_sending)
2955
2956
return KVConnectorTransferResults(
2957
finished_sending=done_sending,
@@ -2996,6 +2979,34 @@ class NixlBaseConnectorWorker:
2979
"""
2980
raise NotImplementedError
2981
2982
+ def _reap_expired_send_leases(self, done_sending: set[str]) -> None:
2983
+ """Reclaim expired send-side KV leases into ``done_sending``.
2984
+
2985
+ ``_reqs_to_send`` is not ordered by expiry: heartbeats update the
2986
+ deadline in place, and mixed TTLs share the map, so a live head
2987
+ entry can sit in front of already-expired ones. Scan every entry
2988
+ rather than stopping at the first still-live request.
2989
+ """
2990
+ if not self._reqs_to_send:
2991
+ return
2992
+ now = time.perf_counter()
2993
+ expired = [
2994
+ req_id for req_id, expires in self._reqs_to_send.items() if now >= expires
2995
+ ]
2996
+ for req_id in expired:
2997
+ count = self.consumer_notification_counts_by_req.pop(req_id, 0)
2998
+ self.expected_consumer_notifications_by_req.pop(req_id, None)
2999
+ self.xfer_stats.record_kv_expired_req()
3000
+ logger.warning(
3001
+ "Releasing expired KV blocks for request %s which were "
3002
+ "retrieved by %d remote worker(s) before lease expired.",
3003
+ req_id,
3004
+ count,
3005
+ )
3006
+ self._reqs_to_process.remove(req_id)
3007
+ del self._reqs_to_send[req_id]
3008
+ done_sending.add(req_id)
3009
+
3010
def _handle_heartbeat(self, payload: str) -> None:
3011
"""Extend leases for requests referenced in a heartbeat.
3012