add SageMaker Neo support

Author: wweic

src/sagemaker/amazon/amazon_estimator.py

@@ -353,6 +353,13 @@ def registry(region_name, algorithm=None):
             "eu-west-2": "644912444149",
             "us-west-1": "632365934929",
         }[region_name]
+    elif algorithm in ['image-classification-neo', 'xgboost-neo']:
+        account_id = {
+            'us-west-2': '301217895009',
+            'us-east-1': '785573368785',
+            'eu-west-1': '802834080501',
+            'us-east-2': '007439368137'
+        }[region_name]
     else:
         raise ValueError("Algorithm class:{} doesn't have mapping to account_id with images".format(algorithm))
     return "{}.dkr.ecr.{}.amazonaws.com".format(account_id, region_name)

src/sagemaker/estimator.py

@@ -26,7 +26,7 @@
                                 validate_source_dir)
 from sagemaker.job import _Job
 from sagemaker.local import LocalSession
-from sagemaker.model import Model
+from sagemaker.model import Model, NEO_ALLOWED_TARGET_INSTANCE_FAMILY, NEO_ALLOWED_FRAMEWORKS
 from sagemaker.model import (SCRIPT_PARAM_NAME, DIR_PARAM_NAME, CLOUDWATCH_METRICS_PARAM_NAME,
                              CONTAINER_LOG_LEVEL_PARAM_NAME, JOB_NAME_PARAM_NAME, SAGEMAKER_REGION_PARAM_NAME)
 from sagemaker.predictor import RealTimePredictor
