Pyo3 bindings to tskit-rust

[currocam]

Hello,
Thanks a lot for putting this together! I’m finding it extremely useful for my PhD work. At some point, I wanted to pass Rust tables into a Python runtime, and I couldn’t figure out how, so I searched in the issues here looking for references.

I didn’t find any so I made (mostly for my future self) a minimal project for exposing rust-tskit into a Python interface. Perhaps someone else finds it useful. Sorry if I missed it and it’s already covered in the book (or there’s already a better reference! Please let me know).

https://github.com/currocam/tskit-rust-pyo3-starter

[molpopgen]

Thanks @currocam -- I have done tricks like this myself. While nothing involving pyo3/python will end up in this repo (the maintenance hassle is too high), there is room for a second repo that handles these abstractions. I think you can modify your code to be zero copy -- at the C level, a tree sequence can "borrow" a table collection rather than take ownership of it. (This is something that a standalone repo would definitely want to do!) Once you go down this road, though, you have the issue where the rust lifetimes need to be maintained via the Python side but the rust lifetime info has been lost due to working through the raw pointers!

[molpopgen]

If the current rust-side TableCollection had the following function, then zero-copy transfer becomes "trivial":

pub fn into_raw(self) -> *mut bindings::tsk_table_collection_t {
todo!()
}

Your "read pointer" function needs to return the *mut *mut pointer instead of dereferencing it. Then, it can be rebound to the pointer generated from rust. (You also need to free the pointer from non-rust tskit, etc.).

The challenge with this is that the safety requirements become many and nuanced. In addition to a bunch of proximal concerns, you also require that the tskit-python has been compiled with a tskit-c that is ABI compatible with that used for tskit-rust. There is no way to guarantee that other than manual inspection / having knowledge of the ecosystem. Fortunately, tables and tree sequences are very stable right now.

ABI stability includes precise layout stability!

These concerns are why I have always preferred dump/load to get data from rust to Python. Many of the same concerns are involved, but there is a lot less code to maintain!

[molpopgen]

See #830 for steps towards better support for this kind of thing.

[currocam]

Thanks!

That looks great indeed. I agree it's easier to dump the table, but I was hoping to reduce I/O for some applications. Thanks for the zero copy suggestion, I will definitely look into this. I can't imagine having shared access from the Rust and Python side (something like a mutex around the pointer?). I'll work on passing the "TSK_TAKE_OWNERSHIP".

[molpopgen]

It looks like the "take ownership" is nowhere in the Python API! I have a workaround for that -- I'll try to mock something up shortly.

[molpopgen]

See here. Once I polish this up I can transfer it over to tskit-dev and push to crates.io.

[molpopgen]

@currocam -- I have a solution in the works. As you noted, truly shared access is impossible. But tskit-rust can simply pass the buck on memory safety and let client code worry about things.

[currocam]

That's amazing! Thank you. I'll take a look soon

[molpopgen]

I am feeling the need to add some Python-side tests to this in tskit-rust. (There are some there already for metadata transfer.) On my local machine, I learned that mismatching PyMem_Malloc with C's free only leads to runtime errors when you are building and running everything against a Python built with debugging symbols! It is relatively tricky to set up Python+debug in GitHub actions.

[currocam]

PyMem_Malloc with C's free only leads to runtime errors when you are building and running everything against a Python built with debugging symbols!

This sounds terribly complicated to debug! 😆

I'm happy to contribute to the test suite (and this and other issues for that matter) when I'm back from a conference and teaching in a couple of weeks.

[molpopgen]

PyMem_Malloc with C's free only leads to runtime errors when you are building and running everything against a Python built with debugging symbols!

This sounds terribly complicated to debug! 😆

Fortunately, it is easy to do locally -- Fedora Linux is great about having many possible Python builds to work with. And uv facilitates things otherwise.

I'm happy to contribute to the test suite (and this and other issues for that matter) when I'm back from a conference and teaching in a couple of weeks.

I'm going to merge the tskit-rust PR "soon" so that the other/new repo can stop depending on a random branch. I'll return to tests soon after that but I'm off camping and visiting family next week.

[molpopgen]

@currocam -- I have published an alpha version of a crate to do one-way data exchange. See here. I think we can close this now?