End-to-end workflow
After you've set up Amazon Bedrock Marketplace, you can use the following example code in your end-to-end workflow. If you need more context, you can read the sections that follow the code.
from botocore.exceptions import ClientError import pprint from datetime import datetime import json import time import sys import boto3 import argparse SM_HUB_NAME = 'SageMakerPublicHub' DELIMITER = "\n\n\n\n================================================================================================" class Bedrock: def __init__(self, region_name) -> None: self.region_name = region_name self.boto3_session = boto3.session.Session() self.sagemaker_client = self.boto3_session.client( service_name='sagemaker', region_name=self.region_name, ) self.bedrock_client = self.boto3_session.client( service_name='bedrock', region_name=self.region_name ) self.endpoint_paginator = self.bedrock_client.get_paginator('list_marketplace_model_endpoints') self.bedrock_runtime_client = self.boto3_session.client( service_name='bedrock-runtime', region_name=self.region_name) def list_models(self): SM_RESPONSE_FIELD_NAME = 'HubContentSummaries' SM_HUB_CONTENT_TYPE = 'Model' response = self.sagemaker_client.list_hub_contents( MaxResults=100, HubName=SM_HUB_NAME, HubContentType=SM_HUB_CONTENT_TYPE ) all_models = Bedrock.extract_bedrock_models(response[SM_RESPONSE_FIELD_NAME]) while ("NextToken" in response) and response["NextToken"]: response = self.sagemaker_client.list_hub_contents( MaxResults=100, HubName=SM_HUB_NAME, HubContentType=SM_HUB_CONTENT_TYPE, NextToken=response['NextToken'] ) extracted_models = Bedrock.extract_bedrock_models(response[SM_RESPONSE_FIELD_NAME]) if not extracted_models: # Bedrock enabled models always appear first, therefore can return when results are empty. return all_models all_models.extend(extracted_models) time.sleep(1) return all_models def describe_model(self, hub_name: str, hub_content_name: str): return self.sagemaker_client.describe_hub_content( HubName=hub_name, HubContentType='Model', HubContentName=hub_content_name ) def list_endpoints(self): for response in self.endpoint_paginator.paginate(): for endpoint in response['marketplaceModelEndpoints']: yield endpoint def list_endpoints_for_model(self, hub_content_arn: str): for response in self.endpoint_paginator.paginate( modelSourceEquals=hub_content_arn): for endpoint in response['marketplaceModelEndpoints']: yield endpoint # acceptEula needed only for gated models def create_endpoint(self, model, endpoint_config, endpoint_name: str, tags = []): model_arn = model['HubContentArn'] if self._requires_eula(model=model): return self.bedrock_client.create_marketplace_model_endpoint( modelSourceIdentifier=model_arn, endpointConfig=endpoint_config, endpointName=endpoint_name, acceptEula=True, tags=tags ) else: return self.bedrock_client.create_marketplace_model_endpoint( modelSourceIdentifier=model_arn, endpointConfig=endpoint_config, endpointName=endpoint_name, tags=tags ) def delete_endpoint(self, endpoint_arn: str): return self.bedrock_client.delete_marketplace_model_endpoint(endpointArn=endpoint_arn) def describe_endpoint(self, endpoint_arn: str): return self.bedrock_client.get_marketplace_model_endpoint(endpointArn=endpoint_arn)['marketplaceModelEndpoint'] def update_endpoint(self, endpoint_arn: str, endpoint_config): return self.bedrock_client.update_marketplace_model_endpoint(endpointArn=endpoint_arn, endpointConfig=endpoint_config) def register_endpoint(self, endpoint_arn: str, model_arn: str): return self.bedrock_client.register_marketplace_model_endpoint(endpointIdentifier=endpoint_arn, modelSourceIdentifier=model_arn)['marketplaceModelEndpoint']['endpointArn'] def deregister_endpoint(self, endpoint_arn: str): return self.bedrock_client.deregister_marketplace_model_endpoint(endpointArn=endpoint_arn) def invoke(self, endpoint_arn: str, body): response = self.bedrock_runtime_client.invoke_model(modelId=endpoint_arn, body=body, contentType='application/json') return json.loads(response["body"].read()) def invoke_with_stream(self, endpoint_arn: str, body): return self.bedrock_runtime_client.invoke_model_with_response_stream(modelId=endpoint_arn, body=body) def converse(self, endpoint_arn: str, conversation): return self.bedrock_runtime_client.converse(modelId=endpoint_arn, messages=conversation) def converse_with_stream(self, endpoint_arn: str, conversation): return self.bedrock_runtime_client.converse_stream(modelId=endpoint_arn, messages=conversation, inferenceConfig={"maxTokens": 4096, "temperature": 0.5, "topP": 0.9}) def wait_for_endpoint(self, endpoint_arn: str): endpoint = self.describe_endpoint(endpoint_arn=endpoint_arn) while endpoint['endpointStatus'] in ['Creating', 'Updating']: print( f"Endpoint {endpoint_arn} status is still {endpoint['endpointStatus']}. Waiting 10 seconds before continuing...") time.sleep(10) endpoint = self.describe_endpoint(endpoint_arn=endpoint_arn) print(f"Endpoint status: {endpoint['status']}") def _requires_eula(self, model): if 'HubContentDocument' in model: hcd = json.loads(model['HubContentDocument']) if ('HostingEulaUri' in hcd) and hcd['HostingEulaUri']: return True return False @staticmethod def extract_bedrock_models(hub_content_summaries): models = [] for content in hub_content_summaries: if ('HubContentSearchKeywords' in content) and ( '@capability:bedrock_console' in content['HubContentSearchKeywords']): print(f"ModelName: {content['HubContentDisplayName']}, modelSourceIdentifier: {content['HubContentArn']}") models.append(content) return models def run_script(sagemaker_execution_role: str, region: str): # Script params model_arn = 'arn:aws:sagemaker:AWS Region:aws:hub-content/SageMakerPublicHub/Model/example-model-name/hub-content-arn' model_name = 'example-model-name' sample_endpoint_name = f'test-ep-{datetime.now().strftime("%Y-%m-%d%H%M%S")}' sagemaker_execution_role = sagemaker_execution_role conversation = [ { "role": "user", "content": [ { "text": "whats the best park in the US?" } ] } ] bedrock = Bedrock(region_name=region) ### ### Model discovery ### # List all models - no new Bedrock Marketplace API here. Uses existing SageMaker APIs print(DELIMITER) print("All models:") all_models = bedrock.list_models() # Describe a model - no new Bedrock Marketplace API here. Uses existing SageMaker APIs # Examples: # bedrock.describe_model("SageMakerPublicHub", "huggingface-llm-amazon-mistrallite") # bedrock.describe_model("SageMakerPublicHub", "huggingface-llm-gemma-2b-instruct") print(DELIMITER) print(f'Describing model: {model_name}') model = bedrock.describe_model(SM_HUB_NAME, model_name) pprint.pprint(model) ## If customer wants to use a proprietary model, they need to subscribe to it first ## If customer wants to use a gated model, they need to accept EULA. Note: EULA Acceptance is on-creation, and needs ## to be provided on every call. Cannot un-accept a EULA ## If customer wants to use an open weight model, they can proceed to deploy ### ### Model deployment to create endpoints ### # # Create endpoint - uses Bedrock Marketplace API endpoint_arn = bedrock.create_endpoint( endpoint_name=sample_endpoint_name, endpoint_config={ "sageMaker": { "initialInstanceCount": 1, "instanceType": "ml.g5.2xlarge", "executionRole": sagemaker_execution_role # Other fields: # kmsEncryptionKey: KmsKeyId # vpc: VpcConfig } }, # Optional: # tags: TagList model=model )['marketplaceModelEndpoint']['endpointArn'] # # Describe endpoint - uses Bedrock Marketplace API endpoint = bedrock.describe_endpoint(endpoint_arn=endpoint_arn) print(DELIMITER) print('Created endpoint:') pprint.pprint(endpoint) # Wait while endpoint is being created print(DELIMITER) bedrock.wait_for_endpoint(endpoint_arn=endpoint_arn) ### ### Currently, customers cannot use self-hosted endpoints with Bedrock Runtime APIs and tools. They can only pass a model ID to the APIs. ### Bedrock Marketplace will enable customers to use self-hosted endpoints through existing Bedrock Runtime APIs and tools ### See below examples of calling invoke_model, invoke_model_with_response_stream, converse and converse_stream ### Customers will be able to use the endpoints with Bedrock dev tools also (Guardrails, Model eval, Agents, Knowledge bases, Prompt flows, Prompt management) - examples not shown below ### # Prepare sample data for invoke calls by getting default payload in model metadata model_data = json.loads(bedrock.describe_model('SageMakerPublicHub', model_name)['HubContentDocument']) payload = list(model_data["DefaultPayloads"].keys())[0] invoke_body = model_data["DefaultPayloads"][payload]["Body"] invoke_content_field_name = 'generated_text' # Invoke model (text) - without stream - uses existing Bedrock Runtime API print(DELIMITER) print(f'Invoking model with body: {invoke_body}') invoke_generated_response = bedrock.invoke(endpoint_arn=endpoint_arn, body=json.dumps(invoke_body)) print(f'Generated text:') print(invoke_generated_response[invoke_content_field_name]) sys.stdout.flush() # Converse with model (chat) - without stream - uses existing Bedrock Runtime API print(DELIMITER) print(f'Converse model with conversation: {conversation}') print(bedrock.converse(endpoint_arn=endpoint_arn, conversation=conversation)['output']) ### ## Other endpoint management operations ### # List all endpoints - uses Bedrock Marketplace API print(DELIMITER) print('Listing all endpoints') for endpoint in bedrock.list_endpoints(): pprint.pprint(endpoint) # List endpoints for a model # Example: bedrock.list_endpoints_for_model(hub_content_arn='arn:aws:sagemaker:us-west-2:aws:hub-content/SageMakerPublicHub/Model/huggingface-textgeneration1-mpt-7b-storywriter-bf16/3.2.0') print(DELIMITER) print(f"Listing all endpoints for model: {model_arn}") for endpoint in bedrock.list_endpoints_for_model(hub_content_arn=model_arn): pprint.pprint(endpoint) # # Update endpoint - uses new API provided by Bedrock Marketplace updated_endpoint_arn = bedrock.update_endpoint( endpoint_arn=endpoint_arn, endpoint_config={ "sageMaker": { "initialInstanceCount": 2, # update to increase instance count "instanceType": "ml.g5.2xlarge", "executionRole": sagemaker_execution_role # Other fields: # kmsEncryptionKey: KmsKeyId # vpc: VpcConfig } # Optional: # tags: TagList } )['marketplaceModelEndpoint']['endpointArn'] # Wait while endpoint is being updated print(DELIMITER) bedrock.wait_for_endpoint(endpoint_arn=updated_endpoint_arn) # Confirm endpoint update updated_endpoint = bedrock.describe_endpoint(endpoint_arn=updated_endpoint_arn) print(f'Updated endpoint: {updated_endpoint}') assert updated_endpoint['endpointConfig']['sageMaker']['initialInstanceCount'] == 2 print(DELIMITER) print(f'Confirmed that updated endpoint\'s initialInstanceCount config changed from 1 to 2') # Wait while endpoint is being updated print(DELIMITER) bedrock.wait_for_endpoint(endpoint_arn=updated_endpoint_arn) # Deregister endpoint - uses Bedrock Marketplace API print(DELIMITER) print(f'De-registering endpoint: {updated_endpoint_arn}') bedrock.deregister_endpoint(endpoint_arn=updated_endpoint_arn) try: pprint.pprint(bedrock.describe_endpoint(endpoint_arn=updated_endpoint_arn)) except ClientError as err: assert err.response['Error']['Code'] == 'ResourceNotFoundException' print(f"Confirmed that endpoint {updated_endpoint_arn} was de-registered") # Re-register endpoint - uses Bedrock Marketplace API print(DELIMITER) print(f'Registered endpoint: {bedrock.register_endpoint(endpoint_arn=updated_endpoint_arn, model_arn=model_arn)}') pprint.pprint(bedrock.describe_endpoint(endpoint_arn=updated_endpoint_arn)) # Delete endpoint - uses Bedrock Marketplace API print(DELIMITER) print(f'Deleting endpoint: {updated_endpoint_arn}') bedrock.delete_endpoint(endpoint_arn=updated_endpoint_arn) try: pprint.pprint(bedrock.describe_endpoint(endpoint_arn=updated_endpoint_arn)) except ClientError as err: assert err.response['Error']['Code'] == 'ResourceNotFoundException' print(f"Confirmed that endpoint {updated_endpoint_arn} was deleted") if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument('--sagemaker-execution-role', required=True) parser.add_argument('--region', required=True) args = parser.parse_args() run_script(args.sagemaker_execution_role, args.region)
