[Fix] fix naming style and remove registry

Author: flesher0813

MANIFEST.in

@@ -1,8 +1,4 @@
-include LICENSE
-include pyproject.toml
 include CMakeLists.txt
-include requirements.txt
-include setup.py
-
-recursive-include examples *
-recursive-include benchmarks *
+graft ucm
+graft examples
+graft benchmarks

docs/source/developer-guide/add_metrics.md

@@ -3,8 +3,8 @@ UCM supports custom metrics with bidirectional updates from both Python and C++
 
 ## Architecture Overview
 The metrics consists of these components below:
-- **metrics** : Central stats registry that manages all metric lifecycle operations (registration, creation, updates, queries)
-- **observability.py** : Prometheus integration layer that handles metric exposition and multiprocess collection
+- **monitor** : Central stats registry that manages all metric lifecycle operations (registration, creation, updates, queries)
+- **observability.py** : Prometheus integration layer that handles metric exposition
 - **metrics_config.yaml** : Declarative configuration that defines which custom metrics to register and their properties
 
 ## Getting Started
@@ -31,7 +31,7 @@ prometheus:
 
   # Gauge metrics configuration
   gauges:
-    - name: "lookup_hit_rate"
+    - name: "external_lookup_hit_rate"
       documentation: "Hit rate of ucm lookup requests"
       multiprocess_mode: "livemostrecent"
 
@@ -43,15 +43,14 @@ prometheus:
 ```
 
 ### Use Monitor APIs to Update Stats
-The monitor provides a unified interface for metric operations. Note that the workflow requires registering a stats class before creating an instance.
+The monitor provides a unified interface for metric operations. Users only need to create stats and update them, while the observability component is responsible for fetching the stats and pushing them to Prometheus.
 :::::{tab-set}
 :sync-group: install
 
 ::::{tab-item} Python side interfaces
 :selected:
 :sync: py
 **Lifecycle Methods**
-- `register_istats(name, py::object)`: Register a new stats class implementation.
 - `create_stats(name)`: Create and initialize a registered stats object.
 
 **Operation Methods**
@@ -64,11 +63,8 @@ The monitor provides a unified interface for metric operations. Note that the wo
 
 **Example:** Using built-in ConnStats
 ```python
-from ucm.integration.vllm.conn_stats import ConnStats
 from ucm.shared.metrics import ucmmonitor
 
-conn_stats = ConnStats(name="ConnStats")
-ucmmonitor.register_stats("ConnStats", conn_stats) # Register stats
 ucmmonitor.create_stats("ConnStats") # Create a stats obj
 
 # Update stats
