Workload YAML reference pour le CLI Python héritée

important

Cette documentation est obsolète et ne sera peut-être pas mise à jour.

Le CLI air basé sur Python, installé avec le package databricks-air, est désormais obsolète et n’est plus activement maintenu.

Utilisez la CLI de Databricks pour les nouvelles charges de travail. Consultez Utiliser la CLI Databricks avec AI Runtime.

Define a training job's experiment name, compute, command, environment, and code source in the workload YAML config you pass to air run --file. Cette page documente chaque champ.

remarque

La vérité de terrain pour la configuration YAML est l'aide intégrée à la CLI. Exécutez air -h config pour la vue globale et air -h config.<section> (par exemple, air -h config.environment) pour le détail par section.

Configuration minimale

YAML

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

Soumettre avec :

Bash

air run --file train.yaml -p profile

Concepts fondamentaux

Core fields

La plupart des configurations de formation comprennent cinq composants :

1. experiment_name (Obligatoire) : crée ou ajoute à une expérimentation MLflow.
2. environment (Facultatif) : 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 bash utilisées pour lancer l'entraînement.
5. code_source (Facultatif) : chemin d'accès à votre code d'entraînement, rendu disponible à distance.

Pour connaître les valeurs prises en charge et les contraintes de champ, consultez la référence.

Votre premier job d'entraînement

YAML

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érimentation MLflow nommée simple-training (ou ajoute une nouvelle exécution si elle existe déjà).
• environment utilise l'environnement par default et installe torch et transformers.
• compute alloue un nœud H100 (8 GPU H100).
• code_source upload le dossier repo sur le nœud, disponible sur $CODE_SOURCE_PATH.
• command exécute train.py via torchrun sur les 8 GPU H100. Le fichier se trouve localement à l’emplacement /home/username/repo/train.py.

Cas d'utilisation courants

Ajouter des variables d’environnement

YAML

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)

YAML

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 page Gestion des secrets pour la configuration.

Lorsque vous partagez un template YAML, les autres utilisateurs doivent créer leurs propres secrets ou avoir accès au secret référencé.

Environnement

Utilisez le bloc environment pour sélectionner un environnement GPU Serverless et installer les dépendances Python. Par exemple, la configuration suivante sélectionne la version d'environnement Standard 4 et installe PyTorch et Transformers :

YAML

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

Version de l’environnement

environment.version est optionnel et sélectionne la version de l'environnement managé pour la charge de travail.

Quelques exemples :

• "4" ou "5" pour utiliser la version d’environnement Standard correspondante.
• "databricks_ai_v5" d'utiliser l'environnement Databricks AI version 5, qui inclut des packages spécifiques au machine learning préinstallés. (Liste complète des packages)

L'exemple suivant sélectionne la version 5 de l'environnement Databricks AI :

YAML

environment:
  version: 'databricks_ai_v5'
  dependencies: []

Si vous spécifiez environment.version, vous devez également fournir environment.dependencies sous forme de liste in-line. Utilisez une liste vide si vous n’avez pas besoin d’installer de package supplémentaire.

Pour plus d’information sur les environnements disponibles pour AI Runtime, consultez Configurer votre environnement.

Dépendances Python

Listez les dépendances Python de votre workload sous forme de liste en ligne sous environment.dependencies.

Dependency format

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

• Fichiers de requirements : référence à un requirements.txt existant utilisant -r, par exemple -r '/Workspace/Shared/requirements.txt'. Les variables d'environnement telles que $HOME sont développées.
• Wheels : un chemin absolu vers un fichier .whl, par exemple /Workspace/Shared/path/to/simplejson-3.19.3-py3-none-any.whl.
• URL d’index : une URL d’index, par exemple --index-url https://pypi.org/simple.

YAML

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 type pip suivants sont pris en charge en tant qu'éléments de liste :

• Appliqué à l'ensemble de l'installation : --index-url, --extra-index-url et --find-links (-f) définissent ou étendent les index de packages.
• 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 les spécifications), suivi de la dépendance à laquelle il s'applique.

Par exemple, pour installer flash-attn par rapport à torch déjà installé (sans isolation de build) et sans résoudre ses propres dépendances :

YAML

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

remarque

--trusted-host n’est pas pris en charge. Puisque uv configure la confiance par URL d’index, utilisez --index-url ou --extra-index-url à la place.

Images Docker personnalisées

Comme alternative à environment.dependencies, vous pouvez spécifier une image de conteneur Docker personnalisée à l’aide de environment.docker_image.url. environment.docker_image.url s’exclut mutuellement avec environment.dependencies et environment.version : vous ne pouvez pas utiliser l’un ou l’autre dans la même charge de travail.

YAML

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, enregistrez-la auprès de air register image. Pour tous les détails, y compris les exigences relatives aux images, les images de base Databricks et les modèles de Dockerfile, consultez la section Utiliser des images Docker personnalisées avec l’ancienne CLI Python.

Travailler avec des codes sources

Le bloc code_source upload du code local afin que le job d’entraînement puisse l’exécuter.

• root_path est le répertoire local dont il faut faire un snapshot. By default, air packages l'arborescence de travail telle quelle (y compris les modifications non validées) sous forme de simple archive tar.
• Pour faire un snapshot d'une version Git pin à la place, ajoutez un bloc git: avec un branch ou commit. Ceci requiert que root_path soit un repository Git et active la création de snapshots tenant compte des versions (caching, git archive).
• Pour les grands repository, include_paths vous permet de prendre un instantané d’un sous-ensemble.

Exemple minimal

YAML

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 la machine distante, le code est placé à /databricks/code_source/<directory_name>, où <directory_name> est le composant de chemin final de root_path. $CODE_SOURCE_PATH est défini sur ce chemin absolu, utilisez-le donc dans votre commande plutôt que de coder l’emplacement en dur.

Git repositories: pin by branch or commit

Pour les repositories Git, ajoutez un bloc git: pour pin la version du code par Branch ou par SHA de commit. branch et commit s'excluent mutuellement : spécifiez-en exactement un dans le bloc.

Pin à une Branch (utilise le HEAD local de cette Branch) :

YAML

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

Pin to a commit SHA (reproductibilité exacte) :

YAML

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

Champs de clé :

• root_path (Obligatoire) : chemin local vers la racine de votre repository Git.
• git.branch (Optionnel) : Nom de la branch. Utilise le HEAD local ; aucun fetch distant. S'exclut mutuellement avec git.commit.
• git.commit (Facultatif) : commit SHA spécifique. S'exclut mutuellement avec git.branch.
• git.remote (Optionnel) : utilisez le HEAD distant de la branch au lieu du HEAD local. Définissez sur true pour détecter automatiquement la ressource distante, ou sur un nom de ressource distante (par exemple, upstream) pour effectuer la récupération à partir d'une ressource distante spécifique. Uniquement valide avec git.branch.

Si vous omettez le bloc git:, air conditionne l'arborescence de travail sous la forme d'une simple archive tar, incluant toutes les modifications non validées. Aucun champ supplémentaire n'est requis.

Répertoires non Git

Vous pouvez créer des instantanés de repository qui ne sont pas des repository git. Omettez le bloc git:, qui exige que root_path soit un repository git. Sans cela, il n’y a pas de mise en cache des versions ; une archive tar fraîche est upload pour chaque exécution.

YAML

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

Filtrage de dossier avec include_paths

Pour les grands monorepos, ne faites un snapshot que de dossiers spécifiques afin de réduire le temps d'upload et de download ainsi que la taille du snapshot :

YAML

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 n'est pas précisé, l'ensemble du repository est inclus par default.
• Les chemins d'accès doivent être relatifs à la racine du repository (sans / initial).
• .. n'est pas autorisé ; vous ne pouvez pas faire référence à des répertoires parents.

Fonctionnalités avancées

Custom hyperparameters

Transmettez une configuration structurée à votre script d'entraînement via HYPERPARAMETERS_PATH:

YAML

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 :

Python

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é des jobs

YAML

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 le workload échoue, il est relancé deux fois. Chaque tentative dispose de 90 minutes pour s'exécuter, de sorte que le budget total en temps réel est de 90 × 3 = 270 minutes.

Cost attribution

Associer une charge de travail à une politique budgétaire existante via usage_policy_name. Le nom est résolu en identifiant de la politique lors du lancement de la charge de travail. Pour la configuration, voir Attribuer l’utilisation avec les politiques d’utilisation serverless.

YAML

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 des champs de base

Champ	Type	Description	Exemple
experiment_name	chaîne	Nom de l'Experimentation pour MLflow.	"my-training-job"
mlflow_artifact_location	chaîne	Emplacement racine des artefacts MLflow consignés par l'exécution. Facultatif.	/Volumes/main/default/mlflow-artifacts/my-training
environment.dependencies	liste	Liste en ligne des spécifications de dépendance pip.	["torch", "transformers"]
environment.version	chaîne	Version de l’environnement Serverless GPU. Facultatif. Utilise l’environnement default en cas d’omission. Consultez Version de l’environnement.	"4", "5", "databricks_ai_v5"
compute.num_accelerators	int	Nombre de GPU. Doit être un multiple des GPU par nœud pour le compute.accelerator_type sélectionné.	1, 4, 8
compute.accelerator_type	chaîne	Configuration des accélérateurs, y compris le type de GPU et la forme des nœuds. Voir Configurations GPU prises en charge.	"GPU_1xA10", "GPU_1xH100", "GPU_8xH100"
code_source	dict	Configuration du code source.	Consultez la section Travailler avec des codes sources.
command	chaîne	Commandes Bash pour lancer l'entraînement.	torchrun --nproc_per_node=8 train.py

