Anunciamos imágenes indexadas con SOCI para trabajos de SageMaker AI más rápidos

Cada minuto que tu imagen de contenedor de aprendizaje profundo pasa descargándose es un minuto que pagas y un minuto que tu equipo espera. La descarga de imágenes de contenedor es puro sobrecosto: tiempo facturable del trabajo que no produce ningún avance en el modelo. Para las imágenes modernas de aprendizaje profundo, ese sobrecosto son varios minutos en cada trabajo. Hoy, nos complace anunciar la indexación Seekable OCI (SOCI) para Amazon SageMaker AI Training y los trabajos de procesamiento de Amazon SageMaker AI, que elimina la descarga completa de la imagen de contenedor de la ruta crítica de cada trabajo.

Con la indexación SOCI, SageMaker AI carga tu imagen de contenedor de forma diferida. Solo obtiene los archivos necesarios para iniciar el contenedor y transmite el resto en segundo plano mientras se ejecuta tu código, de modo que tu trabajo comienza en segundos en lugar de minutos. Lo adoptas apuntando a un AWS Deep Learning Container (DLC) preindexado o agregando un índice a tu propia imagen. Ninguna de las dos opciones requiere cambios en la definición de tu trabajo, el punto de entrada o los canales de datos. El resultado son inicios de trabajo más rápidos, menor costo del trabajo y más tiempo dedicado a construir modelos en lugar de esperar a la infraestructura.

Las imágenes indexadas con SOCI inician los contenedores de entrenamiento de Amazon SageMaker AI en segundos, sin impacto en el rendimiento del entrenamiento. Como mostramos en la sección Lo que medimos , esto redujo el tiempo de inactividad de la GPU durante la descarga de imágenes hasta en un 90 por ciento en nuestros benchmarks.

Esta publicación muestra cómo adoptar la indexación SOCI: qué es, los resultados que medimos y cómo activarla para tus propios trabajos.

El desafío

Los trabajos de entrenamiento y procesamiento de SageMaker AI se ejecutan dentro de imágenes de contenedor almacenadas en Amazon Elastic Container Registry (Amazon ECR). Antes de que un trabajo pueda ejecutar tu código, el servicio descarga la imagen completa en la instancia. Descomprime y escribe cada capa en el disco.

Para imágenes pequeñas esto toma segundos. Pero las imágenes de aprendizaje profundo superan rutinariamente los 10 GB o más. Esa descarga añade minutos a cada trabajo, y tú la pagas como tiempo del trabajo. El impacto es mayor exactamente donde la velocidad de iteración importa más: barridos de ajuste de hiperparámetros con cientos de trabajos, experimentos rápidos y trabajos de procesamiento cortos, donde la descarga de la imagen puede dominar el tiempo total de ejecución.

Cómo funciona la indexación SOCI

Seekable OCI (SOCI) es una tecnología de código abierto que carga imágenes de contenedor de forma diferida. En lugar de descargar toda la imagen antes de iniciar, el runtime usa un índice SOCI para obtener solo los archivos necesarios para lanzar el contenedor de inmediato. Luego transmite el resto en segundo plano a medida que tu código lo accede.

Un índice SOCI son metadatos que mapean cada archivo dentro de tu imagen. Lo almacenas en Amazon ECR como un artefacto compatible con OCI vinculado a tu imagen original. Como el índice es un artefacto separado, el digest de tu imagen y cualquier firma existente siguen siendo válidos.

SageMaker AI busca automáticamente un índice

Cuando creas un trabajo de entrenamiento o procesamiento, SageMaker AI verifica si la imagen que especificaste tiene un índice SOCI y actúa en consecuencia:

Como los trabajos siempre retroceden a una descarga estándar cuando es necesario, agregar un índice SOCI nunca impide que tu trabajo se ejecute. Solo ganas un inicio más rápido. Nunca arriesgas romper un flujo de trabajo existente.

Lo que medimos

Descargamos las mismas imágenes de contenedor de dos maneras en trabajos de entrenamiento reales de SageMaker AI: descarga estándar y carga diferida con índice SOCI. Luego ejecutamos trabajos de entrenamiento completos desde ambas para medir el efecto en la velocidad de entrenamiento.

Configuración

• Ocho imágenes, 1.8–35.4 GB comprimidos.
• instancias de GPU y CPU ml.g4dn, ml.g5, ml.g6, ml.g6e, ml.g7, ml.g7e, ml.p4d, ml.p5 y ml.m7i, desde xlarge hasta 48xlarge, en la Región US West (Oregon) (us-west-2) y la Región US East (N. Virginia) (us-east-1).
• Tiempo de descarga medido desde el inicio de la extracción de la imagen hasta que el contenedor está listo para ejecutarse. Todas las cifras son medianas de ejecuciones repetidas.

Resultado 1: el tiempo de descarga cae de minutos a segundos en cada tamaño de imagen

El siguiente gráfico compara el tiempo de descarga de imágenes de contenedor de un Docker pull estándar frente a la carga diferida de SOCI en ocho tamaños de imagen, desde 1.8 GB hasta 35.4 GB.

Cada tamaño de imagen muestra dos barras. La barra del tiempo de descarga de Docker (azul) crece constantemente con el tamaño de la imagen, desde 48 segundos con 1.8 GB hasta 7 minutos 17 segundos con 35.4 GB.

Una etiqueta de aceleración sobre cada par muestra la relación entre el tiempo de Docker y el tiempo de SOCI. Las aceleraciones van desde aproximadamente 10x para la imagen más pequeña hasta más de 35x para la más grande.

Figura 1: Tiempo de descarga de imágenes de contenedor, un Docker pull estándar en comparación con la carga diferida de SOCI, en ocho tamaños de imagen

El tiempo de descarga estándar crece con el tamaño de la imagen porque cada archivo se descomprime y se escribe. La carga diferida se mantuvo en 3–11 segundos independientemente del tamaño, una mejora de 8x–46x. Cuanto mayor sea la imagen, más minutos ahorras.

Resultado 2: El ahorro se mantiene en cada tipo y tamaño de instancia

Al ejecutar la imagen de 13.2 GB en 13 combinaciones de tipo/tamaño de instancia, la descarga estándar varió entre 3 min 16 s y 4 min 20 s, mientras que la carga diferida se mantuvo en 6-8 s. Eso es aproximadamente 34 veces más rápido en cada configuración. No necesitas cambiar el tipo ni el tamaño de instancia para beneficiarte.

Figura 2: Tiempo de descarga de una imagen de 13.2 GB en 13 combinaciones de tipo y tamaño de instancia

Resultado 3: El rendimiento del entrenamiento no se ve afectado más allá de una penalización en la primera época

Ejecutamos dos tipos de trabajos de entrenamiento completos desde la imagen obtenida con Docker y desde la imagen indexada con SOCI, y comparamos cada época:

• Clasificación de imágenes en ResNet-50 sobre Imagenette (9,469 imágenes)
  • Imágenes de contenedor: PyTorch (DLC de 13.2 GB) y TensorFlow (DLC de 9.9 GB)
  • Instancias: ml.p5.48xlarge, ml.p4d.24xlarge, ml.g7e.48xlarge, ml.g7.48xlarge, ml.g6e.12xlarge, ml.g6e.24xlarge y ml.g6e.48xlarge.
• Ajuste fino de modelos de lenguaje de gran tamaño (LLM): ajuste fino supervisado (SFT) de parámetros completos de Qwen2.5-0.5B sobre alpaca-cleaned (51,760 ejemplos) con el Hugging Face Trainer.
  • Imagen de contenedor: DLC de Hugging Face PyTorch de 9.4 GB.
  • Instancias: ml.p5.48xlarge, ml.p4d.24xlarge, ml.g6.48xlarge.

El siguiente gráfico desglosa el inicio del trabajo en el tiempo de descarga de la imagen y el tiempo de la primera época de entrenamiento, comparando una extracción de Docker estándar con una imagen indexada con SOCI en configuraciones de GPU y CPU para ambas cargas de trabajo.

Figura 3: Inicio del trabajo dividido en el tiempo de descarga de la imagen y el tiempo de la primera época de entrenamiento, un Docker pull estándar en comparación con una imagen indexada con SOCI

• Desde la segunda época en adelante, el tiempo por época es el mismo (dentro de la variación de una ejecución a otra) en cada instancia, para ambas cargas de trabajo.
• La primera época tarda entre 3 y 16 s más desde la imagen de carga diferida, mientras se obtienen los archivos que el marco de trabajo toca por primera vez. Este costo de la primera época es el mismo en cada tamaño de instancia.
• Resultado neto: cada trabajo terminó entre 1 y 6 min antes. Ese es el ahorro en la descarga menos el costo único de la primera época, un ahorro fijo por trabajo.

Prueba esto con tu propia carga de trabajo

Estos números provienen de las dos cargas de trabajo descritas anteriormente. El costo de la primera época depende de tu imagen y de lo que tu código toca al inicio:

• La carga diferida gana cuando el primer paso de entrenamiento necesita solo una fracción de la imagen, y el resto se transmite mientras entrenas.
• Si tu punto de entrada lee la mayor parte de la imagen antes de que comience el entrenamiento (por ejemplo, escaneándola, calculando sumas de verificación o precargando su contenido), la primera época esperará aproximadamente los mismos bytes que habría obtenido un pull completo. La mejora neta sería pequeña.
• Los inicios con carga anticipada también resultaron favorables en nuestras pruebas:
  • torch.compile y una extensión CUDA compilada justo a tiempo (JIT) añadieron entre 7 y 15 s a la primera época (el ahorro neto sigue siendo de 3 min 25 s a 3 min 56 s)
  • Un modelo de 6 GB integrado en la imagen transmitió sus pesos al momento de la carga, añadiendo entre 2 y 3.5 min, pero el trabajo en general aún terminó entre 2 y 2.5 mins antes en la imagen indexada.

Cómo empezar

Puedes adoptar la indexación con SOCI de dos maneras: usar un AWS Deep Learning Container preindexado o crear un índice para tu propia imagen personalizada.

Opción 1: Usar un AWS Deep Learning Container preindexado

AWS publica variantes indexadas con SOCI de varios Deep Learning Containers populares, identificadas por el sufijo -soci en la etiqueta de la imagen:

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

Las variantes preindexadas están disponibles para PyTorch Training, TensorFlow Training, vLLM y SGLang, incluidas variantes arm64 para algunos marcos de trabajo. Para adoptar la carga diferida, actualiza la etiqueta de la imagen en tu definición de trabajo a la variante -soci. Consulta la lista completa en el AWS Deep Learning Containers Gallery.

Nota: Si construyes tu propia imagen FROM una imagen base de DLC, debes crear un índice SOCI para tu imagen final. El índice de la imagen base no se traslada a las imágenes derivadas.

Opción 2: Crear un índice SOCI para tu propia imagen

Para indexar una imagen personalizada, necesitas un host Linux (como una instancia de Amazon Elastic Compute Cloud (Amazon EC2) que ejecute Amazon Linux 2023) con containerd, el SOCI snapshotter, y Finch o nerdctl instalado.

1. Configura tus variables de entorno.

   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. Instala las herramientas necesarias.

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

3. Autentícate en Amazon ECR.

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

4. Extraiga su imagen original.

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

5. Cree la imagen indexada con SOCI.

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

6. Publique (push) la imagen indexada.

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

Esto crea dos artefactos en su repositorio: el índice SOCI y un manifiesto de índice de imagen compatible con OCI que vincula su imagen con el índice. Su imagen original y su digest permanecen sin cambios.

Automatice la creación de índices en CI/CD

Para generar índices automáticamente en cada push, implemente el CFN ECR SOCI Index Builder. Este utiliza una regla de Amazon EventBridge para detectar eventos de push de imágenes en Amazon ECR y construir el índice. No se requieren cambios en su pipeline de compilación.

Utilice la imagen indexada en su trabajo de SageMaker AI

Referencie su imagen indexada con SOCI exactamente como lo haría con cualquier otra imagen. No se necesitan cambios en la API. Solo cambia la etiqueta (tag).

Ejemplo de entrenamiento (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()

Verifique que su imagen incluya un índice SOCI

Antes de lanzar, confirme que su etiqueta resuelve a una imagen 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}'

Busque un manifiesto secundario cuyo artifactType comience con application/vnd.amazon.soci.index. junto a su manifiesto de plataforma. Si no está allí, la etiqueta no apunta a una imagen indexada.

Aspectos a tener en cuenta

• Solo repositorios privados de Amazon ECR. Las imágenes en otros registros, o extraídas a través de un registro Docker privado (TrainingRepositoryAccessMode establecido en Vpc), no son aptas para la carga diferida.
• La mejora escala con el tamaño de la imagen. Las imágenes menores a aproximadamente 4 GB ven mejoras relativamente pequeñas.
• Verifique las suposiciones específicas de Docker. Las imágenes indexadas con SOCI se ejecutan en un tiempo de ejecución basado en containerd. Revise los puntos de entrada y los scripts de configuración en busca de código que llame a la CLI de Docker, lea /var/run/docker.sock o derive identificadores de rutas de cgroup. SageMaker AI monta un marcador /.dockerenv por compatibilidad con versiones anteriores, pero evite agregar nuevas dependencias en él.
• La indexación SOCI para trabajos de SageMaker AI está disponible en todas las Regiones de AWS.

Conclusión

Con la indexación SOCI, puede minimizar el tiempo de descarga de imágenes de contenedor en sus trabajos de entrenamiento y procesamiento de SageMaker AI. Agregue un índice SOCI a sus imágenes, o cambie a un AWS Deep Learning Container preindexado. De cualquier manera, reduce el tiempo de inicio de minutos a segundos, disminuye el costo total del trabajo y acelera su ciclo de iteración de aprendizaje automático (ML). Y sus definiciones de trabajo, llamadas de API y pipelines de datos existentes permanecen exactamente iguales.

Para obtener más información sobre Amazon SageMaker AI y comenzar a usar la indexación SOCI, explore los siguientes recursos:

1. Elija un Deep Learning Container indexado con SOCI. Explore la galería de Deep Learning Containers y agregue -soci a la etiqueta de su imagen. Para obtener orientación específica del framework, consulte índice SOCI para AWS Deep Learning Containers.
2. Indexe su propia imagen personalizada. Siga el tutorial paso a paso en la Guía para desarrolladores de SageMaker AI, o consulte la guía de inicio de SOCI snapshotter para opciones avanzadas.
3. Automatice la indexación en su pipeline de CI/CD. Implemente el CFN ECR SOCI Index Builder para generar índices automáticamente en cada envío de imagen. No se requieren cambios en el pipeline.
4. Pruebe localmente con Finch. Use el Guía de lazy-loading con Finch para validar tu imagen en un runtime basado en containerd antes de enviar un trabajo.
5. Profundiza en la tecnología. Lee Reducing container cold start times using SOCI index on DLAMI and DLC en el AWS Machine Learning Blog para ver benchmarks, o explora Bajo el capó: Carga diferida de imágenes de contenedor con Seekable OCI y AWS Fargate en el AWS Containers Blog para entender en detalle el mecanismo de lazy-loading.

Un inicio más rápido de contenedores está a solo una etiqueta de imagen de distancia. Pruébalo en tu próximo trabajo de SageMaker AI y nota la diferencia.

Sobre los autores