Programar cargas de trabalho de GPU e compor tarefas

info

Visualização

This recurso is in Public Preview. To use it, a workspace admin must enable the AI Runtime preview from the Previews page. See Gerenciar Databricks previews.

Use Declarative Automation Bundles para programar uma carga de trabalho de GPU do AI Runtime e combiná-la com outro trabalho, como a preparação de dados upstream. Defina a carga de trabalho e seu programar em YAML e, em seguida, implante-os como um job. Para um pipeline de treinamento, faça a execução do pré-processamento no compute de CPU e comece o treinamento de GPU depois que os dados estiverem prontos.

Este guia começa com uma carga de trabalho que imprime "Hello world" e, em seguida, cria um pipeline programado de pré-processamento e treinamento. Se você já tiver um arquivo YAML de carga de trabalho do AI Runtime, consulte Converter uma carga de trabalho do AI Runtime em um pacote.

Como funciona

Um bundle contém seu código e uma configuração databricks.yml.
Um Job agrupa tarefas e define quando elas estão em execução.
Um ai_runtime_task faz a execução de seu comando em compute de GPU Serverless.

O campo depends_on conecta as tarefas em um gráfico acíclico direcionado (DAG), e o campo schedule programa seu Job para execução em uma cadência.

databricks bundle deploy faz o upload do código e cria ou atualiza o job. databricks bundle run começa uma execução de job imediatamente. Em cada execução, o AI Runtime provisiona o compute de GPU, executa o comando e registra a execução no experimento do MLflow especificado.

Requisitos

Um workspace em uma região suportada. Consulte Requisitos.
A versão mais recente da CLI do Databricks, autenticada no seu workspace. Atualize uma instalação existente antes de usar estes exemplos.
Para o exemplo de pré-processamento, um volume do Unity Catalog existente. A identidade de execução do job precisa dos privilégios USE CATALOG, USE SCHEMA, READ VOLUME e WRITE VOLUME para esse volume. Consulte Privilégios para volumes do Unity Catalog.

Exemplo de "Olá, Mundo" com `ai_runtime_task`

Este exemplo executa um comando shell em uma GPU A10.

Criar um diretório para o bundle. Nesse diretório, crie command.sh com o seguinte conteúdo:
Bash
```
#!/usr/bin/env bash
set -euo pipefail
echo "Hello world"
```

Crie databricks.yml no mesmo diretório:

YAML
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

No diretório do bundle, valide, implante e execute o job:

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

Abra a URL de execução impressada pela CLI e view a saída da tarefa hello. Ele contém Hello world. A execução também aparece no experimento do MLflow hello-ai-runtime.

Exemplo mais complexo: programar um pipeline de pré-processamento e treinamento

This example prepares a small dataset on CPU compute, then ensinar a linear model on one A10 GPU. Both tarefas use the same volume file to pass data between them. The Job is programar for 09:00 UTC daily, with the programar em pausa until you test it.

Criar o projeto

Crie um diretório separado com a seguinte disposição:

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

Crie prep.py com o seguinte conteúdo. Substitua <catalog>, <schema> e <volume> pelos nomes do seu volume. Este notebook de origem normaliza os valores de entrada e grava os dados preparados no volume:

Python
# 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}")

Crie src/train.py com o seguinte conteúdo: Use o mesmo caminho de volume que em prep.py:

Python
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}")

Crie command.sh para a execução do código de treinamento. O pacote abaixo empacota src/, portanto, CODE_SOURCE_PATH aponta para o diretório extraído src:
Bash
```
#!/usr/bin/env bash
set -euo pipefail
cd "$CODE_SOURCE_PATH"
python train.py
```

Crie databricks.yml com o seguinte conteúdo. O pacote empacota src/ para a tarefa de GPU. O notebook prep é executado no CPU compute serverless, e depends_on começa train somente após prep ser bem-sucedido. max_concurrent_runs: 1 impede que execuções deste job substituam o arquivo de entrada umas das outras:

YAML
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

Teste e ative o programar

A partir de scheduled-training/, implante o pacote e comece uma execução manual:

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

Abra o URL de execução impresso pela CLI. Confirme se prep é bem-sucedido antes que train comece. A saída prep relata 10 linhas preparadas e a saída train relata Trained on cuda:0 e a perda.
Em databricks.yml, altere o(a) pause_status do programar de PAUSED para UNPAUSED. O bloco de programar agora é:
YAML
```
schedule:
  quartz_cron_expression: '0 0 9 * * ?'
  timezone_id: UTC
  pause_status: UNPAUSED
```
Implante a alteração para ativar o programar diário:
Bash
```
databricks bundle deploy --target dev
```
Definir pause_status: UNPAUSED explicitamente habilita este programar mesmo para um alvo de desenvolvimento. Em Jobs & Pipelines , abra o job implantado e confirme se o programar está ativo para 09:00 UTC. Consulte execução de Job em programar.

Para pausa de execuções programadas, defina pause_status: PAUSED e faça o deploy novamente. Para remover o Job de exemplo e os arquivos upload pelo bundle, execução o seguinte a partir do diretório do bundle:

Bash
databricks bundle destroy --target dev

Os dados preparados no volume permanecem. Exclua o diretório de dados scheduled-training quando não precisar mais dele.

Recursos adicionais

Convert an existing AI Runtime workload YAML into a bundle. See Convert an AI Runtime workload to a bundle.
Configure o empacotamento de código, o compute, os ambientes e os parâmetros de tarefa. Consulte Configurar tarefas do pacote do AI Runtime.
Programe um Notebooks existente na GPU. Consulte Programar com a API de Jobs e os Pacotes de Automação Declarativa.
Track execuções de treinamento. See Experiment acompanhamento and observability.
Checkpoint model, optimizer, and pipeline de dados state. See Improve treinamento desempenho and resiliency on AI Runtime.

Como funciona​

Requisitos​

Exemplo de "Olá, Mundo" com ai_runtime_task​

Exemplo mais complexo: programar um pipeline de pré-processamento e treinamento​

Criar o projeto​

Teste e ative o programar​

Recursos adicionais​