multiprocess-tests-schedule-run #1585
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| name: multiprocess-tests-schedule-run | |
| on: | |
| # continuous | |
| schedule: | |
| # Run every 4 hour | |
| - cron: "0 */4 * * *" | |
| permissions: | |
| contents: read | |
| actions: write # to cancel previous workflows | |
| issues: write # to create failure alerts | |
| statuses: write | |
| concurrency: | |
| group: ${{ github.workflow }}-${{ github.head_ref || github.ref }} | |
| cancel-in-progress: true | |
| jobs: | |
| multiprocess-checkpoint-benchmarks: | |
| name: "multiprocess-checkpoint-benchmarks (Python ${{ matrix.python-version }}, jax=${{ matrix.jax-version }})" | |
| runs-on: linux-g2-16-l4-1gpu-x4 | |
| # runs-on: linux-x86-ct5lp-4tpu-x4 | |
| container: us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:infrastructure-public-image-2d2a7b1e6e2e | |
| defaults: | |
| run: | |
| working-directory: checkpoint | |
| strategy: | |
| matrix: | |
| python-version: ["3.12"] | |
| jax-version: ["0.6.0"] | |
| steps: | |
| - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 | |
| - name: Set up Python ${{ matrix.python-version }} | |
| uses: actions/setup-python@0b93645e9fea7318ecaed2b359559ac225c90a2b # v5.3.0 | |
| with: | |
| python-version: ${{ matrix.python-version }} | |
| - name: Install dependencies | |
| run: | | |
| pip install -e . | |
| pip install -e .[testing,gcs] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html | |
| pip uninstall -y orbax | |
| if [[ "${{ matrix.jax-version }}" == "newest" ]]; then | |
| pip install -U jax[k8s,cuda12] jaxlib | |
| elif [[ "${{ matrix.jax-version }}" == "nightly" ]]; then | |
| pip install -U --pre jax[k8s,cuda12] jaxlib --extra-index-url https://us-python.pkg.dev/ml-oss-artifacts-published/jax-public-nightly-artifacts-registry/simple/ | |
| else | |
| pip install "jax[k8s,cuda12]>=${{ matrix.jax-version }}" "jaxlib>=${{ matrix.jax-version }}" | |
| fi | |
| pip install gcsfs | |
| pip install portpicker tensorboard | |
| - name: Run correctness integration tests | |
| env: | |
| GCS_BUCKET_PATH: gs://orbax-benchmarks/benchmark-results/${{ github.run_id }} | |
| TF_FORCE_GPU_ALLOW_GROWTH: true | |
| XLA_PYTHON_CLIENT_PREALLOCATE: false | |
| run: | | |
| cd orbax/checkpoint/_src/testing/oss | |
| failed_benchmarks="" | |
| benchmark_configs_file="multiprocess_benchmark_configs.txt" | |
| echo "Running benchmarks specified in $benchmark_configs_file" | |
| benchmark_configs_file_path="$PWD/$benchmark_configs_file" | |
| cd ../benchmarks | |
| while IFS= read -r entry || [ -n "$entry" ]; do | |
| if [ -n "$entry" ]; then | |
| echo "Running benchmark for $entry" | |
| if ! python -c "import sys; import jax; jax.distributed.initialize(); print(jax.devices()); from absl import app; import run_benchmarks; sys.argv = ['run_benchmarks.py', '--config_file="$entry"', '--output_directory=$GCS_BUCKET_PATH']; app.run(run_benchmarks.main)"; then | |
| echo "Benchmark $entry failed" | |
| failed_benchmarks="$failed_benchmarks $entry" | |
| fi | |
| fi | |
| done < "$benchmark_configs_file_path" | |
| cd ../../../../.. | |
| if [ -n "$failed_benchmarks" ]; then | |
| echo "The following benchmarks failed:$failed_benchmarks" | |
| exit 1 | |
| fi | |
| # python -m pytest orbax/checkpoint/_src/handlers/array_checkpoint_handler_test.py | |
| # cd orbax/checkpoint/_src/testing/benchmarks && python run_benchmarks.py --config_file=configs/pytree_checkpoint_benchmark.yaml --output_directory=$GCS_BUCKET_PATH | |
| # The below step just reports the success or failure of tests as a "commit status". | |
| # This is needed for copybara integration. | |
| - name: Create Issue on Continuous Failure | |
| if: failure() | |
| uses: actions/github-script@v7 | |
| with: | |
| script: | | |
| const runs = await github.rest.actions.listWorkflowRuns({ | |
| owner: context.repo.owner, | |
| repo: context.repo.repo, | |
| workflow_id: 'multiprocess_tests.yml', | |
| per_page: 5 | |
| }); | |
| const previousRuns = runs.data.workflow_runs.filter(run => run.id !== context.runId).slice(0, 2); | |
| console.log(`Previous runs conclusions: ${previousRuns.map(r => r.conclusion).join(', ')}`); | |
| const allFailed = previousRuns.every(run => run.conclusion === 'failure'); | |
| if (allFailed && previousRuns.length === 2) { | |
| // Search for existing open issues with the same title | |
| const issues = await github.rest.issues.listForRepo({ | |
| owner: context.repo.owner, | |
| repo: context.repo.repo, | |
| labels: 'continuous-integration', | |
| state: 'open' | |
| }); | |
| const duplicate = issues.data.find(issue => issue.title === `🚨 Continuous Failure: ${context.workflow}`); | |
| if (duplicate) { | |
| console.log("An open issue already exists for this failure. Skipping creation."); | |
| } else { | |
| console.log("Previous 2 runs also failed and no open issue found. Creating issue."); | |
| await github.rest.issues.create({ | |
| owner: context.repo.owner, | |
| repo: context.repo.repo, | |
| title: `🚨 Continuous Failure: ${context.workflow}`, | |
| body: `The workflow has failed 3 times consecutively. \n\nLatest failing run: https://github.com/${context.repo.owner}/${context.repo.repo}/actions/runs/${context.runId}`, | |
| labels: ['bug', 'continuous-integration'] | |
| }); | |
| } | |
| } else { | |
| console.log("Did not meet continuous failure criteria."); | |
| } | |
| multiprocess-checkpoint-benchmarks-summary: | |
| needs: multiprocess-checkpoint-benchmarks | |
| runs-on: ubuntu-latest | |
| if: always() | |
| steps: | |
| - name: Report success or failure as github status | |
| shell: bash | |
| run: | | |
| if [ "${{ needs.multiprocess-checkpoint-benchmarks.result }}" = "success" ]; then | |
| status="success" | |
| else | |
| status="failure" | |
| fi | |
| curl -sS --request POST \ | |
| --url https://api.github.com/repos/${{ github.repository }}/statuses/${{ github.sha }} \ | |
| --header 'authorization: Bearer ${{ secrets.GITHUB_TOKEN }}' \ | |
| --header 'content-type: application/json' \ | |
| --data '{ | |
| "state": "'$status'", | |
| "target_url": "https://github.com/${{ github.repository }}/actions/runs/${{ github.run_id }}", | |
| "description": "All matrix jobs '$status'", | |
| "context": "github-actions/multiprocess-checkpoint-benchmarks" | |
| }' | |
| if [ "$status" = "failure" ]; then | |
| exit 1 | |
| fi | |
| multiprocess-unit-tests: | |
| name: "multiprocess-unit-tests (Python ${{ matrix.python-version }}, jax=${{ matrix.jax-version }})" | |
| runs-on: linux-x86-ct6e-180-8tpu | |
| container: us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:infrastructure-public-image-2d2a7b1e6e2e | |
| defaults: | |
| run: | |
| working-directory: checkpoint | |
| strategy: | |
| matrix: | |
| python-version: ["3.11"] | |
| jax-version: ["newest"] | |
| steps: | |
| - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 | |
| - name: Set up Python ${{ matrix.python-version }} | |
| uses: actions/setup-python@0b93645e9fea7318ecaed2b359559ac225c90a2b # v5.3.0 | |
| with: | |
| python-version: ${{ matrix.python-version }} | |
| - name: Install dependencies | |
| run: | | |
| pip install -e . | |
| pip install -e .[testing,gcs] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html | |
| pip uninstall -y orbax | |
| pip install gcsfs | |
| pip install portpicker pytest chex pyyaml tensorboard | |
| if [ "${{ matrix.jax-version }}" = "newest" ]; then | |
| pip install -U "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html | |
| elif [ "${{ matrix.jax-version }}" = "nightly" ]; then | |
| pip install -U --pre "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html --extra-index-url https://us-python.pkg.dev/ml-oss-artifacts-published/jax-public-nightly-artifacts-registry/simple/ | |
| else | |
| pip install "jax[tpu]==${{ matrix.jax-version }}" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html | |
| fi | |
| - name: Run 2 multiprocess tests | |
| env: | |
| TEST_TMPDIR: /tmp/orbax_test | |
| run: | | |
| python orbax/checkpoint/_src/testing/oss/run_multihost.py --num_processes=2 --tpu_chips_per_process=4 orbax/checkpoint/_src/testing/oss/run_tests.py --filename=orbax/checkpoint/_src/testing/oss/tagged_tests_whole_suite.yaml --processes=2 | |
| - name: Run 4 multiprocess tests | |
| run: | | |
| python orbax/checkpoint/_src/testing/oss/run_multihost.py --num_processes=4 --tpu_chips_per_process=2 orbax/checkpoint/_src/testing/oss/run_tests.py --filename=orbax/checkpoint/_src/testing/oss/tagged_tests_whole_suite.yaml --processes=4 | |
| - name: Run single process tests | |
| run: | | |
| python orbax/checkpoint/_src/testing/oss/run_multihost.py --num_processes=1 --tpu_chips_per_process=8 orbax/checkpoint/_src/testing/oss/run_tests.py --filename=orbax/checkpoint/_src/testing/oss/tagged_tests_whole_suite.yaml --processes=1 | |
| - name: Create Issue on Continuous Failure | |
| if: failure() | |
| uses: actions/github-script@v7 | |
| with: | |
| script: | | |
| const runs = await github.rest.actions.listWorkflowRuns({ | |
| owner: context.repo.owner, | |
| repo: context.repo.repo, | |
| workflow_id: 'multiprocess_tests.yml', | |
| per_page: 5 | |
| }); | |
| const previousRuns = runs.data.workflow_runs.filter(run => run.id !== context.runId).slice(0, 2); | |
| console.log(`Previous runs conclusions: ${previousRuns.map(r => r.conclusion).join(', ')}`); | |
| const allFailed = previousRuns.every(run => run.conclusion === 'failure'); | |
| if (allFailed && previousRuns.length === 2) { | |
| // Search for existing open issues with the same title | |
| const issues = await github.rest.issues.listForRepo({ | |
| owner: context.repo.owner, | |
| repo: context.repo.repo, | |
| labels: 'continuous-integration', | |
| state: 'open' | |
| }); | |
| const duplicate = issues.data.find(issue => issue.title === `🚨 Continuous Failure: ${context.workflow}`); | |
| if (duplicate) { | |
| console.log("An open issue already exists for this failure. Skipping creation."); | |
| } else { | |
| console.log("Previous 2 runs also failed and no open issue found. Creating issue."); | |
| await github.rest.issues.create({ | |
| owner: context.repo.owner, | |
| repo: context.repo.repo, | |
| title: `🚨 Continuous Failure: ${context.workflow}`, | |
| body: `The workflow has failed 3 times consecutively. \n\nLatest failing run: https://github.com/${context.repo.owner}/${context.repo.repo}/actions/runs/${context.runId}`, | |
| labels: ['bug', 'continuous-integration'] | |
| }); | |
| } | |
| } else { | |
| console.log("Did not meet continuous failure criteria."); | |
| } | |
| multiprocess-unit-tests-summary: | |
| needs: multiprocess-unit-tests | |
| runs-on: ubuntu-latest | |
| if: always() | |
| steps: | |
| - name: Report success or failure as github status | |
| shell: bash | |
| run: | | |
| if [ "${{ needs.multiprocess-unit-tests.result }}" = "success" ]; then | |
| status="success" | |
| else | |
| status="failure" | |
| fi | |
| curl -sS --request POST \ | |
| --url https://api.github.com/repos/${{ github.repository }}/statuses/${{ github.sha }} \ | |
| --header 'authorization: Bearer ${{ secrets.GITHUB_TOKEN }}' \ | |
| --header 'content-type: application/json' \ | |
| --data '{ | |
| "state": "'$status'", | |
| "target_url": "https://github.com/${{ github.repository }}/actions/runs/${{ github.run_id }}", | |
| "description": "All matrix jobs '$status'", | |
| "context": "github-actions/multiprocess-unit-tests" | |
| }' | |
| if [ "$status" = "failure" ]; then | |
| exit 1 | |
| fi |