Colocar cargas de trabalho de treinamento em produção

Importante

Esse recurso está em Visualização Pública.

Use DABs para definir uma carga de trabalho de treinamento em tempo de execução de IA como código. Mantenha-o no controle de versão, implante-o em vários ambientes, agende-o e combine-o com outras tarefas. Esta página aborda o caminho de trazer seu próprio treinamento, onde um ai_runtime_task executa seu próprio comando em relação a um diretório de código em computação GPU sem servidor.

Dica

  • Use Pacotes de Automação Declarativa para definir cargas de trabalho de treinamento como código, implantá-las em ambientes e agenda-las.
  • O ai_runtime_task executa seu próprio comando em um diretório de código (com seu próprio treinamento).
  • Combine tarefas de GPU e CPU em trabalhos multitarefas.

Essa é uma tarefa diferente de rodar um notebook em GPU serverless usando um bundle. Para o exemplo básico do pacote notebook-on-GPU, veja Agendar com a API de Jobs e Pacotes de Automação Declarativa.

Requirements

  • Um espaço de trabalho com o runtime de IA ativado. Veja Requisitos.
  • A CLI (interface de linha de comando) do Databricks instalada e configurada para implantar pacotes.

Defina uma tarefa de tempo de execução de IA em um pacote

ai_runtime_task nomeia um experimento, aponta para seu código de treinamento com code_source_path e define uma implantação: o comando a ser executado e a GPU em que ele será executado. Adicione isto a uma tarefa no seu pacote:

resources:
  jobs:
    train:
      tasks:
        - task_key: train
          ai_runtime_task:
            experiment: my-experiment
            code_source_path: ./dist/code.tgz
            deployments:
              - command_path: ./command.sh
                compute:
                  accelerator_type: GPU_1xA10
                  accelerator_count: 1

code_source_path aponta para o seu código de treinamento empacotado, e command_path é o script que a tarefa executa. Tentativas, tempos limite e permissões são definidos na tarefa e no job da mesma forma que em qualquer job do Azure Databricks, portanto suas práticas de bundle existentes também se aplicam aqui. Para saber como empacotar e referenciar seu código, veja Enviar seu código de treinamento.

Campos ai_runtime_task

Campo Tipo Description
experiment String Required. O nome do experimento do MLflow para a execução. Veja Rastreamento de experimentos e observabilidade.
code_source_path String O código de treinamento a ser executado: o arquivo de saída de um artefato tgz empacotado ou um caminho /Workspace ou /Volumes para o código que já foi carregado. Confira Enviar seu código de treinamento.
deployments Sequência Required. Uma única implantação que descreve o comando e os recursos computacionais para executá-lo. Cada entrada contém command_path, compute, e um opcional name.
deployments[].command_path String Required. O script que a tarefa executa em cada nó.
deployments[].compute.accelerator_type String Required. O tipo de GPU, por exemplo GPU_1xA10, GPU_1xH100, ou GPU_8xH100.
deployments[].compute.accelerator_count Número inteiro Required. O número total de GPUs em todos os nós — um múltiplo da quantidade por nó codificada em accelerator_type.
deployments[].name String Um nome opcional para a implantação, usado nos logs e na interface.
docker_image_url String Uma imagem Docker personalizada opcional para executar o comando, em vez do ambiente gerenciado. Consulte Usar imagens personalizadas do Docker.
mlflow_run String Um nome de exibição opcional para a execução do MLflow.
mlflow_experiment_directory String Um diretório opcional de espaço de trabalho sob o qual o experimento é criado. Deve começar com /Workspace. Defina isso ao rodar como um principal de serviço que não tem diretório de usuário padrão.
mlflow_artifact_location String Uma localização raiz opcional para artefatos do MLflow, por exemplo, um caminho /Volumes/<catalog>/<schema>/<volume>/…. Deve corresponder à localização do artefato de um experimento existente ou ser omitido.

