Menjalankan pipeline dengan TPU

Halaman ini menjelaskan cara menjalankan pipeline Apache Beam di Dataflow dengan TPU. Tugas yang menggunakan TPU dikenai biaya seperti yang ditentukan di halaman harga Dataflow.

Untuk mengetahui informasi selengkapnya tentang penggunaan TPU dengan Dataflow, lihat Dukungan Dataflow untuk TPU.

Opsional: Membuat reservasi tertentu untuk menggunakan akselerator

Meskipun Anda dapat menggunakan TPU on-demand, sebaiknya gunakan TPU Dataflow dengan target khusus Google Cloud pemesanan. Hal ini membantu memastikan Anda memiliki akses ke akselerator yang tersedia dan waktu mulai pekerja yang cepat. Pipeline yang menggunakan pemesanan TPU tidak memerlukan kuota TPU tambahan.

Jika Anda tidak membuat reservasi dan memilih untuk menggunakan TPU sesuai permintaan, sediakan kuota TPU sebelum menjalankan pipeline.

Opsional: Sediakan kuota TPU

Anda dapat menggunakan TPU dalam kapasitas sesuai permintaan atau menggunakan reservasi. Jika ingin menggunakan TPU sesuai permintaan, Anda harus menyediakan kuota TPU sebelum melakukannya. Jika menggunakan pemesanan yang ditargetkan secara khusus, Anda dapat melewati bagian ini.

Untuk menggunakan TPU on-demand tanpa reservasi, periksa batas dan penggunaan saat ini kuota API Compute Engine untuk TPU sebagai berikut:

Konsol

  1. Buka halaman Quotas di konsol Google Cloud :

    Buka Kuota

  2. Di kotak Filter , lakukan tindakan berikut:

    1. Gunakan tabel berikut untuk memilih dan menyalin properti kuota berdasarkan versi TPU dan jenis mesin. Misalnya, jika Anda berencana membuat node TPU v5e sesuai permintaan yang jenis mesinnya diawali dengan ct5lp-, masukkan Name: TPU v5 Lite PodSlice chips.

      Versi TPU, jenis mesin dimulai dengan Properti dan nama kuota untuk instance on-demand
      TPU v5e,
      ct5lp- 
      Name:
      TPU v5 Lite PodSlice chips
      TPU v5p,
      ct5p- 
      Name:
      TPU v5p chips
      TPU v6e,
      ct6e- 
      Dimensions (e.g. location):
      tpu_family:CT6E

    2. Pilih properti Dimensi (mis. lokasi) dan masukkan region:, diikuti dengan nama region tempat Anda berencana memulai pipeline. Misalnya, masukkan region:us-west4 jika Anda berencana menggunakan zona us-west4-a. Kuota TPU bersifat regional, sehingga semua zona dalam region yang sama menggunakan kuota TPU yang sama.

Mengonfigurasi image container kustom

Untuk berinteraksi dengan TPU di pipeline Dataflow, Anda harus menyediakan software yang dapat beroperasi pada perangkat XLA di lingkungan runtime pipeline Anda. Hal ini memerlukan penginstalan library TPU berdasarkan kebutuhan pipeline Anda dan mengonfigurasi variabel lingkungan berdasarkan perangkat TPU yang Anda gunakan.

Untuk menyesuaikan image container, instal Apache Beam ke dalam image dasar siap pakai yang memiliki library TPU yang diperlukan. Atau, instal software TPU ke dalam image yang dipublikasikan dengan rilis Apache Beam SDK.

Untuk menyediakan image container kustom, gunakan opsi pipeline sdk_container_image. Untuk mengetahui informasi selengkapnya, lihat Menggunakan container kustom di Dataflow.

Saat menggunakan akselerator TPU, Anda perlu menetapkan variabel lingkungan berikut di image container.

ENV TPU_SKIP_MDS_QUERY=1 # Don't query metadata
ENV TPU_HOST_BOUNDS=1,1,1 # There's only one host
ENV TPU_WORKER_HOSTNAMES=localhost
ENV TPU_WORKER_ID=0 # Always 0 for single-host TPUs

