Update _sandbox.yaml

NVIDIA · Nov 17, 2023 · d04c2d1 · d04c2d1
1 parent 6bd14fb
commit d04c2d1
Showing 1 changed file with 12 additions and 13 deletions.
diff --git a/.github/workflows/_sandbox.yaml b/.github/workflows/_sandbox.yaml
@@ -116,19 +116,19 @@ jobs:
       REF_XLA: ${{ needs.metadata.outputs.REF_XLA }}
     secrets: inherit
 
-  build-t5x:
+  build-te:
     needs: [metadata, build-jax]
-    uses: ./.github/workflows/_build_t5x.yaml
+    uses: ./.github/workflows/_build_te.yaml
     with:
       BUILD_DATE: ${{ needs.metadata.outputs.BUILD_DATE }}
       BASE_IMAGE: ${{ needs.build-jax.outputs.DOCKER_TAGS }}
-      REPO_T5X: ${{ needs.metadata.outputs.REPO_T5X }}
-      REF_T5X: ${{ needs.metadata.outputs.REF_T5X }}
+      REPO_TE: ${{ needs.metadata.outputs.REPO_TE }}
+      REF_TE: ${{ needs.metadata.outputs.REF_TE }}
     secrets: inherit
 
 
   build-summary:
-    needs: [build-base, build-jax, build-t5x]
+    needs: [build-base, build-jax, build-te]
     # needs: [build-base, build-jax, build-te, build-t5x, build-pax, build-pax-aarch64, build-rosetta-t5x, build-rosetta-pax]
     if: always()
     runs-on: ubuntu-22.04
@@ -143,12 +143,12 @@ jobs:
           | ------------ | -------------------------------------------------- |
           | Base         | ${{ needs.build-base.outputs.DOCKER_TAGS }}        |
           | JAX          | ${{ needs.build-jax.outputs.DOCKER_TAGS }}         |
-          | T5X          | ${{ needs.build-t5x.outputs.DOCKER_TAGS }}         |
+          | JAX-TE       | ${{ needs.build-te.outputs.DOCKER_TAGS }}          |
 
           EOF
 
   retrofit-containers:
-    needs: [build-base, build-jax, build-t5x]
+    needs: [build-base, build-jax, build-te]
     if: always()
     runs-on: ubuntu-22.04
     env:
@@ -227,18 +227,17 @@ jobs:
       JAX_IMAGE: ${{ needs.build-jax.outputs.DOCKER_TAGS }}
     secrets: inherit
 
-  test-t5x:
-    needs: build-t5x
-    uses: ./.github/workflows/_test_t5x.yaml
+  test-te:
+    needs: build-te
+    uses: ./.github/workflows/_test_te.yaml
     with:
-      T5X_IMAGE: ${{ needs.build-t5x.outputs.DOCKER_TAGS }}
+      JAX_TE_IMAGE: ${{ needs.build-te.outputs.DOCKER_TAGS }}
     secrets: inherit
-
 
   finalize:
     if: always()
     # TODO: use dynamic matrix to make dependencies self-updating
-    needs: [build-summary, test-distribution, test-jax, test-t5x]
+    needs: [build-summary, test-distribution, test-jax, test-te]
     uses: ./.github/workflows/_finalize.yaml
     with:
       PUBLISH_BADGE: false