Orquestrar cargas de trabalho do Multislice usando o JobSet e Kueue

Neste tutorial, mostramos como orquestrar várias cargas de trabalho de várias frações no Google Kubernetes Engine (GKE) para melhorar a utilização de recursos. Você implanta uma carga de trabalho Jax como exemplo, executa ela na TPU Multislice e implementa o enfileiramento de jobs com JobSet e Kueue. O Kueue determina quando os jobs devem ser executados com base nos recursos disponíveis, nas cotas e em uma hierarquia para compartilhamento justo entre as equipes.

Este tutorial é destinado a engenheiros de machine learning (ML) e administradores e operadores de plataforma interessados nos recursos de orquestração de contêineres do Kubernetes para treinar LLMs. Para saber mais sobre papéis comuns e tarefas de exemplo referenciados no conteúdo do Google Cloud , consulte Tarefas e funções de usuário comuns do GKE.

Antes de ler esta página, confira se você conhece os seguintes conceitos:

Prepare o ambiente

  1. No Google Cloud console, inicie uma instância do Cloud Shell:
    Abrir o Cloud Shell

  2. Defina as variáveis de ambiente padrão usando o comando gcloud config set:

    gcloud config set project PROJECT_ID
    

    Substitua PROJECT_ID pelo Google Cloud ID do projeto.

Os clusters do Autopilot que executam a versão 1.29.2-gke.1521000 ou mais recente ativa as TPUs por padrão. As TPUs em clusters do Autopilot são configuradas na especificação da carga de trabalho. Para mais informações, consulte a seção Definir cargas de trabalho de multislice com JobSets.

Crie um cluster do GKE

No Cloud Shell, crie um cluster do GKE:

Piloto automático

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

Nesse comando:

  • A flag --location especifica a região do Compute Engine do plano de controle do cluster.
  • A flag --cluster-version especifica a versão do Kubernetes para o cluster.
  • A flag --release-channel especifica o canal de lançamento do seu cluster. Nesse caso, o canal rápido oferece suporte às versões mais recentes disponíveis no GKE.

Padrão

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

Substitua CONTROL_PLANE_LOCATION pelo local em que você quer criar o cluster. Verifique se ele tem capacidade para o tipo de máquina ct5lp-hightpu-4t. A criação do cluster pode levar vários minutos.

Se você usar o modo Autopilot do GKE, pule para a seção Criar os recursos do Kueue. Os clusters do Autopilot que executam a versão 1.29.2-gke.1521000 ou mais recente ativa as TPUs por padrão.

Criar três pools de nós de fração de TPU no modo Standard

Nesta seção, você cria pools de nós de TPU usando o comando gcloud beta container node-pools create.

  1. Crie o primeiro pool de nós chamado 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
    

    Substitua NODE_LOCATION por uma ou mais zonas na região do cluster em que você quer criar os nós.

  2. Crie o segundo pool de nós chamado 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. Crie o terceiro pool de nós chamado 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
    

O GKE cria três pools de nós. Cada pool de nós é uma fração de TPU separada.

Nas etapas anteriores, você usou o comando gcloud beta container node-pools create para criar os pools de nós. Esses comandos usam as seguintes flags:

  • --node-locations: a lista separada por vírgulas de uma ou mais zonas em que o GKE cria os pools de nós.
  • --machine-type: o tipo de máquina a ser usado para nós. Nesse caso, você usou ct5lp-hightpu-4t. Para mais informações sobre os tipos de máquinas compatíveis com TPU, use a tabela em Escolher a versão da TPU.
  • --tpu-topology: a topologia de TPU a ser usada para o pool de nós. Nesse caso, você usou 2x4. Para mais informações sobre topologias de TPU, consulte Escolher a topologia de TPU.

Criar os recursos do Kueue

  1. Crie o seguinte manifesto 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. Aplique o manifesto kueue.yaml:

    kubectl apply -f kueue.yaml
    

O GKE cria os seguintes recursos do Kueue:

  • ResourceFlavor: uma abstração dos recursos em um cluster. Neste exemplo, o GKE cria três frações com a topologia 2x4. Cada fração de TPU tem uma topologia 2x4 com 8 chips (24 chips de TPU no total).
  • ClusterQueue: uma fila global que gerencia cargas de trabalho e recursos de cluster.
  • LocalQueue: agrupa cargas de trabalho estreitamente relacionadas que normalmente são executadas por um único locatário (usuário). Cada LocalQueue aponta para uma ClusterQueue, de onde os recursos são alocados para executar as cargas de trabalho. Uma carga de trabalho do Kueue é uma abstração que representa uma carga de trabalho em lote. Nesse caso, cada carga de trabalho é um JobSet.

Definir cargas de trabalho Multislices com JobSets

Nesta seção, você vai criar três JobSets. Um Jobset é uma API de carga de trabalho que permite gerenciar um grupo de jobs do Kubernetes como uma unidade. O caso de uso mais comum para um JobSet é o treinamento distribuído, mas também é possível usá-lo para executar cargas de trabalho em lote.

Os JobSets a seguir executam uma carga de trabalho Jax que gera o número global de chips de TPU na fração, fica em suspensão por 60 segundos para simular algum tempo de treinamento de modelo e, em seguida, é encerrada.

  1. Instale a API JobSet no cluster:

    VERSION=v0.8.1
    kubectl apply --server-side -f https://github.com/kubernetes-sigs/jobset/releases/download/$VERSION/manifests.yaml
    
  2. Crie o seguinte manifesto jobsets-multislice.yaml:

    Piloto automático

    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
    

    Padrão

    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. Aplique o manifesto jobsets-multislice.yaml:

    kubectl apply -f jobsets-multislice.yaml
    

