Executar aplicativos do Cloud TPU no GKE

Aviso: estas instruções se aplicam somente a TPUs que usam a arquitetura de nós de TPU. A arquitetura do nó da TPU é compatível apenas com o GKE nas TPUs v2 e v3. Estas instruções se aplicam somente à TPU v2 e v3. Para informações sobre o uso de TPUs no GKE, incluindo TPU v4, v5e ou v5p, consulte Implantar cargas de trabalho de TPU.

Neste guia, descrevemos como:

• Definir uma configuração do Cloud TPU como preparação para executá-la no Google Kubernetes Engine
• Crie um cluster do GKE com suporte ao Cloud TPU.
• Use o TensorBoard para visualizar métricas do Cloud TPU e analisar o desempenho.
• Criar e conteinerizar o modelo no Docker

Para mais informações sobre arquiteturas de VM de TPU, consulte Arquitetura do sistema. Este guia só pode ser usado com a arquitetura de nós de TPU.

Benefícios da execução de aplicativos do Cloud TPU no GKE

Os aplicativos de treinamento do Cloud TPU podem ser configurados para execução em contêineres do GKE (em inglês). Quando são eles, você tem os seguintes benefícios:

• Melhorias na configuração e no gerenciamento do fluxo de trabalho: o GKE gerencia o ciclo de vida da TPU. Depois que a inicialização e o treinamento do Cloud TPU forem configurados com o GKE, as cargas de trabalho poderão ser repetidas e gerenciadas pelo GKE, incluindo a recuperação de falhas de jobs.

• Custo otimizado: você só paga pela TPU enquanto o job estiver ativo. O GKE cria e exclui TPUs automaticamente de acordo com os requisitos de recurso de um pod.

• Uso flexível: basta uma pequena alteração na especificação do pod para solicitar um acelerador de hardware diferente (CPU, GPU ou TPU):

  kind: Pod
  metadata:
    name: example-tpu
    annotations:
      # The Cloud TPUs that will be created for this Job will support
      # TensorFlow 2.12.1. This version MUST match the
      # TensorFlow version that your model is built on.
      tf-version.cloud-tpus.google.com: "2.12.1"
  spec:
    containers:
    - name: example-container
      resources:
        limits:
          cloud-tpus.google.com/v2: 8
          # See the line above for TPU, or below for CPU / GPU.
          # cpu: 2
          # nvidia.com/gpu: 1
  

• Escalonabilidade:o GKE fornece APIs (Job e Deployment) que podem ser escalonadas para centenas de pods e nós de TPU do GKE.

• Tolerância a falhas: a API Job do GKE e o mecanismo de checkpoint do TensorFlow fornecem a semântica da execução até a conclusão. Em caso de falha nas instâncias de VM ou nos nós do Cloud TPU, os jobs de treinamento são automaticamente executados de novo com o estado mais recente lido no checkpoint.

Requisitos e limitações de configuração da Cloud TPU e do GKE

Ao definir sua configuração do GKE, observe o seguinte:

• O Cloud TPU não é compatível com pools de nós do Windows Server.
• É preciso criar o cluster do GKE e os pools de nós em uma zona em que a Cloud TPU esteja disponível. Também é possível criar os buckets do Cloud Storage para manter seus dados e modelos de treinamento na mesma região que o cluster do GKE. Consulte o documento de tipos e zonas para ver uma lista das zonas disponíveis.
• Você precisa usar endereços IP compatíveis com RFC 1918 para seus clusters do GKE. Para mais informações, consulte Rede do GKE.
• Cada contêiner pode solicitar, no máximo, uma Cloud TPU. No entanto, é permitido que vários contêineres em um pod solicitem uma.

Antes de começar

1. Faça login na sua conta do Google Cloud. Se você começou a usar o Google Cloud agora, crie uma conta para avaliar o desempenho de nossos produtos em situações reais. Clientes novos também recebem US$ 300 em créditos para executar, testar e implantar cargas de trabalho.

2. No console do Google Cloud, na página do seletor de projetos, selecione ou crie um projeto do Google Cloud.

   Observação: se você não pretende manter os recursos criados neste procedimento, crie um projeto novo em vez de selecionar um que já existe. Depois de concluir essas etapas, é possível excluir o projeto. Para fazer isso, basta remover todos os recursos associados a ele.

   Acessar o seletor de projetos

3. Verifique se a cobrança está ativada para o seu projeto do Google Cloud.

4. No console do Google Cloud, na página do seletor de projetos, selecione ou crie um projeto do Google Cloud.

   Observação: se você não pretende manter os recursos criados neste procedimento, crie um projeto novo em vez de selecionar um que já existe. Depois de concluir essas etapas, é possível excluir o projeto. Para fazer isso, basta remover todos os recursos associados a ele.

   Acessar o seletor de projetos

5. Verifique se a cobrança está ativada para o seu projeto do Google Cloud.

Ao usar o Cloud TPU com o GKE, são utilizados componentes faturáveis do Google Cloud Platform no projeto. Para fazer uma estimativa de custos, verifique os preços da Cloud TPU e os preços do GKE. Siga as instruções para limpar os recursos quando não forem mais necessários.

6. Ative as seguintes APIs no console do Google Cloud:

   • API da Cloud TPU
   • API Compute Engine
   • API GKE

Como criar um cluster compatível com o Cloud TPU

Use as instruções a seguir para configurar o ambiente e criar um cluster do GKE compatível com o Cloud TPU usando a CLI gcloud:

1. Instale os componentes gcloud, necessários para executar o GKE com o Cloud TPU:

   $ gcloud components install kubectl 

2. Configure gcloud om o ID do projeto do Google Cloud:

   $ gcloud config set project project-name
   

   Substitua project-name pelo nome do projeto do Google Cloud.

   Na primeira vez que você executar esse comando em uma nova VM do Cloud Shell, será exibida uma página Authorize Cloud Shell. Clique em Authorize na parte inferior da página para permitir que gcloud faça chamadas de API do Google Cloud com suas credenciais.

3. Configure gcloud com a zona em que você planeja usar um recurso do Cloud TPU. Este exemplo usa us-central1-b, mas é possível usar uma TPU em qualquer zona compatível.

   $ gcloud config set compute/zone us-central1-b
   

4. Use o comando gcloud container clusters create para criar um cluster no GKE com suporte para o Cloud TPU.

   Observação: se for necessário especificar um intervalo CIDR atual com uma VPC compartilhada, use o comando mostrado em Como usar um intervalo CIDR atual com a VPC compartilhada.

   $ gcloud container clusters create cluster-name \
     --release-channel=stable \
     --scopes=cloud-platform \
     --enable-ip-alias \
     --enable-tpu
   

   Descrições de sinalizações de comando

   release-channel

   Os canais de lançamento oferecem uma maneira de gerenciar upgrades automáticos dos clusters. Ao criar um novo cluster, escolha o canal de lançamento. Seu cluster só terá upgrade para as versões oferecidas no canal.

   scopes

   Garante que todos os nós no cluster tenham acesso ao bucket do Cloud Storage. O cluster e o bucket de armazenamento precisam estar no mesmo projeto para isso funcionar. Por padrão, os pods do Kubernetes herdam os escopos dos nós em que são implantados. Portanto, scopes=cloud-platform fornece o escopo cloud-platform a todos os pods do Kubernetes em execução no cluster. Se você quiser limitar o acesso por pod, veja como autenticar com contas de serviço no guia do GKE.

   enable-ip-alias

   Indica que o cluster usa intervalos de IP do alias. Isso é necessário para usar a Cloud TPU no GKE.

   enable-tpu

   Indica que o cluster precisa ser compatível com a Cloud TPU.

   tpu-ipv4-cidr (opcional, não especificado acima)

   Indica o intervalo CIDR a ser usado para a Cloud TPU. Especifique IP_RANGE na forma de IP/20, por exemplo, 10.100.0.0/20. Se essa sinalização não for especificada, um intervalo CIDR de tamanho /20 será alocado e atribuído automaticamente.

Quando o cluster for criado, você verá uma mensagem semelhante a esta:

NAME             LOCATION       MASTER_VERSION    MASTER_IP     MACHINE_TYPE   NODE_VERSION      NUM_NODES  STATUS
cluster-name  us-central1-b  1.16.15-gke.4901  34.71.245.25  n1-standard-1  1.16.15-gke.4901  3          RUNNING

Solicite uma Cloud TPU na especificação do pod do Kubernetes

Na especificação de pod do Kubernetes:

• É necessário criar os modelos nos contêineres usando a mesma versão do TensorFlow. Veja as versões compatíveis.

• Na especificação do contêiner, especifique o recurso do Cloud TPU na seção limits, no campo resource.

  Observe que a unidade do recurso da Cloud TPU corresponde ao número de núcleos da Cloud TPU. A tabela a seguir lista exemplos de solicitações de recursos válidas. Consulte Tipos e zonas de TPU para ver uma lista completa de recursos de TPU válidos.

  Se o recurso destinado ao uso for um pod do Cloud TPU, solicite uma cota, porque a cota padrão do pod do Cloud TPU é zero.

  Solicitação de recurso	Tipo de Cloud TPU
  cloud-tpus.google.com/v2: 8	Um dispositivo Cloud TPU v2 (8 núcleos)
  cloud-tpus.google.com/preemptible-v2: 8	Um dispositivo preemptivo da Cloud TPU v2 (8 núcleos)
  cloud-tpus.google.com/v3: 8	Um dispositivo da Cloud TPU v3 (8 núcleos)
  cloud-tpus.google.com/preemptible-v3: 8	Um dispositivo preemptivo da Cloud TPU v3 (8 núcleos)
  cloud-tpus.google.com/v2: 32	Um pod do Cloud TPU v2-32 (32 núcleos)
  cloud-tpus.google.com/v3: 32	Um pod do Cloud TPU v3-32 (32 núcleos)

  Para saber mais sobre como especificar recursos e limites na especificação de Pods, consulte a documentação do Kubernetes.

A especificação de pod de amostra a seguir solicita uma TPU preemptiva do Cloud TPU v2-8 com o TensorFlow 2.12.0.

A vida útil dos nós da Cloud TPU é vinculada aos pods do Kubernetes que os solicitam. A Cloud TPU é criada sob demanda quando o pod do Kubernetes é programado e reciclada quando o pod é excluído.

apiVersion: v1
kind: Pod
metadata:
  name: gke-tpu-pod
  annotations:
     # The Cloud TPUs that will be created for this Job will support
     # TensorFlow 2.12.1. This version MUST match the
     # TensorFlow version that your model is built on.
     tf-version.cloud-tpus.google.com: "2.12.1"
spec:
  restartPolicy: Never
  containers:
  - name: gke-tpu-container
    # The official TensorFlow 2.12.1 image.
    # https://hub.docker.com/r/tensorflow/tensorflow
    image: tensorflow/tensorflow:2.12.1
    command:
    - python
    - -c
    - |
      import tensorflow as tf
      print("Tensorflow version " + tf.__version__)

      tpu = tf.distribute.cluster_resolver.TPUClusterResolver('$(KUBE_GOOGLE_CLOUD_TPU_ENDPOINTS)')
      print('Running on TPU ', tpu.cluster_spec().as_dict()['worker'])

      tf.config.experimental_connect_to_cluster(tpu)
      tf.tpu.experimental.initialize_tpu_system(tpu)
      strategy = tf.distribute.TPUStrategy(tpu)

      @tf.function
      def add_fn(x,y):
          z = x + y
          return z

      x = tf.constant(1.)
      y = tf.constant(1.)
      z = strategy.run(add_fn, args=(x,y))
      print(z)
    resources:
      limits:
        # Request a single Preemptible v2-8 Cloud TPU device to train the model.
        cloud-tpus.google.com/preemptible-v2: 8

Como criar o job

Siga estas etapas para criar o job no cluster do GKE e instalar o kubectl.

1. Usando um editor de texto, crie uma especificação de pod, example-job.yaml, e copie e cole a especificação mostrada anteriormente.

2. Execute o job:

   $ kubectl create -f example-job.yaml
   

   pod "gke-tpu-pod" created

   Esse comando cria o job que programa o pod automaticamente.

3. Verifique se o pod do GKE foi programado e se os nós da Cloud TPU foram provisionados. Um pod do GKE que solicita nós do Cloud TPU pode ficar pendente por até cinco minutos antes da execução. Você verá uma saída semelhante à seguinte até que o pod do GKE seja programado.

   $ kubectl get pods -w
   
   NAME          READY     STATUS    RESTARTS   AGE
   gke-tpu-pod   0/1       Pending   0          1m
   

   Depois de aproximadamente cinco minutos, você verá algo assim:

   NAME          READY     STATUS              RESTARTS   AGE
   gke-tpu-pod   0/1       Pending             0          21s
   gke-tpu-pod   0/1       Pending             0          2m18s
   gke-tpu-pod   0/1       Pending             0          2m18s
   gke-tpu-pod   0/1       ContainerCreating   0          2m18s
   gke-tpu-pod   1/1       Running             0          2m48s
   gke-tpu-pod   0/1       Completed           0          3m8s
   

   Você precisa usar Ctrl-C para sair do comando "kubectl get".

   É possível imprimir informações de registro e recuperar informações mais detalhadas sobre cada pod do GKE usando os comandos kubectl a seguir. Por exemplo, para ver a saída do registro do pod do GKE, use:

   $ kubectl logs gke-tpu-pod

   A resposta será semelhante a esta:

   2021-09-24 18:55:25.400699: I tensorflow/core/platform/cpu_feature_guard.cc:142]
   This TensorFlow binary is optimized with oneAPI Deep Neural Network Library
   (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 FMA
   To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.
   2021-09-24 18:55:25.405947: I tensorflow/core/distributed_runtime/rpc/grpc_channel.cc:272]
   Initialize GrpcChannelCache for job worker -> {0 -> 10.0.16.2:8470}
   2021-09-24 18:55:25.406058: I tensorflow/core/distributed_runtime/rpc/grpc_channel.cc:272]
   Initialize GrpcChannelCache for job localhost -> {0 -> localhost:32769}
   2021-09-24 18:55:28.091729: I tensorflow/core/distributed_runtime/rpc/grpc_channel.cc:272]
   Initialize GrpcChannelCache for job worker -> {0 -> 10.0.16.2:8470}
   2021-09-24 18:55:28.091896: I tensorflow/core/distributed_runtime/rpc/grpc_channel.cc:272]
   Initialize GrpcChannelCache for job localhost -> {0 -> localhost:32769}
   2021-09-24 18:55:28.092579: I tensorflow/core/distributed_runtime/rpc/grpc_server_lib.cc:427]
   Started server with target: grpc://localhost:32769
   Tensorflow version 2.12.1
   Running on TPU  ['10.0.16.2:8470']
   PerReplica:{
     0: tf.Tensor(2.0, shape=(), dtype=float32),
     1: tf.Tensor(2.0, shape=(), dtype=float32),
     2: tf.Tensor(2.0, shape=(), dtype=float32),
     3: tf.Tensor(2.0, shape=(), dtype=float32),
     4: tf.Tensor(2.0, shape=(), dtype=float32),
     5: tf.Tensor(2.0, shape=(), dtype=float32),
     6: tf.Tensor(2.0, shape=(), dtype=float32),
     7: tf.Tensor(2.0, shape=(), dtype=float32)
   }
   

   Para ver uma descrição completa do pod do GKE, use:

   $ kubectl describe pod gke-tpu-pod
   

   Consulte Introspecção e depuração de aplicativos para mais detalhes.

Crie e coloque o modelo em um contêiner na imagem do Docker

Para detalhes sobre esse processo, consulte compilar e conteinerizar seu próprio modelo.

Ativar o suporte do Cloud TPU em um cluster atual

Esse recurso está sujeito aos "Termos de ofertas pré-GA" na seção "Termos de Serviço gerais" dos Termos específicos de serviço. Os recursos pré-GA estão disponíveis "no estado em que se encontram" e podem ter suporte limitado. Para mais informações, consulte as descrições da fase de lançamento.

Para ativar o suporte do Cloud TPU em um cluster atual do GKE, execute as seguintes etapas na Google Cloud CLI:

1. Ative o suporte do Cloud TPU:

   gcloud beta container clusters update cluster-name --enable-tpu
   

   Substitua cluster-name pelo nome do cluster.

   Importante: o endpoint original do servidor da API pode não estar disponível após a atualização do cluster.

2. Atualize a entrada kubeconfig:

   gcloud container clusters get-credentials cluster-name
   

Como definir um intervalo CIDR personalizado

Esse recurso está sujeito aos "Termos de ofertas pré-GA" na seção "Termos de Serviço gerais" dos Termos específicos de serviço. Os recursos pré-GA estão disponíveis "no estado em que se encontram" e podem ter suporte limitado. Para mais informações, consulte as descrições da fase de lançamento.

Por padrão, o GKE aloca um bloco CIDR com o tamanho de /20 para as TPUs provisionadas pelo cluster. Especifique um intervalo CIDR personalizado para o Cloud TPU executando o seguinte comando:

gcloud beta container clusters update cluster-name \
  --enable-tpu \
  --tpu-ipv4-cidr 10.100.0.0/20

Substitua:

• cluster-name: o nome do cluster atual.
• 10.100.0.0/20: o intervalo CIDR personalizado.

Como usar intervalos CIDR atuais com a VPC compartilhada

Observação: esta seção é destinada a usuários que estão implantando um cluster em uma rede VPC compartilhada que precisam usar intervalos de CIDR atuais. Se esse não for seu caso de uso, pule esta seção.

Siga o guia sobre TPU em clusters do GKE usando uma VPC compartilhada para verificar a configuração correta dessa VPC.

• Cluster do GKE com o Cloud TPU usando uma VPC compartilhada

Como desativar o Cloud TPU em um cluster

Importante: desativar o suporte do Cloud TPU exclui todas as TPUs provisionadas pelo cluster do GKE. Qualquer carga de trabalho que exija o Cloud TPU ficará indisponível quando a operação for iniciada.

Para desativar o suporte do Cloud TPU em um cluster atual do GKE, execute as seguintes etapas na Google Cloud CLI:

1. Verifique se nenhuma de suas cargas de trabalho está usando o Cloud TPU:

   $ kubectl get tpu
   

2. Desative o suporte do Cloud TPU no cluster:

   $ gcloud beta container clusters update cluster-name --no-enable-tpu
   

   Substitua cluster-name pelo nome do cluster.

   Para clusters zonais, essa operação leva cerca de 5 minutos e, para clusters regionais, ela leva cerca de 15 minutos, dependendo da região do cluster.

3. Depois que as operações forem concluídas sem erros, será possível verificar se as TPUs provisionadas pelo cluster foram removidas:

   $ gcloud compute tpus list
   

   Os nomes das TPUs criadas pelo Cloud TPU têm o seguinte formato:

   $ gke-cluster-name-cluster-id-tpu-tpu-id
   

   Substitua:

   • cluster-name: o nome do cluster atual.
   • cluster-id: o ID do cluster atual.
   • tpu-id: ID do Cloud TPU.

   Se alguma TPU aparecer, será possível excluí-las manualmente executando:

   $ gcloud compute tpus delete gke-cluster-name-cluster-id-tpu-tpu-id
   

Limpar

Quando terminar de usar o Cloud TPU no GKE, limpe os recursos para evitar cobranças extras na conta do Cloud Billing.

1. Execute o seguinte comando para excluir o cluster do GKE, substituindo cluster-name pelo nome do cluster e project-name pelo nome do projeto do Google Cloud:

   $ gcloud container clusters delete cluster-name \
   --project=project-name --zone=us-central1-b
   

2. Quando terminar de examinar os dados, use o comando gsutil para excluir o bucket do Cloud Storage que você criou. Substitua bucket-name pelo nome do seu bucket do Cloud Storage:

   $ gsutil rm -r gs://bucket-name