Mengonfigurasi penskalaan otomatis untuk workload LLM di TPU


Halaman ini menunjukkan cara menyiapkan infrastruktur penskalaan otomatis menggunakan Horizontal Pod Autoscaler (HPA) GKE untuk men-deploy model bahasa besar (LLM) Gemma menggunakan JetStream satu host.

Untuk mempelajari lebih lanjut cara memilih metrik untuk penskalaan otomatis, lihat Praktik terbaik untuk menskalakan otomatis workload LLM dengan TPU di GKE.

Sebelum memulai

Sebelum memulai, pastikan Anda telah menjalankan tugas berikut:

  • Aktifkan Google Kubernetes Engine API.
  • Aktifkan Google Kubernetes Engine API
  • Jika ingin menggunakan Google Cloud CLI untuk tugas ini, instal lalu lakukan inisialisasi gcloud CLI. Jika sebelumnya Anda telah menginstal gcloud CLI, dapatkan versi terbaru dengan menjalankan gcloud components update.

Penskalaan otomatis menggunakan metrik

Anda dapat menggunakan metrik performa khusus beban kerja yang dikeluarkan oleh server inferensi JetStream atau metrik performa TPU untuk mengarahkan penskalaan otomatis untuk Pod Anda.

Untuk menyiapkan penskalaan otomatis dengan metrik, ikuti langkah-langkah berikut:

  1. Ekspor metrik dari server JetStream ke Cloud Monitoring. Anda menggunakan Google Cloud Managed Service for Prometheus, yang menyederhanakan deployment dan konfigurasi kolektor Prometheus. Google Cloud Managed Service for Prometheus diaktifkan secara default di cluster GKE Anda; Anda juga dapat mengaktifkannya secara manual.

    Contoh manifes berikut menunjukkan cara menyiapkan definisi resource PodMonitoring untuk mengarahkan Google Cloud Managed Service for Prometheus untuk mengambil metrik dari Pod Anda dengan interval berulang 15 detik:

    Jika Anda perlu meng-scrap metrik server, gunakan manifes berikut. Dengan metrik server, interval pengambilan data yang seringnya 5 detik didukung.

    apiVersion: monitoring.googleapis.com/v1
    kind: PodMonitoring
    metadata:
      name: jetstream-podmonitoring
    spec:
      selector:
        matchLabels:
          app: maxengine-server
      endpoints:
      - interval: 15s
        path: "/"
        port: PROMETHEUS_PORT
      targetLabels:
        metadata:
        - pod
        - container
        - node
    

    Jika Anda perlu meng-scrap metrik TPU, gunakan manifes berikut. Dengan metrik sistem, interval pengambilan data yang seringnya 15 detik didukung.

    apiVersion: monitoring.googleapis.com/v1
    kind: PodMonitoring
    metadata:
      name: tpu-metrics-exporter
      namespace: kube-system
      labels:
        k8s-app: tpu-device-plugin
    spec:
      endpoints:
        - port: 2112
          interval: 15s
      selector:
        matchLabels:
          k8s-app: tpu-device-plugin
    
  2. Instal Adaptor Metrik. Adaptor ini membuat metrik server yang Anda ekspor ke Monitoring terlihat oleh pengontrol HPA. Untuk mengetahui detail selengkapnya, lihat Penskalaan otomatis pod horizontal dalam dokumentasi Google Cloud Managed Service for Prometheus.

    Adaptor Stackdriver Metrik Kustom

    Adaptor Stackdriver Metrik Kustom mendukung metrik kueri dari Google Cloud Managed Service for Prometheus, mulai dari adaptor versi v0.13.1.

    Untuk menginstal Adaptor Stackdriver Metrik Kustom, lakukan hal berikut:

    1. Siapkan koleksi terkelola di cluster Anda.

    2. Instal Adaptor Stackdriver Metrik Kustom di cluster Anda.

      kubectl apply -f https://raw.githubusercontent.com/GoogleCloudPlatform/k8s-stackdriver/master/custom-metrics-stackdriver-adapter/deploy/production/adapter_new_resource_model.yaml
      
    3. Jika Anda mengaktifkan Workload Identity Federation for GKE di cluster Kubernetes dan menggunakan Workload Identity Federation for GKE, Anda juga harus memberikan peran Monitoring Viewer ke akun layanan tempat adaptor berjalan. Ganti PROJECT_ID dengan project ID Anda.

    export PROJECT_NUMBER=$(gcloud projects describe PROJECT_ID --format 'get(projectNumber)')
    gcloud projects add-iam-policy-binding projects/PROJECT_ID \
      --role roles/monitoring.viewer \
      --member=principal://iam.googleapis.com/projects/$PROJECT_NUMBER/locations/global/workloadIdentityPools/PROJECT_ID.svc.id.goog/subject/ns/custom-metrics/sa/custom-metrics-stackdriver-adapter
    

    Adaptor Prometheus

    Perhatikan pertimbangan berikut saat menggunakan prometheus-adapter untuk melakukan penskalaan menggunakan Google Cloud Managed Service for Prometheus:

    • Rutekan kueri melalui proxy UI frontend Prometheus, seperti saat membuat kueri Google Cloud Managed Service for Prometheus menggunakan Prometheus API atau UI. Frontend ini diinstal di langkah berikutnya.
    • Secara default, argumen prometheus-url Deployment prometheus-adapter ditetapkan ke --prometheus-url=http://frontend.default.svc:9090/, dengan default adalah namespace tempat Anda men-deploy frontend. Jika Anda men-deploy frontend di namespace lain, konfigurasikan argumen ini dengan benar.
    • Di kolom .seriesQuery pada konfigurasi aturan, Anda tidak dapat menggunakan pencocok ekspresi reguler (regex) pada nama metrik. Sebagai gantinya, tentukan nama metrik sepenuhnya.

    Karena data dapat memerlukan waktu sedikit lebih lama untuk tersedia dalam Google Cloud Managed Service for Prometheus dibandingkan dengan Prometheus upstream, mengonfigurasi logika penskalaan otomatis yang terlalu cepat dapat menyebabkan perilaku yang tidak diinginkan. Meskipun tidak ada jaminan tentang keaktualan data, data biasanya tersedia untuk dikueri 3-7 detik setelah dikirim ke Google Cloud Managed Service for Prometheus, tidak termasuk latensi jaringan.

    Semua kueri yang dikeluarkan oleh prometheus-adapter memiliki cakupan global. Artinya, jika Anda memiliki aplikasi di dua namespace yang memunculkan metrik dengan nama yang sama, konfigurasi HPA yang menggunakan metrik tersebut akan diskalakan menggunakan data dari kedua aplikasi. Untuk menghindari penskalaan menggunakan data yang salah, selalu gunakan filter namespace atau cluster di PromQL Anda.

    Untuk menyiapkan contoh konfigurasi HPA menggunakan prometheus-adapter dan koleksi terkelola, ikuti langkah-langkah berikut:

    1. Siapkan koleksi terkelola di cluster Anda.
    2. Deploy proxy UI frontend Prometheus di cluster Anda. Buat manifes berikut bernama prometheus-frontend.yaml:

        apiVersion: apps/v1
        kind: Deployment
        metadata:
          name: frontend
        spec:
          replicas: 2
          selector:
            matchLabels:
              app: frontend
          template:
            metadata:
              labels:
                app: frontend
            spec:
              automountServiceAccountToken: true
              affinity:
                nodeAffinity:
                  requiredDuringSchedulingIgnoredDuringExecution:
                    nodeSelectorTerms:
                    - matchExpressions:
                      - key: kubernetes.io/arch
                        operator: In
                        values:
                        - arm64
                        - amd64
                      - key: kubernetes.io/os
                        operator: In
                        values:
                        - linux
              containers:
              - name: frontend
                image: gke.gcr.io/prometheus-engine/frontend:v0.8.0-gke.4
                args:
                - "--web.listen-address=:9090"
                - "--query.project-id=PROJECT_ID"
                ports:
                - name: web
                  containerPort: 9090
                readinessProbe:
                  httpGet:
                    path: /-/ready
                    port: web
                securityContext:
                  allowPrivilegeEscalation: false
                  capabilities:
                    drop:
                    - all
                  privileged: false
                  runAsGroup: 1000
                  runAsNonRoot: true
                  runAsUser: 1000
                livenessProbe:
                  httpGet:
                    path: /-/healthy
                    port: web
        ---
        apiVersion: v1
        kind: Service
        metadata:
          name: prometheus
        spec:
          clusterIP: None
          selector:
            app: frontend
          ports:
          - name: web
            port: 9090
      

      Kemudian, terapkan manifes:

      kubectl apply -f prometheus-frontend.yaml
      
    3. Pastikan prometheus-adapter diinstal di cluster Anda dengan menginstal diagram helm prometheus-community/prometheus-adapter. Buat file values.yaml berikut:

      rules:
        default: false
        external:
        - seriesQuery: 'jetstream_prefill_backlog_size'
          resources:
            template: <<.Resource>>
          name:
            matches: ""
            as: "jetstream_prefill_backlog_size"
          metricsQuery: avg(<<.Series>>{<<.LabelMatchers>>,cluster="CLUSTER_NAME"})
        - seriesQuery: 'jetstream_slots_used_percentage'
          resources:
            template: <<.Resource>>
          name:
            matches: ""
            as: "jetstream_slots_used_percentage"
          metricsQuery: avg(<<.Series>>{<<.LabelMatchers>>,cluster="CLUSTER_NAME"})
        - seriesQuery: 'memory_used'
          resources:
            template: <<.Resource>>
          name:
            matches: ""
            as: "memory_used_percentage"
          metricsQuery: avg(memory_used{cluster="CLUSTER_NAME",exported_namespace="default",container="jetstream-http"}) / avg(memory_total{cluster="CLUSTER_NAME",exported_namespace="default",container="jetstream-http"})
      

      Kemudian, gunakan file ini sebagai file nilai untuk men-deploy diagram helm Anda:

      helm repo add prometheus-community https://prometheus-community.github.io/helm-charts && helm repo update && helm install example-release prometheus-community/prometheus-adapter -f values.yaml
      

    Jika menggunakan Workload Identity Federation untuk GKE, Anda juga perlu mengonfigurasi dan memberikan otorisasi ke akun layanan dengan menjalankan perintah berikut:

    1. Pertama, buat akun layanan dalam cluster dan Google Cloud:

      gcloud iam service-accounts create prom-frontend-sa && kubectl create sa prom-frontend-sa
      
    2. Kemudian, tautkan kedua akun layanan, pastikan untuk mengganti PROJECT_ID dengan project ID Anda:

      gcloud iam service-accounts add-iam-policy-binding \
        --role roles/iam.workloadIdentityUser \
        --member "serviceAccount:PROJECT_ID.svc.id.goog[default/prom-frontend-sa]" \
        jetstream-iam-sa@PROJECT_ID.iam.gserviceaccount.com \
      &&
      kubectl annotate serviceaccount \
        --namespace default \
        prom-frontend-sa \
        iam.gke.io/gcp-service-account=jetstream-iam-sa@PROJECT_ID.iam.gserviceaccount.com
      
    3. Selanjutnya, berikan peran monitoring.viewer ke akun layanan Google Cloud:

      gcloud projects add-iam-policy-binding PROJECT_ID \
        --member=serviceAccount:jetstream-iam-sa@PROJECT_ID.iam.gserviceaccount.com \
        --role=roles/monitoring.viewer
      
    4. Terakhir, tetapkan akun layanan deployment frontend Anda sebagai akun layanan dalam cluster baru:

      kubectl set serviceaccount deployment frontend prom-frontend-sa
      
  3. Siapkan resource HPA berbasis metrik. Deploy resource HPA yang didasarkan pada metrik server pilihan Anda. Untuk mengetahui detail selengkapnya, lihat Penskalaan otomatis pod horizontal dalam dokumentasi Google Cloud Managed Service for Prometheus. Konfigurasi HPA tertentu bergantung pada jenis metrik (server atau TPU) dan adaptor metrik yang diinstal.

    Beberapa nilai diperlukan di semua konfigurasi HPA dan harus ditetapkan untuk membuat resource HPA:

    • MIN_REPLICAS: Jumlah minimum replika pod JetStream yang diizinkan. Jika tidak mengubah manifes deployment JetStream dari langkah Deploy JetStream, sebaiknya tetapkan ke 1.
    • MAX_REPLICAS: Jumlah maksimum replika pod JetStream yang diizinkan. Contoh deployment JetStream memerlukan 8 chip per replika dan kumpulan node berisi 16 chip. Jika Anda ingin mempertahankan latensi penskalaan tetap rendah, tetapkan ke 2. Nilai yang lebih besar akan memicu Autoscaler Cluster untuk membuat node baru di node pool, sehingga meningkatkan latensi penskalaan.
    • TARGET: Rata-rata yang ditargetkan untuk metrik ini di semua instance JetStream. Lihat Dokumentasi Kubernetes untuk Penskalaan Otomatis untuk mengetahui informasi selengkapnya tentang cara menentukan jumlah replika dari nilai ini.

    Adaptor Stackdriver Metrik Kustom

    Adaptor Stackdriver Metrik Kustom mendukung penskalaan beban kerja Anda dengan nilai rata-rata setiap kueri metrik dari Google Cloud Managed Service for Prometheus di semua Pod. Saat menggunakan Adaptor Stackdriver Metrik Kustom, sebaiknya lakukan penskalaan dengan metrik server jetstream_prefill_backlog_size dan jetstream_slots_used_percentage serta metrik TPU memory_used.

    Untuk membuat manifes HPA guna melakukan penskalaan dengan metrik server, buat file hpa.yaml berikut:

    apiVersion: autoscaling/v2
    kind: HorizontalPodAutoscaler
    metadata:
      name: jetstream-hpa
      namespace: default
    spec:
      scaleTargetRef:
        apiVersion: apps/v1
        kind: Deployment
        name: maxengine-server
      minReplicas: MIN_REPLICAS
      maxReplicas: MAX_REPLICAS
      metrics:
      - type: Pods
        pods:
          metric:
            name: prometheus.googleapis.com|jetstream_METRIC|gauge
          target:
            type: AverageValue
            averageValue: TARGET
    

    Saat menggunakan Adaptor Stackdriver Metrik Kustom dengan metrik TPU, sebaiknya hanya gunakan metrik kubernetes.io|node|accelerator|memory_used untuk penskalaan. Untuk membuat manifes HPA guna melakukan penskalaan dengan metrik ini, buat file hpa.yaml berikut:

    apiVersion: autoscaling/v2
    kind: HorizontalPodAutoscaler
    metadata:
      name: jetstream-hpa
      namespace: default
    spec:
      scaleTargetRef:
        apiVersion: apps/v1
        kind: Deployment
        name: maxengine-server
      minReplicas: MIN_REPLICAS
      maxReplicas: MAX_REPLICAS
      metrics:
      - type: External
        external:
          metric:
            name: prometheus.googleapis.com|memory_used|gauge
            selector:
              matchLabels:
                metric.labels.container: jetstream-http
                metric.labels.exported_namespace: default
          target:
            type: AverageValue
            averageValue: TARGET
    

    Adaptor Prometheus

    Prometheus Adapter mendukung penskalaan workload Anda dengan nilai kueri PromQL dari Google Cloud Managed Service for Prometheus. Sebelumnya, Anda telah menentukan metrik server jetstream_prefill_backlog_size dan jetstream_slots_used_percentage yang mewakili nilai rata-rata di semua Pod.

    Untuk membuat manifes HPA guna melakukan penskalaan dengan metrik server, buat file hpa.yaml berikut:

    apiVersion: autoscaling/v2
    kind: HorizontalPodAutoscaler
    metadata:
      name: jetstream-hpa
      namespace: default
    spec:
      scaleTargetRef:
        apiVersion: apps/v1
        kind: Deployment
        name: maxengine-server
      minReplicas: MIN_REPLICAS
      maxReplicas: MAX_REPLICAS
      metrics:
      - type: External
        external:
          metric:
            name: jetstream_METRIC
          target:
            type: AverageValue
            averageValue: TARGET
    

    Untuk membuat manifes HPA guna melakukan penskalaan dengan metrik TPU, sebaiknya hanya gunakan memory_used_percentage yang ditentukan dalam file nilai helm prometheus-adapter. memory_used_percentage adalah nama yang diberikan ke kueri PromQL berikut yang mencerminkan memori rata-rata saat ini yang digunakan di semua akselerator:

    avg(kubernetes_io:node_accelerator_memory_used{cluster_name="CLUSTER_NAME"}) / avg(kubernetes_io:node_accelerator_memory_total{cluster_name="CLUSTER_NAME"})
    

    Untuk membuat manifes HPA guna melakukan penskalaan dengan memory_used_percentage, buat file hpa.yaml berikut:

    apiVersion: autoscaling/v2
    kind: HorizontalPodAutoscaler
    metadata:
      name: jetstream-hpa
      namespace: default
    spec:
      scaleTargetRef:
        apiVersion: apps/v1
        kind: Deployment
        name: maxengine-server
      minReplicas: MIN_REPLICAS
      maxReplicas: MAX_REPLICAS
      metrics:
      - type: External
        external:
          metric:
            name: memory_used_percentage
          target:
            type: AverageValue
            averageValue: TARGET
    

