AWS Batch 配额管理如何让 SageMaker Training 任务跨团队借还 GPU 资源

With the launch of AWS Batch support for Amazon SageMaker Training jobs, customers gained a new fully managed way to queue, schedule, and retry machine learning (ML) training workloads without building their own infrastructure. Instead, they use the queuing and scheduling features of AWS Batch. As organizations scale their artificial intelligence and machine learning (AI/ML) operations and onboard more teams, coordinating access to shared graphics processing unit (GPU) capacity can introduce additional communication, coordination, and configuration overhead for administrators. This overhead includes setting up separate queues, adjusting capacity limits, and carefully orchestrating the right workloads to get the right resources at the right time. A core value proposition of AWS Batch is helping customers increase utilization of their compute, and as training environments grow, so does the opportunity to do more with existing capacity.

Working backward from customers managing multi-team training environments, we heard a consistent request: give us a way to define allocations per team or project, share idle capacity automatically, and run our highest-priority workloads first. This request led us to build quota management, a new scheduling capability for AWS Batch. Administrators can use it to divide compute capacity into quota shares, organized by team, project, or workload. Each quota share uses configurable resource-sharing strategies that maximize GPU utilization while honoring each team’s allocation when it matters most.

This post shows how quota management works, including how to configure quota shares with lending and borrowing strategies, set up in-share preemption for priority workloads, and get started using the SageMaker Python SDK.

Quota shares and preemption

Administrators can allocate and share constrained GPU resources by defining a quota share for separate entities, such as individual teams or projects. Each quota share includes compute quotas (capacity limits) and resource-sharing strategies for idle compute. These strategies determine whether idle compute is reserved until jobs arrive (RESERVE), lent to other quota shares that have work available (LEND), or both lent when idle and borrowed when the share has work that might benefit (LEND_AND_BORROW).

A quota share operates as a virtual queue nested hierarchically within an associated job queue. When scheduling jobs for a job queue, AWS Batch prioritizes quota shares that are under their capacity limits as it iterates through the associated quota shares. If jobs arrive in a quota share that has lent out its capacity, AWS Batch performs cross-share preemption. This reclaims resources for the lending quota share, so administrators don’t need to track down the jobs borrowing capacity and terminate them on their own. If jobs arrive in a quota share that is at or above its capacity limit, AWS Batch checks whether there’s available capacity to borrow from another quota share. If so, workloads can run using otherwise-idle capacity.

You can turn on priority-based in-share preemption when you configure a quota share. Where cross-share preemption reclaims lent capacity as jobs arrive in a lending quota share, in-share preemption preempts lower priority jobs that are in the SCHEDULED, STARTING, or RUNNING state. This allows higher priority jobs at the front of the queue to start running sooner.

Prerequisites

Before you proceed with this example, you must have an AWS account with permissions to create and manage AWS Batch resources. For this example, see the sample IAM permissions needed to invoke AWS Batch and SageMaker APIs.

The example uses the SageMaker Python SDK v3. Complete the setup required to install a compatible version of the SageMaker Python SDK:

!pip install sagemaker

Configuring quota management resources

Using the new quota management capabilities requires creating a Scheduling Policy that contains a quotaSharePolicy, a Service Environment that specifies its capacity limits in terms of specific instances, and a SageMaker Training job queue that is connected to both. After that, you can associate quota shares to a job queue. The user guide has AWS Command Line Interface (AWS CLI) and console examples for configuring quota management resources and quota shares. An end-to-end example of quota management integrated with the aws_batch module in the SageMaker Python SDK v3 is contained in a sample notebook. In this example, we step through code snippets that summarize the key components, and use the helper functions from the notebook to create a Scheduling Policy, Service Environment, Job Queue, and three Quota Shares.

The three quota shares configured here represent three different projects. Each has the same capacity limit set: a single ml.g5.xlarge instance. The QS1 quota share represents background jobs. For this project, you want to run the most important background job first and borrow as much spare capacity as is available. Within this quota share, no jobs are urgent, and preemptions are acceptable because all jobs use checkpointing. This quota share uses the LEND_AND_BORROW resource-sharing strategy and has in-share preemption turned on. Higher priority jobs within the same quota share can then preempt running jobs that are lower priority.

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 represents medium priority workloads. Teams choosing this configuration would prefer to wait for capacity to become available rather than endure additional preemptions. Those preemptions come from borrowed capacity being reclaimed or from in-share preemptions because of higher priority job submission. This is appropriate for shorter jobs, or when most submitted jobs have similar priorities. The quota share still benefits from job priority being taken into consideration at job dispatch. Submitted jobs trigger AWS Batch to reclaim any lent capacity needed to run the jobs.

qs2 = QuotaShareConfig(
    name="QS2",
    capacity_unit=INSTANCE_TYPE,
    max_capacity=1,
    in_share_preemption=False,
    sharing_strategy="LEND"
)

QS3 represents workloads that need to start as quickly as possible with minimal interruptions. It’s configured to reserve capacity and neither lends nor borrows to other quota shares. It doesn’t experience cross-share preemptions, because its workloads don’t exceed their capacity limits. By also turning off in-share preemption, workloads submitted to this quota share are not preempted by AWS Batch as long as the capacity limit is unchanged.

qs3 = QuotaShareConfig(
    name="QS3",
    capacity_unit=INSTANCE_TYPE,
    max_capacity=1,
    in_share_preemption=False,
    sharing_strategy="RESERVE"
)

These quota share configurations are inputs to the create_quota_management_resources function in the sample notebook, which also configures the corresponding Scheduling Policy, Service Environment, and Job Queue with the constant names chosen.

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]
)

Submit SageMaker Training jobs to QS1

The aws_batch module within the SageMaker Python SDK has been extended to allow programmatic submission of SageMaker Training jobs to a quota-management-enabled AWS Batch job queue. To demonstrate these new capabilities, we use a Hello World ModelTrainer that’s submitted to the job queue. The following code snippet identifies an image URI to use based on the Region and instance type selected.

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",
)

We use this image URI within a basic hello world model trainer example:

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),
)

To use the job queue within the SageMaker Python SDK, provide the queue name to the TrainingQueue constructor within the aws_batch module.

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}")

We can then submit three jobs to QS1, our background jobs quota share. The first job submitted has the highest priority, followed by the medium priority job, and then the lowest priority job. As nothing is currently running in QS1 or QS2, two instances’ worth of jobs should start, corresponding to the first and second jobs submitted (as they had the highest relative priority). Note that because this is only a hello world training job, it doesn’t require data, but typically a ModelTrainer requires data to perform training.

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())

Submitting jobs to QS2 triggers cross-share preemption to reclaim lent capacity

The team that owns QS2 might discover they have work to perform and submit a job to QS2. Doing so triggers a preemption of capacity back from QS1, which is currently borrowing that capacity (QS3’s capacity isn’t configured to be borrowed). This is a cross-share preemption.

# 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())

Updating job priorities in QS1 triggers in-share preemption to run highest priority jobs

If a scientist realizes that the qs1_job_low job in QS1, which is still RUNNABLE, needs to begin running now because of changing priorities, they can use the new UpdateServiceJob API to increase the priority of an already submitted job in QS1. Because there’s no available capacity to borrow, this kicks off an in-share preemption, so that the job currently running is preempted and qs1_job_low can begin.

# 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())

(Optional) Clean up created resources

To clean up the resources created in this post, run the following code:

from utils.aws_batch_resource_management import delete_resources

delete_resources(resource_manager, resources)

Conclusion

In this post, we covered quota management, a new scheduling capability in AWS Batch, and how to use it with SageMaker Training jobs and get started setting up quota shares for your own projects. For more details, see the AWS Batch quota management user guide.

AWS Batch with quota management can help your organization increase compute utilization and more effectively prioritize your teams and workloads, freeing up time for your infrastructure admins and ML scientists. By implementing quota management, your teams can focus on their workloads and reduce time spent managing and coordinating infrastructure. This capability is especially powerful with Amazon SageMaker AI training plans, so that your organization can reserve capacity in the quantity you need, during the time you need it. By using AWS Batch with SageMaker AI, you can make full use of the training plan. We encourage you to try out this new capability so it can make a meaningful impact in your operations.

About the authors

Michelle Goodstein

Michelle is a Principal Engineer on Amazon Athena. While working on AWS Batch, she led scheduling improvements for AI/ML, including quota management for SageMaker Training jobs, to drive utilization, efficiency, and cost optimization, as well as improved observability into the job execution lifecycle. She enjoys building innovative solutions to distributed systems problems spanning data, compute, and AI/ML.

David Lindskog

David is a Senior Software Engineer at AWS Batch. He was project lead for AWS Batch support for SageMaker Training jobs and most recently led the launch of AWS Batch’s first free CloudWatch metrics.  Across his time at Amazon, David has enjoyed designing and building highly scalable and operationally robust backend systems for compute-heavy workloads, robotics telemetry, and AI/ML training.

Tom Burggraf

Tom is the Head of Software Strategy for US Federal Supercomputing at AWS, where he shapes strategy across AWS services for federal research and scientific computing customers. While leading product for AWS Batch, he drove features that help research platform builders scale efficiently, including AWS Batch support for SageMaker Training jobs and quota management. Prior to AWS, he was a product leader in FinTech and a consultant for product organizations across multiple industries.