Implementa cargas de trabajo de TPU en Autopilot para GKE

En esta página, se describe cómo implementar cargas de trabajo que usan aceleradores de CloudTPU (TPUs) en clústeres Autopilot de Google Kubernetes Engine (GKE). Ya deberías estar familiarizado con los siguientes conceptos:

  1. Introducción a Cloud TPU
  2. Arquitectura de sistemas de Cloud TPU
  3. Acerca de las TPUs en GKE

Cómo funcionan las TPUs en Autopilot

Si deseas usar TPUs en las cargas de trabajo de Autopilot, solicita una versión de TPU y una topología compatible para esa versión de TPU en el manifiesto de tu carga de trabajo. Luego, usa los campos resources.requests y resources.limits de Kubernetes a fin de especificar la cantidad de chips TPU que usará la carga de trabajo. 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 los recursos completos 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ífica
  3. Pods de tiempo de ejecución extendido

Planifica tu configuración de TPU

Antes de solicitar TPUs, decide la configuración que deseas según los requisitos de CPU y memoria de la carga de trabajo. Debes decidir lo siguiente:

  • Versión de TPU: la versión específica de Cloud TPU, como v5e.
  • Topología para la versión de TPU seleccionada: la disposición y la cantidad de TPUs

La versión de TPU y la topología que seleccionas determinan si GKE aprovisiona los nodos como porciones de host único o porciones de varios hosts. En porciones de host único, cada nodo es independiente de otros nodos TPU. En las porciones de varios hosts, GKE crea un grupo de nodos que tienen VMs de TPU interconectadas. Las porciones de varios hosts son atómicas, lo que significa que GKE escala todo el grupo interconectado de nodos como una sola unidad.

Para obtener información sobre las versiones de TPU disponibles, las topologías correspondientes, la capacidad de CPU y memoria y el tipo de porción resultante, consulta Elige una configuración de TPU de Autopilot.

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 Kubernetes Engine de Google.
    • Habilitar la API de Kubernetes Engine de Google
  • Si deseas usar Google Cloud CLI para esta tarea, instala y, luego, inicializa gcloud CLI. Si ya instalaste gcloud CLI, ejecuta gcloud components update para obtener la versión más reciente.
  • Asegúrate de tener un clúster de Autopilot que ejecute la versión 1.29.2-gke.1521000 o posterior de GKE.
  • Para usar las TPUs reservadas, asegúrate de tener una reserva de capacidad específica existente. Para obtener instrucciones, consulta Consume recursos zonales reservados.

Asegúrate de tener una cuota de TPU

Para crear nodos TPU, debes tener una 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 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 Spot Pods.

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. Selecciona la propiedad Servicio, ingresa API de Compute Engine y presiona Intro.

    2. Selecciona la propiedad Tipo y elige Cuota.

    3. Selecciona la propiedad Nombre y, luego, ingresa un nombre de cuota según el tipo de TPU que desees, de la siguiente manera:

      • TPU v5p (vista previa): Chips TPU v5p
      • TPU v5e (tpu-v5-lite-podslice): Chips PodSlice de TPU v5 Lite
      • TPU v5e (tpu-v5-lite-device): Chips de dispositivo TPU v5 Lite
      • TPU v4 (tpu-v4-podslice): Chips PodSlice de TPU v4

      Para los Spot Pods, selecciona la cuota "interrumpible" correspondiente.

    4. 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 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 deseada al proyecto y debes solicitar un aumento de la cuota de TPU.

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 usar las TPUs en GKE, asegúrate de usar las siguientes versiones:
    Tipo de TPU Versión de libtpu.so
    TPU v5e
    tpu-v5-lite-podslice
    tpu-v5-lite-device
    TPU v5p
    • 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
  2. Configura las siguientes variables de entorno para el contenedor que solicita los recursos TPU:
    • TPU_WORKER_ID: un número entero único para cada Pod. Este ID denota un ID de trabajador único en la porción de TPU. Los valores admitidos para este campo varían de cero a la cantidad de pods menos uno.
    • TPU_WORKER_HOSTNAMES: una lista separada por comas de nombres de host de VM de TPU o direcciones IP que deben comunicarse entre sí dentro de la porción. Debe haber un nombre de host o una dirección IP para cada VM de TPU en la porción. El TPU_WORKER_ID ordena y cero indexa la lista de direcciones IP o nombres de host.

      • GKE inserta de forma automática estas variables de entorno mediante un webhook de mutación cuando se crea un objeto Job con las propiedades completionMode: Indexed, subdomain, parallelism > 1 y requesting google.com/tpu. GKE agrega un Service sin interfaz gráfica a fin de que se agreguen los registros DNS para los Pods que respaldan el Service.

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

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
  • La cantidad de chips TPU para un contenedor en tu carga de trabajo

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

Consideraciones para las solicitudes de TPU en las cargas de trabajo

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

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

Como práctica recomendada para maximizar la rentabilidad, siempre consume todos los chips de la porción que solicites. Si solicitas una porción de varios hosts de dos nodos con 4 chips cada uno, deberías implementar una carga de trabajo que se ejecute en ambos nodos y consuma los 8 chips 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 Service sin interfaz gráfica que seleccione tu carga de trabajo por el nombre. Este Service sin interfaz gráfica permite que los Pods en diferentes nodos de la porción de varios hosts se comuniquen entre sí mediante la actualización de 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:
      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: python:3.10
        ports:
        - containerPort: 8471 # Default port using which TPU VMs communicate
        - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
        command:
        - bash
        - -c
        - |
          pip install 'jax[tpu]' -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
          python -c 'import jax; print("TPU cores:", jax.device_count())'
        resources:
          requests:
            cpu: 10
            memory: 500Gi
            google.com/tpu: NUMBER_OF_CHIPS
          limits:
            cpu: 10
            memory: 500Gi
            google.com/tpu: NUMBER_OF_CHIPS

    Reemplaza lo siguiente:

    • TPU_TYPE: el tipo de TPU que se usará, como tpu-v4-podslice. Debe ser un valor admitido por GKE.
    • TOPOLOGY: la disposición de chips TPU en la porción, como 2x2x4. Debe ser una topología compatible para 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 el trabajo:

    kubectl create -f tpu-autopilot.yaml

Cuando creas este trabajo, GKE hace lo siguiente de manera automática:

  1. Aprovisiona los 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 cualquiera de tus otras cargas de trabajo se ejecute en los mismos nodos que las cargas de trabajo de TPU.

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

La siguiente carga de trabajo muestra 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 selecció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
        cloud.google.com/gke-tpu-topology: 2x2x4
      containers:
      - name: tpu-job
        image: python:3.10
        ports:
        - containerPort: 8471 # Default port using which TPU VMs communicate
        - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
        command:
        - bash
        - -c
        - |
          pip install 'jax[tpu]' -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
          python -c 'import jax; print("TPU cores:", jax.device_count())'
        resources:
          requests:
            cpu: 10
            memory: 500Gi
            google.com/tpu: 4
          limits:
            cpu: 10
            memory: 500Gi
            google.com/tpu: 4
  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 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 al siguiente: 

    TPU cores: 16

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 selecció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
    cloud.google.com/gke-tpu-topology: 2x4
  containers:
  - name: tpu-job
    image: python:3.10
    ports:
    - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
    command:
    - bash
    - -c
    - |
      pip install 'jax[tpu]' -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
      python -c 'import jax; print("Total TPU chips:", jax.device_count())'
    resources:
      requests:
        google.com/tpu: 8
      limits:
        google.com/tpu: 8
  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 VM de TPU tiene ocho chips (porción de TPU de host único).

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

    El resultado es similar al siguiente: 

    Total TPU chips: 8

Observabilidad y métricas

Panel

En la página Clústeres de Kubernetes de la consola de Google Cloud, la pestaña Observabilidad muestra las métricas de observabilidad de TPU. Para obtener más información, consulta Métricas de observabilidad de GKE.

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 una versión de JAX 0.4.14 o posterior 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.
  • Uso de memoria: cantidad de memoria del acelerador asignada en bytes. Se tomaron muestras cada 60 segundos.
  • Memoria total: el total de memoria 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

Métricas de host

En la versión 1.28.1-gke.1066000 o posterior de GKE, la VM de TPU exporta 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 realizaron 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 del 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/container/node/tensorcore_utilization
  • kubernetes.io/container/node/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 las 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 extendidos para retrasar aún más las actualizaciones automáticas de los 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.