[build] rewrite setup.py

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

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,85 @@
 #
 
 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:
     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"
-
-
 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)
 
         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()
-
-        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)}")
+            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"]
 
-        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/shared/metrics/CMakeLists.txt

@@ -12,4 +12,9 @@ file(GLOB_RECURSE BINDINGS_SRCS CONFIGURE_DEPENDS "${CMAKE_CURRENT_SOURCE_DIR}/c
 pybind11_add_module(ucmmonitor ${BINDINGS_SRCS})
 target_link_libraries(ucmmonitor PRIVATE -Wl,--whole-archive monitor_static -Wl,--no-whole-archive)
 target_include_directories(ucmmonitor PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/cc)
-set_target_properties(ucmmonitor PROPERTIES LIBRARY_OUTPUT_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR})
+
+file(RELATIVE_PATH INSTALL_REL_PATH
+     ${CMAKE_SOURCE_DIR}
+     ${CMAKE_CURRENT_SOURCE_DIR}
+)
+install(TARGETS ucmmonitor LIBRARY DESTINATION ${INSTALL_REL_PATH} COMPONENT ucm)

ucm/shared/trans/CMakeLists.txt

@@ -16,4 +16,9 @@ target_link_libraries(trans PUBLIC infra_status)
 file(GLOB_RECURSE UCMTRANS_CPY_SOURCE_FILES "./cpy/*.cc")
 pybind11_add_module(ucmtrans ${UCMTRANS_CPY_SOURCE_FILES})
 target_link_libraries(ucmtrans PRIVATE trans)
-set_target_properties(ucmtrans PROPERTIES LIBRARY_OUTPUT_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR})
+
+file(RELATIVE_PATH INSTALL_REL_PATH
+     ${CMAKE_SOURCE_DIR}
+     ${CMAKE_CURRENT_SOURCE_DIR}
+)
+install(TARGETS ucmtrans LIBRARY DESTINATION ${INSTALL_REL_PATH} COMPONENT ucm)

ucm/sparse/esa/retrieval/CMakeLists.txt

@@ -2,7 +2,11 @@
 pybind11_add_module(retrieval_backend cpy/retrieval_backend.cpp)
 
 # 设置输出库的目录
