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:
- Disponibilidade da versão atual da TPU com a arquitetura do sistema do Cloud TPU
- Multislice de TPU no GKE
Prepare o ambiente
No Google Cloud console, inicie uma instância do Cloud Shell:
Abrir o Cloud ShellDefina 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
.
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.
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
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ê usouct5lp-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ê usou2x4
. Para mais informações sobre topologias de TPU, consulte Escolher a topologia de TPU.
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, o GKE cria três
frações com a topologia
2x4
. Cada fração de TPU tem uma topologia2x4
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.
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
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
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
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.
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: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.