Organiza cargas de trabajo de Multislice con JobSet y Kueue

En este instructivo, se muestra cómo organizar varias cargas de trabajo de varias porciones en Google Kubernetes Engine (GKE) para mejorar el uso de recursos. Implementarás una carga de trabajo de Jax como ejemplo, la ejecutarás en TPU Multislice y, luego, implementarás el encolamiento de trabajos con JobSet y Kueue. Kueue determina cuándo deben ejecutarse los Jobs según los recursos disponibles, las cuotas y una jerarquía para el uso compartido equitativo entre los equipos.

Este instructivo está dirigido a ingenieros de aprendizaje automático (AA) y administradores y operadores de plataformas que estén interesados en las capacidades de organización de contenedores de Kubernetes para entrenar LLM. Para obtener más información sobre los roles comunes y las tareas de ejemplo a las que hacemos referencia en el contenido de Google Cloud , consulta Roles y tareas comunes de los usuarios de GKE.

Antes de leer esta página, asegúrate de estar familiarizado con lo siguiente:

• Disponibilidad de la versión actual de TPU con la arquitectura del sistema de Cloud TPU
• TPU Multislice en GKE

Prepare el entorno

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

2. Configura las variables de entorno predeterminadas con el comando gcloud config set:

   gcloud config set project PROJECT_ID
   

   Reemplaza PROJECT_ID por tu Google Cloud ID del proyecto.

Los clústeres de Autopilot que ejecutan la versión 1.29.2-gke.1521000 o posterior habilitan las TPUs de forma predeterminada. Las TPU en los clústeres de 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=CONTROL_PLANE_LOCATION \
    --cluster-version 1.29.2-gke.1521000 \
    --release-channel rapid

Este comando incluye lo siguiente:

• La marca --location especifica la región de Compute Engine del plano de control del clúster.
• La marca --cluster-version especifica la versión de Kubernetes para tu clúster.
• La marca --release-channel especifica el canal de versiones de tu clúster. En este caso, el canal rápido admite las versiones más recientes disponibles en GKE.

Estándar

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

Reemplaza CONTROL_PLANE_LOCATION por la ubicación en la que deseas crear el 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 en modo Standard

En esta sección, crearás grupos de nodo TPU con el comando gcloud beta container node-pools create.

1. Crea el primer grupo de nodos llamado nodepool1:

   gcloud beta container node-pools create nodepool1 \
       --location=CONTROL_PLANE_LOCATION \
       --cluster=multislice-cluster \
       --node-locations=NODE_LOCATION \
       --machine-type=ct5lp-hightpu-4t \
       --tpu-topology=2x4 \
       --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=CONTROL_PLANE_LOCATION \
       --cluster=multislice-cluster \
       --node-locations=NODE_LOCATION \
       --machine-type=ct5lp-hightpu-4t \
       --tpu-topology=2x4 \
       --project=PROJECT_ID
   

3. Crea el tercer grupo de nodos llamado nodepool3:

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

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

En los pasos anteriores, usaste el comando gcloud beta container node-pools create para crear los grupos de nodos. Estos comandos usan las siguientes marcas:

• --node-locations: Es la lista separada por comas de una o más zonas en las que GKE crea los grupos de nodos.
• --machine-type: el tipo de máquina que se usará para los nodos. En este caso, usaste ct5lp-hightpu-4t. Para obtener más información sobre los tipos de máquinas compatibles con TPU, usa la tabla en Elige la versión de TPU.
• --tpu-topology: Es la topología de TPU que se usará para el grupo de nodos. En este caso, usaste 2x4. Para obtener más información sobre las topologías de TPU, consulta Cómo elegir la topología de TPU.

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. Un Jobset es una API de cargas de trabajo que te permite administrar un grupo de objetos Job de Kubernetes como una unidad. El caso de uso más común de un JobSet es el entrenamiento distribuido, pero también puedes usarlo para ejecutar cargas de trabajo por lotes.

Los siguientes 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. Instala la API de JobSet en tu clúster:

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

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

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

Nota: Para mitigar cualquier problema, asegúrate de que tu clúster cumpla con todos los requisitos previos que se indican en la sección Antes de comenzar de la documentación de instalación de Jobset. Si tienes problemas después de asegurarte de que tu clúster cumple con los requisitos, visita Solución de problemas de Jobset.

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

Las métricas y los paneles de observabilidad de JobSet y del grupo de nodos en la Google Cloud consola están disponibles para el público en general.

Paneles

Para ver el estado de tus grupos de nodos de TPU de varios hosts en GKE, ve al panel Estado del grupo de nodos de TPU de GKE que proporciona Cloud Monitoring:

Ir al estado del grupo de nodos de TPU de GKE

Para obtener más información, consulta Supervisa las métricas de estado de los nodos y los grupos de nodos de TPU.

En la página Kubernetes Engine AI/ML de la consola deGoogle Cloud , la pestaña Implementación de IA > Trabajos muestra un panel de supervisión de JobSet con información integral sobre el estado y el rendimiento de los JobSets y su infraestructura subyacente, como el estado de JobSet, la disponibilidad de réplicas y el estado de las réplicas. El panel también incluye métricas de infraestructura, como las de CPU, GPU, TPU, memoria y almacenamiento. Para obtener más información, consulta Supervisa el estado de JobSet con métricas.

Supervisa qué Pods se están ejecutando

kubectl get pods

El resultado es similar a este:

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.

Supervisa el estado de JobSet con métricas

Para comprender si un JobSet se ejecuta según lo previsto o inferir si se interrumpió, puedes usar las métricas de Prometheus del paquete de métricas de JobSet, como kube_jobset_succeeded_replicas.

Ten en cuenta que las métricas de estado de Jobset solo son compatibles con la versión 1.32.1-gke.135700 de GKE o versiones posteriores. Las métricas de estado de JobSet están habilitadas de forma predeterminada en los clústeres creados recientemente con versiones compatibles. En el caso de los clústeres existentes que se actualizan a versiones compatibles, los clientes deben habilitar manualmente el paquete de métricas de JobSet. Para obtener más información, consulta la documentación.

Para este instructivo, verifica la finalización de JobSet con esta consulta de PromQL:

kube_jobset_succeeded_replicas{
  cluster="multislice-cluster",
  jobset_name=~"mulitslice-.*"}

Supervisa el tiempo de actividad, los tiempos de recuperación (TTR) y los tiempos entre interrupciones (TBI) de JobSet

Las siguientes métricas son útiles para supervisar la disponibilidad de un JobSet:

• kubernetes.io/jobset/uptime: Es el tiempo total en el que el JobSet estuvo disponible.
• kubernetes.io/jobset/times_to_recover: Es la distribución del período de recuperación de un JobSet. Cada muestra indica un solo evento de recuperación de un período de inactividad para el JobSet.
• kubernetes.io/jobset/times_between_interruptions: Es la distribución del intervalo entre el final de la interrupción anterior y el comienzo de la interrupción actual para un JobSet. Cada muestra indica una sola duración entre la interrupción anterior y la actual.

Estas métricas se aplican a los JobSets que tienen exactamente un trabajo replicado de GPU o TPU. El cálculo de las métricas se basa solo en la disponibilidad de ese único trabajo replicado. Las métricas son compatibles con todas las versiones de GKE.

Para ver el tiempo de actividad de los JobSets que usaste en este instructivo, ejecuta la siguiente consulta de PromQL:

avg_over_time(
  kubernetes_io:jobset_uptime{
    monitored_resource="k8s_entity", entity_type="jobset",
    entity_name=~"multislice-.*",cluster_name="multislice-cluster"}[${__interval}])

Para ver las distribuciones del TBI de los JobSets de este instructivo, ejecuta la siguiente consulta en PromQL:

histogram_quantile(0.50,
  sum_over_time(
    kubernetes_io:jobset_times_between_interruptions_bucket{
      monitored_resource="k8s_entity",entity_type="jobset",
      entity_name=~"multislice-.*",cluster_name="multislice-cluster"}[${__interval}]))

Puedes extender el intervalo de la consulta a un horizonte temporal más largo, como 7 días, y calcular el tiempo promedio entre interrupciones (MTBI) durante este período:

sum(sum_over_time(
  kubernetes_io:jobset_times_between_interruptions_sum{
    monitored_resource="k8s_entity",entity_type="jobset",
    entity_name=~"multislice-.*",cluster_name="multislice-cluster"}[${__interval}]))
/
sum(sum_over_time(
  kubernetes_io:jobset_times_between_interruptions_count{
    monitored_resource="k8s_entity",entity_type="jobset",
    entity_name=~"multislice-.*",cluster_name="multislice-cluster"}[${__interval}]))

Para ver las distribuciones del TTR, puedes ejecutar las siguientes consultas de PromQL:

histogram_quantile(0.50,
  sum_over_time(
    kubernetes_io:jobset_times_to_recover_bucket{
      monitored_resource="k8s_entity",entity_type="jobset",
      entity_name=~"multislice-.*",cluster_name="multislice-cluster"}[${__interval}]))

Después de aumentar el intervalo de consultas a un horizonte temporal más largo, como 7 días, puedes calcular el tiempo medio de recuperación (MTTR) durante este período:

sum(sum_over_time(
  kubernetes_io:jobset_times_to_recover_sum{
    monitored_resource="k8s_entity",entity_type="jobset",
    entity_name=~"multislice-.*",cluster_name="multislice-cluster"}[${__interval}]))
/
sum(sum_over_time(
  kubernetes_io:jobset_times_to_recover_count{
    monitored_resource="k8s_entity",entity_type="jobset",
    entity_name=~"multislice-.*",cluster_name="multislice-cluster"}[${__interval}]))

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

GKE incluye una política de preferencia, que define cómo Kueue asigna los recursos disponibles. La política especifica que se puede interrumpir una carga de trabajo si una carga de trabajo de mayor prioridad necesita los recursos. Es más probable que las cargas de trabajo con un valor de prioridad más bajo se interrumpan por cargas de trabajo con una prioridad más alta.