Add support for async fit()

when calling fit(wait=False) it will return immediately. The training
job will carry on even if the process exits. by using attach() the
estimator can be retrieved by providing the training job name.

Author: Ignacio Quintero

src/sagemaker/amazon/amazon_estimator.py

@@ -64,6 +64,20 @@ def data_location(self, data_location):
             data_location = data_location + '/'
         self._data_location = data_location
 
+    @classmethod
+    def _from_training_job(cls, init_params, hyperparameters, image, sagemaker_session):
+
+        # The hyperparam names may not be the same as the class attribute that holds them,
+        # for instance: local_lloyd_init_method is called local_init_method. We need to map these
+        # and pass the correct name to the constructor.
+
+        for attribute, value in cls.__dict__.items():
+            if isinstance(value, hp):
+                if value.name in hyperparameters:
+                    init_params[attribute] = hyperparameters[value.name]
+
+        return cls(sagemaker_session=sagemaker_session, **init_params)
+
     def fit(self, records, mini_batch_size=None, **kwargs):
         """Fit this Estimator on serialized Record objects, stored in S3.
 

src/sagemaker/estimator.py

@@ -152,8 +152,47 @@ def fit(self, inputs, wait=True, logs=True, job_name=None):
         self.latest_training_job = _TrainingJob.start_new(self, inputs)
         if wait:
             self.latest_training_job.wait(logs=logs)
+
+
+    @classmethod
+    def _from_training_job(cls, init_params, hyperparameters, image, sagemaker_session):
+        raise NotImplementedError()
+
+    @classmethod
+    def attach(cls, training_job_name, sagemaker_session=None, **kwargs):
+        """Attach to an existing training job.
+
+        Create an Estimator bound to an existing training job. After attaching, if
+        the training job has a Complete status, it can be ``deploy()`` ed to create
+        a SageMaker Endpoint and return a ``Predictor``.
+
+        If the training job is in progress, attach will block and display log messages
+        from the training job, until the training job completes.
+
+        Args:
+            training_job_name (str): The name of the training job to attach to.
+            sagemaker_session (sagemaker.session.Session): Session object which manages interactions with
+                Amazon SageMaker APIs and any other AWS services needed. If not specified, the estimator creates one
+                using the default AWS configuration chain.
+            **kwargs: Additional kwargs passed to the :class:`~sagemaker.estimator.Estimator` constructor.
+
+        Returns:
+            sagemaker.estimator.Framework: ``Estimator`` with the attached training job.
+        """
+        sagemaker_session = sagemaker_session or Session()
+
+        if training_job_name is not None:
+            job_details = sagemaker_session.sagemaker_client.describe_training_job(TrainingJobName=training_job_name)
+            init_params, hp, image = cls._prepare_estimator_params_from_job_description(job_details)
+
         else:
-            raise NotImplemented('Asynchronous fit not available')
+            raise ValueError('must specify training_job name')
+
+        estimator = cls._from_training_job(init_params, hp, image, sagemaker_session)
+        estimator.latest_training_job = _TrainingJob(sagemaker_session=sagemaker_session,
+                                                     training_job_name=init_params['base_job_name'])
+        estimator.latest_training_job.wait()
+        return estimator
 
     def deploy(self, initial_instance_count, instance_type, endpoint_name=None, **kwargs):
         """Deploy the trained model to an Amazon SageMaker endpoint and return a ``sagemaker.RealTimePredictor`` object.
@@ -528,56 +567,24 @@ def hyperparameters(self):
         return self._json_encode_hyperparameters(self._hyperparameters)
 
     @classmethod
-    def attach(cls, training_job_name, sagemaker_session=None, **kwargs):
-        """Attach to an existing training job.
-
-        Create an Estimator bound to an existing training job. After attaching, if
-        the training job has a Complete status, it can be ``deploy()`` ed to create
-        a SageMaker Endpoint and return a ``Predictor``.
-
-        If the training job is in progress, attach will block and display log messages
-        from the training job, until the training job completes.
-
-        Args:
-            training_job_name (str): The name of the training job to attach to.
-            sagemaker_session (sagemaker.session.Session): Session object which manages interactions with
-                Amazon SageMaker APIs and any other AWS services needed. If not specified, the estimator creates one
-                using the default AWS configuration chain.
-            **kwargs: Additional kwargs passed to the :class:`~sagemaker.estimator.Estimator` constructor.
-
-        Returns:
-            sagemaker.estimator.Framework: ``Estimator`` with the attached training job.
-        """
-        sagemaker_session = sagemaker_session or Session()
-
-        if training_job_name is not None:
-            job_details = sagemaker_session.sagemaker_client.describe_training_job(TrainingJobName=training_job_name)
-            init_params, hp, _ = cls._prepare_estimator_params_from_job_description(job_details)
-
-        else:
-            # this case is only valid when called from inheriting class and then the class must declare framework
-            if not hasattr(cls, '__framework_name__'):
-                raise ValueError('must specify training_job name')
-            init_params = dict(kwargs)
-            hp = init_params.pop('hyperparameters')
+    def _from_training_job(cls, init_params, hyperparameters, image, sagemaker_session):
 
         # parameters for framework classes
         framework_init_params = dict()
