TLDR; Рассказываем, как мы интегрировали CodeBERT-based модель классификации секретов в production-продукт с жёсткими ограничениями по железу, сократив время инференса с 320 до 90 секунд и размер модели с ~600 до ~130 МБ — без дискретных ускорителей и тяжёлых зависимостей.
В предыдущей статье мы рассказали о том, как готовили и собирали нашу модель на основе модели CodeBERT для классификации срабатываний движков поиска секретов в коде. Там речь шла о датасетах, токенизаторе, метриках и о том, как нам удалось поднять PR AUC* с 0.70 до 0.90.
*PR AUC — Area Under the Precision-Recall Curve, буквально “площадь под кривой “точность-полнота””, простыми словами: уверенность модели при вероятностной оценке истинно положительного редкого класса в общей выборке.
Обо мне
Я Натан, техлид модуля Secrets в CodeScoring. Работаю над Secrets уже почти 2 года, раньше работал в FinTech и EdTech, сейчас пишу на Python и Go.
О компании
В CodeScoring мы занимаемся созданием решения для безопасной работы с open source, проверки совместимости лицензий, поиска секретов и оценки качества кода в разрезе команды.
О чём статья?
Как мы подходили к вопросу интеграции тяжело вычисляемой модели в AppSec-продукт в рамках ограниченной по мощности целевой машины.
Статья будет полезна в первую очередь MLOps’ам, которые хотят совместить машинное обучение и информационную безопасность, а также тем, кто интегрирует ресурсоемкие решения в условиях ограниченных вычислительных мощностей без просадки производительности продукта.
Контекст: что такое on-premise и какие вводные в целом?
Наш продукт запускается on-premise, то есть на мощностях клиента. Целевой сценарий — помощь в верификации результатов сканирования на секреты в коде в рамках CI/CD пайплайнов.
В прошлой реализации модели у нас был отлаженный, привычный для пользователя процесс, поэтому мы хотели вписать в него новую версию, чтобы не вызывать лишнего трения у пользователя.
Трансформеры — это тяжёлая история, и прежде чем двигаться дальше, нужно было честно ответить: можем ли мы запустить это в условиях ограниченных ресурсов по железу клиента и требованиями безопасности с приемлемым качеством UX?
После нескольких раундов обсуждений и отсеивания реального и нереального, мы пришли к следующим требованиям:
Модель работает строго на машине клиента.
Инференс в рамках анализа должен занимать не более 5 минут для гипотетического проекта, который содержит 1000 секретов, поскольку работа модели над секретами — это самая емкая по времени часть. На практике такие объемы почти не встречаются. Тут также стоит учитывать то, что Секреты не одиноки и будут конкурировать за ресурсы с другими частями продукта и системы.
Минимизируем объем зависимостей, это необходимо для:
Минимизации веса продуктивной сборки.
Уменьшения количества потенциальных уязвимостей.
Облегчения процесса сборки из исходников, это обязательный этап при сертификации ФСТЭК России.
В качестве целевой машины клиента считаем некий Linux сервер, с архитектурой процессора x86 с 4-мя логическими потоками, тактовой частотой 2 ГГц, 16 ГБ ОЗУ и способностью запускать Docker.
Модель живет на диске и запускается контейнером с потребителем задач. Она осуществляет инференс в рамках задачи анализа исходников на наличие секретов, где сначала секреты находятся при помощи движков типа gitleaks или trufflehog, после чего их оценивает модель (отвечая на вопрос “является ли секрет истинно положительным?”) по ряду свойств, где ключевыми являются сам секрет и его окружение. Пользователь запускает анализ через интерфейс (вручную или по расписанию), после чего задача выполняется в фоне.
Кроме того, пользователь может дообучать модель на своих размеченных данных прямо в интерфейсе продукта, без их передачи за пределы контура — это принципиально для наших клиентов с требованиями к безопасности данных.
Итерация первая: PyTorch, тяжелый и неповоротливый
Первый инстинкт был очевидным: взять модель, которая была реализована изначально на PyTorch, и встроить в продукт.
Структура самой модели на PyTorch выглядит следующим образом:
То есть модель — это последовательность из нескольких модулей: препроцессинг, токенизатор, трансформер, выходной слой.
Вес и зависимости
PyTorch тянет за собой огромное дерево транзитивных зависимостей. Для продукта в сфере AppSec очень важно самому быть безопасным. Каждая лишняя зависимость — потенциальная поверхность атаки, которую нужно отслеживать, патчить и аудировать.
Бенчмарк
У нас есть свой датасет, который представляет из себя набор записей, в каждой из которых есть: строка с вхождением секрета, само вхождение секрета, несколько связанных атрибутов и разметка, которую мы считаем истиной и которая нам сообщает, является ли секрет истинно положительным. Таких записей мы взяли 1000 штук и использовали в качестве “линейки”.
Скорость инференса
Первые замеры на нашем бенчмарке, на Debian образе, дали 320 секунд. Инференс является частью процесса анализа, который пользователь запускает интерактивно или по графику. Поэтому максимально допустимым временем на 1000 секретов взяли 4-5 минут, с учетом того, что остальные части анализа будут проходить достаточно быстро.
Alpine Linux и musl libc
Наши production-образы построены на Alpine Linux, который использует musl libc вместо glibc.
Небольшое пояснение, что такое libc и в чем разница между musl и glibc: libc — это базовый набор компонентов в Linux-системах, который позволяет программам на C и C++ работать нормально, самая популярная реализация этой библиотеки — это glibc. В ней есть очень много полезных функций, которые позволяют легче писать и запускать различное ПО на glibc Linux, но как часто бывает, за всё хорошее нужно чем-то платить, и тут точно так же. Linux-системы на glibc часто носят в себе большое количество уязвимостей и вообще имеют больший объем, поэтому, если вам очень принципиальна легкость и безопасность, то стоит обращаться к иной реализации libc. В нашем случае мы используем образы Alpine, которые реализуют libc через свой musl-набор, он легкий и безопасный, но часто не имеет “сахарного” функционала.
PyTorch официально не распространяет колеса под musl. Это означало либо отказ от Alpine, либо сборку PyTorch из исходников. Отказ от Alpine для нас будет довольно тяжелым, поэтому мы решили попробовать собрать PyTorch под Alpine.
Однако это привело нас к тому, что эти колеса нужно собирать определенным образом, и не было четкой уверенности, что конкретные колеса будут работать на разных платформах. Помимо этого, хоть и функционально musl и glibc являются (почти) тождественными, всё равно возникали опасения насчёт того, что PyTorch будет работать на них по-разному — и не напрасно.
Основной проблемой была библиотека с реализацией потоков libpthread у glibc, которая имеет функцию pthread_attr_setaffinity_np() и отсутствует в musl’овой реализации. Это лечилось отключением USE_DISTRIBUTED флага при сборке, однако это убивало напрочь скорость инференса модели. Она, конечно, работала на CPU, но насладиться результатами её вычислений уже могли бы, скорее всего, наши потомки. Поэтому этот вариант нам тоже не подошел.
Итерация вторая: Wolfi+PyTorch, неоправданная сложность
Прежде чем прийти к финальному решению, мы попробовали обойти проблему с musl иначе. Поскольку мы используем alpine из-за его малого footprint’а и низкого количества уязвимостей, переход на стандартные образы с glibc нам не подходил — они слишком тяжелые и часто содержат уязвимости. Однако мы нашли исключение: базовый образ Wolfi, который работает на glibc, но при этом по своим свойствам (весу и безопасности) очень близок к Alpine. Так как модель запускается в контейнере с потребителем задач, идея была простой: заменить его базовый образ на Wolfi и таким образом закрыть все предметные и бизнесовые требования. Мы дошли до реальной сборки образа — и даже собрали его. Но вылезло сразу две проблемы:
Wolfi не позволяет бесплатно стягивать версии пакетов, кроме последней. Для production-продукта, где воспроизводимость сборки критична, это неприемлемо: нельзя зафиксировать конкретную версию и быть уверенным, что через месяц образ соберётся идентично.
PyTorch внутри Wolfi-контейнера вёл себя непредсказуемо: колеса с разными параметрами сборки вели себя по-разному на разных платформах, у коллеги могло быть все хорошо, у меня же инференс мог залипать на каком-то из этапов и дебажить PyTorch — это не самое благодарное занятие с точки зрения времени и усилий, поэтому нужно было что-то другое, что-то простое и элегантное.
Отсюда появилась идея о том, что нам стоит совсем избежать PyTorch в продуктивной среде.
Последняя итерация: ONNX, удачный эксперимент
ONNX (Open Neural Network Exchange) — формат и среда исполнения для нейронных сетей, изначально созданный Microsoft как платформенно-нейтральная альтернатива фреймворкам вроде PyTorch и TensorFlow. Он гораздо легче, модульнее и не тянет за собой ворох зависимостей.
ONNX, в отличие от PyTorch, создаёт статичный граф вычислений из модели, что и является основным секретом его быстродействия. Также, имея такой граф, вы можете буквально взглянуть на скелет конвертированной в ONNX-формат модели.
При разработке нам нужно было превратить PyTorch модель в ONNX формат и затем заставить это работать в среде нашего продукта с учётом целевой машины. У PyTorch есть API, которое позволяет конвертировать модели в нужный нам формат, поэтому через серию тестов различных подходов мы пришли к определенному скрипту конверсии, который нас удовлетворял, и мы получили нашу модель в нужном формате.
Правда, и здесь был подводный камень: официальных колес onnxruntime под Alpine / musl также не существует. Но, в отличие от PyTorch, ONNX устроен иначе — он разбит на отдельные пакеты, из которых можно взять только нужную функциональность. Это делало задачу сборки из исходников гораздо проще.
Мы форкнули репозиторий microsoft/onnxruntime в наш внутренний GitLab, внесли необходимые патчи для совместимости с musl и нашими требованиями, а также настроили CI-пайплайн, который собирает колеса и публикует их в наш внутренний репозиторий пакетов. Итоговый результат: два пакета onnxruntime и onnxruntime-extensions весом около 25 МБ.
При сборке основного колеса мы убрали сборку и исполнение юнит-тестов после сборки при помощи --cmake_extra_defines onnxruntime_BUILD_UNIT_TESTS=OFF и --skip_tests, поскольку это требовало дополнительных зависимостей и усложняло пайплайн.
Также стоит собираться в Release конфигурации при помощи --config Release, и нет нужды в том, чтобы подтягивать все сабмодули из репозитория — их мы сами подкладываем в свой репозиторий, поэтому добавляем --skip_submodule_sync.
Помимо этого были пропатчены два места при копировании репозитория:
onnxruntime/core/platform/posix/stacktrace.cc, в котором включается резолв стектрейса при падении библиотеки — это нам в продакшене не нужно.version.txt— тут просто была неправильная версия, которую мы заменили на фактическую
Ещё можно подключить сборщик Ninja и кэш для CMAKE, которые вместе должны ускорять сборку в целом и её повторение.
Сам факт конвертации модели из PyTorch в ONNX формат дал немедленный бонус: ~15% прирост скорости инференса без каких-либо дополнительных оптимизаций.
Помимо этого, для onnx существуют библиотеки на других языках, и сам он превращает модель в универсальный формат .onnx, который дружит и с другими инструментами и фреймворками, например, OpenVINO.
Оптимизация: от 600 МБ и 320 секунд к 130 МБ и 90 секундам
Конвертации в ONNX оказалось недостаточно — 320 секунд превратились в 270, что всё ещё неприемлемо, поскольку оставляет маленький зазор по времени для всего остального. Поэтому мы применили ряд дополнительных техник.
Квантизация модели
Квантизация — это снижение точности представления данных в модели для уменьшения размера и ускорения вычислений. После нескольких серий тестирований довольно большого количества кандидатов, выведенных квантизацией, мы остановились на int8-квантизации на уровне отдельных операций внутри графа модели, совместив её с рядом структурных оптимизаций самой модели и её обертки в коде. В итоге это позволило нам:
Снизить размер модели с ~500 МБ до ~130 МБ
Существенно сократить время инференса с ~270 секунд до ~90 секунд
Бенчмарк-методология
Для всех замеров мы использовали один подход: вышеупомянутый проприетарный размеченный датасет из 1000 секретов, запущенный на машине с 4 потоками CPU / 2 ГГц / 16 ГБ RAM. Это наше определение «средней машины клиента»: мы хотели быть уверены, что это работает именно у клиента, а не только локально у разработчиков.
Итоговый результат: 90 секунд для 1000 секретов. Это в 3.5 раза быстрее первоначального варианта, пускай и с минимальной потерей точности (около 2%). С учётом того, что реальный датасет клиента редко содержит 1000 секретов единовременно, на практике время работы значительно меньше. Сравнительная таблица основных параметров каждого подхода:
Вариант | Вес библиотеки | Размер модели | Время инференса (1000 секретов) |
|---|---|---|---|
PyTorch (baseline) | 1 ГБ (cpu-only) | ~600 МБ | ~320 сек |
ONNX без оптимизаций | 25 МБ | ~500 МБ | ~250-230 сек |
ONNX + квантизация + оптимизации | 25 МБ | ~130 МБ | ~90 сек |
Итог
У нас получилось интегрировать новую модель в продукт с соблюдением всех условий, которые являются пересечением нескольких необходимостей в виде: удобства пользователя, быстродействия, экономии ресурсов и соответствия требованиям безопасности.
Главный урок
Если коротко: не бойтесь экспериментировать. На момент, когда мы столкнулись с этой задачей, информации о практическом применении ONNX в продуктивной среде было немного — особенно применительно к musl-окружениям. Часть решений пришлось находить самостоятельно методом проб и ошибок. Но результат оправдал усилия.
Если говорить шире: внедрение трансформеров в on-premise продукт — это задача не только для команды data-инженеров. Это комплексная инженерная задача, в которой нужно не только создать гипотезу о модели, проверить её, оптимизировать точность, но и затем собрать это всё в функциональный модуль остальной системы, которая помимо того, что имеет свой стек и закономерности работы, работает не в облаке, а на машине клиента (которая, обычно, не будет иметь у себя на борту дискретные ускорители).
Благодарность
Большая благодарность моему коллеге Никите Бесперстову за помощь со сборкой библиотек из исходников и в прокладывании пути в темном лесу системных библиотек Linux.
Также хочу поблагодарить Антона Володченко за помощь с формулированием данной статьи и наставлениями по процессу публикации в целом.
Читателям
Если у вас есть опыт запуска трансформеров в жёстко ограниченных on-premise окружениях — будем рады обсудить в комментариях.