TPU에서 LLM 워크로드 자동 확장 구성


이 페이지에서는 JetStream을 사용하여 Gemma 대규모 언어 모델(LLM)을 배포하기 위해 GKE 수평형 포드 자동 확장 처리(HPA)를 사용하여 단일 호스트 자동 확장 인프라를 설정하는 방법을 보여줍니다.

자동 확장을 위한 측정항목 선택에 대한 자세한 내용은 GKE에서 TPU로 LLM 워크로드를 자동 확장하기 위한 권장사항을 참조하세요.

시작하기 전에

시작하기 전에 다음 태스크를 수행했는지 확인합니다.

  • Google Kubernetes Engine API를 사용 설정합니다.
  • Google Kubernetes Engine API 사용 설정
  • 이 태스크에 Google Cloud CLI를 사용하려면 gcloud CLI를 설치한 후 초기화합니다. 이전에 gcloud CLI를 설치한 경우 gcloud components update를 실행하여 최신 버전을 가져옵니다.

측정항목을 사용하여 자동 확장

JetStream 추론 서버에서 내보낸 워크로드별 성능 측정항목 또는 TPU 성능 측정항목을 사용하여 포드의 직접 자동 확장을 수행할 수 있습니다.

측정항목으로 자동 확장을 설정하려면 다음 단계를 수행합니다.

  1. JetStream 서버에서 Cloud Monitoring으로 측정항목 내보내기 Prometheus 수집기의 배포 및 구성을 간소화하는 Google Cloud Managed Service for Prometheus를 사용합니다. Google Cloud Managed Service for Prometheus는 GKE 클러스터에서 기본적으로 사용 설정되어 있습니다. 수동으로 사용 설정할 수도 있습니다.

    다음 매니페스트 예에서는 Google Cloud Managed Service for Prometheus가 15초 간격으로 반복적으로 포드에서 측정항목을 스크래핑하도록 PodMonitoring 리소스 정의를 설정하는 방법을 보여줍니다.

    서버 측정항목을 스크래핑해야 하는 경우 다음 매니페스트를 사용합니다. 서버 측정항목을 사용하면 최대 5초 간격으로 스크래핑 간격이 지원됩니다.

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

    TPU 측정항목을 스크래핑해야 하는 경우 다음 매니페스트를 사용합니다. 시스템 측정항목을 사용하면 최대 15초 간격으로 스크래핑 간격이 지원됩니다.

    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. 측정항목 어댑터를 설치합니다. 이 어댑터는 Monitoring으로 내보낸 서버 측정항목이 HPA 컨트롤러에 표시되도록 합니다. 자세한 내용은 Google Cloud Managed Service for Prometheus 문서의 수평형 포드 자동 확장을 참고하세요.

    커스텀 측정항목 Stackdriver 어댑터

    커스텀 측정항목 Stackdriver 어댑터는 어댑터 버전 v0.13.1로 시작하는 Google Cloud Managed Service for Prometheus에서 측정항목을 쿼리하도록 지원합니다.

    커스텀 측정항목 Stackdriver 어댑터를 설치하려면 다음을 수행합니다.

    1. 클러스터의 관리 컬렉션을 설정합니다.

    2. 클러스터에 커스텀 측정항목 Stackdriver 어댑터를 설치합니다.

      kubectl apply -f https://raw.githubusercontent.com/GoogleCloudPlatform/k8s-stackdriver/master/custom-metrics-stackdriver-adapter/deploy/production/adapter_new_resource_model.yaml
      
    3. Kubernetes 클러스터에서 GKE용 워크로드 아이덴티티 제휴를 사용 설정하고 GKE용 워크로드 아이덴티티 제휴를 사용하는 경우 어댑터가 실행되는 서비스 계정에 모니터링 보기 권한도 부여해야 합니다.

    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
    

    Prometheus Adapter

    Google Cloud Managed Service for Prometheus를 사용하여 prometheus-adapter로 확장할 때는 다음 고려사항에 유의하세요.

    • Prometheus API 또는 UI를 사용하여 Google Cloud Managed Service for Prometheus를 쿼리할 때와 마찬가지로 쿼리를 Prometheus 프런트엔드 UI 프록시를 통해 라우팅합니다. 이 프런트엔드는 이후 단계에서 설치됩니다.
    • 기본적으로 prometheus-adapter 배포의 prometheus-url 인수는 --prometheus-url=http://frontend.default.svc:9090/로 설정되며 여기서 default는 프런트엔드를 배포한 네임스페이스입니다. 프런트엔드를 다른 네임스페이스에 배포한 경우 이 인수를 적절하게 구성합니다.
    • 규칙 구성의 .seriesQuery 필드에서 측정항목 이름에 정규 표현식(정규식) 일치자를 사용할 수 없습니다. 대신 측정항목 이름을 완전히 지정합니다.

    업스트림 Prometheus에 비해 Google Cloud Managed Service for Prometheus 내에서 데이터를 사용할 수 있는 데 시간이 조금 더 걸릴 수 있으므로 과도하게 빠른 자동 확장 로직을 구성하면 원치 않는 동작이 발생할 수 있습니다. 데이터 최신 상태에 대한 보장은 없지만 데이터가 Google Cloud Managed Service for Prometheus로 전송된 후 일반적으로 3~7초가 지나면 데이터를 쿼리할 수 있습니다(네트워크 지연 시간 제외).

    prometheus-adapter에서 실행하는 모든 쿼리는 전역 범위에 있습니다. 즉, 이름이 같은 측정항목을 내보내는 네임스페이스 두 개에 애플리케이션이 있는 경우 해당 측정항목을 사용하는 HPA 구성은 두 애플리케이션 모두의 데이터로 확장됩니다. 잘못된 데이터를 사용하여 확장되지 않도록 항상 PromQL에서 namespace 또는 cluster 필터를 사용하세요.

    prometheus-adapter 및 관리형 컬렉션을 사용하여 HPA 구성 예시를 설정하려면 다음 단계를 따르세요.

    1. 클러스터의 관리 컬렉션을 설정합니다.
    2. 클러스터에 Prometheus 프런트엔드 UI 프록시를 배포합니다. 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
      

      그런 후 다음 매니페스트를 적용합니다.

      kubectl apply -f prometheus-frontend.yaml
      
    3. prometheus-community/prometheus-adapter Helm 차트를 설치하여 클러스터에 prometheus-adapter가 설치되어 있는지 확인합니다. 다음 values.yaml 파일을 만듭니다.

      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"})
      

      그런 후 이 파일을 Helm 차트를 배포하기 위한 값 파일로 사용합니다.

      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
      

    GKE용 워크로드 아이덴티티 제휴를 사용하는 경우 다음 명령어를 실행하여 서비스 계정을 구성하고 승인해야 합니다.

    1. 먼저 클러스터 내 서비스 계정과 Google Cloud 서비스 계정을 만듭니다.

      gcloud iam service-accounts create prom-frontend-sa && kubectl create sa prom-frontend-sa
      
    2. 그런 다음 두 서비스 계정을 결합합니다.

      gcloud iam service-accounts add-iam-policy-binding \
        --role roles/iam.workloadIdentityUser \
        --member "serviceAccount:tpu-vm-gke-testing.svc.id.goog[default/prom-frontend-sa]" \
        jetstream-iam-sa@tpu-vm-gke-testing.iam.gserviceaccount.com \
      &&
      kubectl annotate serviceaccount \
        --namespace default \
        prom-frontend-sa \
        iam.gke.io/gcp-service-account=jetstream-iam-sa@tpu-vm-gke-testing.iam.gserviceaccount.com
      
    3. 그런 후 Google Cloud 서비스 계정에 monitoring.viewer 역할을 부여합니다.

      gcloud projects add-iam-policy-binding tpu-vm-gke-testing \
        --member=serviceAccount:jetstream-iam-sa@tpu-vm-gke-testing.iam.gserviceaccount.com \
        --role=roles/monitoring.viewer
      
    4. 마지막으로 프런트엔드 배포 서비스 계정을 새 클러스터 내 서비스 계정으로 설정합니다.

      kubectl set serviceaccount deployment frontend prom-frontend-sa
      
  3. 측정항목 기반 HPA 리소스를 설정합니다. 원하는 서버 측정항목을 기반으로 HPA 리소스를 배포합니다. 자세한 내용은 Google Cloud Managed Service for Prometheus 문서의 수평형 포드 자동 확장을 참고하세요. 특정 HPA 구성은 측정항목 유형(서버 또는 TPU)과 설치된 측정항목 어댑터에 따라 다릅니다.

    모든 HPA 구성에 몇 가지 값이 필요하며 HPA 리소스를 만들려면 다음 값을 설정해야 합니다.

    • MIN_REPLICAS: 허용되는 최소 JetStream 포드 복제본 수입니다. JetStream 배포 단계에서 JetStream 배포 매니페스트를 수정하지 않는 경우 1로 설정하는 것이 좋습니다.
    • MAX_REPLICAS: 허용되는 최대 JetStream 포드 복제본 수입니다. 이 JetStream 배포 예시에서는 복제본당 8개의 칩이 필요하며 노드 풀에는 16개의 칩이 포함되어 있습니다. 수직 확장 지연 시간을 짧게 유지하려면 이 값을 2로 설정하세요. 값이 클수록 클러스터 자동 확장 처리가 노드 풀에 새 노드를 만들어 수직 확장 지연 시간이 늘어납니다.
    • TARGET: 모든 JetStream 인스턴스에 걸친 이 측정항목의 타겟 평균입니다. 이 값에서 복제본 수가 결정되는 방식에 관한 자세한 내용은 자동 확장용 Kubernetes 문서를 참고하세요.

    커스텀 측정항목 Stackdriver 어댑터

    커스텀 측정항목 Stackdriver 어댑터는 모든 포드에서 Google Cloud Managed Service for Prometheus의 개별 측정항목 쿼리의 평균 값으로 워크로드 확장을 지원합니다. 커스텀 측정항목 Stackdriver 어댑터를 사용하는 경우 jetstream_prefill_backlog_sizejetstream_slots_used_percentage 서버 측정항목과 memory_used TPU 측정항목을 사용하여 확장하는 것이 좋습니다.

    서버 측정항목으로 확장할 HPA 매니페스트를 만들려면 다음 hpa.yaml 파일을 만듭니다.

    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
    

    TPU 측정항목과 함께 커스텀 측정항목 Stackdriver 어댑터를 사용하는 경우 확장에 kubernetes.io|node|accelerator|memory_used 측정항목만 사용하는 것이 좋습니다. 이 측정항목으로 확장하기 위해 HPA 매니페스트를 만들려면 다음 hpa.yaml 파일을 만듭니다.

    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
    

    Prometheus Adapter

    Prometheus 어댑터는 Google Cloud Managed Service for Prometheus의 PromQL 쿼리 값으로 워크로드 확장을 지원합니다. 앞에서 모든 포드의 평균값을 나타내는 jetstream_prefill_backlog_sizejetstream_slots_used_percentage 서버 측정항목을 정의했습니다.

    서버 측정항목으로 확장할 HPA 매니페스트를 만들려면 다음 hpa.yaml 파일을 만듭니다.

    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
    

    TPU 측정항목으로 확장하기 위해 HPA 매니페스트를 만들려면 prometheus-adapter helm 값 파일에 정의된 memory_used_percentage만 사용하는 것이 좋습니다. memory_used_percentage는 모든 가속기에서 사용되는 현재 평균 메모리를 반영하는 다음 PromQL 쿼리에 지정된 이름입니다.

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

    memory_used_percentage로 확장하기 위해 HPA 매니페스트를 만들려면 다음 hpa.yaml 파일을 만듭니다.

    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
    

여러 측정항목을 사용하여 확장

여러 측정항목을 기준으로 확장을 구성할 수도 있습니다. 여러 측정항목을 사용하여 복제본 수를 결정하는 방법을 알아보려면 자동 확장에 관한 Kubernetes 문서를 참고하세요. 이 유형의 HPA 매니페스트를 빌드하려면 각 HPA 리소스의 spec.metrics 필드에서 모든 항목을 단일 HPA 리소스로 수집합니다. 다음 스니펫은 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

자동 확장 모니터링 및 테스트

HPA 구성에 따라 JetStream 워크로드가 확장되는 방식을 관찰할 수 있습니다.

복제본 수를 실시간으로 관찰하려면 다음 명령어를 실행합니다.

kubectl get hpa --watch

이 명령어의 출력은 다음과 비슷하게 표시됩니다.

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

HPA의 확장 기능을 테스트하려면 모델 엔드포인트에 100개의 요청을 일괄 전송하는 다음 명령어를 사용하세요. 이렇게 하면 사용 가능한 디코딩 슬롯이 소진되고 미리 채우기 큐에 요청 백로그가 발생하여 HPA가 트리거되어 모델 배포 크기가 증가합니다.

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 }'

다음 단계