Se usó la API de Cloud Translation para traducir esta página.

Implementa cargas de trabajo de TPU en GKE Autopilot

Autopilot

En esta página, se describe cómo acelerar las cargas de trabajo de aprendizaje automático (AA) mediante el uso de aceleradores de Cloud TPU (TPU) 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 manera óptima en GKE y supervisarlas después de la implementación.

Esta página está dirigida a administradores y operadores de la plataforma, especialistas en datos y AI, y desarrolladores de aplicaciones que deseen 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 las que hacemos referencia en el contenido de Google Cloud , consulta Tareas y roles comunes de los usuarios de GKE Enterprise.

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, especifica 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.
Es 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 los 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:

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 el 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 gcloud components update para obtener la versión más reciente.
Nota: Para las instalaciones de gcloud CLI existentes, asegúrate de configurar las propiedades compute/region y compute/zone. Cuando configuras las ubicaciones predeterminadas, puedes evitar errores en la CLI de gcloud como el siguiente: One of [--zone, --region] must be supplied: Please specify location.

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 cuota para las TPU y otros recursos de GKE

En las siguientes secciones, te ayudaremos 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 reservación 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:

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

Ir a Cuotas

En la casilla Filtro, haz lo siguiente:

Selecciona la propiedad Servicio, ingresa API de Compute Engine y presiona Intro.
Selecciona la propiedad Tipo y elige Cuota.

Selecciona la propiedad Nombre y, luego, ingresa el nombre 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 TPU v5 Lite PodSlice chips.

Versión de TPU	`cloud.google.com/gke-tpu-accelerator`	Nombre de la cuota para instancias bajo demanda	Nombre de la cuota para las instancias Spot²
TPU v3	`tpu-v3-device`	`TPU v3 Device chips`	`Preemptible TPU v3 Device chips`
TPU v3	`tpu-v3-slice`	`TPU v3 PodSlice chips`	`Preemptible TPU v3 PodSlice chips`
TPU v4	`tpu-v4-podslice`	`TPU v4 PodSlice chips`	`Preemptible TPU v4 PodSlice chips`
TPU v5e	`tpu-v5-lite-device`	`TPU v5 Lite Device chips`	`Preemptible TPU v5 Lite Device chips`
TPU v5e	`tpu-v5-lite-podslice`	`TPU v5 Lite PodSlice chips`	`Preemptible TPU v5 Lite PodSlice chips`
TPU v5p	`tpu-v5p-slice`	`TPU v5p chips`	`Preemptible TPU v5p chips`
TPU Trillium	`tpu-v6e-slice`	`TPU v6e Slice chips`	`Preemptible TPU v6e Lite PodSlice chips`

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 aumento 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 relevante 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 subred: Cada nodo de Kubernetes usa rangos 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 rango no se debe compartir con ningún otro clúster. Para evitar que se agote 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 de 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 una cuota más alta.

Opciones para aprovisionar TPUs en GKE

GKE Autopilot te permite usar TPU directamente en cargas de trabajo individuales con nodeSelectors de Kubernetes.

Como alternativa, puedes solicitar TPU con clases de procesamiento personalizadas. Las clases de procesamiento personalizadas permiten que los administradores de la plataforma definan una jerarquía de configuraciones de nodos para que GKE priorice durante las decisiones de escalamiento de nodos, de modo que las cargas de trabajo se ejecuten en el hardware seleccionado.

Para obtener instrucciones, consulta la sección Aprovisiona TPU de forma centralizada con clases de procesamiento personalizadas.

Prepara tu aplicación de TPU

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

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	`libtpu.so` versión
TPU Trillium (v6e) `tpu-v6e-slice`	Versión recomendada de jax[tpu]: v0.4.9 o posterior Versión recomendada de torchxla[tpuvm]: v2.1.0 o posterior
TPU v5e `tpu-v5-lite-podslice` `tpu-v5-lite-device`	Versión recomendada de jax[tpu]: v0.4.9 o posterior Versión recomendada de torchxla[tpuvm]: v2.1.0 o posterior
TPU v5p `tpu-v5p-slice`	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`	Versión recomendada de jax[tpu]: v0.4.4 o posterior Versión recomendada de torchxla[tpuvm]: v2.0.0 o posterior
TPU v3 `tpu-v3-slice` `tpu-v3-device`	Versión recomendada de jax[tpu]: v0.4.4 o posterior Versión recomendada de torchxla[tpuvm]: v2.0.0 o posterior

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.

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
Es 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 y nodos de TPU compatibles en una porción, consulta Elige una configuración de TPU de Autopilot.

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 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 todas las TPUs en la porción que solicitas. Si solicitas una porción de varios hosts de dos nudos con 4 chips de TPU cada uno, debes implementar una carga de trabajo que se ejecute en ambos nudos y consuma los 8 chips de 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 sin cabeza que seleccione tu carga de trabajo por nombre. Este servicio sin interfaz gráfica permite que los Pods en diferentes nodos de la porción de varios hosts se comuniquen entre sí actualizando la configuración de DNS de Kubernetes para que apunte a los Pods de la carga de trabajo.

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

Implementa el trabajo:
```
kubectl create -f tpu-autopilot.yaml
```

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

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.
Agrega taints a los Pods y tolerancias a los nodos para evitar que ninguna de tus otras cargas de trabajo se ejecute en los mismos nodos que las cargas de trabajo de TPU.

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

En TPU Trillium, puedes usar la programación de colecciones para agrupar nodos de porción de TPU. Agrupar estos nodos de fragmento 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 fragmentos dentro de la colección disponibles para entregar tráfico.

TPU Trillium admite la programación de colecciones para grupos de nodos de host único y varios hosts que ejecutan cargas de trabajo de inferencia. A continuación, se describe cómo el comportamiento de programación de la recopilación depende del tipo de porción de TPU que usas:

Porción de TPU de varios hosts: GKE agrupa 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 información sobre las limitaciones 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 colecciones 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 versiones 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 una 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 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 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 colecciones en nodos de Slice de TPU de un solo host está disponible para los clústeres de Autopilot en la versión 1.31.2-gke.1088000 y versiones 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
  ...

Aprovisiona TPU de forma centralizada con clases de procesamiento personalizadas

Para aprovisionar TPUs con una clase de procesamiento personalizada, haz lo siguiente:

Asegúrate de que tu clúster tenga una clase de procesamiento personalizado disponible que seleccione las TPU. Para obtener información sobre cómo especificar TPU en clases de procesamiento personalizadas, consulta Reglas de TPU.

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_NAME
      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_CLASS_NAME: Es el nombre de la clase de procesamiento personalizado existente que especifica las TPU.
NUMBER_OF_CHIPS: Es la cantidad de chips TPU que usará el contenedor. Debe ser el mismo valor para limits y requests, igual que el valor del campo tpu.count en la clase de procesamiento personalizado seleccionada.

Implementa el trabajo:
```
kubectl create -f tpu-workload.yaml
```

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

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, GKE podría recurrir a prioridades más bajas para maximizar la disponibilidad.
Agrega taints a los Pods y tolerancias a los nodos para evitar que ninguna de tus otras cargas de trabajo se ejecute en los mismos nodos que las cargas de trabajo de TPU.

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

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.

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

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.

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

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 elección de versión y topología da como resultado una porción de host único.

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

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).
Obtén los registros del pod:
```
kubectl logs tpu-job-jax-v5
```
El resultado es similar a este:
```
Total TPU chips: 8
```

Observa y supervisa las TPU

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, 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/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 cargas de trabajo de TPU en Autopilot

Las siguientes recomendaciones podrían 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.