맞춤 스케줄러로 동적 슬라이싱 사용

이 문서에서는 슬라이스 커스텀 리소스와 직접 상호작용하여 동적 슬라이싱을 사용하는 방법을 설명합니다. 슬라이스를 만들고, 파티션 상태를 모니터링하고, 슬라이스 상태를 확인할 수 있습니다.

이 안내를 따르기 전에 동적 슬라이싱 개념을 이해해야 합니다.

맞춤 스케줄러로 동적 슬라이싱을 사용해야 하는 이유

복잡한 일정 요구사항이 있거나 동적 슬라이싱을 기존 일정 인프라와 통합하려는 경우 자체 스케줄러를 사용하여 슬라이스 커스텀 리소스를 관리하세요.

슬라이스 커스텀 리소스를 직접 관리하는 대신 스케줄러를 사용하려면 GKE에서 Kueue 및 토폴로지 인식 스케줄링 (TAS)과의 통합을 제공합니다. 자세한 내용은 Kueue 및 TAS로 동적 슬라이스 예약을 참고하세요.

워크플로 개요

맞춤 스케줄러로 동적 슬라이싱을 사용하려면 이 문서에서 다음 작업을 실행하세요.

  1. 슬라이스 컨트롤러를 사용 설정합니다.
  2. 증분 프로비저닝으로 노드 풀 만들기
  3. 워크로드 요구사항에 따라 슬라이스 커스텀 리소스 만들기 슬라이스 커스텀 리소스를 클러스터에 적용합니다.
  4. 파티션 상태 및 슬라이스 상태를 모니터링합니다.
  5. 완료되면 슬라이스를 삭제합니다.

슬라이스 커스텀 리소스의 필드 및 상태에 대한 자세한 내용은 슬라이스 커스텀 리소스 참조 정보를 참고하세요.

요구사항

GKE에서 동적 슬라이싱을 사용하려면 다음 요구사항을 충족해야 합니다.

  • 신속 채널에서 버전 1.35.2-gke.1842000 이상의 Standard 클러스터를 사용합니다.
  • Ironwood (TPU7x) 버전을 사용합니다.
  • 노드에 Container-Optimized OS 이미지를 사용합니다.
  • 증분 프로비저닝을 사용하려면 모든 용량 모드 예약을 사용하세요. 모든 용량 모드는 TPU Cluster Director에서 사용 설정하는 기능입니다.

시작하기 전에

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

  • Google Kubernetes Engine API를 사용 설정합니다.
    • Google Kubernetes Engine API 사용 설정
  • 이 태스크에 Google Cloud CLI를 사용하려면 gcloud CLI를 설치한 후 초기화합니다. 이전에 gcloud CLI를 설치했으면 gcloud components update 명령어를 실행하여 최신 버전을 가져옵니다. 이전 gcloud CLI 버전에서는 이 문서의 명령어를 실행하지 못할 수 있습니다.

슬라이스 컨트롤러 사용 설정

동적 슬라이싱을 사용하려면 클러스터에서 슬라이스 컨트롤러를 사용 설정하세요.

  1. 클러스터를 업데이트합니다.

    gcloud container clusters update CLUSTER_NAME \
    --location=LOCATION \
    --enable-slice-controller

    다음을 바꿉니다.

  2. kubectl 명령어로 클러스터와 통신할 수 있도록 사용자 인증 정보를 가져옵니다.

    gcloud config set container/cluster CLUSTER_NAME
gcloud container clusters get-credentials CLUSTER_NAME \
    --location=LOCATION

  3. 다음 명령어의 출력에서 slices.accelerator.gke.io 값이 있는지 확인합니다.

    kubectl get crd slices.accelerator.gke.io

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

    slices.accelerator.gke.io                2026-01-09T23:58:02Z

증분 프로비저닝으로 노드 풀 만들기

이 섹션에서는 증분 프로비저닝을 사용하여 TPU 노드 풀을 만드는 방법을 설명합니다. GKE는 모든 TPU 용량을 TPU VM 또는 하위 블록의 16노드 그룹의 노드 풀로 변환합니다. GKE는 호스트 머신의 정상 부분에 노드를 배치하고 비정상 머신이 복구되는 동안 점진적으로 프로비저닝하여 정상 VM 16개를 모두 찾을 수 없는 경우에도 이러한 노드 풀을 프로비저닝합니다.

다음 중 하나에 속하도록 노드 풀을 타겟팅할 수 있습니다.

  • 모든 용량 모드 예약에 노출되는 특정 TPU 블록입니다. 차단 타겟팅을 사용하면 GKE가 지정된 블록 내에서 사용 가능한 모든 하위 블록에 노드 풀을 만들 수 있습니다.
  • 더 세부적인 제어를 위한 TPU의 특정 하위 블록 또는 특정 16노드 TPU VM 그룹입니다.

워크로드 정책 만들기

Ironwood (TPU7x)로 TPU 슬라이스 노드 풀을 만들려면 먼저 accelerator-topology-mode 필드가 provision_only로 설정된 워크로드 정책을 만들어야 합니다. 이 설정은 증분 프로비저닝 프로세스를 트리거합니다.

워크로드 정책을 만듭니다.

gcloud compute resource-policies create workload-policy WORKLOAD_POLICY_NAME \
        --project=PROJECT_ID \
        --region=REGION  \
        --type=HIGH_THROUGHPUT \
        --accelerator-topology=4x4x4 \
        --accelerator-topology-mode=provision_only

다음을 바꿉니다.

  • WORKLOAD_POLICY_NAME: 워크로드 정책의 이름입니다.
  • PROJECT_ID: Google Cloud 프로젝트 ID입니다.
  • REGION: 워크로드 정책의 리전입니다.

이 명령어에서 다음을 수행합니다.

  • 단일 하위 블록 내의 총 칩 수와 일치하도록 항상 accelerator-topology 필드를 4x4x4로 설정합니다.
  • 증분 프로비저닝 프로세스가 트리거되도록 항상 accelerator-topology-mode 필드를 provision_only로 설정하세요. provision_only 필드가 설정되면 노드 풀은 ICI 또는 OCS 링크를 형성하지 않고 TPU 노드를 프로비저닝합니다.

블록 또는 하위 블록에 속하도록 노드 풀 타겟팅

'전체 용량' 모드 예약 내에서 특정 하위 블록 또는 블록을 타겟팅할 수 있습니다.

  • 블록 타겟팅: 각 노드 풀은 지정된 블록의 용량을 사용합니다. GKE는 해당 블록의 사용 가능한 하위 블록 내에 노드 풀을 배치합니다. 사용하려는 블록에 하위 블록이 있는 만큼 노드 풀을 만들어야 합니다.

  • 하위 블록 타겟팅: 각 노드 풀은 사용 가능한 특정 하위 블록에 매핑됩니다. 하위 블록 타겟팅을 사용하는 경우 하나 이상의 VM이 정상이면 GKE가 노드 풀을 만듭니다. 증분 프로비저닝은 모든 노드가 지정된 하위 블록 내에 배치되도록 하는 데 도움이 됩니다.

Block

  1. 예약의 블록 이름과 블록에서 사용 가능한 하위 블록 수를 가져오려면 모든 용량 모드 예약의 토폴로지 및 상태 보기 문서에서 다음 단계를 완료하세요.

    1. 모든 예약 블록을 나열하고 name: 필드의 값을 복사하여 블록의 이름을 식별합니다. 이 값은 이 문서에 있는 블록 또는 BLOCK_NAME의 이름입니다.

    2. 예약 블록을 설명하고 reservationSubBlockCount 필드의 값을 식별하여 만들 노드 풀 수를 확인합니다. 이 값은 사용 가능한 하위 블록 수입니다. 예를 들어 reservationSubBlockCount: 4 값은 블록에 사용할 수 있는 하위 블록이 4개 있으며 별도의 노드 풀 4개를 만들어야 함을 나타냅니다.

  2. 예약 경로를 설정합니다.

    export RESERVATION_PATH="projects/PROJECT_ID/reservations/RESERVATION_NAME/reservationBlocks/BLOCK_NAME"

    다음을 바꿉니다.

    • RESERVATION_NAME: TPU 예약의 이름입니다.
    • BLOCK_NAME: 블록의 이름입니다.

  3. 이전 단계에서 식별된 각 하위 블록에 대해 노드 풀을 만듭니다. 예를 들어 개수가 4이면 이 명령어를 네 번 실행합니다. 각 노드 풀에 고유한 이름을 사용합니다.

    gcloud container node-pools create NODE_POOL_NAME \
      --cluster=CLUSTER_NAME \
      --node-locations=ZONE \
      --machine-type=tpu7x-standard-4t \
      --num-nodes=16 \
      --placement-policy=WORKLOAD_POLICY_NAME \
      --reservation-affinity=specific \
      --reservation=${RESERVATION_PATH}

    다음을 바꿉니다.

    • NODE_POOL_NAME: 새 노드 풀의 이름입니다.
    • CLUSTER_NAME: GKE 클러스터의 이름입니다.
    • WORKLOAD_POLICY_NAME: 생성한 워크로드 정책의 이름입니다.
    • ZONE: 노드 풀의 영역입니다(예: us-central1-a).

하위 블록

  1. 블록 이름과 사용 가능한 하위 블록의 ID를 가져오려면 모든 용량 모드 예약의 토폴로지 및 상태 보기 문서에서 다음 단계를 완료하세요.

    1. 차단 이름을 확인하려면 모든 예약 차단을 나열하고 name: 필드의 값을 복사합니다. 이 값은 이 문서에 있는 블록 또는 BLOCK_NAME의 이름입니다.

    2. 하위 블록의 이름을 식별하려면 블록의 모든 하위 블록을 나열하고 reservationSubBlocks 아래 각 항목의 name: 필드에 있는 값을 복사합니다. 이 값은 이 문서에 나오는 하위 블록 또는 SUBBLOCK_NAME의 이름입니다.

  2. 예약 경로를 설정합니다.

    export RESERVATION_PATH="projects/PROJECT_ID/reservations/RESERVATION_NAME/reservationBlocks/BLOCK_NAME/reservationSubBlocks/SUBBLOCK_NAME"

    다음을 바꿉니다.

    • RESERVATION_NAME: TPU 예약의 이름입니다.
    • BLOCK_NAME: 블록의 이름입니다.
    • SUBBLOCK_NAME: 하위 블록의 이름입니다.

  3. 노드 풀을 만듭니다.

    gcloud container node-pools create NODE_POOL_NAME \
        --project=PROJECT_ID \
        --cluster=CLUSTER_NAME \
        --node-locations=ZONE \
        --machine-type=tpu7x-standard-4t \
        --num-nodes=16 \
        --placement-policy=WORKLOAD_POLICY_NAME \
        --reservation-affinity=specific \
        --reservation=${RESERVATION_PATH}

    다음을 바꿉니다.

    • NODE_POOL_NAME: 새 노드 풀의 고유한 이름(예: sub-block-pool-1)
    • PROJECT_ID: Google Cloud 프로젝트 ID입니다.
    • CLUSTER_NAME: GKE 클러스터의 이름입니다.
    • ZONE: 노드 풀의 영역입니다(예: us-central2-b).
    • WORKLOAD_POLICY_NAME: 생성한 워크로드 정책의 이름입니다.

이 단계에서는 노드가 생성되지만 칩 간 상호 연결 (ICI) 링크는 아직 활성화되지 않습니다. 따라서 이러한 노드 풀에서 직접 워크로드를 실행할 수 없습니다.

슬라이스를 형성하고 워크로드를 예약할 수 있도록 필요한 모든 ICI 링크를 사용 설정하려면 다음 방법 중 하나를 사용하여 동적 슬라이스를 만드세요.

  • 슬라이스 커스텀 리소스를 만듭니다. 포드 대신 슬라이스 컨트롤러가 활성화하는 지정된 토폴로지를 정의하는 슬라이스 커스텀 리소스를 사용합니다.
  • Kueue 및 TAS로 GKE 워크로드를 예약합니다. Kueue는 Slice 커스텀 리소스의 생성 및 삭제를 자동으로 처리합니다. Kueue에서 만든 슬라이스 커스텀 리소스를 수동으로 수정하지 마세요.

동적 슬라이스 형성

노드 풀을 만든 후 Slice 커스텀 리소스를 만들어 더 큰 동적 슬라이스를 형성할 수 있습니다. 포드 대신 슬라이스 컨트롤러가 활성화하는 지정된 토폴로지를 정의하는 슬라이스 커스텀 리소스를 사용합니다.

노드 및 파티션 상태 확인

  1. 노드 풀에서 노드 이름을 가져오려면 다음 명령어를 실행합니다.

    kubectl get nodes -l cloud.google.com/gke-nodepool=${NODE_POOL_NAME}

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

    NAME                                 STATUS   ROLES    AGE    VERSION
