GKE AI/ML 워크로드 우선순위 최적화

이 문서에서는 특히 예약에 용량이 없거나 온디맨드 리소스를 통해 용량이 없는 경우 Google Kubernetes Engine (GKE)에서 이기종 AI/ML 워크로드의 리소스 사용률을 극대화하고 다운타임을 최소화하기 위한 도구와 권장사항을 설명합니다. 이종 워크로드는 동일한 GKE 클러스터에서 동시에 실행되는 다양한 유형의 AI/ML 워크로드를 의미합니다. 예를 들어 지연 시간에 민감한 온라인 추론 서비스를 중단 가능한 일괄 학습 작업과 함께 실행할 수 있습니다.

이 가이드에서는 플랫폼 관리자, 운영자, 데이터 및 AI 전문가를 위한 권장사항을 제공합니다.

AI/ML 워크로드 우선순위 지정의 이점

이종 워크로드는 우선순위가 다르며 제한된 용량과 리소스를 공유합니다. 이 페이지의 권장사항에서는 다음과 같은 이점을 얻을 수 있도록 GKE 및 오픈소스 도구를 구성하는 방법을 설명합니다.

  • 우선순위가 높은 워크로드의 다운타임을 최소화합니다.
  • 우선순위가 높은 워크로드를 빠르게 실행합니다.
  • 리소스 소비를 최적화합니다.

배경

GKE는 리소스 사용률을 최적화하기 위해 다음 오픈소스 도구를 지원합니다.

  • Kueue: 일괄 처리, AI, 고성능 컴퓨팅 워크로드를 위해 설계된 Kubernetes 기반 워크로드 큐 시스템입니다. Kueue는 leaderworkerset와 같은 커스텀 리소스 정의로 정의된 다른 워크로드 유형을 관리하도록 확장할 수 있습니다. Kueue는 Kubernetes 클러스터에서 할당량과 워크로드가 이를 소비하는 방식을 관리합니다. Kueue는 워크로드의 대기 시기, 워크로드의 시작 시기 (예: 포드 생성), 워크로드에 속한 포드의 선점 시기를 결정합니다.

    Kueue에 대한 자세한 내용은 Kueue 개념 문서를 참고하세요.

  • 핫스왑: 평균 복구 시간 (MTTR)을 줄이는 기술입니다. 핫스왑을 사용하면 클러스터 리소스가 완전히 사용되고 주문형 인스턴스나 기존 예약에서 추가 용량을 사용할 수 없는 경우 워크로드 우선순위에 따라 선점이 가능합니다.

    • 워크로드를 호스팅하는 노드가 비정상 상태가 되면 적합한 예비 노드에서 워크로드가 다시 예약됩니다. 여유 노드가 없는 경우 핫스왑은 복구 중인 워크로드를 위한 공간을 확보하기 위해 우선순위가 낮은 워크로드를 선점할 수 있습니다.
    • PriorityClass로 포드를 구성하면 우선순위가 더 높게 구성된 워크로드가 실행 중인 낮은 우선순위 워크로드를 제거하여 리소스를 획득합니다. 이 제거 프로세스를 선점이라고 합니다.

사용 사례

다음 표를 참고하여 각 사용 사례에 대한 권장사항을 알아보세요.

사용 사례 권장사항 설명
우선순위가 다른 여러 워크로드 Kueue를 사용하여 대기열을 정의하고 중요도에 따라 워크로드에 우선순위를 할당합니다. Kueue는 특정 팀이나 프로젝트가 일정량의 리소스에 액세스할 수 있도록 할당량을 관리할 수 있습니다.

Kueue를 사용하면 다음 구성을 적용할 수 있습니다.

  • 높은 Kueue WorkloadPriority를 할당하여 우선순위가 높은 작업을 우선 처리합니다.
  • 우선순위가 낮은 워크로드도 결국 리소스를 수신할 수 있도록 Kueue의 공정 공유 대기열을 사용 설정합니다.

권장사항 구성을 테스트하려면 이 문서의 Kueue 예시를 참고하세요.

현재 MTTR을 줄여야 합니다. 중단이 발생하면 핫스왑을 사용하여 정상 리소스에서 워크로드를 재예약하고 우선순위가 높은 워크로드를 위해 우선순위가 낮은 워크로드를 선점합니다.

핫스왑을 사용하면 다음 구성을 적용할 수 있습니다.

  • PriorityClasses를 구성하여 워크로드의 우선순위 수준을 정의합니다.
  • 중요한 워크로드에 더 높은 PriorityClasses를 할당합니다.
  • 중단이 발생하면 정상 노드에서 워크로드의 일정을 자동으로 다시 조정합니다.

권장사항 구성을 테스트하려면 이 문서의 핫스왑 예시를 참고하세요.

제한된 리소스를 두고 경쟁하는 여러 AI 워크로드 Kueue와 Hotswap 결합 이 조합은 초기 예약 중과 런타임 중에 모두 중요한 워크로드의 우선순위를 지정하는 강력한 시스템을 제공합니다.

