Merge pull request #350 from torchmd/fix_windows

Fix windows builds

.github/workflows/CI.yml

@@ -12,7 +12,14 @@ on:
 
 jobs:
   test:
-    runs-on: "ubuntu-latest"
+    runs-on: ${{ matrix.os }}
+
+    strategy:
+      fail-fast: false
+      matrix:
+        os: ["ubuntu-latest", "macos-latest", "macos-13", "windows-2022"]
+        python-version: ["3.10"]
+
     defaults: # Needed for conda
       run:
         shell: bash -l {0}
@@ -21,18 +28,50 @@ jobs:
       - name: Check out
         uses: actions/checkout@v4
 
-      - name: Create a conda environment
-        uses: conda-incubator/setup-miniconda@v3
+      - uses: conda-incubator/setup-miniconda@v3
+        with:
+          python-version: ${{ matrix.python-version }}
+          channels: conda-forge
+          conda-remove-defaults: "true"
+        if: matrix.os != 'macos-13'
+
+      - uses: conda-incubator/setup-miniconda@v3
         with:
-          activate-environment: torchmd-net
-          environment-file: environment.yml
+          python-version: ${{ matrix.python-version }}
+          channels: conda-forge
+          mamba-version: "*"
+          conda-remove-defaults: "true"
+        if: matrix.os == 'macos-13'
 
-      - name: Install the package
-        run: pip -vv install .
+      - name: Install OS-specific conda dependencies
+        run: |
+          if [[ "${{ runner.os }}" == "Linux" ]]; then
+            conda install --file conda_deps_linux.txt --channel conda-forge --override-channels
+          elif [[ "${{ runner.os }}" == "macos-13" ]]; then
+            mamba install --file conda_deps_osx.txt --channel conda-forge --override-channels
+          elif [[ "${{ runner.os }}" == "macOS" ]]; then
+            conda install --file conda_deps_osx.txt --channel conda-forge --override-channels
+          elif [[ "${{ runner.os }}" == "Windows" ]]; then
+            conda install --file conda_deps_win.txt --channel conda-forge --override-channels
+          fi
+
+      - name: Install testing packages
+        run: conda install -y -c conda-forge flake8 pytest psutil
 
       - name: List the conda environment
         run: conda list
 
+      - name: Build and install the package
+        run: |
+          if [[ "${{ runner.os }}" == "Windows" ]]; then
+            export LIB="C:/Miniconda/envs/test/Library/lib"
+            pip -vv install .
+          else
+            pip -vv install .
+          fi
+        env:
+          CPU_ONLY: 1
+
       - name: Lint with flake8
         run: |
           # stop the build if there are Python syntax errors or undefined names
@@ -42,6 +81,10 @@ jobs:
 
       - name: Run tests
         run: pytest -v -s
+        env:
+          CPU_ONLY: 1
+          WINDOWS_TESTS: ${{ runner.os == 'Windows' && 'true' || 'false' }}
+          OMP_PREFIX: ${{ runner.os == 'macOS' && '/Users/runner/miniconda3/envs/test' || '' }}
 
       - name: Test torchmd-train utility
         run: torchmd-train --help

conda_deps_linux.txt

@@ -0,0 +1,10 @@
+h5py
+nnpops
+pip
+libtorch >=2.5.1
+pytorch-cpu >=2.5.1
+pytorch_geometric
+lightning
+torchmetrics
+tqdm
+gxx

conda_deps_osx.txt

@@ -0,0 +1,12 @@
+h5py
+nnpops
+pip
+libtorch>=2.5.1
+pytorch-cpu>=2.5.1
+pytorch_geometric
+lightning
+torchmetrics
+tqdm
+clangxx
+llvm-openmp
+pybind11

conda_deps_win.txt

@@ -0,0 +1,13 @@
+h5py
+pip
+libtorch >=2.5.1
+pytorch-cpu >=2.5.1
+pytorch_geometric
+lightning
+torchmetrics
+tqdm
+vc
+vc14_runtime
+vs2015_runtime
+vs2019_win-64
+sleef

setup.py