Champ	Type	Description	Exemple
experiment_name	chaîne	Nom de l'Experimentation pour MLflow.	"my-training-job"
mlflow_artifact_location	chaîne	Emplacement racine des artefacts MLflow consignés par l'exécution. Facultatif.	/Volumes/main/default/mlflow-artifacts/my-training
environment.dependencies	liste	Liste en ligne des spécifications de dépendance pip.	["torch", "transformers"]
environment.version	chaîne	Version de l’environnement Serverless GPU. Facultatif. Utilise l’environnement default en cas d’omission. Consultez Version de l’environnement.	"4", "5", "databricks_ai_v5"
compute.num_accelerators	int	Nombre de GPU. Doit être un multiple des GPU par nœud pour le compute.accelerator_type sélectionné.	1, 4, 8
compute.accelerator_type	chaîne	Configuration des accélérateurs, y compris le type de GPU et la forme des nœuds. Voir Configurations GPU prises en charge.	"GPU_1xA10", "GPU_1xH100", "GPU_8xH100"
code_source	dict	Configuration du code source.	Consultez la section Travailler avec des codes sources.
command	chaîne	Commandes Bash pour lancer l'entraînement.	torchrun --nproc_per_node=8 train.py

Configurations GPU prises en charge

accelerator_type	GPU par nœud	num_accelerators exigence	Notes
GPU_1xA10	1	Tout entier positif	A10 unique, idéal pour le développement et les petites charges de travail.
GPU_1xH100	1	1	H100 unique.
GPU_8xH100	8	Un multiple positif de 8	Nœud H100 complet, typique pour l'entraînement distribué.

accelerator_type	GPU par nœud	num_accelerators exigence	Notes
GPU_1xA10	1	Tout entier positif	A10 unique, idéal pour le développement et les petites charges de travail.
GPU_1xH100	1	1	H100 unique.
GPU_8xH100	8	Un multiple positif de 8	Nœud H100 complet, typique pour l'entraînement distribué.

Pour en savoir plus sur les fonctionnalités des accélérateurs et les cas d'usage recommandés, consultez la page Options matérielles.

compute.num_accelerators correspond au nombre total de GPU pour le workload. Il doit s'agir d'un multiple des GPU par nœud pour le compute.accelerator_type sélectionné.

Champs optionnels

Configuration de l’environnement

YAML

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

Pour connaître les versions de l’environnement, le format des dépendances et les indicateurs d’installation pris en charge, consultez la section Environnement.

Configuration de l’image Docker personnalisée

YAML

environment:
  docker_image:
    url: myorg/myrepo:mytag

S’exclut mutuellement avec environment.dependencies et environment.version. Enregistrez l’image avec air register image avant utilisation. Voir Utiliser des images Docker personnalisées avec l’ancienne interface CLI Python.

Configuration du code source

YAML

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 sont mutuellement exclusifs : spécifiez-en exactement un dans le bloc git:.
• git.remote requiert git.branch (sans effet avec git.commit).
• Si vous omettez le bloc git:, l'arborescence de travail est empaquetée sous forme d'archive tar simple, incluant toutes les modifications non validées.

Paramètres personnalisés

Transmis à la charge de travail via HYPERPARAMETERS_PATH:

YAML

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

Nom de l’exécution MLflow

YAML

mlflow_run_name: 'experiment-001-baseline'

MLflow artifact location

Définissez mlflow_artifact_location pour stocker les artefacts d’une expérimentation MLflow dans un emplacement racine personnalisé. Si vous omettez ce champ, une nouvelle expérimentation utilise l’emplacement DBFS default, tel que dbfs:/databricks/mlflow-tracking/<experiment-id>/....

YAML

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

Si l'accès à DBFS est restreint ou si vous préférez Unity Catalog, indiquez un chemin /Volumes/<catalog>/<schema>/<volume>/... ou l'URI dbfs:/Volumes/<catalog>/<schema>/<volume>/... équivalente. La CLI air convertit un chemin /Volumes en l'URI dbfs: utilisée par MLflow.

Utilisez un emplacement unique pour chaque Experimentation. L'emplacement des artefacts d'une expérimentation MLflow est fixe lors de la création de l'expérimentation. Si experiment_name identifie une expérimentation existante, mlflow_artifact_location doit correspondre à son emplacement d'artefact ou être omis. Pour utiliser un autre emplacement, veuillez spécifier un nouveau nom d'Experimentation.

Résolution de chemin d'accès

Tous les chemins du fichier YAML de la charge de travail sont relatifs au fichier YAML de la charge de travail, sauf s'il s'agit de chemins absolus.

Structure des dossiers :

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

Configuration YAML :

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