JobSet 및 Kueue를 사용하여 멀티슬라이스 워크로드 조정


이 튜토리얼에서는 Google Kubernetes Engine (GKE)에서 여러 멀티슬라이스 워크로드를 조정하여 리소스 사용률을 개선하는 방법을 보여줍니다. 예로 Jax 워크로드를 배포하고 TPU 멀티슬라이스에서 실행한 후 JobSet 및 Kueue로 작업 큐잉을 구현합니다. Kueue는 사용 가능한 리소스, 할당량, 팀 간의 공정한 공유를 위한 계층 구조를 기반으로 작업 실행 시기를 결정합니다.

이 튜토리얼은 LLM을 학습하기 위한 Kubernetes의 컨테이너 조정 기능에 관심이 있는 머신러닝 (ML) 엔지니어, 플랫폼 관리자, 운영자를 대상으로 합니다. Google Cloud 콘텐츠에서 참조하는 일반적인 역할 및 예시 작업에 대해 자세히 알아보려면 일반 GKE Enterprise 사용자 역할 및 작업을 참고하세요.

이 페이지를 읽기 전에 다음 사항을 숙지해야 합니다.

목표

  1. 3개의 v5e TPU 슬라이스가 있는 GKE 클러스터로 환경을 준비합니다. 각 TPU 슬라이스에는 칩이 8개 있는 2x4 토폴로지가 있습니다. 따라서 총 24개의 TPU v5e TPU 칩이 있습니다.
  2. 할당량이 워크로드 간에 공정하게 공유되도록 Kueue 리소스를 만듭니다.
  3. 멀티슬라이스 워크로드를 실행합니다.

시작하기 전에

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

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

환경 준비

  1. Google Cloud 콘솔에서 Cloud Shell 인스턴스를 시작합니다.
    Cloud Shell 열기

  2. 기본 환경 변수를 설정합니다.

    gcloud config set project PROJECT_ID
    gcloud config set compute/region COMPUTE_REGION
    

    다음 값을 바꿉니다.

버전 1.29.2-gke.1521000 이상을 실행하는 Autopilot 클러스터는 기본적으로 TPU를 사용 설정합니다. Autopilot 클러스터의 TPU는 워크로드 사양에서 구성됩니다. 자세한 내용은 JobSet으로 멀티슬라이스 워크로드 정의 섹션을 참조하세요.

GKE 클러스터 만들기

Cloud Shell에서 GKE 클러스터를 만드세요.

Autopilot

gcloud container clusters create-auto multislice-cluster \
    --location=LOCATION \
    --cluster-version 1.29.2-gke.1521000 \
    --release-channel rapid

표준

gcloud container clusters create multislice-cluster \
    --location=LOCATION

LOCATION을 클러스터를 만들려는 위치로 바꿉니다. ct5lp-hightpu-4t 머신 유형에 대한 용량이 있는지 확인합니다. 클러스터 생성에는 몇 분 정도 걸릴 수 있습니다.

GKE Autopilot 모드를 사용하는 경우 Kueue 리소스 만들기 섹션으로 건너뜁니다. 버전 1.29.2-gke.1521000 이상을 실행하는 Autopilot 클러스터는 기본적으로 TPU를 사용 설정합니다.

Standard 모드 TPU 슬라이스 노드 풀 3개 만들기

  1. nodepool1이라는 첫 번째 노드 풀을 만듭니다.

    gcloud beta container node-pools create nodepool1 \
        --location=LOCATION \
        --cluster=multislice-cluster \
        --node-locations=NODE_LOCATION \
        --machine-type=ct5lp-hightpu-4t \
        --tpu-topology=2x4 \
        --num-nodes=2 \
        --project=PROJECT_ID
    

    NODE_LOCATION을 노드를 만들려는 클러스터 리전의 하나 이상의 영역으로 바꿉니다.

  2. nodepool2라는 두 번째 노드 풀을 만듭니다.

    gcloud beta container node-pools create nodepool2 \
        --location=LOCATION \
        --cluster=multislice-cluster \
        --node-locations=NODE_LOCATION \
        --machine-type=ct5lp-hightpu-4t \
        --tpu-topology=2x4 \
        --num-nodes=2 \
        --project=PROJECT_ID
    
  3. nodepool3라는 세 번째 노드 풀을 만듭니다.

    gcloud beta container node-pools create nodepool3 \
        --location=LOCATION \
        --cluster=multislice-cluster \
        --node-locations=NODE_LOCATION \
        --machine-type=ct5lp-hightpu-4t \
        --tpu-topology=2x4 \
        --num-nodes=2 \
        --project=PROJECT_ID
    

GKE에서 3개의 노드 풀을 만듭니다. 각 노드 풀은 별도의 TPU 슬라이스입니다.

Kueue 리소스 만들기

  1. 다음 kueue.yaml 매니페스트를 만듭니다.

    apiVersion: kueue.x-k8s.io/v1beta1
    kind: ResourceFlavor
    metadata:
      name: "vlp-24"
    spec:
      nodeLabels:
        cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
        cloud.google.com/gke-tpu-topology: 2x4
    ---
    apiVersion: kueue.x-k8s.io/v1beta1
    kind: ClusterQueue
    metadata:
      name: "cluster-queue"
    spec:
      namespaceSelector: {}
      queueingStrategy: BestEffortFIFO
      resourceGroups:
      - coveredResources: ["google.com/tpu"]
        flavors:
        - name: "vlp-24"
          resources:
          - name: "google.com/tpu"
            nominalQuota: 24
    
    ---
    apiVersion: kueue.x-k8s.io/v1beta1
    kind: LocalQueue
    metadata:
      namespace: default
      name: multislice-queue
    spec:
      clusterQueue: cluster-queue
    
  2. kueue.yaml 매니페스트를 적용합니다.

    kubectl apply -f kueue.yaml
    

    GKE는 다음과 같은 Kueue 리소스를 만듭니다.

  • ResourceFlavor: 클러스터의 리소스를 추상화한 것입니다. 이 예시에서 GKE는 2x4 토폴로지로 TPU 슬라이스 3개를 만듭니다. 각 TPU 슬라이스에는 칩이 8개(총 TPU 칩 24개) 있는 2x4 토폴로지가 있습니다.
  • ClusterQueue: 워크로드와 클러스터 리소스를 관리하는 전역 대기열입니다.
  • LocalQueue: 일반적으로 단일 테넌트(사용자)가 실행하는 밀접하게 관련된 워크로드를 그룹화합니다. 각 LocalQueue는 워크로드를 실행하기 위해 리소스가 할당된 ClusterQueue를 가리킵니다. Kueue 워크로드는 일괄 워크로드를 나타내는 추상화입니다. 이 경우 각 워크로드는 JobSet입니다.

JobSet로 멀티슬라이스 워크로드 정의

