fused-ops: Python 3.12 support, manylinux retagging, TestPyPI toggle, wheel caching

.github/workflows/publish_fusedops.yml

@@ -5,6 +5,13 @@ on:
     types:
       - published
   workflow_dispatch:
+    inputs:
+      use_testpypi:
+        description: 'Upload to TestPyPI instead of PyPI'
+        required: false
+        default: 'false'
+        type: choice
+        options: ['true', 'false']
 
 jobs:
   build:
@@ -37,22 +44,49 @@ jobs:
             cuda_short: "124"
             torch: "2.5.0"
             python: "3.11"
+          - cuda: "11.8.0"
+            cuda_short: "118"
+            torch: "2.2.0"
+            python: "3.12"
+          - cuda: "12.1.0"
+            cuda_short: "121"
+            torch: "2.3.0"
+            python: "3.12"
+          - cuda: "12.4.0"
+            cuda_short: "124"
+            torch: "2.5.0"
+            python: "3.12"
 
     container:
       image: nvidia/cuda:${{ matrix.cuda }}-devel-ubuntu22.04
 
+    defaults:
+      run:
+        shell: bash -euo pipefail {0}
+
     steps:
       - name: Checkout
         uses: actions/checkout@v4
 
+      - name: Restore wheel cache
+        id: cache-wheel
+        uses: actions/cache@v4
+        with:
+          path: fused_ops/dist
+          key: wheel-cu${{ matrix.cuda_short }}-torch${{ matrix.torch }}-py${{ matrix.python }}-${{ hashFiles('fused_ops/src/**', 'fused_ops/include/**', 'fused_ops/setup.py', 'fused_ops/pyproject.toml') }}
+
       - name: Install system deps
+        if: steps.cache-wheel.outputs.cache-hit != 'true'
         run: |
           apt-get update -q
+          DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends software-properties-common tzdata ca-certificates
+          if [ "${{ matrix.python }}" = "3.12" ]; then
+            add-apt-repository -y ppa:deadsnakes/ppa
+            apt-get update -q
+          fi
           apt-get install -y --no-install-recommends \
             python${{ matrix.python }} \
             python${{ matrix.python }}-dev \
-            python${{ matrix.python }}-distutils \
-            python3-pip \
             curl \
             git
           update-alternatives --install /usr/bin/python python \
@@ -62,21 +96,32 @@ jobs:
           curl -sS https://bootstrap.pypa.io/get-pip.py | python
 
       - name: Install PyTorch
+        if: steps.cache-wheel.outputs.cache-hit != 'true'
         run: |
-          pip install --no-cache-dir \
+          python -m pip install --no-cache-dir \
             torch==${{ matrix.torch }} \
             --index-url https://download.pytorch.org/whl/cu${{ matrix.cuda_short }}
 
       - name: Install build tools
-        run: pip install --no-cache-dir wheel setuptools ninja
+        if: steps.cache-wheel.outputs.cache-hit != 'true'
+        run: python -m pip install --no-cache-dir wheel setuptools ninja
 
       - name: Build wheel
+        if: steps.cache-wheel.outputs.cache-hit != 'true'
         working-directory: fused_ops
         env:
           CUDA_VERSION: ${{ matrix.cuda_short }}
           TORCH_CUDA_ARCH_LIST: "7.0;7.5;8.0;8.6;8.9;9.0"
           FORCE_CUDA: "1"
-        run: python setup.py bdist_wheel
+        run: python -m pip wheel . --no-deps --no-build-isolation -w dist
+
+      - name: Retag wheel as manylinux
+        if: steps.cache-wheel.outputs.cache-hit != 'true'
+        working-directory: fused_ops/dist
+        run: |
+          for f in *.whl; do
+            mv "$f" "${f/linux_x86_64/manylinux2014_x86_64}"
+          done
 
       - name: Upload wheel as artifact
         uses: actions/upload-artifact@v4
@@ -87,7 +132,7 @@ jobs:
   publish:
     needs: build
     runs-on: ubuntu-latest
-    if: startsWith(github.ref, 'refs/tags/')
+    if: startsWith(github.ref, 'refs/tags/') || github.event_name == 'workflow_dispatch'
 
     steps:
       - name: Download all wheels
@@ -99,12 +144,12 @@ jobs:
       - name: List wheels
         run: find all_wheels/ -name "*.whl" | sort
 
-      - name: Publish to PyPI
+      - name: Publish to PyPI / TestPyPI
+        env:
+          TWINE_USERNAME: __token__
+          TWINE_PASSWORD: ${{ github.event.inputs.use_testpypi == 'true' && secrets.TEST_PYPI_API_TOKEN || secrets.PYPI_API_TOKEN }}
+          TWINE_REPOSITORY_URL: ${{ github.event.inputs.use_testpypi == 'true' && 'https://test.pypi.org/legacy/' || 'https://upload.pypi.org/legacy/' }}
         run: |
+          set -euo pipefail
           pip install twine
-          # twine will route each wheel to the correct PyPI project based on
-          # the package name baked into the wheel by setup.py (fireants-fused-ops-cu118 etc.)
-          # All projects must be pre-created on pypi.org under your account.
-          twine upload all_wheels/*.whl \
-            --username __token__ \
-            --password ${{ secrets.PYPI_API_TOKEN }}
+          twine upload --verbose --skip-existing all_wheels/*.whl

fused_ops/pyproject.toml

@@ -1,14 +1,3 @@
 [build-system]
 requires = ["setuptools", "wheel", "torch>=2.3.0"]
 build-backend = "setuptools.build_meta"
-
-[project]
-name = "fireants_fused_ops"
-version = "1.2.0"
-description = "Fused CUDA operations for FireANTs"
-authors = [{name = "Rohit Jena"}]
-requires-python = ">=3.8"
-dependencies = ["torch>=2.3.0"]
-
-[project.urls]
-Homepage = "https://github.com/rohitrango/FireANTs"

fused_ops/setup.py

@@ -20,11 +20,16 @@
 
 include_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'include')
 
+cuda_short = os.environ.get('CUDA_VERSION', '')
+package_name = f'fireants-fused-ops-cu{cuda_short}' if cuda_short else 'fireants-fused-ops'
+
 setup(
-    name='fireants_fused_ops',
-    version='1.0.0',
+    name=package_name,
+    version='1.2.0',
     description='Fused CUDA operations for FireANTs',
     author='Rohit Jena',
+    url='https://github.com/rohitrango/FireANTs',
+    python_requires='>=3.8',
     ext_modules=[
         cpp_extension.CUDAExtension(
             name='fireants_fused_ops',
@@ -51,5 +56,5 @@
         )
     ],
     cmdclass={'build_ext': cpp_extension.BuildExtension},
-    install_requires=['torch>=2.3.0'],
+    install_requires=['torch>=2.1.0'],
 )

scripts/test_fusedops_import_from_pypi.sh

@@ -0,0 +1,213 @@
+#!/usr/bin/env bash
+# test_fusedops.sh — verify fused-ops wheels from TestPyPI actually work
+# Installs fireants-fused-ops-cu<X> from test.pypi.org into isolated conda
+# envs, then runs the test_fusedops_*.py test suite against each variant.
+# Assumes: CUDA-capable machine with compatible driver, conda on PATH.
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+PROJECT_ROOT="$(cd "$SCRIPT_DIR/.." && pwd)"
+cd "$PROJECT_ROOT"
+
+# ---------- full CI matrix (cuda_short:torch:python) ----------
+ALL_ENTRIES=(
+  "118:2.1.0:3.10"
+  "118:2.1.0:3.11"
+  "118:2.2.0:3.12"
+  "121:2.3.0:3.10"
+  "121:2.3.0:3.11"
+  "121:2.3.0:3.12"
+  "124:2.5.0:3.10"
+  "124:2.5.0:3.11"
+  "124:2.5.0:3.12"
+)
+# ---------------------------------------------------------------
+
+PYPI_INDEX="https://test.pypi.org/simple/"
+ENTRIES=()
+VERSION=""   # optional: pin a specific package version
+IMPORT_ONLY=false  # if true, only verify install + import (skip pytest)
+
+usage() {
+  cat <<EOF
+Usage: $0 [OPTIONS]
+
+Test fireants-fused-ops wheels from TestPyPI in isolated conda environments.
+
+Options:
+  --entry CU:TORCH:PY  Run a single matrix entry, e.g. "124:2.5.0:3.10"
+  --all                 Run all 9 matrix entries (default if no --entry given)
+  --version VER         Pin fused-ops package version (e.g. "1.2.0")
+  --pypi                Install from real PyPI instead of TestPyPI
+  --import-only         Only verify install + import (skip running tests)
+  -h, --help            Show this help
+
+Environment variable overrides (single-entry shortcut):
+  CUDA_SHORT, TORCH, PYTHON   e.g.  CUDA_SHORT=124 TORCH=2.5.0 PYTHON=3.10 $0
+
+Examples:
+  $0                                  # full matrix from TestPyPI
+  $0 --entry 124:2.5.0:3.10           # single entry
+  $0 --version 1.2.0                  # pin version
+  $0 --pypi                           # test real PyPI release
+  $0 --import-only                    # quick: just verify install + import
+  CUDA_SHORT=121 TORCH=2.3.0 PYTHON=3.11 $0
+EOF
+  exit 0
+}
+
+# Parse CLI args
+while [[ $# -gt 0 ]]; do
+  case "$1" in
+    --entry)     ENTRIES+=("$2"); shift 2 ;;
+    --all)       shift ;;
+    --version)   VERSION="$2"; shift 2 ;;
+    --pypi)      PYPI_INDEX="https://pypi.org/simple/"; shift ;;
+    --import-only) IMPORT_ONLY=true; shift ;;
+    -h|--help)   usage ;;
+    *)           echo "Unknown option: $1"; usage ;;
+  esac
+done
+
+# If env vars are set and no --entry was given, build a single entry from them
+if [[ ${#ENTRIES[@]} -eq 0 && -n "${CUDA_SHORT:-}" && -n "${TORCH:-}" && -n "${PYTHON:-}" ]]; then
+  ENTRIES+=("${CUDA_SHORT}:${TORCH}:${PYTHON}")
+fi
+
+# Default: full matrix
+if [[ ${#ENTRIES[@]} -eq 0 ]]; then
+  ENTRIES=("${ALL_ENTRIES[@]}")
+fi
+
+# ---------- conda init helper ----------
+setup_conda() {
+  if ! command -v conda &>/dev/null; then
+    echo "ERROR: conda not found on PATH" >&2
+    exit 1
+  fi
+  local conda_base
+  conda_base="$(conda info --base)"
+  # shellcheck source=/dev/null
+  source "${conda_base}/etc/profile.d/conda.sh"
+
+  # Ensure pip/requests use the system CA bundle (handles corporate proxies
+  # that re-sign TLS with a custom root CA not in Python's certifi bundle).
+  if [[ -z "${SSL_CERT_FILE:-}" ]]; then
+    for ca in /etc/ssl/certs/ca-certificates.crt \
+              /etc/pki/tls/certs/ca-bundle.crt \
+              /etc/ssl/ca-bundle.pem; do
+      if [[ -f "$ca" ]]; then
+        export SSL_CERT_FILE="$ca"
+        export REQUESTS_CA_BUNDLE="$ca"
+        echo "Using system CA bundle: $ca"
+        break
+      fi
+    done
+  fi
+}
+
+# ---------- per-entry install & test (runs in a subshell) ----------
+run_entry() {
+  set -euo pipefail
+  local cuda_short="$1" torch_ver="$2" py_ver="$3"
+  local env_name="fusedops_test_cu${cuda_short}_torch${torch_ver}_py${py_ver}"
+  local pkg_name="fireants-fused-ops-cu${cuda_short}"
+  local pkg_spec="$pkg_name"
+  if [[ -n "$VERSION" ]]; then
+    pkg_spec="${pkg_name}==${VERSION}"
+  fi
+
+  echo ""
+  echo "============================================================"
+  echo "  CUDA_SHORT=$cuda_short  TORCH=$torch_ver  PYTHON=$py_ver"
+  echo "  package: $pkg_spec  (from $PYPI_INDEX)"
+  echo "  conda env: $env_name"
+  echo "============================================================"
+
+  echo ">>> Creating conda env (python=${py_ver})"
+  conda create -n "$env_name" python="$py_ver" -y -q
+
+  echo ">>> Activating env"
+  conda activate "$env_name"
+
+  echo ">>> Installing PyTorch ${torch_ver} (cu${cuda_short})"
+  python -m pip install --no-cache-dir -q \
+    torch=="${torch_ver}" \
+    --index-url "https://download.pytorch.org/whl/cu${cuda_short}"
+
+  echo ">>> Installing ${pkg_spec} from ${PYPI_INDEX}"
+  # Use --no-deps to avoid the fused-ops 'torch>=2.1.0' requirement pulling
+  # a different torch build (e.g. CPU-only from PyPI) over our cu-specific one.
+  python -m pip install --no-cache-dir --no-deps \
+    --index-url "$PYPI_INDEX" \
+    --extra-index-url "https://pypi.org/simple/" \
+    "$pkg_spec"
+
+  # Install fireants as editable WITHOUT pulling deps — torch is already
+  # installed at the exact version we need and SimpleITK==2.2.1 may not
+  # exist for this Python.  Only the fireants source tree is needed for
+  # the test imports.
+  echo ">>> Installing fireants (local, no-deps)"
+  python -m pip install --no-cache-dir -q --no-deps -e .
+
+  # Install minimal test deps that aren't already present
+  echo ">>> Installing test dependencies"
+  python -m pip install --no-cache-dir -q pytest numpy scipy scikit-image \
+    nibabel matplotlib tqdm pandas hydra-core SimpleITK 2>/dev/null \
+    || python -m pip install --no-cache-dir -q pytest numpy scipy scikit-image \
+         nibabel matplotlib tqdm pandas hydra-core
+
+  echo ">>> Verifying import"
+  python -c "import torch; print('torch', torch.__version__); import fireants_fused_ops; print('fireants_fused_ops loaded OK')"
+
+  if [[ "$IMPORT_ONLY" == true ]]; then
+    echo ">>> --import-only: skipping tests"
+  else
+    echo ">>> Running fused-ops tests"
+    python -m pytest -v tests/test_fusedops*.py
+  fi
+
+  echo ">>> Entry cu${cuda_short}/torch${torch_ver}/py${py_ver} — PASS"
+}
+
+# ---------- main ----------
+setup_conda
+
+declare -A RESULTS
+overall_rc=0
+
+for entry in "${ENTRIES[@]}"; do
+  IFS=: read -r cuda_short torch_ver py_ver <<< "$entry"
+  env_name="fusedops_test_cu${cuda_short}_torch${torch_ver}_py${py_ver}"
+
+  # Temporarily disable errexit so the outer script doesn't abort on failure,
+  # but the *subshell* still has set -e active (re-enabled inside run_entry).
+  set +e
+  ( run_entry "$cuda_short" "$torch_ver" "$py_ver" )
+  rc=$?
+  set -e
+
+  if [[ $rc -eq 0 ]]; then
+    RESULTS["$entry"]="PASS"
+  else
+    RESULTS["$entry"]="FAIL"
+    overall_rc=1
+  fi
+
+  # Always clean up the conda env
+  echo ">>> Cleaning up conda env: $env_name"
+  conda env remove -n "$env_name" -y 2>/dev/null || true
+done
+
+# ---------- summary ----------
+echo ""
+echo "======================== SUMMARY ========================"
+printf "%-8s  %-8s  %-6s  %s\n" "CUDA" "TORCH" "PYTHON" "RESULT"
+printf "%-8s  %-8s  %-6s  %s\n" "------" "------" "------" "------"
+for entry in "${ENTRIES[@]}"; do
+  IFS=: read -r cuda_short torch_ver py_ver <<< "$entry"
+  printf "%-8s  %-8s  %-6s  %s\n" "cu$cuda_short" "$torch_ver" "$py_ver" "${RESULTS[$entry]}"
+done
+echo "========================================================="
+
+exit "$overall_rc"