import base64 import json import logging import os import uuid from dataclasses import dataclass from datetime import datetime from typing import Optional import boto3 from aws_lambda_powertools import Tracer from sagemaker import Predictor from sagemaker.deserializers import JSONDeserializer from sagemaker.predictor_async import AsyncPredictor from sagemaker.serializers import JSONSerializer from common.ddb_service.client import DynamoDbUtilsService from common.excepts import BadRequestException from common.response import ok, created from common.util import s3_scan_files, generate_presigned_url_for_keys, \ record_latency_metrics, record_count_metrics, get_workflow_name from libs.comfy_data_types import ComfyExecuteTable, InferenceResult from libs.enums import ComfyExecuteType, EndpointStatus, ServiceType from libs.utils import get_endpoint_by_name, response_error, get_endpoint_name_by_workflow_name tracer = Tracer() region = os.environ.get('AWS_REGION') bucket_name = os.environ.get('S3_BUCKET_NAME') execute_table = os.environ.get('EXECUTE_TABLE') logger = logging.getLogger(__name__) logger.setLevel(os.environ.get('LOG_LEVEL') or logging.ERROR) endpoint_instance_id = os.environ.get('ENDPOINT_INSTANCE_ID') ddb_service = DynamoDbUtilsService(logger=logger) sqs_url = os.environ.get('MERGE_SQS_URL') index_name = "endpoint_name-startTime-index" predictors = {} multi_gpu_instance_type_list = ['ml.p5.48xlarge', 'ml.p4d.24xlarge', 'ml.p3.8xlarge', 'ml.p3.16xlarge', 'ml.p3dn.24xlarge', 'ml.p2.8xlarge', 'ml.p2.16xlarge', 'ml.g4dn.12xlarge', 'ml.g5.12xlarge', 'ml.g5.24xlarge', 'ml.g5.48xlarge'] @dataclass class PrepareProps: need_reboot: Optional[bool] = False prepare_type: Optional[str] = "inputs" s3_source_path: Optional[str] = None local_target_path: Optional[str] = None sync_script: Optional[str] = None @dataclass class ExecuteEvent: prompt_id: str prompt: dict endpoint_name: Optional[str] = '' need_sync: bool = True number: Optional[str] = None front: Optional[bool] = None extra_data: Optional[dict] = None client_id: Optional[str] = None workflow: Optional[str] = None need_prepare: Optional[bool] = False prepare_props: Optional[PrepareProps] = None multi_async: Optional[bool] = False workflow_name: Optional[str] = None def sen_sqs_msg(message_body, endpoint_name): sqs_client = boto3.client('sqs', region_name=region) response = sqs_client.send_message( QueueUrl=sqs_url, MessageBody=json.dumps(message_body), MessageGroupId=endpoint_name ) message_id = response['MessageId'] return message_id def build_s3_images_request(prompt_id, bucket_name, s3_path): s3 = boto3.client('s3', region_name=region) response = s3.list_objects_v2(Bucket=bucket_name, Prefix=s3_path) image_video_dict = {} for obj in response.get('Contents', []): object_key = obj['Key'] file_name = object_key.split('/')[-1] if object_key.lower().endswith(('.png', '.jpg', '.jpeg', '.gif', '.bmp', '.mp4', '.mov', '.avi')): response_data = s3.get_object(Bucket=bucket_name, Key=object_key) object_data = response_data['Body'].read() encoded_data = base64.b64encode(object_data).decode('utf-8') image_video_dict[file_name] = encoded_data return {'prompt_id': prompt_id, 'image_video_data': image_video_dict} @tracer.capture_method def invoke_sagemaker_inference(event: ExecuteEvent): if not event.endpoint_name and not event.workflow_name: raise Exception(f"Cannot match an available environment，please check your EndpointName or your WorkflowVersion.") if event.workflow_name: endpoint_name = get_endpoint_name_by_workflow_name(name=event.workflow_name) else: endpoint_name = event.endpoint_name try: ep = get_endpoint_by_name(endpoint_name) except Exception as e: raise Exception(f"Please create an endpoint first, then try again.") if ep.endpoint_status not in [EndpointStatus.IN_SERVICE.value, EndpointStatus.UPDATING.value]: raise Exception(f"Endpoint {endpoint_name} is {ep.endpoint_status} status, not InService or Updating.") if event.workflow: event.workflow = get_workflow_name(event.workflow, ep.instance_type) logger.info(f"endpoint: {ep}") payload = event.__dict__ logger.info('inference payload: {}'.format(payload)) inference_id = str(uuid.uuid4()) job_status = ComfyExecuteType.CREATED.value inference_job = ComfyExecuteTable( prompt_id=event.prompt_id, endpoint_name=event.endpoint_name, inference_type=ep.endpoint_type, instance_id=endpoint_instance_id, need_sync=event.need_sync, status=job_status, prompt_params={'prompt': str(event.prompt), 'number': event.number, 'front': event.front, 'extra_data': str(event.extra_data), 'client_id': event.client_id}, prompt_path='', create_time=datetime.now().isoformat(), sagemaker_raw={}, output_path='', temp_path='', output_files=[], temp_files=[], multi_async=event.multi_async, batch_id='', workflow=event.workflow, ) logger.info(f"inference job: {inference_job.__dict__}") record_count_metrics(ep_name=ep.endpoint_name, metric_name='InferenceTotal', service=ServiceType.Comfy.value, workflow=inference_job.workflow ) if event.multi_async and ep.endpoint_type == 'Async' and ep.instance_type in multi_gpu_instance_type_list: logger.info(f"executing multi-gpu inference {ep.instance_type} {ep.endpoint_type} {event.multi_async}") save_item = inference_job.__dict__ ddb_service.put_items(execute_table, entries=save_item) sen_sqs_msg({"event": payload, "save_item": save_item, "inference_id": inference_id}, endpoint_name) # just for test multi gpu # payload1 = payload # payload1['prompt_id'] = payload['prompt_id']+"1" # save_item1 = save_item # save_item1['prompt_id'] = payload['prompt_id']+"1" # sen_sqs_msg({"event": payload1, "save_item": save_item1, "inference_id": inference_id+"1"}, endpoint_name) # payload2 = payload # payload2['prompt_id'] = payload['prompt_id'] + "1" # save_item2 = save_item # save_item2['prompt_id'] = payload['prompt_id'] + "1" # sen_sqs_msg({"event": payload2, "save_item": save_item2, "inference_id": inference_id+"2"}, endpoint_name) return created(data=response_schema(inference_job), decimal=True) elif ep.endpoint_type == 'Async': ddb_service.put_items(execute_table, entries=inference_job.__dict__) resp = async_inference([payload], inference_id, ep.endpoint_name) # TODO status check and save logger.info(f"async inference {ep.instance_type} {ep.endpoint_type} {event.multi_async} response: {resp}") return created(data=response_schema(inference_job), decimal=True) inference_job.start_time = datetime.now().isoformat() ddb_service.put_items(execute_table, entries=inference_job.__dict__) resp = real_time_inference([payload], inference_id, ep.endpoint_name) logger.info(f"real time inference response: {resp}") if resp and len(resp) > 0: resp = InferenceResult(**resp[0]) resp = s3_scan_files(resp) inference_job.status = resp.status inference_job.sagemaker_raw = resp.__dict__ inference_job.output_path = resp.output_path inference_job.output_files = resp.output_files inference_job.temp_path = resp.temp_path inference_job.temp_files = resp.temp_files inference_job.complete_time = datetime.now().isoformat() if resp.status == 'fail' or resp.status != 'Completed' or resp.status != 'success': inference_job.message = resp.message ddb_service.put_items(execute_table, entries=inference_job.__dict__) if ep.endpoint_type == 'Real-time': inference_job.output_files = generate_presigned_url_for_keys(inference_job.output_path, inference_job.output_files) inference_job.temp_files = generate_presigned_url_for_keys(inference_job.temp_path, inference_job.temp_files) if resp.status == 'fail' or resp.status != 'Completed' or resp.status != 'success': inference_job.message = resp.message record_count_metrics(ep_name=ep.endpoint_name, metric_name='InferenceFailed', service=ServiceType.Comfy.value, workflow=inference_job.workflow ) else: record_count_metrics(ep_name=ep.endpoint_name, metric_name='InferenceSucceed', service=ServiceType.Comfy.value, workflow=inference_job.workflow ) record_latency_metrics(start_time=inference_job.start_time, ep_name=ep.endpoint_name, metric_name='InferenceLatency', workflow=inference_job.workflow, service=ServiceType.Comfy.value) else: logger.info(f"inference error by sg resp none!{resp}") record_count_metrics(ep_name=ep.endpoint_name, metric_name='InferenceFailed', service=ServiceType.Comfy.value, workflow=inference_job.workflow ) return ok(data=response_schema(inference_job), decimal=True) def response_schema(inference_job: ComfyExecuteTable): if not inference_job.output_files: inference_job.output_files = [] if not inference_job.temp_files: inference_job.temp_files = [] data = { 'prompt_id': inference_job.prompt_id, 'status': inference_job.status, 'create_time': inference_job.create_time, 'endpoint_name': inference_job.endpoint_name, 'inference_type': inference_job.inference_type, 'need_sync': inference_job.need_sync, 'start_time': inference_job.start_time, 'complete_time': inference_job.complete_time, 'output_path': inference_job.output_path, 'output_files': inference_job.output_files, 'temp_path': inference_job.temp_path, 'temp_files': inference_job.temp_files, } return data @tracer.capture_lambda_handler def handler(raw_event, ctx): try: logger.info(f"execute start... Received event: {raw_event}") logger.info(f"Received ctx: {ctx}") event = ExecuteEvent(**json.loads(raw_event['body'])) if not event.prompt: raise BadRequestException("Prompt is required") return invoke_sagemaker_inference(event) except Exception as e: return response_error(e) @tracer.capture_method def async_inference(payload: any, inference_id, endpoint_name): tracer.put_annotation(key="inference_id", value=inference_id) initial_args = {"InvocationTimeoutSeconds": 3600} return get_async_predict_client(endpoint_name).predict_async(data=payload, initial_args=initial_args, inference_id=inference_id) @tracer.capture_method def real_time_inference(data: any, inference_id, endpoint_name): tracer.put_annotation(key="inference_id", value=inference_id) return get_real_time_predict_client(endpoint_name).predict(data=data, inference_id=inference_id) @tracer.capture_method def get_real_time_predict_client(endpoint_name): tracer.put_annotation(key="endpoint_name", value=endpoint_name) if endpoint_name in predictors: return predictors[endpoint_name] predictor = Predictor(endpoint_name) predictor.serializer = JSONSerializer() predictor.deserializer = JSONDeserializer() predictors[endpoint_name] = predictor return predictor @tracer.capture_method def get_async_predict_client(endpoint_name): tracer.put_annotation(key="endpoint_name", value=endpoint_name) if endpoint_name in predictors: return predictors[endpoint_name] predictor = Predictor(endpoint_name) predictor = AsyncPredictor(predictor, name=endpoint_name) predictor.serializer = JSONSerializer() predictor.deserializer = JSONDeserializer() predictors[endpoint_name] = predictor return predictor