Deploy endpoint in machine learning pipeline which will be integrated with AWS Step Functions

ml-pipeline-deploy-lambda.py

import time
import json
import boto3

def create_endpoint(event, job_name_prefix='machine-learning-pipeline'):
    sagemaker_client = boto3.client('sagemaker')
    
    # create model
    timestamp = time.strftime('-%Y-%m-%d-%H-%M-%S', time.gmtime())
    model_name= job_name_prefix + timestamp
    print(model_name)
    model_data = event['jobInfo']['ModelArtifacts']['S3ModelArtifacts']
    print(model_data)
    training_image = event['jobInfo']['AlgorithmSpecification']['TrainingImage']
    print(training_image)
    role = event['jobInfo']['RoleArn']
    
    primary_container = {
        'Image': training_image,
        'ModelDataUrl': model_data,
    }
    
    create_model_response = sagemaker_client.create_model(
        ModelName = model_name,
        ExecutionRoleArn = role,
        PrimaryContainer = primary_container)
    
    print(create_model_response['ModelArn'])
    
    # create endpoint config
    timestamp = time.strftime('-%Y-%m-%d-%H-%M-%S', time.gmtime())
    endpoint_config_name = job_name_prefix + '-epc' + timestamp
    
    endpoint_config_response = sagemaker_client.create_endpoint_config(
        EndpointConfigName = endpoint_config_name,
        ProductionVariants=[{
            'InstanceType':'ml.m4.xlarge',
            'InitialInstanceCount':1,
            'ModelName':model_name,
            'VariantName':'AllTraffic'}])
    
    print('Endpoint configuration name: {}'.format(endpoint_config_name))
    print('Endpoint configuration arn:  {}'.format(endpoint_config_response['EndpointConfigArn']))
    
    # create endpoint 
    timestamp = time.strftime('-%Y-%m-%d-%H-%M-%S', time.gmtime())
    endpoint_name = job_name_prefix + '-ep' + timestamp
    print('Endpoint name: {}'.format(endpoint_name))
    
    endpoint_params = {
        'EndpointName': endpoint_name,
        'EndpointConfigName': endpoint_config_name,
    }
    endpoint_response = sagemaker_client.create_endpoint(**endpoint_params)
    print('EndpointArn = {}'.format(endpoint_response['EndpointArn']))
    
    return endpoint_response

def lambda_handler(event, context):
    
    endpoint_response = create_endpoint(event)
    
    return {
        'statusCode': 200,
        'body': json.dumps(endpoint_response)
    }