Add 'finalize' stage to CI workflow (#258)

Closes #239

.github/workflows/_finalize.yaml

@@ -0,0 +1,29 @@
+name: ~finalize workflow run with job report and badge updates
+
+on:
+  workflow_call:
+    inputs:
+      PUBLISH_BADGE:
+        type: boolean
+        description: Update the landing page badges with run outcomes
+        default: false
+        required: false
+
+jobs:
+  upload-badge:
+    uses: ./.github/workflows/_upload_badge.yaml
+    secrets: inherit
+
+  report:
+    uses: ./.github/workflows/_summary.yaml
+    secrets: inherit
+
+  publish-badge:
+    needs: [upload-badge]
+    if: inputs.PUBLISH_BADGE == true
+    uses: ./.github/workflows/_copy_gist.yaml
+    with:
+      SRC_GIST_ID: ${{ needs.upload-badge.outputs.GIST_ID }}
+      DST_GIST_ID: ${{ vars.BADGE_ENDPOINT_GIST_ID }}
+      FILES: '.*badge.*.json'
+    secrets: inherit

.github/workflows/ci.yaml

@@ -236,3 +236,12 @@ jobs:
     with:
       PAX_IMAGE: ${{ needs.build-pax.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-te, test-t5x, test-pax]
+    uses: ./.github/workflows/_finalize.yaml
+    with:
+      PUBLISH_BADGE: false
+    secrets: inherit

.github/workflows/nightly-jax-test-unit.yaml

@@ -59,23 +59,11 @@ jobs:
     with:
       JAX_IMAGE: ${{ needs.metadata.outputs.JAX_IMAGE }}
     secrets: inherit
-
-  upload-badge:
-    needs: [run]
-    uses: ./.github/workflows/_upload_badge.yaml
-    secrets: inherit
-
-  report:
+
+  finalize:
+    if: always()
     needs: [run]
-    uses: ./.github/workflows/_summary.yaml
-    secrets: inherit
-
-  publish-badge:
-    needs: [metadata, upload-badge]
-    if: github.event_name == 'workflow_run' || needs.metadata.outputs.PUBLISH == 'true'
-    uses: ./.github/workflows/_copy_gist.yaml
+    uses: ./.github/workflows/_finalize.yaml
     with:
-      SRC_GIST_ID: ${{ needs.upload-badge.outputs.GIST_ID }}
-      DST_GIST_ID: ${{ vars.BADGE_ENDPOINT_GIST_ID }}
-      FILES: '.*badge.*.json'
+      PUBLISH_BADGE: ${{ github.event_name == 'workflow_run' || needs.metadata.outputs.PUBLISH == 'true' }}
     secrets: inherit