xref: /aosp_15_r20/external/pytorch/.ci/docker/manywheel/Dockerfile_cuda_aarch64 (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1FROM quay.io/pypa/manylinux_2_28_aarch64 as base
2
3# Cuda ARM build needs gcc 11
4ARG DEVTOOLSET_VERSION=11
5
6# Language variables
7ENV LC_ALL=en_US.UTF-8
8ENV LANG=en_US.UTF-8
9ENV LANGUAGE=en_US.UTF-8
10
11# Installed needed OS packages. This is to support all
12# the binary builds (torch, vision, audio, text, data)
13RUN yum -y install epel-release
14RUN yum -y update
15RUN yum install -y \
16  autoconf \
17  automake \
18  bison \
19  bzip2 \
20  curl \
21  diffutils \
22  file \
23  git \
24  make \
25  patch \
26  perl \
27  unzip \
28  util-linux \
29  wget \
30  which \
31  xz \
32  yasm \
33  less \
34  zstd \
35  libgomp \
36  sudo \
37  gcc-toolset-${DEVTOOLSET_VERSION}-toolchain
38
39# Ensure the expected devtoolset is used
40ENV PATH=/opt/rh/gcc-toolset-${DEVTOOLSET_VERSION}/root/usr/bin:$PATH
41ENV LD_LIBRARY_PATH=/opt/rh/gcc-toolset-${DEVTOOLSET_VERSION}/root/usr/lib64:/opt/rh/gcc-toolset-${DEVTOOLSET_VERSION}/root/usr/lib:$LD_LIBRARY_PATH
42
43# git236+ would refuse to run git commands in repos owned by other users
44# Which causes version check to fail, as pytorch repo is bind-mounted into the image
45# Override this behaviour by treating every folder as safe
46# For more details see https://github.com/pytorch/pytorch/issues/78659#issuecomment-1144107327
47RUN git config --global --add safe.directory "*"
48
49
50FROM base as openssl
51# Install openssl (this must precede `build python` step)
52# (In order to have a proper SSL module, Python is compiled
53# against a recent openssl [see env vars above], which is linked
54# statically. We delete openssl afterwards.)
55ADD ./common/install_openssl.sh install_openssl.sh
56RUN bash ./install_openssl.sh && rm install_openssl.sh
57ENV SSL_CERT_FILE=/opt/_internal/certs.pem
58
59FROM openssl as final
60# remove unncessary python versions
61RUN rm -rf /opt/python/cp26-cp26m /opt/_internal/cpython-2.6.9-ucs2
62RUN rm -rf /opt/python/cp26-cp26mu /opt/_internal/cpython-2.6.9-ucs4
63RUN rm -rf /opt/python/cp33-cp33m /opt/_internal/cpython-3.3.6
64RUN rm -rf /opt/python/cp34-cp34m /opt/_internal/cpython-3.4.6
65
66FROM base as cuda
67ARG BASE_CUDA_VERSION
68# Install CUDA
69ADD ./common/install_cuda_aarch64.sh install_cuda_aarch64.sh
70RUN bash ./install_cuda_aarch64.sh ${BASE_CUDA_VERSION} && rm install_cuda_aarch64.sh
71
72FROM base as magma
73ARG BASE_CUDA_VERSION
74# Install magma
75ADD ./common/install_magma.sh install_magma.sh
76RUN bash ./install_magma.sh ${BASE_CUDA_VERSION} && rm install_magma.sh
77
78FROM base as nvpl
79# Install nvpl
80ADD ./common/install_nvpl.sh install_nvpl.sh
81RUN bash ./install_nvpl.sh && rm install_nvpl.sh
82
83FROM final as cuda_final
84ARG BASE_CUDA_VERSION
85RUN rm -rf /usr/local/cuda-${BASE_CUDA_VERSION}
86COPY --from=cuda     /usr/local/cuda-${BASE_CUDA_VERSION}  /usr/local/cuda-${BASE_CUDA_VERSION}
87COPY --from=magma    /usr/local/cuda-${BASE_CUDA_VERSION}  /usr/local/cuda-${BASE_CUDA_VERSION}
88COPY --from=nvpl /opt/nvpl/lib/  /usr/local/lib/
89COPY --from=nvpl /opt/nvpl/include/  /usr/local/include/
90RUN ln -sf /usr/local/cuda-${BASE_CUDA_VERSION} /usr/local/cuda
91ENV PATH=/usr/local/cuda/bin:$PATH
92