Referência yaml de carga de trabalho

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

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

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).
• environmentinstala as dependências de Python listadas (aqui torch e transformers).
• compute aloca um nó H100 (8 GPUs H100).
• code_source carrega a pasta repo no nó, disponível em $CODE_SOURCE_PATH.
• command é executado train.py por meio torchrun das 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.

dependências de Python

Liste as dependências de Python da carga de trabalho como uma lista environment.dependenciesembutida em:

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

environment.version seleciona a versão do ambiente de GPU sem servidor. Ele é opcional e usa "4"como padrão .

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.
• Rodas: um caminho absoluto para um .whl arquivo, 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

Sinalizadores de instalação com suporte

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

• Aplicado à instalação inteira: --index-url, --extra-index-urle --find-links (-f) defina ou estenda os índices do pacote.
• Aplicado à dependência que os segue: --no-deps, , --no-build-isolatione --no-cache-dir--force-reinstall. Coloque o sinalizador em sua própria linha (ou antes da especificação), seguido pela dependência à qual ele se aplica.

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

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

Imagens personalizadas do Docker

Como alternativa, environment.dependenciesvocê pode especificar uma imagem de contêiner personalizada do Docker usando environment.docker_image.url. environment.docker_image.url é mutuamente exclusivo com ambos environment.dependencies e environment.version não é possível usar nenhum dos dois 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 obter detalhes completos, incluindo requisitos de imagem, imagens base do Databricks e padrões do Dockerfile, consulte Usar imagens personalizadas do Docker.

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 a ser instantâneo. Por padrão, air empacota a árvore de trabalho as-is (incluindo quaisquer alterações não confirmadas) como uma tarball simples.
• Em vez disso, para instantâneo de uma versão do Git fixada, adicione um git: bloco com um branch ou commit. Isso requer root_path ser um repositório git e habilita o instantâneo com reconhecimento de versão (cache). git archive
• Para repositórios grandes, include_paths permite que você 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 git: bloco para fixar a versão de código por branch ou confirmar SHA. branch e commit são mutuamente exclusivos: especifique exatamente um dentro do bloco.

Fixar em um branch (usa o HEAD local desse branch):

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 repositório git.
• git.branch (Opcional): nome do branch. Usa o HEAD local; sem busca remota. Mutuamente exclusivo com git.commit.
• git.commit (Opcional): SHA de confirmação específica. Mutuamente exclusivo com git.branch.
• git.remote (Opcional): use o HEAD remoto do branch em vez do local. Defina para true 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 git: bloco, air empacota a árvore de trabalho como uma tarball simples, incluindo quaisquer alterações não confirmadas. Nenhum campo extra é necessário.

Diretórios não git

Você pode instantâneo de diretórios que não são repositórios git. Omita o git: bloco, que requer root_path ser 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, instantâneo apenas pastas específicas para reduzir o tempo de carregamento e download e o tamanho do instantâneo:

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. Se omitido, todo o repositório será incluído por padrão.
• Os caminhos devem ser relativos à raiz do repositório (sem leading /).
• .. não é permitido; você não pode referenciar 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 a 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

Campos principais

Tipos de GPU com suporte

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

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 o formato de dependência, sinalizadores de instalação com suporte eenvironment.version, consulte Python dependências.

Configuração de imagem personalizada do Docker

environment:
  docker_image:
    url: myorg/myrepo:mytag

Mutuamente exclusivo com environment.dependencies e environment.version. Registre a imagem com air register image antes do uso. Consulte Usar imagens personalizadas do Docker.

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 git: bloco, a árvore de trabalho será empacotada como uma tarball simples, incluindo quaisquer alterações não confirmadas.

Parâmetros personalizados

Passado para a 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'

Resolução de caminho

Todos os caminhos na carga de trabalho YAML 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 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