lmcache-cli 0.4.8.dev42__py3-none-any.whl → 0.5.1.dev0__py3-none-any.whl

lmcache/__init__.py

@@ -74,6 +74,14 @@ def _get_backend() -> Any:
     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",

lmcache/_version.py

@@ -18,7 +18,7 @@ version_tuple: tuple[int | str, ...]
 commit_id: str | None
 __commit_id__: str | None
 
-__version__ = version = '0.4.8.dev42'
-__version_tuple__ = version_tuple = (0, 4, 8, 'dev42')
+__version__ = version = '0.5.1.dev0'
+__version_tuple__ = version_tuple = (0, 5, 1, 'dev0')
 
-__commit_id__ = commit_id = 'g9ffb104f3'
+__commit_id__ = commit_id = 'g389b9cfcb'

lmcache/cli/commands/__init__.py

@@ -8,7 +8,7 @@ automatically — no edits to this file are required.
 """
 
 # First Party
-from lmcache.cli.commands.base import BaseCommand
+from lmcache.cli.commands.base import BaseCommand, CompositeCommand
 from lmcache.v1.utils.subclass_discovery import discover_subclasses
 
 
@@ -37,4 +37,4 @@ def _discover_commands() -> list[BaseCommand]:
 
 ALL_COMMANDS: list[BaseCommand] = _discover_commands()
 
-__all__ = ["ALL_COMMANDS", "BaseCommand"]
+__all__ = ["ALL_COMMANDS", "BaseCommand", "CompositeCommand"]

lmcache/cli/commands/base.py

@@ -1,9 +1,13 @@
 # SPDX-License-Identifier: Apache-2.0
 """Abstract base class and shared helpers for CLI subcommands."""
 
+# Future
+from __future__ import annotations
+
 # Standard
 import abc
 import argparse
+import sys
 
 # First Party
 from lmcache.cli.metrics import (
@@ -12,6 +16,9 @@ from lmcache.cli.metrics import (
     StreamHandler,
     get_formatter,
 )
+from lmcache.logging import init_logger
+
+logger = init_logger(__name__)
 
 
 class BaseCommand(abc.ABC):
@@ -125,6 +132,96 @@ class BaseCommand(abc.ABC):
         return metrics
 
 
+class CompositeCommand(BaseCommand):
+    """Base class for commands that contain auto-discovered sub-subcommands.
+
+    Subclasses only need to implement :meth:`name` and :meth:`help`.
+    Sub-subcommands are discovered automatically by scanning the package
+    for concrete :class:`BaseCommand` subclasses.
+
+    Example::
+
+        class QueryCommand(CompositeCommand):
+            def name(self) -> str:
+                return "query"
+
+            def help(self) -> str:
+                return "Run one inference request and report metrics."
+    """
+
+    def add_arguments(self, _parser: argparse.ArgumentParser) -> None:
+        """No top-level arguments; all args are registered by subcommands."""
+
+    def register(self, subparsers: argparse._SubParsersAction) -> None:
+        """Register this command and auto-discover all sub-subcommands.
+
+        Scans the package where this class is defined for concrete
+        :class:`BaseCommand` subclasses and registers each one as a
+        nested subcommand.
+
+        Args:
+            subparsers: The subparsers action from the root parser.
+        """
+        # Deferred import to avoid circular dependency
+        # First Party
+        from lmcache.v1.utils.subclass_discovery import discover_subclasses
+
+        parser = subparsers.add_parser(
+            self.name(),
+            help=self.help(),
+            description=self.help(),
+        )
+        inner = parser.add_subparsers(
+            dest=f"{self.name()}_target",
+            required=True,
+        )
+
+        # Discover subcommands in the package where the concrete
+        # CompositeCommand subclass is defined. Exclude the module
+        # that defines the composite command itself (the __init__.py).
+        package = self.__class__.__module__
+
+        def _raise(module_name: str, exc: Exception) -> None:
+            raise exc
+
+        self._subcmds: dict[str, BaseCommand] = {}
+        for cls in discover_subclasses(
+            package,
+            BaseCommand,  # type: ignore[type-abstract]
+            module_filter=lambda name: not name.startswith("_"),
+            require_defined_in_module=True,
+            on_import_error=_raise,
+        ):
+            # Skip the composite command class itself
+            if cls is self.__class__:
+                continue
+            inst = cls()
+            self._subcmds[inst.name()] = inst
+            inst.register(inner)
+
+        logger.debug(
+            "CompositeCommand[%s] discovered subcommands: %s",
+            self.name(),
+            list(self._subcmds.keys()),
+        )
+
+    def execute(self, args: argparse.Namespace) -> None:
+        """Dispatch to the appropriate sub-subcommand.
+
+        Args:
+            args: Parsed CLI arguments.
+        """
+        target = getattr(args, f"{self.name()}_target", None)
+        subcmd = self._subcmds.get(target) if target else None
+        if subcmd is None:
+            print(
+                f"Unknown {self.name()} target: {target}",
+                file=sys.stderr,
+            )
+            sys.exit(1)
+        subcmd.execute(args)
+
+
 def _add_output_args(parser: argparse.ArgumentParser) -> None:
     """Add the common ``--format`` and ``--output`` flags.
 

lmcache/cli/commands/bench/__init__.py

@@ -1,30 +1,17 @@
 # SPDX-License-Identifier: Apache-2.0
-"""``lmcache bench`` command — sustained performance benchmarking."""
+"""``lmcache bench`` command — sustained performance benchmarking.
 
-# Standard
-import argparse
-import sys
+Sub-subcommands are auto-discovered from sub-packages in this package.
+To add a new benchmark, create a sub-package with a concrete
+:class:`BaseCommand` subclass in its ``__init__.py`` — no edits to
+this file are required.
+"""
 
 # First Party
-from lmcache.cli.commands.base import BaseCommand
-from lmcache.cli.commands.bench.engine_bench.command import (
-    register_engine_parser,
-    run_engine_bench,
-)
-from lmcache.cli.commands.bench.l2_adapter_bench.command import (
-    register_l2_parser,
-    run_l2_adapter_bench,
-)
-from lmcache.cli.commands.bench.server_bench.command import (
-    register_server_parser,
-    run_server_bench,
-)
-from lmcache.logging import init_logger
+from lmcache.cli.commands.base import CompositeCommand
 
-logger = init_logger(__name__)
 
-
-class BenchCommand(BaseCommand):
+class BenchCommand(CompositeCommand):
     """CLI command for sustained performance benchmarking."""
 
     def name(self) -> str:
@@ -32,36 +19,3 @@ class BenchCommand(BaseCommand):
 
     def help(self) -> str:
         return "Run sustained performance benchmarks."
-
-    def add_arguments(self, _parser: argparse.ArgumentParser) -> None:
-        pass  # args registered in register() via subparsers
-
-    def register(self, subparsers: argparse._SubParsersAction) -> None:
-        parser = subparsers.add_parser(
-            self.name(),
-            help=self.help(),
-            description="Run sustained performance benchmarks.",
-        )
-        inner = parser.add_subparsers(
-            dest="bench_target",
-            required=True,
-            metavar="{engine,server,l2}",
-        )
-        register_engine_parser(inner, self.execute)
-        register_server_parser(inner, self.execute)
-        register_l2_parser(inner, self.execute)
-
-    def execute(self, args: argparse.Namespace) -> None:
-        handlers = {
-            "engine": lambda a: run_engine_bench(self, a),
-            "server": lambda a: run_server_bench(self, a),
-            "l2": lambda a: run_l2_adapter_bench(self, a),
-        }
-        handler = handlers.get(args.bench_target)
-        if handler is None:
-            print(
-                f"Unknown bench target: {args.bench_target}",
-                file=sys.stderr,
-            )
-            sys.exit(1)
-        handler(args)

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

@@ -1 +1,49 @@
 # SPDX-License-Identifier: Apache-2.0
+"""``lmcache bench engine`` subpackage.
+
+Exposes :class:`EngineBenchCommand` for auto-discovery by
+:class:`~lmcache.cli.commands.base.CompositeCommand`.
+"""
+
+# Standard
+import argparse
+
+# First Party
+from lmcache.cli.commands.base import BaseCommand
+
+
+class EngineBenchCommand(BaseCommand):
+    """Benchmark an inference engine."""
+
+    def name(self) -> str:
+        return "engine"
+
+    def help(self) -> str:
+        return "Benchmark an inference engine."
+
+    def add_arguments(self, parser: argparse.ArgumentParser) -> None:
+        # First Party
+        from lmcache.cli.commands.bench.engine_bench.command import (
+            add_engine_arguments,
+        )
+
+        add_engine_arguments(parser)
+
+    def register(self, subparsers: argparse._SubParsersAction) -> None:
+        """Register without the common output args (--format/--output).
+
+        The engine bench manages its own output lifecycle (CSV/JSON export,
+        --quiet) so the generic _add_output_args flags are not applicable
+        and would only clutter ``-h`` output.
+        """
+        parser = subparsers.add_parser(self.name(), help=self.help())
+        self.add_arguments(parser)
+        parser.set_defaults(func=self.execute)
+
+    def execute(self, args: argparse.Namespace) -> None:
+        # First Party
+        from lmcache.cli.commands.bench.engine_bench.command import (
+            run_engine_bench,
+        )
+
+        run_engine_bench(self, args)

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

