agentic-kvc/microbench/fresh_setup/instrument_mooncake.py

#!/usr/bin/env python3
"""Instrument vLLM's shipped MooncakeConnector for per-stage MB2 timing.

vLLM 0.18.1 ships its OWN MooncakeConnector at
  vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py
which is the one --kv-transfer-config '{"kv_connector":"MooncakeConnector"}'
loads. The mooncake-package's mooncake_connector_v1.py is a separate
implementation and not used by vLLM 0.18.1.

We add structured JSONL events at two key sites:

  _send_blocks (P-side, the actual batch_transfer_sync_write):
      {event=send_blocks, remote_session, total_bytes, duration_s,
       t_start_unix, ret, tp_rank, t_log_unix}

  receive_kv_from_single_worker (D-side, async ZMQ-driven pull):
      {event=receive_kv_enter / receive_kv_finish, worker_addr,
       req_ids, duration_s, t_start_unix, tp_rank, t_log_unix}

Logs land in ``$MB2_LOG_DIR/mb2_transfer_pid<pid>.jsonl`` (default
MB2_LOG_DIR=/tmp). One file per PID so multi-worker vllm doesn't fight.

Apply/revert markers: ``# MB2_INSTRUMENT_START/END``.

Usage:
  python instrument_mooncake.py --apply  [--venv PATH]
  python instrument_mooncake.py --revert [--venv PATH]
  python instrument_mooncake.py --check  [--venv PATH]
"""
from __future__ import annotations

import argparse
import re
from pathlib import Path

DEFAULT_VENV = Path("/home/admin/cpfs/wjh/agentic-kv-fresh/.venv")
TARGET_REL = (
    "lib/python3.12/site-packages/vllm/"
    "distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py"
)

START_MARK = "# MB2_INSTRUMENT_START"
END_MARK = "# MB2_INSTRUMENT_END"

HEADER_ANCHOR = "import time\n"

HEADER_INSERT = f"""

{START_MARK}
import json as _mb2_json
import os as _mb2_os
import threading as _mb2_threading
import time as _mb2_time
_MB2_LOG_DIR = _mb2_os.environ.get("MB2_LOG_DIR", "/tmp")
try:
    _mb2_os.makedirs(_MB2_LOG_DIR, exist_ok=True)
except Exception:
    pass
_MB2_LOG_PATH = _mb2_os.path.join(
    _MB2_LOG_DIR, f"mb2_transfer_pid{{_mb2_os.getpid()}}.jsonl"
)
_MB2_LOG_FILE = None
_MB2_LOG_LOCK = _mb2_threading.Lock()
def _mb2_log_event(d):
    global _MB2_LOG_FILE
    d.setdefault("t_log_unix", _mb2_time.time())
    if _MB2_LOG_FILE is None:
        _MB2_LOG_FILE = open(_MB2_LOG_PATH, "a", buffering=1)
    with _MB2_LOG_LOCK:
        _MB2_LOG_FILE.write(_mb2_json.dumps(d) + "\\n")
{END_MARK}
"""

# ---- Patch 2: _send_blocks ----

SEND_TARGET = """    def _send_blocks(
        self,
        remote_session: str,
        src_ptrs: list[int],
        dst_ptrs: list[int],
        lengths: list[int],
    ) -> int:
        start_time = time.perf_counter()
        ret_value = self.engine.batch_transfer_sync_write(
            remote_session, src_ptrs, dst_ptrs, lengths
        )"""

SEND_REPLACE = f"""    def _send_blocks(
        self,
        remote_session: str,
        src_ptrs: list[int],
        dst_ptrs: list[int],
        lengths: list[int],
    ) -> int:
        {START_MARK}
        _mb2_total_bytes = sum(lengths)
        _mb2_t_start_unix = time.time()
        {END_MARK}
        start_time = time.perf_counter()
        ret_value = self.engine.batch_transfer_sync_write(
            remote_session, src_ptrs, dst_ptrs, lengths
        )
        {START_MARK}
        _mb2_log_event({{"event": "send_blocks",
                          "remote_session": remote_session,
                          "total_bytes": _mb2_total_bytes,
                          "duration_s": time.perf_counter() - start_time,
                          "t_start_unix": _mb2_t_start_unix,
                          "ret": ret_value,
                          "tp_rank": getattr(self, "tp_rank", -1)}})
        {END_MARK}"""

