Как работают Model Serving инструменты изнутри. Пишем свой на Python и Docker с нуля

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

В этом обсуждении мы сосредоточимся на концептуальной структуре, которую будем называть `servingml`, которая демонстрирует базовые элементы и рабочий процесс типичного сервинга модели.

Ключевые компоненты servingml Framework

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

1. Core. В основе `servingml` лежит базовый класс для экземпляров модели, который формирует основу фреймворка. Это ядро является универсальным и мы добавим интеграцию с scikit-learn в качестве примера, и возможность добавления других, что будет гарантировать, что широкий спектр моделей машинного обучения можно легко адаптировать и обслуживать с помощью нашей платформы.

2. CLI and Build Tools. Чтобы облегчить переход от обученной модели к развертываемому артефакту, мы создадим интерфейс командной строки (CLI) для `servingml`. Задача этого CLI — собрать необходимый код и зависимости в одном каталоге. Эта консолидация является важным шагом, упрощающим процесс контейнеризации модели и ее среды.

3. Project Configuration via YAML. Примечательной особенностью многих Model Serving инструментов является использование файла YAML для конфигурации проекта. В этом файле указаны важные сведения, такие как требования к модели, имя модели и путь к сервису. Этот метод настройки упрощает процесс установки, обеспечивая четкое и удобное определение параметров модели и сервиса.

4. Docker Integration. Центральное место в процессе развертывания занимает использование Docker. Платформа `servingml` использует шаблон Dockerfile (Dockerfile.j2), в котором описано, как упаковать модель вместе с ее зависимостями в контейнер Docker. Этот процесс гарантирует, что модель может выполняться последовательно в любой среде, устраняя проблемы, возникающие из-за различных зависимостей или конфигураций.

5. ServingML Server and Deployment. Последней частью головоломки является сервер ServingML. Этот сервер, работающий внутри контейнера Docker, отвечает за получение каталога, содержащего модель, файл Dockerfile и любой дополнительный необходимый код. Затем на основе этого создается образ Docker и контейнер из этого образа. После чего модель сможет предоставлять прогнозы через REST API для обычных пользователей.

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

Реализация Model Serving класса для кастомной модели

Изначально, мы должны реализировать базовый класс ModelServer, который будет наследоваться кастомным классом для создания load, predict, preprocess и postprocess методов. ModelServer на основе этих методов будет создавать функцию для обработки api запросов приложением Starlette.

service.py

  import numpy as np

import Starlette

from typing import Dict, List

from servingml import ModelServer
from servingml.frameworks.sklearn import load_model


class ModelClass(ModelServer):
    def load(self):
        self.model = load_model("iris_clf")

    def predict(self, body: Dict) -> List:
        """Generate model predictions from sample"""
        sample_input = np.asarray(body['inputs'])
        result: np.ndarray = self.model.predict(sample_input)
        return result.tolist()

    def postprocess(self, sample_output: List) -> str:
        """Make postprocessing for returning class name"""
        target_names = ['setosa', 'versicolor', 'virginica']
        # Convert numeric prediction to species name
        predicted_species = target_names[sample_output[0]]
        return predicted_species


model_class = ModelClass(
    model_name="iris_clf",
)

app: Starlette = model_class.asgi_app

Это наш кастомный код. Здесь мы загружаем обученную модель, которая хранится в нашем model store, создаем функции predict, для создания предсказания и postprocess, которая возвращает текстовый класс, вместо вероятностей.

Переменная app здесь, это Starlette приложение, которое затем будет использоваться uvicorn для создания asgi сервера.

Это ModelServer класс

servingml._model_server.py

from typing import Dict, Any
from starlette.applications import Starlette
from starlette.requests import Request
from starlette.responses import JSONResponse
from starlette.routing import Route
import logging

