xref: /aosp_15_r20/external/pytorch/.ci/docker/ubuntu-rocm/Dockerfile (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1ARG UBUNTU_VERSION
2
3FROM ubuntu:${UBUNTU_VERSION}
4
5ARG UBUNTU_VERSION
6
7ENV DEBIAN_FRONTEND noninteractive
8
9# Set AMD gpu targets to build for
10ARG PYTORCH_ROCM_ARCH
11ENV PYTORCH_ROCM_ARCH ${PYTORCH_ROCM_ARCH}
12
13# Install common dependencies (so that this step can be cached separately)
14COPY ./common/install_base.sh install_base.sh
15RUN bash ./install_base.sh && rm install_base.sh
16
17# Install clang
18ARG LLVMDEV
19ARG CLANG_VERSION
20COPY ./common/install_clang.sh install_clang.sh
21RUN bash ./install_clang.sh && rm install_clang.sh
22
23# Install user
24COPY ./common/install_user.sh install_user.sh
25RUN bash ./install_user.sh && rm install_user.sh
26
27# Install conda and other packages (e.g., numpy, pytest)
28ARG ANACONDA_PYTHON_VERSION
29ARG CONDA_CMAKE
30ENV ANACONDA_PYTHON_VERSION=$ANACONDA_PYTHON_VERSION
31ENV PATH /opt/conda/envs/py_$ANACONDA_PYTHON_VERSION/bin:/opt/conda/bin:$PATH
32COPY requirements-ci.txt /opt/conda/requirements-ci.txt
33COPY ./common/install_conda.sh install_conda.sh
34COPY ./common/common_utils.sh common_utils.sh
35RUN bash ./install_conda.sh && rm install_conda.sh common_utils.sh /opt/conda/requirements-ci.txt
36
37# Install gcc
38ARG GCC_VERSION
39COPY ./common/install_gcc.sh install_gcc.sh
40RUN bash ./install_gcc.sh && rm install_gcc.sh
41
42# (optional) Install protobuf for ONNX
43ARG PROTOBUF
44COPY ./common/install_protobuf.sh install_protobuf.sh
45RUN if [ -n "${PROTOBUF}" ]; then bash ./install_protobuf.sh; fi
46RUN rm install_protobuf.sh
47ENV INSTALLED_PROTOBUF ${PROTOBUF}
48
49# (optional) Install database packages like LMDB and LevelDB
50ARG DB
51COPY ./common/install_db.sh install_db.sh
52RUN if [ -n "${DB}" ]; then bash ./install_db.sh; fi
53RUN rm install_db.sh
54ENV INSTALLED_DB ${DB}
55
56# (optional) Install vision packages like OpenCV
57ARG VISION
58COPY ./common/install_vision.sh ./common/cache_vision_models.sh ./common/common_utils.sh ./
59RUN if [ -n "${VISION}" ]; then bash ./install_vision.sh; fi
60RUN rm install_vision.sh cache_vision_models.sh common_utils.sh
61ENV INSTALLED_VISION ${VISION}
62
63# Install rocm
64ARG ROCM_VERSION
65COPY ./common/install_rocm.sh install_rocm.sh
66RUN bash ./install_rocm.sh
67RUN rm install_rocm.sh
68COPY ./common/install_rocm_magma.sh install_rocm_magma.sh
69RUN bash ./install_rocm_magma.sh
70RUN rm install_rocm_magma.sh
71ADD ./common/install_miopen.sh install_miopen.sh
72RUN bash ./install_miopen.sh ${ROCM_VERSION} && rm install_miopen.sh
73ENV ROCM_PATH /opt/rocm
74ENV PATH /opt/rocm/bin:$PATH
75ENV PATH /opt/rocm/hcc/bin:$PATH
76ENV PATH /opt/rocm/hip/bin:$PATH
77ENV PATH /opt/rocm/opencl/bin:$PATH
78ENV PATH /opt/rocm/llvm/bin:$PATH
79ENV MAGMA_HOME /opt/rocm/magma
80ENV LANG C.UTF-8
81ENV LC_ALL C.UTF-8
82
83# Install amdsmi
84COPY ./common/install_amdsmi.sh install_amdsmi.sh
85RUN bash ./install_amdsmi.sh
86RUN rm install_amdsmi.sh
87
88# (optional) Install non-default CMake version
89ARG CMAKE_VERSION
90COPY ./common/install_cmake.sh install_cmake.sh
91RUN if [ -n "${CMAKE_VERSION}" ]; then bash ./install_cmake.sh; fi
92RUN rm install_cmake.sh
93
94# (optional) Install non-default Ninja version
95ARG NINJA_VERSION
96COPY ./common/install_ninja.sh install_ninja.sh
97RUN if [ -n "${NINJA_VERSION}" ]; then bash ./install_ninja.sh; fi
98RUN rm install_ninja.sh
99
100ARG TRITON
101# Install triton, this needs to be done before sccache because the latter will
102# try to reach out to S3, which docker build runners don't have access
103COPY ./common/install_triton.sh install_triton.sh
104COPY ./common/common_utils.sh common_utils.sh
105COPY ci_commit_pins/triton.txt triton.txt
106COPY triton_version.txt triton_version.txt
107RUN if [ -n "${TRITON}" ]; then bash ./install_triton.sh; fi
108RUN rm install_triton.sh common_utils.sh triton.txt triton_version.txt
109
110# Install AOTriton
111COPY ./aotriton_version.txt aotriton_version.txt
112COPY ./common/common_utils.sh common_utils.sh
113COPY ./common/install_aotriton.sh install_aotriton.sh
114RUN ["/bin/bash", "-c", "./install_aotriton.sh /opt/rocm && rm -rf install_aotriton.sh aotriton_version.txt common_utils.sh"]
115ENV AOTRITON_INSTALLED_PREFIX /opt/rocm/aotriton
116
117# Install ccache/sccache (do this last, so we get priority in PATH)
118COPY ./common/install_cache.sh install_cache.sh
119ENV PATH /opt/cache/bin:$PATH
120RUN bash ./install_cache.sh && rm install_cache.sh
121
122# Include BUILD_ENVIRONMENT environment variable in image
123ARG BUILD_ENVIRONMENT
124ENV BUILD_ENVIRONMENT ${BUILD_ENVIRONMENT}
125
126# Install LLVM dev version (Defined in the pytorch/builder github repository)
127COPY --from=pytorch/llvm:9.0.1 /opt/llvm /opt/llvm
128
129USER jenkins
130CMD ["bash"]
131