Как ускорить Python с помощью C-расширений

Привет, Хабр! Я – Игорь Алимов, ведущий разработчик группы Python в МТС Digital, работаю над продуктами Smart Rollout и B2B-портал. В этой статье я расскажу о том, как писать быстрый код на Python с использованием C-расширений и победить GIL.

На мысли об ускорении Python меня натолкнула статья о языках программирования будущего. Список перспективных (по мнению автора) языков я приводить не буду, но скажу, что языку Python в будущем было отказано. В числе недостатков Python автор выделил низкую производительность и наличие GIL. Действительно, Python не слишком быстрый и имеет блокировку, которая разрешает одновременное выполнение только одного потока инструкций. Что с этим делать? Переучиваться на Java/Go/Rust/____(нужное подчеркнуть/вписать)? Погодите, есть другие способы ускорить Python и нивелировать его недостатки.

Что это за метод?

• Пишем первую реализацию на Python. На производительность не обращаем внимания, наша цель – получить результат, чтобы в дальнейшем было, с чем сравнивать. Если на этом этапе производительность нас устраивает – задача выполнена, в противном случае переходим ко второму пункту

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

• Производим рефакторинг кода так, чтобы выделить в виде отдельной функции, класса или модуля код, на который уходит больше всего времени. На этом этапе производим оптимизацию кода, используя параллельное исполнение и другие приемы, помогающие уменьшить время выполнения. При этом контролируем правильность результата и время. Если результат нас по-прежнему не устраивает – переходим к следующему пункту

• Переписываем проблемный код на C

Почему именно на C? Для Python существуют порядка десяти различных способов использования нативного кода, но такой метод обеспечит наибольшую производительность при переключении между Python и нативным кодом. Язык C – это interlingua всех языков программирования. Ядро Linux, большинство системных библиотек, GTK и еще много чего написаны на нем.

Даже если библиотека написана на каком-то другом языке программирования, двоичный интерфейс (так называемое ABI) у нее C-подобный. Эталонная реализация Python также написана на C, о чем говорит ее название – CPython. Главный минус – особенности языка C: арифметика указателей, ручное управление памятью, отсутствие средств метапрограммирования. С другой стороны, С – это высокая производительность, прямой доступ к памяти и аппаратуре, стандартный интерфейс. Применение обоих языков программирования позволит использовать их сильные стороны и значительно уменьшит влияние недостатков.

Конкретный пример использования С-расширений

Постановка задачи

Есть некоторая начальная строка, нам необходимо добавить к ней один или несколько случайно выбранных символов из заданного набора, чтобы хеш полной строки имел особый вид. В нашем конкретном случае будем использовать sha256 и хеш полной строки должен начинаться с 8 нулей.

Начальная строка: 'Начальное значение!'

Набор символов для создания случайной строки: punctuation + digits + ascii_letters из модуля string

Хеш: sha256

Ожидаемое начала хеша: '00000000'

Написание первого варианта

Криптофункция хеш по определению является односторонней, поэтому единственный практический способ – это полный перебор всех сочетаний символов из заданного набора, сначала по одному символу, потом по два и так далее. Это позволит пронумеровать все возможные сочетания и написать функцию, которая по номеру возвращает сочетание символов. Первая реализация находится в файле prototype.py. Это простая однопоточная реализация, вот ее особенности:

• Использование лога для вывода

• Использование argparse для парсинга аргументов командной строки

• Вводится понятие раунда, по умолчанию 100000000 хешей, по окончании раунда выводится производительность раунда в kH/s (кило хешей в секунду)

• Основную работу выполняет функция mining, которая обрабатывает раунд целиком

• Функция get_value по номеру сочетания символов возвращает строку с этими символами

• Для расчета хешей используется библиотечный модуль hashlib.sha256

Лог работы первого варианта находится в файле prototype.log.

Хеш 00000000331cb4111b0fb7fff9a9014aa45376e25b59516ff57e0789f86d98ce от строки 'Начальное значение![JBYW' был получен в 71 раунде, время 3 часа 46 минут 56.715 секунд, средняя производительность 527.490 kH/s

Анализ первого варианта, рефакторинг

И без профилирования понятно, что больше всего времени занимает функция mining. Для ускорения работы используем параллельное вычисление. Помня про GIL, применяем для этого модуль multiprocessing. Каждый раунд делим на несколько параллельно действующих обработчиков (workers), причем новый раунд начинаем, не дожидаясь момента, когда закончат работу все workers предыдущего раунда. Исходный код находится в файле prototype_multi.py. Особенности этой версии:

• Добавлена функция form_task, которая формирует задание для очередного раунда

• Функция mining отправляет свой результат в очередь

• В функции main строки 123-133 ставят задание на исполнение

• Строка 134 в функции main ожидает завершения workers

• Строки 135-150 функции main производит анализ возвращаемых результатов

Результат выполнения второго варианта находится в файле prototype-multi.log. Итоги те же, время выполнения 1 час 1 минута 55.01 секунда, средняя производительность 1933.423 kH/s.

Рекомендации для написания C-функций

Написанное ниже относится к реализации Python на CPython версии 3.9. В других версиях CPython возможны изменения в API. Библиотека собиралась компилятором GCC 11.2.0, сборка и тестирование проводились в Linux. В сборке нативных расширений под Windows есть особенности, которые не отражены в этой статье. Основная документация – Extending and Embedding the Python Interpreter и Python/C API Reference Manual.

Прототип C-функции

Если нужно, чтобы C-функция была видна в Python, она должна иметь следующий прототип:

#define PY_SSIZE_T_CLEAN
#include <Python.h>

static PyObject *
foo_bar(PyObject *self, PyObject *args)
{
    ...
}

В первой строке мы определяем макрос PY_SSIZE_T_CLEAN, затем подключаем заголовочный файл Python.h – в нем содержатся прототипы функций, макросы, определения типов данных, которые потребуются при написании кода. Ключевое слово static указывает на область видимости функции, которой является единица компиляции. Проще говоря, функция не видна за пределами текущего файла. Возвращает функция указатель на PyObject, в нем будут возвращены результаты работы функции.

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

У функции два аргумента, это указатель на модуль, в котором расположена функция, и указатель на аргументы функции. Мы рассматриваем функции с позиционными аргументами. Да, есть возможность задать функцию с ключевыми аргументами, но в этой статье мы о ней говорить не будем. Если функция является методом объекта, то первый указатель будет ссылаться на объект.

Поговорим немного о структуре PyObject. Все сущности (числа, строки, составные типы, объекты, функции, модули) в Python представлены расширениями этой структуры. Правильнее было бы говорить, что они наследуют эту структуру, но в языке C нет наследования. Любые Python-сущности – расширения структуры PyObject и расположены в heap (свободная область памяти между кодом программы и стеком). CPython оперирует указателями на эти объекты – как ссылки, которых в языке C нет. Это порождает особенности работы Python c переменными и данными: мы можем записать в переменную число и строку, но нам не удастся сложить число и строку. А все потому, что переменная имеет тип ссылки на PyObject и в нее можно поместить любое значение. Объекты Python типа int и str не имеют операции +, они обладают различными типами. Эти типы и определяют операции, которые с ними можно проводить, а переменные всегда имеют один и тот же тип – ссылка на PyObject.

Работа с исключениями

При работе C-функции могут возникать невосстановимые ошибки: выделения памяти, открытия файла и прочие. Кажется логичным решением возбудить соответствующее исключение и немедленно прекратить работу функции. Но тут проявляется проблема – в языке C нет исключений.

При написании программ на языке C используется другой подход: каждая функция возвращает признак (обычно это int), если функция выполнилась успешно – возвращается 0, в противном случае – 1. Из этого правила есть исключения, чтобы узнать о них, необходимо изучить документацию на функцию или системный вызов. Также устанавливается специальная переменная errno, локальная для каждого потока выполнения. С ее помощью можно уточнить, что именно пошло не так. В нашем случае мы возвращаем указатель на объект и признаком аварийного завершения будет возврат NULL.

Нужно понимать разницу между Python None и C NULL. Любая функция Python, если мы не возвращаем из нее значения, возвращает None. Это нормальный Python-объект, который имеет адрес, а C-выражение NULL означает отсутствие указателя. Кроме возвращения NULL, наша функция должна установить объект исключения. Делается это следующим образом:

void PyErr_SetString(PyObject *type, const char *message)

type – указатель на тип исключения

message – текст сообщения об ошибке

Есть и общая форма установки исключения:

void PyErr_SetObject(PyObject *type, PyObject *value)

value - указатель на произвольное значение

Стоит заметить, что не всегда эти функции вызываются напрямую. Например, при ошибке выделения памяти проще будет сделать:

if (value == NULL)
    return PyErr_NoMemory();

Функция PyErr_NoMemory сама установит соответствующий объект исключения и вернет NULL, который отдадим вызывающей функции.

