Referência de YAML de Workload para a CLI legada do Python

importante

Esta documentação foi descontinuada e pode não ser atualizada.

A CLI do air baseada em Python, instalada com o pacote databricks-air, agora está obsoleta e não é mais mantida ativamente.

Use a CLI do Databricks para novas cargas de trabalho. Consulte Usar a CLI do Databricks com o AI Runtime.

Defina o nome do experimento do Job de treinamento, compute, comando, ambiente e código-fonte no arquivo YAML de configuração de carga de trabalho que você passa para air run --file. Esta página documenta cada campo.

nota

A verdade fundamental para a configuração YAML é a ajuda in-CLI. Execução air -h config para a view de nível superior e air -h config.<section> (por exemplo, air -h config.environment) para detalhes por seção.

Minimal configuration

YAML

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

Enviar com:

Bash

air run --file train.yaml -p profile

Conceitos principais

Campos principais

A maioria das configurações de treinamento inclui cinco componentes:

1. experiment_name (Obrigatório): Cria um experimento do MLflow ou adiciona a ele.
2. environment (Opcional): Dependências do Python e versão do ambiente base.
3. compute (Required): recursos de GPU (tipo e contagem).
4. command (Obrigatório): o comando ou os comandos bash usados para iniciar o treinamento.
5. code_source (Opcional): Caminho para o seu código de treinamento, disponibilizado remotamente.

Para ver os valores compatíveis e as restrições de campo, consulte a Referência.

Seu primeiro Job de treinamento

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

Nesta configuração:

• experiment_name cria um experimento do MLflow chamado simple-training (ou anexa uma nova execução se ele já existir).
• environment usa o ambiente default e instala torch e transformers.
• compute aloca um nó H100 (8 GPUs H100).
• code_source faz o upload da pasta repo para o nó, disponível em $CODE_SOURCE_PATH.
• command executa train.py via torchrun nas 8 GPUs H100. O arquivo está em /home/username/repo/train.py localmente.

Casos de uso comuns

Adicionar variáveis de ambiente

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

Usar segredos (chaves de API, tokens)

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

Os segredos usam o formato scope/key e devem ser configurados nos Segredos do Databricks. Consulte Gerenciamento de segredos para obter informações sobre a configuração.

Ao compartilhar um template YAML, outros usuários devem criar seus próprios segredos ou ter acesso ao segredo referenciado.

Ambiente

Use o bloco environment para selecionar um ambiente Serverless GPU e instalar dependências do Python. Por exemplo, a seguinte configuração seleciona a versão 4 do ambiente Standard e instala o PyTorch e o Transformers:

YAML

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

Versão do ambiente

environment.version é opcional e seleciona a versão do ambiente gerenciado para a carga de trabalho.

Alguns exemplos incluem:

• "4" ou "5" para usar a versão de ambiente Standard correspondente.
• "databricks_ai_v5" para usar a versão 5 do ambiente do Databricks AI, que inclui pacotes específicos de ML pré-instalados. (Lista completa de pacotes)

O exemplo a seguir seleciona a versão 5 do ambiente do Databricks AI:

YAML

environment:
  version: 'databricks_ai_v5'
  dependencies: []

Se você especificar environment.version, também deverá fornecer environment.dependencies como uma lista embutida. Use uma lista vazia se não precisar instalar pacotes adicionais.

Para obter informações sobre ambientes disponíveis para o AI Runtime, consulte Configurar seu ambiente.

Dependências do Python

Liste as dependências Python da sua workload como uma lista embutida sob environment.dependencies.

Formato de dependência

A lista de dependências segue a Especificação de Ambiente Base do Databricks. Cada entrada é uma especificação de pacote no estilo pip (por exemplo, my-library==6.1). A lista também aceita as seguintes entradas:

• Arquivos de requisitos : uma referência a um requirements.txt existente usando -r, por exemplo, -r '/Workspace/Shared/requirements.txt'. Variáveis de ambiente, como $HOME, são expandidas.
• Wheels : um caminho absoluto para um arquivo .whl, por exemplo, /Workspace/Shared/path/to/simplejson-3.19.3-py3-none-any.whl.
• URLs de índice : um URL de índice, por exemplo --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

Sinalizadores de instalação suportados

As dependências são instaladas com uv. Os seguintes sinalizadores no estilo pip são suportados como entradas de lista:

• Aplicado a toda a instalação : --index-url, --extra-index-url e --find-links (-f) definem ou estendem os índices de pacotes.
• Aplicado à dependência que as sucede : --no-deps, --no-build-isolation, --no-cache-dir e --force-reinstall. Coloque a flag em sua própria linha (ou antes da especificação), seguida da dependência à qual ela se aplica.

Por exemplo, para instalar flash-attn em relação ao torch já instalado (sem isolamento de build) e sem resolver suas próprias dependências:

YAML

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

nota

--trusted-host não é suportado. Como o uv configura a confiança por URL de índice, use --index-url ou --extra-index-url em vez disso.

Docker Images personalizados

Como alternativa a environment.dependencies, você pode especificar uma imagem de container Docker personalizada usando environment.docker_image.url. environment.docker_image.url é mutuamente exclusivo com environment.dependencies e environment.version — você não pode usar nenhum dos dois na mesma carga de trabalho.

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

Antes de usar uma imagem personalizada, registre-a com air register image. Para obter todos os detalhes, incluindo requisitos de imagem, imagens base do Databricks e padrões de Dockerfile, consulte Usar imagens personalizadas do Docker com a CLI do Python herdada.

Trabalhar com fontes de código

O bloco code_source faz o upload do código local para que o job de treinamento possa executá-lo.

• root_path é o diretório local para o snapshot. Por default, air empacota a árvore de trabalho tal como está (incluindo quaisquer alterações sem commit) como um tarball simples.
• Para fazer o snapshot de uma versão do git pin em vez disso, adicione um bloco git: com um branch ou commit. Isso requer que root_path seja um repository git e habilita o snapshot com reconhecimento de versão (caching, git archive).
• Para repositories grandes, o include_paths permite criar um snapshot de um subconjunto.

Exemplo mínimo

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

On the remote machine, the code is placed at /databricks/code_source/<directory_name>, where <directory_name> is the final path component of root_path. $CODE_SOURCE_PATH is set to that absolute path, so use it in your comando rather than hard-coding the location.

Repositórios do Git: pin por branch ou commit

Para git repository, adicione um bloco git: para pin a versão do código por branch ou por SHA de commit. branch e commit são mutuamente exclusivos: especifique exatamente um dentro do bloco.

