run_tests_against_package.yml

# Copyright 2023-2026 Google LLC

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

#     https://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# This workflow runs tests against the built MaxText package.

name: Run Tests Against MaxText Package

on:
  workflow_call:
    inputs:
      flavor:
        required: true
        type: string
      device_type:
        required: true
        type: string
      device_name:
        required: true
        type: string
      repository_name:
        description: 'Artifact Registry repository name'
        required: false
        type: string
        default: 'maxtext-images'
      base_image:
        required: true
        type: string
      pytest_marker:
        required: true
        type: string
      pytest_extra_args:
        required: false
        type: string
      pytest_addopts:
        required: false
        type: string
        default: ''
      is_scheduled_run:
        required: true
        type: boolean
      xla_python_client_mem_fraction:
        required: true
        type: string
      tf_force_gpu_allow_growth:
        required: true
        type: string
      container_resource_option:
        required: true
        type: string
      cloud_runner:
        required: false
        type: string
      worker_group:
        required: false
        type: number
        default: 1
      total_workers:
        required: false
        type: number
        default: 1
      maxtext_sha:
        description: 'Git SHA to checkout if MaxText is not pre-installed'
        required: false
        type: string
      # Flag to skip source checkout and wheel installation
      maxtext_installed:
        description: 'If false, maxtext_sha must be provided for checkout'
        type: boolean
        default: false
      is_update_hlo:
        required: false
        type: boolean
        default: false
      install_torch_cpu:
        description: 'Install CPU-only torch for layer-wise reference tests'
        required: false
        type: boolean
        default: false
permissions:
  contents: read
