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.requestsspec.containers.resources.limits 필드의 TPU 칩 수입니다.

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

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

  1. 스팟 포드
  2. 특정 용량 예약
  3. 연장된 실행 시간 포드
  4. 유연한 시작

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
      Ironwood (TPU7x),
      tpu7x
      Dimensions (e.g. location):
      tpu_family:tpu7x
      Name:
      Preemptible TPU slices tpu7x
    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 노드 수 이상으로 설정해야 합니다.

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

TPU 애플리케이션 준비

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

  1. JAX, PyTorch, TensorFlow와 같은 프레임워크는 libtpu 공유 라이브러리를 사용하여 TPU VM에 액세스합니다. libtpu에는 XLA 컴파일러, TPU 런타임 소프트웨어, TPU 드라이버가 포함됩니다. PyTorch 및 JAX의 각 출시 버전에는 특정 libtpu.so 버전이 필요합니다. 패키지 버전 충돌을 방지하려면 JAX AI 이미지를 사용하는 것이 좋습니다. GKE에서 TPU를 사용하려면 다음 버전을 사용해야 합니다. tpu7x
    TPU 유형 libtpu.so 버전
    Ironwood (TPU7x)
    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. 워크로드 매니페스트에서 Kubernetes 노드 선택기를 추가하여 GKE가 정의된 TPU 머신 유형 및 TPU 토폴로지로 TPU 워크로드를 예약하도록 합니다.

      nodeSelector:
        cloud.google.com/gke-tpu-accelerator: TPU_ACCELERATOR
        cloud.google.com/gke-tpu-topology: TPU_TOPOLOGY
        cloud.google.com/placement-policy-name: WORKLOAD_POLICY # Required only for Ironwood (TPU7x)
      

    다음을 바꿉니다.

    • TPU_ACCELERATOR: TPU 가속기 이름입니다. 예를 들어 tpu7x-standard-4t을 사용합니다.
    • TPU_TOPOLOGY: TPU 슬라이스의 물리적 토폴로지입니다. 토폴로지 형식은 TPU 버전에 따라 달라집니다. 예를 들어 2x2x2를 사용하세요. 자세한 내용은 GKE에서 TPU 계획을 참고하세요.
    • WORKLOAD_POLICY: TPU 포드를 배치하는 데 사용할 워크로드 정책의 이름입니다. 이 노드 선택기는 Ironwood (TPU7x)에만 필요합니다.

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

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

