Referência YAML de carga de trabalho para a CLI Python legada

Important

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

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

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

Defina o nome do experimento, o recurso de computação, o comando, o ambiente e o código-fonte de um trabalho de treinamento na configuração YAML da carga de trabalho que você passa para air run --file. Esta página documenta todos os campos.

Note

A verdade básica para a configuração do YAML é a ajuda da CLI. Execute air -h config para a exibição de nível superior e air -h config.<section> (por exemplo, air -h config.environment) para obter detalhes por seção.

Configuração mínima

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

Envie com:

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

Para valores suportados e restrições de campo, veja Referência.

Seu primeiro trabalho de treinamento

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 acrescenta uma nova execução se ela já existir).
  • environment usa o ambiente padrão e instala torch e transformers.
  • compute aloca um nó H100 (8 GPUs H100).
  • code_source carrega a pasta repo para o nó, que fica disponível em $CODE_SOURCE_PATH.
  • command é executado train.py via torchrun nas 8 GPUs H100. O arquivo reside localmente /home/username/repo/train.py .

Casos de uso comuns

Adicionar variáveis de ambiente

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)

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 em Segredos do Databricks. Consulte o gerenciamento de segredo para configurar.

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

Ambiente

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

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.

Os exemplos incluem:

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

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

environment:
  version: 'databricks_ai_v5'
  dependencies: []

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

Para informações sobre ambientes disponíveis para o tempo de execução de IA, veja Configurar seu ambiente.

dependências de Python

Liste as dependências de Python da sua carga de trabalho como uma lista inline em 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 existente requirements.txt usando -r, por exemplo -r '/Workspace/Shared/requirements.txt'. Variáveis de ambiente, como, por exemplo $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: uma URL de índice, por exemplo --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

Flags de instalação suportados

As dependências são instaladas com uv. Os seguintes sinalizadores no estilo pip são aceitos como itens 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 os segue: --no-deps, --no-build-isolation, --no-cache-dir e --force-reinstall. Coloque o sinalizador em sua própria linha (ou antes da especificação), seguido pela dependência à qual se aplica.

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

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

Note

Não há suporte para --trusted-host. Como o uv configura a confiança por URL de índice, use --index-url ou --extra-index-url em vez disso.

Imagens personalizadas do Docker

Como alternativa a environment.dependencies, você pode especificar uma imagem de contêiner Docker personalizada usando environment.docker_image.url. environment.docker_image.url é mutuamente exclusivo em relação a environment.dependencies e environment.version — não é possível usar nenhum deles na mesma carga de trabalho.

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 detalhes completos, incluindo requisitos de imagem, imagens base Databricks e padrões Dockerfile, veja Usar imagens Docker personalizadas com a CLI legada de Python.

Trabalhar com fontes de código

O code_source bloco carrega o código local para que o trabalho de treinamento possa executá-lo.

  • root_path é o diretório local do qual será feito um snapshot. Por padrão, air empacota a árvore de trabalho como está (incluindo quaisquer alterações não commitadas) em um arquivo tar simples.
  • Em vez disso, para criar um instantâneo de uma versão do git fixa, adicione um bloco git: com um branch ou commit. Isso exige que root_path seja um repositório Git e habilita a criação de snapshots sensíveis à versão (cache, git archive).
  • Para repositórios grandes, include_paths permite criar um instantâneo de um subconjunto.

Exemplo mínimo

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

No computador remoto, o código é colocado em /databricks/code_source/<directory_name>, onde <directory_name> está 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.

Repositórios Git: fixar por ramificação ou confirmação

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

Fixar em uma ramificação (usa o HEAD local dessa ramificação):

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

Fixar em um SHA de confirmação (reprodutibilidade exata):

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

Campos-chave:

  • root_path (Obrigatório): caminho local para a raiz do seu repositório Git.
  • git.branch (Opcional): nome da ramificação. Usa o HEAD local; sem busca remota. Mutuamente exclusivo com git.commit.
  • git.commit (Opcional): SHA específico do commit. Mutuamente exclusivo com git.branch.
  • git.remote (Opcional): Use o HEAD remoto da ramificação em vez do HEAD local. Defina como true para detectar automaticamente o repositório remoto ou como um nome de repositório remoto (por exemplo, upstream) para buscar de um repositório remoto específico. Somente válido com git.branch.

