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，并安装用于 gcloud alpha mldiagnostics 的 alpha 组件）和 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 平台）。 工作负载监控默认启用，支持 jobset 和 job GKE 作业类型，并兼容 1.36.0-gke.4681000 及更高版本的 GKE （参见 为 ML Diagnostics 配置 GKE， 其中也介绍了第 4 步路径 B 中按需性能分析所需的集群设置）。如果工作负载使用其他 框架（例如 PyTorch）或其他自定义资源类型，gcloud alpha mldiagnostics 将不会 列出该工作负载的 ML 运行或受监控事件。Megascale XLA 挂起分析器和 Megascale XLA 指标要求 LibTPU 版本为 0.40.0 或更高版本（参见 使用 ML Diagnostics 进行工作负载监控 中的“开始使用”）。
• 必需的 IAM 角色：
  • Cluster Director Editor（roles/hypercomputecluster.editor），这是 ML Diagnostics 平台 的“IAM 权限”部分列出的角色，用于此技能中的 ML Diagnostics CLI 和 API 调用 （ML 运行、受监控事件和按需分析器会话）
  • Monitoring Viewer（roles/monitoring.viewer），用于第 2 步中的 PromQL 查询
  • Logs Viewer（roles/logging.viewer），用于第 1 步中的 Cloud Logging 查询
  • Kubernetes Engine Viewer（roles/container.viewer），用于第 3 步中的 kubectl get nodes 查询
  • 执行修复操作时（[High Risk] 步骤）：Kubernetes Engine Cluster Admin （roles/container.clusterAdmin）
• 参考文档：
  • ML Diagnostics 平台 （章节：“IAM 权限”）
  • 为 ML Diagnostics 配置 GKE （章节：“使用 gcloud CLI、Google Cloud 控制台或 Terraform 进行设置”、 “手动安装”、“Connection-operator”）
  • 使用 ML Diagnostics 进行工作负载监控 （章节：“开始使用”、“Megascale XLA (MXLA) 挂起分析器”、“通过 API 访问工作负载 监控信息”、“系统指标”）
  • 开始使用 ML Diagnostics CLI （章节：“列出机器学习运行”、“受监控事件命令”、“列出分析器目标”、“捕获按需 分析器会话”）
  • 获取支持 （章节：“联系支持前”）
  • 自动修复节点 （章节：“修复条件”、“验证是否已为 Standard 节点池启用节点自动修复”、 “为现有 Standard 节点池启用自动修复”、“TPU 切片节点中的节点自动修复”）
  • 在 GKE Standard 中部署 TPU 工作负载 （章节：“为 TPU 切片节点配置自动修复”）
  • 转储 HLO 计算（OpenXLA）

只读规则：仅运行只读诊断命令。切勿驱逐、删除或重新创建节点，也不要运行任何其他会更改集群的命令。 请向用户提供可供他们自行应用的修复方案。

建议修复方案时，请链接到介绍该方案的文档章节。

---

MXLA Hang Analyzer 路由概览

当多切片工作器在 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：分析器未检测到或原因未知。

---

诊断工作流

步骤 0：收集上下文并设置调查时间范围 [Low Risk]

收集目标参数。默认查询 {issue_time} 前后各 30 分钟的 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 访问信息 Workload Monitoring 信息 检查 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 时，请使用以下只读 kubectl 查询检查 container.googleapis.com/instance_id，将这些数字 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： 不要运行 gcloud compute instance-groups managed 命令（例如 delete）来操作节点池的托管实例组。应按照 路径 C：硬件或网络故障中的说明，通过 GKE 管理节点。
2. 切勿因编译器分歧或输入停滞而隔离或替换节点：如果分析器报告编译器或 HLO 分歧、 程序排队问题或数据输入停滞，不要隔离或替换 TPU 节点。请改为遵循 路径 A：编译器或 HLO 分歧 或 路径 B：主机程序排队或数据输入停滞。