欢迎光临
我们一直在努力

怎样为云端AI训练配置和优化S3/OSS存储的读写性能?

在现代AI训练工作流中,数据通常存储在对象存储服务(如AWS S3、阿里云OSS)中。然而,由于模型规模和数据集大小的爆炸式增长,存储I/O往往成为训练过程中的首要瓶颈,尤其是在涉及数百万个小文件(例如图像、文本片段)的情况下。对象存储的高延迟特性,使得频繁的小文件读取操作,极大地拖慢了GPU的利用率。

本文将深入探讨如何通过基础设施配置、客户端优化以及数据格式转换,来最大化S3/OSS的读写吞吐量,确保数据能够及时喂饱昂贵的计算资源。

1. 性能瓶颈分析:小文件问题与连接限制

对象存储(Object Storage)是为高吞吐量和大对象设计,但其首次字节延迟(Time-to-First-Byte Latency)相对较高。当训练数据集包含数百万个KB级别的小文件时,数据加载器需要发起数百万次单独的HTTP请求(List、Get操作)。每次请求都需要握手、认证和建立连接,这些开销累积起来远超实际数据传输所需的时间,导致性能急剧下降。

2. 基础设施层面优化

2.1 启用VPC Endpoint/Private Link

如果训练集群部署在云供应商的VPC内部,确保数据流量通过专用的VPC Endpoint(如AWS的Gateway Endpoint或Interface Endpoint)访问S3/OSS。这避免了数据绕行互联网或通过NAT网关,显著降低了延迟并提高了带宽一致性。

2.2 存储区域和计算区域一致性

确保存储桶与训练实例位于同一区域(Region)甚至同一可用区(Availability Zone),以实现最低的网络延迟。

3. 客户端配置与并行度调整

Python生态系统通常使用

1
s3fs

1
ossfs

(基于

1
fsspec

)来访问云存储。优化这些客户端的连接参数是提升性能的关键。

3.1 增加连接池大小

默认的HTTP连接池大小可能不足以支撑高并发的数据加载。我们可以通过调整连接配置来允许更多并发连接。


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
import s3fs

# 假设您正在配置s3fs访问AWS S3
S3_ACCESS_KEY = "..."
S3_SECRET_KEY = "..."

# 关键优化参数:设置更大的max_pool_connections
# 默认值通常为10,增加到50或更高,取决于您的集群网络限制
fs = s3fs.S3FileSystem(
    key=S3_ACCESS_KEY,
    secret=S3_SECRET_KEY,
    client_kwargs={
        'config': {
            'max_pool_connections': 64,
            'connect_timeout': 10,
            'read_timeout': 60
        }
    }
)

print(f"S3 FileSystem configured with max connections: 64")
# 现在可以使用这个fs对象进行文件操作,例如:
# data = fs.glob("s3://my-bucket/dataset/*.jpg")

3.2 PyTorch/TensorFlow DataLoader 并行优化

数据加载器(如PyTorch

1
DataLoader

)通过

1
num_workers

参数实现了进程级并行。当读取远程存储时,增加

1
num_workers

可以显著提高并发请求S3/OSS的能力。

注意: 盲目增加

1
num_workers

会消耗更多CPU和内存。对于I/O密集型任务,通常需要更高的

1
num_workers

(例如,8到16),但必须与您的连接池大小保持协调。


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
import torch
from torch.utils.data import DataLoader, Dataset

class RemoteDataset(Dataset):
    # 假设这个Dataset内部使用了上文配置的fs对象或直接使用s3fs/fsspec
    def __init__(self, file_list, fs):
        self.file_list = file_list
        self.fs = fs

    def __len__(self):
        return len(self.file_list)

    def __getitem__(self, idx):
        path = self.file_list[idx]
        # 这是一个模拟的读取操作,实际中应使用fs对象读取
        with self.fs.open(path, 'rb') as f:
            data = f.read()
        # 模拟数据处理和返回
        return torch.tensor(len(data))

# --- 训练配置 ---

# 强烈建议将num_workers设置得较高以应对网络延迟
# pin_memory=True 有助于加速数据传输到GPU
BATCH_SIZE = 128
NUM_WORKERS = 12 # I/O Bound任务的核心优化点

training_dataset = RemoteDataset(file_list=['...'], fs=fs)
trainer_loader = DataLoader(
    training_dataset,
    batch_size=BATCH_SIZE,
    num_workers=NUM_WORKERS,
    pin_memory=True
)

# 开始训练循环...

4. 根本解决方案:数据格式聚合

解决小文件问题的最根本方法是数据聚合(Data Aggregation)。将数千个小文件打包成少量、体积较大的文件,可以大幅减少I/O请求次数。

4.1 推荐的优化数据格式

  • WebDataset (推荐): 专门为深度学习设计,将图像及其标签打包成大的
    1
    .tar

    1
    .zip

    文件。它支持高效的并行读取和流式传输。

  • TFRecord/RecordIO: TensorFlow生态系统的标准格式,用于序列化数据。
  • Parquet/Zarr: 适用于结构化数据或高维数组。

4.2 WebDataset 转换示例(概念性)

以下是一个概念性的数据打包过程,将小文件转换为一个大型

1
.tar

文件,减少S3/OSS的GET请求数量。


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
# 假设我们有1000张图片和对应的标签文件

# 依赖: webdataset
# !pip install webdataset

import webdataset as wds

# 定义打包函数
def create_webdataset(input_files, output_tar_path):
    with wds.TarWriter(output_tar_path) as sink:
        for i, (image_path, label_path) in enumerate(input_files):
            # 将数据写入tar sink
            with open(image_path, 'rb') as image_file,
                 open(label_path, 'r') as label_file:

                key = f"sample_{i:06d}"
                sink.write({
                    "__key__": key,
                    "jpg": image_file.read(),
                    "txt": label_file.read() # 标签或其他元数据
                })
    print(f"Successfully created aggregated dataset at: {output_tar_path}")

# 实际操作:将这个大型.tar文件上传到S3/OSS
# fs.put('local_dataset.tar', 's3://my-bucket/optimized/dataset.tar')

通过将数据集从上百万个小文件转换为几十个大文件,数据加载器只需要发起几十次高吞吐量的

1
GET

请求,而不是数百万次高延迟的连接建立请求,从而彻底消除S3/OSS的I/O瓶颈。

【本站文章皆为原创,未经允许不得转载】:汤不热吧 » 怎样为云端AI训练配置和优化S3/OSS存储的读写性能?
分享到: 更多 (0)

评论 抢沙发

  • 昵称 (必填)
  • 邮箱 (必填)
  • 网址