Référence YAML de la charge de travail

Définissez le nom de l’expérience, le calcul, la commande, l’environnement et la source de code d’un travail d’apprentissage dans la configuration YAML de la charge de travail que vous transmettez.air run --file Cette page documente chaque champ.

Configuration minimale

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

Envoyer avec :

air run --file train.yaml -p profile

Concepts de base

Champs principaux

La plupart des configurations d’apprentissage incluent cinq composants :

1. experiment_name (Obligatoire) : Crée ou ajoute à une expérience MLflow.
2. environment(Optionnel) : dépendances Python et version de l’environnement de base.
3. compute (Obligatoire) : ressources GPU (type et nombre).
4. command (Obligatoire) : La ou les commandes de découp utilisées pour lancer l’entraînement.
5. code_source (Optionnel) : Chemin vers votre code d’entraînement, disponible à distance.

Pour les valeurs prises en charge et les contraintes de champ, voir Référence.

Votre premier travail de formation

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

Dans cette configuration :

• experiment_name crée une expérience MLflow nommée simple-training (ou ajoute une nouvelle exécution s’il existe déjà).
• environment utilise l’environnement par défaut et installe torch et transformers.
• compute alloue un nœud H100 (8 GPU H100).
• code_source charge le dossier repo sur le nœud, disponible à l’adresse $CODE_SOURCE_PATH.
• command s’exécute train.py via torchrun les 8 GPU H100. Le fichier se trouve /home/username/repo/train.py localement.

Cas d’utilisation courants

Ajouter des variables d’environnement

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

Utiliser des secrets (clés API, jetons)

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

Les secrets utilisent le format scope/key et doivent être configurés dans Databricks Secrets. Consultez la gestion des secrets pour la configuration.

Lors du partage d’un modèle YAML, d’autres utilisateurs doivent créer leurs propres secrets ou avoir accès au secret référencé.

Environnement

Utilisez ce environment bloc pour sélectionner un environnement GPU serverless et installer des dépendances Python. Par exemple, la configuration suivante sélectionne la version 4 de l’environnement standard et installe PyTorch et Transformers :

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

Version de l’environnement

environment.version est optionnel et sélectionne la version de l’environnement géré pour la charge de travail.

Voici quelques exemples :

• "4" ou "5" d’utiliser la version correspondante de l’environnement standard.
• "databricks_ai_v5" pour utiliser l’environnement IA de Databricks version 5, qui inclut des paquets préinstallés spécifiques au ML. (Liste complète du paquet)

L’exemple suivant sélectionne l’environnement IA de Databricks version 5 :

environment:
  version: 'databricks_ai_v5'
  dependencies: []

Si vous spécifiez environment.version, vous devez également fournir environment.dependencies une liste en ligne. Utilisez une liste vide si vous n’avez pas besoin d’installer des paquets supplémentaires.

Pour des informations sur les environnements disponibles pour l’exécution IA, voir Configurer votre environnement.

dépendances Python

Listez les dépendances Python de votre charge de travail sous une liste en ligne sous environment.dependencies.

Format de dépendance

La liste de dépendances suit la spécification de l’environnement de base Databricks. Chaque entrée est une spécification de package de style pip (par exemple, my-library==6.1). La liste accepte également les entrées suivantes :

• Fichiers de configuration requise : référence à une utilisation requirements.txtexistante-r, par exemple -r '/Workspace/Shared/requirements.txt'. Les variables d’environnement telles que $HOME celles-ci sont développées.
• Roues : chemin absolu d’un .whl fichier, par exemple /Workspace/Shared/path/to/simplejson-3.19.3-py3-none-any.whl.
• URL d’index : URL d’index, par exemple --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

Indicateurs d’installation pris en charge

Les dépendances sont installées avec uv. Les indicateurs de style pip suivants sont pris en charge en tant qu’entrées de liste :

• Appliqué à l’ensemble de l’installation : --index-url, --extra-index-urlet --find-links (-f) définir ou étendre les index de package.
• Appliqué à la dépendance qui les suit : --no-deps, , --no-build-isolation--no-cache-dir, et --force-reinstall. Placez l’indicateur sur sa propre ligne (ou avant la spécification), suivi de la dépendance à laquelle il s’applique.

Par exemple, pour effectuer une installation flash-attn sur l’élément déjà installé torch (aucune isolation de build) et sans résoudre ses propres dépendances :

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

Images Docker personnalisées

En guise d’alternative, environment.dependenciesvous pouvez spécifier une image conteneur Docker personnalisée à l’aide environment.docker_image.urlde . environment.docker_image.url s’exclue mutuellement des deux environment.dependencies et environment.version vous ne pouvez pas utiliser l’une ou l’autre dans la même charge de travail.

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

Avant d’utiliser une image personnalisée, inscrivez-la avec air register image. Pour plus d’informations, notamment les exigences d’image, les images de base Databricks et les modèles Dockerfile, consultez Utiliser des images Docker personnalisées.

Utiliser des sources de code

Le code_source bloc charge le code local afin que le travail d’entraînement puisse l’exécuter.

• root_path est le répertoire local à instantané. Par défaut, air empaquete l’arborescence de travail as-is (y compris les modifications non validées) en tant que tarball brut.
• Pour instantanéner une version git épinglée à la place, ajoutez un git: bloc avec un branch ou commit. Il doit root_path s’agir d’un référentiel Git et d’activer la capture instantanée prenant en charge la version (mise en cache, git archive).
• Pour les dépôts volumineux, include_paths vous permet d’instantanér un sous-ensemble.

Exemple minimal

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

Sur l’ordinateur distant, le code est placé à /databricks/code_source/<directory_name>, où <directory_name> est le composant final du chemin d’accès .root_path $CODE_SOURCE_PATH est défini sur ce chemin absolu, donc utilisez-le dans votre commande plutôt que de coder en dur l’emplacement.

Référentiels Git : épingler par branche ou commit

Pour les référentiels Git, ajoutez un git: bloc pour épingler la version du code par branche ou par validation SHA. branch et commit s’excluent mutuellement : spécifiez exactement un dans le bloc.

Épingler à une branche (utilise le HEAD local de cette branche) :

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

Épingler à une sha de validation (reproductibilité exacte) :

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

Champs clés :

• root_path (Obligatoire) : chemin local à la racine de votre dépôt Git.
• git.branch (Facultatif) : Nom de la branche. Utilise HEAD local ; aucune extraction à distance. Mutuellement exclusif avec git.commit.
• git.commit (Facultatif) : sha de validation spécifique. Mutuellement exclusif avec git.branch.
• git.remote (Facultatif) : utilisez le head distant de la branche au lieu de celui local. Affectez la valeur à true la détection automatique de la distance ou à un nom distant (par exemple upstream) pour extraire à partir d’un distant spécifique. Valide uniquement avec git.branch.

Si vous omettez le bloc, git: empaquette l’arborescence air de travail en tant que tarball brut, y compris les modifications non validées. Aucun champ supplémentaire n’est requis.

Répertoires non git

Vous pouvez créer des répertoires d’instantanés qui ne sont pas des référentiels Git. Omettez le git: bloc, qui doit root_path être un dépôt Git. Sans cela, il n’y a pas de mise en cache de version ; un tarball frais est chargé pour chaque exécution.

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

Filtrage de dossiers avec include_paths

Pour les monorepos volumineux, les dossiers spécifiques d’instantanés sont uniquement destinés à réduire le temps de chargement et la taille du téléchargement et de l’instantané :

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

Points clés :

• Le champ est facultatif. S’il est omis, l’ensemble du référentiel est inclus par défaut.
• Les chemins d’accès doivent être relatifs à la racine du référentiel (pas de début /).
• .. n’est pas autorisé ; vous ne pouvez pas référencer les répertoires parents.

Fonctions avancées

Hyperparamètres personnalisés

Passez une configuration structurée à votre script d’apprentissage via 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

Lisez-les dans votre script :

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']

Fiabilité du travail

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

Si la charge de travail échoue, elle est retentée deux fois. Chaque tentative a 90 minutes pour se terminer, donc le budget total de l’horloge murale est de 90 × 3 = 270 minutes.

Attribution des coûts

Attachez une charge de travail à une stratégie budgétaire existante via usage_policy_name. Le nom est résolu en ID de la stratégie lors du lancement de la charge de travail. Pour la configuration, consultez l’utilisation des attributs avec des stratégies d’utilisation serverless.

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

Référence

Référence du champ de base

Configurations GPU prises en charge

Pour connaître les fonctionnalités de l’accélérateur et les cas d’usage recommandés, consultez les options matérielles.

compute.num_accelerators est le nombre total de GPU pour la charge de travail. Il doit s’agir d’un multiple des GPU par nœud pour le fichier sélectionné compute.accelerator_type.

Champs facultatifs

Configuration de l’environnement

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

Pour les versions de l’environnement, le format de dépendance et les drapeaux d’installation pris en charge, voir Environnement.

Configuration personnalisée de l’image Docker

environment:
  docker_image:
    url: myorg/myrepo:mytag

Mutuellement exclusif avec environment.dependencies et environment.version. Inscrivez l’image avant air register image l’utilisation. Consultez Utiliser des images Docker personnalisées.

Configuration de la source du code

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/

Contraintes de champ :

• git.branch et git.commit s’excluent mutuellement : spécifiez exactement un dans le git: bloc.
• git.remote nécessite git.branch (il n’a aucun effet avec git.commit).
• Si vous omettez le git: bloc, l’arborescence de travail est empaquetée en tant que tarball simple, y compris les modifications non validées.

Paramètres personnalisés

Passé à la charge de travail via HYPERPARAMETERS_PATH:

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

Nom de l’exécution MLflow

mlflow_run_name: 'experiment-001-baseline'

Résolution du chemin d’accès

Tous les chemins d’accès de la charge de travail YAML sont relatifs à la charge de travail YAML, sauf s’ils sont des chemins absolus.

Structure de dossiers :

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

Configuration YAML :

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