Add LambdaModel and LambdaPredictor

Author: bveeramani

src/sagemaker/serverless/__init__.py

@@ -0,0 +1,15 @@
+# Copyright 2017-2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+"""Classes for performing machine learning on serverless compute."""
+from sagemaker.serverless.model import LambdaModel  # noqa: F401
+from sagemaker.serverless.predictor import LambdaPredictor  # noqa: F401

src/sagemaker/serverless/model.py

@@ -0,0 +1,90 @@
+# Copyright 2017-2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+"""Models that can be deployed to serverless compute."""
+from __future__ import absolute_import
+
+import time
+from typing import Optional
+
+import boto3
+import botocore
+
+from sagemaker.model import ModelBase
+
+from .predictor import LambdaPredictor
+
+
+class LambdaModel(ModelBase):
+    """A model that can be deployed to Lambda."""
+
+    def __init__(
+        self, image_uri: str, role: str, client: Optional[botocore.client.BaseClient] = None
+    ) -> None:
+        """Initialize instance attributes.
+
+        Arguments:
+            image_uri: URI of a container image in the Amazon ECR registry. The image
+                should contain a handler that performs inference.
+            role: The Amazon Resource Name (ARN) of the IAM role that Lambda will assume
+                when it performs inference
+            client: The Lambda client used to interact with Lambda.
+        """
+        self._client = client or boto3.client("lambda")
+        self._image_uri = image_uri
+        self._role = role
+
+    def deploy(
+        self, function_name: str, timeout: int, memory_size: int, wait: bool = False
+    ) -> LambdaPredictor:
+        """Create a Lambda function using the image specified in the constructor.
+
+        Arguments:
+            function_name: The name of the function.
+            timeout: The number of seconds that the function can run for before being terminated.
+            memory_size: The amount of memory in MB that the function has access to.
+            wait: If true, wait until the deployment completes (default: True).
+
+        Returns:
+            A LambdaPredictor instance that performs inference using the specified image.
+        """
+        response = self._client.create_function(
+            FunctionName=function_name,
+            PackageType="Image",
+            Role=self._role,
+            Code={
+                "ImageUri": self._image_uri,
+            },
+            Timeout=timeout,
+            MemorySize=memory_size,
+        )
+
+        if not wait:
+            return LambdaPredictor(function_name, client=self._client)
+
+        # Poll function state.
+        polling_interval = 5
+        while response["State"] == "Pending":
+            time.sleep(polling_interval)
+            response = self._client.get_function_configuration(FunctionName=function_name)
+
+        if response["State"] != "Active":
+            raise RuntimeError("Failed to deploy model to Lambda: %s" % response["StateReason"])
+
+        return LambdaPredictor(function_name, client=self._client)
+
+    def destroy(self) -> None:
+        """Destroy resources associated with this model.
+
+        This method does not delete the image specified in the constructor. As
+        a result, this method is a no-op.
+        """

src/sagemaker/serverless/predictor.py

@@ -0,0 +1,81 @@
+# Copyright 2017-2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+"""Predictors that are hosted on serverless compute."""
+from __future__ import absolute_import
+
+from typing import Optional, Tuple
+
+import boto3
+import botocore
+
+from sagemaker import deserializers, serializers
+from sagemaker.predictor import PredictorBase
+
+
+class LambdaPredictor(PredictorBase):
+    """A deployed model hosted on Lambda."""
+
+    def __init__(
+        self, function_name: str, client: Optional[botocore.client.BaseClient] = None
+    ) -> None:
+        """Initialize instance attributes.
+
+        Arguments:
+            function_name: The name of the function.
+            client: The Lambda client used to interact with Lambda.
+        """
+        self._client = client or boto3.client("lambda")
+        self._function_name = function_name
+        self._serializer = serializers.JSONSerializer()
+        self._deserializer = deserializers.JSONDeserializer()
+
+    def predict(self, data: dict) -> dict:
+        """Invoke the Lambda function specified in the constructor.
+
+        This function is synchronous. It will only return after the function
+        has produced a prediction.
+
+        Arguments:
+            data: The data sent to the Lambda function as input.
+
+        Returns:
+            The data returned by the Lambda function.
+        """
+        response = self._client.invoke(
+            FunctionName=self._function_name,
+            InvocationType="RequestResponse",
+            Payload=self._serializer.serialize(data),
+        )
+        return self._deserializer.deserialize(
+            response["Payload"],
+            response["ResponseMetadata"]["HTTPHeaders"]["content-type"],
+        )
+
+    def destroy(self) -> None:
+        """Destroy the Lambda function specified in the constructor."""
+        self._client.delete_function(FunctionName=self._function_name)
+
+    @property
+    def content_type(self) -> str:
+        """The MIME type of the data sent to the Lambda function."""
+        return self._serializer.CONTENT_TYPE
+
+    @property
+    def accept(self) -> Tuple[str]:
+        """The content type(s) that are expected from the Lambda function."""
+        return self._deserializer.ACCEPT
+
+    @property
+    def function_name(self) -> str:
+        """The name of the Lambda function this predictor invokes."""
+        return self._function_name