Amazon SageMaker로 LLM 응답 Streaming 서빙하기

Cyu 24.09.11

2,010 8 0

facebook twitter kakao link

8 0

facebook twitter kakao link

DEVOTEE 요약: Amazon SageMaker는 머신러닝 모델 구축, 훈련, 배포를 쉽게 할 수 있도록 지원하며, 모델 응답을 스트리밍하기 위해 Lambda Function URL과 CloudFront를 활용할 수 있습니다. SageMaker로 배포된 모델은 API Gateway 대신 Lambda를 사용해 응답을 토큰 단위로 스트리밍할 수 있으며, 이를 위해 Node.js 기반의 Lambda Function을 작성해야 합니다. 또한, CloudFront와 Lambda Function URL을 연동해 sigV4 인증을 우회하여 모델 응답을 스트리밍할 수 있습니다.

DEVOTEE 추천 블로그

SageMaker Endpoint Serving

Amazon SageMaker는 머신러닝 모델을 더 쉽게 구축, 훈련, 배포할 수 있도록 지원하며 데이터 준비부터 모델 배포까지 일련의 과정을 효율적으로 처리할 수 있게 도와주는 서비스입니다.

일반적으로 Sagemaker로 배포한 모델은 API Gateway와 Lambda를 이용해서 외부에 serverless로 서빙할 수 있는데요, S

ageMaker Endpoint를 서빙하는 자세한 방법은 여기 를 참고하시기 바랍니다.

Lambda Response Streaming

위 방법으로 모델을 서빙하면 모델의 응답이 끝날 때까지 기다렸다가 응답을 돌려주게 됩니다.

LLM을 서빙할 경우 모델의 특성상 응답 생성 시간이 오래 걸릴 수도 있는데, 만약 모델의 응답을 토큰 단위로 streaming 하려면 어떻게 해야할까요?

AWS 에 따르면 현재 API Gateway는 응답 스트리밍 기능을 지원하지 않고 있습니다.

You can not use Amazon API Gateway and Application Load Balancer to progressively stream response payloads, but you can use the functionality to return larger payloads with API Gateway.

따라서 응답을 스트리밍하기 위해 API Gateway 대신 Lambda의 Function URL을 사용해야 하며, 람다함수는 Node.js로 구현해야 합니다.

추가로, Lambda Function URL은 호출 시 sigV4 인증이 필요하며, 이를 우회하려면 CloudFront에 Origin Access Control을 붙여서 사용해야 합니다.

시스템 구성

SageMaker Endpoint + Lambda Function URL + CloudFront

SageMaker 엔드포인트 배포

본 블로그에서는 예시로 S3에 저장된 모델을 배포하는 코드를 작성해보았습니다.

아래 코드는 SageMaker Studio 환경에서 돌아가는 것을 가정합니다. (로컬 python 환경에서 돌리려면 몇 가지 패키지만 더 import 해주시면 됩니다. )

이 외에도 HuggingFace Hub등 다양한 source에서 모델을 가져와 배포할 수 있습니다.

### 필요 파이썬 패키지 설치
!pip install datasets openpyxl transformers sagemaker --quiet
!pip install -U sagemaker --quiet

### Sagemaker 환경 설정
# 세션정보, 업로드용 버킷, role, region 정보 받아오기
import sagemaker
import boto3

sess = sagemaker.Session()

sagemaker_session_bucket=None
if sagemaker_session_bucket is None and sess is not None:
    sagemaker_session_bucket = sess.default_bucket()

try:
    role = sagemaker.get_execution_role()
except ValueError:
    iam = boto3.client('iam')
    role = iam.get_role(RoleName='sagemaker_execution_role')['Role']['Arn']

sess = sagemaker.Session(default_bucket=sagemaker_session_bucket)

print(f"sagemaker role arn: {role}")
print(f"sagemaker bucket: {sess.default_bucket()}")
print(f"sagemaker session region: {sess.boto_region_name}")

### 적합한 LLM Deep Learning Container 검색
# 본 blog에서는 huggingface에서 제공하는 tgi 서빙 컨테이너 활용
# SageMaker SDK에서 제공하는 get_huggingface_llm_image_uri 메소드를 통해 지정된 백엔드, 세션, 지역, 버전에 기반한 원하는 Hugging Face LLM DLC의 URI를 가져올 수 있음
from sagemaker.huggingface import get_huggingface_llm_image_uri

