Menjalankan aplikasi Cloud TPU di container Docker

Container Docker mempermudah konfigurasi aplikasi dengan menggabungkan kode Anda dan semua dependensi yang diperlukan dalam satu paket yang dapat didistribusikan. Anda dapat menjalankan container Docker dalam VM TPU untuk menyederhanakan konfigurasi dan berbagi aplikasi Cloud TPU. Dokumen ini menjelaskan cara menyiapkan container Docker untuk setiap framework ML yang didukung oleh Cloud TPU.

Penting: Anda dapat menjalankan maksimum 256 pod Docker di VM TPU.

Melatih model TensorFlow di container Docker

Perangkat TPU

1. Buat file bernama Dockerfile di direktori Anda saat ini dan tempel teks berikut

   FROM python:3.8
   RUN pip install https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/tensorflow/tf-2.12.0/tensorflow-2.12.0-cp38-cp38-linux_x86_64.whl
   RUN curl -L https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/libtpu/1.6.0/libtpu.so -o /lib/libtpu.so
   RUN git clone https://github.com/tensorflow/models.git
   WORKDIR ./models
   RUN pip install -r official/requirements.txt
   ENV PYTHONPATH=/models

2. Membuat bucket Cloud Storage

   gsutil mb -c standard -l europe-west4 gs://your-bucket-name

3. Membuat VM TPU

   gcloud compute tpus tpu-vm create your-tpu-name \
   --zone=europe-west4-a \
   --accelerator-type=v2-8 \
   --version=tpu-vm-tf-2.16.1-pjrt

4. Menyalin Dockerfile ke VM TPU Anda

   gcloud compute tpus tpu-vm scp ./Dockerfile your-tpu-name:

5. SSH ke VM TPU

   gcloud compute tpus tpu-vm ssh your-tpu-name \
   --zone=europe-west4-a

6. Membangun gambar Docker

   sudo docker build -t your-image-name .

7. Memulai container Docker

   sudo docker run -ti --rm --net=host --name your-container-name --privileged your-image-name bash

8. Menetapkan variabel lingkungan

   export STORAGE_BUCKET=gs://your-bucket-name
   export DATA_DIR=gs://cloud-tpu-test-datasets/fake_imagenet
   export MODEL_DIR=${STORAGE_BUCKET}/resnet-2x

9. Melatih ResNet

   python3 official/vision/train.py \
   --tpu=local \
   --experiment=resnet_imagenet \
   --mode=train_and_eval \
   --config_file=official/vision/configs/experiments/image_classification/imagenet_resnet50_tpu.yaml \
   --model_dir=${MODEL_DIR} \
   --params_override="task.train_data.input_path=${DATA_DIR}/train*, task.validation_data.input_path=${DATA_DIR}/validation*,trainer.train_steps=100"

Setelah skrip pelatihan selesai, pastikan Anda membersihkan resource-nya.

1. Ketik exit untuk keluar dari container Docker
2. Ketik exit untuk keluar dari VM TPU
3. Menghapus VM TPU

    $ gcloud compute tpus tpu-vm delete your-tpu-name --zone=europe-west4-a

Pod TPU

1. Buat file bernama Dockerfile di direktori Anda saat ini dan tempel teks berikut

   FROM python:3.8
   RUN pip install https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/tensorflow/tf-2.12.0/tensorflow-2.12.0-cp38-cp38-linux_x86_64.whl
   RUN curl -L https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/libtpu/1.6.0/libtpu.so -o /lib/libtpu.so
   RUN git clone https://github.com/tensorflow/models.git
   WORKDIR ./models
   RUN pip install -r official/requirements.txt
   ENV PYTHONPATH=/models

2. Membuat VM TPU

   gcloud compute tpus tpu-vm create your-tpu-name \
   --zone=europe-west4-a \
   --accelerator-type=v3-32 \
   --version=tpu-vm-tf-2.16.1-pod-pjrt

3. Menyalin Dockerfile ke VM TPU Anda

   gcloud compute tpus tpu-vm scp ./Dockerfile your-tpu-name:

4. SSH ke VM TPU

   gcloud compute tpus tpu-vm ssh your-tpu-name \
   --zone=europe-west4-a

5. Membangun gambar Docker

   sudo docker build -t your-image-name .

6. Memulai container Docker

   sudo docker run -ti --rm --net=host --name your-container-name --privileged your-image-name bash

7. Melatih ResNet

   python3 official/vision/train.py \
   --tpu=local \
   --experiment=resnet_imagenet \
   --mode=train_and_eval \
   --config_file=official/vision/configs/experiments/image_classification/imagenet_resnet50_tpu.yaml \
   --model_dir=${MODEL_DIR} \
   --params_override="task.train_data.input_path=${DATA_DIR}/train*, task.validation_data.input_path=${DATA_DIR}/validation*,task.train_data.global_batch_size=2048,task.validation_data.global_batch_size=2048,trainer.train_steps=100"

Setelah skrip pelatihan selesai, pastikan Anda membersihkan resource-nya.

1. Ketik exit untuk keluar dari container Docker
2. Ketik exit untuk keluar dari VM TPU
3. Menghapus VM TPU

     $ gcloud compute tpus tpu-vm delete your-tpu-name --zone=europe-west4-a

Melatih model PyTorch di container Docker

Perangkat TPU

1. Membuat VM Cloud TPU

   gcloud compute tpus tpu-vm create your-tpu-name \
   --zone=europe-west4-a \
   --accelerator-type=v2-8 \
   --version=tpu-ubuntu2204-base

2. SSH ke VM TPU

   gcloud compute tpus tpu-vm ssh your-tpu-name \
   --zone=europe-west4-a

3. Mulai container di VM TPU menggunakan image PyTorch/XLA yang berfungsi malam hari.

   sudo docker run -ti --rm --name your-container-name --privileged gcr.io/tpu-pytorch/xla:r2.0_3.8_tpuvm bash

   Catatan: Setelah menjalankan perintah ini, command prompt Anda akan berubah yang menunjukkan terminal terhubung ke container yang sedang berjalan.

4. Mengonfigurasi runtime TPU

   Ada dua opsi runtime PyTorch/XLA: PJRT dan XRT. Sebaiknya gunakan PJRT kecuali jika Anda memiliki alasan untuk menggunakan XRT. Untuk mempelajari lebih lanjut konfigurasi runtime yang berbeda, lihat alasan Anda menggunakan XRT. Untuk mempelajari lebih lanjut berbagai konfigurasi runtime, lihat dokumentasi runtime PJRT.

   musik

   export PJRT_DEVICE=TPU

   XRT

   export XRT_TPU_CONFIG="localservice;0;localhost:51011"

5. Meng-clone repo PyTorch XLA

   git clone --recursive https://github.com/pytorch/xla.git

6. Melatih ResNet50

   python3 xla/test/test_train_mp_imagenet.py --fake_data --model=resnet50 --num_epochs=1

Setelah skrip pelatihan selesai, pastikan Anda membersihkan resource-nya.

1. Ketik exit untuk keluar dari container Docker
2. Ketik exit untuk keluar dari VM TPU
3. Menghapus VM TPU

    $ gcloud compute tpus tpu-vm delete your-tpu-name --zone=europe-west4-a

Pod TPU

Saat menjalankan kode PyTorch di Pod TPU, Anda harus menjalankan kode di semua pekerja TPU secara bersamaan. Salah satu cara untuk melakukannya adalah dengan menggunakan perintah gcloud compute tpus tpu-vm ssh dengan flag --worker=all dan --command. Prosedur berikut menunjukkan cara membuat image Docker untuk mempermudah penyiapan setiap pekerja TPU.

1. Membuat VM TPU

   gcloud compute tpus tpu-vm create your-tpu-name \
   --zone=us-central2-b \
   --accelerator-type=v4-32 \
   --version=tpu-ubuntu2204-base

2. Menambahkan pengguna saat ini ke grup Docker

   gcloud compute tpus tpu-vm ssh your-tpu-name \
   --zone=us-central2-b \
   --worker=all \
   --command="sudo usermod -a -G docker $USER"

3. Menjalankan skrip pelatihan dalam container pada semua pekerja TPU.

   gcloud compute tpus tpu-vm ssh your-tpu-name --worker=all \
   --zone=us-central2-b \
   --command="docker run --rm --privileged --net=host  -e PJRT_DEVICE=TPU gcr.io/tpu-pytorch/xla:r2.0_3.8_tpuvm python /pytorch/xla/test/test_train_mp_imagenet.py --fake_data --model=resnet50 --num_epochs=1"

   Flag perintah Docker:

   • --rm menghapus penampung setelah prosesnya dihentikan.
   • --privileged mengekspos perangkat TPU ke container.
   • --net=host mengikat semua port container ke VM TPU untuk memungkinkan komunikasi antara host di Pod.
   • -e menetapkan variabel lingkungan.

Setelah skrip pelatihan selesai, pastikan Anda membersihkan resource-nya.

Hapus VM TPU menggunakan perintah berikut:

$ gcloud compute tpus tpu-vm delete your-tpu-name \
  --zone=us-central2-b

Melatih model JAX di container Docker

Perangkat TPU

1. Membuat VM TPU

   gcloud compute tpus tpu-vm create your-tpu-name \
   --zone=europe-west4-a \
   --accelerator-type=v2-8 \
   --version=tpu-ubuntu2204-base

