-
Notifications
You must be signed in to change notification settings - Fork 59
/
setup.py
99 lines (76 loc) · 3.07 KB
/
setup.py
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
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
"""Build the package with Rust binary."""
import os
import shutil
import subprocess
from io import open
from setuptools import Extension, find_packages, setup
from setuptools.command.build_ext import build_ext as _build_ext
here = os.path.abspath(os.path.dirname(__file__))
PACKAGE_NAME = "mosec"
with open(os.path.join(here, "README.md"), encoding="utf-8") as f:
readme = f.read()
with open(os.path.join(here, "requirements/dev.txt"), encoding="utf-8") as f:
dev_requirements = f.read().splitlines()
with open(os.path.join(here, "requirements/doc.txt"), encoding="utf-8") as f:
doc_requirements = f.read().splitlines()
with open(os.path.join(here, "requirements/mixin.txt"), encoding="utf-8") as f:
mixin_requirements = f.read().splitlines()
class RustExtension(Extension):
# pylint: disable=too-few-public-methods
"""Custom Extension class for rust."""
ext_modules = []
if os.getenv("PRODUCTION_MODE"):
ext_modules.append(RustExtension(name="mosec.bin", sources=["src/*"]))
class RustBuildExt(_build_ext):
"""Custom build extension class for rust."""
def run(self):
"""Run the build command."""
for ext in self.extensions:
self.build_extension(ext)
def build_extension(self, ext: Extension):
"""Build the extension."""
if not isinstance(ext, RustExtension):
return super().build_extension(ext)
libpath = ext.name.replace(".", os.sep)
build_libpath = os.path.join(self.build_lib, libpath)
rust_target = os.getenv("RUST_TARGET")
build_cmd = ["cargo", "build", "--release"]
if rust_target is not None:
build_cmd += ["--target", rust_target]
print(f"running rust cargo package build: {build_cmd}")
errno = subprocess.call(build_cmd)
assert errno == 0, "Error occurred while building rust binary"
# package the binary
if rust_target is not None:
target_dir = os.path.join("target", rust_target, "release", PACKAGE_NAME)
else:
target_dir = os.path.join("target", "release", PACKAGE_NAME)
os.makedirs(build_libpath, exist_ok=True)
shutil.copy(target_dir, build_libpath)
if self.inplace:
os.makedirs(os.path.dirname(libpath), exist_ok=True)
shutil.copy(build_libpath, libpath)
return None
setup(
name=PACKAGE_NAME,
author="Keming Yang",
author_email="kemingy94@gmail.com",
description="Model Serving made Efficient in the Cloud.",
long_description=readme,
long_description_content_type="text/markdown",
url="https://github.com/mosecorg/mosec",
license="Apache License 2.0",
use_scm_version=True,
packages=find_packages(exclude=["examples*", "tests*"]),
include_package_data=True,
python_requires=">=3.8",
setup_requires=["setuptools_scm>=7.0"],
extras_require={
"dev": dev_requirements,
"mixin": mixin_requirements,
"doc": doc_requirements,
},
zip_safe=False,
ext_modules=ext_modules, # type: ignore
cmdclass={"build_ext": RustBuildExt}, # type: ignore
)