Ray Serve-Anwendung mit einem Stable Diffusion-Modell in Google Kubernetes Engine (GKE) mit TPUs bereitstellen

In diesem Leitfaden wird beschrieben, wie Sie ein Stable Diffusion-Modell mithilfe von TPUs, Ray Serve und dem Ray Operator-Add-on in Google Kubernetes Engine (GKE) bereitstellen und ausführen.

Dieser Leitfaden richtet sich an Kunden von generativer KI, neue oder bestehende Nutzer von GKE, ML-Entwickler, MLOps-Entwickler (DevOps) oder Plattformadministratoren, die daran interessiert sind, Funktionen zur Kubernetes-Containerorchestrierung für die Bereitstellung von Modellen mit Ray zu nutzen.

Ray und Ray Serve

Ray ist ein skalierbares Open-Source-Computing-Framework für KI/ML-Anwendungen. Ray Serve ist eine Modellbereitstellungsbibliothek für Ray, die zur Skalierung und Bereitstellung von Modellen in einer verteilten Umgebung verwendet wird. Weitere Informationen finden Sie unter Ray Serve in der Ray-Dokumentation.

TPUs

Tensor Processing Units (TPUs) sind spezielle Hardwarebeschleuniger, die das Training und die Inferenz von großen Modellen für maschinelles Lernen erheblich beschleunigen. Mit Rays und TPUs können Sie leistungsstarke ML-Anwendungen nahtlos skalieren. Weitere Informationen zu TPUs finden Sie unter Einführung in Cloud TPU in der Cloud TPU-Dokumentation.

KubeRay-Webhook zur TPU-Initialisierung

Als Teil des Ray Operator-Add-ons bietet GKE das Validieren und Mutieren von Webhooks, die die TPU-Pod-Planung und bestimmte TPU-Umgebungsvariablen verarbeiten, die für Frameworks wie JAX für die Containerinitialisierung erforderlich sind. Der KubeRay TPU-Webhook mutiert Pods mit dem Label app.kubernetes.io/name: kuberay und fordert TPUs mit den folgenden Attributen an:

  • TPU_WORKER_ID: Eine eindeutige Ganzzahl für jeden Worker-Pod im TPU-Slice.
  • TPU_WORKER_HOSTNAMES: Eine Liste von DNS-Hostnamen für alle TPU-Worker, die innerhalb des Slices miteinander kommunizieren müssen. Diese Variable wird nur für TPU-Pods in einer Gruppe mit mehreren Hosts eingefügt.
  • replicaIndex: Ein Pod-Label, das eine eindeutige Kennung für das Replikat der Worker-Gruppe enthält, zu der der Pod gehört. Dies ist nützlich für Arbeitsgruppen mit mehreren Hosts, bei denen mehrere Worker-Pods zum selben Replikat gehören können. Ray verwendet diese Funktion, um das Multi-Host-Autoscaling zu ermöglichen.
  • TPU_NAME: Ein String, der den GKE TPU-PodSlice darstellt, zu dem dieser Pod gehört, und für den gleichen Wert wie das Label replicaIndex festgelegt.
  • podAffinity: Sorgt dafür, dass GKE TPU-Pods mit übereinstimmenden replicaIndex-Labels im selben Knotenpool plant. Dadurch kann GKE TPUs mit mehreren Hosts atomar nach Knotenpools statt nach einzelnen Knoten skalieren.

Ziele

  • Erstellen Sie einen GKE-Cluster mit einem TPU-Knotenpool.
  • Stellen Sie einen Ray-Cluster mit TPUs bereit.
  • Benutzerdefinierte RayService bereitstellen
  • Interagieren Sie mit dem Stable Diffusion-Modellserver.

Kosten

In diesem Dokument verwenden Sie die folgenden kostenpflichtigen Komponenten von Google Cloud:

Mit dem Preisrechner können Sie eine Kostenschätzung für Ihre voraussichtliche Nutzung vornehmen. Neuen Google Cloud-Nutzern steht möglicherweise eine kostenlose Testversion zur Verfügung.

Nach Abschluss der in diesem Dokument beschriebenen Aufgaben können Sie weitere Kosten vermeiden, indem Sie die erstellten Ressourcen löschen. Weitere Informationen finden Sie unter Bereinigen.

Hinweise

Die Software, die Sie für diese Anleitung benötigen, ist in Cloud Shell vorinstalliert, einschließlich kubectl und der gcloud CLI. Wenn Sie Cloud Shell nicht verwenden, installieren Sie die 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_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.

Ausreichendes Kontingent sicherstellen

Prüfen Sie, ob Ihr Google Cloud-Projekt über ein ausreichendes TPU-Kontingent in Ihrer Compute Engine-Region oder -Zone verfügt. Weitere Informationen finden Sie in der Cloud TPU-Dokumentation unter Ausreichende TPU- und GKE-Kontingente bereitstellen. Möglicherweise müssen Sie Ihre Kontingente auch für Folgendes erhöhen:

  • Nichtflüchtiger SSD-Speicher (GB):
  • IP-Adressen in Verwendung

Umgebung vorbereiten

So bereiten Sie die Umgebung vor:

  1. Starten Sie eine Cloud Shell-Sitzung über die Google Cloud Console. Klicken Sie dazu in der Google Cloud Console auf Symbol für die Cloud Shell-Aktivierung Cloud Shell aktivieren. Dadurch wird im unteren Bereich der Google Cloud Console eine Sitzung gestartet.

  2. Legen Sie Umgebungsvariablen fest:

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

    Ersetzen Sie Folgendes:

    • PROJECT_ID: Ihre Google Cloud-Projekt-ID.
    • CLUSTER_VERSION ist die zu verwendende GKE-Version. Es muss 1.30.1 oder höher sein.

  3. Klonen Sie das GitHub-Repository:

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

  4. Wechseln Sie in das Arbeitsverzeichnis:

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

Cluster mit TPU-Knotenpool erstellen

Erstellen Sie einen Standard-GKE-Cluster mit einem TPU-Knotenpool:

  1. Erstellen Sie einen Cluster im Standardmodus mit aktiviertem Ray-Operator:

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

  2. So erstellen Sie einen TPU-Knotenpool mit einem 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

Wenn Sie TPUs mit dem Standardmodus verwenden möchten, müssen Sie Folgendes auswählen:

  • Ein Compute Engine-Standort mit Kapazität für TPU-Beschleuniger
  • Einen kompatiblen Maschinentyp für die TPU und
  • Die physische Topologie des TPU-Pod-Slices

RayCluster-Ressource mit TPUs konfigurieren

Konfigurieren Sie das RayCluster-Manifest, um Ihre TPU-Arbeitslast vorzubereiten:

TPU nodeSelector konfigurieren

GKE verwendet Kubernetes-nodeSelectors, um dafür zu sorgen, dass TPU-Arbeitslasten auf der entsprechenden TPU-Topologie und dem entsprechenden Beschleuniger geplant werden. Weitere Informationen zum Auswählen von TPU nodeSelectors finden Sie unter TPU-Arbeitslasten in GKE Standard bereitstellen.

Aktualisieren Sie das Manifest ray-cluster.yaml, um Ihren Pod auf einem v4-TPU-Podslice mit einer 2x2x1-Topologie zu planen:

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

TPU-Containerressource konfigurieren

Wenn Sie einen TPU-Beschleuniger verwenden möchten, müssen Sie die Anzahl der TPU-Chips angeben, die GKE jedem Pod zuweisen soll. Konfigurieren Sie dazu die google.com/tpu-Ressource limits und requests im TPU-Containerfeld Ihres RayCluster-Manifests workerGroupSpecs.

Aktualisieren Sie das ray-cluster.yaml-Manifest mit Ressourcenlimits und ‑anfragen:

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"

Worker-Gruppe numOfHosts konfigurieren

KubeRay v1.1.0 fügt der benutzerdefinierten RayCluster-Ressource das Feld numOfHosts hinzu, das die Anzahl der TPU-Hosts angibt, die pro Worker-Gruppen-Replikat erstellt werden sollen. Bei Worker-Gruppen mit mehreren Hosts werden Replikate als PodSlices und nicht als einzelne Worker behandelt. Pro Replikate werden numOfHosts Worker-Knoten erstellt.

Aktualisieren Sie das Manifest ray-cluster.yaml so:

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

Benutzerdefinierte RayService-Ressource erstellen

Erstellen Sie eine benutzerdefinierte RayService-Ressource:

  1. Prüfen Sie das folgende 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

    Dieses Manifest beschreibt eine benutzerdefinierte RayService-Ressource, die eine RayCluster-Ressource mit einem Hauptknoten und einer TPU-Worker-Gruppe mit einer 2x2x1-Topologie erstellt. Dies bedeutet, dass jeder Worker-Knoten 4 v4-TPU-Chips hat.

    Der TPU-Knoten gehört zu einem einzelnen v4-TPU-Podslice mit einer 2x2x1-Topologie. Wenn Sie eine Multi-Host-Worker-Gruppe erstellen möchten, ersetzen Sie die gke-tpu nodeSelector-Werte, die google.com/tpu-Containerlimits und ‑anfragen sowie die numOfHosts-Werte durch Ihre Multi-Host-Konfiguration. Weitere Informationen zu TPU-Topologien mit mehreren Hosts finden Sie in der Cloud TPU-Dokumentation unter Systemarchitektur.

  2. Wenden Sie das Manifest auf Ihren Cluster an:

    kubectl apply -f ray-service-tpu.yaml

  3. Prüfen Sie, ob die RayService-Ressource ausgeführt wird:

    kubectl get rayservices

    Die Ausgabe sieht in etwa so aus:

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

    In dieser Ausgabe gibt Running in der Spalte SERVICE STATUS an, dass die RayService-Ressource bereit ist.

(Optional) Ray-Dashboard aufrufen

Sie können Ihre Ray Serve-Bereitstellung und relevante Logs im Ray-Dashboard aufrufen.

  1. Richten Sie eine Portweiterleitungssitzung vom Ray-Head-Dienst zum Ray-Dashboard ein:

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

  2. Rufe in einem Webbrowser folgende Seite auf:http://localhost:8265/

  3. Klicken Sie auf den Tab Bereitstellen.

Prompts an den Modellserver senden

  1. Richten Sie eine Portweiterleitungssitzung zum Serve-Endpunkt aus dem Ray-Head-Dienst ein:

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

  2. Öffnen Sie eine neue Cloud Shell-Sitzung.

  3. Senden Sie einen Text-zu-Bild-Prompt an den Server des Stable Diffusion-Modells:

    python stable_diffusion_tpu_req.py  --save_pictures

    Die Ergebnisse der stabilen Diffusionsinferenz werden in einer Datei mit dem Namen diffusion_results.png gespeichert.

    Von Stable Diffusion das mit acht Abschnitten generierte Bild: ein grüner Stuhl, ein Mann vor einem Haus, ein Roboter auf der Straße, eine Familie an einem Tisch, ein Dokument, das im Park läuft, ein fließender Drain, ein Porträt von Bären und ein Wasserfall.

Bereinigen

Projekt löschen

    Delete a Google Cloud project:

    gcloud projects delete PROJECT_ID

Einzelne Ressourcen löschen

Zum Löschen des Clusters geben Sie Folgendes ein:

gcloud container clusters delete ${CLUSTER_NAME}

Nächste Schritte