[Перевод] Как при помощи Rust в 150 раз ускорить код на Python26.02.2023 00:46

Python — довольно простой в освоении язык, по сравнению с некоторыми другими языками код на нём пишется очень быстро. Но в жертву приносится скорость выполнения кода.

Перепишем часть Python-кода в Rust и импортируем этот код в виде пакета Python впроект. Получим сверхбыстрый пакет Python, который сможем импортировать и использовать, как любой другой пакет. В качестве бонуса добавим многопроцессорность и в итоге получим функцию, которая примерно в 150 раз быстрее обычного кода на Python.

Обзор

Проблему решим в 6 шагов:

1. Решим вопрос о том, почему функция медленная.
2. Подготовим проект.
3. Перепишем функцию в Rust.
4. Скомпилируем код на Rust и разместим его в пакете Python.
5. Импортируем пакет Python в проект.
6. Выполним бенчмарк чистого Python и функции на Rust.

Пакет maturin скомпилирует Rust-код и преобразует его в готовый к работе пакет Python.

1. Решим вопрос о том, почему функция медленная

Важно понять, почему функция работает медленно. Давайте представим, что проекту требуется функция подсчёта количества простых чисел в диапазоне между двумя другими числами:

def primecounter_py(range_from:int, range_til:int) -> (int, int):
  """ Returns the number of found prime numbers using range"""
  check_count = 0
  prime_count = 0
  range_from = range_from if range_from >= 2 else 2
  for num in range(range_from, range_til + 1):
    for divnum in range(2, num):
      check_count += 1
      if ((num % divnum) == 0):
        break
    else:
      prime_count += 1
  return prime_count, check_count

Пожалуйста, обратите внимание на то, что:

1. На самом деле указывать количество проверок не обязательно, но это позволяет сравнить Python и Rust далее.
2. Код обеих функций далёк от оптимизированного. Здесь важно показать, что с помощью Rust можно ускорить небольшие фрагменты Python, а ещё сравнить производительность языков.

Функция primecounter_py (10, 20) вернёт 4 (11, 13, 17 и 19 — простые числа) и количество выполненных функцией проверок на простое число. Небольшие диапазоны выполняются очень быстро, но на больших вы увидите снижение производительности:

range      milliseconds
1-1K                  4
1-10K               310
1-25K              1754
1-50K              6456
1-75K             14019
1-100K            24194

Чем больше диапазон, тем меньше скорость функции.

Почему primecounter_py работает медленно?

Код может быть медленным по многим причинам: из-за ввода-вывода, ожидания API, оборудования или архитектуры Python как языка. У нас последний случай. Подход к обработке переменных в Python делает язык очень простым, но он страдает небольшим снижением скорости, которое очевидно, когда приходится выполнять много вычислений. С другой стороны, эта функция очень подходит для оптимизации с помощью Rust.

Проблема в конкурентности?

Множество проблем со скоростью можно решить выполнением нескольких задач одновременно. Для разделения всех задач на несколько ядер можно использовать несколько процессоров, но мы придерживаемся оптимизации на Rust, ведь так мы сможем добиться распределения выполнения более быстрой функции.

Задачи с большим количеством операций ввода-вывода, например ожидание API, можно оптимизировать при помощи потоков. Чтобы увеличить скорость выполнения таких задач, ознакомьтесь с этой статьей или статьёй ниже о том, как задействовать несколько процессоров.

2. Подготовим проект

Ниже устанавливаем зависимости и создаём все файлы и папки для кода на Rust и его сборки в пакет.

a) создание venv

Создайте виртуальную среду и активируйте её. Затем установите maturin; он поможет преобразовать код Rust в пакет Python:

python -m venv venv
source venv/bin/activate
pip install maturin

б) файлы и папки Rust

Каталог my_rust_module будет содержать код на Rust:

mkdir my_rust_module
cd my_rust_module

в) инициализация maturin

Вызовем метод .init. Вы увидите несколько вариантов на выбор, из них выберите pyo3. Пакет maturin создаст файлы и папки, структуру которых вы увидите ниже:

my_folder
 |- venv
 |- my_rust_module
   |- .github
   |- src
    |- lib.rs
   |- .gitignore
   |- Cargo.toml
   |- pyproject.toml

Самый важный файл здесь — это /my_rust_module/src/lib.rs. Именно этот файл будет содержать код, который мы собираемся превратить в пакет Python.

Обратите внимание, что maturin создал конфигурацию проекта — файл Cargo.toml. Кроме конфигурации, этот файл содержит все зависимости (например, requirements.txt). Я отредактировал его, чтобы он выглядел вот так:

[package]
name = "my_rust_module"
version = "0.1.0"
edition = "2021"
# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
[lib]
name = "my_rust_module"
crate-type = ["cdylib"]
[dependencies]
pyo3 = { version = "0.17.3", features = ["extension-module"] }

