Implantar um aplicativo Ray Serve com um modelo de difusão estável no Google Kubernetes Engine (GKE) com TPUs

Standard

Este guia demonstra como implantar e disponibilizar um modelo de difusão estável no Google Kubernetes Engine (GKE) usando TPUs, Ray Serve e o complemento Ray Operator.

Este guia é destinado a clientes de IA generativa, usuários novos ou atuais do GKE, engenheiros de ML, engenheiros de MLOps (DevOps) ou administradores de plataformas interessados em usar os recursos de orquestração de contêineres do Kubernetes para disponibilizar modelos usando o Ray.

Sobre o Ray e o Ray Serve

O Ray é um framework de computação escalonável e de código aberto para aplicativos de IA/ML. O Ray Serve é uma biblioteca de disponibilização de modelos para Ray usado para dimensionar e exibir modelos em um ambiente distribuído. Para mais informações, consulte Ray Serve na documetação do Ray.

Sobre TPUs

As unidades de processamento de tensor (TPUs) são aceleradores de hardware especializados projetados para acelerar significativamente o treinamento e a inferência de modelos de machine learning em grande escala. O uso do Ray com TPUs permite escalonar continuamente o alto desempenho aplicativos de ML. Para mais informações sobre TPUs, consulte Introdução ao Cloud TPU na documentação do Cloud TPU.

Sobre o webhook de inicialização do TPU do KubeRay

Como parte do complemento do Operador do Ray, o GKE fornece webhooks de validação e mutação que processam a programação de pods de TPU e algumas variáveis de ambiente de TPU necessárias por frameworks como o JAX para a inicialização do contêiner. O webhook do KubeRay para TPUs transforma pods com o rótulo app.kubernetes.io/name: kuberay que solicitam TPUs com as seguintes propriedades:

• TPU_WORKER_ID: um número inteiro exclusivo para cada pod de worker na fração de TPU.
• TPU_WORKER_HOSTNAMES: uma lista de nomes de host DNS para todos os workers da TPU que precisam se comunicar entre si na fatia. Essa variável é injetada apenas em pods de TPU em um grupo de vários hosts.
• replicaIndex: um rótulo de pod que contém um identificador exclusivo para o réplica do grupo de workers ao qual o pod pertence. Isso é útil para workers de vários hosts grupos de trabalho, em que vários pods de worker podem pertencer à mesma réplica, e é usada pelo Ray para ativar o escalonamento automático de vários hosts.
• TPU_NAME: uma string que representa o PodSlice do TPU do GKE ao qual este pod pertence, definido com o mesmo valor do rótulo replicaIndex.
• podAffinity: garante que o GKE programe pods de TPU com correspondência replicaIndex rótulos no mesmo pool de nós. Isso permite que o GKE escalonar TPUs de vários hosts atomicamente por pools de nós, em vez de nós únicos.

Objetivos

• Criar um cluster do GKE com um pool de nós de TPU.
• Implante um cluster do Ray com TPUs.
• Implantar um recurso personalizado do RayService.
• Interagir com o servidor de modelo de difusão estável.

Custos

Neste documento, você usará os seguintes componentes faturáveis do Google Cloud:

• GKE

Para gerar uma estimativa de custo baseada na projeção de uso deste tutorial, use a calculadora de preços.

Novos usuários do Google Cloud podem estar qualificados para um teste sem custo financeiro.

Ao concluir as tarefas descritas neste documento, é possível evitar o faturamento contínuo excluindo os recursos criados. Para mais informações, consulte Limpeza.

Antes de começar

O Cloud Shell vem pré-instalado com o software necessário para este tutorial, incluindo kubectl e a CLI gcloud. Se você não usa o Cloud Shell, instale a CLI gcloud.

Faça login na sua conta do Google Cloud . Se você começou a usar o Google Cloudagora, crie uma conta para avaliar o desempenho dos nossos produtos em situações reais. Clientes novos também recebem US $300 em créditos para executar, testar e implantar cargas de trabalho.

Instale a CLI do Google Cloud.

Observação:se você já instalou a CLI gcloud, execute o comando gcloud components update para verificar se tem a versão mais recente.

Ao usar um provedor de identidade (IdP) externo, primeiro faça login na CLI gcloud com sua identidade federada.

Para inicializar a CLI gcloud, execute o seguinte comando:

gcloud init

Crie ou selecione um Google Cloud projeto.

Funções necessárias para selecionar ou criar um projeto

• Selecionar um projeto: não é necessário ter um papel específico do IAM para selecionar um projeto. É possível selecionar qualquer projeto em que você tenha um papel.
• Criar um projeto: para criar um projeto, é preciso ter o papel de Criador de projetos (roles/resourcemanager.projectCreator), que contém a permissão resourcemanager.projects.create. Saiba como conceder papéis.

Observação: se você não pretende manter os recursos criados neste procedimento, crie um projeto novo em vez de selecionar um que já existe. Depois de concluir essas etapas, é possível excluir o projeto. Para fazer isso, basta remover todos os recursos associados a ele.

• Crie um projeto do Google Cloud :

  gcloud projects create PROJECT_ID

  Substitua PROJECT_ID por um nome para o projeto do Google Cloud que você está criando.

• Selecione o projeto Google Cloud que você criou:

  gcloud config set project PROJECT_ID

  Substitua PROJECT_ID pelo nome do projeto do Google Cloud .

Verifique se o faturamento está ativado para o projeto do Google Cloud .

Ative a API GKE, se ainda não tiver feito isso:

Funções necessárias para ativar APIs

Para ativar APIs, você precisa da permissão serviceusage.services.enable. Se você criou o projeto, provavelmente já tem essa permissão pelo papel de Proprietário (roles/owner). Caso contrário, é possível receber essa permissão pelo papel de Administrador do Service Usage (roles/serviceusage.serviceUsageAdmin). Saiba como conceder papéis.

gcloud services enable container.googleapis.com

Instale a CLI do Google Cloud.

Observação:se você já instalou a CLI gcloud, execute o comando gcloud components update para verificar se tem a versão mais recente.

Ao usar um provedor de identidade (IdP) externo, primeiro faça login na CLI gcloud com sua identidade federada.

Para inicializar a CLI gcloud, execute o seguinte comando:

gcloud init

Crie ou selecione um Google Cloud projeto.

Funções necessárias para selecionar ou criar um projeto

• Selecionar um projeto: não é necessário ter um papel específico do IAM para selecionar um projeto. É possível selecionar qualquer projeto em que você tenha um papel.
• Criar um projeto: para criar um projeto, é preciso ter o papel de Criador de projetos (roles/resourcemanager.projectCreator), que contém a permissão resourcemanager.projects.create. Saiba como conceder papéis.

Observação: se você não pretende manter os recursos criados neste procedimento, crie um projeto novo em vez de selecionar um que já existe. Depois de concluir essas etapas, é possível excluir o projeto. Para fazer isso, basta remover todos os recursos associados a ele.

• Crie um projeto do Google Cloud :

  gcloud projects create PROJECT_ID

  Substitua PROJECT_ID por um nome para o projeto do Google Cloud que você está criando.

• Selecione o projeto Google Cloud que você criou:

  gcloud config set project PROJECT_ID

  Substitua PROJECT_ID pelo nome do projeto do Google Cloud .

Verifique se o faturamento está ativado para o projeto do Google Cloud .

Ative a API GKE, se ainda não tiver feito isso:

Funções necessárias para ativar APIs

Para ativar APIs, você precisa da permissão serviceusage.services.enable. Se você criou o projeto, provavelmente já tem essa permissão pelo papel de Proprietário (roles/owner). Caso contrário, é possível receber essa permissão pelo papel de Administrador do Service Usage (roles/serviceusage.serviceUsageAdmin). Saiba como conceder papéis.

gcloud services enable container.googleapis.com

1. Conceda papéis à sua conta de usuário. Execute uma vez o seguinte comando para cada um dos seguintes papéis do IAM: roles/container.clusterAdmin, roles/container.admin

   gcloud projects add-iam-policy-binding PROJECT_ID --member="user:USER_IDENTIFIER" --role=ROLE

   Substitua:

   • PROJECT_ID: o ID do projeto.
   • USER_IDENTIFIER: o identificador da sua conta de usuário . Por exemplo, [e-mail].
   • ROLE: o papel do IAM que você concede à sua conta de usuário.

Verificar se há cota suficiente

Verifique se o projeto do Google Cloud tem cota suficiente de TPU na região ou zona do Compute Engine. Para mais informações, consulte Garantir cotas suficientes de TPU e GKE na documentação do Cloud TPU. Talvez também seja necessário aumentar cotas para:

• Persistent disk SSD (GB)
• Endereços IP em uso

Prepare o ambiente

Para preparar o ambiente, siga estas etapas:

1. Inicie uma sessão do Cloud Shell no console do Google Cloud clicando em Ativar o Cloud Shell no console doGoogle Cloud . Isso inicia uma sessão no painel inferior do console Google Cloud .

