O problema de coordenar GPUs entre múltiplas equipes
Quem já gerenciou ambientes de treinamento de modelos de aprendizado de máquina (Machine Learning — ML) em escala sabe o quanto é trabalhoso dividir recursos de unidades de processamento gráfico (Graphics Processing Unit — GPU) entre times diferentes. À medida que as organizações crescem e mais equipes passam a usar a mesma infraestrutura de inteligência artificial e aprendizado de máquina (Artificial Intelligence/Machine Learning — AI/ML), surgem problemas clássicos: filas separadas para cada time, ajustes manuais de limites de capacidade, e uma quantidade considerável de comunicação e coordenação para garantir que o workload certo receba o recurso certo na hora certa.
Foi exatamente esse cenário que motivou a AWS a desenvolver o gerenciamento de cotas (quota management), uma nova capacidade de agendamento dentro do AWS Batch. A proposta é clara: permitir que administradores definam alocações por equipe ou projeto, compartilhem capacidade ociosa de forma automática e executem os workloads de maior prioridade primeiro — sem precisar gerenciar isso manualmente.
Como funciona o gerenciamento de cotas
O gerenciamento de cotas introduz o conceito de quota shares (cotas compartilhadas), que funcionam como filas virtuais aninhadas dentro de uma job queue do AWS Batch. Cada quota share representa uma entidade — uma equipe, um projeto ou um tipo de workload — e possui dois elementos principais:
- Cotas de computação: os limites de capacidade que aquela cota pode usar.
- Estratégias de compartilhamento de recursos: regras que determinam o que acontece com a capacidade ociosa.
Existem três estratégias de compartilhamento disponíveis:
- RESERVE: a capacidade fica reservada para a própria cota, mesmo quando ociosa. Não empresta nem toma emprestado de outras cotas.
- LEND: a capacidade ociosa é emprestada para outras cotas que tenham trabalho disponível. Quando um job chega na cota que emprestou, o AWS Batch recupera automaticamente os recursos.
- LEND_AND_BORROW: além de emprestar quando ociosa, a cota também pode tomar capacidade emprestada de outras cotas quando tiver jobs aguardando.
Preempção entre cotas e dentro de cotas
Dois mecanismos de preempção tornam o sistema ainda mais inteligente:
Preempção entre cotas (cross-share preemption): quando uma cota que emprestou sua capacidade recebe novos jobs, o AWS Batch automaticamente recupera esses recursos — sem que o administrador precise identificar e encerrar manualmente os jobs que estavam usando a capacidade emprestada.
Preempção dentro da cota (in-share preemption): quando ativada em uma quota share, permite que jobs de maior prioridade submetidos àquela cota interrompam jobs de menor prioridade que já estejam nos estados SCHEDULED, STARTING ou RUNNING. Isso garante que trabalhos urgentes comecem a executar mais rapidamente.
Pré-requisitos
Para seguir o exemplo descrito pela AWS, é necessário ter uma conta AWS com permissões para criar e gerenciar recursos do AWS Batch. As permissões de IAM de exemplo necessárias para invocar as APIs do AWS Batch e do SageMaker estão disponíveis no repositório oficial. O exemplo utiliza o SageMaker Python SDK v3, que pode ser instalado com:
!pip install sagemakerConfigurando os recursos de gerenciamento de cotas
Para usar o gerenciamento de cotas, a AWS requer a criação de três componentes principais:
- Uma Política de Agendamento (Scheduling Policy) contendo a
quotaSharePolicy. - Um Ambiente de Serviço (Service Environment) que define os limites de capacidade em termos de instâncias específicas.
- Uma fila de jobs de treinamento do SageMaker conectada aos dois anteriores.
O guia de usuário do AWS Batch traz exemplos via Interface de Linha de Comando (AWS Command Line Interface — AWS CLI) e console para configurar os recursos de gerenciamento de cotas e quota shares. Um exemplo completo integrado ao módulo aws_batch do SageMaker Python SDK v3 está disponível em um notebook de exemplo.
Exemplo com três projetos distintos
O exemplo oficial da AWS demonstra três quota shares representando três projetos diferentes, cada um com o mesmo limite de capacidade: uma instância ml.g5.xlarge.
QS1 — Jobs em segundo plano (background jobs): configurada com a estratégia LEND_AND_BORROW e preempção interna ativada. Isso permite que ela tome capacidade emprestada de outras cotas e que jobs de maior prioridade dentro dela interrompam os de menor prioridade. O uso de checkpointing é assumido, tornando preempções aceitáveis.
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 — Workloads de prioridade média: configurada com a estratégia LEND e preempção interna desativada. Ideal para jobs mais curtos ou quando a maioria dos jobs submetidos tem prioridades similares. A cota ainda empresta sua capacidade ociosa, e quando um job chega, o AWS Batch recupera automaticamente o que foi emprestado.
qs2 = QuotaShareConfig(
name="QS2",
capacity_unit=INSTANCE_TYPE,
max_capacity=1,
in_share_preemption=False,
sharing_strategy="LEND"
)QS3 — Workloads críticos com mínima interrupção: configurada com a estratégia RESERVE e preempção interna desativada. Não empresta nem toma emprestado, garantindo que os workloads submetidos a essa cota não sofram preempções pelo AWS Batch enquanto o limite de capacidade não for alterado.
qs3 = QuotaShareConfig(
name="QS3",
capacity_unit=INSTANCE_TYPE,
max_capacity=1,
in_share_preemption=False,
sharing_strategy="RESERVE"
)Essas configurações são então passadas para a função create_quota_management_resources, que cria a Scheduling Policy, o Service Environment e a Job Queue de forma integrada:
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]
)Submetendo jobs de treinamento do SageMaker
O módulo aws_batch do SageMaker Python SDK foi estendido para permitir a submissão programática de jobs de treinamento do SageMaker para filas habilitadas com gerenciamento de cotas. O exemplo usa um ModelTrainer simples do tipo “Hello World”. Primeiro, identifica-se a URI da imagem com base na região e no tipo de instância:
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",
)Em seguida, configura-se o ModelTrainer:
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 a fila dentro do SDK, basta fornecer o nome da fila ao construtor TrainingQueue:
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}")Submetendo jobs à QS1 com diferentes prioridades
O exemplo submete três jobs à QS1 com prioridades distintas. Como nenhum job está rodando em QS1 ou QS2 no momento, dois jobs começam imediatamente — o de maior prioridade usa a capacidade da própria QS1, e o segundo toma emprestado a capacidade ociosa da QS2:
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())Submissão à QS2 dispara preempção entre cotas
Quando a equipe dona da QS2 submete um novo job, o AWS Batch automaticamente recupera a capacidade que estava emprestada para a QS1 — sem nenhuma intervenção manual do administrador:
# 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())Atualizando prioridades dispara preempção dentro da cota
Se um cientista percebe que o qs1_job_low precisa executar imediatamente por mudança de prioridades, ele pode usar a nova API UpdateServiceJob para aumentar a prioridade do job já submetido. Como não há capacidade disponível para tomar emprestado, o AWS Batch inicia uma preempção dentro da própria cota, interrompendo o job em execução para que o de maior prioridade comece:
# 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())Limpeza dos recursos criados
Para remover todos os recursos criados durante o exemplo, basta executar:
from utils.aws_batch_resource_management import delete_resources
delete_resources(resource_manager, resources)Conclusão
O gerenciamento de cotas do AWS Batch representa um avanço importante para organizações que operam ambientes de treinamento de ML com múltiplas equipes. Em vez de depender de comunicação manual e ajustes constantes de configuração, administradores podem definir regras claras de alocação e deixar o próprio serviço gerenciar o compartilhamento e a recuperação de capacidade de forma automática.
A funcionalidade é especialmente poderosa quando combinada com os planos de treinamento do Amazon SageMaker AI, que permitem reservar capacidade na quantidade e no período necessários. Segundo a AWS, isso garante que toda a capacidade reservada seja efetivamente utilizada. Para mais detalhes técnicos, o guia de usuário do gerenciamento de cotas do AWS Batch traz a documentação completa sobre como configurar e operar essa funcionalidade.
Fonte
Lend and borrow resources across teams with AWS Batch’s quota management for SageMaker Training jobs (https://aws.amazon.com/blogs/machine-learning/lend-and-borrow-resources-across-teams-with-aws-batchs-quota-management-for-sagemaker-training-jobs/)
Leave a Reply