Reworked setup system using both pyproject.toml and setup.py. Reworked CI

.github/workflows/deploy_pypi.yml

@@ -7,27 +7,31 @@ on:
   workflow_dispatch:
 
 jobs:
-  build:
-
+  pypi-publish:
+    name: Publish release to PyPI
     runs-on: ubuntu-latest
+    environment:
+      name: pypi
+      url: https://pypi.org/p/torch-harmonics
 
     steps:
       - uses: actions/checkout@v4
+
       - name: Set up Python 3.9
         uses: actions/setup-python@v5
         with:
-          python-version: 3.9
-      - name: Install dependencies
-        run: |
-          python -m pip install --upgrade pip
-          python -m pip install torch --index-url https://download.pytorch.org/whl/cpu
-          python -m pip install setuptools wheel
-      - name: Install package
+          python-version: "3.9"
+
+      - name: Install build dependencies
         run: |
-          python -m pip install -e .
-      - name: Build a binary wheel and a source tarball
+          python3 -m pip install --upgrade pip setuptools build wheel
+          python3 -m pip install numpy
+          python3 -m pip install torch --extra-index-url https://download.pytorch.org/whl/cpu
+
+      - name: Build distribution
         run: |
-          python setup.py sdist bdist_wheel
+          python3 -m build --sdist
+
       - name: Publish package to PyPI
         uses: pypa/gh-action-pypi-publish@release/v1
         with:

.github/workflows/tests.yml

@@ -1,25 +1,31 @@
-name: tests
+name: Run local tests
 
-on: [push]
+on:
+  push:
+  workflow_dispatch:
 
 jobs:
   build:
-
+    name: Run local tests
     runs-on: ubuntu-latest
 
     steps:
     - uses: actions/checkout@v4
+
     - name: Set up Python 3.9
       uses: actions/setup-python@v5
       with:
-        python-version: '3.10'
+        python-version: "3.9"
+
     - name: Install dependencies
       run: |
         python3 -m pip install --upgrade pip setuptools wheel
-        python3 -m pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cpu
+        python3 -m pip install torch --extra-index-url https://download.pytorch.org/whl/cpu
+
     - name: Install package
       run: |
         python3 -m pip install -e .
+
     - name: Test with pytest
       run: |
         python3 -m pip install pytest pytest-cov parameterized

README.md

@@ -80,9 +80,11 @@ Download directly from PyPI:
 ```bash
 pip install torch-harmonics
 ```
-If you would like to have accelerated CUDA extensions for the discrete-continuous convolutions, please use the '--cuda_ext' flag:
+If you would like to enforce the compilation of CUDA extensions for the discrete-continuous convolutions, you can do so by setting the `FORCE_CUDA_EXTENSION` flag. You may also want to set appropriate architectures with the `TORCH_CUDA_ARCH_LIST` flag.
 ```bash
-pip install --global-option --cuda_ext torch-harmonics
+export FORCE_CUDA_EXTENSION=1
+export TORCH_CUDA_ARCH_LIST="7.0 7.2 7.5 8.0 8.6 8.7 9.0+PTX"
+pip install torch-harmonics
 ```
 :warning: Please note that the custom CUDA extensions currently only support CUDA architectures >= 7.0.
 

pyproject.toml

@@ -0,0 +1,51 @@
+[build-system]
+requires = [ "setuptools", "setuptools-scm", "torch>=2.4.0"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "torch_harmonics"
+authors = [
+    { name="Boris Bonev", email = "bbonev@nvidia.com" },
+    { name="Thorsten Kurth", email = "tkurth@nvidia.com" },
+    { name="Mauro Bisson" },
+    { name="Massimiliano Fatica" },
+    { name="Jean Kossaifi" },
+    { name="Nikola Kovachki" },
+    { name="Christian Hundt" },
+]
+
+maintainers = [
+    { name="Thorsten Kurth", email = "tkurth@nvidia.com" },
+    { name="Boris Bonev", email = "bbonev@nvidia.com" },
+]
+
+readme = "README.md"
+
+dynamic = ["version"]
+
+description = "Differentiable signal processing on the sphere for PyTorch."
+requires-python = ">=3.9"
+
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Programming Language :: Python :: 3.9",
+    "License :: OSI Approved :: BSD License",
+    "Operating System :: OS Independent",
+]
+
+dependencies = [
+    "torch>=2.4.0",
+    "numpy>=1.22.4,<1.25",
+]
+
+[tool.setuptools.dynamic]
+version = {attr = "torch_harmonics.__version__"}
+
+[tool.setuptools.packages.find]
+    include = ["torch_harmonics*"]
+
+[project.optional-dependencies]
+dev = [
+    "pytest>=6.0.0",
+    "coverage>=6.5.0",
+]

setup.py

@@ -29,67 +29,44 @@
 # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 #
 
