Use a CLI do Databricks com o Runtime de IA

A databricks air CLI executa cargas de trabalho de GPU de nó único no AI Runtime, a plataforma de computação de GPU serverless sob demanda. Defina a computação, ambiente e comando de uma carga de trabalho em YAML, e então envie-os com databricks air run.

O AI Runtime cuida do provisionamento da GPU e configuração do ambiente, além de fornecer status de execução, logs, métricas do sistema e rastreamento do MLflow. O fluxo de trabalho YAML e CLI funciona naturalmente com agentes de codificação, que podem definir, enviar e monitorar cargas de trabalho.

Use a CLI para:

  • Desenvolva e itere em aplicações Python do seu ambiente local.
  • Defina cargas de trabalho repetíveis em arquivos YAML que você possa controlar por versão.
  • Submeta e gerencie cargas de trabalho a partir de um terminal ou fluxo de trabalho automatizado de desenvolvimento.

As cargas de trabalho típicas incluem ajuste fino e pós-treinamento em LLM, treinamento distribuído, inferência em lote, aplicações Ray e busca de hiperparâmetros.

Para requisitos de espaço de trabalho, veja Requisitos. Para enviar uma carga de trabalho, siga o quickstart.

Os comandos da CLI do Databricks para o Runtime de IA atendem aos mesmos padrões de conformidade que as cargas de trabalho do Runtime de IA. Para obter os padrões e exceções com suporte, consulte o perfil de segurança de conformidade.

Observação

Os databricks air comandos e a configuração da carga de trabalho diferem da CLI legada baseada em Python air no pacote databricks-air. Veja migração da CLI em Python.

Explore os exemplos executáveis

O AI Runtime Cookbook contém cargas de trabalho de CLI Databricks executáveis para treinamento de modelos, ajuste fino, inferência em lote, Ray e busca por hiperparâmetros. Cada receita inclui um arquivo YAML da carga de trabalho, código-fonte e instruções para executá-la.

Como funciona uma carga de trabalho

Você fornece um arquivo YAML de workload, o código da aplicação quando necessário e o comando que inicia sua aplicação. Quando você envia um workload:

  1. A definição da carga de trabalho é validada em relação ao schema suportado e às restrições de configuração.
  2. Se code_source estiver configurado, o código selecionado é empacotado como um snapshot e carregado.
  3. Os recursos de computação da GPU e o ambiente solicitados são preparados em cada nó.
  4. O comando fornecido no YAML da carga de trabalho roda em cada nó.

O envio retorna um ID da execução do job que identifica a carga de trabalho. Use isso para verificar o status, inspecionar logs e a configuração, ou cancelar a carga de trabalho. Você também pode consultar a página de execução de Jobs e o experimento MLflow para detalhes de execução, métricas do sistema, logs e artefatos.

Esse fluxo de trabalho oferece provisionamento gerenciado de GPU, configuração de ambiente reproduzível e observabilidade embutida sem exigir que você configure ou mantenha a infraestrutura da GPU.

Configuração da carga de trabalho

Um arquivo YAML de carga de trabalho descreve os seguintes componentes:

  • Experimento: O experimento e a execução do MLflow que rastreiam a carga de trabalho.
  • Ambiente: O ambiente gerenciado, dependências ou imagem Docker personalizada para usar.
  • Compute: O tipo de acelerador e o número de aceleradores a serem provisionados.
  • Fonte do código: Um snapshot opcional do código da aplicação para disponibilizar à carga de trabalho.
  • Comando: O comando shell que inicia sua aplicação.
  • Configurações de runtime: Opções como tentativas, timeouts, variáveis de ambiente, segredos e permissões.

Para campos de configuração e exemplos, veja Referência YAML de Workload.

Use uma imagem Docker personalizada

Em vez de um ambiente gerenciado, uma carga de trabalho pode rodar na sua própria imagem Docker. Use uma imagem personalizada quando sua carga de trabalho precisar de bibliotecas do sistema, extensões CUDA ou outras dependências que você não possa instalar com environment.dependencies. Imagens personalizadas do Docker estão no Public Preview. Para usá-los, um administrador do espaço de trabalho deve ativar a pré-visualização do Databricks Artifact Registry na Previews page do workspace.

