run_pathways_tests.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 unit tests with Pathways backend.

name: Run Pathways Tests

on:
  workflow_call:
    inputs:
      device_type:
        required: true
        type: string
      device_name:
        required: true
        type: string
      base_image:
        required: true
        type: string
      pytest_marker:
        required: true
        type: string
      pytest_addopts:
        required: false
        type: string
        default: ''
      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
      maxtext_sha:
        required: true
        type: string
      total_workers:
        required: false
        type: string
        default: '1'
      worker_group:
        required: false
        type: string
        default: '1'


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/${{ inputs.base_image }}
      env:
        XLA_PYTHON_CLIENT_MEM_FRACTION: ${{ inputs.xla_python_client_mem_fraction }}
        TF_FORCE_GPU_ALLOW_GROWTH: ${{ inputs.tf_force_gpu_allow_growth }}
        IFRT_PROXY_USE_INSECURE_GRPC_CREDENTIALS: true
        JAX_PLATFORMS: "proxy"
        JAX_BACKEND_TARGET: "grpc://localhost:29000"
        TPU_VISIBLE_DEVICES: ""
      options: ${{ inputs.container_resource_option }}
    steps:
      - name: Checkout MaxText
        run: |
          git config --global --add safe.directory /__w/maxtext/maxtext
          git clone https://github.com/google/maxtext.git .
          git fetch origin ${{ inputs.maxtext_sha }}
          git checkout FETCH_HEAD
      - name: Download the maxtext wheel
        uses: actions/download-artifact@v4
        with:
          name: maxtext-wheel
      - name: Install the maxtext wheel
        shell: bash
        run: |
          python3 -m uv venv --seed
          source .venv/bin/activate
          maxtext_wheel=$(ls maxtext-*-py3-none-any.whl 2>/dev/null)
          uv pip install ${maxtext_wheel}[tpu] --resolution=lowest
          install_tpu_pre_train_extra_deps
          python3 --version
          python3 -m pip freeze
      - name: Copy test assets files
        run : gcloud storage cp gs://maxtext-test-assets/* tests/assets
      - name: Run Tests
        run: |
          if [ "${{ inputs.total_workers }}" -gt 1 ]; then
            .venv/bin/python3 -m pip install --quiet pytest-split
            SPLIT_ARGS="--splits ${{ inputs.total_workers }} --group ${{ inputs.worker_group }}"
          else
            SPLIT_ARGS=""
          fi
          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
          export MAXTEXT_REPO_ROOT=$(pwd)
          export MAXTEXT_ASSETS_ROOT=$(pwd)/src/maxtext/assets
          export MAXTEXT_TEST_ASSETS_ROOT=$(pwd)/tests/assets
          export MAXTEXT_PKG_DIR=$(pwd)/src/maxtext
          # TODO(b/454659463): Enable test_default_hlo_match after volume mount is supported.
          .venv/bin/python3 -m pytest ${{ inputs.pytest_addopts }} -v -m "${FINAL_PYTEST_MARKER}" -k "not AotHloIdenticalTest and not CompileThenLoad" --durations=0 ${SPLIT_ARGS}
        env:
          PYTHONPATH: "${{ github.workspace }}/src"
    services:
      resource_manager:
        image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:latest
        ports:
          - "29001:29001"
          - "29002:29002"
        options:
          --entrypoint=[/usr/pathways/run/cloud_pathways_server_sanitized, --server_port=29001, --node_type=resource_manager, --enforce_kernel_ipv6_support=false, --instance_count=1, --instance_type=tpuv6e:2x2, --gcs_scratch_location=gs://cloud-pathways-staging/tmp]
        env:
          TPU_SKIP_MDS_QUERY: true

      worker:
        image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:latest
        ports:
          - "29005:29005"
          - "29006:29006"
          - "8471:8471"
          - "8080:8080"
        options:
          --entrypoint=[/usr/pathways/run/cloud_pathways_server_sanitized, --server_port=29005, --resource_manager_address=localhost:29001, --enforce_kernel_ipv6_support=false, --gcs_scratch_location=gs://cloud-pathways-staging/tmp]
          --tpu=4

      proxy:
        image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server:latest
        ports:
          - "29000:29000"
        env:
          IFRT_PROXY_USE_INSECURE_GRPC_CREDENTIALS: true
          XLA_FLAGS: "--xla_dump_to=/tmp/aot_test_dump --xla_dump_hlo_as_text --xla_dump_hlo_module_re=jit_train_step"
        options:
          --entrypoint=[/usr/pathways/run/cloud_proxy_server_sanitized, --server_port=29000, --resource_manager_address=localhost:29001, --enforce_kernel_ipv6_support=false, --gcs_scratch_location=gs://cloud-pathways-staging/tmp, --xla_tpu_scoped_vmem_limit_kib=65536, --xla_tpu_spmd_rng_bit_generator_unsafe=true]