diff --git a/.github/workflows/_build_pax.yaml b/.github/workflows/_build_pax.yaml
index 62bc175ad..de2b3dafd 100644
--- a/.github/workflows/_build_pax.yaml
+++ b/.github/workflows/_build_pax.yaml
@@ -163,7 +163,7 @@ jobs:
           # bring in utility functions
           source .github/workflows/scripts/to_json.sh
 
-          badge_label='PAX ${{ inputs.ARCHITECTURE }} build'
+          badge_label='Upstream PAX ${{ inputs.ARCHITECTURE }} build'
           tags="${{ steps.final-metadata.outputs.tags }}"
           digest="${{ steps.final-build.outputs.digest }}"
           outcome="${{ steps.final-build.outcome }}"
@@ -171,11 +171,11 @@ jobs:
           if [[ ${outcome} == "success" ]]; then
             badge_message="pass"
             badge_color=brightgreen
-            summary="PAX build on ${{ inputs.ARCHITECTURE }}: $badge_message"
+            summary="Upstream PAX build on ${{ inputs.ARCHITECTURE }}: $badge_message"
           else
             badge_message="fail"
             badge_color=red
-            summary="PAX build on ${{ inputs.ARCHITECTURE }}: $badge_message"
+            summary="Upstream PAX build on ${{ inputs.ARCHITECTURE }}: $badge_message"
           fi
 
           to_json \
diff --git a/.github/workflows/_build_rosetta.yaml b/.github/workflows/_build_rosetta.yaml
index f30a0d588..971f60c47 100644
--- a/.github/workflows/_build_rosetta.yaml
+++ b/.github/workflows/_build_rosetta.yaml
@@ -3,6 +3,10 @@ name: ~build Rosetta container
 on:
   workflow_call:
     inputs:
+      ARCHITECTURE:
+        type: string
+        description: 'CPU architecture to build the image for, e.g. amd64, arm64'
+        required: true
       BASE_LIBRARY:
         type: string
         description: 'Choice of base library to build on:'
@@ -17,15 +21,23 @@ on:
         description: "Build date in YYYY-MM-DD format"
         required: false
         default: 'NOT SPECIFIED'
-      PLATFORMS:
+      ARTIFACT_NAME:
+        type: string
+        description: 'Name of the artifact zip file'
+        required: false
+        default: 'artifact-rosetta-build'
+      BADGE_FILENAME:
         type: string
-        description: 'JSON list of platforms. Ex: ["amd64"]'
+        description: 'Name of the endpoint JSON file for shields.io badge (w/o .json || arch || library)'
         required: false
-        default: '["arm64", "amd64"]'
+        default: 'badge-rosetta-build'
     outputs:
-      DOCKER_TAGS:
-        description: "Tags of the image built"
-        value: ${{ jobs.merge.outputs.DOCKER_TAGS }}
+      DOCKER_TAG_MEALKIT:
+        description: "Tags of the 'mealkit' image built"
+        value: ${{ jobs.build-rosetta.outputs.DOCKER_TAG_MEALKIT }}
+      DOCKER_TAG_FINAL:
+        description: "Tags of the complete image built"
+        value: ${{ jobs.build-rosetta.outputs.DOCKER_TAG_FINAL }}
 
 env:
   UPLD_IMAGE: ghcr.io/nvidia/jax-toolbox-internal
@@ -38,12 +50,14 @@ permissions:
 
 jobs:
 
-  build:
-    strategy:
-      fail-fast: false
-      matrix:
-        PLATFORM: ${{ fromJSON(inputs.PLATFORMS) }}
-    runs-on: [self-hosted, "${{ matrix.PLATFORM }}", small]
+  build-rosetta:
+    runs-on: [self-hosted, "${{ inputs.ARCHITECTURE }}", small]
+    env:
+      BADGE_FILENAME_FULL: ${{ inputs.BADGE_FILENAME }}-${{ inputs.BASE_LIBRARY }}-${{ inputs.ARCHITECTURE }}.json
+      ARTIFACT_NAME_FULL: ${{ inputs.ARTIFACT_NAME}}-${{ inputs.BASE_LIBRARY }}-${{ inputs.ARCHITECTURE }}
+    outputs:
+      DOCKER_TAG_MEALKIT: ${{ steps.mealkit-metadata.outputs.tags }}
+      DOCKER_TAG_FINAL: ${{ steps.final-metadata.outputs.tags }}
     steps:
       - name: Set default BASE_IMAGE
         id: defaults
@@ -67,101 +81,112 @@ jobs:
           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 }}-${{ 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.12.1
 
-      - name: Build docker images
+      - name: Set docker metadata - mealkit
+        id: mealkit-metadata
+        uses: docker/metadata-action@v4
+        with:
+          images: |
+            ${{ env.UPLD_IMAGE }}
+          flavor: |
+            latest=false
+          tags: |
+            type=raw,value=${{ github.run_id }}-${{ inputs.BASE_LIBRARY }}-${{ inputs.ARCHITECTURE }}-mealkit
+          labels:
+            org.opencontainers.image.created=${{ inputs.BUILD_DATE }}
+
+      - name: Build mealkit image
+        uses: docker/build-push-action@v4
+        with:
+          context: rosetta/
+          push: true
+          file: rosetta/Dockerfile.${{ inputs.BASE_LIBRARY }}
+          platforms: linux/${{ inputs.ARCHITECTURE }}
+          tags: ${{ steps.mealkit-metadata.outputs.tags }}
+          labels: ${{ steps.mealkit-metadata.outputs.labels }}
+          target: mealkit
+          build-args: |
+            BASE_IMAGE=${{ steps.defaults.outputs.BASE_IMAGE }}
+
+      - name: Set docker metadata - final
+        id: final-metadata
+        uses: docker/metadata-action@v4
+        with:
+          images: |
+            ${{ env.UPLD_IMAGE }}
+          flavor: |
+            latest=false
+          tags: |
+            type=raw,value=${{ github.run_id }}-${{ inputs.BASE_LIBRARY }}-${{ inputs.ARCHITECTURE }}
+          labels:
+            org.opencontainers.image.created=${{ inputs.BUILD_DATE }}
+
+      - name: Build final image
+        id: final-build
         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 }}
+          platforms: linux/${{ inputs.ARCHITECTURE }}
+          tags: ${{ steps.final-metadata.outputs.tags }}
+          labels: ${{ steps.final-metadata.outputs.labels }}
           target: final
           build-args: |
             BASE_IMAGE=${{ steps.defaults.outputs.BASE_IMAGE }}
 
       - name: Extract patches
-        run: rosetta/scripts/extract-patches.sh ${{ steps.meta.outputs.tags }}
+        run: rosetta/scripts/extract-patches.sh ${{ steps.final-metadata.outputs.tags }}
 
       - name: Archive generated patches
         uses: actions/upload-artifact@v3
         with:
-          name: patches-${{ inputs.BASE_LIBRARY }}-${{ github.run_id }}-${{ inputs.BUILD_DATE }}-${{ matrix.PLATFORM }}
+          name: patches-${{ inputs.BASE_LIBRARY }}-${{ github.run_id }}-${{ inputs.BUILD_DATE }}-${{ inputs.ARCHITECTURE }}
           path: rosetta/patches
 
-      # 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
+      - name: Generate sitrep
+        if: success() || failure()
         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
+          # bring in utility functions
+          source .github/workflows/scripts/to_json.sh
+
+          badge_label='${{ inputs.BASE_LIBRARY }} ${{ inputs.ARCHITECTURE }} build'
+          tags="${{ steps.final-metadata.outputs.tags }}"
+          digest="${{ steps.final-build.outputs.digest }}"
+          outcome="${{ steps.final-build.outcome }}"
+
+          if [[ ${outcome} == "success" ]]; then
+            badge_message="pass"
+            badge_color=brightgreen
+            summary="${{ inputs.BASE_LIBRARY }} build on ${{ inputs.ARCHITECTURE }}: $badge_message"
+          else
+            badge_message="fail"
+            badge_color=red
+            summary="${{ inputs.BASE_LIBRARY }} build on ${{ inputs.ARCHITECTURE }}: $badge_message"
+          fi
 
-  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
+          to_json \
+            summary \
+            badge_label tags digest outcome \
+          > sitrep.json
 
-      - name: Login to GitHub Container Registry
-        uses: docker/login-action@v2
-        with:
-          registry: ghcr.io
-          username: ${{ github.repository_owner }}
-          password: ${{ secrets.GITHUB_TOKEN }}
+          schemaVersion=1 \
+          label="${badge_label}" \
+          message="${badge_message}" \
+          color="${badge_color}" \
+          to_json schemaVersion label message color \
+          > ${{ env.BADGE_FILENAME_FULL }}
 
-      - name: Set docker metadata
-        id: meta
-        uses: docker/metadata-action@v4
+      - name: Upload sitrep and badge
+        uses: actions/upload-artifact@v3
         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 }}
+          name: ${{ env.ARTIFACT_NAME_FULL }}
+          path: |
+            sitrep.json
+            ${{ env.BADGE_FILENAME_FULL }}
diff --git a/.github/workflows/_build_t5x.yaml b/.github/workflows/_build_t5x.yaml
index 606840b71..2a6b9966e 100644
--- a/.github/workflows/_build_t5x.yaml
+++ b/.github/workflows/_build_t5x.yaml
@@ -162,7 +162,7 @@ jobs:
           # bring in utility functions
           source .github/workflows/scripts/to_json.sh
 
-          badge_label='T5X ${{ inputs.ARCHITECTURE }} build'
+          badge_label='Upstream T5X ${{ inputs.ARCHITECTURE }} build'
           tags="${{ steps.final-metadata.outputs.tags }}"
           digest="${{ steps.final-build.outputs.digest }}"
           outcome="${{ steps.final-build.outcome }}"
@@ -170,11 +170,11 @@ jobs:
           if [[ ${outcome} == "success" ]]; then
             badge_message="pass"
             badge_color=brightgreen
-            summary="T5X build on ${{ inputs.ARCHITECTURE }}: $badge_message"
+            summary="Upstream T5X build on ${{ inputs.ARCHITECTURE }}: $badge_message"
           else
             badge_message="fail"
             badge_color=red
-            summary="T5X build on ${{ inputs.ARCHITECTURE }}: $badge_message"
+            summary="Upstream T5X build on ${{ inputs.ARCHITECTURE }}: $badge_message"
           fi
 
           to_json \
diff --git a/.github/workflows/_ci.yaml b/.github/workflows/_ci.yaml
index 1447f5076..932631059 100644
--- a/.github/workflows/_ci.yaml
+++ b/.github/workflows/_ci.yaml
@@ -139,28 +139,31 @@ jobs:
       REF_PRAXIS: ${{ needs.metadata.outputs.REF_PRAXIS }}
     secrets: inherit
 
-  # build-rosetta-t5x:
-  #   uses: ./.github/workflows/_build_rosetta.yaml
-  #   needs: [metadata, build-t5x]
-  #   with:
-  #     BUILD_DATE: ${{ needs.metadata.outputs.BUILD_DATE }}
-  #     BASE_IMAGE: ${{ needs.build-t5x.outputs.DOCKER_TAG_MEALKIT }}
-  #     BASE_LIBRARY: t5x
-  #   secrets: inherit
-
-  # build-rosetta-pax:
-  #   uses: ./.github/workflows/_build_rosetta.yaml
-  #   needs: [metadata, build-pax]
-  #   with:
-  #     BUILD_DATE: ${{ needs.metadata.outputs.BUILD_DATE }}
-  #     BASE_IMAGE: ${{ needs.build-pax.outputs.DOCKER_TAG_MEALKIT }}
-  #     BASE_LIBRARY: pax
-  #   secrets: inherit
-
-  # test-distribution:
-  #   needs: metadata
-  #   uses: ./.github/workflows/_test_distribution.yaml
-  #   secrets: inherit
+  build-rosetta-t5x:
+    uses: ./.github/workflows/_build_rosetta.yaml
+    if: inputs.ARCHITECTURE == 'amd64'  # T5X arm64 build is wip in PR 252
+    needs: [metadata, build-t5x]
+    with:
+      ARCHITECTURE: amd64
+      BUILD_DATE: ${{ needs.metadata.outputs.BUILD_DATE }}
+      BASE_IMAGE: ${{ needs.build-t5x.outputs.DOCKER_TAG_MEALKIT }}
+      BASE_LIBRARY: t5x
+    secrets: inherit
+
+  build-rosetta-pax:
+    uses: ./.github/workflows/_build_rosetta.yaml
+    needs: [metadata, build-pax]
+    with:
+      ARCHITECTURE: ${{ inputs.ARCHITECTURE }}
+      BUILD_DATE: ${{ needs.metadata.outputs.BUILD_DATE }}
+      BASE_IMAGE: ${{ needs.build-pax.outputs.DOCKER_TAG_MEALKIT }}
+      BASE_LIBRARY: pax
+    secrets: inherit
+
+  test-distribution:
+    needs: metadata
+    uses: ./.github/workflows/_test_distribution.yaml
+    secrets: inherit
 
   test-jax:
     needs: build-jax
@@ -194,10 +197,4 @@ jobs:
       PAX_IMAGE: ${{ needs.build-pax.outputs.DOCKER_TAG_FINAL }}
     secrets: inherit
 
-  # test-vit:
-  #   needs: build-rosetta-t5x
-  #   uses: ./.github/workflows/_test_vit.yaml
-  #   with:
-  #     ROSETTA_T5X_IMAGE: ${{ needs.build-rosetta-t5x.outputs.DOCKER_TAGS }}
-  #   secrets: inherit 
  
