Pular para o conteúdo principal
Última atualização em

Referência de Carga de Trabalho YAML

info

Beta

O AI Runtime CLI está em Beta.

Esta página é a referência para as configurações YAML de carga de trabalho passadas para air run --file.

nota

A verdade fundamental para a configuração YAML é a ajuda da 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: requirements.yaml
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 essenciais​

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

  1. experiment_name: Obrigatório. Cria ou anexa a um experimento do MLflow.
  2. environment: Opcional. Dependências do Python e ambiente base
  3. compute: Obrigatório. Recursos da GPU (tipo e quantidade).
  4. command: Obrigatório. O comando ou comandos bash usados para iniciar o treinamento.
  5. code_source: Opcional. Caminho do código de treinamento, disponibilizado remotamente.

Seu primeiro Job de treinamento​

YAML
experiment_name: simple-training
environment:
  dependencies: requirements.yaml
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 nomeado simple-training (ou anexa uma nova execução se ele já existir).
  • environment instala dependências de requirements.yaml.
  • compute aloca um nó H100 (8 GPUs H100).
  • code_source Faz 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á localizado localmente em /home/username/repo/train.py.

Casos de uso comuns​

Adicionar variáveis de ambiente​

YAML
experiment_name: training-with-env
environment:
  dependencies: requirements.yaml
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: requirements.yaml
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 no Databricks Secrets. Consulte Gerenciamento de segredos para configuração.

Ao compartilhamento um padrão YAML, outros usuários devem criar seus próprios segredos ou ter acesso ao segredo referenciado.

Trabalhar com fontes de código​

O bloco code_source faz 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 como está (incluindo quaisquer alterações sem commit) como um tarball simples.
  • Para tirar um Snapshot de uma versão git fixada em vez disso, adicione um bloco git: com um branch ou commit. Isso exige que root_path seja um repositório Git e permite a criação de snapshots com reconhecimento de versão (cache, git archive).
  • Para repositórios grandes, include_paths permite que você faça um Snapshot de um subconjunto.

Exemplo mínimo​

YAML
experiment_name: simple-training
environment:
  dependencies: requirements.yaml
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 final do caminho de root_path. $CODE_SOURCE_PATH está definido para esse caminho absoluto — use-o em seu comando em vez de codificar o local.

Git repositórios: pin por branch ou commit​

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

Pin a 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 a um SHA de commit (reprodutibilidade exata):

YAML
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 do branch. Usa HEAD local; sem busca remota. Mutuamente exclusivo com git.commit.
  • git.commit (Opcional) — SHA de commit específico. Mutuamente exclusivo com git.branch.
  • git.remote (Opcional) — Usar o HEAD remoto do branch em vez do local. Defina como true para detectar automaticamente o remoto, ou para 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:, air empacota a árvore de trabalho como um tarball simples, incluindo quaisquer alterações sem commit, sem que nenhum campo extra seja necessário.

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

É possível fazer um snapshot de diretórios que não sejam repositórios Git. Omitir o bloco git: — ele requer que root_path seja um repositório Git. Sem ele, não há cache de versão; um novo tarball é carregado a 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, crie 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, o repositório inteiro é incluído por default.
  • Os caminhos devem ser relativos à raiz do repositório (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 script de treinamento por meio de HYPERPARAMETERS_PATH:

YAML
experiment_name: parameterized-training
environment:
  dependencies: requirements.yaml
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

Leiam no 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: requirements.yaml
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á reexecutada duas vezes. Cada tentativa tem 90 minutos para ser concluída; o orçamento total de tempo real é de 90 x 3 = 270 minutos.

Atribuição de custos​

Vincule uma carga de trabalho a uma política de orçamento existente via usage_policy_id. Para a configuração, consulte Uso de Atributos com Políticas de Uso serverless.

YAML
experiment_name: my-training
environment:
  dependencies: requirements.yaml
compute:
  num_accelerators: 1
  accelerator_type: GPU_1xA10
command: echo "Hello World"
usage_policy_id: abcd123-25b8-3e87-9a2c-f86eb19d101c

Referência​

Campos essenciais​

campo

Tipo

Descrição

Exemplo

experiment_name

string

Nome da experiência do MLflow.

"my-training-job"

environment.dependencies

string

Caminho para requirements.yaml.

"requirements.yaml"

compute.num_accelerators

int

Número de GPUs.

1, 4, 8

compute.accelerator_type

string

Tipo de GPU.

"GPU_1xA10", "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

Tipos de GPU compatíveis​

accelerator_type

GPUs por nó

Notas

GPU_1xA10

1

Single A10 — bom para desenvolvimento e cargas de trabalho pequenas.

GPU_1xH100

1

Único H100.

GPU_8xH100

8

Nó H100 completo — típico para treinamento distribuído.

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

Campos opcionais​

Configuração do ambiente

YAML
environment:
  dependencies: requirements.yaml
env_variables:
  BATCH_SIZE: '32'
secrets:
  HF_TOKEN: 'my_scope/hf_token'

Para o formato de arquivo de dependências, consulte requirements.yaml reference.

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

Resolução de Caminho​

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

Estrutura de pastas:

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

Configuração YAML:

YAML
experiment_name: my-training
environment:
  dependencies: requirements.yaml # Relative to train.yaml
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
Nesta página