# ---- Patch 3: receive_kv_from_single_worker entry ----

RECV_ENTRY_TARGET = """    async def receive_kv_from_single_worker(
        self,
        worker_addr: str,
        pull_metas: dict[ReqId, PullReqMeta],
    ):
        req_ids = set(pull_metas)"""

RECV_ENTRY_REPLACE = f"""    async def receive_kv_from_single_worker(
        self,
        worker_addr: str,
        pull_metas: dict[ReqId, PullReqMeta],
    ):
        req_ids = set(pull_metas)
        {START_MARK}
        _mb2_recv_t_start = time.perf_counter()
        _mb2_recv_t_start_unix = time.time()
        _mb2_log_event({{"event": "receive_kv_enter",
                          "worker_addr": worker_addr,
                          "req_ids": [str(r) for r in req_ids],
                          "t_start_unix": _mb2_recv_t_start_unix,
                          "tp_rank": getattr(self, "tp_rank", -1)}})
        {END_MARK}"""

# ---- Patch 4: receive_kv_from_single_worker FINISH ----

RECV_FINISH_TARGET = """                    if response.status == MooncakeXferResponseStatus.FINISH:
                        break"""

RECV_FINISH_REPLACE = f"""                    if response.status == MooncakeXferResponseStatus.FINISH:
                        {START_MARK}
                        _mb2_log_event({{"event": "receive_kv_finish",
                                          "worker_addr": worker_addr,
                                          "req_ids": [str(r) for r in req_ids],
                                          "duration_s": time.perf_counter() - _mb2_recv_t_start,
                                          "t_start_unix": _mb2_recv_t_start_unix,
                                          "tp_rank": getattr(self, "tp_rank", -1)}})
                        {END_MARK}
                        break"""

PATCHES = [
    ("header",               HEADER_ANCHOR,      HEADER_ANCHOR + HEADER_INSERT),
    ("_send_blocks",         SEND_TARGET,        SEND_REPLACE),
    ("receive_kv (entry)",   RECV_ENTRY_TARGET,  RECV_ENTRY_REPLACE),
    ("receive_kv (FINISH)",  RECV_FINISH_TARGET, RECV_FINISH_REPLACE),
]


def find_target(venv_or_path: Path) -> Path:
    candidates = [venv_or_path, DEFAULT_VENV / TARGET_REL]
    for c in candidates:
        if c.is_file():
            return c
        if c.is_dir():
            sub = c / TARGET_REL
            if sub.is_file():
                return sub
    raise FileNotFoundError(f"cannot find vllm MooncakeConnector at {venv_or_path}")


def is_patched(text: str) -> bool:
    return START_MARK in text


def apply(target: Path) -> None:
    text = target.read_text()
    if is_patched(text):
        print(f"[mb2-instr] already patched: {target}")
        return
    new = text
    for name, src, dst in PATCHES:
        if src not in new:
            raise RuntimeError(
                f"patch {name!r}: anchor not found in {target}. "
                "File may have been refactored."
            )
        new = new.replace(src, dst, 1)
    target.write_text(new)
    print(f"[mb2-instr] applied {len(PATCHES)} patches -> {target}")


def revert(target: Path) -> None:
    text = target.read_text()
    if not is_patched(text):
        print(f"[mb2-instr] not patched (nothing to revert): {target}")
        return
    pat = re.compile(
        r"[ \t]*" + re.escape(START_MARK) + r".*?" + re.escape(END_MARK) + r"\n?",
        flags=re.DOTALL,
    )
    new = pat.sub("", text)
    new = re.sub(r"\n{3,}", "\n\n", new)
    target.write_text(new)
    print(f"[mb2-instr] reverted: {target}")


def main() -> None:
    p = argparse.ArgumentParser()
    p.add_argument("--apply", action="store_true")
    p.add_argument("--revert", action="store_true")
    p.add_argument("--check", action="store_true")
    p.add_argument("--venv", type=Path, default=DEFAULT_VENV)
    args = p.parse_args()
    target = find_target(args.venv)
    if args.apply:
        apply(target)
    elif args.revert:
        revert(target)
    elif args.check:
        text = target.read_text()
        print(f"[mb2-instr] {'PATCHED' if is_patched(text) else 'CLEAN'}: {target}")
    else:
        p.error("specify --apply / --revert / --check")


if __name__ == "__main__":
    main()