Техника старой школы для современного мира

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

Python — красивый язык программирования, но он не известен своей скоростью.

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

Сегодня я собираюсь проиллюстрировать один из самых необычных методов для этого: динамическая генерация кода, или сокращенно codegen. Это классический метод, который был распространен во многих языках программирования вплоть до 1980-х годов, потерял популярность по причинам, которые мы обсудим ниже, но сегодня его стоит пересмотреть.

Поскольку это нюанс, я собираюсь сначала обсудить, когда эта техника уместна, а затем проиллюстрировать ее подробным реальным примером из Humu, где она ускорила производственный стек почти в 2 раза.

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

Что такое динамическая генерация кода?

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

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

Когда вы должны использовать его?

Таким образом, у вас есть хороший кандидат для этого подхода, если:

  1. У вас есть внутренний цикл, в котором одна функция вызывается много, и
  2. Вы знаете из профилирования, что этот вызов функции чрезвычайно дорог, и
  3. Вызов функции обходится дорого, потому что он должен быть универсальным, но при любом частном выполнении более крупной функции с внутренним циклом он вообще не должен быть универсальным, и
  4. Вы не можете устранить это с помощью какого-либо более простого метода, и
  5. Вы работаете на языке высокого уровня, который запускает байт-код на собственной виртуальной машине. (например, Python или Java)

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

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

Это нарушили две вещи: растущее разнообразие машинных архитектур (поэтому вам нужно написать один генератор кода для каждого возможного типа ЦП!) и изменения в архитектуре ЦП, которые делают практически невозможным написание эффективного машинного кода вручную без дублирования половины логика LLVM. В какой-то момент это стало слишком сложно, и появились более простые способы решения проблемы.

Но в таких языках, как Python или Java, у которых есть собственная виртуальная машина, ни одна из этих проблем не применима!

Реальный пример

Я собираюсь проиллюстрировать это реальным примером из моей собственной работы несколько лет назад, когда codegen ускорил производственную систему почти в 2 раза.

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

Настройка проблемы:

Хуму приходится манипулировать множеством наборов данных со структурированными схемами, а также уделять чрезвычайно пристальное внимание контролю доступа и защите данных. (Людские данные очень конфиденциальны!) В результате у него есть собственная система хранения X-уровня, которая управляет всем этим, которая, в свою очередь, находится поверх системы хранения W-уровня от облачного провайдера.

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

  • Каждая ячейка данных, возвращаемая базовым стеком хранилища, имеет ключ Tuple[bytes, ...], и все ключи, возвращаемые во время одной операции чтения, имеют одинаковый формат;
  • Для данной операции чтения у нас есть фиксированный критерий фильтрации для каждого поля этого кортежа, который представляет собой некоторую комбинацию отсутствия фильтра или равенства, больше или равно или меньше некоторого фиксированного значения.

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

def _match(self, *encodedKeys: bytes) -> bool:
  return all(
    keyInfo.matches(encodedKey)
    for keyInfo, encodedKey in zip(self._keys, encodedKeys)
 )

Здесь self — это одноразовый объект, который охватывает время чтения, self._keys содержит именованные кортежи в массиве, параллельном ключам, которые мы получим, а encodedKeys содержит ключевой кортеж из фактической ячейки данных.

Проблема:

В зависимости от аргументов операции чтения эта функция будет вызываться от одного раза (для получения фиксированных ключей) до десятков миллионов раз. Профилирование производственных нагрузок с помощью cprofile показало, что до 50% всего процессорного времени уходит на _match.

Это особенно расстраивает, потому что то, что делает _match, очень просто: для любого фиксированного запроса его логика едва ли превышает несколько if операторов. Но накладные расходы на вызов zip и вызовы функций на matches (и требуемая логика ветвления для разных типов компараторов) действительно складываются.

Интуиция:

Если бы у вас была волшебная палочка, вам бы эта функция вообще не понадобилась; вместо этого объект для каждого запроса будет просто иметь функцию, которая выполняет логику, необходимую только для этого одного запроса, и делает это очень быстро. Конструктор этого объекта для каждого запроса может просто создать функцию и использовать ее.

Это идеальный момент для codegen.

Неправильный способ сделать это: генерация Python

