Configurer des TPU avec le découpage dynamique GKE

Ce document explique comment utiliser le slicing dynamique dans Google Kubernetes Engine (GKE). Le partitionnement dynamique vous permet de configurer des sous-blocs de TPU provisionnés dans différentes topologies. Cette fonctionnalité réduit la nécessité de recréer des pools de nœuds, améliore la tolérance aux pannes en permettant la récupération automatique en cas de défaillance et optimise l'utilisation des ressources.

Ce document est destiné aux ingénieurs en IA/ML et aux administrateurs de plate-forme qui souhaitent optimiser l'utilisation des TPU, réduire le temps de provisionnement et améliorer la tolérance aux pannes pour les charges de travail d'entraînement et d'inférence à grande échelle.

Avant de lire ce document, assurez-vous de connaître les éléments suivants :

Qu'est-ce que le slicing dynamique ?

Le slicing dynamique offre une flexibilité dans la gestion de la capacité Cloud TPU en vous permettant de dissocier le provisionnement des TPU. Le découpage dynamique implique le processus suivant :

  1. Vous provisionnez des ressources sous forme d'unités plus petites appelées sous-blocs. Un sous-bloc est l'unité de base logique de la capacité Ironwood (TPU7x). Pour Ironwood (TPU7x), il s'agit d'un groupe de 16 VM TPU avec une topologie 4x4x4 de puces TPU interconnectées. Dans le contexte du mode "Toute la capacité" des TPU et du découpage dynamique, un pool de nœuds correspond directement à un sous-bloc.
  2. Le découpage dynamique assemble ensuite ces sous-blocs pour former des tranches plus grandes.

Avantages du slicing dynamique

Le découpage dynamique vous permet d'effectuer les opérations suivantes :

  • Réduisez le temps de provisionnement : le provisionnement individuel des sous-blocs permet un provisionnement global plus rapide, car il minimise l'impact de toute défaillance individuelle.
  • Réduisez le délai de reprise : en cas de défaillance d'une puce TPU, la plus petite unité de défaillance est un sous-bloc. Le découpage dynamique isole les sous-blocs défectueux afin que les charges de travail puissent être reprogrammées sur des sous-blocs opérationnels plus rapidement que si vous deviez reprovisionner un grand bloc entier.
  • Remodeler la capacité : si vous avez des exigences de charge de travail diverses, vous n'avez pas besoin de supprimer et de recréer des pools de nœuds pour les modifications de topologie, ce qui serait nécessaire sans le partitionnement dynamique. À la place, vous pouvez reconfigurer dynamiquement les pools de nœuds provisionnés pour qu'ils correspondent aux formes spécifiées.

Éléments clés du découpage dynamique

Le slicing dynamique introduit les concepts clés suivants :

  • Provisionnement incrémentiel des pools de nœuds : le slicing dynamique utilise le provisionnement incrémentiel, qui est un modèle de provisionnement tolérant aux pannes pour les pools de nœuds. Ce modèle convertit toute votre capacité TPU en pools de nœuds de groupes de 16 VM TPU.
  • Contrôleur de tranche : contrôleur de ressources personnalisées Kubernetes s'exécutant dans le plan de contrôle GKE et gérant le découpage dynamique. Le contrôleur de tranche gère le cycle de vie d'une ressource personnalisée Slice, qui représente une tranche dynamique. Le contrôleur de tranche gère la création, la surveillance continue et la suppression de la tranche. Lorsque vous utilisez un planificateur, celui-ci dirige la création et la suppression de la ressource personnalisée Slice.
  • Ressource personnalisée de tranche : assemble dynamiquement ces sous-blocs en fonction de la topologie TPU demandée. Ce processus repose sur la reconfiguration dynamique du réseau OCS pour connecter les pools de nœuds TPU, ce qui permet d'assurer des performances optimisées. Vous pouvez inspecter la progression ou l'état de la formation de tranches dynamiques en inspectant les champs d'état de la ressource personnalisée Slice.

Conditions requises

Pour utiliser le slicing dynamique dans GKE, vous devez répondre aux exigences suivantes :

  • Utilisez un cluster Standard en version 1.35.0-gke.274500 ou ultérieure, dans le canal rapide.
  • Utilisez la version Ironwood (TPU7x).
  • Utilisez l'image Container-Optimized OS pour vos nœuds.
  • Pour utiliser le provisionnement incrémentiel, utilisez les réservations en mode "Toute la capacité". Le mode "Toute capacité" est une fonctionnalité activée par TPU Cluster Director.

Avant de commencer

Avant de commencer, effectuez les tâches suivantes :

  • Activez l'API Google Kubernetes Engine.
    • Activer l'API Google Kubernetes Engine
  • Si vous souhaitez utiliser la Google Cloud CLI pour cette tâche, installez et initialisez la gcloud CLI. Si vous avez déjà installé la gcloud CLI, obtenez la dernière version en exécutant la commande gcloud components update. Il est possible que les versions antérieures de la gcloud CLI ne permettent pas d'exécuter les commandes de ce document.

Limites

  • Une seule tranche doit utiliser des sous-blocs au sein du même bloc TPU sous une réservation. Pour utiliser des sous-blocs dans des blocs TPU, utilisez TPU Multislice.
  • Le découpage dynamique n'est pas compatible avec les topologies inférieures à 4x4x4.

Utiliser le slicing dynamique dans GKE avec Kueue

Cette section décrit le workflow d'utilisation du slicing dynamique dans GKE.

  1. Affichez la topologie et l'état des réservations en mode "Toute capacité".
  2. Activez le contrôleur de tranche dans votre cluster.
  3. Créez des pools de nœuds TPU.
  4. Configurez Kueue pour créer une ressource personnalisée Slice.
  5. Exécutez des charges de travail sur le slicing dynamique avec Kueue.
  6. Effectuez un nettoyage.

Activer le contrôleur de segments

Pour utiliser le slicing dynamique, activez le contrôleur de tranche dans votre cluster.

  1. Mettez à jour votre cluster :

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

    Remplacez les éléments suivants :

  2. Obtenez des identifiants pour pouvoir communiquer avec votre cluster à l'aide des commandes kubectl :

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

  3. Dans le résultat de la commande suivante, vérifiez que la valeur slices.accelerator.gke.io est présente :

    kubectl get crd slices.accelerator.gke.io

    Le résultat ressemble à ce qui suit :

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

Créer des pools de nœuds avec provisionnement incrémentiel

Cette section explique comment créer les pools de nœuds TPU avec provisionnement incrémental. GKE convertit toute votre capacité TPU en pools de nœuds de groupes de 16 VM TPU, ou sous-blocs. GKE provisionne ces pools de nœuds même s'il ne parvient pas à trouver les 16 VM opérationnelles. Pour ce faire, il place les nœuds sur les parties opérationnelles de la machine hôte et provisionne progressivement les machines non opérationnelles à mesure qu'elles sont réparées.

Vous pouvez cibler votre pool de nœuds pour qu'il appartienne à l'un des éléments suivants :

  • Bloc spécifique de TPU, exposé dans les réservations en mode "Toute la capacité". Le ciblage par bloc permet à GKE de créer le pool de nœuds dans n'importe quel sous-bloc disponible du bloc spécifié.
  • Un sous-bloc spécifique ou un groupe de 16 nœuds de VM TPU spécifiques pour un contrôle plus précis.

Créer une règle de charge de travail

Pour créer un pool de nœuds de tranche TPU avec Ironwood (TPU7x), vous devez d'abord créer une règle de charge de travail avec le champ accelerator-topology-mode défini sur provision_only. Ce paramètre déclenche le processus de provisionnement incrémentiel.

Créez une règle de charge de travail :

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

Remplacez les éléments suivants :

  • WORKLOAD_POLICY_NAME : nom de votre règle de charge de travail.
  • PROJECT_ID : ID de votre projet Google Cloud .
  • REGION : région de la règle de charge de travail.

Dans cette commande, procédez comme suit :

  • Définissez toujours le champ accelerator-topology sur 4x4x4 pour qu'il corresponde au nombre total de jetons dans un même sous-bloc.
  • Définissez toujours le champ accelerator-topology-mode sur provision_only pour vous assurer que le processus de provisionnement incrémentiel est déclenché. Lorsque le champ provision_only est défini, le pool de nœuds provisionne des nœuds TPU sans former de liens ICI ni OCS.

Cibler votre pool de nœuds pour qu'il appartienne à un bloc ou à un sous-bloc

Vous pouvez cibler des sous-blocs ou des blocs spécifiques dans votre réservation en mode "Toute capacité".

  • Cibler un bloc : chaque pool de nœuds utilise la capacité d'un bloc spécifié. GKE place le pool de nœuds dans un sous-bloc disponible de ce bloc. Vous devez créer autant de pools de nœuds que de sous-blocs dans le bloc que vous souhaitez utiliser.

  • Cibler un sous-bloc : chaque pool de nœuds est associé à un sous-bloc spécifique et disponible. Lorsque vous utilisez le ciblage par sous-bloc, GKE crée le pool de nœuds si au moins une VM est opérationnelle. Le provisionnement incrémentiel garantit que tous les nœuds sont placés dans le sous-bloc spécifié.

Bloquer

  1. Pour récupérer le nom du bloc dans une réservation et le nombre de sous-blocs disponibles dans le bloc, suivez les étapes décrites dans le document Afficher la topologie et l'état des réservations en mode "Toute la capacité" :

    1. Identifiez le nom du bloc en listant tous les blocs de réservation et en copiant la valeur du champ name:. Cette valeur correspond au nom du bloc ou de BLOCK_NAME dans ce document.

    2. Déterminez le nombre de pools de nœuds à créer en décrivant un bloc de réservation et en identifiant la valeur dans le champ reservationSubBlockCount. Cette valeur correspond au nombre de sous-blocs disponibles. Par exemple, la valeur reservationSubBlockCount: 4 indique que le bloc dispose de quatre sous-blocs et que vous devez créer quatre pools de nœuds distincts.

  2. Définissez le chemin de réservation :

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

    Remplacez les éléments suivants :

    • RESERVATION_NAME : nom de votre réservation de TPU.
    • BLOCK_NAME : nom du bloc.

  3. Créez un pool de nœuds pour chaque sous-bloc identifié à l'étape précédente. Par exemple, si le nombre est 4, exécutez cette commande quatre fois. Attribuez un nom unique à chaque pool de nœuds.

    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}

    Remplacez les éléments suivants :

    • NODE_POOL_NAME : nom de votre nouveau pool de nœuds.
    • CLUSTER_NAME : nom de votre cluster GKE.
    • WORKLOAD_POLICY_NAME : nom de la règle de charge de travail que vous avez créée.
    • ZONE : zone du pool de nœuds, par exemple us-central1-a.

Sous-bloc

  1. Pour récupérer le nom du bloc et les ID des sous-blocs disponibles, suivez les étapes décrites dans le document Afficher la topologie et l'état des réservations en mode Toute capacité :

    1. Pour identifier le nom du bloc, listez tous les blocs de réservation et copiez la valeur du champ name:. Cette valeur correspond au nom du bloc ou de BLOCK_NAME dans ce document.

    2. Pour identifier le nom des sous-blocs, listez tous les sous-blocs d'un bloc et copiez la valeur du champ name: pour chaque entrée sous reservationSubBlocks. Cette valeur correspond au nom du sous-bloc ou SUBBLOCK_NAME dans ce document.

  2. Définissez le chemin de réservation :

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

    Remplacez les éléments suivants :

    • RESERVATION_NAME : nom de votre réservation de TPU.
    • BLOCK_NAME : nom du bloc.
    • SUBBLOCK_NAME : nom du sous-bloc.

  3. Créez le pool de nœuds :

    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}

    Remplacez les éléments suivants :

    • NODE_POOL_NAME : nom unique du nouveau pool de nœuds, par exemple sub-block-pool-1.
    • PROJECT_ID : ID de votre projet Google Cloud .
    • CLUSTER_NAME : nom de votre cluster GKE.
    • ZONE : zone du pool de nœuds, par exemple us-central2-b.
    • WORKLOAD_POLICY_NAME : nom de la règle de charge de travail que vous avez créée.

À ce stade, les nœuds sont créés, mais leurs liens Inter-Chip Interconnect (ICI) ne sont pas encore actifs. Par conséquent, vous ne pouvez pas exécuter de charges de travail directement sur ces pools de nœuds.

Pour activer tous les liens ICI nécessaires à la formation du slice et permettre la planification des charges de travail, créez un slice dynamique à l'aide de l'une des méthodes suivantes :

  • Créez une ressource personnalisée Slice. Au lieu d'utiliser des pods, vous utilisez une ressource personnalisée Slice pour définir la topologie spécifiée, que le contrôleur de tranche active.
  • Planifiez des charges de travail GKE avec Kueue et TAS. Kueue gère automatiquement la création et la suppression des ressources personnalisées Slice. Évitez de modifier manuellement les ressources personnalisées Slice créées par Kueue.

Créer une tranche dynamique avec Kueue et TAS

Dans cette section, vous allez planifier des charges de travail GKE avec Kueue et TAS.

Installer les ressources JobSet et Kueue pour le slicing dynamique

  1. Installez 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. Installez 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. Installez le contrôleur de tranche Kueue :

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

  4. Pour configurer Kueue pour le slicing dynamique, enregistrez le fichier manifeste suivant sous le nom 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. Appliquez le fichier manifeste dynamic-slice-topology.yaml :

    kubectl apply -f dynamic-slice-topology.yaml

    Dans ce fichier manifeste, vous configurez Kueue pour le découpage dynamique en définissant les ressources suivantes :

    • Topologie de tranche dynamique Ironwood (TPU7x) (superslice-topology) : la topologie définit les niveaux que Kueue prend en compte lorsqu'il planifie des charges de travail de découpage dynamique. Voici les différents niveaux :
      • Libellé cloud.google.com/gce-topology-block : ce niveau est nécessaire pour comprendre quels sous-blocs appartiennent à quels blocs, car seules les sous-blocs du même bloc peuvent former un segment.
      • Libellé cloud.google.com/gke-tpu-partition-4x4x4-id : ce niveau représente des sous-blocs Ironwood (TPU7x) individuels (topologie 4x4x4).
      • Libellé kubernetes.io/hostname : ce niveau est requis pour attribuer des pods à des VM spécifiques et pour observer leurs libellés et leurs taints.
    • ResourceFlavor SuperSlice Ironwood (TPU7x) (superslice-rf) : le resource flavor pour les sous-blocs Ironwood (TPU7x) inclut le libellé cloud.google.com/gke-tpu-accelerator: tpu7x pour correspondre aux nœuds avec des machines Ironwood (TPU7x).
    • SuperSlice AdmissionCheck (superslice-ac) : ce contrôle d'admission indique à Kueue de ne pas planifier de charge de travail tant que le contrôleur de tranche GKE n'a pas confirmé que la tranche est devenue active. Le contrôle d'admission est d'abord défini, puis ajouté au ClusterQueue qui gère les charges de travail de segmentation dynamique.
    • ClusterQueue (cq) et LocalQueue (lq) : ces champs gèrent les ressources google.com/tpu. La ressource ClusterQueue cq inclut la vérification de l'admission superslice-ac. Le nominalQuota pour google.com/tpu peut être configuré de deux manières :
      • Quota spécifique : définissez nominalQuota pour qu'il corresponde à la capacité existante pour la gestion du partage équitable et des quotas.
      • Quota illimité : définissez nominalQuota sur une valeur très élevée, telle que "999999", pour modéliser un quota illimité. Pour se concentrer sur TAS et le découpage dynamique, cette configuration contourne la fonctionnalité de gestion des quotas de Kueue.

Définir la sélection de l'état du sous-bloc

Au-delà de l'état et de la disponibilité standard des nœuds, GKE expose l'état spécifique de chaque sous-bloc à l'aide du libellé cloud.google.com/gke-tpu-partition-4x4x4-state. Cette étiquette permet à GKE de tenir compte des facteurs qui influencent la formation des tranches, tels que l'état des liens TPU.

Vous pouvez définir la valeur du libellé cloud.google.com/gke-tpu-partition-4x4x4-state comme suit :

  • HEALTHY : l'infrastructure est opérationnelle.
  • DEGRADED : l'infrastructure du sous-bloc est dans un état dégradé, par exemple en raison de la dégradation du lien OCS. Le sous-bloc peut toujours former une tranche, mais les performances globales peuvent être inférieures à celles des sous-blocs sains. Si vous pouvez tolérer une dégradation potentielle des performances, vous pouvez configurer votre charge de travail pour qu'elle utilise des sous-blocs DEGRADED à l'aide de l'affinité de nœuds, comme indiqué dans l'exemple 3.
  • UNHEALTHY : le sous-bloc est défectueux et ne peut pas former de tranche.

Le webhook du contrôleur de tranche Kueue valide si une charge de travail inclut une exigence d'état de sous-bloc spécifique. Si aucune préférence n'est indiquée, le webhook injecte une affinité de nœud par défaut.

Voici le comportement à suivre :

  • Si un nodeSelector ou un nodeAffinity ciblant le libellé cloud.google.com/gke-tpu-partition-4x4x4-state est présent, il reste inchangé.

  • Si aucune configuration de libellé de ce type n'existe, le webhook injecte l'affinité de nœud par défaut suivante pour s'assurer que seuls les sous-blocs disponibles sont utilisés :

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

La section suivante inclut des exemples dans lesquels le libellé cloud.google.com/gke-tpu-partition-4x4x4-state est configuré pour spécifier les différentes configurations d'état des sous-blocs.

Exécuter des charges de travail de test sur le partitionnement dynamique avec Kueue

Cette section explique comment déployer des charges de travail sur le découpage dynamique avec Kueue et TAS. Il inclut trois exemples qui montrent comment créer une charge de travail de tranche dynamique et une charge de travail composée de plusieurs tranches. Les charges de travail sont envoyées en tant que JobSets.

Exemple 1 : Une seule charge de travail utilise un seul slice dynamique

L'exemple suivant décrit comment créer une charge de travail à l'aide d'une tranche avec une topologie 4x12x16, qui est composée de 12 sous-blocs. Le nombre de pods a été calculé comme suit : (4 * 12 * 16) / 4 puces par nœud = 192 pods.

  1. Enregistrez le manifeste suivant sous le nom 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

    Dans ce fichier manifeste, les annotations suivantes indiquent à Kueue les caractéristiques et la topologie de la tranche pour configurer les éléments suivants :

    • cloud.google.com/gke-tpu-slice-topology : spécifie "4x12x16" comme topologie de tranche dynamique. Les exigences concernant la topologie de l'accélérateur tpu7x incluent les règles suivantes :
      • La topologie minimale est 4x4x4.
      • La topologie doit être une chaîne tridimensionnelle au format AxBxC. Exemple : 4x8x8.
      • Chacune des dimensions (A, B et C) doit être un multiple de quatre.
      • Les dimensions doivent être triées par ordre croissant : A <= B <= C. Par exemple, 4x8x4 n'est pas valide et doit être remplacé par 4x4x8.
      • Le produit des dimensions (ABC) ne doit pas dépasser 9 216.
      • Les topologies de tranche les plus grandes acceptées peuvent inclure jusqu'à 32 sous-blocs. Par exemple, 8x16x16 avec 32 sous-blocs, 8x12x20 avec 30 sous-blocs ou 12x12x12 avec 27 sous-blocs sont dans les limites acceptées.
    • cloud.google.com/gke-tpu-accelerator: tpu7x : planifie les pods sur les VM qui exécutent Ironwood (TPU7x).
    • kueue.x-k8s.io/queue-name : attribue le JobSet à une LocalQueue Kueue.

  2. Appliquez le fichier manifeste big-super-slice.yaml :

    kubectl apply -f big-super-slice.yaml

    Une fois le fichier manifeste appliqué, Kueue crée un JobSet nommé big-super-slice. Kueue tente ensuite de former une tranche dynamique unique avec une topologie 4x12x16. Une fois le slice actif, Kueue accepte la charge de travail et les 192 pods sont planifiés sur les nœuds pour former le slice dynamique qui exécute vos charges de travail.

Exemple 2 : Charge de travail avec plusieurs répliques

L'exemple suivant montre comment créer une charge de travail qui utilise deux tranches dynamiques, chacune composée de quatre sous-blocs.

  1. Enregistrez le manifeste suivant sous le nom 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. Appliquez le fichier manifeste two-super-slices.yaml :

    kubectl apply -f two-super-slices.yaml

Dans ce fichier manifeste, vous définissez replicas: 2 dans le champ replicatedJobs. Une fois le fichier manifeste appliqué, Kueue tente de former deux tranches distinctes avec une topologie 4x8x8. Kueue crée une tranche dynamique pour chaque réplica défini dans jobset.spec.replicatedJobs[].replicas. Si n répliques sont spécifiées, Kueue crée n tranches dynamiques pour la charge de travail et attend que toutes les tranches deviennent actives avant d'admettre la charge de travail.

Exemple 3 : Charge de travail avec une seule tranche dynamique et NodeAffinity

À partir de Kueue 0.15, Kueue est compatible avec NodeAffinity pour la sélection des nœuds TAS. Cette fonctionnalité peut être utilisée pour permettre aux nœuds HEALTHY et DEGRADED de faire partie d'une tranche dynamique. L'exemple suivant montre comment configurer une charge de travail avec une seule tranche dynamique et NodeAffinity :

  1. Enregistrez le manifeste suivant sous le nom 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. Appliquez le fichier manifeste slice-8x8x8-na.yaml :

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

    Une fois le fichier manifeste appliqué, Kueue crée un JobSet nommé slice-8x8x8-na. Kueue tente ensuite de former une tranche dynamique unique avec une topologie 8x8x8, ce qui permet d'inclure à la fois les nœuds HEALTHY et DEGRADED en raison de l'affinité de nœud spécifiée. Une fois la tranche active, Kueue admet la charge de travail et les 128 pods sont planifiés sur les nœuds qui forment la tranche dynamique.

Surveiller l'état de la tranche

Pour vérifier l'état de vos tranches dynamiques, exécutez la commande suivante :

kubectl describe slice SLICE_NAME

Remplacez SLICE_NAME par le nom de votre tranche. Le nom du slice est généralement dérivé du nom JobSet et de l'index de réplica. Dans l'exemple 1, une tranche créée par Kueue aurait un nom semblable à default-jobset-big-super-slice-yyyyy-job-jax-0.

Le résultat ressemble à ce qui suit :

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>

Le nom de la tranche respecte les règles suivantes pour assurer la compatibilité avec les conventions d'attribution de noms des ressources Compute Engine sous-jacentes :

  • Modèle : {namespace}-jobset-{jobset.metadata.name}-kueueHash[5-character]-{jobset.spec.replicatedJobs[].name}-sliceIndex.
  • Longueur : le nom ne comporte pas plus de 54 caractères. L'application de contrôleur ajoute un tiret et un hachage de cluster de huit caractères pour créer des noms de ressources Compute Engine, qui sont limités à 63 caractères.
  • Format : le nom correspond à l'expression régulière ^[a-z]([-a-z0-9]*[a-z0-9])?$. Le nom présente les caractéristiques suivantes :
    • Commence par une lettre minuscule.
    • Ne contient que des lettres minuscules, des chiffres et des traits d'union (-).
    • Se termine par une lettre minuscule ou un chiffre (il ne peut pas se terminer par un trait d'union).

Effectuer un nettoyage

Pour éviter des frais inattendus, supprimez vos tranches avant de supprimer les pools de nœuds.

  1. Supprimez le JobSet. Cette action déclenche la suppression des ressources personnalisées Slice associées par Kueue.

    kubectl delete jobset JOBSET_NAME

    Remplacez JOBSET_NAME par le nom de votre JobSet, par exemple big-super-slice.

  2. Supprimez le pool de nœuds TPU :

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

(Facultatif) Utiliser le slicing dynamique avec votre propre programmateur

Ce document se concentre sur l'utilisation de Kueue et TAS. Toutefois, vous pouvez également gérer le découpage dynamique avec votre propre planificateur personnalisé. Si vous choisissez d'utiliser un autre planificateur, suivez les informations de référence sur la ressource personnalisée Slice.

Étapes suivantes