gke-np-status-update-7b4c890c-0jhp   Ready    <none>   2d1h   v1.35.1-gke.1396002
gke-np-status-update-7b4c890c-377r   Ready    <none>   2d1h   v1.35.1-gke.1396002
gke-np-status-update-7b4c890c-gb51   Ready    <none>   2d1h   v1.35.1-gke.1396002

  2. 노드의 프로비저닝 모델을 확인합니다.

    kubectl describe node NODE_NAME | grep "cloud.google.com/gke-accelerator-topology-mode"

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

    cloud.google.com/gke-accelerator-topology-mode: PROVISION_ONLY

    이 값은 워크로드 정책을 만들 때 정의된 accelerator-topology-mode=provision_only 설정과 일치합니다.

  3. 노드 라벨 정보를 검색합니다.

    kubectl describe node NODE_NAME | grep "cloud.google.com/gke-tpu-partition-4x4x4-id"

    NODE_NAME을 노드 풀에 있는 노드 중 하나의 이름으로 바꿉니다.

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

    cloud.google.com/gke-tpu-partition-4x4x4-id=fba785f80d18552357dcdef6d3d16c27

    cloud.google.com/gke-tpu-partition-4x4x4-state 주석은 노드를 사용하여 동적 슬라이스를 형성할 수 있는지 여부를 나타냅니다. 이 라벨은 다음 값을 지원합니다.

    • HEALTHY: 노드가 정상이며 완전히 작동합니다.
    • DEGRADED: 노드가 손상되었지만 동적 슬라이스 형성에 여전히 사용할 수 있습니다.
    • UNHEALTHY: 노드가 오작동하여 슬라이스를 형성하는 데 사용할 수 없습니다.
    • UNSET: 노드 풀의 노드가 부족하여 상태가 정의되지 않았습니다.
    • INCOMPLETE: 파티션 내의 일부 노드가 프로비저닝되지 않았습니다.

  4. 노드에 node.gke.io/created-by-mig 주석이 포함되어 있는지 확인합니다.

    kubectl describe node NODE_NAME | grep "node.gke.io/created-by-mig"

    NODE_NAME을 노드 풀에 있는 노드 중 하나의 이름으로 바꿉니다.

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

    node.gke.io/created-by-mig: projects/735972712744/zones/us-central1-ai1a/team/string

    출력에는 GKE 컨트롤 플레인이 Kubernetes 노드를 기본 Compute Engine 리소스에 연결할 수 있도록 하는 node.gke.io/created-by-mig 라벨이 포함됩니다.

슬라이스 커스텀 리소스 만들기

  1. 슬라이스 커스텀 리소스를 정의합니다.

    apiVersion: accelerator.gke.io/v1beta1
kind: Slice
metadata:
  # Name of the slice resource
  name: SLICE_NAME
spec:
  # Specify the type of accelerator for this slice
  type: "tpu7x"
  # Define the desired topology for the accelerator slice
  topology: TOPOLOGY
  partitionIds:
    - PARTITION_ID # Example: a9476d1b02bd4f4e75ffffae3bd23c01
    - PARTITION_ID_2
    # ... add more partition IDs as needed

    다음을 바꿉니다.

    • SLICE_NAME: 슬라이스의 이름입니다. 이름은 metadata.name 조건을 충족해야 합니다.
    • TOPOLOGY: 동적 슬라이스의 토폴로지입니다. 토폴로지는 다음 조건을 충족해야 합니다.
      • 요청된 토폴로지의 각 차원은 4의 배수여야 합니다(예: 4A x 4B x 4C).
      • 토폴로지 차원의 세 값 AxBxC은 비감소 순서 (A ≤ B ≤ C)여야 합니다. 예를 들어 4x4x8는 유효하지만 4x8x4는 유효하지 않습니다. 이 순서는 일관된 슬라이스 형성을 보장하고 예기치 않은 동작을 방지하는 데 도움이 됩니다.
      • 토폴로지 차원의 세 값의 곱인 A × B × C은 9,216을 초과해서는 안 됩니다.
    • PARTITION_ID: 슬라이스를 구성하는 4x4x4 파티션을 식별하는 문자열 목록입니다. 각 파티션이 64개의 칩으로 구성된 총 칩 수를 기반으로 파티션 수를 계산합니다. spec.partitionIds 목록의 항목 수는 계산된 파티션 수 ((A × B × C) / 64)와 정확히 일치해야 합니다. partitionIds는 다음 조건을 충족해야 합니다.
      • 각 파티션은 예약 하위 블록에 매핑되어야 합니다.
      • 연결된 모든 하위 블록은 동일한 예약에 속해야 합니다.
      • 연결된 모든 블록은 동일한 예약 내에 있어야 합니다.
      • 연결된 노드 풀의 모든 노드가 ready 상태여야 합니다.
    • type 필드의 값은 tpu7x이어야 합니다.
    • 슬라이스 컨트롤러가 슬라이스 형성 중에 자동으로 재시도하도록 하려면 슬라이스 커스텀 리소스에 slice.gke.io/retry-on-failure: "true" 주석을 추가하면 됩니다. SliceCreationFailed 상태 이유로 인해 슬라이스가 생성되지 않으면 컨트롤러는 슬라이스가 성공적으로 형성될 때까지 재시도합니다.

    예를 들어 4x8x8 슬라이스를 만들려면 고유한 파티션 ID 4개를 제공해야 합니다.

    apiVersion: accelerator.gke.io/v1beta1