Se você omitir o bloco git:, air empacota a árvore de trabalho como um arquivo tar simples, incluindo quaisquer alterações ainda não commitadas. Nenhum campo extra é necessário.

Diretórios sem Git

Você pode criar snapshots de diretórios que não são repositórios Git. Omita o bloco git:, que requer que root_path seja um repositório Git. Sem ele, não há cache de versão; um tarball novo é carregado para cada execução.

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, faça snapshot apenas de pastas específicas para reduzir o tempo de upload e download e o tamanho do snapshot:

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 principais:

  • O campo é opcional. Quando omitido, todo o repositório será incluído por padrão.
  • Os caminhos devem ser relativos à raiz do repositório (sem / no início).
  • .. não é permitido; você não pode fazer referência a diretórios pai.

Recursos avançados

Hiperparâmetros personalizados

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

Leia-os em seu script:

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 trabalho

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, ela será repetida duas vezes. Cada tentativa tem 90 minutos para ser concluída, portanto, o orçamento total do relógio de parede é de 90 × 3 = 270 minutos.

Atribuição de custo

Anexe uma carga de trabalho a uma política de orçamento existente por meio de usage_policy_name. O nome é resolvido para o ID da política quando a carga de trabalho é iniciada. Para configurar, consulte o uso de atributos com políticas de uso sem servidor.

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

Referência do campo principal

Campo Tipo Description Exemplo
experiment_name cadeia Nome do experimento para MLflow. "my-training-job"
mlflow_artifact_location cadeia Localização raiz para os artefatos do MLflow registrados durante a execução. Optional. /Volumes/main/default/mlflow-artifacts/my-training
environment.dependencies list Lista em linha de especificações de dependências do pip. ["torch", "transformers"]
environment.version cadeia Versão do ambiente de GPU sem servidor. Optional. Usa o ambiente padrão se for omitido. Veja a versão 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 cadeia Configuração do acelerador, incluindo o tipo de GPU e o formato do nó. Veja 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 cadeia 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 Exigência Notes
GPU_1xA10 1 Qualquer inteiro positivo Uma A10, ideal para desenvolvimento e pequenas cargas de trabalho.
GPU_1xH100 1 1 Um único H100.
GPU_8xH100 8 Múltiplo positivo de 8 Nó H100 completo, típico para treinamento distribuído.

Para capacidades de aceleradores e casos de uso recomendados, veja 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_typeselecionado.

Campos opcionais

Configuração do ambiente

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

Para versões do ambiente, formato de dependência e flags de instalação suportados, veja Ambiente.

Configuração de imagem personalizada do Docker

environment:
  docker_image:
    url: myorg/myrepo:mytag

Exclusivo mutuamente com environment.dependencies e environment.version. Registre a imagem com air register image antes do uso. Veja Usar imagens Docker personalizadas na CLI legada do Python.

Configuração da fonte de código

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 git: bloco.
  • git.remote requer git.branch (não tem nenhum efeito com git.commit).
  • Se você omitir o bloco git:, a árvore de trabalho será empacotada em um arquivo tar simples, incluindo quaisquer alterações não commitadas.

Parâmetros personalizados

Fornecido à carga de trabalho por meio de HYPERPARAMETERS_PATH:

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

Nome da execução do MLflow

mlflow_run_name: 'experiment-001-baseline'

Localização do artefato do MLflow

Defina mlflow_artifact_location para armazenar artefatos de um experimento do MLflow em um local raiz personalizado. Se você omitir esse campo, um novo experimento usa a localização padrão do DBFS, como dbfs:/databricks/mlflow-tracking/<experiment-id>/....

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

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

Use um local único 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 deve corresponder à localização do artefato ou ser omitido. Para usar um local diferente, especifique um novo nome de experimento.

Resolução de caminho

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

Estrutura de pastas:

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

Configuração do 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
  • Última atualização em