Men-deploy workload TPU di GKE Autopilot

Halaman ini menjelaskan cara men-deploy workload yang menggunakan akselerator (TPU) Cloud TPU di cluster Autopilot Google Kubernetes Engine (GKE). Anda seharusnya sudah memahami konsep berikut:

  1. Pengantar Cloud TPU
  2. Arsitektur sistem Cloud TPU
  3. Tentang TPU di GKE

Cara kerja TPU di Autopilot

Untuk menggunakan TPU dalam workload Autopilot, Anda harus meminta versi TPU dan topologi yang didukung untuk versi TPU tersebut dalam manifes beban kerja Anda. Anda kemudian menggunakan kolom resources.requests dan resources.limits Kubernetes untuk menentukan jumlah chip TPU yang akan digunakan oleh workload. Saat Anda men-deploy beban kerja, GKE akan menyediakan node yang memiliki konfigurasi TPU yang diminta dan menjadwalkan Pod Anda pada node tersebut. GKE menempatkan setiap beban kerja pada node-nya sendiri, sehingga setiap Pod dapat mengakses seluruh resource node dengan risiko gangguan yang seminimal mungkin.

TPU di Autopilot kompatibel dengan kemampuan berikut:

  1. Pod Spot
  2. Reservasi kapasitas khusus
  3. Pod runtime yang diperpanjang

Merencanakan konfigurasi TPU

Sebelum meminta TPU, tentukan konfigurasi yang diinginkan berdasarkan persyaratan memori dan CPU workload. Anda harus menentukan hal berikut:

  • Versi TPU: versi Cloud TPU spesifik, seperti v5e.
  • Topologi untuk versi TPU yang dipilih: pengaturan dan jumlah TPU.

Versi TPU dan topologi yang Anda pilih menentukan apakah GKE menyediakan node sebagai slice host tunggal atau slice multi-host. Dalam potongan host tunggal, setiap node tidak bergantung pada node TPU lainnya. Dalam slice multi-host, GKE membuat grup node yang memiliki VM TPU yang saling terhubung. Slice multi-host bersifat atomik, yang berarti GKE menskalakan seluruh grup node yang saling terhubung sebagai satu unit.

Untuk mengetahui informasi tentang versi TPU yang tersedia, topologi terkait, kapasitas CPU dan memori, serta jenis slice yang dihasilkan, lihat Memilih konfigurasi TPU Autopilot.

Harga

Untuk informasi harga, lihat Harga Autopilot.

Sebelum memulai

Sebelum memulai, pastikan Anda telah menjalankan tugas berikut:

  • Aktifkan Google Kubernetes Engine API.
    • Aktifkan Google Kubernetes Engine API
  • Jika ingin menggunakan Google Cloud CLI untuk tugas ini, instal lalu initialize gcloud CLI. Jika sebelumnya Anda telah menginstal gcloud CLI, dapatkan versi terbaru dengan menjalankan gcloud components update.
  • Pastikan Anda memiliki cluster Autopilot yang menjalankan GKE versi 1.29.2-gke.1521000+ atau yang lebih baru.
  • Untuk menggunakan TPU yang dicadangkan, pastikan Anda sudah memiliki reservasi kapasitas tertentu. Untuk mengetahui petunjuknya, lihat Menggunakan resource zona yang dicadangkan.

Pastikan Anda memiliki kuota TPU

Untuk membuat TPU node, Anda harus memiliki kuota TPU yang tersedia kecuali menggunakan reservasi kapasitas yang sudah ada. Jika Anda menggunakan TPU yang dicadangkan, lewati bagian ini.

Pembuatan node TPU di GKE memerlukan kuota Compute Engine API (compute.googleapis.com), bukan kuota Cloud TPU API (tpu.googleapis.com). Nama kuota ini berbeda pada Pod Autopilot reguler dan di Pod Spot.

Untuk memeriksa batas dan penggunaan kuota Compute Engine API Anda saat ini untuk TPU, ikuti langkah-langkah berikut:

  1. Buka halaman Kuota di Konsol Google Cloud:

    Buka Quotas

  2. Di kotak Filter , lakukan hal berikut:

    1. Pilih properti Layanan, masukkan Compute Engine API, lalu tekan Enter.

    2. Pilih properti Jenis, lalu pilih Kuota.

    3. Pilih properti Name dan masukkan nama kuota berdasarkan jenis TPU yang Anda inginkan, sebagai berikut:

      • TPU v5p (Pratinjau): Chip TPU v5p
      • TPU v5e (tpu-v5-lite-podslice): Chip PodSlice TPU v5 Lite
      • TPU v5e (tpu-v5-lite-device): Chip Perangkat TPU v5 Lite
      • TPU v4 (tpu-v4-podslice): Chip PodSlice TPU v4

      Untuk Spot Pod, pilih kuota "Preemptible" yang sesuai.

    4. Pilih properti Dimension (e.g. locations) dan masukkan region: diikuti dengan nama region tempat Anda berencana membuat TPU di GKE. Misalnya, masukkan region:us-west4 jika Anda berencana untuk membuat node TPU di zona us-west4-a. Kuota TPU bersifat regional, sehingga semua zona dalam region yang sama menggunakan kuota TPU yang sama.

Jika tidak ada kuota yang cocok dengan filter yang Anda masukkan, berarti project belum diberi kuota yang ditentukan untuk region yang diinginkan, dan Anda harus meminta peningkatan kuota TPU.

Menyiapkan aplikasi TPU

Workload TPU memiliki persyaratan persiapan berikut.

  1. Framework seperti JAX, PyTorch, dan TensorFlow mengakses VM TPU menggunakan library bersama libtpu. libtpu mencakup compiler XLA, software runtime TPU, dan driver TPU. Setiap rilis PyTorch dan JAX memerlukan versi libtpu.so tertentu. Untuk menggunakan TPU di GKE, pastikan Anda menggunakan versi berikut:
    Jenis TPU Versi libtpu.so
    TPU v5e
    tpu-v5-lite-podslice
    tpu-v5-lite-device
    TPU v5p
    • Versi jax[tpu] yang direkomendasikan: 0.4.19 atau yang lebih baru.
    • Versi torchxla[tpuvm] yang direkomendasikan: disarankan untuk menggunakan build versi malam pada 23 Oktober 2023.
    TPU v4
    tpu-v4-podslice
  2. Tetapkan variabel lingkungan berikut untuk penampung yang meminta resource TPU:
    • TPU_WORKER_ID: Bilangan bulat unik untuk setiap Pod. ID ini menunjukkan worker-id unik di slice TPU. Nilai yang didukung untuk kolom ini berkisar dari nol hingga jumlah Pod dikurangi satu.
    • TPU_WORKER_HOSTNAMES: Daftar yang dipisahkan koma dari nama host VM TPU atau alamat IP yang perlu berkomunikasi satu sama lain dalam slice. Harus ada nama host atau alamat IP untuk setiap VM TPU dalam slice. Daftar alamat IP atau nama host diurutkan dan nol diindeks oleh TPU_WORKER_ID.

      • GKE otomatis memasukkan variabel lingkungan ini menggunakan webhook yang berubah saat Tugas dibuat dengan completionMode: Indexed, subdomain, parallelism > 1, dan meminta properti google.com/tpu. GKE menambahkan Layanan headless sehingga data DNS ditambahkan untuk Pod yang mendukung Layanan.

Setelah menyelesaikan persiapan beban kerja, Anda dapat menjalankan Tugas yang menggunakan TPU.

Meminta TPU dalam workload

Bagian ini menunjukkan cara membuat Tugas yang meminta TPU di Autopilot. Dalam workload apa pun yang memerlukan TPU, Anda harus menentukan hal berikut:

  • Pemilih node untuk versi TPU dan topologi
  • Jumlah chip TPU untuk container dalam workload Anda

Untuk daftar versi TPU yang didukung, topologi, serta jumlah chip TPU dan node yang sesuai dalam satu slice, lihat Memilih konfigurasi TPU Autopilot.

Pertimbangan untuk permintaan TPU dalam workload

Hanya satu container dalam Pod yang dapat menggunakan TPU. Jumlah chip yang diminta container harus sama dengan jumlah chip yang terpasang ke node dalam slice. Misalnya, jika Anda meminta TPU v5e (tpu-v5-lite-podslice) dengan topologi 2x4, Anda dapat meminta salah satu hal berikut:

  • 4 chip, yang membuat dua node multi-host dengan masing-masing 4 chip
  • 8 chip, yang membuat satu node host tunggal dengan 8 chip

Sebagai praktik terbaik untuk memaksimalkan efisiensi biaya, selalu gunakan semua chip dalam slice yang Anda minta. Jika meminta slice multi-host yang terdiri dari dua node yang terdiri dari 4 chip, Anda harus men-deploy beban kerja yang berjalan di kedua node dan menggunakan ke-8 chip dalam slice tersebut.

Membuat workload yang meminta TPU

Langkah-langkah berikut akan membuat Tugas yang meminta TPU. Jika memiliki workload yang berjalan di Slice TPU multi-host, Anda juga harus membuat Layanan headless yang memilih beban kerja berdasarkan nama. Service headless ini memungkinkan Pod pada node yang berbeda pada slice multi-host berkomunikasi satu sama lain dengan mengupdate konfigurasi DNS Kubernetes agar mengarah ke Pod dalam beban kerja.

  1. Simpan manifes berikut sebagai 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:
      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: python:3.10
        ports:
        - containerPort: 8471 # Default port using which TPU VMs communicate
        - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
        command:
        - bash
        - -c
        - |
          pip install 'jax[tpu]' -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
          python -c 'import jax; print("TPU cores:", jax.device_count())'
        resources:
          requests:
            cpu: 10
            memory: 500Gi
            google.com/tpu: NUMBER_OF_CHIPS
          limits:
            cpu: 10
            memory: 500Gi
            google.com/tpu: NUMBER_OF_CHIPS

    Ganti kode berikut:

    • TPU_TYPE: jenis TPU yang akan digunakan, seperti tpu-v4-podslice. Harus berupa nilai yang didukung oleh GKE.
    • TOPOLOGY: susunan chip TPU dalam irisan, seperti 2x2x4. Harus berupa topologi yang didukung untuk jenis TPU yang dipilih.
    • NUMBER_OF_CHIPS: jumlah chip TPU yang akan digunakan oleh penampung. Harus bernilai sama untuk limits dan requests.

  2. Deploy Job:

    kubectl create -f tpu-autopilot.yaml

Saat Anda membuat Tugas ini, GKE otomatis melakukan hal berikut:

  1. Menyediakan node untuk menjalankan Pod. Bergantung pada jenis TPU, topologi, dan permintaan resource yang Anda tentukan, node ini dapat berupa slice host tunggal atau slice multi-host.
  2. Menambahkan taint ke Pod dan toleransi ke node untuk mencegah workload lain berjalan di node yang sama dengan workload TPU.

Contoh: Menampilkan total chip TPU dalam slice multi-host

Workload berikut menampilkan jumlah chip TPU di semua node dalam slice TPU multi-host. Untuk membuat slice multi-host, beban kerja memiliki parameter berikut:

  • Versi TPU: TPU v4
  • Topologi: 2x2x4

Pemilihan versi dan topologi ini menghasilkan slice multi-host.

  1. Simpan manifes berikut sebagai 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
        cloud.google.com/gke-tpu-topology: 2x2x4
      containers:
      - name: tpu-job
        image: python:3.10
        ports:
        - containerPort: 8471 # Default port using which TPU VMs communicate
        - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
        command:
        - bash
        - -c
        - |
          pip install 'jax[tpu]' -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
          python -c 'import jax; print("TPU cores:", jax.device_count())'
        resources:
          requests:
            cpu: 10
            memory: 500Gi
            google.com/tpu: 4
          limits:
            cpu: 10
            memory: 500Gi
            google.com/tpu: 4
  2. Deploy manifes: 
    kubectl create -f available-chips-multihost.yaml

    GKE menjalankan slice TPU v4 dengan empat VM TPU (slice TPU multi-host). Slice ini memiliki 16 chip yang saling terhubung.

  3. Pastikan Tugas membuat empat Pod: 
    kubectl get pods

    Outputnya mirip dengan yang berikut ini:

    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. Dapatkan log salah satu Pod: 
    kubectl logs POD_NAME

    Ganti POD_NAME dengan nama salah satu Pod yang dibuat. Contoh, tpu-job-podslice-0-5cd8r.

    Outputnya mirip dengan hal berikut ini: 

    TPU cores: 16

Contoh: Menampilkan chip TPU dalam satu node

Beban kerja berikut adalah Pod statis yang menampilkan jumlah chip TPU yang terpasang ke node tertentu. Untuk membuat node host tunggal, beban kerja memiliki parameter berikut:

  • Versi TPU: TPU v5e
  • Topologi: 2x4

