Контейнеризация моделей машинного обучения на GPU: архитектура, производительность и трёхступенчатый build pipeline с Docker+NVIDIA

Категория: Статьи • 24 June 2025 • 5 мин чтения

Контейнеризация ML-моделей давно перестала быть экзотикой. Однако в 2025 году настоящая индустриальная зрелость достигается только тогда, когда ML-продуктов становится десятки, inference идёт на GPU, обновления катятся в production ежедневно, а инфраструктура автосборки и логирования не ломается от одного слова “cuda”. Контейнеризация таких систем — это не просто docker build и FROM pytorch/pytorch. Это архитектурное ремесло.

Ниже — подробная, практическая архитектура контейнеризации ML-инференса для GPU, форматируемая через современный stack: CUDA-aware images, mulitstage Docker build, оптимизация размера образов, reproducible builds и runtime стратегии.

Ключевая задача: reproduce + run на GPU + обновляемо + быстро собирается

Цель — собрать стабильный, предсказуемый, лёгкий контейнер, поддерживающий:

  1. Воспроизводимость (на разных CI runner’ах)
  2. GPU-инференс (CUDA Runtime)
  3. Горячее обновление модели без перекомпиляции всего контейнера
  4. Минимальный размер финального образа
  5. Возможность хостинга в регистри без платной тарификации за jumbo images

Подход: трёхстадийная сборка образа

Стадия 1: Builder — Python deps + CUDA headers из исходников

Dockerfile
# Builder stage — полный dev stack (CUDA toolkits, compilers, headers)
FROM nvidia/cuda:12.2.0-devel-ubuntu22.04 AS builder

ENV DEBIAN_FRONTEND=noninteractive

RUN apt-get update && apt-get install -y \
    python3.10 python3.10-venv python3-pip build-essential \
    git wget curl cmake \
    && rm -rf /var/lib/apt/lists/*

RUN python3.10 -m venv /opt/venv
ENV PATH="/opt/venv/bin:$PATH"

# pip с последним wheel
RUN pip install --upgrade pip wheel

# Копируем requirements.txt и устанавливаем в venv
COPY requirements.txt .
RUN pip install -r requirements.txt

Здесь собирается виртуальное окружение со всеми необходимыми зависимостями. Модули с C-расширениями (например, torchvision с компиляцией CUDA-расширений) важно собирать именно на этой стадии — они часто требуют заголовочных файлов cuda.h, которые на runtime не нужны.

Стадия 2: Base runtime GPU — минимальный CUDA runtime + env

Dockerfile
# Runtime stage — минимальный CUDA runtime + system deps
FROM nvidia/cuda:12.2.0-runtime-ubuntu22.04 AS base

ENV DEBIAN_FRONTEND=noninteractive

RUN apt-get update && apt-get install -y \
    python3.10 libgl1 libglib2.0-0 \
    && rm -rf /var/lib/apt/lists/*

COPY --from=builder /opt/venv /opt/venv
ENV PATH="/opt/venv/bin:$PATH"

Контейнер на этой стадии содержит минимальный runtime — только CUDA runtime и нужные libs. Отдельно копируется виртуальное окружение Python из стадии builder. Это уменьшает final image size в 3-5 раз.

Стадия 3: Final (optional) — прокладка модели, API endpoint

Dockerfile
FROM base AS final

WORKDIR /app

# entrypoint
COPY ./serve.py .
COPY ./custom_module.py .

# runtime model weights (горячее обновление возможно через volume mount)
COPY ./models /models

CMD ["python", "serve.py"]

В этой стадии добавляются бизнес-специфичные части: API-обёртка, логика inference, логеры и, при желании, модель. На практике модель можно монтировать как volume для избавления от rebuild при hot-swap проверках в staging-средах.

Выгоды от мультиярусной сборки

  • Финальный образ ~400MB вместо типичных 3-4GB (в 2025 году это уже важно для edge inference)
  • Повторное использование builder stage в CI. Можно собирать множество моделей против одного builder snapshot
  • Компиляция зависимостей torchvision, xformers, clip etc. один раз, не при каждом docker build context

Работа с GPU: запуск inference в контейнере с nvidia-container-runtime

Чтобы контейнер мог использовать GPU, требуется:

  1. Установленный NVIDIA Container Toolkit (nvidia-container-runtime)
  2. docker run с параметром:
bash
docker run --gpus=all --rm my-inference-container

Проверка GPU:

python
import torch
print(torch.cuda.is_available())           # True
print(torch.cuda.get_device_name(0))       # NVIDIA A100-SXM4-40GB

Если is_available() возвращает False, проблема почти всегда в mismatch между host NVIDIA driver version и CUDA runtime в контейнере. Пример:

  • Host driver: 535.x
  • Container CUDA runtime: 12.2
  • Соответствие: OK (главное — driver >= runtime)

Под капотом: host драйвер обеспечивает реализацию ядра, контейнеру требуется лишь API runtime. Совпадения “точь в точь” не обязательно.

Хранение модели: монтировать или bake в слой?

Есть два подхода к размещению модельных весов:

  1. Монтировать через Docker volume:

    • размонтируемо (удобно в staging)
    • минимальный размер образа
    • гибко для A/B тестов

  2. Bake внутрь образа:

    • пригодно для offline deployment (например, edge на клиентской машине)
    • проще “провести SDLC как артефакт”
    • можно подписывать и кэшировать как immutable content

Обычно в prodоритет bake + версии через CI. В staging и early debug достаточно монтирования:

bash
docker run -v $(pwd)/models:/models --gpus=all myimage:latest

Что ломается чаще всего

1. Ошибки при запуске: failed to initialize CUDA context

Причины:

  • Контейнер собран на CUDA 12.x, а драйвер на хосте старее (например, 515). Обновить драйвер.

  • Контейнер не запущен с --gpus=all. Требуется явное указание GPU-доступа.

2. Не работает torchvision с CUDA

torchvision должен быть собран против конкретной сборки torch и поддерживать ту же версию CUDA. Легче всего — ставить их из одного wheel source.

Плохо:

Dockerfile
RUN pip install torch torchvision

Хорошо:

Dockerfile
RUN pip install torch==2.2.0+cu122 torchvision==0.17.0+cu122 -f https://download.pytorch.org/whl/torch_stable.html

3. Слишком большой образ (3+ ГБ)

Обычно из-за:

  • полных dev toolkits (gcc, компиляторы, cmake) в финальном образе
  • pip install в runtime, не в builder
  • кэширования локальных моделей в ~/.cache/torch, huggingface/transformers

Fix: использовать .dockerignore, вынести pip в фазу builder, удалять кэш перед final stage:

Dockerfile
RUN rm -rf ~/.cache

Рекомендации для CI / CD

  • Builder образ (stage 1) можно публиковать отдельно и переиспользовать с тегом ci::torch-2.2-cu122-builder
  • Final image публикуется с тегом версии модели ml-service:resnet50-v3.12
  • Используйте digest-based pull вместо latest
  • Проверяйте хэш моделей перед bake (например, через SHA256 в models/manifest.json)

Выводы

Контейнеризация inference-сервисов для GPU — деталезависимая инженерная задача. Без избыточности правильное построение образов требует разумного разделения стадий (build/runtime), контроля над CUDA runtime совместимостью, и осознанного выбора: baked модель vs. mounted. Архитектура, приведённая выше, масштабируется, ускоряет CI/CD, и сокращает время rollout'а с часов до минут — с полной поддержкой GPU в продакшн-средах.

Сложное поведение читается проще, когда инженеры рассматривают не только технические зависимости, но и свойства жизненного цикла моделей.

Будьте тактичны с CUDA, и она отплатит вам снижением latency на миллисекунды.