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