Gunakan image Docker kustom dengan AI Runtime

Important

Fitur ini ada di Beta. Untuk menggunakannya, admin ruang kerja harus mengaktifkan pratinjau AI Runtime Beta Features dan Databricks Artifact Registry dari halaman Pratinjau ruang kerja.

AI Runtime dapat menjalankan citra kontainer Docker kustom yang disimpan di Artifact Registry. Gunakan gambar kustom saat Anda membutuhkan:

  • Perpustakaan sistem atau dependensi kompleks yang tidak dapat Anda instal dengan environment.dependencies.
  • Lingkungan yang dapat direproduksi di seluruh pengembangan, penelitian, dan produksi.
  • Image privat yang disetujui organisasi dan dibuat oleh platform atau tim keamanan Anda.

Prerequisites

Dorong gambar ke Registri Artefak

Sebelum menggunakan gambar kustom dengan AI Runtime, simpan di Artifact Registry di workspace yang sama dengan tempat Anda mengirimkan beban kerja.

  1. Buat atau pilih katalog dan skema Unity Catalog untuk gambar tersebut.

  2. Ikuti Mulai dengan Artifact Registry untuk mengonfigurasi autentikasi Docker, memberikan hak istimewa yang diperlukan, dan mendorong gambar.

  3. Catat nama image di Unity Catalog dalam format berikut:

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

    Contohnya, main.ml.training:v1. Jangan sertakan nama host registry dalam konfigurasi beban kerja.

Tip

Sebagai alternatif, gunakan databricks air images push perintah helper di Databricks CLI.

Gunakan citra Docker dalam beban kerja

Tentukan nama Unity Catalog untuk image tersebut di YAML workload Anda pada bagian 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

Contoh ini menjalankan train.py dari /app pada gambar. Untuk mengunggah kode aplikasi secara terpisah tanpa membangun ulang gambar, lihat Jalankan kode aplikasi yang diunggah.

Saat membawa gambar Docker Anda sendiri, environment.dependencies dan environment.version tidak didukung. Menentukan environment.unity_catalog_image dengan salah satu bidang memicu kesalahan. Jika Anda memiliki dependensi tambahan, instal paket di Dockerfile sebagai gantinya.

Kirimkan beban kerja:

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

Profil harus mengautentikasi ke workspace yang sama tempat gambar disimpan.

Variabel lingkungan yang disuntikkan ke dalam kontainer Anda

Runtime AI menyuntikkan variabel lingkungan berikut ke setiap kontainer saat runtime:

  • CODE_SOURCE_PATH: jalur ke kode aplikasi yang diunggah, saat code_source dikonfigurasi.
  • NUM_NODES: jumlah total node.
  • LOCAL_WORLD_SIZE: GPU per node.
  • WORLD_SIZE: jumlah total proses.
  • POD_RANK: peringkat node saat ini (0-indeks). Juga diinjeksikan sebagai NODE_RANK.
  • MASTER_ADDR: alamat koordinasi untuk rank-0 (hanya untuk multi-node).
  • MASTER_PORT: port koordinasi rank-0 (hanya multi-node).

Examples

Contoh berikut menunjukkan cara menjalankan kode aplikasi yang diunggah dan pelatihan terdistribusi dengan gambar kustom.

Jalankan kode aplikasi yang diunggah

Gunakan code_source untuk mengunggah kode aplikasi secara terpisah dari gambar. Anda dapat mengedit dan mengirim ulang kode Anda tanpa membangun ulang gambar. Instal dependensi Python dan sistem untuk kode tersebut di dalam image.

Tempatkan train.py di direktori lokal src di sebelah workload.yaml. Konfigurasi berikut mengunggah src dan menjalankan train.py-nya di dalam image kustom:

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 diuraikan relatif terhadap workload.yaml. AI Runtime menetapkan $CODE_SOURCE_PATH ke jalur direktori yang diunggah di dalam kontainer. Lihat code_source untuk opsi sumber kode.

H100 multinode dengan RDMA

Untuk job H100 multinode yang memerlukan bandwidth jaringan penuh pada instans AWS p5, gunakan salah satu image dasar Databricks sebagai basis image Anda dengan NCCL dan EFA yang telah diprakonfigurasi:

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

Bangun gambar Anda sendiri

Saat membangun image Anda sendiri, Databricks merekomendasikan penggunaan databricks-ai-runtime skill dengan agen coding atau memulai dari image dasar Databricks.

Menggunakan agen pengodian

Instal keterampilan Claude Code databricks-ai-runtime untuk panduan langkah demi langkah tentang Dockerfile, termasuk membangun dari nol, kompatibilitas CUDA/NCCL/EFA, masalah umum, dan daftar periksa sebelum proses pembangunan. Keterampilan ini membutuhkan Databricks CLI versi 1.0.0 atau yang lebih baru.

databricks aitools install --skills databricks-ai-runtime

Gambar dasar Databricks

Databricks menerbitkan gambar dasar pada Docker Hub di databricksruntime/air dengan CUDA, NCCL, dan jaringan khusus cloud (AWS EFA atau Azure InfiniBand) yang telah dikonfigurasi sebelumnya.

Tag Variant CUDA Gunakan ketika diperlukan
dcs-base-azure-runtime Runtime 12 Menginstal roda bawaan saja
dcs-base-azure-devel Devel 12 Mengkompilasi ekstensi CUDA (memerlukan nvcc)

Dockerfile berikut menambahkan PyTorch ke citra dasar Databricks. Gambar dasar menyediakan Python di /opt/venv, dikelola oleh uv. uv pip install Secara default menargetkan lingkungan tersebut. Untuk menggunakan lingkungan yang berbeda, buat dan aktifkan venv sebelum menjalankan uv pip install.

Untuk menyertakan skrip pelatihan Anda ke dalam gambar, letakkan train.py di sebelah Dockerfile. Dockerfile menyalinnya ke /app/train.py. Jika Anda mengunggah kode aplikasi dengan code_source, abaikan COPY instruksi dan simpan train.py di direktori sumber Anda.

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

Bangun gambar secara lokal:

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

Kemudian ikuti Mulai dengan Artifact Registry untuk menandai dan mendorong gambar ke Artifact Registry. Gunakan nama yang dihasilkan <catalog>.<schema>.<image>:<tag> seperti environment.unity_catalog_image pada workload YAML.

Tip

Sebagai alternatif, gunakan databricks air images push perintah helper di Databricks CLI dan ikuti petunjuk interaktif.

Limitations

  • Gambar harus disimpan di Registri Artefak di ruang kerja tempat Anda mengirimkan beban kerja.
  • Ukuran gambar harus di bawah 20 GB.
  • WORKDIR tidak diterapkan saat dijalankan. Gunakan path absolut untuk file yang sudah disertakan di dalam image. Misalnya, gunakan python /app/train.py, bukan python train.py.
  • Anda tidak dapat menggunakan environment.dependencies atau environment.version dengan environment.unity_catalog_image. Jika Anda membutuhkan paket tambahan di luar apa yang ada dalam gambar, Anda harus menambahkannya ke Dockerfile.

Troubleshooting

Untuk kesalahan autentikasi terkait registry, izin, dorongan gambar, atau penemuan gambar, lihat Troubleshoot Artifact Registry.

ssl. SSLError saat memuat ketergantungan

Gambar kustom dapat gagal pada runtime dengan kesalahan OpenSSL saat pustaka mencoba membuat konteks SSL, misalnya:

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

Kesalahan muncul saat mengimpor pustaka yang membuka koneksi jaringan, seperti huggingface_hub, dan mencegah pustaka tersebut dimuat.

Hal ini terjadi karena beban kerja AI Runtime berjalan pada host dengan Federal Information Processing Standards (FIPS) diaktifkan. Ketika pustaka kriptografi gambar tidak sesuai FIPS, OpenSSL gagal menginisialisasi dalam mode FIPS, sehingga pembuatan konteks SSL gagal.

Solusi yang direkomendasikan:

Beban kerja perusahaan, pemerintah, layanan kesehatan, dan keuangan sering bergantung pada kepatuhan FIPS 140-2 atau 140-3 untuk audit FedRAMP, CMMC, atau HIPAA. Jika beban kerja Anda harus tetap sesuai dengan FIPS, buat image Anda dengan pustaka kriptografis yang sesuai dengan FIPS.

Jika beban kerja Anda tidak memerlukan kepatuhan FIPS, Anda dapat menonaktifkan mode FIPS dengan mengatur OPENSSL_FORCE_FIPS_MODE variabel lingkungan ke 0. Melakukannya dapat secara diam-diam melanggar persyaratan kepatuhan.

Untuk menonaktifkan mode FIPS, atur di YAML beban kerja Anda pada bagian env_variables:

env_variables:
  OPENSSL_FORCE_FIPS_MODE: '0'

Atau, atur variabel di Dockerfile Anda sehingga berlaku untuk setiap beban kerja yang menggunakan gambar:

ENV OPENSSL_FORCE_FIPS_MODE=0

Kirim ulang beban kerja dan konfirmasikan kesalahan SSL tidak lagi muncul saat dependensi dimuat.

Sumber daya tambahan

  • Terakhir diperbarui pada