厭倦了C++，CS&ML博士用Rust重寫(xiě)Python擴(kuò)展，還總結(jié)了九條規(guī)則

作者：機(jī)器之心編譯 2022-01-04 05:51:03

擁有 CS 與機(jī)器學(xué)習(xí)博士學(xué)位的 Carl M. Kadie，通過(guò)更新 Python 中生物信息學(xué)軟件包 Bed-Reader，為研究者帶來(lái)了在 Rust 中編寫(xiě) Python 擴(kuò)展的九個(gè)規(guī)則。

Python 是數(shù)據(jù)科學(xué)家最流行的編程語(yǔ)言之一，其內(nèi)部集成了高質(zhì)量分析庫(kù)，包括 NumPy、SciPy、自然語(yǔ)言工具包等，這些庫(kù)中的許多都是用 C 和 C++ 實(shí)現(xiàn)的。

然而，C 和 C++ 兼容性差，且本身不提供線程安全。有研究者開(kāi)始轉(zhuǎn)向 Rust，重寫(xiě) C++ 擴(kuò)展。

擁有 CS 與機(jī)器學(xué)習(xí)博士學(xué)位的 Carl M. Kadie，通過(guò)更新 Python 中生物信息學(xué)軟件包 Bed-Reader，為研究者帶來(lái)了在 Rust 中編寫(xiě) Python 擴(kuò)展的九個(gè)規(guī)則。以下是原博客的主要內(nèi)容。

一年前，我厭倦了我們軟件包 Bed-Reader 的 C++ 擴(kuò)展，我用 Rust 重寫(xiě)了它，令人高興的是，得到的新擴(kuò)展和 C/C++ 一樣快，但具有更好的兼容性和安全性。一路走來(lái)，我學(xué)會(huì)了這九條規(guī)則，可以幫助你創(chuàng)建更好的擴(kuò)展代碼，這九條規(guī)則包括：

創(chuàng)建一個(gè)包含 Rust 和 Python 項(xiàng)目的單獨(dú)存儲(chǔ)庫(kù)
使用 maturin & PyO3 在 Rust 中創(chuàng)建 Python-callable translator 函數(shù)
讓 Rust translator 函數(shù)調(diào)用 nice Rust 函數(shù)
在 Python 中預(yù)分配內(nèi)存
將 nice Rust 錯(cuò)誤處理翻譯 nice Python 錯(cuò)誤處理
多線程與 Rayon 和 ndarray::parallel，返回任何錯(cuò)誤
允許用戶控制并行線程數(shù)
將 nice 動(dòng)態(tài)類(lèi)型 Python 函數(shù)翻譯成 nice Rust 泛型函數(shù)
創(chuàng)建 Rust 和 Python 測(cè)試

其中，文中提到的 nice 這個(gè)詞是指使用最佳實(shí)踐和原生類(lèi)型創(chuàng)建。換句話說(shuō)：在代碼頂部，編寫(xiě) nice Python 代碼;在中間，用 Rust 編寫(xiě) translator 代碼;在底部，編寫(xiě) nice Rust 代碼。結(jié)構(gòu)如下圖所示：

上述策略看似顯而易見(jiàn)，但遵循它可能會(huì)很棘手。本文提供了有關(guān)如何遵循每條規(guī)則的實(shí)用建議和示例。

我在 Bed-Reader 進(jìn)行了實(shí)驗(yàn)，Bed-Reader 是一個(gè) Python 包，用于讀取和寫(xiě)入 PLINK Bed Files，這是一種在生物信息學(xué)中用于存儲(chǔ) DNA 數(shù)據(jù)的二進(jìn)制格式。Bed 格式的文件可以達(dá)到 TB。Bed-Reader 讓用戶可以快速、隨機(jī)地訪問(wèn)數(shù)據(jù)的子集。它在用戶選擇的 int8、float32 或 float64 中返回一個(gè) NumPy 數(shù)組。

