|
23 | 23 | # |
24 | 24 |
|
25 | 25 | import os |
26 | | -import shutil |
27 | 26 | import subprocess |
28 | 27 | import sys |
29 | | -import sysconfig |
30 | | -import warnings |
31 | | -from glob import glob |
32 | 28 |
|
33 | | -import pybind11 |
34 | | -import torch |
35 | | -import torch.utils.cpp_extension |
36 | 29 | from setuptools import Extension, find_packages, setup |
37 | 30 | from setuptools.command.build_ext import build_ext |
38 | 31 |
|
39 | | -# Suppress warnings about packages absent from packages configuration |
40 | | -# These are expected for C++ source directories, test directories, etc. |
41 | | -warnings.filterwarnings( |
42 | | - "ignore", message=".*Package.*is absent from the `packages` configuration.*" |
43 | | -) |
44 | | - |
45 | 32 | ROOT_DIR = os.path.abspath(os.path.dirname(__file__)) |
46 | 33 | PLATFORM = os.getenv("PLATFORM") |
47 | | - |
48 | 34 | ENABLE_SPARSE = os.getenv("ENABLE_SPARSE") |
49 | 35 |
|
50 | 36 |
|
51 | 37 | def _enable_sparse() -> bool: |
52 | 38 | return ENABLE_SPARSE is not None and ENABLE_SPARSE.lower() == "true" |
53 | 39 |
|
54 | 40 |
|
55 | | -def _is_cuda() -> bool: |
56 | | - return PLATFORM == "cuda" or (hasattr(torch, "cuda") and torch.cuda.is_available()) |
57 | | - |
58 | | - |
59 | | -def _is_maca() -> bool: |
60 | | - return PLATFORM == "maca" |
61 | | - |
62 | | - |
63 | 41 | class CMakeExtension(Extension): |
64 | | - def __init__(self, name: str, sourcedir: str = ""): |
| 42 | + def __init__(self, name: str, source_dir: str = ""): |
65 | 43 | super().__init__(name, sources=[]) |
66 | | - self.sourcedir = os.path.abspath(sourcedir) |
| 44 | + self.cmake_file_path = os.path.abspath(source_dir) |
67 | 45 |
|
68 | 46 |
|
69 | 47 | class CMakeBuild(build_ext): |
70 | 48 | def run(self): |
| 49 | + build_dir = os.path.abspath(self.build_temp) |
| 50 | + os.makedirs(build_dir, exist_ok=True) |
| 51 | + |
71 | 52 | for ext in self.extensions: |
72 | 53 | self.build_cmake(ext) |
73 | 54 |
|
74 | | - self._copy_so_files_to_build_lib() |
75 | | - |
76 | 55 | def build_cmake(self, ext: CMakeExtension): |
77 | | - build_dir = self.build_temp |
78 | | - os.makedirs(build_dir, exist_ok=True) |
| 56 | + build_dir = os.path.abspath(self.build_temp) |
| 57 | + install_dir = os.path.abspath(self.build_lib) |
79 | 58 |
|
80 | 59 | cmake_args = [ |
81 | | - "cmake", |
82 | 60 | "-DCMAKE_BUILD_TYPE=Release", |
83 | 61 | f"-DPYTHON_EXECUTABLE={sys.executable}", |
| 62 | + f"-DCMAKE_INSTALL_PREFIX={install_dir}", |
84 | 63 | ] |
85 | 64 |
|
86 | | - torch_cmake_prefix = torch.utils.cmake_prefix_path |
87 | | - pybind11_cmake_dir = pybind11.get_cmake_dir() |
88 | | - |
89 | | - cmake_prefix_paths = [torch_cmake_prefix, pybind11_cmake_dir] |
90 | | - cmake_args.append(f"-DCMAKE_PREFIX_PATH={';'.join(cmake_prefix_paths)}") |
91 | | - |
92 | | - torch_includes = torch.utils.cpp_extension.include_paths() |
93 | | - python_include = sysconfig.get_path("include") |
94 | | - pybind11_include = pybind11.get_include() |
95 | | - |
96 | | - all_includes = torch_includes + [python_include, pybind11_include] |
97 | | - cmake_include_string = ";".join(all_includes) |
98 | | - cmake_args.append(f"-DEXTERNAL_INCLUDE_DIRS={cmake_include_string}") |
99 | | - |
100 | | - if _is_cuda(): |
101 | | - cmake_args.append("-DRUNTIME_ENVIRONMENT=cuda") |
102 | | - else: |
103 | | - cmake_args.append("-DRUNTIME_ENVIRONMENT=ascend") |
104 | | - |
105 | 65 | if _enable_sparse(): |
106 | | - cmake_args.append("-DBUILD_UCM_SPARSE=ON") |
107 | | - |
108 | | - cmake_args.append(ext.sourcedir) |
109 | | - |
110 | | - print(f"[INFO] Building {ext.name} module with CMake") |
111 | | - print(f"[INFO] Source directory: {ext.sourcedir}") |
112 | | - print(f"[INFO] Build directory: {build_dir}") |
113 | | - print(f"[INFO] CMake command: {' '.join(cmake_args)}") |
| 66 | + cmake_args += ["-DBUILD_UCM_SPARSE=ON"] |
| 67 | + |
| 68 | + match PLATFORM: |
| 69 | + case "cuda": |
| 70 | + cmake_args += ["-DRUNTIME_ENVIRONMENT=cuda"] |
| 71 | + case "ascend": |
| 72 | + cmake_args += ["-DRUNTIME_ENVIRONMENT=ascend"] |
| 73 | + case "musa": |
| 74 | + cmake_args += ["-DRUNTIME_ENVIRONMENT=musa"] |
| 75 | + case "maca": |
| 76 | + cmake_args += ["-DRUNTIME_ENVIRONMENT=maca"] |
| 77 | + cmake_args += ["-DBUILD_UCM_SPARSE=OFF"] |
| 78 | + case _: |
| 79 | + # default ascend |
| 80 | + cmake_args += ["-DRUNTIME_ENVIRONMENT=ascend"] |
114 | 81 |
|
115 | | - subprocess.check_call(cmake_args, cwd=build_dir) |
| 82 | + subprocess.check_call( |
| 83 | + ["cmake", *cmake_args, ext.cmake_file_path], cwd=build_dir |
| 84 | + ) |
116 | 85 | subprocess.check_call( |
117 | 86 | ["cmake", "--build", ".", "--config", "Release", "--", "-j8"], |
118 | 87 | cwd=build_dir, |
119 | 88 | ) |
120 | 89 |
|
121 | | - def _copy_so_files_to_build_lib(self): |
122 | | - """Copy .so files from source directories to build_lib for installation.""" |
123 | | - if not hasattr(self, "build_lib") or not self.build_lib: |
124 | | - return |
125 | | - |
126 | | - packages = _get_packages() |
127 | | - copied_count = 0 |
128 | | - |
129 | | - for package in packages: |
130 | | - # Source directory where CMake outputs .so files |
131 | | - source_package_dir = os.path.join(ROOT_DIR, package.replace(".", os.sep)) |
132 | | - |
133 | | - # Destination in build_lib |
134 | | - build_package_dir = os.path.join( |
135 | | - self.build_lib, package.replace(".", os.sep) |
136 | | - ) |
137 | | - |
138 | | - # Find all .so files in the source package directory |
139 | | - so_files = glob(os.path.join(source_package_dir, "*.so")) |
140 | | - |
141 | | - if so_files: |
142 | | - # Ensure destination directory exists |
143 | | - os.makedirs(build_package_dir, exist_ok=True) |
144 | | - |
145 | | - # Copy each .so file |
146 | | - for so_file in so_files: |
147 | | - dest_file = os.path.join( |
148 | | - build_package_dir, os.path.basename(so_file) |
149 | | - ) |
150 | | - shutil.copy2(so_file, dest_file) |
151 | | - copied_count += 1 |
152 | | - print( |
153 | | - f"[INFO] Copied {os.path.basename(so_file)} to {build_package_dir}" |
154 | | - ) |
155 | | - |
156 | | - if copied_count > 0: |
157 | | - print(f"[INFO] Successfully copied {copied_count} .so file(s) to build_lib") |
158 | | - else: |
159 | | - print( |
160 | | - "[WARNING] No .so files found to copy. Extensions may not have been built." |
161 | | - ) |
162 | | - |
163 | | - |
164 | | -def _get_packages(): |
165 | | - """Discover Python packages, optionally filtering out sparse-related ones.""" |
166 | | - sparse_enabled = _enable_sparse() |
167 | | - exclude_patterns = [] |
168 | | - if not sparse_enabled: |
169 | | - exclude_patterns.append("ucm.sparse*") |
170 | | - |
171 | | - packages = find_packages(exclude=exclude_patterns) |
172 | | - return packages |
173 | | - |
174 | | - |
175 | | -ext_modules = [] |
176 | | -ext_modules.append(CMakeExtension(name="ucm", sourcedir=ROOT_DIR)) |
| 90 | + subprocess.check_call( |
| 91 | + ["cmake", "--install", ".", "--config", "Release", "--component", "ucm"], |
| 92 | + cwd=build_dir, |
| 93 | + ) |
177 | 94 |
|
178 | | -packages = _get_packages() |
179 | 95 |
|
180 | 96 | setup( |
181 | 97 | name="uc-manager", |
182 | 98 | version="0.1.2", |
183 | 99 | description="Unified Cache Management", |
184 | 100 | author="Unified Cache Team", |
185 | | - packages=packages, |
| 101 | + packages=find_packages(), |
186 | 102 | python_requires=">=3.10", |
187 | | - ext_modules=ext_modules, |
| 103 | + ext_modules=[CMakeExtension(name="ucm", source_dir=ROOT_DIR)], |
188 | 104 | cmdclass={"build_ext": CMakeBuild}, |
189 | 105 | zip_safe=False, |
190 | 106 | ) |
0 commit comments