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

Padrão

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:

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 uma avaliação gratuita.

Ao concluir as tarefas descritas neste documento, é possível evitar o faturamento contínuo excluindo os recursos criados. Saiba mais em 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.

To initialize the gcloud CLI, run the following command:

gcloud init

Note: If you installed the gcloud CLI previously, make sure you have the latest version by running

gcloud components
      update

Create or select a Google Cloud 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.

To initialize the gcloud CLI, run the following command:

gcloud init

Note: If you installed the gcloud CLI previously, make sure you have the latest version by running

gcloud components
      update

Create or select a Google Cloud 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

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_IDENTIFIER" --role=ROLE
```
- Replace PROJECT_ID with your project ID.
- Replace USER_IDENTIFIER with the identifier for your user account. For example, user:myemail@example.com.
- Replace ROLE with each individual role.

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:

Inicie uma sessão do Cloud Shell no Console do Google Cloud clicando em Ativar o Cloud Shell no Console do Google Cloud. Isso inicia uma sessão no painel inferior do Cloud Console.
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: seuID de projeto no Google Cloud.
- CLUSTER_VERSION: a versão do GKE a ser usada. Precisa ser 1.30.1 ou mais recente.

Clone o repositório do GitHub:

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

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:

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}

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 \
    --tpu-topology=2x2x1

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:

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.

Aplique o manifesto ao cluster:
```
kubectl apply -f ray-service-tpu.yaml
```
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.

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
```
Em um navegador da Web, acesse http://localhost:8265/.
Clique na guia Exibir.

Enviar comandos ao servidor do modelo

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
```
Abra uma nova sessão do Cloud Shell.
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 existente para as tarefas neste documento, a exclusão dele incluirá a exclusão de quaisquer outros trabalhos feitos no projeto.
Os IDs do projeto personalizados são perdidos. Ao criar o projeto, você pode ter criado um código 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 passar por várias arquiteturas, tutoriais ou guias de início rápido, a reutilização de projetos pode evitar que você exceda os limites da cota do projeto.

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.