Usa immagini Docker personalizzate con AI Runtime

Important

Questa funzionalità è in versione beta. Per utilizzarlo, un amministratore dello spazio di lavoro deve abilitare le funzionalità beta di Runtime AI e le anteprime del registro artefatti Databrick dalla pagina Anteprime dello spazio di lavoro.

AI Runtime può eseguire un'immagine container Docker personalizzata memorizzata in Artefact Registry. Usare un'immagine personalizzata quando è necessario:

  • librerie di sistema o dipendenze complesse che non puoi installare con environment.dependencies.
  • Un ambiente riproducibile attraverso sviluppo, ricerca e produzione.
  • Immagini private, approvate dall'organizzazione, create dalla tua piattaforma o dal team di sicurezza.

Prerequisites

Invia un'immagine al Registro degli Artefatti

Prima di usare un'immagine personalizzata con AI Runtime, archivila in Artefact Registry nello stesso workspace che usi per inviare il carico di lavoro.

  1. Crea o seleziona un catalogo e uno schema del Catalogo Unity per l'immagine.

  2. Segui Inizia con Artefact Registry per configurare l'autenticazione Docker, concedere i privilegi necessari e spingere l'immagine.

  3. Nota il nome del Catalogo Unity dell'immagine nel seguente formato:

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

    Ad esempio: main.ml.training:v1. Non includere il nome host del registro nella configurazione del carico di lavoro.

Tip

In alternativa, usa il databricks air images push comando helper nella CLI di Databricks.

Usa un'immagine Docker in un workload

Specifica il nome del Catalogo Unity dell'immagine nel tuo workload YAML sotto 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

Questo esempio viene eseguito train.py da /app nell'immagine. Per caricare il codice applicativo separatamente senza ricostruire l'immagine, vedi Esegui codice applicativo caricato.

Quando si utilizza una propria immagine Docker, environment.dependencies e environment.version non sono supportati. Specificare environment.unity_catalog_image in uno dei due campi genera un errore. Se si hanno dipendenze aggiuntive, installare invece i pacchetti nel Dockerfile.

Inviare il carico di lavoro:

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

Il profilo deve autenticarsi nello stesso spazio di lavoro dove è memorizzata l'immagine.

Variabili di ambiente inserite nel contenitore

Il runtime di intelligenza artificiale inserisce le variabili di ambiente seguenti in ogni contenitore in fase di esecuzione:

  • CODE_SOURCE_PATH: percorso verso il codice applicativo caricato, quando code_source è configurato.
  • NUM_NODES: numero totale di nodi.
  • LOCAL_WORLD_SIZE: GPU per nodo.
  • WORLD_SIZE: numero totale di processi.
  • POD_RANK: posizione corrente del nodo (con indice iniziale 0). Viene anche iniettato come NODE_RANK.
  • LOCAL_ADDR: indirizzo IP locale del nodo (solo in configurazioni multi-nodo).
  • MASTER_ADDR: indirizzo di coordinamento di rango 0 (solo per configurazioni multi-nodo).
  • MASTER_PORT: porta di coordinamento di grado 0 (solo per configurazioni multi-nodo).

Examples

I seguenti esempi mostrano come eseguire codice applicativo caricato e formazione distribuita con un'immagine personalizzata.

Esegui il codice applicativo caricato

Usalo code_source per caricare il codice dell'applicazione separatamente dall'immagine. Puoi modificare e riinviare il tuo codice senza ricostruire l'immagine. Installa il codice in Python e le dipendenze di sistema nell'immagine.

Inserisci train.py in una directory locale src accanto a workload.yaml. La seguente configurazione carica src ed esegue il relativo train.py all'interno dell'immagine personalizzata:

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 si risolve rispetto a workload.yaml. AI Runtime imposta $CODE_SOURCE_PATH sul percorso della directory caricata nel container. Vedi code_source per le opzioni di sorgente del codice.

H100 multinodo con RDMA

Per i processi H100 multinodo che richiedono una larghezza di banda di rete completa nelle istanze di AWS p5, basare l'immagine su una delle immagini di base di Databricks con NCCL ed EFA preconfigurate:

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

Creare un'immagine personalizzata

Quando si crea un'immagine personalizzata, Databricks consiglia di usare la competenza databricks-ai-runtime con un agente di coding oppure di partire da un'immagine di base Databricks.

Usare un agente di codifica