Управление памятью

Если не брать в расчет глобальные и локальные переменные потока, в языке C есть два места, где располагаются переменные:

#include <stdlib.h>

void
foo()
{
    int c;                       // Переменная расположена на стеке
    char *buffer = malloc(4096); // Выделяется буфер размером 4096 байт, расположенный в heap, при этом указатель
                                 // на буфер расположен на стеке
    ...
    free(buffer);                // Освобождение памяти 
}

Все переменные, созданные на стеке, автоматически уничтожаются при выходе из функции, при этом их значение теряется. Если в указанном примере опустить вызов free, то при выходе из функции будут утеряно значение buffer, которое является указателем на память в heap и до завершения программы никто не сможет обратиться к этой области. Это явление называется утечкой памяти.

Для программы, которая недолго работает, это не так страшно – при завершении программы вся память будет возвращена операционной системе. Но если у нас утечка памяти в продолжительно работающем сервисе и на каждом запросе к этому сервису мы теряем пару килобайт – через несколько часов память в системе будет полностью израсходована. Тогда к нам в гости придет oom killer и убьет наш процесс. Двойное освобождение памяти даже более критично, чем утечка, оно проводит к немедленному завершению процесса.

Попытка разыменовывать (обратиться к области памяти) неверный указатель также приводит к немедленному завершению процесса (здесь я немного драматизирую ситуацию, так как, возможно, процесс останется жив, но мы точно считаем мусор, который никак не соответствует данным).

Использование ранее освобожденной памяти в языке C носит название неопределенное поведение (undefined behavior, UB). Если повезет – мы считаем данные. Если нет – получим мусор, а попытки записать данные в такую область памяти, скорее всего, приведут к загадочным ошибкам, неправильным результатам, аварийным завершением процесса. Причем ошибки эти проявятся позже, в процессе выполнения программы.

Как в таких, откровенно говоря, спартанских условиях, Python реализует управление памятью? Для этого используются две стратегии: подсчет ссылок и сборщик мусора.

Подсчет ссылок (Reference Counts).

Рассмотрим структуру PyObject, базовую для всех Python-сущностей:

typedef struct _object {
    _PyObject_HEAD_EXTRA
    Py_ssize_t ob_refcnt;
    PyTypeObject *ob_type;
} PyObject;

Если не брать в расчет _PyObject_HEAD_EXTRA, то в структуре PyObject всего два поля: ob_type – указатель на тип объекта и ob_refcnt – счетчик ссылок. В поле ob_refcnt указано количество ссылок на объект. Действуют следующие правила при работе со счетчиком ссылок:

• Объект никому не принадлежит и находится в общем доступе

• Все, кто заинтересован в использовании этого объекта, должны явно увеличить счетчик ссылок

• После того, как объект становится не нужен, следует уменьшить счетчик ссылок

• Не следует напрямую обращаться к полю ob_refcnt, для этого используются макросы Py_INCREF(x) и Py_DECREF(x)

• Если после макроса Py_DECREF(x) счетчик ссылок становится равным 0, объект уничтожается и память освобождается

• Действует правило заимствования – считается, что аргументы функции удерживает вызывающая функция. В возвращаемом результате увеличивает счетчик ссылок создающая его функция и в таком виде его возвращает

Стратегия с подсчетом ссылок не работает в случае циклических ссылок между объектами. Тогда у обоих (или более) объектов счетчик ссылок никогда не будет равен 0, хотя оба объекта никому не нужны. Для разрешения такой ситуации используется сборщик мусора (garbage collection gc), он находит циклические ссылки и удаляет объекты, входящие в цикл, предварительно проверив, что на них никто больше не ссылается.

Если нашей C-функции потребуется память для собственных нужд, можно использовать стандартные функции malloc/free, но Python/C API Reference Manual/Memory Management рекомендует использовать PyMem_RawMalloc/PyMem_RawCalloc/PyMem_RawRealloc/PyMem_RawFree, которые по своим входным и выходным параметрам соответствуют malloc/calloc/realloc/free. Аргументация такая – сборщику мусора проще выполнять свою работу, если он знает о всей используемой памяти.

На этом я заканчиваю теоретическую часть о написании С-расширений в Python. Во второй части этой статьи продемонстрирую конкретный пример реализации функции mining на C. Подписывайтесь на блог МТС, чтобы не пропустить продолжение.

Если у вас есть вопросы или замечания к первой части – пишите их в комментариях, с радостью ответим на них!