@@ -131,6 +131,8 @@ def __init__(self, role, train_instance_count, train_instance_type,
         self.output_kms_key = output_kms_key
         self.latest_training_job = None
 
+        self._compiled_models = {}
+
         # VPC configurations
         self.subnets = subnets
         self.security_group_ids = security_group_ids
@@ -216,6 +218,57 @@ def fit(self, inputs=None, wait=True, logs=True, job_name=None):
         if wait:
             self.latest_training_job.wait(logs=logs)
 
+    def _compilation_job_name(self):
+        base_name = self.base_job_name or base_name_from_image(self.train_image())
+        return name_from_base('compilation-' + base_name)
+
+    def compile_model(self, target_instance_family, input_shape, output_path, framework=None, framework_version=None,
+                      compile_max_run=5 * 60, tags=None, **kwargs):
+        """Compile a Neo model using the input model.
+
+        Args:
+            target_instance_family (str): Identifies the device that you want to run your model after compilation, for
+                example: ml_c5. Allowed strings are: ml_c5, ml_m5, ml_c4, ml_m4, jetsontx1, jetsontx2, ml_p2, ml_p3,
+                deeplens, rasp3b
+            input_shape (dict): Specifies the name and shape of the expected inputs for your trained model in json
+                dictionary form, for example: {‘data’:[1,3,1024,1024]}, or {‘var1’: [1,1,28,28], ‘var2’:[1,1,28,28]}
+            output_path (str): Specifies where to store the compiled model
+            framework (str): The framework that is used to train the original model. Allowed values: 'mxnet',
+                'tensorflow', 'pytorch', 'onnx', 'xgboost'
+            framework_version (str): The version of the framework
+            compile_max_run (int): Timeout in seconds for compilation (default: 3 * 60).
+                After this amount of time Amazon SageMaker Neo terminates the compilation job regardless of its
+                current status.
+            tags (list[dict]): List of tags for labeling a compilation job. For more, see
+                https://docs.aws.amazon.com/sagemaker/latest/dg/API_Tag.html.
+            **kwargs: Passed to invocation of ``create_model()``. Implementations may customize
+                ``create_model()`` to accept ``**kwargs`` to customize model creation during deploy.
+                For more, see the implementation docs.
+        Returns:
+            sagemaker.model.Model: A SageMaker ``Model`` object. See :func:`~sagemaker.model.Model` for full details.
+        """
+        if target_instance_family not in NEO_ALLOWED_TARGET_INSTANCE_FAMILY:
+            raise ValueError("Please use valid target_instance_family,"
+                             "allowed values: {}".format(NEO_ALLOWED_TARGET_INSTANCE_FAMILY))
+        if framework and framework not in NEO_ALLOWED_FRAMEWORKS:
+            raise ValueError("Please use valid framework, allowed values: {}".format(NEO_ALLOWED_FRAMEWORKS))
+
+        if (framework is None) != (framework_version is None):
+            raise ValueError("You should provide framework and framework_version at the same time.")
+
+        model = self.create_model(**kwargs)
+
+        self._compiled_models[target_instance_family] = model.compile(target_instance_family,
+                                                                      input_shape,
+                                                                      output_path,
+                                                                      self.role,
+                                                                      tags,
+                                                                      self._compilation_job_name(),
+                                                                      compile_max_run,
+                                                                      framework=framework,
+                                                                      framework_version=framework_version)
+        return self._compiled_models[target_instance_family]
+
     @classmethod
     def attach(cls, training_job_name, sagemaker_session=None, model_channel_name='model'):
         """Attach to an existing training job.
@@ -257,7 +310,7 @@ def attach(cls, training_job_name, sagemaker_session=None, model_channel_name='m
         estimator.latest_training_job.wait()
         return estimator
 
-    def deploy(self, initial_instance_count, instance_type, endpoint_name=None, **kwargs):
+    def deploy(self, initial_instance_count, instance_type, endpoint_name=None, use_compiled_model=False, **kwargs):
         """Deploy the trained model to an Amazon SageMaker endpoint and return a ``sagemaker.RealTimePredictor`` object.
 
         More information:
@@ -269,6 +322,7 @@ def deploy(self, initial_instance_count, instance_type, endpoint_name=None, **kw
                 for example, 'ml.c4.xlarge'.
             endpoint_name (str): Name to use for creating an Amazon SageMaker endpoint. If not specified, the name of
                 the training job is used.
+            use_compiled_model (bool): Flag to select whether to use compiled (optimized) model. Default: False.
             **kwargs: Passed to invocation of ``create_model()``. Implementations may customize
                 ``create_model()`` to accept ``**kwargs`` to customize model creation during deploy.
                 For more, see the implementation docs.
@@ -280,7 +334,15 @@ def deploy(self, initial_instance_count, instance_type, endpoint_name=None, **kw
         self._ensure_latest_training_job()
         endpoint_name = endpoint_name or self.latest_training_job.name
         self.deploy_instance_type = instance_type
-        return self.create_model(**kwargs).deploy(
+        if use_compiled_model:
+            family = '_'.join(instance_type.split('.')[:-1])
+            if family not in self._compiled_models:
+                raise ValueError("No compiled model for {}. "
+                                 "Please compile one with compile_model before deploying.".format(family))
+            model = self._compiled_models[family]
+        else:
+            model = self.create_model(**kwargs)
+        return model.deploy(
             instance_type=instance_type,
             initial_instance_count=initial_instance_count,
             endpoint_name=endpoint_name)

src/sagemaker/model.py

@@ -12,11 +12,23 @@
 # language governing permissions and limitations under the License.
 from __future__ import absolute_import
 
+import json
 import logging
 
 import sagemaker
 from sagemaker import fw_utils, local, session, utils
 
+NEO_ALLOWED_TARGET_INSTANCE_FAMILY = set(['ml_c5', 'ml_m5', 'ml_c4', 'ml_m4', 'jetson_tx1', 'jetson_tx2', 'ml_p2',
+                                          'ml_p3', 'deeplens', 'rasp3b'])
+NEO_ALLOWED_FRAMEWORKS = set(['mxnet', 'tensorflow', 'pytorch', 'onnx', 'xgboost'])
+
+NEO_IMAGE_ACCOUNT = {
+    'us-west-2': '301217895009',
+    'us-east-1': '785573368785',
+    'eu-west-1': '802834080501',
+    'us-east-2': '007439368137'
+}
+
 
 class Model(object):
     """A SageMaker ``Model`` that can be deployed to an ``Endpoint``."""
@@ -53,6 +65,7 @@ def __init__(self, model_data, image, role=None, predictor_cls=None, env=None, n
         self.vpc_config = vpc_config
         self.sagemaker_session = sagemaker_session
         self._model_name = None
+        self._is_compiled_model = False
 
     def prepare_container_def(self, instance_type):  # pylint: disable=unused-argument
         """Return a dict created by ``sagemaker.container_def()`` for deploying this model to a specified instance type.
@@ -68,6 +81,93 @@ def prepare_container_def(self, instance_type):  # pylint: disable=unused-argume
         """
         return sagemaker.container_def(self.image, self.model_data, self.env)
 
+    def _framework(self):
+        return getattr(self, '__framework_name__', None)
+
+    def _get_framework_version(self):
+        return getattr(self, 'framework_version', None)
+
+    def _compilation_job_config(self, target_instance_type, input_shape, output_path, role, compile_max_run,
+                                job_name, framework, tags):
+        input_model_config = {
+            'S3Uri': self.model_data,
+            'DataInputConfig': input_shape if type(input_shape) != dict else json.dumps(input_shape),
+            'Framework': framework
+        }
+        role = self.sagemaker_session.expand_role(role)
+        output_model_config = {
+            'TargetDevice': target_instance_type,
+            'S3OutputLocation': output_path
+        }
+
+        return {'input_model_config': input_model_config,
+                'output_model_config': output_model_config,
+                'role': role,
+                'stop_condition': {
+                    'MaxRuntimeInSeconds': compile_max_run
+                },
+                'tags': tags,
+                'job_name': job_name}
+
+    def _neo_image_account(self, region):
+        if region not in NEO_IMAGE_ACCOUNT:
+            raise ValueError("Neo is not currently supported in {}, "
+                             "valid regions: {}".format(region, NEO_IMAGE_ACCOUNT.keys()))
+        return NEO_IMAGE_ACCOUNT[region]
+
+    def _neo_image(self, region, target_instance_type, framework, framework_version):
+        return fw_utils.create_image_uri(region,
+                                         'neo-' + framework.lower(),
+                                         target_instance_type.replace('_', '.'),
+                                         framework_version,
+                                         py_version='py3',
+                                         account=self._neo_image_account(region))
+
+    def compile(self, target_instance_family, input_shape, output_path, role,
+                tags=None, job_name=None, compile_max_run=5 * 60, framework=None, framework_version=None):
+        """Compile this ``Model`` with SageMaker Neo.
+
+        Args:
+            target_instance_family (str): Identifies the device that you want to run your model after compilation, for
+                example: ml_c5. Allowed strings are: ml_c5, ml_m5, ml_c4, ml_m4, jetsontx1, jetsontx2, ml_p2, ml_p3,
+                deeplens, rasp3b
+            input_shape (dict): Specifies the name and shape of the expected inputs for your trained model in json
+                dictionary form, for example: {‘data’:[1,3,1024,1024]}, or {‘var1’: [1,1,28,28], ‘var2’:[1,1,28,28]}
+            output_path (str): Specifies where to store the compiled model
+            role (str): Execution role
+            tags (list[dict]): List of tags for labeling a compilation job. For more, see
+                https://docs.aws.amazon.com/sagemaker/latest/dg/API_Tag.html.
+            job_name (str): The name of the compilation job
+            compile_max_run (int): Timeout in seconds for compilation (default: 3 * 60).
+                After this amount of time Amazon SageMaker Neo terminates the compilation job regardless of its
+                current status.
+            framework (str): The framework that is used to train the original model. Allowed values: 'mxnet',
+                'tensorflow', 'pytorch', 'onnx', 'xgboost'
+            framework_version (str)
+        Returns:
+            sagemaker.model.Model: A SageMaker ``Model`` object. See :func:`~sagemaker.model.Model` for full details.
+        """
+        framework = self._framework() or framework
+        if framework is None:
+            raise ValueError("You must specify framework, allowed values {}".format(NEO_ALLOWED_FRAMEWORKS))
+        if framework not in NEO_ALLOWED_FRAMEWORKS:
+            raise ValueError("You must provide valid framework, allowed values {}".format(NEO_ALLOWED_FRAMEWORKS))
+        if job_name is None:
+            raise ValueError("You must provide a compilation job name")
+
+        framework = framework.upper()
+        framework_version = self._get_framework_version() or framework_version
+
+        config = self._compilation_job_config(target_instance_family, input_shape, output_path, role,
+                                              compile_max_run, job_name, framework, tags)
+        self.sagemaker_session.compile_model(**config)
+        job_status = self.sagemaker_session.wait_for_compilation_job(job_name)
+        self.model_data = job_status['ModelArtifacts']['S3ModelArtifacts']
+        self.image = self._neo_image(self.sagemaker_session.boto_region_name, target_instance_family, framework,
+                                     framework_version)
+        self._is_compiled_model = True
+        return self
+
     def deploy(self, initial_instance_count, instance_type, endpoint_name=None, tags=None):
         """Deploy this ``Model`` to an ``Endpoint`` and optionally return a ``Predictor``.
 
@@ -98,13 +198,21 @@ def deploy(self, initial_instance_count, instance_type, endpoint_name=None, tags
             else:
                 self.sagemaker_session = session.Session()
 
+        compiled_model_suffix = '-'.join(instance_type.split('.')[:-1])
         container_def = self.prepare_container_def(instance_type)
         self.name = self.name or utils.name_from_image(container_def['Image'])
         if self.role is None:
             raise ValueError("Role can not be null for deploying a model")
+        if self._is_compiled_model:
+            self.name += compiled_model_suffix
         self.sagemaker_session.create_model(self.name, self.role, container_def, vpc_config=self.vpc_config)
         production_variant = sagemaker.production_variant(self.name, instance_type, initial_instance_count)
-        self.endpoint_name = endpoint_name or self.name
+        if endpoint_name:
+            self.endpoint_name = endpoint_name
+        else:
+            self.endpoint_name = self.name
+            if self._is_compiled_model and not self.endpoint_name.endswith(compiled_model_suffix):
+                self.endpoint_name += compiled_model_suffix
         self.sagemaker_session.endpoint_from_production_variants(self.endpoint_name, [production_variant], tags)
         if self.predictor_cls:
             return self.predictor_cls(self.endpoint_name, self.sagemaker_session)

src/sagemaker/session.py

@@ -283,6 +283,40 @@ def train(self, image, input_mode, input_config, role, job_name, output_config,
         LOGGER.debug('train request: {}'.format(json.dumps(train_request, indent=4)))
         self.sagemaker_client.create_training_job(**train_request)
 
+    def compile_model(self, input_model_config, output_model_config, role,
+                      job_name, stop_condition, tags):
+        """Create an Amazon SageMaker Neo compilation job.
+
+        Args:
+            input_model_config (dict): the trained model and the Amazon S3 location where it is stored.
+            output_model_config (dict): - Identifies the Amazon S3 location where you want Amazon SageMaker Neo to save
+                the results of compilation job
+            role (str): An AWS IAM role (either name or full ARN). The Amazon SageMaker Neo compilation jobs use this
+                role to access model artifacts. You must grant sufficient permissions to this role.
+            job_name (str): Name of the compilation job being created.
+            stop_condition (dict): Defines when compilation job shall finish. Contains entries that can be understood
+                by the service like ``MaxRuntimeInSeconds``.
+            tags (list[dict]): List of tags for labeling a compile model job. For more, see
+                https://docs.aws.amazon.com/sagemaker/latest/dg/API_Tag.html.
+
+        Returns:
+            str: ARN of the compile model job, if it is created.
+        """
+
+        compilation_job_request = {
+            'InputConfig': input_model_config,
+            'OutputConfig': output_model_config,
+            'RoleArn': role,
+            'StoppingCondition': stop_condition,
+            'CompilationJobName': job_name
+        }
+
+        if tags is not None:
+            compilation_job_request['Tags'] = tags
+
+        LOGGER.info('Creating compilation-job with name: {}'.format(job_name))
+        self.sagemaker_client.create_compilation_job(**compilation_job_request)
+
     def tune(self, job_name, strategy, objective_type, objective_metric_name,
              max_jobs, max_parallel_jobs, parameter_ranges,
              static_hyperparameters, image, input_mode, metric_definitions,
@@ -613,6 +647,23 @@ def wait_for_job(self, job, poll=5):
         self._check_job_status(job, desc, 'TrainingJobStatus')
         return desc
 
+    def wait_for_compilation_job(self, job, poll=5):
+        """Wait for an Amazon SageMaker Neo compilation job to complete.
+
+        Args:
+            job (str): Name of the compilation job to wait for.
+            poll (int): Polling interval in seconds (default: 5).
+
+        Returns:
+            (dict): Return value from the ``DescribeCompilationJob`` API.
+
+        Raises:
+            ValueError: If the compilation job fails.
+        """
+        desc = _wait_until(lambda: _compilation_job_status(self.sagemaker_client, job), poll)
+        self._check_job_status(job, desc, 'CompilationJobStatus')
+        return desc
+
     def wait_for_tuning_job(self, job, poll=5):
         """Wait for an Amazon SageMaker hyperparameter tuning job to complete.
 
@@ -1164,6 +1215,28 @@ def _train_done(sagemaker_client, job_name, last_desc):
     return desc, True
 
 
+def _compilation_job_status(sagemaker_client, job_name):
+    compile_status_codes = {
+        'Completed': '!',
+        'InProgress': '.',
+        'Failed': '*',
+        'Stopped': 's',
+        'Stopping': '_'
+    }
+    in_progress_statuses = ['InProgress', 'Stopping', 'Starting']
+
+    desc = sagemaker_client.describe_compilation_job(CompilationJobName=job_name)
+    status = desc['CompilationJobStatus']
+
+    print(compile_status_codes.get(status, '?'), end='')
+    sys.stdout.flush()
+
+    if status in in_progress_statuses:
+        return None
+
+    return desc
+
+
 def _tuning_job_status(sagemaker_client, job_name):
     tuning_status_codes = {
         'Completed': '!',