Você constrói a imagem para linux/amd64, a envia para o Artifact Registry no seu workspace e define environment.unity_catalog_image no YAML da carga de trabalho para o nome do Unity Catalog da imagem. Para os passos, imagens base do Databricks, autenticação Docker e solução de problemas, veja Usar imagens Docker personalizadas com tempo de execução de IA.

Fazer upload do código do aplicativo

Use code_source para fazer upload de arquivos do aplicativo com uma carga de trabalho. O AI Runtime disponibiliza o diretório carregado em cada nó e define $CODE_SOURCE_PATH com essa localização. Por exemplo, o seguinte comando executa train.py a partir do diretório carregado:

command: python "$CODE_SOURCE_PATH/train.py"

Você pode enviar a árvore de trabalho atual ou selecionar um branch Git ou commit. Para um exemplo completo, veja o quickstart.

Desenvolver e gerenciar cargas de trabalho

O desenvolvimento normalmente segue esse processo:

  1. Desenvolva o código da aplicação e o YAML da carga de trabalho localmente.
  2. Valide o YAML localmente com databricks air run --file <file> --dry-run.
  3. Submeta a carga de trabalho usando databricks air run --file <file> --watch para transmitir logs até a execução terminar.
  4. Use databricks air get, databricks air list, e databricks air logs para inspecionar cargas de trabalho e logs.
  5. Opcionalmente, pare uma carga de trabalho em execução com databricks air cancel.

Para enviar em segundo plano e anexar aos logs depois, omita --watch:

databricks air run --file workload.yaml
databricks air get <job-run-id>
databricks air logs <job-run-id>

Use o ID da execução do trabalho retornado pela submissão. Mantenha o mesmo perfil de espaço de trabalho para os três comandos. Para configuração e verificação do perfil padrão, veja Obter informações sobre perfis de configuração.

Para todos os comandos e sinalizadores, consulte air grupo de comandos.

Para execuções agendadas, fluxos de trabalho multitarefa ou promoção entre ambientes, você pode usar databricks air convert-to-dabs como ponto de partida para um bundle. Veja como agendar cargas de trabalho de GPU e compor tarefas.

Acompanhe cargas de trabalho com MLflow

Cada carga de trabalho submetida é representada por uma execução de job do Databricks e acompanhada por meio de um experimento do MLflow:

  • Execução de tarefas em Databricks: Mostra o status da carga de trabalho, tentativas de execução e saída na página de execução de jobs.
  • Experimento MLflow: Contém as execuções do MLflow com métricas do sistema capturadas automaticamente, configuração da carga de trabalho e logs. Sua aplicação também pode usar as APIs do MLflow para registrar parâmetros personalizados, métricas e artefatos.

Use o ID da execução do job com comandos como databricks air get, databricks air logs e databricks air cancel. Veja Rastreamento de experimentos e observabilidade.

Executar cargas de trabalho de nó único e distribuidas

O AI Runtime executa o comando configurado no nó provisionado. Para uma carga de trabalho multi-GPU, ele fornece variáveis de coordenação como LOCAL_WORLD_SIZE e WORLD_SIZE.

Esse modelo de execução é independente do framework. Seu comando deve iniciar e configurar o framework distribuído, por exemplo, invocando torchrun, um script de bootstrap Ray ou outro launcher de framework. Para cargas completas de trabalho, veja o AI Runtime Cookbook.

Usar um agente de IA

A databricks-ai-runtime habilidade do agente fornece instruções que ajudam agentes de IA, como Claude Code, Cursor e Codex, a trabalhar com o Runtime de IA através da linha de comando do Databricks. Use-o para criar o YAML da carga de trabalho, enviar trabalhos de treinamento de GPU, checar o status das execuções, transmitir logs e cancelar execuções. Também inclui orientações para imagens Docker personalizadas.

Após instalar a CLI Databricks e configurar a autenticação, execute o seguinte comando:

databricks aitools install --skills-only --skills databricks-ai-runtime

Siga as instruções para escolher seu agente de IA e o escopo da instalação. Para opções de instalação, veja aitools o grupo de comandos. Para saber o que a habilidade cobre, veja Usar habilidades de agente com IA Runtime.

Recursos adicionais

  • Última atualização em