Menskalakan menggunakan beberapa metrik

Anda juga dapat mengonfigurasi penskalaan berdasarkan beberapa metrik. Untuk mempelajari cara jumlah replika ditentukan menggunakan beberapa metrik, lihat dokumentasi Kubernetes tentang penskalaan otomatis. Untuk mem-build jenis manifes HPA ini, kumpulkan semua entri dari kolom spec.metrics dari setiap resource HPA ke dalam satu resource HPA. Cuplikan berikut menunjukkan contoh cara memaketkan resource HPA:

apiVersion: autoscaling/v2
kind: HorizontalPodAutoscaler
metadata:
  name: jetstream-hpa-multiple-metrics
  namespace: default
spec:
  scaleTargetRef:
    apiVersion: apps/v1
    kind: Deployment
    name: maxengine-server
  minReplicas: MIN_REPLICAS
  maxReplicas: MAX_REPLICAS
  metrics:
  - type: Pods
    pods:
      metric:
        name: jetstream_METRIC
      target:
        type: AverageValue
      averageValue: JETSTREAM_METRIC_TARGET
  - type: External
    external:
      metric:
        name: memory_used_percentage
      target:
        type: AverageValue
      averageValue: EXTERNAL_METRIC_TARGET

Memantau dan menguji penskalaan otomatis

Anda dapat mengamati cara workload JetStream diskalakan berdasarkan konfigurasi HPA.

Untuk mengamati jumlah replika secara real time, jalankan perintah berikut:

kubectl get hpa --watch

Output dari perintah ini akan mirip dengan berikut ini:

NAME            REFERENCE                     TARGETS      MINPODS   MAXPODS   REPLICAS   AGE
jetstream-hpa   Deployment/maxengine-server   0/10 (avg)   1         2         1          1m

Untuk menguji kemampuan HPA Anda dalam melakukan penskalaan, gunakan perintah berikut yang mengirimkan lonjakan 100 permintaan ke endpoint model. Hal ini akan menghabiskan slot dekode yang tersedia dan menyebabkan backlog permintaan di antrean pra-pengisian, yang memicu HPA untuk meningkatkan ukuran deployment model.

seq 100 | xargs -P 100 -n 1 curl --request POST --header "Content-type: application/json" -s localhost:8000/generate --data '{ "prompt": "Can you provide a comprehensive and detailed overview of the history and development of artificial intelligence.", "max_tokens": 200 }'

Langkah selanjutnya