jax-ipu-tests-internal.yaml

name: CI_jax_ipu_tests_internal

env:
  GIT_MAIN_BRANCH: "jax-v0.3.16-ipu"

# Controls when the workflow will run.
on:
  push:
    branches: [ "jax-v0.3.16-ipu" ]
  pull_request:
    branches: [ "jax-v0.3.16-ipu" ]

  # Allows you to run this workflow manually from the Actions tab.
  workflow_dispatch:

jobs:
  lint_and_typecheck:
    if: github.repository != 'graphcore-research/jax-experimental'
    runs-on: [self-hosted, Linux, X64, 20.04, Ubuntu]
    timeout-minutes: 10
    steps:
      - name: Cancel previous
        uses: styfle/cancel-workflow-action@0.11.0
        with:
          access_token: ${{ github.token }}
        if: ${{github.ref != 'refs/head/jax-v0.3.16-ipu'}}
      - uses: actions/checkout@v3
      - name: Set up Python 3.8
        uses: actions/setup-python@v4
        with:
          python-version: 3.8
      - uses: pre-commit/action@v3.0.0

  # JAX unit tests using IPU model
  jax_unit_tests_ipu_model:
    if: github.repository != 'graphcore-research/jax-experimental'
    runs-on: [self-hosted, Linux, X64, 20.04, Ubuntu]
    strategy:
      matrix:
        container: ["graphcore/pytorch:3.2.0-ubuntu-20.04", "graphcore/pytorch:3.3.0-ubuntu-20.04"]
    container:
      image: ${{ matrix.container }}
    timeout-minutes: 10
    steps:
      - name: Cancel previous
        uses: styfle/cancel-workflow-action@0.11.0
        with:
          access_token: ${{ github.token }}
        if: ${{github.ref != 'refs/head/jax-v0.3.16-ipu'}}
      - uses: actions/checkout@v3
      - name: Get pip cache dir
        id: pip-cache
        run: |
          python3 -m pip install --upgrade pip wheel setuptools
      # Query JAX+JAXLIB IPU wheels from latest release.
      - name: Fetch pre-compiled JAXLIB IPU release wheels
        uses: dsaltares/fetch-gh-release-asset@master
        with:
          repo: ${{ github.repository }}
          version: 'latest'
          regex: true
          file: "jax.*.whl"
          target: "wheels/"
          token: ${{ secrets.JAX_IPU_EXPERIMENTAL_PAT }}
      - name: Pip install pre-compiled JAXLIB & build JAX
        run: |
          ls ./wheels/
          export PATH=$PATH:/opt/poplar/bin
          export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/opt/poplar/lib
          export SDK_VERSION_SHORT="$(python ./build/ipu/get_poplar_sdk_version.py)"
          echo $SDK_VERSION_SHORT
          pip3 install -U numpy==1.23.5 scipy etils pytest
          pip3 install ./wheels/jaxlib-0.3.15+ipu.sdk${SDK_VERSION_SHORT}-cp38*.whl
          python3 setup.py bdist_wheel --universal
          pip3 install dist/*.whl
      # Run IPU specific unit tests
      - name: Run JAX IPU model unit tests
        run: |
          export PATH=$PATH:/opt/poplar/bin
          export CPATH=$CPATH:/opt/poplar/include
          export LIBRARY_PATH=$LIBRARY_PATH:/opt/poplar/lib
          export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/opt/poplar/lib
          JAX_IPU_DEVICE_COUNT=2 JAX_IPU_USE_MODEL=true JAX_IPU_MODEL_NUM_TILES=16 pytest --tb=short -vv --log-cli-level=INFO ./tests/ipu/
      # Dockerized workflow known to create issues with self-hosted servers.
      # Solution is to fully cleanup the workspace for the next action.
      # See: https://stackoverflow.com/questions/70483902/how-to-actually-clean-up-the-repository-on-self-hosted-runner-after-github-actio
      - name: Cleanup GITHUB_WORKSPACE folder
        run: |
          ls -la ./
          rm -rf ./* || true
          rm -rf ./.??* || true
          ls -la ./