forked from determined-ai/environments
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Dockerfile-default-gpu
57 lines (44 loc) · 1.81 KB
/
Dockerfile-default-gpu
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
ARG BASE_IMAGE
FROM ${BASE_IMAGE}
COPY dockerfile_scripts /tmp/det_dockerfile_scripts
RUN /tmp/det_dockerfile_scripts/install_google_cloud_sdk.sh
ARG TENSORFLOW_PIP
ARG TORCH_PIP
ARG TORCHVISION_PIP
RUN if [ "$TENSORFLOW_PIP" ]; then \
export HOROVOD_WITH_TENSORFLOW=1 \
&& python -m pip install $TENSORFLOW_PIP; \
else \
export HOROVOD_WITH_TENSORFLOW=0; \
fi
RUN if [ "$TORCH_PIP" ]; then python -m pip install $TORCH_PIP; fi
RUN if [ "$TORCHVISION_PIP" ]; then python -m pip install $TORCHVISION_PIP; fi
ARG TF_CUDA_SYM
RUN if [ "$TF_CUDA_SYM" ]; then ln -s /usr/local/cuda/lib64/libcusolver.so.11 /opt/conda/lib/python3.8/site-packages/tensorflow/python/libcusolver.so.10; fi
ARG TORCH_TB_PROFILER_PIP
RUN if [ "$TORCH_TB_PROFILER_PIP" ]; then python -m pip install $TORCH_TB_PROFILER_PIP; fi
ARG TORCH_CUDA_ARCH_LIST
ARG APEX_GIT
ARG APEX_PATCH
RUN /tmp/det_dockerfile_scripts/install_apex.sh
ARG DET_BUILD_NCCL=1
RUN if [ -n "${DET_BUILD_NCCL}" ]; then /tmp/det_dockerfile_scripts/build_nccl.sh; fi
ARG HOROVOD_NCCL_LINK=${DET_BUILD_NCCL:+STATIC}
ARG HOROVOD_NCCL_HOME=${DET_BUILD_NCCL:+/tmp/det_nccl/build}
ARG HOROVOD_WITH_TENSORFLOW
RUN if [ "$HOROVOD_WITH_TENSORFLOW" ]; then export HOROVOD_WITH_TENSORFLOW=$HOROVOD_WITH_TENSORFLOW; fi
ARG HOROVOD_PIP=""
ARG HOROVOD_WITH_PYTORCH=1
ARG HOROVOD_WITHOUT_MXNET=1
ARG HOROVOD_GPU_OPERATIONS=NCCL
ARG HOROVOD_WITH_MPI
ARG HOROVOD_CPU_OPERATIONS
ARG HOROVOD_WITHOUT_MPI
RUN pip install cmake==3.22.4 protobuf==3.20.1
RUN if [ -n "${HOROVOD_PIP}" ]; then ldconfig /usr/local/cuda/targets/x86_64-linux/lib/stubs && \
pip install "$HOROVOD_PIP" && \
ldconfig; fi
RUN python -m pip install -r /tmp/det_dockerfile_scripts/additional-requirements.txt
ARG DEEPSPEED_PIP
RUN if [ -n "$DEEPSPEED_PIP" ]; then /tmp/det_dockerfile_scripts/install_deepspeed.sh; fi
RUN rm -r /tmp/*