forked from google/jax
-
Notifications
You must be signed in to change notification settings - Fork 2
/
build_wheel_docker_entrypoint.sh
executable file
·63 lines (53 loc) · 1.43 KB
/
build_wheel_docker_entrypoint.sh
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
#!/bin/bash
set -xev
if [ ! -d "/dist" ]
then
echo "/dist must be mounted to produce output"
exit 1
fi
export CC=/dt7/usr/bin/gcc
export GCC_HOST_COMPILER_PATH=/dt7/usr/bin/gcc
export PYENV_ROOT="/pyenv"
export PATH="$PYENV_ROOT/bin:$PATH"
eval "$(pyenv init -)"
PY_VERSION="$1"
echo "Python version $PY_VERSION"
git clone https://github.com/google/jax /build/jax
cd /build/jax/build
mkdir /build/tmp
mkdir /build/root
export TMPDIR=/build/tmp
usage() {
echo "usage: ${0##*/} [py2|py3] [cuda-included|cuda|nocuda]"
exit 1
}
if [[ $# -lt 2 ]]
then
usage
fi
# Builds and activates a specific Python version.
pyenv local "$PY_VERSION"
# Workaround for https://github.com/bazelbuild/bazel/issues/9254
export BAZEL_LINKLIBS="-lstdc++"
export JAX_CUDA_VERSION=$3
case $2 in
cuda-included)
python build.py --enable_cuda --bazel_startup_options="--output_user_root=/build/root"
python include_cuda.py
;;
cuda)
python build.py --enable_cuda --bazel_startup_options="--output_user_root=/build/root"
;;
nocuda)
python build.py --enable_tpu --bazel_startup_options="--output_user_root=/build/root"
;;
*)
usage
esac
if ! python -m auditwheel show dist/jaxlib-*.whl | egrep 'platform tag: "(manylinux2010_x86_64|manylinux_2_12_x86_64)"' > /dev/null; then
# Print output for debugging
python -m auditwheel show dist/jaxlib-*.whl
echo "jaxlib wheel is not manylinux2010 compliant"
exit 1
fi
cp -r dist/* /dist