GKE의 TPU 소개

이제 Google Kubernetes Engine(GKE) 고객은 TPU v4 및 v5e Pod가 포함된 Kubernetes 노드 풀을 만들 수 있습니다. TPU Pod는 고속 상호 연결로 연결된 TPU 기기 모음입니다. 이는 Kubernetes에서 만들고 관리할 수 있는 배포 가능한 최소 컴퓨팅 단위인 Kubernetes 포드와 다릅니다. 전체 TPU Pod가 필요하지 않은 워크로드에서는 TPU 슬라이스라고 하는 전체 TPU Pod의 하위 집합을 사용하면 됩니다. 전체 TPU Pod와 마찬가지로 슬라이스의 각 TPU 기기에 자체 TPU VM이 있습니다. TPU VM 및 연결된 기기를 호스트 또는 TPU 노드라고 합니다. TPU Pod에 대한 자세한 내용은 시스템 아키텍처를 참조하세요.

GKE와 관련하여 사용되는 포드란 용어는 일반적으로 Kubernetes 포드를 의미하므로 혼동을 방지하기 위해 하나 이상의 TPU 기기 모음은 항상 슬라이스로 지칭합니다.

GKE를 사용하는 경우 먼저 GKE 클러스터를 만들어야 합니다.

그런 다음 클러스터에 노드 풀을 추가합니다. GKE 노드 풀은 동일한 속성을 공유하는 VM 모음입니다. TPU 워크로드의 경우 노드 풀은 TPU VM으로 구성됩니다.

노드 풀 유형

GKE는 다음 두 가지 유형의 TPU 노드 풀을 지원합니다.

멀티 호스트 TPU 슬라이스 노드 풀

멀티 호스트 TPU 슬라이스 노드 풀은 상호 연결된 TPU VM이 2개 이상 포함된 노드 풀입니다. 각 VM에는 연결된 TPU 기기가 있습니다. 멀티 호스트 슬라이스의 TPU는 고속 상호 연결(ICI)을 통해 연결됩니다. 멀티 호스트 TPU 슬라이스 노드 풀은 변경할 수 없습니다. 멀티 호스트 슬라이스 노드 풀이 생성된 후에는 노드를 추가할 수 없습니다. 예를 들어 v4-32 노드 풀을 만든 후 나중에 이 노드 풀에 Kubernetes 노드(TPU VM)를 추가할 수 없습니다. GKE 클러스터에 TPU 슬라이스를 더 추가하려면 새 노드 풀을 만들어야 합니다.

멀티 호스트 TPU 슬라이스 노드 풀의 호스트는 단일 원자적 단위로 취급됩니다. GKE가 슬라이스에 하나의 노드를 배포할 수 없으면 슬라이스의 모든 노드를 배포할 수 없습니다.

멀티 호스트 TPU 슬라이스 내의 노드를 복구해야 하는 경우 GKE는 슬라이스의 모든 TPU VM(노드)을 종료하여 워크로드의 모든 Kubernetes 포드를 강제로 제거합니다. 슬라이스의 모든 TPU VM이 시작되어 실행된 후 새 슬라이스의 TPU VM에 Kubernetes 포드를 예약할 수 있습니다.

다음 다이어그램은 v5litepod-16(v5e) 멀티 호스트 TPU 슬라이스의 예시를 보여줍니다. 이 슬라이스에는 4개의 TPU VM이 있습니다. 각 TPU VM에는 고속 상호 연결(ICI)과 연결된 4개의 TPU v5e 칩이 있으며 각 TPU v5e 칩에는 TensorCore 1개가 있습니다.

멀티 호스트 TPU 슬라이스 다이어그램

다음 다이어그램은 TPU v5litepod-16(v5e) 슬라이스 1개(토폴로지: 4x4)와 TPU v5litepod-8(v5e) 슬라이스 1개(토폴로지: 2x4)가 포함된 GKE 클러스터를 보여줍니다.

TPU v5e Pod 다이어그램

멀티 호스트 TPU 슬라이스에서 워크로드를 실행하는 예시는 멀티 호스트 TPU 슬라이스에서 워크로드 실행을 참조하세요.

단일 호스트 TPU 슬라이스 노드 풀

단일 호스트 슬라이스 노드 풀은 하나 이상의 독립적인 TPU VM을 포함하는 노드 풀입니다. 각 VM에는 연결된 TPU 기기가 있습니다. 단일 호스트 슬라이스 노드 풀 내의 VM은 데이터 센터 네트워크(DCN)를 통해 통신할 수 있지만 VM에 연결된 TPU는 상호 연결되지 않습니다. 다음 다이어그램은 7개의 v4-8 머신이 있는 단일 호스트 TPU 슬라이스의 예시를 보여줍니다.

단일 호스트 슬라이스 노드 풀 다이어그램

단일 호스트 TPU 슬라이스에서 워크로드를 실행하는 예시는 TPU 노드에서 워크로드 실행을 참조하세요.

GKE 노드 풀의 TPU 머신 유형

노드 풀을 만들기 전에 워크로드에 필요한 TPU 슬라이스의 TPU 버전과 크기를 선택해야 합니다. TPU v4는 GKE Standard 버전 1.26.1-gke.1500 이상, GKE Standard 버전 1.27.2-gke.2100 이상의 v5e, GKE Standard 버전 1.28.3-gke.1024000 이상의 v5p에서 지원됩니다.

TPU v4, v5e, v5p는 GKE Autopilot 버전 1.29.2-gke.1521000 이상에서 지원됩니다.

여러 TPU 버전의 하드웨어 사양에 대한 자세한 내용은 시스템 아키텍처를 참조하세요. TPU 노드 풀을 만들 때 모델 크기와 필요한 메모리 양에 따라 TPU 슬라이스 크기(TPU 토폴로지)를 선택합니다. 노드 풀을 만들 때 지정하는 머신 유형은 슬라이스의 버전과 크기에 따라 다릅니다.

v5e

다음은 학습 및 추론 사용 사례에 지원되는 TPU v5e 머신 유형 및 토폴로지입니다.

머신 유형 토폴로지 TPU 칩 수 VM 수 권장 사용 사례
ct5lp-hightpu-1t 1x1 1 1 학습, 단일 호스트 추론
ct5lp-hightpu-4t 2x2 4 1 학습, 단일 호스트 추론
ct5lp-hightpu-8t 2x4 8 1 학습, 단일 호스트 추론
ct5lp-hightpu-4t 2x4 8 2 학습, 멀티 호스트 추론
ct5lp-hightpu-4t 4x4 16 4 대규모 학습, 멀티 호스트 추론
ct5lp-hightpu-4t 4x8 32 8 대규모 학습, 멀티 호스트 추론
ct5lp-hightpu-4t 8x8 64 16 대규모 학습, 멀티 호스트 추론
ct5lp-hightpu-4t 8x16 128 32 대규모 학습, 멀티 호스트 추론
ct5lp-hightpu-4t 16x16 256 64 대규모 학습, 멀티 호스트 추론

Cloud TPU v5e는 학습 및 추론 제품 조합입니다. 학습 작업은 처리량과 가용성에 최적화되고 추론 작업은 지연 시간에 최적화됩니다. 자세한 내용은 v5e 학습 가속기 유형 및 v5e 추론 가속기 유형을 참조하세요.

TPU v5e 머신은 us-west4-a, us-east5-b, us-east1-c에서 사용 가능합니다. GKE Standard 클러스터는 제어 영역 버전 1.27.2-gke.2100 이상을 실행해야 합니다. GKE Autopilot은 제어 영역 버전 1.29.2-gke.1521000 이상을 실행해야 합니다. v5e에 대한 자세한 내용은 Cloud TPU v5e 학습을 참조하세요.

머신 유형 비교:

머신 유형 ct5lp-hightpu-1t ct5lp-hightpu-4t ct5lp-hightpu-8t
v5e 칩 수 1 4 8
vCPU 수 24 112 224
RAM(GB) 48 192 384
NUMA 노드 수 1 1 2
선점 가능성 높음 보통 성과자 낮음

칩이 많은 VM의 공간을 확보하기 위해 GKE 스케줄러에서 칩이 적은 VM을 선점하고 다시 예약할 수 있습니다. 따라서 칩이 8개인 VM에서 칩 1개 및 칩 4개로 구성된 VM을 선점할 가능성이 높습니다.

v4 및 v5p

TPU v4 및 v5p 머신 유형은 다음과 같습니다.

머신 유형 vCPU 수 메모리(GB) NUMA 노드 수
ct4p-hightpu-4t 240 407 2
ct5p-hightpu-4t 208 448 2

TPU v4 슬라이스를 만들 때 호스트가 하나 있고 칩이 4개 포함된 ct4p-hightpu-4t 머신 유형을 사용합니다. 자세한 내용은 v4 토폴로지 및 TPU 시스템 아키텍처를 참조하세요. TPU v4 Pod 머신 유형은 us-central2-b에서 사용할 수 있습니다. GKE Standard 클러스터는 제어 영역 버전 1.26.1-gke.1500 이상을 실행해야 합니다. GKE Autopilot 클러스터는 제어 영역 버전 1.29.2-gke.1521000 이상을 실행해야 합니다.

TPU v5p 슬라이스를 만들 때 호스트가 하나 있고 칩이 4개 포함된 ct5p-hightpu-4t 머신 유형을 사용합니다. TPU v5p 포드 머신 유형은 us-west4-a 및 us-east5-a에서 사용 가능합니다. GKE Standard 클러스터는 제어 영역 버전 1.28.3-gke.1024000 이상을 실행해야 합니다. GKE Autopilot은 1.29.2-gke.1521000 이상을 실행해야 합니다. v5p에 대한 자세한 내용은 v5p 학습 소개를 참조하세요.

알려진 문제 및 제한사항

  • 최대 Kubernetes 포드 수: 단일 TPU VM에서 Kubernetes 포드를 최대 256개까지 실행할 수 있습니다.
  • SPECIFIC 예약만 지원: GKE에서 TPU를 사용하는 경우 SPECIFIC이 gcloud container node-pools create 명령어의 --reservation-affinity 플래그에 지원되는 유일한 값입니다.
  • 선점형 TPU의 스팟 VM 변형만 지원: 스팟 VM은 선점형 VM과 유사하며 동일한 가용성 제한이 적용되지만 최대 기간이 24시간입니다.
  • 비용 할당 지원 없음: GKE 비용 할당 및 사용량 측정에는 사용량 또는 TPU 비용에 대한 데이터가 포함되지 않습니다.
  • 자동 확장 처리에서 용량을 잘못 계산할 수 있음: 클러스터 자동 확장 처리에서 새 TPU 노드가 제공되기 전에 해당 노드의 용량을 잘못 계산할 수 있습니다. 이후 클러스터 자동 확장 처리에서 추가 수직 확장을 수행하여 결과적으로 필요한 것보다 많은 노드가 생성될 수 있습니다. 클러스터 자동 확장 처리에서는 일반적인 축소 작업 후 더 이상 필요하지 않은 추가 노드를 축소합니다.
  • 자동 확장 처리에서 확장 취소: 클러스터 자동 확장 처리에서는 대기 상태로 10시간 이상 유지되는 TPU 노드 풀의 수직 확장을 취소합니다. 클러스터 자동 확장 처리는 나중에 이러한 수직 확장 작업을 다시 시도합니다. 이 동작으로 인해 예약을 사용하지 않는 고객의 TPU 획득 가능성이 줄어들 수 있습니다.
  • Taint에서 축소를 막을 수 있음: TPU taint에 대한 톨러레이션(toleration)이 있는 TPU 외의 워크로드는 TPU 노드 풀 드레이닝 중에 다시 생성될 경우 노드 풀의 축소를 막을 수 있습니다.

TPU 및 GKE 할당량이 충분한지 확인

리소스가 생성된 리전에서 특정 GKE 관련 할당량을 늘려야 할 수 있습니다.

다음 할당량의 기본값을 상향 조정해야 할 수 있습니다.

  • Persistent Disk SSD(GB) 할당량: 각 Kubernetes 노드의 부팅 디스크에는 기본적으로 100GB가 필요합니다. 따라서 이 할당량은 (생성될 것으로 예상되는 최대 GKE 노드 수) * 100GB 이상으로 설정해야 합니다.
  • 사용 중인 IP 주소 할당량: 각 Kubernetes 노드는 하나의 IP 주소를 사용합니다. 따라서 이 할당량은 생성될 것으로 예상되는 최대 GKE 노드 수 이상으로 설정해야 합니다.

할당량 상향을 요청하려면 할당량 상향 요청을 참조하세요. TPU 할당량 유형에 대한 자세한 내용은 TPU 할당량을 참조하세요.

할당량 상향 요청이 승인되는 데 며칠이 걸릴 수 있습니다. 며칠 내로 할당량 상향 요청을 승인받는 데 문제가 있는 경우 Google 계정팀에 문의하세요.

TPU 예약 마이그레이션

GKE에서 TPU와 함께 기존 TPU 예약을 사용하지 않으려면 이 섹션을 건너뛰고 Google Kubernetes Engine 클러스터 만들기로 이동합니다.

GKE에서 예약된 TPU를 사용하려면 먼저 TPU 예약을 새 Compute Engine 기반 예약 시스템으로 마이그레이션해야 합니다.

마이그레이션에 대해 알아야 할 몇 가지 중요한 사항은 다음과 같습니다.

  • 새로운 Compute Engine 기반 예약 시스템으로 마이그레이션된 TPU 용량은 Cloud TPU Queued Resource API와 함께 사용할 수 없습니다. TPU 큐에 추가된 리소스를 예약에 사용하려면 TPU 예약의 일부만 새 Compute Engine 기반 예약 시스템으로 마이그레이션해야 합니다.
  • 새로운 Compute Engine 기반 예약 시스템으로 마이그레이션되면 TPU에서 워크로드를 활발하게 실행할 수 없습니다.
  • 마이그레이션을 수행할 시간을 선택하고 Google Cloud 계정팀과 함께 마이그레이션을 예약합니다. 마이그레이션 기간이 운영 시간(월요일~금요일, 오전 9시~오후 5시(태평양 표준시)) 사이여야 합니다.

Google Kubernetes Engine 클러스터 만들기

Google Kubernetes Engine 문서의 클러스터 만들기를 참조하세요.

TPU 노드 풀 만들기

Google Kubernetes Engine 문서에서 노드 풀 만들기를 참조하세요.

권한이 있는 모드 없이 실행

컨테이너의 권한 범위를 줄이려면 TPU 권한 모드를 참조하세요.

TPU 노드에서 워크로드 실행

Google Kubernetes Engine 문서의 TPU 노드에서 워크로드 실행을 참조하세요.

노드 선택기

Kubernetes가 TPU 노드에서 워크로드를 예약하려면 Google Kubernetes Engine 매니페스트에서 각 TPU 노드에 선택기 두 개를 지정해야 합니다.

  • cloud.google.com/gke-accelerator-type를 tpu-v5-lite-podslice 또는 tpu-v4-podslice로 설정합니다.
  • cloud.google.com/gke-tpu-topology를 TPU 노드의 TPU 토폴로지로 설정합니다.

학습 워크로드 및 추론 워크로드 섹션에 이러한 노드 선택기를 사용하는 방법을 보여주는 매니페스트 예시가 포함되어 있습니다.

워크로드 예약 고려사항

TPU에는 Kubernetes에서 특별한 워크로드 예약 및 관리가 필요한 고유한 특성이 있습니다. 자세한 내용은 GKE 문서의 워크로드 예약 고려사항을 참조하세요.

TPU 노드 복구

멀티 호스트 TPU 슬라이스 노드 풀의 TPU 노드가 비정상이면 전체 노드 풀이 다시 생성됩니다. 자세한 내용은 GKE 문서의 노드 자동 복구를 참조하세요.

멀티슬라이스 - 단일 슬라이스 이상으로 확장

더 큰 학습 워크로드를 처리하기 위해 여러 작은 슬라이스를 하나의 멀티슬라이스로 집계할 수 있습니다. 자세한 내용은 Cloud TPU 멀티슬라이스를 참조하세요.

학습 워크로드 튜토리얼

이 튜토리얼은 멀티 호스트 TPU 슬라이스의 학습 워크로드 (예: 4개의 v5e 머신)를 중점적으로 살펴봅니다. 다음 모델을 다룹니다.

  • Hugging Face FLAX 모델: Pokémon에서 확산 학습
  • PyTorch/XLA: WikiText의 GPT2

튜토리얼 리소스 다운로드

다음 명령어를 사용하여 선행 학습된 각 모델의 튜토리얼 Python 스크립트 및 YAML 사양을 다운로드합니다.

git clone https://github.com/GoogleCloudPlatform/ai-on-gke.git

클러스터 만들기 및 연결

Kubernetes 제어 영역이 영역 3개에 복제되어 높은 가용성을 제공하도록 리전 GKE 클러스터를 만듭니다. 사용하는 TPU 버전에 따라 us-west4, us-east1 또는 us-central2에 클러스터를 만듭니다. TPU 및 영역에 대한 자세한 내용은 Cloud TPU 리전 및 영역을 참조하세요.

다음 명령어는 처음에 영역당 노드 하나가 포함된 노드 풀이 있는 빠른 출시 채널을 구독하는 새 GKE 리전 클러스터를 만듭니다. 이 가이드의 예시 추론 워크로드가 Cloud Storage 버킷을 사용해 선행 학습된 모델을 저장하므로 이 명령어는 클러스터에서 워크로드 아이덴티티 및 Cloud Storage FUSE CSI 드라이버 기능을 사용 설정합니다.

gcloud container clusters create cluster-name \
  --region your-region \
  --release-channel rapid \
  --num-nodes=1 \
  --workload-pool=project-id.svc.id.goog \
  --addons GcsFuseCsiDriver

기존 클러스터에 워크로드 아이덴티티 및 Cloud Storage FUSE CSI 드라이버 기능을 사용 설정하려면 다음 명령어를 실행합니다.

gcloud container clusters update cluster-name \
  --region your-region \
  --update-addons GcsFuseCsiDriver=ENABLED \
  --workload-pool=project-id.svc.id.goog

이 예시 워크로드는 다음과 같은 가정으로 구성됩니다.

  • 노드 풀에서 노드 4개와 함께 tpu-topology=4x4를 사용
  • 노드 풀에서 machine-type ct5lp-hightpu-4t를 사용

다음 명령어를 실행하여 새로 만든 클러스터에 연결합니다.

gcloud container clusters get-credentials cluster-name \
--location=cluster-region

Hugging Face FLAX 모델: Pokémon에서 확산 학습

이 예시에서는 Pokémon 데이터 세트를 사용하여 Hugging Face의 안정적인 확산 모델을 학습시킵니다.

안정적인 확산 모델은 모든 텍스트 입력에서 실사 이미지를 생성하는 잠재 텍스트 이미지 변환 모델입니다. 안정적인 확산에 대한 자세한 내용은 다음을 참조하세요.

Docker 이미지 만들기

Dockerfile은 ai-on-gke/tutorials-and-examples/tpu-examples/training/diffusion/ 폴더에 있습니다. 다음 명령어를 실행하여 Docker 이미지를 빌드하고 푸시합니다.

cd ai-on-gke/tutorials-and-examples/tpu-examples/training/diffusion/
docker build -t gcr.io/project-id/diffusion:latest .
docker push gcr.io/project-id/diffusion:latest

워크로드 배포

다음 콘텐츠로 파일을 만들고 이름을 tpu_job_diffusion.yaml로 지정합니다. 이미지 필드에 방금 만든 이미지를 입력합니다.

apiVersion: v1
kind: Service
metadata:
  name: headless-svc
spec:
  clusterIP: None
  selector:
    job-name: tpu-job-diffusion
---
apiVersion: batch/v1
kind: Job
metadata:
  name: tpu-job-diffusion
spec:
  backoffLimit: 0
  # Completions and parallelism should be the number of chips divided by 4.
  # (e.g. 4 for a v5litepod-16)
  completions: 4
  parallelism: 4
  completionMode: Indexed
  template:
    spec:
      subdomain: headless-svc
      restartPolicy: Never
      nodeSelector:
        cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
        cloud.google.com/gke-tpu-topology: 4x4
      containers:
      - name: tpu-job-diffusion
        image: gcr.io/${project-id}/diffusion:latest
        ports:
        - containerPort: 8471 # Default port using which TPU VMs communicate
        - containerPort: 8431 # Port to export TPU usage metrics, if supported
        command:
        - bash
        - -c
        - |
          cd examples/text_to_image
          python3 train_text_to_image_flax.py --pretrained_model_name_or_path=duongna/stable-diffusion-v1-4-flax --dataset_name=lambdalabs/pokemon-blip-captions --resolution=128 --center_crop --random_flip --train_batch_size=4 --mixed_precision=fp16 --max_train_steps=1500 --learning_rate=1e-05 --max_grad_norm=1 --output_dir=sd-pokemon-model
        resources:
          requests:
            google.com/tpu: 4
          limits:
            google.com/tpu: 4

그런 후 다음을 사용하여 배포합니다.

kubectl apply -f tpu_job_diffusion.yaml

삭제

작업 실행이 완료되면 다음을 사용하여 삭제할 수 있습니다.

kubectl delete -f tpu_job_diffusion.yaml

PyTorch/XLA: WikiText의 GPT2

이 튜토리얼에서는 wikitext 데이터 세트를 사용하여 PyTorch/XLA에서 Hugging Face를 사용하여 v5e TPU에서 GPT2를 실행하는 방법을 보여줍니다.

Docker 이미지 만들기

Dockerfile은 ai-on-gke/tutorials-and-examples/tpu-examples/training/gpt/ 폴더에 있습니다. 다음 명령어를 실행하여 Docker 이미지를 빌드하고 푸시합니다.

cd ai-on-gke/tutorials-and-examples/tpu-examples/training/gpt/
docker build -t gcr.io/project-id/gpt:latest .
docker push gcr.io/project-id/gpt:latest

워크로드 배포

다음 YAML을 복사하여 tpu_job_gpt.yaml 파일에 저장합니다. 이미지 필드에 방금 만든 이미지를 입력합니다.

apiVersion: v1
kind: Service
metadata:
  name: headless-svc
spec:
  clusterIP: None
  selector:
    job-name: tpu-job-gpt
---
apiVersion: batch/v1
kind: Job
metadata:
  name: tpu-job-gpt
spec:
  backoffLimit: 0
  # Completions and parallelism should be the number of chips divided by 4.
  # (for example, 4 for a v5litepod-16)
  completions: 4
  parallelism: 4
  completionMode: Indexed
  template:
    spec:
      subdomain: headless-svc
      restartPolicy: Never
      volumes:
      # Increase size of tmpfs /dev/shm to avoid OOM.
      - name: shm
        emptyDir:
          medium: Memory
          # consider adding `sizeLimit: XGi` depending on needs
      nodeSelector:
        cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
        cloud.google.com/gke-tpu-topology: 4x4
      containers:
      - name: tpu-job-gpt
        image: gcr.io/$(project-id)/gpt:latest
        ports:
        - containerPort: 8479
        - containerPort: 8478
        - containerPort: 8477
        - containerPort: 8476
        - containerPort: 8431 # Port to export TPU usage metrics, if supported.
        env:
        - name: PJRT_DEVICE
          value: 'TPU'
        - name: XLA_USE_BF16
          value: '1'
        command:
        - bash
        - -c
        - |
          numactl --cpunodebind=0 python3 -u examples/pytorch/xla_spawn.py   --num_cores 4 examples/pytorch/language-modeling/run_clm.py    --num_train_epochs 3 --dataset_name wikitext     --dataset_config_name wikitext-2-raw-v1 --per_device_train_batch_size 16    --per_device_eval_batch_size 16 --do_train --do_eval  --output_dir /tmp/test-clm     --overwrite_output_dir --config_name my_config_2.json --cache_dir /tmp --tokenizer_name gpt2  --block_size 1024 --optim adafactor --adafactor true --save_strategy no --logging_strategy no --fsdp "full_shard" --fsdp_config fsdp_config.json
        volumeMounts:
        - mountPath: /dev/shm
          name: shm
        resources:
          requests:
            google.com/tpu: 4
          limits:
            google.com/tpu: 4

다음을 사용하여 워크플로를 배포합니다.

kubectl apply -f tpu_job_gpt.yaml

삭제

작업 실행이 완료되면 다음을 사용하여 삭제할 수 있습니다.

kubectl delete -f tpu_job_gpt.yaml

튜토리얼: 단일 호스트 추론 워크로드

이 튜토리얼에서는 JAX, TensorFlow, PyTorch를 사용하여 선행 학습된 모델을 위해 GKE v5e TPU에서 단일 호스트 추론 워크로드를 실행하는 방법을 보여줍니다. 대략적으로 GKE 클러스터에서 수행하는 4가지 단계가 있습니다.

  1. Cloud Storage 버킷을 만들고 버킷에 대한 액세스를 설정합니다. Cloud Storage 버킷은 선행 학습된 모델을 저장하는 데 사용됩니다.

  2. 선행 학습된 모델을 다운로드하여 TPU 호환 모델로 변환합니다. 선행 학습된 모델을 다운로드하고 Cloud TPU 변환기를 사용하며 Cloud Storage FUSE CSI 드라이버를 사용하여 변환된 모델을 Cloud Storage 버킷에 저장하는 Kubernetes 포드를 적용합니다. Cloud TPU 변환기에는 특수 하드웨어가 필요하지 않습니다. 이 튜토리얼에서는 모델을 다운로드하고 CPU 노드 풀에서 Cloud TPU 변환기를 실행하는 방법을 보여줍니다.

  3. 변환된 모델의 서버를 실행합니다. ReadOnlyMany(ROX) 영구 볼륨에 저장된 볼륨으로 지원되는 서버 프레임워크를 사용하여 모델을 제공하는 배포를 적용합니다. 배포 복제본은 노드당 하나의 Kubernetes 포드가 있는 v5e Pod TPU 노드에서 실행되어야 합니다.

  4. 모델 서버를 테스트할 부하 분산기를 배포합니다. LoadBalancer 서비스를 사용하여 서버가 외부 요청에 노출됩니다. 모델 서버 테스트 요청에 대한 예시가 포함된 Python 스크립트가 제공됩니다.

다음 다이어그램은 부하 분산기에서 요청을 라우팅하는 방법을 보여줍니다.

부하 분산기 라우팅을 보여주는 다이어그램

서버 배포 예시

이 예시 워크로드는 다음과 같은 가정으로 구성됩니다.

  • 클러스터를 노드가 3개인 TPU v5 노드 풀에서 실행
  • 노드 풀이 다음과 같은 ct5lp-hightpu-1t 머신 유형을 사용
    • 토폴로지는 1x1임
    • TPU 칩 수는 1개임

다음 GKE 매니페스트는 단일 호스트 서버 배포를 정의합니다.

apiVersion: apps/v1
kind: Deployment
metadata:
  name: bert-deployment
spec:
  selector:
    matchLabels:
      app: tf-bert-server
  replicas: 3 # number of nodes in node pool
  template:
    metadata:
      annotations:
        gke-gcsfuse/volumes: "true"
      labels:
        app: tf-bert-server
    spec:
      nodeSelector:
        cloud.google.com/gke-tpu-topology: 1x1  # target topology
        cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice  # target version
      containers:
      - name: serve-bert
        image: us-docker.pkg.dev/cloud-tpu-images/inference/tf-serving-tpu:2.13.0
        env:
        - name: MODEL_NAME
          value: "bert"
        volumeMounts:
        - mountPath: "/models/"
          name: bert-external-storage
        ports:
        - containerPort: 8500
        - containerPort: 8501
        - containerPort: 8431 # Port to export TPU usage metrics, if supported.
        resources:
          requests:
            google.com/tpu: 1 # TPU chip request
          limits:
            google.com/tpu: 1 # TPU chip request
      volumes:
      - name: bert-external-storage
        persistentVolumeClaim:
          claimName: external-storage-pvc

TPU 노드 풀에서 다른 노드 수를 사용하는 경우 replicas 필드를 해당하는 노드 수로 변경합니다.

Standard 클러스터에서 GKE 버전 1.27 이하를 실행하는 경우에는 매니페스트에 다음 필드를 추가합니다. 

spec:
  securityContext:
    privileged: true

GKE 버전 1.28 이상에서는 권한이 있는 모드에서 Kubernetes 포드를 실행할 필요가 없습니다. 자세한 내용은 권한이 있는 모드 없이 컨테이너 실행을 참조하세요.

다른 머신 유형을 사용하는 경우 다음 안내를 따르세요.

  • cloud.google.com/gke-tpu-topology를 사용 중인 머신 유형의 토폴로지로 설정합니다.
  • resources 아래의 두 google.com/tpu 필드를 해당 머신 유형의 칩 수와 일치하도록 설정합니다.

설정

다음 명령어를 사용하여 튜토리얼 Python 스크립트 및 YAML 매니페스트를 다운로드합니다.

git clone https://github.com/GoogleCloudPlatform/ai-on-gke.git

single-host-inference 디렉터리로 이동합니다.

cd ai-on-gke/gke-tpu-examples/single-host-inference/

Python 환경 설정

이 튜토리얼에서 사용하는 Python 스크립트에는 Python 버전 3.9 이상이 필요합니다. Python 테스트 스크립트를 실행하기 전에 각 튜토리얼에 requirements.txt를 설치해야 합니다.

로컬 환경에서 적절한 Python 설정을 사용할 수 없는 경우 Cloud Shell을 사용해 이 튜토리얼에서 Python 스크립트를 다운로드하고 실행할 수 있습니다.

클러스터 설정

  1. e2-standard-4 머신 유형을 사용하여 클러스터를 만듭니다.

    gcloud container clusters create cluster-name \
--region your-region \
--release-channel rapid \
--num-nodes=1 \
--machine-type=e2-standard-4 \
--workload-pool=project-id.svc.id.goog \
--addons GcsFuseCsiDriver

  2. 단일 호스트 TPU 노드 풀을 만듭니다.

예시 워크로드를 다음과 같이 가정합니다.

  • 클러스터를 노드가 3개인 TPU v5e 노드 풀에서 실행
  • TPU 노드 풀이 머신 유형 ct5lp-hightpu-1t을 사용

앞에서 설명한 것과 다른 클러스터 구성을 사용하는 경우 서버 배포 매니페스트를 수정해야 합니다.

JAX 안정적인 확산 데모의 경우 사용 가능한 메모리가 16Gi 이상인 머신 유형(예: e2-standard-4)을 사용하는 CPU 노드 풀이 필요합니다. 이는 gcloud container clusters create 명령어로 구성하거나 다음 명령어를 사용해 기존 클러스터에 노드 풀을 더 추가하여 구성합니다.

gcloud beta container node-pools create your-pool-name \
  --zone=your-cluster-zone \
  --cluster=your-cluster-name \
  --machine-type=e2-standard-4 \
  --num-nodes=1

다음을 바꿉니다.

  • your-pool-name: 만들 노드 풀의 이름.
  • your-cluster-zone: 클러스터가 생성된 영역입니다.
  • your-cluster-name: 노드 풀을 추가할 클러스터의 이름입니다.
  • your-machine-type: 노드 풀에 만들 노드의 머신 유형입니다.

모델 스토리지 설정

서빙을 위해 모델을 저장할 수 있는 몇 가지 방법이 있습니다. 이 튜토리얼에서는 다음 방법을 사용합니다.

  • 선행 학습된 모델을 TPU에서 작동하도록 변환하기 위해 ReadWriteMany(RWX) 액세스 권한이 있는 Persistent Disk가 지원하는 Virtual Private Cloud를 사용합니다.
  • 여러 단일 호스트 TPU에서 모델을 서빙하기 위해 Cloud Storage 버킷에서 지원되는 동일한 VPC를 사용합니다.

다음 명령어를 실행하여 Cloud Storage 버킷을 만듭니다.

gcloud storage buckets create gs://your-bucket-name \
  --project=your-bucket-project-id \
  --location=your-bucket-location

다음을 바꿉니다.

  • your-bucket-name: Cloud Storage 버킷 이름입니다.
  • your-bucket-project-id: Cloud Storage 버킷을 만든 프로젝트 ID입니다.
  • your-bucket-location: Cloud Storage 버킷 위치입니다. 성능을 향상시키려면 GKE 클러스터가 실행되는 위치를 지정합니다.

다음 단계에 따라 버킷에 GKE 클러스터 액세스 권한을 부여합니다. 다음 예시에서는 기본 네임스페이스 및 기본 Kubernetes 서비스 계정을 사용하여 설정을 간소화합니다. 자세한 내용은 GKE 워크로드 아이덴티티를 사용하여 Cloud Storage 버킷에 대한 액세스 구성을 참조하세요.

  1. 애플리케이션에 대해 IAM 서비스 계정을 만들거나 대신 기존 IAM 서비스 계정을 사용합니다. Cloud Storage 버킷의 프로젝트에서 모든 IAM 서비스 계정을 사용할 수 있습니다.

    gcloud iam service-accounts create your-iam-service-acct \
--project=your-bucket-project-id

    다음을 바꿉니다.

    • your-iam-service-acct: 새 IAM 서비스 계정의 이름입니다.
    • your-bucket-project-id: IAM 서비스 계정을 만든 프로젝트 ID입니다. IAM 서비스 계정은 Cloud Storage 버킷과 동일한 프로젝트에 있어야 합니다.

  2. IAM 서비스 계정에 필요한 스토리지 역할이 있는지 확인합니다.

    gcloud storage buckets add-iam-policy-binding gs://your-bucket-name \
--member "serviceAccount:your-iam-service-acct@your-bucket-project-id.iam.gserviceaccount.com" \
--role "roles/storage.objectAdmin"

    다음을 바꿉니다.

    • your-bucket-name: Cloud Storage 버킷 이름입니다.
    • your-iam-service-acct: 새 IAM 서비스 계정의 이름입니다.
    • your-bucket-project-id: IAM 서비스 계정을 만든 프로젝트 ID입니다.

  3. Kubernetes 서비스 계정과 IAM 서비스 계정 사이에 IAM 정책 바인딩을 추가하여 Kubernetes 서비스 계정이 IAM 서비스 계정을 가장하도록 허용합니다. 이러한 바인딩으로 Kubernetes 서비스 계정이 IAM 서비스 계정 역할을 하도록 허용합니다.

    gcloud iam service-accounts add-iam-policy-binding your-iam-service-acct@your-bucket-project-id.iam.gserviceaccount.com \
  --role roles/iam.workloadIdentityUser \
  --member "serviceAccount:your-project-id.svc.id.goog[default/default]"

    다음을 바꿉니다.

    • your-iam-service-acct: 새 IAM 서비스 계정의 이름입니다.
    • your-bucket-project-id: IAM 서비스 계정을 만든 프로젝트 ID입니다.
    • your-project-id: GKE 클러스터를 만든 프로젝트의 ID입니다. Cloud Storage 버킷과 GKE 클러스터는 동일한 프로젝트 또는 다른 프로젝트에 있을 수 있습니다.

  4. Kubernetes 서비스 계정에 IAM 서비스 계정의 이메일 주소를 주석으로 추가합니다.

    kubectl annotate serviceaccount default \
  --namespace default \
  iam.gke.io/gcp-service-account=your-iam-service-acct@your-bucket-project-id.iam.gserviceaccount.com

    다음을 바꿉니다.

    • your-iam-service-acct: 새 IAM 서비스 계정의 이름입니다.
    • your-bucket-project-id: IAM 서비스 계정을 만든 프로젝트 ID입니다.

  5. 다음 명령어를 실행하여 이 데모의 YAML 파일에 버킷 이름을 입력합니다.

    find . -type f -name "*.yaml" | xargs sed -i "s/BUCKET_NAME/your-bucket-name/g"

    your-bucket-name을 Cloud Storage 버킷의 이름으로 바꿉니다.

  6. 다음 명령어를 사용하여 영구 볼륨 및 영구 볼륨 신청을 만듭니다.

    kubectl apply -f pvc-pv.yaml

JAX 모델 추론 및 서빙

JAX 모델 서비스로 요청을 전송하는 튜토리얼 Python 스크립트를 실행하기 위해 Python 종속 항목을 설치합니다.

pip install -r jax/requirements.txt

JAX BERT E2E 서빙 데모 실행:

이 데모에서는 Hugging Face의 선행 학습된 BERT 모델을 사용합니다.

Kubernetes 포드는 다음 단계를 수행합니다.

  1. 예시 리소스에서 Python 스크립트 export_bert_model.py를 다운로드하고 사용하여 선행 학습된 BERT 모델을 임시 디렉터리에 다운로드합니다.
  2. Cloud TPU 변환기 이미지를 사용하여 선행 학습된 모델을 CPU에서 TPU로 변환하고 설정 중에 만든 Cloud Storage 버킷에 모델을 저장합니다.

이 Kubernetes 포드는 기본 노드 풀 CPU에서 실행되도록 구성되어 있습니다. 다음 명령어를 사용하여 포드를 실행합니다.

kubectl apply -f jax/bert/install-bert.yaml

다음을 사용해 모델이 올바르게 설치되었는지 확인합니다.

kubectl get pods install-bert

STATUS가 Completed를 읽는 데 몇 분 정도 걸릴 수 있습니다.

모델의 TF 모델 서버 실행

이 튜토리얼의 예시 워크로드를 다음과 같이 가정합니다.

  • 클러스터를 노드가 3개인 TPU v5 노드 풀에서 실행
  • 노드 풀이 TPU 칩 1개가 포함된 ct5lp-hightpu-1t 머신 유형을 사용

앞에서 설명한 것과 다른 클러스터 구성을 사용하는 경우 서버 배포 매니페스트를 수정해야 합니다.

배포 적용
kubectl apply -f jax/bert/serve-bert.yaml

다음을 사용하여 서버가 실행 중인지 확인합니다.

kubectl get deployment bert-deployment

AVAILABLE이 3을 읽는 데 1분 정도 걸릴 수 있습니다.

부하 분산기 서비스 적용
kubectl apply -f jax/bert/loadbalancer.yaml

다음을 사용하여 부하 분산기가 외부 트래픽을 지원할 준비가 되었는지 확인합니다.

kubectl get svc tf-bert-service

EXTERNAL_IP가 IP를 나열하는 데 몇 분 정도 걸릴 수 있습니다.

모델 서버로 요청 전송

부하 분산기 서비스에서 외부 IP를 가져옵니다.

EXTERNAL_IP=$(kubectl get services tf-bert-service --output jsonpath='{.status.loadBalancer.ingress[0].ip}')

서버에 요청을 전송하는 스크립트를 실행합니다.

python3 jax/bert/bert_request.py $EXTERNAL_IP

예상 출력:

For input "The capital of France is [MASK].", the result is ". the capital of france is paris.."
For input "Hello my name [MASK] Jhon, how can I [MASK] you?", the result is ". hello my name is jhon, how can i help you?."
삭제

리소스를 삭제하려면 kubectl delete를 역순으로 실행합니다.

kubectl delete -f jax/bert/loadbalancer.yaml
kubectl delete -f jax/bert/serve-bert.yaml
kubectl delete -f jax/bert/install-bert.yaml

JAX 안정적인 확산 E2E 서빙 데모 실행

이 데모에서는 Hugging Face의 선행 학습된 안정적인 확산 모델을 사용합니다.

Flax 안정적인 확산 모델에서 TPU 호환 TF2 저장된 모델 내보내기

안정적인 확산 모델을 내보내려면 클러스터 설정에 설명된 대로 사용 가능한 메모리가 16Gi 이상인 머신 유형을 사용하는 CPU 노드 풀이 클러스터에 있어야 합니다.

Kubernetes 포드는 다음 단계를 실행합니다.

  1. 예시 리소스에서 Python 스크립트 export_stable_diffusion_model.py를 다운로드하고 사용하여 선행 학습된 안정적인 확산 모델을 임시 디렉터리에 다운로드합니다.
  2. Cloud TPU 변환기 이미지를 사용하여 선행 학습된 모델을 CPU에서 TPU로 변환하고 스토리지 설정 중에 만든 Cloud Storage 버킷에 모델을 저장합니다.

이 Kubernetes 포드는 기본 CPU 노드 풀에서 실행되도록 구성되어 있습니다. 다음 명령어를 사용하여 포드를 실행합니다.

kubectl apply -f jax/stable-diffusion/install-stable-diffusion.yaml

다음을 사용해 모델이 올바르게 설치되었는지 확인합니다.

kubectl get pods install-stable-diffusion

STATUS가 Completed를 읽는 데 몇 분 정도 걸릴 수 있습니다.

모델의 TF 모델 서버 컨테이너 실행

예시 워크로드는 다음과 같은 가정으로 구성됩니다.

  • 클러스터를 노드가 3개인 TPU v5 노드 풀에서 실행
  • 노드 풀이 다음과 같은 ct5lp-hightpu-1t 머신 유형을 사용
    • 토폴로지는 1x1임
    • TPU 칩 수는 1개임

앞에서 설명한 것과 다른 클러스터 구성을 사용하는 경우 서버 배포 매니페스트를 수정해야 합니다.

배포를 적용합니다.

kubectl apply -f jax/stable-diffusion/serve-stable-diffusion.yaml

서버가 예상대로 실행되고 있는지 확인합니다.

kubectl get deployment stable-diffusion-deployment

AVAILABLE이 3을 읽는 데 1분 정도 걸릴 수 있습니다.

부하 분산기 서비스를 적용합니다.

kubectl apply -f jax/stable-diffusion/loadbalancer.yaml

다음을 사용하여 부하 분산기가 외부 트래픽을 지원할 준비가 되었는지 확인합니다.

kubectl get svc tf-stable-diffusion-service

EXTERNAL_IP가 IP를 나열하는 데 몇 분 정도 걸릴 수 있습니다.

모델 서버로 요청 전송

부하 분산기에서 외부 IP를 가져옵니다.

EXTERNAL_IP=$(kubectl get services tf-stable-diffusion-service --output jsonpath='{.status.loadBalancer.ingress[0].ip}')

서버에 요청을 전송하는 스크립트를 실행합니다.

python3 jax/stable-diffusion/stable_diffusion_request.py $EXTERNAL_IP

예상 출력:

프롬프트는 Painting of a squirrel skating in New York이며, 출력 이미지는 현재 디렉터리에 stable_diffusion_images.jpg로 저장됩니다.

삭제

리소스를 삭제하려면 kubectl delete를 역순으로 실행합니다.

kubectl delete -f jax/stable-diffusion/loadbalancer.yaml
kubectl delete -f jax/stable-diffusion/serve-stable-diffusion.yaml
kubectl delete -f jax/stable-diffusion/install-stable-diffusion.yaml

TensorFlow ResNet-50 E2E 서빙 데모 실행:

TF 모델 서비스로 요청을 전송하는 튜토리얼 Python 스크립트를 실행하기 위해 Python 종속 항목을 설치합니다.

pip install -r tf/resnet50/requirements.txt
1단계: 모델 변환

모델 변환을 적용합니다.

kubectl apply -f tf/resnet50/model-conversion.yml

다음을 사용해 모델이 올바르게 설치되었는지 확인합니다.

kubectl get pods resnet-model-conversion

STATUS가 Completed를 읽는 데 몇 분 정도 걸릴 수 있습니다.

2단계: TensorFlow Serving으로 모델 서빙

모델 서빙 배포를 적용합니다.

kubectl apply -f tf/resnet50/deployment.yml

다음 명령어를 사용하여 서버가 예상대로 실행되는지 확인합니다.

kubectl get deployment resnet-deployment

AVAILABLE이 3을 읽는 데 1분 정도 걸릴 수 있습니다.

부하 분산기 서비스를 적용합니다.

kubectl apply -f tf/resnet50/loadbalancer.yml

다음을 사용하여 부하 분산기가 외부 트래픽을 지원할 준비가 되었는지 확인합니다.

kubectl get svc resnet-service

EXTERNAL_IP가 IP를 나열하는 데 몇 분 정도 걸릴 수 있습니다.

3단계: 모델 서버에 테스트 요청 전송

부하 분산기에서 외부 IP를 가져옵니다.

EXTERNAL_IP=$(kubectl get services resnet-service --output jsonpath='{.status.loadBalancer.ingress[0].ip}')

테스트 요청(HTTP) 스크립트를 실행하여 모델 서버로 요청을 전송합니다.

python3 tf/resnet50/request.py --host $EXTERNAL_IP

다음과 같은 응답이 표시됩니다.

Predict result: ['ImageNet ID: n07753592, Label: banana, Confidence: 0.94921875',
'ImageNet ID: n03532672, Label: hook, Confidence: 0.0223388672', 'ImageNet ID: n07749582,
Label: lemon, Confidence: 0.00512695312
4단계: 삭제

리소스를 삭제하려면 다음 kubectl delete 명령어를 실행합니다.

kubectl delete -f tf/resnet50/loadbalancer.yml
kubectl delete -f tf/resnet50/deployment.yml
kubectl delete -f tf/resnet50/model-conversion.yml

완료한 후 GKE 노드 풀과 클러스터를 삭제해야 합니다.

PyTorch 모델 추론 및 서빙

PyTorch 모델 서비스로 요청을 전송하는 튜토리얼 Python 스크립트를 실행하기 위해 Python 종속 항목을 설치합니다.

pip install -r pt/densenet161/requirements.txt

TorchServe Densenet161 E2E 서빙 데모 실행:

  1. 모델 보관 파일을 생성합니다.

    1. 모델 보관 파일을 적용합니다.
    kubectl apply -f pt/densenet161/model-archive.yml
    1. 다음을 사용해 모델이 올바르게 설치되었는지 확인합니다.
    kubectl get pods densenet161-model-archive

    STATUS가 Completed를 읽는 데 몇 분 정도 걸릴 수 있습니다.

  2. TorchServe를 사용하여 모델을 서빙합니다.

    1. 모델 서빙 배포를 적용합니다.

      kubectl apply -f pt/densenet161/deployment.yml

    2. 다음 명령어를 사용하여 서버가 예상대로 실행되는지 확인합니다.

      kubectl get deployment densenet161-deployment

      AVAILABLE이 3을 읽는 데 1분 정도 걸릴 수 있습니다.

    3. 부하 분산기 서비스를 적용합니다.

      kubectl apply -f pt/densenet161/loadbalancer.yml

      다음 명령어를 사용하여 부하 분산기가 외부 트래픽을 지원할 준비가 되었는지 확인합니다.

      kubectl get svc densenet161-service

      EXTERNAL_IP가 IP를 나열하는 데 몇 분 정도 걸릴 수 있습니다.

  3. 모델 서버에 테스트 요청을 보냅니다.

    1. 부하 분산기에서 외부 IP를 가져옵니다.

      EXTERNAL_IP=$(kubectl get services densenet161-service --output jsonpath='{.status.loadBalancer.ingress[0].ip}')

    2. 테스트 요청 스크립트를 실행하여 모델 서버에 요청(HTTP)을 전송합니다.

      python3 pt/densenet161/request.py --host $EXTERNAL_IP

      다음과 같은 응답이 표시됩니다.

      Request successful. Response: {'tabby': 0.47878125309944153, 'lynx': 0.20393909513950348, 'tiger_cat': 0.16572578251361847, 'tiger': 0.061157409101724625, 'Egyptian_cat': 0.04997897148132324

  4. 다음 kubectl delete 명령어를 실행하여 리소스를 삭제합니다.

    kubectl delete -f pt/densenet161/loadbalancer.yml
kubectl delete -f pt/densenet161/deployment.yml
kubectl delete -f pt/densenet161/model-archive.yml

    완료한 후 GKE 노드 풀과 클러스터를 삭제해야 합니다.

자주 발생하는 문제 해결하기

GKE 문제 해결 정보는 GKE의 TPU 문제 해결을 참조하세요.

TPU 초기화 실패

다음 오류가 발생하면 TPU 컨테이너를 권한이 있는 모드에서 실행 중인지 확인하거나 컨테이너 내에서 ulimit를 늘렸는지 확인하세요. 자세한 내용은 권한이 있는 모드 없이 실행을 참조하세요.

TPU platform initialization failed: FAILED_PRECONDITION: Couldn't mmap: Resource
temporarily unavailable.; Unable to create Node RegisterInterface for node 0,
config: device_path:      "/dev/accel0" mode: KERNEL debug_data_directory: ""
dump_anomalies_only: true crash_in_debug_dump: false allow_core_dump: true;
could not create driver instance

예약 교착 상태

2개의 작업(작업 A 및 작업 B)이 있는데 둘 다 특정 TPU 토폴로지(예:v4-32)를 사용하는 TPU 슬라이스에 예약 예정이며 GKE 클러스터 내에 2개의 v4-32 TPU 슬라이스(슬라이스 X와 Y)가 있다고 가정하겠습니다. 클러스터에 두 작업을 모두 예약할 수 있는 충분한 용량이 있으므로 이론적으로는 두 작업이 빠르게 예약됩니다. 두 TPU v4-32 슬라이스 각각에 하나씩 작업이 예약됩니다.

하지만 신중하게 계획하지 않으면 예약 교착 상태가 발생할 수 있습니다. Kubernetes 스케줄러가 슬라이스 X에 작업 A의 Kubernetes 포드 하나를 예약한 후 슬라이스 X에서 작업 B의 Kubernetes 포드 하나를 예약합니다. 이 경우 작업 A에 대한 Kubernetes 포드 어피니티 규칙을 고려해 스케줄러가 슬라이스 X에 작업 A의 나머지 모든 Kubernetes 포드를 예약하려고 시도합니다. 작업 B의 경우도 마찬가지입니다. 따라서 작업 A와 작업 B를 단일 슬라이스에 완전하게 예약할 수 없습니다. 그 결과 예약 교착 상태가 발생합니다.

예약 교착 상태를 방지하기 위해 다음 예시와 같이 cloud.google.com/gke-nodepool을 topologyKey로 사용해 Kubernetes 포드 안티-어피니티를 사용할 수 있습니다.

apiVersion: batch/v1
kind: Job
metadata:
 name: pi
spec:
 parallelism: 2
 template:
   metadata:
     labels:
       job: pi
   spec:
     affinity:
       podAffinity:
         requiredDuringSchedulingIgnoredDuringExecution:
         - labelSelector:
             matchExpressions:
             - key: job
               operator: In
               values:
               - pi
           topologyKey: cloud.google.com/gke-nodepool
       podAntiAffinity:
         requiredDuringSchedulingIgnoredDuringExecution:
         - labelSelector:
             matchExpressions:
             - key: job
               operator: NotIn
               values:
               - pi
           topologyKey: cloud.google.com/gke-nodepool
           namespaceSelector:
             matchExpressions:
             - key: kubernetes.io/metadata.name
               operator: NotIn
               values:
               - kube-system
     containers:
     - name: pi
       image: perl:5.34.0
       command: ["sleep",  "60"]
     restartPolicy: Never
 backoffLimit: 4

Terraform으로 TPU 노드 풀 리소스 만들기

Terraform을 사용하여 클러스터 및 노드 풀 리소스를 관리할 수도 있습니다.

기존 GKE 클러스터에 멀티 호스트 TPU 슬라이스 노드 풀 만들기

멀티 호스트 TPU 노드 풀을 만들 기존 클러스터가 있는 경우 다음 Terraform 스니펫을 사용할 수 있습니다.

resource "google_container_cluster" "cluster_multi_host" {
  …
  release_channel {
    channel = "RAPID"
  }
  workload_identity_config {
    workload_pool = "my-gke-project.svc.id.goog"
  }
  addons_config {
    gcs_fuse_csi_driver_config {
      enabled = true
    }
  }
}

resource "google_container_node_pool" "multi_host_tpu" {
  provider           = google-beta
  project            = "${project-id}"
  name               = "${node-pool-name}"
  location           = "${location}"
  node_locations     = ["${node-locations}"]
  cluster            = google_container_cluster.cluster_multi_host.name
  initial_node_count = 2

  node_config {
    machine_type = "ct4p-hightpu-4t"
    reservation_affinity {
      consume_reservation_type = "SPECIFIC_RESERVATION"
      key = "compute.googleapis.com/reservation-name"
      values = ["${reservation-name}"]
    }
    workload_metadata_config {
      mode = "GKE_METADATA"
    }
  }

  placement_policy {
    type = "COMPACT"
    tpu_topology = "2x2x2"
  }
}

다음 값을 바꿉니다.

  • your-project: 워크로드를 실행 중인 Google Cloud 프로젝트입니다.
  • your-node-pool: 만들려는 노드 풀의 이름입니다.
  • us-central2: 워크로드를 실행 중인 리전입니다.
  • us-central2-b: 워크로드를 실행 중인 영역입니다.
  • your-reservation-name: 예약의 이름입니다.

기존 GKE 클러스터에 단일 호스트 TPU 슬라이스 노드 풀 만들기

다음 Terraform 스니펫을 사용합니다.

resource "google_container_cluster" "cluster_single_host" {
  …
  cluster_autoscaling {
    autoscaling_profile = "OPTIMIZE_UTILIZATION"
  }
  release_channel {
    channel = "RAPID"
  }
  workload_identity_config {
  workload_pool = "${project-id}.svc.id.goog"
  }
  addons_config {
    gcs_fuse_csi_driver_config {
      enabled = true
    }
  }
}

resource "google_container_node_pool" "single_host_tpu" {
  provider           = google-beta
  project            = "${project-id}"
  name               = "${node-pool-name}"
  location           = "${location}"
  node_locations     = ["${node-locations}"]
  cluster            = google_container_cluster.cluster_single_host.name
  initial_node_count = 0
  autoscaling {
    total_min_node_count = 2
    total_max_node_count = 22
    location_policy      = "ANY"
  }

  node_config {
    machine_type = "ct4p-hightpu-4t"
    workload_metadata_config {
      mode = "GKE_METADATA"
    }
  }
}

다음 값을 바꿉니다.

  • your-project: 워크로드를 실행 중인 Google Cloud 프로젝트입니다.
  • your-node-pool: 만들려는 노드 풀의 이름입니다.
  • us-central2: 워크로드를 실행 중인 리전입니다.
  • us-central2-b: 워크로드를 실행 중인 영역입니다.