[Refactor][KV Connector][Mooncake] Move connector metadata and wire structs to their own modules

tests/v1/kv_connector/unit/test_mooncake_connector.py

@@ -1117,7 +1117,7 @@ def test_scheduler_request_finished():
     delay_free, _ = scheduler_connector.request_finished(request, block_ids=([12],))
     assert delay_free is False
     assert len(scheduler_connector._reqs_need_send) == 0
-    assert "id-1" in scheduler_connector._reqs_not_processed
+    assert "id-1" in scheduler_connector._transfers_not_processed
 
 
 @contextlib.contextmanager

tests/v1/kv_connector/unit/test_mooncake_wire_format.py

@@ -0,0 +1,100 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import msgspec
+import pytest
+
+from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.protocol import (
+    MooncakeXferMetadata,
+    MooncakeXferResponse,
+    MooncakeXferResponseStatus,
+)
+
+# The wire contract of the push protocol, stated once. Field order, wire name,
+# type and default are what the peer decodes against, so a deliberate change to
+# any of them is meant to show up here as an edit of this table, and an
+# accidental one as a failure. The runtime counterpart is the compatibility
+# hash in metadata.py, which refuses a pair whose wire differs.
+EXPECTED_METADATA_SCHEMA = [
+    ("remote_hostname", "str", "required"),
+    ("remote_port", "int", "required"),
+    ("remote_tp_size", "int", "required"),
+    ("remote_tp_rank", "int", "required"),
+    ("req_blocks", "dict[str, tuple[str, list[list[int]]]]", "required"),
+    ("kv_caches_base_addr", "list[int]", "required"),
+    ("block_lens", "list[int]", "required"),
+    ("kv_block_lens", "list[int]", "required"),
+    ("registered_layer_names", "list[str]", "factory=list"),
+    ("registered_layer_indices", "list[int]", "factory=list"),
+    ("registered_group_indices", "list[int]", "factory=list"),
+    ("registered_shared_group_ids", "list[list[int]]", "factory=list"),
+    ("registered_row_offsets", "list[int]", "factory=list"),
+    ("remote_pp_size", "int", "default=1"),
+    ("layout_id", "int", "default=0"),
+]
+
+EXPECTED_RESPONSE_SCHEMA = [
+    ("status", "MooncakeXferResponseStatus", "required"),
+    ("ok_reqs", "list[str] | None", "default=None"),
+    ("err_reqs", "list[str] | None", "default=None"),
+    ("err_msg", "str | None", "default=None"),
+    ("layout_id", "int", "default=0"),
+]
+
+
+def _default_of(field: msgspec.structs.FieldInfo) -> str:
+    if field.required:
+        return "required"
+    if field.default_factory is not msgspec.NODEFAULT:
+        return f"factory={field.default_factory.__name__}"
+    return f"default={field.default!r}"
+
+
+def _type_name(tp: object) -> str:
+    # Plain classes render as "<class 'str'>"; subscripts and unions read best
+    # as their own repr.
+    return tp.__name__ if isinstance(tp, type) else str(tp)
+
+
+def _schema(struct: type[msgspec.Struct]) -> list[tuple[str, str, str]]:
+    return [
+        (field.name, _type_name(field.type), _default_of(field))
+        for field in msgspec.structs.fields(struct)
+    ]
+
+
+def test_xfer_metadata_schema_is_unchanged():
+    assert _schema(MooncakeXferMetadata) == EXPECTED_METADATA_SCHEMA
+
+
+def test_xfer_response_schema_is_unchanged():
+    assert _schema(MooncakeXferResponse) == EXPECTED_RESPONSE_SCHEMA
+
+
+@pytest.mark.parametrize(
+    "struct", [MooncakeXferMetadata, MooncakeXferResponse], ids=lambda c: c.__name__
+)
+def test_wire_configuration_is_unchanged(struct: type[msgspec.Struct]):
+    config = struct.__struct_config__
+    # Both messages omit defaults on the wire, and neither is tagged or
+    # array-encoded, so decoding depends on the field names in the tables above.
+    assert config.omit_defaults is True
+    assert config.tag is None
+    assert config.tag_field is None
+    assert config.array_like is False
+    # A peer on a newer vLLM sends fields this side does not know yet; rejecting
+    # them would break a mixed-version P/D pair during a rolling upgrade.
+    assert config.forbid_unknown_fields is False
+    # The tables name fields, so a custom wire name would have to be spelled out
+    # there as well.
+    assert all(
+        field.name == field.encode_name for field in msgspec.structs.fields(struct)
+    )
+
+
+def test_xfer_response_status_values_are_unchanged():
+    assert [(status.name, int(status)) for status in MooncakeXferResponseStatus] == [
+        ("FINISH", 0),
+        ("CONTINUE", 1),
+        ("ERROR", 2),
+        ("LAYOUT_MISS", 3),
+    ]

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/metadata.py

