This is an automated email from the ASF dual-hosted git repository.
junrushao 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 37f45c3 Reliability update (#435)
37f45c3 is described below
commit 37f45c3ed834d57de3a3513457ab239c13881272
Author: Yichen Yan <[email protected]>
AuthorDate: Sun Feb 8 09:43:50 2026 +0800
Reliability update (#435)
Closes #433 in a way that do not bailout when compiling in environment
without torch.
Also fixes that dlpack ext's symbols does not match with torch (e.g. in
ngc container)
---
addons/torch_c_dlpack_ext/build_backend.py | 7 +++++--
python/tvm_ffi/_optional_torch_c_dlpack.py | 5 +++++
2 files changed, 10 insertions(+), 2 deletions(-)
diff --git a/addons/torch_c_dlpack_ext/build_backend.py
b/addons/torch_c_dlpack_ext/build_backend.py
index 639977e..3371ec4 100644
--- a/addons/torch_c_dlpack_ext/build_backend.py
+++ b/addons/torch_c_dlpack_ext/build_backend.py
@@ -21,6 +21,7 @@ from __future__ import annotations
import os
import subprocess
import sys
+from contextlib import suppress
from pathlib import Path
from setuptools import build_meta as orig
@@ -62,10 +63,12 @@ def build_wheel(
metadata_directory: orig.StrPath | None = None,
) -> str:
"""Build wheel."""
- if not _is_lib_prebuilt():
- # build wheel from sdist package, compile the torch c dlpack ext
library locally.
+ torch = None
+ with suppress(ModuleNotFoundError):
import torch # noqa: PLC0415
+ if torch is not None and not _is_lib_prebuilt():
+ # build wheel from sdist package, compile the torch c dlpack ext
library locally.
if hasattr(torch.Tensor, "__dlpack_c_exchange_api__") or hasattr(
torch.Tensor, "__c_dlpack_exchange_api__"
):
diff --git a/python/tvm_ffi/_optional_torch_c_dlpack.py
b/python/tvm_ffi/_optional_torch_c_dlpack.py
index 415e151..ca73f16 100644
--- a/python/tvm_ffi/_optional_torch_c_dlpack.py
+++ b/python/tvm_ffi/_optional_torch_c_dlpack.py
@@ -109,6 +109,11 @@ def load_torch_c_dlpack_extension() -> Any: # noqa:
PLR0912, PLR0915
return None
except ImportError:
pass
+ except AttributeError:
+ # When torch_c_dlpack_ext and torch have different ABI
+ # `ctypes.CDLL` will raise an `AttributeError`.
+ # Keep trying JIT
+ pass
try:
# check whether a JIT shared library is built in cache
