YAML-referens för arbetsbelastning

Viktigt!

Den här funktionen finns som allmänt tillgänglig förhandsversion.

Definiera ett träningsjobbs experimentnamn, beräkning, kommando, miljö och kodkälla i arbetsbelastningens YAML-konfiguration som du skickar till air run --file. Den här sidan dokumenterar varje fält.

Note

Grundsanningen för YAML-konfigurationen är hjälpen i CLI. Kör air -h config för vyn på den översta nivån och air -h config.<section> (till exempel air -h config.environment) för information per avsnitt.

Minimal konfiguration

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

Skicka med:

air run --file train.yaml -p profile

Grundläggande begrepp

Kärnfält

De flesta träningskonfigurationer omfattar fem komponenter:

  1. experiment_name:Krävs. Skapar eller lägger till i ett MLflow-experiment.
  2. environment: Valfritt. Python beroenden och basmiljö.
  3. compute:Krävs. GPU-resurser (typ och antal).
  4. command:Krävs. Bash-kommandot eller kommandona som används för att starta träningen.
  5. code_source: Valfritt. Sökväg till träningskoden som görs tillgänglig via fjärranslutning.

Ditt första träningsjobb

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

I den här konfigurationen:

  • experiment_name skapar ett MLflow-experiment med namnet simple-training (eller lägger till en ny körning om den redan finns).
  • environmentinstallerar de listade Python beroenden (här torch och transformers).
  • compute allokerar en H100-nod (8 H100 GPU:er).
  • code_source laddar upp mappen repo till noden, tillgänglig på $CODE_SOURCE_PATH.
  • command körs train.py via torchrun 8 H100 GPU:er. Filen finns lokalt /home/username/repo/train.py .

Vanliga användningsfall

Lägga till miljövariabler

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

Använda hemligheter (API-nycklar, token)

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

Hemligheter använder formatet scope/key och måste konfigureras i Databricks-hemligheter. Se Hemlig hantering för konfiguration.

När du delar en YAML-mall måste andra användare skapa sina egna hemligheter eller ha åtkomst till den refererade hemligheten.

Python beroenden

Lista arbetsbelastningens Python beroenden som en infogad lista under environment.dependencies:

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

environment.version väljer den serverlösa GPU-miljöversionen. Det är valfritt och standardvärdet är "4". För en fullständig lista över tillgängliga miljöversioner, se Miljöversioner.

Versioner 5 och databricks_ai_v5 är också tillgängliga. Versionen 5 är den minimala standardmiljön, som endast innehåller det serverlösa GPU-API:et, Databricks-beroendena och MLflow. Versionen databricks_ai_v5 är Databricks AI-miljön, som innehåller alla paket från Standard-miljön, plus PyTorch och omfattande maskininlärningsbibliotek. Den fullständiga paketlistan finns i Serverlös GPU-miljö version 5.

Beroendeformat

Beroendelistan följer specifikationen för Databricks-basmiljön. Varje post är en paketspecifikation i pip-stil (till exempel my-library==6.1). Listan accepterar också följande poster:

  • Kravfiler: en referens till en befintlig requirements.txt med , -rtill exempel -r '/Workspace/Shared/requirements.txt'. Miljövariabler som $HOME expanderas.
  • Hjul: en absolut sökväg till en .whl fil, till exempel /Workspace/Shared/path/to/simplejson-3.19.3-py3-none-any.whl.
  • Index-URL:er: en index-URL, till exempel --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

Installationsflaggor som stöds

Beroenden installeras med uv. Följande pip-formatflaggor stöds som listposter:

  • Tillämpas på hela installationen: --index-url, --extra-index-urloch --find-links (-f) anger eller utökar paketindexen.
  • Tillämpas på beroendet som följer dem: --no-deps, --no-build-isolation, --no-cache-diroch --force-reinstall. Placera flaggan på en egen rad (eller före specifikationen), följt av det beroende som den gäller för.

Om du till exempel vill installera flash-attn mot den redan installerade torch (ingen byggisolering) och utan att lösa sina egna beroenden:

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

Note

--trusted-host stöds inte. Eftersom uv konfigurerar förtroende per index-URL använder --index-url du eller --extra-index-url i stället.

Anpassade Docker-avbildningar

Som ett alternativ till environment.dependencieskan du ange en anpassad Docker-containeravbildning med hjälp av environment.docker_image.url. environment.docker_image.url är ömsesidigt uteslutande med båda environment.dependencies och environment.version – du kan inte använda någon av dem i samma arbetsbelastning.

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

Innan du använder en anpassad avbildning registrerar du den med air register image. Fullständig information, inklusive bildkrav, Databricks-basavbildningar och Dockerfile-mönster, finns i Använda anpassade Docker-avbildningar.

Arbeta med kodkällor

Blocket code_source laddar upp lokal kod så att träningsjobbet kan köra den.

  • root_path är den lokala katalogen som ska ögonblicksbildas. Som standard air paketeras arbetsträdet as-is (inklusive eventuella ogenomförda ändringar) som en vanlig tarball.
  • Om du vill ögonblicksbilda en fäst git-version i stället lägger du till ett git: block med en branch eller commit. Detta måste root_path vara en git-lagringsplats och möjliggör versionsmedveten ögonblicksbildering (cachelagring, git archive).
  • För stora lagringsplatser include_paths kan du ögonblicksbildera en delmängd.

