Implementa cargas de trabajo de TPU en GKE Autopilot

En esta página, se describe cómo acelerar las cargas de trabajo de aprendizaje automático (AA) con aceleradores de Cloud TPU (TPUs) en clústeres de Autopilot de Google Kubernetes Engine (GKE). Esta guía puede ayudarte a seleccionar las bibliotecas correctas para tus frameworks de aplicaciones de AA, configurar tus cargas de trabajo de TPU para que se ejecuten de forma óptima en GKE y supervisar tus cargas de trabajo después de la implementación.

Esta página está dirigida a los administradores y operadores de plataformas, los especialistas en IA y datos, y los desarrolladores de aplicaciones que desean preparar y ejecutar cargas de trabajo de AA en TPU. Para obtener más información sobre los roles, las responsabilidades y las tareas de ejemplo comunes a los que hacemos referencia en el contenido de Google Cloud , consulta Roles de usuario y tareas comunes de GKE.

Antes de leer esta página, asegúrate de estar familiarizado con los siguientes recursos:

Cómo funcionan las TPUs en Autopilot

Para usar TPUs en las cargas de trabajo de Autopilot, debes especificar lo siguiente en el manifiesto de tu carga de trabajo:

  • La versión de TPU en el campo spec.nodeSelector
  • La topología de TPU en el campo spec.nodeSelector La topología debe ser compatible con la versión de TPU especificada.
  • La cantidad de chips TPU en los campos spec.containers.resources.requests y spec.containers.resources.limits

Cuando implementas la carga de trabajo, GKE aprovisiona nodos que tienen la configuración de TPU solicitada y programa tus Pods en los nodos. GKE coloca cada carga de trabajo en su propio nodo para que cada Pod pueda acceder a todos los recursos del nodo con un riesgo mínimo de interrupción.

Las TPUs en Autopilot son compatibles con las siguientes capacidades:

  1. Spot Pods
  2. Reservas de capacidad específicas
  3. Pods de tiempo de ejecución extendido
  4. Inicio flexible

Planifica tu configuración de TPU

Antes de usar esta guía para implementar cargas de trabajo de TPU, planifica tu configuración de TPU según tu modelo y la cantidad de memoria que requiere. Para obtener más información, consulta Planifica tu configuración de TPU.

Precios

Para obtener información sobre los precios, consulta Precios de Autopilot.

Antes de comenzar

Antes de comenzar, asegúrate de haber realizado las siguientes tareas:

  • Habilita la API de Google Kubernetes Engine.
    • Habilitar la API de Google Kubernetes Engine
  • Si deseas usar Google Cloud CLI para esta tarea, instala y, luego, inicializa gcloud CLI. Si ya instalaste gcloud CLI, ejecuta el comando gcloud components update para obtener la versión más reciente. Es posible que las versiones anteriores de gcloud CLI no admitan la ejecución de los comandos que se describen en este documento.
  • Asegúrate de tener un clúster de Autopilot que ejecute la versión 1.32.3-gke.1927000 o posterior de GKE. Para obtener instrucciones, consulta Crea un clúster de Autopilot.
  • Para usar las TPUs reservadas, asegúrate de tener una reserva de capacidad específica existente. Para obtener instrucciones, consulta Cómo consumir una reserva.

Asegúrate de tener cuota para las TPU y otros recursos de GKE

En las siguientes secciones, se te ayudará a asegurarte de tener suficiente cuota cuando uses TPUs en GKE.

Para crear nodos de porción de TPU, debes tener cuota de TPU disponible, a menos que uses una reserva de capacidad existente. Si usas TPUs reservadas, omite esta sección.

La creación de nodos de porción de TPU en GKE requiere la cuota de la API de Compute Engine (compute.googleapis.com), no la cuota de la API de Cloud TPU (tpu.googleapis.com). El nombre de la cuota es diferente en los Pods de Autopilot normales y en los Pods Spot.

Si deseas verificar el límite y el uso actual de la cuota de la API de Compute Engine para las TPUs, sigue estos pasos:

  1. Dirígete a la página Cuotas en la consola de Google Cloud :

    Ir a Cuotas

  2. En la casilla Filtro, haz lo siguiente:

    1. Usa la siguiente tabla para seleccionar y copiar la propiedad de la cuota según la versión de TPU y el valor en el selector de nodos cloud.google.com/gke-tpu-accelerator. Por ejemplo, si planeas crear nodos TPU v5e a pedido cuyo valor en el selector de nodos cloud.google.com/gke-tpu-accelerator sea tpu-v5-lite-podslice, ingresa Name: TPU v5 Lite PodSlice chips.

      Versión de TPU. cloud.google.com/gke-tpu-accelerator Propiedad y nombre de la cuota para instancias bajo demanda Propiedad y nombre de la cuota para las instancias Spot2
      TPU v3,
      tpu-v3-device 
      Dimensions (e.g. location):
      tpu_family:CT3
      		 No aplicable
      TPU v3,
      tpu-v3-slice 
      Dimensions (e.g. location):
      tpu_family:CT3P
      		 No aplicable
      TPU v4,
      tpu-v4-podslice 
      Name:
      TPU v4 PodSlice chips
      		 
      Name:
      Preemptible TPU v4 PodSlice chips
      TPU v5e,
      tpu-v5-lite-podslice 
      Name:
      TPU v5 Lite PodSlice chips
      		 
      Name:
      Preemptible TPU v5 Lite Podslice
      chips
      TPU v5p,
      tpu-v5p-slice 
      Name:
      TPU v5p chips
      		 
      Name:
      Preemptible TPU v5p chips
      TPU Trillium,
      tpu-v6e-slice 
      Dimensions (e.g. location):
      tpu_family:CT6E
      		 
      Name:
      Preemptible TPU slices v6e
      Ironwood (TPU7x) (vista previa),
      tpu7x 
      Dimensions (e.g. location):
      tpu_family:tpu7x
      		 
      Name:
      Preemptible TPU slices tpu7x

    2. Selecciona la propiedad Dimensiones (p. ej., ubicaciones) y, luego, ingresa region: seguido del nombre de la región en la que planeas crear TPUs en GKE. Por ejemplo, ingresa region:us-west4 si planeas crear nodos de porción TPU en la zona us-west4-a. La cuota de TPU es regional, por lo que todas las zonas dentro de la misma región consumen la misma cuota de TPU.

Si ninguna cuota coincide con el filtro que ingresaste, no se le otorgó ninguna de las cuotas especificadas para la región que necesitas al proyecto y debes solicitar un ajuste de la cuota de TPU.

Cuando se crea una reserva de TPU, los valores de límite y de uso actuales para la cuota correspondiente aumentan según la cantidad de chips en la reserva de TPU. Por ejemplo, cuando se crea una reserva para 16 chips TPU v5e cuyo valor en el selector de nodos cloud.google.com/gke-tpu-accelerator es tpu-v5-lite-podslice, tanto el límite como el uso actual de la cuota de TPU v5 Lite PodSlice chips en la región pertinente aumentan en 16.

Cuotas para recursos adicionales de GKE

Es posible que debas aumentar las siguientes cuotas relacionadas con GKE en las regiones en las que GKE crea tus recursos.

  • 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 el producto de la cantidad máxima de nodos de GKE que esperas crear y 100 GB (nodos * 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.
  • Asegúrate de que max-pods-per-node se alinee con el rango de la subred: Cada nodo de Kubernetes usa rangos de IP secundarios para los Pods. Por ejemplo, un max-pods-per-node de 32 requiere 64 direcciones IP, lo que se traduce en una subred /26 por nodo. Ten en cuenta que este rango no se debe compartir con ningún otro clúster. Para evitar agotar el rango de direcciones IP, usa la marca --max-pods-per-node para limitar la cantidad de Pods que se pueden programar en un nodo. La cuota para max-pods-per-node 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 un ajuste de cuota.

Prepara tu aplicación de TPU

Las cargas de trabajo de TPU tienen los siguientes requisitos de preparación.

  1. Los frameworks como JAX, PyTorch y TensorFlow acceden a las VMs de TPU con la biblioteca compartida libtpu. libtpu incluye el compilador XLA, el software del entorno de ejecución de TPU y el controlador de TPU. Cada versión de PyTorch y JAX requiere una versión determinada de libtpu.so. Para evitar conflictos de versiones de paquetes, te recomendamos que uses una imagen generada por IA de JAX. Para usar las TPUs en GKE, asegúrate de usar las siguientes versiones: tpu7x
    Tipo de TPU libtpu.so versión
    Ironwood (TPU7x) (vista previa)
    TPU Trillium (v6e)
    tpu-v6e-slice
    TPU v5e
    tpu-v5-lite-podslice
    TPU v5p
    tpu-v5p-slice
    • Imagen generada por IA de JAX recomendada: jax0.4.35-rev1 o posterior
    • Versión recomendada de jax[tpu]: 0.4.19 o posterior.
    • Versión recomendada de torchxla[tpuvm]: Se sugiere usar una versión de compilación nocturna del 23 de octubre de 2023.
    TPU v4
    tpu-v4-podslice
    TPU v3
    tpu-v3-slice
    tpu-v3-device
  2. En el manifiesto de la carga de trabajo, agrega un selector de nodos de Kubernetes para asegurarte de que GKE programe la carga de trabajo de TPU en el tipo de máquina de TPU y la topología de TPU que definiste:

      nodeSelector:
    cloud.google.com/gke-tpu-accelerator: TPU_ACCELERATOR
    cloud.google.com/gke-tpu-topology: TPU_TOPOLOGY
    cloud.google.com/placement-policy-name: WORKLOAD_POLICY # Required only for Ironwood (TPU7x)

    Reemplaza lo siguiente:

    • TPU_ACCELERATOR: Es el nombre del acelerador de TPU. Por ejemplo, usa tpu7x-standard-4t.
    • TPU_TOPOLOGY: la topología física para la porción de TPU. El formato de la topología depende de la versión de la TPU. Por ejemplo, usa 2x2x2. Para obtener más información, consulta Planifica las TPU en GKE.
    • WORKLOAD_POLICY: Es el nombre de la política de cargas de trabajo que deseas usar para colocar tus Pods de TPU. Este selector de nodos solo es obligatorio para Ironwood (TPU7x).

Después de completar la preparación de la carga de trabajo, puedes ejecutar un trabajo que use TPUs.

Opciones para aprovisionar TPUs en GKE

Para aprovisionar TPUs en GKE, tienes las siguientes opciones de configuración:
  • Solicitud de carga de trabajo: Especifica la versión y la topología de TPU en el campo spec.nodeSelector y la cantidad de chips TPU en la sección spec.containers.resources. Cuando implementas la carga de trabajo, GKE aprovisiona automáticamente nodos con la configuración de TPU correcta y coloca cada carga de trabajo en su propio nodo dedicado para garantizar el acceso completo a los recursos del nodo. Para obtener instrucciones, consulta Solicita TPUs en una carga de trabajo.

    • Aprovisiona de forma centralizada las TPU con clases de procesamiento personalizadas

    En las siguientes secciones, se muestra cómo crear una ComputeClass personalizada y, luego, crear un trabajo que consuma las TPU definidas en la ComputeClass.

    Crea una ComputeClass personalizada

    Los pasos para crear una ComputeClass personalizada que siga las reglas de TPU difieren según si usas Ironwood (TPUv7) o una versión anterior de TPU.

    Ironwood (TPU7x)

    1. Crea una política de carga de trabajo. Este paso solo es obligatorio si creas un grupo de nodos de varios hosts, que depende de la topología que elijas. Si usas un grupo de nodos de host único, omite este paso.

      gcloud compute resource-policies create workload-policy WORKLOAD_POLICY_NAME \
    --type=HIGH_THROUGHPUT \
    --accelerator-topology=TPU_TOPOLOGY \
    --project=PROJECT_ID \
    --region=REGION

      Reemplaza lo siguiente:

      • WORKLOAD_POLICY_NAME: Es un nombre para tu política de carga de trabajo.
      • TPU_TOPOLOGY: Es la topología de la TPU Ironwood (TPU7x). Por ejemplo, usa 2x2x2. Para obtener más información sobre todas las topologías compatibles con Ironwood (TPU7x), consulta la sección de topología.
      • PROJECT_ID: Es el ID del proyecto de Google Cloud .
      • REGION: Es la región de la política de carga de trabajo. Una política de carga de trabajo es un recurso regional que puedes usar en todos los grupos de nodos.

    2. Guarda el siguiente manifiesto como tpu-compute-class.yaml:

      apiVersion: cloud.google.com/v1
kind: ComputeClass
metadata:
  name: tpu-class
spec:
  priorities:
    - tpu:
        type: tpu7x
        topology: TPU_TOPOLOGY
        count: 4
      placement:
        policyName: WORKLOAD_POLICY_NAME
  nodePoolAutoCreation:
    enabled: true

    3. (Opcional) Puedes consumir una reserva o un subbloque específicos. Por ejemplo, puedes agregar el siguiente specs a tu manifiesto de ComputeClass:

        reservations:
    affinity: Specific
    specific:
      - name: RESERVATION_NAME
        reservationBlock:
          name: RESERVATION_BLOCK_NAME
          reservationSubBlock:
            name: RESERVATION_SUB_BLOCK_NAME

      Reemplaza lo siguiente:

      • RESERVATION_NAME es el nombre de la reserva de capacidad de Compute Engine.
      • RESERVATION_BLOCK_NAME: Es el nombre del bloque de reserva de capacidad de Compute Engine.
      • RESERVATION_SUB_BLOCK_NAME: Es el nombre del sub-bloque de la reserva de capacidad de Compute Engine.

      Para obtener más información, consulta Consume recursos zonales reservados.

    Otras versiones de TPU

    Para aprovisionar TPU v3, v4, v5p, v5e o v6e (Trillium) con una ComputeClass personalizada configurada para TPUs, completa los siguientes pasos:

    1. Guarda el siguiente manifiesto como tpu-compute-class.yaml:

      apiVersion: cloud.google.com/v1
kind: ComputeClass
metadata:
  name: tpu-class
spec:
  priorities:
  - tpu:
      type: TPU_TYPE
      count: NUMBER_OF_CHIPS
      topology: TOPOLOGY
  - spot: true
    tpu:
      type: {"<var>"}}TPU_TYPE
      count: NUMBER_OF_CHIPS
      topology: TOPOLOGY
  - flexStart:
      enabled: true
    tpu:
      type: {"<var>"}}TPU_TYPE
      count: NUMBER_OF_CHIPS
      topology: TOPOLOGY
  nodePoolAutoCreation:
    enabled: true

      Reemplaza lo siguiente:

      • TPU_TYPE: Es el tipo de TPU que se usará, como tpu-v4-podslice. Debe ser un valor compatible con GKE.
      • TOPOLOGY: Es la disposición de los chips de TPU en la porción, como 2x2x4. Debe ser una topología compatible con el tipo de TPU seleccionado.
      • NUMBER_OF_CHIPS: Es la cantidad de chips TPU que usará el contenedor. Debe ser el mismo valor para limits y requests.

    2. Implementa ComputeClass:

      kubectl apply -f tpu-compute-class.yaml

      Para obtener más información sobre las ComputeClasses y las TPU personalizadas, consulta Configuración de TPU.

    Crea un trabajo que consuma TPUs

    1. Guarda el siguiente manifiesto como tpu-job.yaml:

      apiVersion: v1
kind: Service
metadata:
  name: headless-svc
spec:
  clusterIP: None
  selector:
    job-name: tpu-job
---
apiVersion: batch/v1
kind: Job
metadata:
  name: tpu-job
spec:
  backoffLimit: 0
  completions: 4
  parallelism: 4
  completionMode: Indexed
  template:
    spec:
      subdomain: headless-svc
      restartPolicy: Never
      nodeSelector:
        cloud.google.com/compute-class: tpu-class
      containers:
      - name: tpu-job
        image: us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest
        ports:
        - containerPort: 8471 # Default port using which TPU VMs communicate
        - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
        command:
        - bash
        - -c
        - |
          python -c 'import jax; print("TPU cores:", jax.device_count())'
        resources:
          requests:
            cpu: 10
            memory: MEMORY_SIZE
            google.com/tpu: NUMBER_OF_CHIPS
          limits:
            cpu: 10
            memory: MEMORY_SIZE
            google.com/tpu: NUMBER_OF_CHIPS

      Reemplaza lo siguiente:

      • NUMBER_OF_CHIPS: Es la cantidad de chips TPU que usará el contenedor. Debe ser el mismo valor para limits y requests, igual al valor de CHIP_COUNT en la ComputeClass personalizada seleccionada.
      • MEMORY_SIZE: Es la cantidad máxima de memoria que usa la TPU. Los límites de memoria dependen de la versión y la topología de la TPU que uses. Para obtener más información, consulta Valores mínimos y máximos para los aceleradores.
      • NUMBER_OF_CHIPS: Es la cantidad de chips TPU que usará el contenedor. Debe ser el mismo valor para limits y requests.

    2. Implementa el trabajo:

      kubectl create -f tpu-job.yaml

      Cuando creas este trabajo, GKE hace automáticamente lo siguiente:

      • Aprovisiona nodos para ejecutar los Pods. Según el tipo de TPU, la topología y las solicitudes de recursos que especificaste, estos nodos son porciones de host único o de varios hosts. Según la disponibilidad de los recursos de TPU en la prioridad más alta, es posible que GKE recurra a prioridades más bajas para maximizar la capacidad de obtención.
      • Agrega taints a los Pods y tolerancias a los nodos para evitar que otras cargas de trabajo se ejecuten en los mismos nodos que las cargas de trabajo de TPU.

      Para obtener más información, consulta Acerca de las ComputeClasses personalizadas.

    3. Cuando termines esta sección, puedes borrar los recursos que creaste para evitar que se te siga facturando:

      kubectl delete -f tpu-job.yaml

    Solicita TPUs en una carga de trabajo

    En esta sección, se muestra cómo crear un trabajo que solicite TPUs en Autopilot. En cualquier carga de trabajo que necesite TPUs, debes especificar lo siguiente:

    • Selectores de nodos para la versión y la topología de TPU
    • Cantidad de chips TPU para un contenedor en tu carga de trabajo

    Para obtener una lista de las versiones y topologías de TPU compatibles, y la cantidad correspondiente de chips y nodos de TPU en una porción, consulta Elige la versión de TPU.

    Consideraciones para las solicitudes de TPU en cargas de trabajo

    Solo un contenedor en un Pod puede usar TPUs. La cantidad de chips TPU que solicita un contenedor debe ser igual a la cantidad de chips TPU conectados a un nodo en la porción. Por ejemplo, si solicitas una TPU v5e (tpu-v5-lite-podslice) con una topología 2x4, puedes solicitar cualquiera de las siguientes opciones:

    • Chips 4, que crea dos nodos multihost con 4 chips TPU cada uno
    • Chips 8, que crean un nodo de host único con 8 chips de TPU

    Como práctica recomendada para maximizar la eficiencia de costos, siempre consume toda la TPU en la porción que solicitas. Si solicitas una porción de varios hosts de dos nodos con 4 chips TPU cada uno, debes implementar una carga de trabajo que se ejecute en ambos nodos y consuma los 8 chips TPU de la porción.

    Crea una carga de trabajo que solicite TPUs

    Mediante los siguientes pasos, se crea un trabajo que solicita TPUs. Si tienes cargas de trabajo que se ejecutan en porciones de TPU de varios hosts, también debes crear un servicio headless que seleccione tu carga de trabajo por nombre. Este Service sin interfaz gráfica permite que los Pods en diferentes nodos de la segmentación de varios hosts se comuniquen entre sí actualizando la configuración de DNS de Kubernetes para que apunte a los Pods en la carga de trabajo.

    1. Guarda el siguiente manifiesto como tpu-autopilot.yaml:

      apiVersion: v1
kind: Service
metadata:
  name: headless-svc
spec:
  clusterIP: None
  selector:
    job-name: tpu-job
---
apiVersion: batch/v1
kind: Job
metadata:
  name: tpu-job
spec:
  backoffLimit: 0
  completions: 4
  parallelism: 4
  completionMode: Indexed
  template:
    spec:
      # Optional: Run in GKE Sandbox
      # runtimeClassName: gvisor
      subdomain: headless-svc
      restartPolicy: Never
      nodeSelector:
        cloud.google.com/gke-tpu-accelerator: TPU_TYPE
        cloud.google.com/gke-tpu-topology: TOPOLOGY
      containers:
      - name: tpu-job
        image: us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest
        ports:
        - containerPort: 8471 # Default port using which TPU VMs communicate
        - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
        command:
        - bash
        - -c
        - |
          python -c 'import jax; print("TPU cores:", jax.device_count())'
        resources:
          requests:
            cpu: 10
            memory: MEMORY_SIZE
            google.com/tpu: NUMBER_OF_CHIPS
          limits:
            cpu: 10
            memory: MEMORY_SIZE
            google.com/tpu: NUMBER_OF_CHIPS

      Reemplaza lo siguiente:

      • TPU_TYPE: Es el tipo de TPU que se usará, como tpu-v4-podslice. Debe ser un valor compatible con GKE.
      • TOPOLOGY: Es la disposición de los chips de TPU en la porción, como 2x2x4. Debe ser una topología compatible con el tipo de TPU seleccionado.
      • NUMBER_OF_CHIPS: Es la cantidad de chips TPU que usará el contenedor. Debe ser el mismo valor para limits y requests.
      • MEMORY_SIZE: Es la cantidad máxima de memoria que usa la TPU. Los límites de memoria dependen de la versión y la topología de la TPU que uses. Para obtener más información, consulta Valores mínimos y máximos para los aceleradores.

      De manera opcional, también puedes cambiar los siguientes campos:

      • image: Es la imagen generada por IA de JAX que se usará. En el manifiesto de ejemplo, este campo se establece en la imagen generada por IA de JAX más reciente. Para establecer una versión diferente, consulta la lista de imágenes generadas por IA de JAX actuales.
      • runtimeClassname: gvisor: Es el parámetro de configuración que permite ejecutar este Pod en GKE Sandbox. Para usarlo, quita el comentario de esta línea. GKE Sandbox admite las TPU versión v4 y posteriores. Para obtener más información, consulta GKE Sandbox.

    2. Implementa el trabajo:

      kubectl create -f tpu-autopilot.yaml

      Cuando creas este trabajo, GKE hace automáticamente lo siguiente:

      1. Aprovisiona nodos para ejecutar los Pods. Según el tipo de TPU, la topología y las solicitudes de recursos que especificaste, estos nodos son porciones de host único o de varios hosts.
      2. Agrega taints a los Pods y tolerancias a los nodos para evitar que otras cargas de trabajo se ejecuten en los mismos nodos que las cargas de trabajo de TPU.

    3. Cuando termines esta sección, puedes borrar la carga de trabajo que creaste para evitar que se te siga facturando:

      kubectl delete -f tpu-autopilot.yaml

    Crea una carga de trabajo que solicite TPUs y la programación de la recopilación

    En la TPU Trillium, puedes usar la programación de recopilación para agrupar los nodos de porción de TPU. Agrupar estos nodos de segmentación de TPU facilita el ajuste de la cantidad de réplicas para satisfacer la demanda de la carga de trabajo. Google Cloud controla las actualizaciones de software para garantizar que siempre haya suficientes segmentaciones disponibles en la colección para atender el tráfico.

    La TPU Trillium admite la programación de recopilación para grupos de nodos de host único y de varios hosts que ejecutan cargas de trabajo de inferencia. A continuación, se describe cómo se comporta la programación de la recopilación según el tipo de porción de TPU que uses:

    • Porción de TPU de varios hosts: GKE agrupa las porciones de TPU de varios hosts para formar una colección. Cada grupo de nodos de GKE es una réplica dentro de esta colección. Para definir una colección, crea una porción de TPU de varios hosts y asígnale un nombre único. Para agregar más porciones de TPU a la colección, crea otro grupo de nodos de porción de TPU de varios hosts con el mismo nombre de colección y tipo de carga de trabajo.
    • Porción de TPU de host único: GKE considera todo el grupo de nodos de porción de TPU de host único como una colección. Para agregar más porciones de TPU a la colección, puedes cambiar el tamaño del grupo de nodos de porción de TPU de host único.

    Para obtener más información sobre la limitación de la programación de colecciones, consulta Cómo funciona la programación de colecciones.

    Usa una porción de TPU de varios hosts

    La programación de la recopilación en nodos de porción de TPU de varios hosts está disponible para los clústeres de Autopilot en la versión 1.31.2-gke.1537000 y posteriores. Los nodos de porción de TPU de varios hosts con una topología 2x4 solo son compatibles con la versión 1.31.2-gke.1115000 o posterior. Para crear nodos de porción de TPU de varios hosts y agruparlos como una colección, agrega las siguientes etiquetas de Kubernetes a la especificación de tu carga de trabajo:

    • cloud.google.com/gke-nodepool-group-name: Cada colección debe tener un nombre único a nivel del clúster. El valor de la etiqueta cloud.google.com/gke-nodepool-group-name debe cumplir con los requisitos para las etiquetas de clúster.

    • cloud.google.com/gke-workload-type: HIGH_AVAILABILITY

      Por ejemplo, el siguiente bloque de código define una colección con una porción de TPU de varios hosts:

        nodeSelector:
    cloud.google.com/gke-nodepool-group-name: ${COLLECTION_NAME}
    cloud.google.com/gke-workload-type: HIGH_AVAILABILITY
    cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
    cloud.google.com/gke-tpu-topology: 4x4
...

    Usa una porción de TPU de host único

    La programación de la recopilación en nodos de segmentación de TPU de un solo host está disponible para los clústeres de Autopilot en la versión 1.31.2-gke.1088000 y posteriores. Para crear nodos de porción de TPU de host único y agruparlos como una colección, agrega la etiqueta cloud.google.com/gke-workload-type:HIGH_AVAILABILITY en la especificación de tu carga de trabajo.

    Por ejemplo, el siguiente bloque de código define una colección con una porción de TPU de un solo host:

      nodeSelector:
    cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
    cloud.google.com/gke-tpu-topology: 2x2
    cloud.google.com/gke-workload-type: HIGH_AVAILABILITY
  ...

    Usa clases de procesamiento personalizadas para implementar una colección

    Para obtener más información sobre cómo implementar una carga de trabajo que solicita la programación de la carga de trabajo y la recopilación de TPU con clases de procesamiento personalizadas, consulta Recopilación de TPU de varios hosts y Define el tipo de carga de trabajo para el SLO de TPU.

    Ejemplo: Muestra la cantidad total de chips TPU en una porción de varios hosts

    La siguiente carga de trabajo devuelve la cantidad de chips TPU en todos los nodos de una porción de TPU de varios hosts. Para crear una porción de varios hosts, la carga de trabajo tiene los siguientes parámetros:

    • Versión de TPU: TPU v4
    • Topología: 2 x 2 x 4

    Esta elección de versión y topología dan como resultado una porción de varios hosts.

    1. Guarda el siguiente manifiesto como available-chips-multihost.yaml: 
      apiVersion: v1
kind: Service
metadata:
  name: headless-svc
spec:
  clusterIP: None
  selector:
    job-name: tpu-available-chips
---
apiVersion: batch/v1
kind: Job
metadata:
  name: tpu-available-chips
spec:
  backoffLimit: 0
  completions: 4
  parallelism: 4
  completionMode: Indexed
  template:
    spec:
      subdomain: headless-svc
      restartPolicy: Never
      nodeSelector:
        cloud.google.com/gke-tpu-accelerator: tpu-v4-podslice # Node selector to target TPU v4 slice nodes.
        cloud.google.com/gke-tpu-topology: 2x2x4 # Specifies the physical topology for the TPU slice.
      containers:
      - name: tpu-job
        image: us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest
        ports:
        - containerPort: 8471 # Default port using which TPU VMs communicate
        - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
        command:
        - bash
        - -c
        - |
          python -c 'import jax; print("TPU cores:", jax.device_count())' # Python command to count available TPU chips.
        resources:
          requests:
            cpu: 10
            memory: 407Gi
            google.com/tpu: 4 # Request 4 TPU chips for this workload.
          limits:
            cpu: 10
            memory: 407Gi
            google.com/tpu: 4 # Limit to 4 TPU chips for this workload.
    2. Implementa el manifiesto: 
      kubectl create -f available-chips-multihost.yaml

      GKE ejecuta una porción de TPU v4 con cuatro VMs de TPU (porción de TPU de varios hosts). La porción tiene 16 chips TPU interconectados.

    3. Verifica que el trabajo haya creado cuatro pods: 
      kubectl get pods

      El resultado es similar a este:

      NAME                       READY   STATUS      RESTARTS   AGE
tpu-job-podslice-0-5cd8r   0/1     Completed   0          97s
tpu-job-podslice-1-lqqxt   0/1     Completed   0          97s
tpu-job-podslice-2-f6kwh   0/1     Completed   0          97s
tpu-job-podslice-3-m8b5c   0/1     Completed   0          97s
    4. Obtén los registros de uno de los Pods: 
      kubectl logs POD_NAME

      Reemplaza POD_NAME por el nombre de uno de los Pods creados. Por ejemplo, tpu-job-podslice-0-5cd8r

      El resultado es similar a lo siguiente: 

      TPU cores: 16
    5. Opcional: Quita la carga de trabajo: 
      kubectl delete -f available-chips-multihost.yaml

    Ejemplo: Muestra los chips TPU en un solo nodo

    La siguiente carga de trabajo es un Pod estático que muestra la cantidad de chips TPU adjuntos a un nodo específico. Para crear un nodo de host único, la carga de trabajo tiene los siguientes parámetros:

    • Versión de TPU: TPU v5e
    • Topología: 2 x 4

    Esta elección de versión y topología da como resultado una porción de host único.

    1. Guarda el siguiente manifiesto como available-chips-singlehost.yaml: 
      apiVersion: v1
kind: Pod
metadata:
  name: tpu-job-jax-v5
spec:
  restartPolicy: Never
  nodeSelector:
    cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice # Node selector to target TPU v5e slice nodes.
    cloud.google.com/gke-tpu-topology: 2x4 # Specify the physical topology for the TPU slice.
  containers:
  - name: tpu-job
    image: us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest
    ports:
    - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
    command:
    - bash
    - -c
    - |
      python -c 'import jax; print("Total TPU chips:", jax.device_count())'
    resources:
      requests:
        google.com/tpu: 8 # Request 8 TPU chips for this container.
      limits:
        google.com/tpu: 8 # Limit to 8 TPU chips for this container.
    2. Implementa el manifiesto: 
      kubectl create -f available-chips-singlehost.yaml

      GKE aprovisiona nodos con ocho porciones de TPU de host único que usan TPU v5e. Cada nodo TPU tiene ocho chips TPU (porción de TPU de host único).

    3. Obtén los registros del pod: 
      kubectl logs tpu-job-jax-v5

      El resultado es similar a lo siguiente: 

      Total TPU chips: 8
    4. Opcional: Quita la carga de trabajo: 
        kubectl delete -f available-chips-singlehost.yaml

    Observa y supervisa las TPU

    Panel

    La observabilidad del grupo de nodos en la Google Cloud consola está disponible de forma general. Para ver el estado de tus grupos de nodos de TPU de varios hosts en GKE, ve al panel Estado del grupo de nodos de TPU de GKE que proporciona Cloud Monitoring:

    Ir al estado del grupo de nodos de TPU de GKE

    Este panel te brinda estadísticas integrales sobre el estado de tus grupos de nodo TPU de varios hosts. Para obtener más información, consulta Supervisa las métricas de estado de los nodos y los grupos de nodos de TPU.

    En la página Clústeres de Kubernetes de la consola deGoogle Cloud , la pestaña Observabilidad también muestra las métricas de observabilidad de TPU, como el uso de TPU, en el encabezado Aceleradores > TPU. Para obtener más información, consulta Visualiza métricas de observabilidad.

    El panel de TPU solo se propaga si tienes las métricas del sistema habilitadas en tu clúster de GKE.

    Métricas del entorno de ejecución

    En la versión 1.27.4-gke.900 o posterior de GKE, las cargas de trabajo de TPU que usan la versión 0.4.14 o posterior de JAX y especifican containerPort: 8431 exportan las métricas de uso de TPU como métricas del sistema de GKE. Las siguientes métricas están disponibles en Cloud Monitoring para supervisar el rendimiento del entorno de ejecución de tu carga de trabajo de TPU:

    • Ciclo de trabajo: Porcentaje de tiempo durante el último período de muestreo (60 segundos) durante el cual los TensorCores se procesaron de forma activa en un chip TPU. Un porcentaje mayor significa un mejor uso de TPU.
    • Memoria usada: Cantidad de memoria del acelerador asignada en bytes. Se tomaron muestras cada 60 segundos.
    • Memoria total: Memoria total del acelerador en bytes. Se hace un muestreo cada 60 segundos.

    Estas métricas se encuentran en el esquema de nodos de Kubernetes (k8s_node) y de contenedor de Kubernetes (k8s_container).

    Contenedor de Kubernetes:

    • kubernetes.io/container/accelerator/duty_cycle
    • kubernetes.io/container/accelerator/memory_used
    • kubernetes.io/container/accelerator/memory_total

    Nodo de Kubernetes:

    • kubernetes.io/node/accelerator/duty_cycle
    • kubernetes.io/node/accelerator/memory_used
    • kubernetes.io/node/accelerator/memory_total

    Supervisa las métricas de estado de los nodos y grupos de nodos de TPU

    Cuando un trabajo de entrenamiento tiene un error o finaliza con una falla, puedes verificar las métricas relacionadas con la infraestructura subyacente para determinar si la interrupción se debió a un problema con el nodo o el grupo de nodos subyacente.

    Estado del nodo

    En la versión 1.32.1-gke.1357001 o posterior de GKE, la siguiente métrica del sistema de GKE expone el estado de un nodo de GKE:

    • kubernetes.io/node/status_condition

    El campo condition informa las condiciones del nodo, como Ready, DiskPressure y MemoryPressure. El campo status muestra el estado informado de la condición, que puede ser True, False o Unknown. Esta es una métrica con el tipo de recurso supervisado k8s_node.

    Esta consulta de PromQL muestra si un nodo en particular es Ready:

    kubernetes_io:node_status_condition{
    monitored_resource="k8s_node",
    cluster_name="CLUSTER_NAME",
    node_name="NODE_NAME",
    condition="Ready",
    status="True"}

    Para ayudarte a solucionar problemas en un clúster, puedes consultar los nodos que hayan presentado otras condiciones:

    kubernetes_io:node_status_condition{
    monitored_resource="k8s_node",
    cluster_name="CLUSTER_NAME",
    condition!="Ready",
    status="True"}

    Es posible que desees observar específicamente los nodos que no son Ready:

    kubernetes_io:node_status_condition{
    monitored_resource="k8s_node",
    cluster_name="CLUSTER_NAME",
    condition="Ready",
    status="False"}

    Si no hay datos, los nodos están listos. La condición de estado se muestrea cada 60 segundos.

    Puedes usar la siguiente consulta para comprender el estado de los nodos en toda la flota:

    avg by (condition,status)(
  avg_over_time(
    kubernetes_io:node_status_condition{monitored_resource="k8s_node"}[${__interval}]))

    Estado del grupo de nodos

    La siguiente métrica del sistema de GKE para el recurso supervisado k8s_node_pool expone el estado de un grupo de nodos de GKE:

    • kubernetes.io/node_pool/status

    Esta métrica solo se informa para los grupos de nodo TPU de varios hosts.

    El campo status informa el estado del grupo de nodos, como Provisioning, Running, Error, Reconciling o Stopping. Las actualizaciones de estado se realizan después de que se completan las operaciones de la API de GKE.

    Para verificar si un grupo de nodos en particular tiene el estado Running, usa la siguiente consulta de PromQL:

    kubernetes_io:node_pool_status{
    monitored_resource="k8s_node_pool",
    cluster_name="CLUSTER_NAME",
    node_pool_name="NODE_POOL_NAME",
    status="Running"}

    Para supervisar la cantidad de grupos de nodos en tu proyecto agrupados por su estado, usa la siguiente consulta de PromQL:

    count by (status)(
  count_over_time(
    kubernetes_io:node_pool_status{monitored_resource="k8s_node_pool"}[${__interval}]))

    Disponibilidad del grupo de nodos

    La siguiente métrica del sistema de GKE muestra si hay disponible un grupo de nodo TPU de varios hosts:

    • kubernetes.io/node_pool/multi_host/available

    La métrica tiene un valor de True si todos los nodos del grupo de nodos están disponibles y False en caso contrario. La métrica se muestrea cada 60 segundos.

    Para verificar la disponibilidad de los grupos de nodo TPU de varios hosts en tu proyecto, usa la siguiente consulta de PromQL:

    avg by (node_pool_name)(
  avg_over_time(
    kubernetes_io:node_pool_multi_host_available{
      monitored_resource="k8s_node_pool",
      cluster_name="CLUSTER_NAME"}[${__interval}]))

    Recuento de interrupciones de nodos

    La siguiente métrica del sistema de GKE informa el recuento de interrupciones de un nodo de GKE desde la última muestra (la métrica se muestrea cada 60 segundos):

    • kubernetes.io/node/interruption_count

    Los campos interruption_type (como TerminationEvent, MaintenanceEvent o PreemptionEvent) y interruption_reason (como HostError, Eviction o AutoRepair) pueden ayudar a proporcionar el motivo por el que se interrumpió un nodo.

    Para obtener un desglose de las interrupciones y sus causas en los nodos de TPU de los clústeres de tu proyecto, usa la siguiente consulta de PromQL:

      sum by (interruption_type,interruption_reason)(
    sum_over_time(
      kubernetes_io:node_interruption_count{monitored_resource="k8s_node"}[${__interval}]))

    Para ver solo los eventos de mantenimiento del host, actualiza la consulta para filtrar el valor HW/SW Maintenance para el interruption_reason. Usa la siguiente consulta de PromQL:

      sum by (interruption_type,interruption_reason)(
    sum_over_time(
      kubernetes_io:node_interruption_count{monitored_resource="k8s_node", interruption_reason="HW/SW Maintenance"}[${__interval}]))

    Para ver el recuento de interrupciones agregado por grupo de nodos, usa la siguiente consulta de PromQL:

      sum by (node_pool_name,interruption_type,interruption_reason)(
    sum_over_time(
      kubernetes_io:node_pool_interruption_count{monitored_resource="k8s_node_pool", interruption_reason="HW/SW Maintenance", node_pool_name=NODE_POOL_NAME }[${__interval}]))

    Tiempos de recuperación (TTR) del grupo de nodos

    La siguiente métrica del sistema de GKE informa la distribución de las duraciones del período de recuperación para los grupos de nodo TPU de varios hosts de GKE:

    • kubernetes.io/node_pool/accelerator/times_to_recover

    Cada muestra registrada en esta métrica indica un solo evento de recuperación del grupo de nodos a partir de un período de inactividad.

    Esta métrica es útil para hacer un seguimiento del tiempo de recuperación y el tiempo entre interrupciones del grupo de nodo TPU de varios hosts.

    Puedes usar la siguiente consulta de PromQL para calcular el tiempo medio de recuperación (MTTR) de los últimos 7 días en tu clúster:

    sum(sum_over_time(
  kubernetes_io:node_pool_accelerator_times_to_recover_sum{
    monitored_resource="k8s_node_pool", cluster_name="CLUSTER_NAME"}[7d]))
