使用 Amazon SageMaker 异步终端节点优化 Amazon SageMaker JumpStart 基础模型的部署成本 | 亚马逊网络服务

%%time
from sagemaker.jumpstart.model import JumpStartModel, AsyncInferenceConfig
from sagemaker.serializers import JSONSerializer
from sagemaker.deserializers import JSONDeserializer model_id, model_version = "huggingface-llm-falcon-40b-instruct-bf16", "*"
my_model = JumpStartModel(model_id=model_id)
predictor = my_model.deploy( initial_instance_count=0, instance_type="ml.g5.12xlarge", async_inference_config=AsyncInferenceConfig()
)

import boto3 client = boto3.client("application-autoscaling")
resource_id = "endpoint/" + my_model.endpoint_name + "/variant/" + "AllTraffic" # Configure Autoscaling on asynchronous endpoint down to zero instances
response = client.register_scalable_target( ServiceNamespace="sagemaker", ResourceId=resource_id, ScalableDimension="sagemaker:variant:DesiredInstanceCount", MinCapacity=0, # Miminum number of instances we want to scale down to - scale down to 0 to stop incurring in costs MaxCapacity=1, # Maximum number of instances we want to scale up to - scale up to 1 max is good enough for dev
) response = client.put_scaling_policy( PolicyName="Invocations-ScalingPolicy", ServiceNamespace="sagemaker", # The namespace of the AWS service that provides the resource. ResourceId=resource_id, # Endpoint name ScalableDimension="sagemaker:variant:DesiredInstanceCount", # SageMaker supports only Instance Count PolicyType="TargetTrackingScaling", # 'StepScaling'|'TargetTrackingScaling' TargetTrackingScalingPolicyConfiguration={ "TargetValue": 5.0, # The target value for the metric. - here the metric is - SageMakerVariantInvocationsPerInstance "CustomizedMetricSpecification": { "MetricName": "ApproximateBacklogSizePerInstance", "Namespace": "AWS/SageMaker", "Dimensions": [{"Name": "EndpointName", "Value": my_model.endpoint_name}], "Statistic": "Average", }, "ScaleInCooldown": 600, # The amount of time, in seconds, after a scale in activity completes before another scale in activity can start. "ScaleOutCooldown": 300, # ScaleOutCooldown - The amount of time, in seconds, after a scale out activity completes before another scale out activity can start. # 'DisableScaleIn': True|False - indicates whether scale in by the target tracking policy is disabled. # If the value is true, scale in is disabled and the target tracking policy won't remove capacity from the scalable resource. },
)

import time # Invoking the asynchronous endpoint with the SageMaker Python SDK
def query_endpoint(payload): """Query endpoint and print the response""" response = predictor.predict_async( data=payload, input_path="s3://{}/{}".format(bucket, prefix), ) while True: try: response = response.get_result() break except: print("Inference is not ready ...") time.sleep(5) print(f"33[1m Input:33[0m {payload['inputs']}") print(f"33[1m Output:33[0m {response[0]['generated_text']}") query_endpoint(payload)

import json
import time
import boto3
from botocore.exceptions import ClientError s3_client = boto3.client("s3") # Wait until the prediction is generated
def wait_inference_file(bucket, prefix): while True: try: response = s3_client.get_object(Bucket=bucket, Key=prefix) break except ClientError as ex: if ex.response['Error']['Code'] == 'NoSuchKey': print("Waiting for file to be generated...") time.sleep(5) next else: raise except Exception as e: print(e.__dict__) raise return response

# Invoking the asynchronous endpoint with the Boto3 SDK
import boto3 sagemaker_client = boto3.client("sagemaker-runtime") # Query the endpoint function
def query_endpoint_boto3(payload): """Query endpoint and print the response""" response = sagemaker_client.invoke_endpoint_async( EndpointName=my_model.endpoint_name, InputLocation="s3://{}/{}".format(bucket, prefix), ContentType="application/json", Accept="application/json" ) output_url = response["OutputLocation"] output_prefix = "/".join(output_url.split("/")[3:]) # Read the bytes of the file from S3 in output_url with Boto3 output = wait_inference_file(bucket, output_prefix) output = json.loads(output['Body'].read())[0]['generated_text'] # Emit output print(f"33[1m Input:33[0m {payload['inputs']}") print(f"33[1m Output:33[0m {output}") query_endpoint_boto3(payload)

from typing import Dict
from langchain import PromptTemplate
from langchain.llms.sagemaker_endpoint import LLMContentHandler
from langchain.chains import LLMChain
from sagemaker_async_endpoint import SagemakerAsyncEndpoint class ContentHandler(LLMContentHandler): content_type:str = "application/json" accepts:str = "application/json" len_prompt:int = 0 def transform_input(self, prompt: str, model_kwargs: Dict) -> bytes: self.len_prompt = len(prompt) input_str = json.dumps({"inputs": prompt, "parameters": {"max_new_tokens": 100, "do_sample": False, "repetition_penalty": 1.1}}) return input_str.encode('utf-8') def transform_output(self, output: bytes) -> str: response_json = output.read() res = json.loads(response_json) ans = res[0]['generated_text'] return ans chain = LLMChain( llm=SagemakerAsyncEndpoint( input_bucket=bucket, input_prefix=prefix, endpoint_name=my_model.endpoint_name, region_name=sagemaker.Session().boto_region_name, content_handler=ContentHandler(), ), prompt=PromptTemplate( input_variables=["query"], template="{query}", ),
) print(chain.run(payload['inputs']))

predictor.delete_endpoint()