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 pods TPU v4 y v5e. Un Pod de TPU es un grupo de dispositivos de TPU conectados por interconexiones de alta velocidad. Esto es diferente de un Pod de Kubernetes, que es la unidad de procesamiento implementable más pequeña que puedes crear y administrar en Kubernetes. Para las cargas de trabajo que no requieren un pod de TPU completo, puedes usar un subconjunto de un pod de TPU completo llamado porción de TPU. Al igual que los Pods de TPU completos, cada dispositivo de TPU en una porción tiene su propia VM de TPU. Nos referimos a una VM de TPU y a su dispositivo conectado como un host o un nodo TPU. Para obtener más información sobre los pods de TPU, consulta Arquitectura del sistema.

Dado que el término pod que se usa en el contexto de GKE, por lo general, hace referencia a un pod de Kubernetes, para evitar confusiones, siempre nos referiremos a una colección de uno o más dispositivos de TPU como una porción.

Cuando trabajas con GKE, primero debes crear un clúster de GKE.

Luego, agregas grupos de nodos a tu clúster. Los grupos de nodos de GKE son colecciones de VMs que comparten los mismos atributos. En las cargas de trabajo de TPU, los grupos de nodos constan de VMs de TPU.

Tipos de grupos de nodos

GKE admite dos tipos de grupos de nodo TPU:

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

Grupo de nodos de porciones 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 VM de TPU interconectadas. Cada VM tiene un dispositivo de TPU conectado a ella. Las TPU en una porción de varios hosts están conectadas a través de una interconexión de alta velocidad (ICI). Un grupo de nodos de porción de porción de TPU de varios hosts es inmutable. 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 en la porción, ninguno de los nodos de la porción se podrá implementar.

Si es necesario reparar un nodo dentro de una porción de TPU de varios hosts, GKE cerrará todas las VM de TPU (nodos) de la porción y forzará la expulsión de todos los Pods de Kubernetes de 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 de la porción nueva.

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 porción de TPU de varios hosts

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

Diagrama de 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 la carga de trabajo en una porción de TPU de varios hosts.

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 una de estas VMs tiene un dispositivo de TPU conectado a ellas. Si bien las VM 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 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 tus cargas de trabajo en nodos 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 la 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) según el 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.

Versión 5e

A continuación, se indican los tipos y las topologías de máquinas de TPU v5e que se admiten para casos de uso de inferencia y entrenamiento:

Tipo de máquina	Topología	Cantidad de chips TPU	Cantidad de VM	Caso práctico recomendado
`ct5lp-hightpu-1t`	1x1	1	1	Entrenamiento, inferencia de host único
`ct5lp-hightpu-4t`	2x2	4	1	Entrenamiento, inferencia de host único
`ct5lp-hightpu-8t`	2x4	8	1	Entrenamiento, inferencia de 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 y de 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 aceleradores de entrenamiento v5e y Tipos de aceleradores 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 1.29.2-gke.1521000 del plano de control 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	Alta	Media	Baja

A fin de liberar espacio para las VM con más chips, el programador de GKE puede interrumpir y reprogramar las VM con menos chips. Es probable que las VMs de 8 chips interrumpen las VMs de 1 y 4 chips.

v4 y v5p

Los siguientes son los tipos de máquinas de 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 de sistemas de TPU para obtener más información. Los tipos de máquinas de Pods 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 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 Pods 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 una versión posterior. Para obtener más información sobre v5p, consulta Introducción a la capacitación de 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 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 h.
Sin compatibilidad con 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: El escalador automático del clúster puede calcular la capacidad de forma incorrecta para los nodos TPU nuevos 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 reduce verticalmente la escala de los nodos adicionales, si no son necesarios, después de una 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 por más de 10 horas. El escalador automático del clúster volverá a intentar las operaciones de escalamiento vertical más tarde. Este comportamiento puede reducir la disponibilidad de TPU para los clientes que no usan reservas.
Taint puede prevenir la reducción de escala: Las cargas de trabajo que no son de TPU y que tienen tolerancia para el taint de TPU pueden evitar la reducción de escala del grupo de nodos si se vuelven a crear durante el desvío del grupo de nodo TPU.

