GKE Autopilot에서 TPU 워크로드 배포

이 페이지에서는 Google Kubernetes Engine (GKE) Autopilot 클러스터에서 Cloud TPU 가속기 (TPU)를 사용하여 머신러닝 (ML) 워크로드를 가속화하는 방법을 설명합니다. 이 가이드를 통해 ML 애플리케이션 프레임워크에 적합한 라이브러리를 선택하고, GKE에서 최적으로 실행되도록 TPU 워크로드를 설정하고, 배포 후 워크로드를 모니터링할 수 있습니다.

이 페이지는 TPU에서 ML 워크로드를 준비하고 실행하려는 플랫폼 관리자 및 운영자, 데이터 및 AI 전문가, 애플리케이션 개발자를 대상으로 합니다. Google Cloud 콘텐츠에서 참조하는 일반적인 역할, 책임, 예시 태스크에 대해 자세히 알아보려면 일반 GKE 사용자 역할 및 태스크를 참고하세요.

이 페이지를 읽기 전에 다음 리소스를 숙지해야 합니다.

Autopilot에서 TPU 작동 방식

Autopilot 워크로드에서 TPU를 사용하려면 워크로드 매니페스트에서 다음을 지정합니다.

  • spec.nodeSelector 필드의 TPU 버전입니다.
  • spec.nodeSelector 필드의 TPU 토폴로지 토폴로지는 지정된 TPU 버전에서 지원되어야 합니다.
  • spec.containers.resources.requests 및 spec.containers.resources.limits 필드의 TPU 칩 수입니다.

워크로드를 배포하면 GKE에서 요청된 TPU 구성이 있는 노드를 프로비저닝하고 노드에서 포드를 예약합니다. GKE는 각 워크로드를 자체 노드에 배치하므로 중단 위험이 최소화되면서 각 포드에서 노드의 전체 리소스에 액세스할 수 있습니다.

Autopilot의 TPU는 다음 기능과 호환됩니다.

  1. 스팟 포드
  2. 특정 용량 예약
  3. 연장된 실행 시간 포드
  4. Flex-start

TPU 구성 계획

이 가이드를 사용하여 TPU 워크로드를 배포하기 전에 모델과 필요한 메모리 양에 따라 TPU 구성을 계획하세요. 자세한 내용은 TPU 구성 계획을 참고하세요.

가격 책정

가격 책정 정보는 Autopilot 가격 책정을 참조하세요.

시작하기 전에

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

  • Google Kubernetes Engine API를 사용 설정합니다.
    • Google Kubernetes Engine API 사용 설정
  • 이 태스크에 Google Cloud CLI를 사용하려면 gcloud CLI를 설치한 후 초기화합니다. 이전에 gcloud CLI를 설치한 경우 gcloud components update 명령어를 실행하여 최신 버전을 가져옵니다. 이전 gcloud CLI 버전에서는 이 문서의 명령어를 실행하지 못할 수 있습니다.
  • GKE 버전 1.32.3-gke.1927000 이상을 실행하는 Autopilot 클러스터가 있는지 확인합니다. 자세한 내용은 Autopilot 클러스터 만들기를 참고하세요.
  • 예약된 TPU를 사용하려면 기존 특정 용량 예약이 있는지 확인합니다. 자세한 내용은 예약 사용을 참고하세요.

TPU 및 기타 GKE 리소스의 할당량 확인

다음 섹션에서는 GKE에서 TPU를 사용할 때 할당량이 충분한지 확인하는 방법을 설명합니다.

TPU 슬라이스 노드를 만들려면 기존 용량 예약을 사용하지 않는 한 TPU 할당량을 사용할 수 있어야 합니다. 예약된 TPU를 사용하는 경우 이 섹션을 건너뛰세요.

GKE에서 TPU 슬라이스 노드를 만들려면 Cloud TPU API 할당량(tpu.googleapis.com)이 아닌 Compute Engine API 할당량(compute.googleapis.com)이 필요합니다. 할당량 이름은 일반 Autopilot 포드와 스팟 포드에서 서로 다릅니다.

TPU용 Compute Engine API 할당량의 한도 및 현재 사용량을 확인하려면 다음 단계를 따르세요.

  1. Google Cloud 콘솔에서 할당량 페이지로 이동합니다.

    할당량으로 이동

  2. 필터 상자에서 다음을 수행합니다,

    1. 다음 표를 사용하여 TPU 버전 및 cloud.google.com/gke-tpu-accelerator 노드 선택기의 값에 따라 할당량의 속성을 선택하고 복사합니다. 예를 들어 cloud.google.com/gke-tpu-accelerator 노드 선택기의 값이 tpu-v5-lite-podslice인 주문형 TPU v5e 노드를 만들려면 Name: TPU v5 Lite PodSlice chips를 입력합니다.

      TPU 버전, cloud.google.com/gke-tpu-accelerator 온디맨드 인스턴스의 할당량 속성 및 이름 스팟2 인스턴스의 할당량 속성 및 이름
      TPU v3,
      tpu-v3-device 
      Dimensions (e.g. location):
      tpu_family:CT3
      		 해당 사항 없음
      TPU v3,
      tpu-v3-slice 
      Dimensions (e.g. location):
      tpu_family:CT3P
      		 해당 사항 없음
      TPU v4,
      tpu-v4-podslice 
      Name:
      TPU v4 PodSlice chips
      		 
      Name:
      Preemptible TPU v4 PodSlice chips
      TPU v5e,
      tpu-v5-lite-podslice 
      Name:
      TPU v5 Lite PodSlice chips
      		 
      Name:
      Preemptible TPU v5 Lite Podslice
      chips
      TPU v5p,
      tpu-v5p-slice 
      Name:
      TPU v5p chips
      		 
      Name:
      Preemptible TPU v5p chips
      TPU Trillium,
      tpu-v6e-slice 
      Dimensions (e.g. location):
      tpu_family:CT6E
      		 
      Name:
      Preemptible TPU slices v6e

    2. 측정기준(예: 위치) 속성을 선택하고 region:을 입력한 다음 GKE에서 TPU를 생성할 리전의 이름을 입력합니다. 예를 들어 us-west4-a 영역에 TPU 슬라이스 노드를 만들려면 region:us-west4를 입력합니다. TPU 할당량은 리전 기준이므로 동일한 리전 내의 모든 영역에서 동일한 TPU 할당량을 사용합니다.