-        framework_init_params['entry_point'] = json.loads(hp.get(SCRIPT_PARAM_NAME))
-        framework_init_params['source_dir'] = json.loads(hp.get(DIR_PARAM_NAME))
-        framework_init_params['enable_cloudwatch_metrics'] = json.loads(hp.get(CLOUDWATCH_METRICS_PARAM_NAME))
-        framework_init_params['container_log_level'] = json.loads(hp.get(CONTAINER_LOG_LEVEL_PARAM_NAME))
+        framework_init_params['entry_point'] = json.loads(hyperparameters.get(SCRIPT_PARAM_NAME))
+        framework_init_params['source_dir'] = json.loads(hyperparameters.get(DIR_PARAM_NAME))
+        framework_init_params['enable_cloudwatch_metrics'] = json.loads(
+            hyperparameters.get(CLOUDWATCH_METRICS_PARAM_NAME))
+        framework_init_params['container_log_level'] = json.loads(
+            hyperparameters.get(CONTAINER_LOG_LEVEL_PARAM_NAME))
 
         # drop json and remove other SageMaker specific additions
-        hyperparameters = {entry: json.loads(hp[entry]) for entry in hp}
-        framework_init_params['hyperparameters'] = hyperparameters
+        hp_map = {entry: json.loads(hyperparameters[entry]) for entry in hyperparameters}
+        framework_init_params['hyperparameters'] = hp_map
 
         init_params.update(framework_init_params)
 
         estimator = cls(sagemaker_session=sagemaker_session, **init_params)
-        estimator.latest_training_job = _TrainingJob(sagemaker_session=sagemaker_session,
-                                                     training_job_name=init_params['base_job_name'])
-        estimator.latest_training_job.wait()
         estimator.uploaded_code = UploadedCode(estimator.source_dir, estimator.entry_point)
         return estimator
 

src/sagemaker/mxnet/estimator.py

@@ -14,7 +14,6 @@
 from sagemaker.estimator import Framework
 from sagemaker.fw_utils import create_image_uri, framework_name_from_image
 from sagemaker.mxnet.model import MXNetModel
-from sagemaker.session import Session
 
 
 class MXNet(Framework):
@@ -83,42 +82,12 @@ def create_model(self, model_server_workers=None):
                           sagemaker_session=self.sagemaker_session)
 
     @classmethod
-    def attach(cls, training_job_name, sagemaker_session=None):
-        """Attach to an existing training job.
-
-        Create an ``Estimator`` bound to an existing training job. After attaching, if
-        the training job is in a Complete status, it can be ``deploy``ed to create
-        a SageMaker ``Endpoint`` and return a ``Predictor``.
-
-        If the training job is in progress, attach will block and display log messages
-        from the training job, until the training job completes.
-
-        Args:
-            training_job_name (str): The name of the training job to attach to.
-            sagemaker_session (sagemaker.session.Session): Session object which manages interactions with
-                Amazon SageMaker APIs and any other AWS services needed. If not specified, the estimator creates one
-                using the default AWS configuration chain.
-
-        Returns:
-            sagemaker.mxnet.estimator.MXNet: ``Estimator`` with the attached training job.
-
-        Raises:
-            ValueError: If `training_job_name` is None or the image name does not match the framework.
-        """
-        sagemaker_session = sagemaker_session or Session()
-
-        if training_job_name is None:
-            raise ValueError("must specify training_job name")
-
-        job_details = sagemaker_session.sagemaker_client.describe_training_job(TrainingJobName=training_job_name)
-        init_params, hp, image = cls._prepare_estimator_params_from_job_description(job_details)
-
-        init_params.update({'hyperparameters': hp})
-
+    def _from_training_job(cls, init_params, hyperparameters, image, sagemaker_session):
         framework, py_version = framework_name_from_image(image)
         init_params.update({'py_version': py_version})
 
+        training_job_name = init_params['base_job_name']
         if framework != cls.__framework_name__:
             raise ValueError("Training job: {} didn't use image for requested framework".format(training_job_name))
 
-        return super(MXNet, cls).attach(training_job_name=None, sagemaker_session=sagemaker_session, **init_params)
+        return super(MXNet, cls)._from_training_job(init_params, hyperparameters, image, sagemaker_session)

src/sagemaker/tensorflow/estimator.py

@@ -20,7 +20,6 @@
 import sagemaker.tensorflow
 from sagemaker.estimator import Framework
 from sagemaker.fw_utils import create_image_uri, framework_name_from_image
-from sagemaker.session import Session
 from sagemaker.tensorflow.model import TensorFlowModel
 
 logging.basicConfig()
@@ -166,48 +165,19 @@ def fit_super():
             fit_super()
 
     @classmethod
