vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_common ¶

class HandshakeError(MoRIIOError):
    """Exception raised when handshake fails."""

    pass

@dataclass
class LayerTransferPlan:
    """Plan for transferring a single layer."""

    request_id: ReqId
    transfer_id: TransferId
    layer_name: str
    sess_idx: int
    transfer_local_offsets: list[int]
    transfer_remote_offsets: list[int]
    transfer_sizes: list[int]
    use_batch: bool = True

class MoRIIOConnectorMetadata(KVConnectorMetadata):
    def __init__(self):
        self.reqs_to_recv: dict[ReqId, ReqMeta] = {}
        self.reqs_to_save: dict[ReqId, ReqMeta] = {}
        self.reqs_to_send: dict[ReqId, float] = {}
        self.transfer_id_to_request_id: dict[TransferId, ReqId] = {}

    def __repr__(self):
        return (
            f"MoRIIOConnectorMetadata: reqs_to_recv={self.reqs_to_recv}, "
            f"reqs_to_save={self.reqs_to_save}, "
            f"reqs_to_send={self.reqs_to_send}, "
            f"transfer_id_to_request_id={self.transfer_id_to_request_id}"
        )

    def add_new_req(
        self,
        request_id: ReqId,
        local_block_ids: list[int],
        kv_transfer_params: dict[str, Any],
        write_mode=False,
    ):
        """Ingest a peer's ``kv_transfer_params`` into a typed ``ReqMeta``.

        This is the single place peer info enters the connector. The
        ``kv_transfer_params`` contract (as produced by the llm-d sidecar /
        vLLM router, or echoed by the prefill leg's ``request_finished``):

        Required (always):
          * ``transfer_id``        -- stable id shared by both legs.
          * ``remote_engine_id``   -- peer engine id for the handshake table.
          * ``remote_block_ids``   -- peer block ids (may be [] in WRITE mode,
                                      where decode allocates its own blocks).

        Peer address -- ONE of the following two must resolve:
          * embedded in ``request_id`` (vLLM-router PD id form), OR
          * explicit ``remote_host`` + ``remote_handshake_port`` +
            ``remote_notify_port`` (llm-d sidecar / returnable path).
          If neither resolves we raise -- there is no safe default host/port.

        Optional (defaulted):
          * ``tp_size``              (default 1) -- peer TP size.
          * ``remote_dp_size``       (default 1) -- peer GLOBAL DP size.
          * ``remote_dp_size_local`` (default = ``remote_dp_size``) -- per-pod
                                      DP size for Wide-EP multi-pod port/host
                                      folding; 0/absent means single-pod.
          * ``remote_hosts``         (default [remote_host]) -- per-pod IP list
                                      indexed by ``pod_idx = rank // dp_local``.

        Routing keys consumed elsewhere (NOT here): ``remote_dp_rank`` /
        ``remote_dp_rank_override`` gate the decode->prefill notify target in
        MoRIIOConnectorScheduler; they are router-authoritative and never
        self-derived (see that class's request routing contract).
        """
        transfer_id = kv_transfer_params["transfer_id"]

        # Try request_id embedded address first, fallback to explicit params.
        peer_zmq = get_peer_zmq_from_request_id(request_id, is_producer=write_mode)
        if peer_zmq is not None:
            remote_host, remote_handshake_port, remote_notify_port = (
                parse_moriio_zmq_address(peer_zmq)
            )
        else:
            try:
                remote_host = kv_transfer_params["remote_host"]
                if not remote_host:
                    raise ValueError(
                        f"request_id {request_id!r} does not embed a peer "
                        f"zmq_address and kv_transfer_params['remote_host'] is "
                        f"empty; cannot route MoRI-IO transfer"
                    )
                remote_handshake_port = int(kv_transfer_params["remote_handshake_port"])
                remote_notify_port = int(kv_transfer_params["remote_notify_port"])
            except (KeyError, TypeError, ValueError) as e:
                raise ValueError(
                    f"request_id {request_id!r} does not embed a peer "
                    f"zmq_address and kv_transfer_params is missing one or "
                    f"more sidecar-fallback keys (need remote_host, "
                    f"remote_handshake_port, remote_notify_port): {e}"
                ) from e

        # Multi-pod: use multi_pod_hosts list or fallback to single host.
        _pod_hosts = kv_transfer_params.get("remote_hosts") or [remote_host]
        if not isinstance(_pod_hosts, list):
            _pod_hosts = [_pod_hosts]
        _pod_hosts = [str(h) for h in _pod_hosts]
        _remote_dp_size_local = int(
            kv_transfer_params.get(
                "remote_dp_size_local",
                kv_transfer_params.get("remote_dp_size", 1),
            )
        )

        _req = ReqMeta(
            transfer_id=transfer_id,
            local_block_ids=local_block_ids,
            remote_block_ids=kv_transfer_params["remote_block_ids"],
            remote_engine_id=kv_transfer_params["remote_engine_id"],
            remote_host=remote_host,
            remote_port=int(remote_handshake_port),
            remote_handshake_port=int(remote_handshake_port),
            remote_notify_port=int(remote_notify_port),
            # Remote peer TP degree (used as remote_tp_size downstream). The
            # proxy advertises it under "remote_tp_size"; #46332 read "tp_size"
            # which is absent on WRITE producer requests -> defaulted to 1 ->
            # rank collapse. Read the right key; 0 == unknown (== homogeneous).
            tp_size=int(
                kv_transfer_params.get("remote_tp_size")
                or kv_transfer_params.get("tp_size")
                or 0
            ),
            remote_dp_size=kv_transfer_params.get("remote_dp_size", 1),
            remote_dp_rank=kv_transfer_params.get("remote_dp_rank", 0),
            multi_pod_hosts=_pod_hosts,
            remote_dp_size_local=_remote_dp_size_local,
        )
        if write_mode:
            self.reqs_to_save[request_id] = _req
        else:
            self.reqs_to_recv[request_id] = _req

