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

Panduan ini menunjukkan cara men-deploy dan menyajikan model Stable Diffusion di Google Kubernetes Engine (GKE) menggunakan TPU, Ray Serve, dan add-on Ray Operator.

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

Tentang Ray dan Ray Serve

Ray adalah framework komputasi yang skalabel dan open source untuk aplikasi AI/ML. Ray Serve adalah library inferensi model untuk Ray yang digunakan untuk menskalakan dan menyajikan model di lingkungan terdistribusi. Untuk mengetahui informasi selengkapnya, lihat Ray Serve dalam dokumentasi Ray.

Tentang TPU

Tensor Processing Unit (TPU) adalah akselerator hardware khusus yang dirancang untuk mempercepat pelatihan dan inferensi model machine learning berskala besar secara signifikan. Dengan menggunakan Ray bersama TPU, Anda dapat menskalakan aplikasi ML berperforma tinggi dengan lancar. Untuk mengetahui informasi selengkapnya tentang TPU, lihat Pengantar Cloud TPU dalam dokumentasi Cloud TPU.

Tentang webhook inisialisasi TPU KubeRay

Sebagai bagian dari add-on Ray Operator, GKE menyediakan webhook validasi dan mutasi yang menangani penjadwalan Pod TPU dan variabel lingkungan TPU tertentu yang diperlukan oleh framework seperti JAX untuk inisialisasi penampung. Webhook TPU KubeRay mengubah Pod dengan label app.kubernetes.io/name: kuberay yang 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 perlu berkomunikasi satu sama lain dalam slice. Variabel ini hanya disisipkan 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 grup pekerja multi-host, di mana beberapa Pod pekerja mungkin termasuk dalam replika yang sama, dan digunakan oleh Ray untuk mengaktifkan penskalaan otomatis multi-host.
  • TPU_NAME: String yang merepresentasikan PodSlice TPU GKE yang menjadi bagian Pod ini, ditetapkan ke nilai yang sama dengan label replicaIndex.
  • podAffinity: Memastikan GKE menjadwalkan Pod TPU dengan label replicaIndex yang cocok di node pool yang sama. Hal ini memungkinkan GKE menskalakan TPU multi-host secara atomik berdasarkan node pool, bukan node tunggal.

Tujuan

  • Buat cluster GKE dengan node pool TPU.
  • Deploy cluster Ray dengan TPU.
  • Deploy resource kustom RayService.
  • Berinteraksi dengan server model Stable Diffusion.

Biaya

Dalam dokumen ini, Anda akan menggunakan komponen Google Cloudyang dapat ditagih berikut:

Untuk membuat perkiraan biaya berdasarkan proyeksi penggunaan Anda, gunakan kalkulator harga.

Pengguna Google Cloud baru 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, baca bagian Pembersihan.

Sebelum memulai

Cloud Shell telah diinstal dengan software yang Anda perlukan untuk tutorial ini, termasuk kubectl dan gcloud CLI. Jika Anda tidak menggunakan Cloud Shell, instal 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. Instal Google Cloud CLI.

  3. Jika Anda menggunakan penyedia identitas (IdP) eksternal, Anda harus login ke gcloud CLI dengan identitas gabungan Anda terlebih dahulu.

  4. Untuk melakukan inisialisasi gcloud CLI, jalankan perintah berikut:

    gcloud init
  5. Buat atau pilih Google Cloud project.

    Peran yang diperlukan untuk memilih atau membuat project

    • Pilih project: Memilih project tidak memerlukan peran IAM tertentu—Anda dapat memilih project mana pun yang telah diberi peran.
    • Membuat project: Untuk membuat project, Anda memerlukan peran Pembuat Project (roles/resourcemanager.projectCreator), yang berisi izin resourcemanager.projects.create. Pelajari cara memberikan peran.
    • Buat Google Cloud project:

      gcloud projects create PROJECT_ID

      Ganti PROJECT_ID dengan nama untuk Google Cloud project 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.

  6. Verifikasi bahwa penagihan diaktifkan untuk project Google Cloud Anda.

  7. Aktifkan GKE API:

    Peran yang diperlukan untuk mengaktifkan API

    Untuk mengaktifkan API, Anda memerlukan peran IAM Service Usage Admin (roles/serviceusage.serviceUsageAdmin), yang berisi izin serviceusage.services.enable. Pelajari cara memberikan peran.

    gcloud services enable container.googleapis.com
  8. Instal Google Cloud CLI.

  9. Jika Anda menggunakan penyedia identitas (IdP) eksternal, Anda harus login ke gcloud CLI dengan identitas gabungan Anda terlebih dahulu.

  10. Untuk melakukan inisialisasi gcloud CLI, jalankan perintah berikut:

    gcloud init
  11. Buat atau pilih Google Cloud project.

    Peran yang diperlukan untuk memilih atau membuat project

    • Pilih project: Memilih project tidak memerlukan peran IAM tertentu—Anda dapat memilih project mana pun yang telah diberi peran.
    • Membuat project: Untuk membuat project, Anda memerlukan peran Pembuat Project (roles/resourcemanager.projectCreator), yang berisi izin resourcemanager.projects.create. Pelajari cara memberikan peran.
    • Buat Google Cloud project:

      gcloud projects create PROJECT_ID

      Ganti PROJECT_ID dengan nama untuk Google Cloud project 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.

  12. Verifikasi bahwa penagihan diaktifkan untuk project Google Cloud Anda.

  13. Aktifkan GKE API:

    Peran yang diperlukan untuk mengaktifkan API

    Untuk mengaktifkan API, Anda memerlukan peran IAM Service Usage Admin (roles/serviceusage.serviceUsageAdmin), yang berisi izin serviceusage.services.enable. Pelajari cara memberikan peran.

    gcloud services enable container.googleapis.com
  14. Memberikan peran ke akun pengguna 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:USER_IDENTIFIER" --role=ROLE

    Ganti kode berikut:

    • PROJECT_ID: Project ID Anda.
    • USER_IDENTIFIER: ID untuk akun pengguna Anda. Misalnya, myemail@example.com.
    • ROLE: Peran IAM yang Anda berikan ke akun pengguna Anda.

Pastikan kuota mencukupi

Pastikan project Anda memiliki kuota TPU yang cukup di region atau zona Compute Engine Anda. Google Cloud Untuk mengetahui informasi selengkapnya, lihat Pastikan kuota TPU dan GKE mencukupi dalam 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 Google Cloud konsol, dengan mengklik Ikon aktivasi Cloud Shell Activate Cloud Shell di Google Cloud konsol. 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 yang akan digunakan. Harus 1.30.1 atau setelahnya.
  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 node pool TPU

Buat cluster GKE Standard dengan node pool TPU:

  1. Buat cluster mode Standar dengan Ray Operator diaktifkan:

    gcloud container clusters create ${CLUSTER_NAME} \
        --addons=RayOperator \
        --machine-type=n1-standard-8 \
        --cluster-version=${CLUSTER_VERSION} \
        --location=${COMPUTE_REGION}
    
  2. Buat node pool TPU host tunggal:

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

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

Mengonfigurasi resource RayCluster dengan TPU

Konfigurasi manifes RayCluster untuk menyiapkan workload TPU:

Mengonfigurasi TPU nodeSelector

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

Perbarui manifes ray-cluster.yaml untuk menjadwalkan Pod Anda di podslice 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 harus dialokasikan GKE ke setiap Pod dengan mengonfigurasi google.com/tpuresource limits dan requests di kolom container TPU pada manifes RayCluster workerGroupSpecs.

Perbarui manifes ray-cluster.yaml dengan batas dan permintaan 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"

Mengonfigurasi 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 worker. Untuk grup pekerja multi-host, replika diperlakukan sebagai PodSlice, bukan pekerja individual, dengan numOfHosts node pekerja yang dibuat per replika.

Perbarui manifes ray-cluster.yaml dengan kode berikut:

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

Buat 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:
                - 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

    Manifes ini menjelaskan resource kustom RayService yang membuat resource RayCluster dengan 1 node head dan grup pekerja TPU dengan topologi 2x2x1, yang berarti setiap node pekerja akan memiliki 4 chip TPU v4.

    Node TPU termasuk dalam satu podslice TPU v4 dengan topologi 2x2x1. Untuk membuat grup pekerja multi-host, ganti nilai gke-tpu nodeSelector, batas dan permintaan penampung google.com/tpu, dan nilai numOfHosts dengan konfigurasi multi-host Anda. Untuk mengetahui informasi selengkapnya tentang topologi multi-host TPU, lihat Arsitektur sistem dalam 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 bahwa resource RayService sudah siap.

(Opsional) Melihat Dasbor Ray

Anda dapat melihat deployment Ray Serve dan log yang relevan dari Dasbor Ray.

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

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

  3. Klik tab Serve.

Mengirim perintah ke server model

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

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

  3. Kirim perintah teks ke gambar ke server model Stable Diffusion:

    python stable_diffusion_tpu_req.py  --save_pictures
    

    Hasil inferensi stable diffusion 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 duduk di meja, dokter berjalan di taman, naga terbang, potret beruang bergaya Jepang, dan air terjun.

Mengamati workload Ray

Untuk melihat detail RayJob, Anda dapat membuka bagian Kubernetes Engine > AI/ML > Jobs di konsol. Google Cloud

Melihat RayJob di konsol Google Cloud

Pembersihan

Menghapus project

    Menghapus Google Cloud project:

    gcloud projects delete PROJECT_ID

Menghapus resource satu per satu

Untuk menghapus cluster, ketik:

gcloud container clusters delete ${CLUSTER_NAME}

Langkah berikutnya