Aller au contenu principal
Dernière mise à jour le

Référence YAML du Workload

info

Aperçu

Cette fonctionnalité est en aperçu public.

Définissez le nom de l'Experimentation d'un Job de formation, le compute, la commande, l'environnement et le code source dans la configuration YAML de la charge de travail que vous transmettez à air run --file. Cette page documente chaque champ.

remarque

La vérité terrain pour la configuration YAML est l'aide in-CLI. Exécutez air -h config pour la vue de niveau supérieur et air -h config.<section> (par exemple, air -h config.environment) pour les détails 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 clés​

Champs principaux​

La plupart des configurations de formation comprennent cinq composants :

  1. experiment_name: Obligatoire. Crée ou ajoute à une expérimentation MLflow.
  2. environmentFacultatif. Dépendances Python et 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_sourceFacultatif. Chemin d'accès à votre code d'entraînement, mis à disposition à distance.

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érience MLflow nommée simple-training (ou ajoute une nouvelle exécution si elle existe déjà).
  • environment installe les dépendances Python répertoriées (ici, torch et transformers).
  • compute alloue un nœud H100 (8 GPU H100).
  • code_source télécharge le dossier repo vers le nœud, disponible à l'adresse $CODE_SOURCE_PATH.
  • command 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​

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 les 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 les secrets Databricks. Voir Gestion des secrets pour la configuration.

Lors du partage d'un Template YAML, les 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 sous forme de liste en ligne sous environment.dependencies:

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

environment.version sélectionne la version de l'environnement GPU serverless. C'est facultatif et sa valeur par default est "4". Pour une liste complète des versions d'environnement disponibles, consultez Versions de l'environnement serverless.

Les 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 comprend tous les packages de l'environnement Standard, plus PyTorch et des bibliothèques complètes de Machine Learning Library (MLlib). Pour la liste complète des packages, consultez la version 5 de l’environnement GPU Serverless.

Format de dépendance​

La liste des 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 prérequis** : une référence à un existant requirements.txt à -r l'aide de, par -r '/Workspace/Shared/requirements.txt' exemple. Les variables d'environnement telles que $HOME sont développées.
  • Wheels : un chemin d'accès 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 --index-url https://pypi.org/simple exemple.
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'entrées 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 la spécification), suivi de la dépendance à laquelle il s'applique.

Par exemple, pour installer flash-attn avec le 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​

En alternative à environment.dependencies, vous pouvez spécifier une image conteneur Docker personnalisée en utilisant environment.docker_image.url. environment.docker_image.url est mutuellement exclusif avec environment.dependencies et environment.version — vous ne pouvez utiliser ni l'un ni 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 avec air register image. Pour plus de détails, notamment sur les exigences d'image, les images de base Databricks et les modèles Dockerfile, consultez Utiliser des images Docker personnalisées.

Travailler avec les codes sources​

Le bloc code_source peut upload le code local afin que le Job de training puisse l'exécuter.

  • root_path est le répertoire local à capturer. Par default, air met en package l'arborescence de travail telle quelle (y compris les modifications non validées) sous forme de simple archive tar.
  • Pour prendre un instantané d'une version git pin à la place, ajoutez un bloc git: avec un branch ou commit. Cela nécessite que root_path soit un git repository et permet la création d'instantanés sensibles à la version (mise en cache, git archive).
  • Pour les grands repository, include_paths vous permet d'instantanéiser 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 en dur l'emplacement.

Git repository: pin par Branch ou commit​

Pour les repository Git, ajoutez un bloc git: pour pin la version du code par branch ou par SHA de commit. branch et commit sont mutuellement exclusives : spécifiez exactement l'une d'entre elles dans le bloc.

