Add integration test

Author: bveeramani

tests/data/serverless/Dockerfile

@@ -0,0 +1,12 @@
+FROM public.ecr.aws/lambda/python:3.8
+
+COPY requirements.txt /tmp/requirements.txt
+RUN pip3 install -r /tmp/requirements.txt
+
+COPY model.py .
+RUN python model.py
+
+COPY app.py .
+COPY classes.txt .
+
+CMD ["app.handler"]

tests/data/serverless/README.md

@@ -0,0 +1,6 @@
+This folder contains the source code for the image used in the
+`sagemaker.serverless` tests.
+
+The image contains the code for a Lambda handler and all of its dependencies.
+The Lambda handler predicts the class of an image using a pre-trained ResNet-34
+model.

tests/data/serverless/app.py

@@ -0,0 +1,34 @@
+import urllib
+import json
+import os
+
+import torch
+from PIL import Image
+from torchvision import models
+from torchvision import transforms
+
+transform = transforms.Compose(
+    [
+        transforms.Resize(256),
+        transforms.CenterCrop(224),
+        transforms.ToTensor(),
+        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+    ]
+)
+
+with open("classes.txt") as file:
+    classes = [s.strip() for s in file.readlines()]
+
+
+def handler(event, context):
+    data = urllib.request.urlopen(event["url"])
+
+    image = Image.open(data)
+    image = transform(image)
+    image = image.unsqueeze(0)
+
+    model = torch.jit.load("./model.pt")
+    outputs = model(image)
+    target = outputs.argmax().item()
+
+    return {"class": classes[target]}