Introducción a las TPUs en GKE

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

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

Tipos de grupos de nodos

GKE admite dos tipos de grupos de nodo TPU:

Grupo de nodos de 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 VMs de TPU interconectadas. Cada VM tiene un dispositivo 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 slice de varios hosts, no puedes agregar nodos a él. Por ejemplo, no puedes crear un grupo de nodos v4-32 y, luego, agregar un nodo 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 de 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, no se implementará ningún nodo en 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 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 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 TPU para grupos de nodos de GKE

Antes de crear grupos de nodos, debes elegir la versión y el tamaño de la 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 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 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 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 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 los tipos de aceleradores de entrenamiento v5e y los tipos de aceleradores de inferencia v5e.

Las máquinas 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. 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 Medio Low (Baja)

Para liberar espacio para las VMs con más chips, el programador de GKE puede apropiarse de las VMs con menos chips y reprogramarlas. Por lo tanto, las VMs de 8 chips tienen más probabilidades de interrumpir 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 de 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. Tus clústeres de GKE Standard deben ejecutar la versión 1.26.1-gke.1500 o posterior del plano de control. Los clústeres de GKE Autopilot deben ejecutar la versión 1.29.2-gke.1521000 o posterior del plano de control.

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 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 o una posterior del plano de control. GKE Autopilot debe ejecutar 1.29.2-gke.1521000 o versiones posteriores. Para obtener más información sobre v5p, consulta la introducción al entrenamiento 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 las TPU interrumpibles: Las VMs Spot son similares a las VMs interrumpibles y están sujetas a las mismas limitaciones de disponibilidad, pero no tienen una duración máxima de 24 horas.
  • No se admite la asignación de costos: La asignación de costos de GKE y la medición del uso no incluyen datos sobre el uso o los costos de las TPU.
  • El escalador automático puede calcular la capacidad: El escalador automático del clúster podría calcular de forma incorrecta la capacidad de los nodos nuevos que contienen VMs de TPU antes de que esos nodos estén 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. El escalador automático del clúster volverá a intentar esas operaciones de escalamiento vertical más adelante. Este comportamiento puede reducir la capacidad de obtención de TPU para los clientes que no usan reservas.
  • El taint puede impedir la reducción de la escala vertical: Las cargas de trabajo que no sean de TPU y tengan tolerancia al taint de TPU pueden impedir la reducción de la escala vertical del grupo de nodos si se vuelven a crear durante el vaciado 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 donde se crean tus recursos.

Las siguientes cuotas tienen valores predeterminados que es probable que 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 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 Solicita una cuota más alta. Para obtener más información sobre los tipos de cuotas de TPU, consulta Cuota de TPU.

Es posible que las solicitudes de aumento de cuota tarden algunos días en aprobarse. 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 TPUs en GKE, omite esta sección y ve a Cómo crear 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 se pueden ejecutar cargas de trabajo de forma activa en las TPU cuando se migran al nuevo sistema de reserva basado en Compute Engine.
  • Selecciona un horario para realizar la migración y trabaja con tu equipo de la cuenta de Google Cloud para programarla. 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).

Cree un clúster de Google Kubernetes Engine

Consulta Cómo crear 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 deseas reducir el permiso de permisos en tu contenedor, consulta Modo de privilegio de TPU.

Ejecuta cargas de trabajo en grupos de nodo 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 carga de trabajo en el manifiesto de Google Kubernetes Engine:

  • Establece cloud.google.com/gke-accelerator-type como 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 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 nodos

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

Multislice: Ve más allá de un solo fragmento

Puedes hacer una agregación de porciones más pequeñas que formen una porción múltiple para manejar cargas de trabajo de entrenamiento más grandes. Para obtener más información, consulta Multislice de Cloud TPU.

Instructivos de cargas de trabajo de capacitación

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: Entrena la difusión en Pokémon
  • PyTorch/XLA: GPT2 en WikiText

Descarga recursos de instructivos

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

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

El siguiente comando crea un nuevo clúster regional de GKE suscrito al canal de versiones rápidas 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

Para habilitar la federación de identidades para cargas de trabajo para GKE y las funciones del controlador de CSI de Cloud Storage FUSE para 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 creaste recientemente:

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 difusió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 los siguientes recursos:

Crea una imagen de Docker

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

Antes de ejecutar el siguiente comando, asegúrate de que tu cuenta tenga los permisos adecuados para que Docker 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 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éntala con lo siguiente:

kubectl apply -f tpu_job_diffusion.yaml

Limpieza

Después de que se termine de ejecutar el 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 TPUs v5e con HuggingFace en PyTorch/XLA con el conjunto de datos de wikitext.

Crea una imagen de Docker

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

Antes de ejecutar el siguiente comando, asegúrate de que tu cuenta tenga los permisos adecuados para que Docker 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

