forked from ivy-llc/ivy
-
Notifications
You must be signed in to change notification settings - Fork 0
/
DockerfileGPU
146 lines (118 loc) · 5.33 KB
/
DockerfileGPU
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
# BASE CUDA IMAGE #
# ----------------#
ARG UBUNTU_VERSION=20.04
ARG CUDA=11.2
FROM nvidia/cuda:${CUDA}.1-base-ubuntu${UBUNTU_VERSION} as base
# For TensorFlow #
# ---------------#
# Adapted from
# https://github.com/tensorflow/tensorflow/blob/master/tensorflow/tools/dockerfiles/dockerfiles/gpu.Dockerfile
# CUDA is specified again because the FROM directive resets ARGs
# (but their default value is retained if set previously)
ARG CUDA
ARG CUDNN=8.1.0.77-1
ARG CUDNN_MAJOR_VERSION=8
ARG LIB_DIR_PREFIX=x86_64
ARG LIBNVINFER=7.2.2-1
ARG LIBNVINFER_MAJOR_VERSION=7
# Let us install tzdata painlessly
ENV DEBIAN_FRONTEND=noninteractive
# Needed for string substitution
SHELL ["/bin/bash", "-c"]
# taken from https://github.com/Kaggle/docker-python/commit/f1a3cfc6ee71899b1af8d76598b42a2da280448d
RUN \
# Temporarily swap the NVIDIA GPG key. Remove once new base image with new GPG key is released.
rm /etc/apt/sources.list.d/cuda.list && \
apt-key del 7fa2af80 && \
apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/7fa2af80.pub && \
apt-get update
# Refer to https://forums.developer.nvidia.com/t/notice-cuda-linux-repository-key-rotation/212772
RUN apt-get update && apt-get install -y wget
RUN wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/cuda-keyring_1.0-1_all.deb
RUN dpkg -i cuda-keyring_1.0-1_all.deb
RUN apt-get update
# Pick up some TF dependencies
RUN apt-get update && apt-get install -y --no-install-recommends \
build-essential \
cuda-command-line-tools-${CUDA/./-} \
libcublas-${CUDA/./-} \
cuda-nvrtc-${CUDA/./-} \
libcufft-${CUDA/./-} \
libcurand-${CUDA/./-} \
libcusolver-${CUDA/./-} \
libcusparse-${CUDA/./-} \
curl \
libcudnn8=${CUDNN}+cuda${CUDA} \
libfreetype6-dev \
libhdf5-serial-dev \
libzmq3-dev \
pkg-config \
software-properties-common \
unzip
# Install TensorRT if not building for PowerPC
RUN apt-get update && \
apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/machine-learning/repos/ubuntu1804/x86_64/7fa2af80.pub && \
echo "deb https://developer.download.nvidia.com/compute/machine-learning/repos/ubuntu1804/x86_64 /" > /etc/apt/sources.list.d/tensorRT.list && \
apt-get update && \
apt-get install -y --no-install-recommends libnvinfer${LIBNVINFER_MAJOR_VERSION}=${LIBNVINFER}+cuda11.0 \
libnvinfer-plugin${LIBNVINFER_MAJOR_VERSION}=${LIBNVINFER}+cuda11.0 \
&& apt-get clean \
&& rm -rf /var/lib/apt/lists/*
# For CUDA profiling, TensorFlow requires CUPTI.
ENV LD_LIBRARY_PATH /usr/local/cuda-11.0/targets/x86_64-linux/lib:/usr/local/cuda/extras/CUPTI/lib64:/usr/local/cuda/lib64:$LD_LIBRARY_PATH
# Link the libcuda stub to the location where tensorflow is searching for it and reconfigure
# dynamic linker run-time bindings
RUN ln -s /usr/local/cuda/lib64/stubs/libcuda.so /usr/local/cuda/lib64/stubs/libcuda.so.1 \
&& echo "/usr/local/cuda/lib64/stubs" > /etc/ld.so.conf.d/z-cuda-stubs.conf \
&& ldconfig
# See http://bugs.python.org/issue19846
ENV LANG C.UTF-8
RUN apt-get update && apt-get install -y \
python3 \
python3-pip
RUN pip3 install --upgrade pip
RUN python3 -m pip --no-cache-dir install --upgrade \
pip \
setuptools==58.5.3
# Some TF tools expect a "python" binary
RUN ln -s $(which python3) /usr/local/bin/python
# For MXNet #
# ----------#
# Install NCCL
RUN apt-get install -y wget
# leave below commented for now
# RUN wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/cuda-ubuntu1804.pin
# RUN mv cuda-ubuntu1804.pin /etc/apt/preferences.d/cuda-repository-pin-600
# RUN apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/7fa2af80.pub
# RUN add-apt-repository "deb http://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/ /"
RUN apt-get install libnccl2=2.8.4-1+cuda11.0 libnccl-dev=2.8.4-1+cuda11.0
# OpenBLAS
RUN apt-get install -y libopenblas-dev
# CUDA-specific Installations #
# ----------------------------#
RUN pip3 install pytest
RUN apt-get install -y git
RUN apt-get install -y python-opengl
RUN pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_releases.html
# ToDo: work out why most recent jaxlib installs are not working
RUN pip3 install --upgrade "jaxlib[cuda]" -f https://storage.googleapis.com/jax-releases/jax_releases.html
RUN pip3 install --upgrade torch-scatter -f https://pytorch-geometric.com/whl/torch-1.10.0+cu111.html
# Ivy #
# ----#
# Install Ivy Upstream
RUN git clone https://github.com/unifyai/ivy && \
cd ivy && \
cat requirements.txt | grep -v "ivy-" | pip3 install --no-cache-dir -r /dev/stdin && \
cat optional.txt | grep -v "ivy-" | pip3 install --no-cache-dir -r /dev/stdin && \
python3 -m pip install --user -e .
# Install local requirements
COPY requirements.txt .
RUN pip3 install --no-cache-dir -r requirements.txt
# Install local optional
COPY optional.txt .
RUN pip3 install --no-cache-dir -r optional.txt
RUN pip3 install torch==1.10.0+cu111 -f https://download.pytorch.org/whl/torch_stable.html
COPY test_dependencies.py /
RUN python3 test_dependencies.py -fp requirements.txt,optional.txt && \
rm -rf requirements.txt && \
rm -rf optional.txt