Kueue 및 Hotswap을 사용하면 다음 구성을 적용할 수 있습니다.

  • Kueue를 사용하여 우선순위에 따라 워크로드의 초기 예약 및 승인을 관리합니다.
  • 핫스왑을 사용하여 워크로드 중단을 처리하고 신속한 복구를 지원합니다. 핫스왑은 중단이 발생할 때 우선순위가 높은 워크로드의 복구 시간을 줄이는 데 도움이 됩니다.

권장사항 구성을 테스트하려면 이 문서의 Kueue 및 핫스왑 예시를 참고하세요.

권장사항 구현 예

다음 예에서는 Kueue와 Hotswap을 구현하는 방법과 이전 섹션에 설명된 권장사항을 위해 이를 결합하는 방법을 보여줍니다.

Kueue

다음 예시 매니페스트는 Kueue 구성을 보여줍니다.

  apiVersion: kueue.x-k8s.io/v1beta1
  kind: ResourceFlavor
  metadata:
    name: tpu-v6e-slice
  spec:
    nodeLabels:
      cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
  ---
  apiVersion: kueue.x-k8s.io/v1beta1
  kind: ClusterQueue
  metadata:
    name: tpu-training-cq
  spec:
    resourceGroups:
    - flavors:
      - name: tpu-v6e-slice
        resources:
        - name: google.com/tpu
          nominalQuota: 32
    queueingStrategy: BestEffortFIFO
    preemption:
      reclaimWithinCohort: Never
      reclaimOutOfCohort:
        enable: true
        reclaimMoreThanNominalQuota: false
  ---
  apiVersion: kueue.x-k8s.io/v1beta1
  kind: LocalQueue
  metadata:
    name: default-queue
    namespace: default
  spec:
    clusterQueue: tpu-training-cq

이 매니페스트는 다음을 수행합니다.

  • TPU v6e 슬라이스의 노드 라벨을 지정하는 tpu-v6e-slice라는 ResourceFlavor를 정의합니다.
  • TPU 리소스의 할당량을 관리하는 tpu-training-cq라는 ClusterQueue를 정의합니다.
  • default 네임스페이스의 워크로드가 tpu-training-cq 클러스터 대기열을 사용할 수 있도록 하는 default-queue이라는 LocalQueue를 정의합니다.

핫스왑

다음 예시에서는 low-priority-jobhigh-priority-job의 두 우선순위 클래스를 정의하는 핫스왑 구성을 보여줍니다. 이 핫스왑 구성은 우선순위가 높은 JobSet 워크로드를 만들고 MaxText를 사용합니다.

  apiVersion: scheduling.k8s.io/v1
  kind: PriorityClass
  metadata:
    name: low-priority-job
  value: 1000000
  globalDefault: false
  description: "This priority class should be used for low priority pods only."
  ---
  apiVersion: scheduling.k8s.io/v1
  kind: PriorityClass
  metadata:
    name: high-priority-job
  value: 2000000
  globalDefault: false
  description: "This priority class should be used for critical pods only."
  ---
  apiVersion: jobset.x-k8s.io/v1alpha2
  kind: JobSet
  metadata:
    name: high-jax-trillium
    annotations:
      alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
  spec:
    failurePolicy:
      maxRestarts: 10
      restartStrategy: BlockingRecreate
    replicatedJobs:
    - name: slice
      replicas: 2
      template:
        spec:
          backoffLimit: 0
          completions: 4
          parallelism: 4
          template:
            spec:
              nodeSelector:
                cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
                cloud.google.com/gke-tpu-topology: 4x4
              hostNetwork: true
              dnsPolicy: ClusterFirstWithHostNet
              priorityClassName: high-priority-job
              containers:
              - name: jax-program
                image: <IMAGE LOCATION>
                command:
                -   python3
                -   MaxText/train.py
                -   MaxText/configs/base.yml
                -   model_name=llama2-7b
                -   run_name=<UNIQUE RUN NAME>
                -   steps=300
                -   base_output_directory=gs://<OUTPUT BUCKET>
                -   dataset_path=gs://max-datasets-rogue
                -   max_target_length=4096
                -   dataset_type=synthetic
                -   enable_checkpointing=False
                resources:
                  limits:
                    google.com/tpu: 4

이 구성에 따라 Hotswap은 다음 작업을 실행합니다.

  • 인프라 오류로 인해 우선순위가 높은 워크로드가 중단되면 JobSet이 다시 시작합니다. 핫스왑은 인프라가 복구되기 전에 우선순위가 높은 워크로드를 다시 예약하기 위해 우선순위가 낮은 워크로드를 선점합니다. 우선순위가 낮은 워크로드는 실패 상태로 유지됩니다. 이 프로세스는 워크로드 유휴 시간을 크게 줄여줍니다.
  • 인프라가 복구되면 Hotswap은 복구된 노드 풀에서 우선순위가 낮은 워크로드를 다시 예약합니다.

Kueue 및 핫스왑

제한된 리소스가 있는 복잡한 환경에서 운영하는 경우 Kueue와 Hotswap을 결합하세요. 이 조합은 초기 일정 예약 및 런타임 중에 중요한 워크로드의 우선순위를 지정하는 강력한 시스템을 제공합니다.

