AI [Перевод] WebAssembly как платформа расширений для Python: ускорение, встраивание и опасные ловушки API

[AI]

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

В этой роли хорошо себя показывает Lua, а для веб-технологий, разумеется, используется JavaScript. WebAssembly обобщает этот подход: любой язык программирования, способный компилироваться в Wasm, может расширять приложение, выполняющее Wasm. Это требует больше усилий, чем просто передать скрипт в текстовом файле, зато авторы расширений могут писать на удобном им языке и использовать более развитые инструменты разработки — отладки, тестирования и прочего — которые обычно недоступны для традиционных языков расширений.

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

Чтобы понять, какие темы в Python стоит подтянуть перед работой с расширениями, памятью и внешними модулями, можно пройти короткий тест и оценить свой текущий уровень.

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

В качестве предпочтительной среды выполнения Wasm (без поддержки WASI) можно использовать wasm3. Он написан на классическом C и хорошо подходит для встраивания, примерно так же, как, например, SQLite. Производительность у него средняя, однако программа на C, выполняемая в wasm3, всё равно заметно быстрее аналогичной программы на Python. Для него есть Python-привязки — pywasm3, но они распространяются только в виде исходного кода.

Это означает, что на машине должен быть установлен компилятор C, чтобы использовать pywasm3, что противоречит моим целям. Если на системе уже есть такая цепочка инструментов, проще сразу использовать её, минуя Wasm.

Для рассматриваемых в этой статье сценариев лучшим вариантом является wasmtime-py. Дистрибутив включает готовые бинарные файлы для Windows, macOS и Linux на архитектурах x86-64 и ARM64, что покрывает практически все установки Python. На стороне пользователя требуется только интерпретатор Python — никаких нативных инструментов сборки.

Это почти так же удобно, как если бы Wasm был встроен в сам Python. В моих тестах он работает в 3–10 раз быстрее, чем wasm3, поэтому для первого сценария ситуация даже лучше.

Однако есть и недостатки: установленный пакет занимает около 18 МБ и, вероятно, со временем будет сопоставим по размеру с самим интерпретатором Python. Кроме того, его API меняется практически каждый месяц, так что вам придётся постоянно обновляться, иначе через пару лет ваш код может устареть и перестать работать. В этой статье рассматривается версия 40.

Примеры использования и подводные камни​

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

import functools
import wasmtime

store    = wasmtime.Store()
module   = wasmtime.Module.from_file(store.engine, "example.wasm")
instance = wasmtime.Instance(store, module, ())
exports  = instance.exports(store)

memory = exports["memory"].get_buffer_ptr(store)
func1  = functools.partial(exports["func1"], store)
func2  = functools.partial(exports["func2"], store)
func3  = functools.partial(exports["func3"], store)

Store — это область выделения памяти, из которой создаются все объекты Wasm. Освобождать отдельные объекты нельзя, можно только отбросить весь store целиком. Честно говоря, это вполне разумно. Менее разумно другое: чтобы воспользоваться объектами, мне снова и снова приходится передавать store в каждый из них. Эти объекты связаны ровно с одним store и не могут использоваться с другим. Передадите не тот store — и всё аварийно завершится. Хотя он и так уже отслеживается внутри.

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

Объект get_buffer_ptr реализует протокол буфера, и если вы передаёте что-то кроме байтов, то, скорее всего, именно его и стоит использовать для доступа к памяти. Для него действуют обычные оговорки: если вы меняете размер памяти, то, вероятно, стоит заново получить объект буфера. Для байтов, например буферов и строк, я предпочитаю методы read и write.

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

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

malloc = functools.partial(exports["func1"], store)

hello = b"hello"
pointer = malloc(len(hello))
assert pointer
memory = exports["memory"].write(store, hello, pointer)  # НЕПРАВИЛЬНО!

И этого мало: в wasmtime-py есть ещё одна неприятная ловушка. Методы read и write следуют сомнительному соглашению Python, по которому отрицательные индексы отсчитываются от конца. Если malloc вернёт указатель из верхней половины памяти, отрицательный указатель пройдёт проверку границ внутри write, потому что отрицательные индексы там допустимы, и данные будут молча записаны не по тому адресу. Вот так сюрприз.