class ModelServer:
    """Class that creates a Starlette app for serving a machine learning model."""
    
    def __init__(self, model_name: str) -> None:
        self.model_name = model_name
        self.model = self.load()
        logging.basicConfig(level=logging.INFO)

    @property
    def asgi_app(self) -> Starlette:
        return Starlette(routes=[
            Route(f"/v2/models/{self.model_name}/infer", self._predict_fn, methods=["POST"]),
        ])

    async def _predict_fn(self, request: Request) -> JSONResponse:
        """Handler for the Starlette application."""
        try:
            data = await request.json()
            if not data:
                return JSONResponse({"error": "Invalid input data"}, status_code=400)
            processed_data = self.preprocess(data)
            prediction = self.predict(processed_data)
            result = self.postprocess(prediction)
            return JSONResponse({"result": result})
        except Exception as e:
            logging.error(f"Error during prediction: {e}")
            return JSONResponse({"error": "Error processing request"}, status_code=500)

    def load(self) -> Any:
        """Model loading operation."""
        raise NotImplementedError()

    def preprocess(self, body: Dict) -> Any:
        """Preprocess the event body before validation and action."""
        return body

    def postprocess(self, result: Any) -> Any:
        """Postprocess the prediction before returning response."""
        return result

    def predict(self, data: Any) -> Any:
        """Model prediction operation."""
        raise NotImplementedError()

Здесь мы создаем следующий endpoint для модели: /v2/models/{self.model_name}/infer, с использованием V2 Inference Protocol, который является стандартом для всех современных Model Serving инструментов, таких как Seldon Core, Bentoml, Nvidia Triton, Onnx Runtime Server и др.

Интеграция с scikit-learn

Суть интеграции в нашем примере, это возможность простого сохранения обученной scikit-learn модели в локальный model store, и ее загрузки в память для инференса.

servingml.frameworks.sklearn.py

import os
import joblib

from typing import Union

from sklearn.base import BaseEstimator
from sklearn.pipeline import Pipeline

from servingml.constants import MODEL_STORE_SKLEARN, MODEL_STORE_SKLEARN_WORKING


SklearnModel = Union[BaseEstimator, Pipeline]


def save_model(model_name: str, model: SklearnModel) -> None:
    if not os.path.exists(MODEL_STORE_SKLEARN):
        os.makedirs(MODEL_STORE_SKLEARN)
    model_path = os.path.join(MODEL_STORE_SKLEARN, f"{model_name}.pkl")
    joblib.dump(model, model_path)


def load_model(model_name: str) -> SklearnModel:
    print("MODEL_STORE_SKLEARN_WORKING", MODEL_STORE_SKLEARN_WORKING)
    model_path = os.path.join(MODEL_STORE_SKLEARN_WORKING, f"{model_name}.pkl")
    if not os.path.exists(model_path):
        raise ValueError(
            f"Model {model_name} is not found at the sklearn model store. Make sure you saved it first."
        )
    return joblib.load(model_path)

Где константы выглядят следующим образом

servingml.frameworks.constants.py

import os


home_directory = os.path.expanduser("~")
SERVINGML_WORKING_DIR = os.path.join(home_directory, "servingml/svc")
MODEL_STORE_DIR = os.path.join(home_directory, "servingml/model_store")
MODEL_STORE_SKLEARN = os.path.join(MODEL_STORE_DIR, "sklearn")
MODEL_STORE_SKLEARN_WORKING = "./model_store/sklearn"

Это директорая проекта, которая хранит все нужные артифакты для сервинга модели, который мы рассмотрим следующей секции про cli.

Теперь, когда мы имеем интеграцию с scikit-learn мы можем обучить модель и сохранить его в нашем локальном model store

download_model.py

from sklearn import svm
from sklearn import datasets

from servingml.frameworks.sklearn import save_model

# Load training data set
iris = datasets.load_iris()
X, y = iris.data, iris.target

# Train the model
clf = svm.SVC(gamma='scale')
clf.fit(X, y)

# Save model to the local Model Store
saved_model = save_model("iris_clf", clf)

CLI для создания проекта и деплоя

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

В cli у нас есть две команды build, которая создает Dockerfile для модели и собирает все необходимые артефакты для него, и deploy, который отправляет это zip архив на servingml server, который может работать как локально, так и удаленно.

Во-первых у нас есть функция, которая создает Dockerfile из шаблона

def _generate_dockerfile(app_module: str, requirements: List[str], output_path="Dockerfile") -> None:
    # Load the template environment
    env = Environment(loader=FileSystemLoader('.'))
    template = env.get_template('servingml/Dockerfile.j2')

    # Render the template with variables
    rendered_dockerfile = template.render(app_module=app_module, requirements=requirements)

    # Write the rendered template to a file
    with open(os.path.join(SERVINGML_WORKING_DIR, output_path), 'w') as f:
        f.write(rendered_dockerfile)

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

servingml/Dockerfile.j2

# Use an official Python runtime as a parent image
FROM python:3.9-slim

# Set environment varibles
ENV PYTHONDONTWRITEBYTECODE 1
ENV PYTHONUNBUFFERED 1

# Set work directory
WORKDIR /usr/src/app

# Install any needed packages specified directly
RUN pip install --no-cache-dir {{ requirements | join(" ") }}

RUN pip install uvicorn starlette

# Copy the current directory contents into the container at /usr/src/app
COPY . .

# Run the application
CMD ["uvicorn", "{{ app_module }}", "--host", "0.0.0.0", "--port", "8090"]

Комманда build, определенная с помощью click библиотеки, выполняет следующие операции:

1. Настройка рабочего каталога. Проверяет, существует ли SERVINGML_WORKING_DIR. Если нет, то он создает его. Если он существует, каталог удаляется и создается заново. Это обеспечивает чистоту состояния.

2. Чтение конфигурации YAML. Открывает указанный служебный файл (по умолчанию — servingfile.yaml) и загружает его содержимое.
   Это содержимое преобразуется в объект SimpleNamespace для облегчения доступа к атрибутам. Данный файл yaml будет показан ниже

3. Создание Dockerfile. Вызывает _generate_dockerfile() с информацией из файла YAML (именем службы и пакетами) для создания Dockerfile.

4. Логика копирования и включения файлов. Скрипт просматривает файлы в контексте сборки (build_ctx), проверяя каждый файл на соответствие шаблонам включения и исключения из файла YAML.
   Файлы, соответствующие шаблонам включения и не соответствующие шаблонам исключения, копируются в SERVINGML_WORKING_DIR.

5. Обработка модели. Имя модели и структура извлекаются из конфигурации YAML.
   Для моделей scikit-learn (if framework == "sklearn":) скрипт проверяет, существует ли файл модели в MODEL_STORE_SKLEARN.
   Затем файл модели копируется в SERVINGML_WORKING_DIR.

6. Архивирование. Наконец, сценарий создает zip-архив с именем "master.zip" в SERVINGML_WORKING_DIR, который включает файл Dockerfile, необходимые файлы и модель. Этот архив затем можно использовать для развертывания модели.

servingml_cli/build.py

import click
import yaml
import os
import shutil
import fs
import fnmatch

from typing import List
from fs.copy import copy_file
from types import SimpleNamespace

from jinja2 import Environment, FileSystemLoader

from constants import SERVINGML_WORKING_DIR, MODEL_STORE_SKLEARN, MODEL_STORE_SKLEARN_WORKING


def _generate_dockerfile(app_module: str, requirements: List[str], output_path="Dockerfile") -> None:
    # Load the template environment
    env = Environment(loader=FileSystemLoader('.'))
    template = env.get_template('servingml/Dockerfile.j2')

    # Render the template with variables
    rendered_dockerfile = template.render(app_module=app_module, requirements=requirements)

    # Write the rendered template to a file
    with open(os.path.join(SERVINGML_WORKING_DIR, output_path), 'w') as f:
        f.write(rendered_dockerfile)


@click.command()
@click.option("--build_ctx", type=click.Path(), default=".", show_default=True, help="Path to the build context")
@click.option('--servingfile', 
              default="servingfile.yaml")
def build(build_ctx: str, servingfile: str) -> None:
    """Build Dockerfile for model"""
    
    # Create project dir
    if not os.path.exists(SERVINGML_WORKING_DIR):
        os.makedirs(SERVINGML_WORKING_DIR)
    else:
        try:
            shutil.rmtree(SERVINGML_WORKING_DIR)
            os.makedirs(SERVINGML_WORKING_DIR)
        except OSError as e:
            print(f"Error: {SERVINGML_WORKING_DIR} : {e.strerror}")
    with open(servingfile, 'r') as file:
        data = yaml.safe_load(file)
        spec = SimpleNamespace(**data)

    # Generate dockerfile using spec attributes
    _generate_dockerfile(spec.service, spec.packages)

    # Add files to the target directory
    ctx_fs = fs.open_fs(build_ctx)
    target_fs = fs.open_fs(SERVINGML_WORKING_DIR)
    for dir_path, _, files in ctx_fs.walk():
        for f in files:
            path = fs.path.combine(dir_path, f.name).lstrip("/")
            if any(fnmatch.fnmatch(path, pat) for pat in spec.exclude):
                continue  # Skip this file
            # Check if the file matches any of the include patterns
            if any(fnmatch.fnmatch(path, pat) for pat in spec.include):
                target_fs.makedirs(dir_path, recreate=True)
                copy_file(ctx_fs, path, target_fs, path)

    # Copy model to production directory
    framework, model_name = spec.model_name.split(":")
    if framework == "sklearn":
        model_path = os.path.join(MODEL_STORE_SKLEARN, f"{model_name}.pkl")
        if not os.path.exists(model_path):
            raise ValueError(
                f"Model {model_name} is not found at the sklearn model store. Make sure you saved it first."
            )
        target_path = os.path.join(MODEL_STORE_SKLEARN_WORKING, f"{model_name}.pkl")
        if not os.path.exists(MODEL_STORE_SKLEARN_WORKING):
            os.makedirs(MODEL_STORE_SKLEARN_WORKING)
        shutil.copyfile(model_path, target_path)
    
    # Save all as zip archive for further deploying
    shutil.make_archive("master", 'zip', SERVINGML_WORKING_DIR)


if __name__ == '__main__':
    build()

По сути, этот скрипт представляет собой инструмент для автоматизации упаковки моделей в среду Docker, специально адаптированный к потребностям и зависимостям модели. Он инкапсулирует такие процессы, как создание файлов Dockerfile, управление файлами и упаковку моделей, в единый оптимизированный рабочий процесс, что способствует эффективному и безошибочному развертыванию моделей.

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

servingfile.yaml

service: "service:app"  # Starlette application
include:
- "*.py"  # A pattern for matching which files to include in the Bento
exclude:
- "venv/*"
- "__pycache__/*"
- "servingml_server/*"
- "servingml_cli/*"
- "*.pyc"
packages:  # Additional pip packages required by the Service
- scikit-learn
- pandas
model_name: "sklearn:iris_clf" # Saved model name

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

Теперь, когда у нас есть упакованный проект со всем необходимым, а именно с Dockerfile, кодом и моделью, мы хотим отправить этот проект на наш servingml server, который работает локально или на удаленном сервере.

servingml_cli/deploy.py

import requests
import yaml
import click
import json

from types import SimpleNamespace


@click.command()
@click.option('--servingfile', 
              default="servingfile.yaml")
@click.option('--host', 
              default="localhost")
def deploy(servingfile: str, host: str) -> None:
    """Deploy model using Dockerfile local or remote"""
    with open(servingfile, 'r') as file:
            data = yaml.safe_load(file)
            spec = SimpleNamespace(**data)
    # Copy model to production directory
    _, model_name = spec.model_name.split(":")
    url = f"http://{host}:8000/deploy_project"  # Modify with your actual endpoint
    files = {'file': open('master.zip', 'rb')}
    json_data = json.dumps({'modelname': model_name, 'port': 8090})
    data = {'data': json_data}
    response = requests.post(url, data=data, files=files)
    print("response", response.text)


if __name__ == '__main__':
    deploy()

Мы имеем api http://{host}:8000/deploy_project, по которому работает servingml server, в который помимо проекта мы передаем название модели в качестве названия Docker Image, и порт на котором запускать данный образ в виде контейнера. Стоит отметить, port должен быть таким же, как и в Dockerfile.j2, чтобы совпадал внутренний порт работы uvicorn сервера и внешний контейнера.

Теперь давайте настроим cli из этих методов, чтобы иметь возможность вызывать их servingml build и servinml deploy. Для этого нам потребуются два следующих файла

servingml_cli/cli.py

import click
from build import build
from deploy import deploy

@click.group()
def cli() -> None:
    """ServingML Command Line Interface."""
    pass

cli.add_command(build)
cli.add_command(deploy)

if __name__ == "__main__":
    cli()

servingml_cli/setup.py

from setuptools import setup

setup(
    name='servingml',
    version='0.1',
    py_modules=['cli', 'build', 'deploy'],
    install_requires=[
        'Click',
    ],
    entry_points='''
        [console_scripts]
        servingml=cli:cli
    ''',
)

После чего, находясь в директории servingml_cli, вызываем следующую комманду

pip install --editable .

Создаем сервер servingml

Сервер, который принимает упакованный проект и создаем из него docker образ, а затем и контейнер, это ключевая часть нашей платформы. Для создания api endpoint /deploy_project, мы будем использовать fastapi.

Ниже, deploy_project функция сначала распоковывает наш проект во временную директорию, проверяет есть в ней Dockerfile, и отправляет путь к нему с доп. параметрами в backgound task, для создания docker image и docker container.

servingml_server/server.py

@app.post("/deploy_project")
async def deploy_project(
    background_tasks: BackgroundTasks, 
    data: str = Form(...), 
    file: UploadFile = File(...)
) -> Dict[str, str]:
    try:
        build_data = json.loads(data)
        build_data = BuildData(**build_data)
    except (json.JSONDecodeError, ValidationError) as e:
        return {"error": f"Invalid input data: {str(e)}"}
    # Temporary directory to extract files
    temp_dir = "temp_docker_build"
    os.makedirs(temp_dir, exist_ok=True)

    temp_file_path = os.path.join(temp_dir, file.filename)

    # Save the uploaded file to disk
    with open(temp_file_path, "wb") as buffer:
        buffer.write(file.file.read())  # Read from SpooledTemporaryFile and write to disk

    # Extract the ZIP file
    with zipfile.ZipFile(temp_file_path, 'r') as zip_ref:
        zip_ref.extractall(temp_dir)

    # Check if Dockerfile is present
    if not os.path.exists(os.path.join(temp_dir, 'Dockerfile')):
        return {"error": "Dockerfile not found in the zip"}

    # Build Docker image
    image_tag = f"{build_data.modelname}:latest"  # You can customize this
    background_tasks.add_task(
        background_docker_build_task, temp_dir, image_tag, build_data.port, "docker_build_log.txt",
    )
    return {"message": "Docker build started"}

Для создания subprocess мы используем здесь asyncio

async def exec_command(command, log_file_path):
    process = await asyncio.create_subprocess_shell(
        command,
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE,
    )

    stdout, stderr = await process.communicate()

    with open(log_file_path, "a") as log_file:
        if stdout:
            log_file.write(f"[stdout]\n{stdout.decode()}\n")
        if stderr:
            log_file.write(f"[stderr]\n{stderr.decode()}\n")


