Article · AI translation
随着 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 和 3 个 Quota Shares。这里配置的 3 个配额份额代表 3 个不同的项目。每个份额都设置了相同的容量限制：1 个 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" # 创建 SchedulingPolicy、ServiceEnvironment、JobQueue 和 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 模块经过扩展，现可通过编程方式向启用了配额管理的 AWS Batch 作业队列提交 SageMaker Training 作业。为演示这些新功能，我们向作业队列提交一个 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 # 使用 SageMaker Python SDK 构造队列对象 queue = TrainingQueue(JOB_QUEUE_NAME) logger.info(f"Using queue: {queue.queue_name}") 接下来，我们可以向后台作业配额份额 QS1 提交 3 个作业。第 1 个提交的作业优先级最高，其次是中等优先级作业，最后是优先级最低的作业。由于 QS1 和 QS2 当前都没有作业运行，应有相当于 2 个实例容量的作业启动，对应第 1 个和第 2 个提交的作业（因为它们的相对优先级最高）。请注意，由于这只是一个 Hello World 训练作业，因此不需要数据，但 ModelTrainer 通常需要数据来执行训练。from utils.aws_batch_resource_management import JobStatus, await_jobs, list_jobs_by_quota_share # 向 QS1 提交 3 个优先级各不相同的作业，以保证分派顺序。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 首先使用 QS1 的容量被分派，qs1_job_med 随后借用 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 收回容量（QS3 的容量未配置为可供借用）。这就是跨份额抢占。# 向 QS2 提交作业时，它会从 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}) # 优先级最低的 QS1 作业将被抢占 ]) 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 得以开始运行。# 将 qs1_job_low 更新为 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}) # 高优先级作业将在同一配额份额内抢占低优先级作业 ]) 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 配额管理用户指南。采用配额管理的 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 之前，他曾担任金融科技领域的产品负责人，并为多个行业的产品组织提供咨询服务。