forked from mpi4jax/mpi4jax
-
Notifications
You must be signed in to change notification settings - Fork 0
/
setup.py
215 lines (170 loc) · 5.61 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
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
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
import os
import sys
from setuptools import setup, find_packages
from setuptools.extension import Extension
# ensure vendored versioneer is on path
sys.path.append(os.path.dirname(__file__))
import versioneer # noqa: E402
try:
from Cython.Build import cythonize
except ImportError:
HAS_CYTHON = False
else:
HAS_CYTHON = True
try:
import mpi4py
except ImportError:
HAS_MPI4PY = False
else:
HAS_MPI4PY = True
CYTHON_SUBMODULE_NAME = "mpi4jax._src.xla_bridge"
CYTHON_SUBMODULE_PATH = "mpi4jax/_src/xla_bridge"
#######
# Utils
def search_on_path(filenames):
path_dirs = os.getenv("PATH", "").split(os.pathsep)
for p in path_dirs:
for filename in filenames:
full_path = os.path.join(p, filename)
if os.path.exists(full_path):
return os.path.abspath(full_path)
return None
def print_warning(*lines):
print("**************************************************")
for line in lines:
print("*** WARNING: %s" % line)
print("**************************************************")
# /end Utils
############
def mpi_info(cmd):
config = mpi4py.get_config()
cmd_compile = " ".join([config["mpicc"], "-show"])
out_stream = os.popen(cmd_compile)
flags = out_stream.read().strip()
flags = flags.replace(",", " ").split()
if cmd == "compile":
startwith = "-I"
elif cmd == "libdirs":
startwith = "-L"
elif cmd == "libs":
startwith = "-l"
out = []
for flag in flags:
if flag.startswith(startwith):
out.append(flag[2:])
return out
################
# Cuda detection
# Taken from CUPY (MIT License)
def get_cuda_path():
nvcc_path = search_on_path(("nvcc", "nvcc.exe"))
cuda_path_default = None
if nvcc_path is not None:
cuda_path_default = os.path.normpath(
os.path.join(os.path.dirname(nvcc_path), "..")
)
cuda_path = os.getenv("CUDA_PATH", "") # Nvidia default on Windows
if len(cuda_path) == 0:
cuda_path = os.getenv("CUDA_ROOT", "") # Nvidia default on Windows
if len(cuda_path) > 0 and cuda_path != cuda_path_default:
print_warning(
"nvcc path != CUDA_PATH",
"nvcc path: %s" % cuda_path_default,
"CUDA_PATH: %s" % cuda_path,
)
if os.path.exists(cuda_path):
_cuda_path = cuda_path
elif cuda_path_default is not None:
_cuda_path = cuda_path_default
elif os.path.exists("/usr/local/cuda"):
_cuda_path = "/usr/local/cuda"
else:
_cuda_path = None
return _cuda_path
def cuda_info(cmd):
cuda_path = get_cuda_path()
if not cuda_path:
return []
if cmd == "compile":
incdir = os.path.join(cuda_path, "include")
if os.path.isdir(incdir):
return [incdir]
if cmd == "libdirs":
for libdir in ("lib64", "lib"):
full_dir = os.path.join(cuda_path, libdir)
if os.path.isdir(full_dir):
return [full_dir]
if cmd == "libs":
return ["cudart"]
return []
# /end Cuda detection
#####################
def get_extensions():
if not HAS_MPI4PY:
print_warning("mpi4py could not be imported", "(extensions will not be built)")
return []
if not HAS_CYTHON:
print_warning("Cython could not be imported", "(extensions will not be built)")
return []
extensions = [
Extension(
name=f"{CYTHON_SUBMODULE_NAME}.{mod}",
sources=[f"{CYTHON_SUBMODULE_PATH}/{mod}.pyx"],
include_dirs=mpi_info("compile"),
library_dirs=mpi_info("libdirs"),
libraries=mpi_info("libs"),
)
for mod in ("mpi_xla_bridge", "mpi_xla_bridge_cpu")
]
if cuda_info("compile"):
extensions.append(
Extension(
name=f"{CYTHON_SUBMODULE_NAME}.mpi_xla_bridge_gpu",
sources=[f"{CYTHON_SUBMODULE_PATH}/mpi_xla_bridge_gpu.pyx"],
include_dirs=mpi_info("compile") + cuda_info("compile"),
library_dirs=mpi_info("libdirs") + cuda_info("libdirs"),
libraries=mpi_info("libs") + cuda_info("libs"),
)
)
else:
print_warning("CUDA path not found", "(GPU extensions will not be built)")
if HAS_CYTHON:
extensions = cythonize(
extensions,
language_level=3,
)
return extensions
here = os.path.abspath(os.path.dirname(__file__))
# get the long description from the README file
with open(os.path.join(here, "README.rst"), encoding="utf-8") as f:
long_description = f.read()
setup(
name="mpi4jax",
author="Filippo Vicentini",
author_email="filippovicentini@gmail.com",
description=(
"Zero-copy MPI communication of JAX arrays, "
"for turbo-charged HPC applications in Python ⚡"
),
long_description=long_description,
url="https://github.com/mpi4jax/mpi4jax",
license="MIT",
version=versioneer.get_version(),
cmdclass=versioneer.get_cmdclass(),
classifiers=[
"Programming Language :: Python :: 3",
"Programming Language :: Cython",
"Development Status :: 3 - Alpha",
"Environment :: GPU :: NVIDIA CUDA",
"Intended Audience :: Science/Research",
"License :: OSI Approved :: MIT License",
"Operating System :: MacOS :: MacOS X",
"Operating System :: POSIX :: Linux",
"Operating System :: Unix",
"Topic :: Scientific/Engineering",
],
packages=find_packages(),
ext_modules=get_extensions(),
python_requires=">=3.6",
install_requires=["jax", "mpi4py>=3.0.1", "numpy"],
)