forked from DeepLink-org/dlinfer
-
Notifications
You must be signed in to change notification settings - Fork 0
/
setup.py
100 lines (81 loc) · 2.76 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
100
import os
from pathlib import Path
from setuptools import find_packages
import yaml
from skbuild import setup
VERSION = "0.1.0"
vendor_dispatch_key_map = {
"ascend": "PrivateUse1",
}
def gen_vendor_yaml(device):
config = dict()
config["vendor"] = device
assert device in vendor_dispatch_key_map
config["dispatch_key"] = vendor_dispatch_key_map[device]
file_path = Path(__file__).parent / "dlinfer" / "vendor" / "vendor.yaml"
with open(str(file_path), "w") as f:
yaml.safe_dump(config, f)
return str(file_path.name)
def get_device():
device = os.getenv("DEVICE", "").lower()
assert device in vendor_dispatch_key_map
return device
def get_cmake_args():
cmake_args = list()
cmake_device = get_device()
cmake_args.append("-DCMAKE_BUILD_TYPE=Release")
cmake_args.append(f"-DDEVICE={cmake_device}")
return cmake_args
def get_package_data():
cmake_device = get_device()
yaml_file_name = gen_vendor_yaml(cmake_device)
assert cmake_device, "DEVICE shouldn't be empty!"
return {
f"dlinfer.vendor": [
yaml_file_name,
]
}
def get_readme():
with open(str(Path(__file__).parent / "README.md"), "r", encoding="utf-8") as f:
content = f.read()
return content
def get_requirements(file_name):
requirements = []
device_req_root = Path(__file__).parent / "requirements" / get_device()
with open(str(device_req_root / file_name), "r") as f:
for line in f:
line = line.strip()
if not line or line.startswith("#"):
continue
if line.startswith("-r"):
other_file = line.split()[1]
requirements.extend(get_requirements(other_file))
else:
requirements.append(line)
return requirements
def main():
setup(
name=f"dlinfer-{get_device()}",
version=VERSION,
description="DeepLink Inference Extension",
long_description=get_readme(),
long_description_content_type="text/markdown",
url="https://github.com/DeepLink-org/dlinfer",
packages=find_packages(),
package_data=get_package_data(),
exclude_package_data={"": ["tests/*"]},
cmake_args=get_cmake_args(),
cmake_install_target="install",
classifiers=[
"Programming Language :: Python :: 3.8",
"Programming Language :: Python :: 3.9",
"Programming Language :: Python :: 3.10",
"Intended Audience :: Developers",
"Operating System :: POSIX :: Linux",
],
python_requires=">=3.8, <3.11",
setup_requires=get_requirements("build.txt"),
install_requires=get_requirements("runtime.txt"),
)
if __name__ == "__main__":
main()