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

Знакомство с проблемой

В коммерческой разработке многие сценарии использования машинного обучения подразумевают мультитенантную архитектуру и требуют обучения отдельной модели для каждого клиента и/или пользователя.

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

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

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

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

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

Решение

Решение предполагает, что обучение модели машинного обучения выполняется отдельно от обслуживания данной модели. Например, задание Airflow выполняет обучение модели и сохраняет ее в S3, в таком случае единственная ответственность ML сервера — это прогнозы.

Обученная модель ML — это просто файл на диске, поэтому нам нужно сохранить файл и сопоставление: идентификатор пользователя -> идентификатор модели.

Компоненты решения

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

• Model — абстрактная модель, предоставляющая прогноз; его реализация может быть SklearnModel, TensorFlowModel, MyCustomModel и т. д.

• ModelInfoRepository — абстрактный репозиторий, который предоставляет сопоставления userid -> modelid. Например, он может быть реализован как SQAlchemyModelInfoRepository.

• ModelRepository — абстрактный репозиторий, который может возвращать модель по ее ID. Это может быть FileSystemRepository, S3Repository или любая другая реализация репозитория.

from abc import ABC


class Model(ABC):
    @abstractmethod
    def predict(self, data: pd.DataFrame) -> np.ndarray:
        raise NotImplementedError
 

class ModelInfoRepository(ABC):
    @abstractmethod
    def get_model_id_by_user_id(self, user_id: str) -> str:
        raise NotImplementedError
 

class ModelRepository(ABC):
    @abstractmethod
    def get_model(self, model_id: str) -> Model:
        raise NotImplementedError

Реализация

Теперь предположим, что мы обучили модель sklearn, которая хранится в Amazon S3 с сопоставлениями userid -> modelid, определенными в базе данных.

class SklearnModel(Model):
    def __init__(self, model):
        self.model = model
 

    def predict(self, data: pd.DataFrame):
        return self.model.predict(data)
 

class SQAlchemyModelInfoRepository(ModelInfoRepository):
    def __init__(self, sqalchemy_session: Session):
        self.session = sqalchemy_session
 

    def get_model_id_by_user_id(user_id: str) -> str:
        # implementation goes here, query a table in any Database

      
class S3ModelRepository(ModelRepository):
    def __init__(self, s3_client):
        self.s3_client = s3_client
 

    def get_model(self, model_id: str) -> Model:
        # load and deserialize pickle from S3, implementation goes here

Это делает реализацию сервера чрезвычайно простой:

def make_app(model_info_repository: ModelInfoRepository,
    				 model_repsitory: ModelRepository) -> Flask:
    app = Flask("multi-model-server")
    
    @app.predict("/predict/<user_id>")
    def predict(user_id):
        model_id = model_info_repository.get_model_id_by_user_id(user_id)
 
        model = model_repsitory.get_model(model_id)
 
        data = pd.DataFrame(request.json())
 
        predictions = model.predict(data)
 
        return jsonify(predictions.tolist())
 
    return app

Обратите внимание, что благодаря абстракциям сервер Flask полностью независим от конкретной модели и реализации хранилища; мы можем заменить sklearn на TensorFlow и S3 на локальную папку, при этом в коде сервера Flask не меняются строки.

Замечание о кешировании

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

from cachetools import Cache
 
class CachedModelRepository(ModelRepository):
    def __init__(self, model_repository: ModelRepository, cache: Cache):
        self.model_repository = model_repository
        self.cache = cache
 
    @abstractmethod
    def get_model(self, model_id: str) -> Model:
        if model_id not in self.cache:
            self.cache[model_id] = self.model_repository.get_model(model_id)
        return self.cache[model_id]

Пример использования:

from cachetools import LRUCache
 
model_repository = CachedModelRepository(
    S3ModelRepository(s3_client),
    LRUCache(max_size=10)
)

Перед выходом в продакшен

Такой многомодельный сервер - одна из многих частей, необходимых для запуска приложений производственного уровня с возможностями машинного обучения. Безопасность корпоративного уровня, масштабируемость, MLOps и т. д. могут быть даже более важными для успеха и надежности проекта, чем немного более точная модель машинного обучения. Всегда помните о гениальном правиле №4 от Google: пусть первая модель будет простой, а инфраструктура - правильной.