Acerca de las TPUs en GKE


Esta página presenta Cloud TPU y te muestra dónde encontrar información sobre el uso de Cloud TPU con Google Kubernetes Engine (GKE). Las unidades de procesamiento tensorial (TPU) son circuitos integrados personalizados específicos de aplicaciones (ASIC) de Google que se utilizan para acelerar las cargas de trabajo de aprendizaje automático que usan frameworks como TensorFlow ,PyTorch y JAX.

Antes de usar las TPU en GKE, te recomendamos que completes la siguiente ruta de aprendizaje:

  1. Obtén información sobre cómo funcionan los aceleradores de aprendizaje automático con la Introducción a Cloud TPU.
  2. Obtén información sobre la disponibilidad actual de la versión de TPU con la arquitectura del sistema de Cloud TPU.

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

Beneficios de usar las TPUs en GKE

GKE proporciona compatibilidad total con la administración del ciclo de vida de las VM de TPU, que incluye la creación, la configuración y la eliminación de las VM de TPU. GKE también admite VM Spot y el uso de Cloud TPU reservada. Entre los beneficios de usar las TPU en GKE, se incluyen los siguientes:

  • Entorno operativo coherente: Una sola plataforma para todas las cargas de trabajo de aprendizaje automático y otras.
  • Actualizaciones automáticas: GKE automatiza las actualizaciones de versiones, lo que reduce la sobrecarga operativa.
  • Balanceo de cargas: GKE distribuye la carga mediante la reducción de la latencia y la mejora de la confiabilidad.
  • Escalamiento responsivo: GKE escala de forma automática los recursos TPU para satisfacer las necesidades de tus cargas de trabajo.
  • Administración de recursos: ConKueue , un sistema de colas de trabajos nativos de Kubernetes, puedes administrar recursos en varios usuarios dentro de tu organización mediante la puesta en cola, la interrupción, la priorización y el uso compartido equitativo.

Terminología relacionada con TPU en GKE

En este documento, se usa la siguiente terminología relacionada con las TPUs:

  • Tipo de TPU: el tipo de Cloud TPU, como v5e.
  • Porción de TPU: Un conjunto de chips TPU interconectados.
  • Topología de TPU: La cantidad y la disposición física de los chips TPU en una porción de TPU.
  • Nodos de porción de TPU de host único: Uno o más nodos TPU independientes. Las TPUs conectadas a las VMs no están conectadas entre sí mediante interconexiones de alta velocidad.
  • Nodos de porción de TPU de varios hosts: dos o más nodos TPU interconectados. Las TPUs adjuntas a las VMs están conectadas mediante interconexiones de alta velocidad. Los nodos TPU de varios hosts tienen las siguientes características:
    • Atómico: GKE trata todos los nodos interconectados como una sola unidad. Durante las operaciones de escalamiento, GKE escala todo el conjunto de nodos a 0 y crea nodos nuevos. Si una máquina del grupo falla o finaliza, GKE vuelve a crear todo el conjunto de nodos como una unidad nueva.
    • Inmutable: No puedes agregar nodos nuevos de forma manual al conjunto de nodos interconectados.

Cómo funcionan las TPUs en GKE

La prioridad y la administración de recursos de Kubernetes tratan las VM de TPU de la misma manera que otros tipos de VM. Solicitas chips TPU con el nombre de recurso google.com/tpu:

    resources:
        requests:
          google.com/tpu: 4
        limits:
          google.com/tpu: 4

Cuando uses TPU en GKE, debes considerar las siguientes características de TPU:

  • Una VM de TPU puede acceder hasta 8 chips TPU.
  • Una porción de TPU contiene una cantidad fija de chips TPU que depende del tipo de máquina de TPU que elijas.
  • La cantidad de google.com/tpu solicitadas debe ser igual a la cantidad total de chips disponibles en el nodo TPU. Cualquier contenedor en un Pod de GKE que solicite TPU debe consumir todos los chips TPU en el nodo. De lo contrario, la implementación fallará, ya que GKE no puede consumir recursos TPU de forma parcial. Por ejemplo, consulta las siguientes situaciones:
    • El tipo de máquina ct5l-hightpu-8t tiene un único nodo TPU con 8 chips TPU. Un Pod de GKE que requiere ocho chips TPU se puede implementar en el nodo, pero dos Pods de GKE que requieren cuatro chips TPU cada uno no se puede implementar en un nodo.
    • El tipo de máquina ct5lp-hightpu-4t con una topología 2x4 contiene dos nodos TPU con cuatro chips cada uno, lo que da un total de ocho chips TPU. Un Pod de GKE que requiere ocho chips TPU se no puede implementar en ninguno de los nodos de este grupo de nodos, pero se pueden implementar dos Pods que requieran cuatro ocho chips TPU cada uno en los dos nodos del grupo de nodos.
    • TPU v5e con topología 4x4 tiene 16 chips en cuatro nodos. La carga de trabajo de GKE Autopilot que selecciona esta configuración debe solicitar cuatro chips en cada réplica, para un máximo de cuatro réplicas.
  • En los clústeres Standard, se pueden programar varios Pods de Kubernetes en una VM de TPU, pero solo un contenedor de cada Pod puede acceder a los chips TPU.
  • Cada clúster de Standard debe tener al menos un grupo de nodos que no sea TPU para crear Pods del sistema de Kubernetes, como kube-dns.
  • De forma predeterminada, los nodos TPU tienen el google.com/tputaint que evita que se programen Pods que no sean de TPU en ellos. El taint no garantiza que los recursos de TPU se usen por completo. Te permite ejecutar cargas de trabajo que no usan TPU en nodos que no son TPU, lo que libera procesamiento en nodos TPU para el código que usa TPU.
  • GKE recopila los registros emitidos por los contenedores que se ejecutan en las VMs de TPU. Para obtener más información, consulta Logging.
  • Las métricas de uso de TPU, como el rendimiento del entorno de ejecución, están disponibles en Cloud Monitoring. Para obtener más información, consulta Observabilidad y métricas.

Planifica tu configuración de TPU

Para trabajar con TPUs en los clústeres de GKE, debes decidir los siguientes parámetros:

  • Tipo de TPU: Los diferentes tipos de TPU tienen diferentes capacidades, como la relación precio-rendimiento, la capacidad de procesamiento del entrenamiento y la latencia de la entrega. Las VMs a las que están conectadas las TPUs a fin de tener diferentes capacidades de CPU y memoria.
  • Topología: el diseño de los chips en el tipo de TPU. Cada tipo de TPU es compatible con un diseño de chip 2D o 3D y arreglos específicos. Selecciona una topología que coincida con los requisitos de paralelismo de tu modelo.
  • Interconectividad de TPU: El tipo y la topología de TPU determinan si obtienes TPUs o no en varios nodos con interconexiones de alta velocidad. Esto determina si deseas nodos de porción de TPU de host único o de varios hosts.

    • Para modelos a gran escala, usa nodos de porción de TPU de varios hosts
    • Para modelos a pequeña escala, usa nodos de porción de TPU de host único
  • Modo privilegiado: los contenedores que se ejecutan en nodos de GKE en la versión 1.28 y anteriores deben habilitar el modo privilegiado. Las versiones 1.28 o posteriores no necesitan tener habilitado el modo privilegiado para acceder a las TPU.

Elige una configuración de TPU para el modo GKE Autopilot

En Autopilot, debes elegir un tipo de TPU y una topología, y especificarlos en tu manifiesto de Kubernetes. GKE administra el aprovisionamiento de nodos con TPUs y la programación de tus cargas de trabajo.

Disponibilidad de TPU en GKE Autopilot

Las TPUs están disponibles en regiones específicas de Google Cloud. Si quieres usar un tipo de TPU en tu carga de trabajo de GKE, tu clúster debe estar en una región compatible para ese tipo. Para obtener más detalles, consulta Regiones y zonas de TPU en la documentación de Cloud TPU.

Elige un tipo de TPU en Autopilot

Tipo de TPU Cantidad de CPUs virtuales Memoria (GiB) Cantidad de nodos de NUMA Cantidad máxima de chips TPU en una porción
TPU v5p (vista previa)
tpu-v5p-slice
208 448 2 6,144
TPU v5e
tpu-v5-lite-podslice
De 24 a 224 De 48 a 384 1 256
TPU v5e (solo host único)
tpu-v5-lite-device
De 24 a 224 De 48 a 384 De 1 a 2 8
TPU v4
tpu-v4-podslice
240 407 2 4,096

Revisa las especificaciones y los precios del chip en la documentación de Cloud TPU para ayudarte a decidir qué versión usar.

Elige una topología para Autopilot

La topología es el diseño físico de los chips en una porción de TPU. Según el tipo de TPU, la topología es bidimensional o tridimensional. Después de elegir un tipo de TPU, selecciona una topología que sea compatible con ese tipo de TPU. Los requisitos de paralelismo de tu modelo te ayudarán a decidir una topología. El producto de la topología es el número de chips en la porción. Por ejemplo:

  • 2x2x2 es una porción de TPU v4 de varios hosts de 8 chips.
  • 2x2 es una porción de TPU v5e de host único de 4 chips.

La topología que seleccionas para un tipo de TPU determina si obtienes nodos de porción de TPU de un solo host o de varios hosts. Si una topología específica admite porciones de host único y de hosts múltiples, la cantidad de chips TPU que solicita tu carga de trabajo determina el tipo de host que obtienes. Por ejemplo, TPU v5e (tpu-v5-lite-podslice) admite la topología de 2x4 como hosts único y multihost. Si solicitas 4 chips en tu carga de trabajo, obtendrás un nodo de varios hosts que tiene 4 chips. Si solicitas 8 chips en tu carga de trabajo, obtienes un nodo de host único que tiene 8 chips.

En la siguiente tabla, se describen las topologías compatibles en cada tipo de TPU y la cantidad de chips, la cantidad de nodos y el tipo de host:

Tipo de TPU Topología Chips TPU en una porción Cantidad de nodos Tipo de host Notas
TPU v5p (vista previa)
tpu-v5p-slice
2x2x1 4 1 Host único

Se admiten topologías personalizadas para más de 64 chips. Se aplican las siguientes condiciones:

  • Para más de 64 chips, {A}, {B} y {C} deben ser múltiplos de 4.
  • La topología más grande es 16x16x24
  • Los valores deben ser {A}{B}{C}, como 8x12x16.
2x2x2 8 2 Varios hosts
2x2x4 16 4 Varios hosts
2x4x4 32 8 Varios hosts
4x4x4 64 16 Varios hosts
{A}x{B}x{C} A*B*C (A*B*C/4)1 Varios hosts
TPU v5e
tpu-v5-lite-podslice
1x1 1 1 Host único No se admiten las topologías personalizadas.
2x2 4 1
2x4 8 1
2x4 2 1 Varios hosts
4x4 16 4
4x8 32 8
8x8 64 16
8x16 128 32
16x16 256 64
TPU v5e (solo host único)
tpu-v5-lite-device
1x1 1 1 Host único No se admiten las topologías personalizadas
2x2 4 1
2x4 8 1
TPU v4
tpu-v4-podslice
2x2x1 4 1 Host único

Se admiten topologías personalizadas para más de 64 chips. Se aplican las siguientes condiciones:

  • Para más de 64 chips, {A}, {B} y {C} deben ser múltiplos de 4.
  • La topología más grande es 12x16x16
  • Los valores deben ser {A}{B}{C}, como 8x12x16.
2x2x2 8 2 Varios hosts
2x2x4 16 4 Varios hosts
2x4x4 32 8 Varios hosts
4x4x4 64 16 Varios hosts
{A}x{B}x{C} A*B*C (A*B*C/4)1 Varios hosts
  1. Calculado por el producto de topología dividido por cuatro.

Después de elegir un tipo y una topología de TPU, especifícalos en el manifiesto de tu carga de trabajo. Para obtener instrucciones, consulta Implementa cargas de trabajo de TPU en GKE Autopilot.

Elige una configuración de TPU para el modo GKE Standard

En las siguientes secciones, se describen las características de TPU que evalúas cuando planificas y configuras tus cargas de trabajo de TPU en GKE. Para obtener detalles sobre las versiones disponibles, los tipos de máquinas, las topologías válidas y la cantidad de chips, consulta Asignación de configuraciones de TPU en este documento.

Disponibilidad de TPU en el modo GKE Standard

En la siguiente tabla, se muestra la disponibilidad de TPU según el tipo de máquina y la versión:

