在 Amazon SageMaker AI 上引入实时推理的双向流媒体功能

Client: [发送完整问题] → 等待... Model: ...处理... [返回完整答案] Client: [发送下一个问题] → 等待... Model: ...处理... [返回完整答案]

Client: [携带足够上下文的问题开始流动] → Model: ← [答案立即开始流动] Client: → [继续/调整问题] ↓ Model: ← [实时调整答案]

# 该应用程序使用环境变量进行 AWS 身份验证。在运行应用程序之前设置它们： # export AWS_ACCESS_KEY_ID="your-access-key" # export AWS_SECRET_ACCESS_KEY="your-secret-key" # export AWS_DEFAULT_REGION="us-west-2" container_name="sagemaker-bidirectional-streaming" container_tag="latest" cd container account=$(aws sts get-caller-identity --query Account --output text) # 获取当前配置中定义的区域（如果未定义，则默认为 us-west-2） region=$(aws configure get region) region=${region:-us-west-2} container_image_uri="${account}.dkr.ecr.${region}.amazonaws.com/${container_name}:${container_tag}" # 如果 ECR 中不存在该仓库，则创建它。 aws ecr describe-repositories --repository-names "${container_name}" --region "${region}" > /dev/null 2>&1 if [ $? -ne 0 ] then aws ecr create-repository --repository-name "${container_name}" --region "${region}" > /dev/null fi # 获取 ECR 的登录命令并直接执行它 aws ecr get-login-password --region ${region} | docker login --username AWS --password-stdin ${account}.dkr.ecr.${region}.amazonaws.com/${container_name} # 使用镜像名称在本地构建 docker 镜像，然后将其推送到 ECR # 使用完整名称。 docker build --platform linux/amd64 --provenance=false -t ${container_name} . docker tag ${container_name} ${container_image_uri} docker push ${container_image_uri}

com.amazonaws.sagemaker.capabilities.bidirectional-streaming=true

import boto3 from datetime import datetime sagemaker_client = boto3.client('sagemaker', region_name=REGION) # 创建模型 sagemaker_client.create_model( ModelName=MODEL_NAME, PrimaryContainer={'Image': IMAGE_URI, 'Mode': 'SingleModel'}, ExecutionRoleArn=ROLE ) # 创建配置 sagemaker_client.create_endpoint_config( EndpointConfigName=ENDPOINT_CONFIG_NAME, ProductionVariants=[{ 'VariantName': 'AllTraffic', 'ModelName': MODEL_NAME, 'InitialInstanceCount': 1, 'InstanceType': INSTANCE_TYPE, 'InitialVariantWeight': 1.0 }] ) # 创建端点 sagemaker_client.create_endpoint( EndpointName=ENDPOINT_NAME, EndpointConfigName=ENDPOINT_CONFIG_NAME ) print(f"Endpoint '{ENDPOINT_NAME}' creation initiated")

#!/usr/bin/env python3 """ SageMaker AI 双向流媒体 Python SDK 脚本。 此脚本连接到 SageMaker AI 端点以进行双向流媒体通信。 """ import argparse import asyncio import sys from aws_sdk_sagemaker_runtime_http2.client import SageMakerRuntimeHTTP2Client from aws_sdk_sagemaker_runtime_http2.config import Config, HTTPAuthSchemeResolver from aws_sdk_sagemaker_runtime_http2.models import InvokeEndpointWithBidirectionalStreamInput, RequestStreamEventPayloadPart, RequestPayloadPart from smithy_aws_core.identity import EnvironmentCredentialsResolver from smithy_aws_core.auth.sigv4 import SigV4AuthScheme import logging def parse_arguments(): """解析命令行参数。""" parser = argparse.ArgumentParser( description="Connect to SageMaker AI endpoint for bidirectional streaming" ) parser.add_argument( "ENDPOINT_NAME", help="Name of the SageMaker AI endpoint to connect to" ) return parser.parse_args() # 配置 AWS_REGION = "us-west-2" BIDI_ENDPOINT = f"https://runtime.sagemaker.{AWS_REGION}:8443" logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) class SimpleClient: def __init__(self, endpoint_name, region=AWS_REGION): self.endpoint_name = endpoint_name self.region = region self.client = None self.stream = None self.response = None self.is_active = False def _initialize_client(self): config = Config( endpoint_uri=BIDI_ENDPOINT, region=self.region, aws_credentials_identity_resolver=EnvironmentCredentialsResolver(), auth_scheme_resolver=HTTPAuthSchemeResolver(), auth_schemes={"aws.auth#sigv4": SigV4AuthScheme(service="sagemaker")} ) self.client = SageMakerRuntimeHTTP2Client(config=config) async def start_session(self): if not self.client: self._initialize_client() logger.info(f"Starting session with endpoint: {self.endpoint_name}") self.stream = await self.client.invoke_endpoint_with_bidirectional_stream( InvokeEndpointWithBidirectionalStreamInput(endpoint_name=self.endpoint_name) ) self.is_active = True self.response = asyncio.create_task(self._process_responses()) async def send_words(self, words): for i, word in enumerate(words): logger.info(f"Sending payload: {word}") await self.send_event(word.encode('utf-8')) await asyncio.sleep(1) async def send_event(self, data_bytes): payload = RequestPayloadPart(bytes_=data_bytes) event = RequestStreamEventPayloadPart(value=payload) await self.stream.input_stream.send(event) async def end_session(self): if not self.is_active: return await self.stream.input_stream.close() logger.info("Stream closed") async def _process_responses(self): try: output = await self.stream.await_output() output_stream = output[1] while self.is_active: result = await output_stream.receive() if result is None: logger.info("No more responses") break if result.value and result.value.bytes_: response_data = result.value.bytes_.decode('utf-8') logger.info(f"Received: {response_data}") except Exception as e: logger.error(f"Error processing responses: {e}") def main(): """主函数，用于解析参数并运行流式客户端。""" args = parse_arguments() print("=" * 60) print("SageMaker AI 双向流媒体客户端") print("=" * 60) print(f"Endpoint Name: {args.ENDPOINT_NAME}") print(f"AWS Region: {AWS_REGION}") print("=" * 60) async def run_client(): sagemaker_client = SimpleClient(endpoint_name=args.ENDPOINT_NAME) try: await sagemaker_client.start_session() words = ["I need help with", "my account balance", "I can help with that", "and recent charges"] await sagemaker_client.send_words(words) await asyncio.sleep(2) await sagemaker_client.end_session() sagemaker_client.is_active = False if sagemaker_client.response and not sagemaker_client.response.done(): sagemaker_client.response.cancel() logger.info("Session ended successfully") return 0 except Exception as e: logger.error(f"Client error: {e}") return 1 try: exit_code = asyncio.run(run_client()) sys.exit(exit_code) except KeyboardInterrupt: logger.info("Interrupted by user") sys.exit(1) except Exception as e: logger.error(f"Unexpected error: {e}") sys.exit(1) if __name__ == "__main__": main()