Référence YAML du workload pour l’ancienne interface de ligne de commande Python

Important

Cette documentation a été supprimée et peut ne pas être mise à jour.

La interface technique basée sur air Python, installée avec le databricks-air paquet, est désormais obsolète et n’est plus activement maintenue.

Utilisez la ligne de ligne de commande Databricks pour les nouvelles charges de travail. Voir Utiliser l’interface de ligne de commande Databricks avec AI Runtime.

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.

Remarque

La vérité de base pour la configuration YAML est l’aide dans l’interface CLI. Exécutez air -h config pour une vue d’ensemble et air -h config.<section> (par exemple, air -h config.environment) pour 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(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 bash utilisées pour démarrer 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 sur 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é.

Environment

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" pour 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 sous forme de 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 d’exigences : référence à un requirements.txt existant à l’aide de -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 : 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

Remarque

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

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 est mutuellement exclusif avec environment.dependencies et environment.version à la fois — vous ne pouvez utiliser ni l’un ni l’autre au sein de 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 tous les détails, y compris les exigences d’image, les images de base Databricks et les motifs Dockerfile, voir Utiliser des images Docker personnalisées avec l’ancienne CLI Python.

Utiliser des sources de code

Le bloc code_source téléverse le code local afin que la tâche d’entraînement puisse l’exécuter.

  • root_path est le répertoire local dont il faut prendre un instantané. Par défaut, air empaquette l’arborescence de travail telle quelle (y compris les modifications non validées) sous la forme d’une simple archive tar.
  • Pour créer à la place un instantané d’une version Git épinglée, ajoutez un bloc git: avec un branch ou commit. Cela nécessite que root_path soit un dépôt Git et active les instantanés tenant compte des versions (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é dans /databricks/code_source/<directory_name>, où <directory_name> est le dernier élément du chemin 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 dépôts Git, ajoutez un bloc git: pour fixer la version du code en fonction de la branche ou du hachage SHA d’un commit. branch et commit s’excluent mutuellement : spécifiez exactement un dans le bloc.

Épingler sur 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 à un commit SHA (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 commit spécifique. Mutuellement exclusif avec git.branch.
  • git.remote (Facultatif) : utilisez le HEAD distant de la branche au lieu du HEAD local. Définissez la valeur sur true pour détecter automatiquement le dépôt distant, ou sur un nom de dépôt distant (par exemple, upstream) à extraire depuis un dépôt distant spécifique. Valide uniquement avec git.branch.

Si vous omettez le bloc git:, air empaquette l’arborescence de travail sous forme d’archive tar simple, y compris les modifications non validées. Aucun champ supplémentaire n’est requis.

Répertoires non-Git

Vous pouvez créer un instantané de répertoires qui ne sont pas des dépôts Git. Omettez le bloc git:, qui nécessite que root_path soit 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, limitez l’instantané à certains dossiers uniquement afin de réduire le temps de téléversement et de téléchargement ainsi que la taille 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 converti en identifiant de la stratégie au démarrage de la charge de travail. Pour la configuration, consultez l’utilisation des attributs avec des stratégies d’utilisation sans serveur.

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

Champ Catégorie Description Exemple
experiment_name string Nom de l’expérience pour MLflow. "my-training-job"
mlflow_artifact_location string Emplacement racine des artefacts MLflow enregistrés pour l’exécution. Optional. /Volumes/main/default/mlflow-artifacts/my-training
environment.dependencies list Liste incorporée des spécifications de dépendance pip. ["torch", "transformers"]
environment.version string Version de l’environnement GPU sans serveur. Optional. Utilise l’environnement par défaut si c’est omis. Voir la version Environnement. "4", "5", "databricks_ai_v5"
compute.num_accelerators int Nombre de GPU. Doit être un multiple du nombre de GPU par nœud pour le compute.accelerator_type sélectionné. 1, 4, 8
compute.accelerator_type string Configuration de l’accélérateur, y compris le type de GPU et la forme du nœud. Voir Configurations GPU prises en charge. "GPU_1xA10", "GPU_1xH100", "GPU_8xH100"
code_source dict 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

Configurations de GPU prises en charge

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

Pour les capacités d’accélérateur et les cas d’usage recommandés, voir Matériel.

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. Enregistrez l’image avec air register image avant utilisation. Voir Utiliser des images Docker personnalisées avec l’ancienne interface de ligne de commande Python.

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 bloc git:, l’arborescence de travail est empaquetée sous forme d’archive tar simple, y compris toutes les modifications non validées.

Paramètres personnalisés

Transmis à 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'

Emplacement des artefacts MLflow

Configurez mlflow_artifact_location pour stocker des artefacts pour une expérience MLflow dans un emplacement racine personnalisé. Si vous omets ce champ, une nouvelle expérience utilise la position DBFS par défaut, telle que dbfs:/databricks/mlflow-tracking/<experiment-id>/....

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, spécifiez soit un /Volumes/<catalog>/<schema>/<volume>/... chemin, soit l’URI équivalent dbfs:/Volumes/<catalog>/<schema>/<volume>/... . La air CLI convertit un /Volumes chemin vers l’URI dbfs: utilisé par MLflow.

Utilisez un lieu unique à chaque expérience. L’emplacement de l’artefact d’une expérience MLflow est fixé lors de la création de l’expérience. Si experiment_name identifie une expérience existante, mlflow_artifact_location doit correspondre à son emplacement de l’artefact ou être omis. Pour utiliser un autre emplacement, spécifiez un nouveau nom d’expérience.

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