Utilisez des images Docker personnalisées avec AI Runtime

Important

Cette fonctionnalité est en version bêta. Pour l’utiliser, un administrateur d’espace de travail doit activer les fonctionnalités bêta d’exécution IA et les aperçus du registre d’artefacts Databricks depuis la page Aperçu de l’espace de travail.

AI Runtime peut exécuter une image de conteneur Docker personnalisée stockée dans Artifact Registry. Utilisez une image personnalisée lorsque vous avez besoin des éléments suivants :

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

Prerequisites

Pousser une image dans le Registre des Artefacts

Avant d’utiliser une image personnalisée avec AI Runtime, stockez-la dans le Registre d’Artefacts, dans le même espace de travail que vous utilisez pour soumettre la charge de travail.

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

  2. Suivez Commencez avec Artifact Registry pour configurer l’authentification Docker, accorder les privilèges requis et pousser l’image.

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

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

    Par exemple : main.ml.training:v1. N’incluez pas le nom d’hôte du registre dans la configuration de la charge de travail.

Tip

Sinon, utilisez la databricks air images push commande helper dans Databricks CLI.

Utiliser une image Docker dans une charge de travail

Indiquez le nom Unity Catalog de l’image dans votre fichier YAML de charge de travail sous 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

Cet exemple s’exécute train.py à partir de /app sur l’image. Pour télécharger le code de l’application séparément sans reconstruire l’image, voir Exécuter le code de l’application téléchargé.

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

Soumettez la charge de travail :

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

Le profil doit s’authentifier dans le même espace de travail 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 au moment de l’exécution :

  • CODE_SOURCE_PATH: chemin vers le code de l’application téléchargé, lorsque code_source c’est configuré.
  • NUM_NODES: nombre total de nœuds.
  • LOCAL_WORLD_SIZE: GPU par nœud.
  • WORLD_SIZE: nombre total de processus.
  • POD_RANK: rang actuel du nœud (0-indexé). Également injecté sous la forme de NODE_RANK.
  • LOCAL_ADDR : adresse IP du nœud local (en environnement multinœud uniquement).
  • MASTER_ADDR : adresse de coordination de rang 0 (multinœud uniquement).
  • MASTER_PORT: port de coordination de rang 0 (uniquement en mode multinœud).

Examples

Les exemples suivants montrent comment exécuter le code applicatif téléchargé et la formation distribuée avec une image personnalisée.

Exécuter le code applicatif téléchargé

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

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

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. AI Runtime définit $CODE_SOURCE_PATH sur le chemin du répertoire chargé dans le conteneur. Voir code_source pour les options de code source.

H100 à plusieurs nœuds avec RDMA

Pour les travaux H100 à plusieurs nœuds qui nécessitent une bande passante réseau complète sur des instances AWS p5, basez votre image sur l’une des images de base Databricks avec NCCL et EFA préconfigurées :

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éer votre propre image

Pour créer votre propre image, Databricks recommande d’utiliser la compétence databricks-ai-runtime avec un agent de codage ou de partir d’une image de base Databricks.

Utiliser un agent de codage

Installez la compétence Databricks-ai-runtime Claude Code pour obtenir des instructions pas à pas sur Dockerfile, notamment la création à partir de zéro, la compatibilité CUDA/NCCL/EFA, les problèmes courants et une liste de contrôle de pré-build. Cette compétence nécessite Databricks CLI version 1.0.0 ou ultérieure.

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 Variant CUDA À utiliser quand
dcs-base-azure-runtime Runtime 12 Installation des wheels précompilées uniquement
dcs-base-azure-devel Devel 12 Compilation d’extensions CUDA (nécessite nvcc)

Le fichier Docker suivant ajoute PyTorch à une image de base Databricks. Les images de base fournissent Python à la version /opt/venv, géré par uv. uv pip install cible cet environnement par défaut. Pour utiliser un environnement différent, créez et activez un venv avant de lancer uv pip install.

Pour inclure votre script d’entraînement dans l’image, placez-le train.py à côté du fichier Docker. Le Dockerfile le copie dans /app/train.py. Si vous téléversez le code de l’application avec code_source, omettez l’instruction COPY et gardez train.py dans votre répertoire source.

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

Construisez l’image localement :

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

Ensuite, suivez Commencez avec le Registre des Artefacts pour taguer et pousser l’image vers le Registre des Artefacts. Utilisez le nom résultant <catalog>.<schema>.<image>:<tag> comme environment.unity_catalog_image dans la charge de travail YAML.

Tip

Sinon, utilisez la databricks air images push commande helper dans Databricks CLI et suivez les instructions interactives.

Limitations

  • Les images doivent être stockées dans le registre des artefacts, dans l’espace de travail où vous soumettez la charge de travail.
  • La taille de l’image doit être inférieure à 20 Go.
  • WORKDIR n’est pas pris en compte à l’exécution. Utilisez des chemins absolus pour les fichiers insérés dans 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 fichier Dockerfile.

Troubleshooting

Pour les erreurs d’authentification, d’autorisation, d’importation d’image ou de détection d’images liées au registre, voir Dépanner Artifact Registry.

ssl. SSLError lors du chargement des dépendances

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

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

L’erreur s’affiche lors de l’importation de bibliothèques qui ouvrent des connexions réseau, telles que huggingface_hub, et les empêchent de charger.

Cela est dû au fait que les charges de travail d’AI Runtime s’exécutent sur des hôtes sur lesquels les normes fédérales de traitement de l’information (FIPS) sont activées. Lorsque les bibliothèques de chiffrement de l’image ne sont pas conformes à 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 d’entreprise, de gouvernement, de santé et de 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 à FIPS, générez votre image avec des bibliothèques de chiffrement conformes à FIPS.

Si votre charge de travail ne nécessite pas de conformité FIPS, vous pouvez désactiver le mode FIPS en définissant la variable OPENSSL_FORCE_FIPS_MODEd’environnement 0 sur . Cela peut interrompre silencieusement les exigences de conformité.

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

env_variables:
  OPENSSL_FORCE_FIPS_MODE: '0'

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

ENV OPENSSL_FORCE_FIPS_MODE=0

Renvoyez la charge de travail et confirmez que l’erreur SSL n’apparaît plus lorsque les dépendances se chargent.

Ressources additionnelles

  • Dernière mise à jour :