Skip to content

Remove StateDict from torchsim integration #361

Remove StateDict from torchsim integration

Remove StateDict from torchsim integration #361

Workflow file for this run

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:
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v4
- name: Set up Python
uses: actions/setup-python@v5
with:
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:
needs: lint
runs-on: ${{ github.event.repository.private && 'self-hosted' || 'ubuntu-latest' }}
strategy:
matrix:
torch-version: [2.2.0, 2.*]
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.*" ]; 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