jobs:
  run:
    name: ${{ inputs.flavor }}
    runs-on: ${{ inputs.cloud_runner != '' && inputs.cloud_runner || fromJson(format('["self-hosted", "{0}", "{1}"]', inputs.device_type, inputs.device_name)) }}
    timeout-minutes: 90
    container:
      image: us-docker.pkg.dev/${{ vars.PROJECT_NAME || 'tpu-prod-env-multipod' }}/${{ inputs.repository_name || 'maxtext-images' }}/${{ inputs.base_image }} # zizmor: ignore[unpinned-images]
      env:
        XLA_PYTHON_CLIENT_MEM_FRACTION: ${{ inputs.xla_python_client_mem_fraction }}
        TF_FORCE_GPU_ALLOW_GROWTH: ${{ inputs.tf_force_gpu_allow_growth }}
        TPU_SKIP_MDS_QUERY: ${{ inputs.device_type == 'cpu' && '1' || '' }}
        # cpu-unit runs the hardware-agnostic bulk (every test auto-marked cpu_only), and it is the
        # only job that runs them. On one device a test whose batch is smaller than the mesh, or that
        # never enters an axis_rules context, passes for the wrong reason; 4 virtual devices make the
        # sharding real so those bugs surface here instead of on accelerator runners. Also lets the
        # tests that self-skip below 2 devices actually run.
        XLA_FLAGS: ${{ inputs.flavor == 'cpu-unit' && '--xla_force_host_platform_device_count=4' || '' }}
        VLLM_TARGET_DEVICE: "tpu"
        UV_TORCH_BACKEND: "cpu"
        MAXTEXT_PACKAGE_EXTRA: >-
          ${{
            !contains(inputs.pytest_marker, 'not post_training') && 'tpu-post-train'
            || (inputs.device_type == 'cpu' && 'tpu' || inputs.device_type)
          }}
        ALLOW_MULTIPLE_LIBTPU_LOAD: ${{ inputs.device_type == 'cpu' && 'true' || '' }} # bypass /tmp/libtpu_lockfile check for cpu tests, which don't actually use accelerators (to allow concurrency)
      options: ${{ inputs.container_resource_option }}
    steps:
      - name: Checkout MaxText
        if: ${{ !inputs.maxtext_installed }}
        uses: actions/checkout@fbc6f3992d24b796d5a048ff273f7fcc4a7b6c09 # v5.1.0
        with:
          ref: ${{ inputs.maxtext_sha }}
          fetch-depth: 0
          persist-credentials: false
      - name: Download the maxtext wheel
        if: ${{ !inputs.maxtext_installed }}
        uses: actions/download-artifact@d3f86a106a0bac45b974a628896c90dbdf5c8093 # v4.3.0
        with:
          name: maxtext-wheel
      - name: Install the maxtext wheel
        if: ${{ !inputs.maxtext_installed }}
        shell: bash
        run: |
          python3 -m uv venv --seed
          source .venv/bin/activate
          maxtext_wheel=$(ls maxtext-*-py3-none-any.whl 2>/dev/null)
          echo "Installing ${maxtext_wheel} for ${MAXTEXT_PACKAGE_EXTRA}..."
          uv pip install "${maxtext_wheel}[${MAXTEXT_PACKAGE_EXTRA}]" --resolution=lowest
          if [ "${MAXTEXT_PACKAGE_EXTRA}" == "tpu-post-train" ]; then
            if command -v apt-get &> /dev/null; then
              echo "Installing GCC 12 for vLLM C++20 compilation in CI runner..."
              if [ -f /etc/os-release ] && grep -q "bullseye" /etc/os-release; then
                echo "deb http://deb.debian.org/debian bookworm main" | (sudo tee /etc/apt/sources.list.d/bookworm.list 2>/dev/null || tee /etc/apt/sources.list.d/bookworm.list)
                (sudo apt-get update -y && sudo apt-get install -y --no-install-recommends -t bookworm gcc-12 g++-12 build-essential cmake ninja-build) || (apt-get update -y && apt-get install -y --no-install-recommends -t bookworm gcc-12 g++-12 build-essential cmake ninja-build)
              else
                (sudo apt-get update -y && sudo apt-get install -y --no-install-recommends gcc-12 g++-12 build-essential cmake ninja-build) || (apt-get update -y && apt-get install -y --no-install-recommends gcc-12 g++-12 build-essential cmake ninja-build)
              fi
            fi
            install_tpu_post_train_extra_deps
          else
            install_tpu_pre_train_extra_deps
          fi
          if [ "${INPUTS_INSTALL_TORCH_CPU}" == "true" ]; then
            python3 -m pip install --quiet torch==2.11.0 --index-url https://download.pytorch.org/whl/cpu --no-deps
            python3 -c "import torch; print(f'torch {torch.__version__}')"
          fi
          python3 --version
          python3 -m pip freeze
        env:
          INPUTS_INSTALL_TORCH_CPU: ${{ inputs.install_torch_cpu }}
      - name: Copy test assets files
        if: ${{ !inputs.maxtext_installed }}
        run : gcloud storage cp gs://maxtext-test-assets/* tests/assets
      - name: Download Converted Test Durations
        if: inputs.total_workers > 1
        uses: actions/download-artifact@d3f86a106a0bac45b974a628896c90dbdf5c8093 # v4.3.0
        with:
          name: test-durations-${{ github.run_id }}-${{ inputs.flavor }}
          path: ${{ inputs.maxtext_installed && '/deps' || '.' }}
        continue-on-error: true
      - name: Run Tests
        shell: bash
        run: |
          # Cleanly release TPU device locks upon SIGTERM / PR cancellation
          cleanup_tpu_locks() {
            echo "Received cancellation signal; cleaning up /dev/accel* locks and /tmp/libtpu_lockfile..."
            kill -9 "${PYTEST_PID:-}" 2>/dev/null || true
            rm -f /tmp/libtpu_lockfile
            fuser -k -9 /dev/accel* 2>/dev/null || true
            sleep 0.5
            exit 143
          }
          trap cleanup_tpu_locks TERM INT

          # Determine environment and entry directory
          if [ "${INPUTS_MAXTEXT_INSTALLED}" == "true" ]; then
            # Move to the directory where code is baked into the image. See the Dockerfile.
            cd /deps
            PYTHON_EXE="python3"
            # Disable coverage flags when testing against a pre-installed package
            PYTEST_COV_ARGS=""
          else
            source .venv/bin/activate
            # Use the local virtual environment created in Step 3
            PYTHON_EXE=".venv/bin/python3"
            # Ensure pytest-cov is available and enable coverage flags
            uv pip install pytest-cov
            PYTEST_COV_ARGS="--cov=maxtext --cov-report=xml --cov-report=term"
          fi
          export PYTHONPATH="${PWD}/src${PYTHONPATH:+:${PYTHONPATH}}"

          if [ "${INPUTS_IS_SCHEDULED_RUN}" == "true" ]; then
            FINAL_PYTEST_MARKER="${INPUTS_PYTEST_MARKER}"
          else
            if [ -z "${INPUTS_PYTEST_MARKER}" ]; then
              FINAL_PYTEST_MARKER="not scheduled_only or newly_added"
            else
              FINAL_PYTEST_MARKER="${INPUTS_PYTEST_MARKER} and (not scheduled_only or newly_added)"
            fi
          fi
          # TODO: Use package data for testing and remove the env vars
          MAXTEXT_REPO_ROOT=$(pwd)
          echo "MAXTEXT_REPO_ROOT (the current working directory) is ${MAXTEXT_REPO_ROOT}."
          MAXTEXT_ASSETS_ROOT="$(pwd)/src/maxtext/assets"
          MAXTEXT_TEST_ASSETS_ROOT="$(pwd)/tests/assets"
          MAXTEXT_PKG_DIR="$(pwd)/src/maxtext"
          export MAXTEXT_REPO_ROOT MAXTEXT_ASSETS_ROOT MAXTEXT_TEST_ASSETS_ROOT MAXTEXT_PKG_DIR
          # omit this libtpu init args for gpu tests
          if [ "${INPUTS_DEVICE_TYPE}" != "cuda12" ]; then
            export LIBTPU_INIT_ARGS='--xla_tpu_scoped_vmem_limit_kib=65536'
          else
            # For cuda12, explicitly point to the pip-installed CUDA libraries
            # to avoid conflicts with system-level installations on the runner.
            # Dynamically discover the 'nvidia' folder and prepend all its sub-library
            # directories (including nccl, cublas, cudnn) to LD_LIBRARY_PATH to prevent
            # JAX from partially loading incompatible system-level CUDA libraries.
            if [ -d ".venv/lib" ]; then
              NVIDIA_DIR=$(find .venv/lib/ -maxdepth 3 -name "nvidia" -type d 2>/dev/null | head -n 1)
              if [ -n "${NVIDIA_DIR}" ]; then
                for dir in "${NVIDIA_DIR}"/*; do
                  if [ -d "$dir/lib" ]; then
                    LD_LIBRARY_PATH="$(pwd)/$dir/lib:${LD_LIBRARY_PATH}"
                    export LD_LIBRARY_PATH
                  fi
                done
              fi
            fi
          fi
          if [ "${INPUTS_TOTAL_WORKERS}" -gt 1 ] && [ "${INPUTS_DEVICE_TYPE}" = "tpu" ] && [[ "${INPUTS_FLAVOR}" == *-unit ]]; then
            # TPU unit flavors: whole test classes stay on one shard (tests/utils/class_split.py), so the first
            # test of a class pays its XLA compile in scheduled and PR runs alike. The rest keep pytest-split
            # below: one integration class (TrainTests) takes longer than an even shard share, CPU durations
            # are not regression-checked, and GPU has too few test classes to balance.
            if [ -f .test_durations ]; then
              echo "Using downloaded .test_durations to balance the class-level split."
            else
              echo "Warning: .test_durations not found. The class-level split uses equal weights."
            fi
            SPLIT_ARGS="--class-splits=${INPUTS_TOTAL_WORKERS} --class-group=${INPUTS_WORKER_GROUP}"
          elif [ "${INPUTS_TOTAL_WORKERS}" -gt 1 ]; then
            $PYTHON_EXE -m pip install --quiet pytest-split
            if [ -f .test_durations ]; then
              echo "Using downloaded .test_durations with least_duration splitting algorithm."
              SPLIT_ALGO="--splitting-algorithm least_duration"
            else
              echo "Warning: .test_durations not found. Falling back to default split."
              SPLIT_ALGO=""
            fi
            if [ "${INPUTS_DEVICE_TYPE}" = "cpu" ]; then
              $PYTHON_EXE -m pip install --quiet pytest-xdist
              SPLIT_ARGS="--splits ${INPUTS_TOTAL_WORKERS} --group ${INPUTS_WORKER_GROUP} -n auto ${SPLIT_ALGO}"
            else
              SPLIT_ARGS="--splits ${INPUTS_TOTAL_WORKERS} --group ${INPUTS_WORKER_GROUP} ${SPLIT_ALGO}"
            fi
          else
            SPLIT_ARGS=""
          fi

          # Setup substitution: If manually updating HLO, skip tests execution and run only the update script instead!
          if [ "${INPUTS_IS_UPDATE_HLO}" == "true" ]; then
            python3 tests/utils/update_hlo_references.py
          else
            # The GofN suffix lets process_test_results.py notice a shard that left no result file.
            JUNIT_XML="test-results-${INPUTS_FLAVOR}-${INPUTS_WORKER_GROUP}of${INPUTS_TOTAL_WORKERS}.xml"
            # shellcheck disable=SC2086
            $PYTHON_EXE -m pytest ${INPUTS_PYTEST_ADDOPTS} \
              -v \
              -m "${FINAL_PYTEST_MARKER}" \
              --durations=0 \
              --junitxml="${JUNIT_XML}" \
              $PYTEST_COV_ARGS \
              $SPLIT_ARGS \
              ${INPUTS_PYTEST_EXTRA_ARGS} &
            PYTEST_PID=$!
            PYTEST_EXIT_CODE=0
            wait "$PYTEST_PID" || PYTEST_EXIT_CODE=$?
            if [ "$PYTEST_EXIT_CODE" -eq 2 ] || [ "$PYTEST_EXIT_CODE" -eq 3 ]; then
              # An interrupted session (e.g. a collection error) still writes a result file, without the test times.
              # Drop it so process_test_results.py counts this shard as missing instead of saving a near-zero total.
              rm -f "${JUNIT_XML}"
            fi
            if [ "$PYTEST_EXIT_CODE" -eq 5 ]; then
              echo "No tests collected for worker group ${INPUTS_WORKER_GROUP} (pytest exit code 5), treating as success."
            elif [ "$PYTEST_EXIT_CODE" -ne 0 ]; then
              exit "$PYTEST_EXIT_CODE"
            fi
          fi

        env:
          INPUTS_IS_SCHEDULED_RUN: ${{ inputs.is_scheduled_run }}
          INPUTS_PYTEST_MARKER: ${{ inputs.pytest_marker }}
          INPUTS_DEVICE_TYPE: ${{ inputs.device_type }}
          INPUTS_FLAVOR: ${{ inputs.flavor }}
          INPUTS_PYTEST_ADDOPTS: ${{ inputs.pytest_addopts }}
          INPUTS_TOTAL_WORKERS: ${{ inputs.total_workers }}
          INPUTS_WORKER_GROUP: ${{ inputs.worker_group }}
          INPUTS_PYTEST_EXTRA_ARGS: ${{ inputs.pytest_extra_args }}
          INPUTS_MAXTEXT_INSTALLED: ${{ inputs.maxtext_installed }}
          INPUTS_IS_UPDATE_HLO: ${{ inputs.is_update_hlo }}
      # DECOUPLE_GCLOUD=TRUE only exercises the offline code paths when the Google Cloud packages
      # are actually missing, so this runs in its own venv built from the decoupled requirements
      # rather than the venv the other tests use. The suite is small and unsharded, so it runs in
      # the first worker group only.
      - name: Run Targeted Decoupled Tests
        if: ${{ !inputs.maxtext_installed && inputs.flavor == 'cpu-unit' && inputs.worker_group == 1 }}
        shell: bash
        run: |
          python3 -m uv venv --seed .venv_decoupled
          source .venv_decoupled/bin/activate
          uv pip install --resolution=lowest \
            -r src/dependencies/requirements/generated_requirements/decoupled-requirements.txt
          # --no-deps: the maxtext wheel requires the Google Cloud packages this environment must not have.
          uv pip install --no-deps "$(ls maxtext-*-py3-none-any.whl)"
          python3 - <<'PY'
          import importlib.util
          import sys

          def importable(module):
            try:
              return importlib.util.find_spec(module) is not None
            except ModuleNotFoundError:
              return False

          # google.cloud itself resolves through googleapis-common-protos, so check the clients.
          modules = (
              "cloud_accelerator_diagnostics",
              "gcsfs",
              "google.api_core",
              "google.auth",
              "google.cloud.monitoring_v3",
              "google.cloud.storage",
              "google_cloud_mldiagnostics",
              "jetstream",
              "ml_goodput_measurement",
          )
          coupled = [module for module in modules if importable(module)]
          if coupled:
            sys.exit(f"::error::{coupled} importable in the decoupled environment; decoupled mode is not being tested.")
          PY
          export DECOUPLE_GCLOUD=TRUE
          export PYTHONPATH="${PWD}/src${PYTHONPATH:+:${PYTHONPATH}}"
          MAXTEXT_REPO_ROOT=$(pwd)
          MAXTEXT_ASSETS_ROOT="$(pwd)/src/maxtext/assets"
          MAXTEXT_TEST_ASSETS_ROOT="$(pwd)/tests/assets"
          MAXTEXT_PKG_DIR="$(pwd)/src/maxtext"
          export MAXTEXT_REPO_ROOT MAXTEXT_ASSETS_ROOT MAXTEXT_TEST_ASSETS_ROOT MAXTEXT_PKG_DIR
          python3 -m pytest -v \
            -m "decoupled_target" \
            --durations=0 \
            --ignore=tests/post_training \
            --junitxml=test-results-decoupled-targeted.xml
      - name: Upload Reference HLO
        if: ${{ inputs.is_update_hlo }}
        uses: actions/upload-artifact@5d5d22a31266ced268874388b861e4b58bb5c2f3 # v4.3.1
        with:
          name: reference-hlo
          path: ${{ inputs.maxtext_installed && '/deps/tests/utils/reference_hlo_*.txt' || 'tests/utils/reference_hlo_*.txt' }}
          if-no-files-found: ignore
      - name: Upload results to Codecov
        if: ${{ !inputs.maxtext_installed }} # Skip code coverage upload for maxtext image testing
        uses: codecov/codecov-action@0fb7174895f61a3b6b78fc075e0cd60383518dac # v5
        continue-on-error: true
        with:
          token: ${{ secrets.CODECOV_TOKEN }}
          files: ./coverage.xml
          # If scheduled, upload to scheduled flag only. If PR, upload to regular flag only.
          flags: ${{ inputs.is_scheduled_run && 'scheduled' || 'regular' }}
          verbose: true
      - name: Upload Test Results XML
        if: always()
        uses: actions/upload-artifact@5d5d22a31266ced268874388b861e4b58bb5c2f3 # v4.3.1
        with:
          name: test-results-${{ inputs.flavor }}-${{ inputs.worker_group }}-${{ github.run_id }}
          path: ${{ inputs.maxtext_installed && '/deps/test-results-*.xml' || 'test-results-*.xml' }}
          retention-days: 1
          if-no-files-found: ignore