厭倦了C++,CS&ML博士用Rust重寫Python擴充套件,還總結了9條規則

語言: CN / TW / HK

選自medium.com

作者: Carl M. Kadie

機器之心編譯

編輯:杜偉、陳萍

效果好不好,試一試就知道了。

Python 是資料科學家最流行的程式語言之一,其內部集成了高質量分析庫,包括 NumPy、SciPy、自然語言工具包等,這些庫中的許多都是用 C 和 C++ 實現的。

然而,C 和 C++ 相容性差,且本身不提供執行緒安全。有研究者開始轉向 Rust,重寫 C++ 擴充套件。

擁有 CS 與機器學習博士學位的 Carl M. Kadie ,通過更新 Python 中生物資訊學軟體包 Bed-Reader,為研究者帶來了在 Rust 中編寫 Python 擴充套件的九個規則。以下是原部落格的主要內容。

一年前,我厭倦了我們軟體包 Bed-Reader 的 C++ 擴充套件,我用 Rust 重寫了它,令人高興的是,得到的新擴充套件和 C/C++ 一樣快,但具有更好的相容性和安全性。一路走來,我學會了這九條規則,可以幫助你建立更好的擴充套件程式碼,這 九條規則 包括:

1. 建立一個包含 Rust 和 Python 專案的單獨儲存庫

2. 使用 maturin & PyO3 在 Rust 中建立 Python-callable translator 函式

3. 讓 Rust translator 函式呼叫 nice Rust 函式

4. 在 Python 中預分配記憶體

5. 將 nice Rust 錯誤處理翻譯 nice Python 錯誤處理

6. 多執行緒與 Rayon 和 ndarray::parallel,返回任何錯誤

7. 允許使用者控制並行執行緒數

8. 將 nice 動態型別 Python 函式翻譯成 nice Rust 泛型函式

9. 建立 Rust 和 Python 測試

其中, 文中提到的 nice 這個詞是指使用最佳實踐和原生型別建立 。換句話說:在程式碼頂部,編寫 nice Python 程式碼;在中間,用 Rust 編寫 translator 程式碼;在底部,編寫 nice Rust 程式碼。結構如下圖所示:

上述策略看似顯而易見,但遵循它可能會很棘手。本文提供了有關如何遵循每條規則的實用建議和示例。

我在 Bed-Reader 進行了實驗, Bed-Reader 是一個 Python 包 ,用於讀取和寫入 PLINK Bed Files,這是一種在生物資訊學中用於儲存 DNA 資料的二進位制格式。Bed 格式的檔案可以達到 TB。Bed-Reader 讓使用者可以快速、隨機地訪問資料的子集。它在使用者選擇的 int8、float32 或 float64 中返回一個 NumPy 陣列。

我希望 Bed-Reader 擴充套件程式碼具有以下特點:

  • 比 Python 快;

  • 相容 NumPy;

  • 可以進行資料並行多執行緒處理;

  • 與執行資料並行多執行緒的所有其他包相容;

  • 安全。

我們最初的 C++ 擴充套件兼具速度快、與 NumPy 相容,以及使用 OpenMP 進行資料並行多執行緒等特點。遺憾的是,OpenMP 執行時庫 (Runtime library),存在 Python 包相容版本問題。

Rust 提供了 C++ 擴充套件帶來的優勢。除此之外,Rust 通過提供沒有執行時庫的資料並行多執行緒解決了執行時相容性問題。此外,Rust 編譯器還能保證執行緒安全。

在 Rust 中建立 Python 擴充套件需要許多設計決策。根據我使用 Bed-Reader 的經驗,以下是我的使用規則。 

規則 1:建立一個包含 Rust 和 Python 專案的單獨儲存庫

下表顯示瞭如何佈局檔案:

使用 Rust 常用的‘cargo new’命令建立 Cargo.toml 和 src/lib.rs 檔案。Python 沒有 setup.py 檔案。相反,Cargo.toml 包含 PyPi 包資訊,例如包的名稱、版本號、 README 檔案的位置等。要在沒有 setup.py 的情況下工作,pyproject.toml 必須包含:

[build-system]
requires = ["maturin==0.12.5"]
build-backend = "maturin"

一般來說,Python 設定在 pyproject.toml 中(如果不是,則在 pytest.ini 等檔案中)。Python 程式碼位於子資料夾 bed_reader 中。

最後,我們使用 GitHub 操作來構建、測試和準備部署。該指令碼位於 .github/workflows/ci.yml 中。

規則 2:使用 maturin & PyO3 在 Rust 中

建立 Python-callable translator 函式

Maturin 是一個 PyPi 包,可通過 PyO3 構建和釋出 Python 擴充套件。PyO3 是一個 Rust crate,用於在 Rust 中編寫 Python 擴充套件。

在 Cargo.toml 中,包含這些 Rust 依賴項:

[dependencies]
thiserror = "1.0.30"
ndarray-npy = { version = "0.8.1", default-features = false }
rayon = "1.5.1"
numpy = "0.15.0"
ndarray = { version = "0.15.4", features = ["approx", "rayon"] }
pyo3 = { version = "0.15.1", features = ["extension-module"] }
[dev-dependencies]
temp_testdir = "0.2.3"

在 src/lib.rs 底部,包含這兩行:

mod python_module;
mod tests;

規則 3:Rust translator 函式

呼叫 nice Rust 函式

在 src/lib.rs 定義了 nice  Rust 函式,這些函式將完成包的核心工作。 它們能夠輸入和輸出標準 Rust 型別並嘗試遵循 Rust 最佳實踐。例如,對於 Bed-Reader 包, read_no_alloc 是一個 nice Rust 函式,用於從 PLINK Bed 檔案讀取和返回值。

然而,Python 不能直接呼叫這些函式。因此,在檔案 src/python_module.rs 中定義 Python 可以呼叫的 Rust translator 函式,下面為 translator 函式示例:

#[pyfn(m)]
#[pyo3(name = "read_f64")]
fn read_f64_py(
_py: Python<'_>,
filename: &str,
iid_count: usize,
sid_count: usize,
count_a1: bool,
iid_index: &PyArray1<usize>,
sid_index: &PyArray1<usize>,
val: &PyArray2<f64>,
num_threads: usize,
) -> Result<(), PyErr> {
let iid_index = iid_index.readonly();
let sid_index = sid_index.readonly();
let mut val = unsafe { val.as_array_mut() };
let ii = &iid_index.as_slice()?;
let si = &sid_index.as_slice()?;
create_pool(num_threads)?.install(|| {
read_no_alloc(
filename,
iid_count,
sid_count,
count_a1,
ii,
si,
f64::NAN,
&mut val,
)
})?;
Ok(())
}

該函式將檔名、一些與檔案大小相關的整數以及兩個一維 NumPy 陣列作為輸入,這些陣列指示要讀取資料的哪個子集。該函式從檔案中讀取值並填充 val,這是一個預先分配的二維 NumPy 陣列。

注意

將 Python NumPy 1-D 陣列轉換為 Rust slices,通過:

let iid_index = iid_index.readonly();
let ii = &iid_index.as_slice()?;

將 Python NumPy 2d 陣列轉換為 2-D Rust ndarray 物件,通過:

let mut val = unsafe { val.as_array_mut() };

呼叫 read_no_alloc,這是 src/lib.rs 中一個 nice Rust 函式,它將完成核心工作。

規則 4:在 Python 中預分配記憶體

在 Python 中為結果預分配記憶體簡化了 Rust 程式碼。在 Python 端,在 bed_reader/_open_bed.py 中,我們可以匯入 Rust translator 函式:

from .bed_reader import [...] read_f64 [...]

然後定義一個 nice Python 函式來分配記憶體、呼叫 Rust translator 函式並返回結果。

def read([...]):
[...]
val = np.zeros((len(iid_index), len(sid_index)), order=order, dtype=dtype)
[...]
reader = read_f64
[...]
reader(
str(self.filepath),
iid_count=self.iid_count,
sid_count=self.sid_count,
count_a1=self.count_A1,
iid_index=iid_index,
sid_index=sid_index,
val=val,
num_threads=num_threads,
)
[...]
return val

規則 5:將 nice Rust 錯誤處理翻譯 nice Python 錯誤處理

為了瞭解如何處理錯誤,讓我們在 read_no_alloc( src/lib.rs 中 nice Rust 函式)中跟蹤兩個可能的錯誤。

示例錯誤 1:來自標準函式的錯誤。如果 Rust 的標準 File::open 函式找不到檔案或無法開啟檔案, 在這種情況下,如下一行中的? 將導致函式返回一些 std::io::Error 值。 

let mut buf_reader = BufReader::new(File::open(filename)?);

為了定義一個可返回這些值的函式,我們可以給函式一個返回型別 Result<(), BedErrorPlus>。我們定義 BedErrorPlus 時要包含所有 std::io::Error,如下所示:

use thiserror::Error;
...
/// BedErrorPlus enumerates all possible errors
/// returned by this library.
/// Based on https://nick.groenen.me/posts/rust-error-handling/#the-library-error-type
#[derive(Error, Debug)]
pub enum BedErrorPlus {
#[error(transparent)]
IOError(#[from] std::io::Error),
#[error(transparent)]
BedError(#[from] BedError),
#[error(transparent)]
ThreadPoolError(#[from] ThreadPoolBuildError),
}

這是 nice Rust 錯誤處理,但 Python 不理解它。因此,在 src/python_module.rs 中,我們要進行翻譯。首先, 定義 translator 函式 read_f64_py 來返回 PyErr ;其次,實現了一個從 BedErrorPlus 到 PyErr 的轉換器。轉換器使用正確的錯誤訊息建立正確的 Python 錯誤類(IOError、ValueError 或 IndexError)。如下所示:

impl std::convert::From<BedErrorPlus> for PyErr {
fn from(err: BedErrorPlus) -> PyErr {
match err {
BedErrorPlus::IOError(_) => PyIOError::new_err(err.to_string()),
BedErrorPlus::ThreadPoolError(_) => PyValueError::new_err(err.to_string()),
BedErrorPlus::BedError(BedError::IidIndexTooBig(_))
| BedErrorPlus::BedError(BedError::SidIndexTooBig(_))
| BedErrorPlus::BedError(BedError::IndexMismatch(_, _, _, _))
| BedErrorPlus::BedError(BedError::IndexesTooBigForFiles(_, _))
| BedErrorPlus::BedError(BedError::SubsetMismatch(_, _, _, _)) => {
PyIndexError::new_err(err.to_string())
}
_ => PyValueError::new_err(err.to_string()),
}
}
}

示例錯誤 2:特定於函式的錯誤。如果 nice 函式 read_no_alloc 可以開啟檔案,但隨後意識到檔案格式錯誤怎麼辦?它應該引發一個自定義錯誤,如下所示:

if (BED_FILE_MAGIC1 != bytes_vector[0]) || (BED_FILE_MAGIC2 != bytes_vector[1]) {
return Err(BedError::IllFormed(filename.to_string()).into());
}

BedError::IllFormed 型別的自定義錯誤在 src/lib.rs 中定義:

use thiserror::Error;
[...]
// https://docs.rs/thiserror/1.0.23/thiserror/
#[derive(Error, Debug, Clone)]
pub enum BedError {
#[error("Ill-formed BED file. BED file header is incorrect or length is wrong.'{0}'")]
IllFormed(String),
[...]
}

其餘的錯誤處理與示例錯誤 1 中的相同。

最後,對於 Rust 和 Python,標準錯誤和自定義錯誤的結果都屬於帶有資訊性錯誤訊息的特定錯誤型別。

規則 6:多執行緒與 Rayon 和 ndarray::parallel,返回任何錯誤

Rust Rayon crate 提供了簡單且輕量級的資料並行多執行緒。ndarray::parallel 模組將 Rayon 應用於陣列。通常的模式是跨一個或多個 2D 陣列的列(或行)並行化。面對的一個挑戰是從並行執行緒返回任何錯誤訊息。我將重點介紹兩種通過錯誤處理並行化陣列操作的方法。以下兩個示例都出現在 Bed-Reader 的 src/lib.rs 檔案中。

方法 1:par_bridge().try_for_each

Rayon 的 par_bridge 將順序迭代器變成了並行迭代器。如果遇到錯誤,使用 try_for_each 方法可以儘快停止所有處理。

這個例子中,我們遍歷了壓縮(zip)在一起的兩個 things:

  • DNA 位置的二進位制資料;

  • 輸出陣列的列。

然後,按順序讀取二進位制資料,但並行處理每一列的資料。我們停止了任何錯誤。

[... not shown, read bytes for DNA location's data ...]
// Zip in the column of the output array
.zip(out_val.axis_iter_mut(nd::Axis(1)))
// In parallel, decompress the iid info and put it in its column
.par_bridge() // This seems faster that parallel zip
.try_for_each(|(bytes_vector_result, mut col)| {
match bytes_vector_result {
Err(e) => Err(e),
Ok(bytes_vector) => {
for out_iid_i in 0..out_iid_count {
let in_iid_i = iid_index[out_iid_i];
let i_div_4 = in_iid_i / 4;
let i_mod_4 = in_iid_i % 4;
let genotype_byte: u8 = (bytes_vector[i_div_4] >> (i_mod_4 * 2)) & 0x03;
col[out_iid_i] = from_two_bits_to_value[genotype_byte as usize];
}
Ok(())
}
}
})?;

方法 2:par_azip!

ndarray 包的 par_azip!巨集允許並行地通過一個或多個壓縮在一起的陣列或陣列片段。在我看來,這非常具有可讀性。但是,它不直接支援錯誤處理。因此,我們可以通過將任何錯誤儲存到結果列表來新增錯誤處理。

下面是一個效用函式的例子。完整的效用函式從三個計數和總和(count and sum)陣列計算統計量(均值和方差),並且並行工作。如果在資料中發現錯誤,則將該錯誤記錄在結果列表中。在完成所有處理之後,檢查結果列表是否有錯誤。

[...]
let mut result_list: Vec<Result<(), BedError>> = vec![Ok(()); sid_count];
nd::par_azip!((mut stats_row in stats.axis_iter_mut(nd::Axis(0)),
&n_observed in &n_observed_array,
&sum_s in &sum_s_array,
&sum2_s in &sum2_s_array,
result_ptr in &mut result_list)
{
[...some code not shown...]
});
// Check the result list for errors
result_list.par_iter().try_for_each(|x| (*x).clone())?;
[...]

Rayon 和 ndarray::parallel 提供了許多其他不錯的資料並行處理方法。

規則 7:允許使用者控制並行執行緒數

為了更好地使用使用者的其他程式碼,使用者必須能夠控制每個函式可以使用的並行執行緒數。

在下面這個 nice Python read 函式中,使用者可以得到一個可選的 num_threadsargument。如果使用者沒有設定它,Python 會通過這個函式設定它:

def get_num_threads(num_threads=None):
if num_threads is not None:
return num_threads
if "PST_NUM_THREADS" in os.environ:
return int(os.environ["PST_NUM_THREADS"])
if "NUM_THREADS" in os.environ:
return int(os.environ["NUM_THREADS"])
if "MKL_NUM_THREADS" in os.environ:
return int(os.environ["MKL_NUM_THREADS"])
return multiprocessing.cpu_count()

接著在 Rust 端,我們可以定義 create_pool。這個輔助函式從 num_threads 構造一個 Rayon ThreadPool 物件。

pub fn create_pool(num_threads: usize) -> Result<rayon::ThreadPool, BedErrorPlus> {
match rayon::ThreadPoolBuilder::new()
.num_threads(num_threads)
.build()
{
Err(e) => Err(e.into()),
Ok(pool) => Ok(pool),
}
}

最後,在 Rust translator 函式 read_f64_py 中,我們從 create_pool(num_threads)?.install(...) 內部呼叫 read_no_alloc(很好的 Rust 函式)。這將所有 Rayon 函式限制為我們設定的 num_threads。

[...]
create_pool(num_threads)?.install(|| {
read_no_alloc(
filename,
[...]
)
})?;
[...]

規則 8:將 nice 動態型別 Python 函式翻譯成 nice Rust 泛型函式

nice Python read 函式的使用者可以指定返回的 NumPy 陣列的 dtype(int8、float32 或 float64)。從這個選擇中,該函式查詢適當的 Rust translator 函式(read_i8(_py)、read_f32(_py) 或 read_f64(_py)),然後呼叫該函式。

def read(
[...]
dtype: Optional[Union[type, str]] = "float32",
[...]
)
[...]
if dtype == np.int8:
reader = read_i8
elif dtype == np.float64:
reader = read_f64
elif dtype == np.float32:
reader = read_f32
else:
raise ValueError(
f"dtype'{val.dtype}'not known, only"
+ "'int8', 'float32', and 'float64' are allowed."
)
reader(
str(self.filepath),
[...]
)

三個 Rust translator 函式呼叫相同的 Rust 函式,即在 src/lib.rs 中定義的 read_no_alloc。以下是 translator 函式 read_64 (又稱 read_64_py) 的相關部分:

#[pyfn(m)]
#[pyo3(name = "read_f64")]
fn read_f64_py(
[...]
val: &PyArray2<f64>,
num_threads: usize,
) -> Result<(), PyErr> {
[...]
let mut val = unsafe { val.as_array_mut() };
[...]
read_no_alloc(
[...]
f64::NAN,
&mut val,
)
[...]
}

我們在 src/lib.rs 中定義了 niceread_no_alloc 函式。也就是說,該函式適用於具有正確特徵的任何型別的 TOut 。其程式碼的相關部分如下所示:

fn read_no_alloc<TOut: Copy + Default + From<i8> + Debug + Sync + Send>(
filename: &str,
[...]
missing_value: TOut,
val: &mut nd::ArrayViewMut2<'_, TOut>,
) -> Result<(), BedErrorPlus> {
[...]
}

在 nice Python、translator Rust 和 nice Rust 中組織程式碼,可以讓我們為 Python 使用者提供動態型別的程式碼,同時仍能用 Rust 編寫出漂亮的通用程式碼。

規則 9:建立 Rust 和 Python 測試

你可能只想編寫會呼叫 Rust 的 Python 測試。但是,你還應該編寫 Rust 測試。新增 Rust 測試使你可以互動地執行測試和互動地除錯。Rust 測試還為你以後得到 Rust 版本的包提供了途徑。在示例專案中,兩組測試都從  bed_reader/tests/data 讀取測試檔案。

在可行的情況下,我還建議編寫函式的純 Python 版本,然後就可以使用這些慢速 Python 函式來測試快速 Rust 函式的結果。

最後,關於 CI 指令碼,例如 bed-reader/ci.yml,應該同時執行 Rust 和 Python 測試。

原文連結:https://towardsdatascience.com/nine-rules-for-writing-python-extensions-in-rust-d35ea3a4ec29

© THE END 

轉載請聯絡本公眾號獲得授權

投稿或尋求報道:[email protected]