SageMaker AI 엔드포인트에 호스팅된 LLM을 위한 Strands 에이전트용 커스텀 모델 공급자 구축

이 요약은 AI가 원문을 분석해 생성했습니다. 정확한 내용은 원문 기준으로 확인하세요.

핵심 요약

기업들은 제어권 확보와 비용 최적화를 위해 SageMaker AI 엔드포인트에 SGLang이나 vLLM 같은 프레임워크를 사용하여 커스텀 LLM을 배포한다. 하지만 이러한 프레임워크는 주로 OpenAI 호환 형식을 반환하는 반면, Strands 에이전트는 Bedrock Messages API 형식을 요구하여 호환성 문제가 발생한다. 이를 해결하기 위해 SageMakerAIModel 클래스를 확장하여 응답 형식을 변환하는 커스텀 모델 파서를 구현하는 방법을 제시한다. 결과적으로 개발자는 선호하는 추론 엔진을 유지하면서도 Strands 에이전트 SDK의 기능을 온전히 활용할 수 있다.

배경

AWS 계정 및 SageMaker 권한, Python 및 Docker 기본 지식, Strands 에이전트 SDK에 대한 이해, Node.js 및 Yeoman (ml-container-creator 사용 시)

대상 독자

SageMaker에 커스텀 LLM을 배포하고 Strands 에이전트 SDK로 AI 에이전트를 구축하려는 개발자

의미 / 영향

이 방법론은 특정 클라우드 제공업체의 표준 API에 종속되지 않고도 고성능 추론 엔진(SGLang 등)을 자유롭게 선택할 수 있게 한다. 이는 기업이 최신 오픈소스 모델과 최적화 기술을 에이전트 시스템에 즉각적으로 도입할 수 있는 기술적 유연성을 제공한다.

섹션별 상세

SGLang이나 vLLM 같은 커스텀 서빙 프레임워크는 범용성을 위해 OpenAI 호환 형식을 사용하지만, Strands 에이전트는 Bedrock Messages API 형식을 기대한다. 이러한 데이터 구조의 불일치는 기본 SageMakerAIModel 클래스를 사용할 때 필드 접근 오류인 TypeError를 유발하며 시스템 통합을 방해한다.

AWS Labs의 오픈소스 도구인 ml-container-creator를 활용하여 SageMaker BYOC 프로젝트를 자동 생성한다. 이 도구는 Dockerfile, CodeBuild 설정, 배포 스크립트를 포함한 전체 프로젝트 구조를 생성하여 Llama 3.1 모델을 SGLang 서버와 함께 배포하는 과정을 표준화하고 간소화한다.

bash

# Install Yeoman globally
npm install -g yo
# Clone and install ml-container-creator
git clone https://github.com/awslabs/ml-container-creator
cd ml-container-creator
npm install && npm link

# Run the generator
yo ml-container-creator
# Configuration: Framework=transformers, Model Server=sglang, Model=meta-llama/Llama-3.1-8B-Instruct

ml-container-creator를 설치하고 SGLang 기반의 Llama 3.1 배포 프로젝트를 생성하는 과정

SageMakerAIModel을 상속받는 LlamaModelProvider 클래스를 생성하고 stream() 메서드를 오버라이드하여 파싱 로직을 구현한다. 이 메서드는 Strands의 요청을 모델 서버 형식으로 변환하여 전달하고, 모델의 스트리밍 응답(SSE)을 다시 Strands가 이해할 수 있는 contentBlockDelta 및 messageStop 형식으로 변환하여 반환한다.

python

def stream(self, messages: List[Dict[str, Any]], tool_specs: list, system_prompt: Optional[str], **kwargs):
    # ... (중략)
    payload = {
        "messages": payload_messages,
        "stream": True
    }
    response = self.runtime_client.invoke_endpoint_with_response_stream(
        EndpointName=self.endpoint_name,
        Body=json.dumps(payload)
    )
    for event in response['Body']:
        chunk = event['PayloadPart']['Bytes'].decode('utf-8')
        # SSE 포맷 파싱 및 Strands 포맷으로 yield
        if 'content' in delta:
            yield {
                "type": "contentBlockDelta",
                "delta": {"text": delta['content']},
                "contentBlockIndex": 0
            }

SageMakerAIModel을 확장하여 OpenAI 호환 스트리밍 응답을 Strands 포맷으로 변환하는 핵심 로직

에이전트 레이어, 파서 레이어, 모델 배포 레이어로 구성된 시스템 아키텍처 다이어그램이다. — DiagramStrands 에이전트가 커스텀 파서(LlamaModelProvider)를 거쳐 SageMaker 엔드포인트에 도달하는 흐름을 보여준다. 각 레이어의 역할과 데이터 변환 과정을 시각화하여 본문의 구현 구조를 명확히 설명한다.

구현된 커스텀 공급자를 Strands Agent 초기화 시 model 파라미터로 전달하여 통합을 완료한다. 이를 통해 에이전트는 내부 파싱 로직에 구애받지 않고 일관된 인터페이스로 커스텀 호스팅된 Llama 3.1 모델과 대화할 수 있으며, 시스템 프롬프트 및 컨텍스트 관리가 정상적으로 작동한다.

python

from strands.agent import Agent

provider = LlamaModelProvider(
    endpoint_name="llama-31-deployment-endpoint",
    region_name="us-east-1"
)

agent = Agent(
    name="llama-assistant",
    model=provider,
    system_prompt="You are a helpful AI assistant powered by Llama 3.1."
)

response = agent("What are the key benefits of deploying LLMs on SageMaker?")

구현된 커스텀 공급자를 사용하여 Strands 에이전트를 초기화하고 실행하는 예시

실무 Takeaway

ml-container-creator를 사용하면 복잡한 SageMaker BYOC 배포 인프라 코드를 자동화하여 프로덕션급 환경 구축 시간을 단축할 수 있다.
Strands 에이전트 SDK의 stream() 메서드를 커스텀 구현함으로써 표준 API를 따르지 않는 특수한 LLM 엔드포인트도 에이전트 워크플로우에 통합 가능하다.
OpenAI 호환 형식을 Bedrock Messages API 형식으로 변환하는 파서 레이어를 두어 모델 서빙의 유연성과 에이전트 프레임워크의 편의성을 동시에 확보한다.

언급된 리소스

GitHubml-container-creator GitHub Repository