GKE AI/ML 워크로드 우선순위 지정 권장사항

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

이 가이드에서는 플랫폼 관리자, 운영자, 데이터 및 AI 전문가를 위한 권장사항을 제공합니다. 모든 GKE 권장사항의 통합 개요는 GKE 권장사항을 참고하세요.

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을 다시 예약합니다.

다음 단계