pyg-team · akihironitta · Jul 27, 2025 · May 25, 2025 · May 25, 2025 · May 25, 2025
diff --git a/.github/workflows/_build_linux.yml b/.github/workflows/_build_linux.yml
@@ -6,31 +6,71 @@ on:  # yamllint disable-line rule:truthy
       test-matrix:
         type: string
         required: true
+      docker-hub-username:
+        type: string
+        required: true
+    secrets:
+      docker-hub-token:
+        required: true
+      aws-access-key-id:
+        required: true
+      aws-secret-access-key:
+        required: true
 
 defaults:
   run:
     shell: bash
 
 jobs:
   wheel:
-    runs-on: ubuntu-22.04
+    runs-on: ubuntu-latest
     strategy:
       fail-fast: false
       matrix:
         include: ${{ fromJSON(inputs.test-matrix) }}
 
     steps:
+      - name: Login to Docker Hub
+        uses: docker/login-action@v3
+        with:
+          username: ${{ inputs.docker-hub-username }}
+          password: ${{ secrets.docker-hub-token }}
+
+      - name: Configure AWS
+        uses: aws-actions/configure-aws-credentials@v4
+        with:
+          aws-access-key-id: ${{ secrets.aws-access-key-id }}
+          aws-secret-access-key: ${{ secrets.aws-secret-access-key }}
+          aws-region: us-west-1
+
       - name: Checkout repository
         uses: actions/checkout@v4
         with:
           submodules: recursive
 
-      - name: Install dependencies
-        uses: ./.github/actions/setup
+      - name: Free up disk space
+        # Disk space required for CPU builds is usually small enough, so we
+        # save a few minutes by skipping this step.
+        if: ${{ matrix.cuda-version != 'cpu' }}
+        uses: jlumbroso/[email protected]
         with:
-          python-version: ${{ matrix.python-version }}
-          torch-version: ${{ matrix.torch-version }}
-          cuda-version: ${{ matrix.cuda-version }}
+          android: true
+          dotnet: true
+          haskell: true
+          large-packages: true
+          docker-images: true
+          swap-storage: true
+          tool-cache: true
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          # NOTE: This is a version used for cibuildwheel but not for wheels.
+          python-version: 3.11
+
+      - name: Install dependencies
+        run: |
+          pip install cibuildwheel auditwheel
 
       - name: Set version
         run: |
@@ -43,23 +83,4 @@ jobs:
 
       - name: Build wheel
         run: |
-          source ./.github/workflows/cuda/Linux-env.sh ${{ matrix.cuda-version }}
-          echo $TORCH_CUDA_ARCH_LIST
-          pip install build
-          python -m build --wheel --no-isolation --outdir dist
-
-      - name: Install wheel
-        run: |
-          cd dist
-          ls -lah
-          pip install *.whl
-          python -c "import pyg_lib; print('pyg-lib:', pyg_lib.__version__)"
-          python -c "import pyg_lib; print('CUDA:', pyg_lib.cuda_version())"
-          python -c "import pyg_lib.ops"
-          python -c "import pyg_lib.sampler"
-          cd ..
-
-      - name: Test wheel
-        run: |
-          pip install pytest pytest-cov
-          pytest --cov --cov-report=xml
+          bash ./.github/workflows/utils/build_linux.sh ${{ matrix.cuda-version }} ${{ matrix.python-version }} ${{ matrix.torch-version }}
diff --git a/.github/workflows/pull.yml b/.github/workflows/pull.yml
@@ -93,6 +93,11 @@ jobs:
     uses: ./.github/workflows/_build_linux.yml
     with:
       test-matrix: ${{ needs.trigger.outputs.matrix-linux }}
+      docker-hub-username: ${{ vars.DOCKERHUB_USERNAME }}
+    secrets:
+      docker-hub-token: ${{ secrets.DOCKERHUB_TOKEN }}
+      aws-access-key-id: ${{ secrets.AWS_ACCESS_KEY_ID }}
+      aws-secret-access-key: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
 
   macos:
     needs: [trigger]

