Coloque cargas de trabalho de treinamento em produção
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:
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.
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:
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
tgzexplí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.
- Local directory
- Packaged artifact
- Workspace or volume path
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:
code_source_path: ./src
Declare um artefato tgz e aponte code_source_path para o arquivo de saída dele. Em databricks bundle deploy, a CLI cria o tarball, faz o upload dele, e a tarefa o extrai e executa o 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 criar um snapshot de uma branch confirmada ou de um commit.
Para usar o código que já foi feito o upload, defina code_source_path como um caminho /Workspace/… ou /Volumes/…. O Databricks usa o caminho tal como está e não empacota nada.
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:
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.
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:
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:
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:
databricks bundle deploy --target dev
databricks bundle run train_pipeline --target dev
Passos seguintes
- Execute e gerencie as cargas de trabalho do AI Runtime a partir da linha de comando com a AI Runtime CLI.
- Acompanhe as execuções de treinamento e gerencie os pontos de verificação. Consulte Acompanhamento de experimentos e observabilidade.