Aller au contenu principal
Dernière mise à jour le

Utiliser 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 Fonctionnalités bêta de Runtime de l'IA et Registre d'artefacts Databricks sur la page Aperçus du Workspace.

Le Runtime IA peut exécuter une image de conteneur Docker personnalisée stockée dans Artifact Registry. 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.
  • 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​

Pousser une image vers Artifact Registry​

Avant d'utiliser une image personnalisée avec le Runtime IA, stockez-la dans Artifact Registry dans le même workspace que celui utilisé 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, accorder les privilèges requis et pousser l'image.

  3. 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.

astuce

Vous pouvez également utiliser la commande d’aide databricks air images push dans la CLI de Databricks.

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.
  • LOCAL_ADDR: IP du nœud local (nœud multiple uniquement).
  • 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}" \
    /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.

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

Générer l'image localement :

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

Ensuite, suivez start avec Artifact Registry pour taguer et envoyer l’image vers Artifact Registry. Utilisez le nom <catalog>.<schema>.<image>:<tag> résultant comme environment.unity_catalog_image dans le YAML de la charge de travail.

astuce

Vous pouvez également utiliser la databricks air images push commande d'aide dans la CLI de Databricks 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é 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.

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