Introducción a TPU en GKE

Los clientes de Google Kubernetes Engine (GKE) ahora pueden crear grupos de nodos de Kubernetes que contengan porciones de TPU v4 y v5e. Para obtener más información sobre las TPU, consulta Arquitectura del sistema.

Cuando trabajas con GKE, primero debes crear un clúster de GKE. Luego, agrega grupos de nodos a tu clúster. Los grupos de nodos de GKE son colecciones de VMs que comparten los mismos atributos. Para las cargas de trabajo de TPU, los grupos de nodos consisten en VMs de TPU.

Tipos de grupos de nodos

GKE admite dos tipos de grupos de nodo TPU:

Grupo de nodos de porciones de TPU de varios hosts

Un grupo de nodos en porciones de TPU de varios hosts es un grupo de nodos que contiene dos o más VM de TPU interconectadas. Cada VM tiene un dispositivo de TPU conectado. Las TPU de una porción de varios hosts se conectan a través de una interconexión de alta velocidad (ICI). Una vez que se crea un grupo de nodos de porción de varios hosts, no puedes agregarle nodos. Por ejemplo, no puedes crear un grupo de nodos v4-32 y, luego, agregar un nodo de Kubernetes adicional (VM de TPU) al grupo de nodos. Para agregar una porción de TPU adicional a un clúster de GKE, debes crear un grupo de nodos nuevo.

Los hosts en un grupo de nodos de porción de TPU de varios hosts se tratan como una sola unidad atómica. Si GKE no puede implementar un nodo de la porción, no se implementará ningún nodo de la porción.

Si es necesario reparar un nodo dentro de una porción de TPU de varios hosts, GKE apagará todas las VM de TPU de la porción, lo que forzará la expulsión de todos los Pods de Kubernetes en la carga de trabajo. Una vez que todas las VM de TPU de la porción están en funcionamiento, los pods de Kubernetes se pueden programar en las VM de TPU en la nueva porción.

En el siguiente diagrama, se muestra un ejemplo de una porción de TPU de varios hosts v5litepod-16 (v5e). Esta porción tiene cuatro VMs de TPU. Cada VM de TPU tiene cuatro chips TPU v5e conectados con interconexiones de alta velocidad (ICI), y cada chip TPU v5e tiene un TensorCore.

Diagrama de porciones de TPU de varios hosts

En el siguiente diagrama, se muestra un clúster de GKE que contiene una porción v5litepod-16 de TPU (v5e) (topología: 4x4) y una porción v5litepod-8 (v5e) de TPU (topología: 2x4):

Diagrama del Pod de TPU v5e

Para ver un ejemplo de la ejecución de una carga de trabajo en una porción de TPU de varios hosts, consulta Ejecuta tu carga de trabajo en TPU.

Grupos de nodos de porciones de TPU de host único

Un grupo de nodos de porción de host único es un grupo de nodos que contiene una o más VM de TPU independientes. Cada VM tiene un dispositivo de TPU conectado. Si bien las VM dentro de un grupo de nodos de una porción de host único se pueden comunicar a través de la red del centro de datos (DCN), las TPU adjuntas a las VM no están interconectadas.

En el siguiente diagrama, se muestra un ejemplo de una porción de TPU de host único con siete máquinas v4-8:

Diagrama del grupo de nodos de porción de host único

Para ver un ejemplo de la ejecución de una carga de trabajo en una porción de TPU de host único, consulta Ejecuta tu carga de trabajo en TPU.

Tipos de máquinas de TPU para los grupos de nodos de GKE

Antes de crear grupos de nodos, debes elegir la versión de TPU y el tamaño de la porción de TPU que requiere tu carga de trabajo. TPU v4 es compatible con GKE Standard 1.26.1-gke.1500 y versiones posteriores, v5e en GKE Standard 1.27.2-gke.2100 y versiones posteriores, y v5p en GKE Standard 1.28.3-gke.1024000 y versiones posteriores.

TPU v4, v5e y v5p son compatibles con GKE Autopilot versión 1.29.2-gke.1521000 y versiones posteriores.

Para obtener más información sobre las especificaciones de hardware de las diferentes versiones de TPU, consulta Arquitectura del sistema. Cuando crees un grupo de nodo TPU, selecciona un tamaño de porción de TPU (una topología de TPU) en función del tamaño de tu modelo y la cantidad de memoria que requiere. El tipo de máquina que especificas cuando creas los grupos de nodos depende de la versión y el tamaño de las porciones.

v5e

A continuación, se muestran los tipos de máquinas y las topologías de TPU v5e que se admiten en los casos prácticos de entrenamiento y de inferencia:

Tipo de máquina Topología Cantidad de chips TPU Cantidad de VM Caso práctico recomendado
ct5lp-hightpu-1t 1x1 1 1 Inferencia de entrenamiento, host único
ct5lp-hightpu-4t 2x2 4 1 Inferencia de entrenamiento, host único
ct5lp-hightpu-8t 2x4 8 1 Inferencia de entrenamiento, host único
ct5lp-hightpu-4t 2x4 8 2 Inferencia de múltiples hosts y entrenamiento
ct5lp-hightpu-4t 4x4 16 4 Entrenamiento a gran escala, inferencia de varios hosts
ct5lp-hightpu-4t 4x8 32 8 Entrenamiento a gran escala, inferencia de varios hosts
ct5lp-hightpu-4t 8x8 64 16 Entrenamiento a gran escala, inferencia de varios hosts
ct5lp-hightpu-4t 8x16 128 32 Entrenamiento a gran escala, inferencia de varios hosts
ct5lp-hightpu-4t 16x16 256 64 Entrenamiento a gran escala, inferencia de varios hosts

Cloud TPU v5e es un producto combinado de entrenamiento e inferencia. Los trabajos de entrenamiento están optimizados para la capacidad de procesamiento y la disponibilidad, mientras que los trabajos de inferencia están optimizados para la latencia. Para obtener más información, consulta Tipos de acelerador de entrenamiento v5e y Tipos de acelerador de inferencia v5e.

Las máquinas de TPU v5e están disponibles en us-west4-a, us-east5-b y us-east1-c. Los clústeres de GKE Standard deben ejecutar la versión 1.27.2-gke.2100 del plano de control o una posterior. GKE Autopilot debe ejecutar la versión del plano de control 1.29.2-gke.1521000 o una posterior. Para obtener más información sobre v5e, consulta Entrenamiento de Cloud TPU v5e.

Comparación del tipo de máquina:

