Gerenciar recursos da TPU

Nesta página, descrevemos como gerenciar recursos do Cloud TPU usando:

Pré-requisitos

Antes de executar esses procedimentos, instale a Google Cloud CLI, crie um projeto do Google Cloud e ative a API Cloud TPU. Para instruções, consulte Configurar o ambiente do Cloud TPU.

Se você estiver usando a Google Cloud CLI, poderá executar comandos usando o Cloud Shell, uma VM do Compute Engine ou sua máquina local. O Cloud Shell permite interagir com Cloud TPUs sem precisar instalar software. O Cloud Shell é desconectado após um período de inatividade. Se você estiver executando comandos de longa duração, recomendamos instalar a Google Cloud CLI na sua máquina local. Para mais informações sobre a Google Cloud CLI, consulte a Referência do gcloud.

Provisionar Cloud TPUs

É possível provisionar um Cloud TPU usando gcloud, o console do Google Cloud ou a API Cloud TPU.

Há dois métodos para provisionar TPUs usando gcloud:

  1. Como usar recursos na fila: gcloud alpha compute tpus queued-resources create
  2. Usando a API Create Node: gcloud compute tpus tpu-vm create

A prática recomendada é provisionar TPUs usando recursos na fila. Quando você solicita recursos na fila, a solicitação é adicionada a uma fila mantida pelo serviço do Cloud TPU. Quando o recurso solicitado fica disponível, ele é atribuído ao seu projeto do Google Cloud para uso exclusivo imediato. Para mais informações, consulte Recursos em fila gerenciados.

Ao usar o Multislice, você precisa usar recursos na fila e especificar os seguintes parâmetros extras:

export NODE_COUNT=node_count
export NODE_PREFIX=your_tpu_prefix # Optional

onde:

  • ${NODE_COUNT} é o número de frações a serem criadas.
  • ${NODE_PREFIX} é o prefixo que você especifica para gerar nomes para cada fatia. Um número é anexado ao prefixo de cada fatia. Por exemplo, se você definir ${NODE_PREFIX} como mySlice, as fatias serão nomeadas: mySlice-0, mySlice-1, continuando numericamente para cada fração.

Para saber mais sobre o recurso, consulte Introdução ao recurso multislice.

Criar um Cloud TPU usando a API Create Node

Ao criar um Cloud TPU, é preciso especificar a imagem da VM da TPU (também chamada de versão do software da TPU). Para determinar qual imagem de VM usar, consulte Imagens de VM de TPU.

Também é necessário especificar a configuração da TPU em termos de TensorCores ou chips de TPU. Para mais informações, consulte a seção sobre a versão da TPU que você está usando em Arquitetura do sistema.

gcloud

Para criar uma TPU usando a API Create Node, use o comando gcloud compute tpus tpu-vm create.

O comando a seguir usa uma configuração baseada no TensorCore:

$ gcloud compute tpus tpu-vm create tpu-name \
  --zone=us-central2-b \
  --accelerator-type=v4-8 \
  --version=tpu-software-version

Descrições de sinalizações de comando

zone
A zona em que você planeja criar o Cloud TPU.
accelerator-type
O tipo de acelerador especifica a versão e o tamanho do Cloud TPU que você quer criar. Para mais informações sobre os tipos de aceleradores compatíveis com cada versão de TPU, consulte Versões de TPU.
version
A versão do software da TPU.
shielded-secure-boot (opcional)
Especifica que as instâncias de TPU sejam criadas com a inicialização segura ativada. Isso as torna, implicitamente, instâncias de VM protegida. Consulte O que é VM protegida? para obter mais detalhes.

O comando a seguir cria uma TPU com uma topologia específica:

$ gcloud compute tpus tpu-vm create tpu-name \
  --zone=us-central2-b \
  --type=v4 \
  --topology=2x2x1 \
  --version=tpu-software-version

Sinalizações obrigatórias

tpu-name
O nome da VM da TPU que você está criando.
zone
A zona em que você está criando o Cloud TPU.
type
A versão da TPU que você quer usar. Para mais informações, consulte Versões de TPU.
topology
A disposição física dos chips de TPU, especificando o número de chips em cada dimensão. Para mais informações sobre topologias compatíveis com cada versão de TPU, consulte Versões de TPU.
version
A versão do software da TPU que você quer usar. Para mais informações, consulte Versões de software de TPU.

Console

  1. No console do Google Cloud, acesse a página TPUs:

    Acessar TPUs

  2. Clique em Criar TPU.

  3. No campo Nome, insira um nome para a TPU.

  4. Na caixa Zona, selecione a zona em que a TPU será criada.

  5. Na caixa Tipo de TPU, selecione um tipo de acelerador. O tipo de acelerador especifica a versão e o tamanho do Cloud TPU que você quer criar. Para mais informações sobre os tipos de aceleradores compatíveis com cada versão de TPU, consulte Versões de TPU.

  6. Na caixa Versão do software da TPU, selecione uma versão do software. Ao criar uma VM do Cloud TPU, a versão do software da TPU especifica a versão do ambiente de execução da TPU que será instalada. Para mais informações, consulte Imagens de VM da TPU.

  7. Clique em Criar para criar os recursos.

curl

O comando a seguir usa curl para criar uma TPU.

$ curl -X POST -H "Authorization: Bearer $(gcloud auth print-access-token)" -H "Content-Type: application/json" -d "{accelerator_type: 'v4-8', \
runtime_version:'tpu-vm-tf-2.16.1-pjrt', \
network_config: {enable_external_ips: true}, \
shielded_instance_config: { enable_secure_boot: true }}" \
https://tpu.googleapis.com/v2/projects/project-id/locations/us-central2-b/nodes?node_id=node_name

Campos obrigatórios

runtime_version
A versão do ambiente de execução do Cloud TPU que você quer usar.
project
O nome do projeto registrado do Google Cloud.
zone
A zona em que você está criando o Cloud TPU.
node_name
O nome da VM da TPU que você está criando.

Executar um script de inicialização

É possível executar um script de inicialização em cada VM da TPU especificando a sinalização --metadata startup-script ao criar a VM da TPU. O comando a seguir cria uma VM da TPU usando um script de inicialização.

$ gcloud compute tpus tpu-vm create tpu-name \
    --zone=us-central2-b \
    --accelerator-type=tpu-type \
    --version=tpu-vm-tf-2.16.1-pjrt \
    --metadata startup-script='#! /bin/bash
      pip3 install numpy
      EOF'

Conectar-se a um Cloud TPU

gcloud

Conecte-se ao Cloud TPU usando SSH:

$ gcloud compute tpus tpu-vm ssh tpu-name --zone=zone

Quando você solicita uma fração maior que um único host, o Cloud TPU cria uma VM da TPU para cada host. O número de chips de TPU por host depende da versão da TPU.

Para instalar binários ou executar código, conecte-se a cada VM da TPU usando o tpu-vm ssh command.

$ gcloud compute tpus tpu-vm ssh tpu-name

Para se conectar a uma VM TPU específica usando SSH, use a sinalização --worker, que segue um índice baseado em 0:

$ gcloud compute tpus tpu-vm ssh tpu-name --worker=1

Para executar um comando em todas as VMs da TPU com um único comando, use as sinalizações --worker=all e --command:

$ gcloud compute tpus tpu-vm ssh tpu-name \
  --project=your_project_ID \
  --zone=zone \
  --worker=all \
  --command='pip install "jax[tpu]==0.4.20" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'

Para o multislice, é possível executar um comando em uma única VM usando o nome da TPU enumerada, com cada prefixo de fração e o número anexado a ela. Para executar um comando em todas as VMs de TPU em todas as frações, use as sinalizações --node=all, --worker=all e --command, com uma sinalização --batch-size opcional.

$ gcloud compute tpus queued-resources ssh ${QUEUED_RESOURCE_ID} \
  --project=project_ID \
  --zone=zone \
  --node=all \
  --worker=all \
  --command='pip install "jax[tpu]==0.4.20" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html' \
  --batch-size=4

Console

Para se conectar às suas TPUs no console do Google Cloud, use o SSH no navegador:

  1. No console do Google Cloud, acesse a página TPUs:

    Acessar TPUs

  2. Na lista de VMs da TPU, clique em SSH na linha da VM da TPU à qual você quer se conectar.

Listar seus recursos do Cloud TPU

É possível listar todos os Cloud TPUs em uma zona especificada.

gcloud

$ gcloud compute tpus tpu-vm list --zone=zone

Console

No console do Google Cloud, acesse a página TPUs:

Acessar TPUs

Recuperar informações sobre o Cloud TPU

É possível recuperar informações sobre um Cloud TPU especificado.

gcloud

$ gcloud compute tpus tpu-vm describe tpu-name \
  --zone=zone

Console

  1. No console do Google Cloud, acesse a página TPUs:

    Acessar TPUs

  2. Clique no nome do Cloud TPU. O console exibe a página de detalhes do Cloud TPU.

Interrompa seus recursos do Cloud TPU

É possível interromper um único Cloud TPU para interromper as cobranças sem perder a configuração e o software da sua VM.

gcloud

$ gcloud compute tpus tpu-vm stop tpu-name \
  --zone=zone

Console

  1. No console do Google Cloud, acesse a página TPUs:

    Acessar TPUs

  2. Marque a caixa de seleção ao lado do Cloud TPU.

  3. Clique em Interromper.

Inicie seus recursos do Cloud TPU

É possível iniciar um Cloud TPU quando ele está parado.

gcloud

$ gcloud compute tpus tpu-vm start tpu-name \
  --zone=zone

Console

  1. No console do Google Cloud, acesse a página TPUs:

    Acessar TPUs

  2. Marque a caixa de seleção ao lado do Cloud TPU.

  3. Clique em Iniciar.

Excluir um Cloud TPU

Exclua as frações da VM da TPU ao final da sessão.

gcloud

$ gcloud compute tpus tpu-vm delete tpu-name \
  --project=project-id \
  --zone=zone \
  --quiet

Descrições de sinalizações de comando

zone
A zona em que você planeja excluir o Cloud TPU.

Console

  1. No console do Google Cloud, acesse a página TPUs:

    Acessar TPUs

  2. Marque a caixa de seleção ao lado do Cloud TPU.

  3. Clique em Excluir.

Configurações avançadas

Especificar recursos de rede personalizados

Ao criar a TPU, é possível especificar uma rede ou sub-rede.

gcloud

Para especificar a rede ou sub-rede usando a CLI gcloud, use as seguintes sinalizações de comando:

--network NETWORK --subnetwork SUBNETWORK

curl

Para especificar a rede ou sub-rede em uma chamada de curl, adicione os seguintes parâmetros ao corpo da solicitação:

network_config: {network: 'NETWORK', subnet: 'SUBNETWORK', enable_external_ips: true}

Rede

Também é possível especificar a rede a ser usada para a TPU. Se não for especificada, a rede default será usada.

Formatos de rede válidos:

https://www.googleapis.com/compute/{version}/projects/{proj-id}/global/networks/{network}
compute/{version}/projects/{proj-id}/global/networks/{network}
compute/{version}/projects/{proj-##}/global/networks/{network}
projects/{proj-id}/global/networks/{network}
projects/{proj-##}/global/networks/{network}
global/networks/{network}
{network}

Sub-rede

É possível especificar uma sub-rede específica para usar na TPU. A sub-rede especificada precisa estar na mesma região da zona em que a TPU é executada.

Formatos válidos:

https://www.googleapis.com/compute/{version}/projects/{proj-id}/regions/{region}/subnetworks/{subnetwork}
compute/{version}/projects/{proj-id}/regions/{region}/subnetworks/{subnetwork}
compute/{version}/projects/{proj-##}/regions/{region}/subnetworks/{subnetwork}
projects/{proj-id}/regions/{region}/subnetworks/{subnetwork}
projects/{proj-##}/regions/{region}/subnetworks/{subnetwork}
regions/{region}/subnetworks/{subnetwork}
{subnetwork}

Ative o Acesso privado do Google

Para se conectar a VMs da TPU usando SSH, adicione configurações de acesso para as VMs da TPU ou ative o Acesso privado do Google para a sub-rede a que as VMs da TPU estão conectadas.

Para adicionar configurações de acesso, defina enable_external_ips. Quando você cria uma TPU, enable_external_ips é definido por padrão.

Para desativar, ative os IPs internos:

gcloud

Use a sinalização --internal-ips ao criar uma TPU:

--internal-ips

curl

Adicione os parâmetros a seguir ao corpo da solicitação:

network_config: {enable_external_ips: false}

Depois de configurar o Acesso privado do Google, conecte-se à VM usando SSH.

Anexar uma conta de serviço personalizada

Cada VM da TPU tem uma conta de serviço associada que ela usa para fazer solicitações de API em seu nome. As VMs da TPU usam essa conta de serviço para chamar as APIs do Cloud TPU e acessar o Cloud Storage e outros serviços. Por padrão, a VM da TPU usa a conta de serviço padrão do Compute Engine.

A conta de serviço precisa ser definida no mesmo projeto do Google Cloud em que você cria a VM da TPU. As contas de serviço personalizadas usadas para VMs de TPU precisam ter o papel de Leitor de TPU para chamar a API Cloud TPU. Se o código em execução na VM da TPU chamar outros serviços do Google Cloud, ele precisará ter os papéis necessários para acessar esses serviços.

Ao criar uma TPU, é possível especificar uma conta de serviço personalizada usando a sinalização --service-account. Para mais informações sobre contas de serviço, consulte Contas de serviço.

Use os comandos a seguir para especificar uma conta de serviço personalizada.

gcloud

$ gcloud compute tpus tpu-vm create tpu-name \
  --zone=us-central2-b \
  --accelerator-type=tpu-type \
  --version=tpu-vm-tf-2.16.1-pjrt \
  --service-account=your-service-account

curl

$ curl -X POST -H "Authorization: Bearer $(gcloud auth print-access-token)" -H "Content-Type: application/json" -d "{accelerator_type: 'v4-8', \
runtime_version:'tpu-vm-tf-2.16.1-pjrt', \
network_config: {enable_external_ips: true}, \
shielded_instance_config: { enable_secure_boot: true }}" \
service_account: {email: 'your-service-account'} \
https://tpu.googleapis.com/v2/projects/project-id/locations/us-central2-b/nodes?node_id=node_name

Ativar métodos SSH personalizados

  1. Configurar um firewall para SSH.

    A rede padrão é pré-configurada para permitir o acesso SSH a todas as VMs. Se você não usar a rede padrão ou tiver alterado as configurações de rede padrão, talvez seja necessário ativar explicitamente o acesso SSH adicionando uma regra de firewall:

    $ gcloud compute tpus tpu-vm compute firewall-rules create \
      --network=network allow-ssh \
      --allow=tcp:22
    
  2. Conecte-se às VMs da TPU usando SSH.

    $ gcloud compute tpus tpu-vm ssh tpu-name \
      --zone=us-central2-b \
      --project=project-id
    

    Campos obrigatórios

    • tpu-name: nome da VM da TPU.
    • zone: a zona em que você criou a VM da TPU.
    • project-id: o nome do seu projeto do Google Cloud.

    Para uma lista de campos opcionais, consulte a documentação da API gcloud.