3. Перепишем функцию в Rust

Теперь мы готовы воссоздать нашу функцию Python в Rust. Не будем слишком глубоко погружаться в синтаксис Rust, а сосредоточимся на том, как заставить код на Rust работать с Python. Сначала создадим функцию на чистом Rust, затем поместим её в пакет, который сможем импортировать и использовать в Python.

Если вы никогда не видели код Rust, то приведённый ниже код может немного сбить с толку. Самое главное, что функция primecounter — это чистый Rust; она не имеет ничего общего с Python.

Откройте /my_rust_module/src/lib.rs и вставьте в него этот код:

use pyo3::prelude::*;

#[pyfunction]
fn primecounter(range_from:u64, range_til:u64) -> (u32, u32) {
 /* Returns the number of found prime numbers between [range_from] and [range_til] """ */
 let mut prime_count:u32 = 0;
 let mut check_count:u32 = 0;
 let _from:u64 = if range_from < 2 { 2 } else { range_from };
 let mut prime_found:bool;

  for num in _from..=range_til {
    prime_found = false;
    for divnum in 2..num {
      check_count += 1;
      if num % divnum == 0 {
        prime_found = true;
        break;
      }
    }
    if !prime_found {
      prime_count += 1;
    }
  }
  return (prime_count, check_count)
}

/// Put the function in a Python module
#[pymodule]
fn my_rust_module(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(primecounter, m)?)?;
    Ok(())
}

Пройдёмся по самому важному в коде:

1. Функция primecounter — это чистый Rust.
2. Она декорирована #[pyfunction]. Декоратор указывает на то, что мы хотим преобразовать её в функцию Python.
3. Функция my_rust_module упаковывает код на Rust в модуль Python: cоздаётся pymodule.

4. Скомпилируем код на Rust и разместим его в пакете Python

Эта часть может показаться самой сложной, но maturin облегчает работу. Просто выполните эту команду:

maturin build --release

Команда компилирует весь код Rust, оборачивает его в пакет Python и записывает пакет в каталогyour_project_dir/my_rust_module/target/wheels. Далее мы установим наш Wheel-пакет.

5. Импортируем пакет Python в проект

Установить созданный Wheel-пакет можно командой pip install:

pip install target/wheels/my_rust_module-0.1.0-cp39-cp39-manylinux_2_28_x86_64.whl

Теперь можно рабоать с модулем:

import my_rust_module

primecount, eval_count = my_rust_module.primecounter(range_from=0, range_til=500)
# returns 95 22279

6. Выполним бенчмарк чистого Python и функции на Rust

Давайте вызовем обе версии функций с несколькими аргументами:

range   Py ms   py e/sec    rs ms   rs e/sec
1-1K        4      17.6M     0.19       417M
1-10K     310      18.6M       12       481M
1-25K    1754      18.5M       66       489M
1-50K    6456      18.8M      248       488M
1-75K   14019      18.7M      519       505M
1-100K  24194      18.8M      937       485M

Они возвращают результат и количество вычисленных ими чисел. В приведённом выше обзоре вы видите, что, когда дело доходит до вычислений в секунду, Rust превосходит Python в 27 раз.

7. Многопроцессорность

Код ниже распределяет числа, которые нам нужно вычислить, на все ядра процессора:

# batch size is determined by the range divided over the amount of available CPU's 
batch_size = math.ceil((range_til - range_from) / mp.cpu_count())

# The lines below divide the ranges over all available CPU's. 
# A range of 0 - 10 will be divided over 4 cpu's like:
# [(0, 2), (3, 5), (6, 8), (9, 9)]
number_list = list(range(range_from, range_til))
number_list = [number_list[i * batch_size:(i + 1) * batch_size] for i in range((len(number_list) + batch_size - 1) // batch_size)]
number_list_from_til = [(min(chunk), max(chunk)) for chunk in number_list]

primecount = 0
eval_count = 0
with mp.Pool() as 
    results = mp_pool.starmap(my_rust_module.primecounter, number_list_from_til)
    for _count, _evals in results:
        primecount += _count
        eval_count += _evals

Снова попробуем найти все простые числа в диапазоне от 0 до 100 000. Теперь это означает, что нужно выполнить почти 500 млн. проверок. Как видите, Rust выполняет эти проверки за 0,88 секунды. При многопроцессорной обработке процесс завершается за 0,16 секунды — это в 5,5 раз быстрее, то есть 2,8 млрд. вычислений в секунду.

            calculations     duration    calculations/sec
rust:            455.19M    882.03 ms          516.1M/sec
rust MP:         455.19M    160.62 ms            2.8B/sec

В сравнении с первоначальной функцией Python с единственным процессором мы увеличили количество вычислений с 18,8 млн. до 2,8 млрд. в секунду. Это означает, что наша функция теперь примерно в 150 раз быстрее.

© Habrahabr.ru