nightly-jax-test-unit.yaml

name: Nightly JAX unit test

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:

  metadata:
    runs-on: ubuntu-22.04
    outputs:
      JAX_IMAGE: ${{ steps.meta.outputs.JAX_IMAGE }}
      PUBLISH: ${{ steps.meta.outputs.PUBLISH }}
    steps:
      - name: Set metadata
        id: meta
        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
          echo "PUBLISH=${{ inputs.PUBLISH }}" >> $GITHUB_OUTPUT

  run-jobs:
    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

  publish:
    needs: [metadata, run-jobs]
    strategy:
      fail-fast: false
      matrix:
        GPU_ARCH: [V100, A100]
    uses: ./.github/workflows/_publish_badge.yaml
    if: (github.event_name == 'workflow_run' && github.event.workflow_run.conclusion == 'success') || github.event_name == 'workflow_dispatch'
    secrets: inherit
    with:
      ENDPOINT_FILENAME: 'jax-unit-test-status-${{ matrix.GPU_ARCH }}.json'
      PUBLISH: ${{ github.event_name == 'workflow_run' || needs.metadata.outputs.PUBLISH == 'true' }}
      SCRIPT: |
        ARTIFACTS="${{ needs.run-jobs.outputs.ARTIFACT_NAME }}-${{ matrix.GPU_ARCH }}/*"
        FAILED_TESTS=$(cat $ARTIFACTS | grep -c 'FAILED in' || true)
        PASSED_TESTS=$(cat $ARTIFACTS | grep -c 'PASSED in' || true)
        TOTAL_TESTS=$((FAILED_TESTS + PASSED_TESTS))
        if [[ $FAILED_TESTS == 0 ]]; then
          BADGE_COLOR=brightgreen
        else
          if [[ $FAILED_TESTS < $TOTAL_TESTS ]]; then
            BADGE_COLOR=yellow
          else
            BADGE_COLOR=red
          fi
        fi
        echo "LABEL='${{ matrix.GPU_ARCH }} Unit'" >> $GITHUB_OUTPUT
        echo "MESSAGE='${PASSED_TESTS}/${TOTAL_TESTS} passed'" >> $GITHUB_OUTPUT
        echo "COLOR='${BADGE_COLOR}'" >> $GITHUB_OUTPUT

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