跳到正文
AWS Machine Learning Blog· Kanwaljit Khurmi·· 4 小时前AI 评分42

AWS 为 SageMaker AI 训练与处理作业推出 SOCI 索引镜像,容器启动从分钟级降至秒级

Announcing SOCI-indexed images for faster SageMaker AI jobs

AI 导读

AWS 宣布为 Amazon SageMaker AI 训练和处理作业提供 Seekable OCI(SOCI)索引,镜像改为懒加载,只拉取启动所需文件、其余后台流式传输。

深度导读
AWS 宣布为 Amazon SageMaker AI 的训练作业和处理作业提供 Seekable OCI(SOCI)索引支持,把容器镜像下载从每个作业的关键路径上移除。此前 SageMaker AI 作业需要先把 Amazon ECR 中的完整镜像下载到实例、解压并写入每一层,深度学习镜像常超过 10 GB,这一过程会给每个作业增加数分钟计费时间,在超参数调优、快速实验和短处理作业中占比尤其高。 SOCI 是一种开源懒加载技术:运行时借助 SOCI 索引只拉取启动容器所需的文件,其余内容在代码运行过程中后台流式传输。索引作为符合 OCI 规范的独立产物存放在 Amazon ECR 中并与原镜像关联,因此镜像 digest 和已有签名保持有效。创建作业时 SageMaker AI 会自动检查镜像是否带索引:有索引则懒加载,无索引则按原方式下载,索引读取失败或懒加载不成功时回退到标准下载,作业不会因此无法运行。 实测覆盖 8 个镜像(压缩后 1.8–35.4 GB)、ml.g4dn 至 ml.p5 及 ml.m7i 等实例、us-west-2 与 us-east-1 两个区域。1.8 GB 镜像下载从 48 秒降至 3–11 秒(约 8 倍),13.2 GB 从 3 分 42 秒降至约 6 秒(约 34 倍),35.4 GB 从 7 分 17 秒降至 3–11 秒(约 38 倍);懒加载在不同镜像大小下稳定在 3–11 秒,整体提升 8–46 倍。13.2 GB 镜像在 13 种实例规格组合上均约快 34 倍。 训练性能方面,团队用 ResNet-50 在 Imagenette 上的图像分类(PyTorch 13.2 GB 与 TensorFlow 9.9 GB DLC)和 Qwen2.5-0.5B 在 alpaca-cleaned 上的全参数 SFT(9.4 GB Hugging Face PyTorch DLC)做了对比。从第二个 epoch 起单 epoch 耗时一致,首个 epoch 因首次读取文件多花 3–16 秒,且该开销与实例规格无关,每个作业净提前 1–6 分钟完成。若入口脚本在训练前读取镜像大部分内容,收益会变小;测试中 torch.compile 与 JIT 编译 CUDA 扩展使首 epoch 增加 7–15 秒,净节省仍有 3 分 25 秒至 3 分 56 秒。 采用方式有两种:一是直接使用带 -soci 后缀标签的预索引 AWS Deep Learning Container,目前覆盖 PyTorch Training、TensorFlow Training、vLLM 和 SGLang,部分框架含 arm64 变体;基于 DLC 基础镜像自行构建的镜像需为最终镜像另建索引。二是为自定义镜像建索引,需在装有 containerd、SOCI snapshotter 及 Finch 或 nerdctl 的 Linux 主机上,通过 nerdctl image convert --soci 转换并推送,也可部署 CFN ECR SOCI Index Builder 在每次推送时自动生成索引。作业侧只需替换镜像标签,无需改动 API。需注意目前仅支持 Amazon ECR 私有仓库,约 4 GB 以下镜像收益相对有限,且索引镜像运行在基于 containerd 的运行时上,需检查对 Docker 的特定假设。

来源:AWS Machine Learning Blog · aws.amazon.com