lambda_function.py

import boto3
import datetime
import botocore
import json

s3 = boto3.client("s3")
ecs = boto3.client("ecs")
ec2 = boto3.client("ec2")
cloudwatch = boto3.client("cloudwatch")
sqs = boto3.client("sqs")

bucket = "BUCKET_NAME"


def killdeadAlarms(fleetId, project):
    checkdates = [
        datetime.datetime.now().strftime("%Y-%m-%d"),
        (datetime.datetime.now() - datetime.timedelta(days=1)).strftime("%Y-%m-%d"),
    ]
    todel = []
    for eachdate in checkdates:
        datedead = ec2.describe_spot_fleet_request_history(
            SpotFleetRequestId=fleetId, StartTime=eachdate
        )
        for eachevent in datedead["HistoryRecords"]:
            if eachevent["EventType"] == "instanceChange":
                if eachevent["EventInformation"]["EventSubType"] == "terminated":
                    todel.append(eachevent["EventInformation"]["InstanceId"])
    todel = [f"{project}_{x}" for x in todel]
    while len(todel) > 100:
        dellist = todel[:100]
        cloudwatch.delete_alarms(AlarmNames=dellist)
        todel = todel[100:]
    if len(todel) <= 100:
        cloudwatch.delete_alarms(AlarmNames=todel)
    print("Old alarms deleted")


def seeIfLogExportIsDone(logExportId):
    while True:
        result = cloudwatch.describe_export_tasks(taskId=logExportId)
        if result["exportTasks"][0]["status"]["code"] != "PENDING":
            if result["exportTasks"][0]["status"]["code"] != "RUNNING":
                print(result["exportTasks"][0]["status"]["code"])
                break
                time.sleep(30)


def downscaleSpotFleet(nonvisible, spotFleetID):
    status = ec2.describe_spot_fleet_instances(SpotFleetRequestId=spotFleetID)
    if nonvisible < len(status["ActiveInstances"]):
        ec2.modify_spot_fleet_request(
            ExcessCapacityTerminationPolicy="noTermination",
            TargetCapacity=str(nonvisible),
            SpotFleetRequestId=spotFleetID,
        )


def check_sqs_queue(queueName):
    response = sqs.get_queue_url(QueueName=queueName)
    queueUrl = response["QueueUrl"]
    response = sqs.get_queue_attributes(
        QueueUrl=queueUrl,
        AttributeNames=[
            "ApproximateNumberOfMessages",
            "ApproximateNumberOfMessagesNotVisible",
        ],
    )
    visible = int(response["Attributes"]["ApproximateNumberOfMessages"])
    nonvisible = int(response["Attributes"]["ApproximateNumberOfMessagesNotVisible"])
    print(
        f"Found {visible} visible messages and {nonvisible} nonvisible messages in queue."
    )
    return visible, nonvisible


def lambda_handler(event, lambda_context):
    # Triggered any time SQS queue ApproximateNumberOfMessagesVisible = 0
    # OR ApproximateNumberOfMessagesNotVisible = 0
    messagestring = event["Records"][0]["Sns"]["Message"]
    messagedict = json.loads(messagestring)
    queueName = messagedict["Trigger"]["Dimensions"][0]["value"]
    project = queueName.rsplit("_", 1)[0]

    # Download monitor file
    monitor_file_name = f"{queueName.split('Queue')[0]}SpotFleetRequestId.json"
    monitor_local_name = f"/tmp/{monitor_file_name}"
    monitor_on_bucket_name = f"monitors/{monitor_file_name}"

    with open(monitor_local_name, "wb") as f:
        try:
            s3.download_fileobj(bucket, monitor_on_bucket_name, f)
        except botocore.exceptions.ClientError as error:
            print("Error retrieving monitor file.")
            return
    with open(monitor_local_name, "r") as input:
        monitorInfo = json.load(input)

    monitorcluster = monitorInfo["MONITOR_ECS_CLUSTER"]
    monitorapp = monitorInfo["MONITOR_APP_NAME"]
    fleetId = monitorInfo["MONITOR_FLEET_ID"]
    loggroupId = monitorInfo["MONITOR_LOG_GROUP_NAME"]
    CLEAN_DASHBOARD = monitorInfo["CLEAN_DASHBOARD"]
    print(f"Monitor triggered for {monitorcluster} {monitorapp} {fleetId} {loggroupId}")

    visible, nonvisible = check_sqs_queue(queueName)

    # If no visible messages, downscale machines
    if visible == 0 and nonvisible > 0:
        print("No visible messages. Tidying as we go.")
        killdeadAlarms(fleetId, project)
        downscaleSpotFleet(nonvisible, fleetId)

    # If no messages in progress, cleanup
    if visible == 0 and nonvisible == 0:
        print("No messages in progress. Cleaning up.")
        ecs.update_service(
            cluster=monitorcluster,
            service=f"{monitorapp}Service",
            desiredCount=0,
        )
        print("Service has been downscaled")

        # Delete the alarms from active machines and machines that have died.
        active_dictionary = ec2.describe_spot_fleet_instances(
            SpotFleetRequestId=fleetId
        )
        active_instances = []
        for instance in active_dictionary["ActiveInstances"]:
            active_instances.append(instance["InstanceId"])
        while len(active_instances) > 100:
            dellist = active_instances[:100]
            cloudwatch.delete_alarms(AlarmNames=dellist)
            active_instances = active_instances[100:]
        if len(active_instances) <= 100:
            cloudwatch.delete_alarms(AlarmNames=active_instances)
        killdeadAlarms(fleetId, monitorapp, project)

        # Read spot fleet id and terminate all EC2 instances
        ec2.cancel_spot_fleet_requests(
            SpotFleetRequestIds=[fleetId], TerminateInstances=True
        )
        print("Fleet shut down.")

        # Remove SQS queue, ECS Task Definition, ECS Service
        ECS_TASK_NAME = monitorapp + "Task"
        ECS_SERVICE_NAME = monitorapp + "Service"

        print("Deleting existing queue.")
        queueoutput = sqs.list_queues(QueueNamePrefix=queueName)
        try:
            if len(queueoutput["QueueUrls"]) == 1:
                queueUrl = queueoutput["QueueUrls"][0]
            else:  # In case we have "AnalysisQueue" and "AnalysisQueue1" and only want to delete the first of those
                for eachUrl in queueoutput["QueueUrls"]:
                    if eachUrl.split("/")[-1] == queueName:
                        queueUrl = eachUrl
            sqs.delete_queue(QueueUrl=queueUrl)
        except KeyError:
            print("Can't find queue to delete.")

        print("Deleting service")
        try:
            ecs.delete_service(cluster=monitorcluster, service=ECS_SERVICE_NAME)
        except:
            print("Couldn't delete service.")

        print("De-registering task")
        taskArns = ecs.list_task_definitions(familyPrefix=ECS_TASK_NAME)
        for eachtask in taskArns["taskDefinitionArns"]:
            fulltaskname = eachtask.split("/")[-1]
            ecs.deregister_task_definition(taskDefinition=fulltaskname)

        print("Removing cluster if it's not the default and not otherwise in use")
        if monitorcluster != "default":
            result = ecs.describe_clusters(clusters=[monitorcluster])
        if (
            sum(
                [
                    result["clusters"][0]["pendingTasksCount"],
                    result["clusters"][0]["runningTasksCount"],
                    result["clusters"][0]["activeServicesCount"],
                ]
            )
            == 0
        ):
            ecs.delete_cluster(cluster=monitorcluster)

        # Remove alarms that triggered monitor
        print("Removing alarms that triggered Monitor")
        cloudwatch.delete_alarms(
            AlarmNames=[
                f"ApproximateNumberOfMessagesVisibleisZero_{monitorapp}",
                f"ApproximateNumberOfMessagesNotVisibleisZero_{monitorapp}",
            ]
        )

        # Remove Cloudwatch dashboard if created and cleanup desired
        if CLEAN_DASHBOARD.lower() == "true":
            dashboard_list = cloudwatch.list_dashboards()
            for entry in dashboard_list["DashboardEntries"]:
                if monitorapp in entry["DashboardName"]:
                    cloudwatch.delete_dashboards(
                        DashboardNames=[entry["DashboardName"]]
                    )

        # Delete monitor file
        s3.delete_object(Bucket=bucket, Key=monitor_on_bucket_name)