Se usó la API de Cloud Translation para traducir esta página.

Introducción a las TPU en GKE

Los clientes de Google Kubernetes Engine (GKE) ahora pueden crear grupos de nodos de Kubernetes que contengan porciones TPU v4 y v5e. 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 nodos de TPU:

Porción de TPU de varios hosts
Porción de TPU de host único

Grupo de nodos de porción de TPU de varios hosts

Un grupo de nodos de porción de TPU de varios hosts es un grupo de nodos que contiene dos o más con VMs de TPU interconectadas. Cada VM tiene un dispositivo de TPU conectado. Las TPUs 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 TPU de varios hosts, no puedes agregar nodos a él. Por ejemplo: no puedes crear un grupo de nodos v4-32 y, luego, agregar un Kubernetes (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 los nodos de la porción.

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

En el siguiente diagrama, se muestra un ejemplo de una TPU de varios hosts v5litepod-16 (v5e) porción. 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 la porción de TPU de varios hosts

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

Diagrama del pod de TPU v5e

Para ver un ejemplo de cómo ejecutar una carga de trabajo en una porción de TPU de varios hosts, consulta Ejecuta tu carga de trabajo en TPUs.

Grupos de nodos de porción 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 VMs de TPU independientes. Cada VM tiene un dispositivo TPU conectado. Si bien las VMs dentro de un grupo de nodos de porción de host único pueden comunicarse a través de la red del centro de datos (DCN), las TPU conectadas a las VMs 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 cómo ejecutar una carga de trabajo en una porción de TPU de host único, consulta Ejecuta tu carga de trabajo en TPUs.

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 del porción de TPU que requiere tu carga de trabajo. TPU v4 es compatible con la versión estándar de GKE 1.26.1-gke.1500 y versiones posteriores, v5e en la versión estándar de GKE 1.27.2-gke.2100 y versiones posteriores, y v5p en la versión estándar de GKE 1.28.3-gke.1024000 y versiones posteriores.

Las 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 creas un grupo de nodos de TPU, selecciona un tamaño de porción de TPU (una topología de TPU) según el tamaño de tu modelo y la cantidad de memoria que requiere. El tipo de máquina que especifiques cuando crees tus grupos de nodos depende de la versión y el tamaño de tus slices.

v5e

A continuación, se muestran los tipos de máquinas y las topologías de TPU v5e que son compatibles con los casos de uso 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	Entrenamiento, inferencia de varios hosts
`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 del plano de control 1.27.2-gke.2100 o una posterior. GKE Autopilot debe ejecutar la versión 1.29.2-gke.1521000 o posterior del plano de control. Más información acerca de 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	Alta	Medio	Low (Baja)

Para liberar espacio para las VMs con más chips, el programador de GKE puede adelantar y reprogramar las VMs con menos chips. Por lo tanto, las VMs de 8 chips tienen más probabilidades de usurpar las VMs de 1 y 4 chips.

v4 y v5p

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

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 TPU v4 están disponibles en us-central2-b. Tu Los clústeres de GKE Standard deben ejecutar el plano de control versión 1.26.1-gke.1500 o 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 TPU v5p están disponibles en us-west4-a y us-east5-a. GKE Standard Los clústeres 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 versiones posteriores. 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.
Sin asistencia para la asignación de costos: Asignación de costos de GKE y medición de uso no incluyen datos sobre el uso ni los costos de las TPU.
El escalador automático puede calcular la capacidad: El escalador automático del clúster podría calcular. incorrecta para los nodos nuevos que contengan VMs de TPU antes de que disponibles. Luego, 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á la escala verticalmente de los nodos adicionales, si no son necesarios, después de la operación de reducción de escala vertical regular.
El escalador automático cancela el escalamiento vertical: El escalador automático de clústeres cancela el escalamiento vertical de los grupos de nodos TPU que permanecen en estado de espera durante más de 10 horas. Clúster El escalador automático reintentará realizar esas operaciones de escalamiento vertical más adelante. Es posible que este comportamiento y reduce la capacidad de obtener TPU para los clientes que no usan reservas.
El taint puede evitar la reducción de escala: Cargas de trabajo no basadas en TPU que tienen tolerancia para El taint de TPU puede evitar la reducción de escala vertical del grupo de nodos si se vuelven a crear 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 los recursos.

Las siguientes cuotas tienen valores predeterminados que probablemente deban aumentarse:

Cuota de SSD de Persistent Disk (GB): El disco de arranque de cada nodo de Kubernetes requiere 100 GB de forma predeterminada. Por lo tanto, esta cuota debe establecerse al menos como alta como (la cantidad máxima de nodos de GKE que prevés creando) * 100 GB.
Cuota de direcciones IP en uso: Cada nodo de Kubernetes consume una dirección IP. Por lo tanto, esta cuota debe establecerse al menos tan alta como la cantidad máxima de nodos de GKE que esperas 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 dificultades para que se aprueben tus solicitudes de aumento de cuota en un plazo de unos días, comunícate con tu 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 dirígete a Crea un clúster de Google Kubernetes Engine.

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

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

La capacidad de TPU migrada al nuevo sistema de reserva basado en Compute Engine no se puede usar con la API de recursos en cola de Cloud TPU. Si deseas usar recursos de TPU en cola con tu reserva, deberás migrar una parte de tu reserva de TPU al nuevo sistema de reserva basado en Compute Engine.
No puede haber cargas de trabajo en ejecución activa 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 cuenta de Google Cloud. con tu equipo de cuentas para programar la migración. El período de migración se debe durante el horario de atención (de lunes a viernes, de 9:00 a.m. a 5:00 p.m., hora del Pacífico).

Cree 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 Cómo crear un grupo de nodos en la documentación de Google Kubernetes Engine.

Ejecución sin modo de privilegio

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

Ejecuta cargas de trabajo en grupos de nodos TPU

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

Selectores de nodos

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

Establecer cloud.google.com/gke-accelerator-type en tpu-v5-lite-podslice, tpu-v5p-slice 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 manifiestos de ejemplo 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 una programación especial de la carga de trabajo y la administración de cambios 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 múltiples 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: Ve más allá de un solo fragmento

Puedes agregar porciones más pequeñas en una multislice para controlar cargas de trabajo de entrenamiento. 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 en pokémones
PyTorch/XLA: GPT2 en WikiText

Descarga recursos de instructivos

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

Crear y conectarse al clúster

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

El siguiente comando crea un nuevo clúster regional de GKE suscrito al canal de versiones rápido con un grupo de nodos que inicialmente contiene un nodo por zona. El comando también habilita la federación de identidades para cargas de trabajo para GKE y las funciones del controlador de CSI de Cloud Storage FUSE en tu 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

Habilitar la Federación de identidades para cargas de trabajo para GKE y las funciones del controlador de CSI del FUSE de Cloud Storage para en 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: Entrena la difusión en Pokémon

En este ejemplo, se entrena el modelo de difusión estable de HuggingFace con el conjunto de datos de Pokémon.

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

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 la configuración permisos para que Docker 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 TPUs v5e con HuggingFace en PyTorch/XLA con el conjunto de datos de 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 la configuración permisos para que Docker 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 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

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

Crea un bucket de Cloud Storage y configura el acceso a él. Se usa un El bucket de Cloud Storage se usa para almacenar el modelo previamente entrenado.
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 convertidor 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.
Inicia el servidor para el modelo convertido. Aplica una implementación que entregue el modelo con un framework de servidor respaldado por el volumen almacenado en el volumen persistente ReadOnlyMany (ROX). Las réplicas de implementación se deben ejecutar en un grupo de nodos de slice v5e con un Pod de Kubernetes por nodo.
Implementar un balanceador de cargas para probar el servidor de modelos El servidor se expone a solicitudes externas con el servicio de LoadBalancer. Se proporcionó una secuencia de comandos de Python con una solicitud de ejemplo para probar la en el servidor de modelos.

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 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 implementación de servidor de host único.

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 la replicas a 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 es necesario ejecutar Pods de Kubernetes en modo privilegiado en la versión 1.28 de GKE o versiones posteriores. Para obtener más información, consulta Ejecuta contenedores sin modo privilegiado.

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

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

Configuración

Descarga las secuencias de comandos de Python y los manifiestos YAML del instructivo 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 usas en este instructivo requieren la versión 3.9 o una posterior. 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 en este instructivo.

Configura el clúster

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

Crea el nodo TPU de host único.

En las cargas de trabajo de ejemplo, se supone 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 que se describió 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 agregando 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: El tipo de máquina de los nodos que se crearán en tu grupo de nodos.

Configura el almacenamiento de modelos

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 un 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 respaldadas 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 de tu bucket de Cloud Storage. Para mejorar el rendimiento, especifica la ubicación donde de que tu clúster de GKE esté en ejecución.

Sigue estos pasos para otorgarle a tu clúster de GKE acceso en el bucket. Para simplificar la configuración, en los siguientes ejemplos se usa la configuración espacio de nombres y la cuenta de servicio predeterminada de Kubernetes. Si deseas obtener más información, consulta Configura el acceso a los buckets de Cloud Storage a través de la federación de identidades para cargas de trabajo de GKE para GKE.

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.

Nota: Si usas una cuenta de servicio de IAM existente, omite este paso.
```
gcloud iam service-accounts create your-iam-service-acct \
--project=your-bucket-project-id
```
Reemplaza lo siguiente:
- your-iam-service-acct: Es el nombre del servicio de IAM nuevo. de servicio predeterminada.
- your-bucket-project-id: Es el ID del proyecto en el que creaste tu cuenta de servicio de IAM. La cuenta de servicio de IAM debe estar en el mismo proyecto que tu bucket de Cloud Storage.
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 del servicio de IAM nuevo. de servicio predeterminada.
- your-bucket-project-id: El ID del proyecto en el que creaste el Cuenta de servicio de IAM.
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. Esta vinculación permite que la cuenta de servicio de Kubernetes actúe 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 del servicio de IAM nuevo. de servicio predeterminada.
- your-bucket-project-id: El ID del proyecto en el que creaste el Cuenta de servicio de IAM.
- your-project-id: Es el ID del proyecto en el que creaste tu clúster de GKE. Tus buckets de Cloud Storage y tu clúster de GKE pueden estar en el mismo proyecto o en proyectos diferentes.
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 del servicio de IAM nuevo. de servicio predeterminada.
- your-bucket-project-id: Es el ID del proyecto en el que creaste tu cuenta de servicio de IAM.
Ejecuta el siguiente comando para completar el nombre de tu 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 depósito de Cloud Storage.
Crea el volumen persistente y la reclamación de volumen persistente con el siguiente comando:
```
kubectl apply -f pvc-pv.yaml
```

Inferencia y entrega de modelos de JAX

Instala las dependencias de Python para ejecutar secuencias de comandos de Python del instructivo que envían solicitudes al servicio de modelos de JAX.

pip install -r jax/requirements.txt

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

Esta demostración usa un modelo BERT previamente entrenado. de Hugging Face.

El Pod de Kubernetes realiza los siguientes pasos:

Descarga y usa la secuencia de comandos Python export_bert_model.py del ejemplo. para descargar el modelo BERT previamente entrenado en un directorio temporal.
Usa la imagen del convertidor de Cloud TPU para convertir el modelo previamente entrenado de CPU a TPU y almacenarlo en el bucket de Cloud Storage que creaste durante la configuración.

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

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 que se describió 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 mostrar una IP.

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 los 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 dispersión estable previamente entrenado de Hugging Face.

Exporta un modelo guardado de TF2 compatible con TPU desde el modelo de difusión estable de Flax

Para exportar los modelos de difusión estables, 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 Cómo configurar un clúster.

El Pod de Kubernetes ejecuta los siguientes pasos:

Descarga y usa la secuencia de comandos de Python export_stable_diffusion_model.py de los recursos de ejemplo para descargar el modelo de difusión estable previamente entrenado en un directorio temporal.
Usa la imagen del convertidor de Cloud TPU para convertir el modelo previamente entrenado de CPU a TPU y almacenarlo en el bucket de Cloud Storage que creaste durante la configuración de 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 que se describió 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 según lo esperado:

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 lo siguiente:

kubectl get svc tf-stable-diffusion-service

Es posible que EXTERNAL_IP tarde unos minutos en mostrar una IP.

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 los 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 instructivos que envíen al servicio tf.model.

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

Verifica que el modelo se haya instalado correctamente con lo siguiente:

kubectl get pods resnet-model-conversion

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

Verifica que el servidor se ejecute según lo esperado con el siguiente comando:

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 lo siguiente:

kubectl get svc resnet-service

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

Paso 3: Envía la 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 ser similar a la siguiente:

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 las dependencias de Python para ejecutar secuencias de comandos de Python del instructivo que envían solicitudes al servicio de modelos de PyTorch:

pip install -r pt/densenet161/requirements.txt

Ejecuta la demostración de entrega de extremo a extremo de TorchServe Densenet161:

Generar archivo de modelo.
1. Aplica el archivo del 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
```
STATUS puede tardar unos minutos en leer Completed.
Entrega el modelo con TorchServe:
1. Aplica la implementación de la entrega de modelos:
```
kubectl apply -f pt/densenet161/deployment.yml
```
2. Verifica que el servidor se ejecute según lo esperado con el siguiente comando:
```
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 mostrar una IP.

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

Obtén la IP externa del balanceador de cargas:

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

Ejecuta la secuencia de comandos de solicitud de prueba para enviar una 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

Ejecuta los siguientes comandos kubectl delete para limpiar los recursos:
```
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 de GKE. y clúster cuando no hayan terminado con ellos.

Cómo solucionar problemas comunes

Puedes encontrar información para solucionar problemas de GKE en Soluciona problemas de TPU en GKE.

No se pudo inicializar la TPU

Si encuentras el siguiente error, asegúrate de que estás ejecutando tu TPU contenedor en modo privilegiado o aumentaste el ulimit en tu contenedor. Para obtener más información, consulta Ejecución sin modo de privilegio.

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 la TPU. porciones con una topología de TPU determinada (por ejemplo, v4-32). También supongamos que tienes dos porciones de TPU v4-32 dentro del clúster de GKE; vamos a llamamos a esas porción X y a la porción Y. Dado que tu clúster tiene una capacidad suficiente para programar ambos trabajos, en teoría, ambos trabajos deberían programarse rápidamente: un trabajo en cada una de las dos porciones de TPU v4-32.

Sin embargo, sin una planificación cuidadosa, es posible que se produzca un bloqueo de programación. Supongamos que el programador de Kubernetes programa un Pod de Kubernetes desde Job. 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 sucede con 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 en la programación, puedes usar el Pod de Kubernetes antiafinidad 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 los recursos del clúster y el grupo de nodos.

Crea un grupo de nodos de porción 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 de TPU de varios hosts, 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 tu carga de trabajo.
us-central2-b: Es la zona en la que ejecutas tu carga de trabajo.
your-reservation-name: el nombre de tu 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 tu carga de trabajo.
us-central2-b: Es la zona en la que ejecutas tu carga de trabajo.