SageMaker中遍历PyTorch DataLoader触发KeyError:'error'报错排查

• 核心触发点：你将S3路径下的文件列表查询逻辑写在了自定义Dataset的__getitem__方法中。__getitem__会在DataLoader每次读取单个样本时执行，若开启多worker加载，会短时间内并发发起海量list_objects请求，直接触发S3服务端限流（返回429/503类响应）。
• KeyError: 'error'的本质：触发S3限流时，服务端返回的异常响应结构不符合botocore的标准错误格式，缺失初始化HTTPClientError所需的error字段，导致异常解析阶段先抛出KeyError，把真实的S3限流错误完全掩盖。
• 额外逻辑错误：list_objects是分页类枚举接口，单次调用最多返回1000条对象记录，在单样本读取环节反复调用该接口本身就不符合API设计用法，会让数据读取效率下降1-2个数量级，完全无法支撑正常训练的IO需求。
• 外层SageMaker抛出的UnexpectedStatusException、AlgorithmError: ExecuteUserScriptError只是训练任务失败的外层包装，根源就是上述数据读取环节的S3调用错误。

1. 重构Dataset的S3调用逻辑 #

将全量文件枚举逻辑从__getitem__迁移到Dataset的__init__方法中，训练初始化阶段一次性拉取所有训练样本的S3路径列表缓存到内存，__getitem__仅负责根据索引读取单张对应图片、做数据增强，禁止在单样本读取环节调用全量文件列表接口。
参考实现：

from torch.utils.data import Dataset
import boto3
from botocore.config import Config
from PIL import Image
import io
import torchvision.transforms as T

class Pix2PixS3Dataset(Dataset):
    def __init__(self, s3_prefix, transform=None):
        # 配置S3客户端自适应重试，降低限流概率
        s3_config = Config(
            retries = {
                'max_attempts': 10,
                'mode': 'adaptive'
            }
        )
        self.s3_client = boto3.client('s3', config=s3_config)
        self.transform = transform

        # 初始化阶段一次性拉取全量样本路径
        self.bucket, self.prefix = s3_prefix.replace("s3://", "").split("/", 1)
        self.sample_keys = []
        paginator = self.s3_client.get_paginator('list_objects_v2')
        for page in paginator.paginate(Bucket=self.bucket, Prefix=self.prefix):
            if "Contents" not in page:
                continue
            for obj in page["Contents"]:
                key = obj["Key"]
                # 过滤目录占位符、非图片文件
                if not key.endswith("/") and key.lower().endswith((".png", ".jpg", ".jpeg")):
                    self.sample_keys.append(key)

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

    def __getitem__(self, idx):
        # 仅读取当前索引对应的单张图片，不做全量枚举
        obj_key = self.sample_keys[idx]
        resp = self.s3_client.get_object(Bucket=self.bucket, Key=obj_key)
        img = Image.open(io.BytesIO(resp["Body"].read())).convert("RGB")
        
        # 原有pix2pix样本切分、数据增强逻辑
        if self.transform:
            img = self.transform(img)
        w = img.shape[-1]
        x, y = img[..., :w//2], img[..., w//2:]
        return x, y

2. 大训练集场景优化IO逻辑 #

如果训练样本总量超过10G，不要在训练过程中实时通过boto3拉取S3文件，选择以下两种方案之一：

• 启动SageMaker训练任务时，将S3数据输入模式设置为Fast File，直接以POSIX文件系统方式挂载S3路径，无需提前下载全量数据，也不需要手动编写boto3读文件逻辑。
• 在训练脚本最开头执行aws s3 sync命令，将全量训练数据同步到训练实例的本地NVMe存储盘，后续直接从本地路径读取文件，彻底避免训练过程中反复调用S3 API。

3. DataLoader参数适配 #

• num_workers参数不要设置过高，单张G4dn/G5规格的训练实例设置4-8个worker即可，worker数过高会放大S3并发请求压力，提升限流概率。
• 可以设置prefetch_factor=2，控制每个worker预加载的样本数量，避免短时间内发起过多S3请求。

修改完成后先单独实例化Dataset，随机抽取10-20个索引调用__getitem__验证读取逻辑正常，再启动短周期训练跑10-20个batch确认无S3相关报错，再启动全量训练任务即可。

内容的提问来源于stack exchange，提问作者CoderMath