Add Maxtext tests to NGC test run (#743)

NVIDIA · Apr 22, 2024 · 99a71ea · 99a71ea
1 parent 4152544
commit 99a71ea
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 2 deletions.
diff --git a/.github/workflows/_test_maxtext.yaml b/.github/workflows/_test_maxtext.yaml
@@ -100,7 +100,7 @@ jobs:
           #SBATCH --time=00:10:00
           #SBATCH --output=${{ steps.meta.outputs.LOG_FILE }}
           #SBATCH --export="VOCAB_PATH=gs://t5-data/vocabs/cc_all.32000.100extra/sentencepiece.model,ENROOT_PASSWORD=${{ secrets.GITHUB_TOKEN }},NVCR_TOKEN=${{ secrets.NVCR_TOKEN }}"
-          
+
           # preload enroot container using one task per node
           time srun \
             --ntasks-per-node=1 \

diff --git a/.github/workflows/ngc-release-testing.yaml b/.github/workflows/ngc-release-testing.yaml
@@ -13,6 +13,12 @@ on:
         description: "Rosetta Pax image to run tests on"
         required: false
         default: ''
+      MAXTEXT_IMAGE:
+        type: string
+        description: "MaxText image to run tests on"
+        required: false
+        default: ''
+
 
 concurrency:
   group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }}
@@ -59,9 +65,16 @@ jobs:
     with:
       PAX_IMAGE: ${{ inputs.PAX_IMAGE }}
     secrets: inherit
+
+  test-maxtext:
+    if: inputs.MAXTEXT_IMAGE != ''
+    uses: ./.github/workflows/_test_maxtext.yaml
+    with:
+      MAXTEXT_IMAGE: ${{ inputs.MAXTEXT_IMAGE }}
+    secrets: inherit
 
   finalize:
-    needs: [ test-jax, test-rosetta-pax ]
+    needs: [ test-jax, test-rosetta-pax, test-maxtext ]
     if: "!cancelled()"
     uses: ./.github/workflows/_finalize.yaml
     secrets: inherit