Мне стало интересно, насколько распространена эта ошибка, и я поискал примеры в сети. Мне удалось найти только один нетривиальный случай использования wasmtime-py в реальном проекте — в изолированном просмотрщике PDF. И, как я и ожидал, там попались в ловушку с отрицательным указателем. Более того, это ещё и переполнение буфера с записью в область памяти Python:

buf_ptr = malloc(store, len(pdf_data))
mem_data = memory.data_ptr(store)

for i, byte in enumerate(pdf_data):
    mem_data[buf_ptr + i] = byte

Метод data_ptr возвращает сырой указатель ctypes без проверки границ, так что на деле здесь сразу две ошибки. Во-первых, если разработчика вообще заботит работа в песочнице, доверять указателям, пришедшим из Wasm, не следует. Во-вторых, остаётся риск отрицательного указателя, а в таком случае запись произойдёт за пределами памяти Wasm, уже в памяти Python, что, будем надеяться, хотя бы приведёт к segfault.

Что с этим делать? Каждый указатель, приходящий из Wasm, нужно усекать маской:

pointer = malloc(...) & 0xffffffff # правильно для wasm32!

Так результат интерпретируется как беззнаковый. Для 64-битного Wasm нужна 64-битная маска, хотя на практике получить корректный отрицательный указатель из 64-битного Wasm вы не сможете. Это правило относится и к JavaScript, где обычно используют такой приём:

let pointer = malloc(...) >>> 0

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

Когда у вас уже есть корректный адрес, можно использовать его с представлением памяти через протокол буфера. Если вы работаете с NumPy, есть разные способы взаимодействовать с этой памятью, оборачивая её в типы NumPy, но только если хост использует порядок байтов little-endian. Если же у вас машина с big-endian, то от запуска Wasm, пожалуй, лучше сразу отказаться.

Первый сценарий, который я имею в виду, обычно сводится к копированию обычных значений Python туда и обратно. Здесь очень удобен пакет struct:

vec2   = malloc(...) & 0xffffffff
memory = exports["memory"].get_buffer_ptr(store)
struct.pack_into("<ii", memory, vec2, x, y)

Он играет примерно ту же роль, что и DataView в JavaScript. Если нужно копировать много чисел, то в CPython быстрее сформировать собственную строку формата, чем использовать цикл:

nums: list[int] = ...
struct.pack_into(f"<{len(nums)}i", memory, buf, *nums)

Чтобы извлечь структуры обратно, используйте struct.unpack_from. Если вы передаёте строки, придётся вызывать .encode() и .decode(), чтобы преобразовывать их в байты и обратно, а байты хорошо подходят для работы через read и write.

На практике в реальных Wasm-программах вы, скорее всего, будете взаимодействовать с гостевым аллокатором (guest allocator) извне, запрашивая память, в которую затем копируются входные данные для функции. В моих примерах используется malloc, потому что он не требует дополнительных пояснений, но, как это часто бывает, bump аллокатор решает задачу гораздо лучше, особенно потому, что не требует тащить в Wasm-программу полноценный аллокатор общего назначения.

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

WebAssembly как более быстрый Python​

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

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

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

Мой основной экспериментальный тест — это вариация моего решения задачи «Two Sum», которое я изначально написал для JavaScript, затем адаптировал под pywasm3, а позже и под wasmtime-py. Пример простой, но достаточно показательный и хорошо отражает тот тип Wasm-встраивания, который я имею в виду. Интерфейс остаётся тем же, но реализация выполняется через Wasm.

# Исходный интерфейс в стиле Python
def twosum(nums: list[int], target: int) -> tuple[int, int] | None:
   ...

# Интерфейс Wasm с состоянием
class TwoSumWasm():
   def __init__(self):
       store    = wasmtime.Store()
       module   = wasmtime.Module.from_file(store.engine, ...)
       instance = wasmtime.Instance(store, module, ())
       ...

   def twosum(self, nums, target):
       # ... использовать экземпляр wasm ...

