In a previous post, I said I’d like to follow up and try creating some python bindings for rust.

In a microservice-inspired era, it seems the default way to offload performance-critical modules from high level languages would be through some type of json-rpc (maybe even protocol-buffers, capn-proto, etc).

But I still think there’s something nice about having a synchronous way to invoke optimized modules from another language! I’ve been meaning to test rust-python bindings since reading Armin Ronacher’s post about Fixing Python Performance with Rust at Sentry.

Since Sentry open-sourced their utility for integrating rust with python’s setuptools (see: Evolving Our Rust With Milksnake), I though I’d test it out.

(Please browse the source if interested!)

Milksnake

After getting the example up and running with milksnake, I started tinkering with the CFFI exports to get them closer to the interfaces I would need to handle game_of_life structs and functions. It wasn’t long before I started seeing my Python console littered with cryptic errors:

TypeError: result type 'struct Point' is opaque

TypeError: initializer for ctype 'Point *' must be a pointer or array, not cdata 'Point'

TypeError: invalid result type: 'uintptr_t[10][10]'

cffi.error.CDefError: cannot parse "void double(Point *foo);"

And of course the ubiquitous Segmentation fault (core dumped)!

Now, there were good reasons for these errors. And they are undoubtedly solvable. But I found them pretty difficult to approach because they were:

  • far more finicky and low-level than typical Python Exceptions
  • far less clear and helpful than rust compilation errors

Also, I found that the build.rs used by milksnake usually garbled the compiler-errors. Eg always rendering: thread 'main' panicked at 'called `Result::unwrap()` on an `Err` value: ParseSyntaxError.

PyO3

Since I want to focus my energy on working with rust and python rather than debugging C code, I decided to try the PyO3 instead. Since setuptools_rust is implemented with PyO3, I used it instead of milksnake, though I would imagine milksnake would have worked as well.

The overview is very similar with setuptools_rust as with milksnake:

  • build Rust code into a reusable crate
  • build a new crate that wraps the rust types with PyO3 types
  • use setuptools_rust to automatically generate low-level Python bindings
  • write high-level Python wrappers around the low-level bindings

For game_of_life, my goal is two simple functions:

  • setup: return a Python set representing the inital board state (Acorn config)
  • step: calculate the next board state (optionally the state after ntimes iterations)

While PyO3 has the capability to deal with tasks like allocating PyObjects and handling variable-args/keywords-args, I found that it was simplest to let the high level python wrappers do this instead.

from . import _gol

def setup():
    container = set()
    _gol.setup(container)
    return container

def step(board, ntimes=1):
    _gol.step(board, ntimes)

The expectation is that _gol (the auto-generated low level bindings) will drop all objects allocated internally during computation, and fill the output into the objects container and board, respectively.

And here is how that is done in rust:

fn do_setup(container: &PySet) {
    // load HashSet values into PySet
    for elem in game::setup() {
        container.add(elem).unwrap();
    }
}

fn do_step(py_gil: Python, board: &PySet, ntimes: usize) {
    // empty PySet into HashSet
    let mut curr = game::Board::new();
    for _ in 0..board.len() {
        let pt: (isize, isize) = board.pop().unwrap().extract(py_gil).unwrap();
        curr.insert(pt);
    }

    let mut next = curr;
    for _ in 0..ntimes {
        next = game::next_generation(&next);
    }

    // load HashSet values back into PySet
    for elem in next {
        board.add(elem).unwrap();
    }
}

Relatively straightfoward! Lots of unwraps though…

Now to create the bindings we need to use PyO3’s modinit macro to wrap the functions:

#[py::modinit(_gol)]
fn init(py: Python, m: &PyModule) -> PyResult<()> {

    #[pyfn(m, "setup")]
    fn setup(_py: Python, container: &PySet) -> PyResult<()> {
        do_setup(container);
        return Ok(());
    }

    #[pyfn(m, "step")]
    fn step(py: Python, board: &PySet, ntimes: usize) -> PyResult<()> {
        do_step(py, board, ntimes);
        return Ok(());
    }

    return Ok(());
}

Note that the only reason I defined do_setup and do_step elsewhere is so that the compiler doesn’t blame all errors on the line #[py::modinit(_gol)] (due to macro expansion).

Now all that’s left is to build the project using setuptools!

from setuptools import setup
from setuptools_rust import Binding, RustExtension

setup(
    name='gol_py',
    version='1.0',
    rust_extensions=[
        RustExtension('gol_py._gol', 'Cargo.toml', binding=Binding.PyO3)
    ],
    packages=['gol_py'],
    zip_safe=False)

Since the module is output to build/lib/, a tiny wrapper makes testing build correctness easier:

import os

os.system('python3 setup.py build')
os.chdir('build/lib')

import gol_py

board = gol_py.setup()
gol_py.step(board, ntimes=1000)

assert(len(board) == 457)

print('Computation is correct!')

Which we can invoke as a script:

$ py -m test
running build
running build_py
running build_ext
running build_rust
cargo rustc --lib --manifest-path Cargo.toml --features pyo3/extension-module pyo3/python3 --release -- --crate-type cdylib
    Finished release [optimized] target(s) in 0.23 secs
Computation is correct!

Conclusions

It worked surprisingly well! PyO3 uses procedural macros, which makes compilation errors in modinit harder to debug, and also forces builds to use rust nightly, but otherwise having PyObject structs to work with makes the bindings pretty easy to write.

Followup posts may include: