82 lines
2.4 KiB
Python
82 lines
2.4 KiB
Python
import logging
|
|
import os
|
|
import time
|
|
from datetime import datetime, timedelta
|
|
|
|
import boto3
|
|
from aws_lambda_powertools import Tracer
|
|
|
|
from common.ddb_service.client import DynamoDbUtilsService
|
|
from common.response import ok, not_found
|
|
from libs.comfy_data_types import ComfyExecuteTable
|
|
from libs.utils import response_error
|
|
|
|
tracer = Tracer()
|
|
logger = logging.getLogger(__name__)
|
|
logger.setLevel(os.environ.get('LOG_LEVEL') or logging.ERROR)
|
|
ddb = boto3.resource('dynamodb')
|
|
client = boto3.client('logs')
|
|
ddb_service = DynamoDbUtilsService(logger=logger)
|
|
|
|
execute_table = os.environ.get('EXECUTE_TABLE')
|
|
|
|
|
|
@tracer.capture_lambda_handler
|
|
def handler(event, ctx):
|
|
try:
|
|
logger.info(f"Received ctx: {ctx}")
|
|
prompt_id = event['pathParameters']['id']
|
|
|
|
item = ddb_service.get_item(table=execute_table, key_values={"prompt_id": prompt_id})
|
|
logger.info(item)
|
|
|
|
if not item:
|
|
return not_found(f"execute not found for prompt_id: {prompt_id}")
|
|
|
|
job = ComfyExecuteTable(**item)
|
|
logger.info(job)
|
|
|
|
query = f"""
|
|
fields @timestamp, @message, @logStream
|
|
| filter @message like /{prompt_id}/
|
|
| sort @timestamp asc
|
|
| limit 1000
|
|
"""
|
|
|
|
dt = datetime.strptime(job.start_time, '%Y-%m-%dT%H:%M:%S.%f')
|
|
response = client.start_query(
|
|
logGroupName=f'/aws/sagemaker/Endpoints/{job.endpoint_name}',
|
|
startTime=int(dt.timestamp()),
|
|
endTime=int((dt + timedelta(hours=1)).timestamp()),
|
|
queryString=query,
|
|
limit=1000
|
|
)
|
|
|
|
query_id = response['queryId']
|
|
|
|
# Wait for the query to complete
|
|
while True:
|
|
query_status = client.get_query_results(
|
|
queryId=query_id
|
|
)
|
|
status = query_status['status']
|
|
if status == 'Complete':
|
|
break
|
|
elif status == 'Failed' or status == 'Cancelled':
|
|
print("Query failed or was cancelled")
|
|
break
|
|
time.sleep(1)
|
|
|
|
results = []
|
|
for result in query_status['results']:
|
|
item = {field['field']: field['value'] for field in result}
|
|
results.append({
|
|
'timestamp': item["@timestamp"],
|
|
'message': item["@message"],
|
|
'logStream': item["@logStream"],
|
|
})
|
|
|
|
return ok(data={'results': results}, decimal=True)
|
|
except Exception as e:
|
|
return response_error(e)
|