Usa el segmentado dinámico con un programador personalizado

En este documento, se describe cómo usar el segmentado dinámico interactuando directamente con los recursos personalizados de Slice. Puedes crear segmentos, supervisar los estados de las particiones y verificar el estado de los segmentos.

Antes de seguir estas instrucciones, asegúrate de comprender los conceptos de la segmentación dinámica.

¿Por qué usar el segmentado dinámico con un programador personalizado?

Usa tu propio programador para administrar los recursos personalizados de Slice si tienes requisitos de programación complejos o si deseas integrar el segmentado dinámico con tu infraestructura de programación existente.

En caso de que prefieras usar un programador en lugar de administrar los recursos personalizados de Slice directamente, GKE proporciona integración con Kueue y Topology Aware Scheduling (TAS). Para obtener más información, consulta Cómo programar segmentos dinámicos con Kueue y TAS.

Descripción general del flujo de trabajo

Para usar la segmentación dinámica con un programador personalizado, realiza las siguientes tareas que se describen en este documento:

  1. Habilita el controlador de segmentos.
  2. Crea grupos de nodos con aprovisionamiento incremental.
  3. Crea recursos personalizados de Slice según los requisitos de tu carga de trabajo. Aplica el recurso personalizado de Slice a tu clúster.
  4. Supervisa los estados de las particiones y el estado de los segmentos.
  5. Borra el segmento cuando termines.

Para obtener más información sobre los campos y el estado del recurso personalizado Slice, consulta la información de referencia del recurso personalizado Slice.

Requisitos

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

  • Usa un clúster estándar en la versión 1.35.2-gke.1842000 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 la gcloud CLI, ejecuta el comando gcloud components update para obtener la versión más reciente. Es posible que las versiones anteriores de la gcloud CLI no admitan la ejecución de los comandos que se describen en este documento.

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 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 del destino 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.

Cómo segmentar 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.

  • 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.
  • Segmentación por 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 en el 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 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.

Cómo formar una segmentación dinámica

Después de crear los grupos de nodos, puedes formar una porción dinámica más grande creando un recurso personalizado de Slice. En lugar de Pods, usas un recurso personalizado de Slice para definir la topología especificada, que luego activa el controlador de segmentación.

Verifica el estado de los nodos y las particiones

  1. Para obtener los nombres de los nodos del grupo de nodos, ejecuta el siguiente comando:

    kubectl get nodes -l cloud.google.com/gke-nodepool=${NODE_POOL_NAME}
    

    El resultado es similar al siguiente:

    NAME                                 STATUS   ROLES    AGE    VERSION
    gke-np-status-update-7b4c890c-0jhp   Ready    <none>   2d1h   v1.35.1-gke.1396002
    gke-np-status-update-7b4c890c-377r   Ready    <none>   2d1h   v1.35.1-gke.1396002
    gke-np-status-update-7b4c890c-gb51   Ready    <none>   2d1h   v1.35.1-gke.1396002
    
  2. Verifica el modelo de aprovisionamiento del nodo:

    kubectl describe node NODE_NAME | grep "cloud.google.com/gke-accelerator-topology-mode"
    

    El resultado es similar al siguiente:

    cloud.google.com/gke-accelerator-topology-mode: PROVISION_ONLY
    

    Este valor coincide con el parámetro de configuración accelerator-topology-mode=provision_only definido cuando creaste la política de carga de trabajo.

  3. Recupera la información de la etiqueta del nodo:

    kubectl describe node NODE_NAME | grep "cloud.google.com/gke-tpu-partition-4x4x4-id"
    

    Reemplaza NODE_NAME por el nombre de uno de los nodos del grupo de nodos.

    El resultado es similar al siguiente:

    cloud.google.com/gke-tpu-partition-4x4x4-id=fba785f80d18552357dcdef6d3d16c27
    

    La anotación cloud.google.com/gke-tpu-partition-4x4x4-state indica si el nodo está disponible para formar un segmento dinámico. Esta etiqueta admite los siguientes valores:

    • HEALTHY: El nodo está en buen estado y es completamente funcional.
    • DEGRADED: El nodo está dañado, pero aún se puede usar para la formación de segmentos dinámicos.
    • UNHEALTHY: El nodo no funciona correctamente y no se puede usar para formar una segmentación.
    • UNSET: El estado no está definido debido a que no hay suficientes nodos en el grupo de nodos.
    • INCOMPLETE: No se aprovisionaron todos los nodos dentro de la partición.
  4. Verifica que el nodo incluya la anotación node.gke.io/created-by-mig:

    kubectl describe node NODE_NAME | grep "node.gke.io/created-by-mig"
    

    Reemplaza NODE_NAME por el nombre de uno de los nodos del grupo de nodos.

    El resultado es similar al siguiente:

    node.gke.io/created-by-mig: projects/735972712744/zones/us-central1-ai1a/team/string
    

    El resultado incluye la etiqueta node.gke.io/created-by-mig, que permite que el plano de control de GKE vincule los nodos de Kubernetes a sus recursos subyacentes de Compute Engine.

Crea un recurso personalizado de Slice

  1. Define el recurso personalizado de Slice:

    apiVersion: accelerator.gke.io/v1beta1
    kind: Slice
    metadata:
      # Name of the slice resource
      name: SLICE_NAME
    spec:
      # Specify the type of accelerator for this slice
      type: "tpu7x"
      # Define the desired topology for the accelerator slice
      topology: TOPOLOGY
      partitionIds:
        - PARTITION_ID # Example: a9476d1b02bd4f4e75ffffae3bd23c01
        - PARTITION_ID_2
        # ... add more partition IDs as needed
    

    Reemplaza lo siguiente:

    • SLICE_NAME: Es un nombre para tu segmento. El nombre debe cumplir con las condiciones de metadata.name.
    • TOPOLOGY: Es la topología de la porción dinámica. La topología debe cumplir con las siguientes condiciones:
      • Cada dimensión de la topología solicitada debe ser un múltiplo de cuatro, por ejemplo, 4A x 4B x 4C.
      • Los tres valores de las dimensiones de la topología, AxBxC, deben estar en orden no decreciente (A ≤ B ≤ C). Por ejemplo, 4x4x8 es válido, pero 4x8x4 no lo es. Este orden ayuda a garantizar una formación de segmentos coherente y evita comportamientos inesperados.
      • El producto de los tres valores en las dimensiones de la topología, A × B × C, no debe exceder los 9,216.
    • PARTITION_ID: Es una lista de cadenas que identifican las particiones de 4x4x4 que componen el segmento. Calcula la cantidad de particiones según la cantidad total de chips, donde cada partición consta de 64 chips. La cantidad de elementos de tu lista spec.partitionIds debe coincidir exactamente con la cantidad calculada de particiones ((A × B × C) / 64). El partitionIds debe cumplir con las siguientes condiciones:
      • Cada partición debe asignarse a un subbloque de reserva.
      • Todos los subbloques asociados deben pertenecer a la misma reserva.
      • Todos los bloques asociados deben residir en la misma reserva.
      • Los grupos de nodos asociados deben tener todos los nodos en estado ready.
    • El valor del campo type debe ser tpu7x.
    • De manera opcional, para permitir que el controlador de segmentos vuelva a intentar automáticamente la formación de segmentos, puedes agregar la anotación slice.gke.io/retry-on-failure: "true" al recurso personalizado de segmentos. Si no se crea el segmento debido al motivo de estado SliceCreationFailed, el controlador volverá a intentarlo hasta que se forme correctamente.

    Por ejemplo, para crear un segmento 4x8x8, debes proporcionar cuatro IDs de partición únicos.

    apiVersion: accelerator.gke.io/v1beta1
    kind: Slice
    metadata:
        name: test-slice-example
        annotations:
          slice.gke.io/retry-on-failure: "true" # Optional annotation to retry slice formation
    spec:
        type: "tpu7x"
        topology: "4x8x8" # (4*8*8)/64 = 4 partitions
        partitionIds:
            - "p0"
            - "p1"
            - "p2"
            - "p3"
    
  2. Aplica el recurso personalizado de Slice:

    kubectl apply -f test-slice-example.yaml
    

    En este punto, GKE intenta crear la segmentación. Si ocurre uno de los siguientes problemas, falla la creación de la segmentación y el motivo del estado en el recurso personalizado Slice se actualiza a SliceCreationFailed o FAILED:

    • Si los nodos seleccionados en el recurso personalizado no existen, el motivo del estado es SliceCreationFailed.
    • Si otro segmento usa algún nodo del recurso personalizado, el motivo del estado es SliceCreationFailed.
    • Si los nodos del recurso personalizado no forman parte del mismo bloque de reserva, el motivo del estado es FAILED.
    • Si los nodos no están en la misma reserva, el motivo del estado es FAILED.
    • Si la topología no coincide con la cantidad de particiones, el motivo del estado es SliceCreationFailed.

    Para obtener más información sobre el estado del recurso personalizado Slice, consulta Estado de Slice.

Supervisa el estado del recurso personalizado de Slice

Para verificar el estado del recurso personalizado Slice, ejecuta el siguiente comando:

kubectl describe slice SLICE_NAME

Reemplaza SLICE_NAME por el nombre de la segmentación.

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-01-11T23:45:15Z
  Finalizers:
    accelerator.gke.io/slice-finalizer
  Generation:        1
  Resource Version:  1768175347356335006
  UID:               d0b71e5c-be3f-4788-aead-930c7afec4f2
Spec:
  Partition Ids:
    2c79463990ff67c4e3c2648666bfedfa
    ba898ffcac0ad0946e8ff036d771ee53
    [more partition IDs]
  Topology:  8x16x16
  Type:      tpu7x
Status:
  Conditions:
    Last Transition Time:  2026-01-11T23:45:38Z
    Message:               ""
    
    Reason:                FAILED
    
    Status:                False
    Type:                  Ready
Events:

El campo reason en el estado del recurso personalizado de Slice indica el estado actual de la segmentación. El ciclo de vida de un recurso personalizado de Slice sigue esta progresión:

  • SliceNotCreated: El controlador realiza la inicialización y las verificaciones de recursos.
    • Si no se cumplen los requisitos previos, el estado cambia a SliceCreationFailed.
    • Si la validación se aprueba, el estado cambia a ACTIVATING.
  • ACTIVATING: GKE está formando la segmentación.
    • Si se ejecuta de forma correcta, el estado cambia a ACTIVE.
    • Si los subbloques se degradan, pero el segmento es utilizable, el estado cambia a ACTIVE_DEGRADED.
    • Si falla la formación, el estado pasa a FAILED.
  • DEACTIVATING: Si se borra el recurso personalizado de Slice o se produce una falla crítica en un estado activo o con errores, se comienza a desmantelar el segmento.
  • INCOMPLETE: Es el paso final antes de que se borre el recurso por completo.

Para obtener más información sobre el estado del recurso personalizado Slice, consulta Estado de Slice.

Ejecuta cargas de trabajo en el segmentado dinámico

Cuando el recurso personalizado Slice está en estado ACTIVE, puedes ejecutar cargas de trabajo en él. En la siguiente sección, se incluyen ejemplos de cargas de trabajo que usan la segmentación dinámica. Las cargas de trabajo se envían como Jobs o JobSets.

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

En el siguiente ejemplo, se muestra una carga de trabajo que usa un solo segmento de subbloque.

  1. Guarda el siguiente manifiesto de muestra como tpu-job-jax-v7x-64.yaml:

    apiVersion: v1
    kind: Service
    metadata:
    name: headless-svc
    spec:
    clusterIP: None
    selector:
        job-name: tpu-job-jax-v7x-64
    ---
    apiVersion: batch/v1
    kind: Job
    metadata:
    name: tpu-job-jax-v7x-64
    spec:
    backoffLimit: 0
    completions: 16
    parallelism: 16
    completionMode: Indexed
    template:
        metadata:
        annotations:
            cloud.google.com/gke-tpu-slice-topology: 4x4x4
        spec:
        nodeSelector:
            cloud.google.com/gke-tpu-topology: 4x4x4
            cloud.google.com/gke-tpu-accelerator: tpu7x
            cloud.google.com/gke-tpu-slice: test-slice
        subdomain: headless-svc
        restartPolicy: Never
        containers:
        - name: tpu-job-jax
            env:
            - name: TPU_ACCELERATOR_TYPE
              value: tpu7x-128
            image: python:3.12
            securityContext:
            privileged: false
            command:
            - bash
            - -c
            - |
            set -ex
            pip install -U --pre jax jaxlib libtpu requests -i https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/ -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
            pip list
            python -c 'import jax; print("Total TPU devices (cores):", jax.device_count())'
            resources:
            requests:
                google.com/tpu: 4
            limits:
                google.com/tpu: 4
    

    En el manifiesto se muestra lo siguiente:

    • cloud.google.com/gke-tpu-slice-topology y cloud.google.com/gke-tpu-topology definen la topología del segmento dinámico.
    • env.value: tpu7x-128 es el tipo de acelerador de TPU y la cantidad total de núcleos en la porción. La cantidad de núcleos se calcula multiplicando las dimensiones de la topología por la cantidad de núcleos por chip. Por ejemplo, para una topología 4x4x4, el cálculo es 4 × 4 × 4 × 2 = 128, donde 2 es la cantidad de núcleos por chip para tpu7x (Ironwood [TPU7x]). Por lo tanto, el TPU_ACCELERATOR_TYPE es tpu7x-128.
  2. Aplica el manifiesto tpu-job-jax-v7x-64.yaml:

    kubectl apply -f tpu-job-jax-v7x-64.yaml
    

Ejemplo 2: Implementa una carga de trabajo en grupos de nodos de Multislice con JobSet

En este ejemplo, se muestra cómo implementar una carga de trabajo en grupos de nodos de Multislice con JobSet.

  1. Instala JobSet:

    kubectl apply --server-side -f https://github.com/kubernetes-sigs/jobset/releases/download/v0.10.1/manifests.yaml
    
  2. Guarda el siguiente manifiesto de muestra como tpu-multislice-jax.yaml:

    apiVersion: jobset.x-k8s.io/v1alpha2
    kind: JobSet
    metadata:
      name: tpu-multislice-jax
      annotations:
        alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-tpu-slice
    spec:
      failurePolicy:
        maxRestarts: 3
      replicatedJobs:
      - name: slice-job
        replicas: 2
        template:
          spec:
            parallelism: 16
            completions: 16
            backoffLimit: 0
            completionMode: Indexed
            template:
              metadata:
                annotations:
                  # The shape of the slice
                  cloud.google.com/gke-tpu-slice-topology: 4x4x4
              spec:
                hostNetwork: true
                dnsPolicy: ClusterFirstWithHostNet
                nodeSelector:
                  cloud.google.com/gke-tpu-topology: 4x4x4
                  cloud.google.com/gke-tpu-accelerator: tpu7x
                  # IMPORTANT: Do NOT put 'cloud.google.com/gke-tpu-slice' here manually.
                  # The exclusive-topology annotation handles the slice assignment automatically.
                containers:
                - name: jax-worker
                  image: python:3.12
                  securityContext:
                    privileged: true
                  ports:
                  - containerPort: 8471
                  command:
                  - bash
                  - -c
                  - |
                    set -ex
                    pip install -U --pre jax jaxlib libtpu requests -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
                    # Verify JobSet injected the specific slice ID for this worker
                    echo "JobSet Index: $JOB_COMPLETION_INDEX"
                    python -c 'import jax; print("Total TPU devices:", jax.device_count())'
                  resources:
                    requests:
                      google.com/tpu: 4
                    limits:
                      google.com/tpu: 4
    
  3. Aplica el manifiesto tpu-multislice-jax.yaml:

    kubectl apply -f tpu-multislice-jax.yaml
    

    En el manifiesto se muestra lo siguiente:

    • El campo replicas: 2 en replicatedJobs indica que JobSet crea dos Jobs separados, cada uno correspondiente a una porción de TPU 4x4x4.
    • La anotación alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-tpu-slice ayuda a garantizar que cada trabajo se asigne a una porción de TPU única.
    • La anotación cloud.google.com/gke-tpu-slice-topology: 4x4x4 define la topología de cada segmento dinámico.
    • En este ejemplo, la variable de entorno TPU_ACCELERATOR_TYPE no se configura de forma explícita, ya que JobSet controla la asignación de segmentos. El código de JAX detecta automáticamente los dispositivos TPU disponibles dentro de su segmento asignado.

Borra el segmento

  1. Borra el segmento:

    kubectl patch slice $SLICE_NAME --type json \
      -p='[{"op": "remove", "path": "/metadata/finalizers"}]'
    
  2. Verifica que se haya borrado el segmento:

    kubectl get slices
    

Inhabilita el Slice Controller

Para inhabilitar el controlador de segmentación, quítalo del clúster.

  1. Verifica que los recursos personalizados de Slice estén vacíos:

    kubectl get slice -A
    
  2. Actualiza el clúster para inhabilitar el controlador de segmentación:

    gcloud container clusters update ${CLUSTER_NAME} \
        --location=${REGION} \
        --no-enable-slice-controller
    
  3. Borra el recurso personalizado de Slice :

    kubectl delete crd slices.accelerator.gke.io
    
  4. Verifica que se haya borrado el recurso personalizado de Slice:

    kubectl get crd | grep slices.accelerator.gke.io
    
  5. Quita las etiquetas agregadas por el controlador de segmentación. Debes quitar estas etiquetas:

    • cloud.google.com/gke-tpu-slice
    • cloud.google.com/gke-tpu-topology
    1. Para quitarlo de un nodo específico, actualiza el nombre del nodo.
    export NODE_NAME="gke-tpu-bdac9600-3bdg"
    kubectl label node $NODE_NAME cloud.google.com/gke-tpu-slice- cloud.google.com/gke-tpu-slice-topology-
    
    1. Si quieres quitar estas etiquetas de todos los nodos de tu clúster, haz lo siguiente:
    kubectl label nodes --all cloud.google.com/gke-tpu-slice- cloud.google.com/gke-tpu-slice-topology-
    
    1. Verifica las etiquetas del nodo y confirma que estén vacías:
    export NODE_NAME="gke-tpu-bdac9600-3bdg"
    kubectl describe node $NODE_NAME | grep "cloud.google.com/gke-tpu-slice"
    

¿Qué sigue?