agentic-pd-hybrid/src/agentic_pd_hybrid/snapshot_link.py

"""Minimal D→P snapshot link over Mooncake RDMA.

This module provides a thin wrapper around mooncake.engine.TransferEngine
for one-sided RDMA writes of KV bytes from a Decode worker (sender) to a
Prefill worker (receiver). It deliberately does NOT use the heavyweight
MooncakeKVManager pipeline (which is tied to PREFILL/DECODE roles and
chunked transfer protocols): we want a simple, testable byte transport
that can be reused by SGLang and by stand-alone smoke tests.

Layout:
    SnapshotPeer        — engine + pre-registered receive buffer (receiver)
                          or sender handle (sender)
    SnapshotEndpoint    — what the receiver advertises so the sender can
                          target it: (session_id, base_ptr, length)
    SnapshotPusher      — sender-side: holds a target endpoint, calls
                          batch_transfer_sync_write

All transfers are SYNCHRONOUS, single-shot, in-memory.

Higher layers add: control plane (how D learns P's endpoint), per-session
slot allocation, KV format/layout, hand-off into SGLang scheduler.
"""

from __future__ import annotations

import ctypes
import logging
import os
import threading
from dataclasses import dataclass
from typing import Optional

logger = logging.getLogger(__name__)


@dataclass(frozen=True)
class SnapshotEndpoint:
    """What the receiver advertises so the sender can reach it.

    Attributes
    ----------
    session_id : str
        ``"host:rpc_port"`` string identifying the receiver's mooncake
        TransferEngine. Returned by ``TransferEngine.get_rpc_port()``
        joined with the host the engine was initialized with.
    base_ptr : int
        Address of the registered receive buffer on the receiver side.
    capacity_bytes : int
        Length of the registered region.
    """

    session_id: str
    base_ptr: int
    capacity_bytes: int


def _import_transfer_engine():
    try:
        from mooncake.engine import TransferEngine
    except ImportError as e:  # pragma: no cover
        raise ImportError(
            "mooncake.engine.TransferEngine is required for snapshot_link. "
            "Make sure mooncake-transfer-engine is installed in the venv."
        ) from e
    return TransferEngine


