Add checkpoint/_src/checkpointers/BUILD, add tests and update depende… #15675
Workflow file for this run
This file contains 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: build | |
| on: | |
| # continuous | |
| schedule: | |
| # Run every hour | |
| - cron: "0 * * * *" | |
| push: | |
| branches: | |
| - main | |
| - 'test_*' | |
| pull_request: | |
| branches: | |
| - main | |
| permissions: | |
| contents: read | |
| actions: write # to cancel previous workflows | |
| concurrency: | |
| group: ${{ github.workflow }}-${{ github.head_ref || github.ref }} | |
| cancel-in-progress: true | |
| jobs: | |
| build-checkpoint: | |
| name: "build-checkpoint (Python ${{ matrix.python-version }}, jax=${{ matrix.jax-version }})" | |
| runs-on: ubuntu-latest | |
| defaults: | |
| run: | |
| working-directory: checkpoint | |
| strategy: | |
| matrix: | |
| python-version: ["3.10", "3.11", "3.12"] | |
| jax-version: ["newest"] | |
| include: | |
| - python-version: "3.10" | |
| jax-version: "0.4.34" # keep in sync with minimum version in checkpoint/pyproject.toml | |
| - python-version: "3.13" | |
| jax-version: "nightly" | |
| 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 | |
| # TODO(b/275613424): remove `pip install -e .` and `pip uninstall -y orbax`. | |
| # Currently in place to override remote orbax import due to flax dependency. | |
| run: | | |
| pip install -e . | |
| pip install -e .[testing] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html | |
| pip uninstall -y orbax | |
| if [[ "${{ matrix.jax-version }}" == "newest" ]]; then | |
| pip install -U jax jaxlib | |
| elif [[ "${{ matrix.jax-version }}" == "nightly" ]]; then | |
| pip install -U --pre jax jaxlib -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html | |
| else | |
| pip install "jax==${{ matrix.jax-version }}" "jaxlib==${{ matrix.jax-version }}" | |
| fi | |
| - name: Test with pytest | |
| run: | | |
| python -m pytest | |
| # The below step just reports the success or failure of tests as a "commit status". | |
| # This is needed for copybara integration. | |
| - name: Report success or failure as github status | |
| if: always() | |
| shell: bash | |
| run: | | |
| status="${{ job.status }}" | |
| lowercase_status=$(echo $status | tr '[:upper:]' '[:lower:]') | |
| 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": "'$lowercase_status'", | |
| "target_url": "https://github.com/${{ github.repository }}/actions/runs/${{ github.run_id }}", | |
| "description": "'$status'", | |
| "context": "github-actions/build" | |
| }' | |
| build-export: | |
| name: "build-export (Python ${{ matrix.python-version }}, jax=${{ matrix.jax-version }})" | |
| runs-on: ubuntu-latest | |
| defaults: | |
| run: | |
| working-directory: export | |
| strategy: | |
| matrix: | |
| python-version: ["3.10", "3.11", "3.12"] | |
| jax-version: ["newest"] | |
| include: | |
| - python-version: "3.10" | |
| jax-version: "0.4.34" # keep in sync with minimum version in export/pyproject.toml | |
| - python-version: "3.12" # TODO(jakevdp): update to 3.13 when tf supports it. | |
| jax-version: "nightly" | |
| 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: Extract branch name | |
| shell: bash | |
| run: echo "branch=${GITHUB_HEAD_REF:-${GITHUB_REF#refs/heads/}}" >> $GITHUB_OUTPUT | |
| id: extract_branch | |
| - name: Install dependencies | |
| run: | | |
| pip install . | |
| pip install .[testing] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html | |
| if [[ "${{ matrix.jax-version }}" == "newest" ]]; then | |
| pip install -U jax jaxlib | |
| elif [[ "${{ matrix.jax-version }}" == "nightly" ]]; then | |
| pip install -U --pre jax jaxlib -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html | |
| else | |
| pip install "jax==${{ matrix.jax-version }}" "jaxlib==${{ matrix.jax-version }}" | |
| fi | |
| - name: Test with pytest | |
| run: | | |
| test_dir=$(mktemp -d) | |
| cp orbax/export/conftest.py ${test_dir} | |
| for t in $(find orbax/export -maxdepth 1 -name '*_test.py'); do | |
| cp ${t} ${test_dir} | |
| XLA_FLAGS=--xla_force_host_platform_device_count=8 pytest ${test_dir}/$(basename ${t}) | |
| done | |
| # The below step just reports the success or failure of tests as a "commit status". | |
| # This is needed for copybara integration. | |
| - name: Report success or failure as github status | |
| if: always() | |
| shell: bash | |
| run: | | |
| status="${{ job.status }}" | |
| lowercase_status=$(echo $status | tr '[:upper:]' '[:lower:]') | |
| 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": "'$lowercase_status'", | |
| "target_url": "https://github.com/${{ github.repository }}/actions/runs/${{ github.run_id }}", | |
| "description": "'$status'", | |
| "context": "github-actions/build" | |
| }' |