Referencia YAML de carga de trabajo para la antigua CLI de Python

En el archivo YAML de configuración de la carga de trabajo que pasa a air run --file, defina el nombre del experimento, los recursos de proceso, el comando, el entorno y el origen del código de un trabajo de entrenamiento. 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 junto 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 los comandos de Bash utilizados para iniciar el 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 ejecuta train.py a través de torchrun en las 8 GPU 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.

Medio ambiente

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.
• URL de índice: una 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. Se admiten los siguientes indicadores del estilo de pip como elementos de lista:

• Aplicado a toda la instalación: --index-url, --extra-index-url, y --find-links (-f) establecen o amplían los índices de paquetes.
• 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 usando el torch ya instalado (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 environment.dependencies y environment.version: no se puede usar ninguno de los dos en una 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 detalles completos, incluyendo requisitos de imagen, imágenes base de Databricks y patrones Dockerfile, véase Usar imágenes Docker personalizadas con la CLI heredada de Python.

Trabajar con fuentes 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 del que se tomará una instantánea. De forma predeterminada, air empaqueta el árbol de trabajo tal cual (incluidos los cambios sin confirmar) como un archivo tar simple.
• Para crear en su lugar una instantánea de una versión de git fijada, agregue un bloque git: con un branch o un commit. Esto requiere que root_path sea un repositorio de git y habilita la creación de instantáneas con reconocimiento de 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 último componente de la ruta 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. Utiliza el HEAD local; sin recuperación remota. Mutuamente excluyente con git.commit.
• git.commit (Opcional): confirmación específica SHA. Mutuamente excluyente con git.branch.
• git.remote (Opcional): utiliza el HEAD remoto de la rama en lugar del local. Establézcalo en true para detectar automáticamente el repositorio remoto, o en un nombre de remoto (por ejemplo, upstream) para recuperar desde un remoto específico. Solo es válido con git.branch.

Si omite el bloque git:, air empaqueta el árbol de trabajo como un archivo tar simple, incluidos los cambios sin confirmar. 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, cree instantáneas solo de 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 a su script de entrenamiento mediante 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 dispone de 90 minutos para completarlo, por lo que el tiempo total disponible es de 90 × 3 = 270 minutos.

Atribución de costos

Adjuntar una carga de trabajo a una política presupuestaria existente mediante usage_policy_name. El nombre se resuelve como el ID 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

Reference

Referencia de campo principal

Configuraciones de GPU compatibles

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 GPU por nodo del compute.accelerator_type seleccionado.

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. Ver Usar imágenes Docker personalizadas con la CLI de Python heredada.

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 bloque git:, el árbol de trabajo se empaqueta como un archivo tar simple, incluidos los cambios sin confirmar.

Parámetros personalizados

Se pasa a la carga de trabajo mediante HYPERPARAMETERS_PATH:

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 los artefactos de MLflow

Configure mlflow_artifact_location para almacenar artefactos de un experimento de 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 interfaz de línea de comandos air convierte una ruta /Volumes en el URI dbfs: 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

Todas las rutas del archivo YAML de la carga de trabajo son relativas al archivo YAML de la carga de trabajo, a menos que sean rutas 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