Updates rosetta CI to build multiarch (Arm64/amd64) for pax

.github/workflows/_build_rosetta.yaml

@@ -17,10 +17,15 @@ on:
         description: "Build date in YYYY-MM-DD format"
         required: false
         default: 'NOT SPECIFIED'
+      PLATFORMS:
+        type: string
+        description: 'JSON list of platforms. Ex: ["amd64"]'
+        required: false
+        default: '["arm64", "amd64"]'
     outputs:
       DOCKER_TAGS:
         description: "Tags of the image built"
-        value: ${{ jobs.build.outputs.DOCKER_TAGS }}
+        value: ${{ jobs.merge.outputs.DOCKER_TAGS }}
 
 env:
   UPLD_IMAGE: ghcr.io/nvidia/jax-toolbox-internal
@@ -34,16 +39,19 @@ permissions:
 jobs:
 
   build:
-    outputs:
-      DOCKER_TAGS: ${{ steps.meta.outputs.tags }}
-    runs-on: [self-hosted, x86, small]
+    strategy:
+      fail-fast: false
+      matrix:
+        PLATFORM: ${{ fromJSON(inputs.PLATFORMS) }}
+    runs-on: [self-hosted, "${{ matrix.PLATFORM }}", small]
     steps:
       - name: Set default BASE_IMAGE
+        id: defaults
         run: |
           if [[ -n "${{ inputs.BASE_IMAGE }}" ]]; then
-            echo "BASE_IMAGE=${{ inputs.BASE_IMAGE }}" >> "$GITHUB_ENV"
+            echo "BASE_IMAGE=${{ inputs.BASE_IMAGE }}" >> "$GITHUB_OUTPUT"
           else
-            echo "BASE_IMAGE=${{ env.DOCKER_REGISTRY }}/upstream-${{ inputs.BASE_LIBRARY }}:latest" >> "$GITHUB_ENV"
+            echo "BASE_IMAGE=${{ env.DOCKER_REGISTRY }}/upstream-${{ inputs.BASE_LIBRARY }}:latest" >> "$GITHUB_OUTPUT"
           fi
 
       - name: Print environment variables
@@ -65,32 +73,86 @@ jobs:
         with:
           images: ${{ env.UPLD_IMAGE }}
           flavor: latest=false
-          tags: type=raw,value=${{ github.run_id }}-${{ inputs.BASE_LIBRARY }}
+          tags: type=raw,value=${{ github.run_id }}-${{ inputs.BASE_LIBRARY }}-${{ matrix.PLATFORM }}
           labels: org.opencontainers.image.created=${{ inputs.BUILD_DATE }}
 
       - name: Set up Docker Buildx
         uses: docker/setup-buildx-action@v2
         with:
           driver-opts: |
-            image=moby/buildkit:v0.10.6
-
-      - name: Pull and validate BASE_IMAGE=${{ env.BASE_IMAGE }} contains BASE_LIBRARY=${{ inputs.BASE_LIBRARY }}
-        shell: bash -x -e {0}
-        run: |
-          PYTHON_PACKAGE=${{ inputs.BASE_LIBRARY }}
-          if [[ $PYTHON_PACKAGE == "pax" ]]; then
-            PYTHON_PACKAGE=paxml
-          fi
-          docker run --rm ${{ env.BASE_IMAGE }} python -c "import importlib,sys; sys.exit(importlib.util.find_spec('$PYTHON_PACKAGE') is None)"
+            image=moby/buildkit:v0.12.1
 
       - name: Build docker images
         uses: docker/build-push-action@v4
         with:
           context: rosetta/
           push: true
           file: rosetta/Dockerfile.${{ inputs.BASE_LIBRARY }}
+          platforms: linux/${{ matrix.PLATFORM }}
           tags: ${{ steps.meta.outputs.tags }}
           labels: ${{ steps.meta.outputs.labels }}
           target: rosetta
           build-args: |
-            BASE_IMAGE=${{ env.BASE_IMAGE }}
+            BASE_IMAGE=${{ steps.defaults.outputs.BASE_IMAGE }}
+
+      # Temporary workaround until the following issues are solved:
+      # https://github.com/orgs/community/discussions/17245
+      # https://github.com/actions/runner/pull/2477
+      # https://github.com/orgs/community/discussions/26639
+      - name: Save image name as text file
+        shell: bash -x -e {0}
+        run: |
+          echo "${{ steps.meta.outputs.tags }}" >> image-name.txt
+
+      - name: Upload image name file as artifact
+        uses: actions/upload-artifact@v3
+        with:
+          name: image-name-${{ inputs.BASE_LIBRARY }}-${{ matrix.PLATFORM }}
+          path: image-name.txt
+
+  merge:
+    runs-on: ubuntu-latest
+    needs: build
+    outputs:
+      DOCKER_TAGS: ${{ steps.meta.outputs.tags }}
+    steps:
+      # TODO: currently downloading all artifacts of the entire workflow
+      # Revise when this request is fulfilled:
+      # https://github.com/actions/download-artifact/issues/214
+      - name: Download image name files into separate folders
+        uses: actions/download-artifact@v3
+
+      - name: Login to GitHub Container Registry
+        uses: docker/login-action@v2
+        with:
+          registry: ghcr.io
+          username: ${{ github.repository_owner }}
+          password: ${{ secrets.GITHUB_TOKEN }}
+
+      - name: Set docker metadata
+        id: meta
+        uses: docker/metadata-action@v4
+        with:
+          images: |
+            ${{ env.UPLD_IMAGE }}
+          flavor: |
+            latest=false
+          tags: |
+            type=raw,value=${{ github.run_id }}-${{ inputs.BASE_LIBRARY }}-multiarch
+          labels:
+            org.opencontainers.image.created=${{ inputs.BUILD_DATE }}
+
+      - name: Combine images into a single multi-arch image
+        shell: bash -x -e {0}
+        run: |
+          docker manifest create ${{ steps.meta.outputs.tags }} $(
+            for IMAGE in $(cat image-name-${{ inputs.BASE_LIBRARY }}-*/image-name.txt); do
+              REPO=$(echo $IMAGE | cut -d: -f1)
+              DIGEST=$(
+                docker manifest inspect $IMAGE |\
+                jq -r '.manifests[] | select(.platform.os == "linux") | .digest'
+              )
+              echo $REPO@${DIGEST}
+            done
+          )
+          docker manifest push ${{ steps.meta.outputs.tags }}

.github/workflows/nightly-rosetta-t5x-build-test.yaml

@@ -61,6 +61,8 @@ jobs:
       BUILD_DATE: ${{ needs.metadata.outputs.BUILD_DATE }}
       BASE_LIBRARY: ${{ needs.metadata.outputs.BASE_LIBRARY }}
       BASE_IMAGE: ${{ needs.metadata.outputs.BASE_IMAGE }}
+      # TODO: Can't build ARM until https://github.com/NVIDIA/JAX-Toolbox/pull/252 is available
+      PLATFORMS: '["amd64"]'
     secrets: inherit
 
   publish-build:

rosetta/Dockerfile.pax

@@ -37,18 +37,4 @@ bash create-distribution.sh \
   -e /opt/praxis-mirror
 rm -rf $(find /opt -name "__pycache__")
 EOF
-
-WORKDIR /opt/rosetta
-RUN <<EOF
-rmdir /opt/*-mirror
-pip install -e /opt/rosetta
-rm -rf ~/.cache/pip/
-EOF
-
-FROM rosetta AS rosetta-devel
 
-WORKDIR /opt/rosetta
-RUN <<EOF
-pip install -e '/opt/rosetta[test,lint]'
-rm -rf ~/.cache/pip/
-EOF

rosetta/Dockerfile.t5x

@@ -40,7 +40,6 @@ EOF
 
 WORKDIR /opt/rosetta
 RUN <<EOF bash -e
-rmdir /opt/*-mirror
 pip install -e /opt/rosetta --extra-index-url https://developer.download.nvidia.com/compute/redist
 rm -rf ~/.cache/pip/
 EOF