Asegúrate de tener cuotas suficientes 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 de disco persistente (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 tan alta como (la cantidad máxima de nodos de GKE que esperas crear) × 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 Cómo solicitar una cuota mayor. Para obtener más información sobre los tipos de cuotas de TPU, consulta Cuota de TPU.

Las solicitudes de aumento de cuota pueden tardar unos días en aprobarse. Si tienes dificultades para que tus solicitudes de aumento de cuota se aprueben dentro de unos 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 las 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 reserva basado en Compute Engine no se puede usar con la API de recursos en cola de Cloud TPU. Si pretendes usar recursos de TPU en cola con tu reserva, solo deberás migrar una parte de tu reserva de TPU al nuevo sistema de reservas basado en Compute Engine.
Ninguna carga de trabajo puede ejecutarse de forma activa en las TPU cuando se migra al nuevo sistema de reserva basado en Compute Engine.
Selecciona un momento para realizar la migración y trabaja con tu equipo de cuentas de Google Cloud a fin de programar la migración. El período de migración debe ser durante el horario de atención (de lunes a viernes, de 9:00 a.m. a 5:00 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.

Crear 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 deseas reducir el alcance del permiso en tu contenedor, consulta el modo privilegio de TPU.

Ejecuta cargas de trabajo en nodos TPU

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

Selectores de nodos

A fin de que Kubernetes programe tu carga de trabajo en nodos TPU, debes especificar dos selectores para cada nodo TPU 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 TPU.

Las secciones Cargas de trabajo de entrenamiento y Cargas de trabajo de inferencia contienen manifiestos de ejemplo en los que se ilustra 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 y administración especiales de las 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 nodo TPU

Si un nodo TPU en un grupo de nodos de porción de TPU de varios hosts está en mal estado, se 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 multiparte para controlar 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 entrenar 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: Entrenamiento de dispersión en Pokémon
PyTorch/XLA: GPT2 en WikiText

Descargar recursos del instructivo

Descarga las secuencias de comandos de Python del instructivo y las especificaciones de 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 a él

Crear un clúster de GKE regional para que el plano de control de Kubernetes se replique en tres zonas y proporcione 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 TPU y 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, en un principio, contiene un nodo por zona. El comando también habilita las funciones del controlador de CSI del FUSE de Cloud Storage y Workload Identity en el clúster porque las cargas de trabajo de inferencia de ejemplo de 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 deseas habilitar las funciones del controlador de CSI del FUSE de Cloud Storage 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 según 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: Entrenamiento de dispersión en Pokémon

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

El modelo de dispersión estable es un modelo latente de texto a imagen que genera imágenes 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/. Ejecuta los comandos siguientes para compilar y enviar 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 termine de ejecutarse tu trabajo, puedes borrarlo con lo siguiente:

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 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/. Ejecuta los siguientes comandos para compilar y enviar 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 el siguiente comando:

kubectl apply -f tpu_job_gpt.yaml

Limpieza

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

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 de GKE v5e para modelos previamente entrenados con JAX, TensorFlow y PyTorch. En un nivel alto, hay cuatro pasos separados para realizar en el clúster de GKE:

Crear un bucket de Cloud Storage y configurar el acceso al bucket Usarás un bucket de Cloud Storage para almacenar el modelo previamente entrenado.
Descarga y convierte un modelo previamente entrenado en un modelo 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.
Inicia el servidor para el modelo convertido. Aplica una Implementación que entregue el modelo mediante un framework de servidor respaldado por el volumen almacenado en el volumen persistente ReadOnly Many (ROX). Las réplicas de implementación deben ejecutarse en un nodo TPU de Pod v5e con un Pod de Kubernetes por nodo.
Implementar un balanceador de cargas para probar el servidor del modelo 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 de modelos.

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

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 ocurre lo siguiente:
- la topología es de 1 x 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 a la cantidad de nodos.

Si tu clúster Standard 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 que ejecutes Pods de Kubernetes en modo de privilegios 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.
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 usarás en este instructivo requieren la versión 3.9 o superior de Python. 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.

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 de la que se describió antes, 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 16 memorias Gi+ disponibles (por ejemplo, e2-standard-4). Esto se configura en el comando gcloud container clusters create o si agregas 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 al que se agregará el grupo de nodos.
your-machine-type: El tipo de máquina de los nodos que se crearán en el 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 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 otorgarle 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 información, consulta Configura el acceso a los buckets de Cloud Storage con Workload Identity de 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 de la cuenta de servicio de IAM nueva.
- your-bucket-project-id: 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.
Asegúrate de que tu cuenta de servicio de IAM tenga las funciones 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: El ID del proyecto en el que creaste la 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 cuentas de servicio. 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 de la cuenta de servicio de IAM nueva.
- your-bucket-project-id: 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 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 de la cuenta de servicio de IAM nueva.
- your-bucket-project-id: El ID del proyecto en el que creaste la cuenta de servicio de IAM.
Ejecuta el siguiente comando para propagar 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 bucket de Cloud Storage.
Crea el volumen persistente y la reclamación del 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 de instructivos que envían solicitudes al servicio del modelo JAX.

pip install -r jax/requirements.txt

Ejecuta la demostración de entrega 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:

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.
Se usa la imagen del convertidor de Cloud TPU para convertir el modelo previamente entrenado de CPU a TPU y se almacena 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 predeterminada. Ejecuta el Pod con el siguiente comando:

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

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

kubectl get pods install-bert

Es posible que STATUS tarde 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 de la que se describió antes, deberás editar el manifiesto de implementación del servidor.

Aplica 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 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

Ejecuta la 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.

Exportar el modelo guardado de TF2 compatible con TPU desde el modelo de Flax Stable Diffusion

La exportación de los modelos de difusión estables requiere que el clúster tenga un grupo de nodos de CPU con un tipo de máquina que tenga más de 16 Gi y memoria disponible, como se describe en Configura el 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.
Se usa la imagen del convertidor de Cloud TPU para convertir el modelo previamente entrenado de CPU a TPU y se almacena 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

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

kubectl get pods install-stable-diffusion

Es posible que STATUS tarde unos minutos en leer Completed.

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

Las cargas de trabajo de ejemplo se configuraron según 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 ocurre lo siguiente:
- la topología es de 1 x 1
- la cantidad de chips TPU es 1

Si usas una configuración de clúster diferente de la que se describió antes, 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 esté ejecutando como se espera:

kubectl get deployment stable-diffusion-deployment

AVAILABLE puede tardar un minuto en leer 3.

Aplica el servicio de 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}')

Ejecutar 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 las dependencias de Python para ejecutar secuencias de comandos de Python de instructivos que envían solicitudes al servicio de modelos de TF.

pip install -r tf/resnet50/requirements.txt

Paso 1: Convierte el modelo

Aplica la conversión del modelo:

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

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

kubectl get pods resnet-model-conversion

Es posible que STATUS tarde unos minutos en leer Completed.

Paso 2: Entrega el modelo con TensorFlow Serving

Aplica la implementación de entrega de modelos:

kubectl apply -f tf/resnet50/deployment.yml

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

kubectl get deployment resnet-deployment

AVAILABLE puede tardar un minuto en leer 3.

Aplica el servicio de 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 del modelo

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 del modelo.

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 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 de instructivos que envían solicitudes al servicio de modelos de PyTorch:

pip install -r pt/densenet161/requirements.txt

Ejecuta la demostración de entrega de TorchServe Densenet161 E2E:

Genera un archivo de modelo.
1. Aplica el archivo de modelo:
```
kubectl apply -f pt/densenet161/model-archive.yml
```
1. Verifica que el modelo se haya instalado correctamente con el siguiente comando:
```
kubectl get pods densenet161-model-archive
```
Es posible que STATUS tarde unos minutos en leer Completed.
Entrega el modelo con TorchServe:
1. Aplica la implementación de entrega de modelos:
```
kubectl apply -f pt/densenet161/deployment.yml
```
2. Verifica que el servidor se esté ejecutando como se espera con el siguiente comando:
```
kubectl get deployment densenet161-deployment
```
  AVAILABLE puede tardar un minuto en leer 3.
3. Aplica el servicio de 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.

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

Obtén una 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 del modelo:

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 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 y el clúster de GKE cuando termines de usarlos.

Cómo solucionar problemas comunes

Puedes encontrar información sobre la solución de problemas de GKE en Soluciona problemas de TPU en GKE.

No se pudo inicializar la TPU

Si encuentras el siguiente error, asegúrate de estar 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 que ambos se deben programar en porciones de TPU 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. Las llamaremos a esas porciones X y Y. Dado que tu clúster tiene una capacidad suficiente para programar ambos trabajos, en teoría ambos deben programarse con rapidez: un trabajo en cada una de las dos porciones v4-32 de TPU.

Sin embargo, sin una planificación cuidadosa, es posible llegar a un interbloqueo de 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 del Pod 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 de grupos de nodo TPU con Terraform

También puedes usar Terraform para administrar los recursos de tu clúster y 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 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 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.