From 3442369981e50a672519df52880c2f376fec97ae Mon Sep 17 00:00:00 2001
From: Vladislav Kozlov <vkozlov@nvidia.com>
Date: Thu, 16 Nov 2023 11:47:26 -0700
Subject: [PATCH] wip: build t5x and rosetta/pax

---
 .github/workflows/_sandbox.yaml | 56 +++++++++++++++++++--------------
 1 file changed, 32 insertions(+), 24 deletions(-)

diff --git a/.github/workflows/_sandbox.yaml b/.github/workflows/_sandbox.yaml
index b664fbda8..e76c3953c 100644
--- a/.github/workflows/_sandbox.yaml
+++ b/.github/workflows/_sandbox.yaml
@@ -11,42 +11,50 @@ 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
+  # 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]
+    uses: ./.github/workflows/_build_t5x.yaml
     with:
       ARCHITECTURE: amd64
-      BASE_IMAGE: ${{ needs.build-jax.outputs.DOCKER_TAG_MEALKIT }}
+      BASE_IMAGE: ghcr.io/nvidia/jax-toolbox-internal:6894053657-jax-amd64
     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
+  # 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]
+    # needs: [build-pax]
     uses: ./.github/workflows/_build_rosetta.yaml
     with:
       ARCHITECTURE: amd64
       BASE_LIBRARY: pax
-      BASE_IMAGE: ${{ needs.build-pax.outputs.DOCKER_TAG_FINAL }}
+      BASE_IMAGE: ghcr.io/nvidia/jax-toolbox-internal:6894053657-upstream-pax-amd64
 
   # test-pax:
   #   needs: [build-pax]