Programar cargas de trabajo de GPU y componer tareas

Utiliza paquetes de automatización declarativa para programar una carga de trabajo de GPU de AI Runtime y combinarla con otras tareas, como la preparación de datos en la fase previa. Define la carga de trabajo y su programación en YAML y, a continuación, impleméntalas como un trabajo. Para una canalización de entrenamiento, ejecuta preprocesamiento en el cómputo de CPU y comienza el entrenamiento con GPU una vez que los datos estén listos.

Esta guía comienza con una carga de trabajo que imprime "Hello world" y luego construye una canalización programada de preprocesamiento y entrenamiento. Si ya tienes una carga de trabajo en tiempo de ejecución de IA en YAML, consulta Convertir una carga de trabajo en tiempo de ejecución de IA a un paquete.

Cómo funciona

Un bundle contiene tu código y una databricks.yml configuración.
Un trabajo agrupa las tareas y define cuándo se ejecutan.
ai_runtime_task ejecuta tu comando en AI Runtime.

El depends_on campo conecta las tareas en un grafo acíclico dirigido (DAG), y el schedule campo programa tu trabajo para que se ejecute con una cadencia.

databricks bundle deploy Sube el código y crea o actualiza la tarea. databricks bundle run inicia una ejecución del trabajo inmediatamente. En cada ejecución, AI Runtime provisiona el cómputo de la GPU, ejecuta tu comando y registra la ejecución en el experimento de MLflow con el nombre indicado.

Requirements

Un espacio de trabajo en una región compatible. Ver Requisitos.
La última CLI de Databricks, autenticada en su espacio de trabajo. Actualiza una instalación existente antes de usar estos ejemplos.
Para el ejemplo de preprocesamiento, un volumen existente de Unity Catalog. La identidad de ejecución del trabajo necesita los privilegios USE CATALOG, USE SCHEMA, READ VOLUME y WRITE VOLUME para ese volumen. Consulte Privilegios para volúmenes del catálogo de Unity.

Ejemplo de Hello World con `ai_runtime_task`

Este ejemplo ejecuta un comando shell en una GPU A10.

Crea un directorio para el paquete. En ese directorio, crea command.sh con el siguiente contenido:
```
#!/usr/bin/env bash
set -euo pipefail
echo "Hello world"
```

Crea databricks.yml en el mismo directorio:

bundle:
  name: hello-ai-runtime

resources:
  jobs:
    hello:
      name: hello-ai-runtime
      tasks:
        - task_key: hello
          environment_key: default
          ai_runtime_task:
            experiment: hello-ai-runtime
            deployments:
              - command_path: ${workspace.file_path}/command.sh
                compute:
                  accelerator_type: GPU_1xA10
                  accelerator_count: 1
      environments:
        - environment_key: default
          spec:
            environment_version: '6'

targets:
  dev:
    mode: development
    default: true

Desde el directorio del paquete, valida, despliega y ejecuta la tarea:

databricks bundle validate --target dev
databricks bundle deploy --target dev
databricks bundle run hello --target dev

Abra la URL de ejecución generada por la CLI y vea la salida de la hello tarea. Contiene Hello world. La ejecución también aparece en el experimento hello-ai-runtime MLflow.

Ejemplo más complejo: programar una cadena de preprocesamiento y entrenamiento

Este ejemplo prepara un pequeño conjunto de datos sobre computación de CPU y luego entrena un modelo lineal en una GPU A10. Ambas tareas utilizan el mismo archivo de volumen para pasar datos entre ellas. El trabajo está programado para las 09:00 UTC todos los días, con el horario pausado hasta que lo pruebes.

Creación del proyecto

Crea un directorio separado con el siguiente diseño:

scheduled-training/
├── databricks.yml
├── prep.py
├── command.sh
└── src/
    └── train.py

Crea prep.py con el siguiente contenido. Sustituye <catalog>, <schema>, y <volume> por los nombres de tus volúmenes. Este cuaderno fuente normaliza los valores de entrada y escribe los datos preparados en el volumen:

# Databricks notebook source
import json
from pathlib import Path

data_path = Path("/Volumes/<catalog>/<schema>/<volume>/scheduled-training/data.json")
inputs = [value / 10 for value in range(10)]
data = {"x": inputs, "y": [2 * value + 1 for value in inputs]}
data_path.parent.mkdir(parents=True, exist_ok=True)
data_path.write_text(json.dumps(data))
print(f"Prepared {len(inputs)} rows at {data_path}")

Crea src/train.py con el siguiente contenido. Utiliza la misma ruta del volumen que en prep.py:

import json
from pathlib import Path

import torch

