def create_infra()

in right_size_your_sagemaker_endpoints/api_helper.py [0:0]
130 lines of code
2 McCabe index (conditional complexity)

def create_infra(project_name, account_id, region):
    """
    Creates AWS Lambda and API Gateway for load testing.
    
    This function is called by the notebook to create an AWS Lambda function
    and an API gateway endpoint that we use to perform the load testing. 
    
    Inputs:
    project_name - unique identifier that's tagged to all resources
    account_id - notebook users's AWS account ID
    region - AWS region to deploy the resources in.
    
    Output:
    API Gateway endpoint URL.
    """    
    ## SETUP THE ROLES AND POLICIES
    
    # Create role that can be assumed by Lambda
    trust_policy = {
      "Version": "2012-10-17",
      "Statement": [
        {
          "Effect": "Allow",
          "Principal": {
            "Service": "lambda.amazonaws.com"
          },
          "Action": "sts:AssumeRole"
        }
      ]
    }

    role = iam_client.create_role(
        RoleName=f"LambdaRoleToInvokeSageMakerEndpoint{project_name}",
        AssumeRolePolicyDocument=json.dumps(trust_policy),
        Description="Role for Lambda function to invoke SageMaker endpoints",
    )

    # Create policy for the role, allowing Lambda to invoke SageMaker endpoint
    policy_doc = {
        "Version": "2012-10-17",
        "Statement": [
            {
                "Action": [
                    "logs:CreateLogGroup",
                    "logs:CreateLogStream",
                    "logs:PutLogEvents",
                    "sagemaker:InvokeEndpoint"
                ],
                "Resource": [
                    "arn:aws:logs:*:*:*",
                    "arn:aws:sagemaker:*:*:*"
                ],
                "Effect": "Allow"
            }
        ]
    }

    policy = iam_client.create_policy(
        PolicyName=f"LambdaSageMakerAccessPolicy{project_name}",
        PolicyDocument=json.dumps(policy_doc)
    )

    # Attach policy to role
    response = iam_resource.Role(role['Role']['RoleName']).attach_policy(PolicyArn=policy['Policy']['Arn'])
    
    LambdaRoleArn = role['Role']['Arn']
    
    time.sleep(10)
    
    ## SETUP AND DEPLOY LAMBDA FUNCTION
    
    # Zip the request prediction code
    with ZipFile('lambda_function.zip', 'w') as zip:
        zip.write('lambda_index.py')
    
    # Create the lambda
    response = lambda_client.create_function(
        FunctionName=f'request-predictions-{project_name}',
        Runtime='python3.8',
        Role=LambdaRoleArn,
        Handler='lambda_index.lambda_handler',
        Code={
            'ZipFile': open('./lambda_function.zip', 'rb').read()
        },
        Description='Function to invoke endpoint and return predictions',
        Timeout=60,
        MemorySize=128
    )
    
    lambda_arn = response['FunctionArn']
    lambda_name = response['FunctionName']
    
    # Create rest api
    rest_api = api_client.create_rest_api(
        name=f'ImageClassifier-{project_name}'
    )
    rest_api_id = rest_api["id"]

    # Get the rest api's root id
    root_resource_id = api_client.get_resources(
        restApiId=rest_api_id
    )['items'][0]['id']

    # Create an api resource
    api_resource = api_client.create_resource(
        restApiId=rest_api_id,
        parentId=root_resource_id,
        pathPart='ImageClassifier'
    )

    api_resource_id = api_resource['id']

    # Add a post method to the rest api resource
    api_method = api_client.put_method(
        restApiId=rest_api_id,
        resourceId=api_resource_id,
        httpMethod='POST',
        authorizationType='NONE'
    )

    lambda_uri = f"arn:aws:apigateway:{region}:lambda:path/2015-03-31/functions/{lambda_arn}/invocations"
    
    
    # Add integrations for mapping SageMaker and Lambda HTTP response codes to API gateway HTTP response codes
    put_integration = api_client.put_integration(
        restApiId=rest_api_id,
        resourceId=api_resource_id,
        httpMethod='POST',
        type='AWS',
        integrationHttpMethod='POST',
        uri=lambda_uri)

    put_method_res = api_client.put_method_response(
        restApiId=rest_api_id,
        resourceId=api_resource_id,
        httpMethod='POST',
        statusCode='200',
        responseModels={'application/json': 'Empty'}
    )

    put_integration_response = api_client.put_integration_response(
        restApiId=rest_api_id,
        resourceId=api_resource_id,
        httpMethod='POST',
        statusCode='200',
        responseTemplates={"application/json": ""}
    )

    put_integration_response = api_client.put_integration_response(
        restApiId=rest_api_id,
        resourceId=api_resource_id,
        httpMethod='POST',
        statusCode='400',
        selectionPattern='Invalid*',
        responseTemplates={"application/json": ""}
    )
        
    put_integration_response = api_client.put_integration_response(
        restApiId=rest_api_id,
        resourceId=api_resource_id,
        httpMethod='POST',
        statusCode='500',
        selectionPattern='Internal*',
        responseTemplates={"application/json": ""}
    )
    
    # Deploy to a stage
    deployment = api_client.create_deployment(
        restApiId=rest_api_id,
        stageName='dev',
      )

    # API gateway needs permissions to invoke the lambda function
    lambda_api_response = lambda_client.add_permission(
        FunctionName=lambda_name,
        StatementId='api-gateway-invoke',
        Action='lambda:InvokeFunction',
        Principal='apigateway.amazonaws.com',
        SourceArn=f"arn:aws:execute-api:{region}:{account_id}:{rest_api_id}/*/POST/ImageClassifier"
    )
    
    api_url = f"https://{rest_api_id}.execute-api.{region}.amazonaws.com/dev/ImageClassifier"
    
    print("API GATEWAY URL: url = "+ f"{api_url}")
    
    return api_url