@@ -0,0 +1,79 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import asyncio
+from collections import defaultdict
+from dataclasses import dataclass
+from typing import Any
+
+from vllm.distributed.kv_transfer.kv_connector.utils import EngineId
+from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorMetadata
+
+ReqId = str  # Internal scheduler request ID
+TransferId = str  # KV transfer coordination ID (shared by P/D)
+
+
+@dataclass
+class PullReqMeta:
+    d_req_id: ReqId
+    transfer_id: TransferId
+    local_block_ids: list[list[int]]
+    remote_engine_id: EngineId
+    remote_bootstrap_addr: str
+    # Set expire time to avoid infinitely sending requests.
+    expire_time: float = float("inf")
+    # Designed for one D pairing to multiple P
+    pull_tasks_count: int = 0
+    # A failure is sticky; still settle later peer results before releasing
+    # destination blocks to prevent write-after-free from delayed peer writes.
+    failed: bool = False
+
+
+@dataclass
+class SendBlockMeta:
+    p_req_id: ReqId
+    transfer_id: TransferId
+    local_block_ids: list[list[int]]
+    ready: asyncio.Event
+    expire_time: float = float("inf")
+    need_send: int = 0
+    sent: int = 0
+    sending: int = 0
+
+
+@dataclass
+class LeaseHeartbeat:
+    """Transfers that one P engine holds for this D engine under a lease."""
+
+    remote_bootstrap_addr: str
+    transfer_ids: list[TransferId]
+
+
+class MooncakeConnectorMetadata(KVConnectorMetadata):
+    def __init__(self):
+        # Use (engine_id, dp_rank) to group reqs with same dp.
+        # See comments in MooncakeBootstrapServer.
+        self.reqs_to_recv: dict[EngineId, dict[ReqId, PullReqMeta]] = defaultdict(dict)
+        self.reqs_to_send: dict[ReqId, tuple[TransferId, list[list[int]]]] = {}
+        self.transfers_not_processed: set[TransferId] = set()
+        # D side: lease renewals to send this step, by remote P engine.
+        self.lease_heartbeats: dict[EngineId, LeaseHeartbeat] = {}
+
+    def add_new_req(
+        self,
+        request_id: ReqId,
+        local_block_ids: list[list[int]],
+        kv_transfer_params: dict[str, Any],
+        load_remote_cache: bool = True,
+    ):
+        transfer_id = kv_transfer_params["transfer_id"]
+        if load_remote_cache:
+            remote_engine_id = kv_transfer_params["remote_engine_id"]
+            self.reqs_to_recv[remote_engine_id][request_id] = PullReqMeta(
+                d_req_id=request_id,
+                local_block_ids=local_block_ids,
+                remote_engine_id=remote_engine_id,
+                remote_bootstrap_addr=kv_transfer_params["remote_bootstrap_addr"],
+                transfer_id=transfer_id,
+            )
+        else:
+            self.reqs_to_send[request_id] = (transfer_id, local_block_ids)