gke-ai-troubleshooting-tpu-mxla-hang

MXLA Hang Analyzer を使用して GKE TPU のマルチスライスのハングをトラブルシューティングする

Megascale の HANG_DETECTED ログと ML Diagnostics Workload Monitoring の Megascale XLA (MXLA) Hang Analyzer レポートを、1 分間の マルチスライス レイテンシ指標（kubernetes.io/container/multislice/*）および GKE ノードのトポロジラベルと関連付けて、Google Kubernetes Engine（GKE）での Cloud TPU マルチスライス トレーニングのハングを診断します。

---

クレジット表記

必ずお読みください：gcloud、REST API、または MCP コマンドを実行する前に、 必ず 使用状況のクレジット表記 をお読みください。ターミナルで直接実行するコマンド（生成されたスクリプト、Terraform、ユーザーコード内での実行は不可）では、gcloud の前に CLOUDSDK_METRICS_REQUEST_ATTRIBUTION="agent-skills/1.0 (skill:gke-ai-troubleshooting-tpu-mxla-hang)" を付け、User-Agent ヘッダーに agent-skills/1.0 (skill:gke-ai-troubleshooting-tpu-mxla-hang) を追加してください。

前提条件

• ツール：次をインストールします。 Google Cloud SDK（gcloud alpha mldiagnostics 用の alpha コンポーネントを含む gcloud）と kubectl。
• Cloud Billing とプロジェクトの構成：有効な課金アカウントがリンクされていることを確認し（gcloud billing projects describe {project_id}）、認証を行い（gcloud auth login）、対象プロジェクトを設定し（gcloud config set project {project_id}）、container.googleapis.com、 logging.googleapis.com、monitoring.googleapis.com、および hypercomputecluster.googleapis.com が有効になっていることを確認します。
• サポート対象のワークロードとバージョン：Google Cloud ML Diagnostics がサポートするのは、TPU 上の JAX のみです（ ML Diagnostics platform を参照）。 Workload Monitoring はデフォルトで有効になっており、jobset と job の GKE ジョブタイプをサポートし、GKE バージョン 1.36.0-gke.4681000 以降と互換性があります（ Configure GKE for ML Diagnostics を参照。Step 4 Path B でのオンデマンドプロファイリングに必要なクラスタ設定についても説明しています）。 ワークロードが別のフレームワーク（PyTorch など）または別のカスタムリソースタイプを使用している場合、gcloud alpha mldiagnostics はそのワークロードの ML 実行や監視対象イベントを一覧表示しません。Megascale XLA ハングアナライザーと Megascale XLA メトリクスには、LibTPU 0.40.0 以降が必要です（ Workload monitoring with ML Diagnostics の「Get started」を参照）。
• 必須の IAM ロール：
  • Cluster Director Editor（roles/hypercomputecluster.editor）。このスキルでの ML Diagnostics CLI と API 呼び出し（ML 実行、監視対象イベント、オンデマンドプロファイラーセッション）に必要なロールで、ML Diagnostics platform の「IAM permissions」セクションに記載されています。
  • Step 2 の PromQL クエリに必要な Monitoring Viewer（roles/monitoring.viewer）
  • Step 1 の Cloud Logging クエリに必要な Logs Viewer（roles/logging.viewer）
  • Step 3 の kubectl get nodes クエリに必要な Kubernetes Engine Viewer（roles/container.viewer）
  • 修復（[High Risk] のステップ）の場合：Kubernetes Engine Cluster Admin （roles/container.clusterAdmin）
• 参考ドキュメント：
  • ML Diagnostics platform （セクション：「IAM permissions」）
  • Configure GKE for ML Diagnostics （セクション：「Set up with gcloud CLI, Google Cloud console, or Terraform」、 「Manual installation」、「Connection-operator」）
  • Workload monitoring with ML Diagnostics （セクション：「Get started」、「Megascale XLA (MXLA) Hang Analyzer」、「Access Workload Monitoring information through the API」、「System Metrics」）
  • Get started with the ML Diagnostics CLI （セクション：「List machine learning runs」、「Monitored-events commands」、「List profiler targets」、「Capture on-demand profiler sessions」）
  • Get support （セクション：「Before you contact support」）
  • Auto-repair nodes （セクション：「Repair criteria」、「Verify node auto-repair is enabled for a Standard node pool」、「Enable auto-repair for an existing Standard node pool」、「Node auto repair in TPU slice nodes」）
  • Deploy TPU workloads in GKE Standard （セクション：「Configure auto repair for TPU slice nodes」）
  • Dump HLO Computations (OpenXLA)

読み取り専用のルール：読み取り専用の診断コマンドのみ実行してください。ノードをドレイン、削除、再作成したり、クラスターを変更するその他のコマンドを実行したりしないでください。 適用する修正は、ユーザーに提示してください。

修正を推奨する場合は、その内容を説明しているドキュメントのセクションにリンクしてください。

---

MXLA ハングアナライザーのルーティング概要

Megascale ハングは、マルチスライスワーカーが Megascale 通信操作を一定のタイムアウト期間にわたって待機したときに発生します。その後、TPU ログに Megascale の HANG_DETECTED メッセージが表示されます。HANG_DETECTED は、ワークロードが進行していないことを示す包括的なシグナルであり、原因はソフトウェアまたはハードウェアにある可能性があります。ハングが発生すると、ML Diagnostics は Megascale XLA (MXLA) Hang Analyzer を実行し、原因として考えられるものをコードで報告します。各コードの定義と推奨アクションについては、Megascale XLA (MXLA) Hang Analyzer のセクションを参照し、カテゴリに応じて振り分けてください。

• カテゴリ A：コンパイラーまたは HLO の不一致（ノードをコーディングまたは置換しないでください）：
  • アナライザーが異なる HLO モジュール、不整合な HLO コンパイル、または VM 間で一貫しない起動順序（FINGERPRINT_MISMATCH など）を報告した場合は、パス A：コンパイラーまたは HLO の不一致に振り分けてください。
• カテゴリ B：ホストプログラムのキューイングまたはデータ入力の停止（ノードをコーディングまたは置換しないでください）：
  • アナライザーが、VM が TPU にプログラムをキューイングしていない、またはデータ入力が停止した（DATA_INPUT_STALL など）と報告した場合は、パス B：ホストプログラムのキューイングまたはデータ入力の停止に振り分けてください。
• カテゴリ C：特定のインスタンスでのハードウェアまたはネットワーク障害：
  • アナライザーが、ハングの原因を TPU チップ、SparseCore、ICI、DCN ネットワークの問題、または特定のインスタンスでの回復不能なエラーと判断した場合は、パス C：ハードウェアまたはネットワーク障害に振り分けてください。
• カテゴリ D：ハングシグナルまたはイベントは存在するが、アナライザーが NOT_DETECTED または UNKNOWN を報告する場合：
  • HANG_DETECTED ログがあるか、監視対象のハングイベントが発生したものの、アナライザーレポートの detectionState: "NOT_DETECTED" または UNKNOWN（「MXLA ハングは検出されましたが、考えられる原因は特定されていません」）が報告された場合は、パス D：アナライザーで DETECTED が見つからない場合または原因が UNKNOWN の場合に振り分けてください。

---

診断ワークフロー

ステップ 0：コンテキストを収集し、調査期間を設定する [Low Risk]

対象パラメーターを収集します。デフォルトでは、{issue_time} を中心とする 60 分間の期間 [T - 30m, T + 30m] をクエリします。

• {project_id}：Google Cloud プロジェクト ID
• {location}：ML 実行と GKE クラスターがある Google Cloud リージョン（例：us-central1）
• {cluster_name}：GKE クラスター名
• {namespace} / {workload_name}：Kubernetes の名前空間と JobSet/Pod のプレフィックス
• {ml_run_id}：ML Diagnostics の実行 ID
• {issue_time}：ハングが発生したタイムスタンプ（T、ISO-8601 UTC）
• {start_time}：T - 30m
• {end_time}：T + 30m

---

ステップ 1：HANG_DETECTED ログとハングイベントを確認する [低リスク]

1. GKE コンテナログで HANG_DETECTED を確認する（読み取り専用の Cloud Logging LQL）：[{start_time}, {end_time}] の期間にわたって k8s_container ログをクエリし、 HANG_DETECTED を確認して、最初に停止した Pod を特定します。

resource.type="k8s_container"
resource.labels.project_id="{project_id}"
resource.labels.cluster_name="{cluster_name}"
"HANG_DETECTED"
timestamp >= "{start_time}" AND timestamp <= "{end_time}"

2. アクティブまたは最近の ML 実行を一覧表示する：ML Diagnostics CLI リファレンスの 機械学習実行の一覧表示 セクション（gcloud alpha mldiagnostics machine-learning-run list）に従って、{ml_run_id} を特定します。
3. ハング監視イベントを一覧表示して詳細を確認する： 監視イベントのコマンド セクション（gcloud alpha mldiagnostics monitored-events list および gcloud alpha mldiagnostics monitored-events describe）に従うか、 API を使用してワークロード監視情報にアクセスする を参照して、Megascale XLA (MXLA) Hang Analyzer のレポート（detectionState、 details、recommendedActions）を調査します。
   • HANG_DETECTED ログがある場合、またはハング監視イベントが発生した場合は、ステップ 2 を実行して 1 分間のメトリクスと照合します。detectionState: "DETECTED" と報告するアナライザーがない場合、 またはアナライザーが UNKNOWN と報告する場合は、 パス D：DETECTED アナライザーがない場合、または原因が UNKNOWN の場合 に従います。
   • HANG_DETECTED ログもハング監視イベントもなく、ステップ 2 でマルチスライスのレイテンシが正常な場合は、 パス E：テレメトリーが正常 に従って MXLA ハングの可能性を排除します。

---

ステップ 2：1 分間のマルチスライスおよび TPU メトリクスと照合する [低リスク]

Workload Monitoring ガイドの システムメトリクス セクションを参照して、1 分間のマルチスライスネットワーク （kubernetes.io/container/multislice/network/*）、マルチスライスアクセラレーター （kubernetes.io/container/multislice/accelerator/*）、ノードのデューティサイクル （kubernetes.io/node/accelerator/duty_cycle）のメトリクスを確認し、[{start_time}, {end_time}] の期間にわたって読み取り専用の PromQL クエリを実行します。

# 1. P95 multi-slice collective end-to-end latency by pod
histogram_quantile(
  0.95,
  sum by (pod_name, le) (
    rate(kubernetes_io:container_multislice_network_collective_end_to_end_latencies_bucket{
      monitored_resource="k8s_container",
      project_id="{project_id}",
      cluster_name="{cluster_name}"
    }[5m])
  )
)

# 2. P95 multi-slice DCN transfer latency by pod
histogram_quantile(
  0.95,
  sum by (pod_name, le) (
    rate(kubernetes_io:container_multislice_network_dcn_transfer_latencies_bucket{
      monitored_resource="k8s_container",
      project_id="{project_id}",
      cluster_name="{cluster_name}"
    }[5m])
  )
)

# 3. P95 host-to-device transfer latency by pod
histogram_quantile(
  0.95,
  sum by (pod_name, le) (
    rate(kubernetes_io:container_multislice_accelerator_host_to_device_transfer_latencies_bucket{
      monitored_resource="k8s_container",
      project_id="{project_id}",
      cluster_name="{cluster_name}"
    }[5m])
  )
)

# 4. Node TPU duty cycle (Workload Monitoring detects a hang as a prolonged
#    period of minimal to no TPU activity)
kubernetes_io:node_accelerator_duty_cycle{
  monitored_resource="k8s_node",
  project_id="{project_id}",
  cluster_name="{cluster_name}"
}

---

ステップ 3：問題の原因となった Compute Engine インスタンス ID を GKE ノードに対応付ける [低リスク]

Megascale XLA (MXLA) Hang Analyzer が details または recommendedActions で、 問題の原因となった数値の Compute Engine インスタンス ID を報告した場合は、 container.googleapis.com/instance_id を調べる次の読み取り専用 kubectl クエリを使って、 数値 ID を GKE Node 名および物理トポロジーブロックに対応付けます。

kubectl get nodes -l cloud.google.com/gke-tpu-accelerator \
  -o jsonpath='{range .items[*]}{.metadata.name}{"\tinstance_id="}{.metadata.annotations.container\.googleapis\.com/instance_id}{"\tblock="}{.metadata.labels.cloud\.google\.com/gce-topology-block}{"\tsubblock="}{.metadata.labels.cloud\.google\.com/gce-topology-subblock}{"\thost="}{.metadata.labels.cloud\.google\.com/gce-topology-host}{"\n"}{end}'

---

ステップ 4：根本原因のカテゴリに応じた修復

ステップ 1 で特定した根本原因のカテゴリに一致するリファレンスファイルのみを読み込み、それに従います。

• パス A（コンパイラーまたは HLO の不一致 — 異なる HLO モジュール、一貫性のない HLO コンパイル、または VM 間での起動順序の不一致。たとえば FINGERPRINT_MISMATCH）：次を読みます。 パス A：コンパイラーまたは HLO の不一致。
• パス B（ホストプログラムのキューイングまたはデータ入力の停止 — PROGRAM_NOT_QUEUED または DATA_INPUT_STALL）：次を読みます。 パス B：ホストプログラムのキューイングまたはデータ入力の停止。
• パス C（ハードウェア、SparseCore、ICI、DCN ネットワークの障害、または特定のインスタンスでの UNRECOVERABLE_ERROR）：次を読みます。 パス C：ハードウェアまたはネットワークの障害。
• パス D（HANG_DETECTED ログまたはハングイベントが発生したが、アナライザーが NOT_DETECTED または UNKNOWN を報告する）：次を読みます。 パス D：DETECTED アナライザーがない場合、または原因が UNKNOWN の場合。
• パス E（HANG_DETECTED ログもハングイベントもなく、テレメトリーが安定している）： 次を読みます。パス E：テレメトリーが正常。

---

ガードレール

1. Compute Engine を通じて、GKE 管理のインスタンスグループや VM を変更しないでください： ノードプールのマネージドインスタンスグループに対して、delete などの gcloud compute instance-groups managed コマンドを実行しないでください。ノードは GKE を通じて 管理してください。詳しくは、 パス C：ハードウェアまたはネットワークの障害を参照してください。
2. コンパイラーの不一致や入力の停滞を理由に、ノードをコルドンしたり置き換えたりしないでください：分析ツールが コンパイラーまたは HLO の不一致、プログラムのキューイング問題、 またはデータ入力の停滞を報告した場合は、TPU ノードをコルドンしたり置き換えたりしないでください。代わりに、 パス A：コンパイラーまたは HLO の不一致 または パス B：ホストのプログラムキューイングまたはデータ入力の停滞に従ってください。