Referência YAML de Carga de Trabalho

Importante

A CLI de Runtime da IA está em Beta.

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

Note

A verdade fundamental para a configuração do YAML é a ajuda dentro da CLI. Execute air -h config para a vista de topo e air -h config.<section> (por exemplo, air -h config.environment) para o detalhe por secção.

Configuração mínima

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

Submeta com:

air run --file train.yaml -p profile

Conceitos-chave

Campos principais

A maioria das configurações de treino inclui cinco componentes:

  1. experiment_name:Necessário. Cria ou acrescenta a um experimento MLflow.
  2. environment: Opcional. Dependências de Python e ambiente base.
  3. compute:Necessário. Recursos da GPU (tipo e contagem).
  4. command:Necessário. O comando ou comandos bash usados para lançar treino.
  5. code_source: Opcional. Caminho para o teu código de treino, disponibilizado remotamente.

O teu primeiro trabalho de treinador

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 experiência MLflow nomeada simple-training (ou adiciona uma nova execução se 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 para o nó, disponível em $CODE_SOURCE_PATH.
  • command Funciona train.pytorchrun através das 8 GPUs H100. O ficheiro fica localizado localmente /home/username/repo/train.py .

Casos comuns de utilização

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 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 Secrets usam o formato scope/key e devem ser configurados no Databricks Secrets. Consulte a gestão de segredos para configuração.

Ao partilhar um modelo YAML, outros utilizadores devem criar os seus próprios segredos ou ter acesso ao segredo referenciado.

Dependências de Python

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

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

environment.version seleciona a versão do ambiente de GPU serverless. É opcional e por defeito é "4".

Formato de dependência

A lista de dependências segue a Especificação do Ambiente Base Databricks. Cada entrada é uma especificação de pacote ao estilo pip (por exemplo, my-library==6.1). A lista também aceita as seguintes entradas:

  • Ficheiros de requisitos: uma referência a um existente requirements.txt usando -r, por exemplo -r '/Workspace/Shared/requirements.txt'. Variáveis do ambiente, como as $HOME são expandidas.
  • Rodas: um caminho absoluto para um .whl ficheiro, por /Workspace/Shared/path/to/simplejson-3.19.3-py3-none-any.whlexemplo.
  • URLs de índice: uma URL de índice, por --index-url https://pypi.org/simpleexemplo .
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. As seguintes flags em 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 dos pacotes.
  • Aplicado à dependência que os segue:--no-deps , --no-build-isolation, --no-cache-dir, e --force-reinstall. Coloque a bandeira na sua própria linha (ou antes da especificação), seguida da dependência a que se aplica.

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

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

Note

--trusted-host não é suportado. 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 a environment.dependencies, pode especificar uma imagem de contentor Docker personalizada usando environment.docker_image.url. environment.docker_image.url é mutuamente exclusivo com ambos environment.dependencies e environment.version — não podes usar nenhum 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, regista-a com air register image. Para detalhes completos, incluindo requisitos de imagem, imagens base Databricks e padrões Dockerfile, consulte Usar imagens Docker personalizadas.

Trabalho com fontes de código

O code_source bloco carrega código local para que o trabalho de treino o possa executar.

  • root_path é o diretório local para snapshot. Por defeito, air empacota a árvore de trabalho as-is (incluindo quaisquer alterações não comprometidas) como um tarball simples.
  • Para fazer snapshot de uma versão git fixada, adicione um git: bloco com um branch ou commit. Isto requer root_path que seja um repositório git e permite snapshots conscientes de versões (caching, git archive).
  • Para repositórios grandes, include_paths permite fazer snapshots 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

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, por isso usa-o no teu comando em vez de codificar a localização de forma fixa.

repositórios git: pin por branch ou commit

Para repositórios git, adiciona um git: bloco para fixar a versão do código por branch ou por commit SHA. branch e commit são mutuamente exclusivas: especificar exatamente uma dentro do bloco.

Pin a um ramo (usa o HEAD local desse ramo):

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 commit SHA (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 filial. Utiliza HEAD local; Nada de buscar remotamente. Mutuamente exclusivo com git.commit.
  • git.commit (Opcional): Commit específico SHA. Mutuamente exclusivo com git.branch.
  • git.remote (Opcional): Use o HEAD remoto da agência em vez do local. Defina para true detetar automaticamente o comando, ou para um nome remoto (por exemplo, upstream) para buscar a partir de um comando específico. Válido apenas com git.branch.

Se omitir o git: bloco, air empacota a árvore de trabalho como uma bola de alcatrão simples, incluindo quaisquer alterações não comprometidas. Não é necessário campo extra.

Diretórios não-git

Podes fazer snapshots de diretórios que não são repositórios git. Omita o git: bloco, que tem root_path de ser um repositório git. Sem ele, não existe cache de versões; Uma bola de alcatrão nova é carregada para cada corrida.

code_source:
  type: snapshot
  snapshot:
    root_path: /home/username/my_project
command: $CODE_SOURCE_PATH/train.py

Filtragem de pastas com include_paths

Para monorepos grandes, faça snapshots apenas em 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:

  • A área é opcional. Se for omitido, todo o repositório é incluído por defeito.
  • Os caminhos devem ser relativos à raiz do repositório (sem abertura /).
  • .. não é permitido; Não podes consultar diretórios principais.

Funcionalidades avançadas

Hiperparâmetros personalizados

Passe a configuração estruturada para o seu script de treino através HYPERPARAMETERS_PATHde:

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

Lê-as no teu guião:

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

Fiabilidade do emprego

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, é tentada novamente duas vezes. Cada tentativa tem 90 minutos para completar, pelo que o orçamento total do relógio de parede é 90 × 3 = 270 minutos.

Atribuição de custos

Anexe uma carga de trabalho a uma política orçamental existente através usage_policy_namede . O nome é definido para o ID da política quando a carga de trabalho é lançada. Para configuração, veja Utilização de atributos com políticas de uso serverless.

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

Campo Tipo Description Example
experiment_name cadeia (de caracteres) Nome do experimento para MLflow. "my-training-job"
environment.dependencies list Lista inline de especificações de dependência de pips. ["torch", "transformers"]
environment.version cadeia (de caracteres) Versão do ambiente de GPU serverless. Optional. O valor padrão é "4". "4"
compute.num_accelerators int Número de GPUs. 1, 4, 8
compute.accelerator_type cadeia (de caracteres) Tipo de GPU. "GPU_1xA10", "GPU_8xH100"
code_source Dicionário Configuração da fonte do código. Ver Trabalho com fontes de código.
command cadeia (de caracteres) O Bash ordena iniciar o treino. torchrun --nproc_per_node=8 train.py

Tipos de GPU suportados

accelerator_type GPUs por nó Notes
GPU_1xA10 1 Um único A10, bom para desenvolvimento e pequenas cargas de trabalho.
GPU_1xH100 1 H100 simples.
GPU_8xH100 8 Nó H100 completo, típico para treino distribuído.

Para capacidades do acelerador e casos de uso recomendados, consulte 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, flags de instalação suportados, e environment.version, veja dependências em Python.

Configuração personalizada da imagem Docker

environment:
  docker_image:
    url: myorg/myrepo:mytag

Mutuamente exclusivo com environment.dependencies e environment.version. Registe a imagem antes air register image de ser usada. Veja Usar imagens Docker personalizadas.

Configuração da fonte do 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 exclusivas: especificar exatamente uma dentro do git: bloco.
  • git.remote requer git.branch (não tem efeito com git.commit).
  • Se omitir o git: bloco, a árvore de trabalho é empacotada como uma bola de alcatrão simples, incluindo quaisquer alterações não comprometidas.

Parâmetros personalizados

Passado para a carga de trabalho através HYPERPARAMETERS_PATHde:

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

Nome da execução MLflow

mlflow_run_name: 'experiment-001-baseline'

Resolução do 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 da pasta:

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

Configuração 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