kind: Slice
metadata:
    name: test-slice-example
    annotations:
      slice.gke.io/retry-on-failure: "true" # Optional annotation to retry slice formation
spec:
    type: "tpu7x"
    topology: "4x8x8" # (4*8*8)/64 = 4 partitions
    partitionIds:
        - "p0"
        - "p1"
        - "p2"
        - "p3"

  2. 슬라이스 커스텀 리소스를 적용합니다.

    kubectl apply -f test-slice-example.yaml

    이 시점에서 GKE는 슬라이스를 만들려고 시도합니다. 다음 문제 중 하나가 발생하면 슬라이스 생성이 실패하고 슬라이스 커스텀 리소스의 상태 이유가 SliceCreationFailed 또는 FAILED로 업데이트됩니다.

    • 맞춤 리소스에서 선택한 노드가 없으면 상태 이유는 SliceCreationFailed입니다.
    • 커스텀 리소스의 노드가 다른 슬라이스에서 사용되는 경우 상태 이유는 SliceCreationFailed입니다.
    • 커스텀 리소스의 노드가 동일한 예약 블록에 속하지 않으면 상태 이유는 FAILED입니다.
    • 노드가 동일한 예약에 있지 않으면 상태 이유는 FAILED입니다.
    • 토폴로지가 파티션 수와 일치하지 않으면 상태 이유가 SliceCreationFailed입니다.

    슬라이스 커스텀 리소스의 상태에 대해 자세히 알아보려면 슬라이스 상태를 참고하세요.

슬라이스 커스텀 리소스 상태 모니터링

슬라이스 커스텀 리소스의 상태를 확인하려면 다음 명령어를 실행하세요.

kubectl describe slice SLICE_NAME

SLICE_NAME을 슬라이스 이름으로 바꿉니다.

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

Name:         test-slice
Namespace:
Labels:       <none>
Annotations:  <none>
API Version:  accelerator.gke.io/v1beta1
Kind:         Slice
Metadata:
  Creation Timestamp:  2026-01-11T23:45:15Z
  Finalizers:
    accelerator.gke.io/slice-finalizer
  Generation:        1
  Resource Version:  1768175347356335006
  UID:               d0b71e5c-be3f-4788-aead-930c7afec4f2
Spec:
  Partition Ids:
    2c79463990ff67c4e3c2648666bfedfa
    ba898ffcac0ad0946e8ff036d771ee53
    [more partition IDs]
  Topology:  8x16x16
  Type:      tpu7x
Status:
  Conditions:
    Last Transition Time:  2026-01-11T23:45:38Z
    Message:               ""
    
    Reason:                FAILED
    
    Status:                False
    Type:                  Ready
Events:

슬라이스 커스텀 리소스의 상태에 있는 reason 필드는 슬라이스의 현재 상태를 나타냅니다. 슬라이스 커스텀 리소스의 수명 주기는 다음 단계를 따릅니다.

  • SliceNotCreated: 컨트롤러가 초기화 및 리소스 검사를 실행합니다.
    • 필수 요건이 충족되지 않으면 상태가 SliceCreationFailed로 전환됩니다.
    • 유효성 검사를 통과하면 상태가 ACTIVATING로 전환됩니다.
  • ACTIVATING: GKE가 슬라이스를 형성하고 있습니다.
    • 성공하면 상태가 ACTIVE로 전환됩니다.
    • 하위 블록이 저하되었지만 슬라이스를 사용할 수 있는 경우 상태가 ACTIVE_DEGRADED로 전환됩니다.
    • 형성에 실패하면 상태가 FAILED로 전환됩니다.
  • DEACTIVATING: 슬라이스 커스텀 리소스가 삭제되거나 활성 또는 실패 상태에서 심각한 오류가 발생하면 슬라이스가 해체되기 시작합니다.
  • INCOMPLETE: 리소스가 완전히 삭제되기 전의 최종 단계입니다.

슬라이스 커스텀 리소스의 상태에 대해 자세히 알아보려면 슬라이스 상태를 참고하세요.

동적 슬라이싱에서 워크로드 실행

슬라이스 커스텀 리소스가 ACTIVE 상태인 경우 슬라이스에서 워크로드를 실행할 수 있습니다. 다음 섹션에는 동적 슬라이싱을 사용하는 워크로드의 예가 포함되어 있습니다. 워크로드는 작업 또는 JobSet으로 제출됩니다.

예 1: 단일 워크로드가 단일 슬라이스를 사용함

다음 예는 단일 하위 블록 슬라이스를 사용하는 워크로드를 보여줍니다.

  1. 다음 샘플 매니페스트를 tpu-job-jax-v7x-64.yaml로 저장합니다.

    apiVersion: v1
kind: Service
metadata:
name: headless-svc
spec:
clusterIP: None
selector:
    job-name: tpu-job-jax-v7x-64
---
apiVersion: batch/v1
kind: Job
metadata:
name: tpu-job-jax-v7x-64
spec:
backoffLimit: 0
completions: 16
parallelism: 16
completionMode: Indexed
template:
    metadata:
    annotations:
        cloud.google.com/gke-tpu-slice-topology: 4x4x4
    spec:
    nodeSelector:
        cloud.google.com/gke-tpu-topology: 4x4x4
        cloud.google.com/gke-tpu-accelerator: tpu7x
        cloud.google.com/gke-tpu-slice: test-slice
    subdomain: headless-svc
    restartPolicy: Never
    containers:
    - name: tpu-job-jax
        env:
        - name: TPU_ACCELERATOR_TYPE
          value: tpu7x-128
        image: python:3.12
        securityContext:
        privileged: false
        command:
        - bash
        - -c
        - |
        set -ex
        pip install -U --pre jax jaxlib libtpu requests -i https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/ -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
        pip list
        python -c 'import jax; print("Total TPU devices (cores):", jax.device_count())'
        resources:
        requests:
            google.com/tpu: 4
        limits:
            google.com/tpu: 4

    이 매니페스트에서 각 항목은 다음을 수행합니다.

    • cloud.google.com/gke-tpu-slice-topology 및 cloud.google.com/gke-tpu-topology은 동적 슬라이스의 토폴로지를 정의합니다.
    • env.value: tpu7x-128은 TPU 가속기 유형과 슬라이스의 총 코어 수입니다. 코어 수는 토폴로지의 차원에 칩당 코어 수를 곱하여 계산됩니다. 예를 들어 4x4x4 토폴로지의 경우 계산은 4 × 4 × 4 × 2 = 128이며, 여기서 2는 tpu7x (Ironwood (TPU7x))의 칩당 코어 수입니다. 따라서 TPU_ACCELERATOR_TYPE은 tpu7x-128입니다.

  2. tpu-job-jax-v7x-64.yaml 매니페스트를 적용합니다.

    kubectl apply -f tpu-job-jax-v7x-64.yaml

예 2: JobSet을 사용하여 멀티슬라이스 노드 풀에 워크로드 배포

이 예에서는 JobSet을 사용하여 멀티슬라이스 노드 풀에 워크로드를 배포하는 방법을 보여줍니다.

  1. JobSet 설치:

    kubectl apply --server-side -f https://github.com/kubernetes-sigs/jobset/releases/download/v0.10.1/manifests.yaml

  2. 다음 샘플 매니페스트를 tpu-multislice-jax.yaml로 저장합니다.

    apiVersion: jobset.x-k8s.io/v1alpha2
kind: JobSet
metadata:
  name: tpu-multislice-jax
  annotations:
    alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-tpu-slice
spec:
  failurePolicy:
    maxRestarts: 3
  replicatedJobs:
  - name: slice-job
    replicas: 2
    template:
      spec:
        parallelism: 16
        completions: 16
        backoffLimit: 0
        completionMode: Indexed
        template:
          metadata:
            annotations:
              # The shape of the slice
              cloud.google.com/gke-tpu-slice-topology: 4x4x4
          spec:
            hostNetwork: true
            dnsPolicy: ClusterFirstWithHostNet
            nodeSelector:
              cloud.google.com/gke-tpu-topology: 4x4x4
              cloud.google.com/gke-tpu-accelerator: tpu7x
              # IMPORTANT: Do NOT put 'cloud.google.com/gke-tpu-slice' here manually.
              # The exclusive-topology annotation handles the slice assignment automatically.
            containers:
            - name: jax-worker
              image: python:3.12
              securityContext:
                privileged: true
              ports:
              - containerPort: 8471
              command:
              - bash
              - -c
              - |
                set -ex
                pip install -U --pre jax jaxlib libtpu requests -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
                # Verify JobSet injected the specific slice ID for this worker
                echo "JobSet Index: $JOB_COMPLETION_INDEX"
                python -c 'import jax; print("Total TPU devices:", jax.device_count())'
              resources:
                requests:
                  google.com/tpu: 4
                limits:
                  google.com/tpu: 4

  3. tpu-multislice-jax.yaml 매니페스트를 적용합니다.

    kubectl apply -f tpu-multislice-jax.yaml

    이 매니페스트에서 각 항목은 다음을 수행합니다.

    • replicatedJobs 아래의 replicas: 2 필드는 JobSet이 각각 4x4x4 TPU 슬라이스에 해당하는 두 개의 별도 작업을 생성함을 나타냅니다.
    • alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-tpu-slice 주석은 각 작업이 고유한 TPU 슬라이스에 할당되도록 하는 데 도움이 됩니다.
    • cloud.google.com/gke-tpu-slice-topology: 4x4x4 주석은 각 동적 슬라이스의 토폴로지를 정의합니다.
    • JobSet이 슬라이스 할당을 처리하므로 이 예에서는 TPU_ACCELERATOR_TYPE 환경 변수가 명시적으로 설정되지 않습니다. JAX 코드는 할당된 슬라이스 내에서 사용 가능한 TPU 기기를 자동으로 감지합니다.

슬라이스 삭제

  1. 슬라이스를 삭제합니다.

    kubectl patch slice $SLICE_NAME --type json \
  -p='[{"op": "remove", "path": "/metadata/finalizers"}]'

  2. 슬라이스가 삭제되었는지 확인합니다.

    kubectl get slices

슬라이스 컨트롤러 사용 중지

슬라이스 컨트롤러를 사용 중지하려면 클러스터에서 삭제하세요.

  1. 슬라이스 커스텀 리소스가 비어 있는지 확인합니다.

    kubectl get slice -A

  2. 클러스터를 업데이트하여 슬라이스 컨트롤러를 사용 중지합니다.

    gcloud container clusters update ${CLUSTER_NAME} \
    --location=${REGION} \
    --no-enable-slice-controller

  3. 슬라이스 커스텀 리소스를 삭제합니다.

    kubectl delete crd slices.accelerator.gke.io

  4. 슬라이스 커스텀 리소스가 삭제되었는지 확인합니다.

    kubectl get crd | grep slices.accelerator.gke.io

  5. 슬라이스 컨트롤러에서 추가한 라벨을 삭제합니다. 다음 라벨을 삭제해야 합니다.

    • cloud.google.com/gke-tpu-slice
    • cloud.google.com/gke-tpu-topology
    1. 특정 노드에서 삭제하려면 노드 이름을 업데이트합니다.
    export NODE_NAME="gke-tpu-bdac9600-3bdg"
kubectl label node $NODE_NAME cloud.google.com/gke-tpu-slice- cloud.google.com/gke-tpu-slice-topology-
    1. 클러스터의 모든 노드에서 이러한 라벨을 삭제하려면 다음 단계를 따르세요.
    kubectl label nodes --all cloud.google.com/gke-tpu-slice- cloud.google.com/gke-tpu-slice-topology-
    1. 노드 라벨을 확인하고 비어 있는지 확인합니다.
    export NODE_NAME="gke-tpu-bdac9600-3bdg"
kubectl describe node $NODE_NAME | grep "cloud.google.com/gke-tpu-slice"

다음 단계