class SnapshotPeer:
    """One Mooncake transfer engine endpoint with a registered receive buffer.

    The engine is dedicated to snapshot traffic — it does NOT share state
    with SGLang's MooncakeKVManager engine. Each SnapshotPeer needs its own
    host:port to listen on.
    """

    def __init__(
        self,
        host: str,
        port: int,
        ib_device: Optional[str] = None,
        receive_capacity_bytes: int = 0,
        protocol: Optional[str] = None,
    ):
        TransferEngine = _import_transfer_engine()
        self.host = host
        self.port = port
        self.ib_device = ib_device
        self.engine = TransferEngine()

        listen = f"{host}:{port}"
        proto = protocol or os.environ.get("MOONCAKE_PROTOCOL", "rdma")
        ret = self.engine.initialize(
            listen,
            "P2PHANDSHAKE",
            proto,
            ib_device or "",
        )
        if ret != 0:
            raise RuntimeError(
                f"snapshot_link: engine.initialize({listen!r}, proto={proto}, "
                f"ib={ib_device}) returned {ret}"
            )

        self._rpc_port = self.engine.get_rpc_port()
        self._session_id = f"{host}:{self._rpc_port}"

        self._recv_buffer = None
        self._recv_ptr = 0
        self._recv_capacity = 0
        if receive_capacity_bytes > 0:
            self._allocate_recv_buffer(receive_capacity_bytes)

        self._lock = threading.Lock()
        logger.info(
            "SnapshotPeer up at %s (rpc=%d, ib=%s, recv=%d B)",
            self._session_id,
            self._rpc_port,
            ib_device,
            receive_capacity_bytes,
        )

    # -- accessors ---------------------------------------------------------

    @property
    def session_id(self) -> str:
        return self._session_id

    @property
    def rpc_port(self) -> int:
        return self._rpc_port

    @property
    def endpoint(self) -> SnapshotEndpoint:
        if self._recv_buffer is None:
            raise RuntimeError(
                "SnapshotPeer has no receive buffer; pass receive_capacity_bytes > 0"
            )
        return SnapshotEndpoint(
            session_id=self._session_id,
            base_ptr=self._recv_ptr,
            capacity_bytes=self._recv_capacity,
        )

    # -- buffer management -------------------------------------------------

    def _allocate_recv_buffer(self, length: int) -> None:
        """Allocate + register a pinned host buffer for receiving."""
        # Use c_ubyte (unsigned) so bytes() conversions of the underlying
        # storage always yield valid byte values.
        buf = (ctypes.c_ubyte * length)()
        addr = ctypes.addressof(buf)
        ret = self.engine.register_memory(addr, length)
        if ret != 0:
            raise RuntimeError(
                f"snapshot_link: register_memory({hex(addr)}, {length}) returned {ret}"
            )
        self._recv_buffer = buf
        self._recv_ptr = addr
        self._recv_capacity = length

    def read_bytes(self, offset: int, length: int) -> bytes:
        """Snapshot the recv buffer at [offset, offset+length) (caller syncs)."""
        if self._recv_buffer is None:
            raise RuntimeError("no recv buffer")
        if offset < 0 or offset + length > self._recv_capacity:
            raise ValueError(
                f"read_bytes({offset}, {length}) out of capacity {self._recv_capacity}"
            )
        # string_at copies via memcpy and yields a proper bytes object — works
        # regardless of signed/unsigned underlying storage.
        return ctypes.string_at(self._recv_ptr + offset, length)

    def register_send_buffer(self, ptr: int, length: int) -> None:
        """Register an externally-allocated send buffer for outbound RDMA writes."""
        with self._lock:
            ret = self.engine.register_memory(ptr, length)
            if ret != 0:
                raise RuntimeError(
                    f"snapshot_link: register send buffer({hex(ptr)}, {length}) returned {ret}"
                )

    def deregister(self, ptr: int) -> None:
        with self._lock:
            try:
                self.engine.unregister_memory(ptr)
            except Exception:
                pass

    # -- transfer ----------------------------------------------------------

    def push(
        self,
        target: SnapshotEndpoint,
        local_ptr: int,
        local_offset: int,
        length: int,
        remote_offset: int = 0,
    ) -> int:
        """Synchronously RDMA-write ``length`` bytes from ``local_ptr+local_offset``
        to ``target.base_ptr+remote_offset`` on the peer identified by
        ``target.session_id``.

        Returns 0 on success, non-zero (or raises) on failure.
        """
        if length <= 0:
            return 0
        if remote_offset < 0 or remote_offset + length > target.capacity_bytes:
            raise ValueError(
                f"push: remote_offset={remote_offset}, length={length} exceeds "
                f"target capacity {target.capacity_bytes}"
            )
        src = local_ptr + local_offset
        dst = target.base_ptr + remote_offset
        try:
            ret = self.engine.transfer_sync_write(
                target.session_id, src, dst, length
            )
        except Exception as e:
            logger.exception("snapshot_link.push transfer_sync_write threw: %s", e)
            return -1
        if ret != 0:
            logger.warning(
                "snapshot_link.push transfer_sync_write returned %d (src=%s, "
                "dst=%s/%s, len=%d)",
                ret,
                hex(src),
                target.session_id,
                hex(dst),
                length,
            )
        return ret

    def batch_push(
        self,
        target: SnapshotEndpoint,
        local_addrs: list[int],
        remote_addrs: list[int],
        lengths: list[int],
    ) -> int:
        """Batched RDMA write (one-shot)."""
        if not local_addrs:
            return 0
        try:
            ret = self.engine.batch_transfer_sync_write(
                target.session_id, local_addrs, remote_addrs, lengths
            )
        except Exception as e:
            logger.exception("snapshot_link.batch_push threw: %s", e)
            return -1
        return ret

    def close(self) -> None:
        """Best-effort shutdown — release the receive buffer registration."""
        if self._recv_ptr:
            try:
                self.engine.unregister_memory(self._recv_ptr)
            except Exception:
                pass
        self._recv_ptr = 0
        self._recv_capacity = 0
        self._recv_buffer = None


def make_session_id(host: str, rpc_port: int) -> str:
    """Build the ``host:port`` form used as mooncake's session id."""
    return f"{host}:{rpc_port}"