Pin to a Branch (uses the local HEAD of that 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 à un 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 clés :

  • root_path (Obligatoire) : chemin local vers la racine de votre repository Git.
  • git.branch (Facultatif) : nom de la Branch. Utilise HEAD local ; aucune récupération à distance. Mutuellement exclusif avec git.commit.
  • git.commit (Facultatif) : hachage de commit spécifique. Exclusif mutuellement avec git.branch.
  • git.remote (Facultatif) : utilisez le HEAD distant de la Branch au lieu du HEAD local. Définissez sur true pour détecter automatiquement la télécommande, ou sur un nom de télécommande (par exemple, upstream) pour récupérer à partir d'une télécommande spécifique. Valide uniquement avec git.branch.

Si vous omettez le bloc git:, air transforme l'arborescence de travail en fichier 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 instantanés de répertoires 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 de version ; une nouvelle archive tarball 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 dossiers avec include_paths​

Pour les grands monoréférentiels, n'instantanez que des dossiers spécifiques afin de réduire le temps d'upload et de download ainsi que la taille de l'instantané :

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 est omis, l'intégralité du repository est incluse par default.
  • Les chemins doivent être relatifs à la racine du repository (pas de / de début).
  • .. n'est pas autorisé ; vous ne pouvez pas référencer de répertoires parents.

Fonctionnalités avancées​

Hyperparamètres personnalisés​

Transmettez la configuration structurée à votre script de formation 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 la charge de travail échoue, elle est relancée deux fois. Chaque tentative dispose de 90 minutes pour être complétée, ainsi le budget total en temps réel est de 90 × 3 = 270 minutes.

Attribution des coûts​

Attachez une charge de travail à une politique budgétaire existante via usage_policy_name. Le nom est résolu à l'ID de la politique lorsque la charge de travail est lancée. Pour la configuration, consultez l'utilisation des attributs 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​

Champs principaux​

Champ

Type

Description

Exemple

experiment_name

chaîne

Nom de l'expérience MLflow.

"my-training-job"

environment.dependencies

Liste

Liste en ligne des spécifications de dépendance pip.

["torch", "transformers"]

environment.version

chaîne

Version de l'environnement GPU Serverless. Facultatif. La valeur par default est "4".

"4"

compute.num_accelerators

int

Nombre de GPU.

1, 4, 8

compute.accelerator_type

chaîne

Type de GPU.

"GPU_1xA10", "GPU_8xH100"

code_source

dict

Configuration du code source.

Consultez Travailler avec les sources de code.

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'expérience MLflow.

"my-training-job"

environment.dependencies

Liste

Liste en ligne des spécifications de dépendance pip.

["torch", "transformers"]

environment.version

chaîne

Version de l'environnement GPU Serverless. Facultatif. La valeur par default est "4".

"4"

compute.num_accelerators

int

Nombre de GPU.

1, 4, 8

compute.accelerator_type

chaîne

Type de GPU.

"GPU_1xA10", "GPU_8xH100"

code_source

dict

Configuration du code source.

Consultez Travailler avec les sources de code.

command

chaîne

Commandes Bash pour lancer l'entraînement.

torchrun --nproc_per_node=8 train.py

Types de GPU pris en charge​

accelerator_type

GPU par nœud

Notes

GPU_1xA10

1

Un seul A10, idéal pour le développement et les petites charges de travail.

GPU_1xH100

1

H100 unique.

GPU_8xH100

8

Nœud H100 complet, typique pour la formation distribuée.

accelerator_type

GPU par nœud

Notes

GPU_1xA10

1

Un seul A10, idéal pour le développement et les petites charges de travail.

GPU_1xH100

1

H100 unique.

GPU_8xH100

8

Nœud H100 complet, typique pour la formation distribuée.

Pour les capacités d'accélérateur et les cas d'utilisation recommandés, consultez Options matérielles.

Champs facultatifs​

Configuration de l’environnement

YAML
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 les dépendances Python.

Configuration d'image Docker personnalisée

YAML
environment:
  docker_image:
    url: myorg/myrepo:mytag

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

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 s’excluent mutuellement : spécifiez-en exactement un dans le bloc git:.
  • git.remote nécessite git.branch (cela n'a aucun effet avec git.commit).
  • Si vous omettez le bloc git:, l'arborescence de travail est empaquetée en tant que tarball simple, y compris toute modification non validée.

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

Résolution de chemin​

Tous les chemins dans le YAML de la charge de travail sont relatifs au YAML de la charge de travail, à moins qu'il ne s'agisse de chemins absolus.

Structure de 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
Sur cette page