Бесконечно можно смотреть на то, как горит огонь, течет вода и как дата сайентисты и ML-инженеры развертывают модели машинного обучения. Примерно у трети ML-инженеров данная задача занимает не менее часа рабочего времени. На хакатоне Data Product Hack от AI Talent Hub мы разработали инструмент для упрощения развертывания моделей машинного обучения MLJET. В статье рассказываем, как он работает. 

Всем привет! Мы – Константин Темплин и Кристина Желтова – магистранты онлайн-программы AI Talent Hub, создателями которой являются AI-компания @NapoleonIT и университет ИТМО В сентябре 2022 года мы участвовали в студенческом хакатоне Data Product Hack, что дало нам хороший опыт в продуктовом мышлении, командной работе и техническом аспекте реализации продукта. За несколько дней мы собрали команду, назвали её AI Wolfs и… победили! И победили мы не потому, что «безумно можно быть первым», а потому, что предложили решение для экономии времени специалиста по данным — разработали минималистичный фреймворк автоматического развертывания моделей машинного обучения MLJET. 

Идея

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

Чтобы проверить актуальность наших суждений и целесообразность создания продукта, мы опросили ML-инженеров. В опросе приняли участие более 30 человек. 67% опрошенных тратили более часа на развертывание модели, при этом 60% сочли этот процесс энергозатратным.

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

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

Анализ существующих решений

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

Оценив сильные и слабые стороны конкурентов, мы определили функционал и основную ценность нашего решения.

Это помогло:

• автоматизировать определение платформы машинного обучения, сбор зависимостей и схемы данных,

• интегрировать множество популярных фреймворков машинного обучения,

• создавать итоговые проекты, не имеющие зависимости от нашего инструмента.

Решение

Спойлер: решение вышло простым и универсальным.

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

В результате использования нашего инструмента пользователь получит директорию с полностью готовым для запуска проектом, а также Docker-образ или PIP-пакет. 

Для того, чтобы пользователь получил готовое решение, наш инструмент проходит следующие шаги:

1. Сборка проектной директории.

2. Сериализация моделей машинного обучения.

3. Генерация скриптового python файла на основе шаблонов.

4. Создание API документации с помощью Swagger.

5. Сборка docker-контейнера или python-пакета.

6. (Опционально) Развертывание.

Сборка проекта

Первый этап работы нашего инструмента состоит в сборке проектной директории. Он учитывает специфику выбранного фреймворка для бэкенда и модели машинного обучения.

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

Все шаблоны для серверной части являются .py файлами. Мы рассматривали также подходы с шаблонизаторами вроде Jinja2, однако использование python скриптов позволяет пользоваться всей мощью статических анализаторов, форматтеров и линтеров. Это также  упрощает тестирование.

Шаблон должен иметь точку входа main, содержать в себе методы-прототипы для замены на конкретную реализацию для ML-модели, а также иметь соответствующие методы-маршруты (например, в Flask это app.route).

Пример шаблона:

import ... 

loaded_model = ...

# список методов для замены
def predict(model, data) -> list: ...
def predict_proba(model, data) -> list: ...

# фреймворка-специфичные функции (или классы) для конечных точек
# они должны быть помечены специальным префиксом (на данный момент - подчеркиванием), для соответствующего метода модели машинного обучения
# например predict -> _predict
@app.post("/predict")
def _predict(...) -> ...: ...

# точка входа
if __name__ == "__main__":
    ...

Шаблоны для оберток различных типов моделей машинного обучения являются функциями. Пример для scikit-learn:

"""Module that contains Scikit-learn model method's wrappers."""
from mljet.contrib.supported import ModelType

# Константа определяющая, для каких моделей используются данные обертки
USED_FOR = [ModelType.SKLEARN_MODEL, ModelType.SKLEARN_PIPE]

def predict(model, data) -> list:
    return model.predict(data).tolist()

def predict_proba(model, data) -> list:
    return model.predict_proba(data).tolist()

Перед сборкой шаблоны проходят следующие проверки:

• проверка на соответствие спецификации,

• проверка статической типизации с помощью mypy,

• замена оберток модели в шаблоне на конкретные реализации,

• добавление необходимых импортов,

• форматирование с использованием black, isort,

• статическая проверка типов с помощью mypy.

Замена функции на конкретную реализацию осуществляется путем поиска и замены по регулярному выражению. 

pyfunc_with_body = re.compile(
   r"(?P<indent>[\t]*)(async def|def)[\t]*(?P<name>\w+)\s*\((?P<params>.*?)\)(?:[ "
   r"\t]*->[ \t]*(?P<return>\w+))?:(?P<body>(?:\n(?P=indent)(?:[ \t]+[^\n]*)|\n)+)"
)

