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 pemesanan tertentu untuk menggunakan akselerator

Meskipun Anda dapat menggunakan TPU sesuai permintaan, sebaiknya gunakan TPU Dataflow dengan pemesanan yang Google Cloudditargetkan secara khusus. Hal ini membantu memastikan Anda memiliki akses ke akselerator yang tersedia dan waktu startup worker yang cepat. Pipeline yang menggunakan pemesanan TPU tidak memerlukan kuota TPU tambahan.

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

Opsional: Menyediakan kuota TPU

Anda dapat menggunakan TPU dalam kapasitas sesuai permintaan atau menggunakan pemesanan. 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 sesuai permintaan tanpa pemesanan, periksa batas dan penggunaan kuota Compute Engine API saat ini untuk TPU sebagai berikut:

Konsol

  1. Buka halaman Kuota di Google Cloud konsol:

    Buka Kuota

  2. Di kotak Filter , lakukan hal 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 diawali dengan Properti dan nama kuota untuk instance sesuai permintaan
      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 Dimensions (e.g. location) 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 di perangkat XLA di lingkungan runtime pipeline. Hal ini mengharuskan Anda menginstal library TPU berdasarkan kebutuhan pipeline 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 harus menetapkan variabel lingkungan berikut dalam 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 harus 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 mungkin 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 berikut:

  • Karena container TPU bisa berukuran besar, untuk menghindari kehabisan ruang disk, tingkatkan ukuran boot disk default menjadi 50 gigabyte atau ukuran yang sesuai seperti yang diperlukan oleh image container Anda menggunakan opsi pipeline --disk_size_gb.
  • Batasi paralelisme intra-worker.

TPU dan paralelisme worker

Dalam konfigurasi default, pipeline Python Dataflow meluncurkan satu proses Apache Beam SDK per VM core. Jenis mesin TPU memiliki sejumlah besar core vCPU, tetapi hanya satu proses yang dapat melakukan komputasi pada perangkat TPU. Selain itu, perangkat TPU mungkin dicadangkan oleh proses selama masa aktif proses. Oleh karena itu, Anda harus membatasi paralelisme intra-worker saat menjalankan pipeline TPU Dataflow. Untuk membatasi paralelisme worker, 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 Beam RunInference API, 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 worker, tetapkan opsi pipeline --experiments=no_use_multiple_sdk_containers.
  • Anda dapat mengurangi lebih lanjut jumlah thread menggunakan opsi pipeline --number_of_worker_harness_threads jika hal tersebut mencapai performa yang lebih baik.

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

Jenis TPU topologi jenis mesin TPU chip 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 Google Cloud project.
  • 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 worker, misalnya, 100.
  • IMAGE: Jalur Artifact Registry untuk image Docker Anda.
  • NUMBER_OF_THREADS: Opsional. Jumlah thread harness worker.

Memverifikasi tugas Dataflow

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

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

    Buka Tugas

  2. Pilih tugas.

  3. Klik tab Metrik tugas.

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

  5. Di panel samping Info tugas, periksa apakah machine_type dimulai dengan ct. Misalnya, 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 Dataflow TPU job.

Langkah berikutnya