durandtibo/coola

name: Test - Nightly Dependencies
on:
  workflow_call:
  workflow_dispatch:  # to trigger manually

jobs:
  jax:
    runs-on: ${{ matrix.os }}
    timeout-minutes: 10
    strategy:
      max-parallel: 8
      fail-fast: false
      matrix:
        os: [ ubuntu-latest ]
        python-version: [ '3.12', '3.11', '3.10', '3.9' ]

    steps:
      - name: Checkout
        uses: actions/checkout@v4

      - name: Set up Python
        uses: actions/setup-python@v5
        with:
          python-version: ${{ matrix.python-version }}

      - name: Install Poetry
        uses: snok/install-poetry@v1

      - name: Set environment variable
        run: |
          echo $(poetry run poetry env info --path)/bin >> $GITHUB_PATH

      - name: Install packages
        run: |
          make install

      - name: Install jax
        run: |
          pip install -U --pre jaxlib -f https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html
          pip install -U --pre jax-cuda12-pjrt jax-cuda12-plugin -f https://storage.googleapis.com/jax-releases/jax_cuda_plugin_nightly_releases.html

      - name: Show installed packages
        run: |
          poetry show
          poetry show --tree

      - name: Run unit tests
        run: |
          make unit-test-cov


  numpy:
    runs-on: ${{ matrix.os }}
    timeout-minutes: 10
    strategy:
      max-parallel: 8
      fail-fast: false
      matrix:
        os: [ ubuntu-latest ]
        python-version: [ '3.12', '3.11', '3.10', '3.9' ]

    steps:
      - name: Checkout
        uses: actions/checkout@v4

      - name: Set up Python
        uses: actions/setup-python@v5
        with:
          python-version: ${{ matrix.python-version }}

      - name: Install Poetry
        uses: snok/install-poetry@v1

      - name: Set environment variable
        run: |
          echo $(poetry run poetry env info --path)/bin >> $GITHUB_PATH

      - name: Install packages
        run: |
          make install

      - name: Install NumPy
        run: |
          pip install -U --pre --only-binary :all: -i https://pypi.anaconda.org/scientific-python-nightly-wheels/simple numpy

      - name: Show installed packages
        run: |
          poetry show
          poetry show --tree

      - name: Run unit tests
        run: |
          make unit-test-cov


  pandas:
    runs-on: ${{ matrix.os }}
    timeout-minutes: 10
    strategy:
      max-parallel: 8
      fail-fast: false
      matrix:
        os: [ ubuntu-latest ]
        python-version: [ '3.12', '3.11', '3.10', '3.9' ]

    steps:
      - name: Checkout
        uses: actions/checkout@v4

      - name: Set up Python
        uses: actions/setup-python@v5
        with:
          python-version: ${{ matrix.python-version }}

      - name: Install Poetry
        uses: snok/install-poetry@v1

      - name: Set environment variable
        run: |
          echo $(poetry run poetry env info --path)/bin >> $GITHUB_PATH

      - name: Install packages
        run: |
          make install

      - name: Install pandas
        run: |
          pip install --pre --extra-index https://pypi.anaconda.org/scientific-python-nightly-wheels/simple pandas

      - name: Show installed packages
        run: |
          poetry show
          poetry show --tree

      - name: Run unit tests
        run: |
          make unit-test-cov


  torch:
    runs-on: ${{ matrix.os }}
    timeout-minutes: 10
    strategy:
      max-parallel: 8
      fail-fast: false
      matrix:
        os: [ ubuntu-latest ]
        python-version: [ '3.12', '3.11', '3.10', '3.9' ]

    steps:
      - name: Checkout
        uses: actions/checkout@v4

      - name: Set up Python
        uses: actions/setup-python@v5
        with:
          python-version: ${{ matrix.python-version }}

      - name: Install Poetry
        uses: snok/install-poetry@v1

      - name: Set environment variable
        run: |
          echo $(poetry run poetry env info --path)/bin >> $GITHUB_PATH

      - name: Install packages
        run: |
          make install

      - name: Install PyTorch
        run: |
          pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu124

      - name: Show installed packages
        run: |
          poetry show
          poetry show --tree

      - name: Run unit tests
        run: |
          make unit-test-cov


  xarray:
    runs-on: ${{ matrix.os }}
    timeout-minutes: 20
    strategy:
      max-parallel: 8
      fail-fast: false
      matrix:
        os: [ ubuntu-latest ]
        python-version: [ '3.12', '3.11', '3.10', '3.9' ]

    steps:
      - name: Checkout
        uses: actions/checkout@v4

      - name: Set up Python
        uses: actions/setup-python@v5
        with:
          python-version: ${{ matrix.python-version }}

      - name: Install Poetry
        uses: snok/install-poetry@v1

      - name: Set environment variable
        run: |
          echo $(poetry run poetry env info --path)/bin >> $GITHUB_PATH

      - name: Install packages
        run: |
          make install

      - name: Install xarray
        run: |
          pip install --index-url https://test.pypi.org/simple --extra-index-url https://pypi.org/simple --pre xarray

      - name: Show installed packages
        run: |
          poetry show
          poetry show --tree

      - name: Run unit tests
        run: |
          make unit-test-cov