Здесь появляется состояние, связанное с самим экземпляром Wasm. Если скрыть его, сделав глобальным, придётся синхронизировать доступ к нему между потоками. В многопоточной программе, возможно, стоит использовать локальные для потока данные, создаваемые по требованию. Мне пока не приходилось это решать.

Однако здесь особенно заметна слабость store в wasmtime: обратите внимание, что компиляция и создание экземпляра связаны в рамках одного store. Я не могу один раз скомпилировать модуль, а затем по мере необходимости создавать временные экземпляры, например как это нужно при каждом запуске программы WASI. Каждый новый экземпляр навсегда расширяет store, в котором происходила компиляция. На практике это означает, что для каждого такого временного экземпляра нам приходится заново и совершенно нерационально перекомпилировать Wasm-программу.

Несмотря на внешний вид API, компиляция и создание экземпляра здесь на самом деле не являются отдельными этапами, как это сделано в Wasm API для JavaScript. wasmtime.Instance принимает store первым аргументом, и это наводит на мысль, что для создания экземпляра можно использовать другой store. Это решило бы проблему, но на момент написания статьи требуется тот же самый store, который использовался при компиляции модуля. Для некоторых реальных сценариев это критический изъян, особенно для WASI.

Обновление: Вольфганг Майер подсказал методы serialize и deserialize, которые позволяют отвязать скомпилированный модуль от его store и тем самым создавать независимые экземпляры. Я попробовал этот подход, и он оказался вполне рабочим обходным решением. Накладные расходы невелики, а при десериализации валидация не выполняется. В моём тесте теперь используется именно этот способ — на будущее, поскольку я ожидаю, что это будет мой типичный сценарий.

WebAssembly как встраиваемые возможности​

Monocypher — замечательная криптографическая библиотека. Она компактная, эффективная и хорошо подходит для встраивания, настолько хорошо, что распространяется в виде единого объединённого файла. Ей не нужны ни libc, ни среда выполнения, поэтому её можно напрямую скомпилировать в Wasm почти любой цепочкой инструментов Clang:

$ clang --target=wasm32 -nostdlib -O2 -Wl,--no-entry -Wl,--export-all
        -o monocypher.wasm monocypher.c

Эта библиотека не «понимает» Wasm, поэтому мне нужен флаг --export-all, чтобы открыть наружу весь интерфейс. И это даже удобно, потому что в рамках одной единицы трансляции всё, что имеет внешнее связывание, и становится интерфейсом.

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

#include "monocypher.c"

extern char  __heap_base[];
static char *heap_used;
static char *heap_high;

void *bump_alloc(ptrdiff_t size)
{
    // ...
}

void bump_reset()
{
    ptrdiff_t len = heap_used - __heap_base;
    __builtin_memset(__heap_base, 0, len);  // стереть ключи и т. п.
    heap_used = __heap_base;
}

О __heap_base я уже писал раньше: это часть ABI. Мы будем складывать ключи, входные данные и прочее в этот «стек», запускать нашу криптографическую процедуру, копировать наружу результат, а затем сбрасывать bump аллокатор, который при этом затирает все чувствительные данные. Нередко одного memset недостаточно: обычно сначала память зануляется, а потом освобождается, и компилятор видит, что время жизни объекта подходит к концу.

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

В этом API много всего, но я остановлюсь только на интерфейсе AEAD. Мы «запираем» данные в зашифрованный контейнер и можем записать снаружи любую незашифрованную метку. Позже этот контейнер можно открыть, но только если ни его содержимое, ни эта метка не были подменены. Очень добротный дизайн API:

void crypto_aead_lock(uint8_t       *cipher_text,
                      uint8_t        mac  [16],
                      const uint8_t  key  [32],
                      const uint8_t  nonce[24],
                      const uint8_t *ad,         size_t ad_size,
                      const uint8_t *plain_text, size_t text_size);