@@ -85,7 +81,6 @@ See more detailed example in [test case](https://github.com/ModelEngine-Group/un
 ::::{tab-item} C++ side interfaces
 :sync: cc
 **Lifecycle Methods**
-- `RegistStats(std::string name, Creator creator)`: Register a new stats class implementation.
 - `CreateStats(const std::string& name)`: Create and initialize a registered stats object.
 
 **Operation Methods**
@@ -102,10 +97,8 @@ UCM supports custom metrics by following steps:
    ```c++
     target_link_libraries(xxxstore PUBLIC storeinfra monitor_static)
     ```
-- Step 2: Inheriting from the IStats class to implement custom stats classes
-- Step 3: Register stats class to monitor
-- Step 4: Create stats object in monitor
-- Step 5: Update or get stats info using operation methods
+- Step 2: Create stats object using function **CreateStats**
+- Step 3: Update using function **UpdateStats**
 
 See more detailed example in [test case](https://github.com/ModelEngine-Group/unified-cache-management/tree/develop/ucm/shared/test/case).
 

pyproject.toml

@@ -1,10 +1,14 @@
 [build-system]
-requires = ["setuptools>=45", "wheel", "cmake", "torch", "pybind11"]
+requires = [
+    "setuptools>=64",
+    "cmake>=3.18",
+    "wheel",
+]
 build-backend = "setuptools.build_meta"
 
 [project]
 name = "uc-manager"
-authors = [{name = "UCM Team"}]
+authors = [{name = "Unified Cache Team"}]
 license = { file="LICENSE" }
 readme = "README.md"
 description = "Persist and reuse KV Cache to speedup your LLM."

setup.py

@@ -23,168 +23,92 @@
 #
 
 import os
-import shutil
 import subprocess
 import sys
-import sysconfig
-import warnings
-from glob import glob
 
-import pybind11
-import torch
-import torch.utils.cpp_extension
 from setuptools import Extension, find_packages, setup
 from setuptools.command.build_ext import build_ext
 
-# Suppress warnings about packages absent from packages configuration
-# These are expected for C++ source directories, test directories, etc.
-warnings.filterwarnings(
-    "ignore", message=".*Package.*is absent from the `packages` configuration.*"
-)
-
 ROOT_DIR = os.path.abspath(os.path.dirname(__file__))
 PLATFORM = os.getenv("PLATFORM")
-
 ENABLE_SPARSE = os.getenv("ENABLE_SPARSE")
 
 
-def _enable_sparse() -> bool:
+def enable_sparse() -> bool:
     return ENABLE_SPARSE is not None and ENABLE_SPARSE.lower() == "true"
 
 
-def _is_cuda() -> bool:
-    return PLATFORM == "cuda" or (hasattr(torch, "cuda") and torch.cuda.is_available())
-
-
-def _is_maca() -> bool:
-    return PLATFORM == "maca"
+def is_editable_mode() -> bool:
+    commands = [arg.lower() for arg in sys.argv]
+    return "develop" in commands or "--editable" in commands or "-e" in commands
 
 
 class CMakeExtension(Extension):
-    def __init__(self, name: str, sourcedir: str = ""):
+    def __init__(self, name: str, source_dir: str = ""):
         super().__init__(name, sources=[])
-        self.sourcedir = os.path.abspath(sourcedir)
+        self.cmake_file_path = os.path.abspath(source_dir)
 
 
 class CMakeBuild(build_ext):
     def run(self):
+        build_dir = os.path.abspath(self.build_temp)
+        os.makedirs(build_dir, exist_ok=True)
+
         for ext in self.extensions:
             self.build_cmake(ext)
 
-        self._copy_so_files_to_build_lib()
-
     def build_cmake(self, ext: CMakeExtension):
-        build_dir = self.build_temp
-        os.makedirs(build_dir, exist_ok=True)
+        build_dir = os.path.abspath(self.build_temp)
+        install_dir = os.path.abspath(self.build_lib)
+        if is_editable_mode:
+            install_dir = ext.cmake_file_path
 
         cmake_args = [
-            "cmake",
             "-DCMAKE_BUILD_TYPE=Release",
             f"-DPYTHON_EXECUTABLE={sys.executable}",
+            f"-DCMAKE_INSTALL_PREFIX={install_dir}",
         ]
 
-        torch_cmake_prefix = torch.utils.cmake_prefix_path
-        pybind11_cmake_dir = pybind11.get_cmake_dir()
+        if enable_sparse():
+            cmake_args += ["-DBUILD_UCM_SPARSE=ON"]
+
+        match PLATFORM:
+            case "cuda":
+                cmake_args += ["-DRUNTIME_ENVIRONMENT=cuda"]
+            case "ascend":
+                cmake_args += ["-DRUNTIME_ENVIRONMENT=ascend"]
+            case "musa":
+                cmake_args += ["-DRUNTIME_ENVIRONMENT=musa"]
+            case "maca":
+                cmake_args += ["-DRUNTIME_ENVIRONMENT=maca"]
+                cmake_args += ["-DBUILD_UCM_SPARSE=OFF"]
+            case _:
+                cmake_args += ["-DRUNTIME_ENVIRONMENT=simu"]
+                cmake_args += ["-DBUILD_UCM_SPARSE=OFF"]
 
-        cmake_prefix_paths = [torch_cmake_prefix, pybind11_cmake_dir]
-        cmake_args.append(f"-DCMAKE_PREFIX_PATH={';'.join(cmake_prefix_paths)}")
-
-        torch_includes = torch.utils.cpp_extension.include_paths()
-        python_include = sysconfig.get_path("include")
-        pybind11_include = pybind11.get_include()
-
-        all_includes = torch_includes + [python_include, pybind11_include]
-        cmake_include_string = ";".join(all_includes)
-        cmake_args.append(f"-DEXTERNAL_INCLUDE_DIRS={cmake_include_string}")
-
-        if _is_cuda():
-            cmake_args.append("-DRUNTIME_ENVIRONMENT=cuda")
-        else:
-            cmake_args.append("-DRUNTIME_ENVIRONMENT=ascend")
-
-        if _enable_sparse():
-            cmake_args.append("-DBUILD_UCM_SPARSE=ON")
-
-        cmake_args.append(ext.sourcedir)
-
-        print(f"[INFO] Building {ext.name} module with CMake")
-        print(f"[INFO] Source directory: {ext.sourcedir}")
-        print(f"[INFO] Build directory: {build_dir}")
-        print(f"[INFO] CMake command: {' '.join(cmake_args)}")
-
-        subprocess.check_call(cmake_args, cwd=build_dir)
+        subprocess.check_call(
+            ["cmake", *cmake_args, ext.cmake_file_path], cwd=build_dir
+        )
         subprocess.check_call(
             ["cmake", "--build", ".", "--config", "Release", "--", "-j8"],
             cwd=build_dir,
         )
 
-    def _copy_so_files_to_build_lib(self):
-        """Copy .so files from source directories to build_lib for installation."""
-        if not hasattr(self, "build_lib") or not self.build_lib:
-            return
-
-        packages = _get_packages()
-        copied_count = 0
-
-        for package in packages:
-            # Source directory where CMake outputs .so files
-            source_package_dir = os.path.join(ROOT_DIR, package.replace(".", os.sep))
-
-            # Destination in build_lib
-            build_package_dir = os.path.join(
-                self.build_lib, package.replace(".", os.sep)
-            )
-
-            # Find all .so files in the source package directory
-            so_files = glob(os.path.join(source_package_dir, "*.so"))
-
-            if so_files:
-                # Ensure destination directory exists
-                os.makedirs(build_package_dir, exist_ok=True)
-
-                # Copy each .so file
-                for so_file in so_files:
-                    dest_file = os.path.join(
-                        build_package_dir, os.path.basename(so_file)
-                    )
-                    shutil.copy2(so_file, dest_file)
-                    copied_count += 1
-                    print(
-                        f"[INFO] Copied {os.path.basename(so_file)} to {build_package_dir}"
-                    )
-
-        if copied_count > 0:
-            print(f"[INFO] Successfully copied {copied_count} .so file(s) to build_lib")
-        else:
-            print(
-                "[WARNING] No .so files found to copy. Extensions may not have been built."
-            )
-
-
-def _get_packages():
-    """Discover Python packages, optionally filtering out sparse-related ones."""
-    sparse_enabled = _enable_sparse()
-    exclude_patterns = []
-    if not sparse_enabled:
-        exclude_patterns.append("ucm.sparse*")
-
-    packages = find_packages(exclude=exclude_patterns)
-    return packages
-
-
-ext_modules = []
-ext_modules.append(CMakeExtension(name="ucm", sourcedir=ROOT_DIR))
+        subprocess.check_call(
+            ["cmake", "--install", ".", "--config", "Release", "--component", "ucm"],
+            cwd=build_dir,
+        )
 
-packages = _get_packages()
 
 setup(
     name="uc-manager",
     version="0.1.2",
     description="Unified Cache Management",
     author="Unified Cache Team",
-    packages=packages,
+    packages=find_packages(),
     python_requires=">=3.10",
-    ext_modules=ext_modules,
+    ext_modules=[CMakeExtension(name="ucm", source_dir=ROOT_DIR)],
     cmdclass={"build_ext": CMakeBuild},
     zip_safe=False,
+    include_package_data=False,
 )

ucm/integration/vllm/conn_stats.py

@@ -553,14 +553,14 @@ def _patch_scheduler() -> None:
 
         from ucm.sparse.base import INVALID_SLOT, UcmSparseRole
         from ucm.sparse.state import ensure_ucm_sparse_initialized, get_ucm_sparse
+        from ucm.utils import Config
 
         def init_ucm_sparse(self):
             self.ucm_sparse = None
             if self.vllm_config.kv_transfer_config is not None:
-                if (
-                    "ucm_sparse_config"
-                    in self.vllm_config.kv_transfer_config.kv_connector_extra_config
-                ):
+                ucm_config = Config(self.vllm_config.kv_transfer_config)
+                ucm_sparse_config = ucm_config.get_config().get("ucm_sparse_config")
+                if ucm_sparse_config:
                     ensure_ucm_sparse_initialized(
                         self.vllm_config, role=UcmSparseRole.SCHEDULER
                     )

ucm/integration/vllm/patch/patch_funcs/v092/vllm_patch.py

@@ -17,7 +17,6 @@
 from vllm.platforms import current_platform
 from vllm.v1.core.sched.output import SchedulerOutput
 
-from ucm.integration.vllm.conn_stats import ConnStats
 from ucm.logger import init_logger
 from ucm.observability import UCMStatsLogger
 from ucm.shared.metrics import ucmmonitor
@@ -173,8 +172,6 @@ def __init__(self, vllm_config: "VllmConfig", role: KVConnectorRole):
 
         self.metrics_config = self.launch_config.get("metrics_config_path", "")
         if self.metrics_config:
-            conn_stats = ConnStats(name="ConnStats")
-            ucmmonitor.register_stats("ConnStats", conn_stats)
             ucmmonitor.create_stats("ConnStats")
             self.stats_logger = UCMStatsLogger(
                 vllm_config.model_config.served_model_name,