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 pubblicare un modello Stable Diffusion Google Kubernetes Engine (GKE) tramite TPU, Ray Serve, e ai Componente aggiuntivo Ray Operator.

Questa guida è rivolta ai clienti dell'IA generativa, agli utenti nuovi o esistenti di GKE, ML engineer, MLOps (DevOps) engineer o piattaforme amministratori interessati a utilizzare l'orchestrazione dei container Kubernetes per fornire i modelli mediante Ray.

Informazioni su Ray and Ray Serve

Ray è un framework di calcolo open source scalabile per applicazioni di AI/ML. Raggio Serve è una libreria di gestione dei modelli per Ray, usata per scalare e distribuire i modelli in in un ambiente distribuito. Per ulteriori informazioni, vedi Ray Serve nel raggio documentazione.

Informazioni sulle TPU

Le TPU (Tensor Processing Unit) sono acceleratori hardware specializzati progettati per velocizzare notevolmente l'addestramento e l'inferenza dei modelli di machine learning. L'utilizzo di Ray con TPU consente di scalare senza problemi ad alte prestazioni per le applicazioni di ML. Per ulteriori informazioni sulle TPU, vedi Introduzione a Cloud TPU in Cloud TPU documentazione.

Informazioni sul webhook di inizializzazione TPU KubeRay

Come parte del componente aggiuntivo Ray Operator, GKE fornisce servizi di mutare webhooks che gestiscono la pianificazione dei pod di TPU e alcune variabili di ambiente TPU necessarie da framework come JAX per l'inizializzazione del container. KubeRay Il webhook TPU muta i pod con l'etichetta app.kubernetes.io/name: kuberay richiedendo 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 per comunicare tra loro all'interno della sezione. Questa variabile è solo inseriti per i pod TPU in un gruppo multi-host.
  • replicaIndex: un'etichetta pod contenente un identificatore univoco per la gruppo di worker a cui appartiene il pod. Questo è utile per il worker multi-host in cui più pod worker potrebbero appartenere alla stessa replica, utilizzata da Ray per abilitare la scalabilità automatica multi-host.
  • TPU_NAME: una stringa che rappresenta la sezione di pod TPU GKE di questo pod , impostato sullo stesso valore dell'etichetta replicaIndex.
  • podAffinity: assicura che GKE pianifichi i pod TPU con corrispondenza replicaIndex etichette nello stesso pool di nodi. In questo modo GKE Scalare le TPU multi-host atomicamente in base ai pool di nodi, invece che 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 un RayService risorsa personalizzata.
  • Interagire con il server del modello Stable Diffusion.

Costi

In questo documento vengono utilizzati 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.

Una volta completate le attività descritte in questo documento, puoi evitare la fatturazione continua eliminando le risorse che hai creato. Per ulteriori informazioni, consulta la pagina Pulizia.

Prima di iniziare

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

  1. 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 gratuiti per l'esecuzione, il test e il deployment dei carichi di lavoro.
  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. Assicurati che la fatturazione sia attivata per il tuo progetto Google Cloud.

  6. Attiva l'API GKE.

    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. Assicurati che la fatturazione sia attivata per il tuo progetto Google Cloud.

  11. Attiva l'API GKE.

    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_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 una quota sufficiente

Assicurati che il tuo progetto Google Cloud abbia una quota TPU sufficiente nella regione o nella zona di Compute Engine. Per ulteriori informazioni, vedi Garantire quote TPU e GKE sufficienti nella documentazione di Cloud TPU. Potresti anche dover aumentare 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. clic su Icona di attivazione di Cloud Shell Attiva Cloud Shell nella Console Google Cloud: Verrà 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 account Google Cloud ID progetto.
    • CLUSTER_VERSION: la versione GKE da per gli utilizzi odierni. Deve essere 1.30.1 o una versione 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 Autopilot o Standard Cluster GKE con un pool di nodi TPU:

Autopilot

Crea un cluster in modalità Autopilot con l'operatore Ray abilitato:

gcloud container clusters create-auto ${CLUSTER_NAME} \
    --enable-ray-operator \
    --cluster-version=${CLUSTER_VERSION} \
    --location=${COMPUTE_REGION}

