使用托管在 SageMaker AI 端点上的 LLM 为 Strands Agent 构建自定义模型提供程序

# 全局安装 Yeoman npm install -g yo # 克隆并安装 ml-container-creator git clone https://github.com/awslabs/ml-container-creator cd ml-container-creator npm install && npm link # 验证安装 yo --generators # 应显示 ml-container-creator

# 运行生成器 yo ml-container-creator # 配置选项: # - Framework: transformers # - Model Server: sglang # - Model: meta-llama/Llama-3.1-8B-Instruct # - Deploy Target: codebuild # - Instance Type: ml.g6.12xlarge (GPU) # - Region: us-east-1

<project-directory>/ ├── Dockerfile # 包含 SGLang 和依赖项的容器 ├── buildspec.yml # CodeBuild 配置 ├── code/ │ └── serve # SGLang 服务器启动脚本 ├── deploy/ │ ├── submit_build.sh # 触发 CodeBuild │ └── deploy.sh # 部署到 SageMaker └── test/ └── test_endpoint.sh # 端点测试脚本

cd llama-31-deployment # 使用 CodeBuild 构建容器 (无需本地 Docker) ./deploy/submit_build.sh # 部署到 SageMaker ./deploy/deploy.sh arn:aws:iam::ACCOUNT:role/SageMakerExecutionRole

import boto3 import json runtime_client = boto3.client('sagemaker-runtime', region_name='us-east-1') payload = { "messages": [ {"user", "content": "Hello, how are you?"} ], "max_tokens": 100, "temperature": 0.7 } response = runtime_client.invoke_endpoint( EndpointName='llama-31-deployment-endpoint', ContentType='application/json', Body=json.dumps(payload) ) result = json.loads(response['Body'].read().decode('utf-8')) print(result['choices'][0]['message']['content'])

{ "id": "cmpl-abc123", "object": "chat.completion", "created": 1704067200, "model": "meta-llama/Llama-3.1-8B-Instruct", "choices": [{ "index": 0, "message": {"role": "assistant", "content": "I'm doing well, thank you for asking!"}, "finish_reason": "stop" }], "usage": { "prompt_tokens": 23, "completion_tokens": 12, "total_tokens": 35 } }

def stream(self, messages: List[Dict[str, Any]], tool_specs: list, system_prompt: Optional[str], **kwargs): # 构建负载 messages payload_messages = [] if system_prompt: payload_messages.append({"role": "system", "content": system_prompt}) # 从 Strands 格式中提取消息内容 for msg in messages: payload_messages.append({"role": "user", "content": msg['content'][0]['text']}) # 构建带有流式传输启用的完整负载 payload = { "messages": payload_messages, "max_tokens": kwargs.get('max_tokens', self.max_tokens), "temperature": kwargs.get('temperature', self.temperature), "top_p": kwargs.get('top_p', self.top_p), "stream": True } try: # 使用流式响应调用 SageMaker 端点 response = self.runtime_client.invoke_endpoint_with_response_stream( EndpointName=self.endpoint_name, ContentType='application/json', Accept='application/json', Body=json.dumps(payload) ) # 处理流式响应 accumulated_content = "" for event in response['Body']: chunk = event['PayloadPart']['Bytes'].decode('utf-8') if not chunk.strip(): continue # 解析 SSE 格式: "data: {json}\n" for line in chunk.split('\n'): if line.startswith('data: '): try: json_str = line.replace('data: ', '').strip() if not json_str: continue chunk_data = json.loads(json_str) if 'choices' in chunk_data and chunk_data['choices']: delta = chunk_data['choices'][0].get('delta', {}) # 以 Strands 格式 Yield 内容块 if 'content' in delta: content_chunk = delta['content'] accumulated_content += content_chunk yield { "type": "contentBlockDelta", "delta": {"text": content_chunk}, "contentBlockIndex": 0 } # 检查完成 finish_reason = chunk_data['choices'][0].get('finish_reason') if finish_reason: yield { "type": "messageStop", "stopReason": finish_reason } # Yield 使用元数据 if 'usage' in chunk_data: yield { "type": "metadata", "usage": chunk_data['usage'] } except json.JSONDecodeError: continue except Exception as e: yield { "type": "error", "error": { "message": f"Endpoint invocation failed: {str(e)}", "type": "EndpointInvocationError" } }

from strands.agent import Agent # 初始化自定义提供程序 provider = LlamaModelProvider( endpoint_name="llama-31-deployment-endpoint", region_name="us-east-1", max_tokens=1000, temperature=0.7 ) # 使用自定义提供程序创建代理 agent = Agent( name="llama-assistant", model=provider, system_prompt=( "You are a helpful AI assistant powered by Llama 3.1, " "deployed on Amazon SageMaker. You provide clear, accurate, " "and friendly responses to user questions." ) ) # 测试代理 response = agent("What are the key benefits of deploying LLMs on SageMaker?") print(response.content)