GKE 和 Filestore - 最多可将 AI / ML 工作负载的训练时间缩短 37%

问题陈述和应用场景

训练 AI /ML 工作负载需要大量数据，而这些数据通常存储在大量的小文件中。试想一下训练无人驾驶汽车，训练大量图片数据或进行蛋白质分析，在这些场景中，训练集通常由许多大小在 100K 到 2MB 之间的小文件组成。在为这些应用场景选择工具时，用户通常会使用 Google 的 Cloud Storage 服务（可提供低延迟和高吞吐量，且性价比合理），并且视情况选用 FUSE 作为文件接口，以实现可移植性。然而，如果数据集是由小文件组成，则会出现延迟问题；在训练工作负载时，每个周期可能都需要批量处理成千上万个小文件，并且可能会有多个工作节点同时访问 Cloud Storage。

为了缩短加载时间，用户需要可提供低延迟和高吞吐量的存储服务。使用 Filestore 作为“加速器”可提供帮助。Filestore 提供可快速访问的文件存储服务（具有多重读写权限的所有优势）以及一个原生 POSIX 接口。您仍然可以将 Cloud Storage 用作主要存储来源，并使用 Filestore 为工作节点提供经济实惠、延迟低的数据访问。

在本博文中，我们重点关注 Filestore 在训练 AI / ML 工作负载方面可发挥的重要作用，帮助您通过做出明智的选择来提升工作负载性能。请阅读下文，了解如何根据角色和职责来使用此解决方案：

使用情况详情

以下屏幕截图重点介绍了如何在 AI / ML 应用中使用 GKE 和 Filestore。您可以在此代码库中找到完整的源代码。

角色 1：Kubernetes 平台管理员，暂存 Filestore 以供数据科学家使用

Kubernetes 平台管理员负责创建供数据科学团队使用的基础设施。在本示例中，平台管理员使用 Kubernetes 永久性卷来设置 Filestore，并通过 Jupyter 笔记本设置让数据科学家可以使用该服务（若涉及多个用户，则通过 JupyterHub 提供访问权限）。这样，数据科学家就可以访问笔记本和编写代码了。

在这个具体示例中，我们使用了现成的 premium-rwx GKE StorageClass，可在后台动态预配 Filestore 基本 SSD 实例。Jupyter Pod 规范使用 GKE Filestore CSI 驱动程序来预配 PersistentVolumeClaim (PVC)，该 PVC 用于将 Filestore 共享挂载到 Pod。已装载卷的路径（用作数据和模型的缓存目录）会作为环境变量公开给数据科学家（笔记本用户）。

屏幕截图 1：使用 Filestore 卷的 Tensorflow 部署

屏幕截图 2：Filestore 永久性卷声明

角色 2：从 Jupyter 笔记本访问数据的数据科学家

数据科学家只想专注于运行实验。在此示例中，我们训练了 Google 的 Vision Transformer 模型 (ViT)，并从 Hugging Face 加载了一个 food101 数据集，该数据集主要由大小约为 100k 的图片组成，共占用 5 GiB 存储空间。我们还使用了 Hugging Face 的缓存功能，该功能可在首次读取数据后自动将其缓存在文件系统上。此文件路径需要作为环境变量共享给您的数据科学团队。通过将 Filestore 路径作为环境变量传递，相关数据将缓存在 Filestore 上。由于数据是缓存的，而不是从 Cloud Storage 提取的，因此我们在 a2-highgpu-1g 机器上运行了两个周期的训练，并与直接从 Cloud Storage 提取数据的训练时间（基准）进行比较，观察到训练时间缩短了 37%！

屏幕截图 3：从 Hugging Face 加载数据集并启用缓存

屏幕截图 4：Jupyter 笔记本文件资源管理器，已装载 Filestore 目录

屏幕截图 5、6、7：下载模型，开始训练并计量训练时间

表 1.  训练结果

使用注意事项

在本博客中，我们重点介绍了在 Cloud Storage 前使用 Filestore 作为加速器的好处，尤其是在数据集由大量小文件组成的情况下。尽管使用 Filestore 实例会产生费用，但如果这样做能缩短训练时间（以及提高 GPU 资源利用率），那么支付一定的存储费用也是值得的。处理较大的文件时，直接使用 Cloud Storage 中的数据可能更为合适。请根据您的应用场景选择最佳架构。