Aller au contenu principal
Dernière mise à jour le

Utiliser des images Docker personnalisées

info

Bêta

Les images Docker personnalisées pour les charges de travail AI Runtime CLI sont en version bêta.

Docker Container Services (DCS) vous permet d'apporter votre propre image de conteneur Docker aux charges de travail air. Utilisez une image personnalisée lorsque vous avez besoin :

  • Versions spécifiques de la bibliothèque système.
  • Dépendances complexes qui ne s'intègrent pas facilement dans environment.dependencies.
  • Un environnement exact pour reproduire les résultats de recherche.
  • Images standard créées par l'équipe plateforme ou sécurité de votre organisation.

Prérequis​

Enregistrer une image​

Avant d'exécuter une charge de travail avec une image personnalisée, enregistrez-la auprès de air register image. L'enregistrement extrait et met en cache l'image dans la plateforme Databricks. Chaque utilisateur doit enregistrer une image une fois par balise d'image. Réenregistrez-vous uniquement lorsque vous transmettez une nouvelle balise ou renouvelez les informations d'identification. L'enregistrement prend 2 à 6 minutes et bloque jusqu'à ce que l'image soit prête.

Images publiques​

Enregistrez des images publiques en fournissant l’URL de l’image Docker et votre profil Databricks :

Shell
air register image docker.io/nvidia/cuda:12.9.0-devel-ubuntu24.04 -p my-databricks-profile

La référence d'image abrégée fonctionne également. Par exemple, library/ubuntu:latest.

Images privées de Docker Hub.​

Pour enregistrer une image Docker Hub privée, générez un jeton d'accès personnel d'abord. Dans les paramètres de votre compte Docker Hub, cliquez sur Jetons d'accès personnels → Générer un nouveau jeton . L'accès en lecture seule est suffisant.

Choisissez l'une des méthodes d'authentification suivantes :

Utilisation de la connexion Docker (recommandé pour une utilisation interactive)​

Connectez-vous à Docker Hub dans le terminal. Il vous sera demandé de fournir votre nom d'utilisateur Docker Hub et votre jeton d'accès personnel :

Shell
docker login

Ceci stocke vos identifiants dans ~/.docker/config.json. Enregistrez ensuite l'image — air lit les identifiants automatiquement :

Shell
air register image myorg/myrepo:mytag -p my-databricks-profile

Utilisation de l'authentification interactive​

Authentifiez-vous et stockez les informations d'identification dans un Secret Scope Databricks en une seule étape :

Shell
air register image myorg/myrepo:mytag --interactive-authenticate -p my-databricks-profile

Vous serez invité à saisir votre nom d'utilisateur Docker Hub et votre jeton d'accès personnel. Les informations d'identification sont stockées dans votre Secret Scope de Workspace pour les enregistrements futurs.

Utilisation d'un secret Databricks pré-enregistré (recommandé pour la CI/les scripts)​

Stockez les informations d’identification dans un secret Databricks et référencez-le directement :

Shell
air register image myorg/myrepo:mytag --scope my-secret-scope --key my-docker-key -p my-databricks-profile

Utiliser une image Docker dans une charge de travail​

Spécifiez l’image Docker dans votre fichier YAML de charge de travail sous environment.docker_image.url:

YAML
experiment_name: my-dcs-training
environment:
  docker_image:
    url: myorg/myrepo:mytag
compute:
  num_accelerators: 1
  accelerator_type: GPU_1xA10
command: python /app/train.py

Lorsque vous apportez votre propre image Docker, environment.dependencies et environment.version ne sont pas pris en charge. La spécification de environment.docker_image.url 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
air run --file workload.yaml -p my-databricks-profile

Variables d'environnement injectées dans votre conteneur​

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

  • NUM_NODES — nombre total de nœuds.
  • LOCAL_WORLD_SIZE — GPU par nœud.
  • WORLD_SIZE — nombre total de processus.
  • POD_RANK — classement de nœud actuel (indexé à partir de 0). Également injecté en tant que NODE_RANK.
  • LOCAL_ADDR — IP du nœud local (multi-nœud seulement).
  • MASTER_ADDR — adresse de coordination rank-0 (multi-nœuds uniquement).
  • MASTER_PORT — port de coordination de rang 0 (multi-nœuds uniquement).

Exemples​

Nœud unique A10​

YAML
experiment_name: my-dcs-single-node
environment:
  docker_image:
    url: myorg/myrepo:mytag
compute:
  num_accelerators: 1
  accelerator_type: GPU_1xA10
command: python3 /app/train.py

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:
  docker_image:
    url: myorg/myrepo:mytag
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 --experimental

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

Cloud

Variante

Quand utiliser

dcs-base-aws-runtime

AWS

Environnement d'exécution

Installation de roues pré-construites uniquement

dcs-base-aws-devel

AWS

Développement

Compilation des extensions CUDA (nécessite nvcc)

dcs-base-azure-runtime

Azure

Environnement d'exécution

Installation de roues pré-construites uniquement

dcs-base-azure-devel

Azure

Développement

Compilation des extensions CUDA (nécessite nvcc)

Tag

Cloud

Variante

Quand utiliser

dcs-base-aws-runtime

AWS

Environnement d'exécution

Installation de roues pré-construites uniquement

dcs-base-aws-devel

AWS

Développement

Compilation des extensions CUDA (nécessite nvcc)

dcs-base-azure-runtime

Azure

Environnement d'exécution

Installation de roues pré-construites uniquement

dcs-base-azure-devel

Azure

Développement

Compilation des extensions CUDA (nécessite nvcc)

Utilisez la variante **runtime** sauf si votre Dockerfile compile des extensions CUDA telles que flash-attn, apex ou des kernels personnalisés.

Exemple de Dockerfile ajoutant PyTorch à une image de base Databricks. Les images de base fournissent Python à l’emplacement /opt/venv, géré par uv. uv pip install cible cet environnement par default ; pour utiliser un environnement différent, créez et activez un environnement virtuel avant d'exécuter uv pip install.

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

Construire, pousser et enregistrer :

Shell
docker build -t myorg/myrepo:mytag .
docker push myorg/myrepo:mytag
air register image myorg/myrepo:mytag --interactive-authenticate -p my-databricks-profile

Exigences​

  • Les images doivent être hébergées sur Docker Hub. Amazon ECR, Google GCR et GitHub GHCR ne sont pas pris en charge.
  • 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.docker_image.url. Si vous avez besoin de packages supplémentaires au-delà de ce qui se trouve dans l'image, vous devez les ajouter au Dockerfile.

Voir aussi​

Sur cette page