diff --git a/.github/workflows/nightly-rosetta-pax-build.yaml b/.github/workflows/nightly-rosetta-pax-build.yaml
index c12cfd8f0..635d122b6 100644
--- a/.github/workflows/nightly-rosetta-pax-build.yaml
+++ b/.github/workflows/nightly-rosetta-pax-build.yaml
@@ -39,6 +39,16 @@ jobs:
       BASE_IMAGE: ${{ steps.meta-vars.outputs.BASE_IMAGE }}
       PUBLISH: ${{ steps.meta-vars.outputs.PUBLISH }}
     steps:
+      - name: Check if the triggering workflow failed
+        id: if-upstream-failed
+        shell: bash -x -e {0}
+        run: |
+          echo "UPSTREAM_FAILED=${{ github.event.workflow_run.conclusion != 'success' }}" >> $GITHUB_OUTPUT
+
+      - name: Cancel workflow if upstream workflow did not success
+        if: ${{ steps.if-upstream-failed.outputs.UPSTREAM_FAILED == 'true' }}
+        uses: styfle/cancel-workflow-action@0.12.0
+
       - name: Set build metadata
         id: meta-vars
         shell: bash -x -e {0}
@@ -52,99 +62,125 @@ jobs:
           echo "BUILD_DATE=${BUILD_DATE}" >> $GITHUB_OUTPUT
           echo "BASE_LIBRARY=${{ env.BASE_LIBRARY }}" >> $GITHUB_OUTPUT
           echo "BASE_IMAGE=${BASE_IMAGE}" >> $GITHUB_OUTPUT
-          echo "PUBLISH=${{ inputs.PUBLISH }}" >> $GITHUB_OUTPUT
+          # A container should be published if:
+          # 1) the workflow is triggered by workflow_dispatch and the PUBLISH input is true, or
+          # 2) the workflow is triggered by workflow_run (i.e., a nightly build)
+          echo "PUBLISH=${{ github.event_name == 'workflow_run' || (github.event_name == 'workflow_dispatch' && inputs.PUBLISH) }}" >> $GITHUB_OUTPUT
 
-  build:
-    if: (github.event_name == 'workflow_run' && github.event.workflow_run.conclusion == 'success') || github.event_name == 'workflow_dispatch'
+  amd64:
     needs: metadata
     uses: ./.github/workflows/_build_rosetta.yaml
     with:
+      ARCHITECTURE: amd64
       BUILD_DATE: ${{ needs.metadata.outputs.BUILD_DATE }}
       BASE_LIBRARY: ${{ needs.metadata.outputs.BASE_LIBRARY }}
       BASE_IMAGE: ${{ needs.metadata.outputs.BASE_IMAGE }}
     secrets: inherit
 
-  publish-build:
-    needs: [metadata, build]
+  arm64:
+    needs: metadata
+    uses: ./.github/workflows/_build_rosetta.yaml
+    with:
+      ARCHITECTURE: arm64
+      BUILD_DATE: ${{ needs.metadata.outputs.BUILD_DATE }}
+      BASE_LIBRARY: ${{ needs.metadata.outputs.BASE_LIBRARY }}
+      BASE_IMAGE: ${{ needs.metadata.outputs.BASE_IMAGE }}
+    secrets: inherit
+
+  publish-build-badge:
+    needs: [metadata, amd64, arm64]
     uses: ./.github/workflows/_publish_badge.yaml
-    if: ( always() )
+    if: always()
     secrets: inherit
     with:
       ENDPOINT_FILENAME: 'rosetta-pax-build-status.json'
-      PUBLISH: ${{ github.event_name == 'workflow_run' || needs.metadata.outputs.PUBLISH == 'true' }}
+      PUBLISH: ${{ needs.metadata.outputs.PUBLISH == 'true' }}
       SCRIPT: |
-        if [[ ${{ needs.build.result }} == "success" ]]; then
+        if [[ ${{ needs.amd64.result }} == "success" && ${{ needs.arm64.result }} == "success" ]]; then
           BADGE_COLOR=brightgreen
           MSG=passing
+          STATUS=success
         else
           BADGE_COLOR=red
           MSG=failing
+          STATUS=failure
         fi
         echo "LABEL='nightly'" >> $GITHUB_OUTPUT
         echo "MESSAGE='${MSG}'" >> $GITHUB_OUTPUT
         echo "COLOR='${BADGE_COLOR}'" >> $GITHUB_OUTPUT
