run_tests_against_package.yml

# Copyright 2025 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 file defines a module for running tests against the built maxtext package.

name: Run Tests Against MaxText Package

on:
  workflow_call:
    inputs:
      device_type:
        required: true
        type: string
      device_name:
        required: true
        type: string
      image_type:
        required: false
        type: string
      pytest_marker:
        required: true
        type: string
      is_scheduled_run:
        required: true
        type: string
      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

permissions:
  contents: read
jobs:
  run:
    runs-on: ${{ inputs.cloud_runner != '' && inputs.cloud_runner || fromJson(format('["self-hosted", "{0}", "{1}"]', inputs.device_type, inputs.device_name)) }}
    container:
      image: gcr.io/tpu-prod-env-multipod/maxtext-unit-test-${{ inputs.device_type == 'cpu' && 'tpu' || inputs.device_type }}:${{ inputs.image_type != '' && inputs.image_type }}
      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' || '' }}
      options: ${{ inputs.container_resource_option }}
    steps:
      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
      - name: Download the maxtext wheel
        uses: actions/download-artifact@634f93cb2916e3fdff6788551b99b062d0335ce0 # v5.0.0
        with:
          name: maxtext-wheel
      - name: Install the maxtext wheel
        run: |
          python3 -m uv pip install maxtext-*-py3-none-any.whl --resolution=lowest --system
          install_maxtext_github_deps
          python3 --version
          python3 -m pip freeze
      - name: Copy test assets files
        run : gcloud storage cp gs://maxtext-test-assets/* src/MaxText/test_assets
      - name: Display contents of workspace # TODO: Remove
        run: ls -R
      - name: Run Tests
        run: |
          if [ "${{ inputs.is_scheduled_run }}" == "true" ]; then
            FINAL_PYTEST_MARKER="${{ inputs.pytest_marker }}"
          else
            FINAL_PYTEST_MARKER="${{ inputs.pytest_marker }} and not scheduled_only"
          fi
          # TODO: Use package data for testing and remove the env vars
          export MAXTEXT_REPO_ROOT=$(pwd)
          export MAXTEXT_ASSETS_ROOT=$(pwd)/src/MaxText/assets
          export MAXTEXT_TEST_ASSETS_ROOT=$(pwd)/src/MaxText/test_assets
          export MAXTEXT_PKG_DIR=$(pwd)/src/MaxText
          # TODO: Fix the skipped tests and remove the deselect flags
          python3 -m pytest -v -m "${FINAL_PYTEST_MARKER}" --durations=0 --deselect "tests/aot_hlo_identical_test.py::AotHloIdenticalTest::test_default_hlo_match" --deselect "tests/tokenizer_test.py::TokenizerTest::test_detokenize"