@@ -7,6 +7,9 @@
 import torch
 from torch.utils.cpp_extension import BuildExtension, CUDAExtension, include_paths, CppExtension
 import os
+import sys
+
+is_windows = sys.platform == 'win32'
 
 try:
     version = (
@@ -43,7 +46,7 @@ def set_torch_cuda_arch_list():
 ExtensionType = CppExtension if not use_cuda else CUDAExtension
 extensions = ExtensionType(
     name='torchmdnet.extensions.torchmdnet_extensions',
-    sources=[os.path.join(extension_root, "extensions.cpp")] + neighbor_sources,
+    sources=[os.path.join(extension_root, "torchmdnet_extensions.cpp")] + neighbor_sources,
     include_dirs=include_paths(),
     define_macros=[('WITH_CUDA', 1)] if use_cuda else [],
 )
@@ -58,5 +61,5 @@ def set_torch_cuda_arch_list():
             'build_ext': BuildExtension.with_options(no_python_abi_suffix=True, use_ninja=False)},
         include_package_data=True,
         entry_points={"console_scripts": ["torchmd-train = torchmdnet.scripts.train:main"]},
-        package_data={"torchmdnet": ["extensions/torchmdnet_extensions.so"]},
+        package_data={"torchmdnet": ["extensions/torchmdnet_extensions.so"] if not is_windows else ["extensions/torchmdnet_extensions.dll"]},
     )

tests/test_dataset_comp6.py

@@ -39,7 +39,7 @@ def test_dataset_s66x8():
             ),
             atol=1e-4,
         )
-        assert pt.allclose(sample.y, pt.tensor([[-5755.7288331]],dtype=pt.float64))
+        assert pt.allclose(sample.y, pt.tensor([[-5755.7288331]], dtype=pt.float64))
         assert pt.allclose(
             sample.neg_dy,
             -pt.tensor(
@@ -65,3 +65,8 @@ def test_dataset_s66x8():
         data_loader = DataLoader(dataset=data_set, batch_size=32, num_workers=2)
         for batch in data_loader:
             pass
+
+        # Extra cleanup for Windows tests which are sensitive to open mmap files
+        del sample
+        del data_loader
+        del data_set

tests/test_module.py

@@ -12,6 +12,7 @@
 from torchmdnet.module import LNNP
 from torchmdnet.data import DataModule
 from torchmdnet import priors
+import os
 
 from utils import load_example_args, DummyDataset
 
@@ -30,6 +31,15 @@ def test_load_model():
 @mark.parametrize("use_atomref", [True, False])
 @mark.parametrize("precision", [32, 64])
 def test_train(model_name, use_atomref, precision, tmpdir):
+    import torch
+    import platform
+
+    accelerator = "auto"
+    if platform.system() == "Darwin" and (os.getenv("CI", None) is not None):
+        # MPS backend runs out of memory on Github
+        torch.set_default_device("cpu")
+        accelerator = "cpu"
+
     args = load_example_args(
         model_name,
         remove_prior=not use_atomref,
@@ -53,6 +63,12 @@ def test_train(model_name, use_atomref, precision, tmpdir):
 
     module = LNNP(args, prior_model=prior)
 
-    trainer = pl.Trainer(max_steps=10, default_root_dir=tmpdir, precision=args["precision"],inference_mode=False)
+    trainer = pl.Trainer(
+        max_steps=10,
+        default_root_dir=tmpdir,
+        precision=args["precision"],
+        inference_mode=False,
+        accelerator=accelerator,
+    )
     trainer.fit(module, datamodule)
     trainer.test(module, datamodule)

tests/test_neighbors.py

@@ -671,6 +671,11 @@ def test_per_batch_box(device, strategy, n_batches, use_forward):
 def test_torch_compile(device, dtype, loop, include_transpose):
     import sys
 
+    # Skip if WINDOWS_TESTS is set to true
+    if os.environ.get("WINDOWS_TESTS", "false") == "true":
+        # torch.compile doesn't detect cl.exe on Windows Github Actions
+        pytest.skip("Skipping test on Windows")
+
     if sys.version_info >= (3, 12):
         pytest.skip("Not available in this version")
     if torch.__version__ < "2.0.0":