lmcache-cli 0.5.2.dev0__py3-none-any.whl → 0.5.2rc2.dev0__py3-none-any.whl

lmcache/__init__.py

@@ -1,14 +1,17 @@
 # SPDX-License-Identifier: Apache-2.0
 
 # Standard
-from typing import Any
-import importlib
 import sys
-import types
 
 # First Party
 from lmcache.logging import init_logger
-from lmcache.v1.platform.device_ext import DeviceExt
+
+# --------------------------
+# Backend instance & Device detection
+# --------------------------
+from lmcache.v1.platform import get_backend
+from lmcache.v1.platform import torch_dev as torch_dev
+from lmcache.v1.platform import torch_device_type as torch_device_type
 
 try:
     # First Party
@@ -17,135 +20,17 @@ except ImportError:
     __version__ = "unknown"
 
 logger = init_logger(__name__)
-# Standard
 
 __all__ = ["__version__", "torch_dev", "torch_device_type"]
 
-
-# --------------------------
-# Device detection
-# --------------------------
-def _detect_device() -> tuple[Any, str]:
-    """
-    Detect the available accelerator and return the corresponding torch
-    device module and device type string.
-
-    Returns:
-        tuple[Any, str]: A tuple of (torch_device_module, device_type_string),
-            e.g. ``(torch.cuda, "cuda")``, ``(torch.musa, "musa")``, or
-            ``(torch.xpu, "xpu")``.
-    """
-    try:
-        # Third Party
-        import torch
-    except ImportError:
-        return None, "cpu"  # fallback，CLI-only
-
-    if hasattr(torch, "musa") and torch.musa.is_available():  # type: ignore[attr-defined]
-        logger.info("MUSA device is available. Using MUSA for LMCache engine.")
-        return torch.musa, "musa"  # type: ignore[attr-defined]
-    elif hasattr(torch, "xpu") and torch.xpu.is_available():
-        return torch.xpu, "xpu"
-    elif hasattr(torch, "hpu") and torch.hpu.is_available():
-        return torch.hpu, "hpu"
-    elif torch.cuda.is_available():
-        return torch.cuda, "cuda"
-    else:
-        # First Party
-        from lmcache.v1.platform.cpu.stub_cpu_device import StubCPUDevice
-
-        # Fallback: always return torch, cpu as stub
-        return StubCPUDevice("cpu"), "cpu"
-
-
-torch_dev, torch_device_type = _detect_device()
-
-logger.info(" torch_dev=%s, torch_device_type=%s", torch_dev, torch_device_type)
-
-
-# Attach the DeviceExt instance as ``torch_dev.ext``.  This monkey-patches a
-# standard torch module (e.g. ``torch.cuda``) with a custom attribute that does
-# not exist in the original module.  The ``# type: ignore[attr-defined]`` suppresses
-# the expected mypy/pyright "attr-defined" error from this intentional extension.
-if torch_dev is not None:
-    torch_dev.ext = DeviceExt(torch_device_type)  # type: ignore[attr-defined]
-else:
-    logger.warning("torch_dev is None, skipping DeviceExt initialization.")
-    pass
-
-
-# --------------------------
-# Dynamic backend selection
-# --------------------------
-def _get_backend() -> Any:
-    """
-    Try backends in order, first successful import wins.
-    """
-    default_module = importlib.import_module("lmcache.python_ops_fallback")
-    # Third Party
-    import torch
-
-    backend_candidates = [
-        # Keep backend priority aligned with _detect_device().
-        # MUSA currently uses a Python adapter under the platform package,
-        # unlike the compiled XPU/CUDA extension modules.
-        (
-            "lmcache.v1.platform.musa.ops",
-            "musa_ops",
-            lambda: hasattr(torch, "musa") and torch.musa.is_available(),  # type: ignore[attr-defined]
-        ),
-        (
-            "lmcache.xpu_ops",
-            "xpu_ops",
-            lambda: torch.xpu.is_available(),
-        ),
-        (
-            "lmcache.c_ops",
-            "cuda_ops",
-            lambda: torch.cuda.is_available(),
-        ),
-        # should extend to more HWs..
-    ]
-
-    for module_name, backend_name, predicate in backend_candidates:
-        # 1 Check whether the backend is available before importing
-        try:
-            if not predicate():
-                logger.info(
-                    "Skipping backend %s: predicate returned False",
-                    module_name,
-                )
-                continue
-        except Exception as e:
-            logger.warning(
-                "Skipping backend %s: predicate raised error: %s",
-                module_name,
-                e,
-            )
-            continue
-        # 2 Run availability check for the backend
-        try:
-            backend_module = importlib.import_module(module_name)
-            merged_module = types.ModuleType("lmcache.c_ops")
-            merged_module.__dict__.update(default_module.__dict__)
-            merged_module.__dict__.update(backend_module.__dict__)
-            logger.info("Using backend: %s", module_name)
-            return merged_module
-        except Exception as e:
-            logger.warning("Failed to import backend %s: %s", module_name, e)
-
-    return default_module
-
-
-# --------------------------
-# Backend instance
-# --------------------------
-try:
-    _ops = _get_backend()
-    # override lmcache.c_ops with merged module,
+_ops = get_backend(torch_device_type)
+if _ops is not None:
+    # Override lmcache.c_ops with merged module,
     # in which:
     #     python_ops_fallback as base,
     #     use backend implementation if exists
     sys.modules["lmcache.c_ops"] = _ops
