-
Notifications
You must be signed in to change notification settings - Fork 0
/
get_tensorflow_container.sh
161 lines (145 loc) · 7.49 KB
/
get_tensorflow_container.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
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
set -euox pipefail
if [ "x$BENCH_DIR" = "x" ]; then
echo "BENCH_DIR is not set. Please set it to the `resnet50` directory of benchmark" >&2
exit 1
fi
export ROOT_DIR=$BENCH_DIR/..
NVIDIA_X86_ACCELERATORS=(A100 H100 WAIH100)
NVIDIA_ARM_ACCELERATORS=(JEDI GH200)
TENSORFLOW_CONTAINER_FILE_NVIDIA_X86=$ROOT_DIR/containers/ngc2301_tf115_cuda1201_nccl2165_py38.sif
TENSORFLOW_CONTAINER_FILE_NVIDIA_ARM=$ROOT_DIR/containers/ngc2301_tf115_cuda1201_nccl2165_py38_arm.sif
TENSORFLOW_CONTAINER_FILE_AMD=$ROOT_DIR/containers/amd_tf27_rocm50_py39-dev.sif
TENSORFLOW_CONTAINER_FILE_IPU=$ROOT_DIR/containers/ipu_tf263_poplar310_py38.sif
TENSORFLOW_CONTAINER_FILE_DONE=$BENCH_DIR/tensorflow_container_done
TENSORFLOW_PACKAGES_AMD=$BENCH_DIR/amd_tensorflow_packages
TENSORFLOW_PACKAGES_IPU=$BENCH_DIR/ipu_tensorflow_packages
TENSORFLOW_PACKAGES_NVIDIA_X86=$BENCH_DIR/nvidia_x86_tensorflow_packages
TENSORFLOW_PACKAGES_NVIDIA_ARM=$BENCH_DIR/nvidia_arm_tensorflow_packages
TENSORFLOW_PACKAGES_FILE_AMD=$BENCH_DIR/amd_tensorflow_packages_installed
TENSORFLOW_PACKAGES_FILE_IPU=$BENCH_DIR/ipu_tensorflow_packages_installed
TENSORFLOW_PACKAGES_FILE_NVIDIA_X86=$BENCH_DIR/nvidia_x86_tensorflow_packages_installed
TENSORFLOW_PACKAGES_FILE_NVIDIA_ARM=$BENCH_DIR/nvidia_arm_tensorflow_packages_installed
if ! [ -d "$ROOT_DIR/containers" ]; then
mkdir -p "$ROOT_DIR/containers"
fi
if ! [ -d "$ROOT_DIR/containers/tmp_dir" ]; then
mkdir "$ROOT_DIR/containers/tmp_dir"
fi
if [ -f $TENSORFLOW_CONTAINER_FILE_DONE]; then
echo "Required containers exists at $ROOT_DIR/containers/" >&2
echo "To rebuild containers delete $TENSORFLOW_CONTAINER_FILE_DONE" >&2
else
export APPTAINER_CACHEDIR=$(mktemp -d -p $ROOT_DIR/containers/tmp_dir)
export APPTAINER_TMPDIR=$(mktemp -d -p $ROOT_DIR/containers/tmp_dir)
fi
##### Installing Containers #####
if [ "$ACCELERATOR" = "GH200" ]; then
export CUDA_VISIBLE_DEVICES=0
if [ -f $TENSORFLOW_CONTAINER_FILE_NVIDIA_ARM ]; then
echo "$TENSORFLOW_CONTAINER_FILE_NVIDIA_ARM" exists >&2
else
# https://docs.nvidia.com/deeplearning/frameworks/tensorflow-release-notes/rel-23-01.html
apptainer pull $TENSORFLOW_CONTAINER_FILE_NVIDIA_ARM docker://nvcr.io/nvidia/tensorflow:23.01-tf1-py3 >&2
echo "Done pulling $TENSORFLOW_CONTAINER_FILE_NVIDIA_ARM" >&2
fi
elif [ "$ACCELERATOR" = "H100" ]; then
export CUDA_VISIBLE_DEVICES=0
if [ -f $TENSORFLOW_CONTAINER_FILE_NVIDIA_X86 ]; then
echo "$TENSORFLOW_CONTAINER_FILE_NVIDIA_X86" exists >&2
else
# https://docs.nvidia.com/deeplearning/frameworks/tensorflow-release-notes/rel-23-01.html
apptainer pull $TENSORFLOW_CONTAINER_FILE_NVIDIA_X86 docker://nvcr.io/nvidia/tensorflow:23.01-tf1-py3 >&2
echo "Done pulling $TENSORFLOW_CONTAINER_FILE_NVIDIA_X86" >&2
fi
elif [ "$ACCELERATOR" = "MI250" ]; then
export ROCM_VISIBLE_DEVICES=0
if [ -f $TENSORFLOW_CONTAINER_FILE_AMD ]; then
echo "$TENSORFLOW_CONTAINER_FILE_AMD" exists >&2
else
# https://hub.docker.com/layers/rocm/tensorflow/rocm5.0-tf2.7-dev/images/sha256-664fbd3e38234f5b4419aa54b2b81664495ed0a9715465678f2bc14ea4b7ae16?context=explore
apptainer pull $TENSORFLOW_CONTAINER_FILE_AMD docker://rocm/tensorflow:rocm5.0-tf2.7-dev >&2
echo "Done pulling $TENSORFLOW_CONTAINER_FILE_AMD" >&2
fi
else
if [ -f $TENSORFLOW_CONTAINER_FILE_IPU ]; then
echo "$TENSORFLOW_CONTAINER_FILE_IPU" exists >&2
else
# https://hub.docker.com/layers/john856/caraml/tensorflow_poplar310_tf263_mpi4py/images/sha256-57cb664cb1e3493657c576b07a0d274363e1097e62820d2f7e03db5e68fe1f0e?context=repo
apptainer pull $TENSORFLOW_CONTAINER_FILE_IPU docker://john856/caraml:tensorflow_poplar310_tf263_mpi4py >&2
echo "Done pulling $TENSORFLOW_CONTAINER_FILE_IPU" >&2
fi
fi
if [ -f $TENSORFLOW_CONTAINER_FILE_NVIDIA_X86 ] && [ -f $TENSORFLOW_CONTAINER_FILE_NVIDIA_ARM ] && [ -f $TENSORFLOW_CONTAINER_FILE_AMD ] && [ -f $TENSORFLOW_CONTAINER_FILE_IPU ]; then
touch $TENSORFLOW_CONTAINER_FILE_DONE
echo "Done pulling TensorFlow Containers!" >&2
fi
rm -rf $APPTAINER_CACHEDIR
rm -rf $APPTAINER_TMPDIR
##### Installing Requirements #####
if ! [ -f $TENSORFLOW_PACKAGES_FILE_AMD ] && [ "$ACCELERATOR" = "MI250" ]; then
mkdir -p $TENSORFLOW_PACKAGES_AMD
apptainer exec $TENSORFLOW_CONTAINER_FILE_AMD \
python -m pip install \
--upgrade pip setuptools distlib \
>&2
export PIP_USER=0
apptainer exec $TENSORFLOW_CONTAINER_FILE_AMD \
python -m pip install \
--prefix=$TENSORFLOW_PACKAGES_AMD \
--ignore-installed --no-deps \
--no-cache-dir \
-r $BENCH_DIR/requirements/amd_tensorflow_requirements.txt\
>&2
touch $TENSORFLOW_PACKAGES_FILE_AMD
echo "Done building additional packages for $ACCELERATOR in $TENSORFLOW_PACKAGES_AMD" >&2
elif ! [ -f $TENSORFLOW_PACKAGES_FILE_IPU ] && [ "$ACCELERATOR" = "GC200" ]; then
mkdir -p $TENSORFLOW_PACKAGES_IPU
apptainer exec $TENSORFLOW_CONTAINER_FILE_IPU \
python -m pip install \
--upgrade pip setuptools distlib \
>&2
export PIP_USER=0
apptainer exec $TENSORFLOW_CONTAINER_FILE_IPU \
python -m pip install \
--prefix=$TENSORFLOW_PACKAGES_IPU \
--ignore-installed --no-deps \
--no-cache-dir \
-r $BENCH_DIR/requirements/ipu_tensorflow_requirements.txt \
>&2
touch $TENSORFLOW_PACKAGES_FILE_IPU
echo "Done building additional packages for $ACCELERATOR in $TENSORFLOW_PACKAGES_IPU" >&2
elif ! [ -f $TENSORFLOW_PACKAGES_FILE_NVIDIA_X86 ] && [[ " ${NVIDIA_X86_ACCELERATORS[@]} " == *" $ACCELERATOR "* ]]; then
mkdir -p $TENSORFLOW_PACKAGES_NVIDIA_X86
apptainer exec $TENSORFLOW_CONTAINER_FILE_NVIDIA_X86 \
python -m pip install \
--upgrade pip setuptools distlib \
>&2
export PIP_USER=0
apptainer exec $TENSORFLOW_CONTAINER_FILE_NVIDIA_X86 \
python -m pip install \
--prefix=$TENSORFLOW_PACKAGES_NVIDIA_X86 \
--ignore-installed --no-deps \
--no-cache-dir \
-r $BENCH_DIR/requirements/nvidia_x86_tensorflow_requirements.txt\
>&2
touch $TENSORFLOW_PACKAGES_FILE_NVIDIA_X86
echo "Done building additional packages for $ACCELERATOR in $TENSORFLOW_PACKAGES_NVIDIA_X86" >&2
elif ! [ -f $TENSORFLOW_PACKAGES_FILE_NVIDIA_ARM ] && [[ " ${NVIDIA_ARM_ACCELERATORS[@]} " == *" $ACCELERATOR "* ]]; then
mkdir -p $TENSORFLOW_PACKAGES_NVIDIA_ARM
apptainer exec $TENSORFLOW_CONTAINER_FILE_NVIDIA_ARM \
python -m pip install \
--upgrade pip setuptools distlib \
>&2
export PIP_USER=0
apptainer exec $TENSORFLOW_CONTAINER_FILE_NVIDIA_ARM \
python -m pip install \
--prefix=$TENSORFLOW_PACKAGES_NVIDIA_ARM \
--ignore-installed --no-deps \
--no-cache-dir \
-r $BENCH_DIR/requirements/nvidia_arm_tensorflow_requirements.txt\
>&2
touch $TENSORFLOW_PACKAGES_FILE_NVIDIA_ARM
echo "Done building additional packages for $ACCELERATOR in $TENSORFLOW_PACKAGES_NVIDIA_ARM" >&2
else
echo "No additional packages required for $ACCELERATOR" >&2
fi