JupyterLite позволяет запускать ядра Python прямо в браузере через WebAssembly. Это удобно для документации, образования и демо — не нужен сервер на каждого пользователя. Но в научном Python-окружении не хватало главного: Numba. Теперь это исправлено. QuantStack представила первую рабочую версию Numba JIT, работающую целиком в браузере, вместе с JupyterLite и emscripten-forge.
В демонстрации Numba разгоняет код примерно в 250 раз в WebAssembly, тогда как нативно — около 90 раз. Обход интерпретатора Python в браузере даёт даже больший относительный выигрыш. Функция превращается в Numba Intermediate Representation, типизируется, опускается до LLVM IR, компилируется в WebAssembly, линкуется и выполняется — без удалённого сервера.
Раньше это было невозможно. Numba — это компилятор, её работа зависит от llvmlite и LLVM. На обычной платформе сгенерированный машинный код кладётся в исполняемую память и вызывается. Браузер запрещает создавать и изменять исполняемую память. Поэтому понадобился WebAssembly-aware execution engine для llvmlite, in-process вызов LLVM-линкера LLD и загрузка скомпилированного кода как Emscripten side module. Такой подход уже использовали в Xeus-Cpp для запуска Clang-Repl в браузере, и теперь его применили к llvmlite.
В JupyterLite llvmlite отдаёт сгенерированный модуль новому движку. Тот компилирует его в WebAssembly-объект, линкует через LLD, загружает как side module в работающий Xeus-Python и добавляет в таблицу функций WebAssembly, чтобы Numba могла его вызывать. Одна функция может требовать несколько модулей: сначала runtime NRT, потом вспомогательные, потом сам пользовательский код — порядок важен. Работает и кэш: @njit(cache=True) сохраняет данные в файловой системе JupyterLite, а при новом сеансе llvmlite переиспользует WebAssembly-объект и перелинковывает его заново. Это настоящая компиляция Numba в браузере, а не эмуляция и не удалённый сервис.
Распространение кода решает emscripten-forge — дистрибутив WebAssembly-пакетов на базе conda/mamba и conda-forge. Numba уже есть в нём, его можно динамически устанавливать в JupyterLite через Mamba.
С появлением Numba открывается и экосистема. В браузере уже работают PyTensor — он компилирует символьные графы через линкер Numba, PyMC на его основе, interpolation.py для ускоренной интерполяции и Dolo.py для численных методов в экономике. Один компилятор открыл дорогу пакетам, которые раньше не могли принести свою модель выполнения в JupyterLite.
В планах — перенос изменений в upstream, расширение тестов, улучшение производительности и кэширования, поддержка SIMD в WebAssembly. Автор работы — Анутош Бхат, научный разработчик QuantStack, мейнтейнер LLVM и соавтор Xeus-Cpp.