forked from arogozhnikov/einops
-
Notifications
You must be signed in to change notification settings - Fork 0
/
test.py
100 lines (82 loc) · 2.87 KB
/
test.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
"""
Usage: python test.py <frameworks>
1. Installs part of dependencies (make sure `which pip` points to correct location)
2. Installs current version of einops in editable mode
3. Runs the tests
"""
import os
import shutil
import sys
from subprocess import Popen, PIPE
from pathlib import Path
__author__ = "Alex Rogozhnikov"
def run(cmd, **env):
# keeps printing output when testing
cmd = cmd.split(" ") if isinstance(cmd, str) else cmd
p = Popen(cmd, cwd=str(Path(__file__).parent), env={**os.environ, **env})
p.communicate()
return p.returncode
# check we have nvidia-smi
have_cuda = False
if shutil.which("nvidia-smi") is not None:
output, _ = Popen("nvidia-smi".split(" "), stdout=PIPE).communicate()
if b"failed because" not in output:
have_cuda = True
def main():
_executable, *frameworks = sys.argv
framework_name2installation = {
"numpy": ["numpy"],
"torch": ["torch --index-url https://download.pytorch.org/whl/cpu"],
"jax": ["jax[cpu]", "jaxlib", "flax"],
"tensorflow": ["tensorflow"],
"chainer": ["chainer"],
"cupy": ["cupy"],
"paddle": ["paddlepaddle==0.0.0 -f https://www.paddlepaddle.org.cn/whl/linux/cpu-mkl/develop.html"],
"oneflow": ["oneflow==0.9.0"],
}
usage = f"""
Usage: python test.py <frameworks>
Example: python test.py numpy pytorch
Available frameworks: {list(framework_name2installation)}
"""
if len(frameworks) == 0:
print(usage)
return
else:
synonyms = {
"tf": "tensorflow",
"pytorch": "torch",
"paddlepaddle": "paddle",
}
frameworks = [synonyms.get(f, f) for f in frameworks]
wrong_frameworks = [f for f in frameworks if f not in framework_name2installation]
if wrong_frameworks:
print(usage)
raise RuntimeError(f"Unrecognized frameworks: {wrong_frameworks}")
other_dependencies = [
"nbformat",
"nbconvert",
"jupyter",
"parameterized",
"pillow",
"pytest",
]
for framework in frameworks:
print(f"Installing {framework}")
pip_instructions = framework_name2installation[framework]
assert 0 == run("pip install {} --progress-bar off".format(" ".join(pip_instructions)))
print("Install testing infra")
assert 0 == run("pip install {} --progress-bar off".format(" ".join(other_dependencies)))
# install einops
assert 0 == run("pip install -e .")
# we need to inform testing script which frameworks to use
# this is done by setting a flag EINOPS_TEST_BACKENDS
from tests import unparse_backends
flag_name, flag_value = unparse_backends(backend_names=frameworks)
return_code = run(
"python -m pytest tests",
**{flag_name: flag_value},
)
assert return_code == 0
if __name__ == "__main__":
main()