Используйте пользовательские образы Docker в AI Runtime

Important

Эта функция доступна в бета-версии. Чтобы использовать его, администратор рабочего пространства должен включить бета-функции AI Runtime и предпросмотр реестра артефактов Databricks на странице предварительного просмотра workspace.

AI Runtime может запускать пользовательский образ контейнера Docker, хранящийся в Artifact Registry. При необходимости используйте пользовательский образ:

  • Системные библиотеки или сложные зависимости, которые нельзя установить с environment.dependencies.
  • Воспроизводимая среда в области разработки, исследований и производства.
  • Приватные, одобренные организацией изображения, созданные вашей платформой или командой безопасности.

Prerequisites

Отправьте изображение в Реестр артефактов

Перед использованием пользовательского изображения с AI Runtime храните его в Реестре артефактов в том же рабочем пространстве, где вы отправляете нагрузку.

  1. Создайте или выберите каталог и схему Unity Catalog для изображения.

  2. Воспользуйтесь руководством «Начало работы с Artifact Registry», чтобы настроить аутентификацию Docker, назначить необходимые разрешения и загрузить образ.

  3. Обратите внимание на название изображения в каталоге Unity в следующем формате:

    <catalog>.<schema>.<image>:<tag>

    Например: main.ml.training:v1. Не включайте имя хоста реестра в конфигурацию рабочей нагрузки.

Tip

В качестве альтернативы используйте databricks air images push команду helper в Databricks CLI.

Используйте образ Docker в рабочей нагрузке

Укажите имя образа в Unity Catalog в YAML-файле рабочей нагрузки в разделе environment.unity_catalog_image:

experiment_name: my-dcs-training
environment:
  unity_catalog_image: main.ml.training:v1
compute:
  num_accelerators: 1
  accelerator_type: GPU_1xA10
command: python /app/train.py

В этом примере train.py запускается из /app на изображении. Чтобы загрузить код приложения отдельно без восстановления изображения, см. раздел «Запуск загруженного кода приложения».

При использовании собственного образа Docker environment.dependencies и environment.version не поддерживаются. Указание environment.unity_catalog_image с помощью любого поля вызывает ошибку. Если у вас есть дополнительные зависимости, установите пакеты в Dockerfile.

Отправьте рабочую нагрузку:

databricks air run -f workload.yaml -p my-databricks-profile

Профиль должен быть аутентифицирован в том же рабочем пространстве, где хранится изображение.

Переменные среды, передаваемые в ваш контейнер

Среда выполнения ИИ внедряет следующие переменные среды в каждый контейнер во время выполнения:

  • CODE_SOURCE_PATH: путь к загруженному коду приложения, когда code_source настроен.
  • NUM_NODES: общее количество узлов.
  • LOCAL_WORLD_SIZE: GPU на узел.
  • WORLD_SIZE: общее количество процессов.
  • POD_RANK: текущий ранг узла (индексирован 0). Также вставляется как NODE_RANK.
  • LOCAL_ADDR: локальный IP узла (только для нескольких узлов).
  • MASTER_ADDR: координатный адрес ранга 0 (только для нескольких узлов).
  • MASTER_PORT: координатный порт ранга 0 (только для нескольких узлов).

Examples

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

Запуск загруженного кода приложения

Используйте code_source для загрузки кода приложения отдельно от изображения. Вы можете редактировать и повторно отправлять код, не собирая изображение заново. Установите в образе Python-зависимости кода и системные зависимости.

Поместите train.py в местный src каталог рядом с workload.yaml. Следующая конфигурация загружает src и запускает его train.py в пользовательском образе:

experiment_name: my-dcs-uploaded-code
environment:
  unity_catalog_image: main.ml.training:v1
compute:
  num_accelerators: 1
  accelerator_type: GPU_1xA10
code_source:
  root_path: ./src
command: |-
  cd "$CODE_SOURCE_PATH"
  python3 train.py

root_path разрешается относительно workload.yaml. AI Runtime задаёт для $CODE_SOURCE_PATH путь к загруженному каталогу в контейнере. См. code_source для вариантов исходного кода.

Многоузловая конфигурация H100 с RDMA

Для многоузловых заданий H100, которым требуется полная пропускная способность сети на инстансах AWS p5, используйте в качестве основы для своего образа один из базовых образов Databricks с преднастроенными NCCL и EFA:

experiment_name: my-dcs-distributed
environment:
  unity_catalog_image: main.ml.training:v1
compute:
  num_accelerators: 16 # 2 nodes × 8 H100
  accelerator_type: GPU_8xH100
command: |-
  torchrun \
    --nnodes="${NUM_NODES}" \
    --nproc_per_node="${LOCAL_WORLD_SIZE}" \
    --node_rank="${POD_RANK}" \
    --rdzv_endpoint="${MASTER_ADDR}:${MASTER_PORT}" \
    /app/train.py

