Use a CLI do Databricks com o Runtime de IA

A databricks air CLI executa cargas de trabalho de GPU de nó único e múltiplos nós 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 de nó único ou de múltiplos nós em arquivos YAML que você possa armazenar em controle de 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.

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. Envie o workload com databricks air run --file <file>. Adicione --watch aos logs do stream 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 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 uma vez em cada nó provisionado. Para uma carga de trabalho distribuída, ele fornece variáveis de coordenação como NUM_NODES, NODE_RANK, LOCAL_WORLD_SIZE, MASTER_ADDR, e MASTER_PORT.

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