Pianificare i carichi di lavoro GPU e creare attività

Importante

Questa funzionalità è in Anteprima Pubblica. Per utilizzarlo, un amministratore dello spazio di lavoro deve abilitare l'anteprima AI Runtime dalla pagina delle Anteprime . Vedere Gestire le anteprime di Azure Databricks.

Usa Declarative Automation Bundles per programmare un carico di lavoro GPU di AI Runtime e combinarlo con altri lavori, come la preparazione dei dati a monte. Definisci il carico di lavoro e il suo programma in YAML, poi esegui il deployment come un job. Per una pipeline di addestramento, esegui la preelaborazione su risorse di calcolo CPU e inizia l'addestramento GPU dopo che i dati sono pronti.

Questa guida inizia con un carico di lavoro che stampa "Hello world", poi costruisce una pipeline programmata di preprocessing e addestramento. Se hai già un file YAML di workload AI Runtime, vedi Conversione di un workload AI Runtime in un bundle.

Come funziona

  • Un bundle contiene il tuo codice e una databricks.yml configurazione.
  • Un lavoro raggruppa i compiti e definisce quando vengono eseguiti.
  • Un ai_runtime_task esegue il tuo comando su una GPU serverless.

Il campo depends_on collega i compiti in un grafo aciclico diretto (DAG) e il campo schedule pianifica l'esecuzione del processo a intervalli regolari.

databricks bundle deploy carica il codice e crea o aggiorna il processo. databricks bundle run avvia immediatamente un'esecuzione del processo. Ad ogni esecuzione, AI Runtime provisiona le risorse di calcolo GPU, esegue il tuo comando e registra l'esecuzione nell'esperimento MLflow con nome specificato.

Requirements

Esempio di Hello World con ai_runtime_task

Questo esempio esegue un comando shell su una GPU A10.

  1. Crea una directory per il bundle. In quella directory, crea command.sh con i seguenti contenuti:

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

  2. Crea databricks.yml nella stessa directory:

    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. Dalla cartella bundle, validare, distribuire ed eseguire il job:

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

  4. Apri l'URL di esecuzione stampato dalla CLI e visualizza l'output dell'attività hello. Contiene Hello world. Il run appare anche nell'esperimento hello-ai-runtime MLflow.

Esempio più complesso: programmare una pipeline di preprocessing e addestramento

Questo esempio prepara un piccolo dataset sul calcolo CPU, poi addestra un modello lineare su una GPU A10. Entrambe le attività utilizzano lo stesso file di volume per trasferire dati tra di esse. Il lavoro è programmato per le 09:00 UTC ogni giorno, con il programma sospeso finché non lo testi.

Creare il progetto

  1. Crea una directory separata con il seguente layout:

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

  2. Crea prep.py con i seguenti contenuti. Sostituisci <catalog>, <schema>, e <volume> con i nomi dei tuoi volumi. Questo notebook di origine normalizza i valori di input e scrive i dati preparati nel 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. Crea src/train.py con i seguenti contenuti. Usa lo stesso percorso di volume 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. Crea command.sh per eseguire il codice di addestramento. Il bundle qui sotto impacchettisce src/, quindi CODE_SOURCE_PATH punta alla directory estratta src :

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

  5. Crea databricks.yml con i seguenti contenuti. Il bundle crea pacchetti src/ per l'attività GPU. Il prep notebook funziona su CPU serverless e depends_on inizia train solo dopo che prep riesce. max_concurrent_runs: 1 impedisce che le esecuzioni di questo lavoro sovrascrivano il file di input dell'altro:

    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

Testare e abilitare la pianificazione

  1. Da scheduled-training/, esegui il deployment del bundle e avvia un'esecuzione manuale:

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

  2. Apri l'URL di esecuzione stampato dalla CLI. Conferma che prep venga completato correttamente prima che train inizi. L'output prep riporta 10 righe preparate, e l'output train riporta Trained on cuda:0 e la perdita.

  3. In databricks.yml, cambia l'pause_status della pianificazione da PAUSED a UNPAUSED. Il blocco di pianificazione è ora:

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

  4. Distribuisci la modifica per attivare la pianificazione giornaliera:

    databricks bundle deploy --target dev

    L'impostazione esplicita di pause_status: UNPAUSED abilita questa pianificazione anche per una destinazione di sviluppo. In Jobs & Pipelines, apri il job distribuito e conferma che la sua pianificazione sia attiva per le 09:00 UTC. Consultare Eseguire lavori secondo una pianificazione.

Per mettere in pausa le esecuzioni programmate, imposta pause_status: PAUSED e distribuisci di nuovo. Per rimuovere il job di esempio e i file caricati dal bundle, esegui quanto segue dalla sua directory bundle:

databricks bundle destroy --target dev

I dati preparati nel tuo volume rimangono. Elimina la scheduled-training directory dati quando non ne hai più bisogno.

Risorse aggiuntive

  • Ultimo aggiornamento: