327 lines
13 KiB
Python
327 lines
13 KiB
Python
import json
|
|
import logging
|
|
import os
|
|
from datetime import datetime
|
|
|
|
import boto3
|
|
from aws_lambda_powertools import Tracer
|
|
|
|
from common.util import record_seconds_metrics, endpoint_clean
|
|
from inferences.inference_libs import update_table_by_pk
|
|
from libs.data_types import Endpoint
|
|
from libs.enums import EndpointStatus, EndpointType, ServiceType
|
|
from libs.utils import get_endpoint_by_name
|
|
|
|
lambda_client = boto3.client('lambda')
|
|
|
|
tracer = Tracer()
|
|
sagemaker_endpoint_table = os.environ.get('ENDPOINT_TABLE_NAME')
|
|
|
|
logger = logging.getLogger(__name__)
|
|
logger.setLevel(os.environ.get('LOG_LEVEL') or logging.ERROR)
|
|
|
|
autoscaling_client = boto3.client('application-autoscaling')
|
|
cw_client = boto3.client('cloudwatch')
|
|
sagemaker = boto3.client('sagemaker')
|
|
|
|
esd_version = os.environ.get("ESD_VERSION")
|
|
cool_down_period = 15 * 60 # 15 minutes
|
|
|
|
s3_bucket_name = os.environ.get('S3_BUCKET_NAME')
|
|
s3 = boto3.resource('s3')
|
|
bucket = s3.Bucket(s3_bucket_name)
|
|
aws_region = os.environ.get('AWS_REGION')
|
|
|
|
|
|
# lambda: handle sagemaker events
|
|
@tracer.capture_lambda_handler
|
|
def handler(event, context):
|
|
logger.info(json.dumps(event))
|
|
endpoint_name = event['detail']['EndpointName']
|
|
endpoint_status = event['detail']['EndpointStatus']
|
|
current_time = datetime.now().isoformat()
|
|
|
|
try:
|
|
endpoint = get_endpoint_by_name(endpoint_name)
|
|
|
|
business_status = get_business_status(endpoint_status)
|
|
|
|
update_endpoint_field(endpoint, 'endpoint_status', business_status)
|
|
|
|
if business_status == EndpointStatus.UPDATING.value:
|
|
update_endpoint_field(endpoint, 'startTime', current_time)
|
|
|
|
if business_status == EndpointStatus.IN_SERVICE.value:
|
|
update_endpoint_field(endpoint, 'endTime', current_time)
|
|
|
|
if endpoint.service_type == 'sd':
|
|
service_type = ServiceType.SD.value
|
|
else:
|
|
service_type = ServiceType.Comfy.value
|
|
|
|
record_seconds_metrics(start_time=endpoint.startTime,
|
|
metric_name='EndpointReadySeconds',
|
|
service=service_type)
|
|
|
|
# if it is the first time in service
|
|
if not endpoint.endTime:
|
|
check_and_enable_autoscaling(endpoint, 'prod')
|
|
|
|
# update the instance count if the endpoint is not deleting or deleted
|
|
if business_status not in [EndpointStatus.DELETING.value, EndpointStatus.DELETED.value]:
|
|
status = sagemaker.describe_endpoint(EndpointName=endpoint_name)
|
|
logger.info(f"Endpoint status: {status}")
|
|
if 'ProductionVariants' in status:
|
|
instance_count = status['ProductionVariants'][0]['CurrentInstanceCount']
|
|
update_endpoint_field(endpoint, 'current_instance_count', instance_count)
|
|
else:
|
|
endpoint_clean(endpoint)
|
|
|
|
if business_status == EndpointStatus.FAILED.value:
|
|
update_endpoint_field(endpoint, 'error', event['FailureReason'])
|
|
|
|
except Exception as e:
|
|
logger.error(e, exc_info=True)
|
|
|
|
return {'statusCode': 200}
|
|
|
|
|
|
def check_and_enable_autoscaling(ep: Endpoint, variant_name):
|
|
if ep.autoscaling:
|
|
enable_autoscaling(ep, variant_name)
|
|
else:
|
|
logger.info(f'no need to enable autoscaling')
|
|
|
|
|
|
@tracer.capture_method
|
|
def enable_autoscaling(ep: Endpoint, variant_name):
|
|
tracer.put_annotation("variant_name", variant_name)
|
|
max_instance_number = int(ep.max_instance_number)
|
|
|
|
min_instance_number = 0
|
|
if ep.endpoint_type == EndpointType.RealTime.value:
|
|
min_instance_number = 1
|
|
|
|
if ep.min_instance_number is not None:
|
|
min_instance_number = int(ep.min_instance_number)
|
|
|
|
# Register scalable target
|
|
response = autoscaling_client.register_scalable_target(
|
|
ServiceNamespace='sagemaker',
|
|
ResourceId='endpoint/' + ep.endpoint_name + '/variant/' + variant_name,
|
|
ScalableDimension='sagemaker:variant:DesiredInstanceCount',
|
|
MinCapacity=min_instance_number,
|
|
MaxCapacity=max_instance_number,
|
|
)
|
|
logger.info(f"Register scalable target response: {response}")
|
|
|
|
if ep.endpoint_type == EndpointType.Async.value:
|
|
enable_autoscaling_async(ep, variant_name)
|
|
|
|
if ep.endpoint_type == EndpointType.RealTime.value:
|
|
enable_autoscaling_real_time(ep, variant_name)
|
|
|
|
|
|
def enable_autoscaling_async(ep: Endpoint, variant_name):
|
|
target_value = 3
|
|
|
|
# Define scaling policy
|
|
response = autoscaling_client.put_scaling_policy(
|
|
PolicyName=f"{ep.endpoint_name}-Invocations-ScalingPolicy",
|
|
ServiceNamespace="sagemaker", # The namespace of the AWS service that provides the resource.
|
|
ResourceId='endpoint/' + ep.endpoint_name + '/variant/' + variant_name, # Endpoint name
|
|
ScalableDimension="sagemaker:variant:DesiredInstanceCount", # SageMaker supports only Instance Count
|
|
PolicyType="TargetTrackingScaling", # 'StepScaling'|'TargetTrackingScaling'
|
|
TargetTrackingScalingPolicyConfiguration={
|
|
"TargetValue": target_value,
|
|
# The target value for the metric. - here the metric is - SageMakerVariantInvocationsPerInstance
|
|
"CustomizedMetricSpecification": {
|
|
"MetricName": "ApproximateBacklogSizePerInstance",
|
|
"Namespace": "AWS/SageMaker",
|
|
"Dimensions": [{"Name": "EndpointName", "Value": ep.endpoint_name}],
|
|
"Statistic": "Average",
|
|
},
|
|
"ScaleInCooldown": 180,
|
|
# The cooldown period helps you prevent your Auto Scaling group from launching or terminating
|
|
"ScaleOutCooldown": 60
|
|
# ScaleOutCooldown - The amount of time, in seconds, after a scale out activity completes before another
|
|
# scale out activity can start.
|
|
},
|
|
)
|
|
logger.info(f"Put scaling policy response")
|
|
logger.info(json.dumps(response))
|
|
alarms = response.get('Alarms')
|
|
for alarm in alarms:
|
|
alarm_name = alarm.get('AlarmName')
|
|
logger.info(f"Alarm name: {alarm_name}")
|
|
response = cw_client.describe_alarms(
|
|
AlarmNames=[alarm_name]
|
|
)
|
|
logger.info(f"Describe alarm response")
|
|
logger.info(response)
|
|
comparison_operator = response.get('MetricAlarms')[0]['ComparisonOperator']
|
|
if comparison_operator == "LessThanThreshold":
|
|
period = cool_down_period # 15 minutes
|
|
evaluation_periods = 4
|
|
datapoints_to_alarm = 4
|
|
target_value = 1
|
|
else:
|
|
period = 30
|
|
evaluation_periods = 1
|
|
datapoints_to_alarm = 1
|
|
target_value = 3
|
|
response = cw_client.put_metric_alarm(
|
|
AlarmName=alarm_name,
|
|
Namespace='AWS/SageMaker',
|
|
MetricName='ApproximateBacklogSizePerInstance',
|
|
Statistic="Average",
|
|
Period=period,
|
|
EvaluationPeriods=evaluation_periods,
|
|
DatapointsToAlarm=datapoints_to_alarm,
|
|
Threshold=target_value,
|
|
ComparisonOperator=comparison_operator,
|
|
AlarmActions=response.get('MetricAlarms')[0]['AlarmActions'],
|
|
Dimensions=[{'Name': 'EndpointName', 'Value': ep.endpoint_name}]
|
|
)
|
|
logger.info(f"Put metric alarm response")
|
|
logger.info(response)
|
|
|
|
step_policy_response = autoscaling_client.put_scaling_policy(
|
|
PolicyName=f"{ep.endpoint_name}-HasBacklogWithoutCapacity-ScalingPolicy",
|
|
ServiceNamespace="sagemaker", # The namespace of the service that provides the resource.
|
|
ResourceId='endpoint/' + ep.endpoint_name + '/variant/' + variant_name,
|
|
ScalableDimension="sagemaker:variant:DesiredInstanceCount", # SageMaker supports only Instance Count
|
|
PolicyType="StepScaling", # 'StepScaling' or 'TargetTrackingScaling'
|
|
StepScalingPolicyConfiguration={
|
|
"AdjustmentType": "ChangeInCapacity",
|
|
# Specifies whether the ScalingAdjustment value in the StepAdjustment property is an absolute number or a
|
|
# percentage of the current capacity.
|
|
"MetricAggregationType": "Average", # The aggregation type for the CloudWatch metrics.
|
|
"Cooldown": 180, # The amount of time, in seconds, to wait for a previous scaling activity to take effect.
|
|
"StepAdjustments": # A set of adjustments that enable you to scale based on the size of the alarm breach.
|
|
[
|
|
{
|
|
"MetricIntervalLowerBound": 0,
|
|
"ScalingAdjustment": 1
|
|
}
|
|
]
|
|
},
|
|
)
|
|
logger.info(f"Put step scaling policy response: {step_policy_response}")
|
|
|
|
cw_client.put_metric_alarm(
|
|
AlarmName=f'{ep.endpoint_name}-HasBacklogWithoutCapacity-Alarm',
|
|
MetricName='HasBacklogWithoutCapacity',
|
|
Namespace='AWS/SageMaker',
|
|
Statistic='Average',
|
|
Period=30,
|
|
EvaluationPeriods=1,
|
|
DatapointsToAlarm=1,
|
|
Threshold=1,
|
|
ComparisonOperator='GreaterThanOrEqualToThreshold',
|
|
TreatMissingData='missing',
|
|
Dimensions=[
|
|
{'Name': 'EndpointName', 'Value': ep.endpoint_name},
|
|
],
|
|
AlarmActions=[step_policy_response['PolicyARN']]
|
|
)
|
|
logger.info(f"Put metric alarm response: {step_policy_response}")
|
|
|
|
logger.info(f"Autoscaling has been enabled for the endpoint: {ep.endpoint_name}")
|
|
|
|
|
|
@tracer.capture_method
|
|
def enable_autoscaling_real_time(ep: Endpoint, variant_name):
|
|
tracer.put_annotation("variant_name", variant_name)
|
|
target_value = 5
|
|
|
|
# Define scaling policy
|
|
response = autoscaling_client.put_scaling_policy(
|
|
PolicyName=f"{ep.endpoint_name}-Invocations-ScalingPolicy",
|
|
ServiceNamespace="sagemaker", # The namespace of the AWS service that provides the resource.
|
|
ResourceId='endpoint/' + ep.endpoint_name + '/variant/' + variant_name, # Endpoint name
|
|
ScalableDimension="sagemaker:variant:DesiredInstanceCount", # SageMaker supports only Instance Count
|
|
PolicyType="TargetTrackingScaling", # 'StepScaling'|'TargetTrackingScaling'
|
|
TargetTrackingScalingPolicyConfiguration={
|
|
"TargetValue": target_value,
|
|
"PredefinedMetricSpecification":
|
|
{
|
|
"PredefinedMetricType": "SageMakerVariantInvocationsPerInstance"
|
|
},
|
|
"ScaleInCooldown": 180,
|
|
# The cooldown period helps you prevent your Auto Scaling group from launching or terminating
|
|
"ScaleOutCooldown": 60
|
|
# ScaleOutCooldown - The amount of time, in seconds, after a scale out activity completes before another
|
|
# scale out activity can start.
|
|
},
|
|
)
|
|
logger.info(f"Put scaling policy response")
|
|
logger.info(json.dumps(response))
|
|
alarms = response.get('Alarms')
|
|
for alarm in alarms:
|
|
alarm_name = alarm.get('AlarmName')
|
|
logger.info(f"Alarm name: {alarm_name}")
|
|
response = cw_client.describe_alarms(
|
|
AlarmNames=[alarm_name]
|
|
)
|
|
logger.info(f"Describe alarm response")
|
|
logger.info(response)
|
|
comparison_operator = response.get('MetricAlarms')[0]['ComparisonOperator']
|
|
if comparison_operator == "LessThanThreshold":
|
|
period = cool_down_period # 15 minutes
|
|
evaluation_periods = 4
|
|
datapoints_to_alarm = 4
|
|
target_value = 1
|
|
else:
|
|
period = 30
|
|
evaluation_periods = 1
|
|
datapoints_to_alarm = 1
|
|
target_value = 5
|
|
response = cw_client.put_metric_alarm(
|
|
AlarmName=alarm_name,
|
|
Namespace='AWS/SageMaker',
|
|
MetricName='InvocationsPerInstance',
|
|
Statistic="Sum",
|
|
Period=period,
|
|
EvaluationPeriods=evaluation_periods,
|
|
DatapointsToAlarm=datapoints_to_alarm,
|
|
Threshold=target_value,
|
|
ComparisonOperator=comparison_operator,
|
|
AlarmActions=response.get('MetricAlarms')[0]['AlarmActions'],
|
|
Dimensions=[
|
|
{'Name': 'EndpointName', 'Value': ep.endpoint_name},
|
|
{'Name': 'VariantName', 'Value': 'prod'},
|
|
]
|
|
)
|
|
logger.info(f"Put metric alarm response")
|
|
logger.info(response)
|
|
|
|
logger.info(f"Autoscaling has been enabled for the endpoint: {ep.endpoint_name}")
|
|
|
|
|
|
def update_endpoint_field(ep: Endpoint, field_name, field_value):
|
|
update_table_by_pk(
|
|
table_name=sagemaker_endpoint_table,
|
|
pk='EndpointDeploymentJobId',
|
|
id=ep.EndpointDeploymentJobId,
|
|
key=field_name,
|
|
value=field_value
|
|
)
|
|
|
|
|
|
def get_business_status(status):
|
|
"""
|
|
Convert SageMaker endpoint status to business status
|
|
:param status: EventBridge event status(upper case)
|
|
:return: business status
|
|
"""
|
|
switcher = {
|
|
"IN_SERVICE": EndpointStatus.IN_SERVICE.value,
|
|
"CREATING": EndpointStatus.CREATING.value,
|
|
"DELETED": EndpointStatus.DELETED.value,
|
|
"FAILED": EndpointStatus.FAILED.value,
|
|
"UPDATING": EndpointStatus.UPDATING.value,
|
|
"DELETING": EndpointStatus.DELETING.value,
|
|
}
|
|
return switcher.get(status, status)
|