Создание собственного образа

При создании собственного образа Databricks рекомендует использовать навык databricks-ai-runtime с агентом для программирования или начать с базового образа Databricks.

Использование агента программирования

Установите навык Claude Code databricks-ai-runtime для получения пошаговых рекомендаций по Dockerfile, включая сборку с нуля, совместимость с CUDA/NCCL/EFA, распространённые проблемы и контрольный список перед сборкой. Для этого навыка требуется Интерфейс командной строки Databricks версии 1.0.0 или более поздней.

databricks aitools install --skills databricks-ai-runtime

Базовые образы Databricks

Databricks публикует базовые образы на Docker Hub по адресу databricksruntime/air с предварительно настроенными CUDA, NCCL и облачными сетевыми настройками (AWS EFA или Azure InfiniBand).

Тег Variant CUDA Используйте, если
dcs-base-azure-runtime Runtime 12 Установка предварительно созданных колес только
dcs-base-azure-devel Devel 12 Компиляция расширений CUDA (требуется nvcc)

Следующий Dockerfile добавляет PyTorch в базовый образ Databricks. Базовые образы предоставляют Python по /opt/venvадресу, управляемому uv. uv pip install По умолчанию нацелена на эту среду. Чтобы использовать другое окружение, создайте и активируйте venv перед запуском uv pip install.

Чтобы включить ваш обучающий скрипт в изображение, разместите train.py его рядом с Dockerfile. Dockerfile копирует его в /app/train.py. Если вы загружаете код приложения с code_source, опустите инструкцию COPY и оставьте train.py в исходном каталоге.

FROM databricksruntime/air:dcs-base-azure-runtime

RUN uv pip install --no-cache \
    torch==2.6.0 torchvision==0.21.0 torchaudio==2.6.0

RUN uv pip install --no-cache \
    transformers==4.45.0 \
    accelerate==0.34.0 \
    'mlflow>=3.6'

COPY ./train.py /app/train.py

Постройте изображение локально:

docker build --platform linux/amd64 -t my-training-image:v1 .

Затем следуйте разделу «Начать с Реестра артефактов », чтобы отметить и отправить изображение в Реестр артефактов. Используйте полученное <catalog>.<schema>.<image>:<tag> имя, как environment.unity_catalog_image в рабочей нагрузке YAML.

Tip

В качестве альтернативы используйте вспомогательную databricks air images pushкоманду в Databricks CLI и следуйте интерактивным подсказкам.

Limitations

  • Изображения должны храниться в Реестре артефактов в рабочем пространстве, куда вы отправляете рабочую нагрузку.
  • Размер изображения должен быть не более 20 ГБ.
  • WORKDIR не учитывается во время выполнения. Используйте абсолютные пути для файлов, включённых в образ. Например, используйте python /app/train.py, а не python train.py.
  • Вы не можете использовать environment.dependencies или environment.version с environment.unity_catalog_image. Если вам нужны дополнительные пакеты, кроме того, что находится в изображении, необходимо добавить их в Dockerfile.

Troubleshooting

Сведения об ошибках, связанных с аутентификацией в реестре, разрешениями, отправкой образов или обнаружением образов, см. в разделе «Устранение неполадок Artifact Registry».

SSL. SSLError при загрузке зависимостей

Пользовательский образ может завершиться с ошибкой во время выполнения из-за ошибки OpenSSL, если библиотека пытается создать контекст SSL, например:

ssl.SSLError: [CRYPTO] unknown error (_ssl.c:3076)

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

Это происходит потому, что рабочие нагрузки AI Runtime выполняются на хостах с включёнными Федеральными стандартами обработки информации (FIPS). Если криптографические библиотеки изображения не соответствуют FIPS, OpenSSL не может инициализироваться в режиме FIPS, поэтому создание контекста SSL завершается ошибкой.

Рекомендуемое решение:

Корпоративные, государственные, медицинские и финансовые нагрузки часто требуют соответствия требованиям FIPS 140-2 или 140-3 для прохождения аудитов FedRAMP, CMMC или HIPAA. Если рабочая нагрузка должна оставаться совместимой с FIPS, создайте образ с помощью криптографических библиотек, совместимых с FIPS.

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

Чтобы отключить режим FIPS, задайте этот параметр в YAML-файле рабочей нагрузки в разделе env_variables:

env_variables:
  OPENSSL_FORCE_FIPS_MODE: '0'

Кроме того, задайте переменную в Dockerfile, чтобы она применялася к каждой рабочей нагрузке, используюшей образ:

ENV OPENSSL_FORCE_FIPS_MODE=0

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

Дополнительные ресурсы

  • Последнее обновление: