Install OpenMPI and mpi4py in AMD Ubuntu images (#8)

Refactor the other commands to cleanup after each docker step to avoid
growing the final layer with unused stuff.

build_tools/install_amdgpu.sh

@@ -28,3 +28,5 @@ fi
 
 wget https://github.com/GPUOpen-Drivers/AMDVLK/releases/download/v-2023.Q3.1/amdvlk_2023.Q3.1_amd64.deb && \
     dpkg -i amdvlk_2023.Q3.1_amd64.deb
+
+apt clean && rm -rf /var/lib/apt/lists/*

build_tools/install_base.sh

@@ -0,0 +1,20 @@
+#!/bin/bash
+
+# Copyright 2024 The IREE Authors
+#
+# Licensed under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+# Basic packages.
+
+set -euo pipefail
+
+apt update
+apt install -y \
+        wget python3.11-dev python3-pip git unzip curl gnupg2 \
+        lsb-release vulkan-tools && \
+    update-alternatives --install /usr/local/bin/python python /usr/bin/python3.11 3 && \
+    update-alternatives --install /usr/local/bin/python3 python3 /usr/bin/python3.11 3
+
+apt clean && rm -rf /var/lib/apt/lists/*

build_tools/install_openmpi.sh

@@ -0,0 +1,17 @@
+#!/bin/bash
+
+# Copyright 2024 The IREE Authors
+#
+# Licensed under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+# Needed by IREE E2E collectives tests.
+
+set -euo pipefail
+
+apt update
+apt install -y libopenmpi-dev
+
+# Needs to run after installing openmpi as it will build against it.
+python -m pip install --no-cache-dir mpi4py

build_tools/install_the_rock.sh

@@ -35,7 +35,7 @@ apt update
 apt install -y \
   repo git-lfs libnuma-dev ninja-build g++ pkg-config libdrm-dev \
   libelf-dev xxd libgl1-mesa-dev
-python -m pip install CppHeaderParser
+python -m pip install --no-cache-dir CppHeaderParser
 
 # Make sure git does not report Committer identity unknown errors.
 # export GIT_COMMITTER_NAME="Noname"
@@ -69,3 +69,5 @@ cmake --build build
 cmake --install build --component amdgpu-runtime
 cmake --install build --component amdgpu-runtime-dev
 cmake --install build --component amdgpu-compiler
+
+apt clean && rm -rf /var/lib/apt/lists/*

dockerfiles/amdgpu_ubuntu_jammy_ghr_x86_64.Dockerfile

@@ -4,11 +4,9 @@
 FROM docker.io/myoung34/github-runner:2.314.1-ubuntu-jammy
 
 # Basic packages.
-RUN apt update && \
-    apt install -y \
-        wget python3.11 python3-pip git unzip curl gnupg2 lsb-release vulkan-tools && \
-    update-alternatives --install /usr/local/bin/python python /usr/bin/python3.11 3 && \
-    update-alternatives --install /usr/local/bin/python3 python3 /usr/bin/python3.11 3
+WORKDIR /install-base
+COPY build_tools/install_base.sh ./
+RUN ./install_base.sh "${CMAKE_VERSION}" && rm -rf /install-base
 
 # CMake
 WORKDIR /install-cmake
@@ -22,17 +20,18 @@ WORKDIR /install-amdgpu
 ARG AMDGPU_VERSION=6.1
 COPY build_tools/install_amdgpu.sh ./
 RUN ./install_amdgpu.sh "${AMDGPU_VERSION}" && rm -rf /install-amdgpu
-WORKDIR /
 
 # TheRock (ROCm)
 WORKDIR /install-the-rock
 COPY build_tools/install_the_rock.sh ./
 RUN ./install_the_rock.sh \
   && rm -rf /install-the-rock
-WORKDIR /
 
-# Clean up.
-RUN apt clean && rm -rf /var/lib/apt/lists/*
+# OpenMPI
+WORKDIR /install-openmpi
+COPY build_tools/install_openmpi.sh ./
+RUN ./install_openmpi.sh \
+  && rm -rf /install-openmpi
 
 # Switch back to the working directory upstream expects.
 WORKDIR /actions-runner

dockerfiles/amdgpu_ubuntu_jammy_x86_64.Dockerfile

@@ -4,11 +4,9 @@
 FROM ubuntu:jammy
 
 # Basic packages.
-RUN apt update && \
-    apt install -y \
-        wget python3.11 python3-pip git unzip curl gnupg2 lsb-release vulkan-tools && \
-    update-alternatives --install /usr/local/bin/python python /usr/bin/python3.11 3 && \
-    update-alternatives --install /usr/local/bin/python3 python3 /usr/bin/python3.11 3
+WORKDIR /install-base
+COPY build_tools/install_base.sh ./
+RUN ./install_base.sh "${CMAKE_VERSION}" && rm -rf /install-base
 
 # CMake
 WORKDIR /install-cmake
@@ -22,14 +20,17 @@ WORKDIR /install-amdgpu
 ARG AMDGPU_VERSION=6.1
 COPY build_tools/install_amdgpu.sh ./
 RUN ./install_amdgpu.sh "${AMDGPU_VERSION}" && rm -rf /install-amdgpu
-WORKDIR /
 
 # TheRock (ROCm)
 WORKDIR /install-the-rock
 COPY build_tools/install_the_rock.sh ./
 RUN ./install_the_rock.sh \
-#  && rm -rf /install-the-rock
-WORKDIR /
+  && rm -rf /install-the-rock
 
-# Clean up.
-RUN apt clean && rm -rf /var/lib/apt/lists/*
+# OpenMPI
+WORKDIR /install-openmpi
+COPY build_tools/install_openmpi.sh ./
+RUN ./install_openmpi.sh \
+  && rm -rf /install-openmpi
+
+WORKDIR /