Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions .buildkite/test_areas/disaggregated.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
33 changes: 33 additions & 0 deletions tests/v1/kv_connector/unit/test_kv_load_failure_recovery.py
Original file line number Diff line number Diff line change
Expand Up @@ -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",
[
Expand Down
47 changes: 47 additions & 0 deletions tests/v1/kv_connector/unit/test_nixl_connector.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
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
Expand Down Expand Up @@ -2822,6 +2823,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,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -64,8 +64,14 @@ def start_load_kv(self, metadata: NixlConnectorMetadata):
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:
Expand Down
Loading