Organiza cargas de trabajo de Multislice con JobSet y Kueue

Autopilot Estándar

En este instructivo, se muestra cómo organizar varias cargas de trabajo de varias porciones en Google Kubernetes Engine (GKE). Ejecuta una carga de trabajo de Jax con TPU Multislice, JobSet 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.

Recomendamos que completes este instructivo si usas cargas de trabajo que requieran recursos TPU para ejecutarse de forma simultánea.

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

Obtén información sobre la disponibilidad actual de la versión de TPU con la arquitectura del sistema de Cloud TPU.
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:

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.
Crea los recursos de Kueue para asegurarte de que las cuotas se compartan de manera equitativa entre las cargas de trabajo.
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, inicializa 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

En la consola de Google Cloud, inicia una instancia de Cloud Shell:
Abrir Cloud Shell
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

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.

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

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

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

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.

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

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

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

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.

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.

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

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."

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

In the Google Cloud console, go to the Manage resources page.
Go to Manage resources
In the project list, select the project that you want to delete, and then click Delete.
In the dialog, type the project ID, and then click Shut down to delete the project.

Borra el recurso individual

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

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

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.