GPU 在空转等数据？PyTorch 训练集进 RustFS 的三种接法

分布式训练里最尴尬的一幕往往是：GPU 算完一批，下一批数据还没从存储拉上来，几十张卡集体空转。把训练集放在对象存储、训练节点按需拉取，能把存储和计算彻底解耦，训练集群扩缩容时存储不用跟着动。很多团队一开始把数据集塞在本地 NVMe，单机训没问题，一上分布式多机就发现数据搬不动、各机看到的还不是同一份；对象存储用同一个 key 空间，所有节点读的是同一份数据，这份一致性是分布式训练的前提。RustFS 是 100% 兼容 S3 的存储，PyTorch、TensorFlow 不用改框架代码，只要把 endpoint 指过来就能当训练数据底座，已有的 boto3、s3fs 那套代码原样可用。不过 s3fs 这类文件系统接口要额外装包，TensorFlow 的 tf.data 走 S3 还依赖 tensorflow-io，后面细说。

训练集为什么适合落在对象存储上

训练数据集常常是两种极端：要么是几十亿张几 KB 的小图、小音频，要么是几百 GB 的大文件；本地盘放不下，NFS 又撑不住高并发随机读。对象存储按 key 取对象，天然适合这种「海量 key、随机读」的模式，训练节点要哪片数据就拉哪片。

RustFS 在这类负载上有一个对训练特别友好的点：官方仓库基准里，4KiB 小对象 PUT 吞吐约为同机 MinIO 的 2.3 倍（注意这是仓库 README 的厂商基准，不是第三方复测）。训练集里大量小文件正是小对象写入和读取最密集的场景，这个差距会直接变成更短的 epoch 等待。RustFS 1.0.0 于 2026-09-16 GA，Apache 2.0，生产就绪。官方集成目录（docs.rustfs.com/zh/developer/integration）的 AI 类目收录的是 MLflow、Ray、vLLM 这类平台，PyTorch、TensorFlow 没有专门的官方指南，但它们走标准 S3 客户端接入，不需要官方指南也能对上，每篇集成指南标注的端点与寻址要求是通用的。

用 boto3 直读，最省事

最朴素的接法就是 boto3，把 endpoint_url 指到 RustFS，其余代码和读 AWS S3 一模一样，框架侧零改造。

python 复制代码

import boto3
from botocore.client import Config

s3 = boto3.client(
    "s3",
    endpoint_url="http://rustfs-host:9000",   # 指向 RustFS S3 API，不是 AWS
    aws_access_key_id="<your-access-key>",
    aws_secret_access_key="<your-secret-key>",
    region_name="us-east-1",                  # RustFS 默认区域
    config=Config(
        signature_version="s3v4",             # RustFS 走 v4 签名
        s3={"addressing_style": "path"},      # path-style 寻址
    ),
)
obj = s3.get_object(Bucket="train-data", Key="imagenet/val/ILSVRC2012_val_00000001.JPEG")
data = obj["Body"].read()

一个要点：region 固定填 us-east-1（RustFS 的默认区域），寻址显式配 path-style 最稳，这是官方 Python SDK 指南的四参数写法，漏掉 signature_version 和 addressing_style 在部分环境会报签名不匹配。内网里可以走明文 HTTP（是否上 TLS 由你自己的网络边界决定），省掉握手开销，高并发拉数据时这一点很实在。桶内对象数随集群容量扩展，训练集从百万级小文件涨到十亿级，同一个桶接着用，不需要提前分库分桶。

接进 PyTorch DataLoader，做分片与缓存

单机 boto3 够用，但分布式训练要的是每个 worker 各取各的分片。fsspec 加 s3fs 能把这个桶挂成类文件系统，直接喂给 torch DataLoader。

python 复制代码

import os
import fsspec
import torch.distributed as dist
from torch.utils.data import DataLoader, IterableDataset

class S3ShardDataset(IterableDataset):
    def __init__(self, prefix, endpoint, key, secret):
        self.fs = fsspec.filesystem(
            "s3",
            client_kwargs={
                "endpoint_url": endpoint,
                # region 与签名配置同 boto3，s3fs 侧这样传
                "config_kwargs": {"s3": {"addressing_style": "path"}},
            },
            key=key, secret=secret,
        )
        self.files = sorted(self.fs.glob(f"{prefix}/shard-*.tar"))

    def __iter__(self):
        # IterableDataset 下 DistributedSampler 不生效，按 rank/world_size 手工切
        rank, world = dist.get_rank(), dist.get_world_size()
        mine = self.files[rank::world]          # 每个 rank 只认领自己的那一份
        worker = torch.utils.data.get_worker_info()
        if worker:                              # 单机 8 个 worker 再切一次
            mine = mine[worker.id::worker.num_workers]
        for path in mine:
            with self.fs.open(path, "rb") as f:
                yield from read_tar_stream(f)   # 流式解 tar，别整段 read 进内存

loader = DataLoader(
    S3ShardDataset("s3://train-data/webdataset",
                   "http://rustfs-host:9000", KEY, SECRET),
    batch_size=32, num_workers=8, prefetch_factor=2,
)

注意切分要自己动手：DistributedSampler 只对 map-style 数据集生效，IterableDataset 里它不介入，不按 rank 切的话每个进程都会把全量分片读一遍------数据重复 N 倍，I/O 也放大 N 倍，正好把本文想解决的问题翻倍。单机内的 worker 切分同理，用 get_worker_info() 自己分。read_tar_stream 换成你们自己的 tar 流式解析（WebDataset 的 tariterator 或自定义），关键是别把单个分片整段 f.read() 拉进内存。

把训练集先打成 WebDataset 或 TFRecord 这种大分片（单分片百 MB 量级），比几 KB 一个对象去读要顺得多：对象数少了，列举和元数据压力都下来。这点对任何 S3 兼容存储都成立，RustFS 上同样受益，配合它小对象吞吐的优势，能把「列举慢」这个瓶颈先挪走。顺带补一句依赖账：PyTorch 侧的 fsspec/s3fs 是 pip install s3fs 的额外依赖；TensorFlow 的 tf.data 要直接读 s3:// 路径，装的是 tensorflow-io，版本还要和 TF 主版本对上。「不用改框架代码」说的是接入方式，不是零依赖。

几十个 worker 同时拉数据时，尾延迟比平均吞吐更要命

分布式训练起 DistributedSampler 后，一个训练进程往往带八到十几个 DataLoader worker，几十张卡同时锤存储。这个时候看的不只是平均吞吐，更是尾延迟的一致性：某一个 worker 偶尔卡一下，整批的步子就被它拖住。RustFS 的存储内核用 Rust 实现，没有垃圾回收带来的停顿抖动。这只排除了 GC 这一种抖动源，尾延迟还取决于纠删码布局、盘和网络，能不能更可控，拿自己的工况压一轮才算数。

落到配置上，DataLoader 的 num_workers 和 prefetch_factor（两者配合用，prefetch_factor 只在 num_workers > 0 时生效）让数据在 GPU 算当前批时就预取下一批；对象侧则把训练集打成大分片，减少单批要列举和打开的对象数。两端一起压，才能把空转时间真正压下来。

预签名 URL 让训练节点绕过密钥直拉

如果训练节点很多、又不想在每个容器里塞长期密钥，可以让调度层用 RustFS 的 Access Key 生成预签名 URL，训练容器拿 URL 直接拉对象，密钥不下发到训练侧。生成时用前面那个带 signature_version="s3v4" 配置的 client，因为 boto3 对自定义 endpoint 的预签名默认可能落到 v2 签名，RustFS 走 v4，不配会有签名不匹配的风险。

python 复制代码

url = s3.generate_presigned_url(
    "get_object",
    Params={"Bucket": "train-data", "Key": "imagenet/train/0001.jpg"},
    ExpiresIn=3600,
)

这种做法适合只读分发、短时有效的数据拉取，密钥只在生成 URL 的一侧存在，符合最小权限原则，训练容器拿到的只是一个有过期时间的链接。

checkpoint 也落 RustFS，容灾不靠本地盘

训练中断要能续上，checkpoint 就不能只写在本地 SSD。直接写到 RustFS 的 s3:// 路径，多机多卡共享同一份最新权重，节点挂了换一台继续。

python 复制代码

torch.save(model.state_dict(), "checkpoint.pt")
s3.upload_file("checkpoint.pt", "ml-checkpoints", f"run-{run_id}/step-{step}.pt")

给 checkpoint 桶开版本控制，回滚时按版本取，不会被一次坏写覆盖。RustFS 的版本化桶和 S3 语义一致，upload_file 的覆盖行为你一样能用版本号收回来，这一层在分布式训练反复写 checkpoint 时很关键。如果训练横跨多个集群，checkpoint 桶还可以用 RustFS 的站点复制异步同步到第二个站点，主集群崩了换站点接着续训，不用从 epoch 0 重来。

落地顺序先排这三步

先把数据集按大小定分片：小文件多的先合并成大分片，大 checkpoint 走 multipart，别用小对象去顶大文件，那样会抵消掉小对象吞吐的优势。单节点压一遍 GET 吞吐，确认 RustFS 端到端口径和客户端 path-style 配置对上，再上分布式 DataLoader。生产凭据走专用 Access Key，不要用管理员密钥。

一个要老实说的中性边界：boto3、s3fs 都是通用 S3 客户端，RustFS 走标准 S3 即可，但公开基准里没有 RustFS 在 PyTorch DataLoader 下的端到端数字，小对象优势在 4KiB 量级最明显，大文件走 multipart 之后差距会收窄，要不要切进来先拿自己的数据集压一轮再定。

想看官方集成清单，集成目录（docs.rustfs.com/zh/developer/integration）的 AI 类目收录了 MLflow、Ray、vLLM 等平台，每篇指南标了 RustFS 端点与寻址要求；PyTorch、TensorFlow 本身没有官方指南页，boto3 的四参数写法以 Python SDK 指南为准。存储引擎本身在 rustfs/rustfs 这个仓库，Apache 2.0，1.0.0 已经 GA，能直接拉镜像起一个单节点先验证上面这些接法。