How to extract a python function into a rust function

[shivamka1]

Okay so I understand so far that we can call rust functions from python and also other way around but how do I take a python function extract it into a rust function to be passed around in rust?

Allow me to elaborate!

I have a function in rust library that I'am building that takes a function which I would like to get from python.

fn call_greet<F>(f: F)
where
    F: Fn(&str) -> (),
{
    let name = "Shivam";
    f(name)
}

Now I wish to write a python function module that takes python function (probably as PyObject?) and extract it into a rust function to be passed around like so:

type O = Box<dyn Fn(&str) -> ()>;

#[pyfunction]
pub fn invoke_passed_func2(func: PyObject) -> PyResult<()> {
    Python::with_gil(|py| -> PyResult<()> {
        let f: O = func.extract(py)?;
        call_greet(f);
        Ok(())
    })
}

But this returns error:

the trait bound `Box<dyn for<'a> Fn(&'a str)>: PyClass` is not satisfied

How do I do it in pyO3? Or if it is even possible?

---

I certainly am able to call python function passed as an argument in a directly like so:

#[pyfunction]
pub fn invoke_passed_func(func: PyObject) -> PyResult<()> {
    Python::with_gil(|py| -> PyResult<()> {
        let name = "Shivam";
        let args = (name,);
        func.call1(py, args)?;
        Ok(())
    })
}

From python:

def greet(name):
    print(f"Hello {name}")

pyo3_example.invoke_passed_func(greet)
# Hello

[mejrs]

You can do something similar to what https://pyo3.rs/main/class/call.html does: store the function as a Py<PyAny> and just .call it.

[shivamka1]

Thanks for your reply @mejrs. I don't think decorator is what I need here. What I need is a rust function extracted from python function which can be forward further up in the rust library. In other words., I intend to leave the python context all together once rust function is extracted. Please see the contrived example I have demonstrated.

Correct me if it is still possible with decorators.

[adamreichold]

I think what @mejrs meant was to invoke .call directly on your PyObject, e.g.

#[pyfunction]
pub fn invoke_passed_func2(py: Python, func: PyObject) {
    call_greet(|name| {
        Python::with_gil(|py| -> PyResult<()> {
           func.call1(py, (name,)).unwrap();
        });
    });
}

You will need to unwrap the result if call_greet does not expect a fallible closure, but that is a separate issue.

Also not that a #[pyfunction] does not need to call Python::with_gil as the GIL is already held and you can get a Python by just including it in the arguments.

[adamreichold]

Similarly, you should be able to get an instance of O using such a closure:

type O = Box<dyn Fn(&str) -> ()>;

#[pyfunction]
pub fn invoke_passed_func2(func: PyObject) {
    let func: O = Box::new(|name| {
        Python::with_gil(|py| -> PyResult<()> {
           func.call1(py, (name,)).unwrap();
        });
    });

    // stash `func` elsewhere or pass it to `call_greet`
}

[shivamka1]

@adamreichold Thanks for your tip on Python::with_gil!

I tried your suggestion as shown below but I still get errors (Sorry pretty new to both rust and pyo3):

type O = Box<dyn Fn(&str) -> ()>;

#[pyfunction]
pub fn invoke_passed_func3(py: Python, func: PyObject) -> PyResult<()> {
    let fun: O = Box::new(|name| {
        Python::with_gil(|py| {
            func.call1(py, (name,)).unwrap();
        });
    });

    call_greet(fun);

    Ok(())
}

fn call_greet<F>(f: F)
where
    F: Fn(&str) -> (),
{
    let name = "Shivam";
    f(name)
}

`func` does not live long enough

[adamreichold]

Sorry, I did not compile the example before posting it. You probably need to move the PyObject into the closure, i.e.

    let fun: O = Box::new(move |name| {
        Python::with_gil(|py| {
            func.call1(py, (name,)).unwrap();
        });
    });

[mejrs]

In other words., I intend to leave the python context all together once rust function is extracted

Well, you can't do that. You can't turn a python function into a rust function. At best you can take a python function and wrap it in a rust function, but you will still need to acquire the GIL to call it.

[qqlearn123]

I am curious to this as well.

I tried calling function supplied from python in rust inside parallel code. It turned out that I cannot do that because the function does not implement Sync trait.

Is there some other way that we can convert user-supplied lambda (or regular) function from python to rust without sacrificing parallelism?

The reason I want to do this is because I am running an expansive function on rust side using within parallel processing but this function takes a lambda as one of its arguments. Ultimately, I want to have python api so that lambda function will be supplied on python side.

[adamreichold]

I tried calling function supplied from python in rust inside parallel code. It turned out that I cannot do that because the function does not implement Sync trait.

You need to take the user-supplied function not as a GIL-bound &'py PyAny, but as a GIL-independent Py<PyAny> a.k.a. PyObject, c.f. https://pyo3.rs/v0.18.3/types

Py<PyAny> is Send and Sync, can be cloned cheaply (just incrementing the reference count), and passed therefore easily passed to other threads. To call this function, you need to obtain a GIL token on the calling thread using e.g. Python::with_gil though and this is where the synchronisation happens. (The GIL token Python itself is obviously not Send since it is bound to the current thread holding the GIL.)

[qqlearn123]

I tried calling function supplied from python in rust inside parallel code. It turned out that I cannot do that because the function does not implement Sync trait.

You need to take the user-supplied function not as a GIL-bound &'py PyAny, but as a GIL-independent Py<PyAny> a.k.a. PyObject, c.f. https://pyo3.rs/v0.18.3/types

Py<PyAny> is Send and Sync, can be cloned cheaply (just incrementing the reference count), and passed therefore easily passed to other threads. To call this function, you need to obtain a GIL token on the calling thread using e.g. Python::with_gil though and this is where the synchronisation happens. (The GIL token Python itself is obviously not Send since it is bound to the current thread holding the GIL.)

Thanks. I tried the following and it actually compiles fine.

use rayon::preclude::*;

#[pyfunction]
fn test(pyfn: PyObject, v: Vec<f64>) -> f64 {
    v.par_iter().map(|&i| {
        Python::with_gil(|py| {
            pyfn.call1(py, (i, i)).unwrap().extract::<f64>(py).unwrap()
        })
    }).sum()
}

But when I run it from python side, it never returns.

import test_module

def f(x, y):
    return x + y

test_module.test(f, [1, 2, 3])

Am I introducing some sort of deadlocks? Or it just cannot work this way (call in parallel code) for user-supplied function from python side?

[adamreichold]

Am I introducing some sort of deadlocks?

Yes, when a #[pyfunction] is called, the GIL is held for the duration of the call. You need to temporarily release it by wrapping your v.par_iter().map(..).sum() invocation in Python::allow_threads as described in https://pyo3.rs/v0.18.3/parallelism

[qqlearn123]

Am I introducing some sort of deadlocks?

Yes, when a #[pyfunction] is called, the GIL is held for the duration of the call. You need to temporarily release it by wrapping your v.par_iter().map(..).sum() invocation in Python::allow_threads as described in https://pyo3.rs/v0.18.3/parallelism

Thanks for the reference!

It turned out that in order to wrap it with allow_threads, I have to add the py: Python<'_> argument in my function and either do Python::allow_thread(py, move || ...) or py.allow_thread(move ||), and also change the name of the inner py inside with_gil.

Now, my function becomes,

#[pyfunction]
fn test(py: Python<'_>, pyfn: PyObject, v: Vec<f64>) -> f64 {
    py.allow_threads(move || {
        v.par_iter().map(|&i| {
            Python::with_gil(|py2| {
                pyfn.call1(py2, (i, i)).unwrap().extract::<f64>(py2).unwrap()
            })
        }).sum()
    })
}

It compiles and runs fine. But, are there some ways to simplify it a little bit or it has to be implemented in this way?

Also, if I run test(f, [1, 2, 3]) in python, will this f truly run on parallel? Or is it still sequential under the hood?

[adamreichold]

It compiles and runs fine. But, are there some ways to simplify it a little bit or it has to be implemented in this way?

I don't think this can be simplified significantly. There is no need to give the inner GIL token a different name though, it can just show the outer one which cannot "enter" the closure passed to allow_threads in any case.

Please not that as written in the example, the parallel iteration is completely useless as the closure passed to map will take the GIL for each item. So this strategy will only help if there is a significant amount of work done outside of the call to with_gil.