Defina tentativas, timeouts, permissões e o ambiente (environment_key) na tarefa e no job — não dentro de ai_runtime_task. Para a referência completa da tarefa, veja a tarefa de execução de IA.

Configure o acelerador de hardware

Defina accelerator_type para a GPU que sua carga de trabalho precisa e accelerator_count para o número total de GPUs. A contagem é um múltiplo do número de GPUs por nó: 1 para GPU_1xA10 e GPU_1xH100, e 8 para GPU_8xH100. Um número maior que o tamanho por nó distribui a tarefa por vários nós — por exemplo, GPU_8xH100 com accelerator_count: 16 é executado em dois nós. Para orientações sobre a escolha de um acelerador, veja Opções de hardware.

Note

Para execuções multinó, o AI Runtime executa seu comando em cada nó e define as variáveis de ambiente padrão de treinamento distribuído no ambiente da tarefa — NUM_NODES, WORLD_SIZE, LOCAL_WORLD_SIZE, MASTER_ADDR e MASTER_PORT. Leia-os a partir do seu comando (por exemplo, uma inicialização de torchrun); você não os define no pacote.

Defina o ambiente e as dependências

Declare um bloco environments no job e faça referência a ele na tarefa com environment_key. O AI Runtime instala as dependências listadas antes que seu comando seja executado:

resources:
  jobs:
    train:
      tasks:
        - task_key: train
          environment_key: default
          ai_runtime_task:
            # experiment, code_source_path, and deployments as above
      environments:
        - environment_key: default
          spec:
            environment_version: '6'
            dependencies:
              - numpy

Para os ambientes disponíveis, veja Configurar seu ambiente.

Envie seu código de treinamento

code_source_path informa à tarefa onde está seu código de treinamento. Ela assume uma de duas formas:

  • Um artefato empacotado tgz — declare um artefato e aponte code_source_path para seu arquivo de saída. O Azure Databricks cria o tarball e faz upload dele em databricks bundle deploy. É assim que você envia código de um diretório local de projeto ou de uma revisão commit do Git.
  • Um caminho para um workspace ou volume — código já carregado, usado como está.

Artefato embalado

Declare um artefato tgz e aponte code_source_path para o seu arquivo de saída. Em databricks bundle deploy, a CLI gera o tarball, faz upload dele, e a tarefa o extrai e executa seu comando nele:

artifacts:
  code:
    type: tgz
    path: .
    include: [src]
    files:
      - source: ./dist/code.tgz
resources:
  jobs:
    train:
      tasks:
        - task_key: train
          ai_runtime_task:
            code_source_path: ./dist/code.tgz

Use include para empacotar arquivos da sua árvore de trabalho, ou git para fazer instantâneos de uma ramificação ou confirmação.

Caminho do espaço de trabalho ou do volume

Para usar código que já foi carregado, defina code_source_path para um caminho /Workspace/… ou /Volumes/…. Azure Databricks usa o caminho as-is e não empacota nada.

Os campos tgz do artefato:

Campo Description
type tgz constrói um tarball com gzip a partir dos arquivos-fonte, em vez de executar um build comando.
path O diretório base para empacotar. include os caminhos e os nomes das entradas do arquivo são relativos a ele.
include Uma lista de subcaminhos de path para o pacote. Omita para empacotar todos os path. Respeita .gitignore; sync.include e sync.exclude do pacote não se aplicam. Alternativa a um build comando.
git Faça um instantâneo de uma referência do Git confirmado em vez da árvore de trabalho. Defina git.branch ou git.commit (commit tem prioridade quando ambos estiverem definidos). Alternativa a um build comando.
files[].source O caminho do tarball criado. Aponte code_source_path para isso.

Note

O Runtime de IA extrai seu código para um diretório e o disponibiliza por meio da variável de ambiente CODE_SOURCE_PATH. Faça referência a ele no seu comando para que os caminhos relativos sejam resolvidos corretamente, por exemplo cd "$CODE_SOURCE_PATH" antes de executar seu script.

Exemplo completo

