AWS Batch adds quota management so teams can share GPU capacity for SageMaker Training

Why AWS built quota management

AWS Batch support for Amazon SageMaker Training jobs already gave customers a fully managed way to queue, schedule and retry machine learning training workloads without building their own infrastructure. As more teams are onboarded, AWS says, coordinating access to shared graphics processing unit (GPU) capacity adds communication, coordination and configuration overhead for administrators. That includes setting up separate queues, adjusting capacity limits and orchestrating which workloads get which resources.

AWS says customers running multi-team training environments repeatedly asked for a way to define allocations per team or project, share idle capacity automatically and run their highest-priority workloads first. Quota management is the company's answer to that request, and AWS ties it to a core goal of AWS Batch: helping customers raise utilization of their compute.

How quota shares, lending and preemption work

Administrators define a quota share for each entity, such as a team or project. Each share holds compute quotas, which are capacity limits, and a resource-sharing strategy for idle compute. There are three strategies:

• RESERVE: idle compute is held until the share's own jobs arrive.
• LEND: idle compute is lent to other quota shares that have work available.
• LEND_AND_BORROW: the share lends when idle and borrows when it has work that might benefit.

A quota share works as a virtual queue nested inside a job queue. When scheduling, AWS Batch prioritizes shares that are under their capacity limits. If jobs arrive in a share that has lent out its capacity, AWS Batch performs cross-share preemption, reclaiming the resources so administrators do not have to find and terminate the borrowing jobs themselves. If jobs arrive in a share that is at or above its limit, AWS Batch checks whether another share has capacity to borrow.

Priority-based in-share preemption can be turned on when a share is configured. It preempts lower-priority jobs in the SCHEDULED, STARTING or RUNNING state so that higher-priority jobs at the front of the queue can start sooner.

Setting it up with the SageMaker Python SDK

Using the feature requires a Scheduling Policy containing a quotaSharePolicy, a Service Environment that specifies capacity limits in terms of specific instances, and a SageMaker Training job queue connected to both. Quota shares are then associated with the job queue. The AWS user guide covers AWS CLI and console configuration, and a sample notebook shows an end-to-end example with the aws_batch module in the SageMaker Python SDK v3.

The example in the post creates three quota shares for three projects, each limited to a single ml.g5.xlarge instance:

• QS1 covers background jobs. It uses LEND_AND_BORROW with in-share preemption on, which suits work that is not urgent and uses checkpointing.
• QS2 covers medium-priority workloads. It uses LEND with in-share preemption off, for teams that would rather wait for capacity than face extra preemptions. AWS Batch still reclaims any lent capacity needed to run its jobs.
• QS3 covers workloads that must start fast with minimal interruption. It uses RESERVE with in-share preemption off, so its jobs are not preempted by AWS Batch as long as the capacity limit is unchanged.

These configurations are passed to a create_quota_management_resources function from the sample notebook, which also sets up the Scheduling Policy, Service Environment and Job Queue. The post then submits a simple Hello World ModelTrainer to the queue to demonstrate programmatic submission, which the aws_batch module has been extended to support.

What it means for platform teams

For organizations that share limited GPU capacity across several teams, the feature moves allocation rules into the scheduler instead of manual coordination. Teams can pick the trade-off that fits their work: guaranteed capacity with no preemption, lending idle capacity while being able to reclaim it, or borrowing spare capacity for interruptible jobs. The prerequisites are an AWS account with permissions to manage AWS Batch resources and the SageMaker Python SDK v3.