/
sum(sum_over_time(
  kubernetes_io:node_pool_accelerator_times_to_recover_count{
    monitored_resource="k8s_node_pool",cluster_name="CLUSTER_NAME"}[7d]))

    Tiempos entre interrupciones (TBI) del grupo de nodos

    Los tiempos entre interrupciones del grupo de nodos miden cuánto tiempo se ejecuta tu infraestructura antes de experimentar una interrupción. Se calcula como el promedio en un período, en el que el numerador mide el tiempo total que estuvo activa tu infraestructura y el denominador mide las interrupciones totales en tu infraestructura.

    En el siguiente ejemplo de PromQL, se muestra el tiempo medio entre interrupciones (MTBI) de 7 días para el clúster determinado:

    sum(count_over_time(
  kubernetes_io:node_memory_total_bytes{
    monitored_resource="k8s_node", node_name=~"gke-tpu.*|gk3-tpu.*", cluster_name="CLUSTER_NAME"}[7d]))
/
sum(sum_over_time(
  kubernetes_io:node_interruption_count{
    monitored_resource="k8s_node", node_name=~"gke-tpu.*|gk3-tpu.*", cluster_name="CLUSTER_NAME"}[7d]))

    Métricas de host

    En la versión 1.28.1-gke.1066000 o posterior de GKE, las VMs de una porción de TPU exportan las métricas de uso de TPU como métricas del sistema de GKE. Las siguientes métricas están disponibles en Cloud Monitoring para supervisar el rendimiento de tu host TPU:

    • Uso de TensorCore: Porcentaje actual del TensorCore que se usa. El valor de TensorCore es igual a la suma de las unidades de multiplicación de matriz (MXUs) más la unidad vectorial. El valor de uso de TensorCore es la división de las operaciones de TensorCore que se llevaron a cabo durante el último período de muestra (60 segundos) por la cantidad compatible de operaciones de TensorCore durante el mismo período. Un valor mayor significa un mejor uso.
    • Uso de ancho de banda de memoria: Porcentaje actual del ancho de banda de memoria del acelerador que se usa. Se calcula dividiendo el ancho de banda de memoria usado durante un período de muestra (60 s) por el ancho de banda máximo admitido durante el mismo período de muestra.

    Estas métricas se encuentran en el esquema de nodos de Kubernetes (k8s_node) y de contenedor de Kubernetes (k8s_container).

    Contenedor de Kubernetes:

    • kubernetes.io/container/accelerator/tensorcore_utilization
    • kubernetes.io/container/accelerator/memory_bandwidth_utilization

    Nodo de Kubernetes:

    • kubernetes.io/node/accelerator/tensorcore_utilization
    • kubernetes.io/node/accelerator/memory_bandwidth_utilization

    Para obtener más información, consulta Métricas de Kubernetes y Métricas del sistema de GKE.

    Logging

    Los registros que emiten los contenedores que se ejecutan en nodos de GKE, incluidas las VMs de TPU, son recopiladas por el agente de Logging de GKE, se envían a Logging y son visibles en Logging.

    Recomendaciones para cargas de trabajo de TPU en Autopilot

    Las siguientes recomendaciones pueden mejorar la eficiencia de tus cargas de trabajo de TPU:

    • Usa Pods de tiempo de ejecución extendido para un período de gracia de hasta siete días antes de que GKE finalice tus Pods para reducir la escala verticalmente o las actualizaciones de nodos. Puedes usar exclusiones y períodos de mantenimiento con Pods de tiempo de ejecución extendido para retrasar aún más las actualizaciones automáticas de nodos.
    • Usa las reservas de capacidad para asegurarte de que tus cargas de trabajo reciban las TPUs solicitadas sin que se coloquen en una cola para consultar la disponibilidad.

    Para obtener información sobre cómo configurar Cloud TPU en GKE, consulta los siguientes Google Cloud recursos: