AWS SageMaker Studio 일배치 실행하는 방법은?!

SageMaker Studio는 쥬피터랩과 같은 환경에서 원하는 인스턴스를 선택해서 실행할 수 있어 편리하다. 하지만, 일반 서버에서 작업하는 것과 거의 차이가 없어서, 신나게 알고리즘을 개발하고 나니 정작 일배치를 적용할 때 큰 어려움을 겪게 됐다. 필자와 같이 어려움을 겪고 있을 다른 사람들을 위해 SageMaker Studio에서 일배치 실행하는 방법에 대해서 다뤄보고자 한다.

 

 

AWS SageMaker Studio

인터넷을 찾아보면 SageMaker Studio에서 일배치로 실행하는 방법을 여럿 볼 수 있다. 그 중에 하나는 노트북 파일을 이용하는 것이다. 관련된 플러그인도 있는데, 이게 노트북 파일을 S3로 옮겨서 한다는 단점이 있다. 필자는 Custom이미지를 사용하고, 미리 작성한 다른 파이썬 파일도 이용해야 하기 때문에 이 방법을 그대로 적용하기에는 어려움이 있었다. 약간의 꼼수(?)를 이용해서 할 수 있는 방법이 있을 줄은 모르겠으나 그보다는 다른 방법을 이용하기로 했다.


필자가 AWS SageMaker Studio에서 일배치를 실행하는 방법은 아래와 같다.

우선 lambda함수와 EventBridge를 이용해 일단위로 lambda함수가 실행되게 한다. 인터넷을 뒤지다보니 lambda함수에서 sagemaker studio의 인스턴스를 켜고, system terminal에서 파이썬을 실행하는 코드를 찾게 됐다.


이에 아래 2가지 방향으로 코드를 분리하였다.

1. 학습에는 GPU 등의 인스턴스가 필요하므로, Container를 띄어서 학습하는 것으로 모두 수정하였다. system terminal에서 파이썬을 실행하면 필요한 인스턴스를 잡기 어렵기 때문에 위와 같이 바꾸었다.

2. container를 띄우기 전까지 필요한 코드는 최소한으로 하고, 별도 모듈이 필요한 부분은 파이썬 가상환경을 만들었다. system terminal에서 파이썬을 실행하면, 필요한 이미지의 파이썬 버전으로 실행할 수 없다는 단점이 있었다. 그래서 system terminal로 들어가, 파이썬 가상환경을 만들고 필요한 최소한의 패키지를 설치하였다.


위와 같이 작업하여 정상적으로 일배치 작업을 할 수 있는 것을 확인하였다. lambda함수 코드를 아래 첨부하니 필요하다면 참고하기 바란다.

 

import json
import boto3
import websocket
import requests
import datetime

def lambda_handler(event, context):
    
    client = boto3.client('sagemaker')

    # https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_CreateApp.html
    # https://towardsdatascience.com/run-setup-scripts-automatically-on-sagemaker-studio-15222b9d2f8c
    # https://docs.aws.amazon.com/sagemaker/latest/dg/nbi-lifecycle-config-install.html
    
    date_str = datetime.datetime.strftime(datetime.datetime.now(), '%Y%m%d%H%M%S')
    app_name = date_str
    userprofilename = 'SageMaker Studio 사용자이름'
    domain_id = '도메인ID'
    
    # instance 생성함
    response = client.create_app(
        DomainId=domain_id,
        UserProfileName = userprofilename,
        AppType='KernelGateway',
        AppName=app_name,
        Tags=[
            {
                'Key': app_name,
                'Value': app_name
            },
        ],
        ResourceSpec={
            'SageMakerImageArn': 'arn:으로 시작하는 주소',
            'SageMakerImageVersionArn': 'arn:~/1 형태의 주소 ',
            'InstanceType': 'ml.t3.medium'
        }
    )
    
    out = 0

    while(out==0):
        import time
        
        client_list = client.list_apps()
        client_list = [c for c in client_list['Apps'] if c['AppName'] == app_name and c['UserProfileName']==userprofilename]
        if client_list[0]['AppName'] == app_name and client_list[0]['Status']=='InService':
            out = 1
        print(client_list[0]['Status'])
        time.sleep(10)
            
    print(client_list)
        
    sagemaker_login_url = client.create_presigned_domain_url(
        DomainId=domain_id,
        UserProfileName = userprofilename,
    )["AuthorizedUrl"]
    
    
    import requests
    reqsess = requests.Session()
    login_resp = reqsess.get(sagemaker_login_url)
    
    base_url = sagemaker_login_url.partition("?")[0].rpartition("/")[0]
    api_base_url = base_url + "/jupyter/default"
    # Using XSRF as a proxy for "was Jupyter ready":
    if "_xsrf" not in reqsess.cookies:
        app_status = "Unknown"
        while app_status not in {"InService", "Terminated"}:
            time.sleep(2)
            app_status = reqsess.get(
                f"{base_url}/app?appType=JupyterServer&appName=default"
            ).text
        ready_resp = reqsess.get(api_base_url)
        
    terminal = reqsess.post(
        f"{api_base_url}/api/terminals",
        params={ "_xsrf": reqsess.cookies["_xsrf"] },
    ).json()
    terminal_name = terminal["name"]
    
    
    import websocket
    
    ws_base_url = \
        f"wss://{api_base_url.partition('://')[2]}/terminals/websocket"
    cookies = reqsess.cookies.get_dict()
    ws = websocket.create_connection(
        f"{ws_base_url}/{terminal_name}",
        cookie="; ".join(["%s=%s" %(i, j) for i, j in cookies.items()]),
    )
    
    
    COMMAND_SCRIPT = [
        "cd /home/sagemaker-user/development/venv && source bin/activate && cd /home/sagemaker-user/development && python3 /home/sagemaker-user/development/run.py"
    ]
    import re
    import json
    import time
    
    prompt_exp = re.compile(r"bash-[\d\.]+\$ $", re.MULTILINE)
    for ix, c in enumerate(COMMAND_SCRIPT):
        print(c)
        ws.send(json.dumps(["stdin", c + "\n"]))    
        while True:
            res = json.loads(ws.recv())
            if res[0] == "stdout" and prompt_exp.search(res[1]):
                # Ready for next command
                break
            time.sleep(0.1)


    # TODO implement
    return {
        'statusCode': 200,
        'body': json.dumps('a')
    }

오늘은 이렇게 AWS SageMaker Studio 일배치 실행하는 방법에 대해서 알아보았다. SageMaker Studio는 좋은 툴이지만 그에 맞는 적합한 방법으로 코드를 작성해야 하는 불편함도 있다. AWS SageMaker Studio에서 개발하는 분들에게 도움이 됐기를 바란다.