Implemente uma aplicação Ray Serve com um modelo Stable Diffusion no Google Kubernetes Engine (GKE) com TPUs

Padrão

Este guia demonstra como implementar e publicar um modelo de Stable Diffusion no Google Kubernetes Engine (GKE) usando TPUs, Ray Serve, e o suplemento Ray Operator.

Este guia destina-se a clientes de IA generativa, utilizadores novos ou existentes do GKE, engenheiros de ML, engenheiros de MLOps (DevOps) ou administradores de plataformas interessados em usar capacidades de orquestração de contentores do Kubernetes para publicar modelos com o Ray.

Acerca do Ray e do Ray Serve

O Ray é uma framework de computação escalável de código aberto para aplicações de IA/AM. O RayServe é uma biblioteca de publicação de modelos para o Ray usada para dimensionar e publicar modelos num ambiente distribuído. Para mais informações, consulte o artigo Ray Serve na documentação do Ray.

Acerca das TPUs

As unidades de processamento tensor (TPUs) são aceleradores de hardware especializados concebidos para acelerar significativamente a preparação e a inferência de modelos de aprendizagem automática de grande escala. A utilização do Ray com TPUs permite-lhe dimensionar facilmente aplicações de ML de elevado desempenho. Para mais informações sobre TPUs, consulte o artigo Introdução ao Cloud TPU na documentação do Cloud TPU.

Acerca do webhook de inicialização do KubeRay TPU

Como parte do suplemento do operador do Ray, o GKE fornece webhooks de validação e mutação que processam o agendamento de pods de TPU e determinadas variáveis de ambiente de TPU necessárias por frameworks como o JAX para a inicialização de contentores. O webhook do KubeRay TPU altera os pods com o pedido de etiqueta app.kubernetes.io/name: kuberay de TPUs com as seguintes propriedades:

• TPU_WORKER_ID: um número inteiro exclusivo para cada pod de processador no fragmento da TPU.
• TPU_WORKER_HOSTNAMES: uma lista de nomes de anfitrião DNS para todos os trabalhadores da TPU que precisam de comunicar entre si na fatia. Esta variável só é injetada para agrupamentos de TPUs num grupo com vários anfitriões.
• replicaIndex: uma etiqueta de agrupamento que contém um identificador exclusivo para a réplica do grupo de trabalhadores à qual o agrupamento pertence. Isto é útil para grupos de trabalho com vários anfitriões, em que vários pods de trabalho podem pertencer à mesma réplica, e é usado pelo Ray para ativar o dimensionamento automático com vários anfitriões.
• TPU_NAME: uma string que representa o GKE TPU PodSlice ao qual este pod pertence, definida com o mesmo valor que a etiqueta replicaIndex.
• podAffinity: garante que o GKE agenda pods de TPU com etiquetas replicaIndexcorrespondentes no mesmo conjunto de nós. Isto permite que o GKE dimensione as TPUs com vários anfitriões de forma atómica por conjuntos de nós, em vez de nós únicos.

Objetivos

• Crie um cluster do GKE com um node pool de TPUs.
• Implemente um cluster do Ray com TPUs.
• Implemente um recurso personalizado RayService.
• Interagir com o servidor do modelo Stable Diffusion.

Custos

Neste documento, usa os seguintes componentes faturáveis do Google Cloud:

• GKE

Para gerar uma estimativa de custos com base na sua utilização prevista, use a calculadora de preços.

Os novos Google Cloud utilizadores podem ser elegíveis para uma avaliação gratuita.

Quando terminar as tarefas descritas neste documento, pode evitar a faturação contínua eliminando os recursos que criou. Para mais informações, consulte o artigo Limpe.

Antes de começar

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

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.

Nota: se instalou a CLI gcloud anteriormente, certifique-se de que tem a versão mais recente executando gcloud components update.

Se estiver a usar um fornecedor de identidade (IdP) externo, tem primeiro de iniciar sessão na CLI gcloud com a sua identidade federada.

Para inicializar a CLI gcloud, execute o seguinte comando:

gcloud init

Create or select a Google Cloud project.

Roles required to select or create a project

• Select a project: Selecting a project doesn't require a specific IAM role—you can select any project that you've been granted a role on.
• Create a project: To create a project, you need the Project Creator role (roles/resourcemanager.projectCreator), which contains the resourcemanager.projects.create permission. Learn how to grant roles.

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.

Verify that billing is enabled for your Google Cloud project.

Enable the GKE API:

Roles required to enable APIs

To enable APIs, you need the Service Usage Admin IAM role (roles/serviceusage.serviceUsageAdmin), which contains the serviceusage.services.enable permission. Learn how to grant roles.

gcloud services enable container.googleapis.com

Install the Google Cloud CLI.

Nota: se instalou a CLI gcloud anteriormente, certifique-se de que tem a versão mais recente executando gcloud components update.

Se estiver a usar um fornecedor de identidade (IdP) externo, tem primeiro de iniciar sessão na CLI gcloud com a sua identidade federada.

Para inicializar a CLI gcloud, execute o seguinte comando:

gcloud init

Create or select a Google Cloud project.

Roles required to select or create a project

• Select a project: Selecting a project doesn't require a specific IAM role—you can select any project that you've been granted a role on.
• Create a project: To create a project, you need the Project Creator role (roles/resourcemanager.projectCreator), which contains the resourcemanager.projects.create permission. Learn how to grant roles.

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.

Verify that billing is enabled for your Google Cloud project.

Enable the GKE API:

Roles required to enable APIs

To enable APIs, you need the Service Usage Admin IAM role (roles/serviceusage.serviceUsageAdmin), which contains the serviceusage.services.enable permission. Learn how to grant roles.

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 the following:

   • PROJECT_ID: Your project ID.
   • USER_IDENTIFIER: The identifier for your user account. For example, [e-mail].
   • ROLE: The IAM role that you grant to your user account.

Certifique-se de que tem quota suficiente

Certifique-se de que o seu Google Cloud projeto tem quota de TPU suficiente na sua região ou zona do Compute Engine. Para mais informações, consulte o artigo Garanta quotas suficientes de TPUs e GKEs na documentação do Cloud TPU. Também pode ter de aumentar as suas quotas para:

• Disco persistente SSD (GB)
• Endereços IP em utilização

Prepare o seu ambiente

Para preparar o seu ambiente, siga estes passos:

1. Inicie uma sessão do Cloud Shell a partir da Google Cloud consola, clicando em Ativar Cloud Shell na Google Cloud consola. Esta ação inicia uma sessão no painel inferior da Google Cloud consola.

2. Defina 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 o seguinte:

   • PROJECT_ID: o seu Google Cloud ID do projeto.
   • CLUSTER_VERSION: a versão do GKE a usar. Tem de ser 1.30.1 ou posterior.

3. Clone o repositório do GitHub:

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

4. Altere para o diretório de trabalho:

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

Crie um cluster com um node pool de TPUs

Crie um cluster do GKE padrão com um node pool de TPUs:

1. Crie um cluster do 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 node pool 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, tem de selecionar:

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

Configure um recurso RayCluster com TPUs

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

Configure a TPU nodeSelector

O GKE usa nodeSelectors do Kubernetes para garantir que as cargas de trabalho da TPU são agendadas na topologia e no acelerador de TPU adequados. Para mais informações sobre a seleção de nodeSelectors de TPUs, consulte o artigo Implemente cargas de trabalho de TPUs no GKE Standard.

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

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

Configure um recurso de contentor de TPU

Para usar um acelerador de TPU, tem de especificar o número de chips de TPU que o GKE deve atribuir a cada pod configurando o google.com/tpurecurso limits e requests no campo do contentor de TPUworkerGroupSpecs do manifesto do RayCluster.

Atualize o manifesto ray-cluster.yaml com os limites e os pedidos 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"

Configure o grupo de trabalhadores numOfHosts

O KubeRay v1.1.0 adiciona um campo numOfHosts ao recurso personalizado RayCluster, que especifica o número de anfitriões de TPU a criar por réplica do grupo de trabalhadores. Para grupos de trabalhadores com vários anfitriões, as réplicas são tratadas como PodSlices em vez de trabalhadores individuais, com numOfHosts nós de trabalhadores a serem 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

Crie um recurso personalizado RayService

Crie um recurso personalizado RayService:

1. Reveja 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

   Este manifesto descreve um recurso personalizado RayService que cria um recurso RayCluster com 1 nó principal e um grupo de trabalhadores de TPU com uma topologia de 2x2x1, o que significa que cada nó de trabalhador terá 4 chips de TPU v4.

   O nó da TPU pertence a um único podslice da TPU v4 com uma topologia de 2x2x1. Para criar um grupo de trabalho com vários anfitriões, substitua os valores gke-tpu nodeSelector, os limites e os pedidos de contentores google.com/tpu e os valores numOfHosts pela sua configuração com vários anfitriões. Para mais informações sobre as topologias de vários anfitriões da TPU, consulte a secção Arquitetura do sistema na documentação da 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 é semelhante ao seguinte:

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

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

(Opcional) Veja o painel de controlo do Ray

Pode ver a implementação do Ray Serve e os registos relevantes no painel de controlo do Ray.

1. Estabeleça uma sessão de encaminhamento de porta para o painel de controlo do Ray a partir do serviço principal do Ray:

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

2. Num navegador de Internet, aceda a http://localhost:8265/.

3. Clique no separador Publicar.

Enviar comandos para o servidor do modelo

1. Estabeleça uma sessão de encaminhamento de portas para o ponto final Serve a partir do serviço principal 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 para o 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 guardados num ficheiro denominado diffusion_results.png.

   

Limpar

Elimine o projeto

Atenção: a eliminação de um projeto tem os seguintes efeitos:

• Tudo no projeto é eliminado. Se usou um projeto existente para as tarefas neste documento, quando o elimina, também elimina qualquer outro trabalho que tenha feito no projeto.
• Os IDs dos projetos personalizados são perdidos. Quando criou este projeto, pode ter criado um ID do projeto personalizado que quer usar no futuro. Para preservar os URLs que usam o ID do projeto, como um URL appspot.com, elimine os recursos selecionados no projeto em vez de eliminar o projeto inteiro.

Se planeia explorar várias arquiteturas, tutoriais ou inícios rápidos, a reutilização de projetos pode ajudar a evitar exceder os limites de quota de projetos.

Delete a Google Cloud project:

gcloud projects delete PROJECT_ID

Elimine recursos individuais

Para eliminar o cluster, escreva:

gcloud container clusters delete ${CLUSTER_NAME}

O que se segue?

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