6

我有一个用 PyO3 用 Rust 编写的 Python 库,它涉及一些昂贵的计算(单个函数调用最多 10 分钟)。从 Python 调用时如何中止执行?

Ctrl+C 好像只在执行结束后才处理,所以本质上是没用的。

最小的可重现示例:

# 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()

输入后立即输入wait.sleep()Ctrl + C字符^C打印到屏幕上,但仅10秒后我终于得到

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

KeyboardInterrupt检测到,但直到调用 Rust 函数结束时才被处理。有没有办法绕过它?

当 Python 代码放入文件并从 REPL 外部执行时,行为是相同的。

4

2 回答 2

3

您的问题与此问题非常相似,只是您的代码是用 Rust 而不是 C++ 编写的。

您没有说您使用的是哪个平台 - 我将假设它类似于 unix。此答案的某些方面对于 Windows 可能不正确。

在类 unix 系统中,Ctrl+C 会导致将SIGINT信号发送到您的进程。在 C 库的最低级别,应用程序可以注册函数,这些函数将在接收到这些信号时调用。有关信号的更详细说明,请参见man signal(7)

因为可以随时调用信号处理程序(甚至在您通常认为是原子的某些操作的一部分中),所以信号处理程序实际上可以做的事情有很大的限制。这与编程语言或环境无关。大多数程序只是在收到信号时设置一个标志然后返回,然后检查该标志并对其采取行动。

Python 也不例外——它为信号设置了一个信号处理程序,该处理程序SIGINT设置了一些标志,它会检查(当这样做是安全的)并采取行动。

这在执行 python 代码时可以正常工作——它会在每个代码语句中至少检查一次标志——但在执行一个用 Rust(或任何其他外语)编写的长时间运行的函数时,情况就不同了。在您的 rust 函数返回之前,不会检查该标志。

您可以通过检查 rust 函数中的标志来改进问题。PyO3公开PyErr_CheckSignals函数,它正是这样做的。这个功能:

检查信号是否已发送到进程,如果是,则调用相应的信号处理程序。如果支持信号模块,则可以调用用 Python 编写的信号处理程序。在所有情况下,SIGINT 的默认效果是引发 KeyboardInterrupt 异常。如果引发异常,则设置错误指示器并且函数返回 -1;否则函数返回 0

因此,您可以在 Rust 函数中以适当的时间间隔调用此函数,并检查返回值。如果是 -1,你应该立即从你的 Rust 函数返回;否则继续。

如果您的 Rust 代码是多线程的,情况会更加复杂。您只能PyErr_CheckSignals从与 Python 解释器调用您的线程相同的线程中调用;如果它返回 -1,您将不得不清理您在返回之前启动的任何其他线程。究竟如何做到这一点超出了这个答案的范围。

于 2020-06-15T01:33:28.470 回答
3

一种选择是生成一个单独的进程来运行 Rust 函数。在子进程中,我们可以设置一个信号处理程序来在中断时退出进程。然后 Python 将能够根据需要引发 KeyboardInterrupt 异常。以下是如何执行此操作的示例:

// 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")

这是按下 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
于 2020-06-15T02:19:00.330 回答