int crypto_aead_unlock(uint8_t       *plain_text,
                       const uint8_t  mac  [16],
                       const uint8_t  key  [32],
                       const uint8_t  nonce[24],
                       const uint8_t *ad,          size_t ad_size,
                       const uint8_t *cipher_text, size_t text_size);

Скомпилировав это в Wasm, мы можем получить доступ к этой функциональности из Python почти так, будто она написана на чистом Python, и взаимодействовать с другими системами, использующими Monocypher.

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

class Monocypher:
    def __init__(self):
        ...
        self._read   = functools.partial(memory.read, store)
        self._write  = functools.partial(memory.write, store)
        self.__alloc = functools.partial(exports["bump_alloc"], store)
        self._reset  = functools.partial(exports["bump_reset"], store)
        self._lock   = functools.partial(exports["crypto_aead_lock"], store)
        self._unlock = functools.partial(exports["crypto_aead_unlock"], store)
        self._csprng = secrets.SystemRandom()

    def _alloc(self, n):
        return self.__alloc(n) & 0xffffffff

    def generate_key(self):
        return self._csprng.randbytes(32)

    def generate_nonce(self):
        return self._csprng.randbytes(24)

    ...

С надёжной основой всё остальное уже делается легко. Блок finally гарантирует, что секретные данные всегда будут удалены из памяти Wasm, а всё остальное сводится к копированию байтов туда и обратно:

def aead_lock(self, text, key, ad = b""):
        assert len(key) == 32
        try:
            macptr   = self._alloc(16)
            keyptr   = self._alloc(32)
            nonceptr = self._alloc(24)
            adptr    = self._alloc(len(ad))
            textptr  = self._alloc(len(text))

            self._write(key, keyptr)
            nonce = self.generate_nonce()
            self._write(nonce, nonceptr)
            self._write(ad,    adptr)
            self._write(text,  textptr)

            self._lock(
                textptr,
                macptr,
                keyptr,
                nonceptr,
                adptr, len(ad),
                textptr, len(text),
            )
            return (
                self._read(macptr, macptr+16),
                nonce,
                self._read(textptr, textptr+len(text)),
            )
        finally:
            self._reset()

А aead_unlock устроен почти так же, только в обратном порядке. Если контейнер не удаётся открыть, например из-за подмены данных, метод выбрасывает исключение:

def aead_unlock(self, text, mac, key, nonce, ad = b""):
        assert len(mac) == 16
        assert len(key) == 32
        assert len(nonce) == 24
        try:
            macptr   = self._alloc(16)
            keyptr   = self._alloc(32)
            nonceptr = self._alloc(24)
            adptr    = self._alloc(len(ad))
            textptr  = self._alloc(len(text))

            self._write(mac, macptr)
            self._write(key, keyptr)
            self._write(nonce, nonceptr)
            self._write(ad, adptr)
            self._write(text, textptr)

            if self._unlock(
                textptr,
                macptr,
                keyptr,
                nonceptr,
                adptr, len(ad),
                textptr, len(text),
            ):
                raise ValueError("AEAD mismatch")
            return self._read(textptr, textptr+len(text))
        finally:
            self._reset()

Использование:

mc = Monocypher()
key = mc.generate_key()
message = "Hello, world!"
mac, nonce, encrypted = mc.aead_lock(message.encode(), key)

Передайте mac, nonce и encrypted другой стороне, или самому себе в будущем, если ключ у вас уже есть:

decrypted = mc.aead_unlock(encrypted, mac, key, nonce)

Полный исходный код можно найти в моём рабочем репозитории.

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

• 
  6 мая в 20:00. «Rust в деле: пишем многопользовательский чат с сервером, клиентом и CLI». Записаться
  
  
• 
  7 мая в 20:00. «Настройка удобного рабочего окружения для Python проекта». Записаться
  
  
• 
  21 мая в 20:00. «Перестаньте бояться указателей: как Go экономит вашу память и CPU». Записаться
  

Еще больше практических демо-уроков по разработке и не только можно найти в календаре мероприятий.