diff --git a/MANIFEST.in b/MANIFEST.in index be4487de7..e5d51923b 100644 --- a/MANIFEST.in +++ b/MANIFEST.in @@ -2,6 +2,7 @@ include LICENSE include pyproject.toml include CMakeLists.txt include requirements.txt +include setup.py recursive-include examples * recursive-include benchmarks * \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index 1e4f65b25..1783f09c0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -5,8 +5,7 @@ build-backend = "setuptools.build_meta" [project] name = "uc-manager" authors = [{name = "UCM Team"}] -license = "MIT" -license-files = ["LICENSE"] +license = { file="LICENSE" } readme = "README.md" description = "Persist and reuse KV Cache to speedup your LLM." requires-python = ">=3.10" diff --git a/setup.py b/setup.py index 217417054..ee8a44f2e 100644 --- a/setup.py +++ b/setup.py @@ -23,9 +23,11 @@ # import os +import shutil import subprocess import sys import sysconfig +import warnings from glob import glob import pybind11 @@ -34,6 +36,12 @@ 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") @@ -45,15 +53,7 @@ def _enable_sparse() -> bool: def _is_cuda() -> bool: - return PLATFORM == "cuda" - - -def _is_npu() -> bool: - return PLATFORM == "ascend" - - -def _is_musa() -> bool: - return PLATFORM == "musa" + return PLATFORM == "cuda" or (hasattr(torch, "cuda") and torch.cuda.is_available()) def _is_maca() -> bool: @@ -71,6 +71,8 @@ def run(self): 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) @@ -97,18 +99,8 @@ def build_cmake(self, ext: CMakeExtension): if _is_cuda(): cmake_args.append("-DRUNTIME_ENVIRONMENT=cuda") - elif _is_npu(): - cmake_args.append("-DRUNTIME_ENVIRONMENT=ascend") - elif _is_musa(): - cmake_args.append("-DRUNTIME_ENVIRONMENT=musa") - elif _is_maca(): - cmake_args.append("-DRUNTIME_ENVIRONMENT=maca") - cmake_args.append("-DBUILD_UCM_SPARSE=OFF") else: - raise RuntimeError( - "No supported accelerator found. " - "Please ensure either CUDA/MUSA or NPU is available." - ) + cmake_args.append("-DRUNTIME_ENVIRONMENT=ascend") if _enable_sparse(): cmake_args.append("-DBUILD_UCM_SPARSE=ON") @@ -126,33 +118,58 @@ def build_cmake(self, ext: CMakeExtension): 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 -def _get_packages(): - """Discover Python packages, optionally filtering out sparse-related ones.""" - packages = find_packages() - if not _enable_sparse(): - packages = [pkg for pkg in packages if not pkg.startswith("ucm.sparse")] - return packages + 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)) -def _get_package_data_with_so(packages=None): - """Automatically discover all packages and include .so files.""" - if packages is None: - packages = _get_packages() - package_data = {} + # Destination in build_lib + build_package_dir = os.path.join( + self.build_lib, package.replace(".", os.sep) + ) - for package in packages: - # Convert package name to directory path - package_dir = os.path.join(ROOT_DIR, 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." + ) - # Check if this package directory contains .so files - so_files = glob(os.path.join(package_dir, "*.so")) - if so_files: - package_data[package] = ["*.so"] - print(f"[INFO] Including .so files for package: {package}") - print(f"[INFO] Package data: {package_data}") - return package_data +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 = [] @@ -169,6 +186,5 @@ def _get_package_data_with_so(packages=None): python_requires=">=3.10", ext_modules=ext_modules, cmdclass={"build_ext": CMakeBuild}, - package_data=_get_package_data_with_so(packages), zip_safe=False, )