Este exemplo treina em uma única GPU A10 de um projeto local, sem nenhuma configuração prévia de CLI além de instalar e configurar a CLI do Azure Databricks. O projeto possui três arquivos:

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

command.sh é o ponto de entrada nomeado por command_path. Ele muda para o diretório de código extraído e executa o script de treinamento:

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

databricks.yml nomeia o pacote, empacota src/ como um artefato tgz, executa-o como um ai_runtime_task, instala numpy no ambiente da tarefa e define alvos de desenvolvimento e produção:

bundle:
  name: my-training

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

resources:
  jobs:
    train:
      name: my-training
      tasks:
        - task_key: train
          environment_key: default
          ai_runtime_task:
            experiment: /Users/me@example.com/my-training
            code_source_path: ./dist/code.tgz
            deployments:
              - command_path: ./command.sh
                compute:
                  accelerator_type: GPU_1xA10
                  accelerator_count: 1
      environments:
        - environment_key: default
          spec:
            environment_version: '6'
            dependencies:
              - numpy

targets:
  dev:
    mode: development
    default: true
  prod:
    mode: production

Implante o pacote e execute a tarefa. databricks bundle deploy constrói e faz upload do tgz artefato e cria o trabalho; databricks bundle run inicia:

databricks bundle deploy --target dev
databricks bundle run train --target dev

Construir fluxos de trabalho multitarefa

Um ai_runtime_task é uma tarefa de trabalho do Azure Databricks; portanto, ela compõe o restante de um trabalho. Você pode executar uma etapa de preparação antes do treinamento, combinar tarefas de GPU e CPU em um único trabalho, e usar aceleradores diferentes por tarefa.

Ordenar tarefas com depends_on

Use depends_on para executar tarefas em sequência. O pipeline a seguir executa um caderno de preparação, depois uma tarefa de treinamento da GPU que começa somente após o sucesso da tarefa de preparação:

resources:
  jobs:
    train_pipeline:
      tasks:
        - task_key: prep
          notebook_task:
            notebook_path: ./prep.py
        - task_key: train
          depends_on:
            - task_key: prep
          ai_runtime_task:
            experiment: my-experiment
            code_source_path: ./dist/code.tgz
            deployments:
              - command_path: ./command.sh
                compute:
                  accelerator_type: GPU_1xA10
                  accelerator_count: 1

Combinar tarefas de GPU e CPU

No pipeline acima, apenas a etapa de treinamento precisa de uma GPU. Manter trabalhos fora da GPU, como preparação de dados, em tarefas separadas mantém o tempo da GPU focado no treinamento.

Note

An ai_runtime_task não suporta valores de tarefa de trabalho do Azure Databricks ({{tasks.<task_key>.values.<name>}} ou dbutils.jobs.taskValues). Para passar dados entre etapas, escreva-os em um local compartilhado que ambas as tarefas possam ler, como um volume do Unity Catalog ou um arquivo de espaço de trabalho, e faça referência a esse caminho de cada tarefa.

Organize a carga de trabalho

Adicione um schedule ao job para executá-lo em intervalos regulares. Envie o cronograma pausado para que a implantação do pacote não inicie as execuções sozinha, e depois despause quando estiver pronto:

resources:
  jobs:
    train_pipeline:
      schedule:
        quartz_cron_expression: '0 0 9 * * ?'
        timezone_id: UTC
        pause_status: PAUSED

Promover do desenvolvimento para a produção

A promoção de desenvolvimento para produção é uma funcionalidade padrão do bundle que a tarefa AI Runtime herda sem nenhuma alteração.

Destinos e modos do pacote

Defina uma meta de produção com mode: production junto com sua meta de desenvolvimento. O alvo controla onde o pacote é implantado e como seus recursos são nomeados:

targets:
  dev:
    mode: development
    default: true
  prod:
    mode: production

Implantar e executar

Implante e execute o pacote contra um alvo com os comandos padrão do bundle:

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

Próximas Etapas