Como depurar VMs do Cloud TPU

Neste documento, descrevemos como usar o pacote PyPI cloud-tpu-diagnostics para gerar stack traces para processos em execução em VMs de TPU. Este pacote despeja os rastros do Python quando ocorre uma falha, por exemplo, falhas de segmentação, exceções de ponto flutuante ou exceções de operação ilegal. Além disso, ele também coleta periodicamente stack traces para ajudar a depurar situações em que o programa não está respondendo.

Para usar o pacote PyPI cloud-tpu-diagnostics, instale-o executando pip install cloud-tpu-diagnostics em todas as VMs de TPU. É possível fazer isso com um comando gcloud compute tpus tpu-vm ssh. Exemplo:

  gcloud compute tpus tpu-vm ssh you-tpu-name \
  --zone=your-zone \
  --project=your-project-name \
  --worker=all \
  --command="pip install cloud-tpu-diagnostics"

Também é preciso adicionar o seguinte código aos seus scripts em execução em todas as VMs da TPU.

from cloud_tpu_diagnostics import diagnostic
from cloud_tpu_diagnostics.configuration import debug_configuration
from cloud_tpu_diagnostics.configuration import diagnostic_configuration
from cloud_tpu_diagnostics.configuration import stack_trace_configuration

stack_trace_config = stack_trace_configuration.StackTraceConfig(
                      collect_stack_trace = True,
                      stack_trace_to_cloud = True)
debug_config = debug_configuration.DebugConfig(
                stack_trace_config = stack_trace_config)
diagnostic_config = diagnostic_configuration.DiagnosticConfig(
                      debug_config = debug_config)

Por padrão, os stack traces são coletados a cada 10 minutos. É possível alterar a duração entre dois eventos de coleta de stack trace para cinco minutos, por exemplo:

stack_trace_config = stack_trace_configuration.StackTraceConfig(
                      collect_stack_trace = True,
                      stack_trace_to_cloud = True,
                      stack_trace_interval_seconds = 300)

Una o método principal com diagnose() para coletar periodicamente os stack traces:

with diagnostic.diagnose(diagnostic_config):
    run_main()

Essa configuração começa a coletar stack traces dentro do diretório /tmp/debugging em cada VM da TPU. Há um agente em execução em todas as VMs de TPU que faz upload dos traces de um diretório temporário para o Cloud Logging.