AWS Batch 推出 quota management 功能，实现 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 的 Scheduling Policy（调度策略）、一个以特定实例为单位指定其容量上限的 Service Environment（服务环境），以及连接两者的 SageMaker Training 作业队列。之后，您就可以将配额份额关联到作业队列。用户指南中提供了有关配置配额管理资源和配额份额的 AWS Command Line Interface (AWS CLI) 和控制台示例。示例笔记本中包含一个将配额管理与 SageMaker Python SDK v3 中的 aws_batch 模块相集成的端到端示例。在本示例中，我们将逐步讲解概括关键组件的代码片段，并使用笔记本中的辅助函数来创建 Scheduling Policy、Service Environment、Job Queue 以及三个 Quota Shares。

此处配置的三个配额份额代表三个不同的项目。每个份额都设置了相同的容量上限：一个 ml.g5.xlarge 实例。QS1 配额份额代表后台作业。对于该项目，您希望最先运行最重要的后台作业，并尽可能借用所有可用的空闲容量。在该配额份额内，没有任何作业是紧急的，因此可以接受抢占，因为所有作业都使用了 checkpointing。该配额份额使用 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、Service Environment 和 Job Queue。

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

向 QS1 提交 SageMaker Training 作业

SageMaker Python SDK 中的 aws_batch 模块已得到扩展，支持以编程方式将 SageMaker Training 任务提交到启用配额管理的 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 提交了一个作业。这样做会触发容量从 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 中仍处于 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 Training 任务中使用它，并开始为您的项目设置配额份额。更多详细信息，请参阅 AWS Batch quota management user guide。

AWS Batch 配额管理可以帮助您的组织提高计算资源利用率，更有效地对团队和工作负载进行优先级排序，从而为基础设施管理员和 ML 科学家腾出时间。通过实施配额管理，您的团队可以专注于自己的工作负载，减少在管理和协调基础设施上花费的时间。该功能与 Amazon SageMaker AI training plans 结合使用尤为强大，使您的组织能够在需要的时间预留所需数量的容量。通过将 AWS Batch 与 SageMaker AI 结合使用，您可以充分利用训练计划。我们鼓励您试用这项新功能，让它为您的运营带来切实的影响。

关于作者