Con el lanzamiento del soporte de AWS Batch para trabajos de entrenamiento de Amazon SageMaker, los clientes obtuvieron una nueva forma completamente administrada de poner en cola, programar y reintentar cargas de trabajo de entrenamiento de aprendizaje automático (ML) sin tener que construir su propia infraestructura. En su lugar, utilizan las características de cola y programación de AWS Batch. A medida que las organizaciones escalan sus operaciones de inteligencia artificial y aprendizaje automático (AI/ML) y contratan a más equipos, coordinar el acceso a la capacidad compartida de procesamiento gráfico (GPU) puede introducir sobrecarga adicional de comunicación, coordinación y configuración para los administradores. Esta sobrecarga incluye la configuración de colas separadas, el ajuste de los límites de capacidad y la orquestación cuidadosa de las cargas de trabajo adecuadas para obtener los recursos adecuados en el momento adecuado. Un valor central de la propuesta de AWS Batch es ayudar a los clientes a aumentar la utilización de su capacidad de cómputo, y a medida que crecen los entornos de entrenamiento, también lo hace la oportunidad de hacer más con la capacidad existente.

Trabajando hacia atrás desde los clientes que administran entornos de entrenamiento multi-equipo, escuchamos una solicitud consistente: dános una forma de definir asignaciones por equipo o proyecto, compartir la capacidad inactiva automáticamente y ejecutar nuestras cargas de trabajo de mayor prioridad primero. Esta solicitud nos llevó a construir administración de cuotas, una nueva capacidad de programación para AWS Batch. Los administradores pueden utilizarla para dividir la capacidad de cómputo en acciones de cuota, organizadas por equipo, proyecto o carga de trabajo. Cada acción de cuota utiliza estrategias de compartición de recursos configurables que maximizan la utilización de GPU mientras honran la asignación de cada equipo cuando más importa.

Esta publicación muestra cómo funciona la administración de cuotas, incluyendo cómo configurar acciones de cuota con estrategias de préstamo y préstamo, configurar la preemption dentro de la acción para cargas de trabajo de prioridad y comenzar a usar el SDK de Python de SageMaker.

Acciones de cuota y preemption

Los administradores pueden asignar y compartir recursos de GPU limitados definiendo una acción de cuota para entidades separadas, como equipos o proyectos individuales. Cada acción de cuota incluye cuotas de cómputo (límites de capacidad) y estrategias de compartición de recursos para cómputo inactivo. Estas estrategias determinan si el cómputo inactivo se reserva hasta que lleguen los trabajos (RESERVE), se presta a otras acciones de cuota que tienen trabajo disponible (LEND), o se presta cuando está inactivo y se solicita cuando la acción tiene trabajo que podría beneficiarse (LEND_AND_BORROW).

Una cuota compartida opera como una cola virtual anidada jerárquicamente dentro de una cola de trabajos asociada. Al programar trabajos para una cola de trabajos, AWS Batch prioriza las cuotas compartidas que están bajo sus límites de capacidad a medida que itera a través de las cuotas compartidas asociadas. Si los trabajos llegan a una cuota compartida que ha prestado su capacidad, AWS Batch realiza una preemption entre cuotas compartidas. Esto recupera recursos para la cuota compartida prestamista, por lo que los administradores no necesitan rastrear los trabajos que prestan capacidad y terminarlos por su cuenta. Si los trabajos llegan a una cuota compartida que está en o por encima de su límite de capacidad, AWS Batch verifica si hay capacidad disponible para prestar de otra cuota compartida. Si es así, las cargas de trabajo pueden ejecutarse utilizando la capacidad que de otro modo estaría inactiva.

Puede activar la preemption basada en prioridad dentro de una cuota compartida cuando configure una cuota compartida. Donde la preemption entre cuotas compartidas recupera la capacidad prestada a medida que llegan trabajos a una cuota compartida prestamista, la preemption dentro de una cuota compartida anula los trabajos de prioridad inferior que están en el estado PROGRAMADO, INICIANDO o EJECUTANDO. Esto permite que los trabajos de prioridad más alta al frente de la cola comiencen a ejecutarse antes.

Requisitos previos

Antes de continuar con este ejemplo, debe tener una cuenta de AWS con permisos para crear y administrar recursos de AWS Batch. Para este ejemplo, consulte los permisos de IAM de ejemplo necesarios para invocar las API de AWS Batch y SageMaker.

El ejemplo utiliza el SDK de Python de SageMaker v3. Complete la configuración necesaria para instalar una versión compatible del SDK de Python de SageMaker:

!pip install sagemaker

Configuración de recursos de administración de cuotas

Usar las nuevas capacidades de administración de cuotas requiere crear una Scheduling Policy que contiene un quotaSharePolicy, un Service Environment que especifica sus límites de capacidad en términos de instancias específicas, y una cola de trabajos de entrenamiento de SageMaker que está conectada a ambos. Después de eso, puede asociar acciones de cuota a una cola de trabajos. La guía del usuario tiene ejemplos de AWS Command Line Interface (AWS CLI) y consola para configurar recursos de administración de cuotas y acciones de cuota. Un ejemplo de extremo a extremo de administración de cuotas integrada con el módulo aws_batch en el SDK de Python de SageMaker v3 se contiene en una libreta de ejemplo. En este ejemplo, pasamos por fragmentos de código que resumen los componentes clave y usamos las funciones de ayuda de la libreta para crear una Scheduling Policy, Service Environment, Job Queue y tres Acciones de Cuota.

Las tres acciones de cuota configuradas aquí representan tres proyectos diferentes. Cada una tiene el mismo límite de capacidad establecido: una sola instancia ml.g5.xlarge. La acción de cuota QS1 representa trabajos en segundo plano. Para este proyecto, desea ejecutar el trabajo en segundo plano más importante primero y tomar tantas capacidades de repuesto como estén disponibles. Dentro de esta acción de cuota, no hay trabajos urgentes y las preempciones son aceptables porque todos los trabajos utilizan puntos de control. Esta acción de cuota utiliza la estrategia de LEND_AND_BORROW de compartición de recursos y tiene la preempción dentro de la acción de cuota activada. Los trabajos de mayor prioridad dentro de la misma acción de cuota pueden entonces preemptions a los trabajos en ejecución que tienen menor prioridad.

from sagemaker.train.aws_batch.boto_client import get_batch_boto_client
from utils.aws_batch_resource_management import AwsBatchResourceManager, QuotaShareConfig, create_quota_management_resources

INSTANCE_TYPE = "ml.g5.xlarge"

# Create SchedulingPolicy, ServiceEnvironment, JobQueue, and QuotaShares
resource_manager = AwsBatchResourceManager(get_batch_boto_client())
qs1 = QuotaShareConfig(
    name="QS1",
    capacity_unit=INSTANCE_TYPE,
    max_capacity=1,
    in_share_preemption=True,
    sharing_strategy="LEND_AND_BORROW",
    borrow_limit=200
)

QS2 representa cargas de trabajo de prioridad media. Los equipos que eligen esta configuración prefieren esperar a que se disponga de capacidad en lugar de soportar preemptions adicionales. Esas preemptions provienen de la capacidad prestada que se reclama o de preemptions dentro de la participación debido a la presentación de trabajos de prioridad más alta. Esto es adecuado para trabajos más cortos, o cuando la mayoría de los trabajos presentados tienen prioridades similares. La participación de cuota aún se beneficia de que se tenga en cuenta la prioridad del trabajo en el momento del despacho del trabajo. Los trabajos presentados desencadenan que AWS Batch recupere cualquier capacidad prestada necesaria para ejecutar los trabajos.

qs2 = QuotaShareConfig(
    name="QS2",
    capacity_unit=INSTANCE_TYPE,
    max_capacity=1,
    in_share_preemption=False,
    sharing_strategy="LEND"
)

QS3 representa cargas de trabajo que necesitan iniciarse lo más rápido posible con mínimas interrupciones. Está configurado para reservar capacidad y no prestar ni tomar prestada capacidad a otras participaciones de cuota. No experimenta preemptions entre participaciones, porque sus cargas de trabajo no exceden los límites de capacidad. Al desactivar también la preemption dentro de la participación, las cargas de trabajo presentadas a esta participación de cuota no son interrumpidas por AWS Batch siempre y cuando el límite de capacidad no cambie.

qs3 = QuotaShareConfig(
    name="QS3",
    capacity_unit=INSTANCE_TYPE,
    max_capacity=1,
    in_share_preemption=False,
    sharing_strategy="RESERVE"
)

Estas configuraciones de participación de cuota son entradas para la función create_quota_management_resources en la libreta de muestra, que también configura la política de programación, el entorno de servicio y la cola de trabajos correspondientes con los nombres constantes elegidos.

SCHEDULING_POLICY_NAME = "my-qm-scheduling-policy"
JOB_QUEUE_NAME = "my-sm-training-qm-jq"
SERVICE_ENVIRONMENT_NAME = "my-sm-training-qm-se"

resources = create_quota_management_resources(
    resource_manager=resource_manager,
    scheduling_policy_name=SCHEDULING_POLICY_NAME,
    job_queue_name=JOB_QUEUE_NAME,
    service_environment_name=SERVICE_ENVIRONMENT_NAME,
    capacity_unit=INSTANCE_TYPE,
    max_capacity=3,
    quota_share_configs=[qs1, qs2, qs3]
)

Presentar trabajos de entrenamiento de SageMaker a QS1

El módulo aws_batch dentro del SDK de Python de SageMaker se ha extendido para permitir la presentación programática de trabajos de entrenamiento de SageMaker a una cola de trabajos de AWS Batch habilitada para la administración de cuotas. Para demostrar estas nuevas capacidades, usamos un ejemplo de entrenador de modelo de "hola mundo" ModelTrainer que se presenta a la cola de trabajos. El siguiente fragmento de código identifica una URI de imagen para usar en función de la región y el tipo de instancia seleccionados.

import logging

logging.basicConfig(
    level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
logging.getLogger("botocore.client").setLevel(level=logging.WARN)
logger = logging.getLogger(__name__)

from sagemaker.core.helper.session_helper import Session
from sagemaker.core import image_uris

session = Session()

image_uri = image_uris.retrieve(
    framework="pytorch",
    region=session.boto_session.region_name,
    version="2.5",
    instance_type=INSTANCE_TYPE,
    image_scope="training",
)

Usamos esta URI de imagen dentro de un ejemplo básico de entrenador de modelo de "hola mundo":

from sagemaker.train.model_trainer import ModelTrainer
from sagemaker.train.configs import SourceCode, Compute, StoppingCondition

TRAINING_JOB_NAME = "hello-world-simple-job"
MAX_RUN_TIME = 300

source_code = SourceCode(command="echo 'Hello World'")

model_trainer = ModelTrainer(
    training_image=image_uri,
    source_code=source_code,
    base_job_name=TRAINING_JOB_NAME,
    compute=Compute(instance_type=INSTANCE_TYPE, instance_count=INSTANCE_COUNT),
    stopping_condition=StoppingCondition(max_runtime_in_seconds=MAX_RUN_TIME),
)

Para usar la cola de trabajos dentro del SDK de Python de SageMaker, proporcione el nombre de la cola al constructor TrainingQueue dentro del módulo aws_batch.

from sagemaker.train.aws_batch.training_queue import TrainingQueue, TrainingQueuedJob

# Construct the queue object using the SageMaker Python SDK
queue = TrainingQueue(JOB_QUEUE_NAME)
logger.info(f"Using queue: {queue.queue_name}")

Podemos enviar tres trabajos a QS1, nuestra cuota de compartición de trabajos en segundo plano. El primer trabajo enviado tiene la mayor prioridad, seguido del trabajo de prioridad media y luego del trabajo de prioridad más baja. Como actualmente no se está ejecutando nada en QS1 o QS2, deberían iniciarse dos instancias de trabajos, correspondientes al primer y segundo trabajo enviado (ya que tenían la prioridad relativa más alta). Tenga en cuenta que, como este es solo un trabajo de entrenamiento de "hola mundo", no requiere datos, pero normalmente un ModelTrainer requiere datos para realizar el entrenamiento.

from utils.aws_batch_resource_management import JobStatus, await_jobs, list_jobs_by_quota_share

# Submit three jobs to QS1, all with different priorities to guarantee dispatch ordering.
qs1_job_high: TrainingQueuedJob = queue.submit(job_name="qs1_job_high", training_job=model_trainer, quota_share_name=qs1.name, priority=3, inputs=None)
qs1_job_med: TrainingQueuedJob = queue.submit(job_name="qs1_job_med", training_job=model_trainer, quota_share_name=qs1.name, priority=2, inputs=None)
qs1_job_low: TrainingQueuedJob = queue.submit(job_name="qs1_job_low", training_job=model_trainer, quota_share_name=qs1.name, priority=1, inputs=None)

logger.info(f"Submitted jobs: {qs1_job_high.job_name}, {qs1_job_med.job_name}, {qs1_job_low.job_name}")
logger.info(f"Waiting for jobs to transition...")

# qs1_job_high is dispatched first using capacity from QS1, and qs1_job_med is dispatched second borrowing idle capacity from QS2
await_jobs([([qs1_job_high, qs1_job_med], JobStatus.dispatched() | JobStatus.terminal())])
list_jobs_by_quota_share(queue, [qs1.name, qs2.name, qs3.name], JobStatus.active())

Enviar trabajos a QS2 desencadena la preemption entre comparticiones para recuperar capacidad prestada

El equipo que posee QS2 puede descubrir que tiene trabajo por realizar y enviar un trabajo a QS2. Al hacerlo, se desencadena una preemption de capacidad desde QS1, que actualmente está prestando esa capacidad (la capacidad de QS3 no está configurada para ser prestada). Esto es una preemption entre comparticiones.

# When a job is submitted to QS2, it will preempt its capacity back from QS1
qs2_job_1: TrainingQueuedJob = queue.submit(job_name="qs2_job_1", training_job=model_trainer, quota_share_name=qs2.name, priority=1, inputs=None)

logger.info(f"Submitted job {qs2_job_1.job_name}")
logger.info(f"Waiting for jobs to transition...")

await_jobs([
    ([qs2_job_1], JobStatus.dispatched() | JobStatus.terminal()),
    ([qs1_job_med], {JobStatus.RUNNABLE}) # Lowest priority QS1 job will be preempted
])
list_jobs_by_quota_share(queue, [qs1.name, qs2.name, qs3.name], JobStatus.active())

Actualizar prioridades de trabajos en QS1 desencadena la preemption dentro de la compartición para ejecutar trabajos de mayor prioridad

Si un científico se da cuenta de que el trabajo qs1_job_low en QS1, que todavía es EJECUTABLE, necesita comenzar a ejecutarse ahora debido a los cambios de prioridad, puede utilizar la nueva API UpdateServiceJob para aumentar la prioridad de un trabajo ya enviado en QS1. Como no hay capacidad disponible para prestar, esto inicia una preemption dentro de la compartición, para que el trabajo que se está ejecutando actualmente sea interrumpido y qs1_job_low pueda comenzar.

# Updating qs1_job_low to be the highest priority job in QS1
qs1_job_low.update(scheduling_priority=4)

logger.info(f"Updated job {qs1_job_low.job_name} to increase its priority")
logger.info(f"Waiting for jobs to transition...")

await_jobs([
    ([qs1_job_low], JobStatus.dispatched() | JobStatus.terminal()),
    ([qs1_job_high], {JobStatus.RUNNABLE}) # High priority job will preempt lower priority jobs within the same quota share
])
list_jobs_by_quota_share(queue, [qs1.name, qs2.name, qs3.name], JobStatus.active())

(Opcional) Limpiar recursos creados

Para limpiar los recursos creados en esta publicación, ejecute el siguiente código:

from utils.aws_batch_resource_management import delete_resources

delete_resources(resource_manager, resources)

Conclusión

En esta publicación, cubrimos la administración de cuotas, una nueva capacidad de programación en AWS Batch, y cómo utilizarla con trabajos de entrenamiento de SageMaker y comenzar a configurar comparticiones de cuotas para sus propios proyectos. Para obtener más detalles, consulte la guía de usuario de administración de cuotas de AWS Batch.

AWS Batch con gestión de cuotas puede ayudar a su organización a aumentar la utilización de cómputo y a priorizar de manera más efectiva a sus equipos y cargas de trabajo, liberando tiempo para sus administradores de infraestructura y científicos de ML. Al implementar la gestión de cuotas, sus equipos pueden centrarse en sus cargas de trabajo y reducir el tiempo dedicado a la gestión y coordinación de infraestructura. Esta capacidad es especialmente poderosa con los planes de entrenamiento de AI de Amazon SageMaker, para que su organización pueda reservar capacidad en la cantidad que necesite, durante el tiempo que la necesite. Al utilizar AWS Batch con SageMaker AI, puede aprovechar al máximo el plan de entrenamiento. Le animamos a probar esta nueva capacidad para que tenga un impacto significativo en sus operaciones.

Sobre los autores