View a markdown version of this page

Speed up training job startup with SOCI-indexed images - Amazon SageMaker AI
How it worksConsiderationsUse an indexed Deep Learning ContainerCreate a SOCI indexVerify an indexUse in a training jobPermissionsPrevent job failuresRelated resources

Speed up training job startup with SOCI-indexed images

A training job downloads its container image before it runs your code. For large images, such as deep learning images, this adds minutes to every training job.

Seekable OCI (SOCI) is an open-source technology from AWS for loading container images on demand (lazy loading). For the project, see soci-snapshotter on the GitHub website.

How it works:

  • Most container runtimes download the whole image before they start the container, even though the container needs only a small part of it to start.

  • A SOCI index maps the files inside each layer of the image, so the runtime can read one file without downloading the whole layer.

  • You create the index from your existing image. You don't rebuild the image. The index is stored next to the image in Amazon ECR.

  • When the training job starts, the runtime downloads only the files needed to start the container. The rest is downloaded in the background as your code uses it.

SOCI-indexed images start Amazon SageMaker AI training containers in seconds, cutting GPU idle time during image pulls by up to 90%, with no impact on training performance.

To use it, reference a SOCI-indexed image in your CreateTrainingJob request. Nothing else in the training job changes.

How it works

When you create a training job, SageMaker AI checks whether the image you specified has a SOCI index and behaves accordingly:

Scenario What happens
The image includes a SOCI index SageMaker AI lazy-loads the image and starts your container before the full download completes
The image has no SOCI index Normal download - no change in behavior
The index can't be read, or lazy loading doesn't succeed The training job falls back to a standard download and continues running

Adding an index never prevents a training job from running. In either case the image download is reported in the training job's Downloading secondary status.

Considerations

  • Where the image is stored

    • Store the image in an Amazon ECR private repository.

    • Reference the image by tag (repository:tag). An image referenced by digest (repository@sha256:…) is downloaded normally.

    • An image pulled through a private Docker registry in your VPC (TrainingRepositoryAccessMode set to Vpc) is downloaded normally.

  • Which index format is detected

    • The tag must point to an image that contains a SOCI index v2 (see SOCI Index Manifest v2 on the GitHub website). nerdctl image convert --soci and the Deep Learning Container -soci tags create this format.

    • A v1 index isn't detected. A v1 index is stored as a separate artifact next to the image, for example by nerdctl push --snapshotter=soci or soci create.

  • When the time saved is smaller

    • Files are downloaded the first time your code opens them.

    • If your code reads a large part of the image before training starts (for example, model weights stored in the image), that time moves from the download to the start of training.

  • Large datasets

    • In File input mode, the training job also waits for the dataset download. With a very large dataset, that wait can take longer than the image download.

    • Use FastFile or Pipe input mode to keep the benefit.

Use a Deep Learning Container that includes a SOCI index

AWS publishes SOCI-indexed variants of several Deep Learning Containers (DLCs), identified by a -soci suffix on the tag. To adopt lazy loading, change only the tag:

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

Create a SOCI index for your own image

You need:

  • A Linux host, such as an Amazon EC2 instance running Amazon Linux 2023.

  • These tools on the host:

    • containerd on the containerd website

    • SOCI snapshotter 0.10.0 or later, on the GitHub website

    • nerdctl 2.1.3 or later, on the GitHub website. Finch, which is built on nerdctl, can be used instead; see Lazy loading on the Finch website.

  • sudo to run the commands. Rootless mode isn't supported.

  • The push permissions in Permissions for the identity on the host.

The procedure creates an indexed copy of an existing image under a new tag; the original image, tag, and digest are unchanged.

To create a SOCI index

  1. Set variables for your account, Region, repository, the tag to index, and the tag for the indexed image:

    ACCOUNT_ID="111122223333"
REGION="us-east-1"
REPOSITORY_NAME="repository-name"
ORIGINAL_IMAGE_TAG="original-image-tag"
SOCI_IMAGE_TAG="soci-indexed-image-tag"
REGISTRY="$ACCOUNT_ID.dkr.ecr.$REGION.amazonaws.com"

  2. Install the tools and start the snapshotter (Amazon Linux 2023; for other distributions see the SOCI snapshotter installation guide on the GitHub website):

    sudo dnf install -y containerd soci-snapshotter nerdctl jq
sudo systemctl start soci-snapshotter
sudo systemctl restart containerd

  3. Confirm that soci is listed under Server → Plugins → Storage:

    sudo nerdctl system info

  4. Authenticate to Amazon ECR and pull the original image:

    aws ecr get-login-password --region "$REGION" | sudo nerdctl login -u AWS --password-stdin "$REGISTRY"
sudo nerdctl pull "$REGISTRY/$REPOSITORY_NAME:$ORIGINAL_IMAGE_TAG"

  5. Create and push the SOCI-indexed image:

    sudo nerdctl image convert --soci \
  "$REGISTRY/$REPOSITORY_NAME:$ORIGINAL_IMAGE_TAG" \
  "$REGISTRY/$REPOSITORY_NAME:$SOCI_IMAGE_TAG"
sudo nerdctl push --platform linux/amd64 "$REGISTRY/$REPOSITORY_NAME:$SOCI_IMAGE_TAG"

    Layers below the minimum layer size aren't indexed and are downloaded normally when the container starts (--soci-min-layer-size, see the nerdctl SOCI documentation on the GitHub website).

After the push:

  • The repository has two new artifacts, both shown on the repository's Images page in the Amazon ECR console:

    • The SOCI index.

    • An OCI image index manifest that links your image and the index under the new tag.

  • Amazon ECR doesn't create SOCI indexes on its own. To index images automatically on every push, deploy the CFN ECR SOCI Index Builder on the GitHub website with SociIndexVersion set to V2. For details, see the Index Builder documentation on the GitHub website. It runs from an Amazon EventBridge rule on Amazon ECR push events.

To run the indexed image on the host the way a training job runs it, and check your entry point for the patterns in Prevent job failures: check your code for Docker-specific assumptions:

sudo nerdctl run --snapshotter soci --rm -it "$REGISTRY/$REPOSITORY_NAME:$SOCI_IMAGE_TAG" /bin/bash

Verify that an image includes a SOCI index

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}'

  • Indexed: the output lists a child manifest whose artifactType begins with application/vnd.amazon.soci.index.

  • Not indexed: no such entry. Jobs that use the tag download the image normally.

Use a SOCI-indexed image in a training job

Set the indexed image as TrainingImage in AlgorithmSpecification:

"AlgorithmSpecification": {
  "TrainingImage": "111122223333.dkr.ecr.us-east-1.amazonaws.com/repository-name:soci-indexed-image-tag",
  "TrainingInputMode": "File"
}

With the SageMaker Python SDK:

from sagemaker.train.model_trainer import ModelTrainer
from sagemaker.core.training.configs import Compute

trainer = ModelTrainer(
    training_image="111122223333.dkr.ecr.us-east-1.amazonaws.com/repository-name:soci-indexed-image-tag",
    role="arn:aws:iam::111122223333:role/execution-role-name",
    compute=Compute(instance_type="ml.g5.12xlarge", instance_count=2),
)
trainer.train()

Permissions

  • To run a training job: the execution role needs the same Amazon ECR actions as for any other image, listed in CreateTrainingJob API: Execution Role Permissions. Lazy loading needs no additional permissions.

    • ecr:GetAuthorizationToken

    • ecr:BatchCheckLayerAvailability

    • ecr:GetDownloadUrlForLayer

    • ecr:BatchGetImage

  • To create and push indexes: the identity on the build host needs pull and push permissions on the repository, for example:

{
  "Version": "2012-10-17",
  "Statement": [
    { "Effect": "Allow", "Action": "ecr:GetAuthorizationToken", "Resource": "*" },
    {
      "Effect": "Allow",
      "Action": [
        "ecr:BatchCheckLayerAvailability",
        "ecr:BatchGetImage",
        "ecr:GetDownloadUrlForLayer",
        "ecr:InitiateLayerUpload",
        "ecr:UploadLayerPart",
        "ecr:CompleteLayerUpload",
        "ecr:PutImage"
      ],
      "Resource": "arn:aws:ecr:us-east-1:111122223333:repository/repository-name"
    }
  ]
}

See Private repository policies in the Amazon Elastic Container Registry User Guide.

Prevent job failures: check your code for Docker-specific assumptions

SOCI-indexed images run on a containerd-based runtime, not on Docker:

  • Container behavior is the same.

  • Code that detects a container by looking for Docker-specific files can take a different path.

  • Review your entry point, framework launchers, and setup scripts for the patterns in the following table.

Pattern Why it matters Recommendation
Checking for the /.dockerenv file Created by the Docker daemon; not part of any container standard Set your own environment variable, or detect containerization without depending on a runtime
Matching docker in /proc/self/cgroup or /proc/1/cgroup Cgroup paths differ between runtimes and cgroup versions Don't branch on the runtime name
Deriving a container ID from a cgroup path or the host name format Runtime-specific formats Pass identifiers in as environment variables
Calling the docker CLI or /var/run/docker.sock inside the container Not available Remove the dependency; SageMaker AI manages the container lifecycle

  • SageMaker AI mounts a /.dockerenv marker in the container so that existing checks keep working. Don't add new dependencies on it.

  • If you use the SageMaker distributed data parallel library (SMDDP) or mpirun launchers, use a current Deep Learning Container. Some earlier launcher versions relied on /.dockerenv.