+        echo "STATUS='${STATUS}'" >> ${GITHUB_OUTPUT}
+  
+  publish-mealkit:
+    needs: [metadata, amd64, arm64]
+    if: needs.metadata.outputs.PUBLISH == 'true'
+    uses: ./.github/workflows/_publish_container.yaml
+    with:
+      SOURCE_IMAGE: |
+        ${{ needs.amd64.outputs.DOCKER_TAG_MEALKIT }}
+        ${{ needs.arm64.outputs.DOCKER_TAG_MEALKIT }}
+      TARGET_IMAGE: pax
+      TARGET_TAGS: |
+        type=raw,value=mealkit,priority=500
+        type=raw,value=mealkit-${{ needs.metadata.outputs.BUILD_DATE }},priority=500
 
-  test-pax:
-    needs: build
+  # TODO: Test ARM when runners available
+  test-amd64:
+    needs: amd64
     uses: ./.github/workflows/_test_pax_rosetta.yaml
-    if: (github.event_name == 'workflow_run' && github.event.workflow_run.conclusion == 'success') || github.event_name == 'workflow_dispatch'
     with:
-      PAX_IMAGE: ${{ needs.build.outputs.DOCKER_TAGS }}
+      PAX_IMAGE: ${{ needs.amd64.outputs.DOCKER_TAG_FINAL }}
     secrets: inherit
 
-  publish-test:
-    needs: [metadata, build, test-pax]
+  # TODO: ARM Tests
+  publish-final:
+    needs: [metadata, amd64, arm64, test-amd64]
+    if: needs.metadata.outputs.PUBLISH == 'true'
+    uses: ./.github/workflows/_publish_container.yaml
+    with:
+      SOURCE_IMAGE: |
+        ${{ needs.amd64.outputs.DOCKER_TAG_FINAL }}
+        ${{ needs.arm64.outputs.DOCKER_TAG_FINAL }}
+      TARGET_IMAGE: pax
+      TARGET_TAGS: |
+        ${{ needs.test-amd64.outputs.TEST_STATUS == 'success' && 'type=raw,value=latest,priority=1000' || '' }}
+        type=raw,value=nightly-${{ needs.metadata.outputs.BUILD_DATE }},priority=900
+
+  publish-test-badge:
+    needs: [metadata, publish-build-badge, test-amd64]
     uses: ./.github/workflows/_publish_badge.yaml
-    if: ( always() )
+    if: always()
     secrets: inherit
     with:
       ENDPOINT_FILENAME: 'rosetta-pax-overall-test-status.json'
-      PUBLISH: ${{ github.event_name == 'workflow_run' || needs.metadata.outputs.PUBLISH == 'true' }}
+      PUBLISH: ${{ needs.metadata.outputs.PUBLISH == 'true' }}
       SCRIPT: |
-        PAX_STATUS=${{ needs.test-pax.outputs.TEST_STATUS }}
+        PAX_STATUS=${{ needs.test-amd64.outputs.TEST_STATUS }}
 
         echo "LABEL='Tests'" >> $GITHUB_OUTPUT
 
-        if [[ ${{ needs.build.result }} == "success" ]]; then
+        STATUS=failure
+        if [[ ${{ needs.publish-build-badge.outputs.STATUS }} == "success" ]]; then
           if [[ $PAX_STATUS == "success" ]]; then
             COLOR=brightgreen
             MESSAGE="MGMN passed"
+            STATUS=success
           else
             COLOR=red
             MESSAGE="MGMN failed"
           fi
         else
-          MESSAGE="n/a"
           COLOR="red"
+          MESSAGE="n/a"
         fi
 
-        echo "MESSAGE='${MESSAGE}'" >> $GITHUB_OUTPUT
         echo "COLOR='${COLOR}'" >> $GITHUB_OUTPUT
+        echo "MESSAGE='${MESSAGE}'" >> $GITHUB_OUTPUT
+        echo "STATUS='${STATUS}'" >> ${GITHUB_OUTPUT}
 
-  publish-latest-container:
-    needs: [metadata, build, test-pax]
-    if: ( ${{ needs.test-pax.outputs.TEST_STATUS == 'success' }} ) && ((github.event_name == 'workflow_run' && github.event.workflow_run.conclusion == 'success') || (github.event_name == 'workflow_dispatch' && inputs.PUBLISH))
-    uses: ./.github/workflows/_publish_container.yaml
-    secrets: inherit
-    with:
-      SOURCE_IMAGE: ${{ needs.build.outputs.DOCKER_TAGS }}
-      TARGET_IMAGE: pax
-      TARGET_TAGS: |
-        type=raw,value=latest,priority=1000
-
-  publish-container:
-    if: (github.event_name == 'workflow_run' && github.event.workflow_run.conclusion == 'success') || (github.event_name == 'workflow_dispatch' && inputs.PUBLISH)
-    needs: [metadata, build]
-    uses: ./.github/workflows/_publish_container.yaml
-    secrets: inherit
+  finalize:
+    if: always()
+    needs: [metadata, amd64, arm64, test-amd64]
+    uses: ./.github/workflows/_finalize.yaml
     with:
-      SOURCE_IMAGE: ${{ needs.build.outputs.DOCKER_TAGS }}
-      TARGET_IMAGE: pax
-      TARGET_TAGS: |
-        type=raw,value=nightly-${{ needs.metadata.outputs.BUILD_DATE }},priority=900
-
-  if-upstream-failed:
-    runs-on: ubuntu-latest
-    if: (github.event_name == 'workflow_run' && github.event.workflow_run.conclusion == 'failure') && github.event_name != 'workflow_dispatch'
-    steps:
-      - run: echo 'Upstream workflow failed, aborting run' && exit 1
+      PUBLISH_BADGE: ${{ needs.metadata.outputs.PUBLISH == 'true' }}
+    secrets: inherit
\ No newline at end of file
diff --git a/.github/workflows/nightly-rosetta-t5x-build-test.yaml b/.github/workflows/nightly-rosetta-t5x-build-test.yaml
index 360f8f586..b8d34d690 100644
--- a/.github/workflows/nightly-rosetta-t5x-build-test.yaml
+++ b/.github/workflows/nightly-rosetta-t5x-build-test.yaml
@@ -39,6 +39,16 @@ jobs:
       BASE_IMAGE: ${{ steps.meta-vars.outputs.BASE_IMAGE }}
       PUBLISH: ${{ steps.meta-vars.outputs.PUBLISH }}
     steps:
+      - name: Check if the triggering workflow failed
+        id: if-upstream-failed
+        shell: bash -x -e {0}
+        run: |
+          echo "UPSTREAM_FAILED=${{ github.event.workflow_run.conclusion != 'success' }}" >> $GITHUB_OUTPUT
+
+      - name: Cancel workflow if upstream workflow did not success
+        if: ${{ steps.if-upstream-failed.outputs.UPSTREAM_FAILED == 'true' }}
+        uses: styfle/cancel-workflow-action@0.12.0
+
       - name: Set build metadata
         id: meta-vars
         shell: bash -x -e {0}
@@ -52,92 +62,134 @@ jobs:
           echo "BUILD_DATE=${BUILD_DATE}" >> $GITHUB_OUTPUT
           echo "BASE_LIBRARY=${{ env.BASE_LIBRARY }}" >> $GITHUB_OUTPUT
           echo "BASE_IMAGE=${BASE_IMAGE}" >> $GITHUB_OUTPUT
-          echo "PUBLISH=${{ inputs.PUBLISH }}" >> $GITHUB_OUTPUT
+          # A container should be published if:
+          # 1) the workflow is triggered by workflow_dispatch and the PUBLISH input is true, or
+          # 2) the workflow is triggered by workflow_run (i.e., a nightly build)
+          echo "PUBLISH=${{ github.event_name == 'workflow_run' || (github.event_name == 'workflow_dispatch' && inputs.PUBLISH) }}" >> $GITHUB_OUTPUT
 
-  build:
-    if: (github.event_name == 'workflow_run' && github.event.workflow_run.conclusion == 'success') || github.event_name == 'workflow_dispatch'
+  amd64:
     needs: metadata
     uses: ./.github/workflows/_build_rosetta.yaml
     with:
+      ARCHITECTURE: amd64
       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:
-    needs: [metadata, build]
+  # TODO: Can't build ARM until https://github.com/NVIDIA/JAX-Toolbox/pull/252 is available
+  arm64:
+    needs: metadata
+    runs-on: ubuntu-22.04
+    outputs:
+      DOCKER_TAG_FINAL: ''
+      DOCKER_TAG_MEALKIT: ''
+    steps:
+      - name: Generate placeholder warning
+        shell: bash -x -e {0}
+        run: |
+          echo "WARNING: arm64 build is not yet supported"
+
+  # TODO: ARM
+  publish-build-badge:
+    needs: [metadata, amd64, arm64]
     uses: ./.github/workflows/_publish_badge.yaml
-    if: ( always() )
+    if: always()
     secrets: inherit
     with:
       ENDPOINT_FILENAME: 'rosetta-t5x-build-status.json'
-      PUBLISH: ${{ github.event_name == 'workflow_run' || needs.metadata.outputs.PUBLISH == 'true' }}
+      PUBLISH: ${{ needs.metadata.outputs.PUBLISH == 'true' }}
       SCRIPT: |
