Referencia de YAML de carga de trabajo

Important

Esta característica está en versión preliminar pública.

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.

Note

La verdad básica para la configuración de YAML es la ayuda en la CLI. Ejecute air -h config para la vista de nivel superior y air -h config.<section> (por ejemplo, air -h config.environment) para obtener detalles por sección.

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 añade a un experimento MLflow.
  2. environment(Opcional): dependencias de Python y versión base del entorno.
  3. compute (Obligatorio): Recursos de la GPU (tipo y número).
  4. command (Obligatorio): El comando o comandos de golpe usados para iniciar entrenamiento.
  5. code_source (Opcional): Ruta hacia tu código de entrenamiento, disponible de forma remota.

Para valores soportados y restricciones de campo, véase Referencia.

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).
  • environment utiliza el entorno predeterminado e instala 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.

Environment

Usa el environment bloque para seleccionar un entorno de GPU Serverless e instalar dependencias de Python. Por ejemplo, la siguiente configuración selecciona la versión 4 del entorno estándar e instala PyTorch y Transformers:

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

Versión del entorno

environment.version es opcional y selecciona la versión del entorno gestionado para la carga de trabajo.

Algunos ejemplos son:

  • "4" o "5" para usar la versión correspondiente del entorno estándar.
  • "databricks_ai_v5" para utilizar el entorno de IA de Databricks versión 5, que incluye paquetes específicos de ML preinstalados. (Lista completa de paquetes)

El siguiente ejemplo selecciona la versión 5 del entorno IA de Databricks:

environment:
  version: 'databricks_ai_v5'
  dependencies: []

Si especificas environment.version, también debes proporcionar environment.dependencies como una lista en línea. Usa una lista vacía si no necesitas instalar paquetes adicionales.

Para información sobre entornos disponibles para IA Runtime, consulta Configurar tu entorno.

dependencias de Python

Lista las dependencias de Python de tu carga de trabajo como una lista en línea bajo environment.dependencies.

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

Banderas de instalación compatibles

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

Note

No se admite --trusted-host. Dado que uv configura la confianza por dirección URL de índice, use --index-url o --extra-index-url en su lugar.

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

Referencia de campo principal

Campo Tipo Description Example
experiment_name string Nombre del experimento para MLflow. "my-training-job"
mlflow_artifact_location string Ubicación raíz para artefactos de MLflow registrados por la ejecución. Optional. /Volumes/main/default/mlflow-artifacts/my-training
environment.dependencies list Lista insertada de especificaciones de dependencia de pip. ["torch", "transformers"]
environment.version string Versión del entorno de GPU sin servidor. Optional. Utiliza el entorno predeterminado si se omite. Ver versión de Entorno. "4", , "5", "databricks_ai_v5"
compute.num_accelerators int Número de GPU. Debe ser un múltiplo de las GPUs por nodo para el seleccionado compute.accelerator_type. 1, , 4, 8
compute.accelerator_type string Configuración del acelerador, incluyendo el tipo de GPU y la forma del nodo. Consulta Configuraciones de GPU soportadas. "GPU_1xA10", , "GPU_1xH100", "GPU_8xH100"
code_source dict Configuración de código fuente. Consulte Trabajar con orígenes de código.
command string Comandos de Bash para iniciar el entrenamiento. torchrun --nproc_per_node=8 train.py

Configuraciones de GPU compatibles

accelerator_type GPU por nodo num_accelerators requisito Notas
GPU_1xA10 1 Cualquier entero positivo Un solo A10, adecuado para el desarrollo y las cargas de trabajo pequeñas.
GPU_1xH100 1 1 H100 único.
GPU_8xH100 8 Un múltiplo positivo de 8 Nodo H100 completo, típico para el entrenamiento distribuido.

Para capacidades de aceleradores y casos de uso recomendados, consulte Opciones de hardware.

compute.num_accelerators es el número total de GPUs para la carga de trabajo. Debe ser un múltiplo de las GPUs por nodo para el seleccionado compute.accelerator_type.

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 versiones del entorno, formato de dependencia y flags de instalación soportados, véase Entorno.

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'

Ubicación de artefactos MLflow

Configurado mlflow_artifact_location para almacenar artefactos para un experimento MLflow en una ubicación raíz personalizada. Si omites este campo, un nuevo experimento utiliza la ubicación predeterminada de DBFS, como dbfs:/databricks/mlflow-tracking/<experiment-id>/....

mlflow_artifact_location: /Volumes/main/default/mlflow-artifacts/my-training

Si el acceso a DBFS está restringido o prefieres Unity Catalog, especifica una /Volumes/<catalog>/<schema>/<volume>/... ruta o el URI equivalente dbfs:/Volumes/<catalog>/<schema>/<volume>/... . La air CLI convierte un /Volumes camino al dbfs: URI que utiliza MLflow.

Utiliza una ubicación única para cada experimento. La ubicación del artefacto de un experimento MLflow se fija cuando se crea el experimento. Si experiment_name identifica un experimento existente, mlflow_artifact_location debe coincidir con la ubicación del artefacto o ser omitido. Para usar una ubicación diferente, especifica un nuevo nombre de experimento.

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
  • Last updated on