@@ -1,10 +1,9 @@
 # SPDX-License-Identifier: Apache-2.0
 """``lmcache bench engine`` subcommand implementation.
 
-This module owns the full registration + execution flow for the
-inference engine benchmark. ``BenchCommand`` only forwards CLI dispatch
-to :func:`run_engine_bench` and parser registration to
-:func:`register_engine_parser`.
+This module provides argument registration via :func:`add_engine_arguments`
+and the execution orchestrator :func:`run_engine_bench` for the inference
+engine benchmark.
 """
 
 # Future
@@ -55,26 +54,12 @@ _LDQA_MAX_OUTPUT_LENGTH_DEFAULT = 128
 # ---------------------------------------------------------------------------
 
 
-def register_engine_parser(
-    subparsers: argparse._SubParsersAction,
-    dispatch_func,
-) -> argparse.ArgumentParser:
-    """Register the ``lmcache bench engine`` subcommand parser.
+def add_engine_arguments(parser: argparse.ArgumentParser) -> None:
+    """Add ``lmcache bench engine`` arguments to *parser*.
 
     Args:
-        subparsers: The ``bench`` subparsers action.
-        dispatch_func: Function to bind via ``set_defaults(func=...)``.
-            Typically ``BenchCommand.execute`` so that the outer
-            dispatcher can route the call back into
-            :func:`run_engine_bench`.
-
-    Returns:
-        The created ``ArgumentParser`` (mostly for testing).
+        parser: The ``ArgumentParser`` for the engine bench subcommand.
     """
-    parser = subparsers.add_parser(
-        "engine",
-        help="Benchmark an inference engine.",
-    )
 
     # --- Config file ---
     parser.add_argument(
@@ -337,9 +322,6 @@ def register_engine_parser(
         help="Number of requests to send (default: 50).",
     )
 
-    parser.set_defaults(func=dispatch_func)
-    return parser
-
 
 # ---------------------------------------------------------------------------
 # Argument resolution helpers

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

@@ -1,2 +1,38 @@
 # SPDX-License-Identifier: Apache-2.0
-"""L2 adapter benchmark subpackage for ``lmcache bench l2``."""
+"""``lmcache bench l2`` subpackage.
+
+Exposes :class:`L2AdapterBenchCommand` for auto-discovery by
+:class:`~lmcache.cli.commands.base.CompositeCommand`.
+"""
+
+# Standard
+import argparse
+
+# First Party
+from lmcache.cli.commands.base import BaseCommand
+
+
+class L2AdapterBenchCommand(BaseCommand):
+    """Benchmark an L2 adapter (store / lookup / load)."""
+
+    def name(self) -> str:
+        return "l2"
+
+    def help(self) -> str:
+        return "Benchmark an L2 adapter (store / lookup / load)."
+
+    def add_arguments(self, parser: argparse.ArgumentParser) -> None:
+        # First Party
+        from lmcache.cli.commands.bench.l2_adapter_bench.command import (
+            add_l2_arguments,
+        )
+
+        add_l2_arguments(parser)
+
+    def execute(self, args: argparse.Namespace) -> None:
+        # First Party
+        from lmcache.cli.commands.bench.l2_adapter_bench.command import (
+            run_l2_adapter_bench,
+        )
+
+        run_l2_adapter_bench(self, args)

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

@@ -1,10 +1,9 @@
 # SPDX-License-Identifier: Apache-2.0
 """``lmcache bench l2`` subcommand implementation.
 
-This module owns the full registration + execution flow for the L2
-adapter benchmark. ``BenchCommand`` only forwards CLI dispatch to
-:func:`run_l2_adapter_bench` and parser registration to
-:func:`register_l2_parser`.
+This module provides argument registration via :func:`add_l2_arguments`
+and the execution orchestrator :func:`run_l2_adapter_bench` for the L2
+adapter benchmark.
 """
 
 # Future
@@ -16,15 +15,6 @@ import argparse
 import os
 import sys
 
-# First Party
-# Reuse the common helper that wires up ``--format / --output /
-# --quiet`` onto a subparser. ``BenchCommand.register`` is overridden
-# and creates inner subparsers manually, bypassing the auto-wiring
-# that ``BaseCommand.register`` normally performs, so we attach those
-# common flags ourselves only on the L2 subparser. The ``engine`` and
-# ``kvcache`` subparsers intentionally stay untouched.
-from lmcache.cli.commands.base import _add_output_args
-
 if TYPE_CHECKING:
     # First Party
     from lmcache.cli.commands.base import BaseCommand
@@ -36,32 +26,12 @@ if TYPE_CHECKING:
 # ---------------------------------------------------------------------------
 
 
-def register_l2_parser(
-    subparsers: argparse._SubParsersAction,
-    dispatch_func,
-) -> argparse.ArgumentParser:
-    """Register the ``lmcache bench l2`` subcommand parser.
+def add_l2_arguments(parser: argparse.ArgumentParser) -> None:
+    """Add ``lmcache bench l2`` arguments to *parser*.
 
     Args:
-        subparsers: The ``bench`` subparsers action.
-        dispatch_func: Function to bind via ``set_defaults(func=...)``.
-            Typically ``BenchCommand.execute`` so that the outer
-            dispatcher can route the call back into
-            :func:`run_l2_adapter_bench`.
-
-    Returns:
-        The created ``ArgumentParser`` (mostly for testing).
+        parser: The ``ArgumentParser`` for the L2 bench subcommand.
     """