async def docker_build_task(temp_dir: str, image_tag: str, port: int, log_file_path: str):
    try:
        # Step 1: Build the Docker image asynchronously
        build_command = f"docker build -t {image_tag} {temp_dir}"
        await exec_command(build_command, log_file_path)

        # Step 2: Run a container from the built image asynchronously
        run_command = f"docker run -d -p {port}:{port} {image_tag}"
        await exec_command(run_command, log_file_path)

    except Exception as e:
        # Log any errors during the build or run process
        with open(log_file_path, "a") as log_file:
            log_file.write(f"Error during Docker operations: {e}\n")
    finally:
        # Cleanup
        shutil.rmtree(temp_dir)


def background_docker_build_task(temp_dir, image_tag, port, log_file_path):
    asyncio.run(docker_build_task(temp_dir, image_tag, port, log_file_path))

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

servingml_server/Dockerfile

# Use an official Python runtime as a parent image
FROM python:3.9-slim

# Set the working directory in the container
WORKDIR /usr/src/app

# Copy the local code to the container
COPY . .

# Install Docker CLI
RUN apt-get update && apt-get install -y docker.io

# Install any needed packages specified in requirements.txt
RUN pip install --no-cache-dir fastapi[all] uvicorn

# Make port 8000 available to the world outside this container
EXPOSE 8000

# Define environment variable
ENV NAME World

# Run app.py when the container launches
CMD ["uvicorn", "server:app", "--host", "0.0.0.0", "--port", "8000"]

Создадим image для servingml server

docker build -t servingml_server .

Для корректной работы сервера, при запуска контейнера нам нужно монтировать сокет Docker хоста в контейнер. Делая это, мы предоставляем контейнеру доступ к daemon Docker на хосте, что позволяет ему запускать другие контейнеры Docker и управлять ими. По сути, контейнер сможет создавать родственные контейнеры и управлять ими в хост-системе.

docker run -d -v /var/run/docker.sock:/var/run/docker.sock -p 8000:8000 servingml_server

Отлично! Теперь у нас есть servingml server готовый принимать наши проекты для создания из них docker conrainer's. Что в нашем случае, производить сервинг и деплой моделей машинного обучения.

Деплой проекта модели через cli и тест api endpoint

Теперь, чтобы задеплоить проект модели, нам нужно сначала создать его, сделав из него архив

servingml build --build-ctx project_path

И сделать деплой архивированной проекта с передав его нашему servingml server

servingml deploy --host your_host

Деплой может занять несколько минут в зависимости от проекта. Для проверки готов ли контейнер мы можем либо смотреть его с помощью docker ps, проверяя используемые image, либо копировать логи.

docker cp container_name:/usr/src/app/docker_build_log.txt ./docker_build_log.txt

Если файла логов в контейнере еще нет, значит процесс создния image все еще идет.

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

invoke.py

import requests

# Endpoint URL
url = "http://localhost:8090/v2/models/iris_clf/infer"

sample_input = [[5.1, 3.5, 1.4, 0.2]]

# Constructing the payload
payload = {
    "inputs": sample_input
}

# Make a POST request
response = requests.post(url, json=payload)

# Checking response
if response.status_code == 200:
    print("Success:", response.json())
else:
    print("Error:", response.text)

В качестве ответа мы должны получить названия класса iris.

Заключение

В основном я опирался на исходный код bentoml и mlrun. Этот код не может быть в production, даже несмотря на то, что он рабочий, стоит использовать уже готовые инструменты с гораздо большим фукнционалом и надежностью. Если вы делаете деплой модели на kubernetes я советую использовать seldon core с его экосистемой, например alibi detect, которая позволяет отслеживать data drift, если же вы делаете деплой на обычный сервер, который не является частью кластера kubernetes используйте bentoml.

Исходный код для проекта можно посмотреть здесь - https://github.com/Nikitala0014/servingml

Телеграм для связи - @NLavrenov00

Надеюсь вам было также интересно, как и мне во время изучения и создания этой статьи :)