Pular para o conteúdo principal

Coloque cargas de trabalho de treinamento em produção

info

Visualização

Esse recurso está em Prévia Pública.

Use DABs para definir uma carga de trabalho de treinamento do AI Runtime como código. Mantenha-o no controle de código-fonte, implante-o em vários ambientes, programe-o e crie-o em conjunto com outras tarefas. Esta página aborda o caminho traga seu próprio treinamento (bring-your-own-training), em que um ai_runtime_task executa seu próprio comando em um diretório de código em compute de GPU Serverless.

Esta é uma tarefa diferente de executar um notebook em GPU serverless por meio de um pacote. Para obter o exemplo básico de pacote de notebook em GPU, consulte API de Jobs e Pacotes de Automação Declarativa.

Requisitos

  • Um workspace com o AI Runtime ativado. Consulte Requisitos.
  • A CLI do Databricks (interface de linha de comando) instalada e configurada para implantar pacotes.

Definir uma tarefa do AI Runtime em um pacote

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

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

command_path é o script que a execução da tarefa. Repetições, tempos limite e permissões são definidos na tarefa e no job da mesma forma que qualquer job do Databricks, portanto, suas práticas de pacote existentes são transferidas.

Configure o acelerador de hardware

Defina accelerator_type como a GPU que sua carga de trabalho precisa e accelerator_count como 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. Uma contagem maior do que o tamanho por nó executa a tarefa em vários nós — por exemplo, GPU_8xH100 com accelerator_count: 16 é executada em dois nós. Para obter orientações sobre como escolher um acelerador, consulte Opções de hardware.

nota

Para execuções com vários nós, o AI Runtime executa seu comando em cada nó e preenche as variáveis de ambiente de treinamento distribuído padrão no ambiente da tarefa —NUM_NODES, WORLD_SIZE, LOCAL_WORLD_SIZE, MASTER_ADDR e MASTER_PORT. Faça a leitura deles a partir do seu comando (por exemplo, uma inicialização de torchrun); eles não são definidos no pacote.

Definir o ambiente e as dependências

Declare um bloco environments no job e faça referência a ele a partir da tarefa com environment_key. O AI Runtime instala as dependências listadas antes de o comando ser executado:

YAML
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: '5'
dependencies:
- numpy

Para os ambientes disponíveis, consulte Configure o seu ambiente.

Envie o seu código de treinamento

code_source_path informa à tarefa onde o seu código de treinamento está. Ele aceita três formulários:

  • Um diretório local — a CLI o empacota e faz o upload em databricks bundle deploy. Esta é a forma mais simples e a que os exemplos acima usam.
  • Um artefato tgz explícito — você declara o artefato por conta própria e aponta para o arquivo de saída dele. Use isto quando precisar empacotar apenas um subconjunto de arquivos ou tirar um snapshot de uma revisão do Git consolidadas em vez da sua árvore de trabalho.
  • Um caminho de Workspace ou volume — código que já teve o upload, usado no estado em que se encontra.

Aponte code_source_path para um diretório. Em databricks bundle deploy, a CLI empacota o diretório, faz o upload dele, e a tarefa o extrai e executa seu comando nele:

YAML
code_source_path: ./src
nota

O AI Runtime extrai seu código para um diretório e o expõe como a variável de ambiente CODE_SOURCE_PATH. Faça referência a ele a partir do seu comando para que os caminhos relativos sejam resolvidos, por exemplo, cd "$CODE_SOURCE_PATH" antes de executar seu script.

Criar fluxos de trabalho multitarefa

Um ai_runtime_task é uma tarefa de job do Databricks, portanto, ele é composto com o restante de um job. Você pode executar uma etapa de preparação antes do treinamento, combinar tarefas de GPU e CPU em um único job e usar diferentes aceleradores por tarefa.

Ordene tarefas com depends_on

Use depends_on para executar tarefas em sequência. O pipeline a seguir executa um notebook de preparação e, em seguida, uma tarefa de treinamento em GPU que é iniciada somente após a conclusão bem-sucedida da tarefa de preparação:

YAML
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: ./src
deployments:
- command_path: ./command.sh
compute:
accelerator_type: GPU_1xA10
accelerator_count: 1

Combinar tarefas de GPU e CPU

No pipeline acima, apenas o passo de treinamento precisa de uma GPU. Manter o trabalho que não usa GPU, como a preparação de dados, em tarefas separadas faz com que o tempo da GPU permaneça focado no treinamento.

nota

Um ai_runtime_task não oferece suporte a valores de tarefa de job do Databricks ({{tasks.<task_key>.values.<name>}} ou dbutils.jobs.taskValues). Para passar dados entre os passos, grave-os em um local compartilhado que ambas as tarefas possam ler, como um volume do Unity Catalog ou um arquivo do workspace, e faça referência a esse caminho a partir de cada tarefa.

Programar a carga de trabalho

Adicione um schedule ao job para executá-lo em uma cadência. Envie a programação pausada para que a implantação do pacote não comece execuções por conta própria, depois retome-a quando estiver pronto:

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

Promover de desenvolvimento para produção

A promoção de desenvolvimento para produção é um recurso padrão de pacote que a tarefa de AI Runtime herda sem alterações.

Destinos e modos de pacote

Defina um destino de produção com mode: production junto com seu destino de desenvolvimento. O destino controla onde o pacote é implantado e como seus recursos são nomeados:

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

Implantar e executar

Implante e execute o pacote em um destino com os comandos de pacote padrão:

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

Passos seguintes