Exécuter une charge de travail de petit lot avec des TPU et le mode de provisionnement de démarrage Flex


Ce guide vous explique comment optimiser le provisionnement de TPU pour les charges de travail d'entraînement de taille moyenne et petite à l'aide du mode de provisionnement à démarrage flexible. Dans ce guide, vous utilisez flex-start pour déployer une charge de travail qui se compose d'un pool de nœuds de tranche TPU.

Ce guide est destiné aux ingénieurs en machine learning (ML), aux administrateurs et opérateurs de plate-forme, ainsi qu'aux spécialistes des données et de l'IA qui souhaitent utiliser les fonctionnalités d'orchestration de conteneurs Kubernetes pour exécuter des charges de travail par lot. Pour en savoir plus sur les rôles courants et les exemples de tâches que nous citons dans le contenu Trusted Cloud by S3NS , consultez Rôles utilisateur et tâches courantes de l'utilisateur dans GKE Enterprise.

Tarification Flex-start

Le démarrage flexible est recommandé si votre charge de travail nécessite des ressources provisionnées de manière dynamique selon les besoins, pendant sept jours maximum avec des réservations à court terme, sans gestion complexe des quotas et avec un accès économique. Le démarrage flexible est optimisé par le planificateur de charges de travail dynamique et facturé selon les tarifs du planificateur de charges de travail dynamique :

  • Remise (jusqu'à 53 %) sur les vCPU, les GPU et les TPU.
  • Vous payez à l'usage.

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 Google Cloud CLI pour cette tâche, installez puis initialisez gcloud CLI. Si vous avez déjà installé gcloud CLI, assurez-vous de disposer de la dernière version en exécutant la commande gcloud components update.
  • Vérifiez que vous disposez d'un cluster Autopilot ou d'un cluster Standard exécutant la version 1.33.0-gke.1712000 ou ultérieure.
  • Assurez-vous de bien comprendre les limites du démarrage flexible.
  • Lorsque vous utilisez un cluster standard, assurez-vous de conserver au moins un pool de nœuds sans activer le démarrage flexible pour que le cluster fonctionne correctement.
  • Vérifiez que vous disposez d'un quota de TPU préemptifs dans les emplacements de vos nœuds.

Créer un pool de nœuds avec démarrage Flex

Pour créer un pool de nœuds avec le démarrage flexible activé sur un cluster Standard existant, utilisez gcloud CLI.

Si vous utilisez un cluster en mode Autopilot, ignorez cette section et passez à la section Exécuter une charge de travail par lot.

Vous pouvez créer un pool de nœuds de tranche TPU à hôte unique ou multi-hôte avec flex-start :

  1. Créez un pool de nœuds avec flex-start :

    Un seul hôte

    gcloud container node-pools create NODE_POOL_NAME \
        --cluster=CLUSTER_NAME \
        --location=LOCATION_NAME \
        --node-locations=NODE_ZONES \
        --machine-type=MACHINE_TYPE \
        --reservation-affinity=none \
        --enable-autoscaling \
        --flex-start \
        --num-nodes 0 \
        --min-nodes=0 \
        --max-nodes=1
    

    Remplacez les éléments suivants :

    • NODE_POOL_NAME : nom que vous avez choisi pour votre pool de nœuds.
    • CLUSTER_NAME : nom du cluster.
    • LOCATION_NAME : région de calcul du plan de contrôle du cluster.
    • NODE_ZONES : liste d'une ou de plusieurs zones, séparées par une virgule, dans lesquelles GKE crée le pool de nœuds.
    • MACHINE_TYPE : type de machine à utiliser pour les nœuds. Pour en savoir plus sur les types de machines compatibles avec les TPU, utilisez le tableau de la section Choisir la version de TPU.

    Hôte multiple

    gcloud container node-pools create NODE_POOL_NAME \
        --cluster=CLUSTER_NAME \
        --location=LOCATION_NAME \
        --node-locations=NODE_ZONES \
        --machine-type=MACHINE_TYPE \
        --tpu-topology=TPU_TOPOLOGY \
        --flex-start \
        --enable-autoscaling \
        --num-nodes=0 \
        --max-nodes=2 \
        --reservation-affinity=none \
        --no-enable-autorepair
    

    Remplacez les éléments suivants :

    • NODE_POOL_NAME : nom que vous avez choisi pour votre pool de nœuds.
    • CLUSTER_NAME : nom du cluster.
    • LOCATION_NAME : région de calcul du plan de contrôle du cluster.
    • NODE_ZONES : liste d'une ou de plusieurs zones, séparées par une virgule, dans lesquelles GKE crée le pool de nœuds.
    • MACHINE_TYPE : type de machine à utiliser pour les nœuds. Par exemple, vous pouvez utiliser ct6e-standard-4t pour TPU Trillium. Pour en savoir plus sur les types de machines disponibles, consultez Choisir la version de TPU.
    • TPU_TOPOLOGY : topologie physique de la tranche de TPU. Le format de la topologie dépend de la version du TPU. Pour en savoir plus sur les topologies TPU, utilisez le tableau de la section Choisir une topologie.

    La commande précédente utilise les options obligatoires suivantes lorsque vous créez un pool de nœuds avec démarrage flexible :

    • --enable-autoscaling : flex-start provisionne uniquement les ressources de calcul nécessaires lorsque votre charge de travail s'exécute. Vous devez définir les paramètres suivants :

      • --num-nodes=0
      • --min-nodes=0
      • --max-nodes défini sur le nombre de machines virtuelles dont votre tranche de TPU a besoin.

        Par exemple, votre commande de création de pool de nœuds peut inclure les paramètres suivants :

        ...
        --machine-type=ct6e-standard-4t \
        --tpu-topology=4x4 \
        --enable-autoscaling \
        --num-nodes=0 \
        --max-nodes=4 \
        

        Cette commande définit le champ --max-nodes sur 4, car une topologie 4x4 se compose de 16 puces et chaque VM ct6e-standard-4t en comporte 4.

      L'autoscaler de cluster augmente le nombre de nœuds jusqu'à ce que votre charge de travail soit satisfaite. Une fois votre charge de travail terminée, l'autoscaler de cluster réduit le nombre de nœuds à zéro.

    • --reservation-affinity=none : le démarrage flexible n'utilise pas et ne nécessite pas de réservations.

  2. Vérifiez l'état du démarrage flexible dans le pool de nœuds :

    gcloud container node-pools describe NODE_POOL_NAME \
        --cluster CLUSTER_NAME \
        --location LOCATION_NAME \
        --format="get(config.flexStart)"
    

    Si le démarrage flexible est activé dans le pool de nœuds, le champ flexStart est défini sur True.

Exécuter une charge de travail par lot

Dans cette section, vous allez créer un job qui planifie un nœud TPU avec un démarrage flexible. Dans Kubernetes, un contrôleur Job crée un ou plusieurs pods et s'assure qu'ils exécutent correctement une tâche spécifique.

  1. Dans la Trusted Cloud console, lancez une session Cloud Shell en cliquant sur Icône d'activation Cloud Shell Activer Cloud Shell. Une session s'ouvre dans le volet inférieur de la console Trusted Cloud .

  2. Créez un fichier nommé dws-flex-start.yaml :

    Un seul hôte

    Utilisez le fichier manifeste suivant pour le fichier dws-flex-start.yaml :

    apiVersion: batch/v1
    kind: Job
    metadata:
      name: job-1
    spec:
      template:
        spec:
          nodeSelector:
            cloud.google.com/gke-flex-start: "true"
            cloud.google.com/gke-tpu-accelerator: ACCELERATOR_TYPE
            cloud.google.com/gke-tpu-topology: TPU_TOPOLOGY
          containers:
          - name: container-1
            image: gcr.io/k8s-staging-perf-tests/sleep:latest
            args: ["3600s"] # Sleep for 1 hour
            resources:
              requests:
                  google.com/tpu: NUM_CHIPS
              limits:
                  google.com/tpu: NUM_CHIPS
          restartPolicy: OnFailure
    

    Hôte multiple

    Utilisez le fichier manifeste suivant pour le fichier dws-flex-start.yaml :

    apiVersion: v1
    kind: Service
    metadata:
      name: headless-svc
    spec:
      clusterIP: None
      selector:
        job-name: job-1
    ---
    apiVersion: batch/v1
    kind: Job
    metadata:
      name: job-1
    spec:
      backoffLimit: 0
      completions: 2
      parallelism: 2
      completionMode: Indexed
      template:
        spec:
          subdomain: headless-svc
          restartPolicy: Never
          nodeSelector:
              cloud.google.com/gke-flex-start: "true"
              cloud.google.com/gke-tpu-accelerator: ACCELERATOR_TYPE
              cloud.google.com/gke-tpu-topology: TPU_TOPOLOGY
          containers:
          - name: tpu-job
            image: python:3.10
            ports:
            - containerPort: 8471 # Default port using which TPU VMs communicate
            - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
            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("TPU cores:", jax.device_count())'
            resources:
              requests:
                google.com/tpu: NUM_CHIPS
              limits:
                google.com/tpu: NUM_CHIPS
    

    Remplacez les éléments suivants :

    • ACCELERATOR_TYPE : type d'accélérateur TPU que vous avez utilisé lors de la création des pools de nœuds. Par exemple, tpu-v4-podslice ou tpu-v5-lite-podslice.
    • TPU_TOPOLOGY : Topologie physique de la tranche de TPU. Par exemple, la valeur peut être 4x4x4 ou 2x2, selon la version du TPU.
    • NUM_CHIPS : le nombre de puces TPU dans chaque VM est de un, quatre ou huit. Pour en savoir plus, consultez Versions de TPU.
  3. Appliquez le fichier manifeste dws-flex-start.yaml :

    kubectl apply -f dws-flex-start.yaml
    
  4. Vérifiez que les jobs s'exécutent sur le même nœud :

    kubectl get pods
    

    Le résultat ressemble à ce qui suit :

    NAME    READY   STATUS      RESTARTS   AGE   IP       NODE               NOMINATED NODE   READINESS GATES
    job-1   0/1     Completed   0          19m   10.(...) gke-flex-zonal-a2  <none>           <none>
    

Effectuer un nettoyage

Pour éviter que les ressources utilisées sur cette page ne soient facturées sur votre compte Trusted Cloud by S3NS , supprimez le projet contenant les ressources, ou conservez le projet et supprimez les ressources individuelles.

Supprimer le projet

  1. In the Trusted Cloud console, go to the Manage resources page.

    Go to Manage resources

  2. In the project list, select the project that you want to delete, and then click Delete.
  3. In the dialog, type the project ID, and then click Shut down to delete the project.

Supprimer la ressource individuelle

  1. Supprimez les jobs :

    kubectl delete job -l "job-name in (job-1,job-2)"
    
  2. Supprimez le pool de nœuds :

    gcloud container node-pools delete NODE_POOL_NAME \
          --location LOCATION_NAME
    
  3. Supprimez le cluster à l'aide de la commande suivante :

    gcloud container clusters delete CLUSTER_NAME
    

Étapes suivantes