class MoRIIOConstants:
    """Constants for MoRIIO connector."""

    # ZMQ message types
    GET_META_MSG = b"get_meta_msg"
    POP_DONE_RECV = b"pop_done_recv"
    OVER = b"OVER"
    COMPLETION_PREFIX = "cmpl"
    TRANSFER_PREFIX = "tx"

    PING_INTERVAL = 3
    MAX_PING_RETRIES = 100
    DEFAULT_HANDSHAKE_PORT = "6301"
    DEFAULT_NOTIFY_PORT = "61005"

    VLLM_MORI_READ_ABORT_REQUEST_TIMEOUT = 3600

    # Timeout (seconds) for waiting_for_transfer_complete before raising TransferError.
    # Overridable via kv_connector_extra_config["transfer_timeout"].
    DEFAULT_TRANSFER_TIMEOUT = 30.0
    # Timeout (seconds) before a deferred send with no finished_sending
    # notification is reaped and its blocks force-freed.
    # Overridable via kv_connector_extra_config["defer_timeout"].
    DEFAULT_DEFER_TIMEOUT = 60.0

class MoRIIOError(Exception):
    """Base exception for MoRIIO operations."""

    pass

@dataclass
class RemoteAllocInfo:
    """Information about remote block allocation."""

    block_ids: list[int]
    writes_done: int = 0
    writes_expected: int | None = None
    decode_dp_rank: int = 0
    completion_request_id: str | None = None
    completion_remote_notify_port: int | None = None
    completion_remote_ip: str | None = None
    completion_notified: bool = False
    transfer_statuses: list[Any] = field(default_factory=list)
    transfer_offsets: dict[
        tuple[tuple[int, ...], tuple[int, ...], torch.dtype],
        tuple[list[int], list[int], list[int]],
    ] = field(default_factory=dict)

@dataclass
class ReqMeta:
    """Metadata for a single request."""

    transfer_id: TransferId
    local_block_ids: list[int]
    remote_block_ids: list[int]
    remote_host: str
    remote_port: int
    remote_handshake_port: int
    remote_notify_port: int
    remote_engine_id: str
    tp_size: int
    remote_dp_size: int
    # Prefill DP rank that owns this request's KV (forwarded by the proxy). The
    # read must target this rank's memory registration; the default 0 preserves
    # the symmetric single-DP behaviour.
    remote_dp_rank: int = 0
    # Multi-pod: list of remote pod IPs indexed by pod_idx.
    multi_pod_hosts: list[str] = field(default_factory=list)
    # Per-pod DP size; 0 means fallback to remote_dp_size.
    remote_dp_size_local: int = 0

class RoleManager:
    """Manages role state across the connector."""

    _instance: "RoleManager | None" = None
    _lock = threading.Lock()

    def __init__(self) -> None:
        self._role: ROLE = ROLE.NOTINIT

    @classmethod
    def get_instance(cls) -> "RoleManager":
        if cls._instance is None:
            with cls._lock:
                if cls._instance is None:
                    cls._instance = cls()
        return cls._instance

    def set_role(self, role: ROLE) -> None:
        """Set the current role."""
        with self._lock:
            self._role = role

    def get_role(self) -> ROLE:
        """Get the current role."""
        return self._role

class TransferError(MoRIIOError):
    """Exception raised when transfer fails."""

    pass