Bergantung pada akselerator yang Anda gunakan, variabel dalam tabel berikut juga perlu ditetapkan.

jenis topologi worker_machine_type Dataflow yang diperlukan variabel lingkungan tambahan
tpu-v5-lite-podslice 1x1 ct5lp-hightpu-1t 
TPU_ACCELERATOR_TYPE=v5litepod-1
TPU_CHIPS_PER_HOST_BOUNDS=1,1,1
tpu-v5-lite-podslice 2x2 ct5lp-hightpu-4t 
TPU_ACCELERATOR_TYPE=v5litepod-4
TPU_CHIPS_PER_HOST_BOUNDS=2,2,1
tpu-v5-lite-podslice 2x4 ct5lp-hightpu-8t 
TPU_ACCELERATOR_TYPE=v5litepod-8
TPU_CHIPS_PER_HOST_BOUNDS=2,4,1
tpu-v6e-slice 1x1 ct6e-standard-1t 
TPU_ACCELERATOR_TYPE=v6e-1
TPU_CHIPS_PER_HOST_BOUNDS=1,1,1
tpu-v6e-slice 2x2 ct6e-standard-4t 
TPU_ACCELERATOR_TYPE=v6e-4
TPU_CHIPS_PER_HOST_BOUNDS=2,2,1
tpu-v6e-slice 2x4 ct6e-standard-8t 
TPU_ACCELERATOR_TYPE=v6e-8
TPU_CHIPS_PER_HOST_BOUNDS=2,4,1
tpu-v5p-slice 2x2x1 ct5p-hightpu-4t 
TPU_ACCELERATOR_TYPE=v5p-8
TPU_CHIPS_PER_HOST_BOUNDS=2,2,1

Contoh Dockerfile untuk image container kustom dapat terlihat seperti contoh berikut:

FROM python:3.11-slim

COPY --from=apache/beam_python3.11_sdk:2.66.0 /opt/apache/beam /opt/apache/beam

# Configure the environment to access TPU device

ENV TPU_SKIP_MDS_QUERY=1
ENV TPU_HOST_BOUNDS=1,1,1
ENV TPU_WORKER_HOSTNAMES=localhost
ENV TPU_WORKER_ID=0

# Configure the environment for the chosen accelerator.
# Adjust according to the accelerator you use.
ENV TPU_ACCELERATOR_TYPE=v5litepod-1
ENV TPU_CHIPS_PER_HOST_BOUNDS=1,1,1

# Install TPU software stack.
RUN pip install jax[tpu] apache-beam[gcp]==2.66.0 -f https://storage.googleapis.com/jax-releases/libtpu_releases.html

ENTRYPOINT ["/opt/apache/beam/boot"]

Menjalankan tugas dengan TPU

Pertimbangan untuk menjalankan tugas Dataflow dengan TPU mencakup hal-hal berikut:

  • Karena container TPU bisa berukuran besar, untuk menghindari kehabisan ruang disk, tingkatkan ukuran disk booting default menjadi 50 gigabyte atau ukuran yang sesuai sebagaimana diperlukan oleh image container Anda menggunakan --disk_size_gb opsi pipeline.
  • Membatasi paralelisme dalam pekerja.

TPU dan paralelisme pekerja