-    parser = subparsers.add_parser(
-        "l2",
-        help="Benchmark an L2 adapter (store / lookup / load).",
-        description=(
-            "Benchmark L2 adapters using the standard LMCache adapter "
-            "configuration mechanism (parse_args_to_l2_adapters_config "
-            "+ create_l2_adapter). Any registered adapter type can be "
-            "tested without code changes."
-        ),
-    )
 
     parser.add_argument(
         "--l2-adapter",
@@ -153,14 +123,6 @@ def register_l2_parser(
         help="Run only the specified operation (default: run all).",
     )
 
-    # Common ``--format / --output / --quiet`` flags. Attached only
-    # to the L2 subparser; the ``engine`` and ``kvcache`` subparsers
-    # intentionally keep their existing arguments unchanged.
-    _add_output_args(parser)
-
-    parser.set_defaults(func=dispatch_func)
-    return parser
-
 
 # ---------------------------------------------------------------------------
 # Core benchmark runner

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

@@ -20,7 +20,7 @@ import time
 
 # First Party
 from lmcache.native_storage_ops import Bitmap
-from lmcache.v1.distributed.api import ObjectKey
+from lmcache.v1.distributed.api import MemoryLayoutDesc, ObjectKey
 from lmcache.v1.distributed.internal_api import L2StoreResult
 from lmcache.v1.memory_management import MemoryObj
 
@@ -37,6 +37,10 @@ LogFn = Callable[[str], None]
 KeyProvider = Callable[[int], list[list[ObjectKey]]]
 ObjProvider = Callable[[int], list[list[MemoryObj]]]
 
+# TODO: bench passes a placeholder layout_desc; a real layout may be
+# required here in the future (e.g. when benchmarking the P2P adapter).
+_PLACEHOLDER_LAYOUT_DESC = MemoryLayoutDesc(shapes=[], dtypes=[])
+
 
 def _bitmap_count(bitmap: Bitmap | None) -> int:
     """Count how many bits are set in *bitmap*. Returns 0 when None."""
@@ -217,6 +221,11 @@ def bench_lookup(
         expected_hit_count=expected_hit_count,
     )
 
+    log(
+        "bench_lookup uses a placeholder MemoryLayoutDesc; this may need a "
+        "real layout for layout-sensitive adapters"
+    )
+
     for r in range(rounds):
         keys_batches = keys_for_round(r)
         assert len(keys_batches) == in_flight
@@ -224,7 +233,11 @@ def bench_lookup(
         t0 = time.perf_counter()
         task_ids: list[int] = []
         for i in range(in_flight):
-            task_ids.append(adapter.submit_lookup_and_lock_task(keys_batches[i]))
+            task_ids.append(
+                adapter.submit_lookup_and_lock_task(
+                    keys_batches[i], _PLACEHOLDER_LAYOUT_DESC
+                )
+            )
 
         results = _wait_lookup_finished(adapter, task_ids, 60.0)
         t1 = time.perf_counter()

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

@@ -1,2 +1,38 @@
 # SPDX-License-Identifier: Apache-2.0
-"""Server benchmark subpackage for ``lmcache bench server``."""
+"""``lmcache bench server`` subpackage.
+
+Exposes :class:`ServerBenchCommand` for auto-discovery by
+:class:`~lmcache.cli.commands.base.CompositeCommand`.
+"""
+
+# Standard
+import argparse
+
+# First Party
+from lmcache.cli.commands.base import BaseCommand
+
+
+class ServerBenchCommand(BaseCommand):
+    """End-to-end test for LMCache MP cache server."""
+
+    def name(self) -> str:
+        return "server"
+
+    def help(self) -> str:
+        return "End-to-end test for LMCache MP cache server (GPU mode)."
+
+    def add_arguments(self, parser: argparse.ArgumentParser) -> None:
+        # First Party
+        from lmcache.cli.commands.bench.server_bench.command import (
+            add_server_arguments,
+        )
+
+        add_server_arguments(parser)
+
+    def execute(self, args: argparse.Namespace) -> None:
+        # First Party
+        from lmcache.cli.commands.bench.server_bench.command import (
+            run_server_bench,
+        )
+
+        run_server_bench(self, args)