2. Defina as variáveis de ambiente:

   export PROJECT_ID=PROJECT_ID
   export CLUSTER_NAME=ray-cluster
   export COMPUTE_REGION=us-central2-b
   export CLUSTER_VERSION=CLUSTER_VERSION
   

   Substitua:

   • PROJECT_ID: o Google Cloud ID do projeto.
   • CLUSTER_VERSION: a versão do GKE a ser usada. Precisa ser 1.30.1 ou mais recente.

3. Clone o repositório do GitHub:

   git clone https://github.com/GoogleCloudPlatform/kubernetes-engine-samples
   

4. Mude para o diretório de trabalho:

   cd kubernetes-engine-samples/ai-ml/gke-ray/rayserve/stable-diffusion
   

Criar um cluster com um pool de nós de TPU

Crie um Padrão Cluster do GKE com um pool de nós da TPU:

1. Crie um cluster no modo padrão com o operador Ray ativado:

   gcloud container clusters create ${CLUSTER_NAME} \
       --addons=RayOperator \
       --machine-type=n1-standard-8 \
       --cluster-version=${CLUSTER_VERSION} \
       --location=${COMPUTE_REGION}
   

2. Crie um pool de nós de TPU de host único:

   gcloud container node-pools create tpu-pool \
       --location=${COMPUTE_REGION} \
       --cluster=${CLUSTER_NAME} \
       --machine-type=ct4p-hightpu-4t \
       --num-nodes=1
   

Para usar TPUs com o modo padrão, selecione:

• Um local do Compute Engine com capacidade para aceleradores de TPU
• Um tipo de máquina compatível com a TPU e
• A topologia física do PodSlice de TPU

Configurar um recurso do RayCluster com TPUs

Configure o manifesto do RayCluster para preparar a carga de trabalho da TPU:

Configurar a TPU nodeSelector

O GKE usa nodeSelectors do Kubernetes para garantir que as cargas de trabalho da TPU sejam programadas na topologia e no acelerador de TPU adequados. Para mais informações sobre como selecionar nodeSelectors de TPU, consulte Implantar cargas de trabalho de TPU no GKE Standard.

Atualize o manifesto ray-cluster.yaml para programar seu pod em um podslice de TPU v4 com uma topologia 2x2x1:

nodeSelector:
  cloud.google.com/gke-tpu-accelerator: tpu-v4-podslice
  cloud.google.com/gke-tpu-topology: 2x2x1

Configurar um recurso de contêiner da TPU

Para usar um acelerador de TPU, especifique o número de chips de TPU que O GKE precisa alocar para cada pod configurando google.com/tpuRecurso limits e requests no campo Contêiner da TPU do manifesto do RayCluster workerGroupSpecs.

Atualize o manifesto ray-cluster.yaml com limites e solicitações de recursos:

resources:
  limits:
    cpu: "1"
    ephemeral-storage: 10Gi
    google.com/tpu: "4"
    memory: "2G"
   requests:
    cpu: "1"
    ephemeral-storage: 10Gi
    google.com/tpu: "4"
    memory: "2G"

Configurar o grupo de workers numOfHosts

O KubeRay v1.1.0 adiciona um campo numOfHosts ao recurso personalizado do RayCluster. que especifica o número de hosts de TPU a serem criados por réplica de grupo de workers. Para grupos de workers de vários hosts, as réplicas são tratadas como podslice, em vez de workers individuais, com nós de trabalho numOfHosts sendo criados por réplica.

Atualize o manifesto ray-cluster.yaml com o seguinte:

workerGroupSpecs:
  # Several lines omitted
  numOfHosts: 1 # the number of "hosts" or workers per replica

Criar um recurso personalizado do RayService

Crie um recurso personalizado do RayService:

1. Analise o seguinte manifesto:

   apiVersion: ray.io/v1
   kind: RayService
   metadata:
     name: stable-diffusion-tpu
   spec:
     serveConfigV2: |
       applications:
         - name: stable_diffusion
           import_path: ai-ml.gke-ray.rayserve.stable-diffusion.stable_diffusion_tpu:deployment
           runtime_env:
             working_dir: "https://github.com/GoogleCloudPlatform/kubernetes-engine-samples/archive/refs/heads/main.zip"
             pip:
               - diffusers==0.7.2
               - flax
               - jax[tpu]==0.4.11
               - -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
               - fastapi
     rayClusterConfig:
       rayVersion: '2.9.0'
       headGroupSpec:
         rayStartParams: {}
         template:
           spec:
             containers:
             - name: ray-head
               image: rayproject/ray-ml:2.9.0-py310
               ports:
               - containerPort: 6379
                 name: gcs
               - containerPort: 8265
                 name: dashboard
               - containerPort: 10001
                 name: client
               - containerPort: 8000
                 name: serve
               resources:
                 limits:
                   cpu: "2"
                   memory: "8G"
                 requests:
                   cpu: "2"
                   memory: "8G"
       workerGroupSpecs:
       - replicas: 1
         minReplicas: 1
         maxReplicas: 10
         numOfHosts: 1
         groupName: tpu-group
         rayStartParams: {}
         template:
           spec:
             containers:
             - name: ray-worker
               image: rayproject/ray-ml:2.9.0-py310
               resources:
                 limits:
                   cpu: "100"
                   ephemeral-storage: 20Gi
                   google.com/tpu: "4"
                   memory: 200G
                 requests:
                   cpu: "100"
                   ephemeral-storage: 20Gi
                   google.com/tpu: "4"
                   memory: 200G
             nodeSelector:
               cloud.google.com/gke-tpu-accelerator: tpu-v4-podslice
               cloud.google.com/gke-tpu-topology: 2x2x1

   Esse manifesto descreve um recurso personalizado do RayService que cria um recurso do RayCluster com um nó principal e um grupo de workers da TPU com uma topologia 2x2x1, ou seja, cada nó de trabalho terá quatro chips de TPU v4.

   O nó de TPU pertence a um único podslice de TPU v4 com topologia 2x2x1. Para criar um grupo de workers com vários hosts, substitua os valores gke-tpu nodeSelector, os limites e solicitações do contêiner google.com/tpu e os valores numOfHosts pela configuração de vários hosts. Para mais informações sobre vários hosts de TPU topologias, consulte Arquitetura do sistema na documentação do Cloud TPU.

2. Aplique o manifesto ao cluster:

   kubectl apply -f ray-service-tpu.yaml
   

3. Verifique se o recurso RayService está em execução:

   kubectl get rayservices
   

   O resultado será assim:

   NAME                   SERVICE STATUS   NUM SERVE ENDPOINTS
   stable-diffusion-tpu   Running          2
   

   Nesta saída, Running na coluna SERVICE STATUS indica que o recurso do RayCluster está pronto.

(Opcional) Acessar o painel do Ray

É possível visualizar a implantação do Ray Serve e os registros relevantes no painel do Ray.

1. Estabeleça uma sessão de encaminhamento de portas para o painel do Ray no cabeçalho do Ray serviço:

   kubectl port-forward svc/stable-diffusion-tpu-head-svc 8265:8265
   

2. Em um navegador da Web, acesse http://localhost:8265/.

3. Clique na guia Exibir.

Enviar comandos ao servidor de modelos

1. Estabeleça uma sessão de encaminhamento de portas para o endpoint do Serve do serviço de cabeça do Ray:

   kubectl port-forward svc/stable-diffusion-tpu-serve-svc 8000
   

2. Abra uma nova sessão do Cloud Shell.

3. Envie um comando de criação de imagens a partir de texto ao servidor de modelos Stable Diffusion:

   python stable_diffusion_tpu_req.py  --save_pictures
   

   Os resultados da inferência de difusão estável são salvos em um arquivo chamado diffusion_results.png.

   

Observar suas cargas de trabalho do Ray

Para conferir detalhes sobre os RayJobs, acesse a seção Kubernetes Engine > IA/ML > Jobs no console do Google Cloud .

Ver RayJobs no console do Google Cloud

Limpar

Excluir o projeto

Cuidado: a exclusão de um projeto tem os seguintes efeitos:

• Tudo no projeto é excluído. Se você tiver usado um projeto atual para as tarefas neste documento, a exclusão dele incluirá a exclusão de quaisquer outros trabalhos feitos nele.
• Os IDs do projeto personalizados são perdidos. Ao criar o projeto, você pode ter criado um ID do projeto personalizado para ser usado no futuro. Para preservar os URLs que usam o ID do projeto, como um URL appspot.com, exclua recursos específicos do projeto, em vez de excluir o projeto inteiro.

Se você planeja explorar várias arquiteturas, tutoriais e guias de início rápido, a reutilização de projetos pode evitar que você exceda limites da cota.

Excluir um projeto do Google Cloud :

gcloud projects delete PROJECT_ID

Excluir recursos individuais

Para excluir o cluster, digite:

gcloud container clusters delete ${CLUSTER_NAME}

A seguir

• Saiba mais sobre o Ray no Kubernetes.
• Confira a documentação do KubeRay.
• Confira arquiteturas de referência, diagramas, tutoriais e práticas recomendadas do Google Cloud. Confira o Centro de arquitetura do Cloud.