Tipo de máquina ct5lp-hightpu-1t ct5lp-hightpu-4t ct5lp-hightpu-8t
Cantidad de chips v5e 1 4 8
Cantidad de CPU virtuales 24 112 224
RAM (GB) 48 192 384
Cantidad de nodos de NUMA 1 1 2
Probabilidad de interrupción Alto Media Baja

Para liberar espacio para las VMs que tienen más chips, el programador de GKE puede interrumpir y reprogramar las VMs con menos chips. Así, las VMs de 8 chips tienen más probabilidades de interrumpir a las VMs de 1 y 4 chips.

v4 y v5p

Los tipos de máquinas TPU v4 y v5p son los siguientes:

Tipo de máquina Cantidad de CPUs virtuales Memoria (GB) Cantidad de nodos de NUMA
ct4p-hightpu-4t 240 407 2
ct5p-hightpu-4t 208 448 2

Cuando crees una porción de TPU v4, usa el tipo de máquina ct4p-hightpu-4t, que tiene un host y contiene 4 chips. Consulta Topologías v4 y arquitectura del sistema de TPU para obtener más información. Los tipos de máquinas de porciones de TPU v4 están disponibles en us-central2-b. Los clústeres de GKE Standard deben ejecutar la versión 1.26.1-gke.1500 del plano de control o una versión posterior. Los clústeres de GKE Autopilot deben ejecutar la versión 1.29.2-gke.1521000 del plano de control o una posterior.

Cuando crees una porción de TPU v5p, usa el tipo de máquina ct5p-hightpu-4t, que tiene un host y contiene 4 chips. Los tipos de máquinas de porciones de TPU v5p están disponibles en us-west4-a y us-east5-a. Los clústeres de GKE Standard deben ejecutar la versión 1.28.3-gke.1024000 del plano de control o una posterior. GKE Autopilot debe ejecutar 1.29.2-gke.1521000 o más tarde. Para obtener más información sobre v5p, consulta Introducción a la capacitación sobre v5p.

Problemas conocidos y limitaciones

  • Cantidad máxima de Pods de Kubernetes: Puedes ejecutar un máximo de 256 Pods de Kubernetes en una sola VM de TPU.
  • Solo reservas específicas: Cuando se usan TPU en GKE, SPECIFIC es el único valor admitido para la marca --reservation-affinity del comando gcloud container node-pools create.
  • Solo se admite la variante de VMs Spot de las TPU interrumpibles: Las VMs Spot son similares a las VMs interrumpibles y están sujetas a las mismas limitaciones de disponibilidad, pero no tienen una duración máxima de 24 horas.
  • No admite la asignación de costos: La asignación de costos de GKE y la medición de uso no incluyen datos sobre el uso ni los costos de las TPU.
  • El escalador automático puede calcular la capacidad: Es posible que el escalador automático del clúster calcule la capacidad de forma incorrecta para los nodos nuevos que contengan VM de TPU antes de que esos nodos estén disponibles. El escalador automático del clúster puede realizar un escalamiento vertical adicional y, como resultado, crear más nodos de los necesarios. El escalador automático del clúster reducirá verticalmente la escala de los nodos adicionales, si no son necesarios, después de la operación regular de reducción de escala.
  • El escalador automático cancela el escalamiento vertical: el escalador automático del clúster cancela el escalamiento vertical de los grupos de nodos TPU que permanecen en estado de espera durante más de 10 horas. El escalador automático del clúster volverá a intentar realizar esas operaciones de escalamiento vertical más adelante. Este comportamiento puede reducir la capacidad de obtener TPU para los clientes que no usan reservas.
  • El taint puede evitar la reducción de escala: Las cargas de trabajo que no son de TPU que tienen una tolerancia para el taint de TPU pueden evitar la reducción de escala del grupo de nodos si se recrean durante el desvío del grupo de nodo TPU.

Asegúrate de que haya suficientes cuotas de TPU y GKE

Es posible que debas aumentar ciertas cuotas relacionadas con GKE en las regiones en las que se crean tus recursos.

Las siguientes cuotas tienen valores predeterminados que probablemente deban aumentarse:

  • Cuota SSD (GB) de Persistent Disk: El disco de arranque de cada nodo de Kubernetes requiere 100 GB de forma predeterminada. Por lo tanto, esta cuota debe establecerse al menos en un valor tan alto como (la cantidad máxima de nodos de GKE que prevés que crearán) × 100 GB.
  • Cuota de direcciones IP en uso: Cada nodo de Kubernetes consume una dirección IP. Por lo tanto, esta cuota debe establecerse en, al menos, un valor tan alto como la cantidad máxima de nodos de GKE que prevés crear.

Para solicitar un aumento de la cuota, consulta Solicita una cuota mayor. Para obtener más información sobre los tipos de cuotas de TPU, consulta Cuota de TPU.

Es posible que la aprobación de las solicitudes de aumento de cuota tarde algunos días. Si tienes algún problema para que se aprueben las solicitudes de aumento de cuota en unos pocos días, comunícate con el equipo de Cuentas de Google.

Migra tu reserva de TPU

Si no planeas usar una reserva de TPU existente con TPU en GKE, omite esta sección y ve a Crea un clúster de Google Kubernetes Engine.

Para usar TPU reservadas con GKE, primero debes migrar tu reserva de TPU a un sistema de reservas nuevo basado en Compute Engine.

Hay varios aspectos importantes que debes saber sobre esta migración:

  • La capacidad de TPU migrada al nuevo sistema de reservas basado en Compute Engine no se puede usar con la API de recursos en cola de Cloud TPU. Si piensas usar recursos en cola de TPU con tu reserva, deberás migrar una parte de tu reserva de TPU al nuevo sistema de reservas basado en Compute Engine.
  • No puede haber cargas de trabajo en ejecución activas en las TPU cuando se migran al nuevo sistema de reservas basado en Compute Engine.
  • Selecciona un momento para realizar la migración y trabaja con tu equipo de cuentas de Google Cloud para programar la migración. El período de migración debe ser durante el horario de atención (de lunes a viernes, de 9 a.m. a 5 p.m., hora del Pacífico).

Crea un clúster de Google Kubernetes Engine

Consulta Crea un clúster en la documentación de Google Kubernetes Engine.

Crea un grupo de nodo TPU

Consulta Crea un grupo de nodos en la documentación de Google Kubernetes Engine.

Ejecución sin modo privilegiado

Si quieres reducir el alcance del permiso en tu contenedor, consulta el modo de privilegio de TPU.

Ejecuta cargas de trabajo en grupos de nodo TPU

