Utiliza imágenes Docker personalizadas con AI Runtime

Importante

Esta característica se encuentra en su versión beta. Para utilizarla, un administrador de espacio de trabajo debe habilitar las funciones beta de AI Runtime y las previsualizaciones del Registro de Artefactos de Databricks desde la página de Previsualizaciones del espacio de trabajo.

AI Runtime puede ejecutar una imagen de contenedor Docker personalizada almacenada en el Registro de Artefactos. Use una imagen personalizada cuando necesite:

  • Bibliotecas del sistema o dependencias complejas que no puedes instalar con environment.dependencies.
  • Un entorno reproducible a través del desarrollo, la investigación y la producción.
  • Imágenes privadas aprobadas por la organización creadas por tu plataforma o equipo de seguridad.

Prerrequisitos

Enviar una imagen al Registro de Artefactos

Antes de usar una imagen personalizada con AI Runtime, guárdala en el Registro de Artefactos en el mismo espacio de trabajo que usas para enviar la carga de trabajo.

  1. Crea o selecciona un catálogo y un esquema de Unity Catalog para la imagen.

  2. Sigue los pasos de Primeros pasos con Artifact Registry para configurar la autenticación de Docker, conceder los privilegios necesarios y subir la imagen.

  3. Fíjate en el nombre del Catálogo de Unity de la imagen en el siguiente formato:

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

    Por ejemplo: main.ml.training:v1. No incluyas el nombre de host del registro en la configuración de la carga de trabajo.

Tip

Alternativamente, usa el databricks air images push comando helper en la CLI de Databricks.

Usa una imagen Docker en una carga de trabajo

Especifica el nombre de la imagen en el Catálogo de Unity en el archivo YAML de tu carga de trabajo, en 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

Este ejemplo ejecuta train.py desde /app en la imagen. Para cargar el código de la aplicación por separado sin reconstruir la imagen, consulte Ejecutar código de aplicación subido.

Si utilizas tu propia imagen de Docker, environment.dependencies y environment.version no son compatibles. Al especificar environment.unity_catalog_image 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:

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

El perfil debe autenticarse en el mismo espacio de trabajo donde se almacena la imagen.

Variables de entorno insertadas en el contenedor

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

  • CODE_SOURCE_PATH: ruta hacia el código de la aplicación subido, cuando code_source está configurado.
  • NUM_NODES: número total de nodos.
  • LOCAL_WORLD_SIZE: GPUs por nodo.
  • WORLD_SIZE: número total de procesos.
  • POD_RANK: rango actual del nodo (0-indexado). También se inyecta como NODE_RANK.
  • LOCAL_ADDR: IP local del nodo (solo para varios nodos).
  • MASTER_ADDR: dirección de coordinación de rango 0 (solo multinodo).
  • MASTER_PORT: puerto de coordinación de rango 0 (solo multinodo).

Examples

Los siguientes ejemplos muestran cómo ejecutar código de aplicación subido y entrenamiento distribuido con una imagen personalizada.

Ejecutar código de aplicación subido

Úsalo code_source para subir el código de la aplicación por separado de la imagen. Puedes editar y volver a enviar tu código sin tener que reconstruir la imagen. Instala el código en Python y las dependencias del sistema en la imagen.

Coloca train.py en un directorio local src junto a workload.yaml. La siguiente configuración carga src y ejecuta su train.py dentro de la imagen personalizada:

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 se resuelve en relación con workload.yaml. AI Runtime establece $CODE_SOURCE_PATH en la ruta del directorio cargado en el contenedor. Consulta code_source para opciones de código fuente.

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:
  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

Creación de su propia imagen

Al crear su propia imagen, Databricks recomienda usar la habilidad databricks-ai-runtime con un agente de código o partir de una imagen base de Databricks.

Uso de un agente de codificación

Instale la aptitud databricks-ai-runtime Claude Code para obtener instrucciones paso a paso de Dockerfile, incluida la compilación desde cero, la compatibilidad de CUDA/NCCL/EFA, los problemas comunes y una lista de comprobación previa a la compilación. Esta aptitud requiere la versión 1.0.0 o posterior de la CLI de Databricks.

databricks aitools install --skills databricks-ai-runtime --experimental

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.

Etiqueta Variant CUDA Se utiliza cuando
dcs-base-azure-runtime Runtime 12 Solo se pueden instalar wheels precompiladas
dcs-base-azure-devel Devel 12 Compilación de extensiones CUDA (requiere nvcc)

El siguiente archivo Dockerfile añade PyTorch a una imagen base de Databricks. Las imágenes base proporcionan Python en /opt/venv, gestionado por uv. uv pip install va dirigido a ese entorno por defecto. Para usar un entorno diferente, crea y activa un venv antes de ejecutar uv pip install.

Para incluir tu script de entrenamiento en la imagen, colócalo train.py junto al archivo Docker. El Dockerfile lo copia a /app/train.py. Si subes el código de la aplicación con code_source, omite la instrucción COPY y conserva train.py en tu directorio fuente.

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

Construye la imagen localmente:

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

Luego sigue Empezar con el Registro de Artefactos para etiquetar y subir la imagen al Registro de Artefactos. Utiliza el nombre resultante <catalog>.<schema>.<image>:<tag> como environment.unity_catalog_image en la carga de trabajo YAML.

Tip

Alternativamente, utiliza el databricks air images push comando helper en la CLI de Databricks y sigue las indicaciones interactivas.

Limitations

  • Las imágenes deben almacenarse en el Registro de Artefactos en el espacio de trabajo donde envías la carga de trabajo.
  • 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.unity_catalog_image. Si necesita paquetes adicionales más allá de lo que se encuentra en la imagen, debe agregarlos al Dockerfile.

Troubleshooting

Para errores relacionados con el registro, como errores de autenticación, permisos, subida de imágenes o detección de imágenes, consulte Solucionar problemas de Artifact Registry.

SSL. SSLError al cargar dependencias

Una imagen personalizada puede fallar durante la ejecución con un error de OpenSSL cuando una biblioteca intenta crear un contexto SSL, por ejemplo:

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

El error aparece al importar bibliotecas que abren conexiones de red, como huggingface_hub, e impide que se carguen.

Esto ocurre porque las cargas de trabajo en tiempo de ejecución de IA se ejecutan en hosts con los Estándares Federales de Procesamiento de Información (FIPS) habilitados. Cuando las bibliotecas criptográficas de la imagen no son compatibles con FIPS, OpenSSL no puede inicializarse en modo FIPS, por lo que se produce un error al crear un contexto SSL.

Solución recomendada:

Las cargas de trabajo empresariales, gubernamentales, sanitarias y financieras suelen depender del cumplimiento de FIPS 140-2 o 140-3 para las auditorías de FedRAMP, CMMC o HIPAA. Si la carga de trabajo debe seguir siendo compatible con FIPS, compile la imagen con bibliotecas criptográficas compatibles con FIPS.

Si la carga de trabajo no requiere el cumplimiento de FIPS, puede deshabilitar el modo FIPS estableciendo la OPENSSL_FORCE_FIPS_MODE variable de entorno en 0. Si lo hace, puede interrumpir silenciosamente los requisitos de cumplimiento.

Para deshabilitar el modo FIPS, configúrelo en el archivo YAML de su carga de trabajo debajo de env_variables:

env_variables:
  OPENSSL_FORCE_FIPS_MODE: '0'

Como alternativa, establezca la variable en el Dockerfile para que se aplique a todas las cargas de trabajo que usen la imagen:

ENV OPENSSL_FORCE_FIPS_MODE=0

Vuelva a enviar la carga de trabajo y confirme que el error SSL ya no aparece cuando se cargan las dependencias.

Recursos adicionales

  • Se actualizó por última vez el