-set_target_properties(retrieval_backend PROPERTIES LIBRARY_OUTPUT_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR})
+file(RELATIVE_PATH INSTALL_REL_PATH
+     ${CMAKE_SOURCE_DIR}
+     ${CMAKE_CURRENT_SOURCE_DIR}
+)
+install(TARGETS retrieval_backend LIBRARY DESTINATION ${INSTALL_REL_PATH} COMPONENT ucm)
 
 # 设置头文件目录，以确保 numaf.h 能找到
 target_include_directories(retrieval_backend PUBLIC

ucm/sparse/gsa/offload_ops/CMakeLists.txt

@@ -96,16 +96,8 @@ if(OpenMP_CXX_FOUND)
     target_link_libraries(gsa_offload_ops PRIVATE OpenMP::OpenMP_CXX)
 endif()
 
-# 设置输出路径
-set(OUTPUT_LIB_DIR ${CMAKE_CURRENT_SOURCE_DIR})
-set_target_properties(gsa_offload_ops PROPERTIES
-    LIBRARY_OUTPUT_DIRECTORY ${OUTPUT_LIB_DIR}
-    RUNTIME_OUTPUT_DIRECTORY ${OUTPUT_LIB_DIR}
-)
-
-# 编译后输出信息
-add_custom_command(TARGET gsa_offload_ops POST_BUILD
-    COMMAND ${CMAKE_COMMAND} -E echo "Built gsa_offload_ops successfully for ${RUNTIME_ENVIRONMENT}"
-    COMMAND ${CMAKE_COMMAND} -E echo "CXX11_ABI=${CXX11_ABI}"
-    COMMAND ${CMAKE_COMMAND} -E echo "Output location: ${OUTPUT_LIB_DIR}"
+file(RELATIVE_PATH INSTALL_REL_PATH
+     ${CMAKE_SOURCE_DIR}
+     ${CMAKE_CURRENT_SOURCE_DIR}
 )
+install(TARGETS gsa_offload_ops LIBRARY DESTINATION ${INSTALL_REL_PATH} COMPONENT ucm)

ucm/sparse/gsa/prefetch/CMakeLists.txt

@@ -94,16 +94,8 @@ if(OpenMP_CXX_FOUND)
     target_link_libraries(gsa_prefetch PRIVATE OpenMP::OpenMP_CXX)
 endif()
 
-# 设置输出路径
-set(OUTPUT_LIB_DIR ${CMAKE_CURRENT_SOURCE_DIR})
-set_target_properties(gsa_prefetch PROPERTIES
-    LIBRARY_OUTPUT_DIRECTORY ${OUTPUT_LIB_DIR}
-    RUNTIME_OUTPUT_DIRECTORY ${OUTPUT_LIB_DIR}
-)
-
-# 编译后输出信息
-add_custom_command(TARGET gsa_prefetch POST_BUILD
-    COMMAND ${CMAKE_COMMAND} -E echo "Built gsa_prefetch successfully for ${RUNTIME_ENVIRONMENT}"
-    COMMAND ${CMAKE_COMMAND} -E echo "CXX11_ABI=${CXX11_ABI}"
-    COMMAND ${CMAKE_COMMAND} -E echo "Output location: ${OUTPUT_LIB_DIR}"
+file(RELATIVE_PATH INSTALL_REL_PATH
+     ${CMAKE_SOURCE_DIR}
+     ${CMAKE_CURRENT_SOURCE_DIR}
 )
+install(TARGETS gsa_prefetch LIBRARY DESTINATION ${INSTALL_REL_PATH} COMPONENT ucm)

ucm/sparse/kvcomp/hash_retrieval/CMakeLists.txt

@@ -1,8 +1,11 @@
 # 添加编译目标
 pybind11_add_module(hash_retrieval_backend cpy/hash_retrieval_backend.cpp)
 
-# 设置输出库的目录
-set_target_properties(hash_retrieval_backend PROPERTIES LIBRARY_OUTPUT_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR})
+file(RELATIVE_PATH INSTALL_REL_PATH
+     ${CMAKE_SOURCE_DIR}
+     ${CMAKE_CURRENT_SOURCE_DIR}
+)
+install(TARGETS hash_retrieval_backend LIBRARY DESTINATION ${INSTALL_REL_PATH} COMPONENT ucm)
 
 # 设置头文件目录，以确保 numaf.h 能找到
 target_include_directories(hash_retrieval_backend PUBLIC

ucm/sparse/kvstar/retrieve/CMakeLists.txt

@@ -1,5 +1,5 @@
 # auto detect cuda ------------------------
-if($ENV{PLATFORM} STREQUAL "cuda")
+if(${RUNTIME_ENVIRONMENT} STREQUAL "cuda")
     if(NOT DEFINED CMAKE_CUDA_COMPILER)
         set(CUDA_ROOT "/usr/local/cuda/" CACHE PATH "Path to CUDA root directory")
         set(CMAKE_CUDA_COMPILER ${CUDA_ROOT}/bin/nvcc)
@@ -106,4 +106,8 @@ endif()
 add_subdirectory(core)
 add_subdirectory(py_intf)
 
-set_target_properties(kvstar_retrieve PROPERTIES LIBRARY_OUTPUT_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR})
+file(RELATIVE_PATH INSTALL_REL_PATH
+     ${CMAKE_SOURCE_DIR}
+     ${CMAKE_CURRENT_SOURCE_DIR}
+)
+install(TARGETS kvstar_retrieve LIBRARY DESTINATION ${INSTALL_REL_PATH} COMPONENT ucm)

ucm/sparse/kvstar/retrieve/core/CMakeLists.txt

@@ -4,7 +4,6 @@ add_library(kvstar_retrieve.core STATIC ${CORE_SRC_FILES})
 
 target_include_directories(kvstar_retrieve.core PUBLIC
         "." "api" "domain" "infra"
-        ${EXTERNAL_INCLUDE_DIRS}
         ${NUMA_INSTALL_DIR}/include
 )
 

ucm/store/localstore/CMakeLists.txt

@@ -9,4 +9,9 @@ target_link_libraries(localstore PUBLIC storeinfra storetask)
 file(GLOB_RECURSE UCMSTORE_LOCAL_CPY_SOURCE_FILES "./cpy/*.cc")
 pybind11_add_module(ucmlocalstore ${UCMSTORE_LOCAL_CPY_SOURCE_FILES})
 target_link_libraries(ucmlocalstore PRIVATE localstore)
-set_target_properties(ucmlocalstore PROPERTIES LIBRARY_OUTPUT_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR})
+
+file(RELATIVE_PATH INSTALL_REL_PATH
+     ${CMAKE_SOURCE_DIR}
+     ${CMAKE_CURRENT_SOURCE_DIR}
+)
+install(TARGETS ucmlocalstore LIBRARY DESTINATION ${INSTALL_REL_PATH} COMPONENT ucm)

ucm/store/nfsstore/CMakeLists.txt

@@ -9,4 +9,9 @@ target_link_libraries(nfsstore PUBLIC storeinfra storedevice storetask)
 file(GLOB_RECURSE UCMSTORE_NFS_CPY_SOURCE_FILES "./cpy/*.cc")
 pybind11_add_module(ucmnfsstore ${UCMSTORE_NFS_CPY_SOURCE_FILES})
 target_link_libraries(ucmnfsstore PRIVATE nfsstore)
-set_target_properties(ucmnfsstore PROPERTIES LIBRARY_OUTPUT_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR})
+
+file(RELATIVE_PATH INSTALL_REL_PATH
+     ${CMAKE_SOURCE_DIR}
+     ${CMAKE_CURRENT_SOURCE_DIR}
+)
+install(TARGETS ucmnfsstore LIBRARY DESTINATION ${INSTALL_REL_PATH} COMPONENT ucm)

ucm/store/pcstore/CMakeLists.txt

@@ -9,4 +9,9 @@ target_link_libraries(pcstore PUBLIC trans storeinfra)
 file(GLOB_RECURSE UCMSTORE_PC_CPY_SOURCE_FILES "./cpy/*.cc")
 pybind11_add_module(ucmpcstore ${UCMSTORE_PC_CPY_SOURCE_FILES})
 target_link_libraries(ucmpcstore PRIVATE pcstore)
-set_target_properties(ucmpcstore PROPERTIES LIBRARY_OUTPUT_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR})
+
+file(RELATIVE_PATH INSTALL_REL_PATH
+     ${CMAKE_SOURCE_DIR}
+     ${CMAKE_CURRENT_SOURCE_DIR}
+)
+install(TARGETS ucmpcstore LIBRARY DESTINATION ${INSTALL_REL_PATH} COMPONENT ucm)