宣布推出支持SOCI索引的镜像，加速SageMaker AI任务

您的深度学习容器镜像每花一分钟下载，就是您付费的一分钟，也是您的团队等待的一分钟。容器镜像下载纯属开销：计入计费任务时间，却对模型进展毫无贡献。对于现代深度学习镜像，每次任务的开销都长达数分钟。今天，我们很高兴地宣布 面向Amazon SageMaker AI训练任务和Amazon SageMaker AI处理任务的Seekable OCI（SOCI）索引，它将完整容器镜像下载从每个任务的关键路径中移除。

借助SOCI索引，SageMaker AI会延迟加载您的容器镜像。它只获取启动容器所需的文件，并在后台随代码运行流式传输其余部分，因此您的任务在几秒而非几分钟内开始。您可以通过指向预索引的AWS Deep Learning Container（DLC），或为自己的镜像添加索引来采用该功能。两种方式都无需更改您的任务定义、入口点或数据通道。其结果是更快的任务启动、更低的任务成本，以及更多时间用于构建模型，而不是等待基础设施。

支持SOCI索引的镜像可在几秒内启动Amazon SageMaker AI训练容器，且对训练性能没有影响。如我们在 我们测量了什么 一节所示，在基准测试中，这将镜像拉取期间的GPU空闲时间最多减少了90 percent。

本文介绍如何采用SOCI索引：它是什么、我们测量的结果，以及如何为您自己的任务开启该功能。

挑战

SageMaker AI训练和处理任务在存储于Amazon Elastic Container Registry（Amazon ECR）中的容器镜像内运行。在任务执行您的代码之前，服务会将整个镜像下载到实例。它会对每一层进行解压并写入磁盘。

对于小镜像，这只需几秒。但深度学习镜像通常超过10 GB或更多。该下载为每个任务增加数分钟，而且您要将其作为任务时间付费。影响最大的恰恰是迭代速度最重要的地方：涉及数百个任务的超参数调优扫描、快速实验以及短处理任务，在这些场景中镜像下载可能主导总运行时间。

SOCI索引的工作原理

Seekable OCI（SOCI） 是一项开源技术，可延迟加载容器镜像。运行时不是在启动前下载整个镜像，而是使用SOCI索引只获取立即启动容器所需的文件。然后在后台随代码访问流式传输其余部分。

SOCI索引是映射镜像内每个文件的元数据。您将其作为与原始镜像关联的符合OCI规范的工件存储在Amazon ECR中。由于索引是单独的工件，您的镜像摘要和任何现有签名仍然有效。

SageMaker AI会自动检查索引

当您创建训练或处理任务时，SageMaker AI会检查您指定的镜像是否具有SOCI索引，并相应地采取行动：

由于任务在需要时始终会回退到标准下载，添加SOCI索引绝不会阻止您的任务运行。您只会获得更快的启动。您绝不会面临破坏现有工作流程的风险。

我们测量了什么

我们在真实的SageMaker AI训练任务上以两种方式拉取相同的容器镜像：标准下载和SOCI索引延迟加载。然后我们用两者运行完整的训练任务，以测量其对训练速度的影响。

设置

• 八张镜像，压缩后大小为 1.8–35.4 GB。
• GPU 和 CPU 实例 ml.g4dn、ml.g5、ml.g6、ml.g6e、ml.g7、ml.g7e、ml.p4d、ml.p5 和 ml.m7i，规格从 xlarge 到 48xlarge，位于美国西部（俄勒冈）区域（us-west-2）和美国东部（弗吉尼亚北部）区域（us-east-1）。
• 下载时间 从开始拉取镜像起计算，直到容器可以运行。所有数字均为多次重复运行的中位数。

结果 1：在每种镜像大小下，下载时间都从数分钟缩短到数秒

下图对比了从 1.8 GB 到 35.4 GB 的八种镜像大小下，标准 Docker 拉取与 SOCI 延迟加载的容器镜像下载时间。

每个镜像大小显示两根柱子。Docker 下载时间柱（蓝色）随镜像大小稳步增长，从 1.8 GB 时的 48 秒增长到 35.4 GB 时的 7 分 17 秒。

每组柱子上方的加速标签显示了 Docker 时间与 SOCI 时间的比值。加速比从最小镜像的约 10 倍到最大镜像的超过 35 倍不等。

图 1：八种镜像大小下，标准 Docker 拉取与 SOCI 延迟加载的容器镜像下载时间对比

标准下载时间随镜像大小增长，因为每个文件都需要解压并写入。延迟加载无论大小均保持在 3–11 秒，提升了 8 倍–46 倍。镜像越大，节省的时间越多。

结果 2：该节省在每种实例类型和规格上都成立

在 13 种实例类型/规格组合上运行 13.2 GB 镜像，标准下载时间介于 3 分 16 秒到 4 分 20 秒之间，而延迟加载保持在 6-8 秒。每种配置都快约 34 倍。无需更改实例类型或规格即可受益。

图 2：13.2 GB 镜像在 13 种实例类型和规格组合上的下载时间

结果 3：除首个 epoch 的代价外，训练性能不受影响

我们分别从 Docker 拉取的镜像和 SOCI 索引的镜像运行了两种完整的训练任务，并比较了每个 epoch：

• 在 ResNet-50 上进行的图像分类，数据集为 Imagenette （9,469 张图像）
  • 容器镜像：PyTorch（13.2 GB DLC）和 TensorFlow（9.9 GB DLC）
  • 实例：ml.p5.48xlarge、ml.p4d.24xlarge、ml.g7e.48xlarge、ml.g7.48xlarge、ml.g6e.12xlarge、ml.g6e.24xlarge 和 ml.g6e.48xlarge。
• 大语言模型（LLM）微调：对 Qwen2.5-0.5B 进行全参数监督微调（SFT） 在 alpaca-cleaned (51,760个示例)，使用Hugging Face Trainer。
  • 容器镜像：9.4 GB Hugging Face PyTorch DLC。
  • 实例类型：ml.p5.48xlarge、ml.p4d.24xlarge、ml.g6.48xlarge。

下图将作业启动分解为镜像下载时间和首个训练 epoch 时间，比较了标准 Docker 拉取与使用 SOCI 索引的镜像在 GPU 和 CPU 配置下针对两种工作负载的表现。

图3:作业启动时间拆分为镜像下载时间和首个训练epoch时间,标准Docker pull与使用SOCI索引的镜像对比

• 从第二个epoch开始,对于两种工作负载,每种任务在每个实例上的单epoch时间都相同(在运行间波动范围内)。
• 从懒加载镜像启动时,首个epoch需要多花3–16秒,因为框架首次触到的文件需要被拉取。这一首个epoch的开销在每个实例规格上都相同。
• 总计:每个作业提前1-6分钟完成。即下载节省的时间减去一次性的首个epoch开销,这是每个作业的固定节省。

在你自己的工作负载上进行测试

这些数字来自前文描述的两种工作负载。首个epoch的开销取决于你的镜像以及代码在启动时触到哪些内容:

• 当第一个训练步骤只需要镜像的一小部分时,懒加载更有优势,其余部分会在训练过程中流式加载。
• 如果你的入口点在训练开始前就读取了镜像的大部分内容(例如扫描、校验或预加载其内容),那么首个epoch就需要等待与完整pull大致相同的字节数。净提升会很有限。
• 在我们的测试中,前置加载的启动仍然更快:
  • torch.compile和即时编译(JIT)构建的CUDA扩展为首个epoch增加了7-15秒(净节省仍为3分25秒–3分56秒)
  • 一个烘焙进镜像的6 GB模型在加载时流式传输其权重,增加了2-3.5分钟,但整个作业在使用索引镜像时仍然提前2-2.5分钟完成。

入门指南

你可以通过两种方式采用SOCI索引:使用预索引的AWS Deep Learning Container,或为自己的自定义镜像创建索引。

方式1:使用预索引的AWS Deep Learning Container

AWS发布了多个热门Deep Learning Container的SOCI索引变体,通过镜像标签上的-soci后缀来标识:

763104351884.dkr.ecr.us-east-1.amazonaws.com/pytorch-training:2.9-gpu-py312-cu130-ubuntu22.04-sagemaker-v1-soci

预索引变体可用于PyTorch Training、TensorFlow Training、vLLM和SGLang,部分框架还提供arm64变体。要采用懒加载,请在你的作业定义中将镜像标签更新为-soci变体。完整列表可在以下位置浏览: AWS Deep Learning Containers Gallery.

注意: 如果你基于DLC基础镜像(FROM)构建自己的镜像,则必须为最终镜像创建SOCI索引。基础镜像的索引不会延续到派生镜像。

方式2:为自己的镜像创建SOCI索引

要为自定义镜像建立索引,你需要一台Linux主机(例如运行Amazon Linux 2023的Amazon Elastic Compute Cloud (Amazon EC2)实例),并安装以下工具: containerd、 SOCI snapshotter,以及 Finch 或 nerdctl 。

1. 设置你的环境变量.

   ACCOUNT_ID="111122223333"
   REGION="us-east-1"
   REPOSITORY_NAME="my-training-image"
   ORIGINAL_IMAGE_TAG="v1.0"
   SOCI_IMAGE_TAG="v1.0-soci"
   REGISTRY="$ACCOUNT_ID.dkr.ecr.$REGION.amazonaws.com"

2. 安装所需的工具.

   sudo dnf install -y containerd fuse soci-snapshotter runfinch-finch
   sudo systemctl start soci-snapshotter
   sudo systemctl restart containerd

3. 对Amazon ECR进行身份验证.

   aws ecr get-login-password --region "$REGION" \
     | sudo nerdctl login -u AWS --password-stdin "$REGISTRY"

4. 拉取您的原始镜像.

   sudo nerdctl pull "$REGISTRY/$REPOSITORY_NAME:$ORIGINAL_IMAGE_TAG"

5. 创建带 SOCI 索引的镜像.

   sudo nerdctl image convert --soci \
     "$REGISTRY/$REPOSITORY_NAME:$ORIGINAL_IMAGE_TAG" \
     "$REGISTRY/$REPOSITORY_NAME:$SOCI_IMAGE_TAG"

6. 推送索引镜像.

   sudo nerdctl push --platform linux/amd64 \
     "$REGISTRY/$REPOSITORY_NAME:$SOCI_IMAGE_TAG"

这会在您的存储库中创建两个工件：SOCI 索引和一个将您的镜像链接到该索引的符合 OCI 标准的镜像索引清单。您的原始镜像及其摘要保持不变。

在 CI/CD 中自动创建索引

要在每次推送时自动生成索引，请部署 CFN ECR SOCI Index Builder。它使用 Amazon EventBridge 规则来检测 Amazon ECR 镜像推送事件并构建索引。无需对您的构建管道进行任何更改。

在您的 SageMaker AI 任务中使用索引镜像

像引用任何其他镜像一样引用您的带 SOCI 索引的镜像。无需进行 API 更改。仅标签不同。

训练示例（SageMaker Python SDK）

from sagemaker.train import ModelTrainer
from sagemaker.train.configs import Compute

trainer = ModelTrainer(
    training_image="111122223333.dkr.ecr.us-east-1.amazonaws.com/my-training-image:v1.0-soci",
    role="arn:aws:iam::111122223333:role/SageMakerExecutionRole",
    compute=Compute(
        instance_type="ml.g5.12xlarge",
        instance_count=2,
    ),
)

trainer.train()

验证您的镜像是否包含 SOCI 索引

启动之前，请确认您的标签指向一个已建立索引的镜像：

aws ecr batch-get-image \
  --repository-name "$REPOSITORY_NAME" \
  --image-ids imageTag="$SOCI_IMAGE_TAG" \
  --accepted-media-types \
  "application/vnd.oci.image.index.v1+json" \
  --query 'images[0].imageManifest' --output text \
  | jq '.manifests[] | {mediaType, artifactType}'

查找一个 artifactType 以 application/vnd.amazon.soci.index. 开头的子清单（child manifest），它应与您的平台清单并列存在。如果没有找到，则该标签并未指向已建立索引的镜像。

需要注意的事项

• 仅限 Amazon ECR 私有仓库。 其他注册表中的镜像，或通过私有 Docker 注册表（TrainingRepositoryAccessMode 设置为 Vpc）拉取的镜像，不支持延迟加载。
• 改进效果与镜像大小成正比。 小于约 4 GB 的镜像改进效果相对较小。
• 检查是否存在针对 Docker 的特定假设。 使用 SOCI 索引的镜像运行在基于 containerd 的运行时上。请检查入口点和设置脚本中是否存在调用 Docker CLI、读取 /var/run/docker.sock 或从 cgroup 路径派生标识符的代码。为保持向后兼容，SageMaker AI 会挂载一个 /.dockerenv 标记文件，但请避免新增对它的依赖。
• SageMaker AI 任务的 SOCI 索引功能已在所有 AWS 区域可用。

结论

借助 SOCI 索引，您可以最大限度地缩短 SageMaker AI 训练和处理作业下载容器镜像的时间。为您的镜像添加 SOCI 索引，或切换到预先建立索引的 AWS Deep Learning Container。无论哪种方式，您都能将启动时间从几分钟缩短到几秒钟，降低作业总成本，并加快机器学习（ML）迭代周期。而且您现有的作业定义、API 调用和数据管道保持完全不变。

要进一步了解 Amazon SageMaker AI 并开始使用 SOCI 索引，请浏览以下资源：

1. 选择一个已建立 SOCI 索引的 Deep Learning Container。 浏览 Deep Learning Containers Gallery 并在您的镜像标签后追加 -soci。有关特定框架的指导，请参阅 SOCI index for AWS Deep Learning Containers.
2. 为您自己的自定义镜像建立索引。 请按照 SageMaker AI Developer Guide中的分步演练操作，或参考 SOCI snapshotter getting-started guide 了解高级选项。
3. 在您的 CI/CD 管道中自动建立索引。 部署 CFN ECR SOCI Index Builder ，以便在每次推送镜像时自动生成索引。无需对管道做任何更改。
4. 使用 Finch 在本地进行测试。 使用 Finch 懒加载教程 以便在提交作业之前在基于 containerd 的运行时上验证您的镜像。
5. 深入了解这项技术。 阅读 Reducing container cold start times using SOCI index on DLAMI and DLC （AWS Machine Learning Blog）了解基准测试结果，或探索 底层原理：使用 Seekable OCI 和 AWS Fargate 实现容器镜像的延迟加载 （AWS Containers Blog）详细了解懒加载机制。

更快的容器启动仅需更换一个镜像标签。在您的下一个 SageMaker AI 作业中试用并亲身体验差异。

关于作者