이 섹션에서는 JobSet 3개를 만듭니다. 이러한 JobSet는 슬라이스에 전역 TPU 칩 수를 출력하는 Jax 워크로드를 실행한 후 60초 동안 대기하여 일부 모델 학습 시간을 시뮬레이션한 후 종료됩니다.

  1. 다음 jobsets-multislice.yaml 매니페스트를 만듭니다.

    Autopilot

    apiVersion: jobset.x-k8s.io/v1alpha2
    kind: JobSet
    metadata:
      name: multislice-1slice
      labels:
        kueue.x-k8s.io/queue-name: multislice-queue
      annotations:
        alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
    spec:
      failurePolicy:
        maxRestarts: 4
      replicatedJobs:
        - name: slice
          replicas: 1
          template:
            spec:
              parallelism: 2
              completions: 2
              backoffLimit: 0
              template:
                spec:
                  nodeSelector:
                    cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
                    cloud.google.com/gke-tpu-topology: 2x4
                  containers:
                  - name: jax-tpu
                    image: python:3.8
                    ports:
                    - containerPort: 8471
                    - containerPort: 8080
                    command:
                    - bash
                    - -c
                    - |
                      pip install "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
                      python -c 'import jax; print("Global device count:", jax.device_count())'
                    resources:
                      limits:
                        google.com/tpu: 4
    
    ---
    apiVersion: jobset.x-k8s.io/v1alpha2
    kind: JobSet
    metadata:
      name: multislice-2slice
      labels:
        kueue.x-k8s.io/queue-name: multislice-queue
      annotations:
        alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
    spec:
      failurePolicy:
        maxRestarts: 4
      replicatedJobs:
        - name: slice
          replicas: 2
          template:
            spec:
              parallelism: 2
              completions: 2
              backoffLimit: 0
              template:
                spec:
                  nodeSelector:
                    cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
                    cloud.google.com/gke-tpu-topology: 2x4
                  containers:
                  - name: jax-tpu
                    image: python:3.8
                    ports:
                    - containerPort: 8471
                    - containerPort: 8080
                    command:
                    - bash
                    - -c
                    - |
                      pip install "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
                      python -c 'import jax; print("Global device count:", jax.device_count())'
                      sleep 60
                    resources:
                      limits:
                        google.com/tpu: 4
    ---
    apiVersion: jobset.x-k8s.io/v1alpha2
    kind: JobSet
    metadata:
      name: multislice-3slice
      labels:
        kueue.x-k8s.io/queue-name: multislice-queue
      annotations:
        alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
    spec:
      failurePolicy:
        maxRestarts: 4
      replicatedJobs:
        - name: slice
          replicas: 3
          template:
            spec:
              parallelism: 2
              completions: 2
              backoffLimit: 0
              template:
                spec:
                  nodeSelector:
                    cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
                    cloud.google.com/gke-tpu-topology: 2x4
                  containers:
                  - name: jax-tpu
                    image: python:3.8
                    ports:
                    - containerPort: 8471
                    - containerPort: 8080
                    command:
                    - bash
                    - -c
                    - |
                      sleep 60
                    resources:
                      limits:
                        google.com/tpu: 4
    

    표준

    apiVersion: jobset.x-k8s.io/v1alpha2
    kind: JobSet
    metadata:
      name: multislice-1slice
      labels:
        kueue.x-k8s.io/queue-name: multislice-queue
      annotations:
        alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
    spec:
      failurePolicy:
        maxRestarts: 4
      replicatedJobs:
        - name: slice
          replicas: 1
          template:
            spec:
              parallelism: 2
              completions: 2
              backoffLimit: 0
              template:
                spec:
                  hostNetwork: true
                  dnsPolicy: ClusterFirstWithHostNet
                  nodeSelector:
                    cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
                    cloud.google.com/gke-tpu-topology: 2x4
                  containers:
                  - name: jax-tpu
                    image: python:3.8
                    ports:
                    - containerPort: 8471
                    - containerPort: 8080
                    securityContext:
                      privileged: true
                    command:
                    - bash
                    - -c
                    - |
                      pip install "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
                      python -c 'import jax; print("Global device count:", jax.device_count())'
                    resources:
                      limits:
                        google.com/tpu: 4
    
    ---
    apiVersion: jobset.x-k8s.io/v1alpha2
    kind: JobSet
    metadata:
      name: multislice-2slice
      labels:
        kueue.x-k8s.io/queue-name: multislice-queue
      annotations:
        alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
    spec:
      failurePolicy:
        maxRestarts: 4
      replicatedJobs:
        - name: slice
          replicas: 2
          template:
            spec:
              parallelism: 2
              completions: 2
              backoffLimit: 0
              template:
                spec:
                  hostNetwork: true
                  dnsPolicy: ClusterFirstWithHostNet
                  nodeSelector:
                    cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
                    cloud.google.com/gke-tpu-topology: 2x4
                  containers:
                  - name: jax-tpu
                    image: python:3.8
                    ports:
                    - containerPort: 8471
                    - containerPort: 8080
                    securityContext:
                      privileged: true
                    command:
                    - bash
                    - -c
                    - |
                      pip install "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
                      python -c 'import jax; print("Global device count:", jax.device_count())'
                      sleep 60
                    resources:
                      limits:
                        google.com/tpu: 4
    ---
    apiVersion: jobset.x-k8s.io/v1alpha2
    kind: JobSet
    metadata:
      name: multislice-3slice
      labels:
        kueue.x-k8s.io/queue-name: multislice-queue
      annotations:
        alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
    spec:
      failurePolicy:
        maxRestarts: 4
      replicatedJobs:
        - name: slice
          replicas: 3
          template:
            spec:
              parallelism: 2
              completions: 2
              backoffLimit: 0
              template:
                spec:
                  hostNetwork: true
                  dnsPolicy: ClusterFirstWithHostNet
                  nodeSelector:
                    cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
                    cloud.google.com/gke-tpu-topology: 2x4
                  containers:
                  - name: jax-tpu
                    image: python:3.8
                    ports:
                    - containerPort: 8471
                    - containerPort: 8080
                    securityContext:
                      privileged: true
                    command:
                    - bash
                    - -c
                    - |
                      sleep 60
                    resources:
                      limits:
                        google.com/tpu: 4
    
  2. jobsets-multislice.yaml 매니페스트를 적용합니다.

    kubectl apply -f jobsets-multislice.yaml
    

GKE는 다음 리소스 요청으로 작업을 만듭니다.

  • multislice-1slice JobSet는 TPU 슬라이스가 총 1개 필요한 작업 1개를 만듭니다.
  • multislice-2slice JobSet는 TPU 슬라이스가 총 2개 필요한 작업 2개를 만듭니다.
  • multislice-3slice JobSet는 TPU 슬라이스가 총 3개 필요한 작업 3개를 만듭니다.

클러스터에 TPU 슬라이스가 3개만 있으므로 모든 JobSet가 한 번에 실행되는 것은 아닙니다. Kueue가 세 가지 multislice-3slice JobSet를 모두 큐에 추가하면 작업은 완료될 때까지 단독으로 실행됩니다. multislice-1slicemultislice-2slice는 나중에 함께 기다렸다가 실행됩니다.

Kueue가 워크로드를 수락했는지 확인

  1. Kueue에서 대기열에 추가된 워크로드를 확인합니다.

    kubectl get workloads
    

    출력은 다음과 비슷합니다.

    NAME                             QUEUE              ADMITTED BY     AGE
    jobset-multislice-1slice-2530a   multislice-queue                   3s
    jobset-multislice-2slice-ffb02   multislice-queue                   4s
    jobset-multislice-3slice-8c695   multislice-queue   cluster-queue   10s
    

Kueue는 필요한 TPU 리소스에 따라 하나 이상의 워크로드를 큐에 추가합니다.

워크로드 모니터링

  1. 실행 중인 포드를 모니터링합니다.

    kubectl get pods
    

    출력은 다음과 비슷합니다.

    NAME                                READY   STATUS      RESTARTS   AGE
    multislice-1slice-slice-0-0-pf2ll   1/1     Running     0          1s
    multislice-1slice-slice-0-1-55g62   1/1     Running     0          1s
    multislice-2slice-slice-0-0-f4hf7   1/1     Running     0          3s
    multislice-2slice-slice-0-1-c8kv7   1/1     Running     0          3s
    multislice-2slice-slice-1-0-7h46t   1/1     Running     0          3s
    multislice-2slice-slice-1-1-lj9hb   1/1     Running     0          3s
    multislice-3slice-slice-0-0-wzq9t   0/1     Completed   0          2m31s
    multislice-3slice-slice-0-1-zf4dp   0/1     Completed   0          2m30s
    multislice-3slice-slice-1-0-hbfn5   0/1     Completed   0          2m31s
    multislice-3slice-slice-1-1-45fgl   0/1     Completed   0          2m30s
    multislice-3slice-slice-2-0-wjbp4   0/1     Completed   0          2m30s
    multislice-3slice-slice-2-1-lwnvs   0/1     Completed   0          2m30s
    

    GKE가 먼저 multislice-3slice의 포드를 예약, 생성, 실행했는지 확인합니다. 그런 다음 GKE가 multislice-1slicemultislice-2slice JobSet에서 포드를 실행했습니다.

Kueue 워크로드 우선순위 및 선점 사용 설정

원하는 경우 대기열에 추가된 워크로드가 Kueue에서 수락되는 순서를 결정하는 Kueue 워크로드 우선순위를 할당할 수 있습니다.

  1. 선점 정책이 포함되도록 ClusterQueue를 업데이트합니다.

    apiVersion: kueue.x-k8s.io/v1beta1
    kind: ResourceFlavor
    metadata:
      name: "vlp-24"
    spec:
      nodeLabels:
        cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
        cloud.google.com/gke-tpu-topology: 2x4
    ---
    apiVersion: kueue.x-k8s.io/v1beta1
    kind: ClusterQueue
    metadata:
      name: "cluster-queue"
    spec:
      namespaceSelector: {}
      resourceGroups:
      - coveredResources: ["google.com/tpu"]
        flavors:
        - name: "vlp-24"
          resources:
          - name: "google.com/tpu"
            nominalQuota: 24
     preemption:
        reclaimWithinCohort: Any
        withinClusterQueue: LowerPriority
    ---
    apiVersion: kueue.x-k8s.io/v1beta1
    kind: LocalQueue
    metadata:
      namespace: default
      name: multislice-queue
    spec:
      clusterQueue: cluster-queue
    
  2. 워크로드에 할당할 각 우선순위 수준에 대해 PriorityClass를 만듭니다.

    apiVersion: scheduling.k8s.io/v1
    kind: PriorityClass
    metadata:
      name: low-priority
    value: 100
    globalDefault: false
    description: "This low priority class should be used for some Pods only."
    
  3. priorityClassName을 JobSet에 할당합니다.

    Autopilot

    apiVersion: jobset.x-k8s.io/v1alpha2
    kind: JobSet
    metadata:
      name: low-priority
      labels:
        kueue.x-k8s.io/queue-name: multislice-queue
      annotations:
        alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
    spec:
      failurePolicy:
        maxRestarts: 4
      replicatedJobs:
        - name: slice
          replicas: 1
          template:
            spec:
              parallelism: 2
              completions: 2
              backoffLimit: 0
              template:
                spec:
                  nodeSelector:
                    cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
                    cloud.google.com/gke-tpu-topology: 2x4
                  priorityClassName: low-priority
                  containers:
                  - name: jax-tpu
                    image: python:3.8
                    ports:
                    - containerPort: 8471
                    - containerPort: 8080
                    command:
                    - bash
                    - -c
                    - |
                      sleep 60
                    resources:
                      limits:
                        google.com/tpu: 4 # Number of TPU chips per worker
    

    표준

    apiVersion: jobset.x-k8s.io/v1alpha2
    kind: JobSet
    metadata:
      name: low-priority
      labels:
        kueue.x-k8s.io/queue-name: multislice-queue
      annotations:
        alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
    spec:
      failurePolicy:
        maxRestarts: 4
      replicatedJobs:
        - name: slice
          replicas: 1
          template:
            spec:
              parallelism: 2
              completions: 2
              backoffLimit: 0
              template:
                spec:
                  hostNetwork: true
                  dnsPolicy: ClusterFirstWithHostNet
                  nodeSelector:
                    cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
                    cloud.google.com/gke-tpu-topology: 2x4
                  priorityClassName: low-priority
                  containers:
                  - name: jax-tpu
                    image: python:3.8
                    ports:
                    - containerPort: 8471
                    - containerPort: 8080
                    securityContext:
                      privileged: true
                    command:
                    - bash
                    - -c
                    - |
                      sleep 60
                    resources:
                      limits:
                        google.com/tpu: 4 # Number of TPU chips per worker
      ```
    

삭제

이 튜토리얼에서 사용된 리소스 비용이 Google Cloud 계정에 청구되지 않도록 하려면 리소스가 포함된 프로젝트를 삭제하거나 프로젝트를 유지하고 개별 리소스를 삭제하세요.

프로젝트 삭제

  1. In the Google Cloud console, go to the Manage resources page.

    Go to Manage resources

  2. In the project list, select the project that you want to delete, and then click Delete.
  3. In the dialog, type the project ID, and then click Shut down to delete the project.

개별 리소스 삭제

  1. Kueue 할당량 시스템을 삭제합니다.

    kubectl delete -n team-a localqueue
    kubectl delete -n team-b localqueue
    kubectl delete clusterqueue
    kubectl delete clusterqueue
    kubectl delete clusterqueue
    kubectl delete resourceflavor
    kubectl delete resourceflavor
    kubectl delete resourceflavor
    
  2. Kueue 매니페스트를 삭제합니다.

    VERSION=kueue.x-k8s.io/v1beta1
    kubectl delete -f \
        https://github.com/kubernetes-sigs/kueue/releases/download/$VERSION/manifests.yaml
    
  3. 다음과 같이 클러스터를 삭제합니다.

    gcloud container clusters delete kueue-cohort --region=COMPUTE_REGION
    

다음 단계