From 4f7db24d1e921afcd47eda71ae17b51bb0b87768 Mon Sep 17 00:00:00 2001 From: Murat Ugur Eminoglu Date: Wed, 22 May 2024 14:51:07 +0300 Subject: [PATCH] Refactor Lambda Script --- aws-custom-cluster/template.yaml | 195 ++++++++++++++++++++++++++----- 1 file changed, 167 insertions(+), 28 deletions(-) diff --git a/aws-custom-cluster/template.yaml b/aws-custom-cluster/template.yaml index 1881d859..3c841262 100644 --- a/aws-custom-cluster/template.yaml +++ b/aws-custom-cluster/template.yaml @@ -350,6 +350,8 @@ Resources: ImageId: ami-05603669082b6ebf0 SecurityGroupIds: - !GetAtt "InstanceSecurityGroup.GroupId" + IamInstanceProfile: + Arn: !GetAtt InstanceProfile.Arn BlockDeviceMappings: - DeviceName: /dev/sda1 Ebs: @@ -360,7 +362,7 @@ Resources: Fn::Base64: !Sub | #!/bin/bash sudo apt-get update -y - sudo apt-get install stunnel python3-pip -y + sudo apt-get install stunnel python3-pip jq -y sudo tee /etc/stunnel/stunnel.conf > /dev/null < 0)] | length > 0' /tmp/curl_output.txt + """ + asg_names = autoscaling_client.describe_auto_scaling_groups() asg_origin_name = [group for group in asg_names['AutoScalingGroups'] if - origin_asg in group['AutoScalingGroupName']] + origin_asg in group['AutoScalingGroupName']] asg_origin_group_names = [group['AutoScalingGroupName'] for group in asg_origin_name][0] - response = autoscaling_client.update_auto_scaling_group( - AutoScalingGroupName=asg_origin_group_names, - MinSize=0, - DesiredCapacity=0 - ) + print(f"Auto Scaling Group name: {asg_origin_group_names}") + origin_calculate_total_instance = autoscaling_client.describe_auto_scaling_groups(AutoScalingGroupNames=[asg_origin_group_names]) + origin_current_capacity = len(origin_calculate_total_instance['AutoScalingGroups'][0]['Instances']) + instance = origin_calculate_total_instance['AutoScalingGroups'][0]['Instances'] + try: + instance_id = instance[0]['InstanceId'] + print("Current capacity and Instance ID", {"origin_current_capacity": origin_current_capacity, "instance_id": instance_id}) + except IndexError: + cloudwatch_set_ok(alarm_name) + print("No instances found in Auto Scaling Group", asg_origin_group_names) + return { + 'statusCode': 200, + 'body': json.dumps({"message": "No instances found in Auto Scaling Group"}) + } - create_rule = elb_client.create_rule( - Actions=[ - { - 'Type': 'forward', - 'TargetGroupArn': target_group_arn - } - ], - Conditions=[ - { - 'Field': 'path-pattern', - 'Values': ['*'] - } - ], - ListenerArn=listener_arn, - Priority=1 - ) + smm_response = ssm_client.send_command( + DocumentName ='AWS-RunShellScript', + Parameters = {'commands': [script]}, + InstanceIds = [instance_id] + ) - print(response) + command_id = smm_response['Command']['CommandId'] + while True: + time.sleep(2) + invocation_response = ssm_client.get_command_invocation( + CommandId=command_id, + InstanceId=instance_id, + ) + if invocation_response['Status'] not in ['Pending', 'InProgress']: + break + + smm_output = invocation_response['StandardOutputContent'].strip() + print(f"SSM command output: {smm_output}") + #debug + #print(invocation_response['StandardOutputContent']) + #error = invocation_response['StandardErrorContent'].strip() + + if origin_current_capacity == 1: + if smm_output == 'false': + print("No live streams found. Updating Auto Scaling Group.") + origin_response = autoscaling_client.update_auto_scaling_group( + AutoScalingGroupName=asg_origin_group_names, + MinSize=0, + DesiredCapacity=0 + ) + + create_rule = elb_client.create_rule( + Actions=[ + { + 'Type': 'forward', + 'TargetGroupArn': target_group_arn + } + ], + Conditions=[ + { + 'Field': 'path-pattern', + 'Values': ['*'] + } + ], + ListenerArn=listener_arn, + Priority=1 + ) + cloudwatch_set_ok(alarm_name) + print(origin_response) + return { + 'statusCode': 200, + 'body': 'Auto Scaling Group updated successfully!' + } + else: + print("Live streams found.") + else: + print(f"Current capacity is not 1, it is {origin_current_capacity}. No updates needed.") + + cloudwatch_set_ok(alarm_name) return { 'statusCode': 200, - 'body': 'Auto Scaling Group updated successfully!' + 'body': 'Auto Scaling Group does not require update.' } + + + def cloudwatch_set_ok(alarm_name): + try: + cloudwatch_client = boto3.client('cloudwatch') + cloudwatch_response = cloudwatch_client.set_alarm_state( + AlarmName=alarm_name, + StateValue='OK', + StateReason='Updating alarm state to OK' + ) + print("Alarm state updated to OK", cloudwatch_response) + return cloudwatch_response + except Exception as e: + error_message = { + "error": str(e), + "function": "cloudwatch_set_ok", + "alarm_name": alarm_name + } + print(json.dumps(error_message)) + raise e + #FunctionName: InstanceDeleteFunction Handler: "index.lambda_handler" Role: !GetAtt LambdaIamRole.Arn @@ -928,7 +1031,7 @@ Resources: Service: "lambda.amazonaws.com" Action: "sts:AssumeRole" Policies: - - PolicyName: "EC2FullAccessPolicy1" + - PolicyName: "AccessPolicy" PolicyDocument: Version: "2012-10-17" Statement: @@ -948,6 +1051,9 @@ Resources: - "elasticloadbalancing:CreateRule" - "apigateway:GET" - "acm:DescribeCertificate" + - "ssm:GetCommandInvocation" + - "ssm:SendCommand" + - "cloudwatch:SetAlarmState" Resource: "*" - PolicyName: "CloudWatchLogsPolicy" PolicyDocument: @@ -977,6 +1083,39 @@ Resources: Action: lambda:InvokeFunction Principal: elasticloadbalancing.amazonaws.com + InstanceProfile: + Type: 'AWS::IAM::InstanceProfile' + Properties: + Roles: + - !Ref EC2Role + + EC2Role: + Type: 'AWS::IAM::Role' + Properties: + AssumeRolePolicyDocument: + Version: '2012-10-17' + Statement: + - Effect: Allow + Principal: + Service: [ec2.amazonaws.com] + Action: ['sts:AssumeRole'] + Policies: + - PolicyName: AllowSSM + PolicyDocument: + Version: '2012-10-17' + Statement: + - Effect: Allow + Action: + - "ssm:*" + - "cloudwatch:PutMetricData" + - "ds:CreateComputer" + - "ds:DescribeDirectories" + - "ec2:DescribeInstanceStatus" + - "logs:*" + - "ssm:*" + - "ec2messages:*" + Resource: '*' + # Cloudwatch rule to set 0 in Autoscale AutoScalingGroupScaleDownAlarm: Type: 'AWS::CloudWatch::Alarm'