diff --git a/.github/workflows/utils/build_linux.sh b/.github/workflows/utils/build_linux.sh
@@ -0,0 +1,34 @@
+#!/bin/bash
+set -ex
+
+CUDA_VERSION="${1:?Specify cuda version, e.g. cpu, cu128}"
+PYTHON_VERSION="${2:?Specify python version, e.g. 3.12}"
+TORCH_VERSION="${3:?Specify torch version, e.g. 2.7.0}"
+echo "CUDA_VERSION: ${CUDA_VERSION}"
+echo "PYTHON_VERSION: ${PYTHON_VERSION//./}"
+echo "TORCH_VERSION: ${TORCH_VERSION}"
+
+source ./.github/workflows/cuda/Linux-env.sh ${CUDA_VERSION}
+echo "FORCE_CUDA: ${FORCE_CUDA}"
+echo "TORCH_CUDA_ARCH_LIST: ${TORCH_CUDA_ARCH_LIST}"
+
+export CIBW_BUILD="cp${PYTHON_VERSION//./}-manylinux_x86_64"
+export CIBW_BEFORE_BUILD="pip install ninja wheel setuptools && pip install torch==${TORCH_VERSION} --index-url https://download.pytorch.org/whl/${CUDA_VERSION}"
+# pyg-lib doesn't have torch as a dependency, so we need to explicitly install it when running tests.
+export CIBW_BEFORE_TEST="pip install pytest && pip install torch==${TORCH_VERSION} --index-url https://download.pytorch.org/whl/${CUDA_VERSION}"
+
+if [[ "${CUDA_VERSION}" == "cu"* ]]; then
+  # Use CUDA-pre-installed image
+  export CIBW_MANYLINUX_X86_64_IMAGE=akihironitta/manylinux:${CUDA_VERSION}
+else
+  export CIBW_MANYLINUX_X86_64_IMAGE=quay.io/pypa/manylinux_2_28_x86_64
+fi
+
+rm -rf Testing libpyg.so build dist outputs  # for local testing
+python -m cibuildwheel --output-dir dist
+ls -ahl dist/
+python -m auditwheel show dist/*.whl
+
+unzip dist/*.whl -d debug/
+ldd debug/libpyg.so
+readelf -d debug/libpyg.so
diff --git a/.github/workflows/utils/full_matrix_linux.json b/.github/workflows/utils/full_matrix_linux.json
@@ -17,31 +17,31 @@
   {
     "torch-version": "2.4.0",
     "python-version": ["3.9", "3.10", "3.11", "3.12"],
-    "cuda-version": ["cpu", "cu118", "cu121", "cu124"]
+    "cuda-version": ["cpu"]
   },
   {
     "torch-version": "2.3.0",
     "python-version": ["3.9", "3.10", "3.11", "3.12"],
-    "cuda-version": ["cpu", "cu118", "cu121"]
+    "cuda-version": ["cpu"]
   },
   {
     "torch-version": "2.2.0",
     "python-version": ["3.9", "3.10", "3.11"],
-    "cuda-version": ["cpu", "cu118", "cu121"]
+    "cuda-version": ["cpu"]
   },
   {
     "torch-version": "2.1.0",
     "python-version": ["3.9", "3.10", "3.11"],
-    "cuda-version": ["cpu", "cu118", "cu121"]
+    "cuda-version": ["cpu"]
   },
   {
     "torch-version": "2.0.0",
     "python-version": ["3.9", "3.10", "3.11"],
-    "cuda-version": ["cpu", "cu117", "cu118"]
+    "cuda-version": ["cpu"]
   },
   {
     "torch-version": "1.13.0",
     "python-version": ["3.9", "3.10", "3.11"],
-    "cuda-version": ["cpu", "cu116", "cu117"]
+    "cuda-version": ["cpu"]
   }
 ]
diff --git a/pyproject.toml b/pyproject.toml
@@ -36,6 +36,26 @@ documentation="https://pyg-lib.readthedocs.io"
 repository="https://github.com/pyg-team/pyg-lib.git"
 changelog="https://github.com/pyg-team/pyg-lib/blob/master/CHANGELOG.md"
 
+[tool.cibuildwheel]
+build-frontend = {name = "build", args = ["--no-isolation"] }
+test-command = "pytest {project}/test"
+
+[tool.cibuildwheel.linux]
+environment-pass = [
+    "FORCE_NINJA",
+    "FORCE_CUDA",
+    "CUDA_HOME",
+    "TORCH_CUDA_ARCH_LIST",
+]
+repair-wheel-command = """
+auditwheel repair \
+  --wheel-dir {dest_dir} {wheel} \
+  --exclude 'libtorch*.so' \
+  --exclude 'libc10*.so' \
+  --exclude 'libgomp*.so*' \
+  --exclude 'libcuda.so*'
+"""
+
 [tool.pytest.ini_options]
 addopts=[
     "--capture=no",