Installare la competenza Claude Code databricks-ai-runtime per istruzioni dettagliate sul Dockerfile, tra cui la compilazione da zero, la compatibilità CUDA/NCCL/EFA, i problemi comuni e un elenco di controllo di pre-compilazione. Questa competenza richiede Databricks CLI versione 1.0.0 o successivo.

databricks aitools install --skills databricks-ai-runtime --experimental

Immagini di base di Databricks

Databricks pubblica immagini di base su Docker Hub in databricksruntime/air con CUDA, NCCL e rete specifica del cloud (AWS EFA o Azure InfiniBand) preconfigurata.

Tag Variant CUDA Usa quando
dcs-base-azure-runtime Runtime 12 Installazione solo ruote predefinite
dcs-base-azure-devel Devel 12 Compilazione di estensioni CUDA (richiede nvcc)

Il seguente Dockerfile aggiunge PyTorch a un'immagine base Databricks. Le immagini di base forniscono Python in /opt/venv, gestito da uv. uv pip install ha come destinazione quell'ambiente per impostazione predefinita. Per usare un ambiente diverso, crea e attiva un venv prima di eseguire uv pip install.

Per includere il tuo script di allenamento nell'immagine, posizionalo train.py accanto al Dockerfile. Il file Docker lo copia su /app/train.py. Se carichi il codice dell'applicazione con code_source, ometti l'istruzione COPY e tieni invece train.py nella directory sorgente.

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

Costruisci l'immagine localmente:

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

Poi segui Inizia con Registro Artefatti per taggere e spingere l'immagine in Registro Artefacti. Usa il nome risultante <catalog>.<schema>.<image>:<tag> come environment.unity_catalog_image nel carico di lavoro YAML.

Tip

In alternativa, usa il databricks air images push comando helper nella CLI di Databricks e segui i prompt interattivi.

Limitations

  • Le immagini devono essere archiviate nel Registro degli Artefatti nello spazio di lavoro dove si invia il carico di lavoro.
  • Le dimensioni dell'immagine devono essere minori di 20 GB.
  • WORKDIR non viene rispettato in fase di esecuzione. Usare percorsi assoluti per i file inseriti nell'immagine. Ad esempio, usare python /app/train.py, non python train.py.
  • Non è possibile usare environment.dependencies o environment.version con environment.unity_catalog_image. Se sono necessari pacchetti aggiuntivi oltre a ciò che si trova nell'immagine, è necessario aggiungerli al Dockerfile.

Troubleshooting

Per gli errori di autenticazione, autorizzazione, push di immagini o individuazione delle immagini relativi ad Artifact Registry, consulta Troubleshoot Artifact Registry.

SSL. SSLError durante il caricamento delle dipendenze

Un'immagine personalizzata può avere esito negativo in fase di esecuzione con un errore OpenSSL quando una libreria tenta di creare un contesto SSL, ad esempio:

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

L'errore viene visualizzato durante l'importazione di librerie che aprono connessioni di rete, ad esempio huggingface_hub, e ne impedisce il caricamento.

Questo avviene perché i carichi di lavoro AI Runtime vengono eseguiti su host con abilitati i Federal Information Processing Standards (FIPS). Quando le librerie di crittografia dell'immagine non sono conformi a FIPS, OpenSSL non riesce a inizializzare in modalità FIPS, quindi la creazione di un contesto SSL non riesce.

Soluzione consigliata:

I carichi di lavoro aziendali, governativi, sanitari e finanziari dipendono spesso dai controlli FIPS 140-2 o 140-3 per i controlli FedRAMP, CMMC o HIPAA. Se il carico di lavoro deve rimanere conforme a FIPS, creare l'immagine con librerie crittografiche conformi a FIPS.

Se il carico di lavoro non richiede la conformità FIPS, è possibile disabilitare la modalità FIPS impostando la OPENSSL_FORCE_FIPS_MODE variabile di ambiente su 0. Facendo ciò, si possono compromettere silenziosamente i requisiti di conformità.

Per disabilitare la modalità FIPS, configurala nel file YAML del carico di lavoro sotto env_variables:

env_variables:
  OPENSSL_FORCE_FIPS_MODE: '0'

In alternativa, impostare la variabile nel Dockerfile in modo che venga applicata a ogni carico di lavoro che usa l'immagine:

ENV OPENSSL_FORCE_FIPS_MODE=0

Invia nuovamente il workload e conferma che l'errore SSL non compaia più durante il caricamento delle dipendenze.

Risorse aggiuntive

  • Ultimo aggiornamento: