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

Referência de Carga de Trabalho YAML

info

Visualização

Este recurso está em Pré-visualização Pública.

Defina o nome do experimento, o compute, o comando, o ambiente e a fonte do código de um Job de treinamento na configuração YAML da 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 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:
    - 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 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 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 valores compatíveis e restrições de campo, consulte Reference.

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 nomeado 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 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:
    - 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 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.

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 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 correspondente do ambiente Standard.
  • "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 você não precisar instalar pacotes adicionais.

Para obter informações sobre os ambientes disponíveis para o AI Runtime, consulte Configure 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 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: uma URL de índice, por --index-url https://pypi.org/simple exemplo.
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 compatíveis

Dependências são instaladas com uv. As seguintes opções no estilo pip são suportadas como entradas de lista:

  • Aplicado a toda a instalação : --index-url, --extra-index-url e --find-links (-f) definem ou estendem os índices do pacote.
  • Aplicado à dependência que os segue : --no-deps, --no-build-isolation, --no-cache-dir e --force-reinstall. A flag deve ser posicionada 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 torch já instalado (sem isolamento de compilação) 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 uv configura a confiança por URL de índice, utilize --index-url ou --extra-index-url em vez disso.

Imagens Docker personalizadas​

Como alternativa a environment.dependencies, é possível especificar uma imagem de contêiner Docker personalizada usando environment.docker_image.url. environment.docker_image.url exclui mutuamente environment.dependencies e environment.version — não é possível usar nenhum deles 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, faça o registro com air register image. Para obter detalhes completos, incluindo requisitos de imagem, imagens base do Databricks e padrões de Dockerfile, consulte Docker Image personalizadas.

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:
    - 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 final do caminho de root_path. $CODE_SOURCE_PATH é definido para esse caminho absoluto, então use-o em seu comando em vez de codificar diretamente 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 branch 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 da branch. Usa HEAD local; sem busca remota. Mutuamente exclusivo com git.commit.
  • git.commit (Opcional): SHA do 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 detectar automaticamente o remoto ou como um nome remoto (por exemplo, upstream) para buscar de um remoto específico. Válido somente com git.branch.

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

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

Você pode tirar Snapshot de diretórios que não são repositórios Git. Omita o bloco git:, que exige que root_path seja um repositório Git. Sem ele, não há cache de versão; um novo tarball é feito 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, 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:
    - 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

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:
    - 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á reexecutada duas vezes. Cada tentativa tem 90 minutos para ser concluída, então 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_name. O nome é resolvido para o ID da política quando a carga de trabalho é iniciada. Para 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 de campo principal​

campo

Tipo

Descrição

Exemplo

experiment_name

string

Nome da experiência do MLflow.

"my-training-job"

mlflow_artifact_location

string

Local raiz para artefatos do 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 de GPU serverless. Opcional. Utiliza 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 o formato do nó. Consulte Configurações de GPU suportadas.

"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 do MLflow.

"my-training-job"

mlflow_artifact_location

string

Local raiz para artefatos do 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 de GPU serverless. Opcional. Utiliza 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 o formato do nó. Consulte Configurações de GPU suportadas.

"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

A10 único, bom para desenvolvimento e cargas de trabalho pequenas.

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

A10 único, bom para desenvolvimento e cargas de trabalho pequenas.

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 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 do 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 sinalizadores de instalação suportados, consulte Ambiente.

Configuração de Docker Image personalizada

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. Veja Use Docker Image personalizadas.

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'

Localização do artefato do MLflow

Defina mlflow_artifact_location para armazenar artefatos de uma execução do MLflow em um local raiz personalizado. Se você omitir este campo, uma nova execução usará o local default do DBFS, 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 o Unity Catalog, especifique um caminho /Volumes/<catalog>/<schema>/<volume>/... ou o URI dbfs:/Volumes/<catalog>/<schema>/<volume>/... equivalente. A CLI air converte um caminho /Volumes para o URI dbfs: que o MLflow usa.

Use um local exclusivo para cada experimento. A localização do artefato de um experimento MLflow é fixada 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 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
└── 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
Nesta página