This is an automated email from the ASF dual-hosted git repository.
tqchen 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 32ab66ee [FIX] Make ffi.Error picklable (#709)
32ab66ee is described below
commit 32ab66ee756997208e7666841eb77b6de5e8420d
Author: Yaxing Cai <[email protected]>
AuthorDate: Tue Aug 25 02:28:50 2026 +0800
[FIX] Make ffi.Error picklable (#709)
## Motivation
`Error.py_error()` attaches the FFI error object to the Python exception
it raises:
```python
py_error.__tvm_ffi_error__ = self
```
`BaseException.__reduce__` includes the instance `__dict__`, so pickling
an FFI-originated exception pickles that `Error` through
`CObject.__getstate__` → `ToJSONGraph`. `ffi.Error` has no native
creator and no `__ffi_new__` type attr, so this fails:
```
TypeError: Type `ffi.Error` does not support ToJSONGraph
(no native creator or __ffi_new__ type attr)
```
Harnesses that ship exceptions across a process boundary lose the
original error and report this instead. It was reported downstream
against a CUDA IMA (error 700), where the replacement message made it
look like tvm-ffi had suppressed a coredump — it had not; the coredump
happens before the error is raised, and the only real defect is that the
exception could not be pickled.
`copy.deepcopy` hits the same path.
## Changes
`ffi.Error` now pickles by value. An error is fully described by its
`(kind, message, backtrace)` strings, so `Error.__reduce__` reconstructs
it through `__init__` rather than the JSON graph. This mirrors `Device`
and `DataType`, which already carry hand-written `__reduce__`
implementations for the same reason.
`extra_context` is intentionally dropped — it may hold arbitrary native
payloads, and preserving it only when it happens to be serializable
would make pickling an error depend on where the error came from.
The attribute stays a live `Error` rather than degrading to `None`,
because `set_last_ffi_error` calls `update_backtrace` on
`__tvm_ffi_error__` when a restored exception propagates back through
the FFI; a `None` there would fail on an unrelated `hasattr` check that
is already in the code.
## Not in scope
An earlier revision of this PR also rewrote the message raised when
pickling types that genuinely cannot round-trip (`ffi.Function`,
`ffi.Module`, `ffi.Tensor`, `ffi.OpaquePyObject`). Per review, that is
dropped here and can be revisited separately — this PR is now only the
`Error` fix. Those types are unchanged and still surface the
serializer's own message.
## Testing
New tests in `tests/python/test_error.py`: exception and bare-`Error`
pickle round trips, `deepcopy`, NULL handle, the `extra_context` drop,
and re-propagating a restored exception through an FFI call (which
exercises `set_last_ffi_error`'s `update_backtrace` on the restored
object).
Full Python suite: 2421 passed, 52 skipped, 2 xfailed. `pre-commit`
clean on all changed files.
---
python/tvm_ffi/core.pyi | 1 +
python/tvm_ffi/cython/error.pxi | 19 +++++++++
tests/python/test_error.py | 91 +++++++++++++++++++++++++++++++++++++++++
3 files changed, 111 insertions(+)
diff --git a/python/tvm_ffi/core.pyi b/python/tvm_ffi/core.pyi
index 668be688..1ebed00d 100644
--- a/python/tvm_ffi/core.pyi
+++ b/python/tvm_ffi/core.pyi
@@ -88,6 +88,7 @@ def _type_index_to_type_info(type_index: int) -> TypeInfo |
None: ...
class Error(Object):
def __init__(self, kind: str, message: str, backtrace: str) -> None: ...
+ def __reduce__(self) -> Any: ...
def update_backtrace(self, backtrace: str) -> None: ...
def py_error(self) -> BaseException: ...
@property
diff --git a/python/tvm_ffi/cython/error.pxi b/python/tvm_ffi/cython/error.pxi
index 16390ed3..224cf4d2 100644
--- a/python/tvm_ffi/cython/error.pxi
+++ b/python/tvm_ffi/cython/error.pxi
@@ -69,6 +69,25 @@ cdef class Error(CObject):
raise MemoryError("Failed to create error object")
(<CObject>self).chandle = out
+ def __reduce__(self):
+ """Pickle by value from ``(kind, message, backtrace)``.
+
+ ``ffi.Error`` has no JSON-graph creator, so the inherited
+ :py:meth:`CObject.__reduce__` cannot serialize it. Because an error is
+ fully described by its three strings, reconstruct it through
+ :py:meth:`__init__` instead.
+
+ Notes
+ -----
+ :py:attr:`extra_context` is intentionally dropped: it may hold
arbitrary
+ native payloads that have no value representation. Preserving it only
+ when it happens to be serializable would make pickling an error depend
+ on where the error came from.
+ """
+ if self.chandle == NULL:
+ return (_new_object, (Error,))
+ return (Error, (self.kind, self.message, self.backtrace))
+
def update_backtrace(self, backtrace: str) -> None:
"""Replace the stored backtrace string with ``backtrace``.
diff --git a/tests/python/test_error.py b/tests/python/test_error.py
index 70d1f0a2..5e48973e 100644
--- a/tests/python/test_error.py
+++ b/tests/python/test_error.py
@@ -16,7 +16,9 @@
# under the License.
+import copy
import gc
+import pickle
import weakref
from typing import NoReturn
@@ -151,3 +153,92 @@ def test_error_no_cyclic_reference() -> None:
finally:
# re-enable gc whenever exception occurs
gc.enable()
+
+
+def _raise_from_cxx() -> BaseException:
+ """Return a Python exception produced by a C++-side FFI error."""
+ test_raise_error = tvm_ffi.get_global_func("testing.test_raise_error")
+ try:
+ test_raise_error("ValueError", "error XYZ")
+ except ValueError as e:
+ return e
+ raise AssertionError("expected the FFI call to raise")
+
+
+def test_exception_pickle_roundtrip() -> None:
+ """Exceptions carrying ``__tvm_ffi_error__`` must survive pickling.
+
+ Regression test: ``ffi.Error`` has no JSON-graph creator, so the inherited
+ ``CObject.__reduce__`` used to fail with a ``ToJSONGraph`` ``TypeError``
+ whenever a test harness pickled an FFI-originated exception.
+ """
+ err = _raise_from_cxx()
+ restored = pickle.loads(pickle.dumps(err))
+
+ assert isinstance(restored, ValueError)
+ assert restored.args == err.args
+ ffi_error = restored.__tvm_ffi_error__ # ty: ignore[unresolved-attribute]
+ assert isinstance(ffi_error, tvm_ffi.core.Error)
+ assert ffi_error.kind == "ValueError"
+ assert ffi_error.message == "error XYZ"
+ assert ffi_error.backtrace.find("TestRaiseError") != -1
+
+
+def test_exception_deepcopy_roundtrip() -> None:
+ """``copy.deepcopy`` goes through ``__reduce_ex__`` and must work too."""
+ err = _raise_from_cxx()
+ restored = copy.deepcopy(err)
+
+ assert isinstance(restored, ValueError)
+ assert restored.__tvm_ffi_error__.kind == "ValueError" # ty:
ignore[unresolved-attribute]
+ assert restored.__tvm_ffi_error__.message == "error XYZ" # ty:
ignore[unresolved-attribute]
+
+
+def test_ffi_error_pickle_roundtrip() -> None:
+ """A bare :class:`tvm_ffi.core.Error` round-trips by value."""
+ error = tvm_ffi.core.Error("TypeError", "boom", 'File "a.py", line 1, in
f\n')
+ restored = pickle.loads(pickle.dumps(error))
+
+ assert isinstance(restored, tvm_ffi.core.Error)
+ assert restored.kind == "TypeError"
+ assert restored.message == "boom"
+ assert restored.backtrace == 'File "a.py", line 1, in f\n'
+ # a fresh object, not the original handle
+ assert not restored.same_as(error)
+
+
+def test_ffi_error_pickle_null_handle() -> None:
+ """An ``Error`` with a NULL handle round-trips without dereferencing it."""
+ error = tvm_ffi.core.Error.__new__(tvm_ffi.core.Error)
+ assert error.__chandle__() == 0
+
+ restored = pickle.loads(pickle.dumps(error))
+ assert isinstance(restored, tvm_ffi.core.Error)
+ assert restored.__chandle__() == 0
+
+
+def test_ffi_error_pickle_drops_extra_context() -> None:
+ """``extra_context`` is intentionally not preserved across pickling.
+
+ It may hold arbitrary native payloads with no value representation, so
+ dropping it keeps pickling an error independent of where it came from.
+ """
+ error = tvm_ffi.core.Error("ValueError", "boom", "")
+ restored = pickle.loads(pickle.dumps(error))
+ assert restored.extra_context is None
+
+
+def test_restored_exception_can_propagate_through_ffi() -> None:
+ """An unpickled exception still works as an FFI error payload.
+
+ ``set_last_ffi_error`` calls ``update_backtrace`` on ``__tvm_ffi_error__``,
+ so the attribute must survive as a live ``Error``, not as ``None``.
+ """
+ restored = pickle.loads(pickle.dumps(_raise_from_cxx()))
+
+ def callback(_: int) -> NoReturn:
+ raise restored
+
+ fapply = tvm_ffi.convert(callback)
+ with pytest.raises(ValueError, match="error XYZ"):
+ fapply(1)
