Configura le TPU con lo slicing dinamico di GKE

Questo documento spiega come utilizzare lo slicing dinamico in Google Kubernetes Engine (GKE). Il sezionamento dinamico ti consente di configurare i sottoblocchi TPU di cui è stato eseguito il provisioning in topologie diverse. Questa funzionalità riduce la necessità di ricreare i pool di nodi, migliora la tolleranza agli errori consentendo il ripristino automatico in caso di errore e ottimizza l'utilizzo delle risorse.

Questo documento è destinato agli ingegneri AI/ML e agli amministratori di piattaforme che vogliono ottimizzare l'utilizzo delle TPU, ridurre i tempi di provisioning e migliorare la tolleranza agli errori per i workload di addestramento e inferenza su larga scala.

Prima di leggere questo documento, assicurati di conoscere quanto segue:

Che cos'è la suddivisione dinamica?

Lo slicing dinamico offre flessibilità nella gestione della capacità di Cloud TPU consentendo di disaccoppiare il provisioning delle TPU. La suddivisione dinamica prevede la seguente procedura:

  1. Esegui il provisioning delle risorse come unità più piccole chiamate sottoblocchi. Un sub-block è l'unità logica fondamentale di base della capacità di Ironwood (TPU7x). Per Ironwood (TPU7x), rappresenta un gruppo di 16 nodi di VM TPU con una topologia 4x4x4 di chip TPU interconnessi. Nel contesto della modalità TPU All Capacity e dello slicing dinamico, un pool di nodi viene mappato direttamente a un sottoblocco.
  2. Il taglio dinamico unisce poi questi sottoblocchi in sezioni più grandi.

Vantaggi della suddivisione dinamica

La suddivisione dinamica ti aiuta a:

  • Ridurre il tempo di provisioning: il provisioning individuale dei sottoblocchi comporta un provisioning complessivo più rapido perché riduce al minimo l'impatto di un singolo errore.
  • Riduzione del tempo di recupero: se si verifica un errore del chip TPU, l'unità più piccola di errore è un sottoblocco. Lo slicing dinamico isola i blocchi secondari difettosi in modo che i workload possano essere riprogrammati sui blocchi secondari integri più rapidamente rispetto al riprovisioning di un'intera slice di grandi dimensioni.
  • Rimodellare la capacità: se hai requisiti di workload diversi, non devi eliminare e ricreare i pool di nodi per le modifiche alla topologia, il che sarebbe necessario senza lo slicing dinamico. In alternativa, puoi riconfigurare dinamicamente i pool di nodi di cui è stato eseguito il provisioning in modo che corrispondano alle forme specificate.

Elementi chiave della suddivisione dinamica

Il sezionamento dinamico introduce i seguenti concetti chiave:

  • Provisioning incrementale dei node pool:lo slicing dinamico utilizza il provisioning incrementale, un modello di provisioning dei node pool tollerante agli errori. Questo modello converte tutta la capacità TPU in node pool di gruppi di 16 nodi di VM TPU.
  • Controller di slice: un controller di risorse personalizzate Kubernetes in esecuzione all'interno del control plane GKE che gestisce lo slicing dinamico. Il controller slice gestisce il ciclo di vita di una risorsa personalizzata slice, che rappresenta una slice dinamica. Il controller di slice gestisce la creazione, il monitoraggio continuo e l'eliminazione dello slice. Quando utilizzi un programma di pianificazione, questo gestisce la creazione e l'eliminazione della risorsa personalizzata Slice.
  • Risorsa personalizzata della sezione: unisce dinamicamente questi sottoblocchi in base alla topologia TPU richiesta. Questo processo si basa sulla riconfigurazione dinamica della rete OCS per connettere i node pool TPU, il che contribuisce a garantire prestazioni ottimizzate. Puoi controllare l'avanzamento o l'integrità della formazione delle sezioni dinamiche controllando i campi di stato della risorsa personalizzata Slice.

Requisiti

Per utilizzare lo slicing dinamico in GKE, devi soddisfare i seguenti requisiti:

  • Utilizza un cluster Standard nella versione 1.35.0-gke.274500 o successive nel canale rapido.
  • Utilizza la versione Ironwood (TPU7x).
  • Utilizza l'immagine Container-Optimized OS per i tuoi nodi.
  • Per utilizzare il provisioning incrementale, utilizza le prenotazioni in modalità Tutta la capacità. La modalità Capacità è una funzionalità abilitata da TPU Cluster Director.

Prima di iniziare

Prima di iniziare, assicurati di aver eseguito le seguenti operazioni:

  • Attiva l'API Google Kubernetes Engine.
    • Attiva l'API Google Kubernetes Engine
  • Se vuoi utilizzare Google Cloud CLI per questa attività, installala e poi inizializza gcloud CLI. Se hai già installato gcloud CLI, scarica l'ultima versione eseguendo il comando gcloud components update. Le versioni precedenti di gcloud CLI potrebbero non supportare l'esecuzione dei comandi in questo documento.

Limitazioni

  • Una singola sezione deve utilizzare i sottoblocchi all'interno dello stesso blocco TPU in una prenotazione. Per utilizzare i sottoblocchi nei blocchi TPU, utilizza TPU Multislices.
  • Lo slicing dinamico non supporta topologie inferiori a 4x4x4.

Utilizzare lo slicing dinamico in GKE con Kueue

Questa sezione descrive il flusso di lavoro per l'utilizzo dello slicing dinamico in GKE.

  1. Visualizza la topologia e lo stato di integrità delle prenotazioni in modalità All Capacity.
  2. Abilita il controller di slice nel cluster.
  3. Crea pool di nodi TPU.
  4. Configura Kueue per creare una risorsa personalizzata Slice.
  5. Esegui i workload sul sezionamento dinamico con Kueue.
  6. Pulisci.

Attiva il controller delle sezioni

Per utilizzare la suddivisione dinamica, abilita il controller di suddivisione nel cluster.

  1. Aggiorna il cluster:

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

    Sostituisci quanto segue:

  2. Recupera le credenziali per comunicare con il cluster con i comandi kubectl:

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

  3. Nell'output del seguente comando, verifica che sia presente il valore slices.accelerator.gke.io:

    kubectl get crd slices.accelerator.gke.io

    L'output è simile al seguente:

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

Crea node pool con provisioning incrementale

Questa sezione descrive come creare i pool di nodi TPU con provisioning incrementale. GKE converte tutta la capacità TPU in node pool di gruppi di 16 nodi di VM TPU o in sottoblocchi. GKE esegue il provisioning di questi pool di nodi anche quando non riesce a trovare tutte le 16 VM integre posizionando i nodi sulle parti integre della macchina host ed eseguendo il provisioning incrementale delle macchine non integre durante la riparazione.

Puoi impostare come target del pool di nodi uno dei seguenti elementi:

  • Un blocco specifico di TPU, esposto nelle prenotazioni in modalità Tutta la capacità. Il targeting a blocchi consente a GKE di creare il pool di nodi in qualsiasi sottoblocco disponibile all'interno del blocco specificato.
  • Un sottoblocco specifico o un gruppo di 16 nodi specifico di VM TPU, per un controllo più granulare.

Crea una policy del workload

Per creare un pool di nodi di slice TPU con Ironwood (TPU7x), devi prima creare una policy del workload con il campo accelerator-topology-mode impostato su provision_only. Questa impostazione attiva il processo di provisioning incrementale.

Crea una policy del workload:

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

Sostituisci quanto segue:

  • WORKLOAD_POLICY_NAME: un nome per la policy del carico di lavoro.
  • PROJECT_ID: il tuo ID progetto Google Cloud .
  • REGION: la regione della policy del workload.

In questo comando, esegui le seguenti operazioni:

  • Imposta sempre il campo accelerator-topology su 4x4x4 in modo che corrisponda al numero totale di chip all'interno di un singolo sottoblocco.
  • Imposta sempre il campo accelerator-topology-mode su provision_only per assicurarti che venga attivato il processo di provisioning incrementale. Quando il campo provision_only è impostato, ilpool di nodil esegue il provisioning dei nodi TPU senza formare link ICI o OCS.

Imposta come target il tuo pool di nodi in modo che appartenga a un blocco o a un sottoblocco

Puoi scegliere come target blocchi o sottoblocchi specifici all'interno della prenotazione in modalità Tutta la capacità.

  • Target di un blocco:ogni pool di nodi utilizza la capacità di un blocco specificato. GKE inserisce ilpool di nodil all'interno di un blocco secondario disponibile in quel blocco. Devi creare tanti pool di nodi quanti sono i sottoblocchi nel blocco che vuoi utilizzare.

  • Target di un sottoblocco:ogni pool di nodi viene mappato a un sottoblocco specifico e disponibile. Quando utilizzi il targeting a livello di sottoblocco, GKE crea il pool di nodi se almeno una VM è integra. Il provisioning incrementale garantisce che tutti i nodi vengano posizionati all'interno del sottoblocco specificato.

Blocca

  1. Per recuperare il nome del blocco in una prenotazione e il conteggio dei sottoblocchi disponibili nel blocco, completa i seguenti passaggi nel documento Visualizzare la topologia e lo stato di integrità delle prenotazioni in modalità Tutta la capacità:

    1. Identifica il nome del blocco elencando tutti i blocchi di prenotazione e copiando il valore nel campo name:. Questo valore è il nome del blocco o di BLOCK_NAME in questo documento.

    2. Determina il numero di pool di nodi da creare descrivendo un blocco di prenotazione e identificando il valore nel campo reservationSubBlockCount. Questo valore è il numero di blocchi secondari disponibili. Ad esempio, il valore reservationSubBlockCount: 4 indica che il blocco ha quattro sottoblocchi disponibili e devi creare quattro pool di nodi separati.

  2. Imposta il percorso di prenotazione:

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

    Sostituisci quanto segue:

    • RESERVATION_NAME: il nome della prenotazione TPU.
    • BLOCK_NAME: il nome del blocco.

  3. Crea un pool di nodi per ogni sottoblocco identificato nel passaggio precedente. Ad esempio, se il conteggio è 4, esegui questo comando quattro volte. Utilizza un nome univoco per ogni pool di nodi.

    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}

    Sostituisci quanto segue:

    • NODE_POOL_NAME: il nome del nuovo pool di nodi.
    • CLUSTER_NAME: il nome del tuo cluster GKE.
    • WORKLOAD_POLICY_NAME: il nome della policy del workload che hai creato.
    • ZONE: la zona del pool di nodi, ad esempio us-central1-a.

Blocco secondario

  1. Per recuperare il nome del blocco e gli ID dei sottoblocchi disponibili, completa i seguenti passaggi nel documento Visualizzare la topologia e lo stato di integrità di tutte le prenotazioni in modalità Capacità:

    1. Per identificare il nome del blocco, elenca tutti i blocchi di prenotazione e copia il valore nel campo name:. Questo valore è il nome del blocco o di BLOCK_NAME in questo documento.

    2. Per identificare il nome dei sottoblocchi, elenca tutti i sottoblocchi di un blocco e copia il valore nel campo name: per ogni voce in reservationSubBlocks. Questo valore è il nome del sottoblocco o SUBBLOCK_NAME in questo documento.

  2. Imposta il percorso di prenotazione:

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

    Sostituisci quanto segue:

    • RESERVATION_NAME: il nome della prenotazione TPU.
    • BLOCK_NAME: il nome del blocco.
    • SUBBLOCK_NAME: il nome del sottoblocco.

  3. Crea il pool di nodi:

    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}

    Sostituisci quanto segue:

    • NODE_POOL_NAME: un nome univoco per il nuovo pool di nodi, ad esempio sub-block-pool-1.
    • PROJECT_ID: il tuo ID progetto Google Cloud .
    • CLUSTER_NAME: il nome del tuo cluster GKE.
    • ZONE: la zona del pool di nodi, ad esempio us-central2-b.
    • WORKLOAD_POLICY_NAME: il nome della policy del carico di lavoro che hai creato.

In questa fase, i nodi vengono creati, ma i relativi link Inter-Chip Interconnect (ICI) non sono ancora attivi. Pertanto, non puoi eseguire direttamente i carichi di lavoro su questi pool di nodi.

Per abilitare tutti i collegamenti ICI necessari per formare la sezione e consentire la pianificazione dei carichi di lavoro, crea una sezione dinamica utilizzando uno dei seguenti metodi:

  • Crea una risorsa personalizzata Slice. Anziché i pod, utilizzi una risorsa personalizzata Slice per definire la topologia specificata, che viene attivata dal controller slice.
  • Pianifica i carichi di lavoro GKE con Kueue e TAS. Kueue gestisce automaticamente la creazione e l'eliminazione delle risorse personalizzate Slice. Evita di modificare manualmente le risorse personalizzate Slice create da Kueue.

Creare una sezione dinamica con Kueue e TAS

In questa sezione, pianifichi i carichi di lavoro GKE con Kueue e TAS.

Installare le risorse JobSet e Kueue per lo slicing dinamico

  1. Installa 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. Installa 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. Installa il controller delle sezioni di Kueue:

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

  4. Per configurare Kueue per lo slicing dinamico, salva il seguente manifest come 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. Applica il manifest dynamic-slice-topology.yaml:

    kubectl apply -f dynamic-slice-topology.yaml

    In questo manifest, configuri Kueue per lo slicing dinamico definendo le seguenti risorse:

    • Topologia di slice dinamico Ironwood (TPU7x) (superslice-topology): la topologia definisce i livelli che Kueue prende in considerazione quando pianifica i carichi di lavoro di slicing dinamico. Questi livelli sono i seguenti:
      • Etichetta cloud.google.com/gce-topology-block: questo livello è necessario per capire quali blocchi secondari appartengono a quali blocchi, perché solo i blocchi secondari dello stesso blocco possono formare una sezione.
      • Etichetta cloud.google.com/gke-tpu-partition-4x4x4-id: questo livello rappresenta i singoli sottoblocchi Ironwood (TPU7x) (topologia 4x4x4).
      • Etichetta kubernetes.io/hostname: questo livello è necessario per assegnare i pod a VM specifiche e per osservare le relative etichette e i relativi taint.
    • ResourceFlavor SuperSlice Ironwood (TPU7x) (superslice-rf): il resource flavor per i sottoblocchi Ironwood (TPU7x) include l'etichetta cloud.google.com/gke-tpu-accelerator: tpu7x per abbinare i nodi alle macchine Ironwood (TPU7x).
    • SuperSlice AdmissionCheck (superslice-ac): questo controllo di ammissione indica a Kueue di non pianificare un carico di lavoro finché il controller delle slice GKE non conferma che la slice è diventata attiva. Il controllo di ammissione viene prima definito e poi aggiunto al ClusterQueue che gestisce i carichi di lavoro di slicing dinamico.
    • ClusterQueue (cq) e LocalQueue (lq): questi campi gestiscono le risorse google.com/tpu. cq ClusterQueue include il controllo di ammissione superslice-ac. Il nominalQuota per google.com/tpu può essere configurato in due modi:
      • Quota specifica: imposta nominalQuota in modo che corrisponda alla capacità esistente per la gestione della quota e la condivisione equa.
      • Quota illimitata: imposta nominalQuota su un valore molto alto, ad esempio "999999", per simulare una quota illimitata. Per concentrarsi su TAS e sulla suddivisione dinamica, questa configurazione ignora la funzionalità di gestione della quota di Kueue.

Definisci la selezione dell'integrità del blocco secondario

Oltre all'integrità e alla disponibilità standard dei nodi, GKE espone lo stato specifico di ogni sottoblocco utilizzando l'etichetta cloud.google.com/gke-tpu-partition-4x4x4-state. Questa etichetta consente a GKE di tenere conto dei fattori che influenzano la formazione delle sezioni, ad esempio lo stato dei link TPU.

Puoi definire il valore dell'etichetta cloud.google.com/gke-tpu-partition-4x4x4-state nel seguente modo:

  • HEALTHY: l'infrastruttura è integra.
  • DEGRADED: l'infrastruttura del blocco secondario è in uno stato di degrado, ad esempio a causa del degrado del collegamento OCS. Il sottoblocco può comunque formare una sezione, ma il rendimento complessivo potrebbe essere inferiore rispetto ai sottoblocchi integri. Se puoi tollerare un potenziale peggioramento delle prestazioni, puoi configurare il workload in modo che utilizzi i sottoblocchi DEGRADED utilizzando l'affinità dei nodi, come mostrato nell'esempio 3.
  • UNHEALTHY: il blocco secondario non è integro e non può formare una sezione.

Il webhook del controller delle sezioni di Kueue verifica se un carico di lavoro include un requisito di integrità specifico per i sottoblocchi. Se non viene indicata alcuna preferenza, il webhook inserisce un'affinità del nodo predefinita.

Il comportamento è il seguente:

  • Se è presente un nodeSelector o un nodeAffinity che ha come target l'etichetta cloud.google.com/gke-tpu-partition-4x4x4-state, rimane invariato.

  • Se non esiste una configurazione di etichetta di questo tipo, il webhook inserisce la seguente affinità dei nodi predefinita per garantire che vengano utilizzati solo i sottoblocchi disponibili:

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

La sezione seguente include esempi in cui l'etichetta cloud.google.com/gke-tpu-partition-4x4x4-state è configurata per specificare le diverse configurazioni di integrità dei sottoblocchi.

Esegui workload di test sullo slicing dinamico con Kueue

Questa sezione descrive come eseguire il deployment dei workload sul dynamic slicing con Kueue e TAS. Include tre esempi che mostrano come creare un workload di slice dinamico e un workload costituito da più slice. I carichi di lavoro vengono inviati come JobSet.

Esempio 1: un singolo workload utilizza una singola sezione dinamica

Il seguente esempio descrive come creare un workload utilizzando una sezione con una topologia 4x12x16, composta da 12 blocchi secondari. Il numero di pod è stato calcolato come: (4 * 12 * 16) / 4 chip per nodo = 192 pod.

  1. Salva il seguente manifest come 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

    In questo manifest, le seguenti annotazioni indicano a Kueue le caratteristiche e la topologia della slice da configurare:

    • cloud.google.com/gke-tpu-slice-topology: specifica "4x12x16" come topologia della sezione dinamica. I requisiti per la topologia dell'acceleratore tpu7x includono le seguenti regole:
      • La topologia minima è 4x4x4.
      • La topologia deve essere una stringa tridimensionale nel formato AxBxC. Ad esempio, 4x8x8.
      • Ogni dimensione (A, B e C) deve essere un multiplo di quattro.
      • Le dimensioni devono essere ordinate in ordine non decrescente: A <= B <= C. Ad esempio, 4x8x4 non è valido; deve essere 4x4x8.
      • Il prodotto delle dimensioni (ABC) non deve superare 9216.
      • Le topologie di slice più grandi supportate possono includere fino a 32 sub-blocchi. Ad esempio, 8x16x16 con 32 blocchi secondari, 8x12x20 con 30 blocchi secondari o 12x12x12 con 27 blocchi secondari rientrano nei limiti accettati.
    • cloud.google.com/gke-tpu-accelerator: tpu7x: pianifica i pod sulle VM che eseguono Ironwood (TPU7x).
    • kueue.x-k8s.io/queue-name: assegna JobSet a una LocalQueue di Kueue.

  2. Applica il manifest big-super-slice.yaml:

    kubectl apply -f big-super-slice.yaml

    Dopo aver applicato il manifest, Kueue crea un JobSet denominato big-super-slice. Kueue tenta quindi di formare una singola sezione dinamica con una topologia 4x12x16. Una volta attiva la sezione, Kueue ammette il workload e i 192 pod vengono pianificati sui nodi per formare la sezione dinamica che esegue i workload.

Esempio 2: workload con più di una replica

Il seguente esempio mostra come creare un workload che utilizza due slice dinamiche, ciascuna composta da quattro blocchi secondari.

  1. Salva il seguente manifest come 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. Applica il manifest two-super-slices.yaml:

    kubectl apply -f two-super-slices.yaml

In questo manifest, imposti replicas: 2 nel campo replicatedJobs. Dopo aver applicato il manifest, Kueue tenta di formare due sezioni separate con una topologia 4x8x8. Kueue crea una sezione dinamica per ogni replica definita in jobset.spec.replicatedJobs[].replicas. Se vengono specificate n repliche, Kueue crea n sezioni dinamiche per il workload e attende che tutte le sezioni diventino attive prima di ammettere il workload.

Esempio 3: workload con una singola sezione dinamica e NodeAffinity

A partire da Kueue 0.15, Kueue supporta NodeAffinity per la selezione dei nodi TAS. Questa funzionalità può essere utilizzata per consentire a entrambi i nodi HEALTHY e DEGRADED di far parte di una sezione dinamica. L'esempio seguente mostra come configurare un carico di lavoro con una singola sezione dinamica e NodeAffinity:

  1. Salva il seguente manifest come 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. Applica il manifest slice-8x8x8-na.yaml:

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

    Dopo aver applicato il manifest, Kueue crea un JobSet denominato slice-8x8x8-na. Kueue tenta quindi di formare una singola sezione dinamica con una topologia 8x8x8, che consente di includere sia i nodi HEALTHY sia quelli DEGRADED grazie all'NodeAffinity specificato. Una volta attiva, la slice ammette il workload e i 128 pod vengono pianificati sui nodi che formano la slice dinamica.

Monitorare lo stato dello slice

Per controllare lo stato delle sezioni dinamiche, esegui questo comando:

kubectl describe slice SLICE_NAME

Sostituisci SLICE_NAME con il nome della tua sezione. Il nome della sezione viene in genere derivato dal nome di JobSet e dall'indice della replica. Per l'esempio 1, una sezione creata da Kueue avrebbe un nome simile a default-jobset-big-super-slice-yyyyy-job-jax-0.

L'output è simile al seguente:

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>

Il nome della sezione rispetta le seguenti regole per garantire la compatibilità con le convenzioni di denominazione delle risorse Compute Engine sottostanti:

  • Modello: {namespace}-jobset-{jobset.metadata.name}-kueueHash[5-character]-{jobset.spec.replicatedJobs[].name}-sliceIndex.
  • Lunghezza: il nome contiene al massimo 54 caratteri. Il controller aggiunge un trattino e un hash del cluster di 8 caratteri per creare i nomi delle risorse Compute Engine, che hanno un limite di 63 caratteri.
  • Format: il nome corrisponde all'espressione regolare ^[a-z]([-a-z0-9]*[a-z0-9])?$. Il nome ha le seguenti caratteristiche:
    • Inizia con una lettera minuscola.
    • Contiene solo lettere minuscole, numeri e trattini (-).
    • Termina con una lettera minuscola o un numero (non può terminare con un trattino).

Esegui la pulizia

Per evitare addebiti imprevisti, elimina le sezioni prima di eliminare i pool di nodi.

  1. Elimina il JobSet. Questa azione attiva Kueue per eliminare le risorse personalizzate Slice associate.

    kubectl delete jobset JOBSET_NAME

    Sostituisci JOBSET_NAME con il nome del JobSet, ad esempio big-super-slice.

  2. Elimina il pool di nodi TPU:

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

(Facoltativo) Utilizzare la suddivisione dinamica con il proprio scheduler

Questo documento è incentrato sull'utilizzo di Kueue e TAS. Tuttavia, puoi anche gestire la suddivisione dinamica con il tuo pianificatore personalizzato. Se scegli di utilizzare uno scheduler diverso, segui le informazioni di riferimento della risorsa personalizzata Slice.

Passaggi successivi