xref: /aosp_15_r20/external/pytorch/.ci/docker/ubuntu/Dockerfile (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1ARG UBUNTU_VERSION
2
3FROM ubuntu:${UBUNTU_VERSION}
4
5ARG UBUNTU_VERSION
6
7ENV DEBIAN_FRONTEND noninteractive
8
9ARG CLANG_VERSION
10
11# Install common dependencies (so that this step can be cached separately)
12COPY ./common/install_base.sh install_base.sh
13RUN bash ./install_base.sh && rm install_base.sh
14
15# Install clang
16ARG LLVMDEV
17COPY ./common/install_clang.sh install_clang.sh
18RUN bash ./install_clang.sh && rm install_clang.sh
19
20# Install user
21COPY ./common/install_user.sh install_user.sh
22RUN bash ./install_user.sh && rm install_user.sh
23
24# Install katex
25ARG KATEX
26COPY ./common/install_docs_reqs.sh install_docs_reqs.sh
27RUN bash ./install_docs_reqs.sh && rm install_docs_reqs.sh
28
29# Install conda and other packages (e.g., numpy, pytest)
30ARG ANACONDA_PYTHON_VERSION
31ARG CONDA_CMAKE
32ARG DOCS
33ENV ANACONDA_PYTHON_VERSION=$ANACONDA_PYTHON_VERSION
34ENV PATH /opt/conda/envs/py_$ANACONDA_PYTHON_VERSION/bin:/opt/conda/bin:$PATH
35ENV DOCS=$DOCS
36COPY requirements-ci.txt requirements-docs.txt /opt/conda/
37COPY ./common/install_conda.sh install_conda.sh
38COPY ./common/common_utils.sh common_utils.sh
39RUN bash ./install_conda.sh && rm install_conda.sh common_utils.sh /opt/conda/requirements-ci.txt /opt/conda/requirements-docs.txt
40RUN if [ -n "${UNINSTALL_DILL}" ]; then pip uninstall -y dill; fi
41
42# Install gcc
43ARG GCC_VERSION
44COPY ./common/install_gcc.sh install_gcc.sh
45RUN bash ./install_gcc.sh && rm install_gcc.sh
46
47# Install lcov for C++ code coverage
48COPY ./common/install_lcov.sh install_lcov.sh
49RUN  bash ./install_lcov.sh && rm install_lcov.sh
50
51# Install cuda and cudnn
52ARG CUDA_VERSION
53COPY ./common/install_cuda.sh install_cuda.sh
54RUN bash ./install_cuda.sh ${CUDA_VERSION} && rm install_cuda.sh
55ENV DESIRED_CUDA ${CUDA_VERSION}
56ENV PATH /usr/local/nvidia/bin:/usr/local/cuda/bin:$PATH
57
58# (optional) Install UCC
59ARG UCX_COMMIT
60ARG UCC_COMMIT
61ENV UCX_COMMIT $UCX_COMMIT
62ENV UCC_COMMIT $UCC_COMMIT
63ENV UCX_HOME /usr
64ENV UCC_HOME /usr
65ADD ./common/install_ucc.sh install_ucc.sh
66RUN if [ -n "${UCX_COMMIT}" ] && [ -n "${UCC_COMMIT}" ]; then bash ./install_ucc.sh; fi
67RUN rm install_ucc.sh
68
69# (optional) Install protobuf for ONNX
70ARG PROTOBUF
71COPY ./common/install_protobuf.sh install_protobuf.sh
72RUN if [ -n "${PROTOBUF}" ]; then bash ./install_protobuf.sh; fi
73RUN rm install_protobuf.sh
74ENV INSTALLED_PROTOBUF ${PROTOBUF}
75
76# (optional) Install database packages like LMDB and LevelDB
77ARG DB
78COPY ./common/install_db.sh install_db.sh
79RUN if [ -n "${DB}" ]; then bash ./install_db.sh; fi
80RUN rm install_db.sh
81ENV INSTALLED_DB ${DB}
82
83# (optional) Install vision packages like OpenCV
84ARG VISION
85COPY ./common/install_vision.sh ./common/cache_vision_models.sh ./common/common_utils.sh ./
86RUN if [ -n "${VISION}" ]; then bash ./install_vision.sh; fi
87RUN rm install_vision.sh cache_vision_models.sh common_utils.sh
88ENV INSTALLED_VISION ${VISION}
89
90# (optional) Install Android NDK
91ARG ANDROID
92ARG ANDROID_NDK
93ARG GRADLE_VERSION
94COPY ./common/install_android.sh ./common/cache_vision_models.sh ./common/common_utils.sh ./
95COPY ./android/AndroidManifest.xml AndroidManifest.xml
96COPY ./android/build.gradle build.gradle
97RUN if [ -n "${ANDROID}" ]; then bash ./install_android.sh; fi
98RUN rm install_android.sh cache_vision_models.sh common_utils.sh
99RUN rm AndroidManifest.xml
100RUN rm build.gradle
101ENV INSTALLED_ANDROID ${ANDROID}
102
103# (optional) Install Vulkan SDK
104ARG VULKAN_SDK_VERSION
105COPY ./common/install_vulkan_sdk.sh install_vulkan_sdk.sh
106RUN if [ -n "${VULKAN_SDK_VERSION}" ]; then bash ./install_vulkan_sdk.sh; fi
107RUN rm install_vulkan_sdk.sh
108
109# (optional) Install swiftshader
110ARG SWIFTSHADER
111COPY ./common/install_swiftshader.sh install_swiftshader.sh
112RUN if [ -n "${SWIFTSHADER}" ]; then bash ./install_swiftshader.sh; fi
113RUN rm install_swiftshader.sh
114
115# (optional) Install non-default CMake version
116ARG CMAKE_VERSION
117COPY ./common/install_cmake.sh install_cmake.sh
118RUN if [ -n "${CMAKE_VERSION}" ]; then bash ./install_cmake.sh; fi
119RUN rm install_cmake.sh
120
121# (optional) Install non-default Ninja version
122ARG NINJA_VERSION
123COPY ./common/install_ninja.sh install_ninja.sh
124RUN if [ -n "${NINJA_VERSION}" ]; then bash ./install_ninja.sh; fi
125RUN rm install_ninja.sh
126
127COPY ./common/install_openssl.sh install_openssl.sh
128RUN bash ./install_openssl.sh
129ENV OPENSSL_ROOT_DIR /opt/openssl
130ENV OPENSSL_DIR /opt/openssl
131RUN rm install_openssl.sh
132
133ARG INDUCTOR_BENCHMARKS
134COPY ./common/install_inductor_benchmark_deps.sh install_inductor_benchmark_deps.sh
135COPY ./common/common_utils.sh common_utils.sh
136COPY ci_commit_pins/huggingface.txt huggingface.txt
137COPY ci_commit_pins/timm.txt timm.txt
138RUN if [ -n "${INDUCTOR_BENCHMARKS}" ]; then bash ./install_inductor_benchmark_deps.sh; fi
139RUN rm install_inductor_benchmark_deps.sh common_utils.sh timm.txt huggingface.txt
140
141ARG TRITON
142# Install triton, this needs to be done before sccache because the latter will
143# try to reach out to S3, which docker build runners don't have access
144COPY ./common/install_triton.sh install_triton.sh
145COPY ./common/common_utils.sh common_utils.sh
146COPY ci_commit_pins/triton.txt triton.txt
147RUN if [ -n "${TRITON}" ]; then bash ./install_triton.sh; fi
148RUN rm install_triton.sh common_utils.sh triton.txt
149
150ARG EXECUTORCH
151# Build and install executorch
152COPY ./common/install_executorch.sh install_executorch.sh
153COPY ./common/common_utils.sh common_utils.sh
154COPY ci_commit_pins/executorch.txt executorch.txt
155RUN if [ -n "${EXECUTORCH}" ]; then bash ./install_executorch.sh; fi
156RUN rm install_executorch.sh common_utils.sh executorch.txt
157
158ARG HALIDE
159# Build and install halide
160COPY ./common/install_halide.sh install_halide.sh
161COPY ./common/common_utils.sh common_utils.sh
162COPY ci_commit_pins/halide.txt halide.txt
163RUN if [ -n "${HALIDE}" ]; then bash ./install_halide.sh; fi
164RUN rm install_halide.sh common_utils.sh halide.txt
165
166ARG ONNX
167# Install ONNX dependencies
168COPY ./common/install_onnx.sh ./common/common_utils.sh ./
169RUN if [ -n "${ONNX}" ]; then bash ./install_onnx.sh; fi
170RUN rm install_onnx.sh common_utils.sh
171
172# (optional) Build ACL
173ARG ACL
174COPY ./common/install_acl.sh install_acl.sh
175RUN if [ -n "${ACL}" ]; then bash ./install_acl.sh; fi
176RUN rm install_acl.sh
177ENV INSTALLED_ACL ${ACL}
178
179# Install ccache/sccache (do this last, so we get priority in PATH)
180ARG SKIP_SCCACHE_INSTALL
181COPY ./common/install_cache.sh install_cache.sh
182ENV PATH /opt/cache/bin:$PATH
183RUN if [ -z "${SKIP_SCCACHE_INSTALL}" ]; then bash ./install_cache.sh; fi
184RUN rm install_cache.sh
185
186# Add jni.h for java host build
187COPY ./common/install_jni.sh install_jni.sh
188COPY ./java/jni.h jni.h
189RUN bash ./install_jni.sh && rm install_jni.sh
190
191# Install Open MPI for CUDA
192COPY ./common/install_openmpi.sh install_openmpi.sh
193RUN if [ -n "${CUDA_VERSION}" ]; then bash install_openmpi.sh; fi
194RUN rm install_openmpi.sh
195
196# Include BUILD_ENVIRONMENT environment variable in image
197ARG BUILD_ENVIRONMENT
198ENV BUILD_ENVIRONMENT ${BUILD_ENVIRONMENT}
199
200# Install LLVM dev version (Defined in the pytorch/builder github repository)
201ARG SKIP_LLVM_SRC_BUILD_INSTALL
202COPY --from=pytorch/llvm:9.0.1 /opt/llvm /opt/llvm
203RUN if [ -n "${SKIP_LLVM_SRC_BUILD_INSTALL}" ]; then set -eu; rm -rf /opt/llvm; fi
204
205# AWS specific CUDA build guidance
206ENV TORCH_CUDA_ARCH_LIST Maxwell
207ENV TORCH_NVCC_FLAGS "-Xfatbin -compress-all"
208ENV CUDA_PATH /usr/local/cuda
209
210USER jenkins
211CMD ["bash"]
212