Привет, Хабр!

Я хочу рассказать об удивительном событии, о котором я узнал пару месяцев назад. Оказывается, одна популярная python-утилита уже более года распространяется в виде бинарных файлов, которые компилируются прямо из python. И речь не про банальную упаковку каким-нибудь PyInstaller-ом, а про честную Ahead-of-time компиляцию целого python-пакета. Если вы удивлены так же как и я, добро пожаловать под кат.

Объясню, почему я считаю это событие по-настоящему удивительным. Существует два вида компиляции:  Ahead-of-time (AOT), когда весь код компилируется до запуска программы и Just in time compiler (JIT), когда непосредственно компиляция программы под требуемую архитектуру процессора осуществляется во время ее выполнения. Во втором случае первоначальный запуск программы осуществляется виртуальной машиной или интерпретатором.

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

• Ahead-of-time compiler: C, C++, Rust, Kotlin, Nim, D, Go, Dart;

• Just in time compiler: Lua, С#, Groovy, Dart.

В python из коробки нет JIT компилятора, но отдельные библиотеки, предоставляющие такую возможность, существуют давно

Смотря на эту таблицу, можно заметить определенную закономерность: статически типизированные языки находятся в обеих строках. Некоторые даже могут распространяться с двумя версиями компиляторов: Kotlin может исполняться как с JIT JavaVM, так и с AOT Kotlin/Native. То же самое можно сказать про Dart (версии 2). A вот динамически типизированные языки компилируются только JIT-ом, что впрочем вполне логично. 

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

При использовании JIT компиляции типы не очень то и нужны, ведь информация о типах собирается во время работы программы. Поэтому все популярные динамически типизированные языки программирования распространяются именно с JIT компилятором. Но как быть с AOT компиляцией кода, в котором нет типов? Меня очень заинтересовал этот вопрос, и я полез разбираться.

Итак, вернемся к утилите, о которой говорилось в начале статьи. Речь про mypy - наиболее популярный синтаксический анализатор python-кода.

С апреля 2019 года эта утилита распространяется в скомпилированном виде, о чем рассказывается в блоге проекта. А для компиляции используется еще одна утилита от тех же авторов — mypyc. Погуглив немного, я нашел достаточно большую статью “Путь к проверке типов 4 миллионов строк Python-кода” про становление и развитие mypy (на Хабре доступен перевод: часть 1, часть 2, часть 3). Там немного рассказывается о целях создания mypyc: столкнувшись с недостаточной производительностью mypy при разборе крупных python-проектов в Dropbox, разработчики добавили кеширование результатов проверки кода, а затем возможность запуска утилиты как сервиса. Но исчерпав очевидные возможности оптимизации, столкнулись с выбором: переписать все на go или на cython. В результате проект пошел по третьему пути — написание своего AOT python-компилятора.

Дело в том, что для правильной работы mypy и так необходимо построить то же синтаксическое дерево, что и интерпретатору во время исполнения кода. То есть mypy уже “понимает” python, но использует эту информацию только для статистического анализа, а вот mypyc может преобразовывать эту информацию в полноценный бинарный код.

Думаю тут многие решили, что разобрались в вопросе того, как скомпилировать динамически типизированный python-код. Python c версии 3.4 поддерживает аннотацию типов, а mypy как раз и используется для проверки корректности аннотаций. Получается, python как бы уже и не динамически типизированный язык, что позволяет применить AOT компиляцию. Но загвоздка в том, что mypyc может компилировать и неаннотированный код!

Функция bubble_sort

Для примера рассмотрим функцию сортировки “пузырьком”. Файл lib.py:

def bubble_sort(data):
   n = len(data)
   for i in range(n - 1):
       for j in range(n - i - 1):
           if data[j] > data[j + 1]:
               buff = data[j]
               data[j] = data[j + 1]
               data[j + 1] = buff
   return data

У типов нет аннотаций, но это не мешает mypyc ее скомпилировать. Чтобы запустить компиляцию, нужно установить mypyc. Он не распространяется отдельным пакетом, но если у вас установлен mypy, то и mypyc уже присутствует в системе! Запускаем mypyc, следующей командой:

> mypyc lib.py

После запуска в проекте будут созданы следующие директории:

• .mypy_cache  — mypy кэш, mypyc неявно запускает mypy для разбора программы и получения AST;

• build — артефакты сборки;