-    def attach(cls, training_job_name, sagemaker_session=None):
-        """Attach to an existing training job.
-
-        Create an ``Estimator`` bound to an existing training job. After attaching, if
-        the training job is in a Complete status, it can be ``deploy``ed to create
-        a SageMaker ``Endpoint`` and return a ``Predictor``.
-
-        If the training job is in progress, attach will block and display log messages
-        from the training job, until the training job completes.
-
-        Args:
-            training_job_name (str): The name of the training job to attach to.
-            sagemaker_session (sagemaker.session.Session): Session object which manages interactions with
-                Amazon SageMaker APIs and any other AWS services needed. If not specified, the estimator creates one
-                using the default AWS configuration chain.
-
-        Returns:
-            sagemaker.tensorflow.estimator.TensorFlow: ``Estimator`` with the attached training job.
-
-        Raises:
-            ValueError: If `training_job_name` is None or the image name does not match the framework.
-        """
-        sagemaker_session = sagemaker_session or Session()
-
-        if training_job_name is None:
-            raise ValueError("must specify training_job name")
-
-        job_details = sagemaker_session.sagemaker_client.describe_training_job(TrainingJobName=training_job_name)
-        init_params, hp, image = cls._prepare_estimator_params_from_job_description(job_details)
-
-        updated_params = cls._update_init_params(hp, ['checkpoint_path', 'training_steps', 'evaluation_steps'])
+    def _from_training_job(cls, init_params, hyperparameters, image, sagemaker_session):
+        updated_params = cls._update_init_params(hyperparameters,
+                                                 ['checkpoint_path', 'training_steps', 'evaluation_steps'])
         init_params.update(updated_params)
 
-        init_params.update({'hyperparameters': hp})
-
         framework, py_version = framework_name_from_image(image)
         init_params.update({'py_version': py_version})
+        training_job_name = init_params['base_job_name']
 
         if framework != cls.__framework_name__:
             raise ValueError("Training job: {} didn't use image for requested framework".format(training_job_name))
 
-        return super(TensorFlow, cls).attach(training_job_name=None, sagemaker_session=sagemaker_session, **init_params)
+        return super(TensorFlow, cls)._from_training_job(init_params, hyperparameters, image, sagemaker_session)
 
     def train_image(self):
         """Return the Docker image to use for training.

tests/integ/test_kmeans.py

@@ -16,14 +16,18 @@
 
 import boto3
 import os
+import time
 
 import sagemaker
 from sagemaker import KMeans, KMeansModel
 from sagemaker.utils import name_from_base
 from tests.integ import DATA_DIR, REGION
 from tests.integ.timeout import timeout, timeout_and_delete_endpoint_by_name
 
+import pytest
 
+
+@pytest.mark.skip(reason="no way of currently testing this")
 def test_kmeans():
 
     with timeout(minutes=15):
@@ -60,3 +64,49 @@ def test_kmeans():
         for record in result:
             assert record.label["closest_cluster"] is not None
             assert record.label["distance_to_cluster"] is not None
+
+
+def test_async_kmeans():
+
+    training_job_name = ""
+    endpoint_name = name_from_base('kmeans')
+
+    with timeout(minutes=15):
+        sagemaker_session = sagemaker.Session(boto_session=boto3.Session(region_name=REGION))
+        data_path = os.path.join(DATA_DIR, 'one_p_mnist', 'mnist.pkl.gz')
+        pickle_args = {} if sys.version_info.major == 2 else {'encoding': 'latin1'}
+
+        # Load the data into memory as numpy arrays
+        with gzip.open(data_path, 'rb') as f:
+            train_set, _, _ = pickle.load(f, **pickle_args)
+
+        kmeans = KMeans(role='SageMakerRole', train_instance_count=1,
+                        train_instance_type='ml.c4.xlarge',
+                        k=10, sagemaker_session=sagemaker_session, base_job_name='test-kmeans')
+
+        kmeans.init_method = 'random'
+        kmeans.max_iterators = 1
+        kmeans.tol = 1
+        kmeans.num_trials = 1
+        kmeans.local_init_method = 'kmeans++'
+        kmeans.half_life_time_size = 1
+        kmeans.epochs = 1
+        kmeans.center_factor = 1
+
+        kmeans.fit(kmeans.record_set(train_set[0][:100]), wait=False)
+        training_job_name = kmeans.latest_training_job.name
+
+        print("Detached from training job. Will re-attach in 20 seconds")
+        time.sleep(20)
+        print("attaching now...")
+
+    with timeout_and_delete_endpoint_by_name(endpoint_name, sagemaker_session, minutes=20):
+        estimator = KMeans.attach(training_job_name=training_job_name, sagemaker_session=sagemaker_session)
+        model = KMeansModel(estimator.model_data, role='SageMakerRole', sagemaker_session=sagemaker_session)
+        predictor = model.deploy(1, 'ml.c4.xlarge', endpoint_name=endpoint_name)
+        result = predictor.predict(train_set[0][:10])
+
+        assert len(result) == 10
+        for record in result:
+            assert record.label["closest_cluster"] is not None
+            assert record.label["distance_to_cluster"] is not None