Интуитивным, но неправильным подходом было бы проанализировать self._keys и написать из него реальный код Python (' and '.join(f'args[{i}] {op} {value}' for i, op, value in comparisons или подобный), а затем передать этот код compile(), чтобы получить функцию, которую мы можем вызвать во внутреннем цикле.

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

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

Правильный способ сделать это: генерация байт-кода

На этом этапе мы собираемся конкретизировать нашу конкретную среду выполнения Python, в данном случае это CPython. (Изменение времени выполнения происходит настолько редко, что нам не нужно явно планировать изменение этого!) CPython работает путем компиляции Python AST в байт-код, который затем выполняет его центральный интерпретатор. Этот байт-код задокументирован как часть библиотеки dis, которая также позволяет дизассемблировать реальные функции и посмотреть, какой у них байт-код; эта страница документирует все инструкции и именно то, что они делают.

Что мы собираемся сделать, так это написать байт-код напрямую — буквально построить полный массив байтов — и превратить полученный байт-код в функциональный объект.

Функция, которая нам понадобится в нашем конкретном случае, довольно проста.

Давайте разберемся, что здесь происходит. В функциях Python все локальные переменные (начиная с аргументов функции) хранятся в массиве; все константы, используемые функцией, хранятся в отдельном массиве. (None всегда является первым элементом этого массива) Стек используется для управления мгновенным состоянием функции.

Инструкции LOAD_FAST и LOAD_CONST соответственно помещают N-ю локальную переменную или константу на вершину стека.

Инструкция COMPARE_OP <cmp_op> выполняет логическое сравнение между двумя верхними элементами стека и заменяет их результатом этого теста. Коды операций соответствуют таким операциям, как == и ›=, а их значения можно найти в dis.cmp_op, проиндексированных фактической строкой операции.

JUMP_IF_FALSE_OR_POP <address> — это инструкция для таких циклов. Он проверяет верхний элемент стека как логическое значение. Если значение равно True, то оно просто извлекается из стека и продолжается; если False, то он оставляет его в стеке и переходит к ячейке <address> в коде.

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

Итак, эта функция состоит из серии N сравнений. Каждый из них берет соответствующий ключ из проверяемой ячейки (помните, что мы будем передавать ключи в качестве N позиционных аргументов функции), константу для сравнения и выполняет соответствующее сравнение. Если результат False, то инструкция JUMP_IF_FALSE_OR_POP немедленно заставляет нас вернуть False; в противном случае переходим к следующему сравнению. Самое последнее сравнение пропускает JUMP_IF_FALSE_OR_POP; все, что находится в стеке, True или False, является правильным возвращаемым значением.

Превратить приведенный выше код на «ассемблере» в реальный байт-код очень просто: в Python 3.6+ байт-код — это просто массив байтов, состоящий из чередующихся кодов операций (один байт) и аргументов (один байт).¹

¹ Многобайтовые аргументы, которые нам не нужны в данном конкретном случае, обрабатываются с помощью кода операции EXTENDED_ARG.

Два нюанса: адреса и номера строк

Одна загвоздка здесь в том, что нам нужно знать значение <return> — адрес в коде, где будет находиться оператор return — чтобы написать более ранние операторы. (В частности, нам нужно знать позицию в результирующем байт-коде, где будет жить эта инструкция возврата). В общем дизайне компилятора есть всевозможные причудливые способы сделать это, обычно включающие разбиение кода на блоки без переходов, анализ каждого один отдельно, а затем связывая их вместе и разрешая все адреса сразу.

Но — и это важно при создании кодов — мы не пишем универсальный компилятор. Мы строим одну конкретную функцию, и быть причудливым и в целом правильным здесь бесполезно; нам просто нужно сгенерировать правильный ответ в этом конкретном случае. А так как каждая инструкция превращается ровно в два байта в сети, мы можем просто вычислить смещение инструкции возврата априори: это 8*(ncmps — 1)+6. (Поскольку каждое сравнение, кроме последнего, использует четыре инструкции, а последнее использует три)

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

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

Как выглядит настоящий код

Собрав все вместе, мы получим следующую функцию:

Это реальный производственный код, не упрощенный для этой статьи!

Как видите, этот код состоит из двух частей. В первой части мы синтезируем таблицы байт-кода и номеров строк самым прямым способом — записывая в объект BytesIO и получая необходимые значения двоичного кода операции из библиотеки dis. Во второй части мы создаем интерфейс, чтобы превратить байт-код во что-то, что Python может вызвать, путем создания объекта CodeType (внутреннее представление объекта кода в Python) и объекта FunctionType из него. Этот последний тип объекта — это именно то, что возвращается оператором def или lambda в Python.

Длинный комментарий к CodeType выделяет фрагменты, которые не очень хорошо документированы. Некоторые важные вещи, которые нужно осознать:

  • Вы должны сообщить объекту, сколько имеется позиционных и kwonly-аргументов и сколько всего локальных переменных. Количество локальных переменных должно включать аргументы.
  • Вам также необходимо указать максимальный размер стека, к которому следует подготовиться. В нашем случае это действительно просто — мы кладем два элемента в стек, сравниваем их и в итоге получаем один элемент в стеке, и продолжаем. В более общих случаях, когда точное значение неочевидно, вы можете захотеть немного послабить это число.
  • Первая константа всегда должна быть None, иначе CPython станет очень недовольным.
  • Вам нужно вставить различные вещи (имена функций, таблицы номеров строк и т. д.), чтобы облегчить отладку, но да, это обязательные аргументы.

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

Как вызывается код

Наконец, мы используем это очень просто: исходный объект запроса имеет локальную переменную _encodedKeyMatcher, которая установлена ​​на вывод этой функции или на «медленную» реализацию по умолчанию, если она по какой-то причине недоступна. Этот объект является вызываемым, как если бы он был установлен в выражение lambda, и просто вызывается для каждой ячейки с *keys в качестве аргумента.

Обратите внимание, что аргумент makeFastComparator — это не массив тех NamedTuples (self._keys), которые мы использовали выше, а скорее список сравнений, которые должны произойти, как кортежи (какой ключ, какая операция, какое значение). Это важно, потому что для каждого ключа требуется от нуля до двух реальных сравнений. Итак, мы добавляем к этому NamedTuple метод для преобразования самого себя в список операций сравнения, извлекая фактические коды сравнения из библиотеки dis.

В итоге

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

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

Все это говорит о том, что codegen — это метод крайней меры: если вы можете найти более простой способ решить проблему путем реструктуризации кода, вы должны обязательно сделать это. Это, на самом деле, достаточно «ковбойская» техника, поэтому, если вы обнаружите, что используете ее, традиция состоит в том, чтобы найти шляпу, помахать в воздухе и кричать «ага».

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