Consulta Ejecuta tus cargas de trabajo de GKE en TPU en la documentación de Google Kubernetes Engine.

Selectores de nodos

Con el fin de que Kubernetes programe tu carga de trabajo en nodos que contengan VM de TPU, debes especificar dos selectores para cada nodo en tu manifiesto de Google Kubernetes Engine:

  • Establece cloud.google.com/gke-accelerator-type en tpu-v5-lite-podslice o tpu-v4-podslice.
  • Establece cloud.google.com/gke-tpu-topology en la topología de TPU del nodo.

Las secciones Cargas de trabajo de entrenamiento y Cargas de trabajo de inferencia contienen ejemplos de manifiestos que ilustran el uso de estos selectores de nodos.

Consideraciones de programación de la carga de trabajo

Las TPU tienen características únicas que requieren programación y administración especiales de cargas de trabajo en Kubernetes. Para obtener más información, consulta Consideraciones sobre la programación de cargas de trabajo en la documentación de GKE.

Reparación de nodos

Si un nodo en un grupo de nodos de porciones de TPU de varios hosts está en mal estado, GKE vuelve a crear todo el grupo de nodos. Para obtener más información, consulta Reparación automática de nodos en la documentación de GKE.

Multislice: Más allá de una sola porción

Puedes agregar porciones más pequeñas en una multislice para manejar cargas de trabajo de entrenamiento más grandes. Para obtener más información, consulta Multislice de Cloud TPU.

Instructivos de cargas de trabajo de entrenamiento

Estos instructivos se enfocan en el entrenamiento de cargas de trabajo en una porción de TPU de varios hosts (por ejemplo, 4 máquinas v5e). Abarcan los siguientes modelos:

  • Modelos FLAX de Hugging Face: difusión de pokémones
  • PyTorch/XLA: GPT2 en WikiText

Descarga los recursos del instructivo

Descarga el instructivo de secuencias de comandos de Python y las especificaciones YAML para cada modelo previamente entrenado con el siguiente comando:

git clone https://github.com/GoogleCloudPlatform/ai-on-gke.git

Crea un clúster y conéctate al clúster

Crear un clúster de GKE regional para que el plano de control de Kubernetes se replique en tres zonas, lo que proporciona una mayor disponibilidad Crea tu clúster en us-west4, us-east1 o us-central2, según la versión de TPU que uses. Para obtener más información sobre las TPU y las zonas, consulta Regiones y zonas de Cloud TPU.

El siguiente comando crea un clúster regional de GKE nuevo suscrito al canal de versiones rápido con un grupo de nodos que, en un principio, contiene un nodo por zona. El comando también habilita las funciones del controlador de CSI de Cloud Storage FUSE y Workload Identity en el clúster porque las cargas de trabajo de inferencia de ejemplo en esta guía usan buckets de Cloud Storage para almacenar modelos previamente entrenados.

gcloud container clusters create cluster-name \
  --region your-region \
  --release-channel rapid \
  --num-nodes=1 \
  --workload-pool=project-id.svc.id.goog \
  --addons GcsFuseCsiDriver

Si quieres habilitar las funciones del controlador de CSI de Cloud Storage FUSE y Workload Identity para los clústeres existentes, ejecuta el siguiente comando:

gcloud container clusters update cluster-name \
  --region your-region \
  --update-addons GcsFuseCsiDriver=ENABLED \
  --workload-pool=project-id.svc.id.goog

Las cargas de trabajo de ejemplo se configuran con las siguientes suposiciones:

  • el grupo de nodos usa tpu-topology=4x4 con cuatro nodos
  • el grupo de nodos usa machine-type ct5lp-hightpu-4t

Ejecuta el siguiente comando para conectarte al clúster que acabas de crear:

gcloud container clusters get-credentials cluster-name \
--location=cluster-region

Modelos FLAX de Hugging Face: difusión de pokémones

En este ejemplo, se entrena el modelo Stable Diffusion de HuggingFace con el conjunto de datos Pokémon.

El modelo de difusión estable es un modelo de texto a imagen latente que genera imágenes fotorrealistas a partir de cualquier entrada de texto. Para obtener más información sobre la difusión estable, consulta:

Crear imagen de Docker

El Dockerfile se encuentra en la carpeta ai-on-gke/tutorials-and-examples/tpu-examples/training/diffusion/.

Antes de ejecutar el siguiente comando, asegúrate de que tu cuenta tenga los permisos adecuados para que Docker la envíe al repositorio.

Compila y envía la imagen de Docker:

cd ai-on-gke/tutorials-and-examples/tpu-examples/training/diffusion/
docker build -t gcr.io/project-id/diffusion:latest .
docker push gcr.io/project-id/diffusion:latest

Implementa la carga de trabajo

Crea un archivo con el siguiente contenido y asígnale el nombre tpu_job_diffusion.yaml. Completa el campo de la imagen con la imagen que acabas de crear.

apiVersion: v1
kind: Service
metadata:
  name: headless-svc
spec:
  clusterIP: None
  selector:
    job-name: tpu-job-diffusion
---
apiVersion: batch/v1
kind: Job
metadata:
  name: tpu-job-diffusion
spec:
  backoffLimit: 0
  # Completions and parallelism should be the number of chips divided by 4.
  # (e.g. 4 for a v5litepod-16)
  completions: 4
  parallelism: 4
  completionMode: Indexed
  template:
    spec:
      subdomain: headless-svc
      restartPolicy: Never
      nodeSelector:
        cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
        cloud.google.com/gke-tpu-topology: 4x4
      containers:
      - name: tpu-job-diffusion
        image: gcr.io/${project-id}/diffusion:latest
        ports:
        - containerPort: 8471 # Default port using which TPU VMs communicate
        - containerPort: 8431 # Port to export TPU usage metrics, if supported
        command:
        - bash
        - -c
        - |
          cd examples/text_to_image
          python3 train_text_to_image_flax.py --pretrained_model_name_or_path=duongna/stable-diffusion-v1-4-flax --dataset_name=lambdalabs/pokemon-blip-captions --resolution=128 --center_crop --random_flip --train_batch_size=4 --mixed_precision=fp16 --max_train_steps=1500 --learning_rate=1e-05 --max_grad_norm=1 --output_dir=sd-pokemon-model
        resources:
          requests:
            google.com/tpu: 4
          limits:
            google.com/tpu: 4

Luego, impleméntalo con el comando siguiente:

kubectl apply -f tpu_job_diffusion.yaml

Limpieza

Cuando el trabajo termine de ejecutarse, puedes borrarlo con los siguientes comandos:

kubectl delete -f tpu_job_diffusion.yaml

PyTorch/XLA: GPT2 en WikiText

En este instructivo, se muestra cómo ejecutar GPT2 en TPU v5e con HuggingFace en PyTorch/XLA mediante el conjunto de datos wikitext.

Crear imagen de Docker

El Dockerfile se encuentra en la carpeta ai-on-gke/tutorials-and-examples/tpu-examples/training/gpt/.

Antes de ejecutar el siguiente comando, asegúrate de que tu cuenta tenga los permisos adecuados para que Docker la envíe al repositorio.

Compila y envía la imagen de Docker:

cd ai-on-gke/tutorials-and-examples/tpu-examples/training/gpt/
docker build -t gcr.io/project-id/gpt:latest .
docker push gcr.io/project-id/gpt:latest

Implementa la carga de trabajo

Copia el siguiente YAML y guárdalo en un archivo llamado tpu_job_gpt.yaml. Completa el campo de la imagen con la imagen que acabas de crear.

apiVersion: v1
kind: Service
metadata:
  name: headless-svc
spec:
  clusterIP: None
  selector:
    job-name: tpu-job-gpt
---
apiVersion: batch/v1
kind: Job
metadata:
  name: tpu-job-gpt
spec:
  backoffLimit: 0
  # Completions and parallelism should be the number of chips divided by 4.
  # (for example, 4 for a v5litepod-16)
  completions: 4
  parallelism: 4
  completionMode: Indexed
  template:
    spec:
      subdomain: headless-svc
      restartPolicy: Never
      volumes:
      # Increase size of tmpfs /dev/shm to avoid OOM.
      - name: shm
        emptyDir:
          medium: Memory
          # consider adding `sizeLimit: XGi` depending on needs
      nodeSelector:
        cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
        cloud.google.com/gke-tpu-topology: 4x4
      containers:
      - name: tpu-job-gpt
        image: gcr.io/$(project-id)/gpt:latest
        ports:
        - containerPort: 8479
        - containerPort: 8478
        - containerPort: 8477
        - containerPort: 8476
        - containerPort: 8431 # Port to export TPU usage metrics, if supported.
        env:
        - name: PJRT_DEVICE
          value: 'TPU'
        - name: XLA_USE_BF16
          value: '1'
        command:
        - bash
        - -c
        - |
          numactl --cpunodebind=0 python3 -u examples/pytorch/xla_spawn.py   --num_cores 4 examples/pytorch/language-modeling/run_clm.py    --num_train_epochs 3 --dataset_name wikitext     --dataset_config_name wikitext-2-raw-v1 --per_device_train_batch_size 16    --per_device_eval_batch_size 16 --do_train --do_eval  --output_dir /tmp/test-clm     --overwrite_output_dir --config_name my_config_2.json --cache_dir /tmp --tokenizer_name gpt2  --block_size 1024 --optim adafactor --adafactor true --save_strategy no --logging_strategy no --fsdp "full_shard" --fsdp_config fsdp_config.json
        volumeMounts:
        - mountPath: /dev/shm
          name: shm
        resources:
          requests:
            google.com/tpu: 4
          limits:
            google.com/tpu: 4

Implementa el flujo de trabajo con lo siguiente:

kubectl apply -f tpu_job_gpt.yaml

Limpieza

Cuando tu trabajo termine de ejecutarse, puedes borrarlo con el siguiente comando:

kubectl delete -f tpu_job_gpt.yaml

Instructivo: Cargas de trabajo de inferencia de un solo host

En este instructivo, se muestra cómo ejecutar una carga de trabajo de inferencia de host único en TPU v5e de GKE para modelos previamente entrenados con JAX, TensorFlow y PyTorch. En un nivel alto, hay cuatro pasos separados que se deben realizar en el clúster de GKE:

  1. Crear un bucket de Cloud Storage y configurar el acceso al bucket Se usa un bucket de Cloud Storage para almacenar el modelo previamente entrenado.

  2. Descarga y convierte un modelo previamente entrenado en uno compatible con TPU. Aplica un pod de Kubernetes que descargue el modelo previamente entrenado, use el convertidor de Cloud TPU y almacene los modelos convertidos en un bucket de Cloud Storage con el controlador de CSI del FUSE de Cloud Storage. El conversor de Cloud TPU no requiere hardware especializado. En este instructivo, se muestra cómo descargar el modelo y ejecutar el convertidor de Cloud TPU en el grupo de nodos de la CPU.

  3. Inicia el servidor para el modelo convertido. Aplica un Deployment que entregue el modelo mediante un marco de trabajo de servidor respaldado por el volumen almacenado en el volumen persistente ReadOnlyOnly (ROX). Las réplicas de implementación deben ejecutarse en un grupo de nodos en porción v5e con un pod de Kubernetes por nodo y en un grupo de nodos en una porción v5e con un pod de Kubernetes por nodo.

  4. Implementa un balanceador de cargas para probar el servidor de modelos. El servidor se expone a solicitudes externas mediante el servicio LoadBalancer. Se proporcionó una secuencia de comandos de Python con una solicitud de ejemplo para probar el servidor del modelo.

En el siguiente diagrama, se muestra cómo enruta las solicitudes el balanceador de cargas.

Un diagrama que muestra el enrutamiento del balanceador de cargas

Ejemplos de implementación de servidores

Estas cargas de trabajo de ejemplo se configuran con las siguientes suposiciones:

  • El clúster se ejecuta con un grupo de nodos de TPU v5 con 3 nodos.
  • El grupo de nodos usa el tipo de máquina ct5lp-hightpu-1t, en el que:
    • la topología es 1 × 1
    • la cantidad de chips TPU es 1

En el siguiente manifiesto de GKE, se define una única Deployment de servidor host.

apiVersion: apps/v1
kind: Deployment
metadata:
  name: bert-deployment
spec:
  selector:
    matchLabels:
      app: tf-bert-server
  replicas: 3 # number of nodes in node pool
  template:
    metadata:
      annotations:
        gke-gcsfuse/volumes: "true"
      labels:
        app: tf-bert-server
    spec:
      nodeSelector:
        cloud.google.com/gke-tpu-topology: 1x1  # target topology
        cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice  # target version
      containers:
      - name: serve-bert
        image: us-docker.pkg.dev/cloud-tpu-images/inference/tf-serving-tpu:2.13.0
        env:
        - name: MODEL_NAME
          value: "bert"
        volumeMounts:
        - mountPath: "/models/"
          name: bert-external-storage
        ports:
        - containerPort: 8500
        - containerPort: 8501
        - containerPort: 8431 # Port to export TPU usage metrics, if supported.
        resources:
          requests:
            google.com/tpu: 1 # TPU chip request
          limits:
            google.com/tpu: 1 # TPU chip request
      volumes:
      - name: bert-external-storage
        persistentVolumeClaim:
          claimName: external-storage-pvc

