How do I use SageMaker Inference Toolkit with Bring Your Own Container

Inference Toolkit

The SageMaker Inference Toolkit is a Python library provided by AWS that can be added to a Docker container to make the machine learning models deployable to SageMaker AI. This library implements a model serving stack that is built on Multi Model Server (MMS). The Inference Toolkit simplifies building your own Docker container and using it in SageMaker AI for inference.

The Inference Toolkit provides the necessary components and functionalities to deploy and run inference workloads for various machine learning models. Key components of the SageMaker Inference Toolkit include:

1. Serving Stack:

It is built upon Multi Model Server (MMS), enabling the serving of an individual model or multiple models from a single endpoint.

2. Inference Handler:

The Toolkit requires the implementation of an inference handler, which defines how your model processes incoming requests and generates predictions.

3. Handler Service:

A handler service, executed by the model server, is also implemented to manage the interaction between the model server and your inference handler.

4. Serving Entrypoint:

A serving entrypoint is required to start the model server and make your model accessible for inference requests.

Create a Docker image and deploy a model

The following section explains how to create your own container with the Inference Toolkit and deploy a RandomForestClassifier model. For the following steps, create all the files in the same folder.

Create Inference handler

1. Implement an inference handler, which is responsible for loading the model and providing input, predict, and output functions.

   • The inference handler must extend the DefaultInferenceHandler class which is available here.

   • The inference handler class must have the following functions:

     1. default_model_fn: Loads the model.
     2. default_input_fn: Deserializes the input data from the request.
     3. default_predict_fn: Uses the model to make predictions.
     4. default_output_fn: Serializes the prediction for the response.

   • The default_model_fn and default_predict_fn functions must be implemented, but default_input_fn and default_output_fn are optional since they are implemented in the DefaultInferenceHandler that is being extended.

   • Below is an example of an inference handler custom_inference_handler.py:

from sagemaker_inference import content_types, decoder, default_inference_handler, encoder, errors
import os
import json
import joblib
import pandas as pd
import logging

logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)

class CustomInferenceHandler(default_inference_handler.DefaultInferenceHandler):

    def default_model_fn(self, model_dir, context=None):
        """
        Load the scikit-learn model from the directory provided by SageMaker.
        The model artifact is saved as a joblib file.
        Args:
            model_dir: a directory where model is saved.
            context (obj): the request context (default: None).

        Returns: A model object.
        """
        logger.info("Loading model.")
        model = joblib.load(os.path.join(model_dir, "model.joblib"))
        return model

    def default_input_fn(self, input_data, content_type, context=None):
        """A default_input_fn responsible for deserializing the input data into an object for prediction.

        Args:
            input_data: the request payload serialized in the content_type format
            content_type: the request content_type
            context (obj): the request context (default: None).

        Returns: input_data deserialized and ready for prediction.
        """
        return decoder.decode(input_data, content_type)

    def default_predict_fn(self, data, model, context=None):
        """
        Perform predictions using the loaded model.
        Args:
            data: input data for prediction deserialized by default_input_fn
            model: model loaded in memory by default_model_fn
            context (obj): the request context (default: None).

        Returns: a prediction
        """
        logger.info("Performing prediction.")
        prediction = model.predict(data)
        return prediction

    def default_output_fn(self, prediction, accept, context=None):
        """A default_output_fn responsible for serializing the prediction result to the desired accept type.

        Args:
            prediction: a prediction result from default_predict_fn
            accept: type which the output data needs to be serialized to.
            context (obj): the request context (default: None).

        Returns: output data serialized
        """
        return encoder.encode(prediction, accept)

Note: passing context as an argument to the handler functions is optional. You can omit context from the function declaration if it is not needed at runtime.

Create Handler Service

2. Implement a handler service that is executed by the model server (MMS). It manages the interaction between the model server and the inference handler.

   • The handler service class must extend the DefaultHandlerService class which is available here.

   • If you implemented your own inference handler, you must implement the __init__ function in the handler service to pass your inference handler.

   • The handler service implements the following important functions, these functions are implemented in the DefaultHandlerService which is being extended, so it is optional to implement them:

     1. initialize: Initialize the model, this function will be called during model loading time.
     2. handle: Invoked for all incoming inference requests to the model server.

   • It is possible to only implement the handler service and not implement the inference handler. In that case, you must implement all the functionality of the inference handler in the handler service. The initialize function will load the model and handle funtion will receive the input data, make predictions using the model, and process the predictions.

   • For more information on how to define the handler service file, see the MMS custom service documentation.

   • Below is an example of a handler service model_handler.py:

from sagemaker_inference.default_handler_service import DefaultHandlerService
from sagemaker_inference.transformer import Transformer
from custom_inference_handler import CustomInferenceHandler

class HandlerService(DefaultHandlerService):
    """Handler service that is executed by the model server.
    Determines specific default inference handlers to use based on model being used.
    This class extends ``DefaultHandlerService``, which define the following:
        - The ``handle`` method is invoked for all incoming inference requests to the model server.
        - The ``initialize`` method is invoked at model server start up.
    Based on: https://github.com/awslabs/multi-model-server/blob/master/docs/custom_service.md
    """
    def __init__(self):
        transformer = Transformer(default_inference_handler=CustomInferenceHandler())
        super(HandlerService, self).__init__(transformer=transformer)

Create Serving Entrypoint

3. Implement a serving entrypoint, which starts the model server. Below is an example of serving entrypoint dockerd-entrypoint.py:

import os
import shlex
import subprocess
import sys
from subprocess import CalledProcessError
from retrying import retry
from sagemaker_inference import model_server

def _retry_if_error(exception):
    return isinstance(exception, CalledProcessError or OSError)

@retry(stop_max_delay=1000 * 50, retry_on_exception=_retry_if_error)
def _start_mms():
    model_server.start_model_server(handler_service="/home/model-server/model_handler.py:handle")

def main():
    if sys.argv[1] == "serve":
        _start_mms()
    else:
        subprocess.check_call(shlex.split(" ".join(sys.argv[1:])))
    # prevent docker exit
    subprocess.call(["tail", "-f", "/dev/null"])

main()

Create Dockerfile

4. Create a Dockerfile with the SageMaker Inference Toolkit.

• To install Inference Toolkit in the Docker image, add the following line to the Dockerfile. Note that “multi-model-server” (MMS) is required when using the Inference Toolkit:

RUN pip install multi-model-server sagemaker-inference

• Add the following two labels in the Dockerfile to advertise multi-model support on the container and to enable the container to use the SAGEMAKER_BIND_TO_PORT environment variable.

LABEL com.amazonaws.sagemaker.capabilities.multi-models=true
LABEL com.amazonaws.sagemaker.capabilities.accept-bind-to-port=true

• Copy the serving entrypoint dockerd-entrypoint.py to the /usr/local/bin/ directory inside the Dockerfile
• Create the /home/model-server/ directory inside the Dockerfile and copy the inference handler custom_inference_handler.py and the handler service model_handler.py to this directory.
• Define an entrypoint script for the Docker image using the following line:

ENTRYPOINT ["python", "/usr/local/bin/dockerd-entrypoint.py"]

• Define the command to be passed to the entrypoint script using the following line:

CMD ["serve"]

• Below is an example of a Dockerfile that uses the Inference Toolkit:

FROM ubuntu:18.04

# Set a docker label to advertise multi-model support on the container
LABEL com.amazonaws.sagemaker.capabilities.multi-models=true
# Set a docker label to enable container to use SAGEMAKER_BIND_TO_PORT environment variable if present
LABEL com.amazonaws.sagemaker.capabilities.accept-bind-to-port=true


# Upgrade installed packages
RUN apt-get update && apt-get upgrade -y && apt-get clean

# Python package management and basic dependencies
RUN apt-get install -y curl python3.7 python3.7-dev python3.7-distutils

# Register the version in alternatives
RUN update-alternatives --install /usr/bin/python python /usr/bin/python3.7 1

# Set python 3 as the default python
RUN update-alternatives --set python /usr/bin/python3.7

# Install necessary dependencies for MMS and SageMaker Inference Toolkit
RUN apt-get -y install --no-install-recommends \
    build-essential \
    ca-certificates \
    openjdk-8-jdk-headless \
    curl \
    vim \
    && rm -rf /var/lib/apt/lists/* \
    && python --version \
    && curl -O https://bootstrap.pypa.io/pip/3.7/get-pip.py \
    && python get-pip.py

RUN update-alternatives --install /usr/bin/python python /usr/bin/python3 1
RUN update-alternatives --install /usr/local/bin/pip pip /usr/local/bin/pip3 1

# Install scikit-learn, pandas, MMS, and SageMaker Inference Toolkit to set up MMS
RUN pip3 --no-cache-dir install scikit-learn pandas \
                                multi-model-server \
                                sagemaker-inference \
                                retrying

# Copy entrypoint script to the image
COPY dockerd-entrypoint.py /usr/local/bin/dockerd-entrypoint.py
RUN chmod +x /usr/local/bin/dockerd-entrypoint.py

RUN mkdir -p /home/model-server/

# Copy the default custom service file to handle incoming data and inference requests
COPY model_handler.py /home/model-server/model_handler.py
COPY custom_inference_handler.py /home/model-server/custom_inference_handler.py

# Define an entrypoint script for the docker image
ENTRYPOINT ["python", "/usr/local/bin/dockerd-entrypoint.py"]

# Define command to be passed to the entrypoint
CMD ["serve"]

Build the Docker image

5. Build the Docker image and push it to ECR.

   • Use the following code to create a shell script build_and_push.sh to build the image and push it to ECR:
   • Give permissions to execute this script using command: chmod +x build_and_push.sh
   • Run the shell script using command: ./build_and_push.sh
   • After the image is built, save the ECR URI of the image, which should look like this: <your-account-id>.dkr.ecr.<your-aws-region>.amazonaws.com/custom-inference-container

# Specify an algorithm name
algorithm_name=custom-inference-container

account=$(aws sts get-caller-identity --query Account --output text)

# Get the region defined in the current configuration (default to us-west-2 if none defined)
region=$(aws configure get region)
region=${region:-us-west-2}

fullname="${account}.dkr.ecr.${region}.amazonaws.com/${algorithm_name}:latest"

# If the repository doesn't exist in ECR, create it.

aws ecr describe-repositories --repository-names "${algorithm_name}" > /dev/null 2>&1
if [ $? -ne 0 ]
then
aws ecr create-repository --repository-name "${algorithm_name}" > /dev/null
fi

# Get the login command from ECR and execute it directly

aws ecr get-login-password --region ${region}|docker login --username AWS --password-stdin ${fullname}

# Build the docker image locally with the image name and then push it to ECR
# with the full name.

docker build -t ${algorithm_name} .
docker tag ${algorithm_name} ${fullname}

docker push ${fullname}

Deploy the SageMaker Endpoint

6. The following code trains a RandomForestClassifier and deploys it to a SageMaker endpoint using the image built above. Substitute image_uri with your ECR image URI from above:

   • For trained model to be compatible with inference image, use scikit-learn version 1.0.2

import joblib
from sklearn.ensemble import RandomForestClassifier
from sklearn.datasets import make_classification
import os
import tarfile
import boto3
import json
import sagemaker
from sagemaker.serializers import JSONSerializer
from sagemaker.model import Model

# Train a simple scikit-learn model
X, y = make_classification(n_samples=1000, n_features=4, random_state=42)
model = RandomForestClassifier(n_estimators=100)
model.fit(X, y)

# Save the model
if not os.path.exists("model"):
    os.makedirs("model")
joblib.dump(model, "model/model.joblib")

# Package the model and inference script into a tar.gz file
with tarfile.open("model.tar.gz", "w:gz") as tar:
    tar.add("model/model.joblib", arcname="model.joblib")
    
# Set up SageMaker session and role
sagemaker_session = sagemaker.Session()
role = sagemaker.get_execution_role()

# Upload your model artifact to S3
model_data_uri = sagemaker_session.upload_data(
    path="model.tar.gz", key_prefix="custom-sklearn-model"
)

# Specify the ECR image URI of your custom container
image_uri ="<your-account-id>.dkr.ecr.<your-aws-region>.amazonaws.com/<your-ecr-repo-name>:<tag>"

# Create a SageMaker Model object
model = Model(
    image_uri=image_uri,
    model_data=model_data_uri,
    role=role,
    sagemaker_session=sagemaker_session,
)

# Deploy model to a SageMaker endpoint
predictor = model.deploy(
    initial_instance_count=1,
    instance_type="ml.m5.xlarge",
    serializer=JSONSerializer()
)

7. Invoke the model, substitute ENDPOINT_NAME with the name of the endpoint created above:

# Initialize the SageMaker runtime client
sagemaker_runtime = boto3.client('sagemaker-runtime')

# Your input data
input_data = [
    [0.5, 0.2, 0.9, 0.1],
    [0.8, 0.7, 0.3, 0.6]
]

# Convert to JSON string
payload = json.dumps(input_data)

# Invoke the endpoint
response = sagemaker_runtime.invoke_endpoint(
    EndpointName='<ENDPOINT_NAME>',
    ContentType='application/json',
    Body=payload
)

# Parse the response
result = json.loads(response['Body'].read().decode())
print(result)

References:

• Default inference handler: https://github.com/aws/sagemaker-inference-toolkit/blob/master/src/sagemaker_inference/default_inference_handler.py
• Default handler service: https://github.com/aws/sagemaker-inference-toolkit/blob/master/src/sagemaker_inference/default_handler_service.py
• Custom Service (MMS): https://github.com/awslabs/multi-model-server/blob/master/docs/custom_service.md
• BYOC: https://github.com/aws/amazon-sagemaker-examples/tree/main/advanced_functionality/multi_model_bring_your_own