排查 Dataflow GPU 作业问题

如果您在使用 GPU 运行 Dataflow 作业时遇到问题，请按以下步骤操作：

遵循使用 Dataflow GPU 的最佳实践中的工作流，确保您的流水线配置正确。
确认您的 Dataflow 作业是否使用了 GPU。如需了解详情，请参阅“使用 GPU 运行流水线”中的验证 Dataflow 作业。
使用独立虚拟机或 Dataflow 调试作业。
如果问题仍然存在，请按照本页面上的其余问题排查步骤操作。

调试作业

如果可能，请使用独立虚拟机调试作业，因为使用独立虚拟机调试通常速度更快。不过，如果组织政策阻止您使用独立虚拟机进行调试，您可以使用 Dataflow 进行调试。

使用独立虚拟机进行调试

在设计和迭代适合您的容器映像时，通过在独立虚拟机上试用容器映像可以更快地缩短反馈环。

您可以在具有 GPU 的独立虚拟机上调试自定义容器，方法是，创建在 Container-Optimized OS 上运行 GPU 的 Compute Engine 虚拟机、安装驱动程序并如下所示启动容器。

创建虚拟机实例。

gcloud compute instances create INSTANCE_NAME \
  --project "PROJECT" \
  --image-family cos-stable \
  --image-project=cos-cloud  \
  --zone=us-central1-f \
  --accelerator type=nvidia-tesla-t4,count=1 \
  --maintenance-policy TERMINATE \
  --restart-on-failure  \
  --boot-disk-size=200G \
  --scopes=cloud-platform

使用 ssh 连接到虚拟机。

gcloud compute ssh INSTANCE_NAME --project "PROJECT"

安装 GPU 驱动程序。使用 ssh 连接到虚拟机后，请在虚拟机上运行以下命令：

# Run these commands on the virtual machine
cos-extensions install gpu
sudo mount --bind /var/lib/nvidia /var/lib/nvidia
sudo mount -o remount,exec /var/lib/nvidia
/var/lib/nvidia/bin/nvidia-smi

启动自定义容器。

Apache Beam SDK 容器使用 /opt/apache/beam/boot 入口点。出于调试目的，您可以使用其他入口点手动启动容器：

docker-credential-gcr configure-docker
docker run --rm \
  -it \
  --entrypoint=/bin/bash \
  --volume /var/lib/nvidia/lib64:/usr/local/nvidia/lib64 \
  --volume /var/lib/nvidia/bin:/usr/local/nvidia/bin \
  --privileged \
  IMAGE

将 IMAGE 替换为 Docker 映像的 Artifact Registry 路径。

验证容器中安装的 GPU 库是否可以访问 GPU 设备。

如果您使用的是 TensorFlow，则可以使用以下命令输出 Python 解释器中的可用设备：
```
>>> import tensorflow as tf
>>> print(tf.config.list_physical_devices("GPU"))
```
如果您使用的是 PyTorch，则可以使用以下命令检查 Python 解释器中的可用设备：
```
>>> import torch
>>> print(torch.cuda.is_available())
>>> print(torch.cuda.device_count())
>>> print(torch.cuda.get_device_name(0))
```

如需迭代流水线，您可以在 Direct Runner 上启动流水线。您还可以在此环境中启动 Dataflow Runner 上的流水线。

使用 Dataflow 进行调试

如果组织限制条件导致您无法在独立虚拟机上进行调试，则可以使用 Dataflow 进行调试。

简化流水线，使其仅检测是否存在 GPU，然后在 Dataflow 上运行该流水线。以下示例展示了此流水线的代码可能是什么样子：

def check_if_gpus_present(element):
  import torch
  import tensorflow as tf

  tensorflow_detects_gpus = tf.config.list_physical_devices("GPU")
  torch_detects_gpus = torch.cuda.is_available()
  if tensorflow_detects_gpus and torch_detects_gpus:
    return element

  if tensorflow_detects_gpus:
    raise Exception('PyTorch failed to detect GPUs with your setup')
  if torch_detects_gpus:
    raise Exception('Tensorflow failed to detect GPUs with your setup')
  raise Exception('Both Tensorflow and PyTorch failed to detect GPUs with your setup')

with beam.Pipeline() as p:
  _ = (p | beam.Create([1,2,3]) # Create a PCollection of the prompts.
         | beam.Map(check_if_gpus_present)
  )

如果流水线成功，您的代码便可访问 GPU。如需识别有问题的代码，请逐步将越来越大的示例插入流水线代码，并在每次更改后运行流水线。

如果您的流水线无法检测到 GPU，请按照本文档无 GPU 用量部分中的步骤操作。

工作器无法启动

如果您的作业卡住，并且 Dataflow 工作器一直不开始处理数据，则可能是使用自定义容器和 Dataflow 时出现问题。如需了解详情，请参阅自定义容器问题排查指南。

如果您是 Python 用户，请验证是否满足以下条件：

容器映像中的 Python 解释器次要版本与您启动流水线时使用的版本相同。如果存在不匹配现象，您可能会看到类似 SystemError: unknown opcode 的错误，并且堆栈轨迹涉及 apache_beam/internal/pickler.py。
如果您使用的是 Apache Beam SDK 2.29.0 或更低版本，则必须能够在 /usr/local/bin/pip 中的映像上访问 pip。

我们建议您在第一次使用自定义映像时，将自定义设置降至最低有效配置。使用本页面的示例中提供的示例自定义容器映像。确保无需请求 GPU 即可使用此容器映像运行简单的 Dataflow 流水线。然后，迭代解决方案。

验证工作器是否有足够的磁盘空间来下载容器映像。如有必要，请调整磁盘大小。下载大型映像需要更长的时间，这会增加工作器启动时间。

作业在启动时立即失败

如果遇到 ZONE_RESOURCE_POOL_EXHAUSTED 或 ZONE_RESOURCE_POOL_EXHAUSTED_WITH_DETAILS 错误，则可以执行以下步骤：

请勿指定工作器可用区，以便 Dataflow 为您选择最佳可用区。
在其他可用区或使用其他加速器类型启动流水线。

作业在运行时失败

如果作业在运行时失败，请检查工作器机器和 GPU 上是否存在内存不足 (OOM) 错误。GPU OOM 错误可能会在工作器日志中表现为 cudaErrorMemoryAllocation out of memory 错误。如果您使用的是 TensorFlow，请验证只使用一个 TensorFlow 进程访问一个 GPU 设备。如需了解详情，请参阅 GPU 和工作器并行。

未使用 GPU

如果您的作业似乎未使用 GPU，请按照本文档调试作业部分中的步骤操作，以验证您的 Docker 映像是否可使用 GPU。

如果 GPU 可用但未使用，问题可能出在流水线代码上。如需调试流水线代码，请从一个成功使用 GPU 的简单流水线开始，然后逐步向流水线添加代码，并在添加每个新代码后测试流水线。如需了解详情，请参阅本文档中的在 Dataflow 中调试部分。

如果您的流水线无法检测到 GPU，请验证以下内容：

容器映像中安装的 NVIDIA 库符合流水线用户代码和它使用的库的要求。
容器映像中已安装的 NVIDIA 库可作为共享库访问。

如果这些设备不可用，则您使用的软件配置可能不兼容。如需验证映像配置，请运行一个简单明确的流水线，该流水线只会检查 GPU 是否可用并可供工作器访问。

排查 TensorFlow 问题

如果 PyTorch 检测到流水线中的 GPU，但 TensorFlow 未检测到，请尝试以下问题排查步骤：

验证您是否使用了 TensorFlow、cuDNN 版本和 CUDA 工具包版本的兼容组合。如需了解详情，请参阅 TensorFlow 文档中的经过测试的build 配置。
如果可能，请升级到最新的兼容 TensorFlow 和 CUDA 版本。
查看 TensorFlow 和 CUDA 的已知问题，以验证是否有已知问题导致您的流水线出现问题。例如，以下已知问题可能会导致 TensorFlow 无法检测到 GPU：TF 2.17.0 RC0 无法与 GPU 配合使用。