Python 的鍵盤中斷不會中止 Rust function (PyO3)

Question

我有一個用 PyO3 用 Rust 編寫的 Python 庫，它涉及一些昂貴的計算（單個 function 調用最多需要 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 function 的調用結束時才進行處理。 有沒有辦法繞過它？

將 Python 代碼放入文件並從 REPL 外部執行時，行為是相同的。

Answer 1

您的問題與此問題非常相似，只是您的代碼是用 Rust 而不是 C++ 編寫的。

您沒有說您使用的是哪個平台 - 我將假設它類似於 unix。 對於 Windows，此答案的某些方面可能不正確。

在類 unix 系統中，Ctrl+C 會導致將SIGINT信號發送到您的進程。 在 C 庫的極低級別，應用程序可以注冊將在接收到這些信號時調用的函數。 有關信號的更詳細說明，請參見man signal(7) 。

因為可以隨時調用信號處理程序（甚至在您通常認為是原子的某些操作的一部分中），所以信號處理程序實際上可以做的事情有很大的限制。 這與編程語言或環境無關。 大多數程序只是在收到信號時設置一個標志然后返回，然后檢查該標志並對其采取行動。

Python 沒有什么不同 - 它為SIGINT信號設置了一個信號處理程序，該信號設置了一些標志，它會檢查（當這樣做是安全的）並采取行動。

這在執行 python 代碼時工作正常 - 每個代碼語句至少檢查一次標志 - 但在執行用 Rust（或任何其他外語）編寫的長時間運行的 function 時，情況就不同了。 在您的 rust function 返回之前，不會檢查該標志。

您可以通過檢查 rust function 中的標志來改進問題。 PyO3公開了PyErr_CheckSignals function 正是這樣做的。 這個 function：

檢查信號是否已發送到進程，如果是，則調用相應的信號處理程序。 如果支持信號模塊，這可以調用用 Python 編寫的信號處理程序。 在所有情況下，SIGINT 的默認效果是引發 KeyboardInterrupt 異常。 如果引發異常，則設置錯誤指示器並且 function 返回 -1； 否則 function 返回 0

因此，您可以在 Rust function 中以適當的時間間隔調用此 function，並檢查返回值。 如果是-1，你應該立即從你的 Rust function 返回； 否則繼續。

如果您的 Rust 代碼是多線程的，則情況會更加復雜。 您只能從與 python 解釋器調用您的線程相同的線程調用PyErr_CheckSignals ； 如果它返回 -1，您將不得不清理您在返回之前啟動的任何其他線程。 究竟如何做到這一點超出了這個答案的 scope 。

Answer 2

一種選擇是生成一個單獨的進程來運行 Rust function。 在子進程中，我們可以設置一個信號處理程序來在中斷時退出進程。 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 后我在我的機器上得到的 output：

$ 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