• lib.cpython-38-x86_64-linux-gnu.so — собственно сборка под целевую платформу. Данный файл представляет из себя готовый CPython Extension.

CPython Extension — встроенный в CPython механизм взаимодействия с кодом, написанным на С/C++. По сути это динамическая библиотека, которую CPython умеет загружать при импорте нашего модуля lib. Через данный механизм осуществляется взаимодействие с модулями, написанными на python.

Компиляция состоит из двух фаз:

1. Компиляция python кода в код С;

2. Компиляция С в бинарный .so файл, для этого mypyc сам запускает gcc (gcc и python-dev также должен быть установлены).

Файл lib.cpython-38-x86_64-linux-gnu.so имеет преимущество перед lib.py при импорте на соответствующей платформе, и исполняться теперь будет именно он.

Ну и давайте сравним производительность модуля до и после компиляции. Для этого создадим файл main.py с кодом запуска сортировки:

import lib


data = lib.bubble_sort(list(range(5000, 0, -1)))

assert data == list(range(1, 5001))

Получим примерно следующие результаты:

Ожидаемо скомпилированный код оказался быстрее (~ в 2 раза), что неплохо, так как для такого результата нам потребовалось запустить лишь одну команду. Хотя от скомпилированного кода привычно ожидаешь большего.

Чтобы ответить на вопрос “как компилируется динамически типизированный код”, придется заглянуть в представление этой функции на С. Но разобрать ее будет достаточно сложно, поэтому давайте попробуем разобраться с примером попроще.  

Функция sum(a, b)

Скомпилируем функцию суммы от двух переменных:

def sum(a, b):
  return a + b

Перед запуском компиляции я ожидал увидеть примерно следующий код на С:

int sum(int a, int b) {
   return a + b;
}

Однако результат оказался cущественно иным (код немного упрощен):

PyObject *CPyDef_sum(PyObject *cpy_r_a, PyObject *cpy_r_b){
    return PyNumber_Add(cpy_r_a, cpy_r_b);
}

Рассмотрим, что тут происходит. Во-первых, так как мы не знаем типы входных переменных, функция в качестве аргументов принимает указатели на объекты класса PyObject, по сути это внутренние CPython структуры. Далее компилятор должен сложить эти объекты, но как, если настоящие типы аргументов неизвестны во время компиляции: это могут быть целые числа, числа с плавающей точкой, списки и вообще не факт, что аргументы можно складывать, тогда нужно вернуть ошибку. И что же делает в этом случае mypyc? 

Как оказалось, все очень просто: он просит CPython самостоятельно сложить эти аргументы. Функция PyNumber_Add — это внутренняя функция СPython, которая доступна из расширения, ведь СPython отлично умеет складывать свои объекты.

Взаимодействие CPython c Extension можно изобразить следующим диалогом:

 — А посчитай-ка мне функцию sum для A, B;

 — Хорошо, но скажи сначала, сколько будет A + B;

 — Будет С;

 — Хорошо, тогда держи ответ - С.

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

Конечно, данный пример выглядит гротескно, но даже несмотря на такую неэффективность, mypyc позволяет добиться существенного прироста производительности, как в примере с сортировкой.

Функция sum(a: int, b: int)

Итак, у нас получилось скомпилировать python, и мы разобрались с тем, как это работает, а также увидели определенную неэффективность полученного результата. Теперь попробуем разобраться в том, как можно это улучшить. Очевидно, что основная проблема заключается во множественном взаимодействии CPython - Extension. Но как это побороть? 

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

def sum(a: int, b: int):
  return a + b

Скомпилированный результат на C (немного очищенный):

PyObject *CPyDef_sum(CPyTagged cpy_r_a, CPyTagged cpy_r_b) {
   CPyTagged cpy_r_r0;
   PyObject *cpy_r_r1;
   cpy_r_r0 = CPyTagged_Add(cpy_r_a, cpy_r_b);
   cpy_r_r1 = CPyTagged_StealAsObject(cpy_r_r0);

   return cpy_r_r1;
}

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

Теперь CPyDef_sum получает на вход не указатели на PyObject, а структуры CPyTagged. Это все еще не int, но уже и не часть CPython, а часть библиотек mypyc, которую он добавляет в скомпилированный код расширения. Для ее инициализации в рантайме сначала проверяется тип, так что теперь функция sum работает только с int и обойти аннотацию не получится.

