Implante cargas de trabalho de TPU no Autopilot do GKE

Nesta página, descrevemos como acelerar as cargas de trabalho de machine learning (ML) usando aceleradores do Cloud TPU (TPUs) em clusters do Autopilot do Google Kubernetes Engine (GKE). Essas orientações podem ajudar você a selecionar as bibliotecas corretas para seus frameworks de aplicativos de ML, configurar as cargas de trabalho de TPU para serem executadas de maneira ideal no GKE e monitorar as cargas de trabalho após a implantação.

Esta página é destinada a administradores e operadores de plataforma, especialistas em dados e IA e desenvolvedores de aplicativos que querem preparar e executar cargas de trabalho de ML em TPUs. Para saber mais sobre as funções, responsabilidades e exemplos de tarefas comuns que mencionamos no conteúdo do Google Cloud , consulte Funções e tarefas comuns do usuário do GKE.

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

Como as TPUs funcionam no Autopilot

Para usar TPUs em cargas de trabalho do Autopilot, especifique o seguinte no manifesto da carga de trabalho:

  • A versão da TPU no campo spec.nodeSelector.
  • A topologia da TPU no campo spec.nodeSelector. A topologia precisa ser compatível com a versão da TPU especificada.
  • O número de chips de TPU nos campos spec.containers.resources.requests e spec.containers.resources.limits.

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:

  1. Pods do Spot
  2. Reservas de capacidade específicas
  3. Pods de tempo de execução estendido
  4. Flex-start

Planejar a configuração da TPU

Antes de usar este guia para implantar cargas de trabalho da TPU, planeje a configuração com base no modelo e na quantidade de memória necessária. Para mais detalhes, consulte Planejar a configuração da TPU.

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 CLI do Google Cloud para essa tarefa, instale e inicialize a gcloud CLI. Se você instalou a CLI gcloud anteriormente, instale a versão mais recente executando o comando gcloud components update. Talvez as versões anteriores da CLI gcloud não sejam compatíveis com a execução dos comandos neste documento.
  • Verifique se você tem um cluster do Autopilot em execução na versão 1.32.3-gke.1927000 ou mais recente do GKE. Para instruções, consulte Criar um cluster do Autopilot.
  • Para usar TPUs reservadas, verifique se você tem uma reserva de capacidade específica. Para instruções, consulte Consumir uma reserva.

Verifique se há cota para TPUs e outros recursos do GKE

As seções a seguir ajudam a garantir que você tenha cota suficiente ao usar TPUs no GKE.

Para criar nós de fração 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 fração 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:

  1. Acesse a página Cotas no console do Google Cloud :

    Acessar "Cotas"

  2. Na caixa Filtro , faça o seguinte:

    1. Use a tabela a seguir para selecionar e copiar a propriedade da cota com base na versão da TPU e no valor no seletor de nós cloud.google.com/gke-tpu-accelerator. Por exemplo, se você planeja criar nós da TPU v5e sob demanda com o valor tpu-v5-lite-podslice no seletor de nós cloud.google.com/gke-tpu-accelerator, insira Name: TPU v5 Lite PodSlice chips.

      Versão da TPU, cloud.google.com/gke-tpu-accelerator Propriedade e nome da cota para instâncias sob demanda Propriedade e nome da cota para instâncias Spot2
      TPU v3,
      tpu-v3-device 
      Dimensions (e.g. location):
      tpu_family:CT3
      		 Não relevante
      TPU v3,
      tpu-v3-slice 
      Dimensions (e.g. location):
      tpu_family:CT3P
      		 Não relevante
      TPU v4,
      tpu-v4-podslice 
      Name:
      TPU v4 PodSlice chips
      		 
      Name:
      Preemptible TPU v4 PodSlice chips
      TPU v5e,
      tpu-v5-lite-podslice 
      Name:
      TPU v5 Lite PodSlice chips
      		 
      Name:
      Preemptible TPU v5 Lite Podslice
      chips
      TPU v5p,
      tpu-v5p-slice 
      Name:
      TPU v5p chips
      		 
      Name:
      Preemptible TPU v5p chips
      TPU Trillium,
      tpu-v6e-slice 
      Dimensions (e.g. location):
      tpu_family:CT6E
      		 
      Name:
      Preemptible TPU slices v6e

    2. 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, insira region:us-west4 se planeja criar nós de fração da TPU na zona us-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 ajuste de cota de TPU.

Quando uma reserva de TPU é criada, os valores de limite e de uso atual da cota correspondente aumentam pelo número de ícones na reserva da TPU. Por exemplo, quando uma reserva é criada para 16 chips de TPU v5e com valor tpu-v5-lite-podslice no seletor de nós cloud.google.com/gke-tpu-accelerator, o Limite e o Uso atual para a cota de TPU v5 Lite PodSlice chips na região relevante aumentam em 16.

Cotas para outros recursos do GKE

Talvez seja necessário aumentar as seguintes cotas relacionadas ao GKE nas regiões em que o GKE cria seus recursos.

  • Cota de SSD de disco permanente (GB): o disco de inicialização de cada nó do Kubernetes requer 100 GB por padrão. Portanto, essa cota precisa ser definida pelo menos tão alta quanto o produto do número máximo de nós do GKE que você espera criar e 100 GB (nós * 100 GB).
  • Cota de endereços IP em uso: cada nó do Kubernetes consome um endereço IP. Portanto, essa cota precisa ser definida pelo menos tão alta quanto o número máximo de nós do GKE que você prevê que serão criados.
  • Verifique se max-pods-per-node está alinhado com o intervalo da sub-rede: cada nó do Kubernetes usa intervalos de IP secundários para pods. Por exemplo, max-pods-per-node de 32 requer 64 endereços IP, o que se traduz em uma sub-rede /26 por nó. Esse intervalo não pode ser compartilhado com nenhum outro cluster. Para evitar o esgotamento do intervalo de endereços IP, use a flag --max-pods-per-node para limitar o número de pods que podem ser programados em um nó. A cota para max-pods-per-node precisa ser definida pelo menos tão alta quanto o número máximo de nós do GKE que você prevê que serão criados.

Para solicitar um aumento de cota, consulte Solicitar um ajuste de cota.

Opções para provisionar TPUs no GKE

O Autopilot do GKE permite usar TPUs diretamente em cargas de trabalho individuais usando nodeSelectors do Kubernetes.

Como alternativa, é possível solicitar TPUs usando classes de computação personalizadas. Com as classes de computação personalizadas, os administradores de plataforma podem definir uma hierarquia de configurações de nós para o GKE priorizar durante as decisões de escalonamento de nós, para que as cargas de trabalho sejam executadas no hardware selecionado.

Para instruções, consulte a seção Fazer o provisionamento centralizado de TPUs com classes de computação personalizadas.

Preparar seu aplicativo de TPU

As cargas de trabalho de TPU têm os requisitos de preparação a seguir.

  1. 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 de libtpu.so. Para evitar conflitos de versão de pacote, recomendamos usar uma imagem de IA do JAX. Para usar TPUs no GKE, use as seguintes versões:
    Tipo de TPU Versão do libtpu.so
    TPU Trillium (v6e)
    tpu-v6e-slice
    TPU v5e
    tpu-v5-lite-podslice
    TPU v5p
    tpu-v5p-slice
    TPU v4
    tpu-v4-podslice
    TPU v3
    tpu-v3-slice
    tpu-v3-device
  2. 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 pelo TPU_WORKER_ID.

      • O GKE injeta essas variáveis de ambiente automaticamente usando um webhook mutável quando um job é criado com completionMode: Indexed, subdomain e parallelism > 1 e solicitando as propriedades google.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 a versão da TPU.

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 chips de TPU que um contêiner solicita precisa ser igual ao número de chips 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:

  • Chips 4, que cria dois nós de vários hosts com quatro chips de TPU cada.
  • Chips 8, que cria um nó de host único com oito chips de TPU

Como prática recomendada para maximizar o custo-benefício, consuma sempre toda a TPU na fração que você solicitou. Se você solicitar uma fração de vários hosts de dois nós com quatro chips de TPU cada, implante uma carga de trabalho que seja executada em ambos os nós e consuma todos os oito chips de TPU 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.

  1. 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:
      # Optional: Run in GKE Sandbox
      # runtimeClassName: gvisor
      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: us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest
        ports:
        - containerPort: 8471 # Default port using which TPU VMs communicate
        - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
        command:
        - bash
        - -c
        - |
          python -c 'import jax; print("TPU cores:", jax.device_count())'
        resources:
          requests:
            cpu: 10
            memory: MEMORY_SIZE
            google.com/tpu: NUMBER_OF_CHIPS
          limits:
            cpu: 10
            memory: MEMORY_SIZE
            google.com/tpu: NUMBER_OF_CHIPS

    Substitua:

    • TPU_TYPE: o tipo de TPU a ser usado, como tpu-v4-podslice. Precisa ser um valor compatível com o GKE.
    • TOPOLOGY: a disposição dos chips de TPU na fração, como 2x2x4. 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 para limits e requests.
    • MEMORY_SIZE: a quantidade máxima de memória usada pela TPU. Os limites de memória dependem da versão e da topologia da TPU que você usa. Para saber mais, consulte Mínimos e máximos para aceleradores.

    Também é possível mudar os seguintes campos:

    • image: a imagem de IA do JAX a ser usada. No manifesto de exemplo, esse campo é definido como a imagem mais recente de IA do JAX. Para definir uma versão diferente, consulte a lista de imagens atuais da IA JAX.
    • runtimeClassname: gvisor: a configuração que permite executar esse pod no GKE Sandbox. Para usar, remova a marca de comentário dessa linha. O GKE Sandbox é compatível com TPUs versão v4 e mais recentes. Para saber mais, consulte GKE Sandbox.

  2. Implante o job:

    kubectl create -f tpu-autopilot.yaml

    Quando você cria esse job, o GKE faz automaticamente o seguinte:

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

  3. Ao concluir esta seção, exclua a carga de trabalho criada para evitar o faturamento contínuo:

    kubectl delete -f tpu-autopilot.yaml

Criar uma carga de trabalho que solicite TPUs e agendamento de coleta

No TPU Trillium, é possível usar o agendamento de coleta para agrupar nós de fração de TPU. Agrupar esses nós de fração de TPU facilita o ajuste do número de réplicas para atender à demanda da carga de trabalho.O Google Cloud controla as atualizações de software para garantir que sempre haja frações suficientes disponíveis na coleção para atender ao tráfego.

A TPU Trillium oferece suporte ao agendamento de coleta para pools de nós de host único e vários hosts que executam cargas de trabalho de inferência. A seguir, descrevemos como o comportamento do agendamento de coleta depende do tipo de fração de TPU usada:

  • Fração de TPU de vários hosts:o GKE agrupa frações de TPU de vários hosts para formar uma coleção. Cada pool de nós do GKE é uma réplica nessa coleção. Para definir uma coleção, crie uma fração de TPU de vários hosts e atribua um nome exclusivo a ela. Para adicionar mais frações de TPU à coleção, crie outro pool de nós de fração de TPU de vários hosts com o mesmo nome de coleção e tipo de carga de trabalho.
  • Fração de TPU de host único:o GKE considera todo o pool de nós de fração de TPU de host único como uma coleção. Para adicionar mais frações de TPU à coleção, redimensione o pool de nós de fração de TPU de host único.

Para saber mais sobre a limitação do agendamento de coleta, consulte Como funciona o agendamento de coleta

Usar uma fração de TPU de vários hosts

O agendamento de coleta em nós de fração de TPU de vários hosts está disponível para clusters do Autopilot na versão 1.31.2-gke.1537000 e mais recente. Os nós de fração de TPU de vários hosts com uma topologia 2x4 são compatíveis apenas com a versão 1.31.2-gke.1115000 ou mais recente. Para criar nós de fração de TPU de vários hosts e agrupá-los como uma coleção, adicione os seguintes rótulos do Kubernetes à especificação da carga de trabalho:

  • cloud.google.com/gke-nodepool-group-name: cada coleção precisa ter um nome exclusivo no nível do cluster. O valor no rótulo cloud.google.com/gke-nodepool-group-name precisa obedecer aos requisitos para rótulos de cluster.

  • cloud.google.com/gke-workload-type: HIGH_AVAILABILITY

    Por exemplo, o bloco de código a seguir define uma coleção com uma fração de TPU de vários hosts:

      nodeSelector:
    cloud.google.com/gke-nodepool-group-name: ${COLLECTION_NAME}
    cloud.google.com/gke-workload-type: HIGH_AVAILABILITY
    cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
    cloud.google.com/gke-tpu-topology: 4x4
...

Usar uma fração de TPU de host único

O agendamento de coleta em nós de fração de TPU de host único está disponível para clusters do Autopilot na versão 1.31.2-gke.1088000 e mais recentes. Para criar nós de fração de TPU de host único e agrupá-los como uma coleção, adicione o rótulo cloud.google.com/gke-workload-type:HIGH_AVAILABILITY na especificação da carga de trabalho.

Por exemplo, o bloco de código a seguir define uma coleção com uma fração de TPU de host único:

  nodeSelector:
    cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
    cloud.google.com/gke-tpu-topology: 2x2
    cloud.google.com/gke-workload-type: HIGH_AVAILABILITY
  ...

Usar classes de computação personalizadas para implantar uma coleção

Para mais informações sobre como implantar uma carga de trabalho que solicita uma carga de trabalho de TPU e o agendamento de coleta usando classes de computação personalizadas, consulte Coleta multihost de TPU e Definir o tipo de carga de trabalho para o SLO de TPU.

Provisionar TPUs de maneira centralizada com classes de computação personalizadas

Para provisionar TPUs com uma classe de computação personalizada que segue as regras de TPU e implantar a carga de trabalho, siga estas etapas:

  1. Salve o seguinte manifesto como tpu-compute-class.yaml:

    apiVersion: cloud.google.com/v1
kind: ComputeClass
metadata:
  name: tpu-class
spec:
  priorities:
  - tpu:
      type: tpu-v5-lite-podslice
      count: 4
      topology: 2x4
  - spot: true
    tpu:
      type: tpu-v5-lite-podslice
      count: 4
      topology: 2x4
  - flexStart:
      enabled: true
    tpu:
      type: tpu-v6e-slice
      count: 4
      topology: 2x4
  nodePoolAutoCreation:
    enabled: true

  2. Implante a classe de computação:

    kubectl apply -f tpu-compute-class.yaml

    Para mais informações sobre classes de computação personalizadas e TPUs, consulte Configuração de TPU.

  3. Salve o seguinte manifesto como tpu-job.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/compute-class: tpu-class
      containers:
      - name: tpu-job
        image: us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest
        ports:
        - containerPort: 8471 # Default port using which TPU VMs communicate
        - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
        command:
        - bash
        - -c
        - |
          python -c 'import jax; print("TPU cores:", jax.device_count())'
        resources:
          requests:
            cpu: 10
            memory: MEMORY_SIZE
            google.com/tpu: NUMBER_OF_CHIPS
          limits:
            cpu: 10
            memory: MEMORY_SIZE
            google.com/tpu: NUMBER_OF_CHIPS

    Substitua:

    • NUMBER_OF_CHIPS: o número de chips de TPU para o contêiner usar. Precisa ser o mesmo valor para limits e requests, igual ao valor no campo tpu.count na classe de computação personalizada selecionada.
    • MEMORY_SIZE: a quantidade máxima de memória que a TPU usa. Os limites de memória dependem da versão e da topologia da TPU que você usa. Para saber mais, consulte Mínimos e máximos para aceleradores.
    • NUMBER_OF_CHIPS: o número de chips de TPU para o contêiner usar. Precisa ser o mesmo valor para limits e requests.

  4. Implante o job:

    kubectl create -f tpu-job.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. Dependendo da disponibilidade de recursos de TPU na prioridade mais alta, o GKE pode voltar a prioridades mais baixas para maximizar a capacidade de obtenção.
    • 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.

    Para saber mais, consulte Sobre as classes de computação personalizadas.

  5. Ao concluir esta seção, exclua os recursos criados para evitar o faturamento contínuo:

    kubectl delete -f tpu-job.yaml

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.

  1. 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 # Node selector to target TPU v4 slice nodes.
        cloud.google.com/gke-tpu-topology: 2x2x4 # Specifies the physical topology for the TPU slice.
      containers:
      - name: tpu-job
        image: us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest
        ports:
        - containerPort: 8471 # Default port using which TPU VMs communicate
        - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
        command:
        - bash
        - -c
        - |
          python -c 'import jax; print("TPU cores:", jax.device_count())' # Python command to count available TPU chips.
        resources:
          requests:
            cpu: 10
            memory: 407Gi
            google.com/tpu: 4 # Request 4 TPU chips for this workload.
          limits:
            cpu: 10
            memory: 407Gi
            google.com/tpu: 4 # Limit to 4 TPU chips for this workload.
  2. Implante o manifesto: 
    kubectl create -f available-chips-multihost.yaml

    O GKE executa uma fatia da TPU v4 com quatro VMs (fatia da TPU de vários hosts). A fração tem 16 chips de TPU interconectados.

  3. 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
  4. Consiga 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
  5. Opcional: remova a carga de trabalho: 
    kubectl delete -f available-chips-multihost.yaml

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.

  1. 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 # Node selector to target TPU v5e slice nodes.
    cloud.google.com/gke-tpu-topology: 2x4 # Specify the physical topology for the TPU slice.
  containers:
  - name: tpu-job
    image: us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest
    ports:
    - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
    command:
    - bash
    - -c
    - |
      python -c 'import jax; print("Total TPU chips:", jax.device_count())'
    resources:
      requests:
        google.com/tpu: 8 # Request 8 TPU chips for this container.
      limits:
        google.com/tpu: 8 # Limit to 8 TPU chips for this container.
  2. 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 nó de TPU tem oito chips de TPU (fração de TPU de host único).

  3. Consiga os registros do pod: 
    kubectl logs tpu-job-jax-v5

    O resultado será assim: 

    Total TPU chips: 8
  4. Opcional: remova a carga de trabalho: 
      kubectl delete -f available-chips-singlehost.yaml

Observar e monitorar TPUs

Painel

A capacidade de observação do pool de nós no consoleGoogle Cloud está em disponibilidade geral. 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

Esse painel oferece insights abrangentes sobre a integridade dos seus pools de nós de TPU de vários hosts. Para mais informações, consulte Monitorar métricas de integridade para nós e pools de nós de TPU.

Na página Clusters do Kubernetes no consoleGoogle Cloud , a guia Observabilidade também mostra métricas de observabilidade da TPU, como o uso da TPU, no cabeçalho Aceleradores > TPU. Para mais informações, consulte Conferir métricas de observabilidade.

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.
  • Memória usada: quantidade de memória do acelerador alocada em bytes. Amostras coletadas a cada 60 segundos.
  • Memória total: 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

Monitorar métricas de integridade para nós e pools de nós de TPU

Quando um job de treinamento tem um erro ou termina com falha, é possível verificar as métricas relacionadas à infraestrutura subjacente para descobrir se a interrupção foi causada por um problema com o nó ou o pool de nós subjacente.

Status do nó

No GKE 1.32.1-gke.1357001 ou mais recente, a seguinte métrica do sistema do GKE expõe a condição de um nó do GKE:

  • kubernetes.io/node/status_condition

O campo condition informa condições no nó, como Ready, DiskPressure e MemoryPressure. O campo status mostra o status informado da condição, que pode ser True, False ou Unknown. Essa é uma métrica com o tipo de recurso monitorado k8s_node.

Essa consulta PromQL mostra se um nó específico está Ready:

kubernetes_io:node_status_condition{
    monitored_resource="k8s_node",
    cluster_name="CLUSTER_NAME",
    node_name="NODE_NAME",
    condition="Ready",
    status="True"}

Para ajudar a resolver problemas em um cluster, confira os nós que apresentaram outras condições:

kubernetes_io:node_status_condition{
    monitored_resource="k8s_node",
    cluster_name="CLUSTER_NAME",
    condition!="Ready",
    status="True"}

Talvez você queira analisar especificamente os nós que não são Ready:

kubernetes_io:node_status_condition{
    monitored_resource="k8s_node",
    cluster_name="CLUSTER_NAME",
    condition="Ready",
    status="False"}

Se não houver dados, os nós estarão prontos. A condição de status é amostrada a cada 60 segundos.

Use a consulta a seguir para entender o status do nó em toda a frota:

avg by (condition,status)(
  avg_over_time(
    kubernetes_io:node_status_condition{monitored_resource="k8s_node"}[${__interval}]))

Status do pool de nós

A seguinte métrica de sistema do GKE para o recurso monitorado k8s_node_pool expõe o status de um pool de nós do GKE:

  • kubernetes.io/node_pool/status

Essa métrica é informada apenas para pools de nós de TPU de vários hosts.

O campo status informa o status do pool de nós, como Provisioning, Running, Error, Reconciling ou Stopping. As atualizações de status acontecem depois que as operações da API do GKE são concluídas.

Para verificar se um pool de nós específico tem o status Running, use esta consulta em PromQL:

kubernetes_io:node_pool_status{
    monitored_resource="k8s_node_pool",
    cluster_name="CLUSTER_NAME",
    node_pool_name="NODE_POOL_NAME",
    status="Running"}

Para monitorar o número de pools de nós no seu projeto agrupados por status, use a seguinte consulta PromQL:

count by (status)(
  count_over_time(
    kubernetes_io:node_pool_status{monitored_resource="k8s_node_pool"}[${__interval}]))

Disponibilidade do pool de nós

A seguinte métrica do sistema do GKE mostra se um pool de nós da TPU de vários hosts está disponível:

  • kubernetes.io/node_pool/multi_host/available

A métrica tem um valor de True se todos os nós do pool de nós estiverem disponíveis e False caso contrário. A métrica é amostrada a cada 60 segundos.

Para verificar a disponibilidade de pools de nós de TPU de vários hosts no seu projeto, use a seguinte consulta PromQL:

avg by (node_pool_name)(
  avg_over_time(
    kubernetes_io:node_pool_multi_host_available{
      monitored_resource="k8s_node_pool",
      cluster_name="CLUSTER_NAME"}[${__interval}]))

Contagem de interrupções de nós

A métrica do sistema do GKE a seguir informa a contagem de interrupções de um nó do GKE desde a última amostra. A métrica é amostrada a cada 60 segundos:

  • kubernetes.io/node/interruption_count

Os campos interruption_type (como TerminationEvent, MaintenanceEvent ou PreemptionEvent) e interruption_reason (como HostError, Eviction ou AutoRepair) podem ajudar a explicar por que um nó foi interrompido.

Para conferir um detalhamento das interrupções e das causas nos nós da TPU nos clusters do seu projeto, use a seguinte consulta do PromQL:

  sum by (interruption_type,interruption_reason)(
    sum_over_time(
      kubernetes_io:node_interruption_count{monitored_resource="k8s_node"}[${__interval}]))

Para ver apenas os eventos de manutenção do host, atualize a consulta para filtrar o valor HW/SW Maintenance de interruption_reason. Use a seguinte consulta em PromQL:

  sum by (interruption_type,interruption_reason)(
    sum_over_time(
      kubernetes_io:node_interruption_count{monitored_resource="k8s_node", interruption_reason="HW/SW Maintenance"}[${__interval}]))

Para ver a contagem de interrupções agregada por pool de nós, use a seguinte consulta do PromQL:

  sum by (node_pool_name,interruption_type,interruption_reason)(
    sum_over_time(
      kubernetes_io:node_pool_interruption_count{monitored_resource="k8s_node_pool", interruption_reason="HW/SW Maintenance", node_pool_name=NODE_POOL_NAME }[${__interval}]))

Tempos de recuperação (TTR) do pool de nós

A métrica do sistema do GKE a seguir informa a distribuição das durações do período de recuperação para pools de nós da TPU de vários hosts do GKE:

  • kubernetes.io/node_pool/accelerator/times_to_recover

Cada amostra registrada nessa métrica indica um único evento de recuperação do pool de nós de um período de inatividade.

Essa métrica é útil para rastrear o tempo de recuperação e o tempo entre interrupções do pool de nós de TPU de vários hosts.

Use a seguinte consulta do PromQL para calcular o tempo médio de recuperação (MTTR) nos últimos sete dias no seu cluster:

sum(sum_over_time(
  kubernetes_io:node_pool_accelerator_times_to_recover_sum{
    monitored_resource="k8s_node_pool", cluster_name="CLUSTER_NAME"}[7d]))
/
sum(sum_over_time(
  kubernetes_io:node_pool_accelerator_times_to_recover_count{
    monitored_resource="k8s_node_pool",cluster_name="CLUSTER_NAME"}[7d]))

Tempos entre interrupções (TBI) do pool de nós

Os tempos entre interrupções do pool de nós medem quanto tempo sua infraestrutura é executada antes de sofrer uma interrupção. Ele é calculado como a média em uma janela de tempo, em que o numerador mede o tempo total em que sua infraestrutura ficou ativa e o denominador mede o total de interrupções na sua infraestrutura.

O exemplo de PromQL a seguir mostra o tempo médio entre interrupções (MTBI, na sigla em inglês) de sete dias para o cluster especificado:

sum(count_over_time(
  kubernetes_io:node_memory_total_bytes{
    monitored_resource="k8s_node", node_name=~"gke-tpu.*|gk3-tpu.*", cluster_name="CLUSTER_NAME"}[7d]))
/
sum(sum_over_time(
  kubernetes_io:node_interruption_count{
    monitored_resource="k8s_node", node_name=~"gke-tpu.*|gk3-tpu.*", cluster_name="CLUSTER_NAME"}[7d]))

Métricas de hospedagem

No GKE 1.28.1-gke.1066000 ou mais recente, as VMs em uma fração da TPU exportam 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/node/accelerator/tensorcore_utilization
  • kubernetes.io/node/accelerator/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.

Para saber como configurar a Cloud TPU no GKE, consulte os seguintes recursos Google Cloud :