Настройка Docker-контейнеров для AI/ML-проектов

Настройка Docker-контейнеров для AI/ML-проектов

Docker для AI/ML решает проблему воспроизводимости окружения: один и тот же контейнер запускается одинаково на ноутбуке разработчика, CI-сервере и production GPU-кластере. Ключевые особенности: использование NVIDIA Container Toolkit для доступа к GPU, многостадийные сборки для минимизации размера образов, кэширование слоёв для ускорения CI.

NVIDIA Container Toolkit

# Установка для доступа GPU внутри контейнеров
distribution=$(. /etc/os-release; echo $ID$VERSION_ID)
curl -s -L https://nvidia.github.io/nvidia-docker/gpgkey | sudo apt-key add -
curl -s -L https://nvidia.github.io/nvidia-docker/$distribution/nvidia-docker.list \
  | sudo tee /etc/apt/sources.list.d/nvidia-docker.list
sudo apt install nvidia-container-toolkit
sudo systemctl restart docker

# Тест
docker run --gpus all nvidia/cuda:12.2.0-base-ubuntu22.04 nvidia-smi

Dockerfile для ML-проекта

# Multi-stage build: меньший финальный образ
FROM nvidia/cuda:12.2.0-cudnn8-devel-ubuntu22.04 AS builder

# Зависимости сборки
RUN apt-get update && apt-get install -y python3.11 python3-pip git \
    && rm -rf /var/lib/apt/lists/*

# Установка зависимостей отдельным слоем (кэшируется)
COPY requirements.txt .
RUN pip install --user --no-cache-dir -r requirements.txt

FROM nvidia/cuda:12.2.0-cudnn8-runtime-ubuntu22.04

RUN apt-get update && apt-get install -y python3.11 \
    && rm -rf /var/lib/apt/lists/*

# Копирование только установленных пакетов
COPY --from=builder /root/.local /root/.local
ENV PATH=/root/.local/bin:$PATH

WORKDIR /app
COPY src/ ./src/

# Non-root пользователь для безопасности
RUN useradd -m -u 1000 mluser
USER mluser

CMD ["python", "src/train.py"]

Docker Compose для локальной разработки

# docker-compose.yml
version: '3.8'
services:
  training:
    build: .
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: all
              capabilities: [gpu]
    volumes:
      - ./data:/app/data:ro          # Данные read-only
      - ./src:/app/src               # Код монтируется для hot-reload
      - ./outputs:/app/outputs       # Результаты
    environment:
      - MLFLOW_TRACKING_URI=http://mlflow:5000
    depends_on:
      - mlflow

  mlflow:
    image: python:3.11-slim
    command: mlflow server --host 0.0.0.0 --port 5000
    ports:
      - "5000:5000"
    volumes:
      - mlflow-data:/mlflow

volumes:
  mlflow-data:

Оптимизация размера образов

Типичные проблемы: образы PyTorch >10GB из-за включения CUDA dev библиотек и лишних пакетов. Решение:

• Использовать runtime тег вместо devel (разница 4-5GB)
• --no-cache-dir при pip install
• Удалять кэш apt после установки пакетов
• Использовать .dockerignore для исключения данных и виртуальных окружений

Итоговый образ для инференса: PyTorch ONNX Runtime + приложение = 2-3GB против 8-10GB при наивном подходе.