-except (ImportError, ModuleNotFoundError):
-    logger.debug("No compute backend loaded; CLI-only mode (torch/numba not installed)")
+else:
+    logger.warning(
+        "No compute backend loaded; CLI-only mode (torch/numba not installed)"
+    )

lmcache/_version.py

@@ -18,7 +18,7 @@ version_tuple: tuple[int | str, ...]
 commit_id: str | None
 __commit_id__: str | None
 
-__version__ = version = '0.5.2.dev0'
-__version_tuple__ = version_tuple = (0, 5, 2, 'dev0')
+__version__ = version = '0.5.2rc2.dev0'
+__version_tuple__ = version_tuple = (0, 5, 2, 'rc2', 'dev0')
 
-__commit_id__ = commit_id = 'g979719d7a'
+__commit_id__ = commit_id = 'gcd2c0d6a6'

lmcache/cli/commands/bench/l2_adapter_bench/__init__.py

@@ -30,9 +30,34 @@ class L2AdapterBenchCommand(BaseCommand):
         add_l2_arguments(parser)
 
     def execute(self, args: argparse.Namespace) -> None:
+        # Standard
+        import copy
+        import sys
+
         # First Party
         from lmcache.cli.commands.bench.l2_adapter_bench.command import (
             run_l2_adapter_bench,
         )
 
-        run_l2_adapter_bench(self, args)
+        # A comma-separated --flamegraph-mode profiles one benchmark run per
+        # mode (the modes cannot share a recording window), each rendered to
+        # its own default path.
+        modes = [args.flamegraph_mode]
+        if getattr(args, "flamegraph", "off") == "on" and "," in args.flamegraph_mode:
+            modes = list(
+                dict.fromkeys(
+                    m.strip() for m in args.flamegraph_mode.split(",") if m.strip()
+                )
+            )
+            if args.flamegraph_output:
+                print(
+                    "Error: --flamegraph-output takes a single path; with "
+                    "several modes each graph uses its own default path.",
+                    file=sys.stderr,
+                )
+                sys.exit(2)
+
+        for mode in modes:
+            run_args = copy.copy(args)
+            run_args.flamegraph_mode = mode
+            run_l2_adapter_bench(self, run_args)

lmcache/cli/commands/bench/l2_adapter_bench/command.py

@@ -122,6 +122,46 @@ def add_l2_arguments(parser: argparse.ArgumentParser) -> None:
         default=None,
         help="Run only the specified operation (default: run all).",
     )
+    parser.add_argument(
+        "--flamegraph",
+        choices=["on", "off"],
+        default="off",
+        help=(
+            "Capture a flame graph of the measured phases if turned on. The benchmark "
+            "profiles itself and renders an SVG"
+        ),
+    )
+    parser.add_argument(
+        "--flamegraph-mode",
+        default="on-cpu",
+        metavar="MODE[,MODE...]",
+        help=(
+            "What to sample when profiling is on (default: on-cpu). Pass "
+            "several comma-separated to profile one benchmark run per mode. "
+            "Modes: on-cpu, off-cpu, wakeup, offwake (perf/bcc), wall, gil "
+            "(py-spy). See the 'lmcache tool flamegraph' docs for detail."
+        ),
+    )
+    parser.add_argument(
+        "--flamegraph-output",
+        default="",
+        metavar="PATH",
+        help=(
+            "SVG output path for '--flamegraph on'. Default: "
+            "/tmp/lmcache_bench_flames/<adapter>.<mode>.svg."
+        ),
+    )
+    parser.add_argument(
+        "--flamegraph-scripts-dir",
+        default="",
+        metavar="DIR",
+        help=(
+            "Directory with the FlameGraph scripts (flamegraph.pl, "
+            "stackcollapse-perf.pl); default ~/FlameGraph (cloned there on "
+            "first use). Unused by --flamegraph-mode wall / gil, which "
+            "render their own SVG."
+        ),
+    )
 
 
 # ---------------------------------------------------------------------------
@@ -152,6 +192,14 @@ def run_l2_adapter_bench(command: "BaseCommand", args: argparse.Namespace) -> No
         bench_lookup,
         bench_store,
     )
+    from lmcache.cli.profiling import (
+        PY_SPY_MODES,
+        FlameProfiler,
+        ProfileError,
+        check_profiling_deps,
+        default_output_path,
+        resolve_flamegraph_dir,
+    )
     from lmcache.v1.distributed.l2_adapters import create_l2_adapter
     from lmcache.v1.distributed.l2_adapters.config import (
         parse_args_to_l2_adapters_config,
@@ -224,6 +272,31 @@ def run_l2_adapter_bench(command: "BaseCommand", args: argparse.Namespace) -> No
     l1_buffer = make_aligned_tensor(2 * keys_per_round * data_size, l1_align_bytes)
     l1_memory_desc = create_l1_memory_desc(l1_buffer, align_bytes=l1_align_bytes)
 
+    # Resolve and validate the flame-graph toolchain up front, before any
+    # adapter (and its worker threads) is created. A user who explicitly
+    # passes ``--flamegraph on`` gets a fast, actionable failure if a
+    # required tool is missing, rather than a benchmark that silently runs
+    # unprofiled. When ``--flamegraph`` is off (default), none of this runs
+    # and the benchmark behaves exactly as before.
+    flamegraph_on = getattr(args, "flamegraph", "off") == "on"
+    flamegraph_dir = ""
+    if flamegraph_on:
+        try:
+            check_profiling_deps(args.flamegraph_mode)
+            # py-spy renders its own SVG; only the perf/bcc modes need
+            # the FlameGraph scripts, so do not clone them otherwise.
+            if args.flamegraph_mode not in PY_SPY_MODES:
+                flamegraph_dir = resolve_flamegraph_dir(
+                    args.flamegraph_scripts_dir, log
+                )
+        except ProfileError as e:
+            print(
+                "Error: --flamegraph on was requested but the profiling "
+                f"toolchain is unavailable:\n  {e}",
+                file=sys.stderr,
+            )
+            sys.exit(2)
+
     log("\n[Init] Creating adapter...")
     try:
         adapter = create_l2_adapter(adapter_cfg, l1_memory_desc=l1_memory_desc)
@@ -232,6 +305,38 @@ def run_l2_adapter_bench(command: "BaseCommand", args: argparse.Namespace) -> No
         print(f"[Init] Failed to create adapter: {e}", file=sys.stderr)
         sys.exit(1)
 
+    # Optional self-profiling: record a flame graph of the measured
+    # phases. Built here, after the adapter (and its worker threads)
+    # exist, so the recorder attaches to a fully-started process. The
+    # toolchain was already validated above, so a ProfileError here is a
+    # late failure (e.g. the output directory is not writable); since the
+    # user explicitly requested a flame graph we fail loudly rather than
+    # downgrade, closing the adapter we just opened first.
+    profiler: FlameProfiler | None = None
+    if flamegraph_on:
+        adapter_name = type(adapter).__name__
+        try:
+            profiler = FlameProfiler(
+                mode=args.flamegraph_mode,
+                output=(
+                    args.flamegraph_output
+                    or default_output_path(adapter_name, args.flamegraph_mode)
+                ),
+                flamegraph_dir=flamegraph_dir,
+                pid=os.getpid(),
+                title=f"{args.flamegraph_mode} ({adapter_name})",
+            )
+        except ProfileError as e:
+            print(
+                f"Error: cannot start flame-graph profiling: {e}",
+                file=sys.stderr,
+            )
+            try:
+                adapter.close()
+            except Exception:
+                pass
+            sys.exit(2)
+
     # ------------------------------------------------------------------
     # Idx layout
     # ------------------------------------------------------------------
@@ -347,6 +452,18 @@ def run_l2_adapter_bench(command: "BaseCommand", args: argparse.Namespace) -> No
     last_load_round_keys: list[list] | None = None
 
     try:
+        if profiler is not None:
+            # Pre-build the payload buffers before the recorder starts so
+            # the one-time tensor allocation + fill (benchmark harness work,
+            # not the adapter) is kept out of the flame graph. The batches
+            # are reused across rounds, so this is the only build; warming
+            # it here keeps the recording focused on adapter I/O.
+            if args.only is None or args.only == "store":
+                _store_objs(0)
+            if args.only is None or args.only == "load":
+                _load_objs(0)
+            profiler.start(log)
+
         # ---- Store ----
         if args.only is None or args.only == "store":
             log(f"[Store] Running {warmup} warmup + {rounds} measurement rounds...")
@@ -398,6 +515,11 @@ def run_l2_adapter_bench(command: "BaseCommand", args: argparse.Namespace) -> No
             last_load_round_keys = _build_round_keys(total_rounds - 1)
             log("")
 
+        # Stop profiling before verification / summary so the flame
+        # graph reflects only the measured store/lookup/load work.
+        if profiler is not None:
+            profiler.stop(log)
+
         # ---- Round-trip verification (last measured round only) ----
         if (
             not args.skip_verify
@@ -433,6 +555,10 @@ def run_l2_adapter_bench(command: "BaseCommand", args: argparse.Namespace) -> No
             results=results,
         )
     finally:
+        # Idempotent: a no-op if profiling already stopped on the normal
+        # path; tears the recorder down if a phase raised.
+        if profiler is not None:
+            profiler.stop(log)
         log("[Cleanup] Closing adapter...")
         try:
             adapter.close()

lmcache/cli/commands/bench/server_bench/__init__.py

@@ -30,9 +30,34 @@ class ServerBenchCommand(BaseCommand):
         add_server_arguments(parser)
 
     def execute(self, args: argparse.Namespace) -> None:
+        # Standard
+        import copy
+        import sys
+
         # First Party
         from lmcache.cli.commands.bench.server_bench.command import (
             run_server_bench,
         )
 
-        run_server_bench(self, args)
+        # A comma-separated --flamegraph-mode drives the load once per mode
+        # (the modes cannot share a recording window), each rendered to its
+        # own default path.
+        modes = [args.flamegraph_mode]
+        if getattr(args, "flamegraph", "off") == "on" and "," in args.flamegraph_mode:
+            modes = list(
+                dict.fromkeys(
+                    m.strip() for m in args.flamegraph_mode.split(",") if m.strip()
+                )
+            )
+            if args.flamegraph_output:
+                print(
+                    "Error: --flamegraph-output takes a single path; with "
+                    "several modes each graph uses its own default path.",
+                    file=sys.stderr,
+                )
+                sys.exit(2)
+
+        for mode in modes:
+            run_args = copy.copy(args)
+            run_args.flamegraph_mode = mode
+            run_server_bench(self, run_args)

lmcache/cli/commands/bench/server_bench/command.py

@@ -63,8 +63,12 @@ from lmcache.cli.commands.bench.server_bench.helpers import (
 )
 
 if TYPE_CHECKING:
+    # Standard
+    from collections.abc import Callable
+
     # First Party
     from lmcache.cli.commands.base import BaseCommand
+    from lmcache.cli.profiling import FlameProfiler
     from lmcache.v1.multiprocess.custom_types import KVCache
 
 
@@ -192,12 +196,145 @@ def add_server_arguments(parser: argparse.ArgumentParser) -> None:
         help=("HTTP base URL for checksum API (default: http://localhost:8080)"),
     )
 
+    prof = parser.add_argument_group(
+        "server profiling",
+        "Flame-graph the MP server process while this benchmark drives "
+        "load into it. The server's store path (hashing, allocation, "
+        "gather, D2H) runs in its own process, not in this client, so "
+        "profiling attaches to --profile-server-pid rather than to the "
+        "benchmark. See 'lmcache tool flamegraph' for the standalone form.",
+    )
+    prof.add_argument(
+        "--flamegraph",
+        choices=["on", "off"],
+        default="off",
+        help="Record a flame graph of the server during the run (default: off).",
+    )
+    prof.add_argument(
+        "--profile-server-pid",
+        type=int,
+        default=0,
+        metavar="PID",
+        help=(
+            "Server process to profile, e.g. $(pgrep -f 'lmcache server'). "
+            "Required when --flamegraph on."
+        ),
+    )
+    prof.add_argument(
+        "--flamegraph-mode",
+        default="gil",
+        metavar="MODE[,MODE...]",
+        help=(
+            "What to sample in the server (default: gil). Pass several "
+            "comma-separated to profile one load run per mode. Modes: on-cpu, "
+            "off-cpu, wakeup, offwake (perf/bcc), wall, gil (py-spy). perf/bcc "
+            "name Python functions only when the server was launched with "
+            "PYTHONPERFSUPPORT=1. See the 'lmcache tool flamegraph' docs."
+        ),
+    )
+    prof.add_argument(
+        "--flamegraph-output",
+        default="",
+        metavar="PATH",
+        help=(
+            "SVG output path. Default: "
+            "/tmp/lmcache_bench_flames/server-pid<PID>.<mode>.svg."
+        ),
+    )
+    prof.add_argument(
+        "--flamegraph-scripts-dir",
+        default="",
+        metavar="DIR",
+        help=(
+            "Directory with the FlameGraph scripts (flamegraph.pl, "
+            "stackcollapse-perf.pl); default ~/FlameGraph (cloned there on "
+            "first use). Unused by --flamegraph-mode wall / gil, which "
+            "render their own SVG."
+        ),
+    )
+
 
 # ---------------------------------------------------------------------------
 # Public entry point
 # ---------------------------------------------------------------------------
 
 
+def _build_server_profiler(
+    args: argparse.Namespace,
+    log: "Callable[[str], None]",
+) -> "FlameProfiler | None":
+    """Build a profiler attached to the server, or ``None`` if disabled.
+
+    Validates the target pid and toolchain eagerly so a misconfigured run
+    fails before any load is sent. The returned profiler is not started;
+    the caller wraps the load loop with ``start`` / ``stop``.
+
+    Args:
+        args: Parsed CLI arguments for ``lmcache bench server``.
+        log: Progress logger.
+
+    Returns:
+        A ready :class:`FlameProfiler` targeting ``--profile-server-pid``,
+        or ``None`` when ``--flamegraph`` is off.
+    """
+    if getattr(args, "flamegraph", "off") != "on":
+        return None
+
+    # First Party
+    from lmcache.cli.profiling import (
+        PY_SPY_MODES,
+        FlameProfiler,
+        ProfileError,
+        check_profiling_deps,
+        default_output_path,
+        resolve_flamegraph_dir,
+    )
+
+    pid = args.profile_server_pid
+    if pid <= 0:
+        print(
+            "Error: --flamegraph on requires --profile-server-pid "
+            "(the pid of the running 'lmcache server').",
+            file=sys.stderr,
+        )
+        sys.exit(2)
+    try:
+        os.kill(pid, 0)
+    except ProcessLookupError:
+        print(f"Error: no such process: --profile-server-pid {pid}", file=sys.stderr)
+        sys.exit(2)
+    except PermissionError:
+        print(
+            f"Error: server pid {pid} belongs to another user; "
+            "profiling it needs root.",
+            file=sys.stderr,
+        )
+        sys.exit(2)
+
+    try:
+        check_profiling_deps(args.flamegraph_mode)
+        flamegraph_dir = ""
+        if args.flamegraph_mode not in PY_SPY_MODES:
+            flamegraph_dir = resolve_flamegraph_dir(args.flamegraph_scripts_dir, log)
+        output = args.flamegraph_output or default_output_path(
+            f"server-pid{pid}", args.flamegraph_mode
+        )
+        return FlameProfiler(
+            mode=args.flamegraph_mode,
+            output=output,
+            flamegraph_dir=flamegraph_dir,
+            pid=pid,
+            title=f"{args.flamegraph_mode} (server pid {pid})",
+        )
+    except ProfileError as e:
+        print(
+            "Error: --flamegraph on was requested but profiling is "
+            f"unavailable:\n  {e}",
+            file=sys.stderr,
+        )
+        sys.exit(2)
+
+
 def run_server_bench(
     command: "BaseCommand",
     args: argparse.Namespace,
@@ -252,6 +389,12 @@ def run_server_bench(
             "using REGISTER_KV_CACHE + STORE/RETRIEVE over POSIX SHM"
         )
 
+    # The profiler targets the server process (--profile-server-pid), not
+    # this benchmark client. Build it before opening any connection so a
+    # bad pid or a missing toolchain fails immediately, not after a full
+    # benchmark has already run. ``None`` when --flamegraph is off.
+    profiler = _build_server_profiler(args, log)
+
     total_requests = 0
     total_checksum_ok = 0
     total_checksum_fail = 0
@@ -475,6 +618,10 @@ def run_server_bench(
         # ``/cache/checksums`` path.
         client_tensors = None if use_handle else client_kv_tensors
 
+        # Record only the steady-state load, not the one-time registration.
+        if profiler is not None:
+            profiler.start(log)
+
         for seq_no in seq_iter:
             log("=== Request seq=%d ===" % seq_no)
 
@@ -557,6 +704,9 @@ def run_server_bench(
     except KeyboardInterrupt:
         log("\nStopping...")
     finally:
+        # Stop recording once load ends, before teardown
+        if profiler is not None:
+            profiler.stop(log)
         # Deregister our context from the server before tearing down the
         # client. Otherwise the server keeps the registration (and the
         # CUDA-IPC / POSIX-SHM mappings it holds) alive forever, leaking

lmcache/cli/commands/coordinator.py

@@ -102,12 +102,23 @@ class CoordinatorCommand(BaseCommand):
             ),
         )
         parser.add_argument(
-            "--blend-chunk-size",
+            "--chunk-size",
             type=int,
             default=None,
             help=(
-                "Tokens per chunk for the global CacheBlend directory; must "
-                "equal the LMCache chunk size the blend servers use (default: 256)."
+                "Tokens per KV chunk: the CacheBlend match unit and the unit used "
+                "to resolve pin token_ids to keys. Must equal the MP servers' "
+                "--chunk-size (default: 256)."
+            ),
+        )
+        parser.add_argument(
+            "--hash-algorithm",
+            type=str,
+            default=None,
+            help=(
+                "Token hash algorithm for pin key resolution; must equal the MP "
+                "servers' --hash-algorithm. 'blake3' (default) is self-contained; "
+                "other algorithms require vLLM importable in the coordinator."
             ),
         )
         parser.add_argument(
@@ -172,7 +183,8 @@ class CoordinatorCommand(BaseCommand):
                 ("eviction_check_interval", args.eviction_check_interval),
                 ("eviction_ratio", args.eviction_ratio),
                 ("trigger_watermark", args.trigger_watermark),
-                ("blend_chunk_size", args.blend_chunk_size),
+                ("chunk_size", args.chunk_size),
+                ("hash_algorithm", args.hash_algorithm),
                 ("blend_probe_stride", args.blend_probe_stride),
                 ("timeout_keep_alive", args.timeout_keep_alive),
             )