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

Вы когда-нибудь пробовали подружить конкретную версию PyTorch с определенной версией CUDA, приправив это нужными системными библиотеками в Dockerfile? Если да, то вы знаете, что это прямой путь в «ад зависимостей». Обычно это заканчивается тем, что дата-сайентист идет к девелопсу, и они вдвоем полдня пытаются заставить модель просто запуститься внутри контейнера.

Инструмент Cog создали люди, которые съели на этом собаку. Один из авторов, Бен Фиршман, когда-то придумал Docker Compose, а Андреас Янссон занимался деплоем ML-моделей в Spotify. Они решили, что пора избавить мир от написания бесконечных и кривых Dockerfile для нейросетей.

Что такое Cog и зачем он вам

Если вкратце, Cog — это опенсорсный инструмент, который берет ваш код на Python и превращает его в готовый к продакшену Docker-контейнер. При этом вам не нужно знать, как писать Dockerfile. Вы описываете окружение в простом YAML-файле, а Cog сам подбирает правильные базовые образы от Nvidia, настраивает CUDA, устанавливает Python и нужные библиотеки.

Главная фишка здесь в стандартизации. На выходе вы получаете не просто «какой-то контейнер», а полноценный HTTP-сервер на базе Rust и Axum, который готов принимать запросы к вашей модели.

Как это работает на практике

Вместо того чтобы возиться с инструкциями RUN apt-get install, вы создаете файл cog.yaml. Выглядит это примерно так:

build:
  gpu: true
  system_packages:
    - "libgl1-mesa-glx"
    - "libglib2.0-0"
  python_version: "3.13"
  python_requirements: requirements.txt
predict: "predict.py:Predictor"

Заметьте, здесь нет ни слова о том, какую версию CUDA ставить. Cog сам знает, какие комбинации CUDA, cuDNN и PyTorch совместимы между собой. Это экономит кучу нервов.

Далее вам нужно обернуть логику предсказания в класс. Это стандартный интерфейс, который позволяет Cog понять, как загружать веса модели и как пробрасывать входные данные.

from cog import BasePredictor, Input, Path
import torch

class Predictor(BasePredictor):
    def setup(self):
        # Модель грузится один раз при старте контейнера
        self.model = torch.load("./weights.pth")

    def predict(self, image: Path = Input(description="Входная картинка")) -> Path:
        # Здесь происходит магия инференса
        output = self.model(preprocess(image))
        return postprocess(output)

Автоматический API и валидация

Интересная деталь: Cog использует стандартные типы Python для определения входов и выходов. На основе вашего кода он автоматически генерирует OpenAPI-схему. Это значит, что если ваша модель ожидает картинку, а ей прислали текст, Cog сам вернет ошибку валидации еще до того, как код модели начнет выполняться.

Когда вы запускаете контейнер, Cog поднимает высокопроизводительный сервер. Вам не нужно прикручивать Flask или FastAPI вручную. Просто отправляете POST-запрос с JSON, и всё работает.

Кому это пригодится

Я часто вижу, как хорошие модели пылятся в репозиториях, потому что их чертовски сложно запустить локально. Cog решает эту проблему.

Исследователям: можно быстро упаковать модель и отдать её инженерам или выложить в общий доступ, не объясняя, как настраивать окружение.
Инженерам: вы получаете предсказуемый артефакт (Docker-образ), который можно деплоить куда угодно — хоть в Kubernetes, хоть на сервис Replicate.
Командам: Cog кэширует зависимости гораздо эффективнее, чем стандартный Docker build, поэтому пересборка образов при изменении кода происходит быстрее.

Стоит ли пробовать

Проект уже достаточно зрелый, у него почти 10 тысяч звезд на GitHub, и за ним стоит живое комьюнити. Из минусов можно отметить, что Cog довольно жестко навязывает свою структуру проекта. Если у вас очень специфический процесс сборки, который не укладывается в их YAML-конфиг, придется попотеть.

Но для 90% задач машинного обучения — от классификации картинок до работы с LLM — это, пожалуй, самый быстрый способ превратить скрипт в рабочий сервис. Если вам надоело воевать с драйверами Nvidia внутри Docker, Cog определенно стоит того, чтобы потратить на него вечер.

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