Uso de imágenes personalizadas de Docker

Docker Container Services (DCS) te permite utilizar tu propia imagen de contenedor Docker en air cargas de trabajo. Use una imagen personalizada cuando necesite:

• Versiones específicas de la biblioteca del sistema.
• Dependencias complejas que no se ajustan limpiamente a environment.dependencies.
• Un entorno exacto para reproducir los resultados de la investigación.
• Imágenes estándar creadas por la plataforma o el equipo de seguridad de su organización.

Prerrequisitos

• Instale la CLI de AI Runtime.
• En el caso de las imágenes privadas, una cuenta de Docker Hub con acceso a la imagen.

Registro de una imagen

Antes de ejecutar una carga de trabajo con una imagen personalizada, regístrela con air register image. El registro extrae y almacena en caché la imagen en la plataforma de Databricks. Cada usuario debe registrar una imagen una vez por etiqueta de imagen. Vuelva a registrarse solo cuando inserte una nueva etiqueta o rote las credenciales. El registro tarda entre 2 y 6 minutos y permanece bloqueado hasta que la imagen esté lista.

Imágenes públicas

Registre imágenes públicas proporcionando la dirección URL de la imagen de Docker y el perfil de Databricks:

air register image docker.io/nvidia/cuda:12.9.0-devel-ubuntu24.04 -p my-databricks-profile

La forma abreviada de la referencia de imagen también funciona. Por ejemplo: library/ubuntu:latest.

Imágenes de Docker Hub privadas

Para registrar una imagen de Docker Hub privada, genere primero un token de acceso personal. En la configuración de la cuenta de Docker Hub, haga clic en Tokens de acceso personal → Generar nuevo token. El acceso de solo lectura es suficiente.

Elija uno de los métodos de autenticación siguientes:

Uso del inicio de sesión de Docker (recomendado para uso interactivo)

Inicie sesión en Docker Hub en el terminal. Se le pedirá el nombre de usuario de Docker Hub y el token de acceso personal:

docker login

Esto almacena las credenciales en ~/.docker/config.json. A continuación, registre la imagen; air lee automáticamente las credenciales:

air register image myorg/myrepo:mytag -p my-databricks-profile

Uso de la autenticación interactiva

Autenticación y almacenamiento de credenciales en un ámbito de secreto de Databricks en un paso:

air register image myorg/myrepo:mytag --interactive-authenticate -p my-databricks-profile

Se le pedirá el nombre de usuario de Docker Hub y el token de acceso personal. Las credenciales se almacenan en el ámbito de secretos de tu espacio de trabajo para futuros registros.

Uso de un secreto de Databricks almacenado previamente (recomendado para CI/scripts)

Almacene las credenciales en un secreto de Databricks y haga referencia a ella directamente:

air register image myorg/myrepo:mytag --scope my-secret-scope --key my-docker-key -p my-databricks-profile

Uso de una imagen de Docker en una carga de trabajo

Especifica la imagen de Docker en el archivo YAML de tu carga de trabajo en environment.docker_image.url:

experiment_name: my-dcs-training
environment:
  docker_image:
    url: myorg/myrepo:mytag
compute:
  num_accelerators: 1
  accelerator_type: GPU_1xA10
command: python /app/train.py

Si utilizas tu propia imagen de Docker, environment.dependencies y environment.version no son compatibles. Al especificar environment.docker_image.url con cualquiera de los campos, se desencadena un error. Si tiene dependencias adicionales, instale los paquetes en el Dockerfile en su lugar.

Envíe la carga de trabajo:

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

Variables de entorno insertadas en el contenedor

AI Runtime inserta las siguientes variables de entorno en cada contenedor en tiempo de ejecución:

• NUM_NODES : número total de nodos.
• LOCAL_WORLD_SIZE : GPU por nodo.
• WORLD_SIZE — número total de procesos.
• POD_RANK — posición actual del nodo (indexado desde 0). También se inyecta como NODE_RANK.
• LOCAL_ADDR : dirección IP del nodo local (solo varios nodos).
• MASTER_ADDR — dirección de coordinación de rank-0 (solo para varios nodos).
• MASTER_PORT — puerto de coordinación de rango 0 (solo para varios nodos).

Examples

A10 de nodo único

experiment_name: my-dcs-single-node
environment:
  docker_image:
    url: myorg/myrepo:mytag
compute:
  num_accelerators: 1
  accelerator_type: GPU_1xA10
command: python3 /app/train.py

H100 de varios nodos con RDMA

Para los trabajos multinodo de H100 que necesiten todo el ancho de banda de red en instancias p5 de AWS, use como base para la imagen una de las imágenes base de Databricks con NCCL y EFA preconfigurados:

experiment_name: my-dcs-distributed
environment:
  docker_image:
    url: myorg/myrepo:mytag
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

Creación de su propia imagen

Imágenes base de Databricks

Databricks publica imágenes base en Docker Hub en databricksruntime/air con CUDA, NCCL y redes específicas de la nube (AWS EFA o Azure InfiniBand) preconfiguradas.

Utiliza la variante runtime a menos que el Dockerfile compile extensiones CUDA como flash-attn, apex o kernels personalizados.

Ejemplo de Dockerfile que agrega PyTorch a una imagen base de Databricks. Las imágenes base proporcionan Python en /opt/venv, gestionado por uv. uv pip install tiene como destino ese entorno de forma predeterminada; para usar un entorno diferente, cree y active una venv antes de ejecutar uv pip install.

FROM databricksruntime/air:dcs-base-aws-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 /app/train

Compilar, enviar y registrar:

docker build -t myorg/myrepo:mytag .
docker push myorg/myrepo:mytag
air register image myorg/myrepo:mytag --interactive-authenticate -p my-databricks-profile

Requisitos

• Las imágenes deben hospedarse en Docker Hub. Amazon ECR, Google GCR y GitHub GHCR no se admiten.
• El tamaño de la imagen debe estar por debajo de 20 GB.
• WORKDIR no se tiene en cuenta en tiempo de ejecución. Utilice rutas absolutas para los archivos incluidos en la imagen. Por ejemplo, use python /app/train.py, no python train.py.
• No puede usar environment.dependencies ni environment.version con environment.docker_image.url. Si necesita paquetes adicionales más allá de lo que se encuentra en la imagen, debe agregarlos al Dockerfile.

Consulte también

• Inicio rápido de la CLI de AI Runtime
• Referencia de YAML para cargas de trabajo
• Referencia de comandos de la CLI de ejecución de IA