Per utilizzare le TPU con la modalità Autopilot, devi selezionare un'opzione Località di Compute Engine con capacità per gli acceleratori TPU. TPU v4 sono disponibili in us-central2-b. Per ulteriori informazioni sulla disponibilità delle TPU per regione/zona, vedi Informazioni sulle TPU in GKE.

Standard

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

    gcloud container clusters create ${CLUSTER_NAME} \
        --addons=RayOperator \
        --cluster-version=${CLUSTER_VERSION} \
        --location=${COMPUTE_REGION}
    
  2. Crea un pool di nodi TPU a singolo host:

    gcloud container node-pools create ${NODEPOOL_NAME} \
        --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à di Compute Engine con capacità per gli acceleratori TPU
  • Un tipo di macchina compatibile con TPU e
  • La topologia fisica della sezione di pod TPU

Configura una risorsa RayCluster con TPU

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

Configura TPU nodeSelector

GKE usa Kubernetes nodeSelectors per assicurare che i carichi di lavoro TPU siano pianificati sull'acceleratore e sulla topologia TPU appropriati. Per ulteriori informazioni sulla selezione di 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 2 x 2 x 1:

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 google.com/tpurisorsa limits e requests nel campo del container TPU del tuo manifest RayCluster workerGroupSpecs.

Aggiorna il manifest ray-cluster.yaml con limiti e richieste per le 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 v1.1.0 aggiunge un campo numOfHosts alla risorsa personalizzata RayCluster, che specifica il numero di host TPU da creare per ogni replica del gruppo di worker. Per gruppi di worker multi-host, le repliche vengono trattate come PodSlice anziché singoli worker, con numOfHosts nodi worker creati per ogni replica.

Aggiorna il manifest 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
    spec:
      serveConfigV2: |
        applications:
          - name: stable_diffusion
            import_path: ai-ml.gke-ray.rayserve.stable-diffusion.stable_diffusion:entrypoint
            runtime_env:
              working_dir: "https://github.com/GoogleCloudPlatform/kubernetes-engine-samples/archive/main.zip"
              pip: ["diffusers==0.12.1"]
      rayClusterConfig:
        rayVersion: '2.9.0'
        headGroupSpec:
          rayStartParams:
            dashboard-host: '0.0.0.0'
          template:
            spec:
              containers:
              - name: ray-head
                image: rayproject/ray-ml:2.9.0
                ports:
                - containerPort: 6379
                  name: gcs
                - containerPort: 8265
                  name: dashboard
                - containerPort: 10001
                  name: client
                - containerPort: 8000
                  name: serve
                resources:
                  limits:
                    cpu: "2"
                    memory: "8Gi"
                  requests:
                    cpu: "2"
                    memory: "8Gi"
        workerGroupSpecs:
        - replicas: 1
          minReplicas: 1
          maxReplicas: 4
          groupName: gpu-group
          rayStartParams: {}
          template:
            spec:
              containers:
              - name: ray-worker
                image: rayproject/ray-ml:2.9.0
                resources:
                  limits:
                    cpu: 4
                    memory: "16Gi"
                    nvidia.com/gpu: 1
                  requests:
                    cpu: 3
                    memory: "16Gi"
                    nvidia.com/gpu: 1
              nodeSelector:
                cloud.google.com/gke-accelerator: nvidia-l4

    Questo file manifest descrive una risorsa personalizzata RayService che crea un'istanza Risorsa RayCluster con 1 nodo head e un gruppo di worker TPU con un nodo 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. A crea un gruppo di worker multi-host, sostituisci i valori gke-tpu nodeSelector, google.com/tpu limiti e richieste di container e numOfHosts valori con la configurazione multi-host. Per ulteriori informazioni su TPU multi-host , vedi le topologie Architettura del sistema 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 la La risorsa RayService è pronta.

(Facoltativo) Visualizzare la dashboard Ray

Puoi visualizzare il deployment di Ray Serve e i log pertinenti da Ray Dashboard.

  1. Stabilire una sessione di port forwarding alla dashboard Ray dall'head Ray servizio:

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

Invia prompt al server del modello

  1. Stabilisci una sessione di port forwarding all'endpoint di servizio dall'intestazione Ray servizio:

    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 Stable Diffusion:

    python stable_diffusion_tpu_req.py  --save_pictures
    

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

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

Esegui la pulizia

Elimina il progetto

    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