Include custom ops in docs #44
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: 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/[email protected] | |
| 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/[email protected] | |
| # 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/[email protected] | |
| 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 ./ |