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

你的深度学习容器镜像每多花一分钟下载，你就多付一分钟的钱，团队也多等一分钟。容器镜像下载是纯粹的开销：可计费的任务时间，却不产生任何模型进展。对于现代深度学习镜像，每个任务都有数分钟这样的开销。今天，我们很高兴宣布推出 面向 Amazon SageMaker AI Training 和 Amazon SageMaker AI 处理任务的 Seekable OCI (SOCI) 索引，它将完整的容器镜像下载从每个任务的关键路径中彻底移除。

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

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

本文介绍如何采用 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：

• Image classification on ResNet-50 on Imagenette (9,469 images)
  • 容器镜像：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。
• Large language model (LLM) fine-tuning: full-parameter supervised fine-tuning (SFT) of Qwen2.5-0.5B on alpaca-cleaned (51,760 examples) with the 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 拉取与 SOCI 索引镜像对比

• 从第二个 epoch 开始，每个 epoch 的时间在每个实例上、两种工作负载下都相同（在运行间差异范围内）。
• 第一个 epoch 因镜像懒加载而多花 3–16 秒，同时框架首次访问的文件会被拉取。这一首轮开销在所有实例规格上都相同。
• 净效果：每个作业提前 1-6 分钟完成。这是下载节省量减去一次性首轮开销后的结果，即每个作业的固定节省量。

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

这些数字来自前面描述的两个工作负载。首轮开销取决于你的镜像以及代码在启动时访问的内容：

• 当第一个训练步骤只需要镜像的一小部分时，懒加载就能胜出，其余部分会在你训练时流式传入。
• 如果你的入口点在训练开始前读取了镜像的大部分内容（例如扫描、校验和或预加载其内容），那么第一个 epoch 等待的字节数大致与完整拉取相同。净提升会很小。
• Front-loaded starts still came out ahead in our tests:
  • 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 基础镜像构建自己的镜像，则必须为最终镜像创建 SOCI 索引。基础镜像的索引不会延续到派生镜像。

选项 2：为自己的镜像创建 SOCI 索引

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

1. Set your environment variables.

   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. Install the required tools.

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

3. Authenticate to Amazon ECR.

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

4. Pull your original image.

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

5. Create the SOCI-indexed image.

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

6. Push the indexed image.

   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 开头，并与你的平台清单并列。如果不存在，则该标签未指向索引镜像。

注意事项

• 仅限 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 lazy-loading walkthrough 在提交作业前于基于 containerd 的运行时上验证镜像。
5. 深入了解该技术。阅读 AWS Machine Learning Blog 上的 Reducing container cold start times using SOCI index on DLAMI and DLC 了解基准测试，或浏览 AWS Containers Blog 上的 Under the hood: Lazy loading container images with Seekable OCI and AWS Fargate 详细了解延迟加载机制。

更快的容器启动只需一个镜像标签。在下一个 SageMaker AI 作业上试用，体验其中的差异。

关于作者

Kanwaljit Khurmi

Kanwaljit 是 Amazon Web Services 的解决方案架构与数据科学家高级经理，专攻 AI 和 ML。他与 AWS 产品团队、工程团队和客户合作，提供指导和技术支持，帮助客户在 AWS 上最大化其混合 GenAI 解决方案的价值。Kanwaljit 擅长帮助客户处理容器化、数据科学和机器学习应用。

Mona Mona

Mona 目前在 Amazon 担任高级 AI/ML 专家解决方案架构师。她已出版三本书，最新著作是《AI Agents on AWS》。她撰写了 20 多篇关于 AI/ML 和云技术的博客，并参与合著了一篇关于 CORD19 神经搜索的研究论文，该论文在著名的 AAAI（美国人工智能促进协会）会议上荣获最佳研究论文奖。

Michael Oguike

Michael 是 Amazon SageMaker AI 的产品经理。他致力于将技术和 AI 应用于解决现实世界的问题。在 AWS，他帮助各行业的客户大规模构建、训练和部署 AI/ML 模型。在加入 AWS 之前，Michael 曾设计和运营海上油气井及设施。工作之余，他喜欢通过书籍和播客探索行为科学与心理学。

Deep Shah

Deep 是 Amazon SageMaker 团队的软件开发工程师，负责为 AWS 上的模型训练提供支持的基础设施。在不扩展训练系统的时候，他喜欢模拟赛车、阅读和旅行。