Utiliser le fractionnement dynamique avec un planificateur personnalisé

Ce document explique comment utiliser le découpage dynamique en interagissant directement avec les ressources personnalisées Slice. Vous pouvez créer des tranches, surveiller l'état des partitions et vérifier l'état des tranches.

Avant de suivre ces instructions, assurez-vous de bien comprendre les concepts du slicing dynamique.

Pourquoi utiliser le slicing dynamique avec un planificateur personnalisé ?

Utilisez votre propre programmateur pour gérer les ressources personnalisées Slice si vous avez des exigences de planification complexes ou si vous souhaitez intégrer le slicing dynamique à votre infrastructure de planification existante.

Si vous préférez utiliser un planificateur au lieu de gérer directement les ressources personnalisées Slice, GKE fournit une intégration avec Kueue et la planification tenant compte de la topologie (TAS, Topology Aware Scheduling). Pour en savoir plus, consultez Planifier des tranches dynamiques avec Kueue et TAS.

Présentation du workflow

Pour utiliser le slicing dynamique avec un programmateur personnalisé, vous devez effectuer les tâches suivantes dans ce document :

  1. Activez le sélecteur de tranche.
  2. Créez des pools de nœuds avec provisionnement incrémentiel.
  3. Créez des ressources personnalisées Slice en fonction des besoins de votre charge de travail. Appliquez la ressource personnalisée Slice à votre cluster.
  4. Surveillez l'état des partitions et l'intégrité des segments.
  5. Supprimez la tranche lorsque vous avez terminé.

Pour en savoir plus sur les champs et l'état de la ressource personnalisée Slice, consultez les informations de référence sur 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.2-gke.1842000 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.

Activer le contrôleur de segments

Pour utiliser le découpage dynamique, activez le contrôleur de découpage 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 des pools de nœuds TPU avec provisionnement incrémentiel. 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 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 puces 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 permet de s'assurer 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 comporte quatre sous-blocs disponibles 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 de santé des réservations en mode "Toute la 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 directement de charges de travail sur ces pools de nœuds.

Pour activer tous les liens ICI nécessaires à la formation de la tranche et permettre la planification des charges de travail, créez une tranche 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.

Former une tranche dynamique

Une fois les pools de nœuds créés, vous pouvez former une tranche dynamique plus grande en créant 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 ensuite.

Vérifier l'état des nœuds et des partitions

  1. Pour obtenir les noms des nœuds du pool de nœuds, exécutez la commande suivante :

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

    Le résultat ressemble à ce qui suit :

    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. Vérifiez le modèle de provisionnement du nœud :

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

    Le résultat ressemble à ce qui suit :

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

    Cette valeur correspond au paramètre accelerator-topology-mode=provision_only défini lors de la création de la règle de charge de travail.

  3. Récupérez les informations sur le libellé du nœud :

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

    Remplacez NODE_NAME par le nom de l'un des nœuds du pool de nœuds.

    Le résultat ressemble à ce qui suit :

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

    L'annotation cloud.google.com/gke-tpu-partition-4x4x4-state indique si le nœud est disponible pour former un segment dynamique. Ce libellé accepte les valeurs suivantes :

    • HEALTHY : le nœud est en bon état et entièrement fonctionnel.
    • DEGRADED : le nœud est défectueux, mais reste utilisable pour la formation de segments dynamiques.
    • UNHEALTHY : le nœud ne fonctionne pas correctement et ne peut pas être utilisé pour former une tranche.
    • UNSET : l'état est indéfini en raison d'un nombre insuffisant de nœuds dans le pool de nœuds.
    • INCOMPLETE : tous les nœuds de la partition ne sont pas provisionnés.
  4. Vérifiez que le nœud inclut l'annotation node.gke.io/created-by-mig :

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

    Remplacez NODE_NAME par le nom de l'un des nœuds du pool de nœuds.

    Le résultat ressemble à ce qui suit :

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

    La sortie inclut le libellé node.gke.io/created-by-mig, qui permet au plan de contrôle GKE d'associer les nœuds Kubernetes à leurs ressources Compute Engine sous-jacentes.

Créer une ressource personnalisée Slice

  1. Définissez la ressource personnalisée 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
    

    Remplacez les éléments suivants :

    • SLICE_NAME : nom de votre tranche. Le nom doit respecter les conditions metadata.name.
    • TOPOLOGY : topologie de la tranche dynamique. La topologie doit répondre aux conditions suivantes :
      • Chaque dimension de la topologie demandée doit être un multiple de quatre (par exemple, 4A x 4B x 4C).
      • Les trois valeurs des dimensions de la topologie, AxBxC, doivent être dans l'ordre croissant (A ≤ B ≤ C). Par exemple, 4x4x8 est valide, mais 4x8x4 ne l'est pas. Cet ordre permet d'assurer une formation cohérente des tranches et d'éviter tout comportement inattendu.
      • Le produit des trois valeurs dans les dimensions de la topologie, A × B × C,ne doit pas dépasser 9 216.
    • PARTITION_ID : liste de chaînes qui identifient les partitions 4x4x4 qui composent le segment. Calculez le nombre de partitions en fonction du nombre total de puces, où chaque partition est constituée de 64 puces. Le nombre d'éléments de votre liste spec.partitionIds doit correspondre exactement au nombre de partitions calculé ((A × B × C) / 64). Le partitionIds doit répondre aux conditions suivantes :
      • Chaque partition doit correspondre à un sous-bloc de réservation.
      • Tous les sous-blocs associés doivent appartenir à la même réservation.
      • Tous les blocs associés doivent se trouver dans la même réservation.
      • Tous les nœuds des pools de nœuds associés doivent être à l'état ready.
    • La valeur du champ type doit être tpu7x.
    • Si vous le souhaitez, vous pouvez ajouter l'annotation slice.gke.io/retry-on-failure: "true" à la ressource personnalisée de tranche pour permettre au contrôleur de tranche de réessayer automatiquement lors de la formation de la tranche. Si la tranche n'est pas créée en raison du motif d'état SliceCreationFailed, le contrôleur réessaie jusqu'à ce qu'elle soit correctement formée.

    Par exemple, pour créer une tranche 4x8x8, vous devez fournir quatre ID de partition uniques.

    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. Appliquez la ressource personnalisée Slice :

    kubectl apply -f test-slice-example.yaml
    

    À ce stade, GKE tente de créer la tranche. Si l'un des problèmes suivants se produit, la création du slice échoue et le motif de l'état de la ressource personnalisée Slice est mis à jour sur SliceCreationFailed ou FAILED :

    • Si les nœuds sélectionnés sur la ressource personnalisée n'existent pas, le motif de l'état est SliceCreationFailed.
    • Si des nœuds de la ressource personnalisée sont utilisés par une autre tranche, le motif de l'état est SliceCreationFailed.
    • Si les nœuds de la ressource personnalisée ne font pas partie du même bloc de réservation, la raison de l'état est FAILED.
    • Si les nœuds ne se trouvent pas dans la même réservation, la raison de l'état est FAILED.
    • Si la topologie ne correspond pas au nombre de partitions, la raison de l'état est SliceCreationFailed.

    Pour en savoir plus sur l'état de la ressource personnalisée Slice, consultez État de Slice.

Surveiller l'état de la ressource personnalisée Slice

Pour vérifier l'état de la ressource personnalisée Slice, exécutez la commande suivante :

kubectl describe slice SLICE_NAME

Remplacez SLICE_NAME par le nom de la tranche.

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

Le champ reason de l'état de la ressource personnalisée Slice indique l'état actuel du slice. Le cycle de vie d'une ressource personnalisée Slice suit cette progression :

  • SliceNotCreated : le contrôleur effectue l'initialisation et vérifie les ressources.
    • Si les conditions préalables ne sont pas remplies, l'état passe à SliceCreationFailed.
    • Si la validation réussit, l'état passe à ACTIVATING.
  • ACTIVATING : GKE est en train de former le slice.
    • En cas de réussite, l'état passe à ACTIVE.
    • Si les sous-blocs sont dégradés, mais que la tranche est utilisable, l'état passe à ACTIVE_DEGRADED.
    • En cas d'échec de la formation, l'état passe à FAILED.
  • DEACTIVATING : si la ressource personnalisée Slice est supprimée ou si une défaillance critique se produit dans un état actif ou en échec, le démantèlement du slice commence.
  • INCOMPLETE : étape finale avant la suppression complète de la ressource.

Pour en savoir plus sur l'état de la ressource personnalisée Slice, consultez État de Slice.

Exécuter des charges de travail sur le découpage dynamique

Lorsque la ressource personnalisée Slice est à l'état ACTIVE, vous pouvez y exécuter des charges de travail. La section suivante inclut des exemples de charges de travail qui utilisent le découpage dynamique. Les charges de travail sont envoyées en tant que Jobs ou JobSets.

Exemple 1 : Une charge de travail unique utilise un seul segment

L'exemple suivant montre une charge de travail qui utilise un seul segment de sous-bloc.

  1. Enregistrez l'exemple de fichier manifeste suivant sous le nom 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
    

    Dans le fichier manifeste :

    • cloud.google.com/gke-tpu-slice-topology et cloud.google.com/gke-tpu-topology définissent la topologie de la tranche dynamique.
    • env.value: tpu7x-128 correspond au type d'accélérateur TPU et au nombre total de cœurs de la tranche. Le nombre de cœurs est calculé en multipliant les dimensions de la topologie par le nombre de cœurs par puce. Par exemple, pour une topologie 4x4x4, le calcul est 4 × 4 × 4 × 2 = 128, où 2 correspond au nombre de cœurs par puce pour tpu7x (Ironwood (TPU7x)). Par conséquent, TPU_ACCELERATOR_TYPE est tpu7x-128.
  2. Appliquez le fichier manifeste tpu-job-jax-v7x-64.yaml :

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

Exemple 2 : Déployer une charge de travail sur des pools de nœuds Multislice à l'aide de JobSet

Cet exemple montre comment déployer une charge de travail sur des pools de nœuds multislices à l'aide de JobSet.

  1. Installez JobSet :

    kubectl apply --server-side -f https://github.com/kubernetes-sigs/jobset/releases/download/v0.10.1/manifests.yaml
    
  2. Enregistrez l'exemple de fichier manifeste suivant sous le nom 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. Appliquez le fichier manifeste tpu-multislice-jax.yaml :

    kubectl apply -f tpu-multislice-jax.yaml
    

    Dans le fichier manifeste :

    • Le champ replicas: 2 sous replicatedJobs indique que JobSet crée deux jobs distincts, chacun correspondant à une tranche de TPU 4x4x4.
    • L'annotation alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-tpu-slice permet de s'assurer que chaque Job est attribué à une tranche de TPU unique.
    • L'annotation cloud.google.com/gke-tpu-slice-topology: 4x4x4 définit la topologie de chaque tranche dynamique.
    • La variable d'environnement TPU_ACCELERATOR_TYPE n'est pas définie explicitement dans cet exemple, car JobSet gère l'attribution des tranches. Le code JAX détecte automatiquement les appareils TPU disponibles dans la tranche qui lui est attribuée.

Supprimer le segment

  1. Supprimez le segment :

    kubectl patch slice $SLICE_NAME --type json \
      -p='[{"op": "remove", "path": "/metadata/finalizers"}]'
    
  2. Vérifiez que la tranche a été supprimée :

    kubectl get slices
    

Désactiver le Slice Controller

Pour désactiver le contrôleur de tranche, supprimez-le du cluster.

  1. Vérifiez que les ressources personnalisées Slice sont vides :

    kubectl get slice -A
    
  2. Mettez à jour le cluster pour désactiver le contrôleur de tranche :

    gcloud container clusters update ${CLUSTER_NAME} \
        --location=${REGION} \
        --no-enable-slice-controller
    
  3. Supprimez la ressource personnalisée Slice :

    kubectl delete crd slices.accelerator.gke.io
    
  4. Vérifiez que la ressource personnalisée Slice est supprimée :

    kubectl get crd | grep slices.accelerator.gke.io
    
  5. Supprimez les libellés ajoutés par le contrôleur de tranche. Vous devez supprimer les libellés suivants :

    • cloud.google.com/gke-tpu-slice
    • cloud.google.com/gke-tpu-topology
    1. Pour supprimer un nœud spécifique, modifiez son nom.
    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 vous souhaitez supprimer ces libellés de tous les nœuds de votre cluster :
    kubectl label nodes --all cloud.google.com/gke-tpu-slice- cloud.google.com/gke-tpu-slice-topology-
    
    1. Vérifiez les libellés de nœuds et assurez-vous qu'ils sont vides :
    export NODE_NAME="gke-tpu-bdac9600-3bdg"
    kubectl describe node $NODE_NAME | grep "cloud.google.com/gke-tpu-slice"
    

Étapes suivantes