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 enviar una carga de trabajo, siga la guía de inicio rápido.

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.

Explora ejemplos ejecutables

El AI Runtime Cookbook contiene cargas de trabajo ejecutables de la CLI de Databricks para entrenamiento de modelos, ajuste fino, inferencia por lotes, Ray y búsqueda de hiperparámetros. Cada receta incluye un archivo YAML de carga de trabajo, código fuente e instrucciones para ejecutarla.

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. Al enviar 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 consultar los campos de configuración y los ejemplos, véase Referencia de YAML de cargas de trabajo.

Subir código de aplicación

Utilice code_source para cargar archivos de aplicación con una carga de trabajo. AI Runtime hace que el directorio cargado esté disponible en todos los nodos y establece $CODE_SOURCE_PATH en su ubicación. Por ejemplo, el siguiente comando ejecuta train.py desde el directorio cargado:

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. Envía la carga de trabajo con databricks air run --file <file>. Añade --watch para transmitir los registros hasta que finalice 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 promoción entre entornos, puedes utilizar 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. Consulta Seguimiento y observabilidad de experimentos.

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 cargas completas de trabajo, consulta el Cookbook de Runtime de IA.

Usa un agente de IA

La databricks-ai-runtimehabilidad de agente proporciona instrucciones que ayudan a los 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

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

  • Última actualización el