GPU-Workloads planen und Aufgaben erstellen

Verwenden Sie Declarative Automation Bundles, um eine KI-Laufzeit-GPU-Workload zu planen und sie mit anderer Arbeit zu kombinieren, wie z. B. der Upstream-Datenvorbereitung. Definiere die Arbeitslast und ihren Zeitplan in YAML und deploye sie dann als Job. Für eine Trainingspipeline führen Sie die Vorverarbeitung auf dem CPU-Compute aus und starten Sie das GPU-Training, nachdem die Daten bereit sind.

Dieser Leitfaden beginnt mit einer Arbeitslast, die „Hello world“ ausgibt, und baut anschließend eine zeitgesteuerte Pipeline für die Vorverarbeitung und das Training auf. Wenn Sie bereits ein KI-Runtime-Workload-YAML haben, lesen Sie KI-Runtime-Workload in ein Bundle konvertieren.

Funktionsweise

  • Ein Bundle enthält deinen Code und eine databricks.yml Konfiguration.
  • Ein Job gruppiert Aufgaben und definiert, wann sie ausgeführt werden.
  • Ein ai_runtime_task führt Ihren Befehl auf der KI-Laufzeit aus.

Das depends_on Feld verbindet Aufgaben zu einem gerichteten azyklischen Graphen (DAG), und das schedule Feld plant den Job so ein, dass er in einem Rhythmus läuft.

databricks bundle deploy Lädt den Code hoch und erstellt oder aktualisiert den Job. databricks bundle run startet sofort eine Jobausführung. Bei jedem Durchlauf stellt AI Runtime die GPU-Rechenleistung bereit, führt Ihren Befehl aus und zeichnet den Run im benannten MLflow-Experiment auf.

Requirements

Hallo-Welt-Beispiel mit ai_runtime_task

Dieses Beispiel führt einen Shell-Befehl auf einer A10-GPU aus.

  1. Erstelle ein Verzeichnis für das Bundle. In diesem Verzeichnis erstellen Sie command.sh mit folgendem Inhalt:

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

  2. Im selben Verzeichnis databricks.yml erstellen:

    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. Aus dem Bundle-Verzeichnis validieren, deployen und führen Sie den Job aus:

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

  4. Öffnen Sie die von der CLI ausgedruckte Run-URL und sehen Sie sich die Ausgabe der hello Aufgabe an. Sie enthält Hello world. Der Lauf taucht auch im hello-ai-runtime MLflow-Experiment auf.

Komplexeres Beispiel: Planen Sie eine Preprocessing- und Trainingspipeline

Dieses Beispiel bereitet einen kleinen Datensatz auf der CPU vor und trainiert dann ein lineares Modell auf einer A10-GPU. Beide Aufgaben verwenden dieselbe Volumendatei, um Daten zwischen ihnen zu übertragen. Der Job ist täglich um 09:00 UTC angesetzt, der Zeitplan ist pausiert, bis du ihn testest.

Erstelle das Projekt

  1. Erstellen Sie ein separates Verzeichnis mit folgendem Layout:

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

  2. Erstellen Sie prep.py mit folgendem Inhalt. Ersetze <catalog>, <schema> und <volume> durch die Namen deiner Volumes. Dieses Quellnotebook normalisiert die Eingabewerte und schreibt die vorbereiteten Daten auf das 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. Erstellen Sie src/train.py mit den folgenden Inhalten. Verwenden Sie denselben Volume-Pfad wie in 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. Erstellen Sie command.sh, um den Trainingscode auszuführen. Das folgende Bundle paketiert src/, sodass CODE_SOURCE_PATH auf das extrahierte src-Verzeichnis verweist:

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

  5. Erstellen Sie databricks.yml mit den folgenden Inhalten. Das Paket enthält src/ für die GPU-Aufgabe. Das prep Notebook läuft mit serverloser CPU-Berechnung und depends_on startet train erst, wenn prep erfolgreich abgeschlossen wird. max_concurrent_runs: 1 verhindert, dass Läufe dieses Jobs gegenseitig die Eingabedatei des anderen überschreiben:

    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

Zeitplan testen und aktivieren

  1. Von scheduled-training/ aus stelle das Bundle bereit und starte einen manuellen Lauf:

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

  2. Öffnen Sie die von der CLI ausgedruckte Run-URL. Bestätigen Sie, dass prep erfolgreich ist, bevor train startet. Die Ausgabe prep meldet 10 vorbereitete Zeilen, und die Ausgabe train meldet Trained on cuda:0 und den Verlust.

  3. In databricks.yml ändere pause_status des Zeitplans von PAUSED zu UNPAUSED. Der Zeitplanblock ist jetzt:

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

  4. Stellen Sie die Änderung bereit, um den täglichen Zeitplan zu aktivieren:

    databricks bundle deploy --target dev

    Das Festlegen von pause_status: UNPAUSED aktiviert diese Zeitplanvorgabe explizit selbst für ein Entwicklungsziel. In Jobs & Pipelines öffnen Sie den bereitgestellten Job und bestätigen, dass sein Zeitplan für 09:00 UTC aktiv ist. Siehe Ausführen von Jobs nach einem Zeitplan.

Um geplante Durchläufe zu pausieren, setzen Sie pause_status: PAUSED und deployen erneut. Um den Beispieljob und die vom Bundle hochgeladenen Dateien zu entfernen, führen Sie Folgendes aus dem Bundle-Verzeichnis aus:

databricks bundle destroy --target dev

Die vorbereiteten Daten in Ihrem Volume bleiben erhalten. Lösche das scheduled-training Datenverzeichnis, wenn du es nicht mehr brauchst.

Weitere Ressourcen

  • Zuletzt aktualisiert um