AWS has described a new scheduling capability for AWS Batch, called quota management, in a post on the AWS Machine Learning Blog. It builds on AWS Batch support for SageMaker Training jobs, which already handles queuing, scheduling and retries for ML training.
What it does
An administrator divides GPU capacity into quota shares, one per team, project or workload. Each share has a capacity limit and a strategy for what happens when it sits idle:
- RESERVE holds the capacity until the share's own jobs arrive.
- LEND offers idle capacity to other shares that have work waiting.
- LEND_AND_BORROW lends when idle and borrows when the share has work that can use extra capacity.
Each share works as a virtual queue inside a job queue. The scheduler favors shares that are under their limits. When a share has lent out capacity and its own jobs arrive, AWS Batch performs cross-share preemption, taking the capacity back without an admin hunting down the borrowing jobs.
A second option is in-share preemption. With it on, lower-priority jobs that are scheduled, starting or running can be pre-empted so higher-priority jobs at the front of the queue start sooner. A new UpdateServiceJob API lets a scientist raise the priority of a job already submitted.
Setup
The setup calls for a scheduling policy that includes a quota share policy, plus a service environment that defines capacity constraints by particular instances, and a SageMaker Training job queue connected to both of these. Quota shares are then attached to the queue. Jobs are submitted through the aws_batch module of the SageMaker Python SDK v3, and the post points to CLI and console examples in the user guide.
The worked example
AWS's example creates three shares, each limited to one ml.g5.xlarge instance:
- A background-work share that lends and borrows, with in-share preemption on, for jobs that checkpoint and tolerate interruption.
- A medium-priority share that only lends and has in-share preemption off, for teams that prefer waiting to being preempted repeatedly.
- A share that only reserves, so its jobs are not preempted while its limit stays unchanged.
In the walkthrough, background jobs borrow the medium share's idle instance. When a job lands in the medium share, capacity is taken back from the background share. Raising a waiting job's priority then preempts a running job inside the same share.
The post also says the capability works with SageMaker AI training plans, so reserved capacity can be shared across teams.