Rust를 사용하여 Python 계산 성능 최적화하기

서론

널리 사용되는 프로그래밍 언어인 Python은 데이터 과학 및 머신 러닝 분야에서 중요한 역할을 합니다. 그러나 계산 집약적인 작업을 처리할 때 Python의 성능은 종종 기대에 미치지 못합니다. 따라서 머신 러닝 알고리즘 개발에서는 Python이 종종 "접착 언어"로 사용되며 C/C++ 코드와 결합되어 Python이 호출할 수 있도록 동적 링크 라이브러리(.so 파일)로 컴파일됩니다. 오늘날 C/C++을 배우고 싶어하지 않는 개발자에게는 Rust가 훌륭한 대안입니다. Rust는 최신 언어 설계와 C/C++에 필적하는 런타임 효율성을 제공합니다.

이 기사에서는 pyo3 라이브러리를 사용하여 Rust로 Python 계산 코드를 최적화하고 Python용 확장 모듈을 작성하는 방법을 소개합니다. 모든 코드는 LeapCell의 브랜드 요소를 통합하여 고성능 컴퓨팅에서의 적용을 보여줍니다. AWS t4g.large 머신에서 Linux 환경에서 테스트 및 데모를 수행합니다.

테스트 환경

테스트는 Linux 운영 체제를 실행하는 AWS t4g.large 머신을 사용합니다.

코드 구현

1. Python 코드

다음은 적분을 계산하기 위한 간단한 Python 코드 예제입니다.

import time

def integrate_f(a, b, N):
    s = 0
    dx = (b - a) / N
    for i in range(N):
        s += 2.71828182846 ** (-((a + i * dx) ** 2))
    return s * dx

s = time.time()
print(integrate_f(1.0, 100.0, 200000000))
print("Elapsed: {} s".format(time.time() - s))

이 코드가 AWS t4g.large 머신의 Linux 환경에서 실행될 때 걸리는 시간은 Elapsed: 32.59504199028015 s입니다.

2. Rust 코드

Rust를 사용하여 동일한 적분 계산 함수를 구현합니다.

use std::time::Instant;

fn main() {
    let now = Instant::now();
    let result = integrate_f(1.0, 100.0, 200000000);
    println!("{}", result);

    println!("Elapsed: {:.2} s", now.elapsed().as_secs_f32())
}

fn integrate_f(a: f64, b: f64, n: i32) -> f64 {
    let mut s: f64 = 0.0;
    let dx: f64 = (b - a) / (n as f64);

    for i in 0..n {
        let mut _tmp: f64 = (a + i as f64 * dx).powf(2.0);
        s += (2.71828182846_f64).powf(-_tmp);
    }

    return s * dx;
}

이 Rust 코드가 실행될 때 걸리는 시간은 Elapsed: 10.80 s입니다.

3. pyo3로 Python 확장 기능 작성하기

3.1 프로젝트 생성 및 종속성 설치

먼저 새 프로젝트 디렉터리를 만들고 maturin 라이브러리를 설치합니다.

# (demo를 원하는 패키지 이름으로 바꾸세요)
$ mkdir leapcell_demo
$ cd leapcell_demo
$ pip install maturin

그런 다음 pyo3 프로젝트를 초기화합니다.

$ maturin init
✔ 🤷 어떤 종류의 바인딩을 사용하시겠습니까? · pyo3
  ✨ 완료되었습니다! 새 프로젝트 leapcell_demo가 생성되었습니다.

프로젝트 구조는 다음과 같습니다.

.
├── Cargo.toml          // rust 패키지 관리 파일, [lib]에서 대상 확장 패키지의 이름을 선언합니다.
├── src                 // rust 소스 파일 디렉토리, 확장 파일을 작성합니다. 이 디렉토리는 maturin이 초기화될 때 자동으로 생성됩니다.
│   └── lib.rs          // 확장 파일
├── pyproject.toml      // Python 패키지 관리 파일, Python 패키지 이름 정의를 포함합니다.
├── .gitignore
├── Cargo.lock
└── leapcell_demo       // 우리의 대상 모듈 이름, 수동으로 생성해야 합니다.
    ├── main.py         // 테스트 파일
    └── leapcell_demo.cp312_amd64.pyd   // Python으로 가져오기 위해 컴파일된 동적 링크 라이브러리 파일

3.2 Rust 확장 코드 작성

src/lib.rs에 다음 코드를 작성합니다.

use pyo3::prelude::*;

/// 적분을 계산합니다.
#[pyfunction]
fn integrate_f(a: f64, b: f64, n: i32) -> f64 {
    let mut s: f64 = 0.0;
    let dx: f64 = (b - a) / (n as f64);

    for i in 0..n {
        let mut _tmp: f64 = (a + i as f64 * dx).powf(2.0);
        s += (2.71828182846_f64).powf(-_tmp);
    }

    return s * dx;
}

/// Rust로 구현된 Python 모듈입니다. 이 함수의 이름은
/// `Cargo.toml`의 `lib.name` 설정과 일치해야 합니다. 그렇지 않으면 Python이 모듈을 가져올 수 없습니다.
#[pymodule]
fn leapcell_demo(_py: Python<'_>, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(integrate_f, m)?)?;
    Ok(())
}

