Refactor LocalSageMakerClient

local sagemakerclient was not very useful beyond creating a single
estimator/endpoint. doing workflows such as training and later creating
and endpoint was really awkward.

These changes make it both resemble the API a bit more and allow
persisting objects across LocalSessions. This is important because most
of the time the SDK classes create sessions behind the scenes.

Author: Ignacio Quintero

src/sagemaker/local/entities.py

@@ -0,0 +1,178 @@
+# Copyright 2017-2018 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.
+from __future__ import absolute_import
+
+import datetime
+import logging
+import time
+import urllib3
+
+from sagemaker.local.image import _SageMakerContainer
+from sagemaker.utils import get_config_value
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.WARNING)
+
+
+class _LocalTrainingJob(object):
+
+    _STARTING = 'Starting'
+    _TRAINING = 'Training'
+    _COMPLETED = 'Completed'
+    _states = ['Starting', 'Training', 'Completed']
+
+    def __init__(self, container):
+        self.container = container
+        self.model_artifacts = None
+        self.state = 'created'
+        self.start_time = None
+        self.end_time = None
+
+    def start(self, input_data_config, hyperparameters):
+        for channel in input_data_config:
+            if channel['DataSource'] and 'S3DataSource' in channel['DataSource']:
+                data_distribution = channel['DataSource']['S3DataSource']['S3DataDistributionType']
+            elif channel['DataSource'] and 'FileDataSource' in channel['DataSource']:
+                data_distribution = channel['DataSource']['FileDataSource']['FileDataDistributionType']
+            else:
+                raise ValueError('Need channel[\'DataSource\'] to have [\'S3DataSource\'] or [\'FileDataSource\']')
+
+            if data_distribution != 'FullyReplicated':
+                raise RuntimeError('DataDistribution: %s is not currently supported in Local Mode' %
+                                   data_distribution)
+
+        self.start = datetime.datetime.now()
+        self.state = self._TRAINING
+
+        self.model_artifacts = self.container.train(input_data_config, hyperparameters)
+        self.end = datetime.datetime.now()
+        self.state = self._COMPLETED
+
+    def describe(self):
+        response = {
+            'ResourceConfig': {
+                'InstanceCount': self.container.instance_count
+            },
+            'TrainingJobStatus': self.state,
+            'TrainingStartTime': self.start_time,
+            'TrainingEndTime': self.end_time,
+            'ModelArtifacts': {
+                'S3ModelArtifacts': self.model_artifacts
+            }
+        }
+        return response
+
+
+class _LocalModel(object):
+
+    def __init__(self, model_name, primary_container):
+        self.model_name = model_name
+        self.primary_container = primary_container
+        self.creation_time = datetime.datetime.now()
+
+    def describe(self):
+        response = {
+            'ModelName': self.model_name,
+            'CreationTime': self.creation_time,
+            'ExecutionRoleArn': 'local:arn-does-not-matter',
+            'ModelArn': 'local:arn-does-not-matter',
+            'PrimaryContainer': self.primary_container
+        }
+        return response
+
+
+class _LocalEndpointConfig(object):
+
+    def __init__(self, config_name, production_variants):
+        self.name = config_name
+        self.production_variants = production_variants
+        self.creation_time = datetime.datetime.now()
+
+    def describe(self):
+        response = {
+            'EndpointConfigName': self.name,
+            'EndpointConfigArn': 'local:arn-does-not-matter',
+            'CreationTime': self.creation_time,
+            'ProductionVariants': self.production_variants
+        }
+        return response
+
+
+class _LocalEndpoint(object):
+
+    _CREATING = 'Creating'
+    _IN_SERVICE = 'InService'
+    _FAILED = 'Failed'
+
+    def __init__(self, endpoint_name, endpoint_config_name):
+        from sagemaker.local import LocalSagemakerClient
+        local_client = LocalSagemakerClient()
+
+        self.name = endpoint_name
+        self.endpoint_config = local_client.describe_endpoint_config(endpoint_config_name)
+        self.production_variant = self.endpoint_config['ProductionVariants'][0]
+
+        model_name = self.production_variant['ModelName']
+        self.primary_container = local_client.describe_model(model_name)['PrimaryContainer']
+
+        self.container = None
+        self.create_time = None
+        self.state = _LocalEndpoint._CREATING
+
+    def serve(self, sagemaker_session):
+        image = self.primary_container['Image']
+        instance_type = self.production_variant['InstanceType']
+        instance_count = self.production_variant['InitialInstanceCount']
+
+        self.create_time = datetime.datetime.now()
+        self.container = _SageMakerContainer(instance_type, instance_count, image, sagemaker_session)
+        self.container.serve(self.primary_container['ModelDataUrl'], self.primary_container['Environment'])
+
+        i = 0
+        http = urllib3.PoolManager()
+        serving_port = get_config_value('local.serving_port', sagemaker_session.config) or 8080
+        endpoint_url = 'http://localhost:%s/ping' % serving_port
+        while True:
+            i += 1
+            if i >= 10:
+                self.state = _LocalEndpoint._FAILED
+                raise RuntimeError('Giving up, endpoint: %s didn\'t launch correctly' % self.name)
+
+            logger.info('Checking if endpoint is up, attempt: %s' % i)
+            try:
+                r = http.request('GET', endpoint_url)
+                if r.status != 200:
+                    logger.info('Container still not up, got: %s' % r.status)
+                else:
+                    # the container is running and it passed the healthcheck status is now InService
+                    self.state = _LocalEndpoint._IN_SERVICE
+                    return
+            except urllib3.exceptions.RequestError:
+                logger.info('Container still not up')
+
+            time.sleep(1)
+
+    def stop(self):
+        if self.container:
+            self.container.stop_serving()
+
+    def describe(self):
+        response = {
+            'EndpointConfigName': self.endpoint_config['EndpointConfigName'],
+            'CreationTime': self.create_time,
+            'ProductionVariants': self.endpoint_config['ProductionVariants'],
+            'EndpointName': self.name,
+            'EndpointArn': 'local:arn-does-not-matter',
+            'EndpointStatus': self.state
+        }
+        return response

src/sagemaker/local/image.py

@@ -134,7 +134,7 @@ def train(self, input_data_config, hyperparameters):
         print('===== Job Complete =====')
         return s3_artifacts
 
-    def serve(self, primary_container):
+    def serve(self, model_dir, environment):
         """Host a local endpoint using docker-compose.
         Args:
             primary_container (dict): dictionary containing the container runtime settings
@@ -148,13 +148,12 @@ def serve(self, primary_container):
         self.container_root = self._create_tmp_folder()
         logger.info('creating hosting dir in {}'.format(self.container_root))
 
-        model_dir = primary_container['ModelDataUrl']
         volumes = self._prepare_serving_volumes(model_dir)
-        env_vars = ['{}={}'.format(k, v) for k, v in primary_container['Environment'].items()]
+        env_vars = ['{}={}'.format(k, v) for k, v in environment.items()]
 
         # If the user script was passed as a file:// mount it to the container.
-        if sagemaker.estimator.DIR_PARAM_NAME.upper() in primary_container['Environment']:
-            script_dir = primary_container['Environment'][sagemaker.estimator.DIR_PARAM_NAME.upper()]
+        if sagemaker.estimator.DIR_PARAM_NAME.upper() in environment:
+            script_dir = environment[sagemaker.estimator.DIR_PARAM_NAME.upper()]
             parsed_uri = urlparse(script_dir)
             if parsed_uri.scheme == 'file':
                 volumes.append(_Volume(parsed_uri.path, '/opt/ml/code'))

src/sagemaker/local/local_session.py

@@ -12,16 +12,15 @@
 # language governing permissions and limitations under the License.
 from __future__ import absolute_import
 
-import datetime
 import logging
 import platform
-import time
 
 import boto3
 import urllib3
 from botocore.exceptions import ClientError
 
 from sagemaker.local.image import _SageMakerContainer
+from sagemaker.local.entities import _LocalEndpointConfig, _LocalEndpoint, _LocalModel, _LocalTrainingJob
 from sagemaker.session import Session
 from sagemaker.utils import get_config_value
 
@@ -37,42 +36,44 @@ class LocalSagemakerClient(object):
 
     Implements the methods with the same signature as the boto SageMakerClient.
     """
+
+    _training_jobs = {}
+    _models = {}
+    _endpoint_configs = {}
+    _endpoints = {}
+
     def __init__(self, sagemaker_session=None):
         """Initialize a LocalSageMakerClient.
 
         Args:
             sagemaker_session (sagemaker.session.Session): a session to use to read configurations
                 from, and use its boto client.
         """
-        self.train_container = None
         self.serve_container = None
         self.sagemaker_session = sagemaker_session or LocalSession()
         self.s3_model_artifacts = None
-        self.model_name = None
-        self.primary_container = None
-        self.role_arn = None
         self.created_endpoint = False
 
-    def create_training_job(self, TrainingJobName, AlgorithmSpecification, RoleArn, InputDataConfig, OutputDataConfig,
-                            ResourceConfig, StoppingCondition, HyperParameters, Tags=None):
-
-        self.train_container = _SageMakerContainer(ResourceConfig['InstanceType'], ResourceConfig['InstanceCount'],
-                                                   AlgorithmSpecification['TrainingImage'], self.sagemaker_session)
-
-        for channel in InputDataConfig:
-
-            if channel['DataSource'] and 'S3DataSource' in channel['DataSource']:
-                data_distribution = channel['DataSource']['S3DataSource']['S3DataDistributionType']
-            elif channel['DataSource'] and 'FileDataSource' in channel['DataSource']:
-                data_distribution = channel['DataSource']['FileDataSource']['FileDataDistributionType']
-            else:
-                raise ValueError('Need channel[\'DataSource\'] to have [\'S3DataSource\'] or [\'FileDataSource\']')
+    def create_training_job(self, TrainingJobName, AlgorithmSpecification, InputDataConfig, OutputDataConfig,
+                            ResourceConfig, HyperParameters, *args, **kwargs):
+        """
+        Create a training job in Local Mode
+        Args:
+            TrainingJobName (str): local training job name.
+            AlgorithmSpecification (dict): Identifies the training algorithm to use.
+            InputDataConfig (dict): Describes the training dataset and the location where it is stored.
+            OutputDataConfig (dict): Identifies the location where you want to save the results of model training.
+            ResourceConfig (dict): Identifies the resources to use for local model traininig.
+            HyperParameters (dict): Specify these algorithm-specific parameters to influence the quality of the final
+                model.
+        """
 
-            if data_distribution != 'FullyReplicated':
-                raise RuntimeError("DataDistribution: %s is not currently supported in Local Mode" %
-                                   data_distribution)
+        container = _SageMakerContainer(ResourceConfig['InstanceType'], ResourceConfig['InstanceCount'],
+                                        AlgorithmSpecification['TrainingImage'], self.sagemaker_session)
+        train_job = _LocalTrainingJob(container)
+        train_job.start(InputDataConfig, HyperParameters)
 
-        self.s3_model_artifacts = self.train_container.train(InputDataConfig, HyperParameters)
+        LocalSagemakerClient._training_jobs[TrainingJobName] = train_job
 
     def describe_training_job(self, TrainingJobName):
         """Describe a local training job.
@@ -83,63 +84,55 @@ def describe_training_job(self, TrainingJobName):
         Returns: (dict) DescribeTrainingJob Response.
 
         """
-        response = {'ResourceConfig': {'InstanceCount': self.train_container.instance_count},
-                    'TrainingJobStatus': 'Completed',
-                    'TrainingStartTime': datetime.datetime.now(),
-                    'TrainingEndTime': datetime.datetime.now(),
-                    'ModelArtifacts': {'S3ModelArtifacts': self.s3_model_artifacts}
-                    }
-        return response
-
-    def create_model(self, ModelName, PrimaryContainer, ExecutionRoleArn):
-        self.model_name = ModelName
-        self.primary_container = PrimaryContainer
-        self.role_arn = ExecutionRoleArn
+        if TrainingJobName not in LocalSagemakerClient._training_jobs:
+            error_response = {'Error': {'Code': 'ValidationException', 'Message': 'Could not find local training job'}}
+            raise ClientError(error_response, 'describe_training_job')
+        else:
+            return LocalSagemakerClient._training_jobs[TrainingJobName].describe()
+
+    def create_model(self, ModelName, PrimaryContainer, *args, **kwargs):
+        """Create a Local Model Object
+
+        Args:
+            ModelName (str): the Model Name
+            PrimaryContainer (dict): a SageMaker primary container definition
+        """
+        LocalSagemakerClient._models[ModelName] = _LocalModel(ModelName, PrimaryContainer)
+
+    def describe_model(self, ModelName):
+        if ModelName not in LocalSagemakerClient._models:
+            error_response = {'Error': {'Code': 'ValidationException', 'Message': 'Could not find local model'}}
+            raise ClientError(error_response, 'describe_model')
+        else:
+            return LocalSagemakerClient._models[ModelName].describe()
 
     def describe_endpoint_config(self, EndpointConfigName):
-        if self.created_endpoint:
-            return True
+        if EndpointConfigName in LocalSagemakerClient._endpoint_configs:
+            return LocalSagemakerClient._endpoint_configs[EndpointConfigName].describe()
         else:
-            error_response = {'Error': {'Code': 'ValidationException', 'Message': 'Could not find endpoint'}}
+            error_response = {'Error': {
+                'Code': 'ValidationException', 'Message': 'Could not find local endpoint config'}}
             raise ClientError(error_response, 'describe_endpoint_config')
 
     def create_endpoint_config(self, EndpointConfigName, ProductionVariants):
-        self.variants = ProductionVariants
+        LocalSagemakerClient._endpoint_configs[EndpointConfigName] = _LocalEndpointConfig(
+            EndpointConfigName, ProductionVariants)
 
     def describe_endpoint(self, EndpointName):
-        return {'EndpointStatus': 'InService'}
+        if EndpointName not in LocalSagemakerClient._endpoints:
+            error_response = {'Error': {'Code': 'ValidationException', 'Message': 'Could not find local endpoint'}}
+            raise ClientError(error_response, 'describe_endpoint')
+        else:
+            return LocalSagemakerClient._endpoints[EndpointName].describe()
 
     def create_endpoint(self, EndpointName, EndpointConfigName):
-        instance_type = self.variants[0]['InstanceType']
-        instance_count = self.variants[0]['InitialInstanceCount']
-        self.serve_container = _SageMakerContainer(instance_type, instance_count,
-                                                   self.primary_container['Image'], self.sagemaker_session)
-        self.serve_container.serve(self.primary_container)
-        self.created_endpoint = True
-
-        i = 0
-        http = urllib3.PoolManager()
-        serving_port = get_config_value('local.serving_port', self.sagemaker_session.config) or 8080
-        endpoint_url = "http://localhost:%s/ping" % serving_port
-        while True:
-            i += 1
-            if i >= 10:
-                raise RuntimeError("Giving up, endpoint: %s didn't launch correctly" % EndpointName)
-
-            logger.info("Checking if endpoint is up, attempt: %s" % i)
-            try:
-                r = http.request('GET', endpoint_url)
-                if r.status != 200:
-                    logger.info("Container still not up, got: %s" % r.status)
-                else:
-                    return
-            except urllib3.exceptions.RequestError:
-                logger.info("Container still not up")
-
-            time.sleep(1)
+        endpoint = _LocalEndpoint(EndpointName, EndpointConfigName)
+        LocalSagemakerClient._endpoints[EndpointName] = endpoint
+        endpoint.serve(self.sagemaker_session)
 
     def delete_endpoint(self, EndpointName):
-        self.serve_container.stop_serving()
+        if EndpointName in LocalSagemakerClient._endpoints:
+            LocalSagemakerClient._endpoints[EndpointName].stop()
 
 
 class LocalSagemakerRuntimeClient(object):