Men-deploy aplikasi Ray Serve dengan model Stable Diffusion di Google Kubernetes Engine (GKE) dengan TPU


Panduan ini menunjukkan cara menerapkan dan menyajikan model Difusi Stabil di Google Kubernetes Engine (GKE) dengan TPU, Ray Serve, dan Add-on Ray Operator.

Panduan ini ditujukan untuk pelanggan AI Generatif, pengguna baru atau lama GKE, ML Engineer, engineer MLOps (DevOps), atau platform administrator yang tertarik menggunakan orkestrasi container Kubernetes untuk menyajikan model yang menggunakan Ray.

Tentang Ray and Ray Serve

Ray adalah framework komputasi open source yang skalabel untuk aplikasi AI/ML. Sinar Serve adalah library penyaluran model untuk Ray yang digunakan untuk penskalaan dan model inferensi di lingkungan yang terdistribusi. Untuk informasi selengkapnya, lihat Penyaluran Sinar di Sinar dokumentasi tambahan.

Tentang TPU

Tensor Processing Unit (TPU) adalah akselerator hardware khusus yang dirancang untuk secara signifikan mempercepat pelatihan dan inferensi mesin skala besar machine learning. Menggunakan Ray dengan TPU memungkinkan Anda menskalakan performa tinggi dengan lancar pada aplikasi ML. Untuk mengetahui informasi selengkapnya tentang TPU, lihat Pengantar Cloud TPU di Cloud TPU dokumentasi tambahan.

Tentang webhook inisialisasi TPU KubeRay

Sebagai bagian dari add-on Ray Operator, GKE menyediakan alat validasi bermutasi webhooks yang menangani penjadwalan Pod TPU dan variabel lingkungan TPU tertentu yang diperlukan oleh kerangka kerja seperti JAX untuk inisialisasi container. KubeRay Webhook TPU mengubah Pod dengan label app.kubernetes.io/name: kuberay meminta TPU dengan properti berikut:

  • TPU_WORKER_ID: Bilangan bulat unik untuk setiap Pod pekerja dalam slice TPU.
  • TPU_WORKER_HOSTNAMES: Daftar nama host DNS untuk semua pekerja TPU yang memerlukan untuk berkomunikasi satu sama lain di dalam slice. Variabel ini hanya diinjeksi untuk Pod TPU dalam grup multi-host.
  • replicaIndex: Label Pod yang berisi ID unik untuk replika grup pekerja tempat Pod berada. Hal ini berguna untuk pekerja multi-host grup, di mana beberapa Pod pekerja mungkin berasal dari replika yang sama, dan yang digunakan Ray untuk mengaktifkan penskalaan otomatis multi-host.
  • TPU_NAME: String yang mewakili GKE TPU PodSlice Pod ini berada di dalamnya, yang ditetapkan ke nilai yang sama dengan label replicaIndex.
  • podAffinity: Memastikan GKE menjadwalkan Pod TPU dengan pencocokan replicaIndex label pada kumpulan node yang sama. Hal ini memungkinkan GKE menskalakan TPU multi-host secara atomik berdasarkan kumpulan node, bukan node tunggal.

Tujuan

  • Membuat cluster GKE dengan kumpulan node TPU.
  • Deploy cluster Ray dengan TPU.
  • Men-deploy RayService resource kustom.
  • Berinteraksi dengan server model Stabil Diffusion.

Biaya

Dalam dokumen ini, Anda menggunakan komponen Google Cloud yang dapat ditagih berikut:

Untuk membuat perkiraan biaya berdasarkan proyeksi penggunaan Anda, gunakan kalkulator harga. Pengguna baru Google Cloud mungkin memenuhi syarat untuk mendapatkan uji coba gratis.

Setelah menyelesaikan tugas yang dijelaskan dalam dokumen ini, Anda dapat menghindari penagihan berkelanjutan dengan menghapus resource yang Anda buat. Untuk mengetahui informasi selengkapnya, lihat Pembersihan.

Sebelum memulai

Cloud Shell diprainstal dengan software yang Anda perlukan untuk ini yang sama, termasuk kubectl, dan gcloud CLI. Jika Anda tidak menggunakan Cloud Shell, menginstal gcloud CLI.

  1. Login ke akun Google Cloud Anda. Jika Anda baru menggunakan Google Cloud, buat akun untuk mengevaluasi performa produk kami dalam skenario dunia nyata. Pelanggan baru juga mendapatkan kredit gratis senilai $300 untuk menjalankan, menguji, dan men-deploy workload.
  2. Menginstal Google Cloud CLI.
  3. Untuk initialize gcloud CLI, jalankan perintah berikut:

    gcloud init
  4. Buat atau pilih project Google Cloud.

    • Membuat project Google Cloud:

      gcloud projects create PROJECT_ID

      Ganti PROJECT_ID dengan nama untuk project Google Cloud yang Anda buat.

    • Pilih project Google Cloud yang Anda buat:

      gcloud config set project PROJECT_ID

      Ganti PROJECT_ID dengan nama project Google Cloud Anda.

  5. Pastikan penagihan telah diaktifkan untuk project Google Cloud Anda.

  6. Enable the GKE API:

    gcloud services enable container.googleapis.com
  7. Menginstal Google Cloud CLI.
  8. Untuk initialize gcloud CLI, jalankan perintah berikut:

    gcloud init
  9. Buat atau pilih project Google Cloud.

    • Membuat project Google Cloud:

      gcloud projects create PROJECT_ID

      Ganti PROJECT_ID dengan nama untuk project Google Cloud yang Anda buat.

    • Pilih project Google Cloud yang Anda buat:

      gcloud config set project PROJECT_ID

      Ganti PROJECT_ID dengan nama project Google Cloud Anda.

  10. Pastikan penagihan telah diaktifkan untuk project Google Cloud Anda.

  11. Enable the GKE API:

    gcloud services enable container.googleapis.com
  12. Berikan peran ke Akun Google Anda. Jalankan perintah berikut satu kali untuk setiap peran IAM berikut: roles/container.clusterAdmin, roles/container.admin

    gcloud projects add-iam-policy-binding PROJECT_ID --member="user:EMAIL_ADDRESS" --role=ROLE
    • Ganti PROJECT_ID dengan project ID Anda.
    • Ganti EMAIL_ADDRESS dengan alamat email Anda.
    • Ganti ROLE dengan setiap peran individual.

Memastikan kuota memadai

Pastikan project Google Cloud Anda memiliki kuota TPU yang cukup di Region atau zona Compute Engine. Untuk informasi selengkapnya, lihat Memastikan kuota TPU dan GKE yang memadai di dokumentasi Cloud TPU. Anda mungkin juga perlu meningkatkan kuota untuk:

  • SSD Persistent Disk (GB)
  • Alamat IP yang digunakan

Menyiapkan lingkungan Anda

Untuk menyiapkan lingkungan Anda, ikuti langkah-langkah berikut:

  1. Luncurkan sesi Cloud Shell dari konsol Google Cloud, dengan mengklik Ikon aktivasi Cloud Shell Aktifkan Cloud Shell di konsol Google Cloud. Tindakan ini akan meluncurkan sesi di panel bawah konsol Google Cloud.

  2. Menetapkan variabel lingkungan:

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

    Ganti kode berikut:

    • PROJECT_ID: project ID Google Cloud Anda.
    • CLUSTER_VERSION: versi GKE untuk gunakan. Harus 1.30.1 atau lebih baru.
  3. Buat clone repositori GitHub:

    git clone https://github.com/GoogleCloudPlatform/kubernetes-engine-samples
    
  4. Ubah ke direktori kerja:

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

Membuat cluster dengan kumpulan node TPU

Buat Autopilot atau Standar Cluster GKE dengan kumpulan node TPU:

Autopilot

Buat cluster mode Autopilot dengan Ray Operator diaktifkan:

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

Untuk menggunakan TPU dengan mode Autopilot, Anda harus memilih Lokasi Compute Engine dengan kapasitas untuk akselerator TPU. TPU v4 tersedia di us-central2-b. Untuk mengetahui informasi selengkapnya tentang ketersediaan TPU menurut wilayah/zona, lihat Tentang TPU di GKE.

Standard

  1. Buat cluster mode Standar dengan Operator Ray yang diaktifkan:

    gcloud container clusters create ${CLUSTER_NAME} \
        --addons=RayOperator \
        --cluster-version=${CLUSTER_VERSION} \
        --location=${COMPUTE_REGION}
    
  2. Buat kumpulan node TPU host tunggal:

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

Untuk menggunakan TPU dengan mode Standar, Anda harus memilih:

  • Lokasi Compute Engine dengan kapasitas untuk akselerator TPU
  • Jenis mesin yang kompatibel untuk TPU dan
  • Topologi fisik TPU PodSlice

Mengonfigurasi resource RayCluster dengan TPU

Konfigurasikan manifes RayCluster untuk menyiapkan beban kerja TPU Anda:

Konfigurasi TPU nodeSelector

GKE menggunakan Kubernetes nodeSelectors untuk memastikan workload TPU dijadwalkan pada topologi dan akselerator TPU yang sesuai. Untuk mengetahui informasi selengkapnya tentang cara memilih TPU nodeSelectors, lihat Men-deploy workload TPU di GKE Standard.

Mengupdate manifes ray-cluster.yaml untuk menjadwalkan Pod Anda di podlice TPU v4 dengan topologi 2x2x1:

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

Mengonfigurasi resource container TPU

Untuk menggunakan akselerator TPU, Anda harus menentukan jumlah chip TPU yang GKE harus dialokasikan ke setiap Pod dengan mengonfigurasi google.com/tpuresource limits dan requests di kolom container TPU dari manifes RayCluster workerGroupSpecs.

Update manifes ray-cluster.yaml dengan permintaan dan batas resource:

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"

Konfigurasi grup pekerja numOfHosts

KubeRay v1.1.0 menambahkan kolom numOfHosts ke resource kustom RayCluster, yang menentukan jumlah host TPU yang akan dibuat per replika grup pekerja. Untuk grup pekerja multi-host, replika diperlakukan sebagai PodSlice, bukan setiap pekerja, dengan numOfHosts node pekerja dibuat per replika.

Update manifes ray-cluster.yaml dengan kode berikut:

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

Membuat resource kustom RayService

Buat resource kustom RayService:

  1. Tinjau manifes berikut:

    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:
                - pydantic<2
                - google-api-python-client
                - pillow
                - diffusers==0.7.2
                - transformers==4.24.0
                - flax
                - ml_dtypes==0.2.0
                - 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: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: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

    Manifes ini menjelaskan sumber daya khusus RayService yang membuat Resource RayCluster dengan 1 node kepala dan kelompok pekerja TPU dengan topologi, yang berarti setiap node pekerja akan memiliki 4 chip TPU v4.

    Node TPU termasuk dalam podlice TPU v4 tunggal dengan topologi 2x2x1. Kepada membuat grup pekerja multi-host, mengganti nilai gke-tpu nodeSelector, Batas dan permintaan container google.com/tpu, serta nilai numOfHosts dengan konfigurasi multi-host Anda. Untuk mengetahui informasi selengkapnya tentang multi-host TPU topologi, lihat Arsitektur sistem di dokumentasi Cloud TPU.

  2. Terapkan manifes ke cluster Anda:

    kubectl apply -f ray-service-tpu.yaml
    
  3. Pastikan resource RayService sedang berjalan:

    kubectl get rayservices
    

    Outputnya mirip dengan hal berikut ini:

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

    Dalam output ini, Running di kolom SERVICE STATUS menunjukkan Resource RayService sudah siap.

(Opsional) Melihat Dasbor Ray

Anda bisa melihat deployment Ray Serve dan log yang relevan dari Ray Dashboard.

  1. Membuat sesi penerusan port ke dasbor Ray dari head Ray layanan:

    kubectl port-forward svc/stable-diffusion-tpu-head-svc 8265:8265
    
  2. Di browser web, buka http://localhost:8265/.

  3. Klik tab Tayangkan.

Mengirim prompt ke server model

  1. Membuat sesi penerusan port ke endpoint Serve dari head Ray layanan:

    kubectl port-forward svc/stable-diffusion-tpu-serve-svc 8000
    
  2. Buka sesi Cloud Shell baru.

  3. Kirimkan perintah text-to-image ke server model Stabil Diffusion:

    python stable_diffusion_tpu_req.py  --save_pictures
    

    Hasil inferensi difusi stabil disimpan ke file bernama diffusion_results.png.

    Gambar yang dihasilkan oleh Stable Diffusion dengan 8 bagian: kursi hijau, seorang pria berdiri di luar rumah, robot di jalan, keluarga yang duduk di meja, dokumen yang berjalan di taman, naga terbang, potret beruang bergaya Jepang, dan air terjun.

Pembersihan

Menghapus project

    Menghapus project Google Cloud:

    gcloud projects delete PROJECT_ID

Menghapus resource satu per satu

Untuk menghapus cluster, ketik:

gcloud container clusters delete ${CLUSTER_NAME}

Langkah selanjutnya