<reponame>cswarth/whylogs<filename>examples/sagemaker_container_example/predictor.py
import datetime
import json
import os

import boto3
import botocore
import flask
import my_model as my_model_code
import numpy as np
import torch
import torchvision.transforms as transforms
from PIL import Image

from whylogs import get_or_create_session

prefix = "/opt/ml/"
# location of you checkpoint in sagemaker container
model_path = os.path.join(prefix, "model")


whylogs_session = get_or_create_session()
logger = whylogs_session.logger(dataset_name="my_deployed_model", dataset_timestamp=datetime.datetime.now(datetime.timezone.utc), with_rotation_time="5m")

# simple custom metric for embedding
embedding_centr = np.linalg.norm(np.ones(1024))


# loads the model into memory from disk and returns it
def model_fn(model_dir):

    model_config = {"backbone": "resnet", "outlayer": "C5", "model_path": os.path.join(model_dir, "my_checkpoint.pth.tar")}
    print("Model Config: {}", format(model_config))
    model = my_model_code.Image_Embeddings(model_config)

    try:
        # try to use gpu if available (there are prob faster ways to check)
        model.to(torch.device("cuda"))
        logger.log({"arch": "cuda"})
    except Exception:
        logger.log({"arch": "cpu"})
        return None

    return model


# pre-processing inputs
def input_fn(request_body):

    # Fetch the input data... assuming your data is in a bucket somewhere
    # it could be also be part of the request as a base64.
    request = request_body

    BUCKET_NAME = request["bucket_name"]
    KEY = os.path.join(request["image_path"])

    # whylog track file paths
    logger.log({"bucket": BUCKET_NAME})
    logger.log({"image_path": KEY})

    s3 = boto3.resource("s3")

    try:
        s3.Bucket(BUCKET_NAME).download_file(KEY, "temp_image.png")
    except botocore.exceptions.ClientError as e:
        if e.response["Error"]["Code"] == "404":
            print("The data does not exist.")
        else:
            raise
    # open image byte data and transform as necessary
    image_ = Image.open("temp_image.png").convert("RGB")

    # logger image
    logger.log_image(image_)

    resize = transforms.Resize(300)
    loader = transforms.Compose([resize, transforms.ToTensor()])
    img = loader(image_).float().unsqueeze(0)

    return img


# predicting the output
def predict_fn(input_object, model):

    # inference call
    try:
        # try to load data to the gpu
        input_object = input_object.to(torch.device("cuda"))
    except Exception:
        pass

    embedding = model(input_object).cpu().detach().numpy().flatten().tolist()

    return embedding


model = model_fn(model_path)

# The flask app for serving predictions
app = flask.Flask(__name__)


@app.route("/ping", methods=["GET"])
def ping():
    """Determine if the container is working and healthy. In this sample container, we declare
    it healthy if we can load the model successfully."""

    status = 200 if model else 404
    return flask.Response(response="\n", status=status, mimetype="application/json")


@app.route("/invocations", methods=["POST"])
def transformation():
    """Do an inference on a single batch of data."""

    if flask.request.content_type == "application/json":
        # load and preprocess image

        image = input_fn(flask.request.json)

    else:
        return flask.Response(
            response="This predictor only supports JSON data",
            status=415,
            mimetype="text/plain",
        )

    # embeded image
    embedding = predict_fn(image, model)
    logger.log({"distance_from_cntr": np.dot(embedding, embedding_centr)})
    result = {}
    result["embedding"] = embedding

    return flask.Response(response=json.dumps(result), status=200, mimetype="application/json")

import datetime
import json
import os

import boto3
import botocore
import flask
import my_model as my_model_code
import numpy as np
import torch
import torchvision.transforms as transforms
from PIL import Image

from whylogs import get_or_create_session

prefix = "/opt/ml/"
# location of you checkpoint in sagemaker container
model_path = os.path.join(prefix, "model")


whylogs_session = get_or_create_session()
logger = whylogs_session.logger(dataset_name="my_deployed_model", dataset_timestamp=datetime.datetime.now(datetime.timezone.utc), with_rotation_time="5m")

# simple custom metric for embedding
embedding_centr = np.linalg.norm(np.ones(1024))


# loads the model into memory from disk and returns it
def model_fn(model_dir):

    model_config = {"backbone": "resnet", "outlayer": "C5", "model_path": os.path.join(model_dir, "my_checkpoint.pth.tar")}
    print("Model Config: {}", format(model_config))
    model = my_model_code.Image_Embeddings(model_config)

    try:
        # try to use gpu if available (there are prob faster ways to check)
        model.to(torch.device("cuda"))
        logger.log({"arch": "cuda"})
    except Exception:
        logger.log({"arch": "cpu"})
        return None

    return model


# pre-processing inputs
def input_fn(request_body):

    # Fetch the input data... assuming your data is in a bucket somewhere
    # it could be also be part of the request as a base64.
    request = request_body

    BUCKET_NAME = request["bucket_name"]
    KEY = os.path.join(request["image_path"])

    # whylog track file paths
    logger.log({"bucket": BUCKET_NAME})
    logger.log({"image_path": KEY})

    s3 = boto3.resource("s3")

    try:
        s3.Bucket(BUCKET_NAME).download_file(KEY, "temp_image.png")
    except botocore.exceptions.ClientError as e:
        if e.response["Error"]["Code"] == "404":
            print("The data does not exist.")
        else:
            raise
    # open image byte data and transform as necessary
    image_ = Image.open("temp_image.png").convert("RGB")

    # logger image
    logger.log_image(image_)

    resize = transforms.Resize(300)
    loader = transforms.Compose([resize, transforms.ToTensor()])
    img = loader(image_).float().unsqueeze(0)

    return img


# predicting the output
def predict_fn(input_object, model):

    # inference call
    try:
        # try to load data to the gpu
        input_object = input_object.to(torch.device("cuda"))
    except Exception:
        pass

    embedding = model(input_object).cpu().detach().numpy().flatten().tolist()

    return embedding


model = model_fn(model_path)

# The flask app for serving predictions
app = flask.Flask(__name__)


@app.route("/ping", methods=["GET"])
def ping():
    """Determine if the container is working and healthy. In this sample container, we declare
    it healthy if we can load the model successfully."""

    status = 200 if model else 404
    return flask.Response(response="\n", status=status, mimetype="application/json")


@app.route("/invocations", methods=["POST"])
def transformation():
    """Do an inference on a single batch of data."""

    if flask.request.content_type == "application/json":
        # load and preprocess image

        image = input_fn(flask.request.json)

    else:
        return flask.Response(
            response="This predictor only supports JSON data",
            status=415,
            mimetype="text/plain",
        )

    # embeded image
    embedding = predict_fn(image, model)
    logger.log({"distance_from_cntr": np.dot(embedding, embedding_centr)})
    result = {}
    result["embedding"] = embedding

    return flask.Response(response=json.dumps(result), status=200, mimetype="application/json")
# location of you checkpoint in sagemaker container # simple custom metric for embedding # loads the model into memory from disk and returns it # try to use gpu if available (there are prob faster ways to check) # pre-processing inputs # Fetch the input data... assuming your data is in a bucket somewhere # it could be also be part of the request as a base64. # whylog track file paths # open image byte data and transform as necessary # logger image # predicting the output # inference call # try to load data to the gpu # The flask app for serving predictions Determine if the container is working and healthy. In this sample container, we declare     it healthy if we can load the model successfully. Do an inference on a single batch of data. # load and preprocess image # embeded image