From 49e0427ca5183a56ccb112e5be1300feb2873ca9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Lucchesi?= Date: Fri, 18 Sep 2026 22:33:02 +0200 Subject: [PATCH] [Bugfix][NIXL] Avoid receive reports for notification-only requests (#57570) Signed-off-by: NickLucche Co-authored-by: OpenAI Codex --- .buildkite/test_areas/disaggregated.yaml | 3 ++ .../unit/test_kv_load_failure_recovery.py | 33 +++++++++++++ .../kv_connector/unit/test_nixl_connector.py | 47 +++++++++++++++++++ .../kv_connector/v1/nixl/pull_worker.py | 10 +++- 4 files changed, 91 insertions(+), 2 deletions(-) diff --git a/.buildkite/test_areas/disaggregated.yaml b/.buildkite/test_areas/disaggregated.yaml index d217090bc5..b187169e76 100644 --- a/.buildkite/test_areas/disaggregated.yaml +++ b/.buildkite/test_areas/disaggregated.yaml @@ -311,6 +311,9 @@ steps: PREFILLER_TP_SIZE: "2" DECODER_TP_SIZE: "2" MODEL_NAMES: "moonshotai/Kimi-Linear-48B-A3B-Instruct" + # Keep TP=2 prefill and TP=1 decode above both hybrid page-size minima. + PREFILL_BLOCK_SIZE: "2048" + DECODE_BLOCK_SIZE: "2048" VLLM_SERVE_EXTRA_ARGS: "--trust-remote-code,--mamba-cache-mode align,--max-model-len 32768" commands: - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh diff --git a/tests/v1/kv_connector/unit/test_kv_load_failure_recovery.py b/tests/v1/kv_connector/unit/test_kv_load_failure_recovery.py index a54a023a13..a8bc2d0e12 100644 --- a/tests/v1/kv_connector/unit/test_kv_load_failure_recovery.py +++ b/tests/v1/kv_connector/unit/test_kv_load_failure_recovery.py @@ -35,6 +35,39 @@ def scheduler(): return create_scheduler(vllm_config) +@pytest.mark.parametrize("policy", ["fail", "recompute"]) +def test_failed_receive_completion_honors_failure_policy(policy): + """Failed receives still need completion to free blocks or resume computation.""" + scheduler = create_scheduler(create_vllm_config(kv_load_failure_policy=policy)) + request = create_request(num_tokens=3 * scheduler.block_size) + scheduler.add_request(request) + scheduler.connector = Mock() + scheduler.connector.get_num_new_matched_tokens.side_effect = [ + (2 * scheduler.block_size, True), + (0, False), + ] + scheduler.connector.request_finished.return_value = (False, None) + scheduler.connector.take_events.return_value = () + block_pool = scheduler.kv_cache_manager.block_pool + free_blocks_before = block_pool.get_num_free_blocks() + scheduler_output = scheduler.schedule() + assert request.status == RequestStatus.WAITING_FOR_REMOTE_KVS + + output = create_model_runner_output([], finished_recving={request.request_id}) + assert output.kv_connector_output is not None + output.kv_connector_output.failed_recving = {request.request_id} + scheduler.update_from_output(scheduler_output, output) + + if policy == "fail": + assert request.status == RequestStatus.FINISHED_ERROR + assert request.request_id not in scheduler.requests + assert block_pool.get_num_free_blocks() == free_blocks_before + else: + assert request.num_computed_tokens == 0 + next_output = scheduler.schedule() + assert request.request_id in next_output.num_scheduled_tokens + + @pytest.mark.parametrize( "num_prompt_blocks,num_external_computed_blocks,invalid_block_idxs", [ diff --git a/tests/v1/kv_connector/unit/test_nixl_connector.py b/tests/v1/kv_connector/unit/test_nixl_connector.py index 3b9f607865..39f8f88d5c 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector.py @@ -10,6 +10,7 @@ import textwrap import time import uuid from collections import defaultdict +from concurrent.futures import Future from types import SimpleNamespace from typing import Any, cast from unittest.mock import MagicMock, patch @@ -2899,6 +2900,52 @@ def test_empty_recv_is_reported_only_when_awaited( assert (request_id in done_recving) is awaiting_kvs +@pytest.mark.parametrize("is_hma", [False, True]) +@pytest.mark.parametrize( + "local_block_ids,awaiting_kvs", + [((), False), (([],), False), (([],), True), (([1, 2, 3],), True)], +) +def test_handshake_failure_reports_only_awaited_recvs( + recv_worker, is_hma, local_block_ids, awaiting_kvs +): + """A failed cleanup handshake must not complete an already-aborted request.""" + worker = recv_worker + worker._is_hma_required = is_hma + worker._recving_metadata.clear() + worker._remote_agents = {} + worker._handshake_lock = contextlib.nullcontext() + worker._ready_requests = queue.Queue() + worker._reqs_to_process = set() + worker.pcp_rank = 0 + metadata = NixlConnectorMetadata() + metadata.add_new_req_to_recv( + request_id="request", + local_block_ids=local_block_ids, + kv_transfer_params={ + "remote_block_ids": ([4, 5, 6],), + "remote_engine_id": "prefill", + "remote_request_id": "prefill-request", + "remote_host": "localhost", + "remote_port": 1234, + }, + awaiting_kvs=awaiting_kvs, + ) + handshake = Future[None]() + handshake.set_exception(RuntimeError("handshake failed")) + with patch.object(worker, "_ensure_handshake", return_value=handshake): + worker.start_load_kv(metadata) + + results = worker.get_transfer_results() + expected = {"request"} if awaiting_kvs else set() + assert results.finished_recving == results.failed_recving == expected + assert worker.get_block_ids_with_load_errors() == ( + set(local_block_ids[0]) if awaiting_kvs and not is_hma else set() + ) + assert not worker._recving_metadata + assert not worker._recv_failures + assert not worker._pending_recv_notifs + + @patch( "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper", FailingNixlWrapper, 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 7711a35954..43fbbe0c3c 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 @@ -64,8 +64,14 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker): len(meta.local_physical_block_ids), len(meta.remote.block_ids), ) - # always store metadata for failure recovery - self._recving_metadata[req_id] = meta + # Full local hits and aborted cleanup only notify P; no recv is awaited. + # On a full local hit: + # - Notification failure must not fail the request: its KV is local. + # - Receive completion must not be reported: num_external_tokens == 0, + # so the scheduler never entered WAITING_FOR_REMOTE_KVS to begin with. + # Aborted cleanup requests have already been removed from the scheduler. + if meta.awaiting_kvs or any(meta.local_block_ids): + self._recving_metadata[req_id] = meta if remote_engine_id not in self._remote_agents: # Initiate handshake with remote engine to exchange metadata. with self._handshake_lock: