Référence YAML de la charge de travail

Important

Cette fonctionnalité est disponible en préversion publique.

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.

Note

La vérité de base pour la configuration YAML est l’aide dans l’interface CLI. Exécutez air -h config la vue de niveau supérieur et air -h config.<section> (par exemple) air -h config.environmentpour obtenir des détails par section.

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: facultatif. Python dépendances et environnement de base.
  3. compute:Obligatoire. Ressources GPU (type et nombre).
  4. command:Obligatoire. Commandes ou commandes bash utilisées pour lancer l’entraînement.
  5. code_source: facultatif. Chemin d’accès à votre code d’entraînement, mis à disposition à distance.

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à).
  • environmentinstalle les dépendances Python répertoriées (ici et torchtransformers).
  • 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é.

dépendances Python

Répertoriez les dépendances Python de votre charge de travail en tant que liste inline sous environment.dependencies:

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

environment.version sélectionne la version de l’environnement GPU serverless. Il est facultatif et est défini par défaut sur "4". Pour obtenir la liste complète des versions d’environnement disponibles, consultez les versions d’environnement serverless.

Versions 5 et databricks_ai_v5 sont également disponibles. La version 5 est l’environnement Standard minimal, qui inclut uniquement l’API GPU serverless, les dépendances Databricks et MLflow. La version databricks_ai_v5 est l’environnement Databricks AI, qui inclut tous les packages de l’environnement Standard, ainsi que pyTorch et des bibliothèques d’apprentissage automatique complètes. Pour obtenir la liste complète des packages, consultez l’environnement GPU serverless version 5.

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

Note

La fonction --trusted-host n'est pas prise en charge. Étant donné que uv configure l’approbation par URL d’index, utilisez --index-url ou --extra-index-url utilisez plutôt.

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

Champs principaux

Champ Catégorie Description Example
experiment_name string Nom de l’expérience pour MLflow. "my-training-job"
environment.dependencies Liste Liste inline des spécifications de dépendance pip. ["torch", "transformers"]
environment.version string Version de l’environnement GPU serverless. facultatif. La valeur par défaut est "4". "4"
compute.num_accelerators int Nombre de GPU. 1, , 48
compute.accelerator_type string Type GPU. "GPU_1xA10", "GPU_8xH100"
code_source Dictionnaire Configuration de la source du code. Consultez Utiliser des sources de code.
command string Commandes Bash pour lancer l’entraînement. torchrun --nproc_per_node=8 train.py

Types GPU pris en charge

accelerator_type GPU par nœud Remarques
GPU_1xA10 1 Un seul A10, adapté au développement et aux petites charges de travail.
GPU_1xH100 1 H100 unique.
GPU_8xH100 8 Nœud H100 complet, typique de l’entraînement distribué.

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

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 le format de dépendance, les indicateurs d’installation pris en charge et environment.version, consultez Python dépendances.

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