Orquestrar cargas de trabalho do Multislice usando o JobSet e Kueue

Autopilot Padrão

Neste tutorial, mostramos como executar uma carga de trabalho Jax usando Multislice TPU no Google Kubernetes Engine (GKE) e no Kueue. O Kueue implementa o enfileiramento de jobs para decidir quando os jobs devem esperar e quando devem começar, com base em cotas e em uma hierarquia para compartilhar recursos entre as equipes.

Neste tutorial, mostramos como orquestrar várias cargas de trabalho Multislice que exigem recursos de TPU para serem executadas simultaneamente.

Antes de usar TPUs no GKE, recomendamos que você conclua o seguinte programa de aprendizado:

Saiba mais sobre a disponibilidade atual da versão da TPU com a arquitetura do sistema do Cloud TPU.
Saiba mais sobre o Multislice TPU no GKE.

Objetivos

O público-alvo deste tutorial são os administradores do GKE que têm clusters no modo GKE Standard e querem executar cargas de trabalho Multislice pela primeira vez.

Este tutorial abrange as etapas a seguir:

Preparar seu ambiente com um cluster do GKE Standard com três frações da TPU v5e. Cada fração de TPU tem uma topologia 2x4 e quatro chips por host. Portanto, são 24 chips da TPU v5e no total.
Criar os recursos do Kueue para garantir que as cotas sejam compartilhadas de maneira equitativa entre as cargas de trabalho.
Executar a carga de trabalho Multislice.

Antes de começar

Antes de começar, verifique se você realizou as tarefas a seguir:

Ativar a API Google Kubernetes Engine.

Ativar a API Google Kubernetes Engine

Se você quiser usar a Google Cloud CLI para essa tarefa, instale e, em seguida, inicialize a CLI gcloud. Se você instalou a CLI gcloud anteriormente, instale a versão mais recente executando gcloud components update.
Observação: para instalações atuais da CLI gcloud, defina as propriedades compute/region e compute/zone. Quando você define os locais padrão, você evita erros como este na gcloud CLI: One of [--zone, --region] must be supplied: Please specify location.

Instale o JobSet v0.2.3 ou mais recente.
Instale o Kueue v0.4.1 ou mais recente.

Preparar o ambiente

No console do Google Cloud, inicie uma instância do Cloud Shell:
Abrir o Cloud Shell
Defina as variáveis de ambiente padrão:
```
gcloud config set project PROJECT_ID
gcloud config set compute/region COMPUTE_REGION
```
Substitua os seguintes valores:
- PROJECT_ID: é o ID do projeto do Google Cloud.
- COMPUTE_REGION: a região do Compute Engine.

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:

Autopilot

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

Padrão

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

Substitua 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 TPU no modo Standard

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

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

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

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

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

Criar os recursos do Kueue

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

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, três frações de TPU com a topologia 2x4 e quatro chips por host, ou seja, 24 chips de TPU.
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. Esses JobSets 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 do modelo e, em seguida, é encerrada.

Crie o seguinte manifesto 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

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

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

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

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

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.

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

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

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

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
  ```

Limpar

Para evitar cobranças na sua conta do Google Cloud pelos recursos usados no tutorial, exclua o projeto que os contém ou mantenha o projeto e exclua os recursos individuais.

Excluir o projeto

No Console do Google Cloud, acesse a página Gerenciar recursos.
Acessar "Gerenciar recursos"
Na lista de projetos, selecione o projeto que você quer excluir e clique em Excluir .
Na caixa de diálogo, digite o ID do projeto e clique em Encerrar para excluí-lo.

Excluir o recurso individual

Excluir o sistema de cotas do 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

Excluir o manifesto do Kueue:

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

Exclua o cluster:

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

A seguir

Saiba mais sobre o Kueue
Saiba como implementar um sistema de enfileiramento de jobs com compartilhamento de cotas entre namespaces no GKE.