| import os |
| import re |
| import sys |
| import time |
| import subprocess |
|
|
| from distutils import sysconfig |
| from setuptools import setup, Extension, find_packages |
| from Cython.Build import build_ext, cythonize |
|
|
|
|
| class CMakeExtension(Extension): |
| def __init__(self, name, sourcedir=''): |
| Extension.__init__(self, name, sources=[]) |
| self.sourcedir = os.path.abspath(sourcedir) |
|
|
|
|
| class CMakeBuild(build_ext): |
| def build_extension(self, ext): |
| if isinstance(ext, CMakeExtension): |
| extdir = os.path.abspath(os.path.dirname(self.get_ext_fullpath(ext.name))) |
| if not extdir.endswith(os.path.sep): |
| extdir += os.path.sep |
|
|
| if not os.path.exists(self.build_temp): |
| os.makedirs(self.build_temp) |
|
|
| subprocess.check_call(['cmake', ext.sourcedir, '-DCMAKE_LIBRARY_OUTPUT_DIRECTORY=' + extdir], |
| cwd=self.build_temp) |
| subprocess.check_call(['cmake', '--build', '.', '--', 'VERBOSE=1', '-j8'], cwd=self.build_temp) |
| else: |
| super().build_extension(ext) |
|
|
|
|
| ext_modules = [CMakeExtension('greedrl_c')] |
|
|
| setup( |
| name='greedrl', |
| version='1.0.0', |
| packages=find_packages(), |
| ext_modules=ext_modules, |
| cmdclass={'build_ext': CMakeBuild}, |
| install_requires=["torch==1.12.1+cu113"], |
| ) |
|
|