Pilihan versi dan topologi ini menghasilkan irisan host tunggal.

  1. Simpan manifes berikut sebagai 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
    cloud.google.com/gke-tpu-topology: 2x4
  containers:
  - name: tpu-job
    image: python:3.10
    ports:
    - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
    command:
    - bash
    - -c
    - |
      pip install 'jax[tpu]' -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
      python -c 'import jax; print("Total TPU chips:", jax.device_count())'
    resources:
      requests:
        google.com/tpu: 8
      limits:
        google.com/tpu: 8
  2. Deploy manifes: 
    kubectl create -f available-chips-singlehost.yaml

    GKE menyediakan node dengan delapan irisan TPU host tunggal yang menggunakan TPU v5e. Setiap VM TPU memiliki delapan chip (slice TPU host tunggal).

  3. Dapatkan log Pod: 
    kubectl logs tpu-job-jax-v5

    Outputnya mirip dengan hal berikut ini: 

    Total TPU chips: 8

Kemampuan observasi dan metrik

Dasbor

Di halaman Cluster Kubernetes di Google Cloud Console, tab Observability menampilkan metrik kemampuan observasi TPU. Untuk informasi selengkapnya, lihat metrik kemampuan observasi GKE.

Dasbor TPU hanya diisi jika Anda mengaktifkan metrik sistem di cluster GKE.

Metrik runtime

Pada GKE versi 1.27.4-gke.900 atau yang lebih baru, workload TPU yang menggunakan JAX versi 0.4.14 atau yang lebih baru dan menentukan containerPort: 8431 akan mengekspor metrik pemakaian TPU sebagai metrik sistem GKE. Metrik berikut tersedia di Cloud Monitoring untuk memantau performa runtime workload TPU Anda:

  • Siklus tugas: Persentase waktu dalam periode pengambilan sampel terakhir (60 detik) saat TensorCore secara aktif memproses TPU chip. Persentase yang lebih besar berarti pemakaian TPU yang lebih baik.
  • Memori yang digunakan: Jumlah memori akselerator yang dialokasikan dalam byte. Sampelnya dibuat setiap 60 detik.
  • Total memori: Total memori akselerator dalam byte. Dibuat sampelnya setiap 60 detik.

Metrik ini terletak di skema node Kubernetes (k8s_node) dan container Kubernetes (k8s_container).

Container Kubernetes:

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

Node Kubernetes:

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

Metrik host

Pada GKE versi 1.28.1-gke.1066000 atau yang lebih baru, VM TPU mengekspor metrik pemakaian TPU sebagai metrik sistem GKE. Metrik berikut tersedia di Cloud Monitoring untuk memantau performa host TPU Anda:

  • Pemakaian TensorCore: Persentase TensorCore yang saat ini dipakai. Nilai TensorCore sama dengan jumlah unit perkalian matriks (MXU) ditambah unit vektor. Nilai pemakaian TensorCore adalah pembagian dari operasi TensorCore yang dilakukan selama periode sampel terakhir (60 detik) dengan jumlah operasi TensorCore yang didukung selama periode yang sama. Nilai yang lebih besar berarti pemanfaatan yang lebih baik.
  • Pemakaian Bandwidth Memori: Persentase dari bandwidth memori akselerator yang saat ini sedang digunakan. Dihitung dengan membagi bandwidth memori yang digunakan selama periode sampel (60 dtk) dengan bandwidth maksimum yang didukung selama periode sampel yang sama.

Metrik ini terletak di skema node Kubernetes (k8s_node) dan container Kubernetes (k8s_container).

Container Kubernetes:

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

Node Kubernetes:

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

Untuk mengetahui informasi selengkapnya, lihat metrik Kubernetes dan metrik sistem GKE.

Logging

Log yang dikeluarkan oleh container yang berjalan di node GKE, termasuk VM TPU, dikumpulkan oleh agen logging GKE, dikirim ke Logging, dan terlihat di Logging.

Rekomendasi untuk workload TPU dalam Autopilot

Rekomendasi berikut dapat meningkatkan efisiensi workload TPU Anda:

  • Gunakan Pod runtime yang diperpanjang selama masa tenggang hingga tujuh hari sebelum GKE menghentikan Pod Anda untuk penurunan skala atau upgrade node. Anda dapat menggunakan masa pemeliharaan dan pengecualian dengan Pod runtime yang diperpanjang untuk menunda upgrade node otomatis lebih lanjut.
  • Gunakan pemesanan kapasitas untuk memastikan bahwa workload Anda menerima TPU yang diminta tanpa ditempatkan dalam antrean ketersediaan.