nightly-jax-test-unit.yaml

name: Nightly JAX unit test
run-name: Nightly JAX unit test (${{ github.event_name == 'workflow_run' && format('nightly {0}', github.event.workflow_run.created_at) || github.event_name }})

on:
  workflow_run:
    workflows: [Nightly JAX build]
    types: [completed]
    branches: [main]
  workflow_dispatch:
    inputs:
      JAX_IMAGE:
        type: string
        description: 'JAX image built by NVIDIA/JAX-Toolbox'
        default: 'ghcr.io/nvidia/jax:latest'
        required: true
      PUBLISH:
        type: boolean
        description: Update status badge?
        default: false
        required: true

permissions:
  contents: read  # to fetch code
  actions:  write # to cancel previous workflows
  packages: write # to upload container

env:
  DEFAULT_JAX_IMAGE: 'ghcr.io/nvidia/jax:latest'

jobs:

  if-upstream-failed:
    if: (github.event_name == 'workflow_run' && github.event.workflow_run.conclusion == 'failure') && github.event_name != 'workflow_dispatch'
    runs-on: ubuntu-latest
    steps:
      - run: echo 'Upstream workflow failed, aborting run' && exit 1

  metadata:
    runs-on: ubuntu-22.04
    outputs:
      JAX_IMAGE: ${{ steps.image.outputs.JAX_IMAGE }}
      PUBLISH: ${{ steps.if-publish.outputs.PUBLISH }}
    steps:
      - name: Determine jax image to use
        id: image
        shell: bash -x -e {0}
        run: |
          if [[ -z "${{ inputs.JAX_IMAGE }}" ]]; then
            JAX_IMAGE=${{ env.DEFAULT_JAX_IMAGE }}
          else
            JAX_IMAGE=${{ inputs.JAX_IMAGE }}
          fi
          echo "JAX_IMAGE=${JAX_IMAGE}" >> $GITHUB_OUTPUT

      - name: Determine whether results will be 'published'
        id: if-publish
        shell: bash -x -e {0}
        run: |
          echo "PUBLISH=${{ github.event_name == 'workflow_run' || (github.event_name == 'workflow_dispatch' && inputs.PUBLISH) }}" >> $GITHUB_OUTPUT

  run:
    needs: metadata
    uses: ./.github/workflows/_test_jax.yaml
    if: (github.event_name == 'workflow_run' && github.event.workflow_run.conclusion == 'success') || github.event_name == 'workflow_dispatch'
    with:
      JAX_IMAGE: ${{ needs.metadata.outputs.JAX_IMAGE }}
    secrets: inherit

  finalize:
    if: always()
    needs: [metadata, run]
    uses: ./.github/workflows/_finalize.yaml
    with:
      PUBLISH_BADGE: ${{ needs.metadata.outputs.PUBLISH == 'true' }}
    secrets: inherit