Referência de YAML de Workload para a CLI Python antiga

importante

Esta documentação foi desativada e pode não estar atualizada.

A CLI air baseada em Python, instalada com o pacote databricks-air, 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 de um job de treinamento, compute, comando, ambiente e fonte de código na configuração YAML de carga de trabalho que você passa para air run --file. Esta página documenta todos os campos.

nota

A verdade fundamental para a configuração YAML é a ajuda na 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.

Configuração mínima

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 ou anexa a um experimento do MLflow.
2. environment (Opcional): Dependências Python e versão do ambiente base.
3. compute (Obrigatório): Recursos de GPU (tipo e contagem).
4. command (Obrigatório): O comando ou comandos bash usados para iniciar o treinamento.
5. code_source (Opcional): Caminho para seu código de treinamento, disponibilizado remotamente.

Para obter valores e restrições de campo compatíveis, consulte 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 uma experiment MLflow chamada simple-training (ou anexa uma nova execução se ela já existir).
• environment usa o ambiente default e instala torch e transformers.
• compute aloca um nó H100 (8 GPUs H100).
• code_source upload a pasta repo para o nó, disponível em $CODE_SOURCE_PATH.
• command faz a execução de 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 configuração.

Durante o compartilhamento de um padrão 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 configuração a seguir seleciona a versão 4 do ambiente Standard e instala o PyTorch e 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 do ambiente Standard correspondente.
• "databricks_ai_v5" para usar a versão 5 do ambiente 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 Databricks AI:

YAML

environment:
  version: 'databricks_ai_v5'
  dependencies: []

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

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

Dependências do Python

Liste as dependências Python da sua carga de trabalho como uma lista em linha em environment.dependencies.

Formato de dependência

A lista de dependências segue a Especificação do Ambiente Base do Databricks. Cada entrada é uma especificação de pacote do tipo 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'. As 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

Flags de instalação suportadas

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

• Aplicado à instalação completa : --index-url, --extra-index-url e --find-links (-f) definem ou estendem os índices de pacotes.
• Aplicado à dependência que os segue : --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 pela dependência à qual ela se aplica.

Por exemplo, para instalar flash-attn em relação ao já instalado torch (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 tem suporte. Como o uv configura a confiança por URL de índice, use --index-url ou --extra-index-url em vez disso.

Imagens Docker personalizadas

Como alternativa ao 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, registro-a com air register image. Para detalhes completos, incluindo requisitos de imagem, imagens base do Databricks e padrões de Dockerfile, consulte Usar Docker Image personalizadas com a CLI legada do Python.

Trabalhar com fontes de código

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

• root_path é o diretório local para snapshot. Por default, air empacota a árvore de trabalho exatamente como está (incluindo quaisquer alterações sem commit) como um tarball simples.
• Para criar um Snapshot de uma versão git pin em vez disso, adicione um bloco git: com um branch ou commit. Isso requer que root_path seja um git repository e habilita o snapshot com reconhecimento de versão (cache, git archive).
• Para repository grandes, include_paths permite criar 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

Na máquina remota, o código é colocado em /databricks/code_source/<directory_name>, onde <directory_name> é o componente de caminho final de root_path. $CODE_SOURCE_PATH é definido como esse caminho absoluto, portanto, use-o em seu comando em vez de codificar o local.

Git repository: pin por Branch ou commit

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

Pin para 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 para 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 git repository.
• git.branch (Opcional): nome da branch. Usa HEAD local; sem fetch remoto. Mutuamente exclusivo com git.commit.
• git.commit (Opcional): SHA de commit específico. Mutuamente exclusivo com git.branch.
• git.remote (Opcional): Use o HEAD remoto do branch em vez do local. Defina como true para autodetectar o remoto ou como um nome remoto (por exemplo, upstream) para buscar de um remoto específico. Somente válido com git.branch.

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

Diretórios que não são do Git

Você pode tirar Snapshots de diretórios que não são repository git. Omita o bloco git:, que exige que root_path seja um repository git. Sem isso, não há cache de versão; um tarball fresco é 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 pasta com include_paths

Para monorepos grandes, tire snapshots 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 key:

• 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 referenciar 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 o workload falhar, ele será reexecutado duas vezes. Cada tentativa tem 90 minutos para ser concluída, portanto o orçamento total é de 90 × 3 = 270 minutos.

Atribuição de custos

Anexe um workload a uma política de orçamento existente via usage_policy_name. O nome é resolvido para o ID da política quando o workload é iniciado. Para 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 de campo central

campo	Tipo	Descrição	Exemplo
experiment_name	string	Nome da experiência para MLflow.	"my-training-job"
mlflow_artifact_location	string	Local raiz para artefatos MLflow registrados pela execução. Opcional.	/Volumes/main/default/mlflow-artifacts/my-training
environment.dependencies	Lista	Lista em linha 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 compatíveis.	"GPU_1xA10", "GPU_1xH100", "GPU_8xH100"
code_source	dicionário	Configuração do código-fonte.	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 MLflow.	"my-training-job"
mlflow_artifact_location	string	Local raiz para artefatos MLflow registrados pela execução. Opcional.	/Volumes/main/default/mlflow-artifacts/my-training
environment.dependencies	Lista	Lista em linha 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 compatíveis.	"GPU_1xA10", "GPU_1xH100", "GPU_8xH100"
code_source	dicionário	Configuração do código-fonte.	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 suportadas

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

Para capacidades 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 versões de ambiente, formato de dependência e flags de instalação compatíveis, consulte Ambiente.

Configuração personalizada da Docker Image

YAML

environment:
  docker_image:
    url: myorg/myrepo:mytag

Mutuamente exclusivo com environment.dependencies e environment.version. Faça o registro da imagem com air register image antes de usar. Consulte Usar imagens personalizadas do Docker com a CLI antiga do Python.

Configuração do código-fonte

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 no bloco git: .
• git.remote requer 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 MLflow em uma localização raiz personalizada. Se você omitir esse 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

Se o acesso ao DBFS estiver restrito ou se você preferir usar o Unity Catalog, especifique um caminho /Volumes/<catalog>/<schema>/<volume>/... ou o URI dbfs:/Volumes/<catalog>/<schema>/<volume>/... equivalente. A CLI do air converte um caminho /Volumes para o URI dbfs: que o MLflow usa.

Use uma localização única para cada experimento. A localização do artefato de um experimento MLflow é fixa quando o experimento é criado. Se experiment_name identificar um experimento existente, mlflow_artifact_location deverá corresponder à sua localização de artefato 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 pasta:

/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