Esta página foi traduzida pela API Cloud Translation.

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

Este guia demonstra como implementar e publicar um modelo 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/AA. O Ray Serve é 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 trabalhadores com vários anfitriões, em que vários pods de trabalhadores 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 os TPUs multi-anfitrião de forma atómica por conjuntos de nós, em vez de nós únicos.

Prepare o seu ambiente

Para preparar o seu ambiente, siga estes passos:

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

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

Crie um cluster com um node pool de TPUs

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

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}

Crie um node pool de TPU de anfitrião ú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/tpurecursolimits 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 trabalho 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:

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.

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 é 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.

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
```
Num navegador de Internet, aceda a http://localhost:8265/.
Clique no separador Publicar.

Enviar comandos para o servidor do modelo

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
```
Abra uma nova sessão do Cloud Shell.
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.

Implemente uma aplicação Ray Serve com um modelo Stable Diffusion no Google Kubernetes Engine (GKE) com TPUs Mantenha tudo organizado com as coleções Salve e categorize o conteúdo com base nas suas preferências.