Programa segmentos dinámicos con Kueue y TAS

En este documento, se explica cómo aprovisionar grupos de nodo TPU y programar segmentos dinámicos en Google Kubernetes Engine (GKE) con Kueue y la programación con reconocimiento de la topología (TAS).

También puedes usar el segmentado dinámico interactuando directamente con los recursos personalizados de segmentación. Para obtener más información, consulta Cómo usar el segmentado dinámico con un programador personalizado.

Antes de seguir estas instrucciones, asegúrate de comprender los conceptos del segmentado dinámico.

Requisitos

Para usar la segmentación dinámica en GKE, debes cumplir con los siguientes requisitos:

  • Usa un clúster estándar en el canal rápido en una de las siguientes versiones:
    • Para la configuración dinámica de supersegmentación (topologías iguales o mayores que 4x4x4), usa la versión 1.35.2-gke.1842000 o posterior.
    • Para la configuración dinámica de división en subsegmentos (topologías más pequeñas que 4x4x4), usa la versión 1.36.0-gke.3712000 o posterior.
  • 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 que habilita Cluster Director de TPU.
  • En el caso de la sublicencia dinámica, asegúrate de que tus nodos tengan eventos de mantenimiento pendientes. Supervisa tus instancias para detectar eventos de mantenimiento pendientes. Si alguno de tus nodos tiene un evento de mantenimiento pendiente con una hora de finalización entre el 18 y el 30 de septiembre de 2026, debes activar manualmente el evento de mantenimiento del host en esos nodos antes de poder usar la sublicencia.

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 la 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 indican en este documento.

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 segmentos en tu clúster.
  3. Instala Kueue, JobSet y LWS.
  4. Crea nodo TPU TPU.
  5. Configura Kueue para crear un recurso personalizado de Slice.
  6. Ejecuta cargas de trabajo en segmentación dinámica con Kueue.
  7. Realiza una limpieza.

Habilita el controlador de segmentos

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

Instala Kueue, JobSet y LWS

Si ya instalaste Kueue, JobSet y LWS, puedes omitir esta sección.

Instala Kueue

Sigue las instrucciones de la documentación de Kueue o ejecuta el siguiente comando:

kubectl apply --server-side -f https://github.com/kubernetes-sigs/kueue/releases/download/KUEUE_VERSION/manifests.yaml

Reemplaza KUEUE_VERSION por la versión de Kueue requerida según tus requisitos de topología. Para el subsegmento dinámico, usa Kueue v0.18.2 o una versión posterior. Para el supersegmentado dinámico, usa Kueue v0.16.6 o una versión posterior.

Instala JobSet

Sigue las instrucciones de la documentación de JobSet o ejecuta el siguiente comando:

kubectl apply --server-side -f https://github.com/kubernetes-sigs/jobset/releases/download/JOBSET_VERSION/manifests.yaml

Reemplaza JOBSET_VERSION por la versión de JobSet requerida según tus requisitos de topología. Para el subsegmento dinámico, usa JobSet v0.12.0 o una versión posterior. Para el supercorte dinámico, usa JobSet v0.11.1 o una versión posterior.

Instala LWS

LeaderWorkerSet (LWS) solo se requiere para la segmentación secundaria dinámica.

Sigue las instrucciones de la documentación de LWS o ejecuta el siguiente comando:

kubectl apply --server-side -f https://github.com/kubernetes-sigs/lws/releases/download/LWS_VERSION/manifests.yaml

Reemplaza LWS_VERSION por la versión de LWS requerida. Usa LWS v0.8.0 o una versión posterior.

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 que comprenden grupos de 16 nodos de VMs de Ironwood (TPU7x) o subbloques. GKE aprovisiona estos grupos de nodos incluso cuando no puede encontrar todas las 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:

  • Es un bloque específico de TPUs que se expone en las reservas del modo All Capacity. El bloqueo de 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 Ironwood (TPU7x) 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 subbloques o bloques específicos dentro de tu reserva en el modo Toda la capacidad.

  • Segmento objetivo: Cada grupo de nodos usa capacidad de un segmento 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 ayuda a garantizar 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 tu nuevo grupo de nodos, 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 el controlador de segmentación activa.
  • 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 de forma manual 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 el controlador de segmentación de Kueue

  1. Para instalar el controlador de segmentación de Kueue, guarda el siguiente manifiesto como slice-controller.yaml:

    # Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

apiVersion: v1
kind: Namespace
metadata:
  labels:
    app.kubernetes.io/managed-by: kustomize
    app.kubernetes.io/name: slice-controller
    control-plane: controller-manager
  name: slice-controller-system
---
apiVersion: v1
kind: ServiceAccount
metadata:
  labels:
    app.kubernetes.io/managed-by: kustomize
    app.kubernetes.io/name: slice-controller
    control-plane: controller-manager
  name: slice-controller-controller-manager
  namespace: slice-controller-system
---
apiVersion: rbac.authorization.k8s.io/v1
kind: Role
metadata:
  labels:
    app.kubernetes.io/managed-by: kustomize
    app.kubernetes.io/name: slice-controller
    control-plane: controller-manager
  name: slice-controller-leader-election-role
  namespace: slice-controller-system
rules:
- apiGroups:
  - ""
  resources:
  - configmaps
  verbs:
  - get
  - list
  - watch
  - create
  - update
  - patch
  - delete
- apiGroups:
  - coordination.k8s.io
  resources:
  - leases
  verbs:
  - get
  - list
  - watch
  - create
  - update
  - patch
  - delete
- apiGroups:
  - ""
  resources:
  - events
  verbs:
  - create
  - patch
---
apiVersion: rbac.authorization.k8s.io/v1
kind: ClusterRole
metadata:
  labels:
    control-plane: controller-manager
  name: slice-controller-manager-role
rules:
- apiGroups:
  - ""
  resources:
  - events
  verbs:
  - create
  - patch
  - update
  - watch
- apiGroups:
  - ""
  resources:
  - nodes
  - pods
  verbs:
  - get
  - list
  - watch
- apiGroups:
  - ""
  resources:
  - secrets
  verbs:
  - get
  - list
  - update
  - watch
- apiGroups:
  - accelerator.gke.io
  resources:
  - slices
  verbs:
  - create
  - delete
  - get
  - list
  - patch
  - update
  - watch
- apiGroups:
  - accelerator.gke.io
  resources:
  - slices/finalizers
  verbs:
  - update
- apiGroups:
  - admissionregistration.k8s.io
  resources:
  - mutatingwebhookconfigurations
  verbs:
  - get
  - list
  - update
  - watch
- apiGroups:
  - batch
  resources:
  - jobs
  verbs:
  - get
  - list
  - patch
  - update
  - watch
- apiGroups:
  - jobset.x-k8s.io
  resources:
  - jobsets
  verbs:
  - get
  - list
  - patch
  - update
  - watch
- apiGroups:
  - leaderworkerset.x-k8s.io
  resources:
  - leaderworkersets
  verbs:
  - get
  - list
  - patch
  - update
  - watch
- apiGroups:
  - kueue.x-k8s.io
  resources:
  - admissionchecks
  - admissionchecks/status
  - workloads/status
  verbs:
  - get
  - patch
  - update
- apiGroups:
  - kueue.x-k8s.io
  resources:
  - workloads
  verbs:
  - create
  - get
  - list
  - patch
  - update
  - watch
---
apiVersion: rbac.authorization.k8s.io/v1
kind: ClusterRole
metadata:
  labels:
    control-plane: controller-manager
  name: slice-controller-metrics-auth-role
rules:
- apiGroups:
  - authentication.k8s.io
  resources:
  - tokenreviews
  verbs:
  - create
- apiGroups:
  - authorization.k8s.io
  resources:
  - subjectaccessreviews
  verbs:
  - create
---
apiVersion: rbac.authorization.k8s.io/v1
kind: ClusterRole
metadata:
  labels:
    control-plane: controller-manager
  name: slice-controller-metrics-reader
rules:
- nonResourceURLs:
  - /metrics
  verbs:
  - get
---
apiVersion: rbac.authorization.k8s.io/v1
kind: RoleBinding
metadata:
  labels:
    app.kubernetes.io/managed-by: kustomize
    app.kubernetes.io/name: slice-controller
    control-plane: controller-manager
  name: slice-controller-leader-election-rolebinding
  namespace: slice-controller-system
roleRef:
  apiGroup: rbac.authorization.k8s.io
  kind: Role
  name: slice-controller-leader-election-role
subjects:
- kind: ServiceAccount
  name: slice-controller-controller-manager
  namespace: slice-controller-system
---
apiVersion: rbac.authorization.k8s.io/v1
kind: ClusterRoleBinding
metadata:
  labels:
    app.kubernetes.io/managed-by: kustomize
    app.kubernetes.io/name: slice-controller
    control-plane: controller-manager
  name: slice-controller-manager-rolebinding
roleRef:
  apiGroup: rbac.authorization.k8s.io
  kind: ClusterRole
  name: slice-controller-manager-role
subjects:
- kind: ServiceAccount
  name: slice-controller-controller-manager
  namespace: slice-controller-system
---
apiVersion: rbac.authorization.k8s.io/v1
kind: ClusterRoleBinding
metadata:
  labels:
    control-plane: controller-manager
  name: slice-controller-metrics-auth-rolebinding
roleRef:
  apiGroup: rbac.authorization.k8s.io
  kind: ClusterRole
  name: slice-controller-metrics-auth-role
subjects:
- kind: ServiceAccount
  name: slice-controller-controller-manager
  namespace: slice-controller-system
---
apiVersion: v1
kind: Secret
metadata:
  labels:
    control-plane: controller-manager
  name: slice-controller-webhook-server-cert
  namespace: slice-controller-system
---
apiVersion: v1
kind: Service
metadata:
  labels:
    app.kubernetes.io/managed-by: kustomize
    app.kubernetes.io/name: slice-controller
    control-plane: controller-manager
  name: slice-controller-controller-manager-metrics-service
  namespace: slice-controller-system
spec:
  ports:
  - name: https
    port: 8443
    protocol: TCP
    targetPort: 8443
  selector:
    app.kubernetes.io/name: slice-controller
    control-plane: controller-manager
---
apiVersion: v1
kind: Service
metadata:
  labels:
    control-plane: controller-manager
  name: slice-controller-webhook-service
  namespace: slice-controller-system
spec:
  ports:
  - port: 443
    protocol: TCP
    targetPort: 9443
  selector:
    control-plane: controller-manager
---
apiVersion: apps/v1
kind: Deployment
metadata:
  labels:
    app.kubernetes.io/managed-by: kustomize
    app.kubernetes.io/name: slice-controller
    control-plane: controller-manager
  name: slice-controller-controller-manager
  namespace: slice-controller-system
spec:
  replicas: 1
  selector:
    matchLabels:
      app.kubernetes.io/name: slice-controller
      control-plane: controller-manager
  template:
    metadata:
      annotations:
        kubectl.kubernetes.io/default-container: manager
      labels:
        app.kubernetes.io/name: slice-controller
        control-plane: controller-manager
    spec:
      containers:
      - args:
        - --metrics-bind-address=:8443
        - --leader-elect
        - --health-probe-bind-address=:8081
        - --zap-log-level=3
        - --feature-gates=UseRetryMechanismForSliceCreation=true
        - --activation-timeout=6m
        command:
        - /manager
        image: tpuongke/kueue-slice-controller:latest
        livenessProbe:
          httpGet:
            path: /healthz
            port: 8081
          initialDelaySeconds: 15
          periodSeconds: 20
        name: manager
        ports:
        - containerPort: 9443
          name: webhook-server
          protocol: TCP
        readinessProbe:
          httpGet:
            path: /readyz
            port: 8081
          initialDelaySeconds: 5
          periodSeconds: 10
        resources:
          limits:
            cpu: 12000m
            memory: 32Gi
          requests:
            cpu: 8000m
            memory: 16Gi
        securityContext:
          allowPrivilegeEscalation: false
          capabilities:
            drop:
            - ALL
        volumeMounts:
        - mountPath: /tmp/k8s-webhook-server/serving-certs
          name: cert
          readOnly: true
      securityContext:
        runAsNonRoot: true
        seccompProfile:
          type: RuntimeDefault
      serviceAccountName: slice-controller-controller-manager
      terminationGracePeriodSeconds: 10
      volumes:
      - name: cert
        secret:
          defaultMode: 420
          secretName: slice-controller-webhook-server-cert
---
apiVersion: admissionregistration.k8s.io/v1
kind: MutatingWebhookConfiguration
metadata:
  labels:
    control-plane: controller-manager
  name: slice-controller-mutating-webhook-configuration
webhooks:
- admissionReviewVersions:
  - v1
  clientConfig:
    service:
      name: slice-controller-webhook-service
      namespace: slice-controller-system
      path: /mutate-batch-v1-job
  failurePolicy: Fail
  name: mjob.kb.io
  rules:
  - apiGroups:
    - batch
    apiVersions:
    - v1
    operations:
    - CREATE
    resources:
    - jobs
  sideEffects: None
- admissionReviewVersions:
  - v1
  clientConfig:
    service:
      name: slice-controller-webhook-service
      namespace: slice-controller-system
      path: /mutate-jobset-x-k8s-io-v1alpha2-jobset
  failurePolicy: Fail
  name: mjobset.kb.io
  rules:
  - apiGroups:
    - jobset.x-k8s.io
    apiVersions:
    - v1alpha2
    operations:
    - CREATE
    resources:
    - jobsets
  sideEffects: None
- admissionReviewVersions:
  - v1
  clientConfig:
    service:
      name: slice-controller-webhook-service
      namespace: slice-controller-system
      path: /mutate-leaderworkerset-x-k8s-io-v1-leaderworkerset
  failurePolicy: Fail
  name: mleaderworkerset.kb.io
  rules:
  - apiGroups:
    - leaderworkerset.x-k8s.io
    apiVersions:
    - v1
    operations:
    - CREATE
    resources:
    - leaderworkersets
  sideEffects: None

  2. Aplica el manifiesto slice-controller.yaml:

    kubectl apply -f slice-controller.yaml

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

    # Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

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

  4. 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 observar sus etiquetas y contaminaciones.
    • Ironwood (TPU7x) SuperSlice ResourceFlavor (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 segmentos de GKE confirme que el segmento 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 campo nominalQuota para google.com/tpu se puede configurar de dos maneras:
      • Cuota específica: Establece el campo nominalQuota para que coincida con la capacidad existente para la administración de cuotas y el uso compartido equitativo.
      • Cuota ilimitada: Establece el campo 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 de la partición

Además del estado y la disponibilidad estándar de los nodos, GKE expone el estado específico de cada forma de partición con la etiqueta cloud.google.com/gke-tpu-partition-[shape]-state (en la que [shape] coincide con la forma del ID de partición, como 2x2x1, 2x2x2, 2x2x4, 2x4x4 o 4x4x4). Esta etiqueta permite que GKE tenga en cuenta los factores que influyen en la formación de la segmentación, como el estado de las vinculaciones de TPU. La configuración dinámica de segmentación secundaria (topologías más pequeñas que 4x4x4) requiere la versión 1.36.0-gke.3712000 o posterior de GKE.

Puedes definir el valor de la etiqueta de estado de la partición de la siguiente manera:

  • HEALTHY: La partición está en buen estado y es completamente funcional.
  • DEGRADED: La infraestructura de la partición se encuentra en un estado degradado, por ejemplo, debido a la degradación del vínculo del OCS. La partición aún puede formar un segmento, pero el rendimiento general podría ser más bajo en comparación con las particiones en buen estado. Este estado solo se aplica a la topología 4x4x4 de nivel superior. Las topologías más pequeñas no tienen un estado degradado.
  • UNHEALTHY: La partición no está en buen estado y no puede formar una división.
  • UNSET: El estado no está definido debido a que no se pudo inicializar el controlador de segmentos de GKE.
  • INCOMPLETE: No se aprovisionaron todos los nodos dentro de la partición.

El webhook del controlador de segmentación de Kueue valida si una carga de trabajo incluye un requisito específico de estado de la partición. 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-[shape]-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 las particiones disponibles:

    nodeAffinity:
  requiredDuringSchedulingIgnoredDuringExecution:
    nodeSelectorTerms:
    - matchExpressions:
      - key: cloud.google.com/gke-tpu-partition-4x4x4-state
        operator: In
        values:
        - "HEALTHY"
        - "DEGRADED"

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 salud de 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 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:

    # Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

apiVersion: jobset.x-k8s.io/v1alpha2
kind: JobSet
metadata:
  name: big-super-slice
  labels:
    kueue.x-k8s.io/queue-name: lq
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
              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:
      • Para la división en subsegmentos dinámica: Puedes especificar topologías más pequeñas que 4x4x4, como 2x2x1, 2x2x2, 2x2x4 o 2x4x4. Estas topologías más pequeñas requieren la versión 1.36.0-gke.3712000 de GKE o una posterior.
      • Para el supercorte dinámico: Puedes especificar topologías iguales o superiores a 4x4x4. Para la configuración dinámica de supersegmentación, cada dimensión de la topología solicitada debe ser un múltiplo de cuatro, por ejemplo, 4A x 4B x 4C.
      • La topología debe ser una cadena tridimensional en el formato AxBxC, por ejemplo, 4x8x8.
      • 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.
    • El webhook inyecta la afinidad de nodo predeterminada para garantizar que se usen los nodos HEALTHY y DEGRADED.

  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 porción está activa, Kueue admite la carga de trabajo y los 192 Pods se programan en los nodos para formar la porción dinámica 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 que segmentan solo los nodos HEALTHY.

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

    # Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

apiVersion: jobset.x-k8s.io/v1alpha2
kind: JobSet
metadata:
  name: two-super-slices
  labels:
    kueue.x-k8s.io/queue-name: lq
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, configura el campo replicas como 2 en la sección 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.

Supervisa la segmentación

Puedes ver el estado de la segmentación y supervisar sus métricas con las métricas del sistema de GKE.

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

Supervisa las métricas de la segmentación

Puedes supervisar las siguientes métricas del sistema de GKE que exponen la condición de una segmentación:

  • kubernetes.io/accelerator/slice/state
  • kubernetes.io/accelerator/partition/state
  • kubernetes.io/accelerator/slice/deformation_durations
  • kubernetes.io/accelerator/slice/formation_durations

Para obtener más información sobre las métricas, consulta Métricas del sistema de GKE.

Realiza una limpieza

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

  1. Borra el JobSet. Esta acción activa Kueue para que 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 el corte por porciones dinámico 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?