입력한 필터와 일치하는 할당량이 없으면 필요한 리전에 대해 지정된 할당량이 프로젝트에 부여되지 않았으므로 TPU 할당량 조정을 요청해야 합니다.

TPU 예약이 생성되면 해당 할당량의 한도 및 현재 사용량 값이 TPU 예약의 칩 수만큼 증가합니다. 예를 들어 cloud.google.com/gke-tpu-accelerator 노드 선택기의 값이 tpu-v5-lite-podslice인 16개의 TPU v5e 칩에 대한 예약이 생성되면 관련 리전에서 TPU v5 Lite PodSlice chips 할당량의 한도 및 현재 사용량이 모두 16개씩 증가합니다.

추가 GKE 리소스의 할당량

GKE에서 리소스를 만드는 리전에서 다음 GKE 관련 할당량을 늘려야 할 수 있습니다.

  • 영구 디스크 SSD(GB) 할당량: 각 Kubernetes 노드의 부팅 디스크에는 기본적으로 100GB가 필요합니다. 따라서 이 할당량은 생성될 것으로 예상되는 최대 GKE 노드 수와 100GB(노드 * 100GB)의 곱 이상으로 설정해야 합니다.
  • 사용 중인 IP 주소 할당량: 각 Kubernetes 노드는 하나의 IP 주소를 사용합니다 따라서 이 할당량은 생성될 것으로 예상되는 최대 GKE 노드 수 이상으로 설정해야 합니다.
  • max-pods-per-node가 서브넷 범위와 일치하는지 확인: 각 Kubernetes 노드는 포드에 보조 IP 범위를 사용합니다. 예를 들어 max-pods-per-node가 32이면 64개의 IP 주소가 필요하며 이는 노드당 /26 서브넷으로 변환됩니다. 이 범위는 다른 클러스터와 공유해서는 안 됩니다. IP 주소 범위가 소진되지 않도록 하려면 --max-pods-per-node 플래그를 사용하여 노드에 예약할 수 있는 포드 수를 제한합니다. max-pods-per-node의 할당량은 생성될 것으로 예상되는 최대 GKE 노드 수 이상으로 설정해야 합니다.

할당량 상향을 요청하려면 할당량 조정 요청을 참고하세요.

GKE에서 TPU를 프로비저닝하는 옵션

GKE Autopilot을 사용하면 Kubernetes nodeSelector를 사용하여 개별 워크로드에서 TPU를 직접 사용할 수 있습니다.

또는 커스텀 컴퓨팅 클래스를 사용하여 TPU를 요청할 수 있습니다. 커스텀 컴퓨팅 클래스를 사용하면 플랫폼 관리자가 노드 확장 결정 중에 GKE가 우선시할 노드 구성의 계층 구조를 정의하여 워크로드가 선택한 하드웨어에서 실행되도록 할 수 있습니다.

자세한 내용은 커스텀 컴퓨팅 클래스로 TPU 중앙에서 프로비저닝 섹션을 참고하세요.

TPU 애플리케이션 준비

TPU 워크로드 준비 요구사항은 다음과 같습니다.

  1. JAX, PyTorch, TensorFlow와 같은 프레임워크는 libtpu 공유 라이브러리를 사용하여 TPU VM에 액세스합니다. libtpu에는 XLA 컴파일러, TPU 런타임 소프트웨어, TPU 드라이버가 포함됩니다. PyTorch 및 JAX의 각 출시 버전에는 특정 libtpu.so 버전이 필요합니다. 패키지 버전 충돌을 방지하려면 JAX AI 이미지를 사용하는 것이 좋습니다. GKE에서 TPU를 사용하려면 다음 버전을 사용해야 합니다.
    TPU 유형 libtpu.so 버전
    TPU Trillium(v6e)
    tpu-v6e-slice
    TPU v5e
    tpu-v5-lite-podslice
    TPU v5p
    tpu-v5p-slice
    • 권장 JAX AI 이미지: jax0.4.35-rev1 이상
    • 권장 jax[tpu] version: 0.4.19 이상
    • 권장 torchxla[tpuvm] 버전: 2023년 10월 23일 나이틀리 버전 빌드를 사용하는 것이 좋습니다
    TPU v4
    tpu-v4-podslice
    TPU v3
    tpu-v3-slice
    tpu-v3-device
  2. TPU 리소스를 요청하는 컨테이너에 다음 환경 변수를 설정합니다.
    • TPU_WORKER_ID: 포드마다 고유한 정수입니다. 이 ID는 TPU 슬라이스의 고유한 worker-id를 나타냅니다. 이 필드에 지원되는 값 범위는 0부터 포드 수에서 1을 뺀 값까지입니다.
    • TPU_WORKER_HOSTNAMES: 슬라이스 내에서 서로 통신해야 하는 TPU VM 호스트 이름이나 IP 주소의 쉼표로 구분된 목록입니다. 슬라이스의 TPU VM마다 호스트 이름이나 IP 주소가 있어야 합니다. IP 주소 또는 호스트 이름 목록은 TPU_WORKER_ID에 의해 정렬되고 색인이 0으로 지정됩니다.

      • completionMode: Indexed, subdomain, parallelism > 1, 요청 google.com/tpu 속성을 사용하여 작업을 만든 경우 GKE에서 변형 웹훅을 사용하여 이러한 환경 변수를 자동으로 삽입합니다. GKE는 서비스를 지원하는 포드에 DNS 레코드가 추가되도록 헤드리스 서비스를 추가합니다.

워크로드 준비가 완료되면 TPU를 사용하는 작업을 실행할 수 있습니다.

워크로드에서 TPU 요청

이 섹션에서는 Autopilot에서 TPU를 요청하는 작업을 만드는 방법을 보여줍니다. TPU가 필요한 모든 워크로드에서 다음을 지정해야 합니다.

  • TPU 버전과 토폴로지의 노드 선택기
  • 워크로드의 컨테이너에 대한 TPU 칩 수

지원되는 TPU 버전, 토폴로지, 슬라이스의 해당 TPU 칩 및 노드 수 목록은 TPU 버전 선택을 참고하세요.

워크로드에서 TPU 요청 고려사항

포드에 있는 컨테이너 하나만 TPU를 사용할 수 있습니다. 컨테이너에서 요청하는 TPU 칩 수는 슬라이스의 노드에 연결된 TPU 칩 수와 같아야 합니다. 예를 들어 2x4 토폴로지로 TPU v5e(tpu-v5-lite-podslice)를 요청하는 경우 다음 중 하나를 요청할 수 있습니다.

  • 각각 TPU 칩 4개가 있는 멀티 호스트 노드 2개를 만드는 4 칩
  • TPU 칩 8개가 있는 단일 호스트 노드 1개를 만드는 8 칩

경제성을 극대화할 수 있는 권장사항은 항상 요청하는 슬라이스의 모든 TPU를 사용하는 것입니다. 각각 TPU 칩 4개가 있는 노드 2개의 멀티 호스트 슬라이스를 요청하는 경우 두 노드 모두에서 실행되고 슬라이스의 TPU 칩 8개를 모두 사용하는 워크로드를 배포해야 합니다.

TPU를 요청하는 워크로드 만들기

다음 단계에서는 TPU를 요청하는 작업을 만듭니다. 멀티 호스트 TPU 슬라이스에서 실행되는 워크로드가 있으면 이름을 기준으로 워크로드를 선택하는 헤드리스 서비스도 만들어야 합니다. 이 헤드리스 서비스를 사용하면 워크로드의 포드를 가리키도록 Kubernetes DNS 구성을 업데이트하여 멀티 호스트 슬라이스의 서로 다른 노드에 있는 포드가 서로 통신할 수 있습니다.

  1. 다음 매니페스트를 tpu-autopilot.yaml로 저장합니다.

    apiVersion: v1
kind: Service
metadata:
  name: headless-svc
spec:
  clusterIP: None
  selector:
    job-name: tpu-job
---
apiVersion: batch/v1
kind: Job
metadata:
  name: tpu-job
spec:
  backoffLimit: 0
  completions: 4
  parallelism: 4
  completionMode: Indexed
  template:
    spec:
      # Optional: Run in GKE Sandbox
      # runtimeClassName: gvisor
      subdomain: headless-svc
      restartPolicy: Never
      nodeSelector:
        cloud.google.com/gke-tpu-accelerator: TPU_TYPE
        cloud.google.com/gke-tpu-topology: TOPOLOGY
      containers:
      - name: tpu-job
        image: us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest
        ports:
        - containerPort: 8471 # Default port using which TPU VMs communicate
        - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
        command:
        - bash
        - -c
        - |
          python -c 'import jax; print("TPU cores:", jax.device_count())'
        resources:
          requests:
            cpu: 10
            memory: MEMORY_SIZE
            google.com/tpu: NUMBER_OF_CHIPS
          limits:
            cpu: 10
            memory: MEMORY_SIZE
            google.com/tpu: NUMBER_OF_CHIPS

    다음을 바꿉니다.

    • TPU_TYPE: 사용할 TPU 유형입니다(예: tpu-v4-podslice). GKE에서 지원하는 값이어야 합니다.
    • TOPOLOGY: 슬라이스의 TPU 칩 배열입니다(예: 2x2x4). 선택한 TPU 유형에 지원되는 토폴로지여야 합니다.
    • NUMBER_OF_CHIPS: 사용할 컨테이너의 TPU 칩 수입니다. limits 및 requests 값이 같아야 합니다.
    • MEMORY_SIZE: TPU가 사용하는 최대 메모리 양입니다. 메모리 제한은 사용 중인 TPU 버전과 토폴로지에 따라 달라집니다. 자세한 내용은 가속기의 최솟값 및 최댓값을 참조하세요.

    선택적으로 다음 필드를 변경할 수도 있습니다.

    • image: 사용할 JAX AI 이미지입니다. 매니페스트 예시에서 이 필드는 최신 JAX AI 이미지로 설정됩니다. 다른 버전을 설정하려면 현재 JAX AI 이미지 목록을 참고하세요.
    • runtimeClassname: gvisor: GKE Sandbox에서 이 포드를 실행할 수 있는 설정입니다. 사용하려면 이 줄의 주석 처리를 삭제하세요. GKE Sandbox는 TPU 버전 v4 이상을 지원합니다. 자세한 내용은 GKE Sandbox를 참고하세요.

  2. 작업을 배포합니다.

    kubectl create -f tpu-autopilot.yaml

    이 작업을 만들면 GKE에서 자동으로 다음을 수행합니다.

    1. 포드를 실행할 노드를 프로비저닝합니다. 지정한 TPU 유형, 토폴로지, 리소스 요청에 따라 이러한 노드는 단일 호스트 슬라이스이거나 멀티 호스트 슬라이스입니다.
    2. 다른 워크로드가 TPU 워크로드와 동일한 노드에서 실행되지 않도록 taint를 포드에, 톨러레이션(toleration)을 노드에 추가합니다.

  3. 이 섹션을 마치면 만든 워크로드를 삭제하여 비용이 계속 청구되지 않도록 할 수 있습니다.

    kubectl delete -f tpu-autopilot.yaml

TPU 및 컬렉션 스케줄링을 요청하는 워크로드 만들기

TPU Trillium에서는 컬렉션 예약을 사용하여 TPU 슬라이스 노드를 그룹화할 수 있습니다. 이러한 TPU 슬라이스 노드를 그룹화하면 워크로드 수요를 충족하도록 복제본 수를 더 쉽게 조정할 수 있습니다 . Google Cloud 는 소프트웨어 업데이트를 제어하여 컬렉션 내에서 항상 충분한 슬라이스를 사용하여 트래픽을 처리할 수 있도록 합니다.

TPU Trillium은 추론 워크로드를 실행하는 단일 호스트 및 멀티 호스트 노드 풀의 컬렉션 예약을 지원합니다. 다음은 컬렉션 예약 동작이 사용하는 TPU 슬라이스의 유형에 따라 어떻게 달라지는지 설명합니다.

  • 멀티 호스트 TPU 슬라이스: GKE는 멀티 호스트 TPU 슬라이스를 그룹화하여 컬렉션을 형성합니다. 각 GKE 노드 풀은 이 컬렉션 내의 복제본입니다. 컬렉션을 정의하려면 멀티 호스트 TPU 슬라이스를 만들고 컬렉션에 고유한 이름을 할당합니다. 컬렉션에 TPU 슬라이스를 더 추가하려면 컬렉션 이름과 워크로드 유형이 동일한 멀티 호스트 TPU 슬라이스 노드 풀을 하나 더 만듭니다.
  • 단일 호스트 TPU 슬라이스: GKE는 전체 단일 호스트 TPU 슬라이스 노드 풀을 컬렉션으로 간주합니다. 컬렉션에 TPU 슬라이스를 더 추가하려면 단일 호스트 TPU 슬라이스 노드 풀의 크기를 조정하면 됩니다.

컬렉션 예약의 제한사항에 대해 알아보려면 컬렉션 예약 작동 방식을 참고하세요.

멀티 호스트 TPU 슬라이스 사용

멀티 호스트 TPU 슬라이스 노드의 컬렉션 예약은 버전 1.31.2-gke.1537000 이상의 Autopilot 클러스터에서 사용할 수 있습니다. 2x4 토폴로지가 있는 멀티 호스트 TPU 슬라이스 노드는 1.31.2-gke.1115000 이상에서만 지원됩니다. 멀티 호스트 TPU 슬라이스 노드를 만들고 컬렉션으로 그룹화하려면 워크로드 사양에 다음 Kubernetes 라벨을 추가하세요.

  • cloud.google.com/gke-nodepool-group-name: 각 컬렉션에는 클러스터 수준에서 고유한 이름이 있어야 합니다. cloud.google.com/gke-nodepool-group-name 라벨의 값은 클러스터 라벨 요구사항을 준수해야 합니다.

  • cloud.google.com/gke-workload-type: HIGH_AVAILABILITY

    예를 들어 다음 코드 블록은 멀티 호스트 TPU 슬라이스가 있는 컬렉션을 정의합니다.

      nodeSelector:
    cloud.google.com/gke-nodepool-group-name: ${COLLECTION_NAME}
    cloud.google.com/gke-workload-type: HIGH_AVAILABILITY
    cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
    cloud.google.com/gke-tpu-topology: 4x4
...

단일 호스트 TPU 슬라이스 사용

단일 호스트 TPU 슬라이스 노드의 컬렉션 예약은 버전 1.31.2-gke.1088000 이상의 Autopilot 클러스터에서 사용할 수 있습니다. 단일 호스트 TPU 슬라이스 노드를 만들고 이를 컬렉션으로 그룹화하려면 워크로드 사양에 cloud.google.com/gke-workload-type:HIGH_AVAILABILITY 라벨을 추가합니다.

예를 들어 다음 코드 블록은 단일 호스트 TPU 슬라이스가 있는 컬렉션을 정의합니다.

  nodeSelector:
    cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
    cloud.google.com/gke-tpu-topology: 2x2
    cloud.google.com/gke-workload-type: HIGH_AVAILABILITY
  ...

맞춤 컴퓨팅 클래스를 사용하여 컬렉션 배포

커스텀 컴퓨팅 클래스를 사용하여 TPU 워크로드 및 컬렉션 스케줄링을 요청하는 워크로드를 배포하는 방법에 대한 자세한 내용은 TPU 멀티 호스트 컬렉션 및 TPU SLO의 워크로드 유형 정의를 참고하세요.

커스텀 컴퓨팅 클래스로 TPU를 중앙에서 프로비저닝

TPU 규칙을 따르는 커스텀 컴퓨팅 클래스로 TPU를 프로비저닝하고 워크로드를 배포하려면 다음 단계를 완료하세요.

  1. 다음 매니페스트를 tpu-compute-class.yaml로 저장합니다.

    apiVersion: cloud.google.com/v1
kind: ComputeClass
metadata:
  name: tpu-class
spec:
  priorities:
  - tpu:
      type: tpu-v5-lite-podslice
      count: 4
      topology: 2x4
  - spot: true
    tpu:
      type: tpu-v5-lite-podslice
      count: 4
      topology: 2x4
  - flexStart:
      enabled: true
    tpu:
      type: tpu-v6e-slice
      count: 4
      topology: 2x4
  nodePoolAutoCreation:
    enabled: true

  2. 컴퓨팅 클래스를 배포합니다.

    kubectl apply -f tpu-compute-class.yaml

    커스텀 컴퓨팅 클래스 및 TPU에 대한 자세한 내용은 TPU 구성을 참조하세요.

  3. 다음 매니페스트를 tpu-job.yaml로 저장합니다.

    apiVersion: v1
kind: Service
metadata:
  name: headless-svc
spec:
  clusterIP: None
  selector:
    job-name: tpu-job
---
apiVersion: batch/v1
kind: Job
metadata:
  name: tpu-job
spec:
  backoffLimit: 0
  completions: 4
  parallelism: 4
  completionMode: Indexed
  template:
    spec:
      subdomain: headless-svc
      restartPolicy: Never
      nodeSelector:
        cloud.google.com/compute-class: tpu-class
      containers:
      - name: tpu-job
        image: us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest
        ports:
        - containerPort: 8471 # Default port using which TPU VMs communicate
        - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
        command:
        - bash
        - -c
        - |
          python -c 'import jax; print("TPU cores:", jax.device_count())'
        resources:
          requests:
            cpu: 10
            memory: MEMORY_SIZE
            google.com/tpu: NUMBER_OF_CHIPS
          limits:
            cpu: 10
            memory: MEMORY_SIZE
            google.com/tpu: NUMBER_OF_CHIPS

    다음을 바꿉니다.

    • NUMBER_OF_CHIPS: 사용할 컨테이너의 TPU 칩 수입니다. limits 및 requests에 대해 동일한 값이어야 하며 선택한 커스텀 컴퓨팅 클래스의 tpu.count 필드 값과 같아야 합니다.
    • MEMORY_SIZE: TPU가 사용하는 최대 메모리 양입니다. 메모리 한도는 사용 중인 TPU 버전과 토폴로지에 따라 다릅니다. 자세한 내용은 가속기의 최솟값 및 최댓값을 참조하세요.
    • NUMBER_OF_CHIPS: 사용할 컨테이너의 TPU 칩 수입니다. limits 및 requests 값이 같아야 합니다.

  4. 작업을 배포합니다.

    kubectl create -f tpu-job.yaml

    이 작업을 만들면 GKE에서 자동으로 다음을 수행합니다.

    • 포드를 실행할 노드를 프로비저닝합니다. 지정한 TPU 유형, 토폴로지, 리소스 요청에 따라 이러한 노드는 단일 호스트 슬라이스이거나 멀티 호스트 슬라이스입니다. 최고 우선순위의 TPU 리소스 가용성에 따라 GKE는 획득 가능성을 최대화하기 위해 낮은 우선순위로 대체할 수 있습니다.
    • 다른 워크로드가 TPU 워크로드와 동일한 노드에서 실행되지 않도록 taint를 포드에, 톨러레이션(toleration)을 노드에 추가합니다.

    자세한 내용은 커스텀 컴퓨팅 클래스 정보를 참조하세요.

  5. 이 섹션을 마치면 만든 리소스를 삭제하여 비용이 계속 청구되지 않도록 할 수 있습니다.

    kubectl delete -f tpu-job.yaml

예시: 멀티 호스트 슬라이스의 총 TPU 칩 표시

다음 워크로드는 멀티 호스트 TPU 슬라이스의 모든 노드에서 TPU 칩 수를 반환합니다. 멀티 호스트 슬라이스를 만들기 위한 워크로드 매개변수는 다음과 같습니다.

  • TPU 버전: TPU v4
  • 토폴로지: 2x2x4

이 버전과 토폴로지를 선택하면 멀티 호스트 슬라이스가 생성됩니다.

  1. 다음 매니페스트를 available-chips-multihost.yaml로 저장합니다. 
    apiVersion: v1
kind: Service
metadata:
  name: headless-svc
spec:
  clusterIP: None
  selector:
    job-name: tpu-available-chips
---
apiVersion: batch/v1
kind: Job
metadata:
  name: tpu-available-chips
spec:
  backoffLimit: 0
  completions: 4
  parallelism: 4
  completionMode: Indexed
  template:
    spec:
      subdomain: headless-svc
      restartPolicy: Never
      nodeSelector:
        cloud.google.com/gke-tpu-accelerator: tpu-v4-podslice # Node selector to target TPU v4 slice nodes.
        cloud.google.com/gke-tpu-topology: 2x2x4 # Specifies the physical topology for the TPU slice.
      containers:
      - name: tpu-job
        image: us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest
        ports:
        - containerPort: 8471 # Default port using which TPU VMs communicate
        - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
        command:
        - bash
        - -c
        - |
          python -c 'import jax; print("TPU cores:", jax.device_count())' # Python command to count available TPU chips.
        resources:
          requests:
            cpu: 10
            memory: 407Gi
            google.com/tpu: 4 # Request 4 TPU chips for this workload.
          limits:
            cpu: 10
            memory: 407Gi
            google.com/tpu: 4 # Limit to 4 TPU chips for this workload.
  2. 매니페스트를 배포합니다. 
    kubectl create -f available-chips-multihost.yaml

    GKE는 4개의 VM(멀티 호스트 TPU 슬라이스)으로 TPU v4 슬라이스를 실행합니다. 슬라이스에는 상호 연결된 TPU 칩 16개가 있습니다.

  3. 작업이 4개의 포드를 만들었는지 확인합니다. 
    kubectl get pods

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

    NAME                       READY   STATUS      RESTARTS   AGE
tpu-job-podslice-0-5cd8r   0/1     Completed   0          97s
tpu-job-podslice-1-lqqxt   0/1     Completed   0          97s
tpu-job-podslice-2-f6kwh   0/1     Completed   0          97s
tpu-job-podslice-3-m8b5c   0/1     Completed   0          97s
  4. 포드 중 하나의 로그를 가져옵니다. 
    kubectl logs POD_NAME

    POD_NAME을 생성된 포드 중 하나의 이름으로 바꿉니다. 예를 들면 tpu-job-podslice-0-5cd8r입니다.

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

    TPU cores: 16
  5. 선택사항: 워크로드 삭제: 
    kubectl delete -f available-chips-multihost.yaml

예시: 단일 노드의 TPU 칩 표시

다음 워크로드는 특정 노드에 연결된 TPU 칩 수를 표시하는 정적 포드입니다. 단일 호스트 노드를 만들기 위한 워크로드 매개변수는 다음과 같습니다.

  • TPU 버전: TPU v5e
  • 토폴로지: 2x4

이 버전과 토폴로지를 선택하면 단일 호스트 슬라이스가 생성됩니다.

  1. 다음 매니페스트를 available-chips-singlehost.yaml로 저장합니다. 
    apiVersion: v1
kind: Pod
metadata:
  name: tpu-job-jax-v5
spec:
  restartPolicy: Never
  nodeSelector:
    cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice # Node selector to target TPU v5e slice nodes.
    cloud.google.com/gke-tpu-topology: 2x4 # Specify the physical topology for the TPU slice.
  containers:
  - name: tpu-job
    image: us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest
    ports:
    - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
    command:
    - bash
    - -c
    - |
      python -c 'import jax; print("Total TPU chips:", jax.device_count())'
    resources:
      requests:
        google.com/tpu: 8 # Request 8 TPU chips for this container.
      limits:
        google.com/tpu: 8 # Limit to 8 TPU chips for this container.
  2. 매니페스트를 배포합니다. 
    kubectl create -f available-chips-singlehost.yaml

    GKE는 TPU v5e를 사용하는 단일 호스트 TPU 슬라이스 8개로 노드를 프로비저닝합니다. 각 TPU 노드에는 8개의 TPU 칩이 있습니다(단일 호스트 TPU 슬라이스).

  3. 포드 로그를 가져옵니다. 
    kubectl logs tpu-job-jax-v5

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

    Total TPU chips: 8
  4. 선택사항: 워크로드 삭제: 
      kubectl delete -f available-chips-singlehost.yaml

TPU 관찰 및 모니터링

대시보드

Google Cloud 콘솔의 노드 풀 모니터링 가능성이 정식 버전으로 출시되었습니다. GKE에서 TPU 멀티 호스트 노드 풀의 상태를 보려면 Cloud Monitoring에서 제공하는 GKE TPU 노드 풀 상태 대시보드로 이동하세요.

GKE TPU 노드 풀 상태로 이동

이 대시보드는 멀티 호스트 TPU 노드 풀의 상태에 대한 포괄적인 통계를 제공합니다. 자세한 내용은 TPU 노드 및 노드 풀의 상태 측정항목 모니터링을 참조하세요.

Google Cloud 콘솔의 Kubernetes 클러스터 페이지에서 관측 가능성 탭에는 가속기 > TPU 헤더 아래에 TPU 사용량과 같은 TPU 관측 가능성 측정항목도 표시됩니다. 자세한 내용은 관측 가능성 측정항목 보기를 참조하세요.

TPU 대시보드는 GKE 클러스터에서 시스템 측정항목을 사용 설정한 경우에만 채워집니다.

런타임 측정항목

GKE 버전 1.27.4-gke.900 이상에서 JAX 버전 0.4.14 이상을 사용하고 containerPort: 8431을 지정하는 TPU 워크로드는 TPU 사용률 측정항목을 GKE 시스템 측정항목으로 내보냅니다. TPU 워크로드의 런타임 성능을 모니터링하기 위해 Cloud Monitoring에서 다음 측정항목을 사용할 수 있습니다.

  • 가동 주기: 이전 샘플링 기간(60초) 동안 TPU 칩에서 TensorCore가 활발하게 처리된 시간의 백분율입니다. 백분율이 높을수록 TPU 사용률이 높은 것입니다.
  • 메모리 사용량: 할당된 가속기 메모리 양(바이트)입니다. 60초마다 샘플링됩니다.
  • 메모리 용량: 총 가속기 메모리 용량(바이트)입니다. 60초마다 샘플링됩니다.

이러한 측정항목은 Kubernetes 노드(k8s_node) 및 Kubernetes 컨테이너(k8s_container) 스키마에 있습니다.

Kubernetes 컨테이너:

  • kubernetes.io/container/accelerator/duty_cycle
  • kubernetes.io/container/accelerator/memory_used
  • kubernetes.io/container/accelerator/memory_total

Kubernetes 노드:

  • kubernetes.io/node/accelerator/duty_cycle
  • kubernetes.io/node/accelerator/memory_used
  • kubernetes.io/node/accelerator/memory_total

TPU 노드 및 노드 풀의 상태 측정항목 모니터링

학습 작업에 오류가 있거나 실패로 종료되면 기본 인프라와 관련된 측정항목을 확인하여 기본 노드 또는 노드 풀의 문제로 인해 중단되었는지 확인할 수 있습니다.

노드 상태

GKE 버전 1.32.1-gke.1357001 이상에서 다음 GKE 시스템 측정항목은 GKE 노드의 상태를 노출합니다.

  • kubernetes.io/node/status_condition

condition 필드는 Ready, DiskPressure, MemoryPressure와 같은 노드의 조건을 보고합니다. status 필드에는 True, False 또는 Unknown일 수 있는 조건의 보고된 상태가 표시됩니다. 이것은 k8s_node 모니터링 리소스 유형이 있는 측정항목입니다.

이 PromQL 쿼리는 특정 노드가 Ready인지 보여줍니다.

kubernetes_io:node_status_condition{
    monitored_resource="k8s_node",
    cluster_name="CLUSTER_NAME",
    node_name="NODE_NAME",
    condition="Ready",
    status="True"}

클러스터의 문제를 해결하려면 다른 조건을 나타낸 노드를 살펴보세요.

kubernetes_io:node_status_condition{
    monitored_resource="k8s_node",
    cluster_name="CLUSTER_NAME",
    condition!="Ready",
    status="True"}

Ready가 아닌 노드를 구체적으로 살펴볼 수 있습니다.

kubernetes_io:node_status_condition{
    monitored_resource="k8s_node",
    cluster_name="CLUSTER_NAME",
    condition="Ready",
    status="False"}

데이터가 없으면 노드가 준비된 것입니다. 상태 조건은 60초마다 샘플링됩니다.

다음 쿼리를 사용하여 Fleet 전체의 노드 상태를 파악할 수 있습니다.

avg by (condition,status)(
  avg_over_time(
    kubernetes_io:node_status_condition{monitored_resource="k8s_node"}[${__interval}]))

노드 풀 상태

k8s_node_pool 모니터링 리소스에 대한 다음 GKE 시스템 측정항목은 GKE 노드 풀의 상태를 노출합니다.

  • kubernetes.io/node_pool/status

이 측정항목은 멀티 호스트 TPU 노드 풀에 대해서만 보고됩니다.

status 필드는 Provisioning, Running, Error, Reconciling, Stopping과 같은 노드 풀의 상태를 보고합니다. GKE API 작업이 완료된 후 상태가 업데이트됩니다.

특정 노드 풀의 상태가 Running인지 확인하려면 다음 PromQL 쿼리를 사용하세요.

kubernetes_io:node_pool_status{
    monitored_resource="k8s_node_pool",
    cluster_name="CLUSTER_NAME",
    node_pool_name="NODE_POOL_NAME",
    status="Running"}

상태별로 그룹화된 프로젝트의 노드 풀 수를 모니터링하려면 다음 PromQL 쿼리를 사용하세요.

count by (status)(
  count_over_time(
    kubernetes_io:node_pool_status{monitored_resource="k8s_node_pool"}[${__interval}]))

노드 풀 가용성

다음 GKE 시스템 측정항목은 멀티 호스트 TPU 노드 풀의 사용 가능 여부를 보여줍니다.

  • kubernetes.io/node_pool/multi_host/available

노드 풀의 모든 노드를 사용할 수 있으면 측정항목의 값이 True이고 그렇지 않으면 False입니다. 측정항목은 60초마다 샘플링됩니다.

프로젝트에서 멀티 호스트 TPU 노드 풀의 가용성을 확인하려면 다음 PromQL 쿼리를 사용하세요.

avg by (node_pool_name)(
  avg_over_time(
    kubernetes_io:node_pool_multi_host_available{
      monitored_resource="k8s_node_pool",
      cluster_name="CLUSTER_NAME"}[${__interval}]))

노드 중단 수

다음 GKE 시스템 측정항목은 마지막 샘플 이후 GKE 노드의 중단 수를 보고합니다(측정항목은 60초마다 샘플링됨).

  • kubernetes.io/node/interruption_count

interruption_type(예: TerminationEvent, MaintenanceEvent, PreemptionEvent) 및 interruption_reason(예: HostError, Eviction, AutoRepair) 필드는 노드가 중단된 이유를 제공하는 데 도움이 될 수 있습니다.

프로젝트의 클러스터에 있는 TPU 노드의 중단 및 원인을 분류하려면 다음 PromQL 쿼리를 사용하세요.

  sum by (interruption_type,interruption_reason)(
    sum_over_time(
      kubernetes_io:node_interruption_count{monitored_resource="k8s_node"}[${__interval}]))

호스트 유지보수 이벤트만 표시하려면 interruption_reason에 대한 HW/SW Maintenance 값을 필터링하도록 쿼리를 업데이트합니다. 다음 PromQL 쿼리를 사용합니다.

  sum by (interruption_type,interruption_reason)(
    sum_over_time(
      kubernetes_io:node_interruption_count{monitored_resource="k8s_node", interruption_reason="HW/SW Maintenance"}[${__interval}]))

노드 풀별로 집계된 중단 수를 확인하려면 다음 PromQL 쿼리를 사용하세요.

  sum by (node_pool_name,interruption_type,interruption_reason)(
    sum_over_time(
      kubernetes_io:node_pool_interruption_count{monitored_resource="k8s_node_pool", interruption_reason="HW/SW Maintenance", node_pool_name=NODE_POOL_NAME }[${__interval}]))

노드 풀 복구 시간(TTR)

다음 GKE 시스템 측정항목은 GKE 멀티 호스트 TPU 노드 풀의 복구 기간 분포를 보고합니다.

  • kubernetes.io/node_pool/accelerator/times_to_recover

이 측정항목에 기록된 각 샘플은 다운타임 기간의 노드 풀에 대한 단일 복구 이벤트를 나타냅니다.

이 측정항목은 멀티 호스트 TPU 노드 풀의 복구 시간 및 중단 사이의 시간을 추적하는 데 유용합니다.

다음 PromQL 쿼리를 사용하여 클러스터의 지난 7일 동안의 평균 복구 시간(MTTR)을 계산할 수 있습니다.

sum(sum_over_time(
  kubernetes_io:node_pool_accelerator_times_to_recover_sum{
    monitored_resource="k8s_node_pool", cluster_name="CLUSTER_NAME"}[7d]))
/
sum(sum_over_time(
  kubernetes_io:node_pool_accelerator_times_to_recover_count{
    monitored_resource="k8s_node_pool",cluster_name="CLUSTER_NAME"}[7d]))

중단 간 노드 풀 시간(TBI)

중단 간 노드 풀 시간은 중단이 발생하기 전까지 인프라가 실행되는 시간을 측정합니다. 이는 일정 기간의 평균으로 계산되며, 분자는 인프라가 가동된 총 시간을 측정하고 분모는 인프라의 총 중단 시간을 측정합니다.

다음 PromQL 예시에서는 지정된 클러스터의 7일 중단 간 평균 시간(MTBI)을 보여줍니다.

sum(count_over_time(
  kubernetes_io:node_memory_total_bytes{
    monitored_resource="k8s_node", node_name=~"gke-tpu.*|gk3-tpu.*", cluster_name="CLUSTER_NAME"}[7d]))
/
sum(sum_over_time(
  kubernetes_io:node_interruption_count{
    monitored_resource="k8s_node", node_name=~"gke-tpu.*|gk3-tpu.*", cluster_name="CLUSTER_NAME"}[7d]))

호스트 측정항목

GKE 버전 1.28.1-gke.1066000 이상에서 TPU 슬라이스의 VM은 TPU 사용률 측정항목을 GKE 시스템 측정항목으로 내보냅니다. TPU 호스트의 성능을 모니터링하기 위해 Cloud Monitoring에서 다음 측정항목을 사용할 수 있습니다.

  • TensorCore 사용률: 현재 사용되는 TensorCore의 백분율입니다. TensorCore 값은 행렬 곱셈 단위(MXU)와 벡터 단위의 합과 동일합니다. TensorCore 사용률은 지난 샘플 기간 (60초) 동안 수행된 TensorCore 작업을 동일 기간 동안 지원된 TensorCore 작업 수로 나눈 값입니다. 값이 클수록 사용률이 높습니다.
  • 메모리 대역폭 사용률: 사용 중인 가속기 메모리 대역폭의 현재 백분율입니다. 샘플 기간(60초) 동안 사용된 메모리 대역폭을 동일한 샘플 기간 동안 지원된 최대 대역폭으로 나누어 계산합니다.

이러한 측정항목은 Kubernetes 노드(k8s_node) 및 Kubernetes 컨테이너(k8s_container) 스키마에 있습니다.

Kubernetes 컨테이너:

  • kubernetes.io/container/accelerator/tensorcore_utilization
  • kubernetes.io/container/accelerator/memory_bandwidth_utilization

Kubernetes 노드:

  • kubernetes.io/node/accelerator/tensorcore_utilization
  • kubernetes.io/node/accelerator/memory_bandwidth_utilization

자세한 내용은 Kubernetes 측정항목 및 GKE 시스템 측정항목을 참조하세요.

로깅

TPU VM을 포함하여 GKE 노드에서 실행되는 컨테이너에서 내보낸 로그는 GKE Logging 에이전트에 의해 수집되어 Logging으로 전송되고 Logging에 표시됩니다.

Autopilot의 TPU 워크로드 권장사항

다음 권장사항으로 TPU 워크로드 효율성이 향상될 수 있습니다.

  • GKE에서 축소 또는 노드 업그레이드를 위해 포드를 종료하기 전의 유예 기간(최대 7일) 동안 연장된 실행 시간 포드를 사용합니다. 연장된 실행 시간 포드와 함께 유지보수 기간과 제외를 사용하여 자동 노드 업그레이드를 더 지연시킬 수 있습니다.
  • 용량 예약을 사용하면 워크로드가 가용성을 위해 큐에 배치되지 않고 요청된 TPU를 수신할 수 있습니다.

GKE에서 Cloud TPU를 설정하는 방법은 다음 Google Cloud 리소스를 참고하세요.