排查 Pathways on Cloud 问题

本页面可帮助机器学习 (ML) 工程师和平台管理员诊断和解决 Google Cloud上的 Pathways 工作负载问题。使用本指南检查组件日志、验证集群和切片就绪情况、监控 Megascale XLA (MXLA) 传输指标,以及修复常见的运行时错误。在每个部分中,选择适用于您环境的 GKE 或 Cluster Director 标签页。

准备工作

准备工作

GKE

在继续操作之前,请先满足以下前提条件:

Cluster Director

在继续操作之前,请先满足以下前提条件:

查看日志

GKE

您可以使用 gcluster job logs 命令直接查看 Pathways 工作负载的日志(传递 --main-only=false 以包含工作器 Pod 日志):

gcluster job logs WORKLOAD_NAME \
    --cluster=CLUSTER \
    --project=PROJECT_ID \
    --location=LOCATION \
    --main-only=false

替换以下内容:

  • WORKLOAD_NAME:使用 Cluster Toolkit 时的工作负载名称,或使用 kubectl 时的 JobSet 名称。
  • CLUSTER:GKE 集群的名称。
  • PROJECT_ID:您的Google Cloud 项目 ID。
  • LOCATION:您在其中创建 GKE 集群的区域或可用区。

如需使用 kubectl 列出工作负载的所有 Pod,请按 JobSet 名称标签进行过滤:

kubectl get pods \
    -l jobset.sigs.k8s.io/jobset-name=WORKLOAD_NAME

如需检查工作负载状态或列出失败的作业,请使用 gcluster job list 命令:

gcluster job list --status Failed \
    --cluster=CLUSTER \
    --project=PROJECT_ID \
    --location=LOCATION

或者,在 Cloud Logging Logs Explorer 中,运行以下查询,并根据您的项目、位置、集群和工作负载进行调整:

resource.type="k8s_container"
resource.labels.project_id="PROJECT_ID"
resource.labels.location="LOCATION"
resource.labels.cluster_name="CLUSTER"
resource.labels.namespace_name="default"
resource.labels.pod_name:"WORKLOAD_NAME"

此查询会匹配多个 Pathways Kubernetes 容器,其名称包括 pathways-rm、pathways-proxy、pathways-worker 和 workload-container。如需确定哪个容器导致了问题,请添加容器名称过滤条件,例如 resource.labels.container_name:"CONTAINER_NAME",其中 CONTAINER_NAME 是容器的名称。

Cluster Director

Pathways 脚本会将日志写入共享文件系统:

  • Pathways 服务作业会在您运行 sbatch 的目录(例如 ~/pathways)中写入 pathways-service-SERVICE_JOB_ID.out 和 pathways-service-SERVICE_JOB_ID.err。
  • Pathways 资源管理器、IFRT 代理服务器和 TPU 工作进程会写入 ~/pathways/service_SERVICE_JOB_ID/logs,包括 rm.log、proxy.log 以及每个工作进程任务的一个 slice-*.out 和 slice-*.err 文件。
  • 示例 MaxText 工作负载会将其日志写入 maxtext-llama31-JOB_ID.err。

在登录节点上,检查 Pathways 作业的状态并找到其作业 ID:

squeue -u ${USER}

监控 Pathways 组件和指标

以下部分适用于 GKE 上的 Pathways。

运行状况监控

您可以通过在容器日志中查找条目来监控各种 Pathways 组件的健康状况。

pathways-proxy 容器在写入以下日志条目时,已准备好处理新的连接请求。如需查看日志,请运行以下命令:

kubectl logs HEAD_POD_NAME --container pathways-proxy

将 HEAD_POD_NAME 替换为 Pathways 头 Pod 的名称。

输出包含类似于以下内容的一行:

I1101 04:51:41.967764       1 proxy_server.cc:125] IFRT proxy server started with status OK

当 pathways-rm 容器写入以下日志条目时,它已准备好处理新的连接请求:

