Organiza cargas de trabajo de Multislice con JobSet y Kueue

Autopilot Estándar

---

En este instructivo, se muestra cómo ejecutar una carga de trabajo de Jax con TPU Multislice en Google Kubernetes Engine (GKE) y Kueue. Kueue implementa la cola de Jobs, decide cuándo deben esperar los Jobs y cuándo deben iniciarse según las cuotas y una jerarquía para compartir recursos de manera equitativa entre los equipos.

En este instructivo, se muestra cómo organizar varias cargas de trabajo de Multislice que requieren recursos de TPU para ejecutarse de forma simultánea.

Antes de usar las TPU en GKE, te recomendamos que completes la siguiente ruta de aprendizaje:

1. Obtén información sobre la disponibilidad actual de la versión de TPU con la arquitectura del sistema de Cloud TPU.
2. Obtén más información sobre TPU Multislice en GKE.

Objetivos

Este instructivo está dirigido a los administradores de GKE que tienen clústeres de GKE existentes y desean ejecutar cargas de trabajo de Multislice por primera vez.

En este instructivo, se abarcan los siguientes pasos:

1. Prepara tu entorno con un clúster de GKE con tres porciones de TPU v5e. Cada porción de TPU tiene una topología 2x4 con 8 chips. Por lo tanto, 24 chips TPU v5e en total.
2. Crea los recursos de Kueue para asegurarte de que las cuotas se compartan de manera equitativa entre las cargas de trabajo.
3. Ejecuta tu carga de trabajo de Multislice.

Antes de comenzar

Antes de comenzar, asegúrate de haber realizado las siguientes tareas:

• Habilita la API de Google Kubernetes Engine.
Habilitar la API de Google Kubernetes Engine
• Si deseas usar Google Cloud CLI para esta tarea, instala y, luego, initialize gcloud CLI. Si ya instalaste gcloud CLI, ejecuta gcloud components update para obtener la versión más reciente.
  Nota: Para las instalaciones de gcloud CLI existentes, asegúrate de configurar las propiedades compute/region y compute/zone. Cuando configuras las ubicaciones predeterminadas, puedes evitar errores en gcloud CLI como el siguiente: One of [--zone, --region] must be supplied: Please specify location.

• Instala JobSet v0.2.3 o una versión posterior.

• Instala Kueue v0.4.1 o una versión posterior.

Prepara el entorno

1. En la consola de Google Cloud, inicia una instancia de Cloud Shell:
   Abrir Cloud Shell

2. Configura las variables de entorno predeterminadas:

   gcloud config set project PROJECT_ID
   gcloud config set compute/region COMPUTE_REGION
   

   Reemplaza los siguientes valores:

   • PROJECT_ID: el ID del proyecto de Google Cloud
   • COMPUTE_REGION: la región de Compute Engine

Los clústeres de Autopilot que ejecutan la versión 1.29.2-gke.1521000 o posterior habilitan las TPUs de forma predeterminada. Las TPUs en los clústeres en modo Autopilot se configuran en la especificación de la carga de trabajo. Para obtener más información, consulta la sección Define tus cargas de trabajo de Multislice con JobSets.

Crea un clúster de GKE

En Cloud Shell, crea un clúster de GKE:

Autopilot

gcloud container clusters create-auto multislice-cluster \
    --location=LOCATION \
    --cluster-version 1.29.2-gke.1521000 \
    --release-channel rapid

Estándar

gcloud container clusters create multislice-cluster \
    --location=LOCATION

Reemplaza LOCATION por la ubicación en la que deseas crear tu clúster. Asegúrate de que tenga capacidad para el tipo de máquina ct5lp-hightpu-4t. La creación del clúster puede tomar varios minutos.

Si usas el modo Autopilot de GKE, ve a la sección Crea los recursos de Kueue. Los clústeres de Autopilot que ejecutan la versión 1.29.2-gke.1521000 o posterior habilitan las TPUs de forma predeterminada.

Crea tres grupos de nodos de porción de TPU de modo Standard

1. Crea el primer grupo de nodos llamado nodepool1:

   gcloud beta container node-pools create nodepool1 \
       --location=LOCATION \
       --cluster=multislice-cluster \
       --node-locations=NODE_LOCATION \
       --machine-type=ct5lp-hightpu-4t \
       --tpu-topology=2x4 \
       --num-nodes=2 \
       --project=PROJECT_ID
   

   Reemplaza NODE_LOCATION por una o más zonas de la región del clúster en la que deseas crear los nodos.

2. Crea el segundo grupo de nodos llamado nodepool2:

   gcloud beta container node-pools create nodepool2 \
       --location=LOCATION \
       --cluster=multislice-cluster \
       --node-locations=NODE_LOCATION \
       --machine-type=ct5lp-hightpu-4t \
       --tpu-topology=2x4 \
       --num-nodes=2 \
       --project=PROJECT_ID
   

3. Crea el tercer grupo de nodos llamado nodepool3:

   gcloud beta container node-pools create nodepool3 \
       --location=LOCATION \
       --cluster=multislice-cluster \
       --node-locations=NODE_LOCATION \
       --machine-type=ct5lp-hightpu-4t \
       --tpu-topology=2x4 \
       --num-nodes=2 \
       --project=PROJECT_ID
   

GKE crea tres grupos de nodos. Cada grupo de nodos es una porción de TPU independiente.

Crea los recursos de Kueue

1. Crea el siguiente manifiesto kueue.yaml:

   apiVersion: kueue.x-k8s.io/v1beta1
   kind: ResourceFlavor
   metadata:
     name: "vlp-24"
   spec:
     nodeLabels:
       cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
       cloud.google.com/gke-tpu-topology: 2x4
   ---
   apiVersion: kueue.x-k8s.io/v1beta1
   kind: ClusterQueue
   metadata:
     name: "cluster-queue"
   spec:
     namespaceSelector: {}
     queueingStrategy: BestEffortFIFO
     resourceGroups:
     - coveredResources: ["google.com/tpu"]
       flavors:
       - name: "vlp-24"
         resources:
         - name: "google.com/tpu"
           nominalQuota: 24
   
   ---
   apiVersion: kueue.x-k8s.io/v1beta1
   kind: LocalQueue
   metadata:
     namespace: default
     name: multislice-queue
   spec:
     clusterQueue: cluster-queue
   

2. Aplica el manifiesto kueue.yaml:

   kubectl apply -f kueue.yaml
   

   GKE crea los siguientes recursos de Kueue:

• ResourceFlavor: una abstracción de los recursos en un clúster. En este ejemplo, GKE crea tres porciones de TPU con la topología 2x4. Cada porción de TPU tiene una topología 2x4 con 8 chips (24 chips TPU en total).
• ClusterQueue: una cola global que administra las cargas de trabajo y los recursos del clúster.
• LocalQueue: agrupa las cargas de trabajo estrechamente relacionadas que suele ejecutar un solo usuario. Cada LocalQueue apunta a una ClusterQueue desde la cual se asignan los recursos para ejecutar sus cargas de trabajo. Una carga de trabajo de Kube es una abstracción que representa una carga de trabajo por lotes, en este caso, cada carga de trabajo es un JobSet.

Define tus cargas de trabajo de Multislice con JobSets

En esta sección, crearás tres JobSets. Estos JobSets ejecutan una carga de trabajo de Jax que da como resultado la cantidad global de chips TPU en la porción y, luego, se suspende durante 60 segundos para simular algún tiempo de entrenamiento de modelos y, luego, se cierra.

1. Crea el siguiente manifiesto jobsets-multislice.yaml:

   Autopilot

   apiVersion: jobset.x-k8s.io/v1alpha2
   kind: JobSet
   metadata:
     name: multislice-1slice  
     labels:
       kueue.x-k8s.io/queue-name: multislice-queue  
     annotations:
       alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
   spec:
     failurePolicy:
       maxRestarts: 4
     replicatedJobs:
       - name: slice
         replicas: 1
         template:
           spec:
             parallelism: 2
             completions: 2
             backoffLimit: 0
             template:
               spec:
                 nodeSelector:
                   cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
                   cloud.google.com/gke-tpu-topology: 2x4
                 containers:
                 - name: jax-tpu
                   image: python:3.8
                   ports:
                   - containerPort: 8471
                   - containerPort: 8080
                   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
   
   ---
   apiVersion: jobset.x-k8s.io/v1alpha2
   kind: JobSet
   metadata:
     name: multislice-2slice
     labels:
       kueue.x-k8s.io/queue-name: multislice-queue
     annotations:
       alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
   spec:
     failurePolicy:
       maxRestarts: 4
     replicatedJobs:
       - name: slice
         replicas: 2
         template:
           spec:
             parallelism: 2
             completions: 2
             backoffLimit: 0
             template:
               spec:
                 nodeSelector:
                   cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
                   cloud.google.com/gke-tpu-topology: 2x4
                 containers:
                 - name: jax-tpu
                   image: python:3.8
                   ports:
                   - containerPort: 8471
                   - containerPort: 8080
                   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())'
                     sleep 60
                   resources:
                     limits:
                       google.com/tpu: 4
   ---
   apiVersion: jobset.x-k8s.io/v1alpha2
   kind: JobSet
   metadata:
     name: multislice-3slice
     labels:
       kueue.x-k8s.io/queue-name: multislice-queue
     annotations:
       alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
   spec:
     failurePolicy:
       maxRestarts: 4
     replicatedJobs:
       - name: slice
         replicas: 3
         template:
           spec:
             parallelism: 2
             completions: 2
             backoffLimit: 0
             template:
               spec:
                 nodeSelector:
                   cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
                   cloud.google.com/gke-tpu-topology: 2x4
                 containers:
                 - name: jax-tpu
                   image: python:3.8
                   ports:
                   - containerPort: 8471
                   - containerPort: 8080
                   command:
                   - bash
                   - -c
                   - |
                     sleep 60
                   resources:
                     limits:
                       google.com/tpu: 4
   

   Estándar

   apiVersion: jobset.x-k8s.io/v1alpha2
   kind: JobSet
   metadata:
     name: multislice-1slice  
     labels:
       kueue.x-k8s.io/queue-name: multislice-queue  
     annotations:
       alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
   spec:
     failurePolicy:
       maxRestarts: 4
     replicatedJobs:
       - name: slice
         replicas: 1
         template:
           spec:
             parallelism: 2
             completions: 2
             backoffLimit: 0
             template:
               spec:
                 hostNetwork: true
                 dnsPolicy: ClusterFirstWithHostNet
                 nodeSelector:
                   cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
                   cloud.google.com/gke-tpu-topology: 2x4
                 containers:
                 - name: jax-tpu
                   image: python:3.8
                   ports:
                   - containerPort: 8471
                   - containerPort: 8080
                   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
   
   ---
   apiVersion: jobset.x-k8s.io/v1alpha2
   kind: JobSet
   metadata:
     name: multislice-2slice
     labels:
       kueue.x-k8s.io/queue-name: multislice-queue
     annotations:
       alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
   spec:
     failurePolicy:
       maxRestarts: 4
     replicatedJobs:
       - name: slice
         replicas: 2
         template:
           spec:
             parallelism: 2
             completions: 2
             backoffLimit: 0
             template:
               spec:
                 hostNetwork: true
                 dnsPolicy: ClusterFirstWithHostNet
                 nodeSelector:
                   cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
                   cloud.google.com/gke-tpu-topology: 2x4
                 containers:
                 - name: jax-tpu
                   image: python:3.8
                   ports:
                   - containerPort: 8471
                   - containerPort: 8080
                   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())'
                     sleep 60
                   resources:
                     limits:
                       google.com/tpu: 4
   ---
   apiVersion: jobset.x-k8s.io/v1alpha2
   kind: JobSet
   metadata:
     name: multislice-3slice
     labels:
       kueue.x-k8s.io/queue-name: multislice-queue
     annotations:
       alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
   spec:
     failurePolicy:
       maxRestarts: 4
     replicatedJobs:
       - name: slice
         replicas: 3
         template:
           spec:
             parallelism: 2
             completions: 2
             backoffLimit: 0
             template:
               spec:
                 hostNetwork: true
                 dnsPolicy: ClusterFirstWithHostNet
                 nodeSelector:
                   cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
                   cloud.google.com/gke-tpu-topology: 2x4
                 containers:
                 - name: jax-tpu
                   image: python:3.8
                   ports:
                   - containerPort: 8471
                   - containerPort: 8080
                   securityContext:
                     privileged: true
                   command:
                   - bash
                   - -c
                   - |
                     sleep 60
                   resources:
                     limits:
                       google.com/tpu: 4
   

2. Aplica el manifiesto jobsets-multislice.yaml:

   kubectl apply -f jobsets-multislice.yaml
   

GKE crea los Jobs con las siguientes solicitudes de recursos:

• El JobSet multislice-1slice crea un Job que requiere una porción de TPU en total.
• El JobSet multislice-2slice crea dos Jobs que requieren dos porciones de TPU en total.
• El JobSet multislice-3slice crea tres Jobs que requieren tres porciones de TPU en total.

Debido a que el clúster solo tiene tres porciones de TPU, no todos los JobSets pueden ejecutarse a la vez. Cuando Kueue pone en cola los tres JobSets multislice-3slice, sus Jobs se ejecutan solos hasta su finalización. multislice-1slice y multislice-2slice esperan y se ejecutan en conjunto después.

Verifica que Kueue admita las cargas de trabajo

1. Verifica las cargas de trabajo en cola en Kueue:

   kubectl get workloads
   

   El resultado es similar al siguiente:

   NAME                             QUEUE              ADMITTED BY     AGE
   jobset-multislice-1slice-2530a   multislice-queue                   3s
   jobset-multislice-2slice-ffb02   multislice-queue                   4s
   jobset-multislice-3slice-8c695   multislice-queue   cluster-queue   10s
   

Kueue pone en cola una o más cargas de trabajo según los recursos TPU que necesiten.

Supervisa las cargas de trabajo

1. Supervisa qué Pods se están ejecutando:

   kubectl get pods
   

   El resultado es similar al siguiente:

   NAME                                READY   STATUS      RESTARTS   AGE
   multislice-1slice-slice-0-0-pf2ll   1/1     Running     0          1s
   multislice-1slice-slice-0-1-55g62   1/1     Running     0          1s
   multislice-2slice-slice-0-0-f4hf7   1/1     Running     0          3s
   multislice-2slice-slice-0-1-c8kv7   1/1     Running     0          3s
   multislice-2slice-slice-1-0-7h46t   1/1     Running     0          3s
   multislice-2slice-slice-1-1-lj9hb   1/1     Running     0          3s
   multislice-3slice-slice-0-0-wzq9t   0/1     Completed   0          2m31s
   multislice-3slice-slice-0-1-zf4dp   0/1     Completed   0          2m30s
   multislice-3slice-slice-1-0-hbfn5   0/1     Completed   0          2m31s
   multislice-3slice-slice-1-1-45fgl   0/1     Completed   0          2m30s
   multislice-3slice-slice-2-0-wjbp4   0/1     Completed   0          2m30s
   multislice-3slice-slice-2-1-lwnvs   0/1     Completed   0          2m30s
   

   Consulta que GKE haya programado, creado y ejecutado los Pods para multislice-3slice primero. Luego, que GKE haya ejecutado los Pods de los JobSets multislice-1slice y multislice-2slice.

Completado correctamente: Ejecutaste correctamente tres cargas de trabajo Jax de Multislice.

Habilita la interrupción y la prioridad de las cargas de trabajo de Kueue

De manera opcional, puedes asignar prioridades de cargas de trabajo de Kueue que determinan el orden en que Kueue admite las cargas de trabajo en cola.

Nota: Según la configuración de la política de interrupción de ClusterQueue, se pueden producir algunas situaciones de interrupción. Por ejemplo, GKE puede interrumpir una carga de trabajo en ejecución para ejecutar una carga de trabajo de mayor prioridad si no hay recursos suficientes para ejecutar ambas cargas de trabajo de forma simultánea. GKE detiene y vuelve a poner en cola una carga de trabajo interrumpida para continuar con el procesamiento cuando sea posible.

1. Actualiza tu ClusterQueue para que tenga una política de interrupción:

   apiVersion: kueue.x-k8s.io/v1beta1
   kind: ResourceFlavor
   metadata:
     name: "vlp-24"
   spec:
     nodeLabels:
       cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
       cloud.google.com/gke-tpu-topology: 2x4
   ---
   apiVersion: kueue.x-k8s.io/v1beta1
   kind: ClusterQueue
   metadata:
     name: "cluster-queue"
   spec:
     namespaceSelector: {}
     resourceGroups:
     - coveredResources: ["google.com/tpu"]
       flavors:
       - name: "vlp-24"
         resources:
         - name: "google.com/tpu"
           nominalQuota: 24
    preemption:
       reclaimWithinCohort: Any
       withinClusterQueue: LowerPriority
   ---
   apiVersion: kueue.x-k8s.io/v1beta1
   kind: LocalQueue
   metadata:
     namespace: default
     name: multislice-queue
   spec:
     clusterQueue: cluster-queue
   

2. Crea una PriorityClass para cada nivel de prioridad distinto que desees asignar a las cargas de trabajo:

   apiVersion: scheduling.k8s.io/v1
   kind: PriorityClass
   metadata:
     name: low-priority
   value: 100
   globalDefault: false
   description: "This low priority class should be used for some Pods only."
   

3. Asigna el priorityClassName a tu JobSet:

   Autopilot

   apiVersion: jobset.x-k8s.io/v1alpha2
   kind: JobSet
   metadata:
     name: low-priority
     labels:
       kueue.x-k8s.io/queue-name: multislice-queue
     annotations:
       alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
   spec:
     failurePolicy:
       maxRestarts: 4
     replicatedJobs:
       - name: slice
         replicas: 1
         template:
           spec:
             parallelism: 2
             completions: 2
             backoffLimit: 0
             template:
               spec:
                 nodeSelector:
                   cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
                   cloud.google.com/gke-tpu-topology: 2x4
                 priorityClassName: low-priority
                 containers:
                 - name: jax-tpu
                   image: python:3.8
                   ports:
                   - containerPort: 8471
                   - containerPort: 8080
                   command:
                   - bash
                   - -c
                   - |
                     sleep 60
                   resources:
                     limits:
                       google.com/tpu: 4 # Number of TPU chips per worker
   

   Estándar

   apiVersion: jobset.x-k8s.io/v1alpha2
   kind: JobSet
   metadata:
     name: low-priority
     labels:
       kueue.x-k8s.io/queue-name: multislice-queue
     annotations:
       alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
   spec:
     failurePolicy:
       maxRestarts: 4
     replicatedJobs:
       - name: slice
         replicas: 1
         template:
           spec:
             parallelism: 2
             completions: 2
             backoffLimit: 0
             template:
               spec:
                 hostNetwork: true
                 dnsPolicy: ClusterFirstWithHostNet
                 nodeSelector:
                   cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
                   cloud.google.com/gke-tpu-topology: 2x4
                 priorityClassName: low-priority
                 containers:
                 - name: jax-tpu
                   image: python:3.8
                   ports:
                   - containerPort: 8471
                   - containerPort: 8080
                   securityContext:
                     privileged: true
                   command:
                   - bash
                   - -c
                   - |
                     sleep 60
                   resources:
                     limits:
                       google.com/tpu: 4 # Number of TPU chips per worker
     ```
   

Limpia

Para evitar que se apliquen cargos a tu cuenta de Google Cloud por los recursos usados en este instructivo, borra el proyecto que contiene los recursos o conserva el proyecto y borra los recursos individuales.

Borra el proyecto

Precaución: Borrar un proyecto tiene estas consecuencias:
• Se borra todo en el proyecto. Si usaste un proyecto existente para las tareas de este documento, cuando lo borres, también se borrará cualquier otro trabajo que hayas realizado en el proyecto.
• Se pierden los ID personalizados de proyectos. Cuando creaste este proyecto, es posible que hayas creado un ID del proyecto personalizado que desees utilizar en el futuro. Para conservar las URL que utilizan el ID del proyecto, como una URL appspot.com, borra los recursos seleccionados dentro del proyecto en lugar de borrar todo el proyecto.
1. En la consola de Google Cloud, ve a la página Administrar recursos.

   Ir a Administrar recursos

2. En la lista de proyectos, elige el proyecto que quieres borrar y haz clic en Borrar.
3. En el diálogo, escribe el ID del proyecto y, luego, haz clic en Cerrar para borrar el proyecto.

Borra el recurso individual

1. Borra el sistema de cuotas de Kueue:

   kubectl delete -n team-a localqueue
   kubectl delete -n team-b localqueue
   kubectl delete clusterqueue
   kubectl delete clusterqueue
   kubectl delete clusterqueue
   kubectl delete resourceflavor
   kubectl delete resourceflavor
   kubectl delete resourceflavor
   

2. Borra el manifiesto de Kueue:

   VERSION=kueue.x-k8s.io/v1beta1
   kubectl delete -f \
       https://github.com/kubernetes-sigs/kueue/releases/download/$VERSION/manifests.yaml
   

3. Borra el clúster:

   gcloud container clusters delete kueue-cohort --region=COMPUTE_REGION
   

¿Qué sigue?

• Más información sobre Kueue
• Obtén más información sobre cómo implementar un sistema de colas de trabajos con uso compartido de cuotas entre espacios de nombres en GKE.