nightly.yaml

name: Nightly Tests

on:
  schedule:
    - cron: "0 12 * * *" # Daily at 12:00 UTC
  workflow_dispatch: # allows triggering the workflow run manually
  pull_request:  # Automatically trigger on pull requests affecting this file
    branches:
      - main
    paths:
      - '**workflows/nightly.yaml'

permissions:
  contents: read
  issues: write  # for failed-build-issue

concurrency:
  group: ${{ github.workflow }}-${{ github.head_ref || github.ref }}
  cancel-in-progress: true

jobs:
  latest-releases:
    name: Test with latest package releases
    runs-on: ubuntu-latest
    timeout-minutes: 10
    strategy:
      fail-fast: false
    steps:
      - uses: actions/checkout@692973e3d937129bcbf40652eb9f2f61becf3332  # v4.1.7
      - uses: actions/setup-python@f677139bbe7f9c59b41e40162b753c062f5d49a3  # v5.2.0
        with:
          python-version: 3.12
      - name: Install dependencies
        run: |
          python -m pip install --upgrade pip
          python -m pip install pytest pytest-xdist
          python -m pip install -U jax flax grain ml_dtypes optax orbax-checkpoint orbax-export tensorflow tensorflow_datasets
      - name: Run tests
        run: |
          pytest -n auto jax_ai_stack
      - name: Notify failed build
        uses: jayqi/failed-build-issue-action@1a893bbf43ef1c2a8705e2b115cd4f0fe3c5649b  # v1.2.0
        if: failure() && github.event.pull_request == null
        with:
          github-token: ${{ secrets.GITHUB_TOKEN }}

  jax-nightly:
    name: Test with jax nightly
    runs-on: ubuntu-latest
    timeout-minutes: 10
    strategy:
      fail-fast: false
    steps:
      - uses: actions/checkout@692973e3d937129bcbf40652eb9f2f61becf3332  # v4.1.7
      - uses: actions/setup-python@f677139bbe7f9c59b41e40162b753c062f5d49a3  # v5.2.0
        with:
          python-version: 3.12
      - name: Install dependencies with jax nightly
        run: |
          python -m pip install --upgrade pip
          python -m pip install .[dev,tfds,grain]
          python -m pip install --upgrade --pre jax jaxlib -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html
      - name: Run tests
        run: |
          pytest -n auto jax_ai_stack
      - name: Notify failed build
        uses: jayqi/failed-build-issue-action@1a893bbf43ef1c2a8705e2b115cd4f0fe3c5649b  # v1.2.0
        if: failure() && github.event.pull_request == null
        with:
          github-token: ${{ secrets.GITHUB_TOKEN }}

  flax-nightly:
    name: Test with flax nightly
    runs-on: ubuntu-latest
    timeout-minutes: 10
    strategy:
      fail-fast: false
    steps:
      - uses: actions/checkout@692973e3d937129bcbf40652eb9f2f61becf3332  # v4.1.7
      - uses: actions/setup-python@f677139bbe7f9c59b41e40162b753c062f5d49a3  # v5.2.0
        with:
          python-version: 3.12
      - name: Install dependencies with flax nightly
        run: |
          python -m pip install --upgrade pip
          python -m pip install .[dev,tfds,grain]
          python -m pip install --upgrade git+https://github.com/google/flax.git
      - name: Run tests
        run: |
          pytest -n auto jax_ai_stack
      - name: Notify failed build
        uses: jayqi/failed-build-issue-action@1a893bbf43ef1c2a8705e2b115cd4f0fe3c5649b  # v1.2.0
        if: failure() && github.event.pull_request == null
        with:
          github-token: ${{ secrets.GITHUB_TOKEN }}

  optax-nightly:
    name: Test with optax nightly
    runs-on: ubuntu-latest
    timeout-minutes: 10
    strategy:
      fail-fast: false
    steps:
      - uses: actions/checkout@692973e3d937129bcbf40652eb9f2f61becf3332  # v4.1.7
      - uses: actions/setup-python@f677139bbe7f9c59b41e40162b753c062f5d49a3  # v5.2.0
        with:
          python-version: 3.12
      - name: Install dependencies with optax nightly
        run: |
          python -m pip install --upgrade pip
          python -m pip install .[dev,tfds,grain]
          python -m pip install --upgrade git+https://github.com/google-deepmind/optax.git
      - name: Run tests
        run: |
          pytest -n auto jax_ai_stack
      - name: Notify failed build
        uses: jayqi/failed-build-issue-action@1a893bbf43ef1c2a8705e2b115cd4f0fe3c5649b  # v1.2.0
        if: failure() && github.event.pull_request == null
        with:
          github-token: ${{ secrets.GITHUB_TOKEN }}

  orbax-nightly:
    name: Test with orbax nightly
    runs-on: ubuntu-latest
    timeout-minutes: 10
    strategy:
      fail-fast: false
    steps:
      - uses: actions/checkout@692973e3d937129bcbf40652eb9f2f61becf3332  # v4.1.7
      - uses: actions/setup-python@f677139bbe7f9c59b41e40162b753c062f5d49a3  # v5.2.0
        with:
          python-version: 3.12
      - name: Install dependencies with orbax-checkpoint and orbax-export nightly
        run: |
          python -m pip install --upgrade pip
          python -m pip install .[dev,tfds,grain]
          python -m pip install --upgrade 'git+https://github.com/google/orbax/#subdirectory=checkpoint' 'git+https://github.com/google/orbax/#subdirectory=export'
      - name: Run tests
        run: |
          pytest -n auto jax_ai_stack
      - name: Notify failed build
        uses: jayqi/failed-build-issue-action@1a893bbf43ef1c2a8705e2b115cd4f0fe3c5649b  # v1.2.0
        if: failure() && github.event.pull_request == null
        with:
          github-token: ${{ secrets.GITHUB_TOKEN }}