[PD][Bugfix] Rebase KV lease deadlines onto worker clock (#50326)
Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: esmeetu <jasonailu87@gmail.com> Co-authored-by: Yifan Qiao <yifanqiao@inferact.ai>
Nick Hill committed
Jul 29, 2026 at 18:25 UTC
4e582c5b54b866a2cc0d44c8f8bbedb52cf5c605
5 files changed
+76
tests/v1/kv_connector/unit/test_nixl_connector.py
+47
@@ -1398,6 +1398,53 @@ def test_kv_connector_stats(default_vllm_config, dist_init):
1398
assert stats_after_reset is None
1399
1400
1401
+def test_reqs_to_send_deadline_rebased_to_worker_clock(default_vllm_config, dist_init):
1402
+ """reqs_to_send deadlines are stamped with the scheduler process's
1403
+ perf_counter, whose epoch differs across processes and (by boot-time
1404
+ deltas) across nodes. Without rebasing, a P worker on a node whose
1405
+ monotonic clock is ahead of the scheduler's by more than the TTL
1406
+ expires the lease on arrival and reports done_sending before D has
1407
+ read the blocks — the freed blocks can then be reallocated and the
1408
+ remote read pulls another request's data (silent accuracy corruption).
1409
+ The worker must anchor the remaining TTL to its own clock.
1410
+ """
1411
+ vllm_config = create_vllm_config()
1412
+ connector = NixlConnector(
1413
+ vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16)
1414
+ )
1415
+ connector.connector_worker = FakeNixlConnectorWorker(
1416
+ vllm_config, connector.engine_id, hand_shake_latency=0
1417
+ )
1418
+ worker = connector.connector_worker
1419
+
1420
+ req_id = "req-lease-clock"
1421
+ ttl = 480.0
1422
+ # Simulate a scheduler whose monotonic clock is 10,000 s behind this
1423
+ # worker's (e.g. its node booted much later): the raw deadline is
1424
+ # then already far in the past in this worker's clock domain.
1425
+ scheduler_clock = time.perf_counter() - 10_000.0
1426
+
1427
+ metadata = NixlConnectorMetadata()
1428
+ metadata.reqs_in_batch = {req_id}
1429
+ metadata.reqs_to_send = {req_id: scheduler_clock + ttl}
1430
+ metadata.scheduler_clock = scheduler_clock
1431
+ connector.bind_connector_metadata(metadata)
1432
+ dummy_ctx = ForwardContext(
1433
+ no_compile_layers={},
1434
+ attn_metadata={},
1435
+ slot_mapping={},
1436
+ )
1437
+ connector.start_load_kv(dummy_ctx)
1438
+
1439
+ remaining = worker._reqs_to_send[req_id] - time.perf_counter()
1440
+ assert ttl - 5.0 < remaining <= ttl + 5.0
1441
+
1442
+ # The expiry sweep must not release the request.
1443
+ done_sending, _ = worker.get_finished()
1444
+ assert req_id not in done_sending
1445
+ assert req_id in worker._reqs_to_process
1446
+
1447
+
1448
def test_kv_connector_stats_aggregation():
1449
"""
1450
Test KV transfer stats aggregation across TP ranks using
vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py
+4
@@ -453,6 +453,10 @@ class NixlBaseConnectorScheduler:
453
self._build_save_meta(meta, scheduler_output)
454
455
meta.reqs_to_send = self._reqs_need_send
456
+ # Clock reference for reqs_to_send: deadlines above are in this
457
+ # process's perf_counter domain; workers (possibly on other nodes,
458
+ # where perf_counter has a different epoch) rebase against this.
459
+ meta.scheduler_clock = time.perf_counter()
460
meta.reqs_in_batch = self._reqs_in_batch
461
meta.reqs_not_processed = self._reqs_not_processed
462
vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py
+6
@@ -180,6 +180,12 @@ class NixlConnectorMetadata(KVConnectorMetadata):
180
self.reqs_to_recv: dict[ReqId, ReqMeta] = {}
181
self.reqs_to_save: dict[ReqId, ReqMeta] = {}
182
self.reqs_to_send: dict[ReqId, float] = {}
183
+ # The scheduler process's time.perf_counter() when this metadata was
184
+ # built. reqs_to_send deadlines are stamped with the scheduler's
185
+ # clock, which is NOT comparable across processes (perf_counter is
186
+ # process/boot-local): workers must rebase the remaining TTL onto
187
+ # their own clock via this reference. 0.0 = unset (legacy metadata).
188
+ self.scheduler_clock: float = 0.0
189
self.reqs_in_batch: set[ReqId] = set()
190
self.reqs_not_processed: set[ReqId] = set()
191
# Heartbeat data grouped by remote engine, sent by D worker to P.
vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py
+12
@@ -92,8 +92,20 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):
92
assert req_id not in self._reqs_to_send
93
94
# Add to requests that are waiting to be read and track expiration.
95
+ # Deadlines are stamped with the scheduler process's perf_counter,
96
+ # which is not comparable to ours when the worker runs in another
97
+ # process on another node (perf_counter epochs differ by boot time).
98
+ # Rebase the remaining TTL onto our clock; broadcast latency only
99
+ # lengthens the lease, which is the safe direction. A cross-node
100
+ # epoch gap larger than the TTL otherwise expires the lease on
101
+ # arrival and the blocks are freed before D reads them.
102
+ now_local = time.perf_counter()
103
for req_id, expiration_time in metadata.reqs_to_send.items():
104
if req_id in self._reqs_to_process:
105
+ if metadata.scheduler_clock:
106
+ expiration_time = now_local + (
107
+ expiration_time - metadata.scheduler_clock
108
+ )
109
self._reqs_to_send[req_id] = expiration_time
110
111
# Send heartbeats to P-side engines to keep KV blocks alive while
vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py
+7
@@ -191,8 +191,15 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker):
191
for req_id in metadata.reqs_not_processed:
192
self._reqs_to_process.discard(req_id)
193
assert req_id not in self._reqs_to_send
194
+ # Rebase scheduler-clock deadlines onto this worker's clock — see the
195
+ # equivalent block in pull_worker.start_load_kv for the rationale.
196
+ now_local = time.perf_counter()
197
for req_id, expiration_time in metadata.reqs_to_send.items():
198
if req_id in self._reqs_to_process:
199
+ if metadata.scheduler_clock:
200
+ expiration_time = now_local + (
201
+ expiration_time - metadata.scheduler_clock
202
+ )
203
self._reqs_to_send[req_id] = expiration_time
204
205
# Heartbeats still leave from the main thread (base worker behaviour).