Implemente TPUs no modo de capacidade total no GKE

Este documento explica como implementar e gerir cargas de trabalho através da funcionalidade modo de capacidade total da TPU no GKE. Todas as reservas no modo de capacidade oferecem um controlo melhorado sobre os seus recursos de TPUs, o que lhe permite colocar as suas cargas de trabalho com um controlo detalhado na sua capacidade reservada.

Este documento destina-se a engenheiros de aprendizagem automática (ML) e administradores e operadores da plataforma que querem usar a orquestração de contentores do Kubernetes com controlo detalhado sobre as implementações de TPUs.

Antes de ler este documento, certifique-se de que conhece o seguinte:

O que é o modo de capacidade total da TPU?

O modo de capacidade total da TPU, ativado pelo TPU Cluster Director, dá-lhe controlo total sobre a capacidade da TPU reservada. O TPU Cluster Director é um serviço de gestão que lhe dá controlo baseado em reservas sobre as suas TPUs.

Ao contrário do modo gerido anterior, em que Google Cloud reserva uma parte da sua capacidade para processar falhas de hardware, o modo de capacidade total concede-lhe acesso à totalidade dos seus recursos de TPU reservados. Este modo oferece visibilidade total do estado do hardware, mas também transfere para si a responsabilidade de gerir falhas de nós e manutenção planeada.

Para mais informações sobre as principais funcionalidades do modo de capacidade total, consulte o artigo Modo de capacidade total na vista geral do diretor do cluster de TPUs.

Terminologia relacionada com o modo de capacidade total no GKE

A tabela seguinte inclui os termos e as equivalências definidos pela dimensão de um bloco, um sub-bloco e um cubo na versão Ironwood (TPU7x). Um cubo é uma 4x4x4topologia de chips de TPU interligados, que só é aplicável a topologias em tuplos de 3 elementos ({A}x{B}x{C}).

Recurso de TPU Núcleos Batatas fritas Anfitriões Cubos
1 chip 2 1 - -
1 anfitrião 8 4 1 -
1 sub-bloco 128 64 16 1
Um bloco contém 144 sub-blocos 18432 9216 2304 144

Para mais informações sobre as topologias permitidas num bloco, consulte o artigo Escolha uma topologia.

Antes de começar

Antes de começar, certifique-se de que realizou as seguintes tarefas:

  • Ative a API Google Kubernetes Engine.
  • Ative a API Google Kubernetes Engine
  • Se quiser usar a CLI gcloud para esta tarefa, instale-a e, em seguida, inicialize-a. Se instalou anteriormente a CLI gcloud, execute o comando gcloud components update para obter a versão mais recente. As versões anteriores da CLI gcloud podem não suportar a execução dos comandos neste documento.
  • Certifique-se de que tem um cluster padrão existente na versão 1.34.0-gke.2201000 ou posterior. Para criar um novo cluster, consulte o artigo Criar um cluster regional.
  • Certifique-se de que tem quota suficiente para as TPUs na região que quer usar.
  • Instale o JobSet v0.2.3 ou posterior.

Limitações

O modo de capacidade total da TPU no GKE suporta apenas versões Ironwood (TPU7x).

Use o modo de capacidade total da TPU no GKE

Esta secção descreve o fluxo de trabalho para usar o modo de capacidade total da TPU no GKE.

  1. Familiarize-se com o TPU Cluster Director.
  2. Peça capacidade de TPU no modo de toda a capacidade.
  3. Ver a topologia e o estado de funcionamento de todas as reservas no modo de capacidade.
  4. Conclua o passo neste documento:
    1. Crie um node pool do GKE
    2. Agende cargas de trabalho
    3. Faça a gestão de falhas de nós com uma reserva no modo de capacidade total
  5. Faça a gestão de eventos de manutenção com as TPUs no modo de capacidade total.
  6. Comunique e repare anfitriões com falhas com TPUs no modo de capacidade total.

Crie um node pool numa reserva no modo de capacidade total

O modo All Capacity no GKE permite-lhe criar pools de nós das seguintes formas:

  • Um conjunto de nós em que o GKE seleciona o bloco ou o sub-bloco na sua reserva de capacidade total da TPU.
  • Um conjunto de nós que segmenta um bloco ou um sub-bloco específico numa reserva no modo de capacidade total da TPU.

O GKE seleciona o bloco ou o sub-bloco na sua reserva de toda a capacidade da TPU

Neste modo, o GKE seleciona o posicionamento do conjunto de nós na sua reserva de capacidade total da TPU. Este processo é semelhante ao da criação de pools de nós com outras opções de aprovisionamento de TPUs, como VMs a pedido ou VMs Spot.

Para criar um node pool, use o comando gcloud container node-pools create com a flag --reservation. Especifique o nome completo do recurso da sua reserva de TPU como o valor da flag --reservation.

Para ver um exemplo de um comando de criação de um node pool, consulte o artigo Crie manualmente um node pool.

Segmente um bloco ou um sub-bloco numa reserva

O modo de capacidade total da TPU permite-lhe segmentar um bloco ou um sub-bloco específico na sua reserva de TPU para cargas de trabalho paralelas. Esta capacidade é útil para cargas de trabalho que requerem proximidade entre os chips de TPU para minimizar a latência.

  1. Veja as configurações de blocos, sub-blocos e anfitriões disponíveis. Conclua os passos no documento Veja a topologia e o estado de funcionamento das reservas no modo de capacidade total.

  2. Crie uma política de carga de trabalho:

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

    Substitua o seguinte:

    • WORKLOAD_POLICY_NAME: um nome para a sua política de carga de trabalho.
    • TPU_TOPOLOGY: a topologia Ironwood da TPU (TPU7x). Por exemplo, 2x2x2. Para ver todas as topologias do Ironwood (TPU7x) suportadas, consulte o artigo Planeie TPUs no GKE.
    • PROJECT_ID: o ID do seu Google Cloud projeto.
    • REGION: a região da política de carga de trabalho. Uma política de carga de trabalho é um recurso regional e pode ser reutilizada em pools de nós que partilham a mesma topologia.
  3. Para criar um conjunto de nós e segmentar um bloco ou um sub-bloco específico da reserva, use a flag --reservation para especificar o nome completo do recurso do bloco ou do sub-bloco de destino na sua reserva.

    Para segmentar um bloco específico na sua reserva, use o seguinte comando:

    gcloud container node-pools create NODE_POOL_NAME \
        --cluster=CLUSTER_NAME \
        --machine-type=tpu7x-standard-4t \
        --placement-policy=WORKLOAD_POLICY_NAME \
        --zone=ZONE \
        --reservation=project/PROJECT/reservation/RESERVATION_NAME/reservationBlocks/BLOCK_NAME
    

    Para segmentar um sub-bloco específico num bloco, use o seguinte comando:

    gcloud container node-pools create NODE_POOL_NAME \
        --cluster=CLUSTER_NAME \
        --machine-type=tpu7x-standard-4t \
        --placement-policy=WORKLOAD_POLICY_NAME \
        --zone=ZONE \
        --reservation=project/PROJECT/reservation/RESERVATION_NAME/reservationBlocks/BLOCK_NAME/reservationSubBlocks/SUB_BLOCK_NAME
    

    Substitua o seguinte:

    • NODE_POOL_NAME: o nome do novo node pool.
    • CLUSTER_NAME: o nome do seu cluster do GKE.
    • WORKLOAD_POLICY_NAME: o nome da política de carga de trabalho que criou.
    • ZONE: a zona do node pool, por exemplo, us-central1-a.
    • PROJECT: o ID do seu Google Cloud projeto.
    • RESERVATION_NAME: o nome da sua reserva de TPU.
    • BLOCK_NAME: o bloco específico na sua reserva.
    • SUB_BLOCK_NAME: o sub-bloco específico na sua reserva.

    Nos comandos anteriores,

Agende cargas de trabalho

Depois de criar um conjunto de nós com VMs de TPU no modo de capacidade total, pode implementar a sua carga de trabalho como qualquer outro conjunto de nós de TPU. Para o modo de capacidade total da TPU, não existem diferenças adicionais no agendamento de cargas de trabalho em comparação com os conjuntos de nós que usam reservas padrão suportadas por SLOs.

Para mais informações e exemplos de cargas de trabalho que usam TPUs, consulte os artigos Execute a sua carga de trabalho em nós de fatia de TPU e Execute uma carga de trabalho de várias fatias.

Faça a gestão das falhas de nós

As reservas no modo de capacidade total da TPU são reservas sem restrições. Sem retenção significa que recebe a capacidade total da TPU, que inclui a parte que Google Cloud normalmente retém para as comutações por falha no modo de capacidade gerida.

No modo de capacidade total da TPU, se uma VM falhar devido a problemas como uma falha de hardware, Google Cloud tenta recuperar a VM no mesmo anfitrião (reparação no local). Por isso, recomendamos que mantenha capacidade adicional para acomodar o reagendamento da carga de trabalho durante falhas de infraestrutura.

Falha e recuperação do nó

Quando um nó falha numa reserva no modo de capacidade total da TPU, ocorrem os seguintes eventos:

  • Google Cloud inicia um evento de reparação para a instância de VM do Compute Engine com falha. Este processo tenta restaurar a VM para um estado RUNNING e devolver o estado do nó do GKE para READY.
  • A VM de TPU entra num estado de reparação e qualquer carga de trabalho em execução nesse nó pode falhar, consoante a respetiva política de comutação por falha. O estado do conjunto de nós não muda para ERROR, mesmo que uma ou mais das respetivas VMs estejam a ter falhas.

Para monitorizar o estado de saúde dos seus nós, siga estes passos:

  1. Liste os nós no node pool:

    kubectl get nodes
    

    Os nós com falhas têm o estado NotReady.

  2. Monitorize o estado do nó do Compute Engine:

    Para VMs de TPU no modo de capacidade total, use gcloud compute instances describe. Este comando também fornece o estado físico da topologia para encontrar detalhes do anfitrião, sub-bloco e bloco.

    gcloud compute instances describe VM_NAME \
            --format="table[box,title=VM-Position](resourceStatus.physical_host_topology:label=location)" \
            --zone=ZONE
    

    Substitua o seguinte:

    • VM_NAME: o nome da instância de VM da TPU.
    • ZONE: a zona da VM, por exemplo, us-central1-a.

    Para ver mais passos sobre como obter informações de topologia e estado sobre a capacidade do modo de capacidade total, consulte o artigo Veja a topologia e o estado de funcionamento das reservas do modo de capacidade total.

Faça a gestão da manutenção

Para lidar com potenciais interrupções e ajudar a garantir que as suas cargas de trabalho permanecem resilientes, pode gerir a manutenção de nós individuais através das políticas de manutenção do GKE. Para mais informações, consulte o artigo Janelas de manutenção e exclusões de manutenção.

O GKE não suporta a manutenção de grupos para VMs de TPU numa reserva no modo de capacidade total. Para realizar a manutenção de grupos ao nível do sub-bloco, do bloco ou da reserva, use as APIs Compute Engine. Para mais informações, consulte o artigo Faça a gestão de eventos de manutenção com as UTPs no modo de capacidade total.

Limpar

Para evitar incorrer em cobranças indesejadas na sua conta Google Cloud , elimine os pools de nós de UTPs que já não tenham cargas de trabalho agendadas. Se as cargas de trabalho em execução tiverem de ser terminadas corretamente, use o comando kubectl drain para limpar as cargas de trabalho antes de eliminar o conjunto de nós.

  1. Elimine um node pool da TPU:

    gcloud container node-pools delete NODE_POOL_NAME \
        --location=LOCATION \
        --cluster=CLUSTER_NAME
    

    Substitua o seguinte:

    • NODE_POOL_NAME: o nome do node pool.
    • CLUSTER_NAME: o nome do cluster.
    • LOCATION: a localização de computação do cluster.

O que se segue?