Programar cargas de trabalho da GPU e compor fluxos

Importante

Este recurso está no Public Preview. Para o utilizar, um administrador de espaço de trabalho deve ativar a pré-visualização do AI Runtime a partir da página de Pré-visualizações . Ver Gerir as pré-visualizações de Azure Databricks.

Utilize Pacotes de Automação Declarativa para agendar uma carga de trabalho de GPU em tempo de execução de IA e combiná-la com outros trabalhos, como a preparação de dados a montante. Defina a carga de trabalho e o seu cronograma em YAML e, em seguida, implemente-os como uma tarefa. Para um pipeline de treino, execute o pré-processamento no CPU compute e inicie o treino GPU depois de os dados estarem prontos.

Este guia começa com uma carga de trabalho que imprime "Hello world", depois constrói um pipeline de pré-processamento e treino programado. Se já tem uma carga de trabalho do AI Runtime (YAML), veja Converter uma carga de trabalho do AI Runtime para um bundle.

Como funciona

  • Um bundle contém o seu código e uma databricks.yml configuração.
  • Um job agrupa as tarefas e define quando elas são executadas.
  • Um(a) ai_runtime_task executa o teu comando em computação GPU serverless.

O campo depends_on liga as tarefas num grafo acíclico dirigido (DAG), e o campo schedule programa a execução da tarefa periodicamente.

databricks bundle deploy Carrega o código e cria ou atualiza a tarefa. databricks bundle run inicia imediatamente uma execução de tarefa. Em cada execução, o AI Runtime aprovisiona os recursos de computação da GPU, executa o seu comando e regista a execução no experimento MLflow nomeado.

Requirements

Exemplo de Hello World com ai_runtime_task

Este exemplo executa um comando shell numa GPU A10.

  1. Cria um diretório para o bundle. Nesse diretório, crie command.sh com o seguinte conteúdo:

    #!/usr/bin/env bash
set -euo pipefail
echo "Hello world"

  2. Crie databricks.yml no mesmo diretório:

    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

  3. A partir do diretório do bundle, valide, faça a implantação e execute a tarefa:

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

  4. Abra o URL de execução impresso pela CLI e veja a saída da tarefa hello. Contém Hello world. A execução também aparece no hello-ai-runtime experimento MLflow.

Exemplo mais complexo: agendar um pipeline de pré-processamento e treino

Este exemplo prepara um pequeno conjunto de dados no CPU e depois treina um modelo linear numa GPU A10. Ambas as tarefas usam o mesmo ficheiro de volume para passar dados entre elas. A tarefa está agendada para as 09:00 UTC diariamente, com o agendamento em pausa até a testar.

Criar o projeto

  1. Crie um diretório separado com o seguinte layout:

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

  2. Crie prep.py com os seguintes conteúdos. Substitua <catalog>, <schema>, e <volume> pelos nomes dos seus volumes. Este caderno de origem normaliza os valores de entrada e escreve os dados preparados no volume:

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

  3. Crie src/train.py com o seguinte conteúdo. Utilize o mesmo caminho de volume que em 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}")

  4. Criar command.sh para correr o código de treino. O conjunto abaixo empacota src/, por isso CODE_SOURCE_PATH aponta para o diretório extraído src :

    #!/usr/bin/env bash
set -euo pipefail
cd "$CODE_SOURCE_PATH"
python train.py

  5. Crie databricks.yml com o seguinte conteúdo. O pacote inclui src/ para a tarefa da GPU. O prep notebook corre em computação CPU serverless e depends_on só inicia train depois de prep ser concluído com êxito. max_concurrent_runs: 1 impede que execuções deste trabalho sobrescrevam o ficheiro de entrada umas das outras:

    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 ativar o agendamento

  1. A partir de scheduled-training/, implemente o pacote e inicie uma execução manual:

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

  2. Abra o URL de execução impresso pela CLI. Confirme que prep tem sucesso antes de train começar. A saída prep reporta 10 linhas preparadas, e a saída train reporta Trained on cuda:0 e a perda.

  3. Em databricks.yml, altere o pause_status da agenda de PAUSED para UNPAUSED. O bloco de horários é agora:

    schedule:
  quartz_cron_expression: '0 0 9 * * ?'
  timezone_id: UTC
  pause_status: UNPAUSED

  4. Implemente a alteração para ativar o horário diário:

    databricks bundle deploy --target dev

    Definir pause_status: UNPAUSED permite explicitamente este agendamento mesmo para um destino de desenvolvimento. Em Jobs & Pipelines, abra o trabalho implementado e confirme que o seu agendamento está ativo para as 09:00 UTC. Consulte Executar trabalhos numa programação.

Para pausar as execuções agendadas, defina pause_status: PAUSED e implemente novamente. Para remover o job de exemplo e os ficheiros carregados pelo bundle, execute o seguinte a partir do seu diretório bundle:

databricks bundle destroy --target dev

Os dados preparados no seu volume permanecem. Apaga a scheduled-training pasta de dados quando já não precisares dela.

Recursos adicionais

  • Última atualização a