Si usas una cantidad diferente de nodos en el grupo de nodo TPU, cambia el campo replicas por la cantidad de nodos.

Si tu clúster estándar ejecuta la versión 1.27 de GKE o una anterior, agrega el siguiente campo a tu manifiesto:

spec:
  securityContext:
    privileged: true

No necesitas ejecutar Pods de Kubernetes en modo privilegiado en GKE 1.28 o versiones posteriores. Para obtener más información, consulta Ejecuta contenedores sin modo privilegiado.

Si usas un tipo de máquina diferente, haz lo siguiente:

  • Establece cloud.google.com/gke-tpu-topology en la topología del tipo de máquina que usas.
  • Configura ambos campos google.com/tpu en resources para que coincidan con la cantidad de chips del tipo de máquina correspondiente.

Configuración

Descarga el instructivo de secuencias de comandos de Python y los manifiestos YAML con el siguiente comando:

git clone https://github.com/GoogleCloudPlatform/ai-on-gke.git

Ve al directorio single-host-inference:

cd ai-on-gke/gke-tpu-examples/single-host-inference/

Configura el entorno de Python

Las secuencias de comandos de Python que usarás en este instructivo requieren Python versión 3.9 o superior. Recuerda instalar requirements.txt para cada instructivo antes de ejecutar las secuencias de comandos de prueba de Python.

Si no tienes la configuración de Python adecuada en tu entorno local, puedes usar Cloud Shell para descargar y ejecutar las secuencias de comandos de Python de este instructivo.

Configura el clúster

  1. Crea un clúster con el tipo de máquina e2-standard-4.

    gcloud container clusters create cluster-name \
--region your-region \
--release-channel rapid \
--num-nodes=1 \
--machine-type=e2-standard-4 \
--workload-pool=project-id.svc.id.goog \
--addons GcsFuseCsiDriver

  2. Crea el nodo TPU de host único.

Las cargas de trabajo de ejemplo suponen lo siguiente:

  • Tu clúster se ejecuta con un grupo de nodos de TPU v5e con 3 nodos.
  • El grupo de nodo TPU usa el tipo de máquina ct5lp-hightpu-1t.

Si usas una configuración de clúster diferente a la descrita anteriormente, deberás editar el manifiesto de implementación del servidor.

Para la demostración de JAX Stable Diffusion, necesitarás un grupo de nodos de CPU con un tipo de máquina que tenga más de 16 Gi de memoria disponible (por ejemplo, e2-standard-4). Esto se configura en el comando gcloud container clusters create o se agrega un grupo de nodos adicional al clúster existente con el siguiente comando:

gcloud beta container node-pools create your-pool-name \
  --zone=your-cluster-zone \
  --cluster=your-cluster-name \
  --machine-type=e2-standard-4 \
  --num-nodes=1

Reemplaza lo siguiente:

  • your-pool-name: el nombre del grupo de nodos que se creará.
  • your-cluster-zone: Es la zona en la que se creó el clúster.
  • your-cluster-name: Es el nombre del clúster en el que se agregará el grupo de nodos.
  • your-machine-type: Es el tipo de máquina de los nodos que se crearán en el grupo de nodos.

Configura el almacenamiento del modelo

Existen varias formas de almacenar tu modelo para la entrega. En este instructivo, usaremos el siguiente enfoque:

  • Para convertir el modelo previamente entrenado de modo que funcione en TPU, usaremos una nube privada virtual respaldada por Persistent Disk con acceso ReadWriteMany (RWX).
  • Para entregar el modelo en varias TPU de host único, usaremos la misma VPC respaldada por el bucket de Cloud Storage.

Ejecuta el siguiente comando para crear un bucket de Cloud Storage.

gcloud storage buckets create gs://your-bucket-name \
  --project=your-bucket-project-id \
  --location=your-bucket-location

Reemplaza lo siguiente:

  • your-bucket-name: Es el nombre del bucket de Cloud Storage.
  • your-bucket-project-id: Es el ID del proyecto en el que creaste el bucket de Cloud Storage.
  • your-bucket-location: Es la ubicación del bucket de Cloud Storage. Para mejorar el rendimiento, especifica la ubicación en la que se ejecuta tu clúster de GKE.

Sigue estos pasos para darle a tu clúster de GKE acceso al bucket. Para simplificar la configuración, en los siguientes ejemplos se usa el espacio de nombres predeterminado y la cuenta de servicio predeterminada de Kubernetes. Para obtener más detalles, consulta Configura el acceso a los buckets de Cloud Storage con Workload Identity de GKE.

  1. Crea una cuenta de servicio de IAM para tu aplicación o usa una cuenta de servicio de IAM existente en su lugar. Puedes usar cualquier cuenta de servicio de IAM en el proyecto de tu bucket de Cloud Storage.

    gcloud iam service-accounts create your-iam-service-acct \
--project=your-bucket-project-id

    Reemplaza lo siguiente:

    • your-iam-service-acct: Es el nombre de la cuenta de servicio de IAM nueva.
    • your-bucket-project-id: Es el ID del proyecto en el que creaste la cuenta de servicio de IAM. La cuenta de servicio de IAM debe estar en el mismo proyecto que tu bucket de Cloud Storage.

  2. Asegúrate de que tu cuenta de servicio de IAM tenga los roles de almacenamiento que necesitas.

    gcloud storage buckets add-iam-policy-binding gs://your-bucket-name \
--member "serviceAccount:your-iam-service-acct@your-bucket-project-id.iam.gserviceaccount.com" \
--role "roles/storage.objectAdmin"

    Reemplaza lo siguiente:

    • your-bucket-name: Es el nombre de tu bucket de Cloud Storage.
    • your-iam-service-acct: Es el nombre de la cuenta de servicio de IAM nueva.
    • your-bucket-project-id: Es el ID del proyecto en el que creaste la cuenta de servicio de IAM.

  3. Para permitir que la cuenta de servicio de Kubernetes actúe en nombre de la cuenta de servicio de IAM, agrega una vinculación de política de IAM entre las dos cuentas de servicio. Esta vinculación permite que la cuenta de servicio de Kubernetes funcione como la cuenta de servicio de IAM.

    gcloud iam service-accounts add-iam-policy-binding your-iam-service-acct@your-bucket-project-id.iam.gserviceaccount.com \
  --role roles/iam.workloadIdentityUser \
  --member "serviceAccount:your-project-id.svc.id.goog[default/default]"

    Reemplaza lo siguiente:

    • your-iam-service-acct: Es el nombre de la cuenta de servicio de IAM nueva.
    • your-bucket-project-id: Es el ID del proyecto en el que creaste la cuenta de servicio de IAM.
    • your-project-id: Es el ID del proyecto en el que creaste el clúster de GKE. Tus buckets de Cloud Storage y el clúster de GKE pueden estar en el mismo proyecto o en proyectos diferentes.

  4. Anota la cuenta de servicio de Kubernetes con la dirección de correo electrónico de la cuenta de servicio de IAM.

    kubectl annotate serviceaccount default \
  --namespace default \
  iam.gke.io/gcp-service-account=your-iam-service-acct@your-bucket-project-id.iam.gserviceaccount.com

    Reemplaza lo siguiente:

    • your-iam-service-acct: Es el nombre de la cuenta de servicio de IAM nueva.
    • your-bucket-project-id: Es el ID del proyecto en el que creaste la cuenta de servicio de IAM.

  5. Ejecuta el siguiente comando para propagar el nombre del bucket en los archivos YAML de esta demostración:

    find . -type f -name "*.yaml" | xargs sed -i "s/BUCKET_NAME/your-bucket-name/g"

    Reemplaza your-bucket-name por el nombre del bucket de Cloud Storage.

  6. Crea la reclamación de volumen persistente y volumen persistente con el siguiente comando:

    kubectl apply -f pvc-pv.yaml

Inferencia y entrega de modelos JAX

Instala dependencias de Python para ejecutar secuencias de comandos de Python instructivos que envíen solicitudes al servicio de modelos de JAX.

pip install -r jax/requirements.txt

Ejecuta la demostración de publicación de JAX BERT E2E:

En esta demostración, se usa un modelo BERT previamente entrenado de Hugging Face.

El Pod de Kubernetes realiza los siguientes pasos:

  1. Descarga y usa la secuencia de comandos de Python export_bert_model.py de los recursos de ejemplo para descargar el modelo bert previamente entrenado en un directorio temporal.
  2. Usa la imagen del convertidor de Cloud TPU para convertir el modelo previamente entrenado de CPU a TPU y almacena el modelo en el bucket de Cloud Storage que creaste durante la setup.

Este Pod de Kubernetes está configurado para ejecutarse en la CPU del grupo de nodos predeterminado. Ejecuta el Pod con el siguiente comando:

kubectl apply -f jax/bert/install-bert.yaml

Sigue estos pasos para verificar que el modelo se haya instalado correctamente:

kubectl get pods install-bert

La STATUS puede tardar unos minutos en leer Completed.

Inicia el servidor de modelos de TF para el modelo

En las cargas de trabajo de ejemplo de este instructivo, se supone lo siguiente:

  • El clúster se ejecuta con un grupo de nodos de TPU v5 con tres nodos.
  • El grupo de nodos usa el tipo de máquina ct5lp-hightpu-1t que contiene un chip TPU.

Si usas una configuración de clúster diferente a la descrita anteriormente, deberás editar el manifiesto de implementación del servidor.

Aplicar la implementación
kubectl apply -f jax/bert/serve-bert.yaml

Verifica que el servidor se esté ejecutando con el siguiente comando:

kubectl get deployment bert-deployment

AVAILABLE puede tardar un minuto en leer 3.

Aplicar el servicio del balanceador de cargas
kubectl apply -f jax/bert/loadbalancer.yaml

Verifica que el balanceador de cargas esté listo para el tráfico externo con el siguiente comando:

kubectl get svc tf-bert-service

Es posible que EXTERNAL_IP tarde unos minutos en tener una IP en la lista.

Envía la solicitud al servidor de modelos

Obtén una IP externa del servicio del balanceador de cargas:

EXTERNAL_IP=$(kubectl get services tf-bert-service --output jsonpath='{.status.loadBalancer.ingress[0].ip}')

Ejecuta una secuencia de comandos para enviar una solicitud al servidor:

python3 jax/bert/bert_request.py $EXTERNAL_IP

Resultado esperado:

For input "The capital of France is [MASK].", the result is ". the capital of france is paris.."
For input "Hello my name [MASK] Jhon, how can I [MASK] you?", the result is ". hello my name is jhon, how can i help you?."
Limpieza

Para limpiar recursos, ejecuta kubectl delete en orden inverso.

kubectl delete -f jax/bert/loadbalancer.yaml
kubectl delete -f jax/bert/serve-bert.yaml
kubectl delete -f jax/bert/install-bert.yaml

Ejecutar demostración de entrega de JAX Stable Diffusion E2E

En esta demostración, se usa el modelo de difusión estable previamente entrenado de Hugging Face.

Exporta el modelo guardado en TF2 compatible con TPU desde el modelo Flax Stable Diffusion

Para exportar los modelos de difusión estable, el clúster debe tener un grupo de nodos de CPU con un tipo de máquina que tenga más de 16 Gi de memoria disponible, como se describe en Clúster de configuración.

El Pod de Kubernetes ejecuta los siguientes pasos:

  1. Descarga y usa la secuencia de comandos de Python export_stable_diffusion_model.py desde los recursos de ejemplo para descargar el modelo de difusión estable previamente entrenado en un directorio temporal.
  2. Usa la imagen del convertidor de Cloud TPU para convertir el modelo previamente entrenado de CPU a TPU y almacena el modelo en el bucket de Cloud Storage que creaste durante la configuración del almacenamiento.

Este Pod de Kubernetes está configurado para ejecutarse en el grupo de nodos de CPU predeterminado. Ejecuta el Pod con el siguiente comando:

kubectl apply -f jax/stable-diffusion/install-stable-diffusion.yaml

Sigue estos pasos para verificar que el modelo se haya instalado correctamente:

kubectl get pods install-stable-diffusion

La STATUS puede tardar unos minutos en leer Completed.

Inicia el contenedor del servidor de modelos de TF para el modelo

Las cargas de trabajo de ejemplo se configuraron con las siguientes suposiciones:

  • El clúster se ejecuta con un grupo de nodos de TPU v5 con tres nodos.
  • el grupo de nodos usa el tipo de máquina ct5lp-hightpu-1t, en el que:
    • la topología es 1 × 1
    • la cantidad de chips TPU es 1

