feature: support Multi-Model endpoints

* feature: support Multi-Model endpoints

* fix: formatting for the Multi-Model endpoint files

* fix: enable local file uploads with add_model, use model parameters during multi-model deploy

* fix: patch multidatamodel unit test to work without default region env variable and add docker start cmd to buildspec

* fix: add validation to MultiDataModel constructor, fix integ test to enable runs in parallel

Author: dlragha

buildspec.yml

@@ -1,6 +1,10 @@
 version: 0.2
 
 phases:
+  pre_build:
+    commands:
+      - start-dockerd
+
   build:
     commands:
       - IGNORE_COVERAGE=-

src/sagemaker/multidatamodel.py

@@ -0,0 +1,315 @@
+# Copyright 2019 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.
+"""This module contains code to create and manage SageMaker ``MultiDataModel``"""
+from __future__ import absolute_import
+
+import os
+from six.moves.urllib.parse import urlparse
+
+import sagemaker
+from sagemaker import s3
+from sagemaker.model import Model
+from sagemaker.session import Session
+
+MULTI_MODEL_CONTAINER_MODE = "MultiModel"
+
+
+class MultiDataModel(Model):
+    """A SageMaker ``MultiDataModel`` that can be used to deploy multiple models to the same
+    SageMaker ``Endpoint``, and also deploy additional models to an existing SageMaker
+    multi-model ``Endpoint``
+    """
+
+    def __init__(
+        self,
+        name,
+        model_data_prefix,
+        model=None,
+        image=None,
+        role=None,
+        sagemaker_session=None,
+        **kwargs
+    ):
+        """Initialize a ``MultiDataModel``. In addition to these arguments, it supports all
+           arguments supported by ``Model`` constructor
+
+        Args:
+            name (str): The model name.
+            model_data_prefix (str): The S3 prefix where all the models artifacts (.tar.gz)
+                in a Multi-Model endpoint are located
+            model (sagemaker.Model): The Model object that would define the
+                SageMaker model attributes like vpc_config, predictors, etc.
+                If this is present, the attributes from this model are used when
+                deploying the ``MultiDataModel``.  Parameters 'image', 'role' and 'kwargs'
+                are not permitted when model parameter is set.
+            image (str): A Docker image URI. It can be null if the 'model' parameter
+                is passed to during ``MultiDataModel`` initialization (default: None)
+            role (str): An AWS IAM role (either name or full ARN). The Amazon
+                SageMaker training jobs and APIs that create Amazon SageMaker
+                endpoints use this role to access training data and model
+                artifacts. After the endpoint is created, the inference code
+                might use the IAM role if it needs to access some AWS resources.
+                It can be null if this is being used to create a Model to pass
+                to a ``PipelineModel`` which has its own Role field or if the 'model' parameter
+                is passed to during ``MultiDataModel`` initialization (default: None)
+            sagemaker_session (sagemaker.session.Session): A SageMaker Session
+                object, used for SageMaker interactions (default: None). If not
+                specified, one is created using the default AWS configuration
+                chain.
+            **kwargs: Keyword arguments passed to the ``Model`` initializer.
+        """
+        # Validate path
+        if not model_data_prefix.startswith("s3://"):
+            raise ValueError(
+                'Expecting S3 model prefix beginning with "s3://". Received: "{}"'.format(
+                    model_data_prefix
+                )
+            )
+
+        if model and (image or role or kwargs):
+            raise ValueError(
+                "Parameters image, role or kwargs are not permitted when model parameter is passed."
+            )
+
+        self.name = name
+        self.model_data_prefix = model_data_prefix
+        self.model = model
+        self.container_mode = MULTI_MODEL_CONTAINER_MODE
+        self.sagemaker_session = sagemaker_session or Session()
+        self.s3_client = self.sagemaker_session.boto_session.client("s3")
+
+        # Set the ``Model`` parameters if the model parameter is not specified
+        if not self.model:
+            super(MultiDataModel, self).__init__(
+                self.model_data_prefix,
+                image,
+                role,
+                name=self.name,
+                sagemaker_session=self.sagemaker_session,
+                **kwargs
+            )
+
+    def prepare_container_def(self, instance_type, accelerator_type=None):
+        """Return a container definition set with MultiModel mode,
+        model data and other parameters from the model (if available).
+
+        Subclasses can override this to provide custom container definitions
+        for deployment to a specific instance type. Called by ``deploy()``.
+
+        Returns:
+            dict[str, str]: A complete container definition object usable with the CreateModel API
+        """
+        # Copy the trained model's image and environment variables if they exist. Models trained
+        # with FrameworkEstimator set framework specific environment variables which need to be
+        # copied over
+        if self.model:
+            container_definition = self.model.prepare_container_def(instance_type, accelerator_type)
+            image = container_definition["Image"]
+            environment = container_definition["Environment"]
+        else:
+            image = self.image
+            environment = self.env
+        return sagemaker.container_def(
+            image,
+            env=environment,
+            model_data_url=self.model_data_prefix,
+            container_mode=self.container_mode,
+        )
+
+    def deploy(
+        self,
+        initial_instance_count,
+        instance_type,
+        accelerator_type=None,
+        endpoint_name=None,
+        update_endpoint=False,
+        tags=None,
+        kms_key=None,
+        wait=True,
+        data_capture_config=None,
+    ):
+        """Deploy this ``Model`` to an ``Endpoint`` and optionally return a ``Predictor``.
+
+        Create a SageMaker ``Model`` and ``EndpointConfig``, and deploy an
+        ``Endpoint`` from this ``Model``. If self.model is not None, then the ``Endpoint``
+        will be deployed with parameters in self.model (like vpc_config,
+        enable_network_isolation, etc).  If self.model is None, then use the parameters
+        in ``MultiDataModel`` constructor will be used. If ``self.predictor_cls`` is not
+        None, this method returns a the result of invoking ``self.predictor_cls`` on
+        the created endpoint name.
+
+        The name of the created model is accessible in the ``name`` field of
+        this ``Model`` after deploy returns
+
+        The name of the created endpoint is accessible in the
+        ``endpoint_name`` field of this ``Model`` after deploy returns.
+
+        Args:
+            initial_instance_count (int): The initial number of instances to run
+                in the ``Endpoint`` created from this ``Model``.
+            instance_type (str): The EC2 instance type to deploy this Model to.
+                For example, 'ml.p2.xlarge', or 'local' for local mode.
+            accelerator_type (str): Type of Elastic Inference accelerator to
+                deploy this model for model loading and inference, for example,
+                'ml.eia1.medium'. If not specified, no Elastic Inference
+                accelerator will be attached to the endpoint. For more
+                information:
+                https://docs.aws.amazon.com/sagemaker/latest/dg/ei.html
+            endpoint_name (str): The name of the endpoint to create (default:
+                None). If not specified, a unique endpoint name will be created.
+            update_endpoint (bool): Flag to update the model in an existing
+                Amazon SageMaker endpoint. If True, this will deploy a new
+                EndpointConfig to an already existing endpoint and delete
+                resources corresponding to the previous EndpointConfig. If
+                False, a new endpoint will be created. Default: False
+            tags (List[dict[str, str]]): The list of tags to attach to this
+                specific endpoint.
+            kms_key (str): The ARN of the KMS key that is used to encrypt the
+                data on the storage volume attached to the instance hosting the
+                endpoint.
+            wait (bool): Whether the call should wait until the deployment of
+                this model completes (default: True).
+            data_capture_config (sagemaker.model_monitor.DataCaptureConfig): Specifies
+                configuration related to Endpoint data capture for use with
+                Amazon SageMaker Model Monitoring. Default: None.
+
+        Returns:
+            callable[string, sagemaker.session.Session] or None: Invocation of
+                ``self.predictor_cls`` on the created endpoint name,
+                if ``self.predictor_cls``
+                is not None. Otherwise, return None.
+        """
+        # Set model specific parameters
+        if self.model:
+            enable_network_isolation = self.model.enable_network_isolation()
+            role = self.model.role
+            vpc_config = self.model.vpc_config
+            predictor = self.model.predictor_cls
+        else:
+            enable_network_isolation = self.enable_network_isolation()
+            role = self.role
+            vpc_config = self.vpc_config
+            predictor = self.predictor_cls
+
+        if role is None:
+            raise ValueError("Role can not be null for deploying a model")
+
+        container_def = self.prepare_container_def(instance_type, accelerator_type=accelerator_type)
+        self.sagemaker_session.create_model(
+            self.name,
+            role,
+            container_def,
+            vpc_config=vpc_config,
+            enable_network_isolation=enable_network_isolation,
+            tags=tags,
+        )
+
+        production_variant = sagemaker.production_variant(
+            self.name, instance_type, initial_instance_count, accelerator_type=accelerator_type
+        )
+        if endpoint_name:
+            self.endpoint_name = endpoint_name
+        else:
+            self.endpoint_name = self.name
+
+        data_capture_config_dict = None
+        if data_capture_config is not None:
+            data_capture_config_dict = data_capture_config._to_request_dict()
+
+        if update_endpoint:
+            endpoint_config_name = self.sagemaker_session.create_endpoint_config(
+                name=self.name,
+                model_name=self.name,
+                initial_instance_count=initial_instance_count,
+                instance_type=instance_type,
+                accelerator_type=accelerator_type,
+                tags=tags,
+                kms_key=kms_key,
+                data_capture_config_dict=data_capture_config_dict,
+            )
+            self.sagemaker_session.update_endpoint(self.endpoint_name, endpoint_config_name)
+        else:
+            self.sagemaker_session.endpoint_from_production_variants(
+                name=self.endpoint_name,
+                production_variants=[production_variant],
+                tags=tags,
+                kms_key=kms_key,
+                wait=wait,
+                data_capture_config_dict=data_capture_config_dict,
+            )
+
+        if predictor:
+            return predictor(self.endpoint_name, self.sagemaker_session)
+        return None
+
+    def add_model(self, model_data_source, model_data_path=None):
+        """Adds a model to the `MultiDataModel` by uploading or copying the model_data_source
+         artifact to the given S3 path model_data_path relative to model_data_prefix
+
+        Args:
+        model_source: Valid local file path or S3 path of the trained model artifact
+        model_data_path: S3 path where the trained model artifact
+                should be uploaded relative to `self.model_data_prefix` path. (default: None).
+                If None, then the model artifact is uploaded to a path relative to model_data_prefix
+
+        Returns:
+            str: S3 uri to uploaded model artifact
+        """
+        parse_result = urlparse(model_data_source)
+
+        # If the model source is an S3 path, copy the model artifact to the destination S3 path
+        if parse_result.scheme == "s3":
+            source_bucket, source_model_data_path = s3.parse_s3_url(model_data_source)
+            copy_source = {"Bucket": source_bucket, "Key": source_model_data_path}
+
+            if not model_data_path:
+                model_data_path = source_model_data_path
+
+            # Construct the destination path
+            dst_url = os.path.join(self.model_data_prefix, model_data_path)
+            destination_bucket, destination_model_data_path = s3.parse_s3_url(dst_url)
+
+            # Copy the model artifact
+            self.s3_client.copy(copy_source, destination_bucket, destination_model_data_path)
+            return os.path.join("s3://", destination_bucket, destination_model_data_path)
+
+        # If the model source is a local path, upload the local model artifact to the destination
+        #  s3 path
+        if os.path.exists(model_data_source):
+            destination_bucket, dst_prefix = s3.parse_s3_url(self.model_data_prefix)
+            if model_data_path:
+                dst_s3_uri = os.path.join(dst_prefix, model_data_path)
+            else:
+                dst_s3_uri = os.path.join(dst_prefix, os.path.basename(model_data_source))
+            self.s3_client.upload_file(model_data_source, destination_bucket, dst_s3_uri)
+            # return upload_path
+            return os.path.join("s3://", destination_bucket, dst_s3_uri)
+
+        # Raise error if the model source is of an unexpected type
+        raise ValueError(
+            "model_source must either be a valid local file path or s3 uri. Received: "
+            '"{}"'.format(model_data_source)
+        )
+
+    def list_models(self):
+        """Generates and returns relative paths to model archives stored at model_data_prefix
+        S3 location.
+
+        Yields: Paths to model archives relative to model_data_prefix path.
+        """
+        bucket, url_prefix = s3.parse_s3_url(self.model_data_prefix)
+        file_keys = self.sagemaker_session.list_s3_files(bucket=bucket, key_prefix=url_prefix)
+        for file_key in file_keys:
+            # Return the model paths relative to the model_data_prefix
+            # Ex: "a/b/c.tar.gz" -> "b/c.tar.gz" where url_prefix = "a/"
+            yield file_key.replace(url_prefix, "")

src/sagemaker/predictor.py

@@ -83,7 +83,7 @@ def __init__(
         self._endpoint_config_name = self._get_endpoint_config_name()
         self._model_names = self._get_model_names()
 
-    def predict(self, data, initial_args=None):
+    def predict(self, data, initial_args=None, target_model=None):
         """Return the inference from the specified endpoint.
 
         Args:
@@ -95,6 +95,9 @@ def predict(self, data, initial_args=None):
             initial_args (dict[str,str]): Optional. Default arguments for boto3
                 ``invoke_endpoint`` call. Default is None (no default
                 arguments).
+            target_model (str): S3 model artifact path to run an inference request on,
+                in case of a multi model endpoint. Does not apply to endpoints hosting
+                single model (Default: None)
 
         Returns:
             object: Inference for the given input. If a deserializer was specified when creating
@@ -103,7 +106,7 @@ def predict(self, data, initial_args=None):
                 as is.
         """
 
-        request_args = self._create_request_args(data, initial_args)
+        request_args = self._create_request_args(data, initial_args, target_model)
         response = self.sagemaker_session.sagemaker_runtime_client.invoke_endpoint(**request_args)
         return self._handle_response(response)
 
@@ -120,11 +123,12 @@ def _handle_response(self, response):
         response_body.close()
         return data
 
-    def _create_request_args(self, data, initial_args=None):
+    def _create_request_args(self, data, initial_args=None, target_model=None):
         """
         Args:
             data:
             initial_args:
+            target_model:
         """
         args = dict(initial_args) if initial_args else {}
 
@@ -137,6 +141,9 @@ def _create_request_args(self, data, initial_args=None):
         if self.accept and "Accept" not in args:
             args["Accept"] = self.accept
 
+        if target_model:
+            args["TargetModel"] = target_model
+
         if self.serializer is not None:
             data = self.serializer(data)