A AWS descreveu um novo recurso de agendamento do AWS Batch, o quota management, em um post no AWS Machine Learning Blog. Ele se apoia no suporte do AWS Batch a jobs do SageMaker Training, que já cuida de fila, agendamento e novas tentativas no treinamento de ML.
O que ele faz
Um administrador divide a capacidade de GPU em quota shares, uma para cada time, projeto ou carga de trabalho. Cada share tem um limite de capacidade e uma estratégia para quando fica ociosa:
- RESERVE guarda a capacidade até chegarem os jobs do próprio share.
- LEND oferece a capacidade ociosa a outros shares que têm trabalho esperando.
- LEND_AND_BORROW empresta quando está ocioso e toma emprestado quando tem trabalho que aproveita capacidade extra.
Cada share funciona como uma fila virtual dentro de uma fila de jobs. O agendador dá preferência aos shares que estão abaixo do limite. Quando um share emprestou capacidade e seus próprios jobs chegam, o AWS Batch faz a preempção entre shares e retoma a capacidade sem que um administrador precise caçar os jobs que a tomaram emprestada.
Há ainda a preempção dentro do share. Com ela ativada, jobs de menor prioridade que estão agendados, iniciando ou em execução podem ser interrompidos para que os de maior prioridade, no início da fila, comecem mais cedo. Uma nova API, a UpdateServiceJob, permite que um cientista aumente a prioridade de um job já enviado.
Configuração
É preciso criar uma política de agendamento que inclua uma política de quota share, um ambiente de serviço que defina as restrições de capacidade por tipo de instância e uma fila de jobs do SageMaker Training ligada aos dois. Depois, as quota shares são associadas à fila. Os jobs são enviados pelo módulo aws_batch do SageMaker Python SDK v3, e o post indica exemplos de CLI e de console no guia do usuário.
O exemplo da AWS
O exemplo cria três shares, cada um limitado a uma instância ml.g5.xlarge:
- Um share de trabalho em segundo plano que empresta e toma emprestado, com preempção interna ativada, voltado a jobs que gravam checkpoints e toleram interrupção.
- Um share de prioridade média que só empresta e tem a preempção interna desativada, pensado para times que preferem esperar a sofrer interrupções repetidas.
- Um share que apenas reserva, de modo que seus jobs não são interrompidos enquanto o limite permanece o mesmo.
No passo a passo, os jobs em segundo plano usam a instância ociosa do share de prioridade média. Quando um job chega a esse share, a capacidade é retomada do share de segundo plano. Depois, ao aumentar a prioridade de um job que aguardava, um job em execução dentro do mesmo share é interrompido.
O post também diz que o recurso funciona com os training plans do SageMaker AI, de modo que a capacidade reservada pode ser compartilhada entre times.