Dalam konfigurasi default, pipeline Python Dataflow meluncurkan satu proses Apache Beam SDK per core VM. Jenis mesin TPU memiliki sejumlah besar core vCPU, tetapi hanya satu proses yang dapat melakukan komputasi pada perangkat TPU. Selain itu, perangkat TPU mungkin dipesan oleh suatu proses selama masa aktif proses tersebut. Oleh karena itu, Anda harus membatasi paralelisme dalam worker saat menjalankan pipeline TPU Dataflow. Untuk membatasi paralelisme pekerja, gunakan panduan berikut:

  • Jika kasus penggunaan Anda melibatkan menjalankan inferensi pada model, gunakan Beam RunInference API. Untuk mengetahui informasi selengkapnya, lihat Inferensi Model Bahasa Besar di Beam.
  • Jika Anda tidak dapat menggunakan API RunInference Beam, gunakan objek bersama multi-proses Beam untuk membatasi operasi tertentu ke satu proses.
  • Jika Anda tidak dapat menggunakan rekomendasi sebelumnya dan lebih memilih untuk meluncurkan hanya satu proses Python per pekerja, tetapkan opsi pipeline --experiments=no_use_multiple_sdk_containers.
  • Untuk pekerja dengan lebih dari 100 vCPU, kurangi jumlah thread dengan menggunakan opsi pipeline --number_of_worker_harness_threads. Gunakan tabel berikut untuk melihat apakah jenis TPU Anda menggunakan lebih dari 100 vCPU.

Tabel berikut mencantumkan total resource komputasi per pekerja untuk setiap konfigurasi TPU.

Jenis TPU topologi jenis mesin Chip TPU vCPU RAM (GB)
tpu-v5-lite-podslice 1x1 ct5lp-hightpu-1t 1 24 48
tpu-v5-lite-podslice 2x2 ct5lp-hightpu-4t 4 112 192
tpu-v5-lite-podslice 2x4 ct5lp-hightpu-8t 8 224 384
tpu-v6e-slice 1x1 ct6e-standard-1t 1 44 176
tpu-v6e-slice 2x2 ct6e-standard-4t 4 180 720
tpu-v6e-slice 2x4 ct6e-standard-8t 8 360 1440
tpu-v5p-slice 2x2x1 ct5p-hightpu-4t 4 208 448

Menjalankan pipeline dengan TPU

Untuk menjalankan tugas Dataflow dengan TPU, gunakan perintah berikut.

python PIPELINE \
  --runner "DataflowRunner" \
  --project "PROJECT" \
  --temp_location "gs://BUCKET/tmp" \
  --region "REGION" \
  --dataflow_service_options "worker_accelerator=type:TPU_TYPE;topology:TPU_TOPOLOGY" \
  --worker_machine_type "MACHINE_TYPE" \
  --disk_size_gb "DISK_SIZE_GB" \
  --sdk_container_image "IMAGE" \
  --number_of_worker_harness_threads NUMBER_OF_THREADS

Ganti kode berikut:

  • PIPELINE: File kode sumber pipeline Anda.
  • PROJECT: Nama project. Google Cloud
  • BUCKET: Bucket Cloud Storage.
  • REGION: Region Dataflow, misalnya, us-central1.
  • TPU_TYPE: Jenis TPU yang didukung, misalnya, tpu-v5-lite-podslice. Untuk mengetahui daftar lengkap jenis dan topologi, lihat Akselerator TPU yang didukung.
  • TPU_TOPOLOGY: Topologi TPU, misalnya, 1x1.
  • MACHINE_TYPE: Jenis mesin yang sesuai, misalnya, ct5lp-hightpu-1t.
  • DISK_SIZE_GB: Ukuran boot disk untuk setiap VM pekerja, misalnya, 100.
  • IMAGE: Jalur Artifact Registry untuk image Docker Anda.
  • NUMBER_OF_THREADS: Opsional. Jumlah thread harness pekerja.

Verifikasi tugas Dataflow Anda

Untuk mengonfirmasi bahwa tugas menggunakan VM pekerja dengan TPU, ikuti langkah-langkah berikut:

  1. Di konsol Google Cloud , buka halaman Dataflow > Jobs.

    Buka Tugas

  2. Pilih lowongan.

  3. Klik tab Metrik tugas.

  4. Di bagian Autoscaling, pastikan ada minimal satu VM Current workers.

  5. Di panel samping Info tugas, periksa untuk memastikan machine_type dimulai dengan ct. Contohnya, ct6e-standard-1t Hal ini menunjukkan penggunaan TPU.

Memecahkan masalah tugas Dataflow

Jika Anda mengalami masalah saat menjalankan tugas Dataflow dengan TPU, lihat Memecahkan masalah tugas TPU Dataflow.

Langkah berikutnya