AWS ha presentado una nueva capacidad de planificación para AWS Batch, la gestión de cuotas, en una entrada del blog de Machine Learning de AWS. Se apoya en la compatibilidad de AWS Batch con los trabajos de SageMaker Training, que ya se encarga de las colas, la planificación y los reintentos del entrenamiento de ML.
Qué hace
Un administrador reparte la capacidad de GPU en cuotas, una por equipo, proyecto o carga de trabajo. Cada cuota tiene un límite de capacidad y una estrategia para cuando queda sin usar:
- RESERVE guarda la capacidad hasta que lleguen los trabajos de la propia cuota.
- LEND ofrece la capacidad ociosa a otras cuotas que tengan trabajo en espera.
- LEND_AND_BORROW presta cuando está ociosa y toma prestado cuando tiene trabajo que puede aprovechar capacidad extra.
Cada cuota funciona como una cola virtual dentro de una cola de trabajos. El planificador da preferencia a las cuotas que están por debajo de su límite. Cuando una cuota ha prestado capacidad y llegan sus propios trabajos, AWS Batch aplica preemption entre cuotas y recupera esa capacidad sin que ningún administrador tenga que localizar los trabajos que la tienen prestada.
Hay una segunda opción, la preemption dentro de la cuota. Con ella activada, los trabajos de menor prioridad que están planificados, arrancando o en ejecución pueden interrumpirse para que los de mayor prioridad, al frente de la cola, empiecen antes. Una nueva API, UpdateServiceJob, permite a un científico de datos subir la prioridad de un trabajo ya enviado.
Configuración
Hace falta una política de planificación que incluya una política de cuotas, un entorno de servicio que fije los límites de capacidad por tipo de instancia y una cola de trabajos de SageMaker Training conectada a ambos. Después se asocian las cuotas a la cola. Los trabajos se envían con el módulo aws_batch del SDK de Python v3 de SageMaker, y la entrada remite a ejemplos de CLI y de consola en la guía de usuario.
El ejemplo práctico
El ejemplo de AWS crea tres cuotas, cada una limitada a una instancia ml.g5.xlarge:
- Una cuota para tareas en segundo plano que presta y toma prestado, con preemption interna activada, pensada para trabajos que guardan puntos de control y toleran interrupciones.
- Una cuota de prioridad media que solo presta y tiene la preemption interna desactivada, para equipos que prefieren esperar a que los interrumpan una y otra vez.
- Una cuota que solo reserva, de modo que sus trabajos no se interrumpen mientras su límite no cambie.
En el recorrido, los trabajos en segundo plano usan la instancia ociosa de la cuota media. Cuando entra un trabajo en la cuota media, se recupera la capacidad de la cuota de segundo plano. Al subir la prioridad de un trabajo en espera, este interrumpe a uno que se está ejecutando dentro de la misma cuota.
La entrada añade que la capacidad funciona con los planes de entrenamiento de SageMaker AI, de modo que la capacidad reservada puede compartirse entre equipos.