Configure seu ambiente AI Runtime
O AI Runtime fornece dois ambientes Python gerenciados. O ambiente padrão inclui cuda. A partir da versão 5 do ambiente, torch e torchvision não são mais pré-instalados. O ambiente Databricks AI vem pré-carregado com PyTorch, Transformers e outros frameworks de ML e de aprendizagem profunda. Escolha o ambiente padrão para controle total sobre sua pilha de dependências ou o ambiente de AI para uma configuração de treinamento pronta para uso. Você pode começar com qualquer ambiente e instalar mais pacotes você mesmo.
Pontos-key nesta página:
- Escolha o ambiente Padrão para controle total sobre as dependências ou o ambiente Databricks AI para uma pilha ML pronta para uso.
- O padrão inclui
cuda. A partir da versão 5 do ambiente,torchetorchvisionnão são pré-instalados. O ambiente de AI adiciona PyTorch, Transformers e muito mais. - Instale mais pacotes com
%uv pip install(versão do ambiente 5 e acima) ou%pip installem um notebook, ou comenvironment.dependenciesna CLI.
Qual ambiente usar
O AI Runtime oferece dois ambientes Python gerenciados, o ambiente Padrão e o ambiente Databricks AI. As cargas de trabalho que você envia com a CLI do Databricks também podem ser executadas em uma Docker Image personalizada.
ambiente | key características | Quando usar |
|---|---|---|
Ambiente padrão | Mínimo; inclui | Você deseja controle total sobre sua pilha de dependências e prefere instalar apenas o que precisa |
Ambiente Databricks AI | Pré-carregado com frameworks de ML populares (PyTorch, Transformers e mais) | Você deseja um ambiente completo para treinamento, ajuste fino e experimentação sem gerenciamento manual de dependências |
Docker Image personalizada (Visualização Pública) | Sua própria imagem armazenada no Artifact Registry, geralmente construída a partir de uma imagem base do Databricks. Disponível para cargas de trabalho da CLI do Databricks. Um admin de workspace deve habilitar a visualização Databricks Artifact Registry . | Você precisa de bibliotecas de sistema, extensões CUDA compiladas ou uma configuração de contêiner existente que não possa reproduzir com |
Você também pode usar um ambiente base de workspace que um administrador de workspace criou para o AI Runtime. Consulte Construir para compute de GPU serverless (AI Runtime).
Para escolher um ambiente em um notebook, abra o painel lateral Ambiente e selecione um ambiente base:
Ambiente padrão (ambiente mínimo)
Um ambiente mínimo e estável contendo apenas os pacotes necessários para a operação do AI Runtime. O ambiente inclui cuda para compatibilidade com GPU. A partir da versão 5 do ambiente, torch e torchvision não são mais pré-instalados. Instale as versões que sua carga de trabalho precisa com %uv pip install ou %pip install. Para os pacotes instalados em cada versão do ambiente, consulte as notas sobre a versão abaixo.
Ideal para: usuários que desejam controle total sobre sua pilha de dependências e preferem instalar apenas o que precisam.
Para selecionar: no painel lateral Ambiente , escolha Standard v6 como seu ambiente base.
Para obter mais detalhes sobre as versões dos pacotes instaladas em versões diferentes, consulte as notas sobre a versão:
Ambiente Databricks AI
Disponível no ambiente 4 e posterior. O ambiente de AI é desenvolvido sobre o ambiente Standard com pacotes de runtime comuns e pacotes específicos para machine learning em GPUs. Os pacotes pré-instalados incluem:
- PyTorch (com suporte a CUDA)
- Transformers (Hugging Face)
- E dependências adicionais de ML/DL
Melhor para: profissionais de ML que desejam um ambiente completo para cargas de trabalho de treinamento, ajuste fino e experimentação sem gerenciamento manual de dependências.
Para selecionar: no painel lateral Ambiente , escolha AI v6 como seu ambiente base.
Para obter mais detalhes sobre as versões dos pacotes instaladas em versões diferentes, consulte as notas sobre a versão:
Instalar pacotes adicionais
Você pode começar com um ambiente gerenciado ou um ambiente base do workspace e instalar seus próprios pacotes Python. Como instalá-los depende de como você se conecta ao AI Runtime.
Em um notebook
Instale pacotes na parte superior do seu Notebook ou script de treinamento com %uv pip, que é instalado com uv e é mais rápido que %pip. Pin versões para que as execuções permaneçam reproduzíveis:
%uv pip install "trl==1.1.0"
%uv pip install "transformers==5.5.4"
Para obter a referência completa de %pip e %uv pip, incluindo quais versões de ambiente oferecem suporte a cada um, consulte Gerenciar bibliotecas com %uv comandos.
Quando você usa o decorador @distributed para cargas de trabalho de multi-GPU, os pacotes instalados antes de chamar .distributed() são automaticamente capturados e propagados para todos os processos distribuídos. O tamanho total dos pacotes instalados não deve exceder 15 GB. Para obter mais detalhes, consulte Adicionar dependências ao notebook.
Com a CLI
A air CLI usa os mesmos ambientes gerenciados. Selecione um com environment.version, depois liste os pacotes para instalar em cima dele sob environment.dependencies no seu YAML de carga de trabalho. A CLI os instala no ambiente selecionado com uv, para que você não precise gerenciar um ambiente virtual separado:
environment:
version: '6'
dependencies:
- trl==1.1.0
- transformers==5.5.4
Para começar a partir do ambiente Databricks AI, defina version como um identificador de ambiente AI, como databricks_ai_v6. Para o formato completo de dependência e sinalizadores de instalação compatíveis, consulte ambiente. Para execução da carga de trabalho em sua própria imagem em vez de um ambiente gerenciado, consulte Usar Docker Images personalizadas com AI Runtime.
Ambientes base do workspace
Um administrador de workspace pode criar um ambiente base de workspace para o AI Runtime, que o torna disponível para todos os usuários do workspace através do menu suspenso Ambiente base . Para obter detalhes, consulte Build for serverless GPU compute (AI Runtime).
Para usar um ambiente base em uma carga de trabalho agendada, faça referência a ele pelo seu identificador no bloco environments de um job ou pacote. Por exemplo, esta tarefa de Pacotes de Automação Declarativa seleciona o ambiente Databricks AI:
resources:
jobs:
train:
tasks:
- task_key: train
environment_key: default
environments:
- environment_key: default
spec:
base_environment: databricks_ai_v6
Para um exemplo completo de job, consulte Programar cargas de trabalho de GPU e compor tarefas.
Armazenamento em cache de ambiente e módulos personalizados
Quando os ambientes são armazenados em cache?
O AI Runtime armazena em cache o ambiente do seu notebook para que, ao reabri-lo, você não precise reinstalar os pacotes. Para obter detalhes, consulte Reset as dependências do ambiente.
O cache se aplica a notebooks interativos. Cargas de trabalho executadas através da air CLI ou agendadas como jobs não usam esse cache; cada execução instala suas dependências do zero.
Como prática recomendada, pin os pacotes em uma versão específica para que seu ambiente em cache permaneça válido e suas execuções permaneçam reproduzíveis.
Como importo módulos personalizados?
As passos a seguir mostram como usar seus próprios módulos Python.
- Sincronize seu código com o workspace. Clone um repository Git no workspace com uma pasta Git ou faça upload dos arquivos diretamente como arquivos de workspace. Armazene o código compartilhado em
/Workspace/Sharedpara projetos de equipe ou em uma pasta de usuário para desenvolvimento pessoal. Manter o código em uma pasta Git permite o controle de versão. - Adicione a pasta a
sys.pathe importe-a na parteacima do seu Notebook ou script de treinamento:
import sys
sys.path.append("/Workspace/Shared/my-project/src")
from my_module import my_function
Limitações
Os seguintes recursos não estão disponíveis no AI Runtime:
- Funções Spark : Você não pode importar ou usar funções PySpark diretamente. O AI Runtime é um ambiente exclusivo para Python; o Spark não está disponível como um runtime local. No entanto, o Spark Connect está disponível para carregamento de dados. Consulte Carregar dados no AI Runtime.
- Bibliotecas do Databricks Runtime ML : Os pacotes pré-instalados não substituem o Databricks Runtime ML. Algumas bibliotecas de ML disponíveis no Databricks Runtime ML podem não estar pré-instaladas no AI Runtime.
- Artefatos privados : O AI Runtime oferece suporte a artefatos privados em determinados casos. Entre em contato com sua equipe de account para obter mais detalhes.