Resolver problemas do job da TPU do Dataflow

Se você tiver problemas ao executar o job do Dataflow com TPUs, siga as etapas de solução de problemas abaixo para resolver o problema.

Resolver problemas da imagem do contêiner

Pode ser útil depurar o contêiner e o software de TPU em uma VM independente. É possível depurar com uma VM criada por um pool de nós do GKE ou em uma VM de worker do Dataflow em execução.

Depurar com uma VM independente

Para depurar seu contêiner em uma VM independente, crie um pool de nós do GKE que use a mesma VM de TPU para experimentação local. Por exemplo, criar um pool de nós do GKE com um dispositivo TPU V5 Lite em us-west1-c seria assim:

  1. Criar um cluster do GKE.

    gcloud container clusters create TPU_CLUSTER_NAME \
      --project PROJECT_ID \
      --release-channel=stable \
      --scopes=cloud-platform \
      --enable-ip-alias \
      --location us-west1-c
    
  2. Crie um pool de nós do GKE.

    gcloud container node-pools create TPU_NODE_POOL_NAME \
      --project PROJECT_ID \
      --location=us-west1-c \
      --cluster=TPU_CLUSTER_NAME \
      --node-locations=us-west1-c \
      --machine-type=ct5lp-hightpu-1t \
      --num-nodes=1 \
      [ --reservation RESERVATION_NAME \
      --reservation-affinity=specific ]
    
  3. Encontre o nome da VM do nó da TPU no pool de nós na interface do GKE ou com o comando a seguir.

    gcloud compute instances list --filter='metadata.kube-labels:"cloud.google.com/gke-nodepool=TPU_NODEPOOL_NAME"'
    
  4. Conecte-se a uma VM criada pelo pool de nós do GKE usando SSH:

    gcloud compute ssh --zone "us-west1-c" "VM_NAME" --project PROJECT_ID
    
  5. Depois de se conectar a uma VM usando SSH, configure o Docker para o Artifact Registry que você está usando.

    docker-credential-gcr configure-docker --registries=us-west1-docker.pkg.dev
    
  6. Em seguida, inicie um contêiner da imagem que você usa.

    docker run --privileged --network=host -it --rm --entrypoint=/bin/bash IMAGE_NAME
    
  7. Dentro do contêiner, teste se as TPUs estão acessíveis.

    Por exemplo, se você tiver uma imagem que usa o PyTorch para utilizar TPUs, abra um interpretador do Python:

    python3
    

    Em seguida, faça um cálculo em um dispositivo TPU:

    import torch
    import torch_xla.core.xla_model as xm
    dev = xm.xla_device()
    t1 = torch.randn(3,3,device=dev)
    t2 = torch.randn(3,3,device=dev)
    print(t1 + t2)
    

    Exemplo de resposta:

    >>> tensor([[ 0.3355, -1.4628, -3.2610],
    >>>        [-1.4656,  0.3196, -2.8766],
    >>>        [ 0.8667, -1.5060,  0.7125]], device='xla:0')
    
  8. Se a computação falhar, talvez a imagem não esteja configurada corretamente.

    Por exemplo, talvez seja necessário definir as variáveis de ambiente necessárias no Dockerfile da imagem. Para confirmar, tente de novo o cálculo depois de definir as variáveis de ambiente manualmente da seguinte forma:

    export TPU_SKIP_MDS_QUERY=1 # Don't query metadata
    export TPU_HOST_BOUNDS=1,1,1 # There's only one host
    export TPU_CHIPS_PER_HOST_BOUNDS=1,1,1 # 1 chips per host
    export TPU_WORKER_HOSTNAMES=localhost
    export TPU_WORKER_ID=0 # Always 0 for single-host TPUs
    export TPU_ACCELERATOR_TYPE=v5litepod-1 # Since we use v5e 1x1 accelerator.
    

    Se as dependências do PyTorch ou do LibTPU estiverem faltando, tente de novo o cálculo depois de instalá-las usando o seguinte comando:

    # Install PyTorch with TPU support
    pip install torch torch_xla[tpu] torchvision -f https://storage.googleapis.com/libtpu-releases/index.html
    

Depurar usando uma VM do Dataflow

Como alternativa, é possível se conectar à instância VM de worker do Dataflow usando SSH enquanto um job está em execução. Como as VMs de worker do Dataflow são desligadas após a conclusão do pipeline, talvez seja necessário aumentar artificialmente o tempo de execução fazendo um cálculo que aguarde por um período prolongado.

Como um dispositivo de TPU não pode ser compartilhado entre vários processos, talvez seja necessário executar um pipeline que não faça cálculos em uma TPU.

  1. Para encontrar uma VM para o job de TPU em execução, pesquise o ID do job do Dataflow na barra de pesquisa do console Google Cloud ou use o comando gcloud a seguir:

    gcloud compute instances list --project PROJECT_ID --filter "STATUS='RUNNING' AND description ~ 'Created for Dataflow job: JOB_ID'"
    
  2. Depois de se conectar a uma VM com TPUs usando SSH, inicie um contêiner da imagem que você usa. Para ver um exemplo, consulte Depurar com uma VM independente.

  3. Dentro do contêiner, reconfigure as configurações da TPU e instale as bibliotecas necessárias para testar a configuração. Para um exemplo, consulte Depurar com uma VM independente.

Os workers não são iniciados

Antes de resolver problemas, verifique se as seguintes opções de pipeline estão definidas corretamente:

  • a opção --dataflow_service_option=worker_accelerator
  • a opção --worker_zone
  • a opção --machine_type

Verifique se os registros do console mostram que os workers estão sendo iniciados, mas o job falha com uma mensagem semelhante a esta:

  Workflow failed. Causes: The Dataflow job appears to be stuck because no worker
  activity has been seen in the last 25m.

A causa desses problemas pode estar relacionada a problemas de capacidade ou de inicialização do worker.

  • Capacidade: se você usar a capacidade de TPU sob demanda ou uma reserva esgotada, novos pipelines poderão não ser iniciados até que a capacidade esteja disponível. Se você usa uma reserva, verifique a capacidade restante na página "Reservas de computação" no console doGoogle Cloud ou com o seguinte comando:

    gcloud compute reservations describe RESERVATION_NAME --zone ZONE
    

    Verifique se o job iniciou alguma VM worker. Quando seu job inicia um worker, loggers como worker, worker_startup, kubelet e outros geralmente fornecem saída. Além disso, na página Métricas do job no console Google Cloud , o número de workers atuais precisa ser maior que zero.

  • Inicialização do worker: verifique os registros job-message e launcher. Se o pipeline iniciar workers, mas eles não puderem ser inicializados, talvez haja erros no contêiner personalizado.

  • Espaço em disco: verifique se há espaço suficiente disponível para seu job. Para aumentar o espaço em disco, use a opção --disk_size_gb.

O job falha com um erro

Use as dicas de solução de problemas a seguir quando o job falhar com um erro.

Falha ao iniciar o pool de workers

Se você receber o erro a seguir, verifique se o pipeline especifica --worker_zone e se a zona corresponde à zona da sua reserva.

JOB_MESSAGE_ERROR: Startup of the worker pool in zone ZONE failed to
bring up any of the desired 1 workers. [...] INVALID_FIELD_VALUE:
Instance 'INSTANCE_NAME' creation failed: Invalid value for field
'resource.reservationAffinity': '{ "consumeReservationType":
"SPECIFIC_ALLOCATION", "key":
"compute.googleapis.com/RESERVATION_NAME...'. Specified reservations
[RESERVATION_NAME] do not exist.

Os grupos gerenciados de instâncias não são compatíveis com as Cloud TPUs

Se você encontrar o seguinte erro, entre em contato com sua equipe de contas para verificar se o projeto foi registrado para usar TPUs ou registre um bug usando o Google Issue Tracker.

apache_beam.runners.dataflow.dataflow_runner.DataflowRuntimeException: Dataflow
pipeline failed. State: FAILED, Error: Workflow failed. Causes: One or more
operations had an error [...]: [INVALID_FIELD_VALUE] 'Invalid value
for field 'resource.instanceTemplate': Managed Instance Groups do not support
Cloud TPUs. '.

Valor inválido para o campo

Se você encontrar o seguinte erro, verifique se a invocação do pipeline define a opção de serviço worker_accelerator do Dataflow.

JOB_MESSAGE_ERROR: Workflow failed. Causes: One or more operations had an error:
'operation-[...]': [INVALID_FIELD_VALUE] 'Invalid value for field
'resource.instanceTemplate': 'projects/[...]-harness'. Regional
Managed Instance Groups do not support Cloud TPUs.'

Dispositivo ou recurso ocupado

Se você encontrar o seguinte erro, um worker do Dataflow que processa seu pipeline provavelmente está executando mais de um processo que acessa a TPU ao mesmo tempo. Isso não é permitido. Para mais informações, consulte TPUs e paralelismo de worker.

RuntimeError: TPU initialization failed: open(/dev/vfio/0): Device or resource
busy: Device or resource busy; Couldn't open iommu group /dev/vfio/0

Se você encontrar o erro anterior ao depurar seu pipeline em uma VM, inspecione e encerre o processo que está retendo a TPU usando os seguintes comandos:

apt update ; apt install lsof
lsof -w /dev/vfio/0
kill -9 PROCESS_ID    # to terminate the process.

Instâncias com aceleradores convidados não são compatíveis com a migração em tempo real

Se você encontrar o seguinte erro, provavelmente o pipeline foi iniciado com um tipo de máquina definido explicitamente que tem aceleradores, mas não especificou a configuração do acelerador corretamente. Verifique se a invocação do pipeline define a opção do serviço worker_accelerator do Dataflow e se o nome da opção não contém erros de digitação.

JOB_MESSAGE_ERROR: Startup of the worker pool in zone ZONE failed to
bring up any of the desired 1 workers. [...] UNSUPPORTED_OPERATION:
Instance INSTANCE_ID creation failed: Instances with guest
accelerators do not support live migration.

O fluxo de trabalho foi rejeitado automaticamente pelo serviço

Os erros a seguir também podem aparecer se algumas das opções de pipeline obrigatórias estiverem ausentes ou incorretas:

The workflow was automatically rejected by the service. The requested
accelerator type tpu-v5-lite-podslice;topology:1x1 requires setting
the worker machine type to ct5lp-hightpu-1t. Learn more at:
https://cloud.google.com/dataflow/docs/guides/configure-worker-vm

O tempo para a atualização do worker expirou

Se você iniciar pipelines em VMs de TPU com muitas vCPUs e não reduzir o número padrão de linhas de execução de worker, o job poderá encontrar erros como os seguintes:

Workflow failed. Causes WORK_ITEM failed.
The job failed because a work item has failed 4 times.
Root cause: Timed out waiting for an update from the worker.

Para evitar esse erro, reduza o número de linhas de execução. Por exemplo, você pode definir: --number_of_worker_harness_threads=50.

Nenhum uso de TPU

Se o pipeline for executado com sucesso, mas os dispositivos TPU não forem usados ou não estiverem acessíveis, verifique se os frameworks que você está usando, como JAX ou PyTorch, podem acessar os dispositivos conectados. Para resolver problemas com sua imagem de contêiner em uma única VM, consulte Depurar com uma VM independente.