다음 예는 Kueue와 Hotswap 구성의 조합을 보여줍니다. 이 예시에서는 MaxText를 사용합니다.

  apiVersion: scheduling.k8s.io/v1
  kind: PriorityClass
  metadata:
    name: low-priority-job
  value: 1000000
  globalDefault: false
  description: "This priority class should be used for low priority pods only."
  ---
  apiVersion: scheduling.k8s.io/v1
  kind: PriorityClass
  metadata:
    name: high-priority-job
  value: 2000000
  globalDefault: false
  description: "This priority class should be used for critical pods only."
  ---
  apiVersion: kueue.x-k8s.io/v1beta1
  kind: ResourceFlavor
  metadata:
    name: tpu-v6e-slice
  spec:
    nodeLabels:
      cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
  ---
  apiVersion: kueue.x-k8s.io/v1beta1
  kind: ClusterQueue
  metadata:
    name: tpu-training-cq
  spec:
    resourceGroups:
    - flavors:
      - name: tpu-v6e-slice
        resources:
        - name: google.com/tpu
          nominalQuota: 32
    queueingStrategy: BestEffortFIFO
    preemption:
      reclaimWithinCohort: Never
      reclaimOutOfCohort:
        enable: true
        reclaimMoreThanNominalQuota: false
  ---
  apiVersion: kueue.x-k8s.io/v1beta1
  kind: LocalQueue
  metadata:
    name: default-queue
    namespace: default
  spec:
    clusterQueue: tpu-training-cq
  ---
  apiVersion: jobset.x-k8s.io/v1alpha2
  kind: JobSet
  metadata:
    name: low-jax-trillium
    annotations:
      kueue.x-k8s.io/queue-name: default-queue
      alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
  spec:
    failurePolicy:
      maxRestarts: 10
      restartStrategy: BlockingRecreate
    replicatedJobs:
    - name: slice
      replicas: 2
      template:
        spec:
          backoffLimit: 0
          completions: 4
          parallelism: 4
          template:
            metadata:
              labels:
                kueue.x-k8s.io/managed-by: kueue
                kueue.x-k8s.io/priority-class: low-priority-job
            spec:
              nodeSelector:
                cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
                cloud.google.com/gke-tpu-topology: 4x4
              hostNetwork: true
              dnsPolicy: ClusterFirstWithHostNet
              priorityClassName: low-priority-job
              containers:
              - name: jax-program
                image: <IMAGE LOCATION>
                command:
                - python3
                - MaxText/train.py
                - MaxText/configs/base.yml
                - model_name=llama2-7b
                - run_name=low-priority-run
                - steps=30000
                - base_output_directory=gs://<OUTPUT BUCKET>
                - dataset_path=gs://max-datasets-rogue
                - max_target_length=4096
                - dataset_type=synthetic
                - enable_checkpointing=False
                resources:
                  limits:
                    google.com/tpu: 4
  ---
  apiVersion: jobset.x-k8s.io/v1alpha2
  kind: JobSet
  metadata:
    name: high-jax-trillium
    annotations:
      kueue.x-k8s.io/queue-name: default-queue
      alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
  spec:
    failurePolicy:
      maxRestarts: 10
      restartStrategy: BlockingRecreate
    replicatedJobs:
    - name: slice
      replicas: 2
      template:
        spec:
          backoffLimit: 0
          completions: 4
          parallelism: 4
          template:
            metadata:
              labels:
                kueue.x-k8s.io/managed-by: kueue
                kueue.x-k8s.io/priority-class: high-priority-job
            spec:
              nodeSelector:
                cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
                cloud.google.com/gke-tpu-topology: 4x4
              hostNetwork: true
              dnsPolicy: ClusterFirstWithHostNet
              priorityClassName: high-priority-job
              containers:
              - name: jax-program
                image: <IMAGE LOCATION>
                command:
                - python3
                - MaxText/train.py
                - MaxText/configs/base.yml
                - model_name=llama2-7b
                - run_name=high-priority-run
                - steps=300
                - base_output_directory=gs://<OUTPUT BUCKET>
                - dataset_path=gs://max-datasets-rogue
                - max_target_length=4096
                - dataset_type=synthetic
                - enable_checkpointing=False
                resources:
                  limits:
                    google.com/tpu: 4

이 구성에 따라 Kueue는 Hotswap과 결합되어 다음 작업을 실행합니다.

  • Kueue는 정의된 우선순위와 사용 가능한 리소스를 기반으로 low-jax-trilliumhigh-jax-trillium JobSet의 클러스터 대기열에 대한 승인을 관리합니다.
  • high-jax-trillium JobSet이 인프라 장애로 인해 중단되면 Hotswap이 low-jax-trillium JobSet을 선점하여 우선순위가 높은 JobSet을 다시 예약합니다.
  • 핫스왑을 사용하면 우선순위가 높은 JobSet이 빠르게 다시 시작되어 유휴 시간이 최소화됩니다.
  • 인프라가 복구되면 Hotswap이 복구된 노드 풀에서 낮은 우선순위의 JobSet을 다시 예약합니다.

다음 단계