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

Standard

Questa guida mostra come eseguire il deployment e pubblicare un modello di diffusione stabile su Google Kubernetes Engine (GKE) utilizzando le TPU, Ray Serve, e il componente aggiuntivo Operatore Ray.

Questa guida è destinata ai clienti di AI generativa, agli utenti nuovi o esistenti di GKE, agli ingegneri ML, agli ingegneri MLOps (DevOps) o agli amministratori di piattaforma interessati a utilizzare le funzionalità di orchestrazione dei container Kubernetes per erogare 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 erogazione del modello per Ray utilizzata per scalare ed erogare i modelli in un ambiente distribuito. Per saperne di più, consulta Ray Serve nella documentazione di Ray.

Informazioni sulle TPU

Le Tensor Processing Unit (TPU) sono acceleratori hardware specializzati progettati per velocizzare 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 saperne di più sulle TPU, consulta la sezione Introduzione a Cloud TPU nella documentazione di Cloud TPU.

Informazioni sul webhook di inizializzazione TPU di KubeRay

Nell'ambito del componente aggiuntivo Ray Operator, GKE fornisce webhook di convalida e mutazione che gestiscono la pianificazione dei pod di TPU e determinate variabili di ambiente TPU richieste da framework come JAX per l'inizializzazione dei container. L'webhook TPU KubeRay modifica 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 worker nella sezione TPU.
• TPU_WORKER_HOSTNAMES: un elenco di nomi host DNS per tutti i worker TPU che devono comunicare tra loro all'interno dello slice. Questa variabile viene inserita solo per i pod TPU in un gruppo multi-host.
• replicaIndex: un'etichetta pod che contiene un identificatore univoco per la replica del gruppo di worker a cui appartiene il pod. Ciò è utile per i gruppi di worker multi-host, in cui più pod worker potrebbero appartenere alla stessa replica, e viene utilizzato da Ray per abilitare la scalabilità automatica multi-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 pianifichi i pod TPU con etichette replicaIndex corrispondenti nello stesso pool di nodi. In questo modo, GKE scala le TPU multi-host in modo atomico per node pool, anziché per singoli nodi.

Obiettivi

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

Costi

In questo documento vengono utilizzati i seguenti componenti fatturabili di Google Cloud:

• GKE

Per generare una stima dei costi in base all'utilizzo previsto, utilizza il calcolatore prezzi.

I nuovi utenti di Google Cloud potrebbero avere diritto a una prova senza costi.

Al termine delle attività descritte in questo documento, puoi evitare l'addebito di ulteriori costi eliminando le risorse che hai creato. Per saperne di più, consulta Esegui la 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.

Accedi al tuo account Google Cloud . Se non conosci Google Cloud, crea un account per valutare le prestazioni dei nostri prodotti in scenari reali. I nuovi clienti ricevono anche 300 $di crediti senza costi per l'esecuzione, il test e il deployment dei workload.

Installa Google Cloud CLI.

Nota: se hai già installato gcloud CLI, assicurati di avere l'ultima versione eseguendo gcloud components update.

Se utilizzi un provider di identità (IdP) esterno, devi prima accedere a gcloud CLI con la tua identità federata.

Per inizializzare gcloud CLI, esegui questo comando:

gcloud init

Crea o seleziona un Google Cloud progetto.

Ruoli richiesti per selezionare o creare un progetto

• Seleziona un progetto: la selezione di un progetto non richiede un ruolo IAM specifico. Puoi selezionare qualsiasi progetto per il quale ti è stato concesso un ruolo.
• Crea un progetto: per creare un progetto, devi disporre del ruolo Autore progetto (roles/resourcemanager.projectCreator), che contiene l'autorizzazione resourcemanager.projects.create. Scopri come concedere i ruoli.

Nota: se non prevedi di conservare le risorse che crei in questa procedura, crea un progetto invece di selezionarne uno già esistente. Una volta completata questa procedura, puoi eliminare il progetto e tutte le relative risorse.

• Creare un progetto Google Cloud :

  gcloud projects create PROJECT_ID

  Sostituisci PROJECT_ID con un nome per il progetto Google Cloud che stai creando.

• Seleziona il progetto Google Cloud che hai creato:

  gcloud config set project PROJECT_ID

  Sostituisci PROJECT_ID con il nome del progetto Google Cloud .

Verifica che la fatturazione sia attivata per il tuo progetto Google Cloud .

Abilita l'API GKE:

Ruoli richiesti per abilitare le API

Per abilitare le API, devi disporre del ruolo IAM Amministratore utilizzo dei servizi (roles/serviceusage.serviceUsageAdmin), che include l'autorizzazione serviceusage.services.enable. Scopri come concedere i ruoli.

gcloud services enable container.googleapis.com

Installa Google Cloud CLI.

Nota: se hai già installato gcloud CLI, assicurati di avere l'ultima versione eseguendo gcloud components update.

Se utilizzi un provider di identità (IdP) esterno, devi prima accedere a gcloud CLI con la tua identità federata.

Per inizializzare gcloud CLI, esegui questo comando:

gcloud init

Crea o seleziona un Google Cloud progetto.

Ruoli richiesti per selezionare o creare un progetto

• Seleziona un progetto: la selezione di un progetto non richiede un ruolo IAM specifico. Puoi selezionare qualsiasi progetto per il quale ti è stato concesso un ruolo.
• Crea un progetto: per creare un progetto, devi disporre del ruolo Autore progetto (roles/resourcemanager.projectCreator), che contiene l'autorizzazione resourcemanager.projects.create. Scopri come concedere i ruoli.

Nota: se non prevedi di conservare le risorse che crei in questa procedura, crea un progetto invece di selezionarne uno già esistente. Una volta completata questa procedura, puoi eliminare il progetto e tutte le relative risorse.

• Creare un progetto Google Cloud :

  gcloud projects create PROJECT_ID

  Sostituisci PROJECT_ID con un nome per il progetto Google Cloud che stai creando.

• Seleziona il progetto Google Cloud che hai creato:

  gcloud config set project PROJECT_ID

  Sostituisci PROJECT_ID con il nome del progetto Google Cloud .

Verifica che la fatturazione sia attivata per il tuo progetto Google Cloud .

Abilita l'API GKE:

Ruoli richiesti per abilitare le API

Per abilitare le API, devi disporre del ruolo IAM Amministratore utilizzo dei servizi (roles/serviceusage.serviceUsageAdmin), che include l'autorizzazione serviceusage.services.enable. Scopri come concedere i ruoli.

gcloud services enable container.googleapis.com

1. Concedi ruoli al tuo account utente. Esegui il seguente comando una volta per ciascuno dei seguenti ruoli IAM: roles/container.clusterAdmin, roles/container.admin

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

   Sostituisci quanto segue:

   • PROJECT_ID: il tuo ID progetto.
   • USER_IDENTIFIER: l'identificatore del tuo account utente . Ad esempio: [e-mail].
   • ROLE: il ruolo IAM che concedi al tuo account utente.

Assicurati di avere una quota sufficiente

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

• SSD Persistent Disk (GB)
• Indirizzi IP in uso

prepara l'ambiente

Per preparare l'ambiente:

1. Avvia una sessione di Cloud Shell dalla console Google Cloud facendo clic su Attiva Cloud Shell nella consoleGoogle 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 Google Cloud ID progetto.
   • 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
   

Crea 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 single-host:

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

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 le TPU

Configura il manifest RayCluster per preparare il workload TPU:

Configura TPU nodeSelector

GKE utilizza i nodeSelectors di Kubernetes per garantire che i workload TPU vengano pianificati sulla topologia e sull'acceleratore TPU appropriati. Per saperne di più sulla selezione dei nodeSelector TPU, consulta Esegui il deployment dei carichi di lavoro TPU in GKE Standard.

Aggiorna il 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 container TPU

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

Aggiorna il manifest ray-cluster.yaml con limiti e richieste 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 gruppo di worker numOfHosts

KubeRay v1.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 worker individuali, con la creazione di numOfHosts nodi worker per replica.

Aggiorna il 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 head 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 multihost, sostituisci i valori gke-tpu nodeSelector, i limiti e le richieste dei container google.com/tpu e i valori numOfHosts con la tua configurazione multihost. Per ulteriori informazioni sulle topologie TPU multi-host, 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 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 Ray dal servizio head Ray:

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

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

3. Fai clic sulla scheda Pubblica.

Inviare prompt al server del modello

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

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

2. Apri una nuova sessione di Cloud Shell.

3. Invia un prompt da 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.

   

Osserva i tuoi carichi di lavoro Ray

Per visualizzare i dettagli dei tuoi RayJob, puoi andare alla sezione Kubernetes Engine > AI/ML > Job della console Google Cloud .

Visualizza RayJobs nella Google Cloud console

Esegui la pulizia

Elimina il progetto

Attenzione: l'eliminazione di un progetto ha i seguenti effetti:

• L'intero contenuto del progetto viene eliminato. Se hai utilizzato un progetto esistente per, le attività descritte in questo documento, eliminerai anche tutto il lavoro che hai svolto nel progetto.
• Gli ID progetto personalizzati non sono più disponibili. Quando hai creato questo progetto, potresti aver creato un ID progetto personalizzato che vuoi utilizzare in futuro. Per conservare gli URL che utilizzano l'ID progetto, ad esempio un URL appspot.com, elimina le risorse selezionate all'interno del progetto anziché eliminare l'intero progetto.

Se intendi esplorare più architetture, tutorial o guide rapide, puoi riutilizzare i progetti ed evitare così di superare i limiti di quota.

Elimina un progetto Google Cloud :

gcloud projects delete PROJECT_ID

Elimina singole risorse

Per eliminare il cluster, digita:

gcloud container clusters delete ${CLUSTER_NAME}

Passaggi successivi

• Scopri di più su Ray su Kubernetes.
• Esplora la documentazione di KubeRay.
• Esplora architetture, diagrammi e best practice di riferimento su Google Cloud. Consulta il nostro Cloud Architecture Center.