Далее происходит вызов CPyTaggetAdd вместо PyNumber_Add. Это уже внутренняя функция mypyc. Если заглянуть в код CPyTaggetAdd, то можно понять, что там происходит проверка диапазонов значений a и b, и если они укладываются в int, то происходит простое суммирование, а также проверка на переполнение:

if (likely(CPyTagged_CheckShort(left) && CPyTagged_CheckShort(right))) {
    CPyTagged sum = left + right;

    if (likely(!CPyTagged_IsAddOverflow(sum, left, right))) {
        return sum;
    }
}

Таким образом, наш диалог CPython - Extension превращается из абсурдного в нормальный:

 — А посчитай-ка мне функцию sum для A, B;

 — Хорошо, тогда держи ответ С.

Функция bubble_sort(data: List[int])

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

def bubble_sort(data: List[int]):
    …

Скомпилируем результат и замерим время сортировки:

Итак, мы получили еще двукратное ускорение относительно скомпилированного, не аннотированного кода, и четырехкратное относительно оригинального!

Пара слов о mypyc

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

Также существуют принципиальные ограничения, накладываемые на компилируемый код или на код, взаимодействующий со скомпилированным:

• Принудительная проверка типов в рантайме;

• В компилируемом коде запрещается monkey patching;

• Mypy хранит классы в С структурах для увеличения скорости доступа к атрибутам, но это приводит к проблемам совместимости.

Эти ограничения носят принципиальный характер и являются следствием архитектуры компилятора. Но из них проистекают другие ограничения, например, невозможность использования модуля стандартной библиотеки abc. Помимо этого, есть большая порция недоработок и багов. Чаще всего они приводят к тому, что код gcc отказывается компилировать полученный С код, при этом, чтобы понять настоящую причину ошибки, приходится прокручивать в голове непростую процедуру реверс инжиниринга. Пока резутльт таков, что при компиляции одного из моих проектов, без проблем компилировалось примерно 20 % модулей, зато каких либо проблем при работе с уже скомпилированными модулями я не заметил.

Тем не менее, большая часть улучшений в их Roadmap закрыта, и проект готовится к публичному релизу.

Nuitka

Уже в процессе работы над статьей, я узнал про еще один проект с аналогичными целями. Механизм работы Nuitka сильно напоминает описанный выше. Разница заключается в том, что Nuitka компилирует Python модуль в С++ код, который также собирается в СPython Extension. Дополнительно существует возможность собрать весь проект в один исполняемый файл, тогда уже сам CPython подключается к проекту как динамическая библиотека libpython.

Nuitka пока не учитывает аннотацию типов, поэтому результирующий код и скорость работы не зависят от наличия аннотаций. Полученная же скорость в моем тесте соответствует результату mypy на не аннотированном коде.

Завершение

Недавно один мой коллега высказал мнение, что mypy сильно усложняет ему жизнь: из текста ошибок невозможно понять, “чего он от меня хочет”, а анализатор из PyCharm немного лучше. Теперь я понимаю, что он недооценивает mypy. Так как он намного большее, чем просто синтаксический анализатор. По сути он реализует подмножество языка, потенциал которого в плане оптимизации сильно превосходит обычный python. Поэтому встранивание mypy в пайплайн проекта — инвестиция не только в поиск ошибок, но и будущий перфоманс приложения. Мне очень понравилось, что взаимодействие с CPython осуществляется через механизм расширений интерпретатора, ведь это позволяет сделать выборочную компиляцию наиболее нагруженных модулей, оставив большую часть кода без изменений. Такой путь представляется мне наиболее безопасным (учитывая, что mypyc до сих пор в альфе). Конечно, использовать ли mypyc на продакшене, решать вам, но если вы уже уперлись в потолок по производительности и подумываете о том, чтобы переписать какие-то части на низкоуровневые языки, то стоит попробовать запустить mypyc, тем более, что сделать это просто, если вы уже используете mypy.

P.S.

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

UPD

В комментариях подсказали, об еще одном python компиляторе - Cython, оказывается теперь он умеет компилировать python код напрямую (минуя ручную фазу преобразования кода в cython-код). Судя по замерам cython пока не учитыват аннотацию типов, но время выполнения (real 1.82) оказалось посередине между результатом mypyc на аннотированном и не аннотированном кода. Но возможно ее добавят в будущем.