Esegui applicazioni Cloud TPU in un container Docker

I container Docker semplificano la configurazione delle applicazioni combinando il codice e tutte le dipendenze necessarie in un unico pacchetto distribuibile. Puoi eseguire container Docker all'interno delle VM TPU per semplificare la configurazione e la condivisione delle applicazioni Cloud TPU. Questo documento descrive come configurare un container Docker per ogni framework ML supportato da Cloud TPU.

Importante: puoi eseguire un massimo di 256 pod Docker su una VM TPU.

Addestra un modello TensorFlow in un container Docker

Dispositivo TPU

1. Crea un file denominato Dockerfile nella directory attuale e incolla il seguente testo

   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. Crea bucket Cloud Storage

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

3. Crea una 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. Copia il Dockerfile nella VM TPU

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

5. SSH nella VM TPU

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

6. Crea l'immagine Docker

   sudo docker build -t your-image-name .
   

7. Avvia il container Docker

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

8. Imposta le variabili di ambiente

   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. Addestra 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"
   

Al termine dello script di addestramento, assicurati di eseguire la pulizia delle risorse.

1. Digita exit per uscire dal container Docker
2. Digita exit per uscire dalla VM TPU
3. Elimina la VM TPU

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

pod di TPU

1. Crea un file denominato Dockerfile nella directory attuale e incolla il seguente testo

   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. Crea una 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. Copia il Dockerfile nella VM TPU

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

4. SSH nella VM TPU

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

5. Crea l'immagine Docker

   sudo docker build -t your-image-name .
   

6. Avvia un container Docker

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

7. Addestra 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"
   

Al termine dello script di addestramento, assicurati di eseguire la pulizia delle risorse.

1. Digita exit per uscire dal container Docker
2. Digita exit per uscire dalla VM TPU
3. Elimina la VM TPU

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

Addestra un modello PyTorch in un container Docker

Dispositivo TPU

1. Crea 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 nella VM TPU

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

3. Avvia un container nella VM TPU utilizzando l'immagine PyTorch/XLA di notte.

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

   Nota: dopo aver eseguito questo comando, il prompt dei comandi cambierà per indicare che il terminale è connesso al container in esecuzione.

4. Configura runtime TPU

   Esistono due opzioni di runtime PyTorch/XLA: PJRT e XRT. Ti consigliamo di usare PJRT, a meno che tu non abbia un motivo per usare XRT. Per saperne di più sulle diverse configurazioni di runtime, vedi che hai un motivo per utilizzare XRT. Per saperne di più sulle diverse configurazioni di runtime, consulta la documentazione sul runtime PJRT.

   PJRT

   export PJRT_DEVICE=TPU
   

   XRT

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

5. Clona il repository XLA di PyTorch

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

6. Addestra ResNet50

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

Al termine dello script di addestramento, assicurati di eseguire la pulizia delle risorse.

1. Digita exit per uscire dal container Docker
2. Digita exit per uscire dalla VM TPU
3. Elimina la VM TPU

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

pod di TPU

Quando esegui il codice PyTorch su un pod di TPU, devi eseguire il codice su tutti i worker TPU contemporaneamente. Un modo per farlo è utilizzare il comando gcloud compute tpus tpu-vm ssh con i flag --worker=all e --command. La procedura seguente mostra come creare un'immagine Docker per semplificare la configurazione di ogni worker TPU.

1. Crea una VM TPU

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

2. Aggiungi l'utente corrente al gruppo Docker

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

3. Esegui lo script di addestramento in un container su tutti i worker 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 dei comandi Docker:

   • --rm rimuove il container al termine del processo.
   • --privileged espone il dispositivo TPU al container.
   • --net=host associa tutte le porte del container alla VM TPU per consentire la comunicazione tra gli host nel pod.
   • -e ha impostato le variabili di ambiente.

Al termine dello script di addestramento, assicurati di eseguire la pulizia delle risorse.

Elimina la VM TPU utilizzando il comando seguente:

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

Addestra un modello JAX in un container Docker

Dispositivo TPU

1. Crea la 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 nella VM TPU

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

3. Avvia daemon Docker nella VM TPU

   sudo systemctl start docker
   

4. Avvia container Docker

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

5. Installa JAX

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

6. Installa FLAX

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

7. Esegui lo script di addestramento 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
   

Al termine dello script di addestramento, assicurati di eseguire la pulizia delle risorse.

1. Digita exit per uscire dal container Docker
2. Digita exit per uscire dalla VM TPU

3. Elimina la VM TPU

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

pod di TPU

Quando esegui codice JAX su un pod di TPU, devi eseguire il codice JAX su tutti i worker TPU contemporaneamente. Un modo per farlo è usare il comando gcloud compute tpus tpu-vm ssh con i flag --worker=all e --command. La procedura riportata di seguito mostra come creare un'immagine Docker per semplificare la configurazione di ogni worker TPU.

1. Crea un file denominato Dockerfile nella directory attuale e incolla il seguente testo

   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. Crea l'immagine Docker

   docker build -t your-image-name .
   

3. Aggiungi un tag all'immagine Docker prima di eseguirne il push ad Artifact Registry. Per saperne di più sull'utilizzo di Artifact Registry, consulta Utilizzare le immagini container.

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

4. Esegui il push dell'immagine Docker ad Artifact Registry

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

5. Crea una VM TPU

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

6. Esegui il pull dell'immagine Docker da Artifact Registry su tutti i worker 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. Esegui il container su tutti i worker 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. Esegui lo script di addestramento su tutti i worker 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"
   

Al termine dello script di addestramento, assicurati di eseguire la pulizia delle risorse.

1. Arresta il container su tutti i worker:

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

2. Elimina la VM TPU utilizzando il comando seguente:

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

Passaggi successivi

• Tutorial su Cloud TPU
• Gestione delle TPU
• Architettura di sistema Cloud TPU
• Esegui il codice TensorFlow sulle pod di TPU TPU
• Esegui codice JAX sulle pod di TPU TPU