Présentation des Cloud TPU multisegments dans GKE [Version preview]

Preview

Ce produit ou cette fonctionnalité est soumis aux "Conditions des offres de pré-DG" dans la section "Conditions d'utilisation générales" des Conditions spécifiques au service. Les produits et fonctionnalités pré-DG sont disponibles "en l'état" et peuvent bénéficier d'une assistance limitée. Pour en savoir plus, consultez les descriptions des étapes de lancement.

Résumé

Ce document explique comment gérer des Cloud TPU dans une configuration multisegment à l'aide de GKE. Nous partons du principe que vous utilisez des segments Cloud TPU uniques sur GKE et que vous disposez d'une expérience générale avec les TPU Cloud Multislice. Nous vous recommandons de modéliser chaque tranche Cloud TPU en tant que pool de nœuds et d'envoyer des tâches impliquant plusieurs pools de nœuds à l'aide de l'API JobSet.

Important :Dans la version 4, l'aperçu privé multisegment n'est compatible qu'avec les charges de travail JAX.

Votre avis

Si vous avez des questions ou des commentaires, veuillez nous envoyer un e-mail à l'adresse [e-mail].

Concepts

Réparation automatique

Lorsqu'une tranche subit un événement de maintenance, une préemption ou une défaillance matérielle, Google Kubernetes Engine en crée une autre. Si la capacité est insuffisante pour créer une tranche, la création ne se termine que lorsque le matériel est disponible. Toutes les autres tranches de l'environnement multisegment sont alors redémarrées afin que l'entraînement puisse continuer. Votre script d'entraînement doit définir un script de démarrage qui vérifie les points de contrôle et les charge avant de redémarrer l'entraînement.

Mise en réseau de centres de données

Un réseau à haut débit et à faible débit qui connecte les segments Cloud TPU dans une configuration multisegment.

Ensemble de données

Données utilisées par un modèle pour l'entraînement ou l'inférence.

L'hôte

Un hôte est un ordinateur physique qui exécute des VM. Un hôte peut s'exécuter sur plusieurs VM à la fois. Chaque VM dispose d'un Cloud TPU dédié.

Inférence

Chargez un modèle de machine learning pré-entraîné sur un hôte et effectuez des prédictions sur les données.

Interchip Interconnect (ICI)

Liens internes à haut débit et à faible latence qui connectent des Cloud TPU dans un pod Cloud TPU.

Pod Kubernetes

Un ou plusieurs conteneurs comportant un stockage partagé et des ressources réseau, ainsi qu'une spécification d'exécution des conteneurs.

Multisegment

Deux segments Cloud TPU ou plus pouvant communiquer via un réseau de centres de données.

Pools de nœuds

Dans le contexte GKE, un pool de nœuds est un groupe de nœuds (VM) pouvant être programmés avec des charges de travail (par exemple, une tâche d'entraînement) par le programmeur Kubernetes. Pour en savoir plus sur l'architecture et les concepts des clusters GKE, consultez la page Architecture d'un cluster GKE.

Ressource en file d'attente

Représentation des ressources Cloud TPU, utilisée pour mettre en file d'attente et gérer une requête pour un environnement Cloud TPU à tranche unique ou multisegment.

Tensor

Structure de données utilisée pour représenter des données multidimensionnelles dans un modèle de machine learning.

Tensor Processing Unit (TPU)

La puce d'accélération du ML développée par Google en interne. Ils sont conçus pour offrir les calculs les plus rapides et les plus économes en énergie pour les tâches de machine learning clés telles que la multiplication matricielle.

pod Cloud TPU

Ensemble de puces Cloud TPU connectées par des interfaces réseau ICI dédiées. Un pod Cloud TPU vous permet de répartir la charge de traitement sur plusieurs Cloud TPU.

Segment Cloud TPU

Sous-section logique d'un pod Cloud TPU constitué de puces Cloud TPU. Toutes les puces d'une tranche communiquent entre elles via le réseau ICI.

VM Cloud TPU

Machine virtuelle exécutant Linux ayant accès aux Cloud TPU sous-jacents. Pour les Cloud TPU v4, chaque VM Cloud TPU dispose d'un accès direct à quatre puces. Les VM Cloud TPU sont parfois appelées nœuds de calcul.

Prérequis

Dans ce document, nous partons du principe que vous connaissez la présentation de Cloud TPU Multislice et la présentation de Cloud TPU dans GKE.

Gérer des Cloud TPU multisegments dans GKE

Provisionner la capacité Cloud TPU pour l'utiliser avec GKE Multislice

Déployez des charges de travail Cloud TPU dans GKE. Vous découvrirez comment réserver un TPU. Toutefois, si vous utilisez la capacité Cloud TPU directement via l'API Cloud TPU et que vous souhaitez utiliser la même capacité pour GKE Multislice, contactez [e-mail] pour déplacer une partie de votre capacité vers GKE.

Remarque :Vous disposerez d'une vue unique de la capacité Cloud TPU en disponibilité générale et vous n'aurez plus besoin de migrer la capacité.

Créer des pools de nœuds dans un environnement multisegment

Avertissement :Votre cluster doit comporter au moins un pool de nœuds non Cloud TPU. Ce pool de nœuds permet de planifier les pods système critiques tels que kube-dns. Tous les nœuds Cloud TPU ont un rejet pour empêcher la planification des pods système sur les nœuds Cloud TPU. Vos charges de travail utilisent ainsi pleinement vos nœuds Cloud TPU.

Pour créer un environnement multisegment, vous devez regrouper plusieurs pools de nœuds GKE dans un groupe multisegment. Vous devez créer un pool de nœuds distinct pour chaque tranche faisant partie de votre environnement multisegment. Vous devez ajouter chaque pool de nœuds à un groupe de segments d'application. Les groupes multisegments permettent de regrouper et de gérer plusieurs pools de nœuds. Les groupes d'applications multiples sont définis en appliquant des libellés (paires nom/valeur) à chaque pool de nœuds qui fera partie de votre environnement multisegment. Dans ce tutoriel, nous utilisons des étiquettes nommées MultisliceGroup et MultisliceGroupSize, mais vous pouvez utiliser le nom de votre choix. Tous les pools de nœuds d'un groupe multisegment doivent avoir les mêmes libellés (paires nom/valeur) et la même topologie Cloud TPU.

Vous pouvez ajouter le libellé Multisegment au moment de la création. Dans cet exemple, nous utilisons MultisliceGroup et MultisliceGroupSize pour les noms des étiquettes:

$ gcloud beta container node-pools create pool-name \
  --region=cluster-region \
  --cluster=cluster-name \
  --node-locations=node-zone \
  --machine-type=machine-type \
  --tpu-topology=tpu-topology \
  --node-labels=MultisliceGroup=multislice-group-name,MultisliceGroupSize=num-slices \
  --num-nodes=num-nodes \
  [--reservation-affinity=specific \
     --reservation=reservation-name]

Remarques :

• Une seule opération de pool de nœuds peut être effectuée à la fois. Vous devez donc attendre que chaque pool soit créé ou supprimé avant de pouvoir en créer et en supprimer.

• Les environnements multisegments composés de tranches d'un seul hôte ne sont pas compatibles. Par exemple, vous ne pouvez pas créer une multisegment avec plusieurs topologies v4-16, v5e avec une topologie (1,4) ou des tranches v5e-8.

Vous pouvez également ajouter les étiquettes MultisliceGroup et MultisliceGroupSize à un pool de nœuds existant:

$ gcloud beta container node-pools update pool-name \
  --region=cluster-region \
  --cluster=cluster-name \
  --node-labels=MultisliceGroup=multislice-group-name,MultisliceGroupSize=num-slices

Vérifier l'état du pool de nœuds

Vous pouvez vérifier l'état de vos pools de nœuds à l'aide de la commande kubectl:

$ kubectl get nodes -l MultisliceGroup=multislice-group-1 | grep "Ready" | wc -l

Remarques :

1. Vous devez vous connecter au cluster GKE avant d'exécuter la commande. Pour en savoir plus, consultez la page Accéder aux clusters.

2. Lorsque vous exécutez la commande get nodes, un message d'erreur peut s'afficher:

   E0608 13:25:01.586645 21293 memcache.go:255] couldn't get resource list for
   http://metrics.k8s.io/v1beta1 : the server is currently unable to handle the
   request
   

   Cette erreur signifie généralement qu'il n'y a pas de pool de nœuds non Cloud TPU dans votre cluster GKE. Ajoutez un pool de nœuds autre que Cloud TPU à votre cluster, puis réessayez la commande.

Exécuter des charges de travail multisegments

API JobSet

L'API JobSet est une API Open Source permettant de gérer un groupe de tâches. Il convient parfaitement aux charges de travail à plusieurs segments. Comme décrit dans l'article Présentation des Cloud TPU dans GKE pour les charges de travail à secteur unique, nous utilisons un seul objet IndexedJob Kubernetes pour l'exécuter sur une seule tranche Cloud TPU. De même, pour les environnements multisegments, nous avons besoin que plusieurs jobs indexés Kubernetes fonctionnent ensemble. JobSet permet de parvenir à ce résultat.

JobSet fournit les éléments suivants:

• Traitement des défaillances: recréation automatique de toutes les tâches enfants en cas d'échec de la tâche enfant.
• Indexation des segments d'application: JobSet indexe chaque tâche avec un ID et l'ajoute en tant qu'annotations/libellés sur la tâche et le pod.
• Emplacement de la charge de travail: se charge automatiquement d'attribuer une seule tâche à un domaine spécifique (dans ce cas, chaque tâche sur un seul pool de nœuds).
• Création sans interface graphique: cycle de vie géré d'un service svc sans interface graphique pour toutes les tâches du JobSet.

Installation

1. Installez l'API JobSet sur votre cluster:

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

2. Vérifiez que le contrôleur JobSet est en cours d'exécution:

   $ kubectl get pods -n jobset-system
   

Exécuter une charge de travail multisegment

Le fichier manifeste YAML suivant montre comment exécuter une charge de travail multisegment sur quatre tranches de v4 à 16. Copiez le fichier YAML suivant, enregistrez-le dans un fichier, puis exécutez kubectl apply -f <file-name.yaml> :

apiVersion: jobset.x-k8s.io/v1alpha2
kind: JobSet
metadata:
  name: multislice-job  # JobSet name (${JOBSET_NAME})
  annotations:
    alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool # 1:1 job replica to node pool assignment
spec:
  failurePolicy:
    maxRestarts: 4  # The set will be restarted on failures up to 4 times.
  replicatedJobs:
    - name: slice    # Part of the name of the child Jobs (<replicateJobName>)
      replicas: 4    # Number of slices
      template:
        spec:
          parallelism: 2   # Must be set to number of nodes in each node pool
          completions: 2   # Must be set to number of nodes in each node pool
          backoffLimit: 0   # Must be set to 0. Fail the job when any pod fails.
          template:
            spec:
              affinity: # The affinity section is to make sure there is only one Multislice job running in a single Multislice Group. More on this below.
                podAffinity:
                  requiredDuringSchedulingIgnoredDuringExecution:
                  - labelSelector:
                      matchExpressions:
                      - key: jobset.sigs.k8s.io/jobset-name
                        operator: In
                        values:
                        - multislice-job # JobSet name
                    topologyKey: MultisliceGroup
                podAntiAffinity:
                  requiredDuringSchedulingIgnoredDuringExecution:
                  - labelSelector:
                      matchExpressions:
                      - key: jobset.sigs.k8s.io/jobset-name
                        operator: NotIn
                        values:
                        - multislice-job # JobSet name
                    topologyKey: MultisliceGroup
                    namespaceSelector:
                      matchExpressions:
                      - key: jobset.sigs.k8s.io/jobset-name
                        operator: Exists
              hostNetwork: true
              dnsPolicy: ClusterFirstWithHostNet
              nodeSelector:
                cloud.google.com/gke-tpu-accelerator: tpu-v4-podslice
                cloud.google.com/gke-tpu-topology: 2x2x2
                MultisliceGroupSize: "4"
              containers:
              - name: jax-tpu
                image: python:3.8
                ports:
                - containerPort: 8471
                - containerPort: 8080 # Port for MXLA coordinator
                securityContext:
                  privileged: true
                command:
                - bash
                - -c
                - |
                  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())'
                resources:
                  limits:
                    google.com/tpu: 4 # Number of Cloud TPU VMs per worker

• Le JobSet crée quatre tâches indexées indexées (nombre de secteurs). Les tâches indexées indexées suivent le modèle de nom suivant : <jobSetName>-<replicateJobName>-<job-index> (dans ce cas, il s'agit de multislice-job-slice-{0,1,2,3}). Les tâches créées apparaissent dans Google Cloud CLI sous Workloads.

• Le JobSet est un service sans adresse IP de cluster dont le nom est identique au nom du jeu de tâches (dans ce cas, il s'agit de multislice-job).

• Chaque tâche indexée crée deux pods Kubernetes (nombre de nœuds de calcul). Les tâches suivent le format de nom suivant: <jobSetName>-<replicateJobName>-<job-index>-<worker-index>-<5 letter suffix> (dans ce cas, il s'agit de multislice-job-slice-{0,1,2,3}-{1,2}-{5 letter suffix}).

• Le multisegment sur Google Kubernetes Engine n'est actuellement compatible qu'avec l'entraînement synchrone multicontrôleur. Pour ce faire, parallelism et completions doivent être définis sur le nombre de nœuds dans chaque pool de nœuds, et backoff doit être défini sur 0. Le nombre de nœuds est divisé par le nombre de cœurs/nœuds. Pour les TPU v4, il s'agit de huit. Pour v4-16, définissez parallelism et completions sur 2.

• La version JAX minimale est: v0.4.9.

Injection de variables d'environnement dans GKE

Le webhook GKE injecte automatiquement des variables d'environnement dans la spécification du pod de la tâche indexée, qui est ensuite héritée par le pod Kubernetes exécutant la charge de travail multisegment. GKE injecte les variables d'environnement suivantes:

Variable d'environnement	Valeur
TPU_WORKER_ID	metadata.annotations['batch.kubernetes.io/job-completion-index']
TPU_WORKER_HOSTNAMES	Calculez la liste d'index du pod séparés par des virgules comme suit: <job-name>-0.<subdomain>,........,<job-name>-<n-1>.<subdomain>
MEGASCALE_NUM_SLICES	metadata.annotations['jobset.sigs.k8s.io/replicatedjob-replicas']
MEGASCALE_SLICE_ID	metadata.annotations['jobset.sigs.k8s.io/job-index']
MEGASCALE_COORDINATOR_ADDRESS	Generate the string using the jobSetName and the replicatedJobName: <jobSetName>-<replicatedJobName>-0-0.<sous-domaine>

Remarque :Si l'une de ces variables d'environnement est déjà spécifiée dans la spécification du pod, GKE n'injectera aucune variable d'environnement (en supposant que vous essayiez de configurer ces variables d'environnement manuellement).

Facultatif: Désactiver hostNetwork sur vos pods GKE

Utilisez hostNetwork: true dans les spécifications de votre pod pour ignorer l'ensemble de la pile réseau Kubernetes et laisser vos pods Kubernetes utiliser le réseau hôte directement pour la communication de VM à VM. Cela améliore les performances du réseau entre les tranches. Pour continuer à utiliser podHostnames pour la détection de nœuds de calcul avec hostNetwork, définissez dnsPolicy: ClusterFirstWithHostNet. C'est important lorsque vous reprenez automatiquement les tâches d'entraînement et que vous devez avoir les mêmes noms pour actualiser les mêmes points de contrôle.

Pour désactiver hostNetworking, supprimez les deux lignes suivantes des spécifications de votre pod:

hostNetwork: true
dnsPolicy: ClusterFirstWithHostNet

Orchestrer des tâches multisegments

Éviter les interblocages

Si vous envoyez plusieurs tâches multisegments (JobSet A et B) d'un même élément MultisliceGroupSize, il est possible que les deux ensembles de tâches soient programmés partiellement (par exemple, cinq tâches enfants sur huit sont programmées à partir de tâches A et trois tâches enfants sur huit sont programmées à partir d'une tâche B). Vos pools de nœuds restent alors inactifs, et les tâches A et B sont toutes les deux bloquées.

Pour éviter cela, vous pouvez utiliser podAffinity et podAntiAffinity. Pour les charges de travail multisegment, remplacez topologyKey par le libellé (MultisliceGroup) représentant un groupe multisegment au lieu du libellé cloud.google.com/gke-nodepool. Exemple :

affinity:
  podAffinity:
    requiredDuringSchedulingIgnoredDuringExecution:
      - labelSelector:
          matchExpressions:
          - key: jobset.sigs.k8s.io/jobset-name
            operator: In
            values:
            - multislice-job # JobSet name
          topologyKey: MultisliceGroup
  podAntiAffinity:
    requiredDuringSchedulingIgnoredDuringExecution:
    - labelSelector:
        matchExpressions:
        - key: jobset.sigs.k8s.io/jobset-name
          operator: NotIn
          values:
          - multislice-job # JobSet name
      topologyKey: MultisliceGroup
      namespaceSelector:
        matchExpressions:
        - key: jobset.sigs.k8s.io/jobset-name
          operator: Exists

Observabilité

Journaux Cloud Logging

Les pods GKE de l'ensemble de tâches sont nommés selon le modèle suivant:

<jobSetName>-<replicateJobName>-<job-index>-<worker-index>-<5 letter suffix>

Vous pouvez afficher vos journaux à l'aide de l'explorateur de journaux Cloud Logging avec le filtre suivant afin d'afficher les journaux de conteneur pour votre charge de travail:

resource.type="k8s_container"
resource.labels.cluster_name=<cluster-name>
labels."k8s-pod/jobset_sigs_k8s_io/jobset-name"=<jobSetName>

Pour filtrer les journaux du segment x et du nœud de calcul y, utilisez le filtre suivant:

resource.type="k8s_container"
resource.labels.cluster_name=<cluster-name>
labels."k8s-pod/jobset_sigs_k8s_io/jobset-name"=<jobSetName>
resource.labels.pod_name:<jobSetName>-<replicateJobName>-<job-index>-<worker-index>