formalize the ROCm image pre-build process

manywheel/Dockerfile

@@ -142,7 +142,7 @@ RUN rpm -e cmake || true
 # cmake-3.18.4 from pip
 RUN yum install -y python3-pip && \
     python3 -mpip install cmake==3.18.4 && \
-    ln -s /usr/local/bin/cmake /usr/bin/cmake
+    ln -sf /usr/local/bin/cmake /usr/bin/cmake
 
 # ninja
 RUN yum install -y http://repo.okay.com.mx/centos/7/x86_64/release/okay-release-1-5.el7.noarch.rpm
@@ -153,19 +153,40 @@ RUN rm -rf /usr/local/cuda-${BASE_CUDA_VERSION}
 COPY --from=cuda     /usr/local/cuda-${BASE_CUDA_VERSION}  /usr/local/cuda-${BASE_CUDA_VERSION}
 COPY --from=magma    /usr/local/cuda-${BASE_CUDA_VERSION}  /usr/local/cuda-${BASE_CUDA_VERSION}
 
-FROM cpu_final as rocm_final
+### Special target for ROCm for building magma and miopen from source ahead of time.
+### For this target, GPU_IMAGE is rocm/dev-centos-7:${ROCM_VERSION}-complete
+FROM ${GPU_IMAGE} as rocm_prebuild
+ENV LC_ALL en_US.UTF-8
+ENV LANG en_US.UTF-8
+ENV LANGUAGE en_US.UTF-8
 ARG ROCM_VERSION=3.7
 ARG PYTORCH_ROCM_ARCH
 ENV PYTORCH_ROCM_ARCH ${PYTORCH_ROCM_ARCH}
-# Install ROCm
-ADD ./common/install_rocm.sh install_rocm.sh
-RUN ROCM_VERSION=${ROCM_VERSION} bash ./install_rocm.sh && rm install_rocm.sh
+# Install ROCm --- not needed, GPU_IMAGE contains full rocm install
+#ADD ./common/install_rocm.sh install_rocm.sh
+#RUN ROCM_VERSION=${ROCM_VERSION} bash ./install_rocm.sh && rm install_rocm.sh
+# Install patched libdrm
 ADD ./common/install_rocm_drm.sh install_rocm_drm.sh
 RUN bash ./install_rocm_drm.sh && rm install_rocm_drm.sh
+# Install magma
+COPY --from=intel /opt/intel /opt/intel
+ADD ./common/install_rocm_magma.sh install_rocm_magma.sh
+RUN bash ./install_rocm_magma.sh && rm install_rocm_magma.sh
+RUN rm -rf /opt/intel
+# Install custom miopen
+# cmake is already installed inside the rocm base image, so remove if present
+RUN rpm -e cmake || true
+# cmake-3.18.4 from pip
+RUN yum install -y python3-pip && \
+    python3 -mpip install cmake==3.18.4 && \
+    ln -sf /usr/local/bin/cmake /usr/bin/cmake
 # cmake3 is needed for the MIOpen build
 RUN ln -sf /usr/local/bin/cmake /usr/bin/cmake3
-### The following is now performed beforehand in a new GPU_IMAGE with magma and miopen preinstalled
-#ADD ./common/install_rocm_magma.sh install_rocm_magma.sh
-#RUN bash ./install_rocm_magma.sh && rm install_rocm_magma.sh
-#ADD ./common/install_miopen.sh install_miopen.sh
-#RUN bash ./install_miopen.sh ${ROCM_VERSION} && rm install_miopen.sh
+RUN yum install -y autoconf aclocal automake make
+ADD ./common/install_miopen.sh install_miopen.sh
+RUN bash ./install_miopen.sh ${ROCM_VERSION} && rm install_miopen.sh
+
+### Final target for ROCm. Assumes ./manywheel/build_docker_rocm_prebuild.sh has already been run.
+### For this target, GPU_IMAGE is rocm/centos-7:${ROCM_VERSION}-magma-miopen-staging
+FROM cpu_final as rocm_final
+

manywheel/build_docker_rocm_prebuild.sh

@@ -0,0 +1,29 @@
+#!/usr/bin/env bash
+
+set -eou pipefail
+
+TOPDIR=$(git rev-parse --show-toplevel)
+
+GPU_ARCH_VERSION=${GPU_ARCH_VERSION:-}
+
+TARGET=rocm_prebuild
+GPU_IMAGE=rocm/dev-centos-7:${GPU_ARCH_VERSION}-complete
+DOCKER_IMAGE=rocm/dev-centos-7:${GPU_ARCH_VERSION}-magma-miopen-staging
+PYTORCH_ROCM_ARCH="gfx900;gfx906;gfx908;gfx90a;gfx1030;gfx1100"
+DOCKER_GPU_BUILD_ARG="--build-arg ROCM_VERSION=${GPU_ARCH_VERSION} --build-arg PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH} --build-arg DEVTOOLSET_VERSION=9"
+
+(
+    set -x
+    DOCKER_BUILDKIT=1 docker build \
+        -t "${DOCKER_IMAGE}" \
+        ${DOCKER_GPU_BUILD_ARG} \
+        --build-arg "GPU_IMAGE=${GPU_IMAGE}" \
+        --target "${TARGET}" \
+        -f "${TOPDIR}/manywheel/Dockerfile" \
+        "${TOPDIR}"
+)
+
+(
+    set -x
+    docker push "${DOCKER_IMAGE}"
+)