Usá AMIs personalizadas en clústeres Slurm de SageMaker HyperPod sin perder seguridad

Introducción

Si tu equipo de ML/AI necesita clústeres Slurm en SageMaker HyperPod con entornos preconfigurados, seguros y auditables, pero no querés lidiar con scripts de lifecycle que ralentizan el despliegue y generan inconsistencias entre nodos, esta guía es para vos. Hasta ahora, la única opción era partir de las AMIs base de HyperPod y customizarlas en runtime con scripts complejos. Eso cambia: ahora podés construir AMIs personalizadas una sola vez, incorporar políticas de seguridad, agentes de monitoreo, drivers propietarios y librerías internas, y luego usarlas para levantar clústeres Slurm en segundos sin sacrificar compatibilidad con las librerías de entrenamiento distribuido.

En este artículo vas a aprender a:

1. Crear una AMI personalizada a partir de las AMIs base de HyperPod.
2. Validar que la AMI sea compatible con SageMaker HyperPod.
3. Desplegar un clúster Slurm usando la AMI personalizada mediante la API de SageMaker.
4. Actualizar clústeres existentes sin downtime.

Qué es y para qué sirve

Amazon SageMaker HyperPod con soporte para AMIs personalizadas en clústeres Slurm permite que los equipos de infraestructura y seguridad definan, una vez, el entorno completo de los nodos del clúster: sistema operativo parcheado, agentes de seguridad (como CrowdStrike, Tanium o Qualys), drivers de GPU específicos, librerías internas de ML y cualquier paquete de cumplimiento (ej: OpenSCAP, CIS Benchmarks). Al crear la AMI con estas dependencias preinstaladas y configuradas, el tiempo de arranque del clúster se reduce de minutos a segundos, ya que SageMaker solo necesita clonar la AMI en lugar de ejecutar scripts de configuración en cada nodo.

• Equipos de seguridad que necesitan incorporar agentes de monitoreo o escaneo de vulnerabilidades antes de que los nodos arranquen.
• Clústeres para entornos regulados (HIPAA, PCI-DSS, SOC2) donde cada nodo debe cumplir con políticas de hardening predefinidas.
• Equipos de ML que usan drivers propietarios de GPU o librerías internas no disponibles en los repositorios públicos de AWS.
• Reducción de la ventana de exposición a vulnerabilidades al evitar scripts de post-arranque que dependen de repositorios externos.

• La AMI personalizada debe basarse en las AMIs base públicas de HyperPod para garantizar compatibilidad con:

– Las capacidades de gestión de clústeres de Slurm (ej: integración con slurmctld).

– Los binarios de HyperPod Agent (hyperpod-agent) y neuron-agent.

• No podés usar AMIs genéricas de EC2 (ej: Amazon Linux 2023) porque perderías la integración con SageMaker HyperPod.
• La AMI debe estar disponible en la misma región donde desplegás el clúster.

Prerequisitos

Antes de empezar, asegurate de tener:

• Si usás Terraform, asegurate de tener el provider aws actualizado a la versión que soporte el recurso aws_sagemaker_hyperpod_custom_image.
• Si tu organización usa VPC endpoints privados, configurá el endpoint com.amazonaws.<región>.sagemaker.api para evitar que la creación de la AMI falle por timeout.
• Costo: La creación de AMIs no tiene costo adicional, pero almacenar AMIs en tu cuenta tiene un costo de $0.05 por GB/mes. Una AMI típica de HyperPod personalizada suele pesar entre 10 y 20 GB.

Guía paso a paso

Paso 1: Identificá la AMI base de HyperPod para tu región

SageMaker HyperPod publica AMIs base en cada región soportada. Usá la última versión disponible para evitar problemas de compatibilidad.

# Listá las AMIs base de HyperPod en us-east-1 (ajustá la región según necesites)
aws ec2 describe-images \
  --owners amazon \
  --filters "Name=name,Values=*sagemaker-hyperpod-base-slurm*" \
  --query "Images[].[ImageId,Name,CreationDate]" \
  --region us-east-1 \
  --output table

------------------------------------------------------------
|                         DescribeImages                   |
+-------------------+--------------------------------------+
|  ImageId          |  Name                                |
+-------------------+--------------------------------------+
|  ami-0abc12345678 |  sagemaker-hyperpod-base-slurm-2.16  |
|  ImageId          |  CreationDate                        |
+-------------------+--------------------------------------+
|  ami-1def2345678 |  2024-07-10                         |
+-------------------+--------------------------------------+

Si no ves AMIs con el filtro sagemaker-hyperpod-base-slurm*, asegurate de que:

• Estés usando la región correcta (ej: ap-southeast-1 en lugar de us-east-1).
• La cuenta AWS tenga acceso a SageMaker HyperPod en esa región.

Paso 2: Creá una AMI personalizada a partir de la AMI base

Usá Packer para automatizar la creación de la AMI con tus personalizaciones. Este ejemplo usa un custom.pkr.hcl que:

• Instala un agente de seguridad (example-security-agent).
• Configura un módulo de cumplimiento (cis-hardening).
• Agrega un driver de GPU propietario (nvidia-550.54.15).
• Parchea el sistema operativo.

# custom.pkr.hcl
packer {
  required_plugins {
    amazon = {
      version = ">= 1.3.0"
      source  = "github.com/hashicorp/amazon"
    }
  }
}

source "amazon-ebs" "hyperpod-custom" {
  region             = "us-east-1"
  source_ami_filter {
    filters = {
      name                = "sagemaker-hyperpod-base-slurm-*"
      root-device-type    = "ebs"
      virtualization-type = "hvm"
      architecture       = "x86_64"
    }
    owners = ["amazon"]
    most_recent = true
  }
  instance_type = "m5.2xlarge"
  ssh_username  = "ec2-user"
  communicator  = "ssh"
  ssh_timeout   = "20m"
  tags = {
    Name        = "hyperpod-slurm-custom-ami"
    Environment = "production"
    Owner       = "infra-team"
  }
}

build {
  sources = ["source.amazon-ebs.hyperpod-custom"]

  provisioner "shell" {
    inline = [
      "sudo yum update -y",
      "sudo yum install -y example-security-agent cis-hardening nvidia-driver-550.54.15",
      "sudo systemctl enable example-security-agent",
      "sudo dnf config-manager --add-repo https://internal-repo.example.com/nvidia.repo",
      "sudo dnf install -y nvidia-driver-550.54.15"
    ]
  }

  provisioner "file" {
    source      = "files/security-policies/"
    destination = "/etc/security/"
  }

  provisioner "shell" {
    inline = [
      "sudo cp /etc/security/policies/*.conf /etc/security/",
      "sudo chmod 644 /etc/security/*.conf"
    ]
  }
}

packer init custom.pkr.hcl
packer validate custom.pkr.hcl
packer build custom.pkr.hcl

==> amazon-ebs.hyperpod-custom: AMIs were created:
==> amazon-ebs.hyperpod-custom: 	us-east-1: ami-abc1234567890

Si el build falla por timeout en ssh, verifica:

• Que tu IP pública esté en la lista de redes permitidas en el security group de la instancia.
• Que el security group permita tráfico SSH (puerto 22) desde tu IP.

Paso 3: Validá la compatibilidad de la AMI con HyperPod

Antes de usar la AMI en producción, verificá que sea compatible con HyperPod. SageMaker HyperPod requiere que la AMI cumpla con:

1. Tener instalado el agente de HyperPod (/opt/aws/hyperpod-agent/bin/hyperpod-agent).
2. Tener los binarios de Slurm 22.05 o superior en /usr/local/slurm.
3. Tener los permisos correctos en los directorios /opt/sagemaker y /var/lib/docker.

# Conectate a la AMI recién creada (reemplazá ami-abc1234567890)
INSTANCE_ID=$(aws ec2 run-instances \
  --image-id ami-abc1234567890 \
  --instance-type m5.large \
  --query "Instances[0].InstanceId" \
  --output text \
  --region us-east-1)

# Esperá a que la instancia esté en estado "running" y con IP asignada
sleep 60

IP=$(aws ec2 describe-instances \
  --instance-ids $INSTANCE_ID \
  --query "Reservations[0].Instances[0].PublicIpAddress" \
  --output text \
  --region us-east-1)

# Verificá la presencia de archivos críticos
ssh -o StrictHostKeyChecking=no ec2-user@$IP << 'EOF'
sudo test -f /opt/aws/hyperpod-agent/bin/hyperpod-agent && echo "✓ hyperpod-agent presente" || echo "✗ hyperpod-agent ausente"
sudo ls -l /usr/local/slurm/bin/slurmctld && echo "✓ slurmctld presente" || echo "✗ slurmctld ausente"
sudo ls -l /etc/systemd/system/hyperpod-agent.service && echo "✓ servicio systemd presente" || echo "✗ servicio systemd ausente"
EOF

# Terminá la instancia
aws ec2 terminate-instances --instance-ids $INSTANCE_ID --region us-east-1

✓ hyperpod-agent presente
✓ slurmctld presente
✓ servicio systemd presente

• Revisá que la AMI base sea la correcta (versión mínima 2.15).
• Verificá que los pasos de provisioning en Packer hayan instalado los componentes necesarios.

Paso 4: Desplegá un clúster Slurm usando la AMI personalizada

Usá la API de SageMaker para crear un clúster con la AMI personalizada. Este ejemplo usa boto3 para Python.

import boto3
import json

client = boto3.client("sagemaker", region_name="us-east-1")

response = client.create_cluster(
    ClusterName="ml-training-cluster",
    ClusterDescription="Clúster Slurm con AMI personalizada y hardening",
    ClusterConfiguration={
        "InstanceGroups": [
            {
                "InstanceGroupName": "gpu-group",
                "InstanceType": "ml.p3.8xlarge",
                "InstanceCount": 2,
                "InstanceRole": "worker",
                "CustomAMIId": "ami-abc1234567890",  # AMI personalizada
                "ResourceTags": [
                    {
                        "Key": "Environment",
                        "Value": "production"
                    }
                ]
            }
        ],
        "SharedStorageConfiguration": {
            "StorageType": "EFS",
            "EfsStorageConfiguration": {
                "FileSystemId": "fs-1234567890abcdef0"
            }
        }
    },
    NetworkConfiguration={
        "EnableInterContainerTrafficEncryption": True,
        "EnableNetworkIsolation": True,
        "VpcConfig": {
            "SecurityGroupIds": ["sg-0abc1234567890ab"],
            "Subnets": ["subnet-0abc1234567890ab"]
        }
    },
    ExecutionRoleArn="arn:aws:iam::123456789012:role/service-role/AmazonSageMaker-ExecutionRole",
    Tags=[
        {
            "Key": "ManagedBy",
            "Value": "hyperpod-custom-ami"
        }
    ]
)

print(json.dumps(response, indent=2, default=str))

{
  "ClusterArn": "arn:aws:sagemaker:us-east-1:123456789012:cluster/ml-training-cluster",
  "ClusterName": "ml-training-cluster",
  "ClusterStatus": "InService"
}

aws sagemaker describe-cluster --cluster-name ml-training-cluster --region us-east-1 --query "ClusterStatus"

• Clúster sin AMI personalizada: ~5-7 minutos.
• Clúster con AMI personalizada: ~1-2 minutos (solo clonación de la AMI).

Paso 5: Actualizá un clúster existente con una AMI personalizada

Si necesitas actualizar un clúster ya existente a una AMI personalizada (ej: para aplicar un parche de seguridad), usá la API UpdateCluster con el campo CustomAMIId.

import boto3

client = boto3.client("sagemaker", region_name="us-east-1")

response = client.update_cluster(
    ClusterName="ml-training-cluster",
    InstanceGroups=[
        {
            "InstanceGroupName": "gpu-group",
            "CustomAMIId": "ami-def9876543210"  # Nueva AMI personalizada
        }
    ]
)

print(response["ClusterArn"])

{
  "ClusterArn": "arn:aws:sagemaker:us-east-1:123456789012:cluster/ml-training-cluster"
}

• La actualización no requiere downtime: SageMaker reemplaza los nodos uno por uno, manteniendo la disponibilidad del clúster.
• Si el clúster usa Spot Instances, la actualización puede fallar si no hay capacidad disponible. En ese caso, configurá InstanceMarketOptions con MarketType: "on-demand" para los nodos críticos.

Consideraciones y buenas prácticas

Limitaciones conocidas

1. Tamaño máximo de AMI: SageMaker HyperPod no soporta AMIs mayores a 100 GB. Si tu AMI personalizada supera este límite, SageMaker fallará con el error ImageSizeExceedsLimit.

– Alternativa: Migrá componentes grandes (ej: datasets) a volúmenes EBS o S3 y montalos en runtime.

1. Dependencias externas: Si tu AMI personalizada depende de paquetes en repositorios externos (ej: nvidia.repo), asegurate de que:

– Los repositorios usen HTTPS (evitá http que pueden fallar por políticas de seguridad).

1. Regiones soportadas: No todas las regiones donde HyperPod está disponible soportan AMIs personalizadas. Verificá la lista actual en la documentación oficial.

Buenas prácticas de seguridad

• Firmá las AMIs: Usá AWS Signer para firmar tus AMIs personalizadas y evitar suplantación.
• Rotá las AMIs: Implementá un pipeline de CI/CD que:

2. Genere una nueva AMI con las correcciones.

3. Actualice los clústeres en staging primero, luego en producción.

• Restringí el acceso: Configurá políticas de IAM que solo permitan a roles específicos crear/usar AMIs personalizadas. Ejemplo:

  {
    "Version": "2012-10-17",
    "Statement": [
      {
        "Effect": "Deny",
        "Action": "ec2:CreateImage",
        "Resource": "*",
        "Condition": {
          "NotIpAddress": {
            "aws:SourceIp": ["192.0.2.0/24"]
          }
        }
      }
    ]
  }
  

Monitoreo y auditoría

• Logs de creación de AMIs: Habilitá los logs de CloudTrail para las APIs CreateImage, CopyImage y CreateCluster para auditar quién creó qué AMI.
• Alertas de cambios: Configurá una regla en CloudWatch para alertar si la AMI personalizada es modificada o eliminada.

  aws events put-rule \
    --name "hyperpod-ami-changes" \
    --event-pattern '{"source":["aws.ec2"],"detail-type":["AWS API Call via CloudTrail"],"detail":{"eventSource":["ec2.amazonaws.com"],"eventName":["CreateImage","CopyImage","DeregisterImage"]}}'
  

Alternativas si no podés usar AMIs personalizadas

Si tu organización no puede construir AMIs personalizadas (ej: por restricciones de compliance), considerá:

1. Scripts de lifecycle: Usá los hooks de OnCreate y OnStart de SageMaker HyperPod para ejecutar scripts de configuración en cada nodo.

1. Golden AMI en cuenta compartida: Crea una AMI personalizada en una cuenta AWS dedicada y compartila con las cuentas de producción usando IAM roles.
2. AWS Systems Manager (SSM): Usá SSM Run Command para aplicar configuraciones después del arranque del clúster.

Conclusión

La capacidad de usar AMIs personalizadas en clústeres Slurm de SageMaker HyperPod cierra una brecha crítica para equipos de infraestructura y seguridad: ahora podés desplegar entornos pre-aprobados, pre-hardened y consistentes en segundos, sin sacrificar la optimización de rendimiento de HyperPod. Al integrar políticas de seguridad, drivers y librerías internas en la AMI, eliminás la dependencia de scripts de post-arranque que ralentizan el despliegue y generan inconsistencias, mientras mantienes la compatibilidad con las librerías de entrenamiento distribuido de SageMaker.

1. Crea tu primera AMI personalizada usando el ejemplo de Packer proporcionado.
2. Valida la compatibilidad con el checklist del Paso 3.
3. Despliega un clúster de prueba en staging y verifica que los nodos arranquen sin errores.
4. Implementá un pipeline de CI/CD para escanear y rotar AMIs periódicamente.

Si tu equipo necesita ayuda para adaptar este flujo a tu entorno (ej: usar Terraform en lugar de Packer, o integrar con herramientas de compliance como Prisma Cloud), contactá a tu Customer Solutions Manager de AWS para revisar opciones de soporte técnico avanzado.