-        if [[ ${{ needs.build.result }} == "success" ]]; then
+        if [[ ${{ needs.amd64.result }} == "success" && ${{ needs.arm64.result }} == "success" ]]; then
           BADGE_COLOR=brightgreen
           MSG=passing
+          STATUS=success
         else
           BADGE_COLOR=red
           MSG=failing
+          STATUS=failure
         fi
         echo "LABEL='nightly'" >> $GITHUB_OUTPUT
         echo "MESSAGE='${MSG}'" >> $GITHUB_OUTPUT
         echo "COLOR='${BADGE_COLOR}'" >> $GITHUB_OUTPUT
+        echo "STATUS='${STATUS}'" >> ${GITHUB_OUTPUT}
+  
+  publish-mealkit:
+    needs: [metadata, amd64, arm64]
+    if: needs.metadata.outputs.PUBLISH == 'true'
+    uses: ./.github/workflows/_publish_container.yaml
+    with:
+      SOURCE_IMAGE: |
+        ${{ needs.amd64.outputs.DOCKER_TAG_MEALKIT }}
+        ${{ needs.arm64.outputs.DOCKER_TAG_MEALKIT }}
+      TARGET_IMAGE: t5x
+      TARGET_TAGS: |
+        type=raw,value=mealkit,priority=500
+        type=raw,value=mealkit-${{ needs.metadata.outputs.BUILD_DATE }},priority=500
 
-  test-unit:
-    if: (github.event_name == 'workflow_run' && github.event.workflow_run.conclusion == 'success') || github.event_name == 'workflow_dispatch'
-    needs: build
+  test-unit-amd64:
+    needs: amd64
     uses: ./.github/workflows/_test_rosetta.yaml
     with:
-      ROSETTA_IMAGE: ${{ needs.build.outputs.DOCKER_TAGS }}
+      ROSETTA_IMAGE: ${{ needs.amd64.outputs.DOCKER_TAG_FINAL }}
     secrets: inherit
 
-  test-t5x:
-    needs: build
+  test-t5x-amd64:
+    needs: amd64
     uses: ./.github/workflows/_test_t5x_rosetta.yaml
-    if: (github.event_name == 'workflow_run' && github.event.workflow_run.conclusion == 'success') || github.event_name == 'workflow_dispatch'
     with:
-      T5X_IMAGE: ${{ needs.build.outputs.DOCKER_TAGS }}
+      T5X_IMAGE: ${{ needs.amd64.outputs.DOCKER_TAG_FINAL }}
     secrets: inherit
 
-  test-vit:
-    needs: build
+  test-vit-amd64:
+    needs: amd64
     uses: ./.github/workflows/_test_vit.yaml
-    if: (github.event_name == 'workflow_run' && github.event.workflow_run.conclusion == 'success') || github.event_name == 'workflow_dispatch'
     with:
-      ROSETTA_T5X_IMAGE: ${{ needs.build.outputs.DOCKER_TAGS }}
+      ROSETTA_T5X_IMAGE: ${{ needs.amd64.outputs.DOCKER_TAG_FINAL }}
     secrets: inherit
 
+  publish-final:
+    needs: [metadata, amd64, arm64, test-t5x-amd64, test-vit-amd64, test-unit-amd64]
+    if: needs.metadata.outputs.PUBLISH == 'true'
+    uses: ./.github/workflows/_publish_container.yaml
+    with:
+      SOURCE_IMAGE: |
+        ${{ needs.amd64.outputs.DOCKER_TAG_FINAL }}
+        ${{ needs.arm64.outputs.DOCKER_TAG_FINAL }}
+      TARGET_IMAGE: t5x
+      TARGET_TAGS: |
+        ${{ ( needs.test-t5x-amd64.outputs.TEST_STATUS == 'success' && needs.test-vit-amd64.outputs.TEST_STATUS == 'success' && needs.test-unit-amd64.outputs.TEST_STATUS == 'success' ) && 'type=raw,value=latest,priority=1000' || '' }}
+        type=raw,value=nightly-${{ needs.metadata.outputs.BUILD_DATE }},priority=900
+
   publish-t5x:
-    needs: [metadata, test-t5x, test-vit]
+    needs: [metadata, test-t5x-amd64, test-vit-amd64]
     uses: ./.github/workflows/_publish_t5x_results.yaml
-    if: (github.event_name == 'workflow_run' && github.event.workflow_run.conclusion == 'success') || github.event_name == 'workflow_dispatch'
     with:
       BUILD_DATE: ${{ needs.metadata.outputs.BUILD_DATE }}
       EXPERIMENT_SUBDIR: ROSETTA_T5X
     secrets: inherit
 
