Esegui il deployment di un'applicazione Ray Serve con un modello Stable Diffusion su Google Kubernetes Engine (GKE) con TPU

Questa guida illustra come eseguire il deployment e il servizio di un modello di diffusione stabile su Google Kubernetes Engine (GKE) utilizzando TPU, Ray Serve e il componente aggiuntivo Ray Operator.

Questa guida è rivolta ai clienti di Generative AI, agli utenti nuovi o esistenti di GKE, agli ingegneri ML, agli ingegneri MLOps (DevOps) o agli amministratori della piattaforma interessati a utilizzare le funzionalità di orchestrazione dei container Kubernetes per la pubblicazione di modelli utilizzando Ray.

Informazioni su Ray e Ray Serve

Ray è un framework di calcolo scalabile open source per applicazioni AI/ML. Ray Serve è una libreria di distribuzione dei modelli per Ray utilizzata per scalare e distribuire i modelli in un ambiente distribuito. Per ulteriori informazioni, consulta Ray Serve nella documentazione di Ray.

Informazioni sulle TPU

Le TPU (Tensor Processing Unit) sono acceleratori hardware specializzati progettati per accelerare notevolmente l'addestramento e l'inferenza di modelli di machine learning su larga scala. L'utilizzo di Ray con le TPU consente di scalare senza problemi le applicazioni ML ad alte prestazioni. Per ulteriori informazioni sulle TPU, consulta la sezione Introduzione a Cloud TPU nella documentazione di Cloud TPU.

Informazioni sull'webhook di inizializzazione TPU di KubeRay

Nell'ambito del componente aggiuntivo Ray Operator, GKE fornisce webhooks di convalida e mutazione che gestiscono la pianificazione dei pod di TPU e alcune variabili di ambiente TPU richieste da framework come JAX per l'inizializzazione del contenitore. L'webhook TPU di KubeRay muta i pod con l'etichetta app.kubernetes.io/name: kuberay che richiedono TPU con le seguenti proprietà:

  • TPU_WORKER_ID: un numero intero univoco per ogni pod di lavoro nella sezione TPU.
  • TPU_WORKER_HOSTNAMES: un elenco di nomi host DNS per tutti i worker TPU che devono comunicare tra loro all'interno del segmento. Questa variabile viene inserita solo per i pod TPU in un gruppo di più host.
  • replicaIndex: un'etichetta del pod che contiene un identificatore univoco per la replica del gruppo di lavoro a cui appartiene il pod. Questo è utile per i gruppi di worker su più host, in cui più pod worker potrebbero appartenere alla stessa replica, ed è utilizzato da Ray per abilitare la scalabilità automatica su più host.
  • TPU_NAME: una stringa che rappresenta la sezione di pod TPU GKE a cui appartiene questo pod, impostata sullo stesso valore dell'etichetta replicaIndex.
  • podAffinity: garantisce che GKE programmi i pod TPU con le etichettereplicaIndex corrispondenti nello stesso pool di nodi. In questo modo, GKE può scalare le TPU multi-host in modo atomico in base ai pool di nodi anziché ai singoli nodi.

Obiettivi

  • Crea un cluster GKE con un pool di nodi TPU.
  • Esegui il deployment di un cluster Ray con TPU.
  • Esegui il deployment di una risorsa personalizzata RayService.
  • Interagire con il server del modello di diffusione stabile.

Costi

In questo documento utilizzi i seguenti componenti fatturabili di Google Cloud:

Per generare una stima dei costi basata sull'utilizzo previsto, utilizza il Calcolatore prezzi. I nuovi utenti di Google Cloud potrebbero essere idonei per una prova gratuita.

Al termine delle attività descritte in questo documento, puoi evitare la fatturazione continua eliminando le risorse che hai creato. Per ulteriori informazioni, consulta la sezione Pulizia.

Prima di iniziare

Cloud Shell è preinstallato con il software necessario per questo tutorial, tra cui kubectl e gcloud CLI. Se non utilizzi Cloud Shell, installa gcloud CLI.

  1. 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.
  2. Install the Google Cloud CLI.

  3. To initialize the gcloud CLI, run the following command: 

    gcloud init

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

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

  6. Enable the GKE API: 

    gcloud services enable container.googleapis.com
  7. Install the Google Cloud CLI.

  8. To initialize the gcloud CLI, run the following command: 

    gcloud init

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

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

  11. Enable the GKE API: 

    gcloud services enable container.googleapis.com

  12. 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:myemail@example.com.

    • Replace ROLE with each individual role.

Assicurati di disporre di una quota sufficiente

Assicurati che il tuo progetto Google Cloud disponga di una quota TPU sufficiente nella regione o nella zona Compute Engine. Per ulteriori informazioni, consulta Garantire quote TPU e GKE sufficienti nella documentazione di Cloud TPU. Potresti anche dover aumentare le quote per:

  • Disco permanente SSD (GB)
  • Indirizzi IP in uso

prepara l'ambiente

Per preparare l'ambiente:

  1. Avvia una sessione Cloud Shell dalla console Google Cloud facendo clic su Icona di attivazione di Cloud Shell Attiva Cloud Shell nella console Google Cloud. Viene avviata una sessione nel riquadro inferiore della console Google Cloud.

  2. Imposta le variabili di ambiente:

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

    Sostituisci quanto segue:

    • PROJECT_ID: il tuo ID progetto Google Cloud.
    • CLUSTER_VERSION: la versione di GKE da utilizzare. Deve essere 1.30.1 o successiva.

  3. Clona il repository GitHub:

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

  4. Passa alla directory di lavoro:

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

Creare un cluster con un pool di nodi TPU

Crea un cluster GKE standard con un pool di nodi TPU:

  1. Crea un cluster in modalità standard con l'operatore Ray abilitato:

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

  2. Crea un pool di nodi TPU con un solo host:

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

Per utilizzare le TPU con la modalità standard, devi selezionare:

  • Una località Compute Engine con capacità per gli acceleratori TPU
  • Un tipo di macchina compatibile per la TPU e
  • La topologia fisica della sezione di pod TPU

Configura una risorsa RayCluster con TPU

Configura il manifest di RayCluster per preparare il tuo carico di lavoro TPU:

Configura TPU nodeSelector

GKE utilizza i nodeSelectors di Kubernetes per garantire che i carichi di lavoro TPU vengano pianificati sull'acceleratore e sulla topologia TPU appropriati. Per ulteriori informazioni sulla selezione di nodeSelectors TPU, consulta Eseguire il deployment dei carichi di lavoro TPU in GKE Standard.

Aggiorna il file manifest ray-cluster.yaml per pianificare il pod su una sezione di pod TPU v4 con una topologia 2x2x1:

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

Configura una risorsa contenitore TPU

Per utilizzare un acceleratore TPU, devi specificare il numero di chip TPU che GKE deve allocare a ogni pod configurando le risorse google.com/tpulimits e requests nel campo del contenitore TPU del manifest RayCluster workerGroupSpecs.

Aggiorna il manifest ray-cluster.yaml con richieste e limiti di risorse:

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"

Configura il gruppo di worker numOfHosts

KubeRay 1.1.0 aggiunge un campo numOfHosts alla risorsa personalizzata RayCluster, che specifica il numero di host TPU da creare per replica del gruppo di worker. Per i gruppi di worker multi-host, le repliche vengono trattate come PodSlice anziché come singoli worker e vengono creati numOfHosts nodi worker per replica.

Aggiorna il file manifest di ray-cluster.yaml con quanto segue:

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

Crea una risorsa personalizzata RayService

Crea una risorsa personalizzata RayService:

  1. Esamina il seguente manifest:

    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

    Questo manifest descrive una risorsa personalizzata RayService che crea una risorsa RayCluster con 1 nodo principale e un gruppo di worker TPU con una topologia 2x2x1, il che significa che ogni nodo worker avrà 4 chip TPU v4.

    Il nodo TPU appartiene a una singola sezione di pod TPU v4 con una topologia 2x2x1. Per creare un gruppo di worker multi-host, sostituisci i valori gke-tpu nodeSelector, le richieste e i limiti dei container google.com/tpu e i valori numOfHosts con la tua configurazione multi-host. Per ulteriori informazioni sulle topologie multi-host TPU, consulta la sezione Architettura di sistema nella documentazione di Cloud TPU.

  2. Applica il manifest al cluster:

    kubectl apply -f ray-service-tpu.yaml

  3. Verifica che la risorsa RayService sia in esecuzione:

    kubectl get rayservices

    L'output è simile al seguente:

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

    In questo output, Running nella colonna SERVICE STATUS indica che la risorsa RayService è pronta.

(Facoltativo) Visualizza la dashboard di Ray

Puoi visualizzare il deployment di Ray Serve e i log pertinenti dalla dashboard di Ray.

  1. Stabilisci una sessione di port forwarding alla dashboard di Ray dal servizio Ray head:

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

  2. In un browser web, vai a http://localhost:8265/.

  3. Fai clic sulla scheda Pubblica.

Invia i prompt al server del modello

  1. Stabilisci una sessione di port forwarding all'endpoint Serve dal servizio Ray head:

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

  2. Apri una nuova sessione Cloud Shell.

  3. Invia un prompt di testo a immagine al server del modello di diffusione stabile:

    python stable_diffusion_tpu_req.py  --save_pictures

    I risultati dell'inferenza di diffusione stabile vengono salvati in un file denominato diffusion_results.png.

    Immagine generata da Diffusione stabile con 8 sezioni: una sedia verde, un uomo in piedi davanti a una casa, un robot per strada, una famiglia seduta a un tavolo, un dottore che cammina in un parco, un drago volante, un ritratto di orsi in stile giapponese e una cascata.

Esegui la pulizia

Elimina il progetto

    Delete a Google Cloud project:

    gcloud projects delete PROJECT_ID

Elimina singole risorse

Per eliminare il cluster, digita:

gcloud container clusters delete ${CLUSTER_NAME}

Passaggi successivi