이 튜토리얼에서는 Google Kubernetes Engine(GKE)에서 여러 멀티슬라이스 워크로드를 조정하는 방법을 보여줍니다. TPU 멀티슬라이스, JobSet, Kueue를 사용하여 Jax 워크로드를 실행합니다. Kueue는 작업을 큐에 추가하여 여러 팀 간에 리소스가 공정하게 공유되도록 할당량과 계층 구조에 따라 작업의 대기 시기 및 시작 시기를 결정합니다.
TPU 리소스를 동시에 실행해야 하는 워크로드를 사용하는 경우 이 튜토리얼을 완료하는 것이 좋습니다.
GKE에서 TPU를 사용하기 전에 다음 학습 과정을 완료하는 것이 좋습니다.
- Cloud TPU 시스템 아키텍처를 사용하는 현재 TPU 버전 가용성 알아보기
- GKE의 TPU 멀티슬라이스 알아보기
목표
이 튜토리얼은 기존 GKE 클러스터가 있고 멀티슬라이스 워크로드를 처음 실행하려는 GKE 관리자를 대상으로 합니다.
이 가이드는 다음 과정을 다룹니다.
- 3개의 v5e TPU 슬라이스가 있는 GKE 클러스터로 환경을 준비합니다. 각 TPU 슬라이스에는 8개의 칩이 있는
2x4토폴로지가 있습니다. 따라서 TPU v5e TPU 칩은 총 24개입니다. - Kueue 리소스를 만들어 할당량이 워크로드 간에 공정하게 공유되도록 합니다.
- 멀티슬라이스 워크로드를 실행합니다.
시작하기 전에
시작하기 전에 다음 태스크를 수행했는지 확인합니다.
- Google Kubernetes Engine API를 사용 설정합니다. Google Kubernetes Engine API 사용 설정
- 이 태스크에 Google Cloud CLI를 사용하려면 gcloud CLI를 설치한 후 초기화합니다. 이전에 gcloud CLI를 설치한 경우
gcloud components update를 실행하여 최신 버전을 가져옵니다.
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_IDNODE_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_IDnodepool3라는 세 번째 노드 풀을 만듭니다.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-queuekueue.yaml매니페스트를 적용합니다.kubectl apply -f kueue.yamlGKE에서 다음 Kueue 리소스를 만듭니다.
- ResourceFlavor: 클러스터에 있는 리소스 추상화입니다. 이 예시에서 GKE는
2x4토폴로지를 사용하여 3개의 TPU 슬라이스를 만듭니다. 각 TPU 슬라이스에는 칩이 8개(총 TPU 칩 24개) 있는2x4토폴로지가 있습니다. - ClusterQueue: 워크로드 및 클러스터 리소스를 관리하는 전역 큐입니다.
- LocalQueue: 일반적으로 단일 테넌트(사용자)가 실행하는 밀접하게 관련된 워크로드를 그룹화합니다. 각 LocalQueue는 워크로드 실행을 위해 리소스가 할당된 ClusterQueue를 가리킵니다. Kueue 워크로드는 일괄 워크로드를 나타내는 추상화이며, 이 경우 각 워크로드는 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: 4jobsets-multislice.yaml매니페스트를 적용합니다.kubectl apply -f jobsets-multislice.yaml
GKE는 다음 리소스 요청으로 작업을 만듭니다.
multislice-1sliceJobSet는 TPU 슬라이스가 총 1개 필요한 작업 1개를 만듭니다.multislice-2sliceJobSet는 TPU 슬라이스가 총 2개 필요한 작업 2개를 만듭니다.multislice-3sliceJobSet는 TPU 슬라이스가 총 3개 필요한 작업 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 2m30sGKE가 먼저
multislice-3slice의 포드를 예약, 생성, 실행했는지 확인합니다. 그런 다음 GKE가multislice-1slice및multislice-2sliceJobSet에서 포드를 실행했습니다.
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 계정에 청구되지 않도록 하려면 리소스가 포함된 프로젝트를 삭제하거나 프로젝트를 유지하고 개별 리소스를 삭제하세요.
프로젝트 삭제
- In the Google Cloud console, go to the Manage resources page.
- In the project list, select the project that you want to delete, and then click Delete.
- In the dialog, type the project ID, and then click Shut down to delete the project.
개별 리소스 삭제
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 resourceflavorKueue 매니페스트를 삭제합니다.
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의 네임스페이스 간 할당량 공유로 작업 큐 시스템 구현 방법 알아보기