Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions MANIFEST.in
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ include LICENSE
include pyproject.toml
include CMakeLists.txt
include requirements.txt
include setup.py

recursive-include examples *
recursive-include benchmarks *
3 changes: 1 addition & 2 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down
100 changes: 58 additions & 42 deletions setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,9 +23,11 @@
#

import os
import shutil
import subprocess
import sys
import sysconfig
import warnings
from glob import glob

import pybind11
Expand All @@ -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")

Expand All @@ -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:
Expand All @@ -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)
Expand All @@ -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")
Expand All @@ -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 = []
Expand All @@ -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,
)