ci: Allow selective triggering of pre-submit CI

By specifying a target-container as input to `TEST_SUBSET` the pre-submit CI will build all required containers and run tests only beginning from that node down to all connected leaves. It assumes that all nodes prior to the target-container are unchanged since their last successful test.
NVIDIA · Apr 30, 2024 · 20161c0 · 20161c0
1 parent e44e507
commit 20161c0
Show file tree

Hide file tree

Showing 2 changed files with 129 additions and 43 deletions.
diff --git a/.github/workflows/_ci.yaml b/.github/workflows/_ci.yaml
@@ -1,5 +1,5 @@
 name: ~CI, single-arch
-run-name: CI-${{ inputs.ARCHITECTURE }}
+run-name: CI-${{ inputs.ARCHITECTURE }}-${{ inputs.TESTSUBSET }}
 on:
   workflow_call:
     inputs:
@@ -16,26 +16,60 @@ on:
         description: Artifact name in current run w/ manifest/patches. Leaving empty uses manifest/patches in current branch
         default: ''
         required: false
+      TEST_SUBSET:
+        type: string
+        description: |
+          Subset of tests to run. Allowed values are one of:
+            - base
+            - jax
+            - levanter
+            - equinox
+            - triton
+            - upstream-t5x
+            - rosetta-t5x
+            - upstream-pax
+            - rosetta-pax
+            - maxtext
+            - grok
+            
+            Will run all downstream-connected nodes and leaves.
+        default: 'base'
+        required: false
     outputs:
       DOCKER_TAGS:
         description: JSON object containing tags of all docker images built
         value: ${{ jobs.collect-docker-tags.outputs.TAGS }}
 
 permissions:
-  contents: read  # to fetch code
-  actions:  write # to cancel previous workflows
+  contents: read # to fetch code
+  actions: write # to cancel previous workflows
   packages: write # to upload container
 
 jobs:
+  pre-flight:
+    runs-on: ubuntu-22.04
+    steps:
+      - name: Validate input `TEST_SUBSET`
+        shell: bash
+        run: |
+          valid_inputs=("base" "core" "levanter" "equinox" "triton" "upstream-t5x" "rosetta-t5x" "upstream-pax" "rosetta-pax" "maxtext" "grok")
+
+          if [[ " ${valid_inputs[*]} " != *" ${{ inputs.TEST_SUBSET }} "* ]]; then
+            echo "Invalid value for \`TEST_SUBSET\` provided. Expected one of: ($valid_inputs), Actual: ${{ inputs.TEST_SUBSET }}"
+            exit 1
+          fi
 
+  # Always
   build-base:
     uses: ./.github/workflows/_build_base.yaml
+    needs: pre-flight
     with:
       ARCHITECTURE: ${{ inputs.ARCHITECTURE }}
       BUILD_DATE: ${{ inputs.BUILD_DATE }}
       MANIFEST_ARTIFACT_NAME: ${{ inputs.MANIFEST_ARTIFACT_NAME }}
     secrets: inherit
 
+  # Always
   build-jax:
     needs: build-base
     uses: ./.github/workflows/_build.yaml
@@ -50,9 +84,10 @@ jobs:
       RUNNER_SIZE: large
     secrets: inherit
 
+  # base, jax, triton
   build-triton:
     needs: build-jax
-    if: inputs.ARCHITECTURE == 'amd64' # Triton does not seem to support arm64
+    if: contains(fromJSON('["base", "jax", "triton"]'), inputs.TEST_SUBSET) && inputs.ARCHITECTURE == 'amd64' # Triton does not seem to support arm64
     uses: ./.github/workflows/_build.yaml
     with:
       ARCHITECTURE: ${{ inputs.ARCHITECTURE }}
@@ -64,9 +99,11 @@ jobs:
       DOCKERFILE: .github/container/Dockerfile.triton
     secrets: inherit
 
+  # base, jax, equinox
   build-equinox:
     needs: build-jax
     uses: ./.github/workflows/_build.yaml
+    if: contains(fromJSON('["base", "jax", "equinox"]'), inputs.TEST_SUBSET)
     with:
       ARCHITECTURE: ${{ inputs.ARCHITECTURE }}
       ARTIFACT_NAME: artifact-equinox-build
@@ -77,9 +114,10 @@ jobs:
       DOCKERFILE: .github/container/Dockerfile.equinox
     secrets: inherit
 
+  # base, jax, maxtext
   build-maxtext:
     needs: build-jax
-    if: inputs.ARCHITECTURE == 'amd64' # Triton does not seem to support arm64
+    if: contains(fromJSON('["base", "jax", "maxtext"]'), inputs.TEST_SUBSET) && inputs.ARCHITECTURE == 'amd64' # Triton does not seem to support arm64
     uses: ./.github/workflows/_build.yaml
     with:
       ARCHITECTURE: ${{ inputs.ARCHITECTURE }}
@@ -91,35 +129,41 @@ jobs:
       DOCKERFILE: .github/container/Dockerfile.maxtext.amd64
     secrets: inherit
 
+  # base, jax, levanter
   build-levanter:
     needs: [build-jax]
     uses: ./.github/workflows/_build.yaml
+    if: contains(fromJSON('["base", "jax", "levanter"]'), inputs.TEST_SUBSET)
     with:
       ARCHITECTURE: ${{ inputs.ARCHITECTURE }}
-      ARTIFACT_NAME: "artifact-levanter-build"
-      BADGE_FILENAME: "badge-levanter-build"
+      ARTIFACT_NAME: 'artifact-levanter-build'
+      BADGE_FILENAME: 'badge-levanter-build'
       BUILD_DATE: ${{ inputs.BUILD_DATE }}
       BASE_IMAGE: ${{ needs.build-jax.outputs.DOCKER_TAG_MEALKIT }}
       CONTAINER_NAME: levanter
       DOCKERFILE: .github/container/Dockerfile.levanter
     secrets: inherit
 
+  # base, jax, upstream-t5x
   build-upstream-t5x:
     needs: build-jax
     uses: ./.github/workflows/_build.yaml
+    if: contains(fromJSON('["base", "jax", "upstream-t5x", "rosetta-t5x"]'), inputs.TEST_SUBSET)
     with:
       ARCHITECTURE: ${{ inputs.ARCHITECTURE }}
-      ARTIFACT_NAME: "artifact-t5x-build"
-      BADGE_FILENAME: "badge-t5x-build"
+      ARTIFACT_NAME: 'artifact-t5x-build'
+      BADGE_FILENAME: 'badge-t5x-build'
       BUILD_DATE: ${{ inputs.BUILD_DATE }}
       BASE_IMAGE: ${{ needs.build-jax.outputs.DOCKER_TAG_MEALKIT }}
       CONTAINER_NAME: upstream-t5x
       DOCKERFILE: .github/container/Dockerfile.t5x.${{ inputs.ARCHITECTURE }}
     secrets: inherit
 
+  # base, jax, upstream-pax
   build-upstream-pax:
     needs: build-jax
     uses: ./.github/workflows/_build.yaml
+    if: contains(fromJSON('["base", "jax", "upstream-pax", "rosetta-pax"]'), inputs.TEST_SUBSET)
     with:
       ARCHITECTURE: ${{ inputs.ARCHITECTURE }}
       ARTIFACT_NAME: artifact-pax-build
@@ -130,42 +174,48 @@ jobs:
       DOCKERFILE: .github/container/Dockerfile.pax.${{ inputs.ARCHITECTURE }}
     secrets: inherit
 
+  # base, jax, upstream-t5x, rosetta-t5x
   build-rosetta-t5x:
     needs: build-upstream-t5x
     uses: ./.github/workflows/_build_rosetta.yaml
+    if: contains(fromJSON('["base", "jax", "upstream-t5x", "rosetta-t5x"]'), inputs.TEST_SUBSET)
     with:
       ARCHITECTURE: ${{ inputs.ARCHITECTURE }}
       BUILD_DATE: ${{ inputs.BUILD_DATE }}
       BASE_IMAGE: ${{ needs.build-upstream-t5x.outputs.DOCKER_TAG_MEALKIT }}
       BASE_LIBRARY: t5x
     secrets: inherit
 
+  # base, jax, upstream-pax, rosetta-pax
   build-rosetta-pax:
     needs: build-upstream-pax
     uses: ./.github/workflows/_build_rosetta.yaml
+    if: contains(fromJSON('["base", "jax", "upstream-pax", "rosetta-pax"]'), inputs.TEST_SUBSET)
     with:
       ARCHITECTURE: ${{ inputs.ARCHITECTURE }}
       BUILD_DATE: ${{ inputs.BUILD_DATE }}
       BASE_IMAGE: ${{ needs.build-upstream-pax.outputs.DOCKER_TAG_MEALKIT }}
       BASE_LIBRARY: pax
     secrets: inherit
 
+  # base, jax, grok
   build-grok:
     needs: [build-jax]
     uses: ./.github/workflows/_build.yaml
+    if: contains(fromJSON('["base", "jax", "grok"]'), inputs.TEST_SUBSET)
     with:
       ARCHITECTURE: ${{ inputs.ARCHITECTURE }}
-      ARTIFACT_NAME: "artifact-grok-build"
-      BADGE_FILENAME: "badge-grok-build"
+      ARTIFACT_NAME: 'artifact-grok-build'
+      BADGE_FILENAME: 'badge-grok-build'
       BUILD_DATE: ${{ inputs.BUILD_DATE }}
       BASE_IMAGE: ${{ needs.build-jax.outputs.DOCKER_TAG_MEALKIT }}
       CONTAINER_NAME: grok
       DOCKERFILE: .github/container/Dockerfile.grok
     secrets: inherit
-    
+
   collect-docker-tags:
     runs-on: ubuntu-22.04
-    if: "!cancelled()"
+    if: '!cancelled()'
     needs:
       - build-base
       - build-jax
@@ -236,9 +286,10 @@ jobs:
       - name: Run integration test ${{ matrix.TEST_SCRIPT }}
         run: bash rosetta/tests/${{ matrix.TEST_SCRIPT }}
 
+  # base, jax
   test-jax:
     needs: build-jax
-    if: inputs.ARCHITECTURE == 'amd64' # arm64 runners n/a
+    if: contains(fromJSON('["base", "jax"]'), inputs.TEST_SUBSET) && inputs.ARCHITECTURE == 'amd64' # arm64 runners n/a
     uses: ./.github/workflows/_test_unit.yaml
     with:
       TEST_NAME: jax
@@ -291,33 +342,37 @@ jobs:
   #       test-equinox.log
   #   secrets: inherit
 
+  # base, jax, upstream-pax
   test-te-multigpu:
     needs: build-upstream-pax
-    if: inputs.ARCHITECTURE == 'amd64' # arm64 runners n/a
+    if: contains(fromJSON('["base", "jax", "upstream-pax"]'), inputs.TEST_SUBSET) && inputs.ARCHITECTURE == 'amd64' # arm64 runners n/a
     uses: ./.github/workflows/_test_te.yaml
     with:
       TE_IMAGE: ${{ needs.build-upstream-pax.outputs.DOCKER_TAG_FINAL }}
     secrets: inherit
 
+  # base, jax, upstream-t5x
   test-upstream-t5x:
     needs: build-upstream-t5x
-    if: inputs.ARCHITECTURE == 'amd64' # arm64 runners n/a
+    if: contains(fromJSON('["base", "jax", "upstream-t5x"]'), inputs.TEST_SUBSET) && inputs.ARCHITECTURE == 'amd64' # arm64 runners n/a
     uses: ./.github/workflows/_test_upstream_t5x.yaml
     with:
       T5X_IMAGE: ${{ needs.build-upstream-t5x.outputs.DOCKER_TAG_FINAL }}
     secrets: inherit
 
+  # base, jax, upstream-t5x, rosetta-t5x
   test-rosetta-t5x:
     needs: build-rosetta-t5x
-    if: inputs.ARCHITECTURE == 'amd64' # no images for arm64
+    if: contains(fromJSON('["base", "jax", "upstream-t5x", "rosetta-t5x"]'), inputs.TEST_SUBSET) && inputs.ARCHITECTURE == 'amd64' # no images for arm64
     uses: ./.github/workflows/_test_t5x_rosetta.yaml
     with:
       T5X_IMAGE: ${{ needs.build-rosetta-t5x.outputs.DOCKER_TAG_FINAL }}
     secrets: inherit
 
+  # base, jax
   test-pallas:
     needs: build-jax
-    if: inputs.ARCHITECTURE == 'amd64' # triton doesn't support arm64(?)
+    if: contains(fromJSON('["base", "jax"]'), inputs.TEST_SUBSET) && inputs.ARCHITECTURE == 'amd64' # triton doesn't support arm64(?)
     uses: ./.github/workflows/_test_unit.yaml
     with:
       TEST_NAME: pallas
@@ -341,9 +396,10 @@ jobs:
         test-pallas.log
     secrets: inherit
 
+  # base, jax, triton
   test-triton:
     needs: build-triton
-    if: inputs.ARCHITECTURE == 'amd64' # no images for arm64
+    if: contains(fromJSON('["base", "jax", "triton"]'), inputs.TEST_SUBSET) && inputs.ARCHITECTURE == 'amd64' # no images for arm64
     uses: ./.github/workflows/_test_unit.yaml
     with:
       TEST_NAME: triton
@@ -367,9 +423,10 @@ jobs:
         test-triton.log
     secrets: inherit
 
+  # base, jax, levanter
   test-levanter:
     needs: build-levanter
-    if: inputs.ARCHITECTURE == 'amd64'  # arm64 runners n/a
+    if: contains(fromJSON('["base", "jax", "levanter"]'), inputs.TEST_SUBSET) && inputs.ARCHITECTURE == 'amd64' # arm64 runners n/a
     uses: ./.github/workflows/_test_unit.yaml
     with:
       TEST_NAME: levanter
@@ -394,9 +451,10 @@ jobs:
         test-levanter.log
     secrets: inherit
 
+  # base, jax, upstream-pax
   test-te:
     needs: build-upstream-pax
-    if: inputs.ARCHITECTURE == 'amd64'  # arm64 runners n/a
+    if: contains(fromJSON('["base", "jax", "upstream-pax"]'), inputs.TEST_SUBSET) && inputs.ARCHITECTURE == 'amd64' # arm64 runners n/a
     uses: ./.github/workflows/_test_unit.yaml
     with:
       TEST_NAME: te
@@ -422,25 +480,28 @@ jobs:
         pytest-report.jsonl
     secrets: inherit
 
+  # base, jax, upstream-pax
   test-upstream-pax:
     needs: build-upstream-pax
-    if: inputs.ARCHITECTURE == 'amd64' # no images for arm64
+    if: contains(fromJSON('["base", "jax", "upstream-pax"]'), inputs.TEST_SUBSET) && inputs.ARCHITECTURE == 'amd64' # no images for arm64
     uses: ./.github/workflows/_test_upstream_pax.yaml
     with:
       PAX_IMAGE: ${{ needs.build-upstream-pax.outputs.DOCKER_TAG_FINAL }}
     secrets: inherit
 
+  # base, jax, upstream-pax, rosetta-pax
   test-rosetta-pax:
     needs: build-rosetta-pax
-    if: inputs.ARCHITECTURE == 'amd64' # no images for arm64
+    if: contains(fromJSON('["base", "jax", "upstream-pax", "rosetta-pax"]'), inputs.TEST_SUBSET) && inputs.ARCHITECTURE == 'amd64' # no images for arm64
     uses: ./.github/workflows/_test_pax_rosetta.yaml
     with:
       PAX_IMAGE: ${{ needs.build-rosetta-pax.outputs.DOCKER_TAG_FINAL }}
     secrets: inherit
 
+  # base, jax, maxtext
   test-maxtext:
     needs: build-maxtext
-    if: inputs.ARCHITECTURE == 'amd64' # no images for arm64
+    if: contains(fromJSON('["base", "jax", "maxtext"]'), inputs.TEST_SUBSET) && inputs.ARCHITECTURE == 'amd64' # no images for arm64
     uses: ./.github/workflows/_test_maxtext.yaml
     with:
       MAXTEXT_IMAGE: ${{ needs.build-maxtext.outputs.DOCKER_TAG_FINAL }}