Aller au contenu principal
Dernière mise à jour le

Utiliser des images Docker personnalisées avec AI Runtime

info

Aperçu

Cette fonctionnalité est en préversion publique. Pour l'utiliser, un administrateur du workspace doit activer la préversion de l'Registre d'artefacts Databricks depuis la page de préversions du workspace.

AI Runtime peut exécuter votre charge de travail dans une image de conteneur Docker personnalisée au lieu d'un environnement géré. Vous construisez l'image pour linux/amd64, vous la poussez vers Artifact Registry dans votre workspace et vous référencez son nom Unity Catalog dans le YAML de charge de travail que vous soumettez avec le Databricks CLI. Utilisez une image personnalisée lorsque vous avez besoin de :

  • Bibliothèques système ou dépendances complexes que vous ne pouvez pas installer avec environment.dependencies.
  • Extensions CUDA que vous compilez à partir de la source, telles que flash-attn ou des noyaux personnalisés.
  • Un environnement reproductible entre le développement, la recherche et la production.
  • Images privées approuvées par l’organisation, construites par votre plateforme ou votre équipe de sécurité.

Prérequis​

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. Pour construire une image en premier lieu, consultez Construire votre propre image.

  1. Créez ou sélectionnez un catalogue et un schéma Unity Catalog pour l'image et obtenez les privilèges Artifact Registry requis.

  2. Créer ou refresh un profil OAuth de workspace pour le workspace :

    Shell
    databricks auth login --profile my-databricks-profile

    L'assistant d'identification Docker nécessite un profil créé par databricks auth login. Il ne prend pas en charge les profils qui utilisent des jetons d'accès personnels ou des informations d'identification machine-à-machine (M2M) OAuth.

  3. Pour authentifier Docker, taguer l'image et l'envoyer en une seule étape, exécutez l'assistant databricks air images push :

    Shell
    databricks air images push \
      --profile my-databricks-profile \
      --source my-training-image:v1 \
      --catalog main \
      --schema ml \
      --artifact training:v1

    Pour plus d'options, y compris la copie d'une image à partir d'un autre registre, consultez Pousser des images avec le Databricks CLI.

    Pour envoyer avec docker push à la place, exécutez databricks auth docker configure --profile my-databricks-profile pour configurer l'assistant d'identification Docker, puis suivez Envoyer une image dans le guide Artifact Registry.

  4. Notez le nom de l'image dans le Unity Catalog 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.

Utiliser une image Docker dans une charge de travail​

Indiquez le nom de l’image dans Unity Catalog 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 upload du code d'application séparément sans reconstruire l'image, voir Exécuter le code d'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 de environment.unity_catalog_image avec l'un ou l'autre des champs Trigger une erreur. Si vous avez des dépendances supplémentaires, installez les packages dans le Dockerfile.

Envoyer 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​

L'environnement Runtime de l'IA injecte les variables d'environnement suivantes dans chaque conteneur au moment de l'exécution :

  • CODE_SOURCE_PATH: chemin d'accès au code de l'application upload, lorsque code_source 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 du nœud actuel (indexé à 0). Également injecté en tant que NODE_RANK.

  • MASTER_ADDR: adresse de coordination du rang 0 (multi-nœuds uniquement).

  • MASTER_PORT: port de coordination de rang 0 (multinœud uniquement).

Exemples​

Les exemples suivants montrent comment exécuter du code d’application upload et un entraînement distribué avec une image personnalisée.

Exécuter le code d'application upload​

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

Placez train.py dans un répertoire local src à côté de workload.yaml. La configuration suivante upload src et exécute son train.py dans 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 résout par rapport à workload.yaml. AI Runtime définit $CODE_SOURCE_PATH sur le chemin du répertoire upload dans le conteneur. Voir code_source pour les options de code source.

H100 multi-nœuds avec RDMA​

Pour les jobs H100 à plusieurs nœuds qui ont besoin d'une bande passante réseau complète 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}" \
    --local_addr="$(hostname -I | awk '{print $1}')" \
    /app/train.py

Créez 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.

Vous pouvez également start à partir d'une autre image de base, comme nvidia/cuda. Dans ce cas, l'image doit inclure un runtime CUDA pris en charge par le Driver GPU de l'hôte et NCCL pour la communication multi-GPU. La compétence databricks-ai-runtime couvre la création d'une image à partir de zéro.

Pour les charges de travail GPU_8xH100 multi-nœuds, une image construite à partir d'une autre image de base doit également inclure EFA.

Utiliser un agent de codage​

Installez la compétence Claude Code databricks-ai-runtime pour obtenir des conseils étape par étape sur Dockerfile, notamment pour la création à partir de zéro, la compatibilité avec CUDA/NCCL/EFA, les problèmes courants et une liste de contrôle de pré-génération. Cette compétence nécessite la version 1.0.0 ou plus de Databricks CLI.

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 la mise en réseau spécifique au cloud (AWS EFA ou Azure InfiniBand) préconfigurés.

Tag

Variante

CUDA

Utilisez quand

dcs-base-aws-runtime

Environnement d'exécution

12

Installation de wheels prédéfinis uniquement

dcs-base-aws-devel

Dével

12

Compilation des extensions CUDA (nécessite nvcc)

dcs-base-aws-runtime-cu13

Environnement d'exécution

13

Installation de wheels prédéfinis uniquement sur CUDA 13

dcs-base-aws-devel-cu13

Dével

13

Compilation des extensions CUDA sur CUDA 13 (nécessite nvcc)

Tag

Variante

CUDA

Utilisez quand

dcs-base-aws-runtime

Environnement d'exécution

12

Installation de wheels prédéfinis uniquement

dcs-base-aws-devel

Dével

12

Compilation des extensions CUDA (nécessite nvcc)

dcs-base-aws-runtime-cu13

Environnement d'exécution

13

Installation de wheels prédéfinis uniquement sur CUDA 13

dcs-base-aws-devel-cu13

Dével

13

Compilation des extensions CUDA sur CUDA 13 (nécessite nvcc)

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

Pour inclure votre script de formation dans l’image, placez train.py à côté du Dockerfile. Le Dockerfile le copie vers /app/train.py. Si vous effectuez un upload de code d’application avec code_source, omettez l’instruction COPY et conservez 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

Les images doivent utiliser la plateforme linux/amd64. Sur un Mac avec Apple Silicon, Docker construit des images arm64 default, donc transmettez toujours --platform linux/amd64. Pour construire l'image localement, exécutez la commande suivante :

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

Ensuite, poussez l'image vers Artifact Registry et utilisez le nom <catalog>.<schema>.<image>:<tag> obtenu en tant que environment.unity_catalog_image dans le YAML de charge de travail.

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é au moment de l’exécution. Utilisez des chemins absolus pour les fichiers intégrés à l’image. Par exemple, utilisez python /app/train.py, pas 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 en plus de ceux de l’image, vous devez les ajouter au Dockerfile.

Dépannage​

Pour les erreurs liées à l’authentification de registre, aux autorisations, à l’envoi d’images ou à la découverte d’images, consultez Dépanner Artifact Registry.

L'assistant d'identification Docker nécessite un profil créé par databricks auth login​

Le push d'une image échoue avec une erreur telle que la suivante :

Text
Error: profile "my-profile" uses auth_type "". Docker credential helper requires a profile created by databricks auth login

Cette erreur se produit lorsque le profil utilise un jeton d'accès personnel ou une autre méthode d'authentification que l'assistant d'identification Docker ne prend pas en charge. Créez un profil OAuth de workspace, puis réexécutez le push avec ce profil :

Shell
databricks auth login --profile my-databricks-profile

Plusieurs profils correspondent à l'ID de workspace​

La configuration de l'authentification Docker ou le push d'une image échoue avec une erreur telle que la suivante :

Text
Error: multiple Databricks profiles match workspace ID 1111111111111111: profile-a and profile-b. Remove duplicate workspace_id entries before using Docker credential helper

Cette erreur se produit lorsque plusieurs profils dans votre fichier de configuration (~/.databrickscfg default) pointent vers le même workspace. Supprimez le profil que vous n'utilisez pas pour Docker, puis réexécutez la commande avec le profil restant. Pour vous déconnecter d'un profil et le supprimer du fichier de configuration, exécutez la commande suivante :

Shell
databricks auth logout --profile profile-a --delete

Pour toutes les options de déconnexion, consultez le groupe de commandes auth.

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 :

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 d’IA Runtime s’exécutent sur des hôtes sur lesquels les FIPS (Federal Information Processing Standards) sont activés. Lorsque les bibliothèques cryptographiques de l’image ne sont pas compatibles FIPS, OpenSSL ne parvient pas à s’initialiser en mode FIPS, donc la création d’un contexte SSL échoue.

Solution recommandée :

Les charges de travail des entreprises, du gouvernement, 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 à FIPS, construisez votre image avec des bibliothèques cryptographiques conformes à FIPS.

Si votre workload 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 le YAML de votre workload sous env_variables:

YAML
env_variables:
  OPENSSL_FORCE_FIPS_MODE: '0'

Alternativement, définissez la variable dans votre Dockerfile pour qu’elle s’applique à chaque charge de travail utilisant l’image :

Dockerfile
ENV OPENSSL_FORCE_FIPS_MODE=0

Soumettez à nouveau la charge de travail et confirmez que l'erreur SSL ne s'affiche plus lors du chargement des dépendances.

Ressources supplémentaires​

Sur cette page