我希望 Bed-Reader 擴(kuò)展代碼具有以下特點(diǎn)：

比 Python 快;
兼容 NumPy;
可以進(jìn)行數(shù)據(jù)并行多線程處理;
與執(zhí)行數(shù)據(jù)并行多線程的所有其他包兼容;
安全。

我們最初的 C++ 擴(kuò)展兼具速度快、與 NumPy 兼容，以及使用 OpenMP 進(jìn)行數(shù)據(jù)并行多線程等特點(diǎn)。遺憾的是，OpenMP 運(yùn)行時(shí)庫(kù) (Runtime library)，存在 Python 包兼容版本問(wèn)題。

Rust 提供了 C++ 擴(kuò)展帶來(lái)的優(yōu)勢(shì)。除此之外，Rust 通過(guò)提供沒(méi)有運(yùn)行時(shí)庫(kù)的數(shù)據(jù)并行多線程解決了運(yùn)行時(shí)兼容性問(wèn)題。此外，Rust 編譯器還能保證線程安全。

在 Rust 中創(chuàng)建 Python 擴(kuò)展需要許多設(shè)計(jì)決策。根據(jù)我使用 Bed-Reader 的經(jīng)驗(yàn)，以下是我的使用規(guī)則。

規(guī)則 1：創(chuàng)建一個(gè)包含 Rust 和 Python 項(xiàng)目的單獨(dú)存儲(chǔ)庫(kù)

下表顯示了如何布局文件：

使用 Rust 常用的‘cargo new’命令創(chuàng)建 Cargo.toml 和 src/lib.rs 文件。Python 沒(méi)有 setup.py 文件。相反，Cargo.toml 包含 PyPi 包信息，例如包的名稱(chēng)、版本號(hào)、 README 文件的位置等。要在沒(méi)有 setup.py 的情況下工作，pyproject.toml 必須包含：

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

一般來(lái)說(shuō)，Python 設(shè)置在 pyproject.toml 中(如果不是，則在 pytest.ini 等文件中)。Python 代碼位于子文件夾 bed_reader 中。

最后，我們使用 GitHub 操作來(lái)構(gòu)建、測(cè)試和準(zhǔn)備部署。該腳本位于 .github/workflows/ci.yml 中。

規(guī)則 2：使用 maturin & PyO3在 Rust 中創(chuàng)建 Python-callable translator 函數(shù)

Maturin 是一個(gè) PyPi 包，可通過(guò) PyO3 構(gòu)建和發(fā)布 Python 擴(kuò)展。PyO3 是一個(gè) Rust crate，用于在 Rust 中編寫(xiě) Python 擴(kuò)展。

在 Cargo.toml 中，包含這些 Rust 依賴項(xiàng)：

[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 類(lèi)型并嘗試遵循 Rust 最佳實(shí)踐。例如，對(duì)于 Bed-Reader 包， read_no_alloc 是一個(gè) 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ù)以及兩個(gè)一維 NumPy 數(shù)組作為輸入，這些數(shù)組指示要讀取數(shù)據(jù)的哪個(gè)子集。該函數(shù)從文件中讀取值并填充 val，這是一個(gè)預(yù)先分配的二維 NumPy 數(shù)組。

注意

將 Python NumPy 1-D 數(shù)組轉(zhuǎn)換為 Rust slices，通過(guò)：

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

將 Python NumPy 2d 數(shù)組轉(zhuǎn)換為 2-D Rust ndarray 對(duì)象，通過(guò)：

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

調(diào)用 read_no_alloc，這是 src/lib.rs 中一個(gè) nice Rust 函數(shù)，它將完成核心工作。

規(guī)則 4：在 Python 中預(yù)分配內(nèi)存

在 Python 中為結(jié)果預(yù)分配內(nèi)存簡(jiǎn)化了 Rust 代碼。在 Python 端，在 bed_reader/_open_bed.py 中，我們可以導(dǎo)入 Rust translator 函數(shù)：

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

