Hyunsoo Lee

[AWS] Lambda를 이용하여 EC2 서버에 명령 전달 및 workflow 자동화

0. INTRO

  • Lambda는 Serverless 컴퓨팅 서비스로 별다른 서버의 구축이나 관리 등이 필요없이 단발성의 코드 실행을 하기에 정말 좋은 서비스이다. 또한 AWS 내부의 다른 서비스나 외부 소스들에서 발생하는 event를 trigger로 받아 더욱더 확장성 있는 활용이 가능하다.
  • 아래 포스팅에서는 EC2의 start 반응을 trigger로 받아 EC2의 재시작시 특정 명령을 서버에 던져주는 Lambda 함수를 만들어보려한다.
  • 해당 기능은 정기적인 배치 작업이 서버의 Docker Container 내부에서 돌 때, 서버 시작 -> Docker Container Start -> Batch Job Start 와 같은 workflow를 자동화 할 수 있다.

1. SSM 권한을 가진 Lmabda 함수 생성

  • AWS SSM(Systems Manager)는 AWS 인프라에 대한 가시성과 제어를 제공하는 서비스로 NAT Gateway가 연결되어 있지 않은 완전 private한 EC2 서버 접속하려 할 때 많이 사용되는 서비스이다. 이 서비스는 서버에 설치된 SSM Agent를 통하여 EC2에 대한 리소스 업데이트 및 관리 업무 등을 수행하게된다.
  • Lambda를 이용하여 제어 명령을 내릴 것이기 때문에 Lambda에 적절한 권한이 설정된 Role 생성 후 연결해준다.

    1. Lambda 권한 부여 및 생성

      image image

      • 다른 Role들이 포함되어 있긴 하지만 기본적으로 아래 두 정책이 포함된 Role을 Lambda 함수와 연결시켜주면 된다

        AWSLambdaBasicExecutionRole
        AmazonSSMFullAccess
        AmazonEC2FullAccess

    2. Lambda 실행 시간 설정

      • EC2 시작 및 shell 명령 전달시 시간이 걸리기 때문에 Lambda의 Default timeout 시간인 3초 보다는 더 긴 실행시간이 필요하다.
      • Configuration > General configuration > Timeout 탭에 들어가 실행시간을 약 1분 가량으로 늘려준다. image

2. EC2에 SSM Role 연결

  • SSM 기능을 활성화하고 접속 허용을 위해 EC2에도 적절한 Role을 연결해주어야 한다.

    1. AmazonSSMManagedInstanceCore 정책을 가진 Role 생성

    2. EC2와 해당 Role 연결

      • EC2 > Actions > Security > Modify IAM role 탭에 들어가서 정책을 연결해준다. image image

3. Lambda Trigger 설정

  1. EC2 Start에 대한 Eventbridge Rule 생성

    • Amazon EventBridge > Rules > Create rule 탭으로 들어가 Event의 이름 설정 및 아래와 같이 특정 EC2가 RUNNING 상태일 때 Event가 생성되도록 규칙을 정한다.

      image

    • 이후 탭에서 Skip to Review and update 버튼을 눌러 Event에 대해서만 생성을 완료한다.

  2. Lambda Trigger에 해당 Event 연결

    • Lambda > Add Trigger 섹션에서 위에서 만든 Eventbridge의 Rule을 Lambda 함수의 Trigger로 설정한다.

      image

4. Lambda 코드 작성

  • EC2가 완전히 Running 상태가 되기까진 약 30초 혹은 그 이상의 시간이 필요한데 그런 state를 고려하지 않고 EC2 Start Trigger를 받아 SSM을 통해 script를 바로 던지면 에러가 발생한다.
  • 이러한 문제를 해결하기 위해서는 Lambda 함수 내 while문을 통해 EC2 state를 몇초마다 체크해주어 서버가 완전히 Running 상태가 되었을 때 Command를 던지도록 하는 로직이 필요하다.
  • boto3 client 선언의 경우 반복적인 선언이 불필요하기 때문에 lambda_handler 함수 밖으로 빼서 선언해주는 것이 시간 및 리소스 절약을 위해서 좋지만 단발성 실행의 경우에는 내부에 선언해도 크게 상관은 없다.
import time
import boto3
from datetime import date, datetime, timezone, timedelta

# EC2 시작, 종료 시간 표시를 위한 timestamp
KST = timezone(timedelta(hours=9))
time_record = datetime.now(KST).strftime('%Y-%m-%d %H:%M:%S')


# EC2 instance의 현재 상태 체크를 위한 함수
def get_instance_status(instance_id):
    ec2 = boto3.client("ec2")
    describeInstance = ec2.describe_instances(InstanceIds=[instance_id])
    status = describeInstance['Reservations'][0]['Instances'][0]['State']['Name']
    return status

# SSM을 통해 ec2 서버에 command를 전달해주는 함수
def run_shell_script_on_ec2(instance_id, shell_script):
    ssm_client = boto3.client('ssm')
    response = ssm_client.send_command(
        InstanceIds=[instance_id],
        DocumentName='AWS-RunShellScript',
        Parameters={
            'commands': [shell_script]
        }
    )
    command_id = response['Command']['CommandId']

    # send_command한 결과를 get_command_invocation으로 받는데 시간차가 있기 때문에 time.sleep을 꼭 써준다.
    time.sleep(5)
    output = ssm_client.get_command_invocation(CommandId=command_id, InstanceId=instance_id)
    return output

shell_script = 'INPUT YOUR CUSTOM SHELL COMMAND OR SCRIPT'

def lambda_handler(event, context):

    INSTANCE_ID = '인스턴트ID 입력'
    status = get_instance_status(INSTANCE_ID)
	
    if status == 'stopped':
        ec2.start_instances(InstanceIds=[INSTANCE_ID])
    
        TRIAL_CNT = 30
        SLEEP_TIME = 10
        cnt = 0

        # EC2 상태가 Running 될 때까지 체크하는 로직
        while cnt < TRIAL_CNT:
            time.sleep(SLEEP_TIME)
            current_status = get_instance_status(INSTANCE_ID)
            if current_status == 'running':
                time.sleep(5)
                try:
                    output = run_shell_script_on_ec2(INSTANCE_ID, shell_script)

                    print(f"EC2 Starts at {time_record}")
                    return {
                        'statusCode': 200,
                        'body': output
                    }
                except: pass
            else:
                cnt += 1

5. OUTRO

  • 회사 업무 수행시 Step Function으로 부터 인자를 넘겨받아 EC2 서버의 docker container에서 작업이 이루어져야하는 정기적인 월배치(Monthly Batch) 작업 구성시 위의 기능을 사용하였고 덕분에 workflow의 많은 부분을 자동화할 수 있었다.
  • 코드 자체는 간단하지만 기능은 파워풀하기 때문에 정기 배치 작업이 도는 EC2가 포함되어 있는 아키텍쳐에 종종 쓰일 수 있을 것이라 생각한다.
  • 서버가 stop에서 start로 변하거나 서버에 command가 던져진 후 수행될때까지 항상 시간차가 있기 때문에 서버의 상태를 체크해주는 로직이나 시간차를 주는 명령어들을 사용함으로써 에러 발생을 방지한다.

Leave a comment

You may also enjoy