Ejecuta aplicaciones de Cloud TPU en GKE

Advertencia: Estas instrucciones se aplican solo a las TPU que usan la arquitectura de nodos TPU. La arquitectura de nodo TPU solo es compatible con GKE en TPU v2 y v3. Estas instrucciones se aplican solo a TPU v2 y v3. Para obtener información sobre el uso de TPU en GKE, incluida TPU v4, v5e o v5p, consulta Implementa cargas de trabajo de TPU.

En esta guía, se describe cómo hacer las siguientes acciones:

• Establece una configuración de Cloud TPU a fin de prepararla para ejecutarla en Google Kubernetes Engine
• Crea un clúster de GKE compatible con Cloud TPU.
• Usa TensorBoard para visualizar las métricas de Cloud TPU y analizar el rendimiento
• Crea y organiza tu modelo en contenedores en la imagen de Docker

Para obtener más información sobre las arquitecturas de las VM de TPU, consulta Arquitectura del sistema. En esta guía, solo se puede usar con la arquitectura de nodos de TPU.

Beneficios de ejecutar aplicaciones de Cloud TPU en GKE

Las aplicaciones de entrenamiento de Cloud TPU se pueden configurar para ejecutarse en contenedores dentro de Pods de GKE de GKE. Cuando lo estén, obtendrás los siguientes beneficios:

• Configuración y administración mejoradas del flujo de trabajo: GKE administra el ciclo de vida de la TPU. Una vez que la inicialización y el entrenamiento de Cloud TPU se configuran con GKE, GKE puede repetir y administrar tus cargas de trabajo, incluida la recuperación ante fallas de trabajos.

• Costo optimizado: Solo pagas por la TPU mientras el trabajo está activo. GKE crea y borra automáticamente las TPU de acuerdo con los requisitos de recursos de un Pod.

• Uso flexible: Un cambio pequeño en la especificación del pod es suficiente para solicitar un acelerador de hardware diferente (CPU, GPU o TPU):

  kind: Pod
  metadata:
    name: example-tpu
    annotations:
      # The Cloud TPUs that will be created for this Job will support
      # TensorFlow 2.12.1. This version MUST match the
      # TensorFlow version that your model is built on.
      tf-version.cloud-tpus.google.com: "2.12.1"
  spec:
    containers:
    - name: example-container
      resources:
        limits:
          cloud-tpus.google.com/v2: 8
          # See the line above for TPU, or below for CPU / GPU.
          # cpu: 2
          # nvidia.com/gpu: 1
  

• Escalabilidad: GKE proporciona APIs (implementación y trabajo) que pueden escalar a cientos de pods de GKE y nodos TPU.

• Tolerancia a errores: La API de trabajo de GKE, junto con el mecanismo de punto de control de TensorFlow, proporcionan la semántica de ejecución completa. Tus trabajos de entrenamiento se volverán a ejecutar de forma automática con la última lectura del estado del punto de control si se producen fallas en las instancias de VM o en los nodos de Cloud TPU.

Requisitos y limitaciones de la configuración de Cloud TPU y GKE

Cuando definas la configuración, ten en cuenta la información que figura a continuación de GKE:

• Cloud TPU no es compatible con los grupos de nodos de Windows Server.
• Debes crear el clúster de GKE y los grupos de nodos en una zona donde tu Cloud TPU esté disponible. También debes crear los depósitos de Cloud Storage para alojar tus modelos y datos de entrenamiento en la misma región que el clúster de GKE. Consulta el documento de tipos y zonas para obtener una lista de las zonas disponibles.
• Debes usar direcciones IP que cumplan con RFC 1918 para tus clústeres de GKE. Para obtener más información, consulta Herramientas de redes de GKE.
• Cada contenedor puede solicitar, como máximo, una Cloud TPU. Sin embargo, varios contenedores de un pod pueden solicitar una Cloud TPU cada uno.

Antes de comenzar

1. Accede a tu cuenta de Google Cloud. Si eres nuevo en Google Cloud, crea una cuenta para evaluar el rendimiento de nuestros productos en situaciones reales. Los clientes nuevos también obtienen $300 en créditos gratuitos para ejecutar, probar y, además, implementar cargas de trabajo.

2. En la página del selector de proyectos de la consola de Google Cloud, selecciona o crea un proyecto de Google Cloud.

   Nota: Si no planeas conservar los recursos que creaste durante este procedimiento, crea un proyecto en lugar de seleccionar uno existente. Cuando termines, puedes borrar el proyecto y quitar todos los recursos asociados con él.

   Ir al selector de proyectos

3. Asegúrate de que la facturación esté habilitada para tu proyecto de Google Cloud.

4. En la página del selector de proyectos de la consola de Google Cloud, selecciona o crea un proyecto de Google Cloud.

   Nota: Si no planeas conservar los recursos que creaste durante este procedimiento, crea un proyecto en lugar de seleccionar uno existente. Cuando termines, puedes borrar el proyecto y quitar todos los recursos asociados con él.

   Ir al selector de proyectos

5. Asegúrate de que la facturación esté habilitada para tu proyecto de Google Cloud.

Cuando usas Cloud TPU con GKE, tu proyecto usa componentes facturables de Google Cloud. Consulta los precios de Cloud TPU y los precios de GKE a fin de calcular los costos, y sigue las instrucciones para limpiar los recursos una vez que termines de usarlos.

6. Habilita las siguientes APIs en la consola de Google Cloud:

   • API de Cloud TPU
   • API de Compute Engine
   • API de GKE

Crea un clúster nuevo compatible con Cloud TPU

Usa las siguientes instrucciones para configurar tu entorno y crear un clúster de GKE compatible con Cloud TPU mediante gcloud CLI:

1. Instala los componentes de gcloud necesarios para ejecutar GKE con Cloud TPU:

   $ gcloud components install kubectl 

2. Reemplaza gcloud por tu ID del proyecto de Google Cloud.

   $ gcloud config set project project-name
   

   Reemplaza project-name por el nombre del proyecto de Google Cloud.

   La primera vez que ejecutes este comando en una VM de Cloud Shell nueva, se mostrará la página Authorize Cloud Shell. Haz clic en Authorize en la parte inferior de la página para permitir que gcloud realice llamadas a la API de Google Cloud con tus credenciales.

3. Configura gcloud por la zona en la que planeas usar un recurso de Cloud TPU. En este ejemplo, se usa us-central1-b, pero puedes usar una TPU en cualquier zona compatible.

   $ gcloud config set compute/zone us-central1-b
   

4. Usa el comando gcloud container clusters create para crear un clúster en GKE compatible con Cloud TPU.

   Nota: Si necesitas especificar un rango de CIDR existente con una VPC compartida, usa el comando que se muestra en Usa un rango de CIDR existente con una VPC compartida.

   $ gcloud container clusters create cluster-name \
     --release-channel=stable \
     --scopes=cloud-platform \
     --enable-ip-alias \
     --enable-tpu
   

   Descripciones de las marcas de comandos

   Canal de versiones

   Los canales de versiones proporcionan una forma de administrar las actualizaciones automáticas para los clústeres. Cuando creas un clúster nuevo, puedes elegir su canal de versiones. El clúster solo se actualizará a las versiones que se ofrecen en ese canal.

   permisos

   Garantiza que todos los nodos del clúster tengan acceso a tu bucket de Cloud Storage. El clúster y el bucket de almacenamiento deben estar en el mismo proyecto para que esto funcione. Ten en cuenta que los Pods de Kubernetes heredan de forma predeterminada los alcances de los nodos en los que se implementan. Por lo tanto, scopes=cloud-platform otorga a todos los Pods de Kubernetes que se ejecutan en el clúster el permiso cloud-platform. Si deseas limitar el acceso según el Pod, consulta la guía de GKE para autenticar con cuentas de servicio.

   enable-ip-alias

   indica que el clúster usa rangos de IP de alias. Esto es necesario para usar Cloud TPU en GKE.

   enable-tpu

   indica que el clúster debe ser compatible con Cloud TPU.

   tpu-ipv4-cidr (opcional, no especificado anteriormente)

   Indica el rango CIDR que se usará para Cloud TPU. Especifica el IP_RANGE con el formato IP/20, como 10.100.0.0/20. Si no especificas esta marca, se asigna y asigna automáticamente un rango CIDR de tamaño /20.

Cuando se cree el clúster, deberías ver un mensaje similar al siguiente:

NAME             LOCATION       MASTER_VERSION    MASTER_IP     MACHINE_TYPE   NODE_VERSION      NUM_NODES  STATUS
cluster-name  us-central1-b  1.16.15-gke.4901  34.71.245.25  n1-standard-1  1.16.15-gke.4901  3          RUNNING

Solicita una Cloud TPU en las especificaciones de tu pod de Kubernetes

En las especificaciones de tu Pod de Kubernetes, haz lo siguiente:

• Debes compilar tus modelos en tus contenedores con la misma versión de TensorFlow. Consulta las versiones compatibles.

• Especifica el recurso de Cloud TPU en la sección limits en el campo resource de la especificación del contenedor.

  Ten en cuenta que la unidad del recurso de Cloud TPU es el número de núcleos de Cloud TPU. En la siguiente tabla, se enumeran ejemplos de solicitudes de recursos válidas. Consulta Tipos y zonas de TPU para obtener una lista completa de los recursos de TPU válidos.

  Si el recurso que está destinado a usarse es un pod de Cloud TPU, debes solicitar cuota, ya que la cuota predeterminada para un pod de Cloud TPU es cero.

  Solicitud de recurso	Tipo de Cloud TPU
  cloud-tpus.google.com/v2: 8	Un dispositivo de Cloud TPU v2 (8 núcleos)
  cloud-tpus.google.com/preemptible-v2: 8	Un dispositivo de Cloud TPU v2 interrumpible (8 núcleos)
  cloud-tpus.google.com/v3: 8	Un dispositivo de Cloud TPU v3 (8 núcleos)
  cloud-tpus.google.com/preemptible-v3: 8	Un dispositivo de Cloud TPU v3 interrumpible (8 núcleos)
  cloud-tpus.google.com/v2: 32	Un pod de Cloud TPU v2-32 (32 núcleos)
  cloud-tpus.google.com/v3: 32	Un pod de Cloud TPU v3-32 (32 núcleos)

  Si deseas obtener más información acerca de cómo especificar recursos y límites en las especificaciones del pod, consulta la documentación de Kubernetes.

En la siguiente especificación de Pod de muestra, se solicita una TPU interrumpible de Cloud TPU v2-8 con TensorFlow 2.12.0.

La vida útil de los nodos de Cloud TPU está vinculada a los pods que los solicitan. Cloud TPU se crea a demanda cuando se programa el pod y se recicla cuando este se borra.

apiVersion: v1
kind: Pod
metadata:
  name: gke-tpu-pod
  annotations:
     # The Cloud TPUs that will be created for this Job will support
     # TensorFlow 2.12.1. This version MUST match the
     # TensorFlow version that your model is built on.
     tf-version.cloud-tpus.google.com: "2.12.1"
spec:
  restartPolicy: Never
  containers:
  - name: gke-tpu-container
    # The official TensorFlow 2.12.1 image.
    # https://hub.docker.com/r/tensorflow/tensorflow
    image: tensorflow/tensorflow:2.12.1
    command:
    - python
    - -c
    - |
      import tensorflow as tf
      print("Tensorflow version " + tf.__version__)

      tpu = tf.distribute.cluster_resolver.TPUClusterResolver('$(KUBE_GOOGLE_CLOUD_TPU_ENDPOINTS)')
      print('Running on TPU ', tpu.cluster_spec().as_dict()['worker'])

      tf.config.experimental_connect_to_cluster(tpu)
      tf.tpu.experimental.initialize_tpu_system(tpu)
      strategy = tf.distribute.TPUStrategy(tpu)

      @tf.function
      def add_fn(x,y):
          z = x + y
          return z

      x = tf.constant(1.)
      y = tf.constant(1.)
      z = strategy.run(add_fn, args=(x,y))
      print(z)
    resources:
      limits:
        # Request a single Preemptible v2-8 Cloud TPU device to train the model.
        cloud-tpus.google.com/preemptible-v2: 8

Crea el trabajo

Sigue estos pasos para crear el trabajo en el clúster de GKE y, luego, instalar kubectl.

1. Con un editor de texto, crea una especificación de Pod, example-job.yaml, y copia y pega la especificación de Pod que se mostró antes.

2. Ejecuta el trabajo:

   $ kubectl create -f example-job.yaml
   

   pod "gke-tpu-pod" created

   Con este comando, se crea el trabajo que programa el pod de forma automática.