llm_image = get_huggingface_llm_image_uri(
  "huggingface",
  version="2.0.1",
  session=sess,
)

print(f"llm image uri: {llm_image}")

### 배포 설정
# 모델 경로 또는 huggingface ID, AWS 인스턴스 타입 등 sagemaker inference toolkit 환경변수 설정  
# HF_MODEL_ID는 필수
# - Toolkits: https://docs.aws.amazon.com/sagemaker/latest/dg/amazon-sagemaker-toolkits.html  
# - TGI관련설정: https://github.com/aws/sagemaker-huggingface-inference-toolkit?tab=readme-ov-file#%EF%B8%8F-environment-variables
import json
from sagemaker.huggingface import HuggingFaceModel

instance_type = "ml.g5.2xlarge"
number_of_gpu = 1
health_check_timeout = 300
llm_endpoint_name = sagemaker.utils.name_from_base("ALADIN-LLM")

# 엔드포인트 구성 설정 정의
config = {
"HF_MODEL_ID": "/opt/ml/model",
"SM_NUM_GPUS": json.dumps(number_of_gpu),
"MAX_INPUT_LENGTH": json.dumps(<input_len>),
"MAX_TOTAL_TOKENS": json.dumps(<total_len>),
"MESSAGES_API_ENABLED": "true", ### OpenAI Chat Completion 호환 설정
}

model_s3_path = "s3://sagemaker-ap-northeast-2-*******/llama2_ft/model/"

llm_model = HuggingFaceModel(
  role=role,
  image_uri=llm_image,
  model_data={'S3DataSource':{'S3Uri': model_s3_path,'S3DataType': 'S3Prefix','CompressionType': 'None'}},
  env=config
)

### 엔드포인트 배포
from sagemaker import serializers, deserializers

llm = llm_model.deploy(
  endpoint_name=llm_endpoint_name,
  serializer=serializers.JSONSerializer(),
  deserializer=deserializers.JSONDeserializer(),
  initial_instance_count=1,
  instance_type=instance_type,
  container_startup_health_check_timeout=health_check_timeout,
)

Lambda Function 생성

이제 위 엔드포인트를 호출하고 클라이언트로 응답을 전송하는 Lambda Function을 생성해야 합니다.

AWS 콘솔에서 아래 단계를 거쳐 스트리밍 처리가 가능한 Lambda Function를 생성합니다.

Lambda > 함수 > 함수 생성

런타임: Node.js 20.x
고급설정 > 함수URL활성화 > 호출모드: RESPONSE_STREAM

Lambda Function URL 복사

생성된 Lambda Function 함수 개요에서 함수 URL을 Copy 합니다.

Lambda Function 작성

이제 실제로 SageMaker Endpoint를 호출하고 응답을 스트리밍 처리하여 클라이언트로 전송하는 Lambda Function 코드를 작성합니다.

응답을 스트리밍하려면 node.js를 사용해야 하고, 일반적인 handler가 아닌 streamifyResponse() 디코더를 사용해 핸들러를 작성해야 합니다.

import { SageMakerRuntimeClient, InvokeEndpointCommand } from "@aws-sdk/client-sagemaker-runtime";
import { PassThrough } from "stream";
 
const client = new SageMakerRuntimeClient({ region: 'us-east-1' });
 
export const handler = awslambda.streamifyResponse(async (event, responseStream) => {
  try {
    const endpointName = 'ALADIN-LMM-1717995738'; // 배포된 sagemaker endpoint 이름
    const inputPayload = JSON.parse(event.body); // 요청된 입력 데이터를 사용
    const jsonString = JSON.stringify(inputPayload); // 객체를 JSON 문자열로 변환
    console.log(jsonString);
 
    const params = {
      EndpointName: endpointName,
      Body: Buffer.from(jsonString, 'utf-8'), // JSON 문자열을 Buffer로 변환
      ContentType: 'application/json'
    };
 
    // SageMaker Endpoint 호출
    const command = new InvokeEndpointCommand(params);
    const response = await client.send(command);
    console.log(response.Body);
 
    // 응답 스트리밍 처리
    const passThrough = new PassThrough();
    passThrough.end(response.Body);
 
    // 스트림 데이터를 클라이언트로 전송
    passThrough.pipe(responseStream);
 
  } catch (error) {
    console.error('Error invoking SageMaker endpoint:', error);
    responseStream.write(JSON.stringify({ error: 'Error invoking SageMaker endpoint' }));
    responseStream.end();
  }
});

CloudFront OAC 생성하기

다음으로 Lambda Function URL 호출 시 sigV4 인증을 우회하기 위해 CloudFront를 생성하고 이를 위에 생성한 Lambda 함수와 연결합니다.

여기를 참고해서 CloudFront OAC 를 생성하고 배포하세요.

배포가 완료된 CloudFront의 도메인을 복사해서 호출 테스트에 활용합니다.

테스트

호출 시에는 request 헤더에 payload의 sha256 해시를 필수로 포함해야 하며, payload에 "stream":true 옵션을 넣어주어야 응답을 스트리밍할 수 있습니다.

호출 결과를 보면 data["res"]["choices"][0]["delta"]["content"] 로 응답 결과를 토큰 단위로 가져오는 것을 볼 수 있습니다.

호출 명령어 예제

payload='{"model": "tgi", "messages": [{"role": "user","content": "플레이리스트 만들어줘"}], "max_tokens": 300, "temperature": 0.1, "top_p": 0.2, "stream": true}'
payload_hash=$(echo -n $payload | openssl dgst -sha256 | awk '{print $2}')
  
curl -N -X POST \
  -H "Content-Type: application/json" \
  -H "x-amz-content-sha256: $payload_hash" \
  -d "$payload" \
  https://d221*****2nvsr.cloudfront.net

호출 결과

curl: (6) Could not resolve host:  
curl: (6) Could not resolve host:  
curl: (6) Could not resolve host:  
data:{"id":"","object":"text_completion","created":1720501473,"model":"/opt/ml/model","system_fingerprint":"2.0.1-native","choices":[{"index":0,"delta":{"role":"assistant","content":"{\""},"logprobs":null,"finish_reason":null}]}
 
data:{"id":"","object":"text_completion","created":1720501473,"model":"/opt/ml/model","system_fingerprint":"2.0.1-native","choices":[{"index":0,"delta":{"role":"assistant","content":"agents"},"logprobs":null,"finish_reason":null}]}
 
data:{"id":"","object":"text_completion","created":1720501473,"model":"/opt/ml/model","system_fingerprint":"2.0.1-native","choices":[{"index":0,"delta":{"role":"assistant","content":"\":"},"logprobs":null,"finish_reason":null}]}
 
data:{"id":"","object":"text_completion","created":1720501473,"model":"/opt/ml/model","system_fingerprint":"2.0.1-native","choices":[{"index":0,"delta":{"role":"assistant","content":" [{\""},"logprobs":null,"finish_reason":null}]}
 
data:{"id":"","object":"text_completion","created":1720501474,"model":"/opt/ml/model","system_fingerprint":"2.0.1-native","choices":[{"index":0,"delta":{"role":"assistant","content":"name"},"logprobs":null,"finish_reason":null}]}
 
data:{"id":"","object":"text_completion","created":1720501474,"model":"/opt/ml/model","system_fingerprint":"2.0.1-native","choices":[{"index":0,"delta":{"role":"assistant","content":"\":"},"logprobs":null,"finish_reason":null}]}
 
data:{"id":"","object":"text_completion","created":1720501474,"model":"/opt/ml/model","system_fingerprint":"2.0.1-native","choices":[{"index":0,"delta":{"role":"assistant","content":" \""},"logprobs":null,"finish_reason":null}]}
 
data:{"id":"","object":"text_completion","created":1720501474,"model":"/opt/ml/model","system_fingerprint":"2.0.1-native","choices":[{"index":0,"delta":{"role":"assistant","content":"a*****"},"logprobs":null,"finish_reason":null}]}
 
data:{"id":"","object":"text_completion","created":1720501474,"model":"/opt/ml/model","system_fingerprint":"2.0.1-native","choices":[{"index":0,"delta":{"role":"assistant","content":".b"},"logprobs":null,"finish_reason":null}]}
 
data:{"id":"","object":"text_completion","created":1720501474,"model":"/opt/ml/model","system_fingerprint":"2.0.1-native","choices":[{"index":0,"delta":{"role":"assistant","content":"uiltin"},"logprobs":null,"finish_reason":null}]}
 
data:{"id":"","object":"text_completion","created":1720501474,"model":"/opt/ml/model","system_fingerprint":"2.0.1-native","choices":[{"index":0,"delta":{"role":"assistant","content":".plugin"},"logprobs":null,"finish_reason":null}]}
 
data:{"id":"","object":"text_completion","created":1720501474,"model":"/opt/ml/model","system_fingerprint":"2.0.1-native","choices":[{"index":0,"delta":{"role":"assistant","content":".music"},"logprobs":null,"finish_reason":null}]}
 
data:{"id":"","object":"text_completion","created":1720501474,"model":"/opt/ml/model","system_fingerprint":"2.0.1-native","choices":[{"index":0,"delta":{"role":"assistant","content":"\","},"logprobs":null,"finish_reason":null}]}
 
data:{"id":"","object":"text_completion","created":1720501474,"model":"/opt/ml/model","system_fingerprint":"2.0.1-native","choices":[{"index":0,"delta":{"role":"assistant","content":" \""},"logprobs":null,"finish_reason":null}]}
 
data:{"id":"","object":"text_completion","created":1720501474,"model":"/opt/ml/model","system_fingerprint":"2.0.1-native","choices":[{"index":0,"delta":{"role":"assistant","content":"intent"},"logprobs":null,"finish_reason":null}]}
 
data:{"id":"","object":"text_completion","created":1720501474,"model":"/opt/ml/model","system_fingerprint":"2.0.1-native","choices":[{"index":0,"delta":{"role":"assistant","content":"\":"},"logprobs":null,"finish_reason":null}]}
 
data:{"id":"","object":"text_completion","created":1720501474,"model":"/opt/ml/model","system_fingerprint":"2.0.1-native","choices":[{"index":0,"delta":{"role":"assistant","content":" \""},"logprobs":null,"finish_reason":null}]}
 
data:{"id":"","object":"text_completion","created":1720501474,"model":"/opt/ml/model","system_fingerprint":"2.0.1-native","choices":[{"index":0,"delta":{"role":"assistant","content":"a*****"},"logprobs":null,"finish_reason":null}]}
 
data:{"id":"","object":"text_completion","created":1720501474,"model":"/opt/ml/model","system_fingerprint":"2.0.1-native","choices":[{"index":0,"delta":{"role":"assistant","content":"Music"},"logprobs":null,"finish_reason":null}]}
 
data:{"id":"","object":"text_completion","created":1720501474,"model":"/opt/ml/model","system_fingerprint":"2.0.1-native","choices":[{"index":0,"delta":{"role":"assistant","content":"_play"},"logprobs":null,"finish_reason":null}]}
 
data:{"id":"","object":"text_completion","created":1720501474,"model":"/opt/ml/model","system_fingerprint":"2.0.1-native","choices":[{"index":0,"delta":{"role":"assistant","content":"\","},"logprobs":null,"finish_reason":null}]}
 
data:{"id":"","object":"text_completion","created":1720501474,"model":"/opt/ml/model","system_fingerprint":"2.0.1-native","choices":[{"index":0,"delta":{"role":"assistant","content":" \""},"logprobs":null,"finish_reason":null}]}
 
data:{"id":"","object":"text_completion","created":1720501474,"model":"/opt/ml/model","system_fingerprint":"2.0.1-native","choices":[{"index":0,"delta":{"role":"assistant","content":"entity"},"logprobs":null,"finish_reason":null}]}
 
data:{"id":"","object":"text_completion","created":1720501474,"model":"/opt/ml/model","system_fingerprint":"2.0.1-native","choices":[{"index":0,"delta":{"role":"assistant","content":"\":"},"logprobs":null,"finish_reason":null}]}
 
data:{"id":"","object":"text_completion","created":1720501474,"model":"/opt/ml/model","system_fingerprint":"2.0.1-native","choices":[{"index":0,"delta":{"role":"assistant","content":" {"},"logprobs":null,"finish_reason":null}]}
 
data:{"id":"","object":"text_completion","created":1720501474,"model":"/opt/ml/model","system_fingerprint":"2.0.1-native","choices":[{"index":0,"delta":{"role":"assistant","content":"}}"},"logprobs":null,"finish_reason":null}]}
 
