Multi-architecture Docker images and parallel builds

.github/workflows/docker.yml

@@ -1,4 +1,4 @@
-name: Build & Push Docker Images
+name: Build and push Docker images to Docker Hub
 
 on:
   workflow_dispatch:
@@ -7,69 +7,80 @@ on:
     - develop
 
 jobs:
-  pre_job:
+  build_docker_images:
+    # This workflow is only of value to PyBaMM and would always be skipped in forks
+    # if: github.repository_owner == 'pybamm-team'
+    name: Image (${{ matrix.build-args }})
     runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        build-args: ["No solvers", "JAX", "ODES", "IDAKLU", "ALL"]
+      fail-fast: true
 
     steps:
       - name: Checkout
         uses: actions/checkout@v4
 
+      - name: Set up QEMU
+        uses: docker/setup-qemu-action@v3
+
       - name: Set up Docker Buildx
         uses: docker/setup-buildx-action@v3
 
-      - name: Login to DockerHub
+      - name: Login to Docker Hub
         uses: docker/login-action@v3
         with:
             username: ${{ secrets.DOCKERHUB_USERNAME }}
             password: ${{ secrets.DOCKERHUB_TOKEN }}
 
-      - name: List built images
-        run: docker images
-
-      - name: Build and Push Docker Image (Without Solvers)
-        uses: docker/build-push-action@v5
-        with:
-          context: .
-          file: scripts/Dockerfile
-          tags: pybamm/pybamm:latest
-          push: true
+      - name: Create tags for Docker images based on build-time arguments
+        id: tags
+        run: |
+          if [ "${{ matrix.build-args }}" = "No solvers" ]; then
+            echo "tag=latest" >> "$GITHUB_OUTPUT"
+          elif [ "${{ matrix.build-args }}" = "JAX" ]; then
+            echo "tag=jax" >> "$GITHUB_OUTPUT"
+          elif [ "${{ matrix.build-args }}" = "ODES" ]; then
+            echo "tag=odes" >> "$GITHUB_OUTPUT"
+          elif [ "${{ matrix.build-args }}" = "IDAKLU" ]; then
+            echo "tag=idaklu" >> "$GITHUB_OUTPUT"
+          elif [ "${{ matrix.build-args }}" = "ALL" ]; then
+            echo "tag=all" >> "$GITHUB_OUTPUT"
+          fi
 
-      - name: Build and Push Docker Image (With JAX Solver)
+      - name: Build and push Docker image to Docker Hub (no solvers)
+        if: matrix.build-args == 'No solvers'
         uses: docker/build-push-action@v5
         with:
           context: .
           file: scripts/Dockerfile
-          tags: pybamm/pybamm:jax
+          tags: pybamm/pybamm:${{ steps.tags.outputs.tag }}
           push: true
-          build-args: |
-            JAX=true
+          platforms: linux/amd64, linux/arm64
 
-      - name: Build and Push Docker Image (With ODES & DAE Solver)
+      - name: Build and push Docker image to Docker Hub (with ODES and IDAKLU solvers)
+        if: matrix.build-args == 'ODES' || matrix.build-args == 'IDAKLU'
         uses: docker/build-push-action@v5
         with:
           context: .
           file: scripts/Dockerfile
-          tags: pybamm/pybamm:odes
+          tags: pybamm/pybamm:${{ steps.tags.outputs.tag }}
           push: true
-          build-args: |
-            ODES=true
+          build-args: ${{ matrix.build-args }}=true
+          platforms: linux/amd64, linux/arm64
 
-      - name: Build and Push Docker Image (With IDAKLU Solver)
+      - name: Build and push Docker image to Docker Hub (with ALL and JAX solvers)
+        if: matrix.build-args == 'ALL' || matrix.build-args == 'JAX'
         uses: docker/build-push-action@v5
         with:
           context: .
           file: scripts/Dockerfile
-          tags: pybamm/pybamm:idaklu
+          tags: pybamm/pybamm:${{ steps.tags.outputs.tag }}
           push: true
-          build-args: |
-            IDAKLU=true
+          build-args: ${{ matrix.build-args }}=true
+          # exclude arm64 for JAX and ALL builds for now, see
+          # https://github.com/google/jax/issues/13608
+          platforms: linux/amd64
 
-      - name: Build and Push Docker Image (With All Solvers)
-        uses: docker/build-push-action@v5
-        with:
-          context: .
-          file: scripts/Dockerfile
-          tags: pybamm/pybamm:latest
-          push: true
-          build-args: |
-            ALL=true
+      - name: List built image(s)
+        run: docker images