Referência YAML da carga de trabalho para a CLI antiga do Python

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.

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

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

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).
• environment usa o ambiente predefinido e instala torch e transformers.
• compute atribui um nó H100 (8 GPUs H100).
• code_source carrega a pasta repo para o nó, disponível em $CODE_SOURCE_PATH.
• command executa train.py através de torchrun nas 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 Gestão de segredos para obter instruções de configuração.

Ao partilhar um modelo YAML, outros utilizadores devem criar os 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 ambiental

environment.version é opcional e seleciona a versão do ambiente gerido para a carga de trabalho.

Os exemplos incluem:

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

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

environment:
  version: 'databricks_ai_v5'
  dependencies: []

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

Para informações sobre ambientes disponíveis para IA Runtime, veja Configurar o 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 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.
• wheel: o caminho absoluto para um ficheiro .whl, por exemplo, /Workspace/Shared/path/to/simplejson-3.19.3-py3-none-any.whl
• URLs de índice: um 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. 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, para instalar flash-attn relativamente ao já instalado torch (sem isolamento de compilação) e sem resolver as suas próprias dependências:

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

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 de environment.dependencies e environment.version — não pode usar qualquer um 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, 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 com a CLI legada de Python.

Trabalho com fontes de código

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

• root_path é o diretório local para snapshot. Por defeito, air empacota a árvore de trabalho tal como está (incluindo quaisquer alterações ainda não consolidadas) como um arquivo tar simples.
• Para criar uma instantânea de uma versão do Git fixada, em vez disso, adicione um bloco git: com um branch ou commit. Isto requer que root_path seja um repositório Git e ativa instantâneos com reconhecimento de versões (armazenamento em cache, 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: afixar pela ramificação ou pelo commit

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

Fixar a 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 a um SHA de commit (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. Usa HEAD local; sem obtenção remota. 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 como true para detetar automaticamente o remoto, ou como o nome de um remoto (por exemplo, upstream) para obter de um remoto específico. Válido apenas com git.branch.

Se omitir o bloco git:, air empacota a árvore de trabalho como um arquivo tar simples, incluindo quaisquer alterações por consolidar. Não é necessário campo extra.

Diretórios que não são Git

Podes fazer snapshots 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 existe cache de versões; um novo arquivo tar é enviado em cada execução.

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, o repositório inteiro é incluído por predefinição.
• Os caminhos devem ser relativos à raiz do repositório (sem / no início).
• .. 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 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-as no 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']

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

Anexar uma carga de trabalho a uma política orçamental existente via usage_policy_name. O nome é resolvido para o ID da política quando a carga de trabalho é iniciada. Para configurar, consulte utilização de atributos com políticas de utilização 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

Configurações de GPU suportadas

Para capacidades do 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. Tem de ser um múltiplo do número de GPUs por nó para o compute.accelerator_type selecionado.

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, consulte Ambiente.

Configuração personalizada da imagem Docker

environment:
  docker_image:
    url: myorg/myrepo:mytag

É mutuamente exclusivo com environment.dependencies e environment.version. Registe a imagem com air register image antes de a utilizar. Consulte Utilizar imagens Docker personalizadas com a CLI legada do Python.

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 bloco git:, a árvore de trabalho é empacotada num arquivo tar simples, incluindo quaisquer alterações não confirmadas.

Parâmetros personalizados

Transmitido para a carga de trabalho através de HYPERPARAMETERS_PATH:

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

Nome da execução MLflow

mlflow_run_name: 'experiment-001-baseline'

Localização do artefacto MLflow

Definir mlflow_artifact_location para armazenar artefactos para uma experiência MLflow numa localização raiz personalizada. Se omitir este 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 preferir utilizar o Unity Catalog, especifique um caminho /Volumes/<catalog>/<schema>/<volume>/... ou o URI dbfs:/Volumes/<catalog>/<schema>/<volume>/... equivalente. A air CLI converte um /Volumes caminho para o dbfs: URI que o MLflow utiliza.

Use um local único para cada experiência. A localização do artefacto de um experimento MLflow é fixa quando o experimento é criado. Se experiment_name identificar um experimento existente, mlflow_artifact_location deve corresponder à localização do seu artefacto ou ser omitido. Para usar uma localização diferente, especifique um novo nome de experiência.

Resolução de percurso

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