kubectl logs HEAD_POD_NAME --container pathways-rm

输出包含类似于以下内容的一行:

I1101 04:50:41.967764       1 server_lib.cc:1473] Pathways Server serving on POD_IP:29001

POD_IP 是 Pod 的 IP 地址,在 Pod 重启或重新调度时可能会发生变化。

如需验证所有 TPU 切片是否已准备就绪,请检查 pathways-rm 容器日志中是否有 Pathways Slices Now Ready:

kubectl logs HEAD_POD_NAME --container pathways-rm

例如,以下日志条目显示 2 个切片中有 2 个已准备就绪:

I1101 04:52:41.967764       1 multi_job_allocator.cc:1063] *** 2/2 Pathways Slices Now Ready ***

Pathways 客户端可以在 IFRT 代理服务器准备就绪后立即连接到该服务器,即使并非所有切片都已准备就绪。在 TPU 切片准备就绪之前,您的作业会使用虚拟切片。虚拟切片可让您的代码在 TPU 不可用时运行。

如需验证 pathways-worker 容器是否已准备就绪,请检查其日志:

kubectl logs WORKER_POD_NAME --container pathways-worker

将 WORKER_POD_NAME 替换为 Pathways 工作器 Pod 的名称。

工作器准备就绪后,输出会包含类似于以下内容的一行:

I1101 04:50:41.967764       1 server_lib.cc:1473] Pathways Server serving on POD_IP:29005

POD_IP 是 Pod 的 IP 地址,在 Pod 重启或重新调度时可能会发生变化。

指标收集

Pathways 可以将低级系统指标写入 Cloud Monitoring 以进行调试。这些指标包括:

  • 数据中心网络 (DCN) 传输延迟时间
  • 集合通信延迟时间
  • 主机到设备传输延迟时间
  • 设备到主机传输延迟时间

您可以在 Google Cloud 运行 GKE 集群的项目的 Monitoring 信息中心内找到这些指标。

如需在 Metrics Explorer 中监控这些指标,请按照以下步骤操作:

  1. 在 Google Cloud 控制台中,前往 Metrics Explorer 页面。
  2. 在选择指标字段中,按指标名称进行过滤。
  3. 如需按工作负载名称和时间范围进行过滤,请点击添加过滤条件,然后按 pod_name 进行过滤。
  4. 对于汇总,请根据您要监控的指标和工作负载选择汇总类型。

DCN 传输延迟时间

此 Megascale XLA (MXLA) 指标用于衡量多切片流量的网络传输延迟时间的累积分布。延迟时间测量从系统发出通过 DCN 传输数据的请求时开始,到收到数据传输已完成的确认时结束。如需监控此指标,请按 dcn_transfer_latencies 指标名称进行过滤。

集合通信延迟时间

此 MXLA 指标用于衡量多切片流量的端到端总体延迟时间的累积分布。延迟时间测量从系统发出集合操作请求时开始,到系统收到数据传输已完成的确认时结束。如需监控此指标,请按 collective_end_to_end_latencies 指标名称进行过滤。

主机到设备传输延迟时间

此 MXLA 指标用于衡量多切片流量的主机到设备传输延迟时间的累积分布。延迟时间测量从系统发出将数据从主机传输到 TPU 设备的请求时开始,到收到数据传输已完成的确认时结束。如需监控此指标,请按 host_to_device_transfer_latencies 指标名称进行过滤。

设备到主机传输延迟时间

此 MXLA 指标用于衡量多切片流量的设备到主机传输延迟时间的累积分布。延迟时间测量从系统发出将数据从 TPU 设备传输到主机的请求时开始,到收到数据传输已完成的确认时结束。如需监控此指标,请按 device_to_host_transfer_latencies 指标名称进行过滤。

调试常见错误

GKE

无法对加速器配置警告进行哈希处理

以下消息只是警告,不会影响 JAX 代码的性能:

INFO:jax._src.cache_key:get (_hash_accelerator_config): unable to hash accelerator config, falling back to hashing devices + platform: UNIMPLEMENTED: GetTopologyForDevices is not supported for the IFRT proxy client. (type <class 'jaxlib.xla_extension.XlaRuntimeError'>)

对资源的权限 logging.logEntries.create 被拒绝

如果您看到以下错误,请确保 Vertex AI Workbench 使用的 Compute Engine 服务账号有权将日志条目写入 Logging:

INFO:absl:Created 'ArrayHandler' with primary_host=0, replica-id=0
WARNING:absl:pathwaysutils: Detected Pathways-on-Cloud backend. Applying changes.
Failed to submit 1 logs.
Traceback (most recent call last):
  File "/opt/conda/lib/python3.10/site-packages/google/api_core/grpc_helpers.py", line 65, in error_remapped_callable
    return callable_(**args, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/grpc/_channel.py", line 1181, in __call__
    return _end_unary_response_blocking(state, call, False, None)
  File "/opt/conda/lib/python3.10/site-packages/google/api_core/grpc_helpers.py", line 1006, in _end_unary_response_blocking
    raise _InactiveRpcError(state) # pytype: disable-not-instantiable
grpc._channel._InactiveRpcError: <_InactiveRpcError of RPC that terminated with:
  status = StatusCode.PERMISSION_DENIED
  details = "Permission 'logging.logEntries.create' denied on resource (or it may not exist)."
  debug_error_string = "UNKNOWN:Error received from peer ipv4:216.239.34.174:443 {created_time:"2024-10-03T20:30:44.820425276+00.00", grpc_status:7, grpc_message:"Permission \'logging.logEntries.create\' denied on resource (or it may not exist)."}"
>

如需解决此问题,请向 Compute Engine 服务账号授予 Logs Writer (roles/logging.logWriter) 角色。

无法连接到 IFRT 代理服务器

当您尝试运行向 Pathways 集群发送 IFRT 命令的第一个 JAX 命令时,该命令可能会停止响应大约一分钟,然后显示类似于以下内容的 RuntimeError:

RuntimeError: Unable to initialize backend 'proxy': UNAVAILABLE: Unable to establish connection to ifrt_proxy server, please check provided address example-workload-proxy-0-0.example-workload.default.svc.example-cluster-domain.:38676'; detailed error: DNS resolution failed (set JAX_PLATFORMS='' to automatically choose an available backend)

如果您看到此错误,则表示您的 IFRT 代理客户端无法连接到 IFRT 代理服务器。请检查以下各项:

  • 验证您是否正确配置了 Virtual Private Cloud (VPC) 网络。
  • 验证您的防火墙规则是否允许连接。
  • 验证 Pathways 集群是否已预配。
  • 在容器日志中查看启动出错提示。

连接到 IFRT 代理服务器时,ALTS 握手失败或套接字已关闭

如果您在 Workload Identity Federation for GKE 节点池运行不含 hostNetwork: true 字段的自定义清单,则可能会看到 Alts handshake failed (TSI_INTERNAL_ERROR) 或 UNAVAILABLE: Unable to establish connection to ifrt_proxy server ... Socket closed。如需解决此问题,请在 pathways-proxy 容器和 JAX 客户端容器上设置 IFRT_PROXY_USE_INSECURE_GRPC_CREDENTIALS: "true" 环境变量(两端必须匹配)。添加此环境变量会停用 ALTS,并改用不安全的 gRPC,因此请仅在受信任的集群网络中使用。

已与 Pathways 集群建立连接

Pathways 集群一次只能与一个客户端保持会话。如果两个单独的笔记本尝试连接到同一个 Pathways 集群,则一个笔记本会连接,另一个笔记本会显示以下错误:

INFO:absl:Created ArrayHandler with primary_host=0, replica_id=0
WARNING:absl:pathwaysutils: Detected Pathways-on-Cloud backend. Applying changes.
E0927 21:19:52.919607   37624 rpc_helper.cc:56] Connection to IFRT proxy server was terminated: CANCELLED: Cancelled
E0927 21:20:03.467547   37719 rpc_helper.cc:56] Connection to IFRT proxy server was terminated: CANCELLED: Cancelled
E0927 21:20:14.011645   37807 rpc_helper.cc:56] Connection to IFRT proxy server was terminated: CANCELLED: Cancelled
E0927 21:20:24.557955   37924 rpc_helper.cc:56] Connection to IFRT proxy server was terminated: CANCELLED: Cancelled
---------------------------------------------------------------------------
XlaRuntimeError                           Traceback (most recent call last)
File /opt/conda/lib/python3.10/site-packages/jax/_src/xla_bridge.py:887, in backends()
    885   continue
--> 887 backend = _init_backend(platform)
    888 _backends[platform] = backend

File /opt/conda/lib/python3.10/site-packages/jax/_src/xla_bridge.py:973, in _init_backend(platform)
    972 logger.debug("Initializing backend '%s'", platform)
--> 973 backend = registration.factory()
    974 # TODO: consider raising more descriptive errors directly from backend
    975 # factories instead of returning None.

File /opt/conda/lib/python3.10/site-packages/pathwaysutils/proxy_backend.py:24, in register_backend_factory.<locals>.<lambda>()
     21 def register_backend_factory():
     22   xla_bridge.register_backend_factory(
     23       "proxy",
---> 24       lambda: ifrt_proxy.get_client(
     25           jax.config.read("jax_backend_target"),
     26           ifrt_proxy.ClientConnectionOptions(),
     27       ),
     28       priority=-1,
     29   )

XlaRuntimeError: UNAVAILABLE: Unable to establish connection to ifrt_proxy server, please check provided address example-workload-proxy-0-0.example-workload.default.svc.example-cluster-domain.:38676'; detailed error: Connection to IFRT proxy server was terminated: CANCELLED: Cancelled

During handling of the above exception, another exception occurred:

RuntimeError                              Traceback (most recent call last)
Cell In[2], line 4
      1 import pathwaysutils
      3 import jax
----> 4 print(jax.devices())

File /opt/conda/lib/python3.10/site-packages/jax/_src/xla_bridge.py:1085, in devices(backend)
   1060 def devices(
   1061     backend: str | xla_client.Client | None = None
   1062 ) -> list[xla_client.Device]:
   1063   """Returns a list of all devices for a given backend.
   1064
   1065   .. currentmodule:: jaxlib.xla_extension
   (...)
   1083     List of Device subclasses.
   1084   """
-> 1085   return get_backend(backend).devices()

File /opt/conda/lib/python3.10/site-packages/jax/_src/xla_bridge.py:1019, in get_backend(platform)
   1015 @lru_cache(maxsize=None)  # don't use util.memoize because there is no X64 dependence.
   1016 def get_backend(
   1017     platform: None | str | xla_client.Client = None
   1018 ) -> xla_client.Client:
-> 1019   return _get_backend_uncached(platform)

File /opt/conda/lib/python3.10/site-packages/jax/_src/xla_bridge.py:998, in _get_backend_uncached(platform)
    994   return platform
    996 platform = (platform or _XLA_BACKEND.value or _PLATFORM_NAME.value or None)
--> 998 bs = backends()
    999 if platform is not None:
   1000   platform = canonicalize_platform(platform)

File /opt/conda/lib/python3.10/site-packages/jax/_src/xla_bridge.py:903, in backends()
    901       else:
    902         err_msg += " (you may need to uninstall the failing plugin package, or set JAX_PLATFORMS=cpu to skip this backend.)"
--> 903       raise RuntimeError(err_msg)
    905 assert _default_backend is not None
    906 if not config.jax_platforms.value:

RuntimeError: Unable to initialize backend 'proxy': UNAVAILABLE: Unable to establish connection to ifrt_proxy server, please check provided address 'example-workload-proxy-0-0.example-workload.default.svc.example-cluster-domain.:38676'; detailed error: Connection to IFRT proxy server was terminated: CANCELLED: Cancelled (set JAX_PLATFORMS='' to automatically choose an available backend)

在原始客户端断开连接后,第二个客户端可以连接。发生意外断开连接后,您可能需要重启 Pathways 集群,以便其他客户端可以连接。

LocalProxy.__init__() 收到意外的关键字参数 'unbound_message'

如果您在导入 pathwaysutils 后看到以下错误,请检查您的环境中是否安装了过时的 Flask 或 Werkzeug 版本(如果需要,请将 pip3 替换为 pip):

pip3 list --outdated

如果列出了 Flask 或 Werkzeug,请升级软件包(请注意,升级可能会影响项目中的其他软件包或依赖项):

pip3 install --upgrade flask Werkzeug

代理服务器出现内部错误

"Internal error from proxy server during Array::IsDeleted(): UNAVAILABLE:Connection to IFRT proxy server was terminated: FAILED_PRECONDITION:
GrpcClientSession: writes no longer allowed."

此错误表示 IFRT 代理服务器与客户端断开连接。如需解决此问题,请重启客户端。对于笔记本,请重启笔记本内核并重新运行笔记本。

SIGTERM 和 HBM OOM

日志中与 RESOURCE_EXHAUSTED 错误关联的 SIGTERM 条目可能表示高带宽内存 (HBM) 内存不足 (OOM) 错误。在这种情况下,请减少 JAX 代码使用的 HBM 量。

INVALID_ARGUMENT

"INVALID_ARGUMENT : Permanent error, with a last message of Lifecycle
matches_prefix cannot specify more than 50 prefixes per config.; Error while
initializing persistent cache storage Cloud Storage"

如果传递给 --pathways-gcs-location 标志的 Cloud Storage 存储桶已达到生命周期政策上限(50 个前缀),则会发生此错误。如需解决此问题,请删除不再使用的 Cloud Storage 生命周期政策。

永久性错误

Permanent error, with a last message of The specified bucket does not exist.; Error while initializing persistent cache storage gcs

当您向 Pathways 容器提供无效的 Cloud Storage 位置时,Pathways 资源管理器或 Pathways 工作器会输出此错误。

守护程序返回错误响应

Error response from daemon: dockerfile parse error line 16: Unknown flag: exclude

如果您使用的是旧版 Docker,则会出现此错误。升级 Docker 版本。

IFRT 代理客户端和服务器未能达成一致

IFRT Proxy client and server failed to agree on the protocol version; supported versions: client = [1, 1], server = [3, 14]

此错误表示您使用的是旧版 MaxText。重新构建最新的 MaxText 映像。

Cluster Director

setup_env.sh 失败并显示 Enroot mktemp 错误

如果 setup_env.sh 失败并显示以下错误,则表示 Enroot 的临时目录不存在:

mktemp: failed to create directory via template '/tmp/UID/tmp/enroot...'

如需解决此问题,请运行 mkdir -p /tmp/$(id -u)/tmp,然后再次运行该脚本。

setup_env.sh 无法拉取 Pathways 映像

如果 setup_env.sh 失败并显示 Couldn't pull ...,则表示没有为请求的 JAX_VERSION 标记 Pathways 映像。使用受支持的 JAX 版本。

集群没有 cpu 分区

如果 sinfo 未显示 cpu 分区,则集群定义缺少 CPU 节点集或其分区名称不同。Pathways 脚本会将作业提交到名为 cpu 的分区。如需解决此问题,请更新集群蓝图以定义名为 cpu 的 CPU 分区并重新部署集群,或者更新 Pathways 作业脚本中的分区名称,使其与 sinfo 显示的 CPU 分区匹配。

TPU 节点未启动

检查预留是否位于集群的可用区中,是否有足够的 ct6e-standard-4t 虚拟机,以及 staticNodeCount 和 acceleratorTopology 是否匹配。

Pathways 服务退出,并显示 Invalid partition name specified

PARTITION 变量有误。从 sinfo 的输出中选择一个 TPU 分区。

Pathways 服务未准备就绪

如果服务日志保持在 Waiting for at least MIN_SLICES=...,则表示运行的切片数量不足。以下命令的 NODELIST(REASON) 列显示了原因,例如,没有可用容量或 SLICE_SHAPE 值与 TPU 节点集不匹配:

squeue -u ${USER} -n pw-slice

Pathways 资源管理器或 IFRT 代理服务器无法启动

如果服务失败并显示 Timed out waiting for Pathways Resource Manager 或 Timed out waiting for Pathways IFRT Proxy Server,则表示 Pathways 容器未启动。检查日志目录中的 rm.log 和 proxy.log。

Pathways 服务意外停止

如果服务因 The Pathways Resource Manager or IFRT proxy exited 而停止,则表示其中一个工作器崩溃,服务停止了其工作器。检查日志目录中的 rm.log 和 proxy.log,然后再次启动服务。

TPU 节点卡在错误状态

如需恢复节点,请在登录节点上运行以下命令:

sudo scontrol update nodename=NODE \
    state=power_down_force

将 NODE 替换为节点的名称,如 sinfo 所示。

MaxText 找不到正在运行的 Pathways 服务

如果 MaxText 失败并显示 No running Pathways service after 60 minutes,则表示该服务未运行或尚未准备就绪。查看服务日志。如需等待更长时间,请在提交 run_maxtext.sh 时设置 SERVICE_WAIT_MINUTES。

MaxText 失败并显示 RESOURCE_EXHAUSTED

模型无法放入 TPU 内存中。使用 2x4 或 4x4 切片,或设置 REDUCED_LAYERS=1。

MaxText 作业保持待处理状态

除了 Pathways 服务使用的资源之外,CPU 节点没有 run_maxtext.sh 请求的 8 个 CPU 和 96 GB 内存。如需解决此问题,请增加 cpu 分区中的 CPU 和内存容量(或节点数量),或减少 run_maxtext.sh 中的资源请求。

无法连接到 IFRT 代理服务器

如果客户端停止响应或失败并显示 UNAVAILABLE,则表示它无法访问代理服务器:

  • 检查服务日志是否显示 READY。
  • 检查 JAX_BACKEND_TARGET 是否与 ~/pathways/cluster.env 中的地址一致。
  • 如果您使用端口转发,请检查 SSH 隧道是否仍在运行。
  • 如果客户端在虚拟机上运行,请检查该虚拟机是否位于集群的子网中。以下命令应该会成功:

    nc -zv PROXY_IP 29000

    将 PROXY_IP 替换为服务日志中的 IFRT 代理 IP 地址。

jax.devices() 列出的是 CPU 设备,而不是 TPU 设备

您在导入 jax 后设置了 JAX_PLATFORMS 和 JAX_BACKEND_TARGET 环境变量。重启笔记本内核,并在导入 jax 之前设置变量。

连接到 IFRT 代理服务器后立即出现的错误

客户端的 jax 或 pathwaysutils 版本与 Pathways 服务器不匹配。安装 ~/pathways/venv 中的版本。

计算失败,并显示 DATA_LOSS 和 PlacementDataLoss

计算使用了未运行的切片上的设备。检查 MIN_SLICES 是否等于 MAX_SLICES,以及每个工作器任务是否都在运行:

squeue -u ${USER} -n pw-slice

Spot 抢占后计算失败

使用被抢占切片的设备的计算会失败。如需从 Spot 抢占中自动恢复,请按照使用 Pathways 进行弹性训练中的说明配置检查点和弹性切片重新扩缩。

后续步骤