Monitor Rust progress from Python

[KCrux]

Hello,

I'm writing a GUI (with PyQt5) for some scientific calculations in Python. The calculations are written in Rust. I would like to display the progress of the calculations in the GUI within a progress bar. But I can't figure out how to do so.

First my Rust code:

#[pyclass]
pub struct Test {
    progress: usize
}

#[pymethods]
impl Test {

    #[new]
    pub fn new() -> Self {
        Self {progress: 0}
    }

    pub fn start(&mut self, py: Python) {

        py.allow_threads(move || {

            for _ in 0..20 {
                std::thread::sleep(std::time::Duration::from_secs_f64(0.1));
                self.progress += 1;
            }
    
        })

    }

    pub fn get_progress(&self) -> usize {
        self.progress
    }

}

Now my Python code:

import rs_module
import time
from threading import Thread

class T(Thread):

    def __init__(self):

        Thread.__init__(self)

        self._test = None

    def run(self):

        print("thread started")

        self._test = rs_module.Test()
        self._test.start()

        print("thread finished")

    def get_progress(self):
        if self._test != None:
            print(self._test.get_progress())
        else:
            print(None)

thread = T()
thread.start()

print("start")

tstart = time.time()
while time.time() - tstart < 10:
    time.sleep(1)
    thread.get_progress()

print("finish")

Running the Python code I get the following error message:

    print(self._test.get_progress())
RuntimeError: Already mutably borrowed

Any suggestions are highly appreciated.

[mejrs]

You have the same re-entrancy problem as in https://pyo3.rs/main/class/call.html#what-is-the-cell-for

Consider changing it to something like this:

#[pyclass]
pub struct Test {
-    progress: usize
+    progress: AtomicUsize
}

#[pymethods]
impl Test {
    #[new]
    pub fn new() -> Self {
-        Self {progress: 0}
+        Self {progress: AtomicUsize::new(0)}
}

-    pub fn start(&mut self, py: Python) {
+    pub fn start(&self, py: Python) {
        py.allow_threads(move || {

            for _ in 0..20 {
                std::thread::sleep(std::time::Duration::from_secs_f64(0.1));
-                self.progress += 1;
+                self.progress.fetch_add(1, Ordering::SeqCst);
            }
    
        })

    }

    pub fn get_progress(&self) -> usize {
-        self.progress
+        self.progress.load(Ordering::SeqCst)
    }
}

[songololo]

I'm not certain what the best strategy would be but what I'm trying to accomplish is to log progress (ideally a progress bar) to report on a long-running parallelised rust function (called from Python via pyo3).

I'm not certain how to proceed because it seems problematic to callback to Python from within a par_iter() to update a progress bar logged from within the containing python function. This because of needing to acquire the GIL from inside the parallel threads which would presumably cause hassles and a major slow-down. Though I suppose if this is theoretically doable then one way to proceed would be to only update (via a callback) every e.g. 100 or 1000 iterations so that this doesn't overly slow-down the parallel processing.

I'm not sure whether a callback is a good or bad approach for this purpose and while researching this I stumbled into the above question, which seems an interesting strategy to get periodic updates on the progress of the parallel iterations without getting them bogged down with calls into Python. That said, I'd ideally be able to call the rust function directly instead of spawning a thread. I'm completely open to ideas.

[birkenfeld]

If the atomic approach is not feasible for you, I would send updates from the parallelized threads into a threadsafe but fast Rust channel/queue, accept those updates in a separate thread, consolidate them and make the Python calls from there.

[songololo]

Thank you for the feedback, this is what I've ended up doing:

From Python:

    def wrapper(queue):
        result = my_func()
        queue.put(result)

    result_queue = Queue()
    thread = threading.Thread(target=wrapper, args=(result_queue,))
    pbar = tqdm(total=my_count)
    thread.start()
    while thread.is_alive():
        time.sleep(1)
        pbar.update(my_struct.progress() - pbar.n)
    pbar.update(my_struct.progress() - pbar.n)
    pbar.close()
    result = result_queue.get()
    thread.join()

From rust:

// I'm using a struct
#[pyclass]
#[derive(Clone)]
pub struct MyStruct {
    pub progress: Arc<AtomicUsize>,
}
// with a progress method
#[pymethods]
impl MyStruct {
    fn progress(&self) -> usize {
        self.progress.as_ref().load(Ordering::Relaxed)
    }
}
// and the function from which it is invoked follows this pattern
// iter
let result = py.allow_threads(move || {
    indices.par_iter().for_each(|src_idx| {
        // progress
        if !pbar_disabled {
            self.progress.fetch_add(1, Ordering::Relaxed);
        }
    });
});

There seems to be a small performance hit but quite manageable in exchange for the feedback...

Curiosity: Would pyo3-log work for this type of scenario, is it better to stick with the above approach?

[adamreichold]

let atomic = Arc::clone(&self.progress);

It should not be necessary to clone the Arc for each item. The closure used by parallel iterators should be able to access stack variables in the outer scope.

As a general design trick, I would suggest starting the background via Rust as its JoinHandle makes it easier to retrieve the computations results (avoiding the need for an extra queue).

[songololo]

Thank you, I've updated to this: self.progress.fetch_add(1, Ordering::Relaxed);