Pin em uma Branch (usa o HEAD local dessa 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 em um commit SHA (reprodutibilidade exata):

YAML

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

Campos de key:

• root_path (Obrigatório): Caminho local para a raiz do seu repository git.
• git.branch (Opcional): nome da Branch. Usa o HEAD local; sem busca remota. Mutuamente exclusivo com git.commit.
• git.commit (Opcional): commit SHA específico. Exclusivo com git.branch.
• git.remote (Opcional): use o HEAD remoto do branch em vez do local. Defina como true para detectar automaticamente o remote ou como um nome remoto (por exemplo, upstream) para buscar de um remote específico. Válido apenas com git.branch.

Se você omitir o bloco git:, o air empacotará a árvore de trabalho como um arquivo tar simples, incluindo quaisquer alterações sem commit. Nenhum campo extra é obrigatório.

Diretórios não Git

É possível criar um Snapshot de diretórios que não sejam git repository. Omitir o bloco git:, que exige que root_path seja um repository Git. Sem isso, não há cache de versões; um novo arquivo tarball é enviado (upload) para cada execução.

YAML

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

Filtragem de pastas com include_paths

Para monorepos grandes, crie um Snapshot apenas de pastas específicas para reduzir o tempo de upload e download e o tamanho do 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

Pontos-chave:

• O campo é opcional. Se omitido, todo o repository é incluído por default.
• Os caminhos devem ser relativos à raiz do repository (sem / inicial).
• .. não é permitido; não é possível fazer referência a diretórios pai.

Recursos avançados

Hiperparâmetros personalizados

Passe a configuração estruturada para o seu script de treinamento 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

Leia-os em seu 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']

Confiabilidade do job

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

Se a carga de trabalho falhar, será feita uma nova tentativa duas vezes. Cada tentativa tem 90 minutos para ser concluída; portanto, o orçamento total de tempo corrido é de 90 x 3 = 270 minutos.

Atribuição de custos

Anexe uma carga de trabalho a uma política orçamentária existente via usage_policy_name. O nome é resolvido para o ID da política quando a carga de trabalho é iniciada. Para obter informações sobre a configuração, consulte Uso de atributos com políticas de uso 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

Referência

Referência do campo principal

campo	Tipo	Descrição	Exemplo
experiment_name	string	Nome da experiência para o MLflow.	"my-training-job"
mlflow_artifact_location	string	Root location for MLflow artifacts logged by the execução. Opcional.	/Volumes/main/default/mlflow-artifacts/my-training
environment.dependencies	Lista	Lista embutida de especificações de dependência do pip.	["torch", "transformers"]
environment.version	string	Versão do ambiente Serverless GPU. Opcional. Usa o ambiente default se omitido. Consulte Versão do ambiente.	"4", "5", "databricks_ai_v5"
compute.num_accelerators	int	Número de GPUs. Deve ser um múltiplo das GPUs por nó para o compute.accelerator_type selecionado.	1, 4, 8
compute.accelerator_type	string	Configuração do acelerador, incluindo o tipo de GPU e a forma do nó. Consulte Configurações de GPU suportadas.	"GPU_1xA10", "GPU_1xH100", "GPU_8xH100"
code_source	dict	Configuração da fonte de código.	Consulte Trabalhar com fontes de código.
command	string	Comandos bash para iniciar o treinamento.	torchrun --nproc_per_node=8 train.py

campo	Tipo	Descrição	Exemplo
experiment_name	string	Nome da experiência para o MLflow.	"my-training-job"
mlflow_artifact_location	string	Root location for MLflow artifacts logged by the execução. Opcional.	/Volumes/main/default/mlflow-artifacts/my-training
environment.dependencies	Lista	Lista embutida de especificações de dependência do pip.	["torch", "transformers"]
environment.version	string	Versão do ambiente Serverless GPU. Opcional. Usa o ambiente default se omitido. Consulte Versão do ambiente.	"4", "5", "databricks_ai_v5"
compute.num_accelerators	int	Número de GPUs. Deve ser um múltiplo das GPUs por nó para o compute.accelerator_type selecionado.	1, 4, 8
compute.accelerator_type	string	Configuração do acelerador, incluindo o tipo de GPU e a forma do nó. Consulte Configurações de GPU suportadas.	"GPU_1xA10", "GPU_1xH100", "GPU_8xH100"
code_source	dict	Configuração da fonte de código.	Consulte Trabalhar com fontes de código.
command	string	Comandos bash para iniciar o treinamento.	torchrun --nproc_per_node=8 train.py

Configurações de GPU compatíveis

accelerator_type	GPUs por nó	num_accelerators requisito	Notas
GPU_1xA10	1	Qualquer número inteiro positivo	A10 único, ideal para desenvolvimento e pequenas cargas de trabalho.
GPU_1xH100	1	1	H100 único.
GPU_8xH100	8	Um múltiplo positivo de 8	Nó H100 completo, típico para treinamento distribuído.

accelerator_type	GPUs por nó	num_accelerators requisito	Notas
GPU_1xA10	1	Qualquer número inteiro positivo	A10 único, ideal para desenvolvimento e pequenas cargas de trabalho.
GPU_1xH100	1	1	H100 único.
GPU_8xH100	8	Um múltiplo positivo de 8	Nó H100 completo, típico para treinamento distribuído.

Para obter recursos de acelerador e casos de uso recomendados, consulte Opções de hardware.

compute.num_accelerators é o número total de GPUs para a carga de trabalho. Deve ser um múltiplo das GPUs por nó para o compute.accelerator_type selecionado.

Campos opcionais

Configuração de ambiente

YAML

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

Para ver as versões de ambiente, o formato de dependência e os sinalizadores de instalação compatíveis, consulte Environment.

Configuração de Docker Image personalizada

YAML

environment:
  docker_image:
    url: myorg/myrepo:mytag

Mutualmente exclusivo com environment.dependencies e environment.version. Registre a imagem com air register image antes de usar. Consulte Usar imagens personalizadas do Docker com a CLI do Python antiga.

Configuração da fonte de código

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/

Restrições de campo:

• git.branch e git.commit são mutuamente exclusivos: especifique exatamente um dentro do bloco git:.
• git.remote exige git.branch (não tem efeito com git.commit).
• Se você omitir o bloco git:, a árvore de trabalho será empacotada como um tarball simples, incluindo quaisquer alterações não confirmadas.

Parâmetros personalizados

Passado para a carga de trabalho via HYPERPARAMETERS_PATH:

YAML

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

Nome da execução do MLflow

YAML

mlflow_run_name: 'experiment-001-baseline'

Localização do artefato do MLflow

Defina mlflow_artifact_location para armazenar artefatos de um experimento no MLflow em uma localização raiz personalizada. Se você omitir este campo, um novo experimento usará a localização DBFS default, como dbfs:/databricks/mlflow-tracking/<experiment-id>/....

YAML

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

If DBFS access is restricted or you prefer Unity Catalog, specify either a /Volumes/<catalog>/<schema>/<volume>/... path or the equivalent dbfs:/Volumes/<catalog>/<schema>/<volume>/... URI. The air CLI converts a /Volumes path to the dbfs: URI that MLflow uses.

Use um local exclusivo para cada experimento. A localização do artefato de um experimento do MLflow é fixa quando o experimento é criado. Se experiment_name identificar um experimento existente, mlflow_artifact_location deverá corresponder à localização do artefato dele ou ser omitido. Para usar um local diferente, especifique um novo nome de experimento.

Resolução de caminho

Todos os caminhos no YAML de carga de trabalho são relativos ao YAML de carga de trabalho, a menos que sejam caminhos absolutos.

Estrutura de pastas:

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

Configuração 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