YAML-referens för arbetsbelastningar för den äldre Python-CLI:n

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.

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 ett MLflow-experiment.
2. environment(Valfritt): Python-beroenden och basversionen av miljön.
3. compute (Krävs): GPU-resurser (typ och antal).
4. command (Krävs): Bash-kommandot eller kommandon som används för att starta utbildning.
5. code_source (Valfritt): Väg till din träningskod, tillgänglig på distans.

För stödda värden och fältbegränsningar, se Referens.

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).
• environment använder standardmiljön och installerar torch och transformers.
• compute tilldelar en H100-nod (8 H100-GPU:er).
• code_source laddar upp mappen repo till noden, tillgänglig på $CODE_SOURCE_PATH.
• command kör train.py via torchrun över de 8 H100-GPU:erna. 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.

Miljö

Använd blocket environment för att välja en serverless GPU-miljö och installera Python-beroenden. Till exempel väljer följande konfiguration Standardmiljö version 4 och installerar PyTorch och Transformers:

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

Miljöversion

environment.version är valfritt och väljer versionen för den hanterade miljön för arbetsbelastningen.

Exempel inkluderar:

• "4" eller "5" för att använda motsvarande version av Standard-miljön.
• "databricks_ai_v5" att använda Databricks AI-miljö version 5, som inkluderar förinstallerade ML-specifika paket. (Fullständig paketlista)

Följande exempel väljer Databricks AI-miljö version 5:

environment:
  version: 'databricks_ai_v5'
  dependencies: []

Om du specificerar environment.version, måste du också ange environment.dependencies som en inline-lista. Använd en tom lista om du inte behöver installera fler paket.

För information om miljöer tillgängliga för AI Runtime, se Sätt upp din miljö.

Python-beroenden

Lista Python-beroendena för arbetsbelastningen som en lista direkt i texten under environment.dependencies.

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 hjälp av -r, till exempel -r '/Workspace/Shared/requirements.txt'. Miljövariabler som $HOME expanderas.
• Wheels: 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

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. För fullständiga detaljer, inklusive bildkrav, Databricks-basavbildningar och Dockerfile-mönster, se Använd anpassade Docker-avbildningar med den äldre Python CLI:n.

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 i stället vill skapa en ögonblicksbild av en låst Git-version, lägger du till ett git:-block med branch eller commit. Detta kräver att root_path är ett Git-repository och möjliggör versionsmedvetna ögonblicksbilder (cachelagring, git archive).
• För stora repositorier låter include_paths dig skapa en ögonblicksbild av 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-repositorier: fäst via gren eller commit

För Git-repositorier lägger du till ett block av typen git: för att låsa kodversionen med gren eller commit-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

Lås till en commit-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 commit-SHA. Ömsesidigt uteslutande med git.branch.
• git.remote (Valfritt): Använd grenens fjärr-HEAD i stället för det lokala. Ställ in på true för att automatiskt identifiera fjärren, eller på ett fjärrnamn (till exempel upstream) för att hämta från en specifik fjärr. Endast giltigt med git.branch.

Om du utelämnar git:-blocket paketerar air arbetsträdet som ett vanligt tar-arkiv, inklusive eventuella ej incheckade ändringar. Inget extra fält krävs.

Kataloger som inte är Git-kataloger

Du kan skapa ögonblicksbilder av kataloger som inte är git-lagringsplatser. Utelämna blocket git:, som kräver att root_path är ett git-arkiv. 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, ta endast snapshots av specifika mappar för att minska uppladdnings- och nedladdningstid samt snapshotstorlek:

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 hänvisa 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 löses upp till policyns ID när arbetsbelastningen startar. 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ältreferens

Stödda GPU-konfigurationer

För acceleratorfunktioner och rekommenderade användningsområden, se Hårdvarualternativ.

compute.num_accelerators är det totala antalet GPU:er för arbetsbelastningen. Det måste vara en multipel av antalet GPU:er per nod för den valda compute.accelerator_type.

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 miljöversioner, beroendeformat och stödda installationsflaggor, se Miljö.

Konfiguration av anpassad Docker-avbildning

environment:
  docker_image:
    url: myorg/myrepo:mytag

Kan inte användas tillsammans med environment.dependencies och environment.version. Registrera avbildningen med air register image före användning. Se Använd anpassade Docker-bilder med den äldre Python CLI:n.

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 ett vanligt tar-arkiv, inklusive eventuella ej incheckade ändringar.

Anpassade parametrar

Vidarebefordras till arbetslasten via HYPERPARAMETERS_PATH:

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

MLflow-körningsnamn

mlflow_run_name: 'experiment-001-baseline'

MLflow-artefaktlokalisering

Ställ mlflow_artifact_location in att lagra artefakter för ett MLflow-experiment på en anpassad rotplats. Om du utelämnar detta fält använder ett nytt experiment standardplatsen för DBFS, såsom dbfs:/databricks/mlflow-tracking/<experiment-id>/....

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

Om DBFS-åtkomsten är begränsad eller du föredrar Unity Catalog, ange antingen en /Volumes/<catalog>/<schema>/<volume>/... sökväg eller motsvarande dbfs:/Volumes/<catalog>/<schema>/<volume>/... URI. CLI:n air konverterar en /Volumes sökväg till den dbfs:-URI som MLflow använder.

Använd en plats som är unik för varje experiment. Ett MLflow-experiments artefaktplats är fastställd när experimentet skapas. Om experiment_name identifierar ett befintligt experiment måste mlflow_artifact_location stämma överens med artefaktens plats eller utelämnas. För att använda en annan plats, ange ett nytt experimentnamn.

Sökvägsupplösning

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