Wheeldist

.github/release.py

@@ -0,0 +1,164 @@
+"""
+Release pipeline steps for .github/workflows/release.yaml.
+
+    python3 .github/release.py resolve-symlinks
+    python3 .github/release.py set-version <path/to/pyproject.toml> [...]
+    python3 .github/release.py install-cuda-toolkits <cuda-root> [...]
+    python3 .github/release.py install-rocm-sdk <path/to/pyproject.toml> <target-dir>
+"""
+
+import argparse
+import platform
+import re
+import shutil
+import subprocess
+import sys
+import tempfile
+import tomllib
+from pathlib import Path
+
+TAG_PATTERN = re.compile(r"v([0-9]+\.[0-9]+\.[0-9]+)")
+GIT_SYMLINK_MODE = "120000"
+CUDA_ROOT_PATTERN = re.compile(r"/usr/local/cuda-([0-9]+)\.([0-9]+)")
+CUDA_REPO_ARCH = {"x86_64": "x86_64", "aarch64": "sbsa"}
+ROCM_INDEX_PATTERN = re.compile(
+    r"https://download\.pytorch\.org/whl/rocm[0-9]+\.[0-9]+"
+)
+ROCM_VERSION_PATTERN = re.compile(r"[0-9]+\.[0-9]+\.[0-9]+")
+SHA256_PATTERN = re.compile(r"[0-9a-f]{64}")
+
+
+def resolve_symlinks(args):
+    """Replaces every symlink tracked by git with a copy of its target."""
+    entries = subprocess.run(
+        ["git", "ls-files", "-s", "-z"],
+        check=True,
+        capture_output=True,
+        text=True,
+    ).stdout.split("\0")
+
+    for entry in filter(None, entries):
+        info, path = entry.split("\t", 1)
+        if info.split()[0] == GIT_SYMLINK_MODE:
+            link = Path(path)
+            target = link.resolve()
+            link.unlink()
+            if target.is_dir():
+                shutil.copytree(target, link)
+            else:
+                shutil.copy2(target, link)
+            print(f"Resolved symlink {link} -> {target}")
+
+
+def get_version_from_tag():
+    tags = subprocess.run(
+        ["git", "tag", "--points-at", "HEAD"],
+        check=True,
+        capture_output=True,
+        text=True,
+    ).stdout.splitlines()
+
+    versions = [m.group(1) for m in map(TAG_PATTERN.fullmatch, tags) if m]
+    assert len(versions) == 1
+    return versions[0]
+
+
+def set_version(args):
+    """Sets the version in each pyproject.toml from the git tag of the form
+    vX.Y.Z on HEAD."""
+    import tomlkit
+
+    version = get_version_from_tag()
+    for pyproject in args.pyprojects:
+        doc = tomlkit.parse(pyproject.read_text())
+        assert "version" in doc["project"]
+        doc["project"]["version"] = version
+        pyproject.write_text(tomlkit.dumps(doc))
+
+
+def install_cuda_toolkits(args):
+    """Installs the CUDA headers and link-time libraries for each toolkit root
+    inside a manylinux container."""
+    subprocess.run(
+        [
+            "dnf",
+            "config-manager",
+            "--add-repo",
+            "https://developer.download.nvidia.com/compute/cuda/repos/rhel8/"
+            f"{CUDA_REPO_ARCH[platform.machine()]}/cuda-rhel8.repo",
+        ],
+        check=True,
+    )
+
+    for root in args.cuda_roots:
+        match = CUDA_ROOT_PATTERN.fullmatch(root)
+        assert match, f"expected a CUDA root like /usr/local/cuda-12.6, got {root!r}"
+        pkg_version = "-".join(match.groups())
+        packages = [
+            "cuda-nvcc",
+            "cuda-cudart-devel",
+            "cuda-nvrtc-devel",
+            "cuda-driver-devel",
+        ]
+        subprocess.run(
+            ["dnf", "install", "-y", *(f"{pkg}-{pkg_version}" for pkg in packages)],
+            check=True,
+        )
+
+
+def install_rocm_sdk(args):
+    """Installs rocm-sdk-core (HIP headers, libamdhip64 and libhiprtc), pinned
+    in [tool.openequivariance.rocm-sdk-core] of pyproject.toml, into a target
+    directory; skipped on non-x86_64 machines."""
+    if platform.machine() == "x86_64":
+        with open(args.pyproject, "rb") as f:
+            pin = tomllib.load(f)["tool"]["openequivariance"]["rocm-sdk-core"]
+        assert ROCM_INDEX_PATTERN.fullmatch(pin["index-url"])
+        assert ROCM_VERSION_PATTERN.fullmatch(pin["version"])
+        assert SHA256_PATTERN.fullmatch(pin["sha256"])
+        with tempfile.TemporaryDirectory() as tmp:
+            requirements = Path(tmp) / "rocm_requirements.txt"
+            requirements.write_text(
+                f"--index-url {pin['index-url']}\n"
+                f"rocm-sdk-core=={pin['version']} "
+                f"--hash=sha256:{pin['sha256']}\n"
+            )
+            subprocess.run(
+                [
+                    sys.executable,
+                    "-m",
+                    "pip",
+                    "install",
+                    "--no-deps",
+                    "--require-hashes",
+                    "--target",
+                    str(args.target),
+                    "-r",
+                    str(requirements),
+                ],
+                check=True,
+            )
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description=__doc__.split("\n\n")[0].strip())
+    steps = parser.add_subparsers(dest="step", required=True)
+
+    step = steps.add_parser("resolve-symlinks", help=resolve_symlinks.__doc__)
+    step.set_defaults(func=resolve_symlinks)
+
+    step = steps.add_parser("set-version", help=set_version.__doc__)
+    step.add_argument("pyprojects", nargs="+", type=Path)
+    step.set_defaults(func=set_version)
+
+    step = steps.add_parser("install-cuda-toolkits", help=install_cuda_toolkits.__doc__)
+    step.add_argument("cuda_roots", nargs="+")
+    step.set_defaults(func=install_cuda_toolkits)
+
+    step = steps.add_parser("install-rocm-sdk", help=install_rocm_sdk.__doc__)
+    step.add_argument("pyproject", type=Path)
+    step.add_argument("target", type=Path)
+    step.set_defaults(func=install_rocm_sdk)
+
+    args = parser.parse_args()
+    args.func(args)

.github/workflows/release.yaml

@@ -4,8 +4,41 @@ on:
   workflow_dispatch:
 
 jobs:
-  build-oeq:
-    name: Build oeq dist 
+  build-oeq-wheels:
+    name: Build oeq wheels (${{ matrix.arch }})
+    runs-on: ${{ matrix.runner }}
+    timeout-minutes: 60
+    strategy:
+      matrix:
+        include:
+          - runner: ubuntu-24.04
+            arch: x86_64
+          - runner: ubuntu-24.04-arm
+            arch: aarch64
+    steps:
+    - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
+      with:
+        fetch-tags: true
+    - name: set up Python
+      uses: actions/setup-python@6d4fd22b20465ccc0f7354476588e8df7c9ca0e0 # v7.1.0
+      with:
+        python-version: '3.12'
+    - name: replace symlinks with copies
+      run: python3 .github/release.py resolve-symlinks
+    - name: set version from tag
+      run: |
+        python3 -m pip install -r .github/workflows/requirements_release_ci.txt
+        python3 .github/release.py set-version openequivariance/pyproject.toml
+    - name: build wheels
+      run: python3 -m cibuildwheel --output-dir wheelhouse openequivariance
+    - name: store the wheels
+      uses: actions/upload-artifact@cf430e030ddbb5b0abf93d22962f4752f3646cd9 # v7.0.2
+      with:
+        name: oeq-wheels-${{ matrix.arch }}
+        path: wheelhouse/*.whl
+
+  build-oeq-sdist:
+    name: Build oeq sdist
     runs-on: ubuntu-latest
     steps:
     - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
@@ -14,27 +47,29 @@ jobs:
     - name: set up Python
       uses: actions/setup-python@6d4fd22b20465ccc0f7354476588e8df7c9ca0e0 # v7.1.0
       with:
-        python-version: '3.10'
+        python-version: '3.12'
+    - name: replace symlinks with copies
+      run: python3 .github/release.py resolve-symlinks
     - name: set version from tag
       run: |
         python3 -m pip install -r .github/workflows/requirements_release_ci.txt
-        python3 .github/autoversion.py openequivariance/pyproject.toml
+        python3 .github/release.py set-version openequivariance/pyproject.toml
     - name: install dependencies, then build source tarball
       run: |
         cd openequivariance
-        python3 -m pip install build 
+        python3 -m pip install build
         python3 -m build --sdist
-    - name: store the distribution packages 
+    - name: store the distribution packages
       uses: actions/upload-artifact@cf430e030ddbb5b0abf93d22962f4752f3646cd9 # v7.0.2
       with:
-        name: oeq-package-distributions
+        name: oeq-sdist
         path: openequivariance/dist/
 
   pypi-publish:
     name: Upload oeq to PyPI
     runs-on: ubuntu-latest
-    needs: build-oeq
-    environment: 
+    needs: [build-oeq-wheels, build-oeq-sdist]
+    environment:
         name: pypi
         url: https://pypi.org/p/openequivariance
     permissions:
@@ -43,8 +78,9 @@ jobs:
       - name: download the distributions
         uses: actions/download-artifact@9000827ccba6bdab643e8b6fd33ac0654aef8333 # v8.0.2
         with:
-          name: oeq-package-distributions
-          path: dist/ 
+          pattern: oeq-*
+          merge-multiple: true
+          path: dist/
       - name: publish package distributions to PyPI
         uses: pypa/gh-action-pypi-publish@7245906fb61fa33a13554f9d5fb856fe61758086 # v1.14.3
 
@@ -61,16 +97,18 @@ jobs:
       uses: actions/setup-python@6d4fd22b20465ccc0f7354476588e8df7c9ca0e0 # v7.1.0
       with:
         python-version: '3.10'
+    - name: replace symlinks with copies
+      run: python3 .github/release.py resolve-symlinks
     - name: set version from tag
       run: |
         python3 -m pip install -r .github/workflows/requirements_release_ci.txt
-        python3 .github/autoversion.py openequivariance_extjax/pyproject.toml
+        python3 .github/release.py set-version openequivariance_extjax/pyproject.toml
     - name: install dependencies, then build source tarball
       run: |
         cd openequivariance_extjax
-        python3 -m pip install build 
-        python3 -m build --sdist 
-    - name: store the distribution packages 
+        python3 -m pip install build
+        python3 -m build --sdist
+    - name: store the distribution packages
       uses: actions/upload-artifact@cf430e030ddbb5b0abf93d22962f4752f3646cd9 # v7.0.2
       with:
         name: oeqjax-package-distributions
@@ -80,7 +118,7 @@ jobs:
     name: Upload extjax to PyPI
     runs-on: ubuntu-latest
     needs: build-oeq-extjax
-    environment: 
+    environment:
         name: pypi
         url: https://pypi.org/p/openequivariance_extjax
     permissions:
@@ -90,6 +128,6 @@ jobs:
         uses: actions/download-artifact@9000827ccba6bdab643e8b6fd33ac0654aef8333 # v8.0.2
         with:
           name: oeqjax-package-distributions
-          path: dist/ 
+          path: dist/
       - name: publish package distributions to PyPI
         uses: pypa/gh-action-pypi-publish@7245906fb61fa33a13554f9d5fb856fe61758086 # v1.14.3

.github/workflows/requirements_cuda_ci.txt

@@ -1,6 +1,7 @@
+--extra-index-url https://download.pytorch.org/whl/cu129
 numpy==2.2.5
-torch==2.13.0 --index-url https://download.pytorch.org/whl/cu129
+torch==2.13.0+cu129
 pytest==9.0.3
 ninja==1.11.1.4
-nanobind==2.10.2
+nanobind==3.1.0
 scikit-build-core==0.11.6

.github/workflows/requirements_release_ci.txt

@@ -1 +1,2 @@
 tomlkit==0.15.2
+cibuildwheel==4.3.5

.github/workflows/requirements_sycl_ci.txt

@@ -3,5 +3,5 @@ numpy==2.2.5
 torch==2.13.0+xpu
 pytest==9.0.3
 ninja==1.11.1.4
-nanobind==2.10.2
+nanobind==3.1.0
 scikit-build-core==0.11.6