Usa la CLI de Databricks con AI Runtime

Important

Los comandos de la CLI de Databricks para el tiempo de ejecución de IA están en Vista previa pública. Un administrador de espacio de trabajo debe activar la vista previa de IA en tiempo de ejecución antes de poder enviar cargas de trabajo. Ver Requisitos.

La databricks air CLI ejecuta cargas de trabajo de GPU de un solo nodo y de varios nodos en AI Runtime, la plataforma de computación de GPU serverless bajo demanda. Define la computación, el entorno y el comando de una carga de trabajo en YAML, y luego envíalo con databricks air run.

AI Runtime se encarga del aprovisionamiento de GPU y la configuración del entorno, y proporciona el estado de ejecución, registros, métricas del sistema y seguimiento de MLflow. El flujo de trabajo YAML y CLI funciona de forma natural con agentes de codificación, que pueden definir, enviar y monitorizar cargas de trabajo.

Utiliza la CLI para:

  • Desarrolla e itera aplicaciones en Python desde tu entorno local.
  • Define cargas de trabajo repetibles de un solo nodo o de varios nodos en archivos YAML que puedes controlar con control de versiones.
  • Envía y gestiona cargas de trabajo desde un terminal o un flujo de trabajo automatizado de desarrollo.

Las cargas de trabajo típicas incluyen ajuste fino y postentrenamiento de LLM, entrenamiento distribuido, inferencia por lotes, aplicaciones Ray y búsqueda de hiperparámetros.

Para entregar una carga de trabajo, sigue el quickstart.

Los comandos de la CLI de Databricks para AI Runtime cumplen con los mismos estándares de cumplimiento que las cargas de trabajo de AI Runtime. Para conocer los estándares y excepciones admitidos, consulte perfil de seguridad de cumplimiento.

Note

Los databricks air comandos y la configuración de la carga de trabajo difieren de la antigua CLI basada air en Python del databricks-air paquete. Consulta la migración de la CLI de Python.

Cómo se ejecuta una carga de trabajo

Proporcionas un archivo YAML de carga de trabajo, el código de la aplicación cuando sea necesario y el comando que inicia tu aplicación. Cuando envías una carga de trabajo:

  1. La definición de la carga de trabajo se valida frente a las restricciones de esquema y configuración soportadas.
  2. Si code_source está configurado, el código seleccionado se empaqueta como una instantánea y se sube.
  3. La computación de GPU y el entorno solicitados se preparan en cada nodo.
  4. El comando proporcionado en la carga de trabajo YAML se ejecuta en cada nodo.

El envío devuelve un ID de ejecución del trabajo que identifica la carga de trabajo. Úsalo para comprobar el estado, inspeccionar registros y configuraciones, o cancelar la carga de trabajo. También puedes consultar la página de Jobs run y el experimento MLflow para detalles de ejecución, métricas del sistema, registros y artefactos.

Este flujo de trabajo proporciona provisión gestionada de GPU, configuración de entornos reproducibles y observabilidad integrada sin necesidad de configurar o mantener la infraestructura de GPU.

Configuración de la carga de trabajo

Un archivo YAML de carga de trabajo describe los siguientes componentes:

  • Experimento: El experimento y la ejecución de MLflow que realizan el seguimiento de la carga de trabajo.
  • Entorno: El entorno gestionado, dependencias o imagen Docker personalizada para usar.
  • Cómputo: El tipo de acelerador y el número de aceleradores a provisionar.
  • Fuente del código: Una instantánea opcional del código de la aplicación para poner a disposición de la carga de trabajo.
  • Comando: El comando shell que inicia tu aplicación.
  • Configuración de ejecución: Opciones como intentos, tiempos de espera, variables de entorno, secretos y permisos.

Para campos de configuración y ejemplos, consulta la referencia de YAML de Workload.

Subir código de aplicación

Usa code_source para subir archivos de la aplicación con una carga de trabajo. AI Runtime pone el directorio subido a disposición en cada nodo y establece $CODE_SOURCE_PATH en esa ubicación. Por ejemplo, el siguiente comando ejecuta train.py desde el directorio subido:

command: python "$CODE_SOURCE_PATH/train.py"

Puedes subir el árbol de trabajo actual o seleccionar una rama o commit de Git. Para un ejemplo completo, consulta el inicio rápido.

Desarrollar y gestionar las cargas de trabajo

El desarrollo suele seguir este proceso:

  1. Desarrolla el código de la aplicación y el archivo YAML de workload localmente.
  2. Valida el YAML localmente con databricks air run --file <file> --dry-run.
  3. Entrega el workload con databricks air run --file <file>. Añade --watch registros al flujo de logs hasta que termine la ejecución.
  4. Usa databricks air get, databricks air list, y databricks air logs para inspeccionar cargas de trabajo y registros.
  5. Opcionalmente, detenga una carga de trabajo en ejecución con databricks air cancel.

Para todos los comandos y opciones, véase air grupo de comandos.

Para ejecuciones programadas, flujos de trabajo multitarea o promociones entre entornos, puedes usar databricks air convert-to-dabs como punto de partida para un paquete. Consulta cómo programar cargas de trabajo de GPU y componer tareas.

Seguimiento de las cargas de trabajo con MLflow

Cada carga de trabajo enviada está representada por una ejecución de trabajo de Databricks y rastreada mediante un experimento de MLflow:

  • Ejecución de trabajos en Databricks: Muestra el estado de la carga de trabajo, los intentos de ejecución y la salida en la página de Trabajos ejecutados.
  • Experimento MLflow: Contiene las ejecuciones de MLflow con métricas del sistema capturadas automáticamente, configuración de carga de trabajo y registros. Tu aplicación también puede usar las APIs de MLflow para registrar parámetros personalizados, métricas y artefactos.

Utiliza el ID de Ejecución de Trabajo con comandos como databricks air get, databricks air logs, y databricks air cancel. Para más detalles sobre seguimiento y observabilidad, consulta Seguir ejecuciones con MLflow y la página de Ejecuciones de trabajos.

Ejecutar cargas de trabajo de nodo único y distribuidas

AI Runtime ejecuta el comando configurado una vez en cada nodo provisionado. Para una carga de trabajo distribuida, proporciona variables de coordinación como NUM_NODES, NODE_RANK, LOCAL_WORLD_SIZE, MASTER_ADDR, y MASTER_PORT.

Este modelo de ejecución es independiente del framework. Tu comando debe iniciar y configurar el framework distribuido, por ejemplo invocando torchrun, un script de arranque Ray u otro launcher de framework. Para ejemplos, véase ejemplos de CLI de Databricks para AI Runtime.

Usa un agente de IA

La databricks-ai-runtime agent skill proporciona instrucciones que ayudan a agentes de IA, como Claude Code, Cursor y Codex, a trabajar con AI Runtime a través de la CLI de Databricks. Úsalo para crear YAML de carga de trabajo, enviar trabajos de entrenamiento de GPU, comprobar el estado de las ejecuciones, transmitir registros y cancelar ejecuciones. También incluye orientación para imágenes personalizadas de Docker.

Después de instalar la CLI de Databricks y configurar la autenticación, ejecuta el siguiente comando:

databricks aitools install --skills-only --skills databricks-ai-runtime --experimental

Sigue las indicaciones para elegir tu agente de IA y el alcance de la instalación. Para opciones de instalación, consulta aitools el grupo de comandos.

Recursos adicionales

  • Se actualizó por última vez el