File size: 1,327 Bytes
db26c81 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 |
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"],
)
|