Hi all,

I have written an extension module for Python in Rust, using pyo3. The data is stored in nested structs which ultimately store their data in ndarrays. The calculations are carried out on GPU by using OpenCL and the results written back to the ndarrays. Scripts using my module run fine for many iterations but then segfault (or sometimes encounter other weird Python errors). The strange thing is the segfault and other errors are not happening in my code, but rather in some numpy methods I'm calling. However these functions are running over the exposed ndarray data.

Since the errors are occuring outside my OpenCL kernel I do not think I have an indexing error there. (I've also double and triple checked that.) The only other thing I can think of this that I'm somehow exposing my array data to Python in an incorrect way.

I need the struct hierarchy to be readable from Python, and the array data to be both readable and writable. Unfortunately pyo3(get) clones data, necessitating wrapping everything in Py references everywhere, which then need to be borrowed at the point of use. Thus, for reasons of ergonomics, I have found it necessary to keep the pyclasses separate from their purely Rust counterparts.

My code is somewhat long, so a minimal example which shows my solution is attached here:

use std::sync::{Arc, Mutex};

use ndarray::Array1;
use numpy::PyArray1;
use pyo3::{
    pyclass, pymethods, pymodule,
    types::{PyModule, PyModuleMethods},
    Bound, PyResult,
};

/// Rust inner struct.
struct Foo {
    array: Array1<f32>,
}

impl Foo {
    /// Some heavy (OpenCL) calculation which outputs its result to array.
    fn calculate(&mut self) {
        todo!();
    }
}

/// Rust outer struct.
struct Bar {
    foo: Foo,
}

/// Python wrapped outer struct, contains ref to the Rust outer struct.
#[pyclass]
pub struct PyBar {
    /// Arc + Mutex necessary to make things threadsafe (in reality `Foo` is not `Send``)
    bar: Arc<Mutex<Bar>>,
}

#[pymethods]
impl PyBar {
    #[getter]
    unsafe fn foo<'py>(this: Bound<'py, PyBar>) -> Bound<'py, PyFoo> {
        let borrow_mut = this.borrow_mut();
        Bound::new(
            this.py(),
            PyFoo {
                bar: borrow_mut.bar.clone(),
            },
        )
        .unwrap()
    }
}

/// Python wrapped inner struct, contains ref to the Rust outer struct containing the relevant inner struct.
#[pyclass]
#[derive(Clone)]
pub struct PyFoo {
    bar: Arc<Mutex<Bar>>,
}

#[pymethods]
impl PyFoo {
    // This works without all the `mut`s too, but I don't know which is correct.
    #[getter]
    pub unsafe fn array<'py>(this: Bound<'py, PyFoo>) -> Bound<'py, PyArray1<f32>> {
        let foo = this.borrow_mut();
        let array = &mut foo.bar.lock().unwrap().foo.array;
        PyArray1::borrow_from_array(array, this.into_any())
    }
}

#[pymodule]
fn my_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_class::<PyFoo>()?;
    m.add_class::<PyBar>()?;

    Ok(())
}

This "solution" feels unsatisfying to me but it seems to work and it fulfills my two requirements of splitting the python and rust structs completely, and allowing writable access to the struct hierarchy.

Two immediate and one auxiliary questions:

1. Is there anything wrong with exposing the ndarray data to Python in this way?
2. Is there a better way, which still maintains the requirements?
3. This is probably an X/Y problem. Can the segfaults be being caused by something else?

Any help is much appreciated!

Thanks,
Dan

1 post - 1 participant

Read full topic