_ci.yaml

name: ~CI, single-arch
run-name: CI-${{ inputs.ARCHITECTURE }}
on:
  workflow_call:
    inputs:
      ARCHITECTURE:
        type: string
        required: true
      BUILD_DATE:
        type: string
        description: 'Build date in YYYY-MM-DD format'
        required: false
        default: NOT SPECIFIED
      CUDA_IMAGE:
        type: string
        description: CUDA image to use as base, e.g. nvidia/cuda:X.Y.Z-devel-ubuntu22.04
        default: 'latest'
        required: false
      MANIFEST_ARTIFACT_NAME:
        type: string
        description: 'Artifact name in current run w/ manifest/patches. Leaving empty uses manifest/patches in current branch'
        default: ''
        required: false
      SOURCE_URLREFS:
        type: string
        description: 'A JSON object containing git url+refs for softwares to be built'
        required: false
        default: '{}'
    outputs:
      DOCKER_TAGS:
        description: 'JSON object containing tags of all docker images built'
        value: ${{ jobs.collect-docker-tags.outputs.TAGS }}

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

jobs:

  build-base:
    uses: ./.github/workflows/_build_base.yaml
    with:
      ARCHITECTURE: ${{ inputs.ARCHITECTURE }}
      BASE_IMAGE: ${{ inputs.CUDA_IMAGE }}
      BUILD_DATE: ${{ inputs.BUILD_DATE }}
      MANIFEST_ARTIFACT_NAME: ${{ inputs.MANIFEST_ARTIFACT_NAME }}
    secrets: inherit

  build-jax:
    needs: build-base
    uses: ./.github/workflows/_build.yaml
    with:
      ARCHITECTURE: ${{ inputs.ARCHITECTURE }}
      ARTIFACT_NAME: artifact-jax-build
      BADGE_FILENAME: badge-jax-build
      BUILD_DATE: ${{ inputs.BUILD_DATE }}
      BASE_IMAGE: ${{ needs.build-base.outputs.DOCKER_TAG }}
      CONTAINER_NAME: jax
      DOCKERFILE: .github/container/Dockerfile.jax
      RUNNER_SIZE: large
      EXTRA_BUILD_ARGS: |
        URLREF_JAX=${{ fromJson(inputs.SOURCE_URLREFS).JAX }}
        URLREF_XLA=${{ fromJson(inputs.SOURCE_URLREFS).XLA }}
        URLREF_FLAX=${{ fromJson(inputs.SOURCE_URLREFS).FLAX }}
        URLREF_TRANSFORMER_ENGINE=${{ fromJson(inputs.SOURCE_URLREFS).TRANSFORMER_ENGINE }}
    secrets: inherit

  build-triton:
    needs: build-jax
    if: inputs.ARCHITECTURE == 'amd64' # Triton does not seem to support arm64
    uses: ./.github/workflows/_build.yaml
    with:
      ARCHITECTURE: ${{ inputs.ARCHITECTURE }}
      ARTIFACT_NAME: artifact-triton-build
      BADGE_FILENAME: badge-triton-build
      BUILD_DATE: ${{ inputs.BUILD_DATE }}
      BASE_IMAGE: ${{ needs.build-jax.outputs.DOCKER_TAG_MEALKIT }}
      CONTAINER_NAME: triton
      DOCKERFILE: .github/container/Dockerfile.triton
      RUNNER_SIZE: large
      EXTRA_BUILD_ARGS: URLREF_JAX_TRITON=${{ fromJson(inputs.SOURCE_URLREFS).JAX_TRITON }}
    secrets: inherit

  build-equinox:
    needs: build-jax
    uses: ./.github/workflows/_build.yaml
    with:
      ARCHITECTURE: ${{ inputs.ARCHITECTURE }}
      ARTIFACT_NAME: artifact-equinox-build
      BADGE_FILENAME: badge-equinox-build
      BUILD_DATE: ${{ inputs.BUILD_DATE }}
      BASE_IMAGE: ${{ needs.build-jax.outputs.DOCKER_TAG_MEALKIT }}
      CONTAINER_NAME: equinox
      DOCKERFILE: .github/container/Dockerfile.equinox
      EXTRA_BUILD_ARGS: |
        URLREF_EQUINOX=${{ fromJson(inputs.SOURCE_URLREFS).EQUINOX }}
    secrets: inherit

  build-maxtext:
    needs: build-jax
    uses: ./.github/workflows/_build.yaml
    with:
      ARCHITECTURE: ${{ inputs.ARCHITECTURE }}
      ARTIFACT_NAME: artifact-maxtext-build
      BADGE_FILENAME: badge-maxtext-build
      BUILD_DATE: ${{ inputs.BUILD_DATE }}
      BASE_IMAGE: ${{ needs.build-jax.outputs.DOCKER_TAG_MEALKIT }}
      CONTAINER_NAME: maxtext
      DOCKERFILE: .github/container/Dockerfile.maxtext.${{ inputs.ARCHITECTURE }}
      EXTRA_BUILD_ARGS: |
        URLREF_MAXTEXT=${{ fromJson(inputs.SOURCE_URLREFS).MAXTEXT }}
    secrets: inherit

  build-levanter:
    needs: [build-jax]
    uses: ./.github/workflows/_build.yaml
    with:
      ARCHITECTURE: ${{ inputs.ARCHITECTURE }}
      ARTIFACT_NAME: "artifact-levanter-build"
      BADGE_FILENAME: "badge-levanter-build"
      BUILD_DATE: ${{ inputs.BUILD_DATE }}
      BASE_IMAGE: ${{ needs.build-jax.outputs.DOCKER_TAG_MEALKIT }}
      CONTAINER_NAME: levanter
      DOCKERFILE: .github/container/Dockerfile.levanter
      EXTRA_BUILD_ARGS: |
        URLREF_LEVANTER=${{ fromJson(inputs.SOURCE_URLREFS).LEVANTER }}
        URLREF_HALIAX=${{ fromJson(inputs.SOURCE_URLREFS).HALIAX }}
    secrets: inherit

  build-upstream-t5x:
    needs: build-jax
    uses: ./.github/workflows/_build.yaml
    with:
      ARCHITECTURE: ${{ inputs.ARCHITECTURE }}
      ARTIFACT_NAME: "artifact-t5x-build"
      BADGE_FILENAME: "badge-t5x-build"
      BUILD_DATE: ${{ inputs.BUILD_DATE }}
      BASE_IMAGE: ${{ needs.build-jax.outputs.DOCKER_TAG_MEALKIT }}
      CONTAINER_NAME: upstream-t5x
      DOCKERFILE: .github/container/Dockerfile.t5x.${{ inputs.ARCHITECTURE }}
      EXTRA_BUILD_ARGS: |
        URLREF_T5X=${{ fromJson(inputs.SOURCE_URLREFS).T5X }}
    secrets: inherit

  build-upstream-pax:
    needs: build-jax
    uses: ./.github/workflows/_build.yaml
    with:
      ARCHITECTURE: ${{ inputs.ARCHITECTURE }}
      ARTIFACT_NAME: artifact-pax-build
      BADGE_FILENAME: badge-pax-build
      BUILD_DATE: ${{ inputs.BUILD_DATE }}
      BASE_IMAGE: ${{ needs.build-jax.outputs.DOCKER_TAG_MEALKIT }}
      CONTAINER_NAME: upstream-pax
      DOCKERFILE: .github/container/Dockerfile.pax.${{ inputs.ARCHITECTURE }}
      EXTRA_BUILD_ARGS: |
        URLREF_PAXML=${{ fromJson(inputs.SOURCE_URLREFS).PAXML }}
        URLREF_PRAXIS=${{ fromJson(inputs.SOURCE_URLREFS).PRAXIS }}
    secrets: inherit

  build-rosetta-t5x:
    needs: build-upstream-t5x
    uses: ./.github/workflows/_build_rosetta.yaml
    with:
      ARCHITECTURE: ${{ inputs.ARCHITECTURE }}
      BUILD_DATE: ${{ inputs.BUILD_DATE }}
      BASE_IMAGE: ${{ needs.build-upstream-t5x.outputs.DOCKER_TAG_MEALKIT }}
      BASE_LIBRARY: t5x
    secrets: inherit

  build-rosetta-pax:
    needs: build-upstream-pax
    uses: ./.github/workflows/_build_rosetta.yaml
    with:
      ARCHITECTURE: ${{ inputs.ARCHITECTURE }}
      BUILD_DATE: ${{ inputs.BUILD_DATE }}
      BASE_IMAGE: ${{ needs.build-upstream-pax.outputs.DOCKER_TAG_MEALKIT }}
      BASE_LIBRARY: pax
    secrets: inherit

  build-gemma:
    needs: build-jax
    uses: ./.github/workflows/_build.yaml
    if: inputs.ARCHITECTURE == 'amd64' # build only amd64
    with:
      ARCHITECTURE: ${{ inputs.ARCHITECTURE }}
      ARTIFACT_NAME: artifact-gemma-build
      BADGE_FILENAME: badge-gemma-build
      BUILD_DATE: ${{ inputs.BUILD_DATE }}
      BASE_IMAGE: ${{ needs.build-jax.outputs.DOCKER_TAG_MEALKIT }}
      CONTAINER_NAME: gemma
      DOCKERFILE: rosetta/Dockerfile.gemma
      DOCKER_CONTEXT: .
      EXTRA_BUILD_ARGS: |
        URLREF_GEMMA=${{ fromJson(inputs.SOURCE_URLREFS).gemma }}
        URLREF_BIG_VISION=${{ fromJson(inputs.SOURCE_URLREFS).big_vision }}
        URLREF_COMMON_LOOP_UTILS=${{ fromJson(inputs.SOURCE_URLREFS).common_loop_utils }}
        URLREF_FLAXFORMER=${{ fromJson(inputs.SOURCE_URLREFS).flaxformer }}
        URLREF_PANOPTICAPI=${{ fromJson(inputs.SOURCE_URLREFS).panopticapi }}
    secrets: inherit

  collect-docker-tags:
    runs-on: ubuntu-22.04
    if: "!cancelled()"
    needs:
      - build-base
      - build-jax
      - build-triton
      - build-equinox
      - build-maxtext
      - build-levanter
      - build-upstream-t5x
      - build-upstream-pax
      - build-rosetta-t5x
      - build-rosetta-pax
      - build-gemma
    outputs:
      TAGS: ${{ steps.collect-tags.outputs.TAGS }}
    steps:
      - name: Save docker tags as a JSON object
        id: collect-tags
        run: |
          TAGS=$(cat <<EOF | jq -c
          [\
            {"flavor": "base",         "stage": "final",   "priority": 800,  "tag": "${{ needs.build-base.outputs.DOCKER_TAG }}"},\
            {"flavor": "jax",          "stage": "final",   "priority": 1000, "tag": "${{ needs.build-jax.outputs.DOCKER_TAG_FINAL }}"},\
            {"flavor": "triton",       "stage": "final",   "priority": 900,  "tag": "${{ needs.build-triton.outputs.DOCKER_TAG_FINAL }}"},\
            {"flavor": "equinox",      "stage": "final",   "priority": 900,  "tag": "${{ needs.build-equinox.outputs.DOCKER_TAG_FINAL }}"},\
            {"flavor": "maxtext",      "stage": "final",   "priority": 900,  "tag": "${{ needs.build-maxtext.outputs.DOCKER_TAG_FINAL }}"},\
            {"flavor": "levanter",     "stage": "final",   "priority": 900,  "tag": "${{ needs.build-levanter.outputs.DOCKER_TAG_FINAL }}"},\
            {"flavor": "upstream-t5x", "stage": "final",   "priority": 900,  "tag": "${{ needs.build-upstream-t5x.outputs.DOCKER_TAG_FINAL }}"},\
            {"flavor": "upstream-pax", "stage": "final",   "priority": 900,  "tag": "${{ needs.build-upstream-pax.outputs.DOCKER_TAG_FINAL }}"},\
            {"flavor": "t5x",          "stage": "final",   "priority": 900,  "tag": "${{ needs.build-rosetta-t5x.outputs.DOCKER_TAG_FINAL }}"},\
            {"flavor": "pax",          "stage": "final",   "priority": 900,  "tag": "${{ needs.build-rosetta-pax.outputs.DOCKER_TAG_FINAL }}"},\
            {"flavor": "gemma",        "stage": "final",   "priority": 900,  "tag": "${{ needs.build-gemma.outputs.DOCKER_TAG_FINAL }}"},\
            {"flavor": "jax",          "stage": "mealkit", "priority": 500,  "tag": "${{ needs.build-jax.outputs.DOCKER_TAG_MEALKIT }}"},\
            {"flavor": "triton",       "stage": "mealkit", "priority": 500,  "tag": "${{ needs.build-triton.outputs.DOCKER_TAG_MEALKIT }}"},\
            {"flavor": "equinox",      "stage": "mealkit", "priority": 500,  "tag": "${{ needs.build-equinox.outputs.DOCKER_TAG_MEALKIT }}"},\
            {"flavor": "maxtext",      "stage": "mealkit", "priority": 500,  "tag": "${{ needs.build-maxtext.outputs.DOCKER_TAG_MEALKIT }}"},\
            {"flavor": "levanter",     "stage": "mealkit", "priority": 500,  "tag": "${{ needs.build-levanter.outputs.DOCKER_TAG_MEALKIT }}"},\
            {"flavor": "upstream-t5x", "stage": "mealkit", "priority": 500,  "tag": "${{ needs.build-upstream-t5x.outputs.DOCKER_TAG_MEALKIT }}"},\
            {"flavor": "upstream-pax", "stage": "mealkit", "priority": 500,  "tag": "${{ needs.build-upstream-pax.outputs.DOCKER_TAG_MEALKIT }}"},\
            {"flavor": "t5x",          "stage": "mealkit", "priority": 500,  "tag": "${{ needs.build-rosetta-t5x.outputs.DOCKER_TAG_MEALKIT }}"},\
            {"flavor": "pax",          "stage": "mealkit", "priority": 500,  "tag": "${{ needs.build-rosetta-pax.outputs.DOCKER_TAG_MEALKIT }}"},\
            {"flavor": "gemma",        "stage": "mealkit", "priority": 500,  "tag": "${{ needs.build-gemma.outputs.DOCKER_TAG_MEALKIT }}"},\

            {}\
          ]
          EOF
          )

          echo "TAGS=${TAGS}" >> $GITHUB_OUTPUT

  test-distribution:
    runs-on: ubuntu-22.04
    strategy:
      matrix:
        TEST_SCRIPT:
          - extra-only-distribution.sh
          - mirror-only-distribution.sh
          - upstream-only-distribution.sh
          - local-patch-distribution.sh
      fail-fast: false
    steps:
      - name: Print environment variables
        run: env
      - name: Set git login for tests
        run: |
          git config --global user.email "jax@nvidia.com"
          git config --global user.name "JAX-Toolbox CI"
      - name: Check out the repository under ${GITHUB_WORKSPACE}
        uses: actions/checkout@v4
      - name: Run integration test ${{ matrix.TEST_SCRIPT }}
        run: bash rosetta/tests/${{ matrix.TEST_SCRIPT }}

  test-jax:
    needs: build-jax
    if: inputs.ARCHITECTURE == 'amd64' # arm64 runners n/a
    uses: ./.github/workflows/_test_unit.yaml
    with:
      TEST_NAME: jax
      EXECUTE: |
        docker run -i --shm-size=1g --gpus all \
        ${{ needs.build-jax.outputs.DOCKER_TAG_FINAL }} \
        bash <<"EOF" |& tee test-backend-independent.log
          test-jax.sh -b backend-independent 
        EOF
        docker run -i --shm-size=1g --gpus all \
        ${{ needs.build-jax.outputs.DOCKER_TAG_FINAL }} \
        bash <<"EOF" |& tee tee test-gpu.log
          nvidia-cuda-mps-control -d
          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-nsys-jax:
    needs: build-jax
    if: inputs.ARCHITECTURE == 'amd64' # arm64 runners n/a
    uses: ./.github/workflows/_test_unit.yaml
    with:
      TEST_NAME: nsys-jax
      EXECUTE: |
        set -o pipefail
        num_tests=0
        num_failures=0
        # Run the pytest-driven tests
        docker run -i --shm-size=1g --gpus all \
          -v $PWD:/opt/output \
          ${{ needs.build-jax.outputs.DOCKER_TAG_FINAL }} \
          bash <<"EOF" |& tee test-nsys-jax.log
            pip install pytest-reportlog -e /opt/jax_nsys/python/jax_nsys[test]
            pytest --report-log=/opt/output/pytest-report.jsonl /opt/jax_nsys_tests
        EOF
        GPUS_PER_NODE=$(nvidia-smi -L | grep -c '^GPU')
        for mode in 1-process 2-process process-per-gpu; do
          DOCKER="docker run --shm-size=1g --gpus all --env XLA_FLAGS=--xla_gpu_enable_command_buffer= --env XLA_PYTHON_CLIENT_MEM_FRACTION=0.9 -v ${PWD}:/opt/output ${{ needs.build-jax.outputs.DOCKER_TAG_FINAL }}"
          if [[ "${mode}" == "1-process" ]]; then
            PROCESS_COUNT=1
            ARGS=""
          elif [[ "${mode}" == "2-process" ]]; then
            # Use two processes with GPUS_PER_NODE/2 GPUs per process in the hope that
            # this will flush out more bugs than process-per-node or process-per-GPU.
            PROCESS_COUNT=2
            ARGS="--process-id RANK --process-count ${PROCESS_COUNT} --coordinator-address 127.0.0.1:12345 --gpus-per-process $((GPUS_PER_NODE/2)) --distributed"
          else
            PROCESS_COUNT=${GPUS_PER_NODE}
            ARGS="--process-id RANK --process-count ${PROCESS_COUNT} --coordinator-address 127.0.0.1:12345 --gpus-per-process 1 --distributed"
          fi
          for collection in full partial; do
            NSYS_JAX="nsys-jax"
            if [[ "${mode}" == "1-process" ]]; then
              # We will not run nsys-jax-combine, so run analyses eagerly
              NSYS_JAX+=" --nsys-jax-analysis communication --nsys-jax-analysis summary"
            fi
            NSYS_JAX+=" --output=/opt/output/${mode}-${collection}-execution-%q{RANK}"
            if [[ "${collection}" == "partial" ]]; then
              NSYS_JAX+=" --capture-range=cudaProfilerApi --capture-range-end=stop"
              # nvbug/4801401
              NSYS_JAX+=" --sample=none"
            fi
            set +e
            ${DOCKER} parallel-launch RANK ${PROCESS_COUNT} ${NSYS_JAX} \
              -- jax-nccl-test ${ARGS} |& tee ${mode}-${collection}-execution.log
            num_failures=$((num_failures + ($? != 0)))
            set -e
            num_tests=$((num_tests + 1))
          done
          if [[ "${mode}" != "1-process" ]]; then
            # Run nsys-jax-combine
            NSYS_JAX_COMBINE="nsys-jax-combine --analysis communication --analysis summary --output=/opt/output/${mode}-${collection}-execution.zip"
            for (( i=0; i<PROCESS_COUNT; i++ )); do
              NSYS_JAX_COMBINE+=" /opt/output/${mode}-${collection}-execution-${i}.zip"
            done
            set +e
            ${DOCKER} ${NSYS_JAX_COMBINE} |& tee ${mode}-${collection}-execution-combine.log
            num_failures=$((num_failures + ($? != 0)))
            set -e
            num_tests=$((num_tests + 1))
          fi
        done
        ls -R .
        echo "NSYS_JAX_TEST_COUNT=${num_tests}" >> $GITHUB_ENV
        echo "NSYS_JAX_FAIL_COUNT=${num_failures}" >> $GITHUB_ENV
        exit $num_failures
      STATISTICS_SCRIPT: |
        summary_line=$(tail -n1 test-nsys-jax.log)
        num_errors=$(echo $summary_line | grep -oE '[0-9]+ error' | awk '{print $1} END { if (!NR) print 0}')
        passed_tests=$(cat pytest-report.jsonl | jq -r 'select(."$report_type" == "TestReport" and .when == "call" and .outcome == "passed") | .outcome' | wc -l)
        failed_tests=$(cat pytest-report.jsonl | jq -r 'select(."$report_type" == "TestReport" and .when == "call" and .outcome == "failed") | .outcome' | wc -l)
        total_tests=$(( NSYS_JAX_TEST_COUNT + passed_tests + failed_tests ))
        num_passed=$(( passed_tests + NSYS_JAX_TEST_COUNT - NSYS_JAX_FAIL_COUNT ))
        num_failed=$(( failed_tests + NSYS_JAX_FAIL_COUNT ))
        echo "TOTAL_TESTS=${total_tests}" >> $GITHUB_OUTPUT
        echo "ERRORS=${num_errors}" >> $GITHUB_OUTPUT
        echo "PASSED_TESTS=${num_passed}" >> $GITHUB_OUTPUT
        echo "FAILED_TESTS=${num_failed}" >> $GITHUB_OUTPUT
      ARTIFACTS: |
        # pytest-driven part
        test-nsys-jax.log
        pytest-report.jsonl
        # nsys-jax logfiles
        *process-*-execution.log
        # nsys-jax output for the case that doesn't use nsys-jax-combine
        1-process-*-execution-0.zip
        # nsys-jax-combine output/logfiles
        *process*-*-execution.zip
        *-execution-combine.log
    secrets: inherit

  # test-equinox:
  #   needs: build-equinox
  #   if: inputs.ARCHITECTURE == 'amd64' # arm64 runners n/a
  #   uses: ./.github/workflows/_test_unit.yaml
  #   with:
  #     IMAGE: ${{ needs.build-equinox.outputs.DOCKER_TAG_FINAL }}
  #     TEST_NAME: equinox
  #     EXECUTE: |
  #       docker run --shm-size=1g --gpus all ${{ needs.build-equinox.outputs.DOCKER_TAG_FINAL }} \
  #       bash -exc -o pipefail \
  #       'pushd /opt/equinox/tests && pip install -r requirements.txt && pytest .' | tee test-equinox.log
  #     STATISTICS_SCRIPT: |
  #       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-equinox.log
  #   secrets: inherit

  test-te-multigpu:
    needs: build-upstream-pax
    if: inputs.ARCHITECTURE == 'amd64' # arm64 runners n/a
    uses: ./.github/workflows/_test_te.yaml
    with:
      TE_IMAGE: ${{ needs.build-upstream-pax.outputs.DOCKER_TAG_FINAL }}
    secrets: inherit

  test-upstream-t5x:
    needs: build-upstream-t5x
    if: inputs.ARCHITECTURE == 'amd64' # arm64 runners n/a
    uses: ./.github/workflows/_test_upstream_t5x.yaml
    with:
      T5X_IMAGE: ${{ needs.build-upstream-t5x.outputs.DOCKER_TAG_FINAL }}
    secrets: inherit

  test-rosetta-t5x:
    needs: build-rosetta-t5x
    if: inputs.ARCHITECTURE == 'amd64' # no images for arm64
    uses: ./.github/workflows/_test_t5x_rosetta.yaml
    with:
      T5X_IMAGE: ${{ needs.build-rosetta-t5x.outputs.DOCKER_TAG_FINAL }}
    secrets: inherit

  test-triton:
    needs: build-triton
    if: inputs.ARCHITECTURE == 'amd64' # no images for arm64
    uses: ./.github/workflows/_test_unit.yaml
    with:
      TEST_NAME: triton
      EXECUTE: |
        docker run -i --shm-size=1g --gpus all --volume $PWD:/output \
        ${{ needs.build-triton.outputs.DOCKER_TAG_FINAL }} \
        bash <<"EOF" |& tee test-triton.log
          # autotuner tests from jax-triton now hit a triton code path that uses utilities from pytorch...
          pip install --no-deps torch --index-url https://download.pytorch.org/whl/cpu
          python /opt/jax-triton/tests/triton_call_test.py --xml_output_file /output/triton_test.xml
        EOF
      STATISTICS_SCRIPT: |
        curl -L -o yq https://github.com/mikefarah/yq/releases/latest/download/yq_linux_$(dpkg --print-architecture) && chmod 777 yq;
        total_tests=$(./yq '.testsuites."+@tests"' triton_test.xml)
        errors=$(./yq '.testsuites."+@errors"' triton_test.xml)
        failed_tests=$(./yq '.testsuites."+@failures"' triton_test.xml)
        passed_tests=$((total_tests - errors - failed_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-triton.log
    secrets: inherit

  test-levanter:
    needs: build-levanter
    if: inputs.ARCHITECTURE == 'amd64'  # arm64 runners n/a
    uses: ./.github/workflows/_test_unit.yaml
    with:
      TEST_NAME: levanter
      EXECUTE: |
        docker run -i --gpus all --shm-size=1g \
        ${{ needs.build-levanter.outputs.DOCKER_TAG_FINAL }} \
        bash <<"EOF" |& tee test-levanter.log
          pip install flake8 pytest soundfile librosa
          PYTHONPATH=/opt/levanter/tests:$PYTHONPATH pytest /opt/levanter/tests -m "not entry and not slow and not ray"
        EOF
      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

  test-te:
    needs: build-upstream-pax
    if: inputs.ARCHITECTURE == 'amd64'  # arm64 runners n/a
    uses: ./.github/workflows/_test_unit.yaml
    with:
      TEST_NAME: te
      EXECUTE: |
        docker run -i --gpus all --shm-size=1g -v $PWD:/log \
        ${{ needs.build-upstream-pax.outputs.DOCKER_TAG_FINAL }} \
        bash <<"EOF" |& tee test-te.log
          pip install pytest-reportlog
          pytest --report-log=log/pytest-report.jsonl ${SRC_PATH_TRANSFORMER_ENGINE}/tests/jax
        EOF
      STATISTICS_SCRIPT: |
        summary_line=$(tail -n1 test-te.log)
        errors=$(echo $summary_line | grep -oE '[0-9]+ error' | awk '{print $1} END { if (!NR) print 0}')
        passed_tests=$(cat pytest-report.jsonl | jq -r 'select(."$report_type" == "TestReport" and .when == "call" and .outcome == "passed") | .outcome' | wc -l)
        failed_tests=$(cat pytest-report.jsonl | jq -r 'select(."$report_type" == "TestReport" and .when == "call" and .outcome == "failed") | .outcome' | wc -l)
        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
      TIMEOUT_MINUTES: 120
      ARTIFACTS: |
        test-te.log
        pytest-report.jsonl
    secrets: inherit

  test-upstream-pax:
    needs: build-upstream-pax
    if: inputs.ARCHITECTURE == 'amd64' # no images for arm64
    uses: ./.github/workflows/_test_upstream_pax.yaml
    with:
      PAX_IMAGE: ${{ needs.build-upstream-pax.outputs.DOCKER_TAG_FINAL }}
    secrets: inherit

  test-rosetta-pax:
    needs: build-rosetta-pax
    if: inputs.ARCHITECTURE == 'amd64' # no images for arm64
    uses: ./.github/workflows/_test_pax_rosetta.yaml
    with:
      PAX_IMAGE: ${{ needs.build-rosetta-pax.outputs.DOCKER_TAG_FINAL }}
    secrets: inherit
  
  test-gemma:
    needs: build-gemma
    uses: ./.github/workflows/_test_unit.yaml  
    if: inputs.ARCHITECTURE == 'amd64'
    with:
      TEST_NAME: gemma
      EXECUTE: |
        docker run --shm-size=1g --gpus all ${{ needs.build-gemma.outputs.DOCKER_TAG_FINAL }} \
        bash -ec \
        "cd /opt/gemma && pip install -e .[dev] && pytest ." | tee test-gemma.log
      STATISTICS_SCRIPT: |
        summary_line=$(tail -n1 test-gemma.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-gemma.log
    secrets: inherit

  test-maxtext:
    needs: build-maxtext
    if: inputs.ARCHITECTURE == 'amd64' # no arm64 gpu runners
    uses: ./.github/workflows/_test_maxtext.yaml
    with:
      MAXTEXT_IMAGE: ${{ needs.build-maxtext.outputs.DOCKER_TAG_FINAL }}
    secrets: inherit