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ê vai usar os seguintes componentes faturáveis do Google Cloud:

• GKE

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

Novos usuários do Google Cloud podem estar qualificados para um teste gratuito.

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 gcloud CLI. Se você não usa o Cloud Shell, instalar a gcloud CLI.

Sign in to your Google Cloud account. If you're new to Google Cloud, create an account to evaluate how our products perform in real-world scenarios. New customers also get $300 in free credits to run, test, and deploy workloads.

Install the Google Cloud CLI.

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

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

Para inicializar a gcloud CLI, execute o seguinte comando:

gcloud init

Create or select a Google Cloud project.

Note: If you don't plan to keep the resources that you create in this procedure, create a project instead of selecting an existing project. After you finish these steps, you can delete the project, removing all resources associated with the project.

• Create a Google Cloud project:

  gcloud projects create PROJECT_ID

  Replace PROJECT_ID with a name for the Google Cloud project you are creating.

• Select the Google Cloud project that you created:

  gcloud config set project PROJECT_ID

  Replace PROJECT_ID with your Google Cloud project name.

Make sure that billing is enabled for your Google Cloud project.

Enable the GKE API:

gcloud services enable container.googleapis.com

Install the Google Cloud CLI.

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

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

Para inicializar a gcloud CLI, execute o seguinte comando:

gcloud init

Create or select a Google Cloud project.

Note: If you don't plan to keep the resources that you create in this procedure, create a project instead of selecting an existing project. After you finish these steps, you can delete the project, removing all resources associated with the project.

• Create a Google Cloud project:

  gcloud projects create PROJECT_ID

  Replace PROJECT_ID with a name for the Google Cloud project you are creating.

• Select the Google Cloud project that you created:

  gcloud config set project PROJECT_ID

  Replace PROJECT_ID with your Google Cloud project name.

Make sure that billing is enabled for your Google Cloud project.

Enable the GKE API:

gcloud services enable container.googleapis.com

1. Grant roles to your user account. Run the following command once for each of the following IAM roles: roles/container.clusterAdmin, roles/container.admin

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

   • Replace PROJECT_ID with your project ID.

   • Replace USER_IDENTIFIER with the identifier for your user account. For example, user:[e-mail].

   • Replace ROLE with each individual role.

Verificar se há cota suficiente

Verifique se o projeto 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 ID do projeto Google Cloud .
   • 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 PodSlices, em vez de workers individuais, com nós de worker 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 de TPU com uma topologia 2x2x1, ou seja, cada nó de worker 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 do modelo

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 texto para imagem ao servidor do modelo 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.

   

Limpar

Excluir o projeto

Cuidado: excluir 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 vai resultar na exclusão de outros trabalhos feitos nele.
• Os IDs de projeto personalizados são perdidos. É possível que, ao criar o projeto, você tenha criado um ID de 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.

Delete a Google Cloud project:

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.