Minimalt exempel

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

På fjärrdatorn placeras koden på /databricks/code_source/<directory_name>, där <directory_name> är den sista sökvägskomponenten för root_path. $CODE_SOURCE_PATH är inställd på den absoluta sökvägen, så använd den i kommandot i stället för att hårdkoda platsen.

Git-lagringsplatser: fäst efter gren eller incheckning

För git-lagringsplatser lägger du till ett git: block för att fästa kodversionen efter gren eller genom att checka in SHA. branch och commit är ömsesidigt uteslutande: ange exakt en i blocket.

Fäst på en gren (använder det lokala HUVUDET för den grenen):

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

Fäst på en inchecknings-SHA (exakt reproducerbarhet):

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

Nyckelfält:

  • root_path (Krävs): Lokal sökväg till roten för git-lagringsplatsen.
  • git.branch (Valfritt): Grennamn. Använder lokalt HEAD; ingen fjärrhämtning. Ömsesidigt uteslutande med git.commit.
  • git.commit (Valfritt): Specifik inchecknings-SHA. Ömsesidigt uteslutande med git.branch.
  • git.remote (Valfritt): Använd grenens fjärr-HEAD i stället för det lokala. Ange till true för att automatiskt identifiera fjärren eller till ett fjärrnamn (till exempel upstream) för att hämta från en specifik fjärranslutning. Endast giltigt med git.branch.

Om du utelämnar git: blocket air paketerade arbetsträdet som en vanlig tarball, inklusive eventuella icke-bakåtkompatibla ändringar. Inget extra fält krävs.

Icke-git-kataloger

Du kan skapa ögonblicksbilder av kataloger som inte är git-lagringsplatser. Utelämna blocket git: , som måste root_path vara en git-lagringsplats. Utan den finns det ingen cachelagring av versioner. en ny tarball laddas upp för varje körning.

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

Mappfiltrering med include_paths

För stora monorepos, ögonblicksbilder endast specifika mappar för att minska uppladdning och nedladdningstid och storlek på ögonblicksbilder:

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

Viktiga punkter:

  • Fältet är valfritt. Om den utelämnas inkluderas hela lagringsplatsen som standard.
  • Sökvägarna måste vara relativa till lagringsplatsens rot (inga inledande /).
  • .. tillåts inte. du kan inte referera till överordnade kataloger.

Avancerade funktioner

Anpassade hyperparametrar

Skicka strukturerad konfiguration till ditt träningsskript 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

Läs dem i skriptet:

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

Tillförlitlighet för jobb

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

Om arbetsbelastningen misslyckas görs ett nytt försök två gånger. Varje försök har 90 minuter att slutföra, så den totala budgeten för väggklockan är 90 × 3 = 270 minuter.

Kostnadstillskrivning

Koppla en arbetsbelastning till en befintlig budgetprincip via usage_policy_name. Namnet matchas till principens ID när arbetsbelastningen startas. Mer information finns i Attributanvändning med serverlösa användningsprinciper.

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

Reference

Kärnfält

Fält Type Description Example
experiment_name string Experimentnamn för MLflow. "my-training-job"
environment.dependencies list Infogad lista över pip-beroendespecifikationer. ["torch", "transformers"]
environment.version string Serverlös GPU-miljöversion. Optional. Standardinställningen är "4". "4"
compute.num_accelerators int Antal GPU:er. 1, , 48
compute.accelerator_type string GPU-typ. "GPU_1xA10", "GPU_8xH100"
code_source Dict Kodkällans konfiguration. Se Arbeta med kodkällor.
command string Bash-kommandon för att starta träningen. torchrun --nproc_per_node=8 train.py

GPU-typer som stöds

accelerator_type GPUs per nod Notes
GPU_1xA10 1 Enkel A10, bra för utveckling och små arbetsbelastningar.
GPU_1xH100 1 Enkel H100.
GPU_8xH100 8 Fullständig H100-nod, typisk för distribuerad träning.

Information om acceleratorfunktioner och rekommenderade användningsfall finns i Maskinvarualternativ.

Valfria fält

Miljökonfiguration

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

För beroendeformatet, installationsflaggor som stöds och environment.version, se Python beroenden.

Konfiguration av anpassad Docker-avbildning

environment:
  docker_image:
    url: myorg/myrepo:mytag

Ömsesidigt uteslutande med environment.dependencies och environment.version. Registrera avbildningen med air register image före användning. Se Använda anpassade Docker-avbildningar.

Konfiguration av kodkälla

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/

Fältbegränsningar:

  • git.branch och git.commit är ömsesidigt uteslutande: ange exakt en i git: blocket.
  • git.remote kräver git.branch (det har ingen effekt med git.commit).
  • Om du utelämnar git: blocket paketeras arbetsträdet som en vanlig tarball, inklusive eventuella icke-bakåtkompatibla ändringar.

Anpassade parametrar

Skickas till arbetsbelastningen via HYPERPARAMETERS_PATH:

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

MLflow-körningsnamn

mlflow_run_name: 'experiment-001-baseline'

Sökvägsmatchning

Alla sökvägar i arbetsbelastningen YAML är relativa till arbetsbelastningen YAML om de inte är absoluta sökvägar.

Mappstrukturen:

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

YAML-konfiguration:

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