AWS Batch 为 SageMaker Training 作业推出配额管理，支持跨团队资源借调

随着 AWS Batch 对 Amazon SageMaker Training 作业的支持正式推出，客户获得了一种全新的全托管方式，可以对机器学习（ML）训练工作负载进行排队、调度和重试，而无需自行构建基础设施。他们转而使用 AWS Batch 的排队和调度功能。随着组织扩展其人工智能与机器学习（AI/ML）运营并引入更多团队，协调对共享图形处理器（GPU）容量的访问可能会给管理员带来额外的沟通、协调和配置开销。这些开销包括设置单独的队列、调整容量限制，以及精心编排合适的工作负载，以便在合适的时间获得合适的资源。AWS Batch 的核心价值主张之一是帮助客户提高计算利用率，而随着训练环境的发展，利用现有容量做更多事情的机会也随之增加。

从管理多团队训练环境的客户出发反向思考，我们听到了一致的诉求：给我们一种方式，按团队或项目定义配额分配，自动共享空闲容量，并优先运行我们最高优先级的工作负载。这一诉求促使我们构建了 配额管理，这是 AWS Batch 的一项新调度能力。管理员可以用它将计算容量划分为配额份额，按团队、项目或工作负载进行组织。每个配额份额都使用可配置的资源共享策略，在最大化 GPU 利用率的同时，在关键时刻尊重每个团队的分配额度。

本文介绍配额管理的工作原理，包括如何配置带有借出和借用策略的配额份额、为优先工作负载设置份额内抢占，以及如何开始使用 SageMaker Python SDK。

配额份额与抢占

管理员可以通过为不同实体（例如各个团队或项目）定义配额份额，来分配和共享受限的 GPU 资源。每个配额份额都包含计算配额（容量限制）以及针对空闲计算的资源共享策略。这些策略决定空闲计算是保留到作业到达（RESERVE）、借给其他有工作可做的配额份额（LEND），还是既在空闲时借出、又在份额有可受益的工作时借用（LEND_AND_BORROW）。

配额份额作为虚拟队列运行，以层级结构嵌套在关联的作业队列中。在为作业队列调度作业时，AWS Batch 会遍历关联的配额份额，并优先考虑未超出容量限制的配额份额。如果作业到达一个已借出其容量的配额份额，AWS Batch 会执行跨份额抢占。这会为借出方配额份额回收资源，因此管理员无需自行查找借用容量的作业并终止它们。如果作业到达一个已达到或超出容量限制的配额份额，AWS Batch 会检查是否有可用容量可从其他配额份额借用。如果有，工作负载就可以使用原本空闲的容量运行。

你可以在配置配额份额时开启基于优先级的份额内抢占。跨份额抢占是在作业到达借出方配额份额时回收借出的容量，而份额内抢占则会抢占处于 SCHEDULED、STARTING 或 RUNNING 状态的较低优先级作业。这使队列前部的高优先级作业能够更早开始运行。

先决条件

在继续此示例之前，您必须拥有一个具有创建和管理 AWS Batch 资源权限的 AWS 账户。对于此示例，请参阅调用 AWS Batch 和 SageMaker API 所需的示例 IAM 权限。

该示例使用 SageMaker Python SDK v3。完成安装兼容版本 SageMaker Python SDK 所需的设置：

!pip install sagemaker

配置配额管理资源

使用新的配额管理功能需要创建一个调度策略，其中包含一个quotaSharePolicy；创建一个服务环境，以特定实例为单位指定其容量限制；以及创建一个 SageMaker 训练作业队列，将其与两者连接。之后，您可以将配额份额关联到作业队列。用户指南提供了用于配置配额管理资源和配额份额的 AWS Command Line Interface (AWS CLI) 和控制台示例。与 SageMaker Python SDK v3 中 aws_batch 模块集成的配额管理端到端示例包含在一个示例笔记本中。在此示例中，我们将逐步介绍概括关键组件的代码片段，并使用笔记本中的辅助函数创建调度策略、服务环境、作业队列和三个配额份额。

此处配置的三个配额份额代表三个不同的项目。每个都具有相同的容量限制设置：单个ml.g5.xlarge实例。QS1 配额份额代表后台作业。对于此项目，您希望首先运行最重要的后台作业，并尽可能借用可用的空闲容量。在此配额份额内，没有作业是紧急的，并且可以接受抢占，因为所有作业都使用检查点。此配额份额使用LEND_AND_BORROW资源共享策略，并开启了份额内抢占。然后，同一配额份额内优先级较高的作业可以抢占正在运行的优先级较低的作业。

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 代表中等优先级工作负载。选择此配置的团队宁愿等待容量可用，也不愿承受额外的抢占。这些抢占来自借用的容量被回收，或由于提交了更高优先级的作业而发生的份额内抢占。这适用于较短的作业，或大多数提交的作业具有相似优先级的情况。该配额份额仍然受益于在作业分派时考虑作业优先级。提交的作业会触发 AWS Batch 回收运行作业所需的任何已借出容量。

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

QS3 代表需要尽可能快速启动且中断最少的工作负载。它配置为预留容量，既不向其他配额份额借出也不借用。它不会经历跨份额抢占，因为其工作负载不会超过其容量限制。通过同时关闭份额内抢占，只要容量限制不变，提交到此配额份额的工作负载就不会被 AWS Batch 抢占。

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

这些配额份额配置是示例笔记本中create_quota_management_resources函数的输入，该函数还使用所选的常量名称配置相应的调度策略、服务环境和作业队列。

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

将 SageMaker 训练作业提交到 QS1

SageMaker Python SDK 中的 aws_batch 模块已扩展，允许以编程方式将 SageMaker 训练作业提交到启用了配额管理的 AWS Batch 作业队列。为了演示这些新功能，我们使用一个 Hello World ModelTrainer 提交到作业队列。以下代码片段根据所选区域和实例类型确定要使用的镜像 URI。

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

我们在一个基础的 hello world 模型训练器示例中使用此镜像 URI：

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

要在 SageMaker Python SDK 中使用作业队列，请在 aws_batch 模块中向 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}")

然后我们可以向 QS1（我们的后台作业配额份额）提交三个作业。第一个提交的作业优先级最高，其次是中等优先级作业，最后是最低优先级作业。由于 QS1 或 QS2 中当前没有运行任何内容，应启动相当于两个实例的作业，对应第一个和第二个提交的作业（因为它们的相对优先级最高）。请注意，由于这只是一个 hello world 训练作业，它不需要数据，但通常 ModelTrainer 需要数据来执行训练。

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

向 QS2 提交作业会触发跨份额抢占以回收借出的容量

拥有 QS2 的团队可能会发现他们有工作要执行并向 QS2 提交作业。这样做会触发从 QS1 抢占回容量，因为 QS1 当前正在借用该容量（QS3 的容量未配置为可借用）。这就是跨份额抢占。

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

更新 QS1 中的作业优先级会触发份额内抢占以运行最高优先级的作业

如果科学家意识到 QS1 中仍处于 RUNNABLE 状态的 qs1_job_low 作业由于优先级变化需要立即开始运行，他们可以使用新的 UpdateServiceJob API 来提高 QS1 中已提交作业的优先级。由于没有可借用的可用容量，这会触发份额内抢占，使当前正在运行的作业被抢占，qs1_job_low 可以开始运行。

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

（可选）清理创建的资源

要清理本文中创建的资源，请运行以下代码：

from utils.aws_batch_resource_management import delete_resources

delete_resources(resource_manager, resources)

结论

在本文中，我们介绍了配额管理——AWS Batch 中的一项新调度功能，以及如何将其与 SageMaker 训练作业配合使用，并开始为您自己的项目设置配额份额。有关更多详细信息，请参阅 AWS Batch 配额管理用户指南。

具有配额管理的 AWS Batch 可以帮助您的组织提高计算利用率，并更有效地确定团队和工作负载的优先级，从而为基础设施管理员和 ML 科学家腾出时间。通过实施配额管理，您的团队可以专注于其工作负载，减少管理和协调基础设施所花费的时间。此功能与 Amazon SageMaker AI 训练计划结合使用时尤为强大，使您的组织能够在需要的时间段内按所需数量预留容量。通过将 AWS Batch 与 SageMaker AI 结合使用，您可以充分利用训练计划。我们鼓励您试用这项新功能，使其在您的运营中发挥有意义的影响。

关于作者

Michelle Goodstein

Michelle 是 Amazon Athena 的首席工程师。在参与 AWS Batch 工作期间，她主导了面向 AI/ML 的调度改进，包括 SageMaker Training 作业的配额管理，以提升利用率、效率和成本优化，并改善作业执行生命周期的可观测性。她乐于为横跨数据、计算和 AI/ML 的分布式系统问题构建创新解决方案。

David Lindskog

David 是 AWS Batch 的高级软件工程师。他曾担任 AWS Batch 支持 SageMaker Training 作业的项目负责人，最近还主导了 AWS Batch 首个免费 CloudWatch 指标的发布。在 Amazon 任职期间，David 一直乐于为计算密集型工作负载、机器人遥测和 AI/ML 训练设计并构建高度可扩展且运行稳健的后端系统。

Tom Burggraf

Tom 是 AWS 美国联邦超级计算软件战略负责人，负责为联邦研究和科学计算客户制定跨 AWS 服务的战略。在领导 AWS Batch 产品期间，他推动了帮助研究平台构建者高效扩展的功能，包括 AWS Batch 对 SageMaker Training 作业的支持和配额管理。在加入 AWS 之前，他曾是金融科技领域的产品负责人，并为多个行业的产品组织担任顾问。