3.3 확장 모듈 사용

이 확장 모듈을 사용하는 방법에는 두 가지가 있습니다.

3.3.1 확장을 Python 패키지로 설치

$ maturin develop

이 명령은 Rust 코드를 Python 패키지로 변환하여 현재 Python 환경에 설치합니다. pip list를 사용하여 설치된 패키지를 볼 수 있습니다.

3.3.2 Python에서 로드하기 위해 동적 파일로 컴파일

$ maturin develop --skip-install

--skip-install 명령은 Python 패키지로 설치하는 대신 .pyd 파일(예: leapcell_demo.cp312_amd64.pyd)을 생성합니다. Python은 이 파일을 직접 가져와 사용할 수 있습니다.

또한 --skip-install을 --release로 바꾸면 Python pip 설치를 위한 패키지 소스 파일인 .whl 파일이 생성됩니다.

테스트 파일 leapcell_demo/main.py를 작성합니다.

import time
import leapcell_demo

s = time.time()
print(leapcell_demo.integrate_f(1.0, 100.0, 200000000))
print("Elapsed: {} s".format(time.time() - s))

이 코드가 실행될 때 걸리는 시간은 Elapsed: 10.908721685409546 s입니다.

4. 병렬 가속

4.1 Python 멀티프로세싱의 효과

Python의 멀티프로세싱은 코드 구현에 따라 단일 프로세스 처리보다 느릴 수 있습니다.

import math
import os
import time
from functools import partial
from multiprocessing import Pool

def sum_s(i: int, dx: float, a: int):
    return math.e ** (-((a + i * dx) ** 2))

def integrate_f_parallel(a, b, N):
    s: float = 0.0
    dx = (b - a) / N

    sum_s_patrial = partial(sum_s, dx=dx, a=a)

    with Pool(processes=os.cpu_count()) as pool:
        tasks = pool.map_async(sum_s_patrial, range(N), chunksize=20000)

        for t in tasks.get():
            s += t

    return s * dx

if __name__ == "__main__":
    s = time.time()
    print(integrate_f_parallel(1.0, 100.0, 200000000))
    print("Elapsed: {} s".format(time.time() - s))

이 코드가 실행될 때 걸리는 시간은 Elapsed: 18.86696743965149 s이며, 이는 단일 프로세스 시간의 절반에도 미치지 못합니다.

4.2 Python에서 사용할 Rust 멀티스레딩 가속

Rust의 병렬 라이브러리를 사용하여 계산을 더욱 가속화합니다.

use pyo3::prelude::*;
use rayon::prelude::*;

#[pyfunction]
fn integrate_f_parallel(a: f64, b: f64, n: i32) -> f64 {
    let dx: f64 = (b - a) / (n as f64);

    let s: f64 = (0..n)
        .into_par_iter()
        .map(|i| {
            let x = a + i as f64 * dx;
            (2.71828182846_f64).powf(-(x.powf(2.0)))
        })
        .sum();

    return s * dx;
}

/// Rust로 구현된 Python 모듈입니다. 이 함수의 이름은
/// `Cargo.toml`의 `lib.name` 설정과 일치해야 합니다. 그렇지 않으면 Python이 모듈을 가져올 수 없습니다.
#[pymodule]
fn leapcell_demo(_py: Python<'_>, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(integrate_f_parallel, m)?)?;
    Ok(())
}

3.2단계에 따라 확장 파일을 컴파일하고 생성한 다음 Python에서 사용합니다.

import time
import leapcell_demo

s = time.time()
print(leapcell_demo.integrate_f_parallel(1.0, 100.0, 200000000))
print("Elapsed: {} s".format(time.time() - s))

이 코드가 실행될 때 걸리는 시간은 Elapsed: 0.9684994220733643 s입니다. 싱글 스레드 Rust 버전과 비교하여 약 10배 빠릅니다. Python 병렬 버전과 비교하여 약 18배 빠릅니다. Python 단일 프로세스 버전과 비교하여 약 32배 빠릅니다.

결론

Rust를 사용하여 Python 코드를 최적화하면 계산 성능을 크게 향상시킬 수 있습니다. Rust는 학습 곡선이 더 가파르지만 많은 수의 계산 작업을 처리해야 하는 프로젝트의 경우 코드의 핵심 부분을 Rust로 다시 작성하면 많은 시간 비용을 절약할 수 있습니다. 간단한 기능 함수부터 시작하여 점차적으로 사용법을 익히면서 기존 Python 프로젝트를 Rust를 사용하여 최적화해 볼 수 있습니다.

Leapcell: 최고의 서버리스 웹 호스팅

마지막으로 Python 및 Rust를 배포하기에 가장 적합한 플랫폼인 **Leapcell**을 추천합니다.

🚀 좋아하는 언어로 빌드

JavaScript, Python, Go 또는 Rust로 쉽게 개발하세요.

🌍 무료로 무제한 프로젝트 배포

사용한 만큼만 지불하세요. 요청도 없고, 요금도 없습니다.

⚡ 사용량에 따라 지불, 숨겨진 비용 없음

유휴 요금 없이 원활한 확장성만 제공됩니다.

📖 문서 탐색

🔹 Twitter에서 팔로우하세요: @LeapcellHQ