-
Cyril Moineau authoredCyril Moineau authored
Code owners
Assign users and groups as approvers for specific file changes. Learn more.
setup.py 4.11 KiB
#!/usr/bin/env python3
import sys
import os
import shutil
import pathlib
import multiprocessing
from math import ceil
import toml
from setuptools import setup, Extension
from setuptools import find_packages
from setuptools.command.build_ext import build_ext
def get_project_name() -> str:
with open(pathlib.Path().absolute() / "pyproject.toml", "r") as file:
project_toml = toml.load(file)
return project_toml["project"]["name"]
def get_project_version() -> str:
aidge_root = pathlib.Path().absolute()
version = open(aidge_root / "version.txt", "r").read().strip()
return version
class CMakeExtension(Extension):
def __init__(self, name):
super().__init__(name, sources=[])
class CMakeBuild(build_ext):
def run(self):
# This lists the number of processors available on the machine
# The compilation will use half of them
max_jobs = str(ceil(multiprocessing.cpu_count() / 2))
max_jobs = os.environ.get("AIDGE_NB_PROC", max_jobs)
cwd = pathlib.Path().absolute()
build_temp = cwd / "build"
if not build_temp.exists():
build_temp.mkdir(parents=True, exist_ok=True)
build_lib = pathlib.Path(self.build_lib)
if not build_lib.exists():
build_lib.mkdir(parents=True, exist_ok=True)
os.chdir(str(build_temp))
install_path = (
os.path.join(sys.prefix, "lib", "libAidge")
if "AIDGE_INSTALL" not in os.environ
else os.environ["AIDGE_INSTALL"]
)
# Read environment variables for CMake options
c_compiler = os.environ.get("AIDGE_C_COMPILER", "gcc")
cxx_compiler = os.environ.get("AIDGE_CXX_COMPILER", "g++")
build_type = os.environ.get("AIDGE_BUILD_TYPE", "Release")
asan = os.environ.get("AIDGE_ASAN", "OFF")
with_cuda = os.environ.get("AIDGE_WITH_CUDA", "OFF")
# using ninja as default build system to build faster and with the same compiler as on windows
build_gen = os.environ.get("AIDGE_BUILD_GEN", "")
build_gen_opts = (
["-G", build_gen]
if build_gen
else []
)
test_onoff = os.environ.get("AIDGE_BUILD_TEST", "OFF")
self.spawn(
[
"cmake",
*build_gen,
str(cwd),
f"-DTEST={test_onoff}",
f"-DCMAKE_INSTALL_PREFIX:PATH={install_path}",
f"-DCMAKE_BUILD_TYPE={build_type}",
f"-DCMAKE_C_COMPILER={c_compiler}",
f"-DCMAKE_CXX_COMPILER={cxx_compiler}",
f"-DENABLE_ASAN={asan}",
f"-DCUDA={with_cuda}",
"-DPYBIND=ON",
"-DPYBIND11_FINDPYTHON=ON",
"-DCMAKE_EXPORT_COMPILE_COMMANDS=ON",
"-DCOVERAGE=OFF",
"-DCMAKE_CUDA_ARCHITECTURES=native",
]
)
if not self.dry_run:
self.spawn(
["cmake", "--build", ".", "--config", build_type, "-j", max_jobs]
)
self.spawn(["cmake", "--install", ".", "--config", build_type])
os.chdir(str(cwd))
aidge_package = build_lib / (get_project_name())
# Get "aidge core" package
# ext_lib = build_temp
print(build_temp.absolute())
# Copy all shared object files from build_temp/lib to aidge_package
for root, _, files in os.walk(build_temp.absolute()):
for file in files:
if (file.endswith(".so") or file.endswith(".pyd")) and (
root != str(aidge_package.absolute())
):
currentFile = os.path.join(root, file)
shutil.copy(currentFile, str(aidge_package.absolute()))
# Copy version.txt in aidge_package
os.chdir(os.path.dirname(__file__))
shutil.copy("version.txt", str(aidge_package.absolute()))
if __name__ == "__main__":
setup(
include_package_data=True,
ext_modules=[CMakeExtension(get_project_name())],
cmdclass={
"build_ext": CMakeBuild,
},
zip_safe=False,
)