-import sys
+import os, sys
+import warnings
 
-try:
-    from setuptools import setup, find_packages
-except ImportError:
-    from distutils.core import setup, find_packages
-
-import re
-from pathlib import Path
-
-import torch
-from torch.utils import cpp_extension
-
-def version(root_path):
-    """Returns the version taken from __init__.py
-
-    Parameters
-    ----------
-    root_path : pathlib.Path
-        path to the root of the package
+from setuptools import setup, find_packages
+from setuptools.command.install import install
 
-    Reference
-    ---------
-    https://packaging.python.org/guides/single-sourcing-package-version/
-    """
-    version_path = root_path.joinpath("torch_harmonics", "__init__.py")
-    with version_path.open() as f:
-        version_file = f.read()
-    version_match = re.search(r"^__version__ = ['\"]([^'\"]*)['\"]", version_file, re.M)
-    if version_match:
-        return version_match.group(1)
-    raise RuntimeError("Unable to find version string.")
+# some code to handle the building of custom modules
+FORCE_CUDA_EXTENSION = os.getenv("FORCE_CUDA_EXTENSION", "0") == "1"
+BUILD_CPP = BUILD_CUDA = False
 
+# try to import torch
+try:
+    import torch
 
-def readme(root_path):
-    """Returns the text content of the README.md of the package
-
-    Parameters
-    ----------
-    root_path : pathlib.Path
-        path to the root of the package
-    """
-    with root_path.joinpath("README.md").open(encoding="UTF-8") as f:
-        return f.read()
+    print(f"setup.py with torch {torch.__version__}")
+    from torch.utils.cpp_extension import BuildExtension, CppExtension
 
+    BUILD_CPP = True
+    from torch.utils.cpp_extension import CUDA_HOME, CUDAExtension
 
-def get_ext_modules(argv):
+    BUILD_CUDA = FORCE_CUDA_EXTENSION or (torch.cuda.is_available() and (CUDA_HOME is not None))
+except (ImportError, TypeError, AssertionError, AttributeError) as e:
+    warnings.warn(f"building custom extensions skipped: {e}")
 
-    compile_cuda_extension = False
+def get_ext_modules():
 
-    if "--cuda_ext" in sys.argv:
-        sys.argv.remove("--cuda_ext")
-        compile_cuda_extension = True
+    ext_modules = []
+    cmdclass = {}
 
-    ext_modules = [
-        cpp_extension.CppExtension("disco_helpers", ["torch_harmonics/csrc/disco/disco_helpers.cpp"]),
-    ]
+    if BUILD_CPP:
+        print(f"Compiling helper routines for torch-harmonics.")
+        ext_modules.append(CppExtension("disco_helpers", ["torch_harmonics/csrc/disco/disco_helpers.cpp"]))
+        cmdclass["build_ext"] = BuildExtension
 
-    if torch.cuda.is_available() or compile_cuda_extension:
+    if BUILD_CUDA:
+        print(f"Compiling custom CUDA kernels for torch-harmonics.")
         ext_modules.append(
-            cpp_extension.CUDAExtension(
+            CUDAExtension(
                 "disco_cuda_extension",
                 [
                     "torch_harmonics/csrc/disco/disco_interface.cu",
@@ -98,37 +75,16 @@ def get_ext_modules(argv):
                 ],
             )
         )
+        cmdclass["build_ext"] = BuildExtension
+
+    return ext_modules, cmdclass
+
+if __name__ == "__main__":
+
+    ext_modules, cmdclass = get_ext_modules()
 
-    return ext_modules
-
-
-root_path = Path(__file__).parent
-README = readme(root_path)
-VERSION = version(root_path)
-
-# external modules
-ext_modules = get_ext_modules(sys.argv)
-
-config = {
-    "name": "torch_harmonics",
-    "packages": find_packages(),
-    "description": "A differentiable spherical harmonic transform for PyTorch.",
-    "long_description": README,
-    "long_description_content_type": "text/markdown",
-    "url": "https://github.com/NVIDIA/torch-harmonics",
-    "author": "Boris Bonev",
-    "author_email": "bbonev@nvidia.com",
-    "version": VERSION,
-    "install_requires": ["torch", "numpy"],
-    "extras_require": {
-        "sfno": ["tensorly", "tensorly-torch"],
-    },
-    "license": "Modified BSD",
-    "scripts": [],
-    "include_package_data": True,
-    "classifiers": ["Topic :: Scientific/Engineering", "License :: OSI Approved :: BSD License", "Programming Language :: Python :: 3"],
-    "ext_modules": ext_modules,
-    "cmdclass": {"build_ext": cpp_extension.BuildExtension} if ext_modules else {},
-}
-
-setup(**config)
+    setup(
+        packages=find_packages(),
+        ext_modules=ext_modules,
+        cmdclass=cmdclass,
+    )