This is an automated email from the ASF dual-hosted git repository.
yaxingcai pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm-ffi.git
The following commit(s) were added to refs/heads/main by this push:
new 75c2a2b [ADDON] Improved github action for torch-c-dlpack-ext (#225)
75c2a2b is described below
commit 75c2a2bb7b8d367f33ade1a1b4f9f14212fc080f
Author: Yaxing Cai <[email protected]>
AuthorDate: Tue Nov 4 20:45:03 2025 -0800
[ADDON] Improved github action for torch-c-dlpack-ext (#225)
This PR improved the github action for torch-c-dlpack-ext
---
.github/workflows/torch_c_dlpack.yml | 122 +++-----------------------
addons/torch_c_dlpack_ext/build_aot_wheels.sh | 109 +++++++++++++++++++++++
addons/torch_c_dlpack_ext/pyproject.toml | 2 +-
3 files changed, 120 insertions(+), 113 deletions(-)
diff --git a/.github/workflows/torch_c_dlpack.yml
b/.github/workflows/torch_c_dlpack.yml
index 8c5a605..93fec23 100644
--- a/.github/workflows/torch_c_dlpack.yml
+++ b/.github/workflows/torch_c_dlpack.yml
@@ -25,12 +25,12 @@ on:
required: true
jobs:
- build_libs:
+ build_wheels:
strategy:
fail-fast: false
matrix:
arch: ["x86_64", "aarch64"]
- python-version: ["cp39", "cp310", "cp311", "cp312", "cp313"]
+ python-version: ["cp39", "cp310", "cp311", "cp312", "cp313", "cp314"]
runs-on: ${{ matrix.arch == 'aarch64' && 'ubuntu-24.04-arm' ||
'ubuntu-latest' }}
steps:
- uses: jlumbroso/free-disk-space@main
@@ -54,120 +54,18 @@ jobs:
&& dnf clean all \
&& dnf -y install cuda-toolkit-13-0 \
&& rm /cuda-repo-rhel8-13-0-local-13.0.2_580.95.05-1.${{ matrix.arch
}}.rpm \
- && dnf clean all \
- && mkdir /libs"
- - name: build torch 2.4 lib
- if: ${{ matrix.python-version != 'cp313' && matrix.python-version !=
'cp314' && matrix.arch != 'aarch64' }}
- run: |
- docker exec build_lib bash -c "uv venv /torch2.4 --python ${{
matrix.python-version }} \
- && source /torch2.4/bin/activate \
- && uv pip install setuptools ninja \
- && uv pip install torch==2.4 --index-url
https://download.pytorch.org/whl/cu124 \
- && cd /tvm-ffi \
- && uv pip install -v . \
- && python -m tvm_ffi.utils._build_optional_torch_c_dlpack
--output-dir /libs \
- && python -m tvm_ffi.utils._build_optional_torch_c_dlpack
--output-dir /libs --build-with-cuda \
- && rm -rf /torch2.4"
- - name: build torch 2.5 lib
- if: ${{ matrix.python-version != 'cp314' && matrix.arch != 'aarch64' }}
- run: |
- docker exec build_lib bash -c "uv venv /torch2.5 --python ${{
matrix.python-version }} \
- && source /torch2.5/bin/activate \
- && uv pip install setuptools ninja \
- && uv pip install torch==2.5 --index-url
https://download.pytorch.org/whl/cu124 \
- && cd /tvm-ffi \
- && uv pip install -v . \
- && python -m tvm_ffi.utils._build_optional_torch_c_dlpack
--output-dir /libs \
- && python -m tvm_ffi.utils._build_optional_torch_c_dlpack
--output-dir /libs --build-with-cuda \
- && rm -rf /torch2.5"
- - name: build torch 2.6 lib
- if: ${{ matrix.python-version != 'cp314' && matrix.arch != 'aarch64' }}
- run: |
- docker exec build_lib bash -c "uv venv /torch2.6 --python ${{
matrix.python-version }} \
- && source /torch2.6/bin/activate \
- && uv pip install setuptools ninja \
- && uv pip install torch==2.6 --index-url
https://download.pytorch.org/whl/cu126 \
- && cd /tvm-ffi \
- && uv pip install -v . \
- && python -m tvm_ffi.utils._build_optional_torch_c_dlpack
--output-dir /libs \
- && python -m tvm_ffi.utils._build_optional_torch_c_dlpack
--output-dir /libs --build-with-cuda \
- && rm -rf /torch2.6"
- - name: build torch 2.7 lib
- if: ${{ matrix.python-version != 'cp314' }}
- run: |
- docker exec build_lib bash -c "uv venv /torch2.7 --python ${{
matrix.python-version }} \
- && source /torch2.7/bin/activate \
- && uv pip install setuptools ninja \
- && uv pip install torch==2.7 --index-url
https://download.pytorch.org/whl/cu128 \
- && cd /tvm-ffi \
- && uv pip install -v . \
- && python -m tvm_ffi.utils._build_optional_torch_c_dlpack
--output-dir /libs \
- && python -m tvm_ffi.utils._build_optional_torch_c_dlpack
--output-dir /libs --build-with-cuda \
- && rm -rf /torch2.7"
- - name: build torch 2.8 lib
- if: ${{ matrix.python-version != 'cp314' }}
- run: |
- docker exec build_lib bash -c "uv venv /torch2.8 --python ${{
matrix.python-version }} \
- && source /torch2.8/bin/activate \
- && uv pip install setuptools ninja \
- && uv pip install torch==2.8 --index-url
https://download.pytorch.org/whl/cu129 \
- && cd /tvm-ffi \
- && uv pip install -v . \
- && python -m tvm_ffi.utils._build_optional_torch_c_dlpack
--output-dir /libs \
- && python -m tvm_ffi.utils._build_optional_torch_c_dlpack
--output-dir /libs --build-with-cuda \
- && rm -rf /torch2.8"
- - name: build torch 2.9 lib
- if: ${{ matrix.python-version != 'cp39' }}
+ && dnf clean all"
+ - name: build torch libs and wheels
run: |
- docker exec build_lib bash -c "uv venv /torch2.9 --python ${{
matrix.python-version }} \
- && source /torch2.9/bin/activate \
- && uv pip install setuptools ninja \
- && uv pip install torch==2.9 --index-url
https://download.pytorch.org/whl/cu129 \
- && cd /tvm-ffi \
- && uv pip install -v . \
- && python -m tvm_ffi.utils._build_optional_torch_c_dlpack
--output-dir /libs \
- && python -m tvm_ffi.utils._build_optional_torch_c_dlpack
--output-dir /libs --build-with-cuda \
- && rm -rf /torch2.9"
- - name: collect built libs
+ docker exec -w /tvm-ffi build_lib bash
./addons/torch_c_dlpack_ext/build_aot_wheels.sh ${{ matrix.arch }} ${{
matrix.python-version }}
+ - name: collect built wheels
run: |
- mkdir libs
- docker cp build_lib:/libs .
- - uses: actions/upload-artifact@v4
- with:
- name: libs-${{ matrix.arch }}-${{ matrix.python-version }}
- path: ./libs/*.so
- build_wheels:
- needs: build_libs
- strategy:
- matrix:
- arch: ["x86_64", "aarch64"]
- python-version: ["cp39", "cp310", "cp311", "cp312", "cp313"]
- runs-on: ${{ matrix.arch == 'aarch64' && 'ubuntu-24.04-arm' ||
'ubuntu-latest' }}
- container: quay.io/pypa/manylinux_2_28_${{ matrix.arch }}
- steps:
- - uses: actions/checkout@v5
- with:
- repository: apache/tvm-ffi
- ref: main
- fetch-depth: 0
- submodules: recursive
- path: tvm-ffi
- - uses: actions/download-artifact@v4
- with:
- pattern: libs-${{ matrix.arch }}-${{ matrix.python-version }}
- path: ./tvm-ffi/addons/torch_c_dlpack_ext/torch_c_dlpack_ext
- merge-multiple: true
- - name: build wheels
- run: |
- /opt/python/${{ matrix.python-version }}-${{ matrix.python-version
}}/bin/pip3 install build wheel auditwheel
- cd ./tvm-ffi/addons/torch_c_dlpack_ext
- /opt/python/${{ matrix.python-version }}-${{ matrix.python-version
}}/bin/python -m build -w
- /opt/python/${{ matrix.python-version }}-${{ matrix.python-version
}}/bin/python -m wheel tags dist/*.whl --python-tag=${{ matrix.python-version
}} --abi-tag=${{ matrix.python-version }} --remove
- auditwheel repair --exclude libtorch.so --exclude libtorch_cpu.so
--exclude libc10.so --exclude libtorch_python.so dist/*.whl -w wheelhouse
+ mkdir wheelhouse
+ docker cp build_lib:/tvm-ffi/addons/torch_c_dlpack_ext/wheelhouse/ .
- uses: actions/upload-artifact@v4
with:
name: pypi-wheels-${{ matrix.arch }}-${{ matrix.python-version }}
- path: ./tvm-ffi/addons/torch_c_dlpack_ext/wheelhouse/*.whl
+ path: ./wheelhouse/*.whl
build_source:
runs-on: ubuntu-latest
steps:
@@ -198,7 +96,7 @@ jobs:
permissions:
id-token: write
attestations: write
- # if: github.event_name == 'workflow_dispatch' # <-- publish only on
manual trigger
+ if: github.event_name == 'workflow_dispatch' # <-- publish only on
manual trigger
steps:
- uses: actions/download-artifact@v4
with:
diff --git a/addons/torch_c_dlpack_ext/build_aot_wheels.sh
b/addons/torch_c_dlpack_ext/build_aot_wheels.sh
new file mode 100755
index 0000000..0361b5f
--- /dev/null
+++ b/addons/torch_c_dlpack_ext/build_aot_wheels.sh
@@ -0,0 +1,109 @@
+#!/bin/bash
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# shellcheck disable=SC1090,1091
+set -eux
+
+arch=$1
+python_version=$2
+
+tvm_ffi="$PWD"
+torch_c_dlpack_ext="$tvm_ffi"/addons/torch_c_dlpack_ext
+
+
+function get_torch_url() {
+ local version="$1"
+ case "$version" in
+ "2.4" | "2.5")
+ echo "https://download.pytorch.org/whl/cu124";
+ ;;
+ "2.6")
+ echo "https://download.pytorch.org/whl/cu126";
+ ;;
+ "2.7")
+ echo "https://download.pytorch.org/whl/cu128";
+ ;;
+ "2.8" | "2.9")
+ echo "https://download.pytorch.org/whl/cu129";
+ ;;
+ *)
+ echo "Unknown or unsupported torch version: $version" >&2
+ return 1
+ ;;
+ esac
+}
+
+
+function check_availability() {
+ local torch_version="$1"
+ case "$torch_version" in
+ "2.4")
+ ! [[ "$arch" == "aarch64" || "$python_version" == "cp313" ||
"$python_version" == "cp314" ]]
+ ;;
+ "2.5" | "2.6")
+ ! [[ "$arch" == "aarch64" || "$python_version" == "cp314" ]]
+ ;;
+ "2.7" | "2.8")
+ ! [[ "$python_version" == "cp314" ]]
+ ;;
+ "2.9")
+ ! [[ "$python_version" == "cp39" ]]
+ ;;
+ *)
+ echo "Unknown or unsupported torch version: $torch_version" >&2
+ return 1
+ ;;
+ esac
+}
+
+
+function build_libs() {
+ local torch_version=$1
+ if check_availability "$torch_version"; then
+ mkdir "$tvm_ffi"/.venv -p
+ uv venv "$tvm_ffi"/.venv/torch"$torch_version" --python
"$python_version"
+ source "$tvm_ffi"/.venv/torch"$torch_version"/bin/activate
+ uv pip install setuptools ninja
+ uv pip install torch=="$torch_version" --index-url "$(get_torch_url
"$torch_version")"
+ uv pip install -v .
+ mkdir "$tvm_ffi"/lib -p
+ python -m tvm_ffi.utils._build_optional_torch_c_dlpack --output-dir
"$tvm_ffi"/lib
+ python -m tvm_ffi.utils._build_optional_torch_c_dlpack --output-dir
"$tvm_ffi"/lib --build-with-cuda
+ ls "$tvm_ffi"/lib
+ deactivate
+ rm -rf "$tvm_ffi"/.venv/torch"$torch_version"
+ else
+ echo "Skipping build for torch $torch_version on $arch with python
$python_version as it is not available."
+ fi
+}
+
+torch_versions=("2.4" "2.5" "2.6" "2.7" "2.8" "2.9")
+for version in "${torch_versions[@]}"; do
+ build_libs "$version"
+done
+
+cp "$tvm_ffi"/lib/*.so "$torch_c_dlpack_ext"/torch_c_dlpack_ext
+uv venv "$tvm_ffi"/.venv/build --python "$python_version"
+source "$tvm_ffi"/.venv/build/bin/activate
+uv pip install build wheel auditwheel
+cd "$torch_c_dlpack_ext"
+python -m build -w
+ls dist
+python -m wheel tags dist/*.whl --python-tag="$python_version"
--abi-tag="$python_version" --remove
+ls dist
+auditwheel repair --exclude libtorch.so --exclude libtorch_cpu.so --exclude
libc10.so --exclude libtorch_python.so dist/*.whl -w wheelhouse
+ls wheelhouse
diff --git a/addons/torch_c_dlpack_ext/pyproject.toml
b/addons/torch_c_dlpack_ext/pyproject.toml
index 3cf0138..343a26c 100644
--- a/addons/torch_c_dlpack_ext/pyproject.toml
+++ b/addons/torch_c_dlpack_ext/pyproject.toml
@@ -23,7 +23,7 @@ description = "torch c dlpack ext"
dependencies = ["torch"]
[build-system]
-requires = ["setuptools>=61.0", "torch"]
+requires = ["setuptools>=61.0"]
build-backend = "build_backend"
backend-path = ["."]
