Lend and borrow resources across teams with AWS Batch’s quota management for SageMaker Training jobs

!pip install sagemaker

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"

# 创建调度策略、服务环境、作业队列和配额共享
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 = QuotaShareConfig(
    name="QS2",
    capacity_unit=INSTANCE_TYPE,
    max_capacity=1,
    in_share_preemption=False,
    sharing_strategy="LEND"
)

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

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

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

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

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

from utils.aws_batch_resource_management import JobStatus, await_jobs, list_jobs_by_quota_share

# 向 QS1 提交三个任务，设置不同优先级以保证调度顺序
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 提交任务时，会从 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_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)