Следующим шагом является сборка необходимых зависимостей — их мы собираем на основе анализа ast дерева сервисных файлов-шаблонов, а также типа используемой модели машинного обучения. Необходимые версии требуемых пакетов мы берем на основе текущего виртуального окружения, используя importlib_metadata.

Пайплайны

При создании решения вдохновлялись принципом работы задач в CI/CD. Например, в Github Actions можно указывать различные задачи (actions) и запускать их последовательно или параллельно.

Так, разработали аналогичный механизм в нашем проекте: основной единицей определили задачу (action).

Получается, для сборки проектной директории и создания docker-контейнера последовательность задач выглядит так:

сборка проекта —> создание docker-контейнера —> отправка на облачную платформу

Для хранения такого списка задач мы использовали DAG — ориентированный ациклический граф, в котором вершинами являются actions, а ребрами — отношение depends_on. Таким образом, для примера выше стадия сборки контейнера опирается на сборку проекта, а отправка на облачную платформу опирается на сборку контейнера.

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

@runtime_checkable
class Stage(Protocol):
   """Stage protocol."""

   # stage name
   name: str

   # a list of the stages names,
   #  on which the stage is based
   depends_on: FrozenSet[str]

   # all stage must be callable
   def __call__(self, *args, **kwargs):
       ...  # fmt: skip

Для того, чтобы функция или класс стали стадией, необходимо использовать декоратор stage, например:

@stage("docker-build", depends_on=["project-build"])
def docker_build(
   model: Estimator,
   tag: Optional[str] = None,
   base_image: Optional[str] = None,
   container_name: Optional[str] = None,
   need_run: bool = True,
   port: int = 5000,
   n_workers: int = 1,
   silent: bool = True,
   remove_project_dir: bool = False,
) -> str:
        …

Декоратор анализирует переданный объект, создает его копию и в зависимости от типа выставляет ему необходимые атрибуты (name, depends_on).

Использование нашего инструмента в простейшем случае для пользователя выглядит следующим образом:

import ...
from mljet import cook
X, y = ...
clf = RandomForestClassifier()
clf.fit(X_train, y_train)
cook(strategy="docker", model=clf, verbose=True)

Что дальше?

Прошло уже более полугода с победы на хакатоне Data Product Hack от проекта AI Talent Hub, а мы продолжаем работать над созданным инструментом автоматизации моделей.
За это время нам удалось: 

1. реализовать возможность генерации итоговых проектов на основе различных web-фреймворков. Cейчас поддерживается Sanic, Flask, Aiohttp, FastAPI (а также множество проверок на корректность),

2. сделать документацию с гайдами для пользователей и разработчиков на ReadTheDocs: https://mljet.readthedocs.io/en/latest/,

3. добавить поддержку XGBoost, Catboost, LightGBM,

4. улучшить workflow: добавили пользовательские сценарии в CI, пакет на Pypi (для версионирования используем семантическое версионирование), интеграцию с Codacy и отслеживание метрик проекта, Renovate,

5. сделать CLI интерфейс.

Сейчас MLJET — это автоматизированный инструмент развертывания ML моделей. Нам очень повезло, потому что наш проект курирует Александр Рыжков, руководитель команды разработки LightAutoML в Sber AI Lab. LightAutoML — это open-source фреймворк от Sber AI Lab, который позволяет автоматически строить модели обучения с учителем.

Автоматизация развертывания отлично дополняет инструмент от Sber AI Lab, так как позволяет представить отлаженную модель в виде готового сервиса. К тому же, в LightAutoML есть возможность генерации отчетов об обучении и работе модели, что впоследствии можно включить отдельным эндпойнтом в развертываемый сервис. Также думаем над усилением инструмента фронтэндом с набором визуализаций.

От первого лица

Александр Рыжков, руководитель команды разработки LightAutoML в Sber AI Lab, ментор-эксперт AI Talent Hub:

«Ребята показывают отличные результаты! В качестве планов по дальнейшему развитию решения планируется более глубокая интеграция с нашим opensource фреймворком для автоматического построения моделей машинного обучения LightAutoML (LAMA). Он доступен в PyPI и на Github для построения системы, способной быстро решать DS задачи с высоким качеством в режиме end-2-end (от идеи до продакшн)». 

Сейчас наш проект находится на стадии активной разработки. Мы хотим развить MLJET, сделать его полноценным продуктом и рассказать об этом всему миру. И это только начало нашего ML-пути! 

Ссылка на Github: github.com/qnbhd/mljet

Ссылка на документацию: mljet.readthedocs.io/en/latest/

P.S. Если хотите делать такие же крутые проекты, то вэлком к нам в онлайн-магистратуру AI Talent Hub