O GKE cria os jobs com as seguintes solicitações de recursos:

  • O JobSet multislice-1slice cria um job que requer uma fração da TPU no total.
  • O JobSet multislice-2slice cria dois jobs que exigem duas frações da TPU no total.
  • O JobSet multislice-3slice cria três jobs que exigem três frações da TPU no total.

Como o cluster tem apenas três frações de TPU, nem todos os JobSets podem ser executados ao mesmo tempo. Quando o Kueue coloca os três JobSets multislice-3slice na fila, os jobs são executados sozinhos até a conclusão. multislice-1slice e multislice-2slice aguardam e são executadas juntas, posteriormente.

Verificar se o Kueue admitiu as cargas de trabalho

  1. Verifique as cargas de trabalho em fila no Kueue:

    kubectl get workloads
    

    O resultado será assim:

    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
    

O Kueue enfileira uma ou mais cargas de trabalho, dependendo dos recursos de TPU necessários.

Monitorar as cargas de trabalho

As métricas e os painéis de observabilidade do JobSet e do pool de nós no Google Cloud console estão em disponibilidade geral.

Painéis

Para conferir o status dos pools de nós de vários hosts da TPU no GKE, acesse o painel Status do pool de nós de TPU do GKE fornecido pelo Cloud Monitoring:

Acessar o status do pool de nós de TPU do GKE

Para mais informações, consulte Monitorar métricas de integridade para nós e pools de nós de TPU.

Na página IA/ML do Kubernetes Engine no console doGoogle Cloud , a guia Implantação de IA > Jobs mostra um painel de monitoramento do JobSet com informações abrangentes sobre a integridade e o desempenho dos JobSets e da infraestrutura subjacente, como status do JobSet, prontidão e estado da réplica. O painel também inclui métricas de infraestrutura, como CPU, GPU, TPU, memória e armazenamento. Para mais informações, consulte Monitorar a integridade do JobSet com métricas.

Monitore quais pods estão em execução

kubectl get pods

O resultado será assim:

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

Observe que o GKE programou, criou e executou os pods para o multislice-3slice primeiro. Na sequência, o GKE executa os pods dos JobSets multislice-1slice e multislice-2slice

Monitorar a integridade do JobSet com métricas

Para entender se um JobSet está sendo executado conforme o esperado ou inferir se ele foi interrompido, use as métricas do Prometheus do pacote de métricas do JobSet, como kube_jobset_succeeded_replicas.

As métricas de integridade do Jobset só são compatíveis com a versão 1.32.1-gke.135700 ou mais recente do GKE. As métricas de integridade do JobSet são ativadas por padrão em clusters recém-criados com versões compatíveis. Para clusters atuais que são atualizados para versões compatíveis, os clientes precisam ativar manualmente o pacote de métricas do JobSet. Para saber mais, consulte a documentação.

Para este tutorial, verifique a conclusão do JobSet com esta consulta PromQL:

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

Monitore o tempo de atividade, os tempos de recuperação (TTR) e os tempos entre interrupções (TBI) do JobSet

As métricas a seguir são úteis para monitorar a disponibilidade de um JobSet:

  • kubernetes.io/jobset/uptime: tempo total em que o JobSet ficou disponível.
  • kubernetes.io/jobset/times_to_recover: distribuição do período de recuperação de um JobSet. Cada amostra indica um único evento de recuperação de um período de inatividade para o JobSet.
  • kubernetes.io/jobset/times_between_interruptions: distribuição do intervalo entre o fim da interrupção anterior e o início da interrupção atual para um JobSet. Cada amostra indica uma única duração entre a interrupção anterior e a atual.

Essas métricas se aplicam a JobSets que têm exatamente um job replicado de GPU ou TPU. O cálculo das métricas se baseia apenas na disponibilidade desse único job replicado. As métricas são compatíveis com todas as versões do GKE.

Para conferir o tempo de atividade dos JobSets usados neste tutorial, execute a seguinte consulta em PromQL:

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

Para conferir as distribuições de TBI dos JobSets deste tutorial, execute a seguinte consulta em 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}]))

Você pode estender o intervalo da consulta para um período mais longo, como 7 dias, e calcular o tempo médio entre interrupções (MTBI) nesse 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 conferir as distribuições de TTR, execute as seguintes consultas em 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}]))

Depois de aumentar o intervalo de consulta para um horizonte de tempo mais longo, como 7 dias, você pode calcular o tempo médio para recuperação (MTTR) nesse 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}]))

Ativar a preempção e as prioridades de carga de trabalho do Kueue

Também é possível atribuir prioridades às cargas de trabalho do Kueue que determinam a ordem em que as cargas de trabalho enfileiradas são admitidas pelo Kueue.

  1. Atualize seu ClusterQueue para que ele tenha uma política de preempção:

    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. Crie um PriorityClass para cada nível de prioridade distinto que você quer atribuir às cargas de trabalho:

    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. Atribua o priorityClassName ao seu JobSet:

    Piloto automático

    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
    

    Padrão

    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
    

O GKE inclui uma política de remoção forçada, que define como o Kueue atribui os recursos disponíveis. A política especifica que uma carga de trabalho pode ser interrompida se uma carga de trabalho de maior prioridade precisar dos recursos. As cargas de trabalho com um valor de prioridade mais baixo têm mais chances de serem interrompidas por cargas de trabalho de maior prioridade.