Como trocar versões de software na Cloud TPU

Visão geral

A versão do software do framework em execução na TPU precisa corresponder à versão em execução na VM local. Essa versão do software agora pode ser ativada no Cloud TPU sem excluir e recriar a TPU. Isso também permite configurar o Cloud TPU com versões noturnas específicas de estruturas de software. Ainda é recomendável selecionar uma versão compatível dessas bibliotecas.

Uso

A maneira recomendada de alternar versões é usar a biblioteca em Python cloud-tpu-client.

Exemplo de uso do TensorFlow.

from cloud_tpu_client import Client
import argparse

parser = argparse.ArgumentParser()
parser.add_argument('--tpu-name',
                    type=str,
                    required=True,
                    help='Name of the TPU Instance')
parser.add_argument('--target-version',
                    type=str,
                    required=True,
                    help='target TPU Runtime version')

args = parser.parse_args()

c = Client(args.tpu_name)
c.configure_tpu_version(args.target_version, restart_type='ifNeeded')
c.wait_for_healthy()

Isso configura o Cloud TPU para corresponder à versão do TensorFlow em execução na sua VM local. Isso inclui versões oficiais, bem como versões noturnas datadas.

O parâmetro restart_type da API configure_tpu_version define o comportamento de reinicialização da TPU ao alternar as versões. As opções são "always" (padrão) e "ifNeeded".

• "always" pode ser usado para corrigir uma TPU com, por exemplo, o status UNHEALTHY_TENSORFLOW ou que retorna erros de falta de memória (OOM, na sigla em inglês) causados por recursos vazados de uma execução anterior. Quando essa opção é ativada, a TPU é reiniciada mesmo quando uma nova versão do framework não está instalada.

• "ifNeeded" pode ser útil porque não reinicia o ambiente de execução se já estiver na versão correta. Portanto, não adiciona nenhum tempo de inicialização significativo a um script de treinamento. Quando essa opção é ativada, a TPU é reiniciada apenas se não tiver a versão correta do framework instalada.

A biblioteca se comunica diretamente com o Cloud TPU para que esse código precise ser executado em uma VM na mesma rede. É recomendável executar isso no código para o restante do modelo.

Outras opções de software

O TensorFlow inclui uma string tf.__version__, que é a maneira mais simples de configurar a versão correta. Outras opções de software incluem:

• PyTorch - pytorch-0.1, pytorch-0.5-dev20200320, pytorch-nightly
• Jax - tpu_driver, tpu_driver0.1-dev20200320, tpu_driver_nightly

Por exemplo, para configurar uma TPU para ser executada com a versão noturna mais recente do PyTorch.

from cloud_tpu_client import Client
c = Client()
c.configure_tpu_version('pytorch-nightly', restart_type='ifNeeded')
c.wait_for_healthy()