Workload YAML-Referenz

Important

Dieses Feature befindet sich in der Public Preview.

Definieren Sie den Namen des Experiments eines Schulungsauftrags, die Berechnung, den Befehl, die Umgebung und die Codequelle in der YaML-Konfiguration der Arbeitsauslastung, an die air run --fileSie übergeben. Diese Seite dokumentiert jedes Feld.

Note

Die grundwahre Wahrheit für die YAML-Konfiguration ist die In-CLI-Hilfe. Führen Sie air -h config für die Ansicht der obersten Ebene und air -h config.<section> (z. B air -h config.environment. ) für Detaildetails pro Abschnitt aus.

Minimale Konfiguration

experiment_name: my-training
environment:
  dependencies:
    - mlflow
compute:
  num_accelerators: 1
  accelerator_type: GPU_1xA10
command: echo "Hello World"

Übermitteln mit:

air run --file train.yaml -p profile

Wichtige Konzepte

Kernfelder

Die meisten Schulungskonfigurationen umfassen fünf Komponenten:

  1. experiment_name (Erforderlich): Erstellt oder fügt ein MLflow-Experiment hinzu.
  2. environment(Optional): Python-Abhängigkeiten und Basisumgebungsversion.
  3. compute (Erforderlich): GPU-Ressourcen (Typ und Anzahl).
  4. command (Erforderlich): Das Bash-Kommando oder die Befehle, die zum Start der Ausbildung verwendet werden.
  5. code_source (Optional): Pfad zu deinem Trainingscode, aus der Ferne verfügbar.

Für unterstützte Werte und Feldbeschränkungen siehe Referenz.

Ihr erster Ausbildungsauftrag

experiment_name: simple-training
environment:
  dependencies:
    - torch
    - transformers
compute:
  num_accelerators: 8
  accelerator_type: GPU_8xH100
code_source:
  type: snapshot
  snapshot:
    root_path: /home/username/repo
command: torchrun --nproc_per_node=8 $CODE_SOURCE_PATH/train.py

In dieser Konfiguration:

  • experiment_name erstellt ein MLflow-Experiment namens simple-training (oder fügt eine neue Ausführung an, wenn es bereits vorhanden ist).
  • environment verwendet die Standardumgebung und installiert torch und transformers.
  • compute weist einen H100-Knoten zu (8 H100 GPUs).
  • code_source lädt den Ordner repo in den Knoten hoch, verfügbar unter $CODE_SOURCE_PATH.
  • command läuft train.py über torchrun die 8 H100 GPUs. Die Datei befindet sich lokal /home/username/repo/train.py .

Häufige Anwendungsfälle

Hinzufügen von Umgebungsvariablen

experiment_name: training-with-env
environment:
  dependencies:
    - torch
    - transformers
env_variables:
  BATCH_SIZE: '32'
  LEARNING_RATE: '0.001'
compute:
  num_accelerators: 8
  accelerator_type: GPU_8xH100
code_source:
  type: snapshot
  snapshot:
    root_path: /home/username/repo
    git:
      branch: main
command: torchrun --nproc_per_node=8 train.py

Verwenden von geheimen Schlüsseln (API-Schlüssel, Token)

experiment_name: training-with-secrets
environment:
  dependencies:
    - torch
    - transformers
secrets:
  HF_TOKEN: 'my_scope/hf_token'
  WANDB_API_KEY: 'my_scope/wandb'
compute:
  num_accelerators: 8
  accelerator_type: GPU_8xH100
code_source:
  type: snapshot
  snapshot:
    root_path: /home/username/repo
    git:
      branch: main
command: torchrun --nproc_per_node=8 train.py

Geheime Schlüssel verwenden das Format scope/key und müssen in Databricks Secrets konfiguriert werden. Informationen zum Einrichten finden Sie unter "Geheime Verwaltung ".

Wenn Sie eine YAML-Vorlage freigeben, müssen andere Benutzer eigene geheime Schlüssel erstellen oder Zugriff auf den geheimen Schlüssel haben.

Environment

Verwenden Sie den Block, environment um eine serverlose GPU-Umgebung auszuwählen und Python-Abhängigkeiten zu installieren. Zum Beispiel wählt die folgende Konfiguration die Standardumgebung Version 4 und installiert PyTorch und Transformers:

environment:
  version: '4'
  dependencies:
    - torch
    - transformers

Umgebungsversion

environment.version ist optional und wählt die verwaltete Umgebungsversion für die Arbeitslast aus.

Dazu gehören:

  • "4" oder "5" die entsprechende Standardumgebungsversion zu verwenden.
  • "databricks_ai_v5" um die Databricks KI-Umgebung Version 5 zu verwenden, die vorinstallierte ML-spezifische Pakete enthält. (Vollständige Paketliste)

Das folgende Beispiel wählt die Databricks-KI-Umgebung Version 5 aus:

environment:
  version: 'databricks_ai_v5'
  dependencies: []

Wenn Sie angeben environment.version, müssen Sie auch eine Inline-Liste angeben environment.dependencies . Verwenden Sie eine leere Liste, wenn Sie keine zusätzlichen Pakete installieren müssen.

Informationen zu Umgebungen, die für AI Runtime verfügbar sind, finden Sie unter "Einrichten Ihrer Umgebung".

Python Abhängigkeiten

Liste die Python-Abhängigkeiten deiner Arbeitslast als Inline-Liste unter environment.dependencies.

Abhängigkeitsformat

Die Abhängigkeitsliste folgt der Databricks Base Environment Specification. Jeder Eintrag ist eine Pip-Style-Paketspezifikation (z. B my-library==6.1. ). Die Liste akzeptiert auch die folgenden Einträge:

  • Anforderungsdateien: eine Referenz auf eine bestehende requirements.txt Verwendung von -r, zum Beispiel -r '/Workspace/Shared/requirements.txt'. Umgebungsvariablen wie $HOME z. B. werden erweitert.
  • Wheels: ein absoluter Pfad zu einer .whl Datei, z. B /Workspace/Shared/path/to/simplejson-3.19.3-py3-none-any.whl. .
  • Index-URLs: eine Index-URL, z. B --index-url https://pypi.org/simple. .
environment:
  version: '4'
  dependencies:
    - --index-url https://pypi.org/simple
    - -r '/Workspace/Shared/requirements.txt'
    - my-library==6.1
    - /Workspace/Shared/path/to/simplejson-3.19.3-py3-none-any.whl

Unterstützte Installationsflags

Abhängigkeiten werden mit uv installiert. Die folgenden Pip-Style-Flags werden als Listeneinträge unterstützt:

  • Wird auf die gesamte Installation angewendet: --index-url, --extra-index-urlund --find-links (-f) legen oder erweitern Sie die Paketindizes.
  • Wird auf die Abhängigkeit angewendet, die ihnen folgt: --no-deps, --no-build-isolation, , , --no-cache-dirund --force-reinstall. Platzieren Sie die Kennzeichnung in einer eigenen Zeile (oder vor der Spezifikation), gefolgt von der Abhängigkeit, auf die sie angewendet wird.

Um beispielsweise eine Installation mit der bereits installierten flash-attn (keine Buildisolation) und ohne Auflösung eigener Abhängigkeiten zu erstellentorch:

environment:
  version: '4'
  dependencies:
    - torch
    - --no-build-isolation
    - --no-deps
    - flash-attn

Note

--trusted-host wird nicht unterstützt. Da uv die Vertrauensstellung pro Index-URL konfiguriert, verwenden --index-url Oder --extra-index-url stattdessen.

Benutzerdefinierte Docker-Images

Alternativ dazu environment.dependencieskönnen Sie ein benutzerdefiniertes Docker-Containerimage mithilfe von environment.docker_image.url. environment.docker_image.url ist sich gegenseitig ausschließend mit beiden environment.dependencies und environment.version — Sie können beides nicht in derselben Workload verwenden.

experiment_name: my-dcs-training
environment:
  docker_image:
    url: myorg/myrepo:mytag
compute:
  num_accelerators: 1
  accelerator_type: GPU_1xA10
command: python /app/train.py

Bevor Sie ein benutzerdefiniertes Bild verwenden, registrieren Sie es mit air register image. Ausführliche Informationen, einschließlich Imageanforderungen, Databricks-Basisimages und Dockerfile-Mustern, finden Sie unter Verwenden von benutzerdefinierten Docker-Images.

Arbeiten mit Codequellen

Der code_source Block lädt lokalen Code hoch, damit der Schulungsauftrag ihn ausführen kann.

  • root_path ist das lokale Verzeichnis zur Momentaufnahme. Packt standardmäßig air die Arbeitsstruktur as-is (einschließlich aller nicht ausgelassenen Änderungen) als einfaches Tarball.
  • Um stattdessen eine angeheftete Git-Version zu snapshotn, fügen Sie einen git: Block mit einem oder einem branch .commit Dies muss root_path ein Git-Repository sein und ermöglicht versionsfähige Momentaufnahmen (Zwischenspeichern, git archive).
  • Bei großen Repositorys include_paths können Sie eine Teilmenge snapshotn.

Minimales Beispiel

experiment_name: simple-training
environment:
  dependencies:
    - torch
    - transformers
compute:
  num_accelerators: 8
  accelerator_type: GPU_8xH100
code_source:
  type: snapshot
  snapshot:
    root_path: /home/username/repo
command: python $CODE_SOURCE_PATH/train.py

Auf dem Remotecomputer wird der Code an der Stelle platziert, an /databricks/code_source/<directory_name>der <directory_name> sich die endgültige Pfadkomponente befindet root_path. $CODE_SOURCE_PATH ist auf diesen absoluten Pfad festgelegt, verwenden Sie ihn also in Ihrem Befehl, anstatt den Speicherort hart zu codieren.

Git-Repositorys: Anheften nach Verzweigung oder Commit

Fügen Sie für Git-Repositorys einen git: Block hinzu, um die Codeversion per Verzweigung oder durch Commit-SHA anzuheften. branch und commit schließen sich gegenseitig aus: Geben Sie genau einen innerhalb des Blocks an.

An eine Verzweigung anheften (verwendet den lokalen HEAD dieser Verzweigung):

code_source:
  type: snapshot
  snapshot:
    root_path: /home/username/repo
    git:
      branch: main # Uses local HEAD of main (no remote fetch)
command: train.sh

An einen Commit SHA anheften (genaue Reproduzierbarkeit):

code_source:
  type: snapshot
  snapshot:
    root_path: /home/username/repo
    git:
      commit: abc1234567 # Pins specific commit
command: train.sh

Schlüsselfelder:

  • root_path (Erforderlich): Lokaler Pfad zum Stammverzeichnis Ihres Git-Repositorys.
  • git.branch (Optional): Verzweigungsname. Verwendet lokale HEAD; kein Remoteabruf. Gegenseitiger Ausschluss mit git.commit
  • git.commit (Optional): Spezifischer Commit-SHA. Gegenseitiger Ausschluss mit git.branch
  • git.remote (Optional): Verwenden Sie den Remote-HEAD der Verzweigung anstelle der lokalen. Legen Sie diese Einstellung fest true , um die Remoteverbindung automatisch zu erkennen, oder auf einen Remotenamen (z. B upstream. auf einen Remotenamen), der von einer bestimmten Remote abgerufen werden soll. Nur gültig mit git.branch.

Wenn Sie den git: Block weglassen, air packt die Arbeitsstruktur als einfache Tarball, einschließlich aller nicht ausgelassenen Änderungen. Es ist kein zusätzliches Feld erforderlich.

Nicht-Git-Verzeichnisse

Sie können Snapshotverzeichnisse erstellen, die keine Git-Repositorys sind. Lassen Sie den git: Block aus, der ein Git-Repository sein muss root_path . Ohne sie gibt es keine Versionszwischenspeicherung; Für jede Ausführung wird ein frischer Tarball hochgeladen.

code_source:
  type: snapshot
  snapshot:
    root_path: /home/username/my_project
command: $CODE_SOURCE_PATH/train.py

Ordnerfilterung mit include_paths

Bei großen Monoreposs werden nur bestimmte Ordner zum Reduzieren der Upload- und Downloadzeit und Momentaufnahmegröße von Snapshots verwendet:

code_source:
  type: snapshot
  snapshot:
    root_path: /home/username/repo
    include_paths:
      - research/models
      - research/common
      - research/configs
command: python $CODE_SOURCE_PATH/research/models/launch_training.py

Wichtige Punkte:

  • Das Feld ist optional. Wenn nicht angegeben, ist das gesamte Repository standardmäßig enthalten.
  • Pfade müssen relativ zum Repositorystamm (keine führende) /sein.
  • .. ist nicht zulässig; Sie können nicht auf übergeordnete Verzeichnisse verweisen.

Erweiterte Funktionen

Benutzerdefinierte Hyperparameter

Übergeben Sie die strukturierte Konfiguration an Ihr Schulungsskript über HYPERPARAMETERS_PATH:

experiment_name: parameterized-training
environment:
  dependencies:
    - torch
    - transformers
compute:
  num_accelerators: 8
  accelerator_type: GPU_8xH100
code_source:
  type: snapshot
  snapshot:
    root_path: /home/username/repo
    git:
      branch: main
command: torchrun --nproc_per_node=8 train.py
parameters:
  model:
    name: 'gpt2'
    hidden_size: 768
  training:
    batch_size: 32
    learning_rate: 0.0001

Lesen Sie sie in Ihrem Skript:

import os
import yaml

with open(os.environ['HYPERPARAMETERS_PATH']) as f:
    params = yaml.safe_load(f)

learning_rate = params['training']['learning_rate']
model_name = params['model']['name']

Auftragssicherheit

experiment_name: reliable-training
environment:
  dependencies:
    - torch
    - transformers
compute:
  num_accelerators: 8
  accelerator_type: GPU_8xH100
code_source:
  type: snapshot
  snapshot:
    root_path: /home/username/repo
    git:
      branch: main
command: torchrun --nproc_per_node=8 train.py
max_retries: 2
timeout_minutes: 90

Wenn die Workload fehlschlägt, wird sie zweimal wiederholt. Jeder Versuch hat 90 Minuten zu erledigen, sodass das Gesamtbudget der Wanduhr 90 × 3 = 270 Minuten beträgt.

Kostenzuordnung

Fügen Sie eine Arbeitsauslastung über eine vorhandene Budgetrichtlinie an usage_policy_name. Der Name wird beim Starten der Workload in die ID der Richtlinie aufgelöst. Informationen zum Einrichten finden Sie unter "Attributverwendung mit serverlosen Verwendungsrichtlinien".

experiment_name: my-training
environment:
  dependencies:
    - mlflow
compute:
  num_accelerators: 1
  accelerator_type: GPU_1xA10
command: echo "Hello World"
usage_policy_name: my team policy

Referenz

Kernfeldreferenz

Feld Typ Description Example
experiment_name string Experimentname für MLflow. "my-training-job"
mlflow_artifact_location string Root-Standort für MLflow-Artefakte, die vom Lauf erfasst werden. Dies ist optional. /Volumes/main/default/mlflow-artifacts/my-training
environment.dependencies Liste Inlineliste der Pip-Abhängigkeitsspezifikationen. ["torch", "transformers"]
environment.version string Serverlose GPU-Umgebungsversion. Dies ist optional. Verwendet die Standardumgebung, wenn sie weggelassen wird. Siehe Umgebungsversion. "4", "5""databricks_ai_v5"
compute.num_accelerators INT Anzahl der GPUs. Es muss ein Vielfaches der GPUs pro Knoten für die ausgewählte compute.accelerator_typegeben. 1, 48
compute.accelerator_type string Beschleunigerkonfiguration, einschließlich GPU-Typ und Knotenform. Siehe Unterstützte GPU-Konfigurationen. "GPU_1xA10", "GPU_1xH100""GPU_8xH100"
code_source dict Codequellkonfiguration. Siehe "Arbeiten mit Codequellen".
command string Bash-Befehle zum Starten der Schulung. torchrun --nproc_per_node=8 train.py

Unterstützte GPU-Konfigurationen

accelerator_type GPUs pro Knoten num_accelerators Anforderung Hinweise
GPU_1xA10 1 Jede positive ganze Zahl Single A10, gut für Entwicklung und kleine Workloads.
GPU_1xH100 1 1 Einzelne H100.
GPU_8xH100 8 Ein positives Vielfaches von 8 Vollständiger H100-Knoten, typisch für verteilte Schulungen.

Für Beschleunigerfähigkeiten und empfohlene Anwendungsfälle siehe Hardware-Optionen.

compute.num_accelerators ist die Gesamtzahl der GPUs für die Arbeitslast. Es muss ein Vielfaches der GPUs pro Knoten für die ausgewählte compute.accelerator_typesein.

Optionale Felder

Umgebungskonfiguration

environment:
  version: '4'
  dependencies:
    - torch
    - transformers
env_variables:
  BATCH_SIZE: '32'
secrets:
  HF_TOKEN: 'my_scope/hf_token'

Für Umgebungsversionen, Abhängigkeitsformat und unterstützte Installationsflags siehe Umgebung.

Benutzerdefinierte Docker-Imagekonfiguration

environment:
  docker_image:
    url: myorg/myrepo:mytag

Sich gegenseitig ausschließend mit environment.dependencies und environment.version. Registrieren Sie das Bild vor air register image der Verwendung. Siehe Verwenden von benutzerdefinierten Docker-Images.

Codequellkonfiguration

code_source:
  type: snapshot
  snapshot:
    root_path: /home/username/repo # REQUIRED — local path to repo or directory
    git: # Optional (git repos only) — pin to a branch or commit
      branch: main # Branch name; uses local HEAD unless 'remote' is set
      # commit: abc1234567 # Mutually exclusive with 'branch'
      remote: false # Optional — true to auto-detect remote HEAD, or a remote name string
    include_paths: # Optional — filter included paths
      - src/
      - configs/

Feldeinschränkungen:

  • git.branch und git.commit schließen sich gegenseitig aus: Geben Sie genau einen innerhalb des git: Blocks an.
  • git.remote erfordert git.branch (es hat keine Wirkung mit git.commit).
  • Wenn Sie den git: Block weglassen, wird der Arbeitsbaum als einfaches Tarball verpackt, einschließlich aller nicht ausgelassenen Änderungen.

Benutzerdefinierte Parameter

Wird an die Workload übergeben über HYPERPARAMETERS_PATH:

parameters:
  model:
    name: 'gpt2'
    hidden_size: 768
  training:
    batch_size: 32

MLflow-Ausführungsname

mlflow_run_name: 'experiment-001-baseline'

MLflow-Artefaktlokalisierung

Stellen Sie mlflow_artifact_location die Artefakte für ein MLflow-Experiment an einem benutzerdefinierten Root-Standort ein. Wenn Sie dieses Feld weglassen, verwendet ein neues Experiment den Standard-DBFS-Standort, wie zum Beispiel dbfs:/databricks/mlflow-tracking/<experiment-id>/....

mlflow_artifact_location: /Volumes/main/default/mlflow-artifacts/my-training

Wenn der DBFS-Zugriff eingeschränkt ist oder Sie Unity Catalog bevorzugen, geben Sie entweder einen /Volumes/<catalog>/<schema>/<volume>/... Pfad oder die entsprechende dbfs:/Volumes/<catalog>/<schema>/<volume>/... URI an. Die air CLI wandelt einen /Volumes Pfad in die dbfs: URI, die MLflow verwendet.

Verwenden Sie einen für jedes Experiment einzigartigen Ort. Der Artefaktstandort eines MLflow-Experiments ist bei der Erstellung des Experiments festgelegt. Wenn ein experiment_name bestehendes Experiment identifiziert wird, mlflow_artifact_location muss es mit seinem Artefaktstandort übereinstimmen oder weggelassen werden. Um einen anderen Ort zu verwenden, geben Sie einen neuen Experimentnamen an.

Pfadauflösung

Alle Pfade in der Workload YAML sind relativ zur Workload YAML, es sei denn, sie sind absolute Pfade.

Ordnerstruktur:

/home/username/my-project/
├── train.yaml
└── scripts/
    └── train.py

YAML-Konfiguration:

experiment_name: my-training
environment:
  dependencies:
    - torch
    - transformers
compute:
  num_accelerators: 8
  accelerator_type: GPU_8xH100
code_source:
  type: snapshot
  snapshot:
    root_path: . # Relative to train.yaml
    git:
      branch: main
command: torchrun --nproc_per_node=8 $CODE_SOURCE_PATH/scripts/train.py
  • Last updated on