Skip to content

[PY] Support pickling FFI String and Bytes#643

Merged
tqchen merged 3 commits into
apache:mainfrom
Ubospica:fix-ffi-string-pickle-main
Jun 23, 2026
Merged

[PY] Support pickling FFI String and Bytes#643
tqchen merged 3 commits into
apache:mainfrom
Ubospica:fix-ffi-string-pickle-main

Conversation

@Ubospica

Copy link
Copy Markdown
Contributor

Summary

  • Add pickle reductions for tvm_ffi.core.String and tvm_ffi.core.Bytes so only the Python payload is serialized.
  • Ensure unpickled values drop the hidden _tvm_ffi_cached_object backing object.
  • Add regression coverage for cached FFI String/Bytes returned from testing.echo.

Notes

Test plan

  • uv pip install --reinstall --verbose . && .venv/bin/pytest -vvs tests/python/test_string.py
  • PYTHONPATH=/raid/user_data/yixind/miniforge3/lib/python3.12/site-packages pre-commit run --files python/tvm_ffi/cython/string.pxi tests/python/test_string.py

Pickled FFI string-like values should only persist their Python payload, avoiding serialization of the hidden cached FFI backing object.

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request implements custom pickling support for the String and Bytes classes in python/tvm_ffi/cython/string.pxi by adding __reduce_ex__ methods, along with corresponding unit tests to verify that cached FFI objects are cleared during roundtrip serialization. The reviewer feedback points out that the current __reduce_ex__ implementations return a 2-tuple, which discards the instance's __dict__ and causes any custom or subclass attributes to be silently lost during pickling. It is recommended to return the state dictionary (excluding _tvm_ffi_cached_object) as the third element of the reduction tuple to preserve these attributes.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

Comment thread python/tvm_ffi/cython/string.pxi Outdated
Comment thread python/tvm_ffi/cython/string.pxi Outdated
Keep ordinary instance state when pickling FFI String and Bytes while still dropping the cached FFI backing object.
@Ubospica

Copy link
Copy Markdown
Contributor Author

Addressed the latest Gemini comments in cb9272a:

  • String.__reduce_ex__ and Bytes.__reduce_ex__ now return state as the third reduction element.
  • The state keeps ordinary instance attributes and drops only _tvm_ffi_cached_object.
  • Added tests for String subclass state and Bytes instance state preservation.

Round-tripping FFI String and Bytes through pickle should produce raw Python str and bytes values instead of preserving FFI wrapper types.
@Ubospica

Copy link
Copy Markdown
Contributor Author

Updated the pickle semantics in a030c1f per requested behavior:

  • pickle.loads(pickle.dumps(tvm_ffi.core.String(...))) now returns raw Python str.
  • pickle.loads(pickle.dumps(tvm_ffi.core.Bytes(...))) now returns raw Python bytes.
  • Tests now assert exact native types for both directly constructed values and cached FFI-returned values.

This intentionally drops FFI wrapper type/subclass state across pickle, matching the desired raw Python payload behavior.

@tqchen tqchen merged commit 88e066e into apache:main Jun 23, 2026
9 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants