Add sitrep report to JAX build, revise JAX unit test sitrep (#281)

NVIDIA · Oct 4, 2023 · 674347b · 674347b
1 parent 858c234
commit 674347b
Show file tree

Hide file tree

Showing 4 changed files with 72 additions and 8 deletions.
diff --git a/.github/workflows/_build_jax.yaml b/.github/workflows/_build_jax.yaml
@@ -33,6 +33,16 @@ on:
         description: Git commit, tag, or branch for XLA
         required: false
         default: main
+      ARTIFACT_NAME:
+        type: string
+        description: 'Name of the artifact zip file'
+        required: false
+        default: 'artifact-jax-build'
+      BADGE_FILENAME:
+        type: string
+        description: 'Name of the endpoint JSON file for shields.io badge'
+        required: false
+        default: 'badge-jax-build'
     outputs:
       DOCKER_TAGS:
         description: "Tags of the image built"
@@ -54,6 +64,8 @@ jobs:
       matrix:
         PLATFORM: [amd64, arm64]
     runs-on: [self-hosted, "${{ matrix.PLATFORM }}", large]
+    env:
+      BADGE_FILENAME_FULL: ${{ inputs.BADGE_FILENAME }}-${{ matrix.PLATFORM }}.json
     steps:
       - name: Print environment variables
         run: env
@@ -103,6 +115,7 @@ jobs:
             image=moby/buildkit:v0.12.1
 
       - name: Build docker images
+        id: build
         uses: docker/build-push-action@v4
         with:
           context: .github/container
@@ -123,6 +136,40 @@ jobs:
             REF_JAX=${{ inputs.REF_JAX }}
             REF_XLA=${{ inputs.REF_XLA }}
 
+      - name: Generate sitrep
+        if: success() || failure()
+        shell: bash -x -e {0}
+        run: |
+          # bring in utility functions
+          source .github/workflows/scripts/to_json.sh
+
+          badge_label='JAX ${{ matrix.PLATFORM }} build'
+          tags="${{ steps.meta.outputs.tags }}"
+          digest="${{ steps.build.outputs.digest }}"
+          outcome="${{ steps.build.outcome }}"
+
+          if [[ ${outcome} == "success" ]]; then
+            badge_message="pass"
+            badge_color=brightgreen
+            summary="JAX build on ${{ matrix.PLATFORM }}: $badge_message"
+          else
+            badge_message="fail"
+            badge_color=red
+            summary="JAX build on ${{ matrix.PLATFORM }}: $badge_message"
+          fi
+
+          to_json \
+            summary \
+            badge_label tags digest outcome \
+          > sitrep.json
+
+          schemaVersion=1 \
+          label="${badge_label}" \
+          message="${badge_message}" \
+          color="${badge_color}" \
+          to_json schemaVersion label message color \
+          > ${{ env.BADGE_FILENAME_FULL }}
+
       # Temporary workaround until the following issues are solved:
       # https://github.com/orgs/community/discussions/17245
       # https://github.com/actions/runner/pull/2477
@@ -138,6 +185,14 @@ jobs:
           name: image-name-jax-${{ matrix.PLATFORM }}
           path: image-name.txt
 
+      - name: Upload sitrep and badge
+        uses: actions/upload-artifact@v3
+        with:
+          name: ${{ inputs.ARTIFACT_NAME }}-${{ matrix.PLATFORM }}
+          path: |
+            sitrep.json
+            ${{ env.BADGE_FILENAME_FULL }}
+
   merge:
     runs-on: ubuntu-latest
     needs: build

diff --git a/.github/workflows/_test_jax.yaml b/.github/workflows/_test_jax.yaml
@@ -60,17 +60,19 @@ jobs:
           docker pull ${{ inputs.JAX_IMAGE }}
 
       - name: Backend-independent tests
-        shell: bash -x {0}
+        shell: bash -x -e {0}
+        # todo: distinguish between test failures and errors
+        continue-on-error: true
         run: |
           docker run --gpus all ${{ inputs.JAX_IMAGE }} test-jax.sh -b backend-independent | tee test-backend-independent.log
 
       - name: GPU-specific tests
-        shell: bash -x {0}
+        shell: bash -x -e {0}
+        continue-on-error: true
         run: |
           docker run --gpus all ${{ inputs.JAX_IMAGE }} test-jax.sh -b gpu | tee test-gpu.log
 
-      - name: Process results and generate report
-        if: success() || failure()
+      - name: Generate sitrep
         shell: bash -x -e {0}
         run: |
           # bring in utility functions
@@ -111,7 +113,6 @@ jobs:
           > ${{ env.BADGE_FILENAME_FULL }}
 
       - name: Upload artifacts
-        if: success() || failure()
         uses: actions/upload-artifact@v3
         with:
           name: ${{ inputs.ARTIFACT_NAME }}-${{ matrix.GPU_ARCH }}

diff --git a/.github/workflows/nightly-jax-build.yaml b/.github/workflows/nightly-jax-build.yaml
@@ -48,3 +48,11 @@ jobs:
       TARGET_TAGS: |
         type=raw,value=latest,priority=1000
         type=raw,value=nightly-${{ needs.metadata.outputs.BUILD_DATE }},priority=900
+
+  finalize:
+    if: always()
+    needs: [build]
+    uses: ./.github/workflows/_finalize.yaml
+    with:
+      PUBLISH_BADGE: ${{ github.event_name == 'schedule' || (github.event_name == 'workflow_dispatch' && inputs.PUBLISH) }}
+    secrets: inherit
diff --git a/.github/workflows/nightly-jax-test-unit.yaml b/.github/workflows/nightly-jax-test-unit.yaml
@@ -50,7 +50,7 @@ jobs:
             JAX_IMAGE=${{ inputs.JAX_IMAGE }}
           fi
           echo "JAX_IMAGE=${JAX_IMAGE}" >> $GITHUB_OUTPUT
-          echo "PUBLISH=${{ inputs.PUBLISH }}" >> $GITHUB_OUTPUT
+          echo "PUBLISH=${{ github.event_name == 'workflow_run' || (github.event_name == 'workflow_dispatch' && inputs.PUBLISH) }}" >> $GITHUB_OUTPUT
 
   run:
     needs: metadata
@@ -59,11 +59,11 @@ jobs:
     with:
       JAX_IMAGE: ${{ needs.metadata.outputs.JAX_IMAGE }}
     secrets: inherit
-  
+
   finalize:
     if: always()
     needs: [run]
     uses: ./.github/workflows/_finalize.yaml
     with:
-      PUBLISH_BADGE: ${{ github.event_name == 'workflow_run' || needs.metadata.outputs.PUBLISH == 'true' }}
+      PUBLISH_BADGE: ${{ needs.metadata.outputs.PUBLISH }}
     secrets: inherit