Utilisez des images Docker personnalisées avec AI Runtime

info

Cette fonctionnalité est en version bêta. Pour l'utiliser, un administrateur de workspace doit activer les aperçus AI Runtime Beta Features et Databricks Artifact Registry depuis la page de prévisualisation du workspace.

AI Runtime can run a custom Docker container image stored in Artifact Registry. Use a custom image when you need:

• Bibliothèques système ou dépendances complexes que vous ne pouvez pas installer avec environment.dependencies.
• Un environnement reproductible dans le cadre du développement, de la recherche et de la production.
• Images privées approuvées par l'organisation et créées par votre plateforme ou votre équipe de sécurité.

Prérequis

• Installez ou mettez à jour vers la version 1,19,0 ou supérieure de la CLI Databricks. Consultez Installer ou mettre à jour la CLI Databricks.
• Demandez à l’administrateur de votre workspace d’activer l’aperçu AI Runtime Beta Features . Pour obtenir des instructions, consultez Gérer les aperçus au niveau du workspace.
• Configurez Artifact Registry dans le même workspace et obtenez l'autorisation d'envoyer (push) et de lire des images. Consultez Prise en main d'Artifact Registry.
• Installez et start Docker sur votre machine locale.

Envoyer une image vers Artifact Registry

Avant d'utiliser une image personnalisée avec AI Runtime, stockez-la dans Artifact Registry dans le même workspace que celui que vous utilisez pour soumettre la charge de travail.

1. Créez ou sélectionnez un catalogue et un schéma Unity Catalog pour l’image.

2. Suivez Prise en main d’Artifact Registry pour configurer l’authentification Docker, octroyer les privilèges requis et envoyer l’image.

3. Notez le nom Unity Catalog de l’image dans le format suivant :

   Text

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

   Par exemple, main.ml.training:v1. N’incluez pas le hostname du registre dans la configuration de la charge de travail.

astuce

Vous pouvez également utiliser la commande d'aidedatabricks air images push dans la CLI Databricks.

Utiliser une image Docker dans une charge de travail

Spécifiez le nom Unity Catalog de l'image dans votre fichier YAML de charge de travail sous environment.unity_catalog_image:

YAML

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

Cet exemple exécute train.py à partir de /app dans l’image. Pour effectuer l’upload du code d’application séparément sans reconstruire l’image, consultez la page Exécuter du code d’application upload.

Lorsque vous apportez votre propre image Docker, environment.dependencies et environment.version ne sont pas pris en charge. La spécification de environment.unity_catalog_image avec l'un ou l'autre champ Trigger une erreur. Si vous avez des dépendances supplémentaires, installez les packages dans le Dockerfile à la place.

Soumettez la charge de travail :

Shell

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

Le profil doit s'authentifier sur le même workspace que celui où l'image est stockée.

Variables d'environnement injectées dans votre conteneur

AI Runtime injecte les variables d'environnement suivantes dans chaque conteneur lors de l'exécution :

• CODE_SOURCE_PATH: chemin d'accès au code d'application upload, lorsque code_source est configuré.
• NUM_NODES: nombre total de nœuds.
• LOCAL_WORLD_SIZE: GPUs per node.
• WORLD_SIZE: nombre total de processus.
• POD_RANK: rang du nœud actuel (indexé à partir de 0). Également injecté en tant que NODE_RANK.
• LOCAL_ADDR: IP du nœud local (multi-nœud uniquement).
• MASTER_ADDR: rank-0 coordination address (multi-node only).
• MASTER_PORT: port de coordination de rang 0 (multinœud uniquement).

Exemples

The following examples show how to run upload application code and distributed training with a custom image.

Exécuter le code d'application upload

Utilisez code_source pour upload le code d’application séparément de l’image. Vous pouvez modifier et soumettre à nouveau votre code sans avoir à reconstruire l'image. Installez les dépendances Python et système du code dans l'image.

Placez train.py dans un répertoire local src à côté de workload.yaml. La configuration suivante effectue l’upload de src et exécute son train.py à l’intérieur de l’image personnalisée :

YAML

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 résout par rapport à workload.yaml. Runtime définit $CODE_SOURCE_PATH sur le chemin d'accès du répertoire upload dans le conteneur. Voir code_source pour les options de code source.

H100 multinœud avec RDMA

Pour les Job H100 multinœuds qui nécessitent une pleine bande passante réseau sur les instances AWS p5, basez votre image sur l'une des images de base Databricks avec NCCL et EFA préconfigurés :

YAML

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

Créez votre propre image

Lors de la création de votre propre image, Databricks recommande d'utiliser la compétence databricks-ai-runtime avec un agent de codage ou de commencer à partir d'une image de base Databricks.

Utilisez un agent de codage

Installez la compétence Claude Code databricks-ai-runtime pour des conseils Dockerfile étape par étape, y compris la construction à partir de zéro, la compatibilité CUDA/NCCL/EFA, les problèmes courants et une liste de contrôle de pré-construction. Cette compétence requiert Databricks CLI version 1.0.0 ou plus récente.

Shell

databricks aitools install --skills databricks-ai-runtime

Images de base Databricks

Databricks publie des images de base sur Docker Hub à l'adresse databricksruntime/air avec CUDA, NCCL et une mise en réseau spécifique au cloud (AWS EFA ou Azure InfiniBand) préconfigurés.

Tag	Variante	CUDA	Quand utiliser
dcs-base-aws-runtime	Environnement d'exécution	12	Installation de roues pré-construites uniquement
dcs-base-aws-devel	Développement	12	Compilation des extensions CUDA (nécessite nvcc)
dcs-base-aws-runtime-cu13	Environnement d'exécution	13	Installation de wheels préconstruits uniquement, sur CUDA 13
dcs-base-aws-devel-cu13	Développement	13	Compilation d'extensions CUDA sur CUDA 13 (nécessite nvcc)

Tag	Variante	CUDA	Quand utiliser
dcs-base-aws-runtime	Environnement d'exécution	12	Installation de roues pré-construites uniquement
dcs-base-aws-devel	Développement	12	Compilation des extensions CUDA (nécessite nvcc)
dcs-base-aws-runtime-cu13	Environnement d'exécution	13	Installation de wheels préconstruits uniquement, sur CUDA 13
dcs-base-aws-devel-cu13	Développement	13	Compilation d'extensions CUDA sur CUDA 13 (nécessite nvcc)

Le Dockerfile suivant ajoute PyTorch à une image de base Databricks. Les images de base fournissent Python sur /opt/venv, géré par uv. uv pip install cible cet environnement par default. Pour utiliser un autre environnement, créez et activez un venv avant d’exécuter uv pip install.

Pour inclure votre script d'entraînement dans l'image, placez train.py à côté du Dockerfile. Le Dockerfile le copie vers /app/train.py. Si vous upload le code de l'application avec code_source, omettez l'instruction COPY et conservez plutôt train.py dans votre répertoire source.

Dockerfile

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.py /app/train.py

Créer l'image localement :

Shell

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

Suivez ensuite Get started with Artifact Registry pour étiqueter et transférer l'image vers Artifact Registry. Utilisez le nom de <catalog>.<schema>.<image>:<tag> résultant comme environment.unity_catalog_image dans le fichier YAML de la charge de travail.

astuce

Vous pouvez également utiliser la commande d’aidedatabricks air images push dans Databricks CLI et suivre les instructions interactives.

Limitations

• Les images doivent être stockées dans Artifact Registry, dans le workspace où vous soumettez la charge de travail.
• La taille de l'image doit être inférieure à 20 Go.
• WORKDIR n'est pas respecté à l'environnement d'exécution. Utilisez des chemins absolus pour les fichiers intégrés à l’image. Par exemple, utilisez python /app/train.py, et non python train.py.
• Vous ne pouvez pas utiliser environment.dependencies ou environment.version avec environment.unity_catalog_image. Si vous avez besoin de packages supplémentaires au-delà de ce qui se trouve dans l'image, vous devez les ajouter au Dockerfile.

Dépannage

For registry-related authentication, permission, image push, or image discovery errors, see Troubleshoot Artifact Registry.

ssl.SSLError lors du chargement des dépendances

Une image personnalisée peut échouer à l'exécution avec une erreur OpenSSL lorsqu'une bibliothèque tente de créer un contexte SSL, par exemple :

Text

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

L'erreur apparaît lors de l'importation de bibliothèques qui ouvrent des connexions réseau, telles que huggingface_hub, et les empêche de se charger.

Cela se produit parce que les workloads AI Runtime s’exécutent sur des hôtes pour lesquels les normes FIPS (Federal Information Processing Standards) sont activées. Lorsque les bibliothèques cryptographiques de l’image ne sont pas conformes aux normes FIPS, OpenSSL ne parvient pas à s’initialiser en mode FIPS, de sorte que la création d’un contexte SSL échoue.

Solution recommandée :

Les charges de travail des entreprises, des gouvernements, du secteur de la santé et de la finance dépendent souvent de la conformité FIPS 140-2 ou 140-3 pour les audits FedRAMP, CMMC ou HIPAA. Si votre charge de travail doit rester conforme à la norme FIPS, créez votre image avec des bibliothèques cryptographiques conformes à la norme FIPS.

Si votre charge de travail ne nécessite pas la conformité FIPS, vous pouvez désactiver le mode FIPS en définissant la variable d'environnement OPENSSL_FORCE_FIPS_MODE sur 0. Cela peut entraîner silencieusement le non-respect des exigences de conformité.

Pour désactiver le mode FIPS, définissez-le dans votre fichier YAML de charge de travail sous env_variables:

YAML

env_variables:
  OPENSSL_FORCE_FIPS_MODE: '0'

Vous pouvez également définir la variable dans votre Dockerfile afin qu'elle s'applique à chaque charge de travail qui utilise l'image :

Dockerfile

ENV OPENSSL_FORCE_FIPS_MODE=0

Soumettez à nouveau la charge de travail et confirmez que l'erreur SSL n'apparaît plus lors du chargement des dépendances.

Ressources supplémentaires

• Guide de démarrage rapide de Databricks CLI pour AI Runtime
• Référence YAML Workload
• air groupe de commandes