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. Início flexível

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 gcloud CLI anteriormente, instale a versão mais recente executando o comando gcloud components update. Talvez as versões anteriores da gcloud CLI 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.

Garantir 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 do 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
      Ironwood (TPU7x) (pré-lançamento),
      tpu7x 
      Dimensions (e.g. location):
      tpu_family:tpu7x
      		 
      Name:
      Preemptible TPU slices tpu7x

    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 até 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, um 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.

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: tpu7x
    Tipo de TPU Versão do libtpu.so
    Ironwood (TPU7x) (pré-lançamento)
    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. No manifesto da carga de trabalho, adicione seletores de nós do Kubernetes para garantir que o GKE programe a carga de trabalho da TPU no tipo de máquina e na topologia de TPU definidos:

      nodeSelector:
    cloud.google.com/gke-tpu-accelerator: TPU_ACCELERATOR
    cloud.google.com/gke-tpu-topology: TPU_TOPOLOGY
    cloud.google.com/placement-policy-name: WORKLOAD_POLICY # Required only for Ironwood (TPU7x)

    Substitua:

    • TPU_ACCELERATOR: o nome do acelerador de TPU. Por exemplo, use tpu7x-standard-4t.
    • TPU_TOPOLOGY: a topologia física da fração da TPU. O formato da topologia depende da versão da TPU. Por exemplo, use 2x2x2. Para saber mais, consulte Planejar TPUs no GKE.
    • WORKLOAD_POLICY: o nome da política de carga de trabalho que você quer usar para colocar os pods de TPU. Esse seletor de nós é necessário apenas para o Ironwood (TPU7x).

Depois de concluir a preparação da carga de trabalho, é possível executar um job que usa TPUs.

Opções para provisionar TPUs no GKE

Para provisionar TPUs no GKE, você tem as seguintes opções de configuração:
  • Solicitação de carga de trabalho: especifique a versão e a topologia da TPU no campo spec.nodeSelector e o número de chips de TPU na seção spec.containers.resources. Quando você implanta a carga de trabalho, o GKE provisiona automaticamente nós com a configuração correta de TPU e coloca cada carga de trabalho no próprio nó dedicado para garantir acesso total aos recursos do nó. Para instruções, consulte Solicitar TPUs em uma carga de trabalho.

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

    As seções a seguir mostram como criar uma ComputeClass personalizada e, em seguida, criar um job que consome as TPUs definidas na ComputeClass.

    Criar uma ComputeClass personalizada

    As etapas para criar uma ComputeClass personalizada que segue as regras de TPU variam dependendo se você usa o Ironwood (TPU7x) ou uma versão anterior da TPU.

    Ironwood (TPU7x)

    1. Crie uma política de carga de trabalho. Esta etapa é necessária apenas se você estiver criando um pool de nós de vários hosts, que depende da topologia escolhida. Se você usa um pool de nós de host único, pule esta etapa.

      gcloud compute resource-policies create workload-policy WORKLOAD_POLICY_NAME \
    --type=HIGH_THROUGHPUT \
    --accelerator-topology=TPU_TOPOLOGY \
    --project=PROJECT_ID \
    --region=REGION

      Substitua:

      • WORKLOAD_POLICY_NAME: um nome para a política de carga de trabalho.
      • TPU_TOPOLOGY: a topologia da TPU Ironwood (TPU7x). Por exemplo, use 2x2x2. Para mais informações sobre todas as topologias do Ironwood (TPU7x) compatíveis, consulte a seção de topologia.
      • PROJECT_ID: o ID do projeto do Google Cloud .
      • REGION: a região da política de carga de trabalho. Uma política de carga de trabalho é um recurso regional que pode ser usado em todos os pools de nós.

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

      apiVersion: cloud.google.com/v1
kind: ComputeClass
metadata:
  name: tpu-class
spec:
  priorities:
    - tpu:
        type: tpu7x
        topology: TPU_TOPOLOGY
        count: 4
      placement:
        policyName: WORKLOAD_POLICY_NAME
  nodePoolAutoCreation:
    enabled: true

    3. (Opcional) Você pode consumir uma reserva ou um sub-bloco específico. Por exemplo, adicione o seguinte specs ao manifesto ComputeClass:

        reservations:
    affinity: Specific
    specific:
      - name: RESERVATION_NAME
        reservationBlock:
          name: RESERVATION_BLOCK_NAME
          reservationSubBlock:
            name: RESERVATION_SUB_BLOCK_NAME

      Substitua:

      • RESERVATION_NAME: o nome da reserva de capacidade do Compute Engine.
      • RESERVATION_BLOCK_NAME: o nome do bloco de reserva de capacidade do Compute Engine.
      • RESERVATION_SUB_BLOCK_NAME: o nome do sub-bloco da reserva de capacidade do Compute Engine.

      Para mais informações, consulte Consumo de recursos zonais reservados.

    Outras versões de TPU

    Para provisionar TPUs v3, v4, v5p, v5e ou v6e (Trillium) usando uma ComputeClass personalizada configurada para TPUs, 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_TYPE
      count: NUMBER_OF_CHIPS
      topology: TOPOLOGY
  - spot: true
    tpu:
      type: {"<var>"}}TPU_TYPE
      count: NUMBER_OF_CHIPS
      topology: TOPOLOGY
  - flexStart:
      enabled: true
    tpu:
      type: {"<var>"}}TPU_TYPE
      count: NUMBER_OF_CHIPS
      topology: TOPOLOGY
  nodePoolAutoCreation:
    enabled: true

      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.

    2. Implante a ComputeClass:

      kubectl apply -f tpu-compute-class.yaml

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

    Criar um job que consuma TPUs

    1. 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 de CHIP_COUNT na ComputeClass 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.

    2. 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 ComputeClasses personalizadas.

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

      kubectl delete -f tpu-job.yaml

    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 de IA do JAX mais recente. Para definir uma versão diferente, consulte a lista de imagens de IA do JAX atuais.
      • 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 programação de coleta

    No TPU Trillium, é possível usar a programação 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 na coleção disponíveis para atender ao tráfego.

    A TPU Trillium oferece suporte à programação 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 à coleta, crie outro pool de nós de fração de TPU de vários hosts com o mesmo nome de coleta 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 da programação de coleta, consulte Como funciona a programação 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ó são compatíveis 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 programação de coleta e carga de trabalho de TPU usando classes de computação personalizadas, consulte Coleta de vários hosts de TPU e Definir o tipo de carga de trabalho para o SLO de 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.

    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á o seguinte: 

      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á o seguinte: 

      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á disponível para todos os usuários. Para conferir o status dos seus pools de nós de vários hosts da TPU no GKE, acesse o painel Status do pool de nós da 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 console doGoogle 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 da 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 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, consulte 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 do 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 de TPU dos clusters no 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 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 seguinte métrica do sistema do GKE 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 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

    O tempo entre interrupções do pool de nós mede 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 :