생성형 AI를 위한 Amazon SageMaker Endpoint 기반 임베딩 모델 배포

서론

생성형 AI의 기술이 하루가 다르게 발전하고 있으며, 그 접근법 또한 다양해지고 있습니다. 그 중, RAG(Retrieval Augmented Generation, 검색 증강 생성) 접근은 도메인의 문제를 효율적으로 풀기 위한 접근법으로, 생성형 AI의 필수적인 접근 중 하나가 되었습니다.

RAG를 활용하는 데에는 여러 방법이 있지만, 가장 널리 사용되는 방법은 임베딩을 사용한 유사도 검색에 기반을 두고 있습니다. 이 접근법은 텍스트를 벡터 형태로 변환하는 과정에서 시작되며, 이를 위해서는 다양한 임베딩 모델들을 실험하고 검토하는 것이 매우 중요합니다.

본 블로그에서는 허깅페이스(Hugging Face)를 기반으로 하는 사용자 정의 임베딩 모델을 Amazon SageMaker Endpoint로 배포하는 두 가지 방법을 살펴볼 예정입니다. 첫 번째는 허깅페이스 추론 툴킷 (HuggingFace Inference Toolkit)을 이용하는 방법이고, 두 번째는 커스텀 이미지를 컨테이너라이즈하여 모델을 활용하는 방법입니다.

샘플코드를 통해, 이 두 방법의 활용 및 내부 구조를 알아보고, 각 방법의 장단점을 바탕으로 어떤 상황에 활용할 수 있을지 살펴보겠습니다. 특히 커스텀 이미지를 통한 방법은 deploying-embedding-model-on-sagemaker-endpoint에 공유된 코드를 참고하시어, 본 블로그의 내용을 보다 잘 이해할 수 있습니다.

공개 임베딩 모델 선택

RAG의 높은 성능을 위해서는 좋은 임베딩 모델이 필요합니다. 허깅페이스 리더보드에서는 공개된 여러 임베딩 모델간의 성능비교를 MTEB: Massive Text Embedding Benchmark 기반으로 하여, 리더보드를 제공하고 있습니다.

참고 : MTEB: Massive Text Embedding Benchmark, EACL 2023

MTEB에서는 각 평가항목의 데이터셋을 기반, 목적에 맞도록 각각 평가된 점수를 종합해 평가하는 방법을 제안합니다. 허깅페이스 리더보드에는 MTEB의 각 점수가 출력되며, 이를 기반하여 목적에 맞는 모델을 선택할 수 있습니다.

한국어 공개 모델

허깅페이스 플랫폼은 다양한 언어를 지원하는 모델들을 제공합니다. 이 중에는 한국어에 특화된 모델과 여러 언어를 지원하는 multilingual 모델들이 포함됩니다. 예를 들어, BM-K/KoSimCSE-roberta는 한국어 데이터에 특화된 모델로, 한국어 텍스트 처리에 최적화된 성능을 제공합니다. 반면, multilingual 모델들은 여러 언어를 지원하며, 이들 중 일부는 한국어를 포함한 다양한 언어에 대한 처리 능력을 가지고 있습니다.

목적에 맞게 모델을 선택하는 것이 중요하며, 이를 위해 MTEB 혹은 자체적인 성능 지표를 선정하여 문제 해결에 가장 적합한 모델을 고르는 것이 필요합니다.
본 블로그에서는 한국어 및 multilingual의 MTEB 지표가 비교적 높은 e5-large-multilingual 모델로 배포하는 방법을 알아보겠습니다.

Amazon SageMaker Endpoint를 활용하여 임베딩 모델 배포하기

Amazon SageMaker Endpoint 활용

SageMaker Endpoint 구동 환경

Amazon SageMaker Endpoint는 스크립트의 실행, 학습 및 모델 배포과정에서 Docker를 기반한 구조에서 동작됩니다. 사용자는 SageMaker 콘솔, AWS CLI, Python 노트북, 또는 Amazon SageMaker Python SDK를 통해 내장 알고리즘과 프레임워크를 직접 활용할 수 있으며, 필요에 따라 사전 구축된 Docker 컨테이너 이미지를 사용할 수 있습니다.

해당 SageMaker Endpoint의 Docker 구동 환경의 시나리오는 크게 3가지로 나눠볼 수 있습니다.

사전 구축된 SageMaker 컨테이너 이미지 사용 : Amazon SageMaker의 사전 구축된 Docker 컨테이너를 사용하는 것은 매우 편리합니다. 이 컨테이너들은 Apache MXNet, TensorFlow, PyTorch, Chainer 등 주요 머신 러닝 프레임워크를 포함하고 있으며, SageMaker 콘솔, AWS CLI, Python 노트북, 또는 Amazon SageMaker Python SDK를 통해 쉽게 사용할 수 있습니다. 이러한 접근 방식은 머신 러닝 개발자들에게 컨테이너 관리에 대한 부담을 줄여주며, 빠르게 프로젝트를 시작할 수 있게 해줍니다.
사전 구축된 SageMaker 컨테이너 이미지 확장: 사전 구축된 SageMaker Docker 컨테이너는 확장이 가능합니다. 예를 들어, PyTorch 컨테이너를 사용자의 특정 필요에 맞게 조정할 수 있습니다. 이는 기본 컨테이너에 없는 추가 기능이나 라이브러리를 포함시키고 싶을 때 매우 유용합니다. 이 방식은 사용자가 보다 세밀한 설정을 할 수 있게 해주며, 더 복잡한 프로젝트 요구 사항을 충족시킬 수 있습니다.
커스텀 이미지 활용 : 가장 많은 유연성을 원하는 사용자들은 자체 Docker 컨테이너를 구축할 수 있습니다. 이 방법을 통해 사용자는 SageMaker와 호환되는 완전히 맞춤화된 모델을 활용할 수 있습니다. 커스텀 모델을 활용하거나 특정한 문제를 해결하는 등에 적합합니다. 하지만 이 접근법은 컨테이너 관리와 구성에 대한 추가적인 이해와 노력을 필요로 합니다.

해당 내용의 자세한 설명은 Use Docker containers to build models 를 참고할 수 있습니다.

이 블로그에서는 커스텀 임베딩 모델을 활용하기 위해, 사전 구축된 SageMaker HuggingFace Inference Toolkit의 사전 구축 이미지 활용 방안 및 커스텀 이미지를 직접 생성하고 배포하는 방안에 대해 알아보겠습니다.

HuggingFace Inference Toolkit을 사용하여 커스텀 코드와 모듈 배포하기

SageMaker HuggingFace Inference Toolkit은 SageMaker에서 트랜스포머 모델을 제공하기 위한 오픈 소스 라이브러리로, 특정 트랜스포머 및 디퓨저 모델 Task에 대해 아래와 같은 기능을 제공합니다.

전처리(Pre-processing)
예측(Predict)
후처리(Post-processing)

이 라이브러리는 추론 요청을 처리하는 모델 서버를 시작하기 위해 SageMaker HuggingFace Inference Toolkit을 활용합니다. 아래 글에서는 다양한 모델 중 하나인 multilingual-e5-large을 사용하여 SageMaker Endpoint를 배포하는 과정을 소개합니다.

SageMaker HuggingFace Inference Toolkit을 사용하여 커스텀 모델 정의하기

HF 추론 툴킷을 사용하면 HuggingFaceHandlerService의 기본 메서드를 재정의할 수 있습니다. 이 과정은 다음과 같이 요약할 수 있습니다.

레포지토리에서 아티팩트 구성 패키지를 내려받습니다.
패키지 내부의 code/ 하위에 inference.py 파일을 작성합니다.
model.tar.gz 파일로 압축합니다.
생성된 모델 아티팩트를 S3에 업로드합니다.
모델 아티팩트를 AWS Sagemaker에서 로드하여 엔드포인트를 생성합니다.

inference.py의 구조

inference.py 파일에는 사용자 정의 추론 모듈을 정의해야 하며, 이를 위해 필요한 추가 종속성은 requirements.txt 파일에 추가합니다. 모델 아티팩트 내 경로는 다음을 참고합니다.

model.tar.gz/
|- pytorch_model.bin
|- tf_model.h5 
|- tokenizer.json 
|- tokenizer_config.json
|- ....
|- code/
  |- inference.py
  |- requirements.txt

사용자 정의 모듈은 다음 메서드를 오버라이드할 수 있습니다.

model_fn(model_dir)

모델을 로드하는 기본 메서드를 재정의합니다. 반환 값 model은 예측을 위해 predict에서 사용됩니다. predict는 인자로 압축 해제된 model.tar.gz의 경로인 model_dir을 받습니다.

transform_fn(model, data, content_type, accept_type)

사용자 정의 구현으로 기본 transform 함수를 재정의합니다. transform_fn에서 고유한 전처리, 예측 및 후처리 단계를 구현해야 합니다. 이 메서드는 아래에 언급된 input_fn, predict_fn 또는 output_fn과 함께 사용할 수 없습니다.

input_fn(input_data, content_type)

전처리를 위한 기본 메서드를 재정의합니다. 반환 값 데이터는 예측을 위해 예측에 사용됩니다. 입력은 다음과 같습니다:

input_data는 요청의 원시 본문입니다.
content_type은 요청 헤더의 콘텐츠 유형입니다.

predict_fn(processed_data, model)

예측에 대한 기본 메서드를 재정의합니다. 반환 값 예측은 후처리에서 사용됩니다. 입력은 전처리 결과인 processed_data입니다.

output_fn(prediction, accept)

후처리에 대한 기본 메서드를 재정의합니다. 반환 값은 요청의 응답(예: JSON)이 됩니다. 입력은 다음과 같습니다:

predict는 예측 결과입니다.
accept는 HTTP 요청의 반환 허용 유형(예: application/json)입니다.

from sagemaker_huggingface_inference_toolkit import decoder_encoder

def model_fn(model_dir):
    loaded_model = ...
    return loaded_model 

def input_fn(input_data, content_type):
    data = decoder_encoder.decode(input_data, content_type)
    return data