Versión de TPU Tipo de máquina que comienza con Versión mínima de GKE Disponibilidad Zona
TPU v4 ct4p- 1.26.1-gke.1500 Disponible de manera general us-central2-b
TPU v5e ct5l- 1.27.2-gke.2100 Disponible de manera general europe-west4-b
us-central1-a
TPU v5e ct5lp- 1.27.2-gke.2100 Disponible de manera general europe-west4-a1
us-central1-a1
us-east1-c
us-east5-b1
us-west1-c
us-west4-a
us-west4-b1
TPU v5p ct5p- 1.28.3-gke.1024000 Vista previa us-east1-d
us-east5-a
us-east5-c
  1. Cuando creas una TPU v5e con un tipo de máquina que comienza con ct5lp- en cualquiera de las zonas europe-west4-a, us-central1-a, us-east5-b o us-west4-b, no se admiten los grupos de nodos TPU v5e de host único. En otras palabras, cuando se crea un grupo de nodos TPU v5e en cualquiera de estas zonas, solo se admite el tipo de máquina ct5lp-hightpu-4t con una topología de al menos 2x4 o superior. Para crear una TPU v5e de host único en us-central1 o europe-west4, elige las zonas us-central1-a o europe-west4-b, respectivamente, y usa tipos de máquina que comiencen con ct5l-, como ct5l-hightpu-1t, ct5l-hightpu-4t o ct5l-hightpu-8t. Para crear una TPU v5e de host único en la región us-west4, elige la zona us-west4-a y usa tipos de máquinas que comienzan con ct5lp-, como ct5lp-hightpu-1t. Ten en cuenta que los tipos de máquinas que comienzan con ct5l- requieren una cuota diferente de los tipos de máquinas que comienzan con ct5lp-.

En las siguientes secciones, se describen las características de TPU que evalúas cuando planificas y configuras tus cargas de trabajo de TPU en GKE. Para obtener detalles sobre las versiones disponibles, los tipos de máquinas, las topologías válidas y la cantidad de chips, consulta Asignación de configuraciones de TPU en este documento.

Tipo de máquina

Los tipos de máquina que admiten recursos TPU siguen una convención de nombres que incluye la versión de TPU y la cantidad de chips por nodo, como ct<version>-hightpu-<node-chip-count>t. Por ejemplo, el tipo de máquina ct5lp-hightpu-1t es compatible con TPU v5e y contiene un chip TPU en total.

Topología

La topología define la disposición física de las TPU dentro de una porción de TPU. GKE aprovisiona una porción de TPU en topologías de dos o tridimensionales, según la versión de la TPU. Debes especificar una topología como la cantidad de chips TPU en cada dimensión:

  • En TPU v4 y v5p, que se programan en grupos de nodos de porciones de TPU de varios hosts, debes definir la topología en 3 tuplas ({A}x{B}x{C}), por ejemplo, 4x4x4. El producto de {A}x{B}x{C} define la cantidad de chips en el grupo de nodos. Por ejemplo, puedes definir topologías pequeñas de menos de 64 chips con formas de topología como 2x2x2, 2x2x4 o 2x4x4. Si usas topologías de más de 64 chips, los valores que asignes a {A}, {B} y {C} deben cumplir las siguientes condiciones:

    • {A}, {B} y {C} son múltiplos de cuatro.
    • La topología más grande compatible con v4 es 12x16x16 y v5p es 16x16x24.
    • Los valores asignados mantienen el patrón A ≤ B ≤ C. Por ejemplo, 4x4x8 o 8x8x8.

Asignación de la configuración de TPU

Usa la siguiente tabla para definir el tipo de máquina de TPU y la topología que se usará según tu caso de uso:

  • Para el entrenamiento o la inferencia de modelos a pequeña escala, usa TPU v4 o TPU v5e con grupos de nodos de porción de TPU de host único.
  • Para el entrenamiento o la inferencia de modelos a gran escala, usa TPU v4 o TPU v5e con grupos de nodos de porción de TPU de varios hosts.
Versión de TPU Tipo de máquina Topología Cantidad de chips TPU Cantidad de VM Tipo de grupo de nodos
TPU v4 ct4p-hightpu-4t 2x2x1 4 1 Host único
2x2x2 8 2 Varios hosts
2x2x4 16 4 Varios hosts
2x4x4 32 8 Varios hosts
{A}x{B}x{C} A*B*C (A*B*C/4)1 Varios hosts
TPU v5p ct5p-hightpu-4t 2x2x1 4 1 Host único
2x2x2 8 2 Varios hosts
2x2x4 16 4 Varios hosts
2x4x4 32 8 Varios hosts
{A}x{B}x{C} A*B*C (A*B*C/4)1 Varios hosts
TPU v5e ct5l-hightpu-1t 1x1 1 1 Host único
ct5l-hightpu-4t 2x2 4 1 Host único
ct5l-hightpu-8t 2x4 8 1 Host único
ct5lp-hightpu-1t 1x1 1 1 Host único
ct5lp-hightpu-4t 2x2 4 1 Host único
ct5lp-hightpu-8t 2x4 8 1 Host único
ct5lp-hightpu-4t 2x4 8 2 Varios hosts
4x4 16 4 Varios hosts
4x8 32 8 Varios hosts
8x8 64 16 Varios hosts
8x16 128 32 Varios hosts
16x16 256 64 Varios hosts
  1. Calculado por el producto de topología dividido por cuatro.

Características de TPU v5e

Las máquinas de TPU v5e tienen las siguientes características técnicas:

Tipo de máquina Cantidad de CPUs virtuales Memoria (GB) Cantidad de nodos de NUMA Probabilidad de ser adelantado
ct5l-hightpu-1t 24 48 1 ROI
ct5l-hightpu-4t 112 192 1 Media
ct5l-hightpu-8t 224 384 2 Inferior
ct5lp-hightpu-1t 24 48 1 ROI
ct5lp-hightpu-4t 112 192 1 Media
ct5lp-hightpu-8t 224 384 1 Bajo

Características de TPU v4 y v5p

Las máquinas TPU v4p y v5p tienen las siguientes características técnicas:

Tipo de máquina Cantidad de CPUs virtuales Memoria (GB) Cantidad de nodos de NUMA
ct4p-hightpu-4t 240 407 2
ct5p-hightpu-4t 208 448 2

Reserva de TPU

Las reservas de TPU están disponibles cuando se compra un compromiso. Cualquier reserva de TPU se puede usar con GKE.

Cuando crees un grupo de nodos TPU, usa las marcas --reservation y --reservation-affinity=specific para consumir una instancia de TPU reservada.

Ajuste de escala automático de TPU en GKE

GKE es compatible con unidades de procesamiento tensorial (TPU) para acelerar las cargas de trabajo de aprendizaje automático. El grupo de nodos de porción de TPU de host único y el grupo de nodos de porción de TPU de varios hosts admiten el ajuste de escala automático y el aprovisionamiento automático.

Con la marca --enable-autoprovisioning en un clúster de GKE, este último crea o borra grupos de nodos de porción de TPU de host único o multihost con una versión de TPU y una topología que cumple con los requisitos de las cargas de trabajo pendientes.

Cuando usas --enable-autoscaling, GKE escala el grupo de nodos según su tipo de la siguiente manera:

  • Grupo de nodos de porción de TPU de host único: GKE agrega o quita nodos TPU en el grupo de nodos existente. El grupo de nodos puede contener cualquier cantidad de nodos TPU entre cero y el tamaño máximo del grupo de nodos según lo determinado por --max-nodes y el valor --total-max-nodes. Cuando el grupo de nodos escala, todos los nodos TPU del grupo tienen el mismo tipo de máquina y topología. Para obtener más información sobre cómo crear un grupo de nodos de porción de TPU de host único, consulta Crea un grupo de nodos.

  • Grupo de nodos de porción de TPU de varios hosts: GKE escala verticalmente el grupo de nodos de forma atómica desde cero hasta la cantidad de nodos necesarios para satisfacer la topología de TPU. Por ejemplo, con un grupo de nodos TPU con un tipo de máquina ct5lp-hightpu-4t y una topología de 16x16, el grupo de nodos contiene 64 nodos. El escalador automático de GKE garantiza que este grupo de nodos tenga exactamente 0 o 64 nodos. Cuando se reduce la escala, GKE expulsa todos los Pods programados y vacía todo el grupo de nodos a cero. Para obtener más información sobre cómo crear un grupo de nodos de porción de TPU de varios hosts, consulta Crea un grupo de nodos.

Limitaciones

  • Para las reservas de capacidad, debes usar una reserva específica.
  • La asignación de costos y la medición de uso de GKE no incluyen ningún dato sobre el uso o los costos de la TPU v4 reservada.
  • TPU v5p y v5e no admiten la transmisión de imágenes o de extracción en us-east5.
  • El ajuste de escala automático de TPU v5p es compatible con los clústeres de GKE con planos de control que ejecutan al menos 1.29.2-gke.1035000 o al menos 1.28.7-gke.1020000.

Consideraciones de programación de la carga de trabajo

Las TPU tienen características únicas que requieren una programación y administración especiales de las cargas de trabajo en Kubernetes. En las siguientes secciones, se describen las prácticas recomendadas de programación.

CPU

Esta sección no se aplica a los clústeres de Autopilot porque GKE coloca cada Pod de TPU en su propio nodo.

Para programar una carga de trabajo en la CPU integrada en una VM de TPU, asegúrate de que tu Pod de GKE pueda tolerar el taint google.com/tpu. Si quieres que la carga de trabajo se implemente en nodos específicos, usa los selectores de nodos.

La prioridad y la administración de recursos de Kubernetes tratan las VM de TPU de la misma manera que otros tipos de VM. Para otorgar a los Pods que requieren prioridad de programación de TPU sobre otros Pods en los mismos nodos, solicita la CPU o la memoria máximas para esos Pods de TPU. Los Pods de prioridad baja deben hacer lo siguiente:

  1. Establecer solicitudes de CPU y memoria bajas a fin de asegurarte de que el nodo tenga suficientes recursos asignables para las cargas de trabajo de TPU. Para obtener más información, consulta Cómo Kubernetes aplica solicitudes y límites de recursos.
  2. No establezcas un límite de CPU (ilimitado) para garantizar que los pods puedan tener aumentos de actividad a fin de usar todos los ciclos sin usar.
  3. Establece un límite de memoria alto para asegurarte de que los Pods puedan usar la mayor parte de la memoria sin usar y, al mismo tiempo, mantener la estabilidad del nodo.

Si un Pod de Kubernetes no solicita CPU y memoria (incluso si solicita TPU), Kubernetes lo considera un Pod de mejor esfuerzo y no hay garantía de que necesitaba CPU y memoria. Solo los pods que solicitan CPU y memoria de forma explícita tienen esas garantías. Si deseas obtener más información, consulta Administración de recursos para Pods y contenedores.

Para obtener más información, consulta las prácticas recomendadas de Kubernetes sobre límites y solicitudes de recursos.

Reduce la interrupción de la carga de trabajo

Si usas TPU para entrenar un modelo de aprendizaje automático y la carga de trabajo se interrumpe, se pierde todo el trabajo realizado desde el último punto de control. Para disminuir la probabilidad de que se interrumpa tu carga de trabajo, haz lo siguiente:

  • Establece una prioridad más alta para este trabajo que para todos los demás trabajos: si los recursos son pocos, el programador de GKE interrumpe los trabajos de menor prioridad a fin de programar un trabajo de mayor prioridad. Esto también garantiza que tu carga de trabajo de mayor prioridad reciba todos los recursos que necesita (hasta el total de recursos disponibles en el clúster). Para obtener más información, consulta Interrupción y prioridad del pod.
  • Configura la exclusión de mantenimiento: una exclusión de mantenimiento es un período no recurrente durante el cual se prohíbe el mantenimiento automático. Para obtener más información, consulta Exclusiones de mantenimiento.
  • Usa Pods de tiempo de ejecución extendido en Autopilot: 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.

Controla las interrupciones debidas al mantenimiento de los nodos

Todos los nodos de GKE, incluidos los que alojan las VM de TPU, están sujetos a eventos de mantenimiento o a otras interrupciones que pueden causar el apagado de nodos. Puedes reducir las interrupciones de las cargas de trabajo que se ejecutan en clústeres de GKE con el plano de control que ejecuta la versión 1.29.1-gke.1425000 y posteriores. GKE alerta a los nodos de un cierre inminente con el envío de una señal SIGTERM al nodo hasta cinco minutos antes de la expulsión. Si tu carga de trabajo usa un framework de AA, como MaxText, Pax o JAX con Orbax, las cargas de trabajo pueden capturar la señal SIGTERM e iniciar un proceso de control.

Puedes configurar GKE para finalizar tus cargas de trabajo del AA de forma correcta con el tiempo máximo de notificación. En el manifiesto del Pod, establece el campo spec.terminationGracePeriodSeconds en 300 segundos (cinco minutos). GKE hace su mejor esfuerzo para finalizar estos Pods de forma correcta y ejecutar la acción de finalización que definiste, por ejemplo, guardar un estado de entrenamiento. GKE respeta cualquier configuración de hasta cinco minutos para la configuración PodDisruptionBudget o terminationGracePeriodSeconds.

El campo spec.terminationGracePeriodSeconds solo controla las interrupciones debidas a eventos de mantenimiento y desfragmentación que ocurren en las VM no interrumpibles. GKE no controla las interrupciones involuntarias, como las fallas de hardware.

Para obtener más información, consulta Configura la finalización correcta de los nodos TPU.

Maximiza el uso de TPU

Si deseas maximizar tu inversión en TPU, programa una combinación de prioridades del trabajo y pásalas en cola para maximizar la cantidad de tiempo que operan tus TPU. Si quieres programar e interrumpir los niveles de un trabajo, debes usar un complemento de Kubernetes que organice los trabajos en colas. Te recomendamos usar Kueue para ese caso de uso.

¿Qué sigue?