Add BYOA implementation and missing docs.

BYO was missing an implementation of _from_training_job(). This adds
that as well as an integration test to verify that.

Also addressed the PR comments and added information to the README.

Author: Ignacio Quintero

README.rst

@@ -97,7 +97,7 @@ SageMaker Python SDK provides several high-level abstractions for working with A
 - **Estimators**: Encapsulate training on SageMaker. Can be ``fit()`` to run training, then the resulting model ``deploy()`` ed to a SageMaker Endpoint.
 - **Models**: Encapsulate built ML models. Can be ``deploy()`` ed to a SageMaker Endpoint.
 - **Predictors**: Provide real-time inference and transformation using Python data-types against a SageMaker Endpoint.
-- **Session**: Provides a collection of convience methods for working with SageMaker resources.
+- **Session**: Provides a collection of convenience methods for working with SageMaker resources.
 
 Estimator and Model implementations for MXNet, TensorFlow, and Amazon ML algorithms are included. There's also an Estimator that runs SageMaker compatible custom Docker containers, allowing you to run your own ML algorithms via SageMaker Python SDK.
 
@@ -1149,7 +1149,8 @@ Optional arguments
 ''''''''''''''''''
 
 -  ``wait (bool)``: Defaults to True, whether to block and wait for the
-   training script to complete before returning.
+   training script to complete before returning. If set to False, it will return immediately, and can later be
+   `attach`ed to.
 -  ``logs (bool)``: Defaults to True, whether to show logs produced by training
    job in the Python session. Only meaningful when wait is True.
 - ``run_tensorboard_locally (bool)``: Defaults to False. Executes TensorBoard in a different
@@ -1178,9 +1179,25 @@ the ``TensorFlow`` estimator parameter ``training_steps`` is finished or when th
 job execution time reaches the ``TensorFlow`` estimator parameter ``train_max_run``.
 
 When the training job finishes, a `TensorFlow serving <https://www.tensorflow.org/serving/serving_basic>`_
-with the result of the training is generated and saved to the S3 location define by
+with the result of the training is generated and saved to the S3 location defined by
 the ``TensorFlow`` estimator parameter ``output_path``.
 
+
+If the ``wait=False`` flag is passed to ``fit``, then it will return immediately. The training job will continue running
+asynchronously. At a later time, a Tensorflow Estimator can be obtained by attaching to the existing training job. If
+the training job is not finished it will start showing the standard output of training and wait until it completes.
+After attaching, the estimator can be deployed as usual.
+
+.. code:: python
+
+    tf_estimator.fit(your_input_data, wait=False)
+    training_job_name = tf_estimator.latest_training_job.name
+
+    # after some time, or in a separate python notebook, we can attach to it again.
+
+    tf_estimator = TensorFlow.attach(training_job_name=training_job_name)
+
+
 The evaluation process
 """"""""""""""""""""""
 
@@ -1244,6 +1261,8 @@ You can access TensorBoard locally at http://localhost:6006 or using your SakeMa
 `https*workspace_base_url*proxy/6006/ <proxy/6006/>`_ (TensorBoard will not work if you forget to put the slash,
 '/', in end of the url). If TensorBoard started on a different port, adjust these URLs to match.
 
+Note that TensorBoard is not supported when passing wait=False to ``fit``.
+
 
 Deploying TensorFlow Serving models
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

src/sagemaker/amazon/amazon_estimator.py

@@ -66,11 +66,21 @@ def data_location(self, data_location):
 
     @classmethod
     def _from_training_job(cls, init_params, hyperparameters, image, sagemaker_session):
+        """Create an Estimator from existing training job data.
+
+        Args:
+            init_params (dict): The init_params the training job was created with.
+            hyperparameters (dict):  The hyperparameters the training job was created with.
+            image (str): Container image (if any) the training job was created with
+            sagemaker_session (sagemaker.session.Session): A sagemaker Session to pass to the estimator.
+
+        Returns: An instance of the calling Estimator Class.
+
+        """
 
         # 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:

src/sagemaker/estimator.py

@@ -155,14 +155,26 @@ def fit(self, inputs, wait=True, logs=True, job_name=None):
 
     @classmethod
     def _from_training_job(cls, init_params, hyperparameters, image, sagemaker_session):
+        """Create an Estimator from existing training job data.
+
+        Args:
+            init_params (dict): The init_params the training job was created with.
+            hyperparameters (dict):  The hyperparameters the training job was created with.
+            image (str): Container image (if any) the training job was created with
+            sagemaker_session (sagemaker.session.Session): A sagemaker Session to pass to the estimator.
+
+        Returns: An instance of the calling Estimator Class.
+
+        """
         raise NotImplementedError()
 
     @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 has a Complete status, it can be ``deploy()`` ed to create
+        Create an Estimator bound to an existing training job, each subclass is responsible to implement
+        ``from_training_job()`` as this method delegates the actual Estimator creation to it. 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
@@ -173,17 +185,22 @@ def attach(cls, training_job_name, sagemaker_session=None):
             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.
+
+        Examples:
+            >>> my_estimator.fit(wait=False)
+            >>> training_job_name = my_estimator.latest_training_job.name
+            Later on:
+            >>> attached_estimator = Estimator.attach(training_job_name)
+            >>> attached_estimator.deploy()
 
         Returns:
-            sagemaker.estimator.Framework: ``Estimator`` with the attached training job.
+            Instance of the calling ``Estimator`` Class with the attached training job.
         """
         sagemaker_session = sagemaker_session or Session()
 