2. SSH ke VM TPU

   gcloud compute tpus tpu-vm ssh your-tpu-name  --zone=europe-west4-a

3. Memulai daemon Docker di VM TPU

   sudo systemctl start docker

4. Memulai container Docker

   sudo docker run -ti --rm --name your-container-name --privileged --network=host python:3.8 bash

5. Instal JAX

   pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html

6. Instal FLAX

   pip install --upgrade clu
   git clone https://github.com/google/flax.git
   pip install --user -e flax

7. Menjalankan skrip pelatihan FLAX MNIST

   cd flax/examples/mnist
   python3 main.py --workdir=/tmp/mnist \
   --config=configs/default.py \
   --config.learning_rate=0.05 \
   --config.num_epochs=5

Setelah skrip pelatihan selesai, pastikan Anda membersihkan resource-nya.

1. Ketik exit untuk keluar dari container Docker
2. Ketik exit untuk keluar dari VM TPU

3. Menghapus VM TPU

   $ gcloud compute tpus tpu-vm delete your-tpu-name --zone=europe-west4-a

Pod TPU

Saat menjalankan kode JAX di Pod TPU, Anda harus menjalankan kode JAX di semua pekerja TPU secara bersamaan. Salah satu cara untuk melakukannya adalah dengan menggunakan perintah gcloud compute tpus tpu-vm ssh dengan tanda --worker=all dan --command. Prosedur berikut menunjukkan cara membuat image Docker untuk mempermudah penyiapan setiap pekerja TPU.

1. Buat file bernama Dockerfile di direktori Anda saat ini dan tempel teks berikut

   FROM python:3.8
   RUN pip install "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
   RUN pip install --upgrade clu
   RUN git clone https://github.com/google/flax.git
   RUN pip install --user -e flax
   WORKDIR ./flax/examples/mnist

2. Membangun gambar Docker

   docker build -t your-image-name .

3. Tambahkan tag ke image Docker Anda sebelum mengirimkannya ke Artifact Registry. Untuk mengetahui informasi selengkapnya tentang cara menggunakan Artifact Registry, lihat Bekerja dengan image container.

   docker tag your-image-name europe-west-docker.pkg.dev/your-project/your-repo/your-image-name:your-tag

4. Mengirim image Docker ke Artifact Registry

   docker push europe-west4-docker.pkg.dev/your-project/your-repo/your-image-name:your-tag

5. Membuat VM TPU

   gcloud compute tpus tpu-vm create your-tpu-name \
   --zone=europe-west4-a \
   --accelerator-type==v2-8 \
   --version=tpu-ubuntu2204-base

6. Ambil image Docker dari Artifact Registry pada semua pekerja TPU.

   gcloud compute tpus tpu-vm ssh your-tpu-name --worker=all \
   --zone=europe-west4-a \
   --command="sudo usermod -a -G docker ${USER}"

   gcloud compute tpus tpu-vm ssh your-tpu-name --worker=all \
   --zone=europe-west4-a \
   --command="gcloud auth configure-docker europe-west4-docker.pkg.dev --quiet"

   gcloud compute tpus tpu-vm ssh your-tpu-name --worker=all \
   --zone=europe-west4-a \
   --command="docker pull europe-west4-docker.pkg.dev/your-project/your-repo/your-image-name:your-tag"

7. Jalankan container di semua pekerja TPU.

   gcloud compute tpus tpu-vm ssh your-tpu-name --worker=all \
   zone=europe-west4-a \
   --command="docker run -ti -d --privileged --net=host --name your-container-name europe-west4-docker.pkg.dev/your-project/your-repo/your-image:your-tag bash"

8. Jalankan skrip pelatihan pada semua pekerja TPU:

   gcloud compute tpus tpu-vm ssh your-tpu-name --worker=all \
   --zone=europe-west4-a \
   --command="docker exec --privileged your-container-name python3 main.py --workdir=/tmp/mnist \
   --config=configs/default.py \
   --config.learning_rate=0.05 \
   --config.num_epochs=5"

Setelah skrip pelatihan selesai, pastikan Anda membersihkan resource-nya.

1. Matikan container bagi semua pekerja:

   gcloud compute tpus tpu-vm ssh your-tpu-name --worker=all \
   --zone=europe-west4-a \
   --command="docker kill your-container-name"

2. Hapus VM TPU menggunakan perintah berikut:

   $ gcloud compute tpus tpu-vm delete your-tpu-name \
   --zone=europe-west4-a

Langkah selanjutnya

• Tutorial Cloud TPU
• Mengelola TPU
• Arsitektur Sistem Cloud TPU
• Menjalankan kode TensorFlow pada slice Pod TPU
• Menjalankan kode JAX pada slice Pod TPU