Skip to content

Basic support of jax.lax.gather in TessellateIPU. #109

Basic support of jax.lax.gather in TessellateIPU.

Basic support of jax.lax.gather in TessellateIPU. #109

Workflow file for this run

name: tests_internal
env:
GIT_MAIN_BRANCH: "main"
# Controls when the workflow will run.
on:
push:
branches: [ "main" ]
pull_request:
branches: [ "main" ]
# Allows you to run this workflow manually from the Actions tab.
workflow_dispatch:
jobs:
lint_and_typecheck:
if: github.repository != 'graphcore-research/tessellate-ipu'
runs-on: [self-hosted, Linux, X64, 20.04, Ubuntu, intel]
timeout-minutes: 10
steps:
- name: Cancel previous
uses: styfle/[email protected]
with:
access_token: ${{ github.token }}
if: ${{github.ref != 'refs/head/main'}}
- 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]
# Tessellate unit tests using IPU model
tessellate_unit_tests_ipu_model:
if: github.repository != 'graphcore-research/tessellate-ipu'
runs-on: [self-hosted, Linux, X64, 20.04, Ubuntu, amd]
container: graphcore/pytorch:3.2.0-ubuntu-20.04
timeout-minutes: 20
steps:
- name: Cancel previous
uses: styfle/[email protected]
with:
access_token: ${{ github.token }}
if: ${{github.ref != 'refs/head/main'}}
- uses: actions/checkout@v3
- name: Update pip
id: pip-cache
run: |
python3 -m pip install --upgrade build cmake nanobind ninja pip scikit_build_core[pyproject]
# Install JAX for IPU.
- name: Install JAX for IPU
run: |
pip3 install --default-timeout=100 -U numpy==1.23.5 scipy etils pytest
pip3 install jax==0.3.16+ipu jaxlib==0.3.15+ipu.sdk320 -f https://graphcore-research.github.io/jax-experimental/wheels.html
# Install other dependencies.
- name: Install dependencies
run: |
pip3 install --default-timeout=100 -r ./test-requirements.txt
# Build & install sdist package.
- name: Install TessellateIPU
run: |
python3 -m build --no-isolation
pip3 install dist/*.tar.gz
# Run repository unit tests.
- name: Run TessellateIPU unit tests (model)
run: |
cd ./tests/
JAX_IPU_USE_MODEL=true JAX_IPU_MODEL_NUM_TILES=16 pytest --tb=short -v --log-cli-level=INFO ./
# 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 ./
# Tessellate unit tests using IPU hardware
tessellate_unit_tests_ipu_hardware:
if: github.repository != 'graphcore-research/tessellate-ipu'
runs-on: [self-hosted, Linux, X64, 20.04, Ubuntu, pod16]
container:
image: graphcore/pytorch:3.2.0-ubuntu-20.04
options: --ulimit memlock=-1:-1 --cap-add=SYS_PTRACE --cap-add=IPC_LOCK --device=/dev/infiniband/ -e IPUOF_VIPU_API_HOST -e IPUOF_VIPU_API_PARTITION_ID --shm-size=128G
timeout-minutes: 20
steps:
- name: Cancel previous
uses: styfle/[email protected]
with:
access_token: ${{ github.token }}
if: ${{github.ref != 'refs/head/main'}}
- uses: actions/checkout@v3
- name: Instal Python base dependencies
id: pip-cache
run: |
python3 -m pip install --upgrade build cmake docker nanobind ninja pip scikit_build_core[pyproject]
# Necessary for attaching IPU hardware.
- name: Attach RDMA Network
run: |
python3 -c "import docker; client=docker.from_env(); client.networks.get('macvlan_rdma_swarm').connect(client.containers.get('${{ job.container.id }}'))"
gc-monitor
# Install JAX for IPU.
- name: Install JAX for IPU
run: |
pip3 install --default-timeout=100 -U numpy==1.23.5 scipy etils pytest
pip3 install jax==0.3.16+ipu jaxlib==0.3.15+ipu.sdk320 -f https://graphcore-research.github.io/jax-experimental/wheels.html
POPLAR_LOG_LEVEL=DEBUG TF_CPP_MIN_LOG_LEVEL=0 JAX_IPU_DEVICE_COUNT=1 python .github/scripts/check_ipu_hardware_device.py
- name: Install repository dependencies
run: |
pip3 install --default-timeout=100 -r ./test-requirements.txt
- name: Install TessellateIPU
run: |
python3 -m build --no-isolation
pip3 install dist/*.whl
# Run IPU hardware unit tests.
- name: Run TessellateIPU unit tests (hardware)
run: |
cd ./tests/
JAX_IPU_DEVICE_COUNT=2 pytest --tb=short -v --log-cli-level=INFO -m ipu_hardware ./