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

Autopilot Standard

이 튜토리얼에서는 Google Kubernetes Engine(GKE) 및 Kueue에서 TPU 멀티슬라이스를 사용하여 Jax 워크로드를 실행하는 방법을 보여줍니다. Kueue는 작업을 큐에 추가하여 여러 팀 간에 리소스가 공정하게 공유되도록 할당량과 계층 구조에 따라 작업의 대기 시기 및 시작 시기를 결정합니다.

이 튜토리얼에서는 TPU 리소스를 동시에 실행해야 하는 여러 멀티슬라이스 워크로드를 조정하는 방법을 보여줍니다.

GKE에서 TPU를 사용하기 전에 다음 학습 과정을 완료하는 것이 좋습니다.

Cloud TPU 시스템 아키텍처에서 현재 TPU 버전 가용성 알아보기
GKE의 TPU 멀티슬라이스 알아보기

목표

이 튜토리얼은 기존 GKE 클러스터가 있고 멀티슬라이스 워크로드를 처음 실행하려는 GKE 관리자를 대상으로 합니다.

이 가이드는 다음 과정을 다룹니다.

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

시작하기 전에

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

Google Kubernetes Engine API를 사용 설정합니다.

Google Kubernetes Engine API 사용 설정

이 태스크에 Google Cloud CLI를 사용하려면 gcloud CLI를 설치한 후 초기화합니다. 이전에 gcloud CLI를 설치한 경우 gcloud components update를 실행하여 최신 버전을 가져옵니다.
참고: 기존 gcloud CLI 설치의 경우 compute/region 및 compute/zone 속성을 설정해야 합니다. 기본 위치를 설정하면 gcloud CLI에서 One of [--zone, --region] must be supplied: Please specify location과 같은 오류를 방지할 수 있습니다.

v0.2.3 이상의 JobSet을 설치합니다.
v0.4.1 이상의 Kueue를 설치합니다.

환경 준비

Google Cloud 콘솔에서 Cloud Shell 인스턴스를 시작합니다.
Cloud Shell 열기
기본 환경 변수를 설정합니다.
```
gcloud config set project PROJECT_ID
gcloud config set compute/region COMPUTE_REGION
```
다음 값을 바꿉니다.
- PROJECT_ID: Google Cloud 프로젝트 ID입니다.
- COMPUTE_REGION: Compute Engine 리전입니다.

버전 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개 만들기

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을 노드를 만들려는 클러스터 리전의 하나 이상의 영역으로 바꿉니다.

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

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 리소스 만들기

다음 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

kueue.yaml 매니페스트를 적용합니다.
```
kubectl apply -f kueue.yaml
```
GKE가 다음과 같은 Kueue 리소스를 만듭니다.

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

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

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

다음 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

jobsets-multislice.yaml 매니페스트를 적용합니다.
```
kubectl apply -f jobsets-multislice.yaml
```

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

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

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

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

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 리소스에 따라 하나 이상의 워크로드를 큐에 추가합니다.

워크로드 모니터링

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

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-1slice 및 multislice-2slice JobSet에서 포드를 실행했습니다.

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

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

선점 정책을 갖도록 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

워크로드에 할당할 각 우선순위 수준에 대해 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."

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 계정에 청구되지 않도록 하려면 리소스가 포함된 프로젝트를 삭제하거나 프로젝트를 유지하고 개별 리소스를 삭제하세요.

프로젝트 삭제

Google Cloud 콘솔에서 리소스 관리 페이지로 이동합니다.
리소스 관리로 이동
프로젝트 목록에서 삭제할 프로젝트를 선택하고 삭제를 클릭합니다.
대화상자에서 프로젝트 ID를 입력한 후 종료를 클릭하여 프로젝트를 삭제합니다.

개별 리소스 삭제

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

Kueue 매니페스트를 삭제합니다.

VERSION=kueue.x-k8s.io/v1beta1
kubectl delete -f \
    https://github.com/kubernetes-sigs/kueue/releases/download/$VERSION/manifests.yaml

다음과 같이 클러스터를 삭제합니다.

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

다음 단계

Kueue 자세히 알아보기
GKE의 네임스페이스 간 할당량 공유로 작업 큐 시스템 구현 방법 알아보기