From db054c387ed35c0b9c2e06b45f70746bed731361 Mon Sep 17 00:00:00 2001 From: Alfred Yu Date: Mon, 3 Nov 2025 17:28:58 +0800 Subject: [PATCH 1/3] [cienet-private] Disable all integration tests --- .github/workflows/CPUTests.yml | 56 ----- .github/workflows/RunTests.yml | 236 ------------------ .github/workflows/UploadDockerImages.yml | 176 ------------- .github/workflows/build_and_test_maxtext.yml | 156 ------------ .github/workflows/build_package.yml | 52 ---- .github/workflows/build_upload_internal.yml | 69 ----- .github/workflows/check_docs_build.yml | 31 --- .../workflows/run_tests_against_package.yml | 114 --------- .github/workflows/run_tests_internal.yml | 84 ------- .github/workflows/stale.yml | 49 ---- 10 files changed, 1023 deletions(-) delete mode 100644 .github/workflows/CPUTests.yml delete mode 100644 .github/workflows/RunTests.yml delete mode 100644 .github/workflows/UploadDockerImages.yml delete mode 100644 .github/workflows/build_and_test_maxtext.yml delete mode 100644 .github/workflows/build_package.yml delete mode 100644 .github/workflows/build_upload_internal.yml delete mode 100644 .github/workflows/check_docs_build.yml delete mode 100644 .github/workflows/run_tests_against_package.yml delete mode 100644 .github/workflows/run_tests_internal.yml delete mode 100644 .github/workflows/stale.yml diff --git a/.github/workflows/CPUTests.yml b/.github/workflows/CPUTests.yml deleted file mode 100644 index 737d1a5f22..0000000000 --- a/.github/workflows/CPUTests.yml +++ /dev/null @@ -1,56 +0,0 @@ -# Copyright 2023–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. - -name: Linter - -on: - pull_request: - -concurrency: - # Dedup pull requests (canceling previous runs of the same workflow for same PR), and scheduled runs but nothing else - group: > - ${{ - github.event_name == 'pull_request' && format('{0}-pr-{1}', github.workflow, github.event.pull_request.number) || - github.event_name == 'schedule' && format('{0}-schedule', github.workflow) || - github.run_id - }} - cancel-in-progress: true - -jobs: - cpu: - name: "CPU tests" - runs-on: ubuntu-latest - strategy: - matrix: - os: [ubuntu-24.04] - python-version: ['3.12'] - steps: - - uses: actions/checkout@v5 - - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v5 - with: - python-version: ${{ matrix.python-version }} - - name: Install Dependencies - run: | - python3 -m pip install --upgrade pip - python3 -m pip install pylint pyink pytype==2024.2.27 - - name: Typecheck the code with pytype - run: | - pytype --jobs auto --disable 'import-error,late-directive,wrong-arg-types,module-attr,unsupported-operands' src/MaxText/ || true - - name: pylint - run: | - pylint --disable=R0401,R0917,W0201,W0613 --ignore-patterns='.pytype,.*pyi$' benchmarks end_to_end src tests - - name: pyink - run: | - pyink --pyink-indentation=2 --line-length=122 --check . diff --git a/.github/workflows/RunTests.yml b/.github/workflows/RunTests.yml deleted file mode 100644 index 0690278891..0000000000 --- a/.github/workflows/RunTests.yml +++ /dev/null @@ -1,236 +0,0 @@ -# Copyright 2023–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 workflow will install Python dependencies, run tests and lint with a variety of Python versions -# For more information see: https://docs.github.com/en/actions/automating-builds-and-tests/building-and-testing-python - -name: Tests - -on: - pull_request: - workflow_dispatch: - schedule: - # Run the job every 4 hours - - cron: '0 */4 * * *' - -concurrency: - # Dedup pull requests (canceling previous runs of the same workflow for same PR), and scheduled runs but nothing else - group: > - ${{ - github.event_name == 'pull_request' && format('{0}-pr-{1}', github.workflow, github.event.pull_request.number) || - github.event_name == 'schedule' && format('{0}-schedule', github.workflow) || - github.run_id - }} - cancel-in-progress: true - -jobs: - prelim: - runs-on: ["self-hosted"] - steps: - - name: Test gsutil installation - run: which gsutil >/dev/null 2>&1 || { echo >&2 "gsutil is required but not installed. Aborting"; exit 24;} - - name: Cleanup old docker images - run: docker system prune --all --force - - tpu_image: - needs: prelim - uses: ./.github/workflows/build_upload_internal.yml - with: - device_type: tpu - device_name: v4-8 - cloud_runner: linux-x86-n2-16-buildkit - build_mode: jax_ai_image - base_image: us-docker.pkg.dev/tpu-prod-env-multipod/jax-stable-stack/candidate/tpu:latest - - gpu_image: - needs: prelim - uses: ./.github/workflows/build_upload_internal.yml - with: - device_type: gpu - device_name: a100-40gb-4 - cloud_runner: linux-x86-n2-16-buildkit - build_mode: jax_ai_image - base_image: us-docker.pkg.dev/tpu-prod-env-multipod/jax-stable-stack/candidate/gpu:latest - - cpu_unit_tests: - needs: tpu_image - strategy: - fail-fast: false - matrix: - worker_group: [1, 2, 3, 4] - uses: ./.github/workflows/run_tests_internal.yml - with: - device_type: cpu - device_name: X64 - image_type: tpu - pytest_marker: 'cpu_only' - xla_python_client_mem_fraction: 0.75 - tf_force_gpu_allow_growth: false - container_resource_option: "--privileged" - is_scheduled_run: ${{ github.event_name == 'schedule' }} - worker_group: ${{ matrix.worker_group }} - total_workers: 4 - - tpu_unit_tests: - needs: tpu_image - uses: ./.github/workflows/run_tests_internal.yml - with: - device_type: tpu - device_name: v4-8 - cloud_runner: linux-x86-ct4p-240-4tpu - pytest_marker: 'not cpu_only and not gpu_only and not integration_test' - xla_python_client_mem_fraction: 0.75 - tf_force_gpu_allow_growth: false - container_resource_option: "--privileged" - is_scheduled_run: ${{ github.event_name == 'schedule' }} - - tpu_integration_tests: - needs: tpu_image - uses: ./.github/workflows/run_tests_internal.yml - with: - device_type: tpu - device_name: v4-8 - cloud_runner: linux-x86-ct4p-240-4tpu - pytest_marker: 'not cpu_only and not gpu_only and integration_test' - xla_python_client_mem_fraction: 0.75 - tf_force_gpu_allow_growth: false - container_resource_option: "--privileged" - is_scheduled_run: ${{ github.event_name == 'schedule' }} - - gpu_unit_tests: - needs: gpu_image - uses: ./.github/workflows/run_tests_internal.yml - with: - device_type: gpu - device_name: a100-40gb-4 - cloud_runner: linux-x86-a2-48-a100-4gpu - pytest_marker: 'not cpu_only and not tpu_only and not integration_test' - pytest_addopts: '--ignore=tests/sft_hooks_test.py' - xla_python_client_mem_fraction: 0.65 - tf_force_gpu_allow_growth: true - container_resource_option: "--shm-size 2g --runtime=nvidia --gpus all --privileged" - is_scheduled_run: ${{ github.event_name == 'schedule' }} - - gpu_integration_tests: - needs: gpu_image - uses: ./.github/workflows/run_tests_internal.yml - with: - device_type: gpu - device_name: a100-40gb-4 - cloud_runner: linux-x86-a2-48-a100-4gpu - pytest_marker: 'not cpu_only and not tpu_only and integration_test' - pytest_addopts: '--ignore=tests/sft_hooks_test.py' - xla_python_client_mem_fraction: 0.65 - tf_force_gpu_allow_growth: true - container_resource_option: "--shm-size 2g --runtime=nvidia --gpus all --privileged" - is_scheduled_run: ${{ github.event_name == 'schedule' }} - - clean_up: - if: ${{ always() }} - needs: [cpu_unit_tests, gpu_unit_tests, gpu_integration_tests, tpu_unit_tests, tpu_integration_tests] - name: "Clean up" - runs-on: ["self-hosted"] - permissions: - contents: read - issues: write - steps: - - name: Authenticate gcloud - run: | - # configure registries as root and as runner - gcloud auth configure-docker --quiet - gcloud auth configure-docker us-docker.pkg.dev --quiet - - name: Delete the tpu image - run: gcloud container images delete "gcr.io/tpu-prod-env-multipod/maxtext_${{ github.run_id }}:tpu" --force-delete-tags --quiet - - name: Delete the gpu image - run: gcloud container images delete "gcr.io/tpu-prod-env-multipod/maxtext_${{ github.run_id }}:gpu" --force-delete-tags --quiet - - notify_failure: - name: Notify failed build # creates an issue or modifies last open existing issue for failed build - needs: [cpu_unit_tests, gpu_unit_tests, gpu_integration_tests, tpu_unit_tests, tpu_integration_tests] - if: ${{ always() }} - runs-on: ubuntu-latest - permissions: - issues: write - steps: - - name: Check whether one of the jobs failed - id: report_failure - if: ${{ contains(needs.*.result, 'failure') && github.event.pull_request == null && github.event_name != 'workflow_dispatch' }} - uses: jayqi/failed-build-issue-action@v1.2 - with: - github-token: ${{ secrets.GITHUB_TOKEN }} - labels: "failed-build" - - - name: Reset consecutive success counter on failure - # This runs only if the previous step actually found or created an issue - if: ${{ steps.report_failure.outputs.issue-number != '' }} - env: - GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - GH_REPO: ${{ github.repository }} - # Use the issue number from the previous step's output - ISSUE_NUMBER: ${{ steps.report_failure.outputs.issue-number }} - run: | - echo "A failure occurred. Resetting success counter on issue #${ISSUE_NUMBER}." - # This command will attempt to remove both labels. - # It will not fail if the labels don't exist. - gh issue remove-label $ISSUE_NUMBER "success-run-1" "success-run-2" --repo $GH_REPO || echo "No success labels to remove." - - notify_success_and_close: - name: Close issue after 3 successful builds - # This job runs only if all the preceding test jobs succeeded - if: ${{ success() && github.event.pull_request == null && github.event_name != 'workflow_dispatch' }} - needs: [cpu_unit_tests, gpu_unit_tests, gpu_integration_tests, tpu_unit_tests, tpu_integration_tests] - runs-on: ubuntu-latest - permissions: - issues: write - steps: - - name: Find existing failure issue - id: find_issue - env: - GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - GH_REPO: ${{ github.repository }} - run: | - ISSUE_NUMBER=$(gh issue list --label "failed-build" --state open --limit 1 --json number -q '.[0].number') - if [[ -z "$ISSUE_NUMBER" ]]; then - echo "No open build failure issue found. Nothing to do." - echo "issue_number=" >> $GITHUB_OUTPUT - else - echo "Found open build failure issue: #${ISSUE_NUMBER}" - echo "issue_number=${ISSUE_NUMBER}" >> $GITHUB_OUTPUT - fi - - - name: Add success label or close issue - if: steps.find_issue.outputs.issue_number != '' - env: - GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - GH_REPO: ${{ github.repository }} - run: | - ISSUE_NUMBER=${{ steps.find_issue.outputs.issue_number }} - LABELS=$(gh issue view $ISSUE_NUMBER --json labels -q '.labels[].name') - - if echo "$LABELS" | grep -q "success-run-2"; then - echo "Third consecutive success. Closing issue #${ISSUE_NUMBER}." - gh issue comment $ISSUE_NUMBER --body "Build succeeded for the third consecutive time. Closing this issue automatically." - gh issue close $ISSUE_NUMBER - # Clean up all tracking labels - gh issue remove-label $ISSUE_NUMBER "failed-build" "success-run-2" --repo $GH_REPO - elif echo "$LABELS" | grep -q "success-run-1"; then - echo "Second consecutive success. Updating label on issue #${ISSUE_NUMBER}." - gh issue comment $ISSUE_NUMBER --body "Build succeeded for the second time. One more successful run will close this issue." - gh issue remove-label $ISSUE_NUMBER "success-run-1" --repo $GH_REPO - gh issue add-label $ISSUE_NUMBER "success-run-2" --repo $GH_REPO - else - echo "First consecutive success since failure. Adding label to issue #${ISSUE_NUMBER}." - gh issue comment $ISSUE_NUMBER --body "Build succeeded. This issue will be auto-closed after two more consecutive successful runs." - gh issue add-label $ISSUE_NUMBER "success-run-1" --repo $GH_REPO - fi diff --git a/.github/workflows/UploadDockerImages.yml b/.github/workflows/UploadDockerImages.yml deleted file mode 100644 index eb4c769678..0000000000 --- a/.github/workflows/UploadDockerImages.yml +++ /dev/null @@ -1,176 +0,0 @@ -# Copyright 2023–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 workflow will install Python dependencies, run tests and lint with a variety of Python versions -# For more information see: https://docs.github.com/en/actions/automating-builds-and-tests/building-and-testing-python - -name: Build Images - -on: - schedule: - # Run the job daily at 12AM UTC - - cron: '0 0 * * *' - workflow_dispatch: - inputs: - target_device: - description: 'Specify target device (all, tpu, or gpu)' - required: true - type: choice - default: 'all' - options: - - all - - tpu - - gpu - -jobs: - build-tpu: - # This job will only run for 'tpu', 'all', schedule, or PR triggers. - if: > - github.event_name == 'schedule' || - github.event_name == 'pull_request' || - github.event.inputs.target_device == 'all' || - github.event.inputs.target_device == 'tpu' - - runs-on: linux-x86-n2-16-buildkit - container: google/cloud-sdk:524.0.0 - - # Use Github Actions matrix to run image builds in parallel - strategy: - fail-fast: false - matrix: - include: - # TPU Image Builds - - image_name: maxtext_jax_stable - dockerfile: ./maxtext_dependencies.Dockerfile - build_args: | - MODE=stable - JAX_VERSION=NONE - LIBTPU_GCS_PATH=NONE - - image_name: maxtext_jax_nightly - dockerfile: ./maxtext_dependencies.Dockerfile - build_args: | - MODE=nightly - JAX_VERSION=NONE - LIBTPU_GCS_PATH=NONE - # TPU Image builds using JAX AI Image - - image_name: maxtext_jax_stable_stack - dockerfile: ./maxtext_jax_ai_image.Dockerfile - base_image: us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest - - image_name: maxtext_stable_stack_nightly_jax - dockerfile: ./maxtext_jax_ai_image.Dockerfile - base_image: us-docker.pkg.dev/tpu-prod-env-multipod/jax-stable-stack/tpu/jax_nightly:latest - - image_name: maxtext_stable_stack_candidate - dockerfile: ./maxtext_jax_ai_image.Dockerfile - base_image: us-docker.pkg.dev/tpu-prod-env-multipod/jax-stable-stack/candidate/tpu:latest - - # Setup for GKE runners per b/412986220#comment82 and b/412986220#comment90 - steps: - - uses: actions/checkout@v5 - - name: Mark git repository as safe - run: git config --global --add safe.directory ${GITHUB_WORKSPACE} - - name: Configure Docker - run: gcloud auth configure-docker us-docker.pkg.dev,gcr.io -q - - name: Set up Docker BuildX - uses: docker/setup-buildx-action@v3.11.1 - with: - driver: remote - endpoint: tcp://localhost:1234 - # Env variables to be passed to Dockerfile - - name: Get short commit hash - id: vars - run: echo "sha_short=$(git rev-parse --short HEAD)" >> $GITHUB_OUTPUT - - name: Get current date - id: date - run: echo "image_date=$(date +%Y-%m-%d)" >> $GITHUB_OUTPUT - # Docker BuildX command config - - name: Build and Push Docker Image - uses: docker/build-push-action@v6 - with: - push: true - context: . - file: ${{ matrix.dockerfile }} - tags: | - gcr.io/tpu-prod-env-multipod/${{ matrix.image_name }}:${{ steps.date.outputs.image_date }} - gcr.io/tpu-prod-env-multipod/${{ matrix.image_name }}:latest - cache-from: type=gha - cache-to: type=gha,mode=max - provenance: false - build-args: | - ${{ matrix.build_args }} - JAX_AI_IMAGE_BASEIMAGE=${{ matrix.base_image }} - COMMIT_HASH=${{ steps.vars.outputs.sha_short }} - DEVICE=tpu - TEST_TYPE=xlml - - # Same as tpu-build step but mirrored for GPUs - build-gpu: - if: > - github.event_name == 'schedule' || - github.event_name == 'pull_request' || - github.event.inputs.target_device == 'all' || - github.event.inputs.target_device == 'gpu' - - runs-on: linux-x86-n2-16-buildkit - container: google/cloud-sdk:524.0.0 - - strategy: - fail-fast: false - matrix: - # GPU Image Builds using JAX AI Image - include: - - image_name: maxtext_gpu_jax_stable_stack - dockerfile: ./maxtext_jax_ai_image.Dockerfile - base_image: us-central1-docker.pkg.dev/deeplearning-images/jax-ai-image/gpu:latest - - image_name: maxtext_gpu_stable_stack_nightly_jax - dockerfile: ./maxtext_jax_ai_image.Dockerfile - base_image: us-docker.pkg.dev/tpu-prod-env-multipod/jax-stable-stack/gpu/jax_nightly:latest - - image_name: maxtext_stable_stack_candidate_gpu - dockerfile: ./maxtext_jax_ai_image.Dockerfile - base_image: us-docker.pkg.dev/tpu-prod-env-multipod/jax-stable-stack/candidate/gpu:latest - - steps: - - uses: actions/checkout@v5 - - name: Mark git repository as safe - run: git config --global --add safe.directory ${GITHUB_WORKSPACE} - - name: Configure Docker - run: gcloud auth configure-docker us-docker.pkg.dev,gcr.io,us-central1-docker.pkg.dev -q - - name: Set up Docker BuildX - uses: docker/setup-buildx-action@v3.11.1 - with: - driver: remote - endpoint: tcp://localhost:1234 - - name: Get short commit hash - id: vars - run: echo "sha_short=$(git rev-parse --short HEAD)" >> $GITHUB_OUTPUT - - name: Get current date - id: date - run: echo "image_date=$(date +%Y-%m-%d)" >> $GITHUB_OUTPUT - - name: Build and Push Docker Image - uses: docker/build-push-action@v6 - with: - push: true - context: . - file: ${{ matrix.dockerfile }} - tags: | - gcr.io/tpu-prod-env-multipod/${{ matrix.image_name }}:${{ steps.date.outputs.image_date }} - gcr.io/tpu-prod-env-multipod/${{ matrix.image_name }}:latest - cache-from: type=gha - cache-to: type=gha,mode=max - provenance: false - build-args: | - ${{ matrix.build_args }} - JAX_AI_IMAGE_BASEIMAGE=${{ matrix.base_image }} - COMMIT_HASH=${{ steps.vars.outputs.sha_short }} - DEVICE=gpu - TEST_TYPE=xlml diff --git a/.github/workflows/build_and_test_maxtext.yml b/.github/workflows/build_and_test_maxtext.yml deleted file mode 100644 index 112e52fb2c..0000000000 --- a/.github/workflows/build_and_test_maxtext.yml +++ /dev/null @@ -1,156 +0,0 @@ -# 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 workflow will build maxtext python package and run tests. - -name: MaxText Package Tests - -on: - pull_request: - workflow_dispatch: - schedule: - # Run the job every 4 hours - - cron: '0 */4 * * *' - -concurrency: - # Dedup pull requests (canceling previous runs of the same workflow for same PR), and scheduled runs but nothing else - group: > - ${{ - github.event_name == 'pull_request' && format('{0}-pr-{1}', github.workflow, github.event.pull_request.number) || - github.event_name == 'schedule' && format('{0}-schedule', github.workflow) || - github.run_id - }} - cancel-in-progress: true - -permissions: - contents: read -jobs: - build_and_upload_maxtext_package: - name: Build and upload maxtext package - uses: ./.github/workflows/build_package.yml - with: - device_type: tpu - device_name: v4-8 - cloud_runner: linux-x86-n2-16-buildkit - - maxtext_cpu_unit_tests: - needs: build_and_upload_maxtext_package - uses: ./.github/workflows/run_tests_against_package.yml - strategy: - fail-fast: false # don't cancel all jobs on failure - matrix: - image_type: ["py312"] - worker_group: [1, 2, 3, 4] - with: - device_type: cpu - device_name: X64 - cloud_runner: linux-x86-n2-16 - image_type: ${{ matrix.image_type }} - pytest_marker: 'cpu_only' - xla_python_client_mem_fraction: 0.75 - tf_force_gpu_allow_growth: false - container_resource_option: "--privileged" - is_scheduled_run: ${{ github.event_name == 'schedule' }} - worker_group: ${{ matrix.worker_group }} - total_workers: 4 - - maxtext_tpu_unit_tests: - needs: build_and_upload_maxtext_package - uses: ./.github/workflows/run_tests_against_package.yml - strategy: - fail-fast: false - matrix: - image_type: ["py312"] - with: - device_type: tpu - device_name: v4-8 - image_type: ${{ matrix.image_type }} - cloud_runner: linux-x86-ct4p-240-4tpu - pytest_marker: 'not cpu_only and not gpu_only and not integration_test' - xla_python_client_mem_fraction: 0.75 - tf_force_gpu_allow_growth: false - container_resource_option: "--privileged" - is_scheduled_run: ${{ github.event_name == 'schedule' }} - - maxtext_tpu_integration_tests: - needs: build_and_upload_maxtext_package - uses: ./.github/workflows/run_tests_against_package.yml - strategy: - fail-fast: false - matrix: - image_type: ["py312"] - with: - device_type: tpu - device_name: v4-8 - image_type: ${{ matrix.image_type }} - cloud_runner: linux-x86-ct4p-240-4tpu - pytest_marker: 'not cpu_only and not gpu_only and integration_test' - xla_python_client_mem_fraction: 0.75 - tf_force_gpu_allow_growth: false - container_resource_option: "--privileged" - is_scheduled_run: ${{ github.event_name == 'schedule' }} - - maxtext_gpu_unit_tests: - needs: build_and_upload_maxtext_package - uses: ./.github/workflows/run_tests_against_package.yml - strategy: - fail-fast: false - matrix: - image_type: ["py312"] - cuda: ["cuda12"] - with: - device_type: ${{ matrix.cuda }} - device_name: a100-40gb-4 - image_type: ${{ matrix.image_type }} - cloud_runner: linux-x86-a2-48-a100-4gpu - pytest_marker: 'not cpu_only and not tpu_only and not integration_test' - pytest_addopts: '--ignore=tests/sft_hooks_test.py' - xla_python_client_mem_fraction: 0.65 - tf_force_gpu_allow_growth: true - container_resource_option: "--shm-size 2g --runtime=nvidia --gpus all --privileged" - is_scheduled_run: ${{ github.event_name == 'schedule' }} - - maxtext_gpu_integration_tests: - needs: build_and_upload_maxtext_package - uses: ./.github/workflows/run_tests_against_package.yml - strategy: - fail-fast: false - matrix: - image_type: ["py312"] - cuda: ["cuda12"] - with: - device_type: ${{ matrix.cuda }} - device_name: a100-40gb-4 - image_type: ${{ matrix.image_type }} - cloud_runner: linux-x86-a2-48-a100-4gpu - pytest_marker: 'not cpu_only and not tpu_only and integration_test' - pytest_addopts: '--ignore=tests/sft_hooks_test.py' - xla_python_client_mem_fraction: 0.65 - tf_force_gpu_allow_growth: true - container_resource_option: "--shm-size 2g --runtime=nvidia --gpus all --privileged" - is_scheduled_run: ${{ github.event_name == 'schedule' }} - - notify_failure: - name: Notify failed build # creates an issue or modifies last open existing issue for failed build - needs: [maxtext_cpu_unit_tests, maxtext_tpu_unit_tests, maxtext_tpu_integration_tests, maxtext_gpu_unit_tests, maxtext_gpu_integration_tests] - if: ${{ always() }} - runs-on: ubuntu-latest - permissions: - issues: write - steps: - - name: Check whether one of the jobs failed - if: ${{ contains(needs.*.result, 'failure') && github.event.pull_request == null && github.event_name != 'workflow_dispatch' }} - uses: jayqi/failed-build-issue-action@1a893bbf43ef1c2a8705e2b115cd4f0fe3c5649b # v1.2.0 - with: - github-token: ${{ secrets.GITHUB_TOKEN }} diff --git a/.github/workflows/build_package.yml b/.github/workflows/build_package.yml deleted file mode 100644 index 15487f1daf..0000000000 --- a/.github/workflows/build_package.yml +++ /dev/null @@ -1,52 +0,0 @@ -# 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 building and uploading a maxtext pacakge -# based on the pyproject.toml at the current github workspace. - -name: Build and Upload MaxText Package - -on: - workflow_call: - inputs: - device_type: - required: true - type: string - device_name: - required: true - type: string - cloud_runner: - required: false - type: string - -permissions: - contents: read -jobs: - build_and_upload: - name: Build and upload maxtext package (${{ inputs.device_name }}) - runs-on: ${{ inputs.cloud_runner != '' && inputs.cloud_runner || fromJson(format('["self-hosted", "{0}", "{1}"]', inputs.device_type, inputs.device_name)) }} - container: python:3.12.3-slim-bullseye - steps: - - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 - - name: Install build tools - run: | - python -m pip install --upgrade pip build uv - - name: Build maxtext wheel - run: | - uv build --wheel - - name: Upload the built maxtext wheel - uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4.6.2 - with: - name: maxtext-wheel - path: dist/* diff --git a/.github/workflows/build_upload_internal.yml b/.github/workflows/build_upload_internal.yml deleted file mode 100644 index 3a073ec5ad..0000000000 --- a/.github/workflows/build_upload_internal.yml +++ /dev/null @@ -1,69 +0,0 @@ -# Copyright 2023–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 building and uploading an image used in UnitTests.yml - -name: Build and Upload Image - -on: - workflow_call: - inputs: - device_type: - required: true - type: string - device_name: - required: true - type: string - build_mode: - required: true - type: string - base_image: - required: false - type: string - cloud_runner: - required: false - type: string - -jobs: - build_and_upload: - name: Build and upload image (${{ inputs.device_name }}) - runs-on: ${{ inputs.cloud_runner != '' && inputs.cloud_runner || fromJson(format('["self-hosted", "{0}", "{1}"]', inputs.device_type, inputs.device_name)) }} - container: google/cloud-sdk:524.0.0 - steps: - - uses: actions/checkout@v5 - - name: Mark git repository as safe - run: git config --global --add safe.directory ${GITHUB_WORKSPACE} - - name: Configure Docker - run: gcloud auth configure-docker us-docker.pkg.dev,gcr.io,us-central1-docker.pkg.dev -q - - name: Set up Docker BuildX - uses: docker/setup-buildx-action@v3.11.1 - with: - driver: remote - endpoint: tcp://localhost:1234 - - name: Get short commit hash - id: vars - run: echo "sha_short=$(git rev-parse --short HEAD)" >> $GITHUB_OUTPUT - - name: Build and push - uses: docker/build-push-action@v6 - with: - push: true - context: . - file: ./maxtext_jax_ai_image.Dockerfile - tags: gcr.io/tpu-prod-env-multipod/maxtext_${{ github.run_id }}:${{ inputs.device_type }} - provenance: false - build-args: | - JAX_AI_IMAGE_BASEIMAGE=${{ inputs.base_image }} - COMMIT_HASH=${{ steps.vars.outputs.sha_short }} - DEVICE=${{ inputs.device_type }} - TEST_TYPE=unit_test diff --git a/.github/workflows/check_docs_build.yml b/.github/workflows/check_docs_build.yml deleted file mode 100644 index 880480c4d4..0000000000 --- a/.github/workflows/check_docs_build.yml +++ /dev/null @@ -1,31 +0,0 @@ -name: Check the Documentation Build - -on: - pull_request: - types: [opened, synchronize] - workflow_dispatch: - -permissions: - contents: read - -jobs: - build-docs: - runs-on: ubuntu-latest - steps: - - name: Checkout repository - uses: actions/checkout@08c6903cd8c0fde910a37f88322edcfb5dd907a8 # v5.0.0 - with: - persist-credentials: false - - - name: Set up Python - uses: actions/setup-python@e797f83bcb11b83ae66e0230d6156d7c80228e7c # v6.0.0 - with: - python-version: '3.12' - cache: 'pip' # caching pip dependencies - - - name: Install dependencies - run: pip install -r requirements_docs.txt - - - name: Build documentation - run: | - sphinx-build -W -b html docs docs/_build/html diff --git a/.github/workflows/run_tests_against_package.yml b/.github/workflows/run_tests_against_package.yml deleted file mode 100644 index 6076efd38a..0000000000 --- a/.github/workflows/run_tests_against_package.yml +++ /dev/null @@ -1,114 +0,0 @@ -# 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 - 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 - worker_group: - required: false - type: number - default: 1 - total_workers: - required: false - type: number - 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/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' || '' }} - MAXTEXT_PACKAGE_EXTRA: ${{ inputs.device_type == 'cpu' && 'tpu' || inputs.device_type }} - 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 - 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}[${MAXTEXT_PACKAGE_EXTRA}] --resolution=lowest - 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 - shell: bash - 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 - # 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' - fi - # TODO: Fix the skipped tests and remove the deselect flags - [ "${{ inputs.total_workers }}" -gt 1 ] && .venv/bin/python3 -m pip install --quiet pytest-split && SPLIT_ARGS="--splits ${{ inputs.total_workers }} --group ${{ inputs.worker_group }}" || SPLIT_ARGS="" - .venv/bin/python3 -m pytest ${{ inputs.pytest_addopts }} -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" $SPLIT_ARGS diff --git a/.github/workflows/run_tests_internal.yml b/.github/workflows/run_tests_internal.yml deleted file mode 100644 index 3647bf83fb..0000000000 --- a/.github/workflows/run_tests_internal.yml +++ /dev/null @@ -1,84 +0,0 @@ -# Copyright 2023–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 used in UnitTests.yml - -name: Run Tests - -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 - 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 - worker_group: - required: false - type: number - default: 1 - total_workers: - required: false - type: number - default: 1 - -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_${{ github.run_id }}:${{ inputs.image_type != '' && inputs.image_type || inputs.device_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.image_type == 'tpu' && inputs.device_type != 'tpu' && '1' || '' }} - options: ${{ inputs.container_resource_option }} - steps: - - uses: actions/checkout@v5 - - 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 - python3 -m pip install -e . --no-dependencies - [ "${{ inputs.total_workers }}" -gt 1 ] && python3 -m pip install --quiet pytest-split && SPLIT_ARGS="--splits ${{ inputs.total_workers }} --group ${{ inputs.worker_group }}" || SPLIT_ARGS="" - export LIBTPU_INIT_ARGS='--xla_tpu_scoped_vmem_limit_kib=65536' - python3 -m pytest ${{ inputs.pytest_addopts }} -v -m "${FINAL_PYTEST_MARKER}" --durations=0 $SPLIT_ARGS diff --git a/.github/workflows/stale.yml b/.github/workflows/stale.yml deleted file mode 100644 index b8430add34..0000000000 --- a/.github/workflows/stale.yml +++ /dev/null @@ -1,49 +0,0 @@ -# Copyright 2023–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. - -name: 'Close stale PRs' -on: - schedule: - - cron: '0 16 * * *' - -jobs: - stale: - runs-on: ubuntu-latest - permissions: - actions: write - issues: write - pull-requests: write - steps: - - uses: actions/stale@v9 - with: - # Disable the action for issues to run on PRs only - days-before-issue-stale: -1 - days-before-issue-close: -1 - # Number of days of inactivity before a PR is marked as stale - days-before-pr-stale: 30 - # Number of days of inactivity before a stale PR is closed - days-before-pr-close: 7 - # The label to apply to stale PRs - stale-pr-label: 'stale' - # This is true by default, adding just to be explicit - remove-stale-when-updated: true - # The comment to post when marking a PR as stale. - stale-pr-message: > - This PR has been automatically marked as stale because it has not had - recent activity. It will be closed soon if no further activity occurs. - Thank you for your contributions. - # The comment to post when closing a stale PR. - close-pr-message: > - This PR was closed because it has been inactive for a while. - Please reopen it if you are still working on it. From e5134ed61a65a2390fb9113d07fe93185031c8a2 Mon Sep 17 00:00:00 2001 From: Alfred Yu Date: Mon, 3 Nov 2025 17:28:58 +0800 Subject: [PATCH 2/3] [cienet-private] Disable all integration tests --- .github/workflows/CPUTests.yml | 56 ---- .github/workflows/RunTests.yml | 249 ------------------ .github/workflows/UploadDockerImages.yml | 176 ------------- .github/workflows/build_and_test_maxtext.yml | 156 ----------- .github/workflows/build_package.yml | 52 ---- .github/workflows/build_upload_internal.yml | 69 ----- .github/workflows/check_docs_build.yml | 31 --- .../workflows/run_pathways_tests_internal.yml | 114 -------- .../workflows/run_tests_against_package.yml | 114 -------- .github/workflows/run_tests_internal.yml | 84 ------ .github/workflows/stale.yml | 49 ---- 11 files changed, 1150 deletions(-) delete mode 100644 .github/workflows/CPUTests.yml delete mode 100644 .github/workflows/RunTests.yml delete mode 100644 .github/workflows/UploadDockerImages.yml delete mode 100644 .github/workflows/build_and_test_maxtext.yml delete mode 100644 .github/workflows/build_package.yml delete mode 100644 .github/workflows/build_upload_internal.yml delete mode 100644 .github/workflows/check_docs_build.yml delete mode 100644 .github/workflows/run_pathways_tests_internal.yml delete mode 100644 .github/workflows/run_tests_against_package.yml delete mode 100644 .github/workflows/run_tests_internal.yml delete mode 100644 .github/workflows/stale.yml diff --git a/.github/workflows/CPUTests.yml b/.github/workflows/CPUTests.yml deleted file mode 100644 index 737d1a5f22..0000000000 --- a/.github/workflows/CPUTests.yml +++ /dev/null @@ -1,56 +0,0 @@ -# Copyright 2023–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. - -name: Linter - -on: - pull_request: - -concurrency: - # Dedup pull requests (canceling previous runs of the same workflow for same PR), and scheduled runs but nothing else - group: > - ${{ - github.event_name == 'pull_request' && format('{0}-pr-{1}', github.workflow, github.event.pull_request.number) || - github.event_name == 'schedule' && format('{0}-schedule', github.workflow) || - github.run_id - }} - cancel-in-progress: true - -jobs: - cpu: - name: "CPU tests" - runs-on: ubuntu-latest - strategy: - matrix: - os: [ubuntu-24.04] - python-version: ['3.12'] - steps: - - uses: actions/checkout@v5 - - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v5 - with: - python-version: ${{ matrix.python-version }} - - name: Install Dependencies - run: | - python3 -m pip install --upgrade pip - python3 -m pip install pylint pyink pytype==2024.2.27 - - name: Typecheck the code with pytype - run: | - pytype --jobs auto --disable 'import-error,late-directive,wrong-arg-types,module-attr,unsupported-operands' src/MaxText/ || true - - name: pylint - run: | - pylint --disable=R0401,R0917,W0201,W0613 --ignore-patterns='.pytype,.*pyi$' benchmarks end_to_end src tests - - name: pyink - run: | - pyink --pyink-indentation=2 --line-length=122 --check . diff --git a/.github/workflows/RunTests.yml b/.github/workflows/RunTests.yml deleted file mode 100644 index c07d7e5ac9..0000000000 --- a/.github/workflows/RunTests.yml +++ /dev/null @@ -1,249 +0,0 @@ -# Copyright 2023–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 workflow will install Python dependencies, run tests and lint with a variety of Python versions -# For more information see: https://docs.github.com/en/actions/automating-builds-and-tests/building-and-testing-python - -name: Tests - -on: - pull_request: - workflow_dispatch: - schedule: - # Run the job every 4 hours - - cron: '0 */4 * * *' - -concurrency: - # Dedup pull requests (canceling previous runs of the same workflow for same PR), and scheduled runs but nothing else - group: > - ${{ - github.event_name == 'pull_request' && format('{0}-pr-{1}', github.workflow, github.event.pull_request.number) || - github.event_name == 'schedule' && format('{0}-schedule', github.workflow) || - github.run_id - }} - cancel-in-progress: true - -jobs: - prelim: - runs-on: ["self-hosted"] - steps: - - name: Test gsutil installation - run: which gsutil >/dev/null 2>&1 || { echo >&2 "gsutil is required but not installed. Aborting"; exit 24;} - - name: Cleanup old docker images - run: docker system prune --all --force - - tpu_image: - needs: prelim - uses: ./.github/workflows/build_upload_internal.yml - with: - device_type: tpu - device_name: v4-8 - cloud_runner: linux-x86-n2-16-buildkit - build_mode: jax_ai_image - base_image: us-docker.pkg.dev/tpu-prod-env-multipod/jax-stable-stack/candidate/tpu:latest - - gpu_image: - needs: prelim - uses: ./.github/workflows/build_upload_internal.yml - with: - device_type: gpu - device_name: a100-40gb-4 - cloud_runner: linux-x86-n2-16-buildkit - build_mode: jax_ai_image - base_image: us-docker.pkg.dev/tpu-prod-env-multipod/jax-stable-stack/candidate/gpu:latest - - cpu_unit_tests: - needs: tpu_image - strategy: - fail-fast: false - matrix: - worker_group: [1, 2, 3, 4] - uses: ./.github/workflows/run_tests_internal.yml - with: - device_type: cpu - device_name: X64 - image_type: tpu - pytest_marker: 'cpu_only' - xla_python_client_mem_fraction: 0.75 - tf_force_gpu_allow_growth: false - container_resource_option: "--privileged" - is_scheduled_run: ${{ github.event_name == 'schedule' }} - worker_group: ${{ matrix.worker_group }} - total_workers: 4 - - tpu_unit_tests: - needs: tpu_image - uses: ./.github/workflows/run_tests_internal.yml - with: - device_type: tpu - device_name: v4-8 - cloud_runner: linux-x86-ct4p-240-4tpu - pytest_marker: 'not cpu_only and not gpu_only and not integration_test' - xla_python_client_mem_fraction: 0.75 - tf_force_gpu_allow_growth: false - container_resource_option: "--privileged" - is_scheduled_run: ${{ github.event_name == 'schedule' }} - - tpu_pathways_unit_tests: - needs: tpu_image - uses: ./.github/workflows/run_pathways_tests_internal.yml - with: - device_type: tpu - device_name: v4-8 - cloud_runner: linux-x86-ct4p-240-4tpu - pytest_marker: 'not cpu_only and not gpu_only and not integration_test' - xla_python_client_mem_fraction: 0.75 - tf_force_gpu_allow_growth: false - container_resource_option: "--privileged" - is_scheduled_run: ${{ github.event_name == 'schedule' }} - - tpu_integration_tests: - needs: tpu_image - uses: ./.github/workflows/run_tests_internal.yml - with: - device_type: tpu - device_name: v4-8 - cloud_runner: linux-x86-ct4p-240-4tpu - pytest_marker: 'not cpu_only and not gpu_only and integration_test' - xla_python_client_mem_fraction: 0.75 - tf_force_gpu_allow_growth: false - container_resource_option: "--privileged" - is_scheduled_run: ${{ github.event_name == 'schedule' }} - - gpu_unit_tests: - needs: gpu_image - uses: ./.github/workflows/run_tests_internal.yml - with: - device_type: gpu - device_name: a100-40gb-4 - cloud_runner: linux-x86-a2-48-a100-4gpu - pytest_marker: 'not cpu_only and not tpu_only and not integration_test' - pytest_addopts: '--ignore=tests/sft_hooks_test.py' - xla_python_client_mem_fraction: 0.65 - tf_force_gpu_allow_growth: true - container_resource_option: "--shm-size 2g --runtime=nvidia --gpus all --privileged" - is_scheduled_run: ${{ github.event_name == 'schedule' }} - - gpu_integration_tests: - needs: gpu_image - uses: ./.github/workflows/run_tests_internal.yml - with: - device_type: gpu - device_name: a100-40gb-4 - cloud_runner: linux-x86-a2-48-a100-4gpu - pytest_marker: 'not cpu_only and not tpu_only and integration_test' - pytest_addopts: '--ignore=tests/sft_hooks_test.py' - xla_python_client_mem_fraction: 0.65 - tf_force_gpu_allow_growth: true - container_resource_option: "--shm-size 2g --runtime=nvidia --gpus all --privileged" - is_scheduled_run: ${{ github.event_name == 'schedule' }} - - clean_up: - if: ${{ always() }} - needs: [cpu_unit_tests, gpu_unit_tests, gpu_integration_tests, tpu_unit_tests, tpu_integration_tests, tpu_pathways_unit_tests] - name: "Clean up" - runs-on: ["self-hosted"] - permissions: - contents: read - issues: write - steps: - - name: Authenticate gcloud - run: | - # configure registries as root and as runner - gcloud auth configure-docker --quiet - gcloud auth configure-docker us-docker.pkg.dev --quiet - - name: Delete the tpu image - run: gcloud container images delete "gcr.io/tpu-prod-env-multipod/maxtext_${{ github.run_id }}:tpu" --force-delete-tags --quiet - - name: Delete the gpu image - run: gcloud container images delete "gcr.io/tpu-prod-env-multipod/maxtext_${{ github.run_id }}:gpu" --force-delete-tags --quiet - - notify_failure: - name: Notify failed build # creates an issue or modifies last open existing issue for failed build - needs: [cpu_unit_tests, gpu_unit_tests, gpu_integration_tests, tpu_unit_tests, tpu_integration_tests, tpu_pathways_unit_tests] - if: ${{ always() }} - runs-on: ubuntu-latest - permissions: - issues: write - steps: - - name: Check whether one of the jobs failed - id: report_failure - if: ${{ contains(needs.*.result, 'failure') && github.event.pull_request == null && github.event_name != 'workflow_dispatch' }} - uses: jayqi/failed-build-issue-action@v1.2 - with: - github-token: ${{ secrets.GITHUB_TOKEN }} - labels: "failed-build" - - - name: Reset consecutive success counter on failure - # This runs only if the previous step actually found or created an issue - if: ${{ steps.report_failure.outputs.issue-number != '' }} - env: - GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - GH_REPO: ${{ github.repository }} - # Use the issue number from the previous step's output - ISSUE_NUMBER: ${{ steps.report_failure.outputs.issue-number }} - run: | - echo "A failure occurred. Resetting success counter on issue #${ISSUE_NUMBER}." - # This command will attempt to remove both labels. - # It will not fail if the labels don't exist. - gh issue remove-label $ISSUE_NUMBER "success-run-1" "success-run-2" --repo $GH_REPO || echo "No success labels to remove." - - notify_success_and_close: - name: Close issue after 3 successful builds - # This job runs only if all the preceding test jobs succeeded - if: ${{ success() && github.event.pull_request == null && github.event_name != 'workflow_dispatch' }} - needs: [cpu_unit_tests, gpu_unit_tests, gpu_integration_tests, tpu_unit_tests, tpu_integration_tests, tpu_pathways_unit_tests] - runs-on: ubuntu-latest - permissions: - issues: write - steps: - - name: Find existing failure issue - id: find_issue - env: - GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - GH_REPO: ${{ github.repository }} - run: | - ISSUE_NUMBER=$(gh issue list --label "failed-build" --state open --limit 1 --json number -q '.[0].number') - if [[ -z "$ISSUE_NUMBER" ]]; then - echo "No open build failure issue found. Nothing to do." - echo "issue_number=" >> $GITHUB_OUTPUT - else - echo "Found open build failure issue: #${ISSUE_NUMBER}" - echo "issue_number=${ISSUE_NUMBER}" >> $GITHUB_OUTPUT - fi - - - name: Add success label or close issue - if: steps.find_issue.outputs.issue_number != '' - env: - GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - GH_REPO: ${{ github.repository }} - run: | - ISSUE_NUMBER=${{ steps.find_issue.outputs.issue_number }} - LABELS=$(gh issue view $ISSUE_NUMBER --json labels -q '.labels[].name') - - if echo "$LABELS" | grep -q "success-run-2"; then - echo "Third consecutive success. Closing issue #${ISSUE_NUMBER}." - gh issue comment $ISSUE_NUMBER --body "Build succeeded for the third consecutive time. Closing this issue automatically." - gh issue close $ISSUE_NUMBER - # Clean up all tracking labels - gh issue remove-label $ISSUE_NUMBER "failed-build" "success-run-2" --repo $GH_REPO - elif echo "$LABELS" | grep -q "success-run-1"; then - echo "Second consecutive success. Updating label on issue #${ISSUE_NUMBER}." - gh issue comment $ISSUE_NUMBER --body "Build succeeded for the second time. One more successful run will close this issue." - gh issue remove-label $ISSUE_NUMBER "success-run-1" --repo $GH_REPO - gh issue add-label $ISSUE_NUMBER "success-run-2" --repo $GH_REPO - else - echo "First consecutive success since failure. Adding label to issue #${ISSUE_NUMBER}." - gh issue comment $ISSUE_NUMBER --body "Build succeeded. This issue will be auto-closed after two more consecutive successful runs." - gh issue add-label $ISSUE_NUMBER "success-run-1" --repo $GH_REPO - fi diff --git a/.github/workflows/UploadDockerImages.yml b/.github/workflows/UploadDockerImages.yml deleted file mode 100644 index 76945ed6e8..0000000000 --- a/.github/workflows/UploadDockerImages.yml +++ /dev/null @@ -1,176 +0,0 @@ -# Copyright 2023–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 workflow will install Python dependencies, run tests and lint with a variety of Python versions -# For more information see: https://docs.github.com/en/actions/automating-builds-and-tests/building-and-testing-python - -name: Build Images - -on: - schedule: - # Run the job daily at 12AM UTC - - cron: '0 0 * * *' - workflow_dispatch: - inputs: - target_device: - description: 'Specify target device (all, tpu, or gpu)' - required: true - type: choice - default: 'all' - options: - - all - - tpu - - gpu - -jobs: - build-tpu: - # This job will only run for 'tpu', 'all', schedule, or PR triggers. - if: > - github.event_name == 'schedule' || - github.event_name == 'pull_request' || - github.event.inputs.target_device == 'all' || - github.event.inputs.target_device == 'tpu' - - runs-on: linux-x86-n2-16-buildkit - container: google/cloud-sdk:524.0.0 - - # Use Github Actions matrix to run image builds in parallel - strategy: - fail-fast: false - matrix: - include: - # TPU Image Builds - - image_name: maxtext_jax_stable - dockerfile: ./dependencies/dockerfiles/maxtext_dependencies.Dockerfile - build_args: | - MODE=stable - JAX_VERSION=NONE - LIBTPU_GCS_PATH=NONE - - image_name: maxtext_jax_nightly - dockerfile: ./dependencies/dockerfiles/maxtext_dependencies.Dockerfile - build_args: | - MODE=nightly - JAX_VERSION=NONE - LIBTPU_GCS_PATH=NONE - # TPU Image builds using JAX AI Image - - image_name: maxtext_jax_stable_stack - dockerfile: ./dependencies/dockerfiles/maxtext_jax_ai_image.Dockerfile - base_image: us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest - - image_name: maxtext_stable_stack_nightly_jax - dockerfile: ./dependencies/dockerfiles/maxtext_jax_ai_image.Dockerfile - base_image: us-docker.pkg.dev/tpu-prod-env-multipod/jax-stable-stack/tpu/jax_nightly:latest - - image_name: maxtext_stable_stack_candidate - dockerfile: ./dependencies/dockerfiles/maxtext_jax_ai_image.Dockerfile - base_image: us-docker.pkg.dev/tpu-prod-env-multipod/jax-stable-stack/candidate/tpu:latest - - # Setup for GKE runners per b/412986220#comment82 and b/412986220#comment90 - steps: - - uses: actions/checkout@v5 - - name: Mark git repository as safe - run: git config --global --add safe.directory ${GITHUB_WORKSPACE} - - name: Configure Docker - run: gcloud auth configure-docker us-docker.pkg.dev,gcr.io -q - - name: Set up Docker BuildX - uses: docker/setup-buildx-action@v3.11.1 - with: - driver: remote - endpoint: tcp://localhost:1234 - # Env variables to be passed to Dockerfile - - name: Get short commit hash - id: vars - run: echo "sha_short=$(git rev-parse --short HEAD)" >> $GITHUB_OUTPUT - - name: Get current date - id: date - run: echo "image_date=$(date +%Y-%m-%d)" >> $GITHUB_OUTPUT - # Docker BuildX command config - - name: Build and Push Docker Image - uses: docker/build-push-action@v6 - with: - push: true - context: . - file: ${{ matrix.dockerfile }} - tags: | - gcr.io/tpu-prod-env-multipod/${{ matrix.image_name }}:${{ steps.date.outputs.image_date }} - gcr.io/tpu-prod-env-multipod/${{ matrix.image_name }}:latest - cache-from: type=gha - cache-to: type=gha,mode=max - provenance: false - build-args: | - ${{ matrix.build_args }} - JAX_AI_IMAGE_BASEIMAGE=${{ matrix.base_image }} - COMMIT_HASH=${{ steps.vars.outputs.sha_short }} - DEVICE=tpu - TEST_TYPE=xlml - - # Same as tpu-build step but mirrored for GPUs - build-gpu: - if: > - github.event_name == 'schedule' || - github.event_name == 'pull_request' || - github.event.inputs.target_device == 'all' || - github.event.inputs.target_device == 'gpu' - - runs-on: linux-x86-n2-16-buildkit - container: google/cloud-sdk:524.0.0 - - strategy: - fail-fast: false - matrix: - # GPU Image Builds using JAX AI Image - include: - - image_name: maxtext_gpu_jax_stable_stack - dockerfile: ./dependencies/dockerfiles/maxtext_jax_ai_image.Dockerfile - base_image: us-central1-docker.pkg.dev/deeplearning-images/jax-ai-image/gpu:latest - - image_name: maxtext_gpu_stable_stack_nightly_jax - dockerfile: ./dependencies/dockerfiles/maxtext_jax_ai_image.Dockerfile - base_image: us-docker.pkg.dev/tpu-prod-env-multipod/jax-stable-stack/gpu/jax_nightly:latest - - image_name: maxtext_stable_stack_candidate_gpu - dockerfile: ./dependencies/dockerfiles/maxtext_jax_ai_image.Dockerfile - base_image: us-docker.pkg.dev/tpu-prod-env-multipod/jax-stable-stack/candidate/gpu:latest - - steps: - - uses: actions/checkout@v5 - - name: Mark git repository as safe - run: git config --global --add safe.directory ${GITHUB_WORKSPACE} - - name: Configure Docker - run: gcloud auth configure-docker us-docker.pkg.dev,gcr.io,us-central1-docker.pkg.dev -q - - name: Set up Docker BuildX - uses: docker/setup-buildx-action@v3.11.1 - with: - driver: remote - endpoint: tcp://localhost:1234 - - name: Get short commit hash - id: vars - run: echo "sha_short=$(git rev-parse --short HEAD)" >> $GITHUB_OUTPUT - - name: Get current date - id: date - run: echo "image_date=$(date +%Y-%m-%d)" >> $GITHUB_OUTPUT - - name: Build and Push Docker Image - uses: docker/build-push-action@v6 - with: - push: true - context: . - file: ${{ matrix.dockerfile }} - tags: | - gcr.io/tpu-prod-env-multipod/${{ matrix.image_name }}:${{ steps.date.outputs.image_date }} - gcr.io/tpu-prod-env-multipod/${{ matrix.image_name }}:latest - cache-from: type=gha - cache-to: type=gha,mode=max - provenance: false - build-args: | - ${{ matrix.build_args }} - JAX_AI_IMAGE_BASEIMAGE=${{ matrix.base_image }} - COMMIT_HASH=${{ steps.vars.outputs.sha_short }} - DEVICE=gpu - TEST_TYPE=xlml diff --git a/.github/workflows/build_and_test_maxtext.yml b/.github/workflows/build_and_test_maxtext.yml deleted file mode 100644 index 112e52fb2c..0000000000 --- a/.github/workflows/build_and_test_maxtext.yml +++ /dev/null @@ -1,156 +0,0 @@ -# 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 workflow will build maxtext python package and run tests. - -name: MaxText Package Tests - -on: - pull_request: - workflow_dispatch: - schedule: - # Run the job every 4 hours - - cron: '0 */4 * * *' - -concurrency: - # Dedup pull requests (canceling previous runs of the same workflow for same PR), and scheduled runs but nothing else - group: > - ${{ - github.event_name == 'pull_request' && format('{0}-pr-{1}', github.workflow, github.event.pull_request.number) || - github.event_name == 'schedule' && format('{0}-schedule', github.workflow) || - github.run_id - }} - cancel-in-progress: true - -permissions: - contents: read -jobs: - build_and_upload_maxtext_package: - name: Build and upload maxtext package - uses: ./.github/workflows/build_package.yml - with: - device_type: tpu - device_name: v4-8 - cloud_runner: linux-x86-n2-16-buildkit - - maxtext_cpu_unit_tests: - needs: build_and_upload_maxtext_package - uses: ./.github/workflows/run_tests_against_package.yml - strategy: - fail-fast: false # don't cancel all jobs on failure - matrix: - image_type: ["py312"] - worker_group: [1, 2, 3, 4] - with: - device_type: cpu - device_name: X64 - cloud_runner: linux-x86-n2-16 - image_type: ${{ matrix.image_type }} - pytest_marker: 'cpu_only' - xla_python_client_mem_fraction: 0.75 - tf_force_gpu_allow_growth: false - container_resource_option: "--privileged" - is_scheduled_run: ${{ github.event_name == 'schedule' }} - worker_group: ${{ matrix.worker_group }} - total_workers: 4 - - maxtext_tpu_unit_tests: - needs: build_and_upload_maxtext_package - uses: ./.github/workflows/run_tests_against_package.yml - strategy: - fail-fast: false - matrix: - image_type: ["py312"] - with: - device_type: tpu - device_name: v4-8 - image_type: ${{ matrix.image_type }} - cloud_runner: linux-x86-ct4p-240-4tpu - pytest_marker: 'not cpu_only and not gpu_only and not integration_test' - xla_python_client_mem_fraction: 0.75 - tf_force_gpu_allow_growth: false - container_resource_option: "--privileged" - is_scheduled_run: ${{ github.event_name == 'schedule' }} - - maxtext_tpu_integration_tests: - needs: build_and_upload_maxtext_package - uses: ./.github/workflows/run_tests_against_package.yml - strategy: - fail-fast: false - matrix: - image_type: ["py312"] - with: - device_type: tpu - device_name: v4-8 - image_type: ${{ matrix.image_type }} - cloud_runner: linux-x86-ct4p-240-4tpu - pytest_marker: 'not cpu_only and not gpu_only and integration_test' - xla_python_client_mem_fraction: 0.75 - tf_force_gpu_allow_growth: false - container_resource_option: "--privileged" - is_scheduled_run: ${{ github.event_name == 'schedule' }} - - maxtext_gpu_unit_tests: - needs: build_and_upload_maxtext_package - uses: ./.github/workflows/run_tests_against_package.yml - strategy: - fail-fast: false - matrix: - image_type: ["py312"] - cuda: ["cuda12"] - with: - device_type: ${{ matrix.cuda }} - device_name: a100-40gb-4 - image_type: ${{ matrix.image_type }} - cloud_runner: linux-x86-a2-48-a100-4gpu - pytest_marker: 'not cpu_only and not tpu_only and not integration_test' - pytest_addopts: '--ignore=tests/sft_hooks_test.py' - xla_python_client_mem_fraction: 0.65 - tf_force_gpu_allow_growth: true - container_resource_option: "--shm-size 2g --runtime=nvidia --gpus all --privileged" - is_scheduled_run: ${{ github.event_name == 'schedule' }} - - maxtext_gpu_integration_tests: - needs: build_and_upload_maxtext_package - uses: ./.github/workflows/run_tests_against_package.yml - strategy: - fail-fast: false - matrix: - image_type: ["py312"] - cuda: ["cuda12"] - with: - device_type: ${{ matrix.cuda }} - device_name: a100-40gb-4 - image_type: ${{ matrix.image_type }} - cloud_runner: linux-x86-a2-48-a100-4gpu - pytest_marker: 'not cpu_only and not tpu_only and integration_test' - pytest_addopts: '--ignore=tests/sft_hooks_test.py' - xla_python_client_mem_fraction: 0.65 - tf_force_gpu_allow_growth: true - container_resource_option: "--shm-size 2g --runtime=nvidia --gpus all --privileged" - is_scheduled_run: ${{ github.event_name == 'schedule' }} - - notify_failure: - name: Notify failed build # creates an issue or modifies last open existing issue for failed build - needs: [maxtext_cpu_unit_tests, maxtext_tpu_unit_tests, maxtext_tpu_integration_tests, maxtext_gpu_unit_tests, maxtext_gpu_integration_tests] - if: ${{ always() }} - runs-on: ubuntu-latest - permissions: - issues: write - steps: - - name: Check whether one of the jobs failed - if: ${{ contains(needs.*.result, 'failure') && github.event.pull_request == null && github.event_name != 'workflow_dispatch' }} - uses: jayqi/failed-build-issue-action@1a893bbf43ef1c2a8705e2b115cd4f0fe3c5649b # v1.2.0 - with: - github-token: ${{ secrets.GITHUB_TOKEN }} diff --git a/.github/workflows/build_package.yml b/.github/workflows/build_package.yml deleted file mode 100644 index 15487f1daf..0000000000 --- a/.github/workflows/build_package.yml +++ /dev/null @@ -1,52 +0,0 @@ -# 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 building and uploading a maxtext pacakge -# based on the pyproject.toml at the current github workspace. - -name: Build and Upload MaxText Package - -on: - workflow_call: - inputs: - device_type: - required: true - type: string - device_name: - required: true - type: string - cloud_runner: - required: false - type: string - -permissions: - contents: read -jobs: - build_and_upload: - name: Build and upload maxtext package (${{ inputs.device_name }}) - runs-on: ${{ inputs.cloud_runner != '' && inputs.cloud_runner || fromJson(format('["self-hosted", "{0}", "{1}"]', inputs.device_type, inputs.device_name)) }} - container: python:3.12.3-slim-bullseye - steps: - - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 - - name: Install build tools - run: | - python -m pip install --upgrade pip build uv - - name: Build maxtext wheel - run: | - uv build --wheel - - name: Upload the built maxtext wheel - uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4.6.2 - with: - name: maxtext-wheel - path: dist/* diff --git a/.github/workflows/build_upload_internal.yml b/.github/workflows/build_upload_internal.yml deleted file mode 100644 index 0fba599642..0000000000 --- a/.github/workflows/build_upload_internal.yml +++ /dev/null @@ -1,69 +0,0 @@ -# Copyright 2023–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 building and uploading an image used in UnitTests.yml - -name: Build and Upload Image - -on: - workflow_call: - inputs: - device_type: - required: true - type: string - device_name: - required: true - type: string - build_mode: - required: true - type: string - base_image: - required: false - type: string - cloud_runner: - required: false - type: string - -jobs: - build_and_upload: - name: Build and upload image (${{ inputs.device_name }}) - runs-on: ${{ inputs.cloud_runner != '' && inputs.cloud_runner || fromJson(format('["self-hosted", "{0}", "{1}"]', inputs.device_type, inputs.device_name)) }} - container: google/cloud-sdk:524.0.0 - steps: - - uses: actions/checkout@v5 - - name: Mark git repository as safe - run: git config --global --add safe.directory ${GITHUB_WORKSPACE} - - name: Configure Docker - run: gcloud auth configure-docker us-docker.pkg.dev,gcr.io,us-central1-docker.pkg.dev -q - - name: Set up Docker BuildX - uses: docker/setup-buildx-action@v3.11.1 - with: - driver: remote - endpoint: tcp://localhost:1234 - - name: Get short commit hash - id: vars - run: echo "sha_short=$(git rev-parse --short HEAD)" >> $GITHUB_OUTPUT - - name: Build and push - uses: docker/build-push-action@v6 - with: - push: true - context: . - file: ./dependencies/dockerfiles/maxtext_jax_ai_image.Dockerfile - tags: gcr.io/tpu-prod-env-multipod/maxtext_${{ github.run_id }}:${{ inputs.device_type }} - provenance: false - build-args: | - JAX_AI_IMAGE_BASEIMAGE=${{ inputs.base_image }} - COMMIT_HASH=${{ steps.vars.outputs.sha_short }} - DEVICE=${{ inputs.device_type }} - TEST_TYPE=unit_test diff --git a/.github/workflows/check_docs_build.yml b/.github/workflows/check_docs_build.yml deleted file mode 100644 index 393bd98fd2..0000000000 --- a/.github/workflows/check_docs_build.yml +++ /dev/null @@ -1,31 +0,0 @@ -name: Check the Documentation Build - -on: - pull_request: - types: [opened, synchronize] - workflow_dispatch: - -permissions: - contents: read - -jobs: - build-docs: - runs-on: ubuntu-latest - steps: - - name: Checkout repository - uses: actions/checkout@08c6903cd8c0fde910a37f88322edcfb5dd907a8 # v5.0.0 - with: - persist-credentials: false - - - name: Set up Python - uses: actions/setup-python@e797f83bcb11b83ae66e0230d6156d7c80228e7c # v6.0.0 - with: - python-version: '3.12' - cache: 'pip' # caching pip dependencies - - - name: Install dependencies - run: pip install -r dependencies/requirements/requirements_docs.txt - - - name: Build documentation - run: | - sphinx-build -W -b html docs docs/_build/html diff --git a/.github/workflows/run_pathways_tests_internal.yml b/.github/workflows/run_pathways_tests_internal.yml deleted file mode 100644 index 079c712b0d..0000000000 --- a/.github/workflows/run_pathways_tests_internal.yml +++ /dev/null @@ -1,114 +0,0 @@ -# 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 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 - image_type: - required: false - 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 - -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_${{ github.run_id }}:${{ inputs.image_type != '' && inputs.image_type || inputs.device_type }} - volumes: - - ${{ github.workspace }}:/tmp - 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" - options: ${{ inputs.container_resource_option }} - steps: - - uses: actions/checkout@v4 - - name: Run - 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 - python3 -m pip install -e . --no-dependencies && - python3 -m pip uninstall -y libtpu && - # TODO(b/454659463): Enable test_default_hlo_match after volume mount is supported. - python3 -m pytest ${{ inputs.pytest_addopts }} -v -m "${FINAL_PYTEST_MARKER}" -k "not test_default_hlo_match" --durations=0 - - 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, --instance_count=1, --instance_type=tpuv4:2x2x1, --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, --gcs_scratch_location=gs://cloud-pathways-staging/tmp] - --tpu=4 - - proxy: - image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server:latest - volumes: - - ${{ github.workspace }}:/tmp - 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, --gcs_scratch_location=gs://cloud-pathways-staging/tmp, --xla_tpu_scoped_vmem_limit_kib=65536, --xla_tpu_spmd_rng_bit_generator_unsafe=true] diff --git a/.github/workflows/run_tests_against_package.yml b/.github/workflows/run_tests_against_package.yml deleted file mode 100644 index b2f1ca6c43..0000000000 --- a/.github/workflows/run_tests_against_package.yml +++ /dev/null @@ -1,114 +0,0 @@ -# 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 - 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 - worker_group: - required: false - type: number - default: 1 - total_workers: - required: false - type: number - 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/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' || '' }} - MAXTEXT_PACKAGE_EXTRA: ${{ inputs.device_type == 'cpu' && 'tpu' || inputs.device_type }} - 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 - 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}[${MAXTEXT_PACKAGE_EXTRA}] --resolution=lowest - uv pip install -r src/install_maxtext_extra_deps/extra_deps_from_github.txt - 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 - shell: bash - 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 - # 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' - fi - # TODO: Fix the skipped tests and remove the deselect flags - [ "${{ inputs.total_workers }}" -gt 1 ] && .venv/bin/python3 -m pip install --quiet pytest-split && SPLIT_ARGS="--splits ${{ inputs.total_workers }} --group ${{ inputs.worker_group }}" || SPLIT_ARGS="" - .venv/bin/python3 -m pytest ${{ inputs.pytest_addopts }} -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" $SPLIT_ARGS diff --git a/.github/workflows/run_tests_internal.yml b/.github/workflows/run_tests_internal.yml deleted file mode 100644 index 3647bf83fb..0000000000 --- a/.github/workflows/run_tests_internal.yml +++ /dev/null @@ -1,84 +0,0 @@ -# Copyright 2023–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 used in UnitTests.yml - -name: Run Tests - -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 - 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 - worker_group: - required: false - type: number - default: 1 - total_workers: - required: false - type: number - default: 1 - -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_${{ github.run_id }}:${{ inputs.image_type != '' && inputs.image_type || inputs.device_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.image_type == 'tpu' && inputs.device_type != 'tpu' && '1' || '' }} - options: ${{ inputs.container_resource_option }} - steps: - - uses: actions/checkout@v5 - - 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 - python3 -m pip install -e . --no-dependencies - [ "${{ inputs.total_workers }}" -gt 1 ] && python3 -m pip install --quiet pytest-split && SPLIT_ARGS="--splits ${{ inputs.total_workers }} --group ${{ inputs.worker_group }}" || SPLIT_ARGS="" - export LIBTPU_INIT_ARGS='--xla_tpu_scoped_vmem_limit_kib=65536' - python3 -m pytest ${{ inputs.pytest_addopts }} -v -m "${FINAL_PYTEST_MARKER}" --durations=0 $SPLIT_ARGS diff --git a/.github/workflows/stale.yml b/.github/workflows/stale.yml deleted file mode 100644 index b8430add34..0000000000 --- a/.github/workflows/stale.yml +++ /dev/null @@ -1,49 +0,0 @@ -# Copyright 2023–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. - -name: 'Close stale PRs' -on: - schedule: - - cron: '0 16 * * *' - -jobs: - stale: - runs-on: ubuntu-latest - permissions: - actions: write - issues: write - pull-requests: write - steps: - - uses: actions/stale@v9 - with: - # Disable the action for issues to run on PRs only - days-before-issue-stale: -1 - days-before-issue-close: -1 - # Number of days of inactivity before a PR is marked as stale - days-before-pr-stale: 30 - # Number of days of inactivity before a stale PR is closed - days-before-pr-close: 7 - # The label to apply to stale PRs - stale-pr-label: 'stale' - # This is true by default, adding just to be explicit - remove-stale-when-updated: true - # The comment to post when marking a PR as stale. - stale-pr-message: > - This PR has been automatically marked as stale because it has not had - recent activity. It will be closed soon if no further activity occurs. - Thank you for your contributions. - # The comment to post when closing a stale PR. - close-pr-message: > - This PR was closed because it has been inactive for a while. - Please reopen it if you are still working on it. From 121f3b4211104e5fa366aa532ee72201153093b8 Mon Sep 17 00:00:00 2001 From: Dora Hsieh Date: Fri, 7 Nov 2025 15:23:30 +0800 Subject: [PATCH 3/3] import llama_or_mistral_ckpt --- .../utils/ckpt_scripts/llama_mistral_mixtral_orbax_to_hf.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/MaxText/utils/ckpt_scripts/llama_mistral_mixtral_orbax_to_hf.py b/src/MaxText/utils/ckpt_scripts/llama_mistral_mixtral_orbax_to_hf.py index 8a9afe34da..a1f525d29a 100644 --- a/src/MaxText/utils/ckpt_scripts/llama_mistral_mixtral_orbax_to_hf.py +++ b/src/MaxText/utils/ckpt_scripts/llama_mistral_mixtral_orbax_to_hf.py @@ -48,7 +48,7 @@ from transformers import LlamaForCausalLM, MistralForCausalLM, AutoModelForCausalLM, AutoConfig from MaxText import checkpointing -from MaxText import llama_or_mistral_ckpt +import llama_or_mistral_ckpt from MaxText import max_logging from MaxText import maxtext_utils from MaxText import pyconfig