Desplegar cargas de trabajo de TPU en Autopilot de GKE

En esta página se describe cómo acelerar las cargas de trabajo de aprendizaje automático (ML) mediante el uso de aceleradores de TPU de Cloud en clústeres de Autopilot de Google Kubernetes Engine (GKE). Estas directrices pueden ayudarte a seleccionar las bibliotecas correctas para tus frameworks de aplicaciones de aprendizaje automático, configurar tus cargas de trabajo de TPU para que se ejecuten de forma óptima en GKE y monitorizar tus cargas de trabajo después de la implementación.

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

Antes de leer esta página, asegúrate de que conoces los siguientes recursos:

Cómo funcionan las TPUs en Autopilot

Para usar las TPUs en cargas de trabajo de Autopilot, especifica lo siguiente en el manifiesto de la carga de trabajo:

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

Cuando despliegas 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 de Autopilot son compatibles con las siguientes funciones:

  1. Spot Pods
  2. Reservas de capacidad específicas
  3. Grupos de anuncios con tiempos de publicación prolongados
  4. Inicio flexible

Planificar la configuración de TPU

Antes de usar esta guía para desplegar cargas de trabajo de TPU, planifica la configuración de TPU en función de tu modelo y de la cantidad de memoria que necesite. Para obtener más información, consulta Planificar la configuración de las TPU.

Precios

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

Antes de empezar

Antes de empezar, asegúrate de que has realizado las siguientes tareas:

  • Habilita la API de Google Kubernetes Engine.
    • Habilitar la API de Google Kubernetes Engine
  • Si quieres usar Google Cloud CLI para esta tarea, instálala y, a continuación, inicialízala. Si ya has instalado la CLI de gcloud, obtén la versión más reciente ejecutando el comando gcloud components update. Es posible que las versiones anteriores de la interfaz de línea de comandos de gcloud no admitan la ejecución de los comandos de este documento.
  • Asegúrate de que tienes un clúster de Autopilot que ejecuta la versión 1.32.3-gke.1927000 de GKE o una posterior. Para obtener instrucciones, consulta Crear un clúster de Autopilot.
  • Para usar las TPUs reservadas, asegúrate de que tienes una reserva de capacidad específica. Para obtener instrucciones, consulta Consumir una reserva.

Asegurar la cuota de TPUs y otros recursos de GKE

En las siguientes secciones se explica cómo asegurarse de que tiene suficiente cuota al usar TPUs en GKE.

Para crear nodos de segmento de TPU, debes tener cuota de TPU disponible, a menos que utilices una reserva de capacidad. Si usas TPUs reservadas, puedes saltarte esta sección.

Para crear nodos de segmento de TPU en GKE, se necesita cuota de la API Compute Engine (compute.googleapis.com), no cuota de la API Cloud TPU (tpu.googleapis.com). El nombre de la cuota es diferente en los pods de Autopilot normales y en los pods de Spot.

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

  1. Ve a la página Cuotas de la Google Cloud consola:

    Ir a Cuotas

  2. En el cuadro Filtrar, haz lo siguiente:

    1. Usa la siguiente tabla para seleccionar y copiar la propiedad de la cuota en función de la versión de TPU y el valor del selector de nodos cloud.google.com/gke-tpu-accelerator. Por ejemplo, si tienes previsto crear nodos de TPU v5e bajo demanda cuyo valor en el selector de nodos cloud.google.com/gke-tpu-accelerator sea tpu-v5-lite-podslice, introduce Name: TPU v5 Lite PodSlice chips.

      Versión de TPU: cloud.google.com/gke-tpu-accelerator Propiedad y nombre de la cuota de instancias bajo demanda Propiedad y nombre de la cuota de 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 (por ejemplo, ubicaciones) e introduce region: seguido del nombre de la región en la que tienes previsto crear las TPUs en GKE. Por ejemplo, introduce region:us-west4 si tienes previsto crear nodos de segmento de TPU en la zona us-west4-a. La cuota de TPU es regional, por lo que todas las zonas de la misma región consumen la misma cuota de TPU.

Si no hay ninguna cuota que coincida con el filtro que has introducido, significa que no se ha concedido al proyecto ninguna de las cuotas especificadas para la región que necesitas y debes solicitar un ajuste de cuota de TPU.

Cuando se crea una reserva de TPU, tanto el límite como los valores de uso actuales de la cuota correspondiente aumentan en el número de chips de la reserva de TPU. Por ejemplo, si se crea una reserva de 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 TPU v5 Lite PodSlice chips en la región correspondiente aumentan en 16.

Cuotas de recursos adicionales de GKE

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

  • Cuota de Persistent Disk SSD (GB): el disco de arranque de cada nodo de Kubernetes requiere 100 GB de forma predeterminada. Por lo tanto, esta cuota debe ser al menos tan alta como el producto del número máximo de nodos de GKE que preveas 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 ser al menos tan alta como el número máximo de nodos de GKE que preveas crear.
  • Asegúrate de que max-pods-per-node se corresponda con el intervalo de la subred: cada nodo de Kubernetes usa intervalos de IP secundarios para los pods. Por ejemplo, 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 intervalo no se debe compartir con ningún otro clúster. Para evitar agotar el intervalo de direcciones IP, usa la marca --max-pods-per-node para limitar el número de pods que se pueden programar en un nodo. La cuota de max-pods-per-node debe ser al menos tan alta como el número máximo de nodos de GKE que preveas crear.

Para solicitar un aumento de la cuota, consulta Solicitar un ajuste de cuota.

Preparar la aplicación de TPU

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

  1. Frameworks como JAX, PyTorch y TensorFlow acceden a las VMs de TPU mediante la biblioteca compartida libtpu. libtpu incluye el compilador XLA, el software de tiempo de ejecución de la TPU y el controlador de la TPU. Cada lanzamiento de PyTorch y JAX requiere una versión libtpu.so concreta. Para evitar conflictos entre versiones de paquetes, te recomendamos que uses una imagen de IA de JAX. Para usar TPUs en GKE, asegúrate de que utilizas las siguientes versiones: tpu7x
    Tipo de TPU Versión de libtpu.so
    Ironwood (TPU7x) (vista previa)
    TPU Trillium (v6e)
    tpu-v6e-slice
    TPU v5e
    tpu-v5-lite-podslice
    TPU v5p
    tpu-v5p-slice
    • Imagen de 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 nightly compilada el 23 de octubre del 2023.
    TPU v4
    tpu-v4-podslice
    TPU v3
    tpu-v3-slice
    tpu-v3-device
  2. En el manifiesto de tu carga de trabajo, añade selectores de nodos de Kubernetes para asegurarte de que GKE programa tu carga de trabajo de TPU en el tipo de máquina y la topología de TPU que has definido:

      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)

    Haz los cambios siguientes:

    • TPU_ACCELERATOR: nombre del acelerador de TPU. Por ejemplo, usa tpu7x-standard-4t.
    • TPU_TOPOLOGY: la topología física de 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 Planificar TPUs en GKE.
    • WORKLOAD_POLICY: el nombre de la política de carga de trabajo que quieras usar para colocar tus pods de TPU. Este selector de nodos solo es necesario para Ironwood (TPU7x).

Una vez que hayas completado la preparación de la carga de trabajo, puedes ejecutar un trabajo que use las 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 la TPU en el campo spec.nodeSelector, así como el número de chips de TPU en la sección spec.containers.resources. Cuando despliegas la carga de trabajo, GKE aprovisiona automáticamente los nodos con la configuración de TPU correcta y coloca cada carga de trabajo en su propio nodo dedicado para asegurar el acceso completo a los recursos del nodo. Para obtener instrucciones, consulta el artículo Solicitar TPUs en una carga de trabajo.

    • Aprovisionar TPUs de forma centralizada con clases de computación personalizadas

    En las siguientes secciones se muestra cómo crear una ComputeClass personalizada y, a continuación, crear un trabajo que consuma las TPUs definidas en la ComputeClass.

    Crear un ComputeClass personalizado

    Los pasos para crear una ComputeClass personalizada que siga las reglas de las TPU varían en función de si usas Ironwood (TPU7x) o una versión anterior de TPU.

    Ironwood (TPU7x)

    1. Crea una política de carga de trabajo. Este paso solo es obligatorio si vas a crear un pool de nodos multihost, que depende de la topología que elijas. Si usas un grupo de nodos de un solo host, sáltate este paso.

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

      Haz los cambios siguientes:

      • WORKLOAD_POLICY_NAME: un nombre para tu política de carga de trabajo.
      • TPU_TOPOLOGY: 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: tu ID de proyecto Google Cloud .
      • REGION: la región de la política de cargas de trabajo. Una política de carga de trabajo es un recurso regional y puedes usarla en varios grupos de nodos.

    2. Guarda el siguiente archivo de 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 añadir lo siguiente specs a tu archivo de manifiesto ComputeClass:

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

      Haz los cambios siguientes:

      • RESERVATION_NAME: el nombre de la reserva de capacidad de Compute Engine.
      • RESERVATION_BLOCK_NAME: el nombre del bloque de reserva de capacidad de Compute Engine.
      • RESERVATION_SUB_BLOCK_NAME: el nombre del subbloque de reserva de capacidad de Compute Engine.

      Para obtener más información, consulta el artículo sobre cómo consumir recursos de zona reservados.

    Otras versiones de TPU

    Para aprovisionar TPUs de las versiones 3, 4, 5p, 5e o 6e (Trillium) mediante una ComputeClass personalizada configurada para TPUs, sigue estos pasos:

    1. Guarda el siguiente archivo de 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

      Haz los cambios siguientes:

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

    2. Despliega el ComputeClass:

      kubectl apply -f tpu-compute-class.yaml

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

    Crear un trabajo que consuma TPUs

    1. Guarda el siguiente archivo de 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

      Haz los cambios siguientes:

      • NUMBER_OF_CHIPS: número de chips de TPU que usará el contenedor. Debe ser el mismo valor para limits y requests, igual al valor de CHIP_COUNT en el elemento ComputeClass personalizado seleccionado.
      • MEMORY_SIZE: 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 utilices. Para obtener más información, consulta Mínimos y máximos de los aceleradores.
      • NUMBER_OF_CHIPS: número de chips de TPU que usará el contenedor. Debe ser el mismo valor para limits y requests.

    2. Desplegar el trabajo:

      kubectl create -f tpu-job.yaml

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

      • Aprovisiona nodos para ejecutar los pods. En función del tipo de TPU, la topología y las solicitudes de recursos que hayas especificado, estos nodos serán slices de un solo host o de varios hosts. En función de la disponibilidad de los recursos de TPU en la prioridad más alta, GKE puede recurrir a prioridades más bajas para maximizar la obtenibilidad.
      • Añade intolerancias 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 evitar que se te siga facturando eliminando los recursos que has creado:

      kubectl delete -f tpu-job.yaml

    Solicitar 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 la TPU
    • El número de chips de TPU de un contenedor de tu carga de trabajo

    Para ver una lista de las versiones y topologías de TPU compatibles, así como el número correspondiente de chips y nodos de TPU en un segmento, consulta Elegir la versión de TPU.

    Consideraciones sobre las solicitudes de TPU en cargas de trabajo

    Solo un contenedor de un pod puede usar TPUs. El número de chips de TPU que solicita un contenedor debe ser igual al número de chips de TPU conectados a un nodo de 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:

    • 4 chips, lo que crea dos nodos de varios hosts con 4 chips de TPU cada uno
    • 8 chips, lo que crea un nodo de un solo host con 8 chips de TPU

    Para maximizar la eficiencia de los costes, te recomendamos que siempre consumas todas las TPUs del segmento que solicites. Si solicitas una porción de varios hosts de dos nodos con 4 chips de TPU cada uno, debes implementar una carga de trabajo que se ejecute en ambos nodos y consuma los 8 chips de TPU de la porción.

    Crear una carga de trabajo que solicite TPUs

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

    1. Guarda el siguiente archivo de 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

      Haz los cambios siguientes:

      • TPU_TYPE: el tipo de TPU que se va a usar, como tpu-v4-podslice. Debe ser un valor admitido por GKE.
      • TOPOLOGY: la disposición de los chips de TPU en el slice, como 2x2x4. Debe ser una topología compatible con el tipo de TPU seleccionado.
      • NUMBER_OF_CHIPS: número de chips de TPU que usará el contenedor. Debe ser el mismo valor para limits y requests.
      • MEMORY_SIZE: 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 utilices. Para obtener más información, consulta Mínimos y máximos de los aceleradores.

      También puede cambiar los siguientes campos:

      • image: la imagen de IA de JAX que se va a usar. En el archivo de manifiesto de ejemplo, este campo se ha definido en la imagen de IA de JAX más reciente. Para definir otra versión, consulta la lista de imágenes de IA de JAX actuales.
      • runtimeClassname: gvisor: el ajuste que te permite ejecutar este pod en GKE Sandbox. Para usarlo, descomenta esta línea. GKE Sandbox es compatible con las TPUs versión v4 y posteriores. Para obtener más información, consulta GKE Sandbox.

    2. Desplegar el trabajo:

      kubectl create -f tpu-autopilot.yaml

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

      1. Aprovisiona nodos para ejecutar los pods. En función del tipo de TPU, la topología y las solicitudes de recursos que hayas especificado, estos nodos serán slices de un solo host o slices de varios hosts.
      2. Añade intolerancias 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 evitar que se te siga facturando eliminando la carga de trabajo que has creado:

      kubectl delete -f tpu-autopilot.yaml

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

    En TPU Trillium, puedes usar la programación de colecciones para agrupar nodos de slices de TPU. Al agrupar estos nodos de segmento de TPU, es más fácil ajustar el número de réplicas para satisfacer la demanda de la carga de trabajo. Google Cloud controla las actualizaciones de software para asegurarse de que siempre haya suficientes segmentos disponibles en la colección para servir el tráfico.

    TPU Trillium admite la programación de recolecciones para grupos de nodos de un solo host 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 recogida en función del tipo de segmento de TPU que utilices:

    • Slice de TPU de varios hosts: GKE agrupa los slices de TPU de varios hosts para formar una colección. Cada grupo de nodos de GKE es una réplica de esta colección. Para definir una colección, crea un segmento de TPU multihost y asigna un nombre único a la colección. Para añadir más porciones de TPU a la colección, crea otro grupo de nodos de porción de TPU multihost con el mismo nombre de colección y tipo de carga de trabajo.
    • Slice de TPU de un solo host: GKE considera todo el pool de nodos de slice de TPU de un solo host como una colección. Para añadir más porciones de TPU a la colección, puedes cambiar el tamaño del grupo de nodos de la porción de TPU de un solo host.

    Para obtener información sobre las limitaciones de la programación de recogidas, consulta Cómo funciona la programación de recogidas.

    Usar un slice de TPU de varios hosts

    La programación de la recogida en nodos de segmentos de TPU de varios hosts está disponible en clústeres de Autopilot con la versión 1.31.2-gke.1537000 y posteriores. Los nodos de las slices de TPU de varios hosts con una topología 2x4 solo se admiten en la versión 1.31.2-gke.1115000 o posteriores. Para crear nodos de segmento de TPU multihost y agruparlos en una colección, añade 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 de clúster. El valor de la etiqueta cloud.google.com/gke-nodepool-group-name debe cumplir los requisitos de 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 un segmento de TPU multihost:

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

    Usar un slice de TPU de un solo host

    La programación de la recogida en nodos de segmentos de TPU de un solo host está disponible en clústeres de Autopilot con la versión 1.31.2-gke.1088000 y posteriores. Para crear nodos de segmento de TPU de un solo host y agruparlos en una colección, añade la etiqueta cloud.google.com/gke-workload-type:HIGH_AVAILABILITY a la especificación de tu carga de trabajo.

    Por ejemplo, el siguiente bloque de código define una colección con un slice 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
  ...

    Usar clases de computación personalizadas para desplegar una colección

    Para obtener más información sobre cómo desplegar una carga de trabajo que solicite la programación de cargas de trabajo y colecciones de TPU mediante clases de computación personalizadas, consulta los artículos Colección multihost de TPU y Definir el tipo de carga de trabajo para el SLO de TPU.

    Ejemplo: Mostrar el total de chips de TPU en un segmento de varias máquinas

    La siguiente carga de trabajo devuelve el número de chips de TPU de todos los nodos de un segmento de TPU de varios hosts. Para crear una porción multihost, la carga de trabajo tiene los siguientes parámetros:

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

    Esta versión y la topología seleccionadas dan como resultado un slice multihost.

    1. Guarda el siguiente archivo de 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 archivo de manifiesto: 
      kubectl create -f available-chips-multihost.yaml

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

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

      El resultado debería ser similar al siguiente:

      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

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

      El resultado debería ser similar al siguiente: 

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

    Ejemplo: Mostrar los chips de TPU en un solo nodo

    La siguiente carga de trabajo es un pod estático que muestra el número de chips de TPU conectados a un nodo específico. Para crear un nodo de un solo host, la carga de trabajo tiene los siguientes parámetros:

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

    Esta versión y topología dan como resultado un segmento de un solo host.

    1. Guarda el siguiente archivo de 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 archivo de manifiesto: 
      kubectl create -f available-chips-singlehost.yaml

      GKE aprovisiona nodos con ocho slices de TPU de un solo host que usan la versión 5e de la TPU. Cada nodo de TPU tiene ocho chips de TPU (slice de TPU de un solo host).

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

      El resultado debería ser similar al siguiente: 

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

    Observar y monitorizar TPUs

    Panel de control

    La observabilidad de los grupos 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 de control Estado de los grupos de nodos de TPU de GKE proporcionado por Cloud Monitoring:

    Ir a Estado del grupo de nodos de TPU de GKE

    Este panel te ofrece estadísticas completas sobre el estado de tus grupos de nodos de TPU multihost. Para obtener más información, consulta Monitorizar 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 métricas de observabilidad de la TPU, como el uso de la TPU, en el encabezado Aceleradores > TPU. Para obtener más información, consulta Ver métricas de observabilidad.

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

    Métricas de tiempo de ejecución

    En GKE 1.27.4-gke.900 o versiones posteriores, las cargas de trabajo de TPU que usen JAX 0.4.14 o versiones posteriores y especifiquen containerPort: 8431 exportan métricas de utilización de la TPU como métricas del sistema de GKE. En Cloud Monitoring están disponibles las siguientes métricas para monitorizar el rendimiento del tiempo de ejecución de tu carga de trabajo de TPU:

    • Ciclo de trabajo: porcentaje del tiempo durante el último periodo de muestreo (60 segundos) en el que los Tensor Cores han estado procesando activamente en un chip de TPU. Un porcentaje mayor significa un mejor uso de la TPU.
    • Memoria usada: cantidad de memoria del acelerador asignada en bytes. Se muestrea cada 60 segundos.
    • Memoria total: memoria total del acelerador en bytes. Se muestrea cada 60 segundos.

    Estas métricas se encuentran en el esquema de nodo 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

    Monitorizar métricas de estado de nodos y grupos de nodos de TPU

    Cuando un trabajo de entrenamiento tiene un error o finaliza con un fallo, puedes consultar las métricas relacionadas con la infraestructura subyacente para determinar si la interrupción se ha debido a un problema con el nodo o el pool de nodos subyacente.

    Estado del nodo

    En la versión 1.32.1-gke.1357001 de GKE o una posterior, 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 sobre las condiciones del nodo, como Ready, DiskPressure y MemoryPressure. El campo status muestra el estado notificado de la condición, que puede ser True, False o Unknown. Se trata de una métrica con el tipo de recurso monitorizado k8s_node.

    Esta consulta de PromQL muestra si un nodo concreto está Ready:

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

    Para solucionar problemas en un clúster, puede que quieras consultar los nodos que hayan mostrado otras condiciones:

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

    Puede que quieras consultar específicamente los nodos que no sean Ready:

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

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

    Puedes usar la siguiente consulta para conocer el estado de los nodos de 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 del recurso monitorizado k8s_node_pool muestra el estado de un grupo de nodos de GKE:

    • kubernetes.io/node_pool/status

    Esta métrica solo se registra en los grupos de nodos de TPU multihost.

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

    Para verificar si un grupo de nodos concreto 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 monitorizar el número de grupos de nodos de 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 nodos de TPU de varios hosts:

    • kubernetes.io/node_pool/multi_host/available

    La métrica tiene el valor 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 comprobar la disponibilidad de los grupos de nodos de 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}]))

    Número de interrupciones de nodos

    La siguiente métrica del sistema de GKE informa del número 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 explicar por qué se ha interrumpido 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 del 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}]))

    Tiempo de recuperación (TTR) de los grupos de nodos

    La siguiente métrica del sistema de GKE informa de la distribución de las duraciones del periodo de recuperación de los grupos de nodos de TPU de varios hosts de GKE:

    • kubernetes.io/node_pool/accelerator/times_to_recover

    Cada muestra registrada en esta métrica indica un evento de recuperación del grupo de nodos tras un periodo de inactividad.

    Esta métrica es útil para monitorizar el tiempo de recuperación y el tiempo entre interrupciones del grupo de nodos de 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]))

    Tiempo entre interrupciones (TBI) del grupo de nodos

    El tiempo entre interrupciones de un grupo de nodos mide el tiempo que se ejecuta tu infraestructura antes de sufrir una interrupción. Se calcula como la media de un periodo, en el que el numerador mide el tiempo total que tu infraestructura ha estado activa y el denominador mide las interrupciones totales de tu infraestructura.

    En el siguiente ejemplo de PromQL se muestra el tiempo medio entre interrupciones (MTBI) de 7 días del clúster en cuestión:

    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 GKE 1.28.1-gke.1066000 o versiones posteriores, las VMs de un segmento de TPU exportan métricas de utilización de TPU como métricas del sistema de GKE. En Cloud Monitoring están disponibles las siguientes métricas para monitorizar el rendimiento de tu host de TPU:

    • Uso de TensorCore: porcentaje actual de TensorCore que se está utilizando. El valor de Tensor Core es igual a la suma de las unidades de multiplicación de matrices (MXUs) más la unidad vectorial. El valor de utilización de Tensor Core es el resultado de dividir las operaciones de Tensor Core que se han realizado durante el periodo de muestreo anterior (60 segundos) entre el número admitido de operaciones de Tensor Core durante el mismo periodo. Cuanto mayor sea el valor, mejor será la utilización.
    • Uso del ancho de banda de la memoria: porcentaje actual del ancho de banda de la memoria del acelerador que se está usando. Se calcula dividiendo el ancho de banda de la memoria usado durante un periodo de muestra (60 s) entre el ancho de banda máximo admitido durante el mismo periodo de muestra.

    Estas métricas se encuentran en el esquema de nodos de Kubernetes (k8s_node) y de contenedores 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 las métricas de Kubernetes y las métricas del sistema de GKE.

    Almacenamiento de registros

    El agente de registro de GKE recoge los registros emitidos por los contenedores que se ejecutan en los nodos de GKE, incluidas las máquinas virtuales de TPU, los envía a Logging y se pueden ver 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 con tiempo de ejecución ampliado para disfrutar de un periodo de gracia de hasta siete días antes de que GKE finalice tus pods por reducciones o actualizaciones de nodos. Puedes usar ventanas de mantenimiento y exclusiones con pods de tiempo de ejecución ampliado para retrasar aún más las actualizaciones automáticas de nodos.
    • Usa reservas de capacidad para asegurarte de que tus cargas de trabajo reciben las TPUs solicitadas sin tener que esperar en una cola para que estén disponibles.

    Para saber cómo configurar las TPUs de Cloud en GKE, consulta los siguientes Google Cloud recursos: