Imagens com índice SOCI aceleram jobs no SageMaker AI

O problema que o SOCI resolve

Quem trabalha com treinamento de modelos no Amazon SageMaker AI conhece bem aquele tempo perdido antes do job sequer começar a processar dados. Antes de executar qualquer linha de código, o serviço precisa baixar a imagem de container armazenada no Amazon Elastic Container Registry (Amazon ECR), descomprimindo e gravando cada camada no disco da instância.

Para imagens pequenas, isso leva alguns segundos. O problema é que imagens de deep learning modernas facilmente ultrapassam 10 GB. Resultado: vários minutos de espera a cada job — e esse tempo é cobrado normalmente, sem gerar nenhum progresso no modelo. O impacto é ainda mais crítico em cenários de alta iteração, como varreduras de hiperparâmetros com centenas de jobs ou experimentos rápidos, onde o download da imagem pode dominar o tempo total de execução.

Como o SOCI funciona

O Seekable OCI (SOCI) é uma tecnologia open source que implementa o conceito de lazy loading (carregamento sob demanda) para imagens de container. Em vez de baixar tudo antes de iniciar, o runtime usa um índice SOCI para buscar apenas os arquivos necessários para subir o container imediatamente. O restante é transmitido em segundo plano conforme o código vai acessando os dados.

Um índice SOCI é basicamente um mapa de todos os arquivos dentro da imagem. Ele é armazenado no Amazon ECR como um artefato compatível com o padrão OCI (Open Container Initiative), vinculado à imagem original. Como o índice é um artefato separado, o digest da imagem e eventuais assinaturas existentes continuam válidos.

Como o SageMaker AI usa o índice

Quando um job de treinamento ou processamento é criado, o SageMaker AI verifica automaticamente se a imagem especificada possui um índice SOCI e age de acordo:

Imagem com índice SOCI: o SageMaker AI faz o lazy loading e inicia o container antes de concluir o download completo.
Imagem sem índice SOCI: comportamento padrão, sem alteração.
Índice ilegível ou falha no lazy loading: o job faz fallback para o download padrão e continua executando normalmente.

Esse mecanismo de fallback garante que adicionar um índice SOCI nunca interrompa um workflow existente. O pior cenário é simplesmente não ter a melhoria de velocidade.

O que foi medido

A AWS realizou benchmarks comparando o download padrão com o lazy loading via SOCI em jobs reais de treinamento no SageMaker AI. Foram utilizadas oito imagens com tamanhos entre 1,8 GB e 35,4 GB comprimidos, em instâncias GPU e CPU das famílias ml.g4dn, ml.g5, ml.g6, ml.g6e, ml.g7, ml.g7e, ml.p4d, ml.p5 e ml.m7i, nas regiões US West (Oregon) e US East (N. Virginia).

Resultado 1: tempo de download cai de minutos para segundos em todos os tamanhos

O download padrão cresce proporcionalmente ao tamanho da imagem, pois cada arquivo precisa ser descomprimido e gravado. Com lazy loading, o tempo ficou entre 3 e 11 segundos independentemente do tamanho — uma melhoria de 8x a 46x. Quanto maior a imagem, mais minutos são economizados por job.

Tamanho da imagem (comprimida)	Download padrão	Lazy loading	Melhoria
1,8 GB	48 s	3–11 s	~8x
13,2 GB	3 min 42 s	~6 s	~34x
35,4 GB	7 min 17 s	3–11 s	~38x

Resultado 2: a melhoria se mantém em todos os tipos e tamanhos de instância

Testando a imagem de 13,2 GB em 13 combinações de tipo e tamanho de instância, o download padrão variou entre 3 min 16 s e 4 min 20 s, enquanto o lazy loading se manteve entre 6 e 8 segundos — cerca de 34x mais rápido em todas as configurações. Não é necessário mudar o tipo de instância para se beneficiar.

Resultado 3: performance de treinamento não é afetada (além de uma penalidade na primeira época)

A AWS também executou jobs de treinamento completos para comparar a performance época a época. Os workloads testados foram:

Classificação de imagens com ResNet-50 no dataset Imagenette (9.469 imagens), usando containers PyTorch (13,2 GB) e TensorFlow (9,9 GB), em instâncias ml.p5.48xlarge, ml.p4d.24xlarge, ml.g7e.48xlarge, ml.g7.48xlarge, ml.g6e.12xlarge, ml.g6e.24xlarge e ml.g6e.48xlarge.
Fine-tuning de modelo de linguagem (LLM): ajuste fino supervisionado completo do Qwen2.5-0.5B no dataset alpaca-cleaned (51.760 exemplos) com o Hugging Face Trainer, usando container PyTorch da Hugging Face (9,4 GB), em instâncias ml.p5.48xlarge, ml.p4d.24xlarge e ml.g6.48xlarge.

A partir da segunda época, o tempo por época é idêntico entre os dois métodos em todas as instâncias e workloads. A primeira época leva entre 3 e 16 segundos a mais com a imagem lazy-loaded, enquanto os arquivos acessados pelo framework pela primeira vez são buscados. Esse custo inicial é fixo e independente do tamanho da instância. O resultado líquido: cada job terminou entre 1 e 6 minutos mais cedo — que é a economia no download menos o custo único da primeira época.

Vale considerar que o lazy loading traz maior ganho quando a primeira etapa de treinamento precisa de apenas uma fração da imagem — o restante é transmitido enquanto o treinamento avança. Se o entry point lê a maior parte da imagem antes de começar (por exemplo, fazendo scan, checksum ou pré-carregamento de conteúdo), a melhoria líquida será menor. Mesmo em cenários mais “pesados” no início — como uso de torch.compile ou extensões CUDA compiladas em tempo real, que adicionaram 7 a 15 segundos à primeira época —, o ganho líquido ainda foi de 3 min 25 s a 3 min 56 s. Um modelo de 6 GB embutido na imagem adicionou 2 a 3,5 minutos à primeira época, mas o job ainda terminou 2 a 2,5 minutos antes.

Como adotar o SOCI indexing

Há duas formas de usar imagens com índice SOCI no SageMaker AI.

Opção 1: usar um AWS Deep Learning Container já indexado

A AWS publica variantes com índice SOCI de vários Deep Learning Containers (DLCs) populares, identificadas pelo sufixo -soci na tag da imagem. Exemplo:

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

Variantes pré-indexadas estão disponíveis para PyTorch Training, TensorFlow Training, vLLM e SGLang, incluindo variantes arm64 para alguns frameworks. Para adotar, basta atualizar a tag da imagem na definição do job para a variante -soci. Consulte a lista completa na AWS Deep Learning Containers Gallery.

Atenção: se você constrói sua própria imagem a partir de uma DLC base, é necessário criar um índice SOCI para a imagem final. O índice da imagem base não é herdado por imagens derivadas.

Opção 2: criar um índice SOCI para sua própria imagem

Para indexar uma imagem customizada, é necessário um host Linux (como uma instância Amazon EC2 com Amazon Linux 2023) com containerd, o SOCI snapshotter e o Finch ou nerdctl instalados.

Defina as variáveis de ambiente:

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"

Instale as ferramentas necessárias:

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

Autentique-se no Amazon ECR:

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

Faça o pull da imagem original:

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

Crie a imagem com índice SOCI:

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

Faça o push da imagem indexada:

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

Esse processo cria dois artefatos no repositório: o índice SOCI e um manifesto de índice OCI vinculando a imagem ao índice. A imagem original e seu digest permanecem inalterados.

Automatizando a criação de índices no pipeline de CI/CD

Para gerar índices automaticamente a cada push, é possível implantar o CFN ECR SOCI Index Builder. Ele usa uma regra do Amazon EventBridge para detectar eventos de push no Amazon ECR e construir o índice automaticamente, sem necessidade de alterações no pipeline de build.

Usando a imagem indexada no job do SageMaker AI

A imagem com índice SOCI é referenciada exatamente como qualquer outra imagem — nenhuma mudança na API é necessária. Apenas a tag difere. Exemplo com o 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()

Verificando se a imagem possui índice SOCI

Antes de executar o job, é possível confirmar se a tag aponta para uma imagem indexada:

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

Procure por um manifesto filho cujo artifactType comece com application/vnd.amazon.soci.index. ao lado do manifesto de plataforma. Se não estiver presente, a tag não aponta para uma imagem indexada.

Pontos de atenção

Somente repositórios privados do Amazon ECR. Imagens em outros registries ou acessadas via Docker registry privado (TrainingRepositoryAccessMode definido como Vpc) não são elegíveis para lazy loading.
A melhoria escala com o tamanho da imagem. Imagens menores que aproximadamente 4 GB têm ganhos relativamente pequenos.
Verifique dependências específicas do Docker. Imagens com índice SOCI rodam em um runtime baseado em containerd. Revise entry points e scripts de setup que chamam o Docker CLI, leem /var/run/docker.sock ou derivam identificadores de caminhos de cgroup. O SageMaker AI monta um marcador /.dockerenv para compatibilidade retroativa, mas evite criar novas dependências nele.
Disponibilidade global. O SOCI indexing para jobs do SageMaker AI está disponível em todas as regiões AWS.

Conclusão

Com o suporte a índices SOCI, o Amazon SageMaker AI passa a oferecer uma forma direta de eliminar o overhead de download de imagens de container em jobs de treinamento e processamento. A adoção é simples — basta apontar para um DLC pré-indexado ou criar o índice para a própria imagem — e não exige nenhuma mudança nas definições de job, chamadas de API ou pipelines de dados existentes. O resultado é um início de job em segundos em vez de minutos, menor custo total e ciclos de iteração de machine learning (ML) mais rápidos.

Para aprofundar o tema, a AWS disponibiliza recursos adicionais: a Deep Learning Containers Gallery com a lista de imagens pré-indexadas, o guia de orientações específicas por framework em SOCI index for AWS Deep Learning Containers, o passo a passo no SageMaker AI Developer Guide, o guia de introdução ao SOCI snapshotter para opções avançadas, o CFN ECR SOCI Index Builder para automação no CI/CD, o walkthrough de lazy loading com Finch para testes locais, e os posts técnicos Reducing container cold start times using SOCI index on DLAMI and DLC e Under the hood: Lazy loading container images with Seekable OCI and AWS Fargate para quem quiser entender o mecanismo em profundidade.

Fonte

Announcing SOCI-indexed images for faster SageMaker AI jobs (https://aws.amazon.com/blogs/machine-learning/announcing-soci-indexed-images-for-faster-sagemaker-ai-jobs/)