_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: '{}'
      MODE:
        type: string
        description: 'Mode selection for running specific tests only'
        required: false
        default: full
    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

  test-nccl:
    if: inputs.ARCHITECTURE == 'amd64' # build only amd64
    needs: build-base
    uses: ./.github/workflows/_test_nccl.yaml
    with:
      CONTAINER: ${{ needs.build-base.outputs.DOCKER_TAG }}
    secrets: inherit

  build-jax:
    needs: build-base
    runs-on: [self-hosted, "${{ inputs.ARCHITECTURE }}", "large"]
    steps:
        - name: Checkout repository
          uses: actions/checkout@v6
        - name: Build JAX container
          id: build-jax
          uses: ./.github/actions/build-container
          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
            ssh-private-key: ${{ secrets.SSH_PRIVATE_KEY }}
            ssh-known-hosts: ${{ vars.SSH_KNOWN_HOSTS }}
            github-token: ${{ secrets.GITHUB_TOKEN }}
            bazel-remote-cache-url: ${{ vars.BAZEL_REMOTE_CACHE_URL }}
            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 }}
    outputs:
      DOCKER_TAG_MEALKIT: ${{ steps.build-jax.outputs.DOCKER_TAG_MEALKIT }}
      DOCKER_TAG_FINAL:   ${{ steps.build-jax.outputs.DOCKER_TAG_FINAL }}

  build-equinox:
    needs: build-jax
    runs-on: [self-hosted, "${{ inputs.ARCHITECTURE }}", "small"]
    outputs:
      DOCKER_TAG_MEALKIT: ${{ steps.build-equinox.outputs.DOCKER_TAG_MEALKIT }}
      DOCKER_TAG_FINAL:   ${{ steps.build-equinox.outputs.DOCKER_TAG_FINAL }}
    steps:
      - name: Checkout repository
        uses: actions/checkout@v6
      - name: Build Equinox container
        id: build-equinox
        uses: ./.github/actions/build-container
        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
          RUNNER_SIZE: small
          ssh-private-key: ${{ secrets.SSH_PRIVATE_KEY }}
          ssh-known-hosts: ${{ vars.SSH_KNOWN_HOSTS }}
          github-token: ${{ secrets.GITHUB_TOKEN }}
          bazel-remote-cache-url: ${{ vars.BAZEL_REMOTE_CACHE_URL }}
          EXTRA_BUILD_ARGS: |
            URLREF_EQUINOX=${{ fromJson(inputs.SOURCE_URLREFS).EQUINOX }}

  build-maxtext:
    needs: build-jax
    runs-on: [self-hosted, "${{ inputs.ARCHITECTURE }}", "small"]
    outputs:
      DOCKER_TAG_MEALKIT: ${{ steps.build-maxtext.outputs.DOCKER_TAG_MEALKIT }}
      DOCKER_TAG_FINAL:   ${{ steps.build-maxtext.outputs.DOCKER_TAG_FINAL }}
    steps:
      - name: Checkout repository
        uses: actions/checkout@v6
      - name: Build MaxText container
        id: build-maxtext
        uses: ./.github/actions/build-container
        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
          RUNNER_SIZE: small
          ssh-private-key: ${{ secrets.SSH_PRIVATE_KEY }}
          ssh-known-hosts: ${{ vars.SSH_KNOWN_HOSTS }}
          github-token: ${{ secrets.GITHUB_TOKEN }}
          bazel-remote-cache-url: ${{ vars.BAZEL_REMOTE_CACHE_URL }}
          EXTRA_BUILD_ARGS: |
            URLREF_MAXTEXT=${{ fromJson(inputs.SOURCE_URLREFS).MAXTEXT }}

  build-torchax:
    needs: build-jax
    runs-on: [self-hosted, "${{ inputs.ARCHITECTURE }}", "small"]
    outputs:
      DOCKER_TAG_MEALKIT: ${{ steps.build-torchax.outputs.DOCKER_TAG_MEALKIT }}
      DOCKER_TAG_FINAL:   ${{ steps.build-torchax.outputs.DOCKER_TAG_FINAL }}
    steps:
      - name: Checkout repository
        uses: actions/checkout@v6
      - name: Build TorchAX container
        id: build-torchax
        uses: ./.github/actions/build-container
        with:
          ARCHITECTURE: ${{ inputs.ARCHITECTURE }}
          ARTIFACT_NAME: artifact-torchax-build
          BADGE_FILENAME: badge-torchax-build
          BUILD_DATE: ${{ inputs.BUILD_DATE }}
          BASE_IMAGE: ${{ needs.build-jax.outputs.DOCKER_TAG_MEALKIT }}
          CONTAINER_NAME: torchax
          DOCKERFILE: .github/container/Dockerfile.torchax
          RUNNER_SIZE: small
          ssh-private-key: ${{ secrets.SSH_PRIVATE_KEY }}
          ssh-known-hosts: ${{ vars.SSH_KNOWN_HOSTS }}
          github-token: ${{ secrets.GITHUB_TOKEN }}
          bazel-remote-cache-url: ${{ vars.BAZEL_REMOTE_CACHE_URL }}
          EXTRA_BUILD_ARGS: |
            URLREF_TORCHAX=${{ fromJson(inputs.SOURCE_URLREFS).TORCHAX }}

  build-axlearn:
    needs: build-jax
    runs-on: [self-hosted, "${{ inputs.ARCHITECTURE }}", "large"]
    outputs:
      DOCKER_TAG_MEALKIT: ${{ steps.build-axlearn.outputs.DOCKER_TAG_MEALKIT }}
      DOCKER_TAG_FINAL:   ${{ steps.build-axlearn.outputs.DOCKER_TAG_FINAL }}
    steps:
      - name: Checkout repository
        uses: actions/checkout@v6
      - name: Build AxLearn container
        id: build-axlearn
        uses: ./.github/actions/build-container
        with:
          ARCHITECTURE: ${{ inputs.ARCHITECTURE }}
          ARTIFACT_NAME: artifact-axlearn-build
          BADGE_FILENAME: badge-axlearn-build
          BUILD_DATE: ${{ inputs.BUILD_DATE }}
          BASE_IMAGE: ${{ needs.build-jax.outputs.DOCKER_TAG_MEALKIT }}
          CONTAINER_NAME: axlearn
          DOCKERFILE: .github/container/Dockerfile.axlearn
          RUNNER_SIZE: large
          ssh-private-key: ${{ secrets.SSH_PRIVATE_KEY }}
          ssh-known-hosts: ${{ vars.SSH_KNOWN_HOSTS }}
          github-token: ${{ secrets.GITHUB_TOKEN }}
          bazel-remote-cache-url: ${{ vars.BAZEL_REMOTE_CACHE_URL }}
          EXTRA_BUILD_ARGS: |
            URLREF_AXLEARN=${{ fromJson(inputs.SOURCE_URLREFS).AXLEARN }}

  collect-docker-tags:
    runs-on: ubuntu-22.04
    if: ${{ !cancelled() }}
    needs:
      - build-base
      - build-jax
      - build-equinox
      - build-maxtext
      - build-axlearn
    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": "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": "axlearn",      "stage": "final",   "priority": 900,  "tag": "${{ needs.build-axlearn.outputs.DOCKER_TAG_FINAL }}"},\
            {"flavor": "jax",          "stage": "mealkit", "priority": 500,  "tag": "${{ needs.build-jax.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": "axlearn",      "stage": "mealkit", "priority": 500,  "tag": "${{ needs.build-axlearn.outputs.DOCKER_TAG_MEALKIT }}"},\

            {}\
          ]
          EOF
          )

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

  test-jax:
    needs: build-jax
    if: >-
      inputs.ARCHITECTURE == 'amd64' &&
      (
        inputs.MODE == 'full' ||
        inputs.MODE == 'jax'
      )
    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 test-single-gpu.log
          nvidia-cuda-mps-control -d
          test-jax.sh -b single-gpu
        EOF
        docker run -i --shm-size=1g --gpus all \
        ${{ needs.build-jax.outputs.DOCKER_TAG_FINAL }} \
        bash <<"EOF" |& tee test-multi-gpu.log
          nvidia-cuda-mps-control -d
          test-jax.sh -b multi-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-multi-gpu.log
        test-single-gpu.log
    secrets: inherit

  test-nsys-jax:
    needs: build-jax
    if: >-
      inputs.ARCHITECTURE == 'amd64' &&
      (
        inputs.MODE == 'full' ||
        inputs.MODE == 'jax'
      )
    uses: ./.github/workflows/_test_unit.yaml
    with:
      TEST_NAME: nsys-jax
      EXECUTE: |
        set -o pipefail
        mkdir -p output-results
        docker run -i --shm-size=1g --gpus all \
          -v $PWD/output-results:/opt/output \
          ${{ needs.build-jax.outputs.DOCKER_TAG_FINAL }} \
          bash <<"EOF" |& tee test-nsys-jax.log
            # nsys-jax is already installed, this is just adding the test dependencies
            pip install pytest-reportlog nsys-jax[test]
            # abuse knowledge that nsys-jax is installed editable, so the tests exist
            test_path=$(python -c 'import importlib.resources; print(importlib.resources.files("nsys_jax").joinpath("..", "tests").resolve())')
            pytest --basetemp=/opt/output/pytest-tmp --report-log=/opt/output/pytest-report.jsonl "${test_path}"
            chmod -R a+rwX /opt/output
        EOF
      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 output-results/pytest-report.jsonl | jq -r 'select(."$report_type" == "TestReport" and .when == "call" and .outcome == "passed") | .outcome' | wc -l)
        failed_tests=$(cat output-results/pytest-report.jsonl | jq -r 'select(."$report_type" == "TestReport" and .when == "call" and .outcome == "failed") | .outcome' | wc -l)
        total_tests=$(( passed_tests + failed_tests ))
        echo "TOTAL_TESTS=${total_tests}" >> $GITHUB_OUTPUT
        echo "ERRORS=${num_errors}" >> $GITHUB_OUTPUT
        echo "PASSED_TESTS=${passed_tests}" >> $GITHUB_OUTPUT
        echo "FAILED_TESTS=${failed_tests}" >> $GITHUB_OUTPUT
      ARTIFACTS: |
        # pytest-driven part
        test-nsys-jax.log
        output-results/pytest-report.jsonl
        output-results/pytest-tmp/
    secrets: inherit

  # test-nsys-jax generates several fresh .zip archive outputs by running nsys-jax with real GPU hardware; this test
  # runs on a regular GitHub Actions runner and checks that offline post-processing works in an environment that does
  # not already have nsys-jax installed
  test-nsys-jax-archive:
    needs: test-nsys-jax
    if: >-
      inputs.ARCHITECTURE == 'amd64' &&
      (
        inputs.MODE == 'full' ||
        inputs.MODE == 'jax'
      )
    strategy:
      matrix:
        os: [ubuntu-22.04, ubuntu-24.04, macOS-latest]
    runs-on: ${{ matrix.os }}
    steps:
    - name: Download nsys-jax output .zip files
      uses: actions/download-artifact@v8
      with:
        name: nsys-jax-unit-test-A100
    - name: Extract archives and execute install scripts
      run: |
        pip install virtualenv # for install.sh
        for zip in $(ls *.zip); do
          ZIP="${PWD}/${zip}"
          pushd $(mktemp -d)
          unzip "${ZIP}"
          ls -l
          # TODO: verify this isn't needed, or make sure it isn't needed
          chmod 755 install.sh
          # Run the notebook with IPython, not Jupyter Lab, so it exits and prints something informative to stdout
          # Skip executing Jupyter lab
          NSYS_JAX_JUPYTER_EXECUTE_NOT_LAB=1 ./install.sh
          popd
        done

  test-nsys-jax-eks:
    needs: build-jax
    if: >-
      inputs.ARCHITECTURE == 'amd64' &&
      (
        inputs.MODE == 'full' ||
        inputs.MODE == 'jax'
      )
    runs-on: eks
    env:
      JAX_DOCKER_IMAGE: ${{ needs.build-jax.outputs.DOCKER_TAG_FINAL }}
      JOB_NAME: ${{ github.run_id }}-nsys-jax
      # Service name cannot start with a number
      SERVICE_NAME: svc-${{ github.run_id}}-nsys-jax
      POSTPROCESS_JOB_NAME: ${{ github.run_id }}-nsys-jax-postprocess
    steps:
    - name: Check out the repository
      uses: actions/checkout@v6
    - name: Login to GitHub Container Registry
      uses: docker/login-action@v4
      with:
        registry: ghcr.io
        username: ${{ github.repository_owner }}
        password: ${{ secrets.GITHUB_TOKEN }}
    - name: K8s GHCR store and delete token
      id: store-token
      uses: ./.github/actions/store-delete-k8s-ghcr
    - name: Configure Kubernetes job
      run: |
        yq -i ea 'select(di == 0).spec.selector.job-name = strenv(JOB_NAME)
          | select(di == 0).metadata.name = strenv(SERVICE_NAME)
          | select(di == 1).metadata.name = strenv(JOB_NAME)
          | select(di == 1).spec.template.spec.subdomain = strenv(SERVICE_NAME)
          | select(di == 1).spec.template.spec.imagePullSecrets[].name = "${{ steps.store-token.outputs.token-name }}"
          | select(di == 1).spec.template.spec.containers[0].image = strenv(JAX_DOCKER_IMAGE)
          | select(di == 1).spec.template.spec.containers[0].env[0].value = strenv(JOB_NAME)
          | select(di == 1).spec.template.spec.containers[0].env[1].value = strenv(SERVICE_NAME)' \
          .github/eks-workflow-files/job.yml
        git diff .github/eks-workflow-files/job.yml
    - name: Submit Kubernetes job
      uses: ./.github/actions/submit-delete-k8s-job
      with:
        job-config-file: .github/eks-workflow-files/job.yml
        job-name: ${{ env.JOB_NAME }}
    - name: Configure post-processing job
      run: |
        export JOB_OUTPUT_PATTERN="${JOB_NAME}-rank*.zip"
        yq -i '.metadata.name = strenv(POSTPROCESS_JOB_NAME)
          | .spec.template.spec.containers[].image = strenv(JAX_DOCKER_IMAGE)
          | .spec.template.spec.imagePullSecrets[].name = "${{ steps.store-token.outputs.token-name }}"
          | .spec.template.spec.initContainers[].command[7] = strenv(JOB_OUTPUT_PATTERN)' \
          .github/eks-workflow-files/post-process-job.yml
        git diff .github/eks-workflow-files/post-process-job.yml
    - name: Submit post process Kubernetes job
      uses: ./.github/actions/submit-delete-k8s-job
      with:
        job-config-file: .github/eks-workflow-files/post-process-job.yml
        job-name: ${{ env.POSTPROCESS_JOB_NAME }}

  test-te-h100:
    needs: build-jax
    if: >-
      inputs.ARCHITECTURE == 'amd64' &&
      (
        inputs.MODE == 'full' ||
        inputs.MODE == 'te'
      )
    uses: ./.github/workflows/_transformer_engine_eks.yaml
    with:
      JAX_IMAGE: ${{ needs.build-jax.outputs.DOCKER_TAG_FINAL }}
      JOB_NAME: transformerengine-${{ github.run_id }}
      S3_BUCKET: jax-toolbox-eks-output
      CI_NAME: transformer-engine
    secrets: inherit

  test-jax-cutlass-h100:
    needs: build-jax
    if: >-
      inputs.ARCHITECTURE == 'amd64' &&
      (
        inputs.MODE == 'full' ||
        inputs.MODE == 'jax-cutlass'
      )
    uses: ./.github/workflows/_jax_cutlass_eks.yaml
    with:
      JAX_IMAGE: ${{ needs.build-jax.outputs.DOCKER_TAG_FINAL }}
      JOB_NAME: jax-cutlass-${{ github.run_id }}
      S3_BUCKET: jax-toolbox-eks-output
      CI_NAME: jax-cutlass
    secrets: inherit

  test-te-a100:
    needs: build-jax
    secrets: inherit
    if: >-
      inputs.ARCHITECTURE == 'amd64' &&
      (
        inputs.MODE == 'full' ||
        inputs.MODE == 'te'
      )
    uses: ./.github/workflows/_test_unit.yaml
    with:
      TEST_NAME: te
      EXECUTE: |
        docker run -i --gpus all --shm-size=1g -v $PWD:/log \
        ${{ needs.build-jax.outputs.DOCKER_TAG_FINAL }} \
        bash <<"EOF" |& tee test-te.log
          set -xu -o pipefail

          LOG_DIR=/log

          pip install pytest-reportlog pytest-xdist
          # Start MPS daemon
          nvidia-cuda-mps-control -d
          # TE's default is slightly different, without the hyphen
          export TE_PATH=${SRC_PATH_TRANSFORMER_ENGINE}
          # 1 GPU per worker, 3 workers per GPU
          pytest-xdist.sh 1 3 ${LOG_DIR}/pytest-report-L0-unittest.jsonl bash ${TE_PATH}/qa/L0_jax_unittest/test.sh
          ## 8 GPUs per worker, 1 worker per GPU. pytest-xdist.sh allows aggregation
          ## into a single .jsonl file of results from multiple pytest invocations
          ## inside the test.sh script, so it's useful even with a single worker per
          ## device.
          pytest-xdist.sh 8 1 ${LOG_DIR}/pytest-report-L0-distributed-unittest.jsonl bash ${TE_PATH}/qa/L0_jax_distributed_unittest/test.sh

          # merge the log files
          cat \
            ${LOG_DIR}/pytest-report-L0-unittest.jsonl \
            ${LOG_DIR}/pytest-report-L0-distributed-unittest.jsonl \
            > ${LOG_DIR}/pytest-report.jsonl

        EOF
      STATISTICS_SCRIPT: |
        report_json=pytest-report.jsonl
        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 $report_json | jq -r 'select(."$report_type" == "TestReport" and .when == "call" and .outcome == "passed") | .outcome' | wc -l)
        failed_tests=$(cat $report_json | 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

        echo "$failed_tests tests failed"
        if [[ $failed_tests -gt 0 ]]; then
            exit 1
        else
            exit 0
        fi

      TIMEOUT_MINUTES: 120
      ARTIFACTS: |
        test-te.log
        pytest-report.jsonl
        pytest-report-L0-unittest.jsonl
        pytest-report-L0-distributed-unittest.jsonl

  test-maxtext:
    needs: build-maxtext
    if: >-
      inputs.ARCHITECTURE == 'amd64' &&
      (
        inputs.MODE == 'full' ||
        inputs.MODE == 'maxtext'
      )
    uses: ./.github/workflows/_test_maxtext.yaml
    with:
      MAXTEXT_IMAGE: ${{ needs.build-maxtext.outputs.DOCKER_TAG_FINAL }}
    secrets: inherit

  test-maxtext-gke:
    needs: build-maxtext
    if: >-
      inputs.ARCHITECTURE == 'amd64' &&
      (
        inputs.MODE == 'full' ||
        inputs.MODE == 'maxtext'
      )
    uses: ./.github/workflows/_test_maxtext_gke_xpk.yaml
    with:
      MAXTEXT_IMAGE: ${{ needs.build-maxtext.outputs.DOCKER_TAG_FINAL }}
    secrets: inherit

  test-axlearn-eks:
    needs: build-axlearn
    if: >-
      inputs.ARCHITECTURE == 'amd64' &&
      (
        inputs.MODE == 'full' ||
        inputs.MODE == 'axlearn'
      )
    runs-on: eks
    env:
      AXLEARN_DOCKER_IMAGE: ${{ needs.build-axlearn.outputs.DOCKER_TAG_FINAL }}
      JOB_NAME: axlearn-${{ github.run_id }}
    steps:
    - name: Check out the repository
      uses: actions/checkout@v6
    - name: Login to GitHub Container Registry
      uses: docker/login-action@v4
      with:
        registry: ghcr.io
        username: ${{ github.repository_owner }}
        password: ${{ secrets.GITHUB_TOKEN }}
    - name: K8s GHCR store and delete token
      id: store-token
      uses: ./.github/actions/store-delete-k8s-ghcr
    - name: Configure axlearn test job
      run: |
        # Replace placeholders in axlearn-job.yml with environment variables
        yq -i ea '
           select(di == 0).metadata.name = strenv(JOB_NAME)
          | select(di == 0).spec.template.spec.containers[0].image = strenv(AXLEARN_DOCKER_IMAGE)
          | select(di == 0).spec.template.spec.containers[0].env[0].value = "${{ github.run_id }}"
          | select(di == 0).spec.template.spec.imagePullSecrets[].name = "${{ steps.store-token.outputs.token-name }}"' \
        .github/eks-workflow-files/axlearn/axlearn-job.yml
        git diff .github/eks-workflow-files/axlearn/axlearn-job.yml
    - name: Submit & delete axlearn test
      uses: ./.github/actions/submit-delete-k8s-job
      with:
        job-config-file: ".github/eks-workflow-files/axlearn/axlearn-job.yml"
        job-name: ${{ env.JOB_NAME }}
    - name: Download logs from S3
      id: log-s3
      if: ${{ !cancelled() }}
      run: |
        mkdir -p axlearn-output
        aws s3 cp s3://jax-toolbox-eks-output/axlearn/${{ github.run_id }}/logs/summary.txt axlearn-output/
        aws s3 cp s3://jax-toolbox-eks-output/axlearn/${{ github.run_id }}/logs/ axlearn-output/ --recursive --exclude "*" --include "*.log"
        aws s3 cp s3://jax-toolbox-eks-output/axlearn/${{ github.run_id }}/logs/ axlearn-output/ --recursive --exclude "*" --include "*.xml"


        passed_tests=$(grep -Eo 'PASSED:[[:space:]]*[0-9]+' axlearn-output/summary.txt | grep -Eo '[0-9]+' )
        failed_tests=$(grep -Eo 'FAILED:[[:space:]]*[0-9]+' axlearn-output/summary.txt | grep -Eo '[0-9]+' )
        skipped_tests=$(grep -Eo 'SKIPPED:[[:space:]]*[0-9]+' axlearn-output/summary.txt | grep -Eo '[0-9]+' )
        total_tests=$((failed_tests + passed_tests + skipped_tests))

        echo "Passed tests: $passed_tests"
        echo "Failed tests: $failed_tests"
        echo "Skipped tests: $skipped_tests"
        echo "Total tests: $total_tests"
        echo "PASSED_TESTS=$passed_tests" >> $GITHUB_OUTPUT
        echo "FAILED_TESTS=$failed_tests" >> $GITHUB_OUTPUT
        echo "TOTAL_TESTS=$total_tests" >> $GITHUB_OUTPUT
    - name: Generate sitrep
      id: sitrep
      if: ${{ !cancelled() }}
      shell: bash -x -e {0}
      run: |
        # bring in utility functions
        source .github/workflows/scripts/to_json.sh

        badge_label='Axlearn EKS Unit'

        total_tests=${{ steps.log-s3.outputs.TOTAL_TESTS }} \
        failed_tests=${{ steps.log-s3.outputs.FAILED_TESTS }} \
        passed_tests=${{ steps.log-s3.outputs.PASSED_TESTS }} \
        errors="0" \
        summary="All tests: $total_tests. Passed: $passed_tests. Failed: $failed_tests." \
        badge_message="Passed $passed_tests out of $total_tests." \
        badge_color="brightgreen"
        if [ "$failed_tests" -gt 0 ]; then
          badge_color="red"
        fi \

        to_json \
          summary \
          errors total_tests passed_tests failed_tests \
          badge_label badge_color badge_message \
        > sitrep.json

        schemaVersion=1 \
        label="${badge_label}" \
        message="Passed $passed_tests out of $total_tests." \
        color=$badge_color \
        to_json schemaVersion label message color \
        > badge-axlearn-test.json

    - name: Upload artifacts
      if: ${{ !cancelled() }}
      uses: actions/upload-artifact@v6
      with:
        name: "artifact-axlearn-test"
        path: |
          sitrep.json
          badge-axlearn-test.json
          axlearn-unittests.jsonl
          axlearn-output/*

  test-axlearn-fuji-models-eks:
    needs: build-axlearn
    if: >-
      inputs.ARCHITECTURE == 'amd64' &&
      (
        inputs.MODE == 'full' ||
        inputs.MODE == 'axlearn'
      )
    runs-on: eks
    env:
      AXLEARN_DOCKER_IMAGE: ${{ needs.build-axlearn.outputs.DOCKER_TAG_FINAL }}
      JOB_NAME: axlearn-fuji-3b-${{ github.run_id }}
    steps:
    - name: Check out the repository
      uses: actions/checkout@v6
    - name: Login to GitHub Container Registry
      uses: docker/login-action@v4
      with:
        registry: ghcr.io
        username: ${{ github.repository_owner }}
        password: ${{ secrets.GITHUB_TOKEN }}
    - name: K8s GHCR store and delete token
      id: store-token
      uses: ./.github/actions/store-delete-k8s-ghcr
    - name: Configure axlearn test job
      run: |
        yq -i ea '
           select(di == 0).metadata.name = strenv(JOB_NAME)
          | select(di == 0).spec.template.spec.containers[0].image = strenv(AXLEARN_DOCKER_IMAGE)
          | select(di == 0).spec.template.spec.imagePullSecrets[].name = "${{ steps.store-token.outputs.token-name }}"' \
        .github/eks-workflow-files/axlearn/axlearn-fuji-model.yml
        git diff .github/eks-workflow-files/axlearn/axlearn-fuji-model.yml

    - name: Submit & delete axlearn fuji model test
      uses: ./.github/actions/submit-delete-k8s-job
      with:
        job-config-file:  ".github/eks-workflow-files/axlearn/axlearn-fuji-model.yml"
        job-name: ${{ env.JOB_NAME }}