Skip to content

Nightly JAX unit test #186

Nightly JAX unit test

Nightly JAX unit test #186

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