Jadwalkan beban kerja GPU dan susun tugas

Gunakan Bundle Otomasi Deklaratif untuk menjadwalkan beban kerja GPU AI Runtime dan menggabungkannya dengan pekerjaan lain, seperti persiapan data hulu. Definisikan beban kerja dan jadwalnya di YAML, lalu deploy keduanya sebagai job. Untuk pipeline pelatihan, jalankan prapemrosesan pada komputasi CPU dan mulai pelatihan GPU setelah data siap.

Panduan ini dimulai dengan beban kerja yang mencetak "Hello world", lalu membangun pipeline pra-pemrosesan dan pelatihan yang terjadwal. Jika Anda sudah memiliki YAML beban kerja AI Runtime, lihat Konversi beban kerja AI Runtime menjadi bundel.

Cara kerjanya

  • Sebuah bundel berisi kode Anda dan databricks.yml konfigurasi.
  • Job mengelompokkan tugas dan menentukan kapan tugas tersebut dijalankan.
  • ai_runtime_task menjalankan perintah Anda di AI Runtime.

Bidang depends_on menghubungkan tugas ke dalam grafik asiklik terarah (DAG), dan bidang schedule menjadwalkan pekerjaan Anda untuk berjalan secara berkala.

databricks bundle deploy mengunggah kode dan membuat atau memperbarui job. databricks bundle run segera memulai eksekusi job. Pada setiap eksekusi, AI Runtime menyediakan komputasi GPU, menjalankan perintah Anda, dan merekam run tersebut dalam eksperimen MLflow bernama.

Requirements

Contoh Hello World dengan ai_runtime_task

Contoh ini menjalankan perintah shell pada satu GPU A10.

  1. Buat direktori untuk bundle. Di direktori tersebut, buat command.sh dengan konten berikut:

    #!/usr/bin/env bash
set -euo pipefail
echo "Hello world"

  2. Buat databricks.yml di direktori yang sama:

    bundle:
  name: hello-ai-runtime

resources:
  jobs:
    hello:
      name: hello-ai-runtime
      tasks:
        - task_key: hello
          environment_key: default
          ai_runtime_task:
            experiment: hello-ai-runtime
            deployments:
              - command_path: ${workspace.file_path}/command.sh
                compute:
                  accelerator_type: GPU_1xA10
                  accelerator_count: 1
      environments:
        - environment_key: default
          spec:
            environment_version: '6'

targets:
  dev:
    mode: development
    default: true

  3. Dari direktori bundle, validasi, sebarkan, dan jalankan pekerjaan:

    databricks bundle validate --target dev
databricks bundle deploy --target dev
databricks bundle run hello --target dev

  4. Buka URL run yang dicetak oleh CLI dan lihat output tugas hello. Ini mengandung Hello world. Run ini juga muncul dalam hello-ai-runtime eksperimen MLflow.

Contoh yang lebih kompleks: Jadwalkan pipeline pra-pemrosesan dan pelatihan

Contoh ini menyiapkan dataset kecil pada komputasi CPU, lalu melatih model linear pada satu GPU A10. Kedua tugas menggunakan file volume yang sama untuk meneruskan data di antara keduanya. Pekerjaan dijadwalkan pada pukul 09:00 UTC setiap hari, dengan jadwal dijeda sampai Anda mengujinya.

Membuat proyek

  1. Buat direktori terpisah dengan tata letak berikut:

    scheduled-training/
├── databricks.yml
├── prep.py
├── command.sh
└── src/
    └── train.py

  2. Buat prep.py dengan konten berikut. Ganti <catalog>, <schema>, dan <volume> dengan nama volume Anda. Notebook sumber ini menormalkan nilai input dan menulis data yang telah disiapkan ke volume:

    # Databricks notebook source
import json
from pathlib import Path

data_path = Path("/Volumes/<catalog>/<schema>/<volume>/scheduled-training/data.json")
inputs = [value / 10 for value in range(10)]
data = {"x": inputs, "y": [2 * value + 1 for value in inputs]}
data_path.parent.mkdir(parents=True, exist_ok=True)
data_path.write_text(json.dumps(data))
print(f"Prepared {len(inputs)} rows at {data_path}")

  3. Buat src/train.py dengan konten berikut. Gunakan jalur volume yang sama seperti di prep.py:

    import json
from pathlib import Path

import torch

data_path = Path("/Volumes/<catalog>/<schema>/<volume>/scheduled-training/data.json")
data = json.loads(data_path.read_text())
x = torch.tensor(data["x"], device="cuda").reshape(-1, 1)
y = torch.tensor(data["y"], device="cuda").reshape(-1, 1)
model = torch.nn.Linear(1, 1).to("cuda")
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)

for _ in range(200):
    optimizer.zero_grad()
    loss = torch.nn.functional.mse_loss(model(x), y)
    loss.backward()
    optimizer.step()

print(f"Trained on {x.device}: loss={loss.item():.4f}")

  4. Buat command.sh untuk menjalankan kode pelatihan. Bundel di bawah ini berisi src/, jadi CODE_SOURCE_PATH menunjuk ke direktori src yang diekstrak:

    #!/usr/bin/env bash
set -euo pipefail
cd "$CODE_SOURCE_PATH"
python train.py

  5. Buat databricks.yml dengan konten berikut. Bundle ini mengemas src/ untuk tugas GPU. prep Notebook berjalan dengan komputasi CPU serverless, dan depends_on hanya mulai train setelah prep berhasil. max_concurrent_runs: 1 mencegah eksekusi pekerjaan ini menimpa file input milik satu sama lain:

    bundle:
  name: scheduled-training

artifacts:
  code:
    type: tgz
    path: .
    include: [src]
    files:
      - source: ./dist/code.tgz

resources:
  jobs:
    train_pipeline:
      name: scheduled-training
      max_concurrent_runs: 1
      schedule:
        quartz_cron_expression: '0 0 9 * * ?'
        timezone_id: UTC
        pause_status: PAUSED
      tasks:
        - task_key: prep
          notebook_task:
            notebook_path: ./prep.py
        - task_key: train
          depends_on:
            - task_key: prep
          environment_key: training
          ai_runtime_task:
            experiment: scheduled-training
            code_source_path: ./dist/code.tgz
            deployments:
              - command_path: ${workspace.file_path}/command.sh
                compute:
                  accelerator_type: GPU_1xA10
                  accelerator_count: 1
      environments:
        - environment_key: training
          spec:
            environment_version: '6'
            dependencies:
              - torch

targets:
  dev:
    mode: development
    default: true

Uji dan aktifkan jadwal

  1. Dari scheduled-training/, deploy bundle dan mulai proses manual:

    databricks bundle validate --target dev
databricks bundle deploy --target dev
databricks bundle run train_pipeline --target dev

  2. Buka URL run yang dicetak oleh CLI. Pastikan prep berhasil sebelum train dimulai. Output prep melaporkan 10 baris yang sudah disiapkan, serta output train melaporkan Trained on cuda:0 dan kerugian.

  3. Di databricks.yml, ubah jadwal pause_status dari PAUSED ke UNPAUSED. Blok jadwal sekarang adalah:

    schedule:
  quartz_cron_expression: '0 0 9 * * ?'
  timezone_id: UTC
  pause_status: UNPAUSED

  4. Terapkan perubahan untuk mengaktifkan jadwal harian:

    databricks bundle deploy --target dev

    Menetapkan pause_status: UNPAUSED secara eksplisit memungkinkan jadwal ini bahkan untuk target pengembangan. Di Jobs & Pipelines, buka pekerjaan yang telah dideploy dan konfirmasi bahwa jadwalnya aktif pada pukul 09:00 UTC. Lihat Menjalankan pekerjaan sesuai jadwal.

Untuk menjeda pelaksanaan terjadwal, atur pause_status: PAUSED lalu deploy lagi. Untuk menghapus pekerjaan contoh dan file yang diunggah oleh bundle, jalankan hal berikut dari direktori bundle-nya:

databricks bundle destroy --target dev

Data yang sudah disiapkan dalam volume Anda tetap ada. Hapus scheduled-training direktori data saat Anda tidak lagi membutuhkannya.

Sumber daya tambahan

  • Terakhir diperbarui pada