GKE에서 TPU를 프로비저닝하는 데는 다음과 같은 구성 옵션이 있습니다.
  • 워크로드 요청: spec.nodeSelector 필드에서 TPU 버전과 토폴로지를 지정하고 spec.containers.resources 섹션에서 TPU 칩 수를 지정합니다. 워크로드를 배포하면 GKE에서 올바른 TPU 구성으로 노드를 자동으로 프로비저닝하고 노드의 리소스에 대한 전체 액세스를 보장하기 위해 각 워크로드를 자체 전용 노드에 배치합니다. 자세한 내용은 워크로드에서 TPU 요청을 참고하세요.
  • 커스텀 컴퓨팅 클래스로 중앙에서 TPU 프로비저닝

    다음 섹션에서는 커스텀 ComputeClass를 만든 다음 ComputeClass에 정의된 TPU를 사용하는 작업을 만드는 방법을 보여줍니다.

    맞춤 ComputeClass 만들기

    TPU 규칙을 따르는 맞춤 ComputeClass를 만드는 단계는 Ironwood (TPU7x)를 사용하는지 아니면 이전 TPU 버전을 사용하는지에 따라 다릅니다.

    Ironwood (TPU7x)

    1. 워크로드 정책을 만듭니다. 이 단계는 선택한 토폴로지에 따라 달라지는 멀티 호스트 노드 풀을 만드는 경우에만 필요합니다. 단일 호스트 노드 풀을 사용하는 경우 이 단계를 건너뜁니다.

      gcloud compute resource-policies create workload-policy WORKLOAD_POLICY_NAME \
          --type=HIGH_THROUGHPUT \
          --accelerator-topology=TPU_TOPOLOGY \
          --project=PROJECT_ID \
          --region=REGION
      

      다음을 바꿉니다.

      • WORKLOAD_POLICY_NAME: 워크로드 정책의 이름입니다.
      • TPU_TOPOLOGY: TPU Ironwood (TPU7x) 토폴로지입니다. 예를 들어 2x2x2를 사용하세요. 지원되는 모든 Ironwood (TPU7x) 토폴로지에 대한 자세한 내용은 토폴로지 섹션을 참고하세요.
      • PROJECT_ID: Google Cloud 프로젝트 ID입니다.
      • REGION: 워크로드 정책의 리전입니다. 워크로드 정책은 리전별 리소스이며 노드 풀 전반에서 사용할 수 있습니다.
    2. 다음 매니페스트를 tpu-compute-class.yaml로 저장합니다.

      apiVersion: cloud.google.com/v1
      kind: ComputeClass
      metadata:
        name: tpu-class
      spec:
        priorities:
          - tpu:
              type: tpu7x
              topology: TPU_TOPOLOGY
              count: 4
            placement:
              policyName: WORKLOAD_POLICY_NAME
        nodePoolAutoCreation:
          enabled: true
      
    3. (선택사항) 특정 예약 또는 하위 블록을 사용할 수 있습니다. 예를 들어 ComputeClass 매니페스트에 다음 specs를 추가할 수 있습니다.

        reservations:
          affinity: Specific
          specific:
            - name: RESERVATION_NAME
              reservationBlock:
                name: RESERVATION_BLOCK_NAME
                reservationSubBlock:
                  name: RESERVATION_SUB_BLOCK_NAME
      

      다음을 바꿉니다.

      • RESERVATION_NAME: Compute Engine 용량 예약의 이름입니다.
      • RESERVATION_BLOCK_NAME: Compute Engine 용량 예약 블록의 이름입니다.
      • RESERVATION_SUB_BLOCK_NAME: Compute Engine 용량 예약 하위 블록의 이름입니다.

      자세한 내용은 예약된 영역별 리소스 소비를 참고하세요.

    기타 TPU 버전

    TPU용으로 구성된 커스텀 ComputeClass를 사용하여 v3, v4, v5p, v5e 또는 v6e (Trillium) TPU를 프로비저닝하려면 다음 단계를 완료하세요.

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

      apiVersion: cloud.google.com/v1
      kind: ComputeClass
      metadata:
        name: tpu-class
      spec:
        priorities:
        - tpu:
            type: TPU_TYPE
            count: NUMBER_OF_CHIPS
            topology: TOPOLOGY
        - spot: true
          tpu:
            type: TPU_TYPE
            count: NUMBER_OF_CHIPS
            topology: TOPOLOGY
        - flexStart:
            enabled: true
          tpu:
            type: TPU_TYPE
            count: NUMBER_OF_CHIPS
            topology: TOPOLOGY
        nodePoolAutoCreation:
          enabled: true
      

      다음을 바꿉니다.

      • TPU_TYPE: 사용할 TPU 유형입니다(예: tpu-v4-podslice). GKE에서 지원하는 값이어야 합니다.
      • TOPOLOGY: 슬라이스의 TPU 칩 배열입니다(예: 2x2x4). 선택한 TPU 유형에 지원되는 토폴로지여야 합니다.
      • NUMBER_OF_CHIPS: 사용할 컨테이너의 TPU 칩 수입니다. limitsrequests 값이 같아야 합니다.
    2. ComputeClass를 배포합니다.

      kubectl apply -f tpu-compute-class.yaml
      

      커스텀 ComputeClass 및 TPU에 대한 자세한 내용은 TPU 구성을 참고하세요.

    TPU를 사용하는 작업 만들기

    1. 다음 매니페스트를 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 칩 수입니다. limitsrequests에 대해 동일한 값이어야 하며 선택한 맞춤 ComputeClass의 CHIP_COUNT 값과 같아야 합니다.
      • MEMORY_SIZE: TPU가 사용하는 최대 메모리 양입니다. 메모리 한도는 사용 중인 TPU 버전과 토폴로지에 따라 다릅니다. 자세한 내용은 가속기의 최솟값 및 최댓값을 참조하세요.
      • NUMBER_OF_CHIPS: 사용할 컨테이너의 TPU 칩 수입니다. limitsrequests 값이 같아야 합니다.
    2. 작업을 배포합니다.

      kubectl create -f tpu-job.yaml
      

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

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

      자세한 내용은 커스텀 ComputeClass 정보를 참고하세요.

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

      kubectl delete -f tpu-job.yaml
      

    워크로드에서 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
        parallelism: NUM_OF_VMS
        completions: NUM_OF_VMS
        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: TPU_TOPOLOGY
              topology.kubernetes.io/zone: LOCATION
            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
      

      다음 순서대로 값을 바꿉니다.

      1. TPU_TYPE: 사용할 TPU 버전을 선택합니다(예: tpu-v5-lite-podslice). 값은 GKE에서 지원해야 합니다.
      2. LOCATION: GKE가 TPU VM을 배치하는 영역을 검증합니다.

        선택사항: us-central1-ai1a와 같은 AI 영역을 사용합니다. AI 영역은 Google Cloud 리전 내에서 AI/ML 워크로드에 최적화된 특수 위치입니다.

      3. TPU_TOPOLOGY: 슬라이스의 TPU 칩 배열을 선택합니다(예: 4x4). 선택한 TPU 유형에 지원되는 토폴로지여야 합니다.
      4. 토폴로지 선택 섹션의 표에서 선택한 TPU_TYPETPU_TOPOLOGY을 찾습니다.
        • NUM_OF_VMS: '기술 사양' 열의 'VM 수' 값을 사용합니다. 작업 사양에서 parallelismcompletions 필드의 값을 설정합니다.
        • NUMBER_OF_CHIPS: VM당 칩 수를 입력합니다. 선택한 행의 '토폴로지의 TPU 칩 수' 값을 'VM 수' 값으로 나누어 계산합니다. google.com/tpu 리소스 요청에 이 값을 설정합니다.
      5. MEMORY_SIZE: TPU가 사용하는 최대 메모리 양입니다. 메모리 한도는 사용 중인 TPU 버전과 토폴로지에 따라 다릅니다. 자세한 내용은 가속기의 최솟값 및 최댓값을 참고하세요.

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

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

      예를 들어 4x4 토폴로지가 있는 tpu-v5-lite-podslice를 생각해 보겠습니다. 이 구성은 다음을 의미합니다.

      • 토폴로지에는 총 16개의 칩이 있습니다 (4*4=16).
      • tpu-v5-lite-podslice는 VM당 4개의 칩이 있는 ct5lp-hightpu-4t 머신에 배포됩니다. 즉, 4x4 토폴로지를 사용하려면 VM이 4개 (16/4=4) 필요합니다. 따라서 parallelismcompletions 필드를 4로 설정합니다.
      • 포드의 각 컨테이너는 4개의 TPU 칩을 요청하며 이는 VM당 칩 수에 해당합니다. 매니페스트는 google.com/tpu: 4 줄로 이를 지정합니다.
    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 칩 수를 반환합니다. 멀티 호스트 슬라이스를 만들기 위한 워크로드 매개변수는 다음과 같습니다.

    • 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)

    노드 풀 TBI는 중단이 발생하기 전까지 TPU 노드가 실행되는 평균 시간입니다.

    다음 PromQL 쿼리는 실행 중인 노드에서 60초 메모리 샘플의 총수를 계산하여 이 TBI를 계산합니다. 각 샘플의 존재는 노드가 60초 동안 작동했음을 나타냅니다. 일정 기간 동안 이러한 샘플을 집계하면 노드의 총 업타임(분)을 추정할 수 있습니다. 예를 들어 10분 동안 10개의 샘플은 노드가 약 10분 동안 작동했음을 의미합니다. 이 개수는 노드 가동시간(분)의 프록시 역할을 합니다.

    그런 다음 선택한 기간의 중단 횟수로 이 수를 나눕니다.

    다음 예시에서는 지정된 클러스터의 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 리소스를 참조하세요.