-  publish-test:
-    needs: [metadata, build, test-unit, test-t5x, test-vit]
+  publish-test-badge:
+    needs: [metadata, publish-build-badge, test-unit-amd64, test-t5x-amd64, test-vit-amd64]
     uses: ./.github/workflows/_publish_badge.yaml
-    if: ( always() )
+    if: always()
     secrets: inherit
     with:
       ENDPOINT_FILENAME: 'rosetta-t5x-overall-test-status.json'
-      PUBLISH: ${{ github.event_name == 'workflow_run' || needs.metadata.outputs.PUBLISH == 'true' }}
+      PUBLISH: ${{ needs.metadata.outputs.PUBLISH == 'true' }}
       SCRIPT: |
-        UNIT_STATUS=${{ needs.test-unit.outputs.TEST_STATUS }}
-        T5X_STATUS=${{ needs.test-t5x.outputs.TEST_STATUS }}
-        VIT_STATUS=${{ needs.test-vit.outputs.TEST_STATUS }}
+        UNIT_STATUS=${{ needs.test-unit-amd64.outputs.TEST_STATUS }}
+        T5X_STATUS=${{ needs.test-t5x-amd64.outputs.TEST_STATUS }}
+        VIT_STATUS=${{ needs.test-vit-amd64.outputs.TEST_STATUS }}
 
         echo "LABEL='Tests'" >> $GITHUB_OUTPUT
 
-        if [[ ${{ needs.build.result }} == "success" ]]; then
+        STATUS=failure
+        if [[ ${{ needs.publish-build-badge.outputs.STATUS }} == "success" ]]; then
           if [[ $UNIT_STATUS == "success" ]] && [[ $T5X_STATUS == "success" ]] && [[ $VIT_STATUS == "success" ]]; then
             COLOR=brightgreen
             MESSAGE="Unit passed / MGMN passed"
+            STATUS=success
           elif [[ $UNIT_STATUS == "success" ]]; then 
             COLOR=yellow
             MESSAGE="Unit passed / MGMN failed"
@@ -149,37 +201,18 @@ jobs:
             MESSAGE="Unit failed / MGMN failed"
           fi
         else
-          MESSAGE="n/a"
           COLOR="red"
+          MESSAGE="n/a"
         fi
 
         echo "MESSAGE='${MESSAGE}'" >> $GITHUB_OUTPUT
         echo "COLOR='${COLOR}'" >> $GITHUB_OUTPUT
+        echo "STATUS='${STATUS}'" >> ${GITHUB_OUTPUT}
 
-  publish-latest-container:
-    needs: [metadata, build, test-t5x, test-unit, test-vit]
-    if: ( needs.test-unit.outputs.TEST_STATUS == 'success' && needs.test-t5x.outputs.TEST_STATUS == 'success' && needs.test-vit.outputs.TEST_STATUS == 'success' ) && ((github.event_name == 'workflow_run' && github.event.workflow_run.conclusion == 'success') || (github.event_name == 'workflow_dispatch' && inputs.PUBLISH))
-    uses: ./.github/workflows/_publish_container.yaml
-    secrets: inherit
-    with:
-      SOURCE_IMAGE: ${{ needs.build.outputs.DOCKER_TAGS }}
-      TARGET_IMAGE: t5x
-      TARGET_TAGS: |
-        type=raw,value=latest,priority=1000
-
-  publish-container:
-    if: (github.event_name == 'workflow_run' && github.event.workflow_run.conclusion == 'success') || (github.event_name == 'workflow_dispatch' && inputs.PUBLISH)
-    needs: [metadata, build]
-    uses: ./.github/workflows/_publish_container.yaml
-    secrets: inherit
+  finalize:
+    if: always()
+    needs: [metadata, amd64, arm64, test-t5x-amd64, test-vit-amd64, test-unit-amd64]
+    uses: ./.github/workflows/_finalize.yaml
     with:
-      SOURCE_IMAGE: ${{ needs.build.outputs.DOCKER_TAGS }}
-      TARGET_IMAGE: t5x
-      TARGET_TAGS: |
-        type=raw,value=nightly-${{ needs.metadata.outputs.BUILD_DATE }},priority=900
-
-  if-upstream-failed:
-    runs-on: ubuntu-latest
-    if: (github.event_name == 'workflow_run' && github.event.workflow_run.conclusion == 'failure') && github.event_name != 'workflow_dispatch'
-    steps:
-      - run: echo 'Upstream workflow failed, aborting run' && exit 1
+      PUBLISH_BADGE: ${{ needs.metadata.outputs.PUBLISH == 'true' }}
+    secrets: inherit
\ No newline at end of file