Dataflow TPU ジョブのトラブルシューティング

TPU を使用した Dataflow ジョブの実行で問題が発生した場合は、次のトラブルシューティング手順に沿って問題を解決してください。

コンテナ イメージのトラブルシューティング

スタンドアロン VM でコンテナと TPU ソフトウェアをデバッグすると便利な場合があります。GKE ノードプールで作成された VM を使用してデバッグすることも、実行中の Dataflow ワーカー VM でデバッグすることもできます。

スタンドアロン VM でデバッグする

スタンドアロン VM でコンテナをデバッグするには、ローカル テストに同じ TPU VM を使用する GKE ノードプールを作成します。たとえば、us-west1-c に 1 つの TPU V5 Lite デバイスを含む GKE ノードプールを作成する場合、次のようになります。

  1. 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. 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. GKE UI または次のコマンドを使用して、ノードプールの TPU ノードの VM 名を見つけます。

    gcloud compute instances list --filter='metadata.kube-labels:"cloud.google.com/gke-nodepool=TPU_NODEPOOL_NAME"'
    
  4. SSH を使用して、GKE ノードプールによって作成された VM に接続します。

    gcloud compute ssh --zone "us-west1-c" "VM_NAME" --project PROJECT_ID
    
  5. SSH を使用して VM に接続したら、使用している Artifact Registry 用に Docker を構成します。

    docker-credential-gcr configure-docker --registries=us-west1-docker.pkg.dev
    
  6. 次に、使用するイメージからコンテナを起動します。

    docker run --privileged --network=host -it --rm --entrypoint=/bin/bash IMAGE_NAME
    
  7. コンテナ内で、TPU にアクセスできることをテストします。

    たとえば、PyTorch を使用して TPU を利用するイメージがある場合は、Python インタープリタを開きます。

    python3
    

    次に、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)
    

    出力例:

    >>> tensor([[ 0.3355, -1.4628, -3.2610],
    >>>        [-1.4656,  0.3196, -2.8766],
    >>>        [ 0.8667, -1.5060,  0.7125]], device='xla:0')
    
  8. 計算が失敗した場合は、イメージが正しく構成されていない可能性があります。

    たとえば、イメージの Dockerfile で必要な環境変数を設定する必要がある場合があります。確認するには、次のように環境変数を手動で設定してから、計算を再試行します。

    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.
    

    PyTorch または LibTPU の依存関係がない場合は、次のコマンドを使用してインストールしてから、計算を再試行できます。

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

Dataflow VM を使用してデバッグする

ジョブの実行中に SSH を使用して Dataflow ワーカー VM インスタンスに接続することもできます。Dataflow ワーカー VM はパイプラインの完了後にシャットダウンするため、長時間待機する計算を行うことで、実行時間を人工的に増やす必要がある場合があります。

TPU デバイスは複数のプロセス間で共有できないため、TPU で計算を行わないパイプラインを実行する必要がある場合があります。

  1. 実行中の TPU ジョブの VM を見つけるには、 Google Cloud コンソールの検索バーで Dataflow ジョブ ID を検索するか、次の gcloud コマンドを使用します。

    gcloud compute instances list --project PROJECT_ID --filter "STATUS='RUNNING' AND description ~ 'Created for Dataflow job: JOB_ID'"
    
  2. SSH を使用して TPU を含む VM に接続したら、使用するイメージからコンテナを起動します。例については、スタンドアロン VM でデバッグするをご覧ください。

  3. コンテナ内で、TPU 設定を再構成し、設定をテストするために必要なライブラリをインストールします。例については、スタンドアロン VM でデバッグするをご覧ください。

ワーカーが起動しない

トラブルシューティングを行う前に、次のパイプライン オプションが正しく設定されていることを確認します。

  • --dataflow_service_option=worker_accelerator オプション
  • --worker_zone オプション
  • --machine_type オプション

コンソールログにワーカーが起動していることが示されているにもかかわらず、ジョブが失敗し、次のようなメッセージが表示されるかどうかを確認します。

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

これらの問題の原因は、容量またはワーカーの起動の問題に関連している可能性があります。

  • 容量: オンデマンド TPU 容量または使用済みの予約を使用している場合、容量が使用可能になるまで新しいパイプラインが開始されないことがあります。予約を使用する場合は、Google Cloud コンソールの [コンピューティングの予約] ページで、または次のコマンドを使用して、残りの容量を確認します。

    gcloud compute reservations describe RESERVATION_NAME --zone ZONE
    

    ジョブがワーカー VM を起動したかどうかを確認します。ジョブがワーカーを起動すると、通常、workerworker_startupkubelet などのロガーが出力を提供します。また、 Google Cloud コンソールの [ジョブ指標] ページで、現在のワーカー数が 0 より大きいことを確認します。

  • ワーカーの起動: job-message ログと launcher ログを確認します。パイプラインがワーカーを起動しても起動できない場合は、カスタム コンテナにエラーがある可能性があります。

  • ディスク容量: ジョブに十分なディスク容量があることを確認します。ディスク容量を増やすには、--disk_size_gb オプションを使用します。

ジョブがエラーで失敗する

ジョブがエラーで失敗した場合は、次のトラブルシューティングのヒントを参考にしてください。

ワーカープールの起動に失敗しました

次のエラーが表示された場合は、パイプラインで --worker_zone が指定されていることと、ゾーンが予約のゾーンと一致していることを確認します。

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.

マネージド インスタンス グループは Cloud TPU をサポートしていない

次のエラーが表示された場合は、アカウント チームに連絡して、プロジェクトが TPU を使用するように登録されているかどうかを確認するか、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. '.

フィールドの値が無効です

次のエラーが表示された場合は、パイプライン呼び出しで worker_accelerator 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.'

デバイスまたはリソースがビジー状態

次のエラーが表示された場合、パイプラインを処理する Dataflow ワーカーが、同時に TPU にアクセスする複数のプロセスを実行している可能性があります。これはサポートされていません。詳細については、TPU とワーカーの並列処理をご覧ください。

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

VM でパイプラインをデバッグしているときに上記のエラーが表示された場合は、次のコマンドを使用して、TPU を保持しているプロセスを検査して終了できます。

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

ゲスト アクセラレータを使用するインスタンスはライブ マイグレーションをサポートしていない

次のエラーが表示された場合は、アクセラレータを含むマシンタイプが明示的に設定されてパイプラインが起動されたものの、アクセラレータ構成が正しく指定されていない可能性があります。パイプライン呼び出しで worker_accelerator Dataflow サービス オプションが設定されていることを確認し、オプション名に誤字脱字がないことを確認します。

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.

ワークフローがサービスによって自動的に拒否されました

必要なパイプライン オプションの一部がないか、正しくない場合、次のエラーも表示されることがあります。

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

ワーカーからの更新待機中にタイムアウトした

多くの vCPU を持つ TPU VM でパイプラインを起動し、ワーカー スレッドのデフォルト数を減らさないと、ジョブで次のようなエラーが発生する可能性があります。

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.

このエラーを回避するには、スレッドの数を減らします。たとえば、--number_of_worker_harness_threads=50 を設定できます。

TPU の使用なし

パイプラインは問題なく実行されていても、TPU デバイスが使用されていないか、アクセスできない場合は、使用しているフレームワーク(JAX や PyTorch など)が接続されたデバイスにアクセスできることを確認します。単一の VM でコンテナ イメージのトラブルシューティングを行うには、スタンドアロン VM でデバッグするをご覧ください。