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
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
|
import contextlib
import unittest
from collections import namedtuple
from mpitestutil import import_MPI
class TestCase(unittest.TestCase):
#
def assertAlmostEqual(self, first, second):
num = complex(second) - complex(first)
den = max(abs(complex(second)), abs(complex(first))) or 1.0
if abs(num / den) > 1e-2:
raise self.failureException(f"{first!r} != {second!r}")
@contextlib.contextmanager
def catchNotImplementedError(self, version=None, subversion=0):
try:
yield
except NotImplementedError:
if version is None:
return
MPI = import_MPI()
mpi_version = MPI.Get_version()
self.assertLess(mpi_version, (version, subversion))
_Version = namedtuple("_Version", ["major", "minor", "patch"]) # noqa: PYI024
def _parse_version(version):
version = (*tuple(map(int, version.split("."))), 0, 0, 0)
return _Version(*version[:3])
class _VersionPredicate:
#
def __init__(self, versionPredicateStr):
import re
re_name = re.compile(r"(?i)^([a-z_]\w*(?:\.[a-z_]\w*)*)(.*)$")
re_pred = re.compile(r"^(<=|>=|<|>|!=|==)(.*)$")
def split(item):
m = re_pred.match(item)
op, version = m.groups()
version = _parse_version(version)
return op, version
vpstr = versionPredicateStr.replace(" ", "")
m = re_name.match(vpstr)
name, plist = m.groups()
if plist:
assert plist[0] == "(" and plist[-1] == ")"
plist = plist[1:-1]
pred = [split(p) for p in plist.split(",") if p]
self.name = name
self.pred = pred
def __str__(self):
if self.pred:
items = [f"{op}{'.'.join(map(str, ver))}" for op, ver in self.pred]
return f"{self.name}({','.join(items)})"
else:
return self.name
def satisfied_by(self, version):
from operator import eq, ge, gt, le, lt, ne
opmap = {"<": lt, "<=": le, ">": gt, ">=": ge, "==": eq, "!=": ne}
version = _parse_version(version)
for op, ver in self.pred:
if not opmap[op](version, ver):
return False
return True
def mpi_predicate(predicate):
MPI = import_MPI()
def key(s):
s = s.replace(" ", "")
s = s.replace("/", "")
s = s.replace("-", "")
s = s.replace("Intel", "I")
s = s.replace("Microsoft", "MS")
return s.lower()
vp = _VersionPredicate(key(predicate))
if vp.name == "mpi":
name, version = "mpi", MPI.Get_version()
version = (*version, 0)
else:
name, version = MPI.get_vendor()
if vp.name == key(name):
x, y, z = version
if vp.satisfied_by(f"{x}.{y}.{z}"):
return vp
return None
def is_mpi(predicate):
return mpi_predicate(predicate)
def is_mpi_gpu(predicate, array):
if array.backend in ("cupy", "numba", "dlpack-cupy"):
if mpi_predicate(predicate):
return True
return False
SkipTest = unittest.SkipTest
skip = unittest.skip
skipIf = unittest.skipIf
skipUnless = unittest.skipUnless
def skipMPI(predicate, *conditions, reason=""):
version = mpi_predicate(predicate)
if version:
if not conditions or any(conditions):
return unittest.skip(str(version))
return unittest.skipIf(False, reason)
def disable(what, reason):
return unittest.skip(reason)(what)
def main(*args, **kwargs):
from main import main
try:
main(*args, **kwargs)
except SystemExit:
pass
|
