_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
  #     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
  #     EXTRA_BUILD_ARGS: |
  #       URLREF_T5X=${{ fromJson(inputs.SOURCE_URLREFS).T5X }}
  #       URLREF_AIRIO=${{ fromJson(inputs.SOURCE_URLREFS).AIRIO }}
  #   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
  #     EXTRA_BUILD_ARGS: |
  #       URLREF_PAXML=${{ fromJson(inputs.SOURCE_URLREFS).PAXML }}
  #       URLREF_PRAXIS=${{ fromJson(inputs.SOURCE_URLREFS).PRAXIS }}
  #       URLREF_LINGVO=${{ fromJson(inputs.SOURCE_URLREFS).LINGVO }}
  #   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

  build-axlearn:
    needs: build-jax
    uses: ./.github/workflows/_build.yaml
    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
      EXTRA_BUILD_ARGS: |
        URLREF_MAXTEXT=${{ fromJson(inputs.SOURCE_URLREFS).MAXTEXT }}
    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
      - 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": "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": "axlearn",      "stage": "mealkit", "priority": 500,  "tag": "${{ needs.build-axlearn.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; failure is explicitly handled below so set +e to
  #       # avoid an early abort here.
  #       set +e
  #       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
  #           # 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 --report-log=/opt/output/pytest-report.jsonl "${test_path}"
  #       EOF
  #       set -e
  #       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-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' # arm64 runners n/a
  #   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@v4
  #     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' # arm64 runners n/a
  #   runs-on: eks
  #   env:
  #     JAX_DOCKER_IMAGE: ${{ needs.build-jax.outputs.DOCKER_TAG_FINAL }}
  #     JOB_NAME: ${{ github.run_id }}-${{ github.run_attempt }}-jax
  #     POSTPROCESS_JOB_NAME: ${{ github.run_id }}-${{ github.run_attempt }}-postprocess
  #     TOKEN_NAME: ${{ github.run_id }}-${{ github.run_attempt }}-token
  #   steps:
  #   - name: Check out the repository
  #     uses: actions/checkout@v4
  #   - name: Login to GitHub Container Registry
  #     uses: docker/login-action@v3
  #     with:
  #       registry: ghcr.io
  #       username: ${{ github.repository_owner }}
  #       password: ${{ secrets.GITHUB_TOKEN }}
  #   - name: Store GitHub Container Registry token as Kubernetes secret
  #     run: |
  #       kubectl create secret generic \
  #         ${{ github.run_id }}-${{ github.run_attempt }}-token \
  #         --from-file=.dockerconfigjson=$HOME/.docker/config.json \
  #         --type=kubernetes.io/dockerconfigjson
  #   - name: Configure Kubernetes job
  #     run: |
  #       yq -i ea 'select(di == 0).spec.selector.job-name = strenv(JOB_NAME)
  #         | select(di == 1).metadata.name = strenv(JOB_NAME)
  #         | select(di == 1).spec.template.spec.imagePullSecrets[].name = strenv(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)' \
  #         .github/eks-workflow-files/job.yml
  #       git diff .github/eks-workflow-files/job.yml
  #   - name: Submit Kubernetes job
  #     run: kubectl apply -f .github/eks-workflow-files/job.yml
  #   - name: Wait for Kubernetes job to start
  #     run: |
  #       while [[ -n $(kubectl get pods --selector=batch.kubernetes.io/job-name=${{ github.run_id }}-${{ github.run_attempt }}-jax --output=jsonpath='{.items[?(@.status.phase == "Pending")].metadata.name}') ]]; do
  #         sleep 2
  #       done
  #   - name: Stream Kubernetes job output
  #     run: kubectl logs --all-containers=true --all-pods=true --follow job/${{ github.run_id }}-${{ github.run_attempt }}-jax
  #   # Clean up in case of errors as well as success
  #   - name: Delete Kubernetes job
  #     if: always()
  #     run: kubectl delete job ${{ github.run_id }}-${{ github.run_attempt }}-jax
  #   - 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 = strenv(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-processing Kubernetes job
  #     run: kubectl apply -f .github/eks-workflow-files/post-process-job.yml
  #   - name: Wait for post-processing Kubernetes job to start
  #     run: |
  #       while [[ -n $(kubectl get pods --selector=batch.kubernetes.io/job-name=${{ github.run_id }}-${{ github.run_attempt }}-postprocess --output=jsonpath='{.items[?(@.status.phase == "Pending")].metadata.name}') ]]; do
  #         sleep 2
  #       done
  #   - name: Stream post-processing Kubernetes job output
  #     run: kubectl logs --all-containers=true --all-pods=true --follow job/${{ github.run_id }}-${{ github.run_attempt }}-postprocess
  #   # Clean up in case of errors as well as success
  #   - name: Delete post-processing Kubernetes job
  #     if: always()
  #     run: kubectl delete job ${{ github.run_id }}-${{ github.run_attempt }}-postprocess
  #   - name: Delete GitHub Container Registry token
  #     if: always()
  #     run: kubectl delete secret ${{ github.run_id }}-${{ github.run_attempt }}-token

  # # 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; this relies on
  #         # actually having a CUDA backend for pytoch
  #         pip install --no-deps torch
  #         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

  # test-axlearn-slurm:
  #   needs: build-axlearn
  #   if: inputs.ARCHITECTURE == 'amd64'
  #   uses: ./.github/workflows/_test_unit.yaml
  #   with: # fix the arguments below
  #     TEST_NAME: axlearn
  #     EXECUTE: |
  #       docker run -i --shm-size=1g --gpus all \
  #       ${{ needs.build-axlearn.outputs.DOCKER_TAG_FINAL }} \
  #       bash <<"EOF" |& tee test-backend-independent.log
  #         test-axlearn.sh --directory "." --output "/opt/output/" --test-files "/opt/axlearn/axlearn/common/*_test.py" 
  #       EOF
  #     STATISTICS_SCRIPT: |
  #       # Parse the summary.txt file to count passed/failed/error tests
  #       # Adjust greps if your output format changes.
  #       passed_tests=$(grep -c ": PASSED" /opt/output/summary.txt || true)
  #       failed_tests=$(grep -c ": FAILED" /opt/output/summary.txt || true)
  #       total_tests=$((failed_tests + passed_tests))

  #       echo "TOTAL_TESTS=${total_tests}" >> $GITHUB_OUTPUT
  #       echo "PASSED_TESTS=${passed_tests}" >> $GITHUB_OUTPUT
  #       echo "FAILED_TESTS=${failed_tests}" >> $GITHUB_OUTPUT
  #     ARTIFACTS: |
  #       test-backend-independent.log
  #   secrets: inherit


  test-axlearn-eks:
    needs: build-axlearn
    if: inputs.ARCHITECTURE == 'amd64'
    runs-on: eks
    env:
      AXLEARN_DOCKER_IMAGE: ${{ needs.build-axlearn.outputs.DOCKER_TAG_FINAL }}
      JOB_NAME: axlearn-${{ github.run_id }}
      TOKEN_NAME: axlearn-${{ github.run_id }}-token
    steps:
    - name: Set date env var for saving files
      run: |
        echo "DATE_TEST_RAN=$(date +'%Y-%m-%d-%H-%M-%S')" >> $GITHUB_ENV
    - name: Check and GHCR Login
      uses: /.github/actions/checkout-ghcr-login
      with: 
          docker-username: ${{ github.repository_owner }}
          password: ${{ secrets.GITHUB_TOKEN }}
          token-name: ${{ env.TOKEN_NAME }}
    - 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[1].env[0].value = strenv(DATE_TEST_RAN)
          | select(di == 0).spec.template.spec.imagePullSecrets[].name = strenv(TOKEN_NAME)' \
        .github/eks-workflow-files/axlearn/axlearn-job.yml
        git diff .github/eks-workflow-files/axlearn/axlearn-job.yml
    - name: Submit & wait for axlearn test job
      uses: ./.github/actions/submit-k8s-job 
      with:
        job-config-file: ".github/eks-workflow-files/axlearn/axlearn-job.yml"
        job-name: ${{ env.JOB_NAME }}

    - name: Delete axlearn test job
      uses: ./.github/actions/delete-k8s-job
      with: 
        job-name: ${{ env.JOB_NAME }}

    - name: Download logs from S3
      run: |
        mkdir -p /tmp/axlearn-output
        aws s3 cp s3://jax-toolbox-eks-output/axlearn/${{ env.DATE_TEST_RAN }}/summary.txt /tmp/axlearn-output/

        passed_tests=$(grep -c ": PASSED" /tmp/axlearn-output/summary.txt || true)
        failed_tests=$(grep -c ": FAILED" /tmp/axlearn-output/summary.txt || true)
        total_tests=$((failed_tests + passed_tests))

        echo "Passed tests: $passed_tests"
        echo "Failed tests: $failed_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.test-stats.outputs.TOTAL_TESTS }} \
        failed_tests=${{ steps.test-stats.outputs.FAILED_TESTS }} \
        passed_tests=${{ steps.test-stats.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"

    - name: Upload artifacts
      if: "!cancelled()"
      uses: actions/upload-artifact@v4
      with:
        name: "artifact-axlearn-test"
        path: |
          sitrep.json
          "badge-axlearn-test"
          summary.txt
    - name: Delete GitHub Container Registry token
      uses: ./.github/actions/delete-ghcr-token
      with: 
        token-name: ${{ env.TOKEN_NAME }}