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. Agora é possível ativar essa versão de software em um Cloud TPU em execução, sem excluir e recriar a TPU. Isso também permite configurar a Cloud TPU com versões noturnas específicas de frameworks de software. Ainda é recomendável selecionar uma versão compatível com esses frameworks.

Uso

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

Exemplo de uso do TensorFlow.

import tensorflow as tf
from cloud_tpu_client import Client

c = Client()
c.configure_tpu_version(tf.__version__, restart_type='ifNeeded')

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

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

  • "sempre" pode ser usado para corrigir uma TPU com, por exemplo, status UNHEALTHY_TENSORFLOW ou que esteja retornando erros de falta de memória (OOM, na sigla em inglês) devido a recursos vazados de uma execução anterior. Quando essa opção é definida, 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, por isso não adicionará qualquer tempo de inicialização significativo a um script de treinamento. Quando essa opção é configurada, a TPU é reiniciada apenas se não tiver a versão correta do framework instalada.

A biblioteca se comunica diretamente com a Cloud TPU. Portanto, esse código precisa 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()