Nesta página, descrevemos como implantar cargas de trabalho que usam aceleradores da Cloud TPU (TPUs) em clusters do Autopilot do Google Kubernetes Engine (GKE). Você já deve estar familiarizado com estes conceitos:
Como as TPUs funcionam no Autopilot
Para usar TPUs em cargas de trabalho do Autopilot, é preciso solicitar uma versão da TPU e uma
topologia compatível para ela no manifesto da carga de trabalho. Em seguida, use os campos resources.requests
e resources.limits
do Kubernetes para especificar o número de chips de TPU para a carga de trabalho. Quando você implanta a carga de trabalho, o GKE provisiona nós que têm a configuração de TPU solicitada e programa seus pods nos nós. O GKE coloca cada carga de trabalho no próprio nó para que cada pod possa acessar todos os recursos do nó com risco minimizado de interrupção.
As TPUs no Autopilot são compatíveis com os seguintes recursos:
Planejar a configuração da TPU
Antes de solicitar TPUs, defina a configuração com base nos requisitos de CPU e memória da carga de trabalho. Você precisa decidir o seguinte:
- Versão do TPU: a versão específica do Cloud TPU, como v5e.
- Topologia da versão selecionada da TPU: a disposição e o número de TPUs.
A versão da TPU e a topologia selecionadas determinam se o GKE provisiona nós como frações de host único ou frações de vários hosts. Em frações de host único, cada nó é independente de outros nós da TPU. Em frações de vários hosts, o GKE cria um grupo de nós que têm VMs de TPU interconectadas. As frações de vários hosts são atômicas, o que significa que o GKE escalona todo o grupo interconectado de nós como uma única unidade.
Para informações sobre as versões de TPU disponíveis, as topologias correspondentes, a CPU e a capacidade de memória e o tipo de fração resultante, consulte Escolher uma configuração de TPU do Autopilot.
Preços
Saiba mais sobre os preços em Preços do Autopilot.
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
.
- Verifique se você tem um cluster do Autopilot em execução na versão 1.29.2-gke.1521000 ou mais recente do GKE.
- Para usar TPUs reservadas, verifique se você tem uma reserva de capacidade específica. Para mais instruções, consulte Como consumir recursos por zona reservados.
Verifique se você tem uma cota de TPU
Para criar nós de TPU, é necessário ter uma cota de TPU disponível, a menos que você esteja usando uma reserva de capacidade atual. Se você estiver usando TPUs reservadas, pule esta seção.Para criar nós de TPU no GKE, é preciso ter uma cota da API Compute Engine (compute.googleapis.com), e não uma cota da API Cloud TPU (tpu.googleapis.com). O nome da cota é diferente nos pods normais do Autopilot e nos pods do Spot.
Para verificar o limite e o uso atual da sua cota da API Compute Engine para TPUs, siga estas etapas:
Acesse a página Cotas no console do Google Cloud.
Na caixa Filtro
, faça o seguinte:Selecione a propriedade Serviço, insira API Compute Engine e pressione Enter.
Selecione a propriedade Tipo e escolha Cota.
Selecione a propriedade Nome e insira um nome de cota com base no tipo de TPU que você quer, da seguinte maneira:
- TPU v5p (pré-lançamento): chips TPU v5p
- TPU v5e (
tpu-v5-lite-podslice
): chips TPU v5 Lite PodSlice - TPU v5e (
tpu-v5-lite-device
): chips de dispositivo TPU v5 Lite - TPU v4 (
tpu-v4-podslice
): chips PodSlice de TPU v4
Para pods do Spot, selecione a cota "Preemptiva" correspondente.
Selecione a propriedade Dimensões (por exemplo, locais) e insira
region:
seguido do nome da região em que você planeja criar TPUs no GKE. Por exemplo, insiraregion:us-west4
se planeja criar nós da TPU na zonaus-west4-a
. A cota de TPU é regional, portanto, todas as zonas na mesma região consomem a mesma cota de TPU.
Se nenhuma cota corresponder ao filtro inserido, isso significa que o projeto não recebeu nenhuma das cotas especificadas para a região desejada. Você precisará solicitar um aumento de cota de TPU.
Preparar seu aplicativo de TPU
As cargas de trabalho de TPU têm os requisitos de preparação a seguir.
- Frameworks como JAX, PyTorch e TensorFlow acessam VMs de TPU usando a biblioteca compartilhada
libtpu
.libtpu
inclui o compilador XLA, o software do ambiente de execução da TPU e o driver da TPU. Cada versão do PyTorch e do JAX requer uma determinada versão delibtpu.so
. Para usar TPUs no GKE, use as seguintes versões:Tipo de TPU Versão do libtpu.so
TPU v5e
tpu-v5-lite-podslice
tpu-v5-lite-device
- Versão recomendada do jax[tpu]: v0.4.9 ou posterior
- Versão recomendada do torchxla[tpuvm]: v2.1.0 ou mais recente
TPU v5p
tpu-v5p-slice
- Versão recomendada do jax[tpu]: 0.4.19 ou mais recente.
- Versão recomendada do torchxla[tpuvm]: sugestão de uso de um build de versão noturno em 23 de outubro de 2023.
TPU v4
tpu-v4-podslice
- jax[tpu]: v0.4.4 ou posterior
- torchxla[tpuvm]: v2.0.0 ou posterior
- Defina as variáveis de ambiente a seguir para o contêiner que solicita os recursos da TPU:
TPU_WORKER_ID
: um número inteiro exclusivo para cada pod. Esse ID denota um código de worker exclusivo na fatia de TPU. Os valores aceitos para esse campo variam de zero ao número de pods menos um.TPU_WORKER_HOSTNAMES
: uma lista separada por vírgulas de nomes de host ou endereços IP de VM da TPU que precisam se comunicar entre si na fatia. É necessário que haja um nome do host ou endereço IP para cada VM da TPU na fatia. A lista de endereços IP ou nomes do host é ordenada e zero indexada peloTPU_WORKER_ID
.
O GKE injeta essas variáveis de ambiente automaticamente usando um webhook mutável quando um job é criado com
completionMode: Indexed
,subdomain
eparallelism > 1
e solicitando as propriedadesgoogle.com/tpu
. O GKE adiciona um serviço sem comando para que os registros DNS sejam adicionados aos pods de apoio do serviço.
Depois de concluir a preparação da carga de trabalho, é possível executar um job que usa TPUs.
Solicitar TPUs em uma carga de trabalho
Nesta seção, mostramos como criar um job que solicita TPUs no Autopilot. Em qualquer carga de trabalho que precise de TPUs, especifique o seguinte:
- Seletores de nós para a versão e topologia da TPU
- O número de chips de TPU para um contêiner na sua carga de trabalho
Para ver uma lista de versões de TPU, topologias e o número correspondente de chips e nós de TPU em uma fração, consulte Escolher uma configuração de TPU do Autopilot.
Considerações sobre solicitações de TPU em cargas de trabalho
Somente um contêiner em um pod pode usar TPUs. O número de ícones que um contêiner
solicita precisa ser igual ao número de ícones anexados a um nó na fração.
Por exemplo, se você solicitar a TPU v5e (tpu-v5-lite-podslice
) com uma topologia 2x4
, poderá solicitar qualquer um dos seguintes itens:
- Ícones
4
, que cria dois nós de vários hosts com quatro ícones cada. - Ícones
8
, que cria um nó de host único com 8 chips
Como prática recomendada para maximizar sua eficiência de custos, consuma sempre todos os ícones na fração que você solicitou. Se você solicitar uma fração de vários hosts de dois nós com quatro chips cada, implante uma carga de trabalho que seja executada em ambos os nós e consuma todos os oito chips na fração.
Criar uma carga de trabalho que solicite TPUs
As etapas a seguir criam um job que solicita TPUs. Se você tiver cargas de trabalho executadas em frações de TPU de vários hosts, também será necessário criar um serviço sem comando que selecione a carga de trabalho por nome. Esse serviço sem comando permite que pods em diferentes nós na fração de vários hosts se comuniquem entre si, atualizando a configuração DNS do Kubernetes para apontar para os pods na carga de trabalho.
Salve o seguinte manifesto como
tpu-autopilot.yaml
:apiVersion: v1 kind: Service metadata: name: headless-svc spec: clusterIP: None selector: job-name: tpu-job --- apiVersion: batch/v1 kind: Job metadata: name: tpu-job spec: backoffLimit: 0 completions: 4 parallelism: 4 completionMode: Indexed template: spec: subdomain: headless-svc restartPolicy: Never nodeSelector: cloud.google.com/gke-tpu-accelerator: TPU_TYPE cloud.google.com/gke-tpu-topology: TOPOLOGY containers: - name: tpu-job image: python:3.10 ports: - containerPort: 8471 # Default port using which TPU VMs communicate - containerPort: 8431 # Port to export TPU runtime metrics, if supported. command: - bash - -c - | pip install 'jax[tpu]' -f https://storage.googleapis.com/jax-releases/libtpu_releases.html python -c 'import jax; print("TPU cores:", jax.device_count())' resources: requests: cpu: 10 memory: 500Gi google.com/tpu: NUMBER_OF_CHIPS limits: cpu: 10 memory: 500Gi google.com/tpu: NUMBER_OF_CHIPS
Substitua:
TPU_TYPE
: o tipo de TPU a ser usado, comotpu-v4-podslice
. Precisa ser um valor compatível com o GKE.TOPOLOGY
: a disposição dos chips de TPU na fração, como2x2x4
. Precisa ser uma topologia compatível com o tipo de TPU selecionado.NUMBER_OF_CHIPS
: o número de chips de TPU para o contêiner usar. Precisa ser o mesmo valor paralimits
erequests
.
Implante o job:
kubectl create -f tpu-autopilot.yaml
Quando você cria esse job, o GKE faz automaticamente o seguinte:
- Provisiona nós para executar os pods. Dependendo do tipo de TPU, topologia e solicitações de recursos especificados, esses nós são frações de host único ou de vários hosts.
- Adiciona taints aos pods e tolerâncias aos nós para evitar que outras cargas de trabalho sejam executadas nos mesmos nós que as cargas de trabalho da TPU.
Exemplo: exibir o total de ícones de TPU em uma fração de vários hosts
A carga de trabalho a seguir retorna o número de chips de TPU em todos os nós em uma fração de TPU de vários hosts. Para criar uma fração de vários hosts, a carga de trabalho tem os seguintes parâmetros:
- Versão do TPU: TPU v4
- Topologia: 2x2x4
Essa seleção de versão e topologia resultam em uma fração de vários hosts.
- Salve o seguinte manifesto como
available-chips-multihost.yaml
:apiVersion: v1 kind: Service metadata: name: headless-svc spec: clusterIP: None selector: job-name: tpu-available-chips --- apiVersion: batch/v1 kind: Job metadata: name: tpu-available-chips spec: backoffLimit: 0 completions: 4 parallelism: 4 completionMode: Indexed template: spec: subdomain: headless-svc restartPolicy: Never nodeSelector: cloud.google.com/gke-tpu-accelerator: tpu-v4-podslice cloud.google.com/gke-tpu-topology: 2x2x4 containers: - name: tpu-job image: python:3.10 ports: - containerPort: 8471 # Default port using which TPU VMs communicate - containerPort: 8431 # Port to export TPU runtime metrics, if supported. command: - bash - -c - | pip install 'jax[tpu]' -f https://storage.googleapis.com/jax-releases/libtpu_releases.html python -c 'import jax; print("TPU cores:", jax.device_count())' resources: requests: cpu: 10 memory: 500Gi google.com/tpu: 4 limits: cpu: 10 memory: 500Gi google.com/tpu: 4
- Implante o manifesto:
kubectl create -f available-chips-multihost.yaml
O GKE executa uma fatia da TPU v4 com quatro VMs da TPU (fatia da TPU de vários hosts). A fração tem 16 chips interconectados.
- Verifique se o job criou quatro pods:
kubectl get pods
O resultado será assim:
NAME READY STATUS RESTARTS AGE tpu-job-podslice-0-5cd8r 0/1 Completed 0 97s tpu-job-podslice-1-lqqxt 0/1 Completed 0 97s tpu-job-podslice-2-f6kwh 0/1 Completed 0 97s tpu-job-podslice-3-m8b5c 0/1 Completed 0 97s
- Receba os registros de um dos pods:
kubectl logs POD_NAME
Substitua
POD_NAME
pelo nome de um dos pods criados. Por exemplo,tpu-job-podslice-0-5cd8r
.O resultado será assim:
TPU cores: 16
Exemplo: exibir os chips da TPU em um único nó
A carga de trabalho a seguir é um pod estático que exibe o número de chips de TPU anexados a um nó específico. Para criar um nó de host único, a carga de trabalho tem os seguintes parâmetros:
- Versão do TPU: TPU v5e
- Topologia: 2x4
Essa seleção de versão e topologia resultam em uma fração de host único.
- Salve o seguinte manifesto como
available-chips-singlehost.yaml
:apiVersion: v1 kind: Pod metadata: name: tpu-job-jax-v5 spec: restartPolicy: Never nodeSelector: cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice cloud.google.com/gke-tpu-topology: 2x4 containers: - name: tpu-job image: python:3.10 ports: - containerPort: 8431 # Port to export TPU runtime metrics, if supported. command: - bash - -c - | pip install 'jax[tpu]' -f https://storage.googleapis.com/jax-releases/libtpu_releases.html python -c 'import jax; print("Total TPU chips:", jax.device_count())' resources: requests: google.com/tpu: 8 limits: google.com/tpu: 8
- Implante o manifesto:
kubectl create -f available-chips-singlehost.yaml
O GKE provisiona nós com oito frações de TPU de host único que usam a TPU v5e. Cada VM de TPU tem oito chips (fração de TPU de host único).
- Receba os registros do pod:
kubectl logs tpu-job-jax-v5
O resultado será assim:
Total TPU chips: 8
Observabilidade e métricas
Painel
Na página Clusters do Kubernetes no Console do Google Cloud, a guia Observabilidade exibe as métricas de observabilidade da TPU. Para mais informações, consulte Métricas de observabilidade do GKE.
O painel da TPU será preenchido apenas se você tiver métricas do sistema ativadas no cluster do GKE.
Métricas do ambiente de execução
No GKE 1.27.4-gke.900 ou mais recente, as cargas de trabalho da TPU
que usam o JAX
0.4.14
ou mais recente e especificam containerPort: 8431
exportam as métricas de utilização da TPU como
métricas do sistema do GKE.
As métricas a seguir estão disponíveis no Cloud Monitoring para monitorar o desempenho do ambiente de execução da carga de trabalho da TPU
- Ciclo de trabalho: porcentagem de tempo durante o período de amostragem anterior (60 segundos) em que os TensorCores estavam ativamente em processamento em um chip de TPU. Uma porcentagem maior significa melhor uso da TPU.
- Uso da memória: quantidade de memória do acelerador alocada em bytes. Amostras coletadas a cada 60 segundos.
- Capacidade de memória: memória total do acelerador em bytes. Amostras coletadas a cada 60 segundos.
Essas métricas estão localizadas no esquema de nó (k8s_node
) e contêiner de contêiner (k8s_container
) do Kubernetes.
Contêiner do Kubernetes:
kubernetes.io/container/accelerator/duty_cycle
kubernetes.io/container/accelerator/memory_used
kubernetes.io/container/accelerator/memory_total
Nó do Kubernetes:
kubernetes.io/node/accelerator/duty_cycle
kubernetes.io/node/accelerator/memory_used
kubernetes.io/node/accelerator/memory_total
Métricas de hospedagem
No GKE 1.28.1-gke.1066000 ou mais recente, a VM da TPU exporta as métricas de utilização da TPU como métricas do sistema do GKE. As seguintes métricas estão disponíveis no Cloud Monitoring para monitorar o desempenho do host da TPU:
- Uso do TensorCore: porcentagem atual do TensorCore que está sendo utilizada. O valor do TensorCore é igual à soma das unidades de multiplicação de matriz (MXUs) mais a unidade vetorial. O valor de utilização do TensorCore é a divisão das operações do TensorCore que foram realizadas no período de amostragem anterior (60 segundos) pelo número de operações do TensorCore com suporte no mesmo período. Um valor maior significa melhor utilização.
- Uso de largura de banda da memória: porcentagem atual da largura de banda da memória do acelerador que está sendo usada. Calculada pela divisão da largura de banda da memória usada durante um período de amostragem (60 segundos) pela largura de banda máxima aceita no mesmo período.
Essas métricas estão localizadas no esquema de nó (k8s_node
) e contêiner de contêiner (k8s_container
) do Kubernetes.
Contêiner do Kubernetes:
kubernetes.io/container/accelerator/tensorcore_utilization
kubernetes.io/container/accelerator/memory_bandwidth_utilization
Nó do Kubernetes:
kubernetes.io/container/node/tensorcore_utilization
kubernetes.io/container/node/memory_bandwidth_utilization
Para mais informações, consulte Métricas do Kubernetes e Métricas do sistema do GKE.
Geração de registros
Os registros emitidos por contêineres em execução em nós do GKE, incluindo VMs de TPU, sãocoletou pelo agente do Logging do GKE, enviado para o Logging e sãovisível no Logging do Google Analytics.
Recomendações para cargas de trabalho de TPU no Autopilot
As recomendações a seguir podem melhorar a eficiência das cargas de trabalho da TPU:
- Use pods de tempo de execução estendido por um período de carência de até sete dias antes que o GKE encerre seus pods para redução de escalonamento ou upgrade de nós. É possível usar janelas de manutenção e exclusões com pods de tempo de execução estendidos para atrasar ainda mais os upgrades automáticos de nós.
- Use reservas de capacidade para garantir que as cargas de trabalho recebam as TPUs solicitadas sem serem colocadas em uma fila para disponibilidade.