Run Tests workflow (mir-group/nequip)

name: Run Tests
 
on:
  push:
    branches:
      - main
      - develop
  pull_request:
    branches:
      - main
      - develop
 
concurrency:
  # github.workflow will be the name of the workflow, e.g. "Run Tests"
  # github.ref will be the branch or tag ref, e.g. "refs/heads/main" or, for a PR, "refs/pull/123/merge"
  group: ${{ github.workflow }}-${{ github.ref }}
  cancel-in-progress: true
 
# yamllint disable rule:line-length
jobs:
  lint:
    timeout-minutes: 30
    runs-on: latchkey-small
    steps:
      - uses: actions/checkout@v4
      - name: Set up Python
        uses: actions/setup-python@v5
        with:
          cache: 'pip'
          python-version: "3.12"
      - name: Install pre-commit
        run: pip install pre-commit
      - name: Run pre-commit
        run: pre-commit run --all-files --show-diff-on-failure
 
  run-tests:
    timeout-minutes: 30
    needs: lint
    runs-on: ${{ github.event.repository.private && 'self-hosted' || 'ubuntu-latest' }}
    strategy:
      matrix:
        torch-version: [2.2.0, 2.10.*]
    container:
      image: nvidia/cuda:12.6.1-cudnn-devel-ubuntu24.04
      options: ${{ github.event.repository.private && '--gpus all' || ' ' }}
    steps:
      - uses: actions/checkout@v4
      - name: Install Python and create venv
        run: |
          apt update && apt install -y python3 python3-pip python3-venv git
          python3 -m venv /opt/venv
          echo "/opt/venv/bin" >> $GITHUB_PATH
      - name: Check Python version
        run: |
          python3 --version
          which python3
      - name: Clear pip cache on self-hosted runner
        if: ${{ github.event.repository.private && contains(runner.labels, 'self-hosted') }}
        run: |
          python3 -m pip cache dir
          python3 -m pip cache purge || true
          rm -rf ~/.cache/pip
      - name: Disable pip cache
        run: |
          echo "PIP_NO_CACHE_DIR=1" >> $GITHUB_ENV
          echo "PIP_NO_CACHE=1" >> $GITHUB_ENV
      - name: Configure torch cache directories
        run: |
          echo "TORCH_EXTENSIONS_DIR=/tmp/torch_extensions_${GITHUB_RUN_ID}_${GITHUB_RUN_ATTEMPT}" >> $GITHUB_ENV
          echo "TORCHINDUCTOR_CACHE_DIR=/tmp/torchinductor_${GITHUB_RUN_ID}_${GITHUB_RUN_ATTEMPT}" >> $GITHUB_ENV
          echo "TRITON_CACHE_DIR=/tmp/triton_${GITHUB_RUN_ID}_${GITHUB_RUN_ATTEMPT}" >> $GITHUB_ENV
      - name: Install PyTorch
        env:
          TORCH: "${{ matrix.torch-version }}"
        run: |
          # use CPU only on GH runner
          if [ "${{ github.event.repository.private }}" = "true" ]; then
              pip install torch==${TORCH} --upgrade
          else
              pip install torch==${TORCH} --index-url https://download.pytorch.org/whl/cpu --upgrade
          fi
      - name: Install test dependencies
        env:
          TORCH: "${{ matrix.torch-version }}"
        run: |
          # install packages that aren't required dependencies but that the tests do need
          pip install h5py packaging
          # for vesin, the following is to avoid
          # `AttributeError: module 'torch' has no attribute 'uint64'`
          if [ "${{ matrix.torch-version }}" = "2.2.0" ]; then
            pip install "vesin==0.4.2"
          else
            pip install vesin
          fi
          # install torch-sim-atomistic only for latest torch version
          if [ "${{ matrix.torch-version }}" = "2.10.*" ]; then
            pip install torch-sim-atomistic
          fi
          #pip install schedulefree  #Disabled to reduce CI runtime; uncomment to test ScheduleFreeLightningModule
          # install CuEquivariance if CUDA is available
          if python3 -c "import torch; exit(0 if torch.cuda.is_available() else 1)"; then
            echo "Installing CuEquivariance (CUDA available)"
            pip install cuequivariance-torch cuequivariance-ops-torch-cu12
          else
            echo "Skipping CuEquivariance installation (CUDA not available)"
          fi
          # install OpenEquivariance only if torch >= 2.7 and CUDA is available
          if python3 -c "import torch; import packaging.version; exit(0 if packaging.version.parse(torch.__version__) >= packaging.version.parse('2.7') else 1)"; then
            if python3 -c "import torch; exit(0 if torch.cuda.is_available() else 1)"; then
              echo "Installing OpenEquivariance (torch >= 2.7 and CUDA available)"
              pip install openequivariance
            else
              echo "Skipping OpenEquivariance installation (CUDA not available)"
            fi
          else
            echo "Skipping OpenEquivariance installation (torch < 2.7)"
          fi
          # Pin numpy<2 and matscipy 1.1.1 for torch 2.2 to avoid compatibility issues
          if [ "${{ matrix.torch-version }}" = "2.2.0" ]; then
            pip install "numpy<2" "matscipy==1.1.1"
          fi
          # Pin lightning <2.6.0 due to test_restarts failures
          pip install "lightning<2.6.0"
          pip install --upgrade-strategy only-if-needed .
          pip install pytest pytest-xdist[psutil]
      #- name: Clear default torch extension cache
      #  run: |
      #    python3 - <<'PY'
      #    import shutil
      #    from torch.utils.cpp_extension import get_default_build_root
 
      #    default_root = get_default_build_root()
      #    print(f"default torch extension cache root: {default_root}")
      #    shutil.rmtree(default_root, ignore_errors=True)
      #    print("deleted default torch extension cache root")
      #    PY
      - name: Test with pytest
        run: |
          # See https://github.com/pytest-dev/pytest/issues/1075
          PYTHONHASHSEED=0 pytest -xv -n auto tests/
      #- name: Clean torch caches
      #  if: always()
      #  run: |
      #    rm -rf "$TORCH_EXTENSIONS_DIR" "$TORCHINDUCTOR_CACHE_DIR" "$TRITON_CACHE_DIR"
# yamllint enable