Programar cargas de trabalho de GPU e compor tarefas

Use Pacotes de Automação Declarativa para programar uma carga de trabalho de GPU do AI Runtime e combiná-la com outros trabalhos, como a preparação de dados upstream. Defina a carga de trabalho e seu programar no YAML, depois 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", em seguida, cria um pipeline de pré-processamento e treinamento agendado. Se você já tiver um YAML de carga de trabalho do AI Runtime, consulte Converter uma carga de trabalho do AI Runtime em um bundle.

Como funciona

Um pacote contém seu código e uma configuração databricks.yml.
Um Job agrupa tarefas e define quando elas estarão em execução.
Uma ai_runtime_task execução seu comando no AI Runtime.

O campo depends_on conecta 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 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 seu comando e registra a execução no experiment MLflow nomeado.

Requisitos

Um workspace em uma região suportada. Consulte Requisitos.
A CLI do Databricks mais recente, 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 faz a execução de um comando shell em uma GPU A10.

Crie 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"
```

Criar 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 impressa pela CLI e view a saída da tarefa hello. Ele contém Hello world. A execução também aparece no experimento MLflow hello-ai-runtime.

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

Este exemplo prepara um pequeno dataset em compute de CPU e, em seguida, ensinar um modelo linear em uma GPU A10. Ambas as tarefas usam o mesmo arquivo de volume para passar dados entre elas. O Job está programado para às 09:00 UTC diariamente, com o programar em pausa até você testá-lo.

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 os seguintes conteúdos. Substitua <catalog>, <schema> e <volume> pelos nomes do seu volume. Este notebook de origem normaliza os valores de entrada e escreve 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 executar o código de treinamento. O pacote abaixo empacota src/, portanto, CODE_SOURCE_PATH aponta para o diretório src extraído:
Bash
```
#!/usr/bin/env bash
set -euo pipefail
cd "$CODE_SOURCE_PATH"
python train.py
```

Crie databricks.yml com o seguinte conteúdo. Os pacotes de pacotes src/ para a tarefa de GPU. O Notebook prep está em execução em compute de CPU Serverless, e depends_on começa train apenas após o sucesso de prep. max_concurrent_runs: 1 impede que execuções desse job sobrescrevam o arquivo de entrada um do outro:

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

Testar e habilitar programar

A partir de scheduled-training/, implante o pacote e começar 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 que prep seja concluído com êxito antes de train começar. A saída de prep informa 10 linhas preparadas, e a saída de train informa Trained on cuda:0 e a perda.
Em databricks.yml, altere pause_status de programar de PAUSED para UNPAUSED. O bloco de programar agora é:
YAML
```
schedule:
  quartz_cron_expression: '0 0 9 * * ?'
  timezone_id: UTC
  pause_status: UNPAUSED
```
Faça o deploy da mudança para ativar o programar diário:
Bash
```
databricks bundle deploy --target dev
```
Definir pause_status: UNPAUSED explicitamente habilita esse agendamento até mesmo para um alvo de desenvolvimento. Em Jobs & Pipelines , abra o Job implantado e confirme se sua programar está ativa para 09:00 UTC. Consulte Executar jobs em uma programar.

Para pausar execuções programadas, defina pause_status: PAUSED e implante novamente. Para remover o Job de exemplo e os arquivos cujo upload foi feito pelo pacote, realize a execução do seguinte a partir do diretório do pacote:

Bash
databricks bundle destroy --target dev

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

Outros recursos

Converta um YAML de carga de trabalho do AI Runtime existente em um pacote. Consulte Converter uma carga de trabalho do AI Runtime em um pacote.
Configure o empacotamento de código, compute, ambientes e parâmetros de tarefa. Veja Configurar tarefas do pacote AI Runtime.
Programe um notebook existente na GPU. Consulte Criar e gerenciar jobs de notebook agendados.
Acompanhe execuções de treinamento. Veja Acompanhamento de experimentos e observabilidade.
Estado do modelo de ponto de verificação, do otimizador e do pipeline de dados. Consulte Melhorar o desempenho e a resiliência do treinamento no 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​

Testar e habilitar programar​

Outros recursos​