data_path = Path("/Volumes/<catalog>/<schema>/<volume>/scheduled-training/data.json")
data = json.loads(data_path.read_text())
x = torch.tensor(data["x"], device="cuda").reshape(-1, 1)
y = torch.tensor(data["y"], device="cuda").reshape(-1, 1)
model = torch.nn.Linear(1, 1).to("cuda")
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)

for _ in range(200):
    optimizer.zero_grad()
    loss = torch.nn.functional.mse_loss(model(x), y)
    loss.backward()
    optimizer.step()

print(f"Trained on {x.device}: loss={loss.item():.4f}")

Cree command.sh para ejecutar el código de entrenamiento. El paquete siguiente empaqueta src/, por lo que CODE_SOURCE_PATH apunta al directorio src extraído:
```
#!/usr/bin/env bash
set -euo pipefail
cd "$CODE_SOURCE_PATH"
python train.py
```

Crea databricks.yml con el siguiente contenido. El paquete incluye src/ para la tarea de GPU. El prep notebook funciona con computación de CPU sin servidor y depends_on solo inicia train después de que prep se complete correctamente. max_concurrent_runs: 1 evita que las ejecuciones de este trabajo se sobrescriban mutuamente el archivo de entrada:

bundle:
  name: scheduled-training

artifacts:
  code:
    type: tgz
    path: .
    include: [src]
    files:
      - source: ./dist/code.tgz

resources:
  jobs:
    train_pipeline:
      name: scheduled-training
      max_concurrent_runs: 1
      schedule:
        quartz_cron_expression: '0 0 9 * * ?'
        timezone_id: UTC
        pause_status: PAUSED
      tasks:
        - task_key: prep
          notebook_task:
            notebook_path: ./prep.py
        - task_key: train
          depends_on:
            - task_key: prep
          environment_key: training
          ai_runtime_task:
            experiment: scheduled-training
            code_source_path: ./dist/code.tgz
            deployments:
              - command_path: ${workspace.file_path}/command.sh
                compute:
                  accelerator_type: GPU_1xA10
                  accelerator_count: 1
      environments:
        - environment_key: training
          spec:
            environment_version: '6'
            dependencies:
              - torch

targets:
  dev:
    mode: development
    default: true

Prueba y activa la programación

Desde scheduled-training/, despliega el paquete y comienza una ejecución manual:

databricks bundle validate --target dev
databricks bundle deploy --target dev
databricks bundle run train_pipeline --target dev

Abre la URL de la ejecución impresa por la CLI. Confirma que prep se complete correctamente antes de que train empiece. La salida de prep muestra 10 filas preparadas, y la salida de train muestra Trained on cuda:0 y la pérdida.
En databricks.yml, cambie el pause_status de la programación de PAUSED a UNPAUSED. El bloque de horarios es ahora:
```
schedule:
  quartz_cron_expression: '0 0 9 * * ?'
  timezone_id: UTC
  pause_status: UNPAUSED
```
Despliega el cambio para activar el horario diario:
```
databricks bundle deploy --target dev
```
Configurar pause_status: UNPAUSED habilita explícitamente esta programación incluso para un objetivo de desarrollo. En Jobs & Pipelines, abre el trabajo desplegado y confirma que su horario está activo a las 09:00 UTC. Consulte Ejecución de trabajos según una programación.

Para pausar las ejecuciones programadas, configura pause_status: PAUSED y vuelve a desplegar. Para eliminar el trabajo de ejemplo y los archivos subidos por el paquete, ejecuta lo siguiente desde su directorio de paquete:

databricks bundle destroy --target dev

Los datos preparados en tu volumen de almacenamiento permanecen. Elimina el scheduled-training directorio de datos cuando ya no lo necesites.

Recursos adicionales

Convierte un YAML de carga de trabajo de IA en tiempo de ejecución existente en un paquete. Consulta Convertir una carga de trabajo de AI Runtime en un paquete.
Configura el empaquetado del código, los recursos de procesamiento, los entornos y los parámetros de la tarea. Consulte Configurar tareas de paquetes de AI Runtime.
Programe un cuaderno existente en la GPU. Consulte Creación y administración de trabajos de cuaderno programados.
Haz seguimiento de las sesiones de entrenamiento. Consulta Seguimiento y observabilidad de experimentos.
Modelo de punto de control, optimizador y estado del canal de datos. Consulta Mejorar el rendimiento y la resiliencia del entrenamiento en AI Runtime.

Comentarios

¿Le ha resultado útil esta página?

Se actualizó por última vez el 2026-10-23

Programar cargas de trabajo de GPU y componer tareas

Cómo funciona

Requirements

Ejemplo de Hello World con ai_runtime_task

Ejemplo más complejo: programar una cadena de preprocesamiento y entrenamiento

Creación del proyecto

Prueba y activa la programación

Recursos adicionales

Comentarios

Recursos adicionales

Ejemplo de Hello World con `ai_runtime_task`