add SageMaker inference example (#152)

Author: silv-io

sagemaker-inference/ReadMe.md

@@ -0,0 +1,9 @@
+# SageMaker Model Inference
+
+This is a small example about how you can use LocalStack to host your PyTorch ML models.
+
+Before using this example you should setup your Docker Client to pull the AWS Deep Learning images ([more info here](https://github.com/aws/deep-learning-containers/blob/master/available_images.md)):
+
+```bash
+    aws ecr get-login-password --region us-east-1 | docker login --username AWS --password-stdin 763104351884.dkr.ecr.us-east-1.amazonaws.com
+```

sagemaker-inference/data/mnist/t10k-images-idx3-ubyte.gz

@@ -0,0 +1,95 @@
+import json
+import random
+import time
+
+import boto3
+import httpx
+import numpy as np
+from mypy_boto3_s3 import S3Client
+from mypy_boto3_sagemaker import SageMakerClient
+from mypy_boto3_sagemaker_runtime import SageMakerRuntimeClient
+
+from mnist import mnist_to_numpy, normalize
+
+LOCALSTACK_ENDPOINT = "http://localhost.localstack.cloud:4566"
+MODEL_BUCKET = "models"
+MODEL_TAR = "./data/model.tar.gz"
+MODEL_NAME = "sample"
+CONFIG_NAME = "sample-cf"
+ENDPOINT_NAME = "sample-ep"
+CONTAINER_IMAGE = "763104351884.dkr.ecr.us-east-1.amazonaws.com/pytorch-inference:1.5.0-cpu-py3"
+EXECUTION_ROLE_ARN = "arn:aws:iam::0000000000000:role/sagemaker-role"
+
+sagemaker: SageMakerClient = boto3.client("sagemaker", endpoint_url=LOCALSTACK_ENDPOINT)
+sagemaker_runtime: SageMakerRuntimeClient = boto3.client("sagemaker-runtime", endpoint_url=LOCALSTACK_ENDPOINT)
+s3: S3Client = boto3.client("s3", endpoint_url=LOCALSTACK_ENDPOINT)
+
+
+def deploy_model(run_id: str = "0"):
+    # Put the Model into the correct bucket
+    s3.create_bucket(Bucket=f"{MODEL_BUCKET}-{run_id}")
+    s3.upload_file(MODEL_TAR, f"{MODEL_BUCKET}-{run_id}", f"{MODEL_NAME}.tar.gz")
+
+    # Create the model in sagemaker
+    sagemaker.create_model(ModelName=f"{MODEL_NAME}-{run_id}", ExecutionRoleArn=EXECUTION_ROLE_ARN,
+                           PrimaryContainer={"Image": CONTAINER_IMAGE,
+                                             "ModelDataUrl": f"s3://{MODEL_BUCKET}-{run_id}/{MODEL_NAME}.tar.gz"})
+    sagemaker.create_endpoint_config(EndpointConfigName=f"{CONFIG_NAME}-{run_id}", ProductionVariants=[{
+        "VariantName": f"var-{run_id}", "ModelName": f"{MODEL_NAME}-{run_id}", "InitialInstanceCount": 1,
+        "InstanceType": "ml.m5.large"
+    }])
+    sagemaker.create_endpoint(EndpointName=f"{ENDPOINT_NAME}-{run_id}", EndpointConfigName=f"{CONFIG_NAME}-{run_id}")
+
+
+def _get_input_dict():
+    X, Y = mnist_to_numpy("data/mnist", train=False)
+    mask = random.sample(range(X.shape[0]), 2)
+    samples = X[mask]
+
+    samples = normalize(samples.astype(np.float32), axis=(1, 2))
+    return {
+        "inputs": np.expand_dims(samples, axis=1).tolist()
+    }
+
+
+def _show_predictions(response):
+    predictions = np.argmax(np.array(response, dtype=np.float32), axis=1).tolist()
+    print(f"Predicted digits: {predictions}")
+
+
+def inference_model_container(run_id: str = "0"):
+    ep = sagemaker.describe_endpoint(EndpointName=f"{ENDPOINT_NAME}-{run_id}")
+    arn = ep["EndpointArn"]
+    tag_list = sagemaker.list_tags(ResourceArn=arn)
+    port = "4510"
+    for tag in tag_list["Tags"]:
+        if tag["Key"] == "_LS_ENDPOINT_PORT_":
+            port = tag["Value"]
+    inputs = _get_input_dict()
+    response = httpx.post(f"http://localhost.localstack.cloud:{port}/invocations", json=inputs,
+                          headers={"Content-Type": "application/json", "Accept": "application/json"})
+    _show_predictions(json.loads(response.text))
+
+
+def inference_model_boto3(run_id: str = "0"):
+    inputs = _get_input_dict()
+    response = sagemaker_runtime.invoke_endpoint(EndpointName=f"{ENDPOINT_NAME}-{run_id}", Body=json.dumps(inputs),
+                                                 Accept="application/json",
+                                                 ContentType="application/json")
+    _show_predictions(json.loads(response["Body"].read()))
+
+
+def _short_uid():
+    import uuid
+
+    return str(uuid.uuid4())[:8]
+
+
+if __name__ == '__main__':
+    test_run = _short_uid()
+    deploy_model(test_run)
+    # wait some time to avoid connection resets in log output
+    # -> not essential as the container spins up quickly enough within the retries of boto
+    time.sleep(2)
+    inference_model_boto3(test_run)
+    inference_model_container(test_run)

sagemaker-inference/data/mnist/t10k-labels-idx1-ubyte.gz

@@ -0,0 +1,97 @@
+import gzip
+import os
+
+import boto3
+import numpy as np
+
+dirname = os.path.dirname(os.path.abspath(__file__))
+
+
+def mnist_to_numpy(data_dir="/tmp/data", train=True):
+    """Download MNIST dataset and convert it to numpy array
+
+    Args:
+        data_dir (str): directory to save the data
+        train (bool): download training set
+
+    Returns:
+        tuple of images and labels as numpy arrays
+    """
+
+    if not os.path.exists(data_dir):
+        os.makedirs(data_dir)
+
+    if train:
+        images_file = "train-images-idx3-ubyte.gz"
+        labels_file = "train-labels-idx1-ubyte.gz"
+    else:
+        images_file = "t10k-images-idx3-ubyte.gz"
+        labels_file = "t10k-labels-idx1-ubyte.gz"
+
+    # download objects
+    s3 = boto3.client("s3")
+    bucket = "sagemaker-sample-files"
+    for obj in [images_file, labels_file]:
+        key = os.path.join("datasets/image/MNIST", obj)
+        dest = os.path.join(data_dir, obj)
+        if not os.path.exists(dest):
+            s3.download_file(bucket, key, dest)
+
+    return _convert_to_numpy(data_dir, images_file, labels_file)
+
+
+def _convert_to_numpy(data_dir, images_file, labels_file):
+    """Byte string to numpy arrays"""
+    with gzip.open(os.path.join(data_dir, images_file), "rb") as f:
+        images = np.frombuffer(f.read(), np.uint8, offset=16).reshape(-1, 28, 28)
+
+    with gzip.open(os.path.join(data_dir, labels_file), "rb") as f:
+        labels = np.frombuffer(f.read(), np.uint8, offset=8)
+
+    return (images, labels)
+
+
+def normalize(x, axis):
+    eps = np.finfo(float).eps
+
+    mean = np.mean(x, axis=axis, keepdims=True)
+    # avoid division by zero
+    std = np.std(x, axis=axis, keepdims=True) + eps
+    return (x - mean) / std
+
+
+def adjust_to_framework(x, framework="pytorch"):
+    """Adjust a ``numpy.ndarray`` to be used as input for specified framework
+
+    Args:
+        x (numpy.ndarray): Batch of images to be adjusted
+            to follow the convention in pytorch / tensorflow / mxnet
+
+        framework (str): Framework to use. Takes value in
+            ``pytorch``, ``tensorflow`` or ``mxnet``
+    Return:
+        numpy.ndarray following the convention of tensors in the given
+        framework
+    """
+
+    if x.ndim == 3:
+        # input is gray-scale
+        x = np.expand_dims(x, 1)
+
+    if framework in ["pytorch", "mxnet"]:
+        # depth-major
+        return x
+    elif framework == "tensorlfow":
+        # depth-minor
+        return np.transpose(x, (0, 2, 3, 1))
+    elif framework == "mxnet":
+        return x
+    else:
+        raise ValueError(
+            "framework must be one of " + "[pytorch, tensorflow, mxnet], got {}".format(framework)
+        )
+
+
+if __name__ == "__main__":
+    X, Y = mnist_to_numpy()
+    X, Y = X.astype(np.float32), Y.astype(np.int8)

sagemaker-inference/data/model.tar.gz

@@ -0,0 +1,25 @@
+anyio==3.6.1
+boto3==1.24.85
+boto3-stubs==1.24.85
+botocore==1.27.85
+botocore-stubs==1.27.85
+certifi==2022.9.24
+charset-normalizer==2.1.1
+h11==0.12.0
+httpcore==0.15.0
+httpx==0.23.0
+idna==3.4
+jmespath==1.0.1
+mypy-boto3-s3==1.24.76
+mypy-boto3-sagemaker==1.24.84
+mypy-boto3-sagemaker-runtime==1.24.84
+numpy==1.23.3
+python-dateutil==2.8.2
+rfc3986==1.5.0
+s3transfer==0.6.0
+six==1.16.0
+sniffio==1.3.0
+types-awscrt==0.14.6
+types-s3transfer==0.6.0.post4
+typing_extensions==4.3.0
+urllib3==1.26.12