[Перевод] Python без Python: как запускать код где угодно

Команда Python for Devs подготовила перевод статьи о том, как можно компилировать Python в быстрые, кроссплатформенные исполняемые файлы без изменения исходного кода. Автор подробно разбирает архитектуру компилятора, объясняет, зачем «понижать» Python до C++, как типы позволяют «приручить» динамику языка и почему эмпирическая оптимизация даёт лучший результат, чем ручной тюнинг.

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

Эту идею уже пробовали реализовать в самых разных формах: рантаймы (Jython, RustPython), DSL (Numba, PyTorch) и даже совершенно новые языки программирования (Mojo). Но по причинам, к которым мы вернёмся позже в этой статье, нам нужно было решение, которое:

1. Компилирует Python полностью заранее, без каких-либо модификаций кода;

2. Работает без Python-интерпретатора или любого другого интерпретатора;

3. Имеет минимальные накладные расходы по сравнению с «чистыми» программами на C или C++;

4. И, что важнее всего, может запускаться где угодно — на сервере, десктопе, мобильных устройствах и в вебе.

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

Контейнеры — неправильный способ распространения ИИ

Мой путь в исследованиях ИИ начался примерно в 2018 году, в те времена, когда это ещё называли глубокое обучение (Deep Learning). Я взял год академического отпуска и только что вышел из своего первого стартапа — венчурного proptech-проекта, где был сооснователем и который позже был приобретён. Одной из самых интересных проблем, с которыми я столкнулся тогда, было редактирование изображений для объявлений о продаже недвижимости. Каждый месяц фотограф по недвижимости отдавал на ручную обработку в Photoshop и Lightroom тысячи фотографий домов, прежде чем они попадали в региональный MLS или на Zillow.

Я объединился со старым другом, и мы решили построить полностью автоматизированный редактор изображений, используя новый класс моделей компьютерного зрения — генеративно-состязательные сети (GAN). Мы обучали собственные архитектуры моделей на наших датасетах, а затем тщательно тестировали их, чтобы убедиться, что всё работает корректно. Но когда пришло время отдать эти ИИ-модели в руки наших дизайн-партнёров, мы просто застряли. Большую часть времени я тратил на попытки упаковать наши модели во что-то, что можно было бы очень легко распространять. И после месяцев борьбы с Dockerfile и сторонними сервисами мне стало абсолютно ясно: контейнеры — это неправильная единица распространения для ИИ-нагрузок.

Чтобы понять почему, нужно заглянуть внутрь контейнера. Контейнеры — это всего лишь самодостаточные файловые системы Linux с изоляцией выполнения и управлением ресурсами. Поэтому, разворачивая нашу ИИ-модель в виде контейнера, мы упаковывали туда код инференса, веса модели, все зависимости Python-пакетов, сам Python-интерпретатор и прочее необходимое ПО — по сути, снимок полноценной операционной системы Linux.

Но что, если вместо самодостаточной операционной системы сделать самодостаточный исполняемый файл, который запускает только нашу ИИ-модель и ничего лишнего? Выгоды здесь были бы существенными. Мы могли бы поставлять гораздо более компактные контейнеры, которые запускаются заметно быстрее, потому что в них не пришлось бы включать ненужные Python-пакеты, сам Python-интерпретатор и весь прочий мусор, обычно попадающий в контейнер. Но, что ещё важнее, такие исполняемые файлы можно было бы запускать не только на наших Linux-серверах — их можно было бы запускать где угодно.

Arm64, Apple и Unity: как всё начиналось

Я начал программировать в одиннадцать лет — благодаря отцу, который категорически отказался покупать мне PlayStation 2, опасаясь, что у меня просядут оценки. Из упрямства, унаследованного и от него, и от мамы, я решил: если он не собирается покупать мне игровую консоль, значит, я просто буду делать игры сам¹. Мне повезло наткнуться на игровой движок, который был интуитивно понятным, позволял один раз собрать проект и запускать его где угодно и, что самое важное, был бесплатным. Это был Unity Engine.

В конце 2013 года Apple представила iPhone 5S — первое устройство компании с тогда ещё относительно новой архитектурой набора инструкций armv8-a. В отличие от предыдущих устройств, это была 64-битная архитектура на ARM. Она позволяла адресовать гораздо больший объём памяти и давала целый ряд приростов производительности. В результате Apple довольно быстро обязала компилировать все новые приложения под arm64.

Unity с его огромной экосистемой разработчиков оказался в состоянии хаоса. Чтобы понять почему, нужен небольшой контекст о том, как работает Unity. Поскольку Unity — это игровой движок, объекты в игре могут быть запрограммированы с собственным поведением. Для описания этого поведения Unity выбрал C#. Но C# не компилируется напрямую в объектный код, поэтому для выполнения во время работы ему нужна виртуальная машина (звучит знакомо?). Unity использовал для этого Mono, но Mono не поддерживал arm64.

Unity отправился в путь, который я до сих пор считаю его величайшим инженерным достижением: IL2CPP. Как следует из названия, компилятор IL2CPP принимает на вход байткод Common Intermediate Language (то есть промежуточное представление, генерируемое компилятором C#), а затем генерирует эквивалентный исходный код на C++. Имея C++-код, его уже можно скомпилировать практически под любую платформу — от GPU Nvidia и WebAssembly до Apple Silicon и всего, что между ними.

Компилятор IL2CPP в Unity преобразует C#-код в C++-код. Затем этот C++-код можно скомпилировать для запуска где угодно.

Мы решили построить ровно то же самое — но для Python.

Набрасываем архитектуру компилятора Python

На высоком уровне компилятор должен был:

1. Принимать на вход обычный Python-код без каких-либо изменений;

2. Трассировать его, чтобы построить граф промежуточного представления (IR);

3. Понижать IR до исходного кода на C++;

4. Компилировать этот C++-код для запуска на разных платформах и архитектурах.

Прежде чем углубляться дальше, у вас может возникнуть вопрос: зачем вообще сначала генерировать C++? Почему не перейти напрямую от IR к объектному коду?

Возвращаясь к тому, с чего мы начинали, основной фокус Muna — вычислительно тяжёлые приложения, прежде всего инференс ИИ. Если вы работали в этой области, то наверняка знакомы с такими технологиями, как CUDA, MLX, TensorRT и им подобными. Но помимо них существует ещё огромное количество фреймворков, библиотек и даже недокументированных ISA, которые приложения могут использовать для ускорения всего — от умножения матриц до задач компьютерного зрения.

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

Создание символического трассировщика для Python

Первый шаг в построении нашего компилятора — сделать символический трассировщик. Задача трассировщика — принять на вход Python-функцию и выдать граф промежуточного представления (IR), который полностью фиксирует поток управления внутри этой функции.

Наши самые первые прототипы опирались на символический трассировщик PyTorch FX, появившийся в PyTorch 2.0. Их трассировщик был построен на базе PEP 523 — возможности CPython, которая позволяла разработчикам на C переопределять то, как интерпретатор вычисляет фреймы байткода. Я не буду углубляться в детали — это само по себе инженерное чудо, — но в двух словах PEP 523 позволил команде PyTorch зарегистрировать хук, который мог записывать каждый вызов функции по мере того, как интерпретатор её выполнял:

К сожалению, у TorchFX было два существенных недостатка, из-за которых нам пришлось написать собственный трассировщик. Во-первых, как только вы подключаетесь к интерпретатору CPython, чтобы записывать выполнение вашей PyTorch-функции, эту функцию нужно действительно запустить. Для PyTorch это не было проблемой, потому что можно вызывать функцию с так называемыми «фейковыми тензорами»: у них корректные типы данных, формы и устройства, но память под данные не выделяется. Более того, такой подход — запускать функцию ради трассировки — отлично укладывался в то, как работали их устаревшие API сериализации (torch.jit и torch.onnx).

Поскольку нам нужна была возможность компилировать произвольные Python-функции, из которых лишь малая часть (или вообще никакая) могла относиться к PyTorch, нам потребовался бы похожий механизм, чтобы разработчики предоставляли нам входные данные для трассировки. Но, в отличие от PyTorch, мы не могли сделать «фейковое изображение», «фейковую строку» или «фейковое что угодно». Для нас это оказалось тупиком.

Вторая проблема заключалась в том, что даже когда мы создавали фейковые данные в качестве входов для трассировщика TorchFX, он умел записывать только операции PyTorch. Нам пришлось бы серьёзно перерабатывать и расширять трассировщик, чтобы он поддерживал трассировку через произвольные функции сотен и тысяч Python-библиотек. Поэтому мы решили строить трассировщик, который будет захватывать Python-функцию, разбирая её абстрактное синтаксическое дерево (AST). Возьмём пример функции:

Наш трассировщик сначала извлёк бы AST примерно так:

Далее он пошагово проходил бы по дереву, разрешал все вызовы функций (то есть определял, к какой исходной библиотеке относится каждый вызов), а затем генерировал проприетарный формат IR. Сейчас наш символический трассировщик поддерживает статический анализ (через разбор AST), частичное вычисление исходного Python-кода, интроспекцию значений во время выполнения (через песочницу) и многое другое. Но, как ни странно, это самая неинтересная часть нашего компиляторного конвейера.

Понижение в C++ через распространение типов

Вот здесь всё становится по-настоящему интересным. Python — динамический язык, поэтому переменные могут иметь любой тип, и эти типы легко меняются:

C++ же, напротив, — строго типизированный язык: у переменных есть чётко определённые, неизменяемые типы, которые должны быть известны в момент объявления. На первый взгляд может показаться, что связать эти два мира — задача практически неразрешимая. Однако есть ключевое наблюдение, которое мы можем использовать:

Когда мы вызываем Python-функцию с конкретными входными данными, мы можем однозначно определить типы всех промежуточных переменных внутри этой функции²:

Если мы знаем, что входные параметры x и y — это экземпляры float, то знаем и то, что результирующий тип их умножения (то есть tmp_1) однозначно определяется тем, что возвращает функция operator.mul. Но как определить operator.mul и получить её возвращаемый тип? Здесь и появляется C++³.

Из этого следует, что tmp_1 обязательно имеет тип float. Тот же процесс можно повторить для операции сложения (tmp_1 + z), чтобы получить финальный результат.

На этом этапе стоит на минуту остановиться и осмыслить, что именно мы уже построили:

1. Мы можем взять Python-функцию и сгенерировать промежуточное представление (IR), которое полностью описывает её поведение;

2. Затем, используя информацию о типах параметров и C++-реализацию Python-оператора (например, operator.mul), мы можем полностью определить тип первой промежуточной переменной в Python-функции;

3. Мы можем повторить пункт (2) для всех последующих промежуточных переменных, пока не распространим типы по всей функции целиком.

Инициализация процесса распространения типов

Один момент, который стоит разобрать подробнее, — откуда мы берём исходную информацию о типах параметров, чтобы запустить процесс распространения типов. В примере выше как мы узнаём, что x, y и z — это экземпляры float?

После прототипирования нескольких разных подходов мы остановились на PEP 484, который добавил в язык Python поддержку аннотаций типов. Сам Python полностью игнорирует эти аннотации, поскольку они не используются во время выполнения⁴. И хотя они решили задачу инициализации распространения типов, у них оказалось два серьёзных недостатка. Во-первых, они конфликтуют с нашей ключевой целью проектирования, потому что их использование требует от разработчиков слегка модифицировать Python-код⁵:

Код выглядит не сильно иначе, и некоторые считают, что аннотации типов в принципе помогают писать более качественный Python-код (в Muna мы делаем их обязательными). Вторая проблема заключалась в том, что для проектирования простого и модульного интерфейса потребления скомпилированных функций нам пришлось бы ограничить количество различных входных типов, которые могут использовать разработчики⁶. В итоге мы решили, что это разумный компромисс с хорошей эргономикой.

Построение библиотеки операторов на C++

К этому моменту вы, возможно, заметили очевидную проблему в нашем дизайне: нам нужно писать реализации на C++ для потенциально десятков или даже сотен тысяч Python-функций из разных библиотек. К счастью, всё это гораздо менее сложно, чем может показаться на первый взгляд. Рассмотрим функцию ниже:

При компиляции cosecant мы видим, что внутри есть вызовы функций sin и reciprocal. Наш компилятор сначала проверяет, может ли он протрассировать каждый вызов. В случае sin у нас нет её определения (есть лишь импорт), поэтому мы не можем пройтись по ней трассировкой⁷. Это образует листовой узел, который нам приходится реализовывать вручную на C++. Вызов reciprocal мы можем протрассировать, поэтому делаем это и получаем для него граф IR. Затем его можно понизить и использовать в других местах вызова.

Ключевое наблюдение здесь в том, что большинство Python-функций, с которыми сталкивается наш компилятор, состоят из сравнительно небольшого набора элементарных функций. То, что создаёт огромное разнообразие кода в реальном мире, — это не уникальное количество таких элементарных функций, а различные способы их композиции.

Тем не менее можно возразить, что таких элементарных функций в разных библиотеках потенциально тысячи — и это будет на 100 % верно. К счастью, у нас теперь есть отличный инструмент, который делает эту задачу простой: генерация кода с помощью ИИ.

Современные LLM способны писать проверяемо корректный и высокопроизводительный код на самых разных языках программирования. Поэтому мы выстраиваем инфраструктуру, которая ограничивает генерируемый ими код, тестирует его на корректность и берёт на себя вспомогательную логику — например, управление зависимостями и условную компиляцию. На данный момент мы уже использовали ИИ для генерации реализаций сотен Python-функций из популярных библиотек, таких как Numpy, OpenCV и PyTorch.

Оптимизация производительности через исчерпывающий перебор

Последняя тема, которую стоит обсудить, — оптимизация производительности. Наиболее распространённые подходы здесь включают написание ручного кода (например, на ассемблере или PTX), использование гетерогенных ускорителей (GPU, NPU), эвристический выбор алгоритмов во время выполнения (например, поиск алгоритма свёртки в ArmCL и cuDNN) или различные комбинации всего перечисленного.

Из нашего прошлого опыта построения сверхнизколатентных конвейеров компьютерного зрения для встраиваемых систем мы усвоили очень горький урок⁸: эффективная оптимизация производительности всегда эмпирична. Задержка конкретной операции на конкретном железе зависит от такого количества факторов, что единственный способ узнать наверняка — просто протестировать каждый возможный подход. Единственная причина, по которой инженерные команды этого не делают, — практическая невозможность: пришлось бы переписывать код десятки или сотни раз, а затем тестировать каждый вариант… но постойте!

Ранее мы разбирали, как распространяем типы по Python-функции с помощью C++-оператора. Чего я тогда не упомянул, так это того, что мы используем не один C++-оператор, а столько, сколько можем написать (кхм, сгенерировать). То есть вместо вот такого:

Мы не просто генерируем одну C++-программу из Python-функции.

Каждый путь от входа к результату — это уникальная программа, гарантированно корректная относительно исходной Python-функции. При этом каждый C++-оператор (цветные прямоугольники) может быть реализован с использованием разных алгоритмов, библиотек и даже аппаратных ускорителей. Давайте разберём конкретный пример:

Функция выше изменяет размер входного изображения до 64×64 с билинейной интерполяцией, используя библиотеку torchvision. При компиляции этой функции под Apple Silicon (macOS, iOS или visionOS) у нас есть целый спектр подходов и библиотек на выбор, включая следующие:

Выше приведён лишь небольшой фрагмент возможных вариантов, поскольку способов реализовать билинейное масштабирование на Apple Silicon существует множество (например, с использованием GPU или Neural Engine). Ключевой момент в том, что мы можем сгенерировать как можно больше таких реализаций (благодаря генерации кода с помощью LLM), а затем выпустить скомпилированные программы, использующие каждую из них — без каких-либо ограничений. Так, в приведённом примере пользовательская Python-функция будет сгенерирована как четыре уникальные программы только для Apple Silicon. В наших тестах на реальных нагрузках мы видели, как одна Python-функция превращалась почти в 200 уникальных программ для 9 целевых платформ компиляции.

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

1. Мы можем оптимизировать код чисто эмпирически — не делая никаких предположений о том, какая реализация будет быстрее, и без отдельного этапа тюнинга производительности после генерации кода; мы просто распространяем все доступные скомпилированные бинарники, собираем телеметрию и на её основе определяем самый быстрый вариант;

2. Мы получаем эффект масштаба сети: поскольку C++-операторы разделяются между тысячами скомпилированных функций и эти функции запускаются на сотнях тысяч уникальных устройств всех наших пользователей, у нас накапливается огромный объём данных, который позволяет оптимизировать каждый фрагмент генерируемого кода.

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

Проектирование пользовательского интерфейса для компилятора

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

Разработчик может просто добавить @compile к своей Python-функции, чтобы указать точку входа для компиляции. После этого он компилирует функцию⁹ и все её зависимости с помощью CLI:

Мы полюбили парадигму декораторов, наблюдая за тем, как разработчики предпочитают описывать сложную инфраструктуру в виде кода¹⁰. Кроме того, это был привычный форм-фактор для Python-экосистемы, что подтверждается его использованием в Numba и PyTorch. Благодаря декоратору наш CLI мог находить функцию — точку входа компиляции — и использовать её как отправную точку для обхода всего остального кода зависимостей (как собственного кода разработчика, так и сторонних пакетов, установленных через pip или uv).

Декоратор @compile также служил основной точкой настройки для разработчиков, компилирующих свои функции. Помимо обязательных параметров — тега (который однозначно идентифицирует функцию на нашей платформе) и описания — разработчики могли указать описание песочницы для воссоздания локального окружения разработки (например, установка Python-пакетов, загрузка файлов), а также метаданные, помогающие компилятору при генерации кода (например, выполнение инференса PyTorch с использованием ONNXRuntime, TensorRT, CoreML, IREE, QNN и других).

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

Заключительные мысли

Откровенно говоря, нам до сих пор трудно поверить, что всё это действительно работает. При этом в компиляторе пока частично реализованы или вовсе отсутствуют многие стандартные возможности Python: исключения, лямбда-выражения, рекурсивные функции и классы. Общая нить, связывающая все эти пробелы, — наша система распространения типов. Хотя распространение типов хорошо работает для простых функций с единичными типами параметров и возвращаемых значений, оно требует дополнительных усилий для составных типов (например, union) и типов более высокого порядка (например, классов и лямбда-выражений).

Ещё одна важная область, над которой мы всё ещё работаем, — это опыт отладки. Хорошая новость в том, что мы гарантируем: после компиляции Python-код разработчиков будет работать так, как ожидается, полностью снимая с них ответственность за отладку во время выполнения. Это похоже на то, как разработчики, использующие Docker или другие технологии контейнеризации, просто ожидают, что всё «будет работать» — почти никто не отлаживает слои Docker-образов. Плохая новость заключается в том, что, поскольку мы позволяем запускать Python-код где угодно, нам приходится решать задачи написания исключительно безопасного кода и сбора детализированных, символизированных трассировок в случае, если какая-то функция выбрасывает исключение. Всё усложняется тем, что ради минимального размера и максимальной скорости скомпилированных бинарников мы компилируем сгенерированный код с полными оптимизациями, неизбежно удаляя ценную отладочную информацию.

Впрочем, не всё было сложно — во многом потому, что развивающийся стандарт C++ стал для нас огромным подспорьем. Muna не существовала бы без C++20: наша генерация кода активно опирается на std::span, concepts и, что особенно важно, корутины. И мы с нетерпением ждём широкой поддержки C++23, потому что используем std::generator для потоковой обработки, <stdfloat> для поддержки float16_t и bfloat16_t, а <stacktrace> — для поддержки исключений Python.

И напоследок: если вы сейчас разворачиваете embedding-модели или модели детекции объектов в своей организации, либо если вам просто показалась интересной эта работа, мы будем рады пообщаться. Нам хотелось бы, чтобы больше разработчиков использовали компилятор для задач и программ, которые мы сами ещё не запускали, и мы всегда рады знакомству с теми, кому по душе приземлённый, низкоуровневый мир Python, C++ и всего, что находится между ними.

Русскоязычное сообщество про Python

Друзья! Эту статью подготовила команда Python for Devs — канала, где каждый день выходят самые свежие и полезные материалы о Python и его экосистеме. Подписывайтесь, чтобы ничего не пропустить!