Después de que se termine de ejecutar el 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 un solo host en las TPUs de GKE v5e para modelos previamente entrenados con JAX, TensorFlow y PyTorch. A grandes rasgos, hay cuatro pasos independientes que se deben realizar en el clúster de GKE:

  1. Crea un bucket de Cloud Storage y configura el acceso a él. Usas un bucket de Cloud Storage para almacenar el modelo previamente entrenado.

  2. Descarga y convierte un modelo previamente entrenado en uno compatible con TPU. Aplica un pod de Kubernetes que descargue el modelo previamente entrenado, use el Convertidor de Cloud TPU y almacene los modelos convertidos en un bucket de Cloud Storage con el controlador de CSI de Cloud Storage FUSE. 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 CPU.

  3. Inicia el servidor del 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.

  4. Implementa 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 el servidor de modelos.

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

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 Deployment 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 tu grupo de nodo TPU, cambia el campo 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 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 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

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

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

  2. Crea el nodo TPU de host único.

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

  • Tu clúster se ejecuta con un grupo de nodos 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 debe 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 publicación. En este instructivo, usaremos el siguiente enfoque:

  • Para convertir el modelo previamente entrenado para que funcione en TPU, usaremos una nube privada virtual respaldada por un 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: La ubicación de tu bucket de Cloud Storage. Para mejorar el rendimiento, especifica la ubicación en la que se ejecuta tu clúster de GKE.

Sigue los pasos que se indican a continuación para otorgarle acceso al bucket al clúster de GKE. Para simplificar la configuración, en los siguientes ejemplos se usa el espacio de nombres predeterminada y la cuenta de servicio de Kubernetes predeterminada. Para 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.

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

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

    Reemplaza lo siguiente:

    • your-iam-service-acct: Es el nombre de la nueva cuenta de servicio de IAM.
    • 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.

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

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

    Reemplaza lo siguiente:

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

  3. Para permitir que la cuenta de servicio de Kubernetes actúe en nombre de la cuenta de servicio de IAM, agrega una vinculación de política de IAM entre las dos. 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 nueva cuenta de servicio de IAM.
    • your-bucket-project-id: Es el ID del proyecto en el que creaste tu 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.

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

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

    Reemplaza lo siguiente:

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

  5. Ejecuta el siguiente comando para propagar el nombre de tu bucket en los archivos YAML de esta demo:

    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.

  6. 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 extremo a extremo de JAX BERT:

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

El Pod de Kubernetes realiza los siguientes pasos:

  1. Descarga y usa la secuencia de comandos de Python export_bert_model.py de los recursos de ejemplo para descargar el modelo BERT previamente entrenado en un directorio temporal.
  2. Usa la imagen del convertidor de Cloud TPU para convertir el modelo previamente entrenado de CPU a TPU y 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

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

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

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

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

kubectl get deployment bert-deployment

AVAILABLE puede tardar un minuto en leer 3.

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

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 la 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 extremo a extremo de Stable Diffusion de JAX

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:

  1. 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.
  2. 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

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

kubectl get pods install-stable-diffusion

Es posible que STATUS tarde 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 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:

La instrucción 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 extremo a extremo de ResNet-50 de TensorFlow:

Instala las dependencias de Python para ejecutar secuencias de comandos de Python del instructivo 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 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

Es posible que STATUS tarde unos minutos en leer Completed.

Paso 2: Entrega el modelo con TensorFlow Serving

Aplica la implementación de la entrega de modelos:

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 mostrar una IP.

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 como 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 los 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 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:

  1. Genera un archivo del modelo.

    1. Aplica el archivo del modelo:
    kubectl apply -f pt/densenet161/model-archive.yml
    1. Verifica que el modelo se haya instalado correctamente con lo siguiente:
    kubectl get pods densenet161-model-archive

    Es posible que STATUS tarde unos minutos en leer Completed.

  2. 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 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.

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

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

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

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

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

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 ejecutar tu contenedor de TPU en modo con privilegios o de haber aumentado el ulimit dentro de 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 se programarán en porciones de TPU con una topología de TPU determinada (por ejemplo, v4-32). Supongamos también que tienes dos porciones de TPU v4-32 dentro del clúster de GKE. Llamaremos a esas porciones X y Y. Dado que tu clúster tiene una capacidad suficiente para programar ambos trabajos, en teoría, ambos 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 del trabajo A en la porción X y, luego, programa un Pod de Kubernetes del trabajo B en la porción X. En este caso, dadas las reglas de afinidad de Pods de Kubernetes para el trabajo A, el programador intentará programar todos los Pods de Kubernetes restantes para el trabajo A en la porción X. Lo mismo 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 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 grupo 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 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 tu 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 tu 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.