Referencia de YAML de carga de trabajo

Defina el nombre del experimento, el proceso, el comando, el entorno y el código fuente de un trabajo de entrenamiento en la configuración de YAML de carga de trabajo que pase a air run --file. Esta página documenta todos los campos.

Configuración mínima

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

Enviar con:

air run --file train.yaml -p profile

Conceptos principales

Campos principales

La mayoría de las configuraciones de entrenamiento incluyen cinco componentes:

1. experiment_name:Obligatorio. Crea o anexa a un experimento de MLflow.
2. environment: opcional. Python dependencias y entorno base.
3. compute:Obligatorio. Recursos de GPU (tipo y recuento).
4. command:Obligatorio. Comando o comandos de Bash usados para iniciar el entrenamiento.
5. code_source: opcional. Ruta de acceso al código de entrenamiento, disponible de forma remota.

Su primer trabajo de entrenamiento

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

En esta configuración:

• experiment_name crea un experimento de MLflow denominado simple-training (o anexa una nueva ejecución si ya existe).
• environmentinstala las dependencias de Python enumeradas (aquí torch y transformers).
• compute asigna un nodo H100 (8 GPU H100).
• code_source carga la carpeta repo en el nodo, disponible en $CODE_SOURCE_PATH.
• command se ejecuta train.py a través torchrun de las GPU 8 H100. El archivo reside en /home/username/repo/train.py localmente.

Casos de uso comunes

Adición de variables de entorno

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

Uso de secretos (claves 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

Los secretos usan el formato scope/key y deben configurarse en Secretos de Databricks. Consulte Administración de secretos para la instalación.

Al compartir una plantilla YAML, otros usuarios deben crear sus propios secretos o tener acceso al secreto al que se hace referencia.

dependencias de Python

Enumere las dependencias de Python de la carga de trabajo como una lista insertada en environment.dependencies:

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

environment.version selecciona la versión del entorno de GPU sin servidor. Es opcional y tiene "4"como valor predeterminado .

Formato de dependencia

La lista de dependencias sigue la especificación del entorno base de Databricks. Cada entrada es una especificación de paquete de estilo pip (por ejemplo, my-library==6.1). La lista también acepta las siguientes entradas:

• Archivos de requisitos: una referencia a un existente requirements.txt mediante -r, por ejemplo -r '/Workspace/Shared/requirements.txt'. Las variables de entorno como $HOME se expanden.
• Ruedas: una ruta de acceso absoluta a un .whl archivo, por ejemplo /Workspace/Shared/path/to/simplejson-3.19.3-py3-none-any.whl.
• Direcciones URL de índice: una dirección URL de índice, por ejemplo --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

Marcas de instalación admitidas

Las dependencias se instalan con uv. Las siguientes marcas de estilo pip se admiten como entradas de lista:

• Se aplica a toda la instalación: --index-url, --extra-index-urly --find-links (-f) establecida o extiende los índices de paquete.
• Se aplica a la dependencia que las sigue: --no-deps, --no-build-isolation, --no-cache-diry --force-reinstall. Coloque la marca en su propia línea (o antes de la especificación), seguida de la dependencia a la que se aplica.

Por ejemplo, para instalar flash-attn en el ya instalado torch (sin aislamiento de compilación) y sin resolver sus propias dependencias:

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

Imágenes personalizadas de Docker

Como alternativa a environment.dependencies, puede especificar una imagen de contenedor de Docker personalizada mediante environment.docker_image.url. environment.docker_image.url es mutuamente excluyente con y environment.dependenciesenvironment.version : no se puede usar en la misma carga de trabajo.

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 una imagen personalizada, regístrela con air register image. Para más información, incluidos los requisitos de imagen, las imágenes base de Databricks y los patrones de Dockerfile, consulte Uso de imágenes personalizadas de Docker.

Trabajar con orígenes de código

El code_source bloque carga código local para que el trabajo de entrenamiento pueda ejecutarlo.

• root_path es el directorio local al que se va a realizar la instantánea. De forma predeterminada, air empaqueta el árbol de trabajo as-is (incluidos los cambios no confirmados) como un tarball sin formato.
• Para crear una instantánea de una versión de Git anclada en su lugar, agregue un git: bloque con o branchcommit. Esto requiere root_path ser un repositorio git y habilitar la creación de instantáneas compatibles con versiones (almacenamiento en caché, git archive).
• En el caso de repositorios de gran tamaño, include_paths permite realizar una instantánea de un subconjunto.

Ejemplo 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

En el equipo remoto, el código se coloca en /databricks/code_source/<directory_name>, donde <directory_name> es el componente de ruta de acceso final de root_path. $CODE_SOURCE_PATH se establece en esa ruta de acceso absoluta, así que úsela en el comando en lugar de codificar de forma rígida la ubicación.

Repositorios de Git: anclar por rama o confirmación

En el caso de los repositorios de Git, agregue un git: bloque para anclar la versión del código por rama o por confirmación SHA. branch y commit son mutuamente excluyentes: especifique exactamente uno dentro del bloque.

Anclar a una rama (usa el HEAD local de esa rama):

code_source:
  type: snapshot
  snapshot:
    root_path: /home/username/repo
    git:
      branch: main # Uses local HEAD of main (no remote fetch)
command: train.sh

Anclar a un SHA de confirmación (reproducibilidad exacta):

code_source:
  type: snapshot
  snapshot:
    root_path: /home/username/repo
    git:
      commit: abc1234567 # Pins specific commit
command: train.sh

Campos clave:

• root_path (Obligatorio): ruta de acceso local a la raíz del repositorio git.
• git.branch (Opcional): Nombre de rama. Usa head local; ninguna captura remota. Mutuamente excluyente con git.commit.
• git.commit (Opcional): confirmación específica SHA. Mutuamente excluyente con git.branch.
• git.remote (Opcional): use head remoto de la rama en lugar de la local. Establézcalo true en para detectar automáticamente el control remoto o en un nombre remoto (por ejemplo, upstream) para capturar desde un remoto específico. Solo es válido con git.branch.

Si omite el git: bloque, air empaqueta el árbol de trabajo como un tarball sin formato, incluidos los cambios no confirmados. No se requiere ningún campo adicional.

Directorios que no son de Git

Puede crear instantáneas de directorios que no sean repositorios de Git. Omita el git: bloque , que requiere root_path ser un repositorio de Git. Sin él, no hay almacenamiento en caché de versiones; se carga un tarball fresco para cada ejecución.

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

Filtrado de carpetas con include_paths

Para monorepos grandes, las instantáneas solo son carpetas específicas para reducir el tiempo de carga y descarga y el tamaño de la instantánea:

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

Puntos clave:

• El campo es opcional. Si se omite, el repositorio completo se incluye de forma predeterminada.
• Las rutas de acceso deben ser relativas a la raíz del repositorio (sin /inicial).
• .. no se permite; no puede hacer referencia a directorios primarios.

Características avanzadas

Hiperparámetros personalizados

Pase la configuración estructurada al script de entrenamiento a travé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

Léelos en el 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']

Confiabilidad del trabajo

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

Si se produce un error en la carga de trabajo, se reintenta dos veces. Cada intento tiene 90 minutos para completarse, por lo que el presupuesto total del reloj es de 90 × 3 = 270 minutos.

Atribución de costos

Adjunte una carga de trabajo a una directiva de presupuesto existente a través de usage_policy_name. El nombre se resuelve en el identificador de la directiva cuando se inicia la carga de trabajo. Para la instalación, consulte Uso de atributos con directivas de uso sin 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

Referencia

Campos principales

Tipos de GPU admitidos

Para conocer las funcionalidades del acelerador y los casos de uso recomendados, consulte Opciones de hardware.

Campos opcionales

Configuración del entorno

environment:
  version: '4'
  dependencies:
    - torch
    - transformers
env_variables:
  BATCH_SIZE: '32'
secrets:
  HF_TOKEN: 'my_scope/hf_token'

Para obtener el formato de dependencia, las marcas de instalación admitidas y environment.version, consulte Python dependencias.

Configuración personalizada de la imagen de Docker

environment:
  docker_image:
    url: myorg/myrepo:mytag

Mutuamente excluyente con environment.dependencies y environment.version. Registre la imagen con air register image antes de usarla. Consulte Uso de imágenes personalizadas de Docker.

Configuración de código fuente

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/

Restricciones de campo:

• git.branch y git.commit son mutuamente excluyentes: especifique exactamente uno dentro del git: bloque.
• git.remote requiere git.branch (no tiene ningún efecto con git.commit).
• Si omite el git: bloque, el árbol de trabajo se empaqueta como un tarball sin formato, incluidos los cambios no confirmados.

Parámetros personalizados

Se pasa a la carga de trabajo a través HYPERPARAMETERS_PATHde :

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

Nombre de ejecución de MLflow

mlflow_run_name: 'experiment-001-baseline'

Resolución de rutas de acceso

Todas las rutas de acceso de la carga de trabajo YAML son relativas a YAML de carga de trabajo, a menos que sean rutas de acceso absolutas.

Estructura de carpetas:

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

Configuración de 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