Skip to content

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

Classes:

Functions:

HandshakeError

Bases: MoRIIOError

Exception raised when handshake fails.

Source code in vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py 
153
154
155
156
class HandshakeError(MoRIIOError):
    """Exception raised when handshake fails."""

    pass

LayerTransferPlan dataclass

Plan for transferring a single layer.

Source code in vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py 
61
62
63
64
65
66
67
68
69
70
71
72
@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

MoRIIOConstants

Constants for MoRIIO connector.

Source code in vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py 
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
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

MoRIIOError

Bases: Exception

Base exception for MoRIIO operations.

Source code in vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py 
147
148
149
150
class MoRIIOError(Exception):
    """Base exception for MoRIIO operations."""

    pass

RemoteAllocInfo dataclass

Information about remote block allocation.

Source code in vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py 
75
76
77
78
79
80
81
82
@dataclass
class RemoteAllocInfo:
    """Information about remote block allocation."""

    block_ids: list[int]
    writes_done: int = 0
    decode_dp_rank: int = 0
    transfer_offset: tuple[list[int], list[int], list[int]] | None = None

ReqMeta dataclass

Metadata for a single request.

Source code in vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py 
385
386
387
388
389
390
391
392
393
394
395
396
397
398
@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

RoleManager

Manages role state across the connector.

Methods:

Source code in vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py 
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
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

get_role()

Get the current role.

Source code in vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py 
127
128
129
def get_role(self) -> ROLE:
    """Get the current role."""
    return self._role

set_role(role)

Set the current role.

Source code in vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py 
122
123
124
125
def set_role(self, role: ROLE) -> None:
    """Set the current role."""
    with self._lock:
        self._role = role

TransferError

Bases: MoRIIOError

Exception raised when transfer fails.

Source code in vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py 
159
160
161
162
class TransferError(MoRIIOError):
    """Exception raised when transfer fails."""

    pass

get_peer_zmq_from_request_id(request_id, is_producer)

Extract the peer's zmq_address from the vLLM router request_id.

The producer (prefill) needs the decode's address; the consumer (decode) needs the prefill's address.

Source code in vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py 
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
def get_peer_zmq_from_request_id(request_id: str, is_producer: bool) -> str:
    """Extract the *peer's* zmq_address from the vLLM router request_id.

    The producer (prefill) needs the decode's address; the consumer (decode)
    needs the prefill's address.
    """
    if is_producer:
        m = _DECODE_ZMQ_RE.search(request_id)
    else:
        m = _PREFILL_ZMQ_RE.search(request_id)
    if m is None:
        raise ValueError(
            f"Cannot parse peer zmq_address from request_id: {request_id!r}"
        )
    return m.group(1)

get_role()

Get the global role.

Source code in vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py 
137
138
139
def get_role() -> ROLE:
    """Get the global role."""
    return RoleManager.get_instance().get_role()

parse_moriio_zmq_address(zmq_address)

Parse the MoRI-IO zmq address into its components.

Parses "host:IP,handshake:PORT,notify:PORT" into (host, handshake_port, notify_port).

Each key-value pair is split on the first colon so that IPv6 addresses (e.g. host:::1) are handled correctly. Raises ValueError if any of host, handshake, or notify keys are absent or if the port values are non-numeric.

Source code in vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py 
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
def parse_moriio_zmq_address(
    zmq_address: str,
) -> tuple[str, int, int]:
    """Parse the MoRI-IO zmq address into its components.

    Parses ``"host:IP,handshake:PORT,notify:PORT"`` into
        (host, handshake_port, notify_port).

    Each key-value pair is split on the *first* colon so that IPv6 addresses
    (e.g. ``host:::1``) are handled correctly.  Raises ``ValueError`` if any
    of ``host``, ``handshake``, or ``notify`` keys are absent or if the port
    values are non-numeric.
    """
    parts: dict[str, str] = {}
    for segment in zmq_address.split(","):
        key, _, val = segment.partition(":")
        parts[key.strip()] = val.strip()
    try:
        host = parts["host"]
        handshake_port = int(parts["handshake"])
        notify_port = int(parts["notify"])
    except (KeyError, ValueError) as e:
        raise ValueError(
            f"Malformed zmq_address {zmq_address!r}: expected "
            f"'host:IP,handshake:PORT,notify:PORT' format"
        ) from e
    return host, handshake_port, notify_port

set_role(role)

Set the global role.

Source code in vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py 
132
133
134
def set_role(role: ROLE):
    """Set the global role."""
    RoleManager.get_instance().set_role(role)

zmq_ctx(socket_type, addr)

Context manager for a ZMQ socket

Source code in vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py 
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
@contextlib.contextmanager
def zmq_ctx(socket_type: Any, addr: str) -> Iterator[zmq.Socket]:
    """Context manager for a ZMQ socket"""

    if socket_type not in (zmq.ROUTER, zmq.REQ, zmq.DEALER):
        raise ValueError(f"Unexpected socket type: {socket_type}")

    ctx: zmq.Context | None = None
    try:
        ctx = zmq.Context()  # type: ignore[attr-defined]
        yield make_zmq_socket(
            ctx=ctx, path=addr, socket_type=socket_type, bind=socket_type == zmq.ROUTER
        )
    finally:
        if ctx is not None:
            ctx.destroy(linger=0)