Refactor attach to remove _from_training_job()

_prepare_init_params_from_job_description() is now a classmethod instead
of being a static method. Each class is responsible to implement their
specific logic to convert a training job description into arguments that
can be passed to its own __init__()

Author: Ignacio Quintero

src/sagemaker/amazon/amazon_estimator.py

@@ -65,28 +65,29 @@ def data_location(self, data_location):
         self._data_location = data_location
 
     @classmethod
-    def _from_training_job(cls, init_params, hyperparameters, image, sagemaker_session):
-        """Create an Estimator from existing training job data.
+    def _prepare_init_params_from_job_description(cls, job_details):
+        """Convert the job description to init params that can be handled by the class constructor
 
         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.
+            job_details: the returned job details from a describe_training_job API call.
 
-        Returns: An instance of the calling Estimator Class.
+        Returns:
+             dictionary: The transformed init_params
 
         """
+        init_params = super(AmazonAlgorithmEstimatorBase, cls)._prepare_init_params_from_job_description(job_details)
 
         # 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]
+                if value.name in init_params['hyperparameters']:
+                    init_params[attribute] = init_params['hyperparameters'][value.name]
 
-        return cls(sagemaker_session=sagemaker_session, **init_params)
+        del init_params['hyperparameters']
+        del init_params['image']
+        return 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

@@ -169,7 +169,7 @@ def _from_training_job(cls, init_params, hyperparameters, image, sagemaker_sessi
         raise NotImplementedError()
 
     @classmethod
-    def attach(cls, training_job_name, sagemaker_session=None):
+    def attach(cls, training_job_name, sagemaker_session=None, job_details=None):
         """Attach to an existing training job.
 
         Create an Estimator bound to an existing training job, each subclass is responsible to implement
@@ -185,6 +185,7 @@ 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.
+            training_job_details (
 
         Examples:
             >>> my_estimator.fit(wait=False)
@@ -198,13 +199,10 @@ def attach(cls, training_job_name, sagemaker_session=None):
         """
         sagemaker_session = sagemaker_session or Session()
 
-        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')
+        job_details = sagemaker_session.sagemaker_client.describe_training_job(TrainingJobName=training_job_name)
+        init_params = cls._prepare_init_params_from_job_description(job_details)
 
-        estimator = cls._from_training_job(init_params, hp, image, sagemaker_session)
+        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()
@@ -257,21 +255,33 @@ def create_model(self, **kwargs):
         """
         pass
 
-    @staticmethod
-    def _prepare_estimator_params_from_job_description(job_details):
-        estimator_params = dict()
+    @classmethod
+    def _prepare_init_params_from_job_description(cls, job_details):
+        """Convert the job description to init params that can be handled by the class constructor
+
+        Args:
+            job_details: the returned job details from a describe_training_job API call.
+
+        Returns:
+             dictionary: The transformed init_params
 
-        estimator_params['role'] = job_details['RoleArn']
-        estimator_params['train_instance_count'] = job_details['ResourceConfig']['InstanceCount']
-        estimator_params['train_instance_type'] = job_details['ResourceConfig']['InstanceType']
-        estimator_params['train_volume_size'] = job_details['ResourceConfig']['VolumeSizeInGB']
-        estimator_params['train_max_run'] = job_details['StoppingCondition']['MaxRuntimeInSeconds']
-        estimator_params['input_mode'] = job_details['AlgorithmSpecification']['TrainingInputMode']
-        estimator_params['base_job_name'] = job_details['TrainingJobName']
-        estimator_params['output_path'] = job_details['OutputDataConfig']['S3OutputPath']
-        estimator_params['output_kms_key'] = job_details['OutputDataConfig']['KmsKeyId']
+        """
+        init_params = dict()
+
+        init_params['role'] = job_details['RoleArn']
+        init_params['train_instance_count'] = job_details['ResourceConfig']['InstanceCount']
+        init_params['train_instance_type'] = job_details['ResourceConfig']['InstanceType']
+        init_params['train_volume_size'] = job_details['ResourceConfig']['VolumeSizeInGB']
+        init_params['train_max_run'] = job_details['StoppingCondition']['MaxRuntimeInSeconds']
+        init_params['input_mode'] = job_details['AlgorithmSpecification']['TrainingInputMode']
+        init_params['base_job_name'] = job_details['TrainingJobName']
+        init_params['output_path'] = job_details['OutputDataConfig']['S3OutputPath']
+        init_params['output_kms_key'] = job_details['OutputDataConfig']['KmsKeyId']
+
+        init_params['hyperparameters'] = job_details['HyperParameters']
+        init_params['image'] = job_details['AlgorithmSpecification']['TrainingImage']
 
-        return estimator_params, job_details['HyperParameters'], job_details['AlgorithmSpecification']['TrainingImage']
+        return init_params
 
     def delete_endpoint(self):
         """Delete an Amazon SageMaker ``Endpoint``.
@@ -388,7 +398,8 @@ class Estimator(EstimatorBase):
 
     def __init__(self, image_name, role, train_instance_count, train_instance_type,
                  train_volume_size=30, train_max_run=24 * 60 * 60, input_mode='File',
-                 output_path=None, output_kms_key=None, base_job_name=None, sagemaker_session=None):
+                 output_path=None, output_kms_key=None, base_job_name=None, sagemaker_session=None,
+                 hyperparameters=None):
         """Initialize an ``Estimator`` instance.
 
         Args:
@@ -420,9 +431,10 @@ def __init__(self, image_name, role, train_instance_count, train_instance_type,
             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.
+            hyperparameters (dict): Dictionary containing the hyperparameters to initialize this estimator with.
         """
         self.image_name = image_name
-        self.hyperparam_dict = {}
+        self.hyperparam_dict = hyperparameters.copy() if hyperparameters else {}
         super(Estimator, self).__init__(role, train_instance_count, train_instance_type,
                                         train_volume_size, train_max_run, input_mode,
                                         output_path, output_kms_key, base_job_name, sagemaker_session)
@@ -478,23 +490,20 @@ def predict_wrapper(endpoint, 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.
+    def _prepare_init_params_from_job_description(cls, job_details):
+        """Convert the job description to init params that can be handled by the class constructor
 
         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.
+            job_details: the returned job details from a describe_training_job API call.
 
-        Returns: An instance of the calling Estimator Class.
+        Returns:
+             dictionary: The transformed init_params
 
         """
+        init_params = super(Estimator, cls)._prepare_init_params_from_job_description(job_details)
 
-        estimator = cls(sagemaker_session=sagemaker_session, **init_params)
-        cls.set_hyperparameters(**hyperparameters)
-
-        return estimator
+        init_params['image_name'] = init_params.pop('image')
+        return init_params
 
 
 class Framework(EstimatorBase):
@@ -602,35 +611,32 @@ def hyperparameters(self):
         return self._json_encode_hyperparameters(self._hyperparameters)
 
     @classmethod
-    def _from_training_job(cls, init_params, hyperparameters, image, sagemaker_session):
-        """Create an Estimator from existing training job data.
+    def _prepare_init_params_from_job_description(cls, job_details):
+        """Convert the job description to init params that can be handled by the class constructor
 
         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.
+            job_details: the returned job details from a describe_training_job API call.
 
-        Returns: An instance of the calling Estimator Class.
+        Returns:
+             dictionary: The transformed init_params
 
         """
+        init_params = super(Framework, cls)._prepare_init_params_from_job_description(job_details)
 
-        # parameters for framework classes
-        framework_init_params = dict()
-        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))
+        init_params['entry_point'] = json.loads(init_params['hyperparameters'].get(SCRIPT_PARAM_NAME))
+        init_params['source_dir'] = json.loads(init_params['hyperparameters'].get(DIR_PARAM_NAME))
+        init_params['enable_cloudwatch_metrics'] = json.loads(
+            init_params['hyperparameters'].get(CLOUDWATCH_METRICS_PARAM_NAME))
+        init_params['container_log_level'] = json.loads(
+            init_params['hyperparameters'].get(CONTAINER_LOG_LEVEL_PARAM_NAME))
 
-        # drop json and remove other SageMaker specific additions
-        deserialized_hps = {entry: json.loads(hyperparameters[entry]) for entry in hyperparameters}
-        framework_init_params['hyperparameters'] = deserialized_hps
+        init_params['hyperparameters'] = {k: json.loads(v) for k, v in init_params['hyperparameters'].items()}
 
-        init_params.update(framework_init_params)
+        return init_params
 
-        estimator = cls(sagemaker_session=sagemaker_session, **init_params)
+    @classmethod
+    def attach(cls, training_job_name, sagemaker_session=None):
+        estimator = super(Framework, cls).attach(training_job_name, sagemaker_session)
         estimator.uploaded_code = UploadedCode(estimator.source_dir, estimator.entry_point)
         return estimator
 

src/sagemaker/mxnet/estimator.py

@@ -82,24 +82,23 @@ def create_model(self, model_server_workers=None):
                           sagemaker_session=self.sagemaker_session)
 
     @classmethod
-    def _from_training_job(cls, init_params, hyperparameters, image, sagemaker_session):
-        """Create an Estimator from existing training job data.
+    def _prepare_init_params_from_job_description(cls, job_details):
+        """Convert the job description to init params that can be handled by the class constructor
 
         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.
+            job_details: the returned job details from a describe_training_job API call.
 
