From 4e582c5b54b866a2cc0d44c8f8bbedb52cf5c605 Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Wed, 29 Jul 2026 18:25:42 -0700 Subject: [PATCH] [PD][Bugfix] Rebase KV lease deadlines onto worker clock (#50326) Signed-off-by: Nick Hill Co-authored-by: esmeetu Co-authored-by: Yifan Qiao --- .../kv_connector/unit/test_nixl_connector.py | 47 +++++++++++++++++++ .../kv_connector/v1/nixl/base_scheduler.py | 4 ++ .../kv_connector/v1/nixl/metadata.py | 6 +++ .../kv_connector/v1/nixl/pull_worker.py | 12 +++++ .../kv_connector/v1/nixl/push_worker.py | 7 +++ 5 files changed, 76 insertions(+) diff --git a/tests/v1/kv_connector/unit/test_nixl_connector.py b/tests/v1/kv_connector/unit/test_nixl_connector.py index c474077ab2..ba9f73f240 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector.py @@ -1398,6 +1398,53 @@ def test_kv_connector_stats(default_vllm_config, dist_init): assert stats_after_reset is None +def test_reqs_to_send_deadline_rebased_to_worker_clock(default_vllm_config, dist_init): + """reqs_to_send deadlines are stamped with the scheduler process's + perf_counter, whose epoch differs across processes and (by boot-time + deltas) across nodes. Without rebasing, a P worker on a node whose + monotonic clock is ahead of the scheduler's by more than the TTL + expires the lease on arrival and reports done_sending before D has + read the blocks — the freed blocks can then be reallocated and the + remote read pulls another request's data (silent accuracy corruption). + The worker must anchor the remaining TTL to its own clock. + """ + vllm_config = create_vllm_config() + connector = NixlConnector( + vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16) + ) + connector.connector_worker = FakeNixlConnectorWorker( + vllm_config, connector.engine_id, hand_shake_latency=0 + ) + worker = connector.connector_worker + + req_id = "req-lease-clock" + ttl = 480.0 + # Simulate a scheduler whose monotonic clock is 10,000 s behind this + # worker's (e.g. its node booted much later): the raw deadline is + # then already far in the past in this worker's clock domain. + scheduler_clock = time.perf_counter() - 10_000.0 + + metadata = NixlConnectorMetadata() + metadata.reqs_in_batch = {req_id} + metadata.reqs_to_send = {req_id: scheduler_clock + ttl} + metadata.scheduler_clock = scheduler_clock + connector.bind_connector_metadata(metadata) + dummy_ctx = ForwardContext( + no_compile_layers={}, + attn_metadata={}, + slot_mapping={}, + ) + connector.start_load_kv(dummy_ctx) + + remaining = worker._reqs_to_send[req_id] - time.perf_counter() + assert ttl - 5.0 < remaining <= ttl + 5.0 + + # The expiry sweep must not release the request. + done_sending, _ = worker.get_finished() + assert req_id not in done_sending + assert req_id in worker._reqs_to_process + + def test_kv_connector_stats_aggregation(): """ Test KV transfer stats aggregation across TP ranks using diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py index e53fb5673e..afa16279ff 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py @@ -453,6 +453,10 @@ class NixlBaseConnectorScheduler: self._build_save_meta(meta, scheduler_output) meta.reqs_to_send = self._reqs_need_send + # Clock reference for reqs_to_send: deadlines above are in this + # process's perf_counter domain; workers (possibly on other nodes, + # where perf_counter has a different epoch) rebase against this. + meta.scheduler_clock = time.perf_counter() meta.reqs_in_batch = self._reqs_in_batch meta.reqs_not_processed = self._reqs_not_processed diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py index f6f3de0a1f..27eab825a1 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py @@ -180,6 +180,12 @@ class NixlConnectorMetadata(KVConnectorMetadata): self.reqs_to_recv: dict[ReqId, ReqMeta] = {} self.reqs_to_save: dict[ReqId, ReqMeta] = {} self.reqs_to_send: dict[ReqId, float] = {} + # The scheduler process's time.perf_counter() when this metadata was + # built. reqs_to_send deadlines are stamped with the scheduler's + # clock, which is NOT comparable across processes (perf_counter is + # process/boot-local): workers must rebase the remaining TTL onto + # their own clock via this reference. 0.0 = unset (legacy metadata). + self.scheduler_clock: float = 0.0 self.reqs_in_batch: set[ReqId] = set() self.reqs_not_processed: set[ReqId] = set() # Heartbeat data grouped by remote engine, sent by D worker to P. diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py index 0c66596c4b..cc32845a4a 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py @@ -92,8 +92,20 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker): assert req_id not in self._reqs_to_send # Add to requests that are waiting to be read and track expiration. + # Deadlines are stamped with the scheduler process's perf_counter, + # which is not comparable to ours when the worker runs in another + # process on another node (perf_counter epochs differ by boot time). + # Rebase the remaining TTL onto our clock; broadcast latency only + # lengthens the lease, which is the safe direction. A cross-node + # epoch gap larger than the TTL otherwise expires the lease on + # arrival and the blocks are freed before D reads them. + now_local = time.perf_counter() for req_id, expiration_time in metadata.reqs_to_send.items(): if req_id in self._reqs_to_process: + if metadata.scheduler_clock: + expiration_time = now_local + ( + expiration_time - metadata.scheduler_clock + ) self._reqs_to_send[req_id] = expiration_time # Send heartbeats to P-side engines to keep KV blocks alive while diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py index 88977e3403..13af0314f0 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py @@ -191,8 +191,15 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker): for req_id in metadata.reqs_not_processed: self._reqs_to_process.discard(req_id) assert req_id not in self._reqs_to_send + # Rebase scheduler-clock deadlines onto this worker's clock — see the + # equivalent block in pull_worker.start_load_kv for the rationale. + now_local = time.perf_counter() for req_id, expiration_time in metadata.reqs_to_send.items(): if req_id in self._reqs_to_process: + if metadata.scheduler_clock: + expiration_time = now_local + ( + expiration_time - metadata.scheduler_clock + ) self._reqs_to_send[req_id] = expiration_time # Heartbeats still leave from the main thread (base worker behaviour).