3. Verifica si se programó el Pod de GKE y si se aprovisionaron los nodos Cloud TPU. Un Pod de GKE que solicita nodos de Cloud TPU puede estar pendiente durante 5 minutos antes de ejecutarse. Verás un resultado similar al siguiente hasta que se programe el Pod de GKE.

   $ kubectl get pods -w
   
   NAME          READY     STATUS    RESTARTS   AGE
   gke-tpu-pod   0/1       Pending   0          1m
   

   Después de unos 5 minutos, deberías ver algo como lo siguiente:

   NAME          READY     STATUS              RESTARTS   AGE
   gke-tpu-pod   0/1       Pending             0          21s
   gke-tpu-pod   0/1       Pending             0          2m18s
   gke-tpu-pod   0/1       Pending             0          2m18s
   gke-tpu-pod   0/1       ContainerCreating   0          2m18s
   gke-tpu-pod   1/1       Running             0          2m48s
   gke-tpu-pod   0/1       Completed           0          3m8s
   

   Debes usar Ctrl-C para salir del comando “kubectl get”.

   Puedes imprimir información de registro y recuperar información más detallada sobre cada pod de GKE con los siguientes comandos de kubectl. Por ejemplo, a fin de ver el resultado del registro del pod de GKE, usa el siguiente comando:

   $ kubectl logs gke-tpu-pod

   Debería ver un resultado similar al siguiente:

   2021-09-24 18:55:25.400699: I tensorflow/core/platform/cpu_feature_guard.cc:142]
   This TensorFlow binary is optimized with oneAPI Deep Neural Network Library
   (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 FMA
   To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.
   2021-09-24 18:55:25.405947: I tensorflow/core/distributed_runtime/rpc/grpc_channel.cc:272]
   Initialize GrpcChannelCache for job worker -> {0 -> 10.0.16.2:8470}
   2021-09-24 18:55:25.406058: I tensorflow/core/distributed_runtime/rpc/grpc_channel.cc:272]
   Initialize GrpcChannelCache for job localhost -> {0 -> localhost:32769}
   2021-09-24 18:55:28.091729: I tensorflow/core/distributed_runtime/rpc/grpc_channel.cc:272]
   Initialize GrpcChannelCache for job worker -> {0 -> 10.0.16.2:8470}
   2021-09-24 18:55:28.091896: I tensorflow/core/distributed_runtime/rpc/grpc_channel.cc:272]
   Initialize GrpcChannelCache for job localhost -> {0 -> localhost:32769}
   2021-09-24 18:55:28.092579: I tensorflow/core/distributed_runtime/rpc/grpc_server_lib.cc:427]
   Started server with target: grpc://localhost:32769
   Tensorflow version 2.12.1
   Running on TPU  ['10.0.16.2:8470']
   PerReplica:{
     0: tf.Tensor(2.0, shape=(), dtype=float32),
     1: tf.Tensor(2.0, shape=(), dtype=float32),
     2: tf.Tensor(2.0, shape=(), dtype=float32),
     3: tf.Tensor(2.0, shape=(), dtype=float32),
     4: tf.Tensor(2.0, shape=(), dtype=float32),
     5: tf.Tensor(2.0, shape=(), dtype=float32),
     6: tf.Tensor(2.0, shape=(), dtype=float32),
     7: tf.Tensor(2.0, shape=(), dtype=float32)
   }
   

   Si deseas ver una descripción completa del Pod de GKE, usa el siguiente comando:

   $ kubectl describe pod gke-tpu-pod
   

   Consulta Inspección y depuración de la aplicación para obtener más detalles.

Crea y organiza tu modelo en contenedores en la imagen de Docker

Consulta cómo compilar y organizar tu propio modelo en contenedores para obtener más detalles sobre este proceso.

Habilita la compatibilidad con Cloud TPU en un clúster existente

Esta función está sujeta a las "Condiciones de las Ofertas de Fase Previa a la DG" de la sección Condiciones Generales del Servicio de las Condiciones Específicas del Servicio. Las funciones de fase previa a la DG están disponibles “tal cual” y podrían tener asistencia limitada. Para obtener más información, consulta las descripciones de la etapa de lanzamiento.