-        Returns: An instance of the calling Estimator Class.
+        Returns:
+             dictionary: The transformed init_params
 
         """
+        init_params = super(MXNet, cls)._prepare_init_params_from_job_description(job_details)
+        framework, py_version = framework_name_from_image(init_params.pop('image'))
 
-        framework, py_version = framework_name_from_image(image)
-        init_params.update({'py_version': py_version})
-
+        init_params['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)._from_training_job(init_params, hyperparameters, image, sagemaker_session)
+        return init_params

src/sagemaker/tensorflow/estimator.py

@@ -11,12 +11,11 @@
 # ANY KIND, either express or implied. See the License for the specific
 # language governing permissions and limitations under the License.
 import logging
+import os
 import subprocess
 import tempfile
 import threading
 
-import os
-
 import sagemaker.tensorflow
 from sagemaker.estimator import Framework
 from sagemaker.fw_utils import create_image_uri, framework_name_from_image
@@ -168,31 +167,32 @@ def fit_super():
             fit_super()
 
     @classmethod
-    def _from_training_job(cls, init_params, hyperparameters, image, sagemaker_session):
-        """Create an Estimator from existing training job data.
+    def _prepare_init_params_from_job_description(cls, job_details):
+        """Convert the job description to init params that can be handled by the class constructor
 
         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.
+            job_details: the returned job details from a describe_training_job API call.
 
-        Returns: An instance of the calling Estimator Class.
+        Returns:
+             dictionary: The transformed init_params
 
         """
+        init_params = super(TensorFlow, cls)._prepare_init_params_from_job_description(job_details)
 
-        updated_params = cls._update_init_params(hyperparameters,
-                                                 ['checkpoint_path', 'training_steps', 'evaluation_steps'])
-        init_params.update(updated_params)
+        # Move some of the tensorflow specific init params from hyperparameters into the main init params.
+        for argument in ['checkpoint_path', 'training_steps', 'evaluation_steps']:
+            value = init_params['hyperparameters'].pop(argument, None)
+            if value is not None:
+                init_params[argument] = value
 
-        framework, py_version = framework_name_from_image(image)
-        init_params.update({'py_version': py_version})
-        training_job_name = init_params['base_job_name']
+        framework, py_version = framework_name_from_image(init_params.pop('image'))
+        init_params['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)._from_training_job(init_params, hyperparameters, image, sagemaker_session)
+        return init_params
 
     def train_image(self):
         """Return the Docker image to use for training.

tests/integ/test_byo_estimator.py

@@ -154,3 +154,5 @@ def test_async_byo_estimator():
         assert len(result['predictions']) == 10
         for prediction in result['predictions']:
             assert prediction['score'] is not None
+
+        assert estimator.train_image() == image_name

tests/unit/test_estimator.py

@@ -70,6 +70,12 @@ def train_image(self):
     def create_model(self):
         return DummyFrameworkModel(self.sagemaker_session)
 
+    @classmethod
+    def _prepare_init_params_from_job_description(cls, job_details):
+        init_params = super(DummyFramework, cls)._prepare_init_params_from_job_description(job_details)
+        init_params.pop("image", None)
+        return init_params
+
 
 class DummyFrameworkModel(FrameworkModel):
 
@@ -251,12 +257,6 @@ def test_attach_framework(sagemaker_session):
     assert framework_estimator.entry_point == 'iris-dnn-classifier.py'
 
 
-def test_attach_no_job_name_framework(sagemaker_session):
-    with pytest.raises(ValueError) as error:
-        Framework.attach(training_job_name=None, sagemaker_session=sagemaker_session)
-    assert 'must specify training_job name' in str(error)
-
-
 def test_fit_then_fit_again(sagemaker_session):
     fw = DummyFramework(entry_point=SCRIPT_PATH, role=ROLE, sagemaker_session=sagemaker_session,
                         train_instance_count=INSTANCE_COUNT, train_instance_type=INSTANCE_TYPE,

tests/unit/test_mxnet.py

@@ -201,9 +201,3 @@ def test_attach_wrong_framework(sagemaker_session):
     with pytest.raises(ValueError) as error:
         MXNet.attach(training_job_name='neo', sagemaker_session=sagemaker_session)
     assert "didn't use image for requested framework" in str(error)
-
-
-def test_attach_no_job_name(sagemaker_session):
-    with pytest.raises(ValueError) as error:
-        MXNet.attach(training_job_name=None, sagemaker_session=sagemaker_session)
-    assert "must specify training_job name" in str(error)

tests/unit/test_tf_estimator.py

@@ -379,9 +379,3 @@ def test_attach_wrong_framework(sagemaker_session):
     with pytest.raises(ValueError) as error:
         TensorFlow.attach(training_job_name='neo', sagemaker_session=sagemaker_session)
     assert "didn't use image for requested framework" in str(error)
-
-
-def test_attach_no_job_name(sagemaker_session):
-    with pytest.raises(ValueError) as error:
-        TensorFlow.attach(training_job_name=None, sagemaker_session=sagemaker_session)
-    assert "must specify training_job name" in str(error)