Configura las TPU con la segmentación dinámica de GKE

En este documento, se explica cómo usar el segmentado dinámico en Google Kubernetes Engine (GKE). El segmentado dinámico te permite configurar subbloques de TPU aprovisionados en diferentes topologías. Esta capacidad reduce la necesidad de volver a crear grupos de nodos, mejora la tolerancia a fallas, ya que permite la recuperación automática cuando se produce una falla, y optimiza el uso de recursos.

Este documento está dirigido a los ingenieros de IA/AA y a los administradores de plataformas que desean optimizar el uso de las TPU, reducir el tiempo de aprovisionamiento y mejorar la tolerancia a fallas para las cargas de trabajo de inferencia y entrenamiento a gran escala.

Antes de leer este documento, asegúrate de estar familiarizado con lo siguiente:

¿Qué es la segmentación dinámica?

El segmentado dinámico ofrece flexibilidad en la administración de la capacidad de Cloud TPU, ya que te permite desacoplar el aprovisionamiento de TPU. El segmentado dinámico implica el siguiente proceso:

  1. Aprovisionas recursos como unidades más pequeñas llamadas sub-bloques. Un subbloque es la unidad lógica fundamental de capacidad de Ironwood (TPU7x). En el caso de Ironwood (TPU7x), representa un grupo de 16 nodos de VMs de TPU con una topología 4x4x4 de chips TPU interconectados. En el contexto del modo All Capacity de TPU y el segmentado dinámico, un grupo de nodos se asigna directamente a un subbloque.
  2. Luego, el corte dinámico une estos subbloques en cortes más grandes.

Beneficios del segmentado dinámico

El segmentado dinámico te ayuda a lograr lo siguiente:

  • Reducir el tiempo de aprovisionamiento: El aprovisionamiento individual de subbloques permite un aprovisionamiento general más rápido, ya que minimiza el impacto de cualquier falla individual.
  • Reducción del tiempo de recuperación: Si se produce una falla en un chip de TPU, la unidad de falla más pequeña es un subbloque. El segmentado dinámico aísla los subbloques defectuosos para que las cargas de trabajo se puedan reprogramar en los subbloques en buen estado más rápido que si se aprovisionara un segmento grande completo.
  • Cambio de forma de la capacidad: Si tienes requisitos de carga de trabajo diversos, no necesitas borrar ni volver a crear grupos de nodos para los cambios de topología, lo que sería necesario sin la segmentación dinámica. En cambio, puedes reconfigurar de forma dinámica los grupos de nodos aprovisionados para que coincidan con las formas especificadas.

Elementos clave del segmentado dinámico

El segmentado dinámico introduce los siguientes conceptos clave:

  • Aprovisionamiento incremental de grupos de nodos: El segmentado dinámico usa el aprovisionamiento incremental, que es un modelo de aprovisionamiento tolerante a fallas de grupos de nodos. Este modelo convierte toda tu capacidad de TPU en grupos de nodos de grupos de 16 nodos de VMs de TPU.
  • Controlador de segmentación: Es un controlador de recursos personalizados de Kubernetes que se ejecuta dentro del plano de control de GKE y administra la segmentación dinámica. El controlador de segmentos administra el ciclo de vida de un recurso personalizado de Slice, que representa un segmento dinámico. El controlador de segmentos se encarga de crear, supervisar de forma continua y borrar el Slice. Cuando usas un programador, este dirige la creación y el borrado del recurso personalizado Slice.
  • Recurso personalizado de Slice: Une dinámicamente estos subbloques según la topología de TPU solicitada. Este proceso se basa en la reconfiguración dinámica de la red de OCS para conectar los grupos de nodo TPU, lo que ayuda a garantizar un rendimiento optimizado. Puedes inspeccionar el progreso o el estado de la formación de segmentos dinámicos a través de los campos de estado del recurso personalizado de Slice.

Requisitos

Para usar el segmentado dinámico en GKE, debes cumplir con los siguientes requisitos:

  • Usa un clúster Standard en la versión 1.35.0-gke.274500 o posterior, en el canal rápido.
  • Usa la versión de Ironwood (TPU7x).
  • Usa la imagen de Container-Optimized OS para tus nodos.
  • Para usar el aprovisionamiento incremental, usa las reservas del modo All Capacity. El modo All Capacity es una función habilitada por Cluster Director de TPU.

Antes de comenzar

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

  • Habilita la API de Google Kubernetes Engine.
    • Habilitar la API de Google Kubernetes Engine
  • Si deseas usar Google Cloud CLI para esta tarea, instala y, luego, inicializa gcloud CLI. Si ya instalaste gcloud CLI, ejecuta el comando gcloud components update para obtener la versión más reciente. Es posible que las versiones anteriores de gcloud CLI no admitan la ejecución de los comandos que se describen en este documento.

Limitaciones

  • Una sola segmentación debe usar subbloques dentro del mismo bloque de TPU en una reserva. Para usar subbloques en bloques de TPU, usa TPU Multislices.
  • El corte dinámico no admite topologías más pequeñas que 4x4x4.

Usa el segmentado dinámico en GKE con Kueue

En esta sección, se describe el flujo de trabajo para usar el segmentado dinámico en GKE.

  1. Consulta la topología y el estado de las reservas en el modo All Capacity.
  2. Habilita el controlador de segmentación en tu clúster.
  3. Crea nodo TPU TPU.
  4. Configura Kueue para crear un recurso personalizado de Slice.
  5. Ejecuta cargas de trabajo en segmentación dinámica con Kueue.
  6. Realiza una limpieza.

Habilita el controlador de segmentación

Para usar el segmentado dinámico, habilita el controlador de segmentación en tu clúster.

  1. Actualiza tu clúster:

    gcloud container clusters update CLUSTER_NAME \
    --location=LOCATION \
    --enable-slice-controller

    Reemplaza lo siguiente:

  2. Obtén credenciales para que puedas comunicarte con tu clúster con los comandos de kubectl:

    gcloud config set container/cluster CLUSTER_NAME
gcloud container clusters get-credentials CLUSTER_NAME \
    --location=LOCATION

  3. En el resultado del siguiente comando, verifica que esté presente el valor slices.accelerator.gke.io:

    kubectl get crd slices.accelerator.gke.io

    El resultado es similar a lo siguiente:

    slices.accelerator.gke.io                2026-01-09T23:58:02Z

Crea grupos de nodos con aprovisionamiento incremental

En esta sección, se describe cómo crear los grupos de nodo TPU con aprovisionamiento incremental. GKE convierte toda tu capacidad de TPU en grupos de nodos de grupos de 16 nodos de VMs de TPU o subbloques. GKE aprovisiona estos grupos de nodos incluso cuando no puede encontrar las 16 VMs en buen estado. Para ello, coloca nodos en las partes en buen estado de la máquina anfitrión y aprovisiona de forma incremental las máquinas en mal estado mientras se reparan.

Puedes segmentar tu grupo de nodos para que pertenezca a cualquiera de los siguientes elementos:

  • Es un bloque específico de TPUs que se expone en las reservas del modo All Capacity. El bloqueo de la segmentación permite que GKE cree el grupo de nodos en cualquier subbloque disponible dentro del bloque especificado.
  • Un subbloque específico o un grupo específico de 16 nodos de VMs de TPU para un control más detallado

Crear una política de cargas de trabajo

Para crear un grupo de nodos de porción de TPU con Ironwood (TPU7x), primero debes crear una política de cargas de trabajo con el campo accelerator-topology-mode establecido en provision_only. Este parámetro de configuración activa el proceso de aprovisionamiento incremental.

Crea una política de cargas de trabajo:

gcloud compute resource-policies create workload-policy WORKLOAD_POLICY_NAME \
        --project=PROJECT_ID \
        --region=REGION  \
        --type=HIGH_THROUGHPUT \
        --accelerator-topology=4x4x4 \
        --accelerator-topology-mode=provision_only

Reemplaza lo siguiente:

  • WORKLOAD_POLICY_NAME: Es un nombre para tu política de carga de trabajo.
  • PROJECT_ID: Es el ID del proyecto de Google Cloud .
  • REGION: Es la región de la política de carga de trabajo.

En este comando, haz lo siguiente:

  • Siempre configura el campo accelerator-topology como 4x4x4 para que coincida con la cantidad total de chips dentro de un solo subbloque.
  • Siempre establece el campo accelerator-topology-mode en provision_only para garantizar que se active el proceso de aprovisionamiento incremental. Cuando se configura el campo provision_only, el grupo de nodos aprovisiona nodos TPU sin formar vínculos de ICI ni de OCS.

Segmenta tu grupo de nodos para que pertenezca a un bloque o a un subbloque

Puedes segmentar tus reservas en el modo Toda la capacidad para que se muestren en bloques o subbloques específicos.

  • Orientación a un bloque: Cada grupo de nodos usa capacidad de un bloque especificado. GKE coloca el grupo de nodos dentro de un subbloque disponible en ese bloque. Debes crear tantos grupos de nodos como subbloques haya en el bloque que quieras usar.

  • Segmenta un subbloque: Cada grupo de nodos se asigna a un subbloque específico y disponible. Cuando se usa la segmentación por subbloque, GKE crea el grupo de nodos si al menos una VM está en buen estado. El aprovisionamiento incremental garantiza que todos los nodos se coloquen dentro del subbloque especificado.

Bloquear

  1. Para recuperar el nombre del bloque en una reserva y el recuento de subbloques disponibles en el bloque, completa los siguientes pasos en el documento Visualiza la topología y el estado de las reservas en el modo All Capacity:

    1. Para identificar el nombre del bloque, enumera todos los bloques de reserva y copia el valor del campo name:. Este valor es el nombre del bloque o BLOCK_NAME en este documento.

    2. Para determinar cuántos grupos de nodos crear, describe un bloque de reserva y, luego, identifica el valor en el campo reservationSubBlockCount. Este valor es la cantidad de subbloques disponibles. Por ejemplo, el valor reservationSubBlockCount: 4 indica que el bloque tiene cuatro subbloques disponibles y que debes crear cuatro grupos de nodos separados.

  2. Establece la ruta de reserva:

    export RESERVATION_PATH="projects/PROJECT_ID/reservations/RESERVATION_NAME/reservationBlocks/BLOCK_NAME"

    Reemplaza lo siguiente:

    • RESERVATION_NAME: Es el nombre de tu reserva de TPU.
    • BLOCK_NAME: Es el nombre del bloque.

  3. Crea un grupo de nodos para cada subbloque identificado en el paso anterior. Por ejemplo, si el recuento es 4, ejecuta este comando cuatro veces. Usa un nombre único para cada grupo de nodos.

    gcloud container node-pools create NODE_POOL_NAME \
      --cluster=CLUSTER_NAME \
      --node-locations=ZONE \
      --machine-type=tpu7x-standard-4t \
      --num-nodes=16 \
      --placement-policy=WORKLOAD_POLICY_NAME \
      --reservation-affinity=specific \
      --reservation=${RESERVATION_PATH}

    Reemplaza lo siguiente:

    • NODE_POOL_NAME: es el nombre de tu grupo de nodos nuevo.
    • CLUSTER_NAME: Es el nombre del clúster de GKE.
    • WORKLOAD_POLICY_NAME: Es el nombre de la política de cargas de trabajo que creaste.
    • ZONE: Es la zona del grupo de nodos, por ejemplo, us-central1-a.

En subbloque

  1. Para recuperar el nombre del bloque y los IDs de los subbloques disponibles, completa los siguientes pasos en el documento Visualiza la topología y el estado de las reservas en el modo All Capacity:

    1. Para identificar el nombre del bloque, enumera todos los bloques de reserva y copia el valor del campo name:. Este valor es el nombre del bloque o BLOCK_NAME en este documento.

    2. Para identificar el nombre de los subbloques, enumera todos los subbloques de un bloque y copia el valor del campo name: para cada entrada en reservationSubBlocks. Este valor es el nombre del subbloque o SUBBLOCK_NAME en este documento.

  2. Establece la ruta de reserva:

    export RESERVATION_PATH="projects/PROJECT_ID/reservations/RESERVATION_NAME/reservationBlocks/BLOCK_NAME/reservationSubBlocks/SUBBLOCK_NAME"

    Reemplaza lo siguiente:

    • RESERVATION_NAME: Es el nombre de tu reserva de TPU.
    • BLOCK_NAME: Es el nombre del bloque.
    • SUBBLOCK_NAME: Es el nombre del subbloque.

  3. Crea el grupo de nodos:

    gcloud container node-pools create NODE_POOL_NAME \
        --project=PROJECT_ID \
        --cluster=CLUSTER_NAME \
        --node-locations=ZONE \
        --machine-type=tpu7x-standard-4t \
        --num-nodes=16 \
        --placement-policy=WORKLOAD_POLICY_NAME \
        --reservation-affinity=specific \
        --reservation=${RESERVATION_PATH}

    Reemplaza lo siguiente:

    • NODE_POOL_NAME: Es un nombre único para el grupo de nodos nuevo, por ejemplo, sub-block-pool-1.
    • PROJECT_ID: Es el ID del proyecto de Google Cloud .
    • CLUSTER_NAME: el nombre del clúster de GKE.
    • ZONE: Es la zona del grupo de nodos, por ejemplo, us-central2-b.
    • WORKLOAD_POLICY_NAME: Es el nombre de la política de carga de trabajo que creaste.

En esta etapa, se crean los nodos, pero sus vínculos de interconexión entre chips (ICI) aún no están activos. Por lo tanto, no puedes ejecutar cargas de trabajo en estos grupos de nodos directamente.

Para habilitar todos los vínculos de ICI necesarios para formar la segmentación y permitir que se programen las cargas de trabajo, crea una segmentación dinámica con uno de los siguientes métodos:

  • Crea un recurso personalizado de Slice. En lugar de Pods, usas un recurso personalizado de Slice para definir la topología especificada, que activa el controlador de segmentación.
  • Programa cargas de trabajo de GKE con Kueue y TAS. Kueue controla automáticamente la creación y eliminación de recursos personalizados de Slice. Evita modificar manualmente los recursos personalizados de Slice creados por Kueue.

Crea un segmento dinámico con Kueue y TAS

En esta sección, programarás cargas de trabajo de GKE con Kueue y TAS.

Instala los recursos de JobSet y Kueue para la segmentación dinámica

  1. Instala JobSet:

    helm install jobset oci://registry.k8s.io/jobset/charts/jobset \
        --version 0.10.1 \
        --namespace jobset-system \
        --create-namespace \
        --set controller.resources.requests.cpu=4 \
        --set controller.resources.requests.memory=16Gi

  2. Instala Kueue:

    helm install kueue oci://registry.k8s.io/kueue/charts/kueue \
        --version 0.16.1 \
        --namespace kueue-system \
        --create-namespace \
        --wait \
        --set controllerManager.replicas=3 \
        --set controllerManager.manager.resources.requests.cpu=16 \
        --set controllerManager.manager.resources.requests.memory=64Gi

  3. Instala el controlador de segmentación de Kueue:

    kubectl apply -f https://gist.githubusercontent.com/mwysokin/cd90010d0d375b3bf57c536905692547/raw/506c36dd070f4ac222ba8a5e58ba28bbfcfa8ed3/kueue-slice-controller-v0.8.0-130.yaml

  4. Para configurar Kueue para el segmentado dinámico, guarda el siguiente manifiesto como dynamic-slice-topology.yaml:

    apiVersion: kueue.x-k8s.io/v1beta1
kind: Topology
metadata:
  name: superslice-topology
spec:
  levels:
  # Label to identify the physical block a sub-block belongs to.
  # Only sub-blocks from the same block can form a slice.
  - nodeLabel: cloud.google.com/gce-topology-block
  # Label to identify individual TPU sub-blocks (4x4x4 topology).
  - nodeLabel: cloud.google.com/gke-tpu-partition-4x4x4-id
  # Standard Kubernetes label for individual nodes.
  # Required to assign Pods to specific VMs.
  - nodeLabel: kubernetes.io/hostname
---
apiVersion: kueue.x-k8s.io/v1beta1
kind: ResourceFlavor
metadata:
  name: superslice-rf
spec:
  nodeLabels:
    cloud.google.com/gke-tpu-accelerator: tpu7x
  topologyName: superslice-topology
---
apiVersion: kueue.x-k8s.io/v1beta1
kind: AdmissionCheck
metadata:
  name: superslice-ac
spec:
  controllerName: accelerator.gke.io/slice
---
apiVersion: kueue.x-k8s.io/v1beta1
kind: ClusterQueue
metadata:
  name: cq
spec:
  namespaceSelector: {}
  admissionChecks:
  - superslice-ac
  resourceGroups:
  - coveredResources:
    - google.com/tpu
    flavors:
    - name: superslice-rf
      resources:
      - name: google.com/tpu
        nominalQuota: "999999"  # modeling unlimited quota
---
apiVersion: kueue.x-k8s.io/v1beta1
kind: LocalQueue
metadata:
  name: lq
  namespace: default
spec:
  clusterQueue: cq

  5. Aplica el manifiesto dynamic-slice-topology.yaml:

    kubectl apply -f dynamic-slice-topology.yaml

    En este manifiesto, configurarás Kueue para el segmentado dinámico definiendo los siguientes recursos:

    • Topología de porción dinámica de Ironwood (TPU7x) (superslice-topology): La topología define los niveles que Kueue considera cuando programa cargas de trabajo de segmentación dinámica. Estos niveles son los siguientes:
      • Etiqueta cloud.google.com/gce-topology-block: Este nivel es necesario para comprender qué subbloques pertenecen a qué bloques, ya que solo los subbloques del mismo bloque pueden formar un segmento.
      • Etiqueta cloud.google.com/gke-tpu-partition-4x4x4-id: Este nivel representa subbloques individuales de Ironwood (TPU7x) (topología 4x4x4).
      • Etiqueta kubernetes.io/hostname: Este nivel es obligatorio para asignar Pods a VMs específicas y para observar sus etiquetas y contaminaciones.
    • ResourceFlavor de SuperSlice de Ironwood (TPU7x) (superslice-rf): El ResourceFlavor para los subbloques de Ironwood (TPU7x) incluye la etiqueta cloud.google.com/gke-tpu-accelerator: tpu7x para que coincida con los nodos con máquinas de Ironwood (TPU7x).
    • SuperSlice AdmissionCheck (superslice-ac): Esta verificación de admisión le indica a Kueue que no programe una carga de trabajo hasta que el controlador de segmentación de GKE confirme que la segmentación se activó. Primero, se define la verificación de admisión y, luego, se agrega al ClusterQueue que controla las cargas de trabajo de segmentación dinámica.
    • ClusterQueue (cq) y LocalQueue (lq): Estos campos administran los recursos de google.com/tpu. La cq ClusterQueue incluye la verificación de admisión superslice-ac. El nominalQuota para google.com/tpu se puede configurar de dos maneras:
      • Cuota específica: Establece nominalQuota para que coincida con la capacidad existente para la administración de cuotas y el uso compartido equitativo.
      • Cuota ilimitada: Establece nominalQuota en un valor muy alto, como "999999", para modelar una cuota ilimitada. Para enfocarse en el TAS y el segmentado dinámico, esta configuración omite la funcionalidad de administración de cuotas de Kueue.

Define la selección del estado del subbloque

Además del estado y la disponibilidad estándar del nodo, GKE expone el estado específico de cada subbloque con la etiqueta cloud.google.com/gke-tpu-partition-4x4x4-state. Esta etiqueta permite que GKE tenga en cuenta los factores que influyen en la formación de la porción, como el estado de los vínculos de TPU.

Puedes definir el valor de la etiqueta cloud.google.com/gke-tpu-partition-4x4x4-state de la siguiente manera:

  • HEALTHY: La infraestructura está en buen estado.
  • DEGRADED: La infraestructura del subbloque está en un estado degradado, por ejemplo, debido a la degradación del vínculo del OCS. El subbloque aún puede formar un segmento, pero el rendimiento general podría ser menor en comparación con los subbloques en buen estado. Si puedes tolerar un rendimiento potencialmente degradado, puedes configurar tu carga de trabajo para que use subbloques DEGRADED con la afinidad de nodos, como se muestra en el ejemplo 3.
  • UNHEALTHY: El subbloque no está en buen estado y no puede formar una porción.

El webhook del controlador de segmentación de Kueue valida si una carga de trabajo incluye un requisito de estado de sub-bloque específico. Si no se indica ninguna preferencia, el webhook inserta una afinidad de nodo predeterminada.

El comportamiento es el siguiente:

  • Si hay un nodeSelector o un nodeAffinity que segmenta la etiqueta cloud.google.com/gke-tpu-partition-4x4x4-state, este no se modifica.

  • Si no existe tal configuración de etiquetas, el webhook inserta la siguiente afinidad de nodo predeterminada para garantizar que solo se usen los subbloques disponibles:

    nodeSelector:
  cloud.google.com/gke-tpu-partition-4x4x4-state: "HEALTHY"

En la siguiente sección, se incluyen ejemplos en los que la etiqueta cloud.google.com/gke-tpu-partition-4x4x4-state está configurada para especificar las diferentes configuraciones de estado de los subbloques.

Ejecuta cargas de trabajo de prueba en el segmentado dinámico con Kueue

En esta sección, se describe cómo implementar cargas de trabajo en el segmentado dinámico con Kueue y TAS. Incluye tres ejemplos que muestran cómo crear una carga de trabajo de segmentación dinámica y una carga de trabajo que consta de varias segmentaciones. Las cargas de trabajo se envían como JobSets.

Ejemplo 1: Una sola carga de trabajo usa un solo segmento dinámico

En el siguiente ejemplo, se describe cómo crear una carga de trabajo con una división que tiene una topología de 4x12x16, que se compone de 12 subbloques. La cantidad de Pods se calculó de la siguiente manera: (4 * 12 * 16) / 4 chips por nodo = 192 Pods.

  1. Guarda el siguiente manifiesto como big-super-slice.yaml:

    apiVersion: jobset.x-k8s.io/v1alpha2
kind: JobSet
metadata:
  name: big-super-slice
  labels:
    kueue.x-k8s.io/queue-name: lq
  annotations:
spec:
  replicatedJobs:
    - name: job-jax
      replicas: 1
      template:
        spec:
          parallelism: 192  # pods per slice calculation: 4*12*16 / 4 = 192
          completions: 192
          backoffLimit: 10
          template:
            metadata:
              annotations:
                cloud.google.com/gke-tpu-slice-topology: 4x12x16
            spec:
              tolerations:
                - key: "google.com/tpu"
                  operator: "Equal"
                  value: "present"
                  effect: "NoSchedule"
              nodeSelector:
                cloud.google.com/gke-tpu-accelerator: tpu7x
                cloud.google.com/gke-tpu-partition-4x4x4-state: "HEALTHY"
              containers:
                - name: jax
                  image: python:latest
                  command:
                    - bash
                    - -c
                    - |
                      printenv
                      pip install "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
                      python -c 'import jax; print("Global device count:", jax.device_count(), "Local device count:", jax.local_device_count())'
                  resources:
                    limits:
                      google.com/tpu: 4
              restartPolicy: Never

    En este manifiesto, las siguientes anotaciones le indican a Kueue las características y la topología de la división para configurar lo siguiente:

    • cloud.google.com/gke-tpu-slice-topology: Especifica "4x12x16" como la topología de la porción dinámica. Los requisitos para la topología del acelerador tpu7x incluyen las siguientes reglas:
      • La topología mínima es 4x4x4.
      • La topología debe ser una cadena tridimensional en el formato AxBxC. Por ejemplo, 4x8x8.
      • Cada dimensión (A, B y C) debe ser un múltiplo de cuatro.
      • Las dimensiones deben ordenarse de forma no decreciente: A <= B <= C. Por ejemplo, 4x8x4 no es válido; debería ser 4x4x8.
      • El producto de las dimensiones (ABC) no debe superar los 9,216.
      • Las topologías de segmentación más grandes admitidas pueden incluir hasta 32 subbloques. Por ejemplo, 8x16x16 con 32 subbloques, 8x12x20 con 30 subbloques o 12x12x12 con 27 subbloques se encuentran dentro de los límites aceptados.
    • cloud.google.com/gke-tpu-accelerator: tpu7x: Programa Pods en VMs que ejecutan Ironwood (TPU7x).
    • kueue.x-k8s.io/queue-name: Asigna el JobSet a una LocalQueue de Kueue.

  2. Aplica el manifiesto big-super-slice.yaml:

    kubectl apply -f big-super-slice.yaml

    Después de aplicar el manifiesto, Kueue crea un JobSet llamado big-super-slice. Luego, Kueue intenta formar una sola porción dinámica con una topología 4x12x16. Una vez que la segmentación está activa, Kueue admite la carga de trabajo y los 192 Pods se programan en los nodos para formar el segmento dinámico que ejecuta tus cargas de trabajo.

Ejemplo 2: Carga de trabajo con más de una réplica

En el siguiente ejemplo, se muestra cómo crear una carga de trabajo que usa dos segmentos dinámicos, cada uno compuesto por cuatro subbloques.

  1. Guarda el siguiente manifiesto como two-super-slices.yaml:

    apiVersion: jobset.x-k8s.io/v1alpha2
kind: JobSet
metadata:
name: two-super-slices
labels:
    kueue.x-k8s.io/queue-name: lq
annotations:
spec:
replicatedJobs:
    - name: job-jax
    replicas: 2
    template:
        spec:
        parallelism: 64  # Pods per slice calculation: (4*8*8) / 4 = 64
        completions: 64
        backoffLimit: 10
        template:
            metadata:
            annotations:
                cloud.google.com/gke-tpu-slice-topology: 4x8x8
            spec:
            tolerations:
                - key: "google.com/tpu"
                operator: "Equal"
                value: "present"
                effect: "NoSchedule"
            nodeSelector:
                cloud.google.com/gke-tpu-accelerator: tpu7x
                cloud.google.com/gke-tpu-partition-4x4x4-state: "HEALTHY"
            containers:
                - name: jax
                image: python:latest
                command:
                    - bash
                    - -c
                    - |
                    printenv
                    pip install "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
                    python -c 'import jax; print("Global device count:", jax.device_count(), "Local device count:", jax.local_device_count())'
                resources:
                    limits:
                    google.com/tpu: 4
            restartPolicy: Never

  2. Aplica el manifiesto two-super-slices.yaml:

    kubectl apply -f two-super-slices.yaml

En este manifiesto, debes configurar replicas: 2 en el campo replicatedJobs. Después de aplicar el manifiesto, Kueue intenta formar dos segmentos separados con una topología 4x8x8. Kueue crea un segmento dinámico para cada réplica definida en jobset.spec.replicatedJobs[].replicas. Si se especifican réplicas de n, Kueue crea n segmentos dinámicos para la carga de trabajo y espera a que todos los segmentos se activen antes de admitir la carga de trabajo.

Ejemplo 3: Carga de trabajo con un solo segmento dinámico y NodeAffinity

A partir de Kueue 0.15, Kueue admite NodeAffinity para la selección de nodos de TAS. Esta funcionalidad se puede usar para permitir que los nodos HEALTHY y DEGRADED formen parte de un segmento dinámico. En el siguiente ejemplo, se muestra cómo configurar una carga de trabajo con un solo segmento dinámico y NodeAffinity:

  1. Guarda el siguiente manifiesto como slice-8x8x8-na.yaml:

    apiVersion: jobset.x-k8s.io/v1alpha2
kind: JobSet
metadata:
  name: slice-8x8x8-na
  labels:
    kueue.x-k8s.io/queue-name: lq
spec:
  replicatedJobs:
    - name: rj1
      replicas: 1
      template:
        spec:
          parallelism: 128
          completions: 128
          backoffLimit: 10
          template:
            metadata:
              annotations:
                cloud.google.com/gke-tpu-slice-topology: 8x8x8
            spec:
              tolerations:
                - key: "google.com/tpu"
                  operator: "Equal"
                  value: "present"
                  effect: "NoSchedule"
              nodeSelector:
                cloud.google.com/gke-tpu-accelerator: tpu7x
              affinity:
                nodeAffinity:
                  requiredDuringSchedulingIgnoredDuringExecution:
                    nodeSelectorTerms:
                      - matchExpressions:
                          - key: cloud.google.com/gke-tpu-partition-4x4x4-state
                            operator: In
                            values:
                              - "HEALTHY"
                              - "DEGRADED"
              containers:
                - name: jax
                  image: python:latest
                  command:
                    - bash
                    - -c
                    - |
                      printenv
                      pip install "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
                      python -c 'import jax; print("Global device count:", jax.device_count(), "Local device count:", jax.local_device_count())'
                  resources:
                    limits:
                      google.com/tpu: 4
              restartPolicy: Never

  2. Aplica el manifiesto slice-8x8x8-na.yaml:

    kubectl apply -f slice-8x8x8-na.yaml

    Después de aplicar el manifiesto, Kueue crea un JobSet llamado slice-8x8x8-na. Luego, Kueue intenta formar una sola porción dinámica con una topología 8x8x8, lo que permite incluir nodos HEALTHY y DEGRADED debido al NodeAffinity especificado. Después de que la segmentación está activa, Kueue admite la carga de trabajo y los 128 Pods se programan en los nodos que forman la segmentación dinámica.

Supervisa el estado de la segmentación

Para verificar el estado de tus segmentos dinámicos, ejecuta el siguiente comando:

kubectl describe slice SLICE_NAME

Reemplaza SLICE_NAME por el nombre de tu segmento. Por lo general, el nombre de la división se deriva del nombre de JobSet y del índice de réplica. En el ejemplo 1, un segmento creado por Kueue tendría un nombre similar a default-jobset-big-super-slice-yyyyy-job-jax-0.

El resultado es similar a lo siguiente:

Name:         test-slice
Namespace:
Labels:       <none>
Annotations:  <none>
API Version:  accelerator.gke.io/v1beta1
Kind:         Slice
Metadata:
  Creation Timestamp:  2026-02-12T23:44:28Z
  Finalizers:
    accelerator.gke.io/slice-finalizer
  Generation:        1
  Resource Version:  1770939905695871008
  UID:               6dbbfe14-4486-4462-864d-e078d0ca8b5b
Spec:
  Partition Ids:
    5eae6a4f59d59cf30a9bf49de618eb2b
  Topology:  4x4x4
  Type:      tpu7x
Status:
  Conditions:
    Last Transition Time:  2026-02-12T23:45:05Z
    Message:
    Reason:                ACTIVE
    Status:                True
    Type:                  Ready
    Last Transition Time:  2026-02-12T23:45:05Z
    Message:               NodeLabelingCompleted
    Reason:                NodeLabelIsAdded
    Status:                True
    Type:                  NodeLabeled
Events:                    <none>

El nombre de la segmentación cumple con las siguientes reglas para garantizar la compatibilidad con las convenciones de nomenclatura de los recursos subyacentes de Compute Engine:

  • Plantilla: {namespace}-jobset-{jobset.metadata.name}-kueueHash[5-character]-{jobset.spec.replicatedJobs[].name}-sliceIndex.
  • Longitud: El nombre tiene 54 caracteres o menos. El controlador agrega un guion y un hash de clúster de 8 caracteres para crear nombres de recursos de Compute Engine, que tienen un límite de 63 caracteres.
  • Formato: El nombre coincide con la expresión regular ^[a-z]([-a-z0-9]*[a-z0-9])?$. El nombre tiene las siguientes características:
    • Comienza con una letra en minúscula.
    • Solo contiene letras minúsculas, números y guiones (-).
    • Termina con una letra minúscula o un número (no puede terminar con un guion).

Realiza una limpieza

Para evitar cargos inesperados, borra tus segmentos antes de borrar los grupos de nodos.

  1. Borra el JobSet. Esta acción hace que Kueue borre los recursos personalizados de Slice asociados.

    kubectl delete jobset JOBSET_NAME

    Reemplaza JOBSET_NAME por el nombre de tu JobSet, por ejemplo, big-super-slice.

  2. Borra el grupo de nodo TPU:

    gcloud container node-pools delete NODE_POOL_NAME \
    --cluster=CLUSTER_NAME \
    --location=LOCATION

Opcional: Usa la segmentación dinámica con tu propio programador

Este documento se centra en el uso de Kueue y TAS. Sin embargo, también puedes administrar la segmentación dinámica con tu propio programador personalizado. Si decides usar otro programador, sigue la información de referencia del recurso personalizado Slice.

¿Qué sigue?