厭倦了C++,CS&ML博士用Rust重寫Python擴展,還總結(jié)了九條規(guī)則
Python 是數(shù)據(jù)科學(xué)家最流行的編程語言之一,其內(nèi)部集成了高質(zhì)量分析庫,包括 NumPy、SciPy、自然語言工具包等,這些庫中的許多都是用 C 和 C++ 實現(xiàn)的。
然而,C 和 C++ 兼容性差,且本身不提供線程安全。有研究者開始轉(zhuǎn)向 Rust,重寫 C++ 擴展。
擁有 CS 與機器學(xué)習(xí)博士學(xué)位的 Carl M. Kadie,通過更新 Python 中生物信息學(xué)軟件包 Bed-Reader,為研究者帶來了在 Rust 中編寫 Python 擴展的九個規(guī)則。以下是原博客的主要內(nèi)容。
一年前,我厭倦了我們軟件包 Bed-Reader 的 C++ 擴展,我用 Rust 重寫了它,令人高興的是,得到的新擴展和 C/C++ 一樣快,但具有更好的兼容性和安全性。一路走來,我學(xué)會了這九條規(guī)則,可以幫助你創(chuàng)建更好的擴展代碼,這九條規(guī)則包括:
- 創(chuàng)建一個包含 Rust 和 Python 項目的單獨存儲庫
- 使用 maturin & PyO3 在 Rust 中創(chuàng)建 Python-callable translator 函數(shù)
- 讓 Rust translator 函數(shù)調(diào)用 nice Rust 函數(shù)
- 在 Python 中預(yù)分配內(nèi)存
- 將 nice Rust 錯誤處理翻譯 nice Python 錯誤處理
- 多線程與 Rayon 和 ndarray::parallel,返回任何錯誤
- 允許用戶控制并行線程數(shù)
- 將 nice 動態(tài)類型 Python 函數(shù)翻譯成 nice Rust 泛型函數(shù)
- 創(chuàng)建 Rust 和 Python 測試
其中,文中提到的 nice 這個詞是指使用最佳實踐和原生類型創(chuàng)建。換句話說:在代碼頂部,編寫 nice Python 代碼;在中間,用 Rust 編寫 translator 代碼;在底部,編寫 nice Rust 代碼。結(jié)構(gòu)如下圖所示:
上述策略看似顯而易見,但遵循它可能會很棘手。本文提供了有關(guān)如何遵循每條規(guī)則的實用建議和示例。
我在 Bed-Reader 進行了實驗,Bed-Reader 是一個 Python 包,用于讀取和寫入 PLINK Bed Files,這是一種在生物信息學(xué)中用于存儲 DNA 數(shù)據(jù)的二進制格式。Bed 格式的文件可以達(dá)到 TB。Bed-Reader 讓用戶可以快速、隨機地訪問數(shù)據(jù)的子集。它在用戶選擇的 int8、float32 或 float64 中返回一個 NumPy 數(shù)組。
我希望 Bed-Reader 擴展代碼具有以下特點:
- 比 Python 快;
- 兼容 NumPy;
- 可以進行數(shù)據(jù)并行多線程處理;
- 與執(zhí)行數(shù)據(jù)并行多線程的所有其他包兼容;
- 安全。
我們最初的 C++ 擴展兼具速度快、與 NumPy 兼容,以及使用 OpenMP 進行數(shù)據(jù)并行多線程等特點。遺憾的是,OpenMP 運行時庫 (Runtime library),存在 Python 包兼容版本問題。
Rust 提供了 C++ 擴展帶來的優(yōu)勢。除此之外,Rust 通過提供沒有運行時庫的數(shù)據(jù)并行多線程解決了運行時兼容性問題。此外,Rust 編譯器還能保證線程安全。
在 Rust 中創(chuàng)建 Python 擴展需要許多設(shè)計決策。根據(jù)我使用 Bed-Reader 的經(jīng)驗,以下是我的使用規(guī)則。
規(guī)則 1:創(chuàng)建一個包含 Rust 和 Python 項目的單獨存儲庫
下表顯示了如何布局文件:
使用 Rust 常用的‘cargo new’命令創(chuàng)建 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 設(shè)置在 pyproject.toml 中(如果不是,則在 pytest.ini 等文件中)。Python 代碼位于子文件夾 bed_reader 中。
最后,我們使用 GitHub 操作來構(gòu)建、測試和準(zhǔn)備部署。該腳本位于 .github/workflows/ci.yml 中。
規(guī)則 2:使用 maturin & PyO3在 Rust 中創(chuàng)建 Python-callable translator 函數(shù)
Maturin 是一個 PyPi 包,可通過 PyO3 構(gòu)建和發(fā)布 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;
規(guī)則 3:Rust translator 函數(shù)調(diào)用 nice Rust 函數(shù)
在 src/lib.rs 定義了 nice Rust 函數(shù),這些函數(shù)將完成包的核心工作。它們能夠輸入和輸出標(biāo)準(zhǔn) Rust 類型并嘗試遵循 Rust 最佳實踐。例如,對于 Bed-Reader 包, read_no_alloc 是一個 nice Rust 函數(shù),用于從 PLINK Bed 文件讀取和返回值。
然而,Python 不能直接調(diào)用這些函數(shù)。因此,在文件 src/python_module.rs 中定義 Python 可以調(diào)用的 Rust translator 函數(shù),下面為 translator 函數(shù)示例:
- #[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_indexiid_index = iid_index.readonly();
- let sid_indexsid_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(())
- }
該函數(shù)將文件名、一些與文件大小相關(guān)的整數(shù)以及兩個一維 NumPy 數(shù)組作為輸入,這些數(shù)組指示要讀取數(shù)據(jù)的哪個子集。該函數(shù)從文件中讀取值并填充 val,這是一個預(yù)先分配的二維 NumPy 數(shù)組。
注意
將 Python NumPy 1-D 數(shù)組轉(zhuǎn)換為 Rust slices,通過:
- let iid_indexiid_index = iid_index.readonly();
- let ii = &iid_index.as_slice()?;
將 Python NumPy 2d 數(shù)組轉(zhuǎn)換為 2-D Rust ndarray 對象,通過:
- let mut val = unsafe { val.as_array_mut() };
調(diào)用 read_no_alloc,這是 src/lib.rs 中一個 nice Rust 函數(shù),它將完成核心工作。
規(guī)則 4:在 Python 中預(yù)分配內(nèi)存
在 Python 中為結(jié)果預(yù)分配內(nèi)存簡化了 Rust 代碼。在 Python 端,在 bed_reader/_open_bed.py 中,我們可以導(dǎo)入 Rust translator 函數(shù):
- from .bed_reader import [...] read_f64 [...]
然后定義一個 nice Python 函數(shù)來分配內(nèi)存、調(diào)用 Rust translator 函數(shù)并返回結(jié)果。
- def read([...]):
- [...]
- val = np.zeros((len(iid_index), len(sid_index)), orderorder=order, dtypedtype=dtype)
- [...]
- reader = read_f64
- [...]
- reader(
- str(self.filepath),
- iid_count=self.iid_count,
- sid_count=self.sid_count,
- count_a1=self.count_A1,
- iid_indexiid_index=iid_index,
- sid_indexsid_index=sid_index,
- valval=val,
- num_threadsnum_threads=num_threads,
- )
- [...]
- return val
規(guī)則 5:將 nice Rust 錯誤處理翻譯 nice Python 錯誤處理
為了了解如何處理錯誤,讓我們在 read_no_alloc( src/lib.rs 中 nice Rust 函數(shù))中跟蹤兩個可能的錯誤。
示例錯誤 1:來自標(biāo)準(zhǔn)函數(shù)的錯誤。如果 Rust 的標(biāo)準(zhǔn) File::open 函數(shù)找不到文件或無法打開文件, 在這種情況下,如下一行中的? 將導(dǎo)致函數(shù)返回一些 std::io::Error 值。
- let mut buf_reader = BufReader::new(File::open(filename)?);
為了定義一個可返回這些值的函數(shù),我們可以給函數(shù)一個返回類型 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 函數(shù) read_f64_py 來返回 PyErr;其次,實現(xiàn)了一個從 BedErrorPlus 到 PyErr 的轉(zhuǎn)換器。轉(zhuǎn)換器使用正確的錯誤消息創(chuàng)建正確的 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:特定于函數(shù)的錯誤。如果 nice 函數(shù) read_no_alloc 可以打開文件,但隨后意識到文件格式錯誤怎么辦?它應(yīng)該引發(fā)一個自定義錯誤,如下所示:
- 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,標(biāo)準(zhǔn)錯誤和自定義錯誤的結(jié)果都屬于帶有信息性錯誤消息的特定錯誤類型。
規(guī)則 6:多線程與 Rayon 和 ndarray::parallel,返回任何錯誤
Rust Rayon crate 提供了簡單且輕量級的數(shù)據(jù)并行多線程。ndarray::parallel 模塊將 Rayon 應(yīng)用于數(shù)組。通常的模式是跨一個或多個 2D 數(shù)組的列(或行)并行化。面對的一個挑戰(zhàn)是從并行線程返回任何錯誤消息。我將重點介紹兩種通過錯誤處理并行化數(shù)組操作的方法。以下兩個示例都出現(xiàn)在 Bed-Reader 的 src/lib.rs 文件中。
方法 1:par_bridge().try_for_each
Rayon 的 par_bridge 將順序迭代器變成了并行迭代器。如果遇到錯誤,使用 try_for_each 方法可以盡快停止所有處理。
這個例子中,我們遍歷了壓縮(zip)在一起的兩個 things:
- DNA 位置的二進制數(shù)據(jù);
- 輸出數(shù)組的列。
然后,按順序讀取二進制數(shù)據(jù),但并行處理每一列的數(shù)據(jù)。我們停止了任何錯誤。
- [... 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!宏允許并行地通過一個或多個壓縮在一起的數(shù)組或數(shù)組片段。在我看來,這非常具有可讀性。但是,它不直接支持錯誤處理。因此,我們可以通過將任何錯誤保存到結(jié)果列表來添加錯誤處理。
下面是一個效用函數(shù)的例子。完整的效用函數(shù)從三個計數(shù)和總和(count and sum)數(shù)組計算統(tǒng)計量(均值和方差),并且并行工作。如果在數(shù)據(jù)中發(fā)現(xiàn)錯誤,則將該錯誤記錄在結(jié)果列表中。在完成所有處理之后,檢查結(jié)果列表是否有錯誤。
- [...]
- 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 提供了許多其他不錯的數(shù)據(jù)并行處理方法。
規(guī)則 7:允許用戶控制并行線程數(shù)
為了更好地使用用戶的其他代碼,用戶必須能夠控制每個函數(shù)可以使用的并行線程數(shù)。
在下面這個 nice Python read 函數(shù)中,用戶可以得到一個可選的 num_threadsargument。如果用戶沒有設(shè)置它,Python 會通過這個函數(shù)設(shè)置它:
- 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。這個輔助函數(shù)從 num_threads 構(gòu)造一個 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 函數(shù) read_f64_py 中,我們從 create_pool(num_threads)?.install(...) 內(nèi)部調(diào)用 read_no_alloc(很好的 Rust 函數(shù))。這將所有 Rayon 函數(shù)限制為我們設(shè)置的 num_threads。
- [...]
- create_pool(num_threads)?.install(|| {
- read_no_alloc(
- filename,
- [...]
- )
- })?;
- [...]
規(guī)則 8:將 nice 動態(tài)類型 Python 函數(shù)翻譯成 nice Rust 泛型函數(shù)
nice Python read 函數(shù)的用戶可以指定返回的 NumPy 數(shù)組的 dtype(int8、float32 或 float64)。從這個選擇中,該函數(shù)查找適當(dāng)?shù)?Rust translator 函數(shù)(read_i8(_py)、read_f32(_py) 或 read_f64(_py)),然后調(diào)用該函數(shù)。
- 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 函數(shù)調(diào)用相同的 Rust 函數(shù),即在 src/lib.rs 中定義的 read_no_alloc。以下是 translator 函數(shù) read_64 (又稱 read_64_py) 的相關(guān)部分:
- #[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 函數(shù)。也就是說,該函數(shù)適用于具有正確特征的任何類型的 TOut 。其代碼的相關(guān)部分如下所示:
- 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 中組織代碼,可以讓我們?yōu)?Python 用戶提供動態(tài)類型的代碼,同時仍能用 Rust 編寫出漂亮的通用代碼。
規(guī)則 9:創(chuàng)建 Rust 和 Python 測試
你可能只想編寫會調(diào)用 Rust 的 Python 測試。但是,你還應(yīng)該編寫 Rust 測試。添加 Rust 測試使你可以交互地運行測試和交互地調(diào)試。Rust 測試還為你以后得到 Rust 版本的包提供了途徑。在示例項目中,兩組測試都從 bed_reader/tests/data 讀取測試文件。
在可行的情況下,我還建議編寫函數(shù)的純 Python 版本,然后就可以使用這些慢速 Python 函數(shù)來測試快速 Rust 函數(shù)的結(jié)果。
最后,關(guān)于 CI 腳本,例如 bed-reader/ci.yml,應(yīng)該同時運行 Rust 和 Python 測試。
原文鏈接:https://towardsdatascience.com/nine-rules-for-writing-python-extensions-in-rust-d35ea3a4ec29
【本文是51CTO專欄機構(gòu)“機器之心”的原創(chuàng)譯文,微信公眾號“機器之心( id: almosthuman2014)”】