def predict_fn(data, model):
    ... 
    outputs = model(data , ... )
    return predictions

def output_fn(prediction, accept):
    ...
    response = decoder_encoder.encode(prediction, accept)
    return response

Inference.py 예시

sentence-transformers/multi-qa-MiniLM-L6-cos-v1 모델을 사용하는 허깅페이스 트랜스포머(HuggingFace Transformers)의 예시는 아래와 같습니다.

from transformers import AutoTokenizer, AutoModel
import torchimport torch.nn.functional as F

EMBEDDING_MODEL_NAME = "sentence-transformers/multi-qa-MiniLM-L6-cos-v1"

#Mean Pooling - Take attention mask into account for correct averaging
def mean_pooling(model_output, attention_mask):
    token_embeddings = model_output[0] #First element of model_output contains all token embeddings
    input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
    return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)
    
def model_fn(model_dir): 
    # Load model from HuggingFace Hub
    tokenizer = AutoTokenizer.from_pretrained(EMBEDDING_MODEL_NAME)
    model = AutoModel.from_pretrained(EMBEDDING_MODEL_NAME)

    return model, tokenizer 

def predict_fn(data, model_and_tokenizer): 
    # destruct model and tokenizer
    model, tokenizer = model_and_tokenizer
    
    # Tokenize sentences
    sentences = data.pop("inputs", data)
    encoded_input = tokenizer(sentences, padding=True, truncation=True, return_tensors='pt')

    # Compute token embeddings
    with torch.no_grad():
        model_output = model(**encoded_input)

    # Perform pooling
    sentence_embeddings = mean_pooling(model_output, encoded_input['attention_mask'])

    # Normalize embeddings
    sentence_embeddings = F.normalize(sentence_embeddings, p=2, dim=1)
    
    # return dictonary, which will be json serializable
    return {"vectors": sentence_embeddings[0].tolist()}

intfloat/multilingual-e5-large 모델의 경우는 아래와 같이 구현합니다.

import torch.nn.functional as F

from torch import Tensorfrom transformers 
import AutoTokenizer, AutoModel

EMBEDDING_MODEL_NAME = "intfloat/multilingual-e5-large"

def average_pool(last_hidden_states: Tensor, attention_mask: Tensor) -> Tensor:
    last_hidden = last_hidden_states.masked_fill(~attention_mask[..., None].bool(), 0.0)
    return last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]

def model_fn(model_dir): 
    # Load model from HuggingFace Hub
    tokenizer = AutoTokenizer.from_pretrained(EMBEDDING_MODEL_NAME)
    model = AutoModel.from_pretrained(EMBEDDING_MODEL_NAME)

    return model, tokenizer 

def predict_fn(data, model_and_tokenizer): 
    tokenizer = AutoTokenizer.from_pretrained(EMBEDDING_MODEL_NAME)
    model = AutoModel.from_pretrained(EMBEDDING_MODEL_NAME)
    
    # Tokenize the input texts
    batch_dict = tokenizer(input_texts, max_length=512, padding=True, truncation=True, return_tensors='pt')

    outputs = model(**batch_dict)
    sentence_embeddings = average_pool(outputs.last_hidden_state, batch_dict['attention_mask'])
    sentence_embeddings = F.normalize(embeddings, p=2, dim=1)

    # return dictonary, which will be json serializable
    return {"vectors": sentence_embeddings[0].tolist()}

배포 과정

이어지는 내용에서는 Sagemaker Jupyter Notebook을 기준으로 위 모델 아티팩트를 배포하는 과정을 설명합니다.

a. code 폴더를 만들고, 사용자 정의 추론을 포함하는 inference.py 파일을 작성합니다.

!mkdir code

%%writefile code/inference.py

from transformers import AutoTokenizer, AutoModel
import torch
import torch.nn.functional as F

def model_fn(model_dir):
  ...(생략)...

def predict_fn(data, model_and_tokenizer):
    ...(생략)...

b. 원하는 모델에 해당하는 추론 툴킷 패키지를 내려받습니다.

!conda install -c conda-forge git-lfs -y
!git lfs install

repository = 'intfloat/multilingual-e5-large'
model_id = repository.split("/")[-1]
s3_bucket = '<임의의 S3 버킷 경로>'
s3_location = f's3://{s3_bucket}/custom_inference/{model_id}/model.tar.gz'

!git clone https://huggingface.co/$repository

c. 앞서 작성한 inference.py 파일을 패키지 내부에 추가합니다.

!cp -r code/ $model_id/code/

d. 패키지를 압축하고, S3 버킷에 업로드합니다.

%cd $model_id
!tar czvf model.tar.gz *
!aws s3 cp model.tar.gz $s3_location

e. 업로드된 모델 아티팩트를 AWS Sagemaker에서 로드하여 엔드포인트로 배포하는 과정은 아래와 같습니다.

from sagemaker import get_execution_role 
from sagemaker.huggingface.model import HuggingFaceModel
from datetime import datetime

role = get_execution_role()
endpoint_name = f'sentence-transform-{datetime.utcnow():%Y-%m-%d-%H%M}'

# create Hugging Face Model Class
huggingface_model = HuggingFaceModel(
   model_data=s3_location,       # path to your model and script
   role=role,                    # iam role with permissions to create an Endpoint
   transformers_version="4.12",  # transformers version used
   pytorch_version="1.9",        # pytorch version used
   py_version='py39',            # python version used
)

# deploying endpoint 
predictor = huggingface_model.deploy(
    initial_instance_count=1,
    instance_type="ml.g4dn.xlarge", 
    endpoint_name=endpoint_name
)

# request inference 
data = {
  "inputs": "Hello, World",
}

res = predictor.predict(data=data)
print(res)

커스텀 컨테이너 이미지 만들어 배포하기

커스텀 코드를 Amazon SageMaker로 가져오기 위해 굳이 컨테이너를 만들 필요가 없을 수도 있습니다. SageMaker에서 직접 지원하는 프레임워크(예: Apache MXNet 또는 TensorFlow)를 사용하는 경우, 해당 프레임워크의 SDK 엔트리포인트를 사용하여 알고리즘을 구현하는 Python 코드를 서빙하기만 하면 됩니다. 이들 프레임워크는 지속적으로 업데이트되고 있으므로, 바퀴의 재발명을 피하기 위해서는 커스텀 컨테이너를 만들기 전에 먼저 해당 프레임워크를 확인하는 것이 좋습니다.

사용 중인 환경이나 프레임워크에 대한 직접적인 SDK 지원이 있더라도 자체 컨테이너를 구축하는 것이 더 효과적일 수 있습니다. 알고리즘을 구현하는 코드가 그 자체로 상당히 복잡하거나 프레임워크에 특별한 추가 사항이 필요한 경우가 그에 해당합니다. 예를 들어, 이 글에서 설명하고자 하는 ‘SentenceTransformer를 FastSentenceTransformer로 교체하는 처리‘ 가 대표적입니다.

본 섹션은 deploying-embedding-model-on-sagemaker-endpoint 의 코드를 참고할 수 있습니다.

커스텀 컨테이너 이미지를 SageMaker Endpoint로 배포하기

커스텀 컨테이너 이미지 빌드와 SageMaker Endpoint 배포 기본 과정의 이해를 돕기 위해, 간단한 더미 컨테이너 이미지를 만들고 배포해 보도록 하겠습니다. 이번 실습에서 다룰 내용은 다음과 같습니다.
컨테이너 이미지는 실제 학습 모델을 구현하는 대신, 추론 요청에 대해 임의의 정적 값을 반환하도록 구현합니다.
SageMaker에서 권장하는 방식인 nginx → gunicorn → wsgi → flask 구조가 아닌, Flask 단독으로 추론 모델을 서빙합니다.
이 구성은 실제 서비스에 적합하지는 않지만, SageMaker가 컨테이너 이미지를 다루는 메커니즘을 이해하는 데에 도움을 줍니다. 실제 서비스에 적절한 구성을 실습하고자 한다면 이어지는 실습 2를 참고하기 바랍니다.

SageMaker Endpoint로 배포하는 컨테이너 이미지는 아래와 같은 두 개의 RESTful API를 제공해야 합니다.

GET /ping
POST /invocations

실습 과정은 가장 간단한 파이썬 웹 프레임워크인 Flask를 이용하여 위의 두 API를 구현하고, 이를 컨테이너 이미지로 빌드하기만 하면 됩니다. 이를 위해, 아래와 같은 구조의 프로젝트를 만들겠습니다.

<PROJECT_ROOT>
├── Dockerfile
├── build_and_push.sh
└── src
    └── predictor.py

STEP-1. 추론 코드 및 컨테이너 빌드 관련 스크립트 작성하기

먼저 src/predictor.py를 살펴봅시다.

# src/predictor.py 실제 추론 모델이 구현되는 곳이다. 

import flask

app = flask.Flask(__name__)

# /ping은 SageMaker Endpoint의 인스턴스 상태를 확인하는 데 사용된다.
@app.route("/ping", methods=["GET"])
def ping():
    response_body = "OK"
    status = 200

    return flask.Response(response=f"{response_body}\n", 
                          status=status, 
                          mimetype="application/json")

# /invocations는 SageMaker Endpoint를 통해 실제 추론을 행하는 API이다. 
# sagemaker.predictor.predict()를 호출하면
# 내부적으로 /invocations가 호출된다.
@app.route("/invocations", methods=["POST"])
def transformation():
    return flask.jsonify([0.00121, -0.0023401]) # 더미 추론 모델을 구현


if __name__ == '__main__':
    app.run('0.0.0.0', port=8080, debug=True)

다음으로, Dockerfile입니다.

# Dockerfile 

FROM ubuntu:18.04

MAINTAINER AWS ProServe GAIA <hijaeeun@amazon.com>

RUN apt-get -y update && apt-get install -y --no-install-recommends \
         wget \
         python3-pip \
         python3-setuptools \
    && rm -rf /var/lib/apt/lists/*

RUN ln -s /usr/bin/python3 /usr/bin/python
RUN ln -s /usr/bin/pip3 /usr/bin/pip

RUN pip --no-cache-dir install flask

ENV PYTHONUNBUFFERED=TRUE
ENV PYTHONDONTWRITEBYTECODE=TRUE
ENV PATH="/opt/program:${PATH}"

# Set up the program in the image
COPY src /opt/program
WORKDIR /opt/program

ENTRYPOINT ["python", "predictor.py"]

마지막으로, 컨테이너 이미지를 빌드하고 ECR에 배포하는 build_and_push.sh 스크립트입니다.

#!/usr/bin/env bash
IMAGE=$1

if [ "$IMAGE" == "" ]
then
    echo "Usage: $0 <image-name>"
    exit 1
fi

chmod +x src/predictor.py

# Get the account number associated with the current IAM credentials
ACCOUNT=$(aws sts get-caller-identity --query Account --output text)

if [ $? -ne 0 ]
then
    exit 255
fi

# Get the region defined in the current configuration (default to us-west-2 if none defined)
REGION=$(aws configure get region)
REGION=${REGION:-us-west-2}


FULLNAME="${ACCOUNT}.dkr.ecr.${REGION}.amazonaws.com/${IMAGE}:latest"

# If the repository doesn't exist in ECR, create it.
aws ecr describe-repositories --repository-names "${IMAGE}" > /dev/null 2>&1

if [ $? -ne 0 ]
then
    aws ecr create-repository --repository-name "${IMAGE}" > /dev/null
fi

# Get the login command from ECR and execute it directly
aws ecr get-login-password --region "${REGION}" | docker login --username AWS --password-stdin "${ACCOUNT}".dkr.ecr."${REGION}".amazonaws.com

# Build the docker image locally with the image name and then push it to ECR
# with the full name.

docker build  -t ${IMAGE} .
docker tag ${IMAGE} ${FULLNAME}

docker push ${FULLNAME}

복잡해 보이지만, 조건절을 제외하면 핵심 로직은 아래와 같은 몇 라인에 불과합니다.

aws ecr get-login-password --region "${REGION}" | docker login --username AWS --password-stdin "${ACCOUNT}".dkr.ecr."${REGION}".amazonaws.com
docker build  -t ${IMAGE} .
docker tag ${IMAGE} ${FULLNAME}
docker push ${FULLNAME}

ECR에 로그인하고, 컨테이너 이미지를 빌드한 후 태깅하고, 이를 ECR 레포지토리에 푸시하는 과정입니다. ECR 이미지 레포지토리가 없으면 새로 만들어주는 코드까지 포함하고 있지만, 레포지토리를 만드는 데에 시간이 소요되므로 그 시간동안 위 스크립트가 반복적으로 실패할 수 있습니다. 가급적 ECR 레포지토리는 미리 만들어 두길 권장합니다.

STEP-2. 컨테이너 이미지 빌드 후 ECR 레포지토리에 푸시하기

이제 실제로 빌드하여 ECR 레포지토리에 푸시해 봅시다. build_and_push.sh 스크립트를 아래와 같이 실행합니다.

$ chmod +x build_and_push.sh 
$ ./build_and_push.sh <ECR 레포지토리에 푸시될 컨테이너 이미지 명>

예제에서는 “dummy-inference-model”이라는 이름을 사용했습니다. 다른 이름을 사용하더라도 실행 결과는 아래와 비슷합니다.

$ ./build_and_push.sh dummy-inference-model  

Login Succeeded
[+] Building 1.0s (12/12) FINISHED                                               
 => [internal] load .dockerignore                                                
 => => transferring context: 2B                                                   
 => [internal] load build definition from Dockerfile                              
 => => transferring dockerfile: 1.16kB                                             
 ...(중략)                                                                          
 => CACHED [5/7] RUN pip --no-cache-dir install flask                               
 => CACHED [6/7] COPY src /opt/program                                               
 => CACHED [7/7] WORKDIR /opt/program                                                 
 => exporting to image                                                                 
 => => exporting layers                                                                 
 => => writing image sha256:6b1cb56f0338c3f13524201fc9d63c1efda574b29353e            
 => => naming to docker.io/library/dummy-inference-model                              

The push refers to repository [98037000.dkr.ecr.ap-northeast-2.amazonaws.com/dummy-inference-model] 
5f70bf18a086: Pushed 
bcdcf07fd681: Pushed 
64eca9b22945: Pushed 
1d9757ae478f: Pushed 
31058fce5202: Pushed 
3915e331ca51: Pushed 
548a79621a42: Pushed 
latest: digest: sha256:3f8cf03020a6698fd1e75b590d178e85afa4feb559b4ea3f size: 1779

STEP-3. Sagemaker Endpoint에 배포하기

해당 파이썬 코드를 커스텀 컨테이너 이미지를 Sagemaker Endpoint로 배포하는 코드는 다음과 같습니다 (참고로 이 단계부터는 Sagemaker Notebook에서 실행할 것을 권장합니다).

import boto3
client = boto3.client('sagemaker')

# Create Sagemaker Model
create_model_response = client.create_model(
    ModelName='sm-dummy-inference-model', 
    PrimaryContainer={
        'Image': '9803771.dkr.ecr.ap-northeast-2.amazonaws.com/dummy-inference-model', 
        'Environment': {}
    }, 
    ExecutionRoleArn='arn:aws:iam::980377:role/service-role/AmazonSageMaker-ExecutionRole-20231127T170565'
)

print("create model response:", create_model_response)

# Create Sagemaker Endpoint Config 
create_endpoint_config_response = client.create_endpoint_config(
    EndpointConfigName='dummy-inference-serve-config', 
    ProductionVariants=[
        {
            'ModelName': 'sm-dummy-inference-model',
            'VariantName': 'variant-sm-dummy-inference-model-1',
            'InitialInstanceCount': 1,
            'InstanceType': 'ml.m4.xlarge'
        },
    ]
)

print('create endpoint config response:', create_endpoint_config_response)

# Create Sagemaker Endpoint 
create_endpoint_response = client.create_endpoint(
    EndpointName='dummy-inference-serve', 
    EndpointConfigName='dummy-inference-serve-config'
)

print('create endpoint response: ', create_endpoint_response)

위 코드는 아래와 같이 3단계로 나눌 수 있습니다.

Sagemaker Model을 생성한다.
Sagemaker Endpoint Config를 생성한다.
Sagemaker Endpoint를 생성한다.

여기서 Sagemaker Model은 허깅페이스나 BERT 등에서 제공하는 모델이 아닌, ECR에 등록된 컨테이너 이미지를 베이스로 하는 Sagemaker의 모델을 의미합니다.

각 단계별 예시 코드를 살펴봅시다.

Sagemaker Model 생성

# Create Sagemaker Model
create_model_response = client.create_model(
    ModelName='sm-dummy-inference-model', 
    PrimaryContainer={
        'Image': '9803771.dkr.ecr.ap-northeast-2.amazonaws.com/dummy-inference-model', 
        'Environment': {}
    }, 
    ExecutionRoleArn='arn:aws:iam::980377:role/service-role/AmazonSageMaker-ExecutionRole-20231127T170565'
)

ModelName
임의의 이름을 입력. 이 모델명은 이후 Endpoint Config 생성 시 참조
Image
ECR 레포지토리 URI를 입력
ExecutionRoleArn
AmazonSagemakerFullAccess를 포함하는 IAM Role Arn. 다른 서비스가 Sagemaker를 실행할 수 있도록 권한을 부여

Sagemaker Endpoint Config 생성

# Create Sagemaker Endpoint Config 
create_endpoint_config_response = client.create_endpoint_config(
    EndpointConfigName='dummy-inference-serve-config', 
    ProductionVariants=[
        {
            'ModelName': 'sm-dummy-inference-model',
            'VariantName': 'variant-sm-dummy-inference-model-1',
            'InitialInstanceCount': 1,
            'InstanceType': 'ml.m4.xlarge'
        },
    ]
)

EndpointConfigName
임의의 엔드포인트 컨피그 이름을 입력한다. 이 컨피그 명은 이후 Endpoint 생성시 참조
ModelName
앞서 생성한 Sagemaker 모델명을 그대로 입력
VariantName
변형할 이름을 임의로 입력

Sagemaker Endpoint 생성

# Create Sagemaker Endpoint 
create_endpoint_response = client.create_endpoint(
    EndpointName='dummy-inference-serve', 
    EndpointConfigName='dummy-inference-serve-config'
)

EndpointName
임의의 엔드포인트 이름을 입력. 이 엔드포인트 명은 이후 엔드포인트를 호출할 모든 클라이언트가 참조
EndpointConfigName
앞서 정의한 엔드포인트 컨피그 이름을 그대로 입력.

STEP-4. 배포된 엔드포인트 호출하기

SageMaker Endpoint로 배포된 임베딩 모델 활용하기 부분을 참고.

실제 허깅페이스 모델 기반 커스텀 이미지 빌드 및 SageMaker Endpoint로 배포하기

이번에는 조금 더 복잡하지만, 실제 배포 사양에 적합한 구조를 다루는 실습을 진행해보겠습니다.

이 실습에서는 Sagemaker에서 권장하는 방식인 nginx – gunicorn – wsgi – flask 형태를 따라 컨테이너를 구성하며, 내부적으로는 허깅페이스 모델을 이용하여 주어진 입력값을 임베딩하는 모델을 호스팅하는 것을 목표로 합니다.

프로젝트 전체 구성은 다음과 같습니다.

<PROJECT ROOT> 
├── Dockerfile
├── Makefile
├── build_and_push.sh
├── deploy_container_image.ipynb
└── src
    ├── nginx.conf
    ├── predictor.py
    ├── train
    ├── serve
    └── wsgi.py

실습 진행 또는 다른 모델을 배포하고자 할 때, predictor.py만 수정해서 사용하면 됩니다.

[Under the hood] Amazon SageMaker가 Docker 컨테이너를 실행하는 방법

동일한 이미지를 트레이닝과 호스팅 양쪽에서 모두 사용할 수 있도록, SageMaker는 train 또는 serve 파라미터를 사용하여 컨테이너를 실행합니다. 컨테이너가 이 인수를 처리하는 방법은 컨테이너에 따라 다릅니다.

Dockerfile에 ENTRYPOINT 항목을 지정하지 않았을 경우

Dockerfile에 ENTRYPOINT 항목을 지정하지 않으면 도커는 아래와 같이 동작합니다.

트레이닝 시점에는 train을 실행함
서빙 시점에는 serve를 실행함

따라서, 이 경우 컨테이너 이미지 내부의 WORKDIR 경로에는 train 및 serve 프로그램이 정의되어 있어야 합니다. 아래는 ENTRYPOINT를 지정하지 않은 경우에 필요한 컨테이너 이미지 내부 구조의 예시입니다.

트레이닝과 호스팅을 위한 컨테이너를 분리하거나, 또는 둘 중 하나의 목적으로만 컨테이너를 빌드해야 하는 경우가 있습니다. BERT 기반 허깅페이스 모델과 같이 사전 훈련된 모델을 사용하여 추론용 호스팅만 배포하는 것이 대표적인 예입니다. 이 때에는 train 및 serve 둘 다 있을 필요는 없습니다. 목적에 맞는 프로그램만 선택적으로 위치하면 됩니다.

Dockerfile에 ENTRYPOINT 항목을 지정했을 경우

Dockerfile에 ENTRYPOINT를 지정하면 컨테이너 시작시 해당 ENTRYPOINT가 실행되고, 그 첫 번째 파라미터로 train 또는 serve 가 전달됩니다. 프로세스는 해당 파라미터를 보고 어떤 작업을 수행해야 할지 결정할 수 있습니다. 수행할 작업을 결정할 수 있습니다.

앞에서와 마찬가지로, 트레이닝과 호스팅을 위한 컨테이너를 분리하거나 혹은 둘 중 하나의 목적으로만 컨테이너를 빌드하는 경우가 있을 수 있습니다. 이 때에도 ENTRYPOINT가 지정되어 있다면 train 또는 serve 가 전달되지만, 사용하지 않아도 됩니다. 다만 안전한 동작 보장을 위해서는 해당 파라미터를 무시하기보다는 올바른 파라미터가 전달되었는지 체크하는 로직을 포함시키는 것이 좋습니다.

[Under the hood] 추론 코드를 커스텀 컨테이너 이미지로 빌드하기

Dockerize를 위한 패키지 구조

다음은 패키지 구성의 예시입니다. Dockerfile에서 ENTRYPOINT를 지정하지 않았으므로, SageMaker는 자동으로 트레이닝 시점에는 train 을, 호스팅 시점에는 serve 를 실행합니다. 따라서 SageMaker가 호출할 수 있는 위치에 해당 프로그램이 존재해야 합니다.

serve 프로그램의 경우, 호스팅을 위한 서버를 가동하도록 처리되어야 합니다. nginx – gunicorn – wsgi로 연결되는 서버 구성 예시는 아래와 같습니다.

def start_server():
    print('Starting the inference server with {} workers.'.format(model_server_workers))


    # link the log streams to stdout/err so they will be logged to the container logs
    subprocess.check_call(['ln', '-sf', '/dev/stdout', '/var/log/nginx/access.log'])
    subprocess.check_call(['ln', '-sf', '/dev/stderr', '/var/log/nginx/error.log'])

    nginx = subprocess.Popen(['nginx', '-c', '/opt/program/nginx.conf'])
    gunicorn = subprocess.Popen(['gunicorn',
                                 '--timeout', str(model_server_timeout),
                                 '-k', 'sync',
                                 '-b', 'unix:/tmp/gunicorn.sock',
                                 '-w', str(model_server_workers),
                                 'wsgi:app'])

    signal.signal(signal.SIGTERM, lambda a, b: sigterm_handler(nginx.pid, gunicorn.pid))

    # If either subprocess exits, so do we.
    pids = set([nginx.pid, gunicorn.pid])
    while True:
        pid, _ = os.wait()
        if pid in pids:
            break

    sigterm_handler(nginx.pid, gunicorn.pid)
    print('Inference server exiting')
    
# The main routine just invokes the start function.
if __name__ == '__main__':
    start_server()

실제 서빙하기 위한 predictor.py 파일의 예시는 아래와 같습니다.

# This is the file that implements a flask server to do inferences. It's the file that you will modify to
# implement the scoring for your own algorithm.

from fast_sentence_transformer
import json
import os

import flask

prefix = "/opt/ml/"

class EmbeddingService(object):
    model = None  # Where we keep the model when it's loaded
    EMBEDDING_MODEL_NAME = "intfloat/multilingual-e5-large"
    
    @classmethod
    def get_model(cls):
        """Get the model object for this instance, loading it if it's not already loaded."""
        if cls.model == None:
            cls.model = FastSentenceTransformer(EMBEDDING_MODEL_NAME)
        return cls.model

    @classmethod
    def predict(cls, input):
        """For the input, do the predictions and return them.
        Args:
            input (a pandas dataframe): The data on which to do the predictions. There will be
                one prediction per row in the dataframe"""
        clf = cls.get_model()
        return clf.predict(input)

# The flask app for serving predictions
app = flask.Flask(__name__)

@app.route("/ping", methods=["GET"])
def ping():
    """Determine if the container is working and healthy. In this sample container, we declare
    it healthy if we can load the model successfully."""
    health = EmbeddingService.get_model() is not None  # You can insert a health check here

    status = 200 if health else 404
    return flask.Response(response="\n", status=status, mimetype="application/json")


@app.route("/invocations", methods=["POST"])
def transformation():
    """Do an inference on a single batch of data. In this sample server, we take data as CSV, convert
    it to a pandas data frame for internal use and then convert the predictions back to CSV (which really
    just means one prediction per line, since there's a single column.
    """
    data = flask.request.data.decode("utf-8")
    result = EmbeddingService.predict(data)

    return flask.Response(response=result, status=200, mimetype="application/json")

SageMaker Endpoint로 배포된 임베딩 모델 활용하기

SageMaker Endpoint를 호출하는 방식은 크게 두가지 방법이 존재합니다. Amazon API Gateway를 기반으로 REST API 형태로 호출하는 방법, 그리고 AWS SDK 라이브러리인 boto3를 사용하여 invoke_endpoint 함수를 사용하는 방법이 있습니다. 이 글 에서는 boto3 라이브러리를 활용하여 invoke_endpoint 기반으로 호출하여, 임베딩값을 구해오는 예시를 설명합니다.

사전 준비 사항

boto3 라이브러리 설치

sagemaker_endpoint_caller.py 예시

boto3를 사용하여 SageMaker Endpoint를 호출하는 Python 코드는 아래와 같이 작성 할 수 있습니다.

import json
import boto3

class SageMakerEndpointCaller:
    def __init__(self, region='us-east-1'):
        self.sagemaker_runtime = boto3.client(
            service_name='sagemaker-runtime',
            region_name=region
        )

    def get_query_embeding(self, query, endpoint_name):
        """
        SageMaker Endpoint를 호출하여 쿼리에 대한 임베딩을 가져옵니다.

        :param query: 임베딩을 생성할 쿼리
        :param endpoint_name: SageMaker Endpoint의 이름
        :return: 임베딩 결과 또는 빈 리스트
        """
        payload = json.dumps({"inputs": query})
        response = self.sagemaker_runtime.invoke_endpoint(
            EndpointName=endpoint_name,
            ContentType='application/json', 
            Body=payload
        )
 
        return self._parse_response(response)
        
    def _parse_response(self, response):
        """
        SageMaker Endpoint의 응답을 파싱합니다.

        :param response: SageMaker Endpoint로부터의 응답
        :return: 임베딩 결과 또는 빈 리스트
        """
        try:
            if response['ResponseMetadata']['HTTPStatusCode'] == 200:
                body = response['Body'].read().decode('utf-8')
                data = json.loads(body)
                return data.get('vectors', [])
        except KeyError:
            pass

        return []

위의 코드를 사용하면 Vector화된 Query를 가져올 수 있습니다. 해당 List를 사용해 Amazon OpenSearch Service의 Vector Database를 검색 할 수 있습니다. 관련된 자세한 내용은 ‘Amazon OpenSearch Service Hybrid Query를 통한 검색 기능 강화’ 블로그에서 설명합니다.

결론

생성형 AI를 활용하기 위해 RAG를 기반한 접근은 필수적인 상황이 되어가고 있으며, 이를 위한 텍스트의 임베딩생성은 필수요소가 되었습니다. AWS위에서 임베딩을 생성하기 위해서는 SageMaker Endpoint를 비교적 쉽게 활용할 수 있으며, 허깅페이스와 오픈소스등을 통한 활용법이 가장 빠르게 잘 구성된 쉽고 다양한 접근을 하는 방법입니다.

본 블로그에서는 허깅페이스의 인퍼런스 툴킷을 활용하는 방법, 그리고 도커라이즈된 커스텀 이미지로 배포하는 방법을 자세히 알아봤으며, 이 두 예시는 모두 SageMaker Endpoint를 기반으로 배포하며 Python을 활용한 사용법 또한 함께 제공하고 있습니다.

위의 방법을 통해, 목적에 맞는 임베딩 모델을 검토하고 서비스에 빠른 적용을 하여 주어진 비즈니스 문제를 해결할 수 있는 참고자료가 되기를 바랍니다.

참고링크

deploying-embedding-model-on-sagemaker-endpoint – 커스텀 이미지를 도커라이징하여, endpoint로 활용하기
Transformers documentation – SageMaker에서 Hugging Face Job 실행하기
intfloat/multilingual-e5-large – 허깅페이스의 multilingual-e5-large 모델 소개 및 사용 가이드
sentence-transformers/multi-qa-MiniLM-L6-cos-v1 – Sentence-Transformer와 HuggingFace Transformers 모델 사이의 여러 구현 예시
embedding 모델을 위한 inference.py 예제
https://github.com/aws/sagemaker-inference-toolkit – Sagemaker Inference Toolkit 샘플코드

AWS 기술 블로그