Si usas una configuración de clúster diferente a la descrita anteriormente, deberás editar el manifiesto de implementación del servidor.

Aplica la implementación:

kubectl apply -f jax/stable-diffusion/serve-stable-diffusion.yaml

Verifica que el servidor se ejecute como se espera:

kubectl get deployment stable-diffusion-deployment

AVAILABLE puede tardar un minuto en leer 3.

Aplica el servicio del balanceador de cargas:

kubectl apply -f jax/stable-diffusion/loadbalancer.yaml

Verifica que el balanceador de cargas esté listo para el tráfico externo con el siguiente comando:

kubectl get svc tf-stable-diffusion-service

Es posible que EXTERNAL_IP tarde unos minutos en tener una IP en la lista.

Envía la solicitud al servidor de modelos

Obtén una IP externa del balanceador de cargas:

EXTERNAL_IP=$(kubectl get services tf-stable-diffusion-service --output jsonpath='{.status.loadBalancer.ingress[0].ip}')

Ejecuta una secuencia de comandos para enviar una solicitud al servidor

python3 jax/stable-diffusion/stable_diffusion_request.py $EXTERNAL_IP

Resultado esperado:

El mensaje es Painting of a squirrel skating in New York y la imagen de salida se guardará como stable_diffusion_images.jpg en tu directorio actual.

Limpieza

Para limpiar recursos, ejecuta kubectl delete en orden inverso.

kubectl delete -f jax/stable-diffusion/loadbalancer.yaml
kubectl delete -f jax/stable-diffusion/serve-stable-diffusion.yaml
kubectl delete -f jax/stable-diffusion/install-stable-diffusion.yaml

Ejecuta la demostración de entrega de TensorFlow ResNet-50 E2E:

Instala dependencias de Python para ejecutar secuencias de comandos de Python de instructivos que envíen solicitudes al servicio de modelo de TF.

pip install -r tf/resnet50/requirements.txt
Paso 1: Convierte el modelo

Aplicar la conversión de modelos:

kubectl apply -f tf/resnet50/model-conversion.yml

Sigue estos pasos para verificar que el modelo se haya instalado correctamente:

kubectl get pods resnet-model-conversion

La STATUS puede tardar unos minutos en leer Completed.

Paso 2: Entregar el modelo con TensorFlow Serving

Aplica la implementación de la entrega del modelo:

kubectl apply -f tf/resnet50/deployment.yml

Usa el siguiente comando para verificar que el servidor se ejecute como se espera:

kubectl get deployment resnet-deployment

AVAILABLE puede tardar un minuto en leer 3.

Aplica el servicio del balanceador de cargas:

kubectl apply -f tf/resnet50/loadbalancer.yml

Verifica que el balanceador de cargas esté listo para el tráfico externo con el siguiente comando:

kubectl get svc resnet-service

Es posible que EXTERNAL_IP tarde unos minutos en tener una IP en la lista.

Paso 3: Envía una solicitud de prueba al servidor de modelos

Obtén la IP externa del balanceador de cargas:

EXTERNAL_IP=$(kubectl get services resnet-service --output jsonpath='{.status.loadBalancer.ingress[0].ip}')

Ejecuta la secuencia de comandos de solicitud de prueba (HTTP) para enviar la solicitud al servidor de modelos.

python3 tf/resnet50/request.py --host $EXTERNAL_IP

La respuesta debería verse de la siguiente manera:

Predict result: ['ImageNet ID: n07753592, Label: banana, Confidence: 0.94921875',
'ImageNet ID: n03532672, Label: hook, Confidence: 0.0223388672', 'ImageNet ID: n07749582,
Label: lemon, Confidence: 0.00512695312
Paso 4: Limpieza

Para limpiar recursos, ejecuta los siguientes comandos de kubectl delete:

kubectl delete -f tf/resnet50/loadbalancer.yml
kubectl delete -f tf/resnet50/deployment.yml
kubectl delete -f tf/resnet50/model-conversion.yml

Asegúrate de borrar el grupo de nodos y el clúster de GKE cuando termines de usarlos.

Inferencia y entrega de modelos de PyTorch

Instala dependencias de Python para ejecutar secuencias de comandos de Python instructivos que envíen solicitudes al servicio de modelos de PyTorch:

pip install -r pt/densenet161/requirements.txt

Ejecuta la demostración de entrega TorchServe Densenet161 E2E:

  1. Generar archivo de modelo.

    1. Aplicar archivo de modelo:
    kubectl apply -f pt/densenet161/model-archive.yml
    1. Sigue estos pasos para verificar que el modelo se haya instalado correctamente:
    kubectl get pods densenet161-model-archive

    La STATUS puede tardar unos minutos en leer Completed.

  2. Entregar el modelo con TorchServe:

    1. Aplica la implementación de la entrega del modelo:

      kubectl apply -f pt/densenet161/deployment.yml

    2. Usa el siguiente comando para verificar que el servidor se ejecute como se espera:

      kubectl get deployment densenet161-deployment

      AVAILABLE puede tardar un minuto en leer 3.

    3. Aplica el servicio del balanceador de cargas:

      kubectl apply -f pt/densenet161/loadbalancer.yml

      Verifica que el balanceador de cargas esté listo para el tráfico externo con el siguiente comando:

      kubectl get svc densenet161-service

      Es posible que EXTERNAL_IP tarde unos minutos en tener una IP en la lista.

  3. Envía una solicitud de prueba al servidor de modelos:

    1. Obtén una IP externa del balanceador de cargas:

      EXTERNAL_IP=$(kubectl get services densenet161-service --output jsonpath='{.status.loadBalancer.ingress[0].ip}')

    2. Ejecuta la secuencia de comandos de solicitud de prueba para enviar la solicitud (HTTP) al servidor de modelos:

      python3 pt/densenet161/request.py --host $EXTERNAL_IP

      Deberías ver una respuesta como la siguiente:

      Request successful. Response: {'tabby': 0.47878125309944153, 'lynx': 0.20393909513950348, 'tiger_cat': 0.16572578251361847, 'tiger': 0.061157409101724625, 'Egyptian_cat': 0.04997897148132324

  4. Para limpiar recursos, ejecuta los siguientes comandos de kubectl delete:

    kubectl delete -f pt/densenet161/loadbalancer.yml
kubectl delete -f pt/densenet161/deployment.yml
kubectl delete -f pt/densenet161/model-archive.yml

    Asegúrate de borrar el grupo de nodos y el clúster de GKE cuando termines.

Cómo solucionar problemas comunes

Puedes encontrar información sobre la solución de problemas de GKE en Cómo solucionar problemas de TPU en GKE.

No se pudo inicializar la TPU

Si te encuentras con el siguiente error, asegúrate de que estás ejecutando el contenedor de TPU en modo privilegiado o de haber aumentado el ulimit dentro del contenedor. Para obtener más información, consulta Ejecución sin modo privilegiado.

TPU platform initialization failed: FAILED_PRECONDITION: Couldn't mmap: Resource
temporarily unavailable.; Unable to create Node RegisterInterface for node 0,
config: device_path:      "/dev/accel0" mode: KERNEL debug_data_directory: ""
dump_anomalies_only: true crash_in_debug_dump: false allow_core_dump: true;
could not create driver instance

Programación de interbloqueo

Supongamos que tienes dos trabajos (trabajo A y trabajo B) y ambos deben programarse en porciones de TPU con una topología de TPU determinada (por ejemplo, v4-32). Supongamos también que tienes dos porciones de TPU v4-32 dentro del clúster de GKE; las llamaremos porción X e Y porción Y. Dado que tu clúster tiene capacidad amplia para programar ambos trabajos, en teoría, ambos deben programarse con rapidez: uno en cada una de las dos porciones v4-32 de TPU.

Sin embargo, sin una planificación cuidadosa, es posible entrar en un interbloqueo de la programación. Supongamos que el programador de Kubernetes programa un Pod de Kubernetes del trabajo A en la porción X y, luego, programa un pod de Kubernetes del trabajo B en la porción X. En este caso, dadas las reglas de afinidad de Pods de Kubernetes para el trabajo A, el programador intentará programar todos los Pods de Kubernetes restantes para el trabajo A en la porción X. Lo mismo para el trabajo B. Por lo tanto, ni el trabajo A ni el trabajo B podrán programarse por completo en una sola porción. El resultado será un interbloqueo de programación.

Para evitar el riesgo de un interbloqueo de programación, puedes usar la antiafinidad del Pod de Kubernetes con cloud.google.com/gke-nodepool como topologyKey, como se muestra en el siguiente ejemplo:

apiVersion: batch/v1
kind: Job
metadata:
 name: pi
spec:
 parallelism: 2
 template:
   metadata:
     labels:
       job: pi
   spec:
     affinity:
       podAffinity:
         requiredDuringSchedulingIgnoredDuringExecution:
         - labelSelector:
             matchExpressions:
             - key: job
               operator: In
               values:
               - pi
           topologyKey: cloud.google.com/gke-nodepool
       podAntiAffinity:
         requiredDuringSchedulingIgnoredDuringExecution:
         - labelSelector:
             matchExpressions:
             - key: job
               operator: NotIn
               values:
               - pi
           topologyKey: cloud.google.com/gke-nodepool
           namespaceSelector:
             matchExpressions:
             - key: kubernetes.io/metadata.name
               operator: NotIn
               values:
               - kube-system
     containers:
     - name: pi
       image: perl:5.34.0
       command: ["sleep",  "60"]
     restartPolicy: Never
 backoffLimit: 4

Crea recursos del grupo de nodo TPU con Terraform

También puedes usar Terraform para administrar tus recursos de clúster y grupo de nodos.

Crear un grupo de nodos de porciones de TPU de varios hosts en un clúster de GKE existente

Si tienes un clúster existente en el que deseas crear un grupo de nodos TPU de hosts múltiples, puedes usar el siguiente fragmento de Terraform:

resource "google_container_cluster" "cluster_multi_host" {
  …
  release_channel {
    channel = "RAPID"
  }
  workload_identity_config {
    workload_pool = "my-gke-project.svc.id.goog"
  }
  addons_config {
    gcs_fuse_csi_driver_config {
      enabled = true
    }
  }
}

resource "google_container_node_pool" "multi_host_tpu" {
  provider           = google-beta
  project            = "${project-id}"
  name               = "${node-pool-name}"
  location           = "${location}"
  node_locations     = ["${node-locations}"]
  cluster            = google_container_cluster.cluster_multi_host.name
  initial_node_count = 2

  node_config {
    machine_type = "ct4p-hightpu-4t"
    reservation_affinity {
      consume_reservation_type = "SPECIFIC_RESERVATION"
      key = "compute.googleapis.com/reservation-name"
      values = ["${reservation-name}"]
    }
    workload_metadata_config {
      mode = "GKE_METADATA"
    }
  }

  placement_policy {
    type = "COMPACT"
    tpu_topology = "2x2x2"
  }
}

Reemplaza los siguientes valores:

  • your-project: Es el proyecto de Google Cloud en el que ejecutas la carga de trabajo.
  • your-node-pool: Es el nombre del grupo de nodos que creas.
  • us-central2: Es la región en la que ejecutas la carga de trabajo.
  • us-central2-b: Es la zona en la que ejecutas la carga de trabajo.
  • your-reservation-name: Es el nombre de la reserva.

Crea un grupo de nodos de porción de TPU de host único en un clúster de GKE existente

Usa el siguiente fragmento de Terraform:

resource "google_container_cluster" "cluster_single_host" {
  …
  cluster_autoscaling {
    autoscaling_profile = "OPTIMIZE_UTILIZATION"
  }
  release_channel {
    channel = "RAPID"
  }
  workload_identity_config {
  workload_pool = "${project-id}.svc.id.goog"
  }
  addons_config {
    gcs_fuse_csi_driver_config {
      enabled = true
    }
  }
}

resource "google_container_node_pool" "single_host_tpu" {
  provider           = google-beta
  project            = "${project-id}"
  name               = "${node-pool-name}"
  location           = "${location}"
  node_locations     = ["${node-locations}"]
  cluster            = google_container_cluster.cluster_single_host.name
  initial_node_count = 0
  autoscaling {
    total_min_node_count = 2
    total_max_node_count = 22
    location_policy      = "ANY"
  }

  node_config {
    machine_type = "ct4p-hightpu-4t"
    workload_metadata_config {
      mode = "GKE_METADATA"
    }
  }
}

Reemplaza los siguientes valores:

  • your-project: Es el proyecto de Google Cloud en el que ejecutas la carga de trabajo.
  • your-node-pool: Es el nombre del grupo de nodos que creas.
  • us-central2: Es la región en la que ejecutas la carga de trabajo.
  • us-central2-b: Es la zona en la que ejecutas la carga de trabajo.