-        if training_job_name is not None:
+        if 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)
-
         else:
             raise ValueError('must specify training_job name')
 
@@ -460,6 +477,25 @@ def predict_wrapper(endpoint, session):
         return Model(self.model_data, image or self.train_image(), self.role, sagemaker_session=self.sagemaker_session,
                      predictor_cls=predictor_cls, **kwargs)
 
+    @classmethod
+    def _from_training_job(cls, init_params, hyperparameters, image, sagemaker_session):
+        """Create an Estimator from existing training job data.
+
+        Args:
+            init_params (dict): The init_params the training job was created with.
+            hyperparameters (dict):  The hyperparameters the training job was created with.
+            image (str): Container image (if any) the training job was created with
+            sagemaker_session (sagemaker.session.Session): A sagemaker Session to pass to the estimator.
+
+        Returns: An instance of the calling Estimator Class.
+
+        """
+
+        estimator = cls(sagemaker_session=sagemaker_session, **init_params)
+        cls.set_hyperparameters(**hyperparameters)
+
+        return estimator
+
 
 class Framework(EstimatorBase):
     """Base class that cannot be instantiated directly.
@@ -567,6 +603,17 @@ def hyperparameters(self):
 
     @classmethod
     def _from_training_job(cls, init_params, hyperparameters, image, sagemaker_session):
+        """Create an Estimator from existing training job data.
+
+        Args:
+            init_params (dict): The init_params the training job was created with.
+            hyperparameters (dict):  The hyperparameters the training job was created with.
+            image (str): Container image (if any) the training job was created with
+            sagemaker_session (sagemaker.session.Session): A sagemaker Session to pass to the estimator.
+
+        Returns: An instance of the calling Estimator Class.
+
+        """
 
         # parameters for framework classes
         framework_init_params = dict()
@@ -578,8 +625,8 @@ def _from_training_job(cls, init_params, hyperparameters, image, sagemaker_sessi
             hyperparameters.get(CONTAINER_LOG_LEVEL_PARAM_NAME))
 
         # drop json and remove other SageMaker specific additions
-        hp_map = {entry: json.loads(hyperparameters[entry]) for entry in hyperparameters}
-        framework_init_params['hyperparameters'] = hp_map
+        deserialized_hps = {entry: json.loads(hyperparameters[entry]) for entry in hyperparameters}
+        framework_init_params['hyperparameters'] = deserialized_hps
 
         init_params.update(framework_init_params)
 

src/sagemaker/mxnet/estimator.py

@@ -83,6 +83,18 @@ def create_model(self, model_server_workers=None):
 
     @classmethod
     def _from_training_job(cls, init_params, hyperparameters, image, sagemaker_session):
+        """Create an Estimator from existing training job data.
+
+        Args:
+            init_params (dict): The init_params the training job was created with.
+            hyperparameters (dict):  The hyperparameters the training job was created with.
+            image (str): Container image (if any) the training job was created with
+            sagemaker_session (sagemaker.session.Session): A sagemaker Session to pass to the estimator.
+
+        Returns: An instance of the calling Estimator Class.
+
+        """
+
         framework, py_version = framework_name_from_image(image)
         init_params.update({'py_version': py_version})
 

src/sagemaker/tensorflow/estimator.py

@@ -169,6 +169,18 @@ def fit_super():
 
     @classmethod
     def _from_training_job(cls, init_params, hyperparameters, image, sagemaker_session):
+        """Create an Estimator from existing training job data.
+
+        Args:
+            init_params (dict): The init_params the training job was created with.
+            hyperparameters (dict):  The hyperparameters the training job was created with.
+            image (str): Container image (if any) the training job was created with
+            sagemaker_session (sagemaker.session.Session): A sagemaker Session to pass to the estimator.
+
+        Returns: An instance of the calling Estimator Class.
+
+        """
+
         updated_params = cls._update_init_params(hyperparameters,
                                                  ['checkpoint_path', 'training_steps', 'evaluation_steps'])
         init_params.update(updated_params)

tests/integ/test_byo_estimator.py

@@ -29,6 +29,13 @@
 from tests.integ.timeout import timeout, timeout_and_delete_endpoint_by_name
 
 
+def fm_serializer(data):
+    js = {'instances': []}
+    for row in data:
+        js['instances'].append({'features': row.tolist()})
+    return json.dumps(js)
+
+
 def test_byo_estimator():
     """Use Factorization Machines algorithm as an example here.
 
@@ -79,12 +86,6 @@ def test_byo_estimator():
 
     endpoint_name = name_from_base('byo')
 
-    def fm_serializer(data):
-        js = {'instances': []}
-        for row in data:
-            js['instances'].append({'features': row.tolist()})
-        return json.dumps(js)
-
     with timeout_and_delete_endpoint_by_name(endpoint_name, sagemaker_session, minutes=20):
         model = estimator.create_model()
         predictor = model.deploy(1, 'ml.m4.xlarge', endpoint_name=endpoint_name)
@@ -97,3 +98,59 @@ def fm_serializer(data):
         assert len(result['predictions']) == 10
         for prediction in result['predictions']:
             assert prediction['score'] is not None
+
+
+def test_async_byo_estimator():
+    image_name = registry(REGION) + "/factorization-machines:1"
+    endpoint_name = name_from_base('byo')
+    training_job_name = ""
+
+    with timeout(minutes=5):
+        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'}
+
+        with gzip.open(data_path, 'rb') as f:
+            train_set, _, _ = pickle.load(f, **pickle_args)
+
+        # take 100 examples for faster execution
+        vectors = np.array([t.tolist() for t in train_set[0][:100]]).astype('float32')
+        labels = np.where(np.array([t.tolist() for t in train_set[1][:100]]) == 0, 1.0, 0.0).astype('float32')
+
+        buf = io.BytesIO()
+        write_numpy_to_dense_tensor(buf, vectors, labels)
+        buf.seek(0)
+
+        bucket = sagemaker_session.default_bucket()
+        prefix = 'test_byo_estimator'
+        key = 'recordio-pb-data'
+        boto3.resource('s3').Bucket(bucket).Object(os.path.join(prefix, 'train', key)).upload_fileobj(buf)
+        s3_train_data = 's3://{}/{}/train/{}'.format(bucket, prefix, key)
+
+        estimator = Estimator(image_name=image_name,
+                              role='SageMakerRole', train_instance_count=1,
+                              train_instance_type='ml.c4.xlarge',
+                              sagemaker_session=sagemaker_session, base_job_name='test-byo')
+
+        estimator.set_hyperparameters(num_factors=10,
+                                      feature_dim=784,
+                                      mini_batch_size=100,
+                                      predictor_type='binary_classifier')
+
+        # training labels must be 'float32'
+        estimator.fit({'train': s3_train_data}, wait=False)
+        training_job_name = estimator.latest_training_job.name
+
+    with timeout_and_delete_endpoint_by_name(endpoint_name, sagemaker_session, minutes=30):
+        estimator = Estimator.attach(training_job_name=training_job_name, sagemaker_session=sagemaker_session)
+        model = estimator.create_model()
+        predictor = model.deploy(1, 'ml.m4.xlarge', endpoint_name=endpoint_name)
+        predictor.serializer = fm_serializer
+        predictor.content_type = 'application/json'
+        predictor.deserializer = sagemaker.predictor.json_deserializer
+
+        result = predictor.predict(train_set[0][:10])
+
+        assert len(result['predictions']) == 10
+        for prediction in result['predictions']:
+            assert prediction['score'] is not None