From f567e2c90b79b3d61a10582059a8cc5b4a466d39 Mon Sep 17 00:00:00 2001
From: Vladislav Kozlov <vkozlov@nvidia.com>
Date: Thu, 16 Nov 2023 11:59:56 -0700
Subject: [PATCH] wip: build all in question

---
 .github/workflows/_build_rosetta.yaml |  8 +--
 .github/workflows/_sandbox.yaml       | 91 ++++++++++++++-------------
 2 files changed, 50 insertions(+), 49 deletions(-)

diff --git a/.github/workflows/_build_rosetta.yaml b/.github/workflows/_build_rosetta.yaml
index b4e0bfa02..816e35ab7 100644
--- a/.github/workflows/_build_rosetta.yaml
+++ b/.github/workflows/_build_rosetta.yaml
@@ -87,7 +87,7 @@ jobs:
         with:
           images: ${{ env.UPLD_IMAGE }}
           flavor: latest=false
-          tags: type=raw,value=${{ github.run_id }}-${{ inputs.BASE_LIBRARY }}-${{ inputs.ARCHITECHTURE }}-mealkit
+          tags: type=raw,value=${{ github.run_id }}-${{ inputs.BASE_LIBRARY }}-${{ inputs.ARCHITECTURE }}-mealkit
           labels: org.opencontainers.image.created=${{ inputs.BUILD_DATE }}
 
       - name: Build docker images - mealkit
@@ -97,7 +97,7 @@ jobs:
           context: rosetta/
           push: true
           file: rosetta/Dockerfile.${{ inputs.BASE_LIBRARY }}
-          platforms: linux/${{ inputs.ARCHITECHTURE }}
+          platforms: linux/${{ inputs.ARCHITECTURE }}
           tags: ${{ steps.mealkit-metadata.outputs.tags }}
           labels: ${{ steps.mealkit-metadata.outputs.labels }}
           target: rosetta
@@ -110,7 +110,7 @@ jobs:
         with:
           images: ${{ env.UPLD_IMAGE }}
           flavor: latest=false
-          tags: type=raw,value=${{ github.run_id }}-${{ inputs.BASE_LIBRARY }}-${{ inputs.ARCHITECHTURE }}-final
+          tags: type=raw,value=${{ github.run_id }}-${{ inputs.BASE_LIBRARY }}-${{ inputs.ARCHITECTURE }}-final
           labels: org.opencontainers.image.created=${{ inputs.BUILD_DATE }}
 
       - name: Build docker images - final
@@ -119,7 +119,7 @@ jobs:
           context: rosetta/
           push: true
           file: rosetta/Dockerfile.${{ inputs.BASE_LIBRARY }}
-          platforms: linux/${{ inputs.ARCHITECHTURE }}
+          platforms: linux/${{ inputs.ARCHITECTURE }}
           tags: ${{ steps.final-metadata.outputs.tags }}
           labels: ${{ steps.final-metadata.outputs.labels }}
           target: rosetta
diff --git a/.github/workflows/_sandbox.yaml b/.github/workflows/_sandbox.yaml
index e76c3953c..af999f836 100644
--- a/.github/workflows/_sandbox.yaml
+++ b/.github/workflows/_sandbox.yaml
@@ -11,63 +11,64 @@ permissions:
 
 jobs:
 
-  # build-base:
-  #   uses: ./.github/workflows/_build_base.yaml
-  #   with:
-  #     ARCHITECTURE: amd64
-  #   secrets: inherit
+  build-base:
+    uses: ./.github/workflows/_build_base.yaml
+    with:
+      ARCHITECTURE: amd64
+    secrets: inherit
 
-  # build-jax:
-  #   needs: [build-base]
-  #   uses: ./.github/workflows/_build_jax.yaml
-  #   with:
-  #     ARCHITECTURE: amd64
-  #     BASE_IMAGE: ${{ needs.build-base.outputs.DOCKER_TAG }}
-  #   secrets: inherit
+  build-jax:
+    needs: [build-base]
+    uses: ./.github/workflows/_build_jax.yaml
+    with:
+      ARCHITECTURE: amd64
+      BASE_IMAGE: ${{ needs.build-base.outputs.DOCKER_TAG }}
+    secrets: inherit
 
-  # build-pax:
-  #   needs: [build-jax]
-  #   uses: ./.github/workflows/_build_pax.yaml
-  #   with:
-  #     ARCHITECTURE: amd64
-  #     BASE_IMAGE: ${{ needs.build-jax.outputs.DOCKER_TAG_MEALKIT }}
-  #   secrets: inherit
+  build-pax:
+    needs: [build-jax]
+    uses: ./.github/workflows/_build_pax.yaml
+    with:
+      ARCHITECTURE: amd64
+      BASE_IMAGE: ${{ needs.build-jax.outputs.DOCKER_TAG_MEALKIT }}
+    secrets: inherit
 
   build-t5x:
-#    needs: [build-jax]
+   needs: [build-jax]
     uses: ./.github/workflows/_build_t5x.yaml
     with:
       ARCHITECTURE: amd64
-      BASE_IMAGE: ghcr.io/nvidia/jax-toolbox-internal:6894053657-jax-amd64
+      BASE_IMAGE: ${{ needs.build-jax.outputs.DOCKER_TAG_MEALKIT }}
+    secrets: inherit
+
+  test-pax:
+    needs: [build-pax]
+    uses: ./.github/workflows/_test_pax.yaml
+    with:
+      PAX_IMAGE: ${{ needs.build-pax.outputs.DOCKER_TAG_FINAL }}
     secrets: inherit
 
-  # test-te:
-  #   needs: [build-pax]
-  #   uses: ./.github/workflows/_test_te.yaml
-  #   with:
-  #     TE_IMAGE: ${{ needs.build-pax.outputs.DOCKER_TAG_FINAL }}
-  #   secrets: inherit
 
-  build-rosetta:
-    # needs: [build-pax]
+  test-te:
+    needs: [build-pax]
+    uses: ./.github/workflows/_test_te.yaml
+    with:
+      TE_IMAGE: ${{ needs.build-pax.outputs.DOCKER_TAG_FINAL }}
+    secrets: inherit
+
+  build-rosetta-pax:
+    needs: [build-pax]
     uses: ./.github/workflows/_build_rosetta.yaml
     with:
       ARCHITECTURE: amd64
       BASE_LIBRARY: pax
-      BASE_IMAGE: ghcr.io/nvidia/jax-toolbox-internal:6894053657-upstream-pax-amd64
-
-  # test-pax:
-  #   needs: [build-pax]
-  #   uses: ./.github/workflows/_test_pax.yaml
-  #   with:
-  #     PAX_IMAGE: ${{ needs.build-pax.outputs.DOCKER_TAG_FINAL }}
-  #   secrets: inherit
+      BASE_IMAGE: ${{ needs.build-pax.outputs.DOCKER_TAG_FINAL }}
 
-  # finalize:
-  #   if: always()
-  #   # TODO: use dynamic matrix to make dependencies self-updating
-  #   needs: [build-jax, build-t5x, build-pax]
-  #   uses: ./.github/workflows/_finalize.yaml
-  #   with:
-  #     PUBLISH_BADGE: false
-  #   secrets: inherit
+  finalize:
+    if: always()
+    # TODO: use dynamic matrix to make dependencies self-updating
+    needs: [build-jax, build-t5x, build-pax]
+    uses: ./.github/workflows/_finalize.yaml
+    with:
+      PUBLISH_BADGE: false
+    secrets: inherit