ngc-release-testing.yaml

name: ~NGC release testing

on:
  workflow_dispatch:
    inputs:
      JAX_IMAGE: 
        type: string
        description: "JAX image to run tests on"
        required: false
        default: ''
      PAX_IMAGE:
        type: string
        description: "Rosetta Pax image to run tests on"
        required: false
        default: ''
      MAXTEXT_IMAGE:
        type: string
        description: "MaxText image to run tests on"
        required: false
        default: ''
      LEVANTER_IMAGE:
        type: string
        description: "Levanter image to run tests on"
        required: false
        default: ''

concurrency:
  group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }}
  cancel-in-progress: ${{ github.ref != 'refs/heads/main' }}

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

jobs:
  test-jax:
    if: inputs.JAX_IMAGE != ''
    uses: ./.github/workflows/_test_unit.yaml
    with:
      TEST_NAME: jax
      EXECUTE: |
        docker run -i --shm-size=1g --gpus all \
        ${{ inputs.JAX_IMAGE }} \
        bash <<"EOF" |& tee test-backend-independent.log
          test-jax.sh -b backend-independent 
        EOF
        docker run -i --shm-size=1g --gpus all \
        ${{ inputs.JAX_IMAGE }} \
        bash <<"EOF" |& tee tee test-gpu.log
          test-jax.sh -b gpu
        EOF
      STATISTICS_SCRIPT: |
        errors=$(cat test-*.log | grep -c 'ERROR:' || true)
        failed_tests=$(cat test-*.log | grep -c 'FAILED in' || true)
        passed_tests=$(cat test-*.log | grep -c 'PASSED in' || true)
        total_tests=$((failed_tests + passed_tests))
        echo "TOTAL_TESTS=${total_tests}" >> $GITHUB_OUTPUT
        echo "ERRORS=${errors}" >> $GITHUB_OUTPUT
        echo "PASSED_TESTS=${passed_tests}" >> $GITHUB_OUTPUT
        echo "FAILED_TESTS=${failed_tests}" >> $GITHUB_OUTPUT
      ARTIFACTS: |
        test-backend-independent.log
        test-gpu.log
    secrets: inherit

  test-rosetta-pax:
    if: inputs.PAX_IMAGE != ''
    uses: ./.github/workflows/_test_pax_rosetta.yaml
    with:
      PAX_IMAGE: ${{ inputs.PAX_IMAGE }}
    secrets: inherit

  test-maxtext:
    if: inputs.MAXTEXT_IMAGE != ''
    uses: ./.github/workflows/_test_maxtext.yaml
    with:
      MAXTEXT_IMAGE: ${{ inputs.MAXTEXT_IMAGE }}
    secrets: inherit

  test-levanter:
    if: inputs.LEVANTER_IMAGE  != ''
    uses: ./.github/workflows/_test_unit.yaml
    with:
      TEST_NAME: levanter
      EXECUTE: |
        docker run --gpus all --shm-size=1g ${{ inputs.LEVANTER_IMAGE }} \
        bash -ec \
        "pip install pytest && PYTHONPATH=/opt/levanter/tests:$PYTHONPATH pytest /opt/levanter/tests" | tee test-levanter.log
      STATISTICS_SCRIPT: |
        summary_line=$(tail -n1 test-levanter.log)
        errors=$(echo $summary_line | grep -oE '[0-9]+ error' | awk '{print $1} END { if (!NR) print 0}')
        failed_tests=$(echo $summary_line | grep -oE '[0-9]+ failed' | awk '{print $1} END { if (!NR) print 0}')
        passed_tests=$(echo $summary_line | grep -oE '[0-9]+ passed' | awk '{print $1} END { if (!NR) print 0}')
        total_tests=$((failed_tests + passed_tests))
        echo "TOTAL_TESTS=${total_tests}" >> $GITHUB_OUTPUT
        echo "ERRORS=${errors}" >> $GITHUB_OUTPUT
        echo "PASSED_TESTS=${passed_tests}" >> $GITHUB_OUTPUT
        echo "FAILED_TESTS=${failed_tests}" >> $GITHUB_OUTPUT
      ARTIFACTS: |
        test-levanter.log
    secrets: inherit

  finalize:
    needs: [ test-jax, test-rosetta-pax, test-maxtext, test-levanter ]
    if: "!cancelled()"
    uses: ./.github/workflows/_finalize.yaml
    secrets: inherit