data:{"id":"","object":"text_completion","created":1720501474,"model":"/opt/ml/model","system_fingerprint":"2.0.1-native","choices":[{"index":0,"delta":{"role":"assistant","content":"]}"},"logprobs":null,"finish_reason":null}]}
 
data:{"id":"","object":"text_completion","created":1720501474,"model":"/opt/ml/model","system_fingerprint":"2.0.1-native","choices":[{"index":0,"delta":{"role":"assistant","content":"<|eot_id|>"},"logprobs":null,"finish_reason":null}]}
 
data:{"id":"","object":"text_completion","created":1720501474,"model":"/opt/ml/model","system_fingerprint":"2.0.1-native","choices":[{"index":0,"delta":{"role":"assistant","content":"<|end_of_text|>"},"logprobs":null,"finish_reason":"eos_token"}]}

마치며

Amazon SageMaker는 강력한 MLOps 기능 등을 제공하며, LLM 이 등장한 이후에도 여러 AWS 서비스와 연동하여 다양한 기능을 지원하고 있습니다.

모델의 응답을 스트리밍하는 기능은 아직 제약사항이 많지만 위와 같이 구현이 가능한 것을 파악할 수 있었습니다.

서빙 성능 및 비용과는 별개로, Serverless 환경에서 빠르게 모델을 서빙할 수 있는 환경을 제공하기에 상당히 유용한 옵션이라고 생각합니다.

#Amazon Web Service #Language Models #A.(에이닷)

등록

DEVOTEE를 활성화 시키면
지금 작성한 댓글에 AI가 댓글을 달아줍니다.

익명

Cyu 님의 최신 블로그

Amazon SageMaker로 LLM 응답 Streaming 서빙하기

SageMaker Endpoint Serving Amazon SageMaker는 머신러닝 모델을 더 쉽게 구축, 훈련, 배포할 수 있도록 지원하며 데이터 준비부터 모델 배포까지 일련의 과정을 효율적으로 처리할 수 있게 도와주는 서비스입니다. 일반적으로 Sagemaker로 배포한 모델은 API Gateway와 Lambda를 이용해서 외부에 serverless로 서빙할 수 있는데요, S ageMaker Endpoint를 서빙하는 자세한 방법은 여기 를 참고하시기 바랍니다. Lambda Response Streaming 위 방법으로 모델을 서빙하면 모델의 응답이 끝날 때까지 기다렸다가 응답을 돌려주게 됩니다. LLM을 서빙할 경우 모델의 특성상 응답 생성 시간이 오래 ...

DEVOTEE 추천 블로그

AWS Bedrock의 imported model 사용 하기

1. AWS Bedrock 기본 정보 AWS Bedrock은 Amazon Web Services(AWS)에서 제공하는 완전 관리형 서비스로, 다양한 고성능 파운데이션 모델(Foundation Model)을 통합 API를 통해 사용할 수 있게 해줍니다. 이 서비스는 AI21 Labs, Anthropic, Cohere, Meta, Mistral AI, Stability AI 및 Amazon과 같은 회사의 모델을 포함하여 다양한 모델을 제공합니다. ［참고］ https://aws.amazon.com/ko/blogs/tech/category/artificial-intelligence/amazon-machine-learning/amazon-bedrock/ ［참고］ https...
AWS CDK로 구축하는 Vision AI 모델 서빙 파이프라인 실전 가이드

이 글은 실제 프로덕션에서 운영을 위해 개발을 진행했던 VAIaaS(Vision AI as a Service) 프로젝트를 기반으로 작성되었습니다. CDK + API Gateway + Lambda + SageMaker 조합으로 Vision AI 모델 서빙 파이프라인을 구축한 경험을 공유합니다. 들어가며 Vision AI 모델을 개발하는 것과 그 모델을 실제 서비스로 제공하는 것은 전혀 다른 문제입니다. 연구 환경에서 높은 정확도를 달성한 모델이라도, 이를 안정적이고 확장 가능한 API 서비스로 전환하는 과정에는 수많은 엔지니어링 과제가 숨어 있습니다. 인프라 프로비저닝, 모델 배포 자동화, 요청 처리 파이프라인 설계, 비용 최적화, 그리고 운영 모니터링까지 — 이 ...

목록으로

DEVOCEAN

로그인 선택

미리보기

커뮤니티

글 등록

임시저장함

데보션 블로그 게재 요청