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]
    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/main'}}
      - 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

  # 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]
    container: graphcore/pytorch:3.2.0-ubuntu-20.04
    timeout-minutes: 20
    steps:
      - name: Cancel previous
        uses: styfle/cancel-workflow-action@0.11.0
        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/cancel-workflow-action@0.11.0
        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
          cd .github/scripts
          POPLAR_LOG_LEVEL=DEBUG TF_CPP_MIN_LOG_LEVEL=0 JAX_IPU_DEVICE_COUNT=1 python 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 ./