然后定義一個(gè) nice Python 函數(shù)來(lái)分配內(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 錯(cuò)誤處理翻譯 nice Python 錯(cuò)誤處理

為了了解如何處理錯(cuò)誤，讓我們?cè)?read_no_alloc( src/lib.rs 中 nice Rust 函數(shù))中跟蹤兩個(gè)可能的錯(cuò)誤。

示例錯(cuò)誤 1：來(lái)自標(biāo)準(zhǔn)函數(shù)的錯(cuò)誤。如果 Rust 的標(biāo)準(zhǔn) File::open 函數(shù)找不到文件或無(wú)法打開(kāi)文件，在這種情況下，如下一行中的? 將導(dǎo)致函數(shù)返回一些 std::io::Error 值。

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

為了定義一個(gè)可返回這些值的函數(shù)，我們可以給函數(shù)一個(gè)返回類(lèi)型 Result<(), BedErrorPlus>。我們定義 BedErrorPlus 時(shí)要包含所有 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 錯(cuò)誤處理，但 Python 不理解它。因此，在 src/python_module.rs 中，我們要進(jìn)行翻譯。首先，定義 translator 函數(shù) read_f64_py 來(lái)返回 PyErr;其次，實(shí)現(xiàn)了一個(gè)從 BedErrorPlus 到 PyErr 的轉(zhuǎn)換器。轉(zhuǎn)換器使用正確的錯(cuò)誤消息創(chuàng)建正確的 Python 錯(cuò)誤類(lèi)(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()), 
        } 
    } 
}

示例錯(cuò)誤 2：特定于函數(shù)的錯(cuò)誤。如果 nice 函數(shù) read_no_alloc 可以打開(kāi)文件，但隨后意識(shí)到文件格式錯(cuò)誤怎么辦?它應(yīng)該引發(fā)一個(gè)自定義錯(cuò)誤，如下所示：

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

BedError::IllFormed 類(lèi)型的自定義錯(cuò)誤在 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), 
[...] 
}

其余的錯(cuò)誤處理與示例錯(cuò)誤 1 中的相同。

最后，對(duì)于 Rust 和 Python，標(biāo)準(zhǔn)錯(cuò)誤和自定義錯(cuò)誤的結(jié)果都屬于帶有信息性錯(cuò)誤消息的特定錯(cuò)誤類(lèi)型。

規(guī)則 6：多線程與 Rayon 和 ndarray::parallel，返回任何錯(cuò)誤

Rust Rayon crate 提供了簡(jiǎn)單且輕量級(jí)的數(shù)據(jù)并行多線程。ndarray::parallel 模塊將 Rayon 應(yīng)用于數(shù)組。通常的模式是跨一個(gè)或多個(gè) 2D 數(shù)組的列(或行)并行化。面對(duì)的一個(gè)挑戰(zhàn)是從并行線程返回任何錯(cuò)誤消息。我將重點(diǎn)介紹兩種通過(guò)錯(cuò)誤處理并行化數(shù)組操作的方法。以下兩個(gè)示例都出現(xiàn)在 Bed-Reader 的 src/lib.rs 文件中。

方法 1：par_bridge().try_for_each

Rayon 的 par_bridge 將順序迭代器變成了并行迭代器。如果遇到錯(cuò)誤，使用 try_for_each 方法可以盡快停止所有處理。

這個(gè)例子中，我們遍歷了壓縮(zip)在一起的兩個(gè) things：

DNA 位置的二進(jìn)制數(shù)據(jù);
輸出數(shù)組的列。

然后，按順序讀取二進(jìn)制數(shù)據(jù)，但并行處理每一列的數(shù)據(jù)。我們停止了任何錯(cuò)誤。

[... 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!宏允許并行地通過(guò)一個(gè)或多個(gè)壓縮在一起的數(shù)組或數(shù)組片段。在我看來(lái)，這非常具有可讀性。但是，它不直接支持錯(cuò)誤處理。因此，我們可以通過(guò)將任何錯(cuò)誤保存到結(jié)果列表來(lái)添加錯(cuò)誤處理。

下面是一個(gè)效用函數(shù)的例子。完整的效用函數(shù)從三個(gè)計(jì)數(shù)和總和(count and sum)數(shù)組計(jì)算統(tǒng)計(jì)量(均值和方差)，并且并行工作。如果在數(shù)據(jù)中發(fā)現(xiàn)錯(cuò)誤，則將該錯(cuò)誤記錄在結(jié)果列表中。在完成所有處理之后，檢查結(jié)果列表是否有錯(cuò)誤。

[...] 
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 提供了許多其他不錯(cuò)的數(shù)據(jù)并行處理方法。

規(guī)則 7：允許用戶控制并行線程數(shù)

為了更好地使用用戶的其他代碼，用戶必須能夠控制每個(gè)函數(shù)可以使用的并行線程數(shù)。

在下面這個(gè) nice Python read 函數(shù)中，用戶可以得到一個(gè)可選的 num_threadsargument。如果用戶沒(méi)有設(shè)置它，Python 會(huì)通過(guò)這個(gè)函數(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。這個(gè)輔助函數(shù)從 num_threads 構(gòu)造一個(gè) Rayon ThreadPool 對(duì)象。

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ù)限制為我們?cè)O(shè)置的 num_threads。

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

規(guī)則 8：將 nice 動(dòng)態(tài)類(lèi)型 Python 函數(shù)翻譯成 nice Rust 泛型函數(shù)

nice Python read 函數(shù)的用戶可以指定返回的 NumPy 數(shù)組的 dtype(int8、float32 或 float64)。從這個(gè)選擇中，該函數(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), 
       [...] 
     )

三個(gè) Rust translator 函數(shù)調(diào)用相同的 Rust 函數(shù)，即在 src/lib.rs 中定義的 read_no_alloc。以下是 translator 函數(shù) read_64 (又稱(chēng) 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, 
     ) 
     [...] 
}

我們?cè)?src/lib.rs 中定義了 niceread_no_alloc 函數(shù)。也就是說(shuō)，該函數(shù)適用于具有正確特征的任何類(lèi)型的 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 用戶提供動(dòng)態(tài)類(lèi)型的代碼，同時(shí)仍能用 Rust 編寫(xiě)出漂亮的通用代碼。

規(guī)則 9：創(chuàng)建 Rust 和 Python 測(cè)試

你可能只想編寫(xiě)會(huì)調(diào)用 Rust 的 Python 測(cè)試。但是，你還應(yīng)該編寫(xiě) Rust 測(cè)試。添加 Rust 測(cè)試使你可以交互地運(yùn)行測(cè)試和交互地調(diào)試。Rust 測(cè)試還為你以后得到 Rust 版本的包提供了途徑。在示例項(xiàng)目中，兩組測(cè)試都從 bed_reader/tests/data 讀取測(cè)試文件。

在可行的情況下，我還建議編寫(xiě)函數(shù)的純 Python 版本，然后就可以使用這些慢速 Python 函數(shù)來(lái)測(cè)試快速 Rust 函數(shù)的結(jié)果。

最后，關(guān)于 CI 腳本，例如 bed-reader/ci.yml，應(yīng)該同時(shí)運(yùn)行 Rust 和 Python 測(cè)試。

原文鏈接：https://towardsdatascience.com/nine-rules-for-writing-python-extensions-in-rust-d35ea3a4ec29

【本文是51CTO專(zhuān)欄機(jī)構(gòu)“機(jī)器之心”的原創(chuàng)譯文，微信公眾號(hào)“機(jī)器之心( id: almosthuman2014)”】

戳這里，看該作者更多好文

責(zé)任編輯：趙寧寧來(lái)源： 51CTO專(zhuān)欄