Python gyorsítása Rusttal: PyO3, maturin és tisztességes mérés

Egy lassú Python-programnál először azt keresd meg, mire megy el az idő. Ha hálózati válaszra vársz, egy számítás Rustba költöztetése kevés segítséget jelenthet. Ha viszont ugyanaz a Python-ciklus sok százezer elemen fut, érdemes egy kis, jól körülhatárolt függvényt kipróbálni natív bővítményként. Ebben a cikkben szomszédos egész számokat hasonlítunk össze, és megszámoljuk az emelkedéseket. Nem írjuk újra az alkalmazást, és nem ígérünk általános gyorsulási szorzót.

A PyO3 biztosítja a Python és Rust közötti kötést; a maturin a buildet és a Python-csomagolást segíti. A kísérőprojekt rögzített PyO3 0.29.3 és maturin 1.12.6 verzióval készült, Rust 1.97.1 és edition 2024 mellett. Ezek a kipróbált verzióink, nem a mindenkori legújabb kiadások ígéretei. A teljes minta az examples/articles/python-pyo3-maturin könyvtárban található.

Először a függvény szerződését határozd meg

A bemenet egész számok listája. Egy növekedés akkor történik, ha az aktuális elem szigorúan nagyobb az előzőnél; az egyenlő érték nem növekedés. Üres vagy egyelemű listánál nulla az eredmény. A Rust-oldal előjeles, 64 bites számokat fogad, ezért a Python tetszőleges pontosságú egész számainak nem minden értéke fér át ezen a határon. A tartományon kívüli értékekre a kötés hibát jelez; ezt külön ellenőrizzük.

A magfüggvény szeletet kér, így önálló Rust-teszttel is vizsgálható. Az összehasonlítás nem kivonással történik: a legnagyobb és legkisebb szám különbsége túlcsordulhatna, miközben a rendezési összehasonlítás biztonságosan elvégzi a feladatot. A windows(2) az egymást követő párokat adja, és a rövid bemenetekhez nem kell külön indexelési ág.

pub fn count_increases(values: &[i64]) -> usize {
    values.windows(2).filter(|pair| pair[1] > pair[0]).count()
}

fn main() {
    assert_eq!(count_increases(&[3, 1, 2, 2, 9]), 2);
    assert_eq!(count_increases(&[]), 0);
}

A Pythonból hívható réteg Vec<i64> értéket vesz át, majd kölcsönadja azt a magfüggvénynek. Ez szándékosan egyszerű, de adatkonverzióval és másolással jár. A kölcsönzésről szóló leckében megértheted, miért elég a belső függvénynek egy szelet. Az ownership és move segít követni, melyik réteg birtokolja az átalakított adatot.

Buildelj elkülönített Python-környezetbe

A projekt Cargo.toml fájljában a könyvtár neve rusthu_kernel, a crate-típus pedig cdylib és rlib. A Python-modul neve ugyanilyen; eltérő nevek helyett érdemes ezt az egyszerű egyezést megtartani. Az rlib lehetővé teszi a Rust-tesztek futtatását is. A Python telepítési környezetét külön virtuális környezetbe helyezzük, hogy ne írjuk felül más projektek függőségeit.

cd examples/articles/python-pyo3-maturin
python3 -m venv .venv
. .venv/bin/activate
python -m pip install maturin==1.12.6
cargo test --locked
maturin develop --release --locked
python verify.py

A fejlesztői build alkalmas hibakeresésre, de nem jó alap teljesítményállításhoz. A fenti parancs ezért optimalizált buildet készít. Kódmódosítás után a bővítményt ismét le kell fordítani; a Python-fájl újraindítása önmagában nem építi újra a natív modult. Telepítési hiba esetén először a használt Python végrehajtható állományát és a virtuális környezetet ellenőrizd, majd a Rust fordítóüzenetét olvasd el.

A helyes eredmény előzze meg a stoppert

A verify.py először 105 determinisztikus példán hasonlítja össze a Python-referenciát és a Rust-eredményt. Szerepel benne üres bemenet, ismétlődő érték, negatív szám és a két szélső egész érték. Ezután külön megpróbál tartományon kívüli számokat átadni, és hibát vár. A Rust-tesztek a magfüggvényt ellenőrzik; a Python-teszt a tényleges nyelvi határt is érinti. A kettő más típusú hibát képes megtalálni.

A mérés tíz, ezer és százezer elemes listákon fut. Mindkét hívási út bemelegítést kap, a bemenet előállítása pedig a mért szakaszon kívül történik. Hét ismétlésből mediánt számítunk, és a kimenet rögzíti az ismétlésenkénti hívásszámot. A Python timeit dokumentációja alapján a mérési környezet és a szóródás értelmezése is számít; egyetlen stopperérték könnyen félrevezet.

A konverzió is része a költségnek

A Rust-ág idejébe a Python-lista Rust-vektorrá alakítása is beleszámít. Ez illik ahhoz a kérdéshez, hogy egy Python-hívó számára megéri-e a csere. A csak Rustban mért ciklusidő más kérdésre válaszolna. A script a Python verzióját és a platformot is kiírja, így az eredmény mellé megmarad a környezet. A gép terhelése, energiamódja és a fordítás módja azonban továbbra is befolyásolhatja a számokat.

A helyi ellenőrzésben a bővítmény valóban lefordult, a Pythonból történő import és az egyezőségi vizsgálat lefutott. Ebből nem következik, hogy egy másik alkalmazás gyorsabb lesz. Kis bemenetnél a hívási és konverziós költség dominálhat; nagy adathalmaznál érdemes a teljes feldolgozási láncot mérni. Ha több függvény között ide-oda másolod ugyanazt a listát, előbb a határfelületet érdemes összevonni.

A következő lépés egy valódi alkalmazásprofil és egy előre megadott elfogadási küszöb: mennyi időt kell megtakarítani, mekkora memóriaköltség fér bele, és mely Python-verziókat kell támogatni. A natív wheel terjesztése platformonként külön feladat. Ez a minta egy helyben működő, ellenőrzött bővítményt és ismételhető mérési eljárást ad; a csomagolási mátrixot és a saját terhelésedet külön kell végigpróbálni.

A példa ellenőrzési környezete

Rust 1.97.1, edition 2024, PyO3 0.29.3, maturin 1.12.6 és Python 3.12.12 környezetben, macOS arm64 gépen ellenőriztük. A két Rust-teszt, az optimalizált bővítmény fordítása, a Python-import, a 105 egyezőségi eset és a két tartományhiba-ellenőrzés sikeres volt. A mérési script hét ismétléssel lefutott; más platformon vagy alkalmazáson mért eredményt ez nem állít.