Para habilitar la compatibilidad con Cloud TPU en un clúster de GKE existente, sigue estos pasos en Google Cloud CLI:

1. Habilita la compatibilidad con Cloud TPU:

   gcloud beta container clusters update cluster-name --enable-tpu
   

   Reemplaza cluster-name por el nombre del clúster.

   Importante: Es posible que el extremo del servidor de la API original no esté disponible después de actualizar el clúster.

2. Actualiza la entrada de kubeconfig:

   gcloud container clusters get-credentials cluster-name
   

Configura un rango CIDR personalizado

Esta función está sujeta a las "Condiciones de las Ofertas de Fase Previa a la DG" de la sección Condiciones Generales del Servicio de las Condiciones Específicas del Servicio. Las funciones de fase previa a la DG están disponibles “tal cual” y podrían tener asistencia limitada. Para obtener más información, consulta las descripciones de la etapa de lanzamiento.

De forma predeterminada, GKE asigna un bloque CIDR con un tamaño de /20 para las TPU aprovisionadas por el clúster. Puedes especificar un rango de CIDR personalizado para Cloud TPU mediante la ejecución del siguiente comando:

gcloud beta container clusters update cluster-name \
  --enable-tpu \
  --tpu-ipv4-cidr 10.100.0.0/20

Reemplaza lo siguiente:

• cluster-name: Es el nombre del clúster existente.
• 10.100.0.0/20: Es rango de CIDR personalizado.

Usa rangos de CIDR existentes con una VPC compartida

Nota: Esta sección está destinada a los usuarios que implementan un clúster en una red de VPC compartida que necesitan usar rangos CIDR existentes. Si este no es tu caso de uso, omite esta sección.

Sigue la guía sobre TPU en clústeres de GKE con una VPC compartida para verificar la configuración correcta de la VPC compartida.

• Clúster de GKE con Cloud TPU mediante una VPC compartida

Inhabilita Cloud TPU en un clúster

Importante: Inhabilitar la compatibilidad con Cloud TPU borra todas las TPU aprovisionadas por el clúster de GKE. Cualquier carga de trabajo que requiera Cloud TPU dejará de estar disponible una vez que se inicie la operación.

Para inhabilitar la compatibilidad con Cloud TPU en un clúster de GKE existente, sigue estos pasos en Google Cloud CLI:

1. Verifica que ninguna de tus cargas de trabajo use Cloud TPU:

   $ kubectl get tpu
   

2. Inhabilita la compatibilidad con Cloud TPU en tu clúster:

   $ gcloud beta container clusters update cluster-name --no-enable-tpu
   

   Reemplaza cluster-name por el nombre del clúster.

   Para los clústeres zonales, esta operación tarda alrededor de 5 minutos, y para los clústeres regionales, esta operación tarda alrededor de 15 minutos, según la región del clúster.

3. Una vez que se completan las operaciones sin errores, puedes verificar que se hayan quitado las TPU aprovisionadas por el clúster:

   $ gcloud compute tpus list
   

   Los nombres de las TPU que crea Cloud TPU tienen el siguiente formato:

   $ gke-cluster-name-cluster-id-tpu-tpu-id
   

   Reemplaza lo siguiente:

   • cluster-name: Es el nombre del clúster existente.
   • cluster-id: Es el ID de tu clúster existente.
   • tpu-id: Es el ID de la Cloud TPU.

   Si aparece alguna TPU, puedes borrarla de forma manual si ejecutas el siguiente comando:

   $ gcloud compute tpus delete gke-cluster-name-cluster-id-tpu-tpu-id
   

Limpia

Una vez que termines con Cloud TPU en GKE, limpia los recursos a fin de evitar que se apliquen cargos adicionales a la cuenta de Facturación de Cloud

1. Ejecuta el siguiente comando para borrar el clúster de GKE; reemplaza cluster-name por el nombre del clúster y project-name por el nombre del proyecto de Google Cloud:

   $ gcloud container clusters delete cluster-name \
   --project=project-name --zone=us-central1-b
   

2. Cuando termines de examinar los datos, usa el comando gsutil para borrar el bucket de Cloud Storage que creaste. Reemplaza bucket-name por el nombre de tu bucket de Cloud Storage.

   $ gsutil rm -r gs://bucket-name