Keyboard Interrupt from Python does not abort Rust function (PyO3)

Question

I have a Python library written in Rust with PyO3, and it involves some expensive calculations (up to 10 minutes for a single function call). How can I abort the execution when calling from Python ?

Ctrl+C seems to only be handled after the end of the execution, so is essentially useless.

Minimal reproducible example:

# Cargo.toml

[package]
name = "wait"
version = "0.0.0"
authors = []
edition = "2018"

[lib]
name = "wait"
crate-type = ["cdylib"]

[dependencies.pyo3]
version = "0.10.1"
features = ["extension-module"]

// src/lib.rs

use pyo3::wrap_pyfunction;

#[pyfunction]
pub fn sleep() {
    std::thread::sleep(std::time::Duration::from_millis(10000));
}

#[pymodule]
fn wait(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_wrapped(wrap_pyfunction!(sleep))
}

$ rustup override set nightly
$ cargo build --release
$ cp target/release/libwait.so wait.so
$ python3
>>> import wait
>>> wait.sleep()

Immediately after having entered wait.sleep() I type Ctrl + C, and the characters ^C are printed to the screen, but only 10 seconds later do I finally get

>>> wait.sleep()
^CTraceback (most recent call last):
  File "<stdin>", line 1, in <module>
KeyboardInterrupt
>>>

The KeyboardInterrupt was detected, but was left unhandled until the end of the call to the Rust function. Is there a way to bypass that ?

The behavior is the same when the Python code is put in a file and executed from outside the REPL.

Answer 1

Your problem is very similar to this one, except that your code is written in Rust instead of C++.

You did not say which platform you are using - I am going to assume it is unix-like. Some aspects of this answer may not be correct for Windows.

In unix-like systems, Ctrl+C results in a SIGINT signal being sent to your process. At the very low level of the C Library, applications can register functions which will be called when these signals are received. See man signal(7) for a more detailed description of signals.

Because signal handlers can be called at any time (even part way through some operation that you would normally consider atomic) there are big limits on what a signal handler can actually do. This is independent of the programming language or environment. Most programs just set a flag when they receive a signal and then return, and then later check that flag and take action on it.

Python is no different - it sets a signal handler for the SIGINT signal which sets some flag, which it that checks (when it is safe to do so) and takes action on.

This works OK when executing python code - it will check the flag at least once per code statement - but it is a different matter when executing a long running function written in Rust (or any other foreign language). The flag does not get checked until your rust function returns.

You can improve matters by checking the flag within your rust function. PyO3 exposes the PyErr_CheckSignals function which does exactly that. This function:

checks whether a signal has been sent to the processes and if so, invokes the corresponding signal handler. If the signal module is supported, this can invoke a signal handler written in Python. In all cases, the default effect for SIGINT is to raise the KeyboardInterrupt exception. If an exception is raised the error indicator is set and the function returns -1; otherwise the function returns 0

So, you could call this function at suitable intervals inside your Rust function, and check the returned value. If it was -1, you should immediately return from your Rust function; otherwise keep going.

The picture is more complex if your Rust code is multi-threaded. You can only call PyErr_CheckSignals from the same thread as the python interpreter called you; and if it returns -1 you would have to clean up any other threads that you had started before returning. Exactly how to do that is beyond the scope of this answer.

Answer 2

One option would be to spawn a separate process to run the Rust function. In the child process, we can set up a signal handler to exit the process on interrupt. Python will then be able to raise a KeyboardInterrupt exception as desired. Here's an example of how to do it:

// src/lib.rs
use pyo3::prelude::*;
use pyo3::wrap_pyfunction;
use ctrlc;

#[pyfunction]
pub fn sleep() {
    ctrlc::set_handler(|| std::process::exit(2)).unwrap();
    std::thread::sleep(std::time::Duration::from_millis(10000));
}

#[pymodule]
fn wait(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_wrapped(wrap_pyfunction!(sleep))
}

# wait.py
import wait
import multiprocessing as mp

def f():
    wait.sleep()

p = mp.Process(target=f)
p.start()
p.join()
print("Done")

Here's the output I get on my machine after pressing CTRL-C:

$ python3 wait.py
^CTraceback (most recent call last):
  File "wait.py", line 9, in <module>
    p.join()
  File "/home/kerby/miniconda3/lib/python3.7/multiprocessing/process.py", line 140, in join
    res = self._popen.wait(timeout)
  File "/home/kerby/miniconda3/lib/python3.7/multiprocessing/popen_fork.py", line 48, in wait
    return self.poll(os.WNOHANG if timeout == 0.0 else 0)
  File "/home/kerby/miniconda3/lib/python3.7/multiprocessing/popen_fork.py", line 28, in poll
    pid, sts = os.waitpid(self.pid, flag)
KeyboardInterrupt