feedback and update tests

Author: shivlaks

src/stepfunctions/steps/sagemaker.py

@@ -12,6 +12,8 @@
 # permissions and limitations under the License.
 from __future__ import absolute_import
 
+import logging
+
 from enum import Enum
 from stepfunctions.inputs import ExecutionInput, StepInput
 from stepfunctions.steps.states import Task
@@ -23,6 +25,8 @@
 from sagemaker.model import Model, FrameworkModel
 from sagemaker.model_monitor import DataCaptureConfig
 
+logger = logging.getLogger('stepfunctions.sagemaker')
+
 SAGEMAKER_SERVICE_NAME = "sagemaker"
 
 
@@ -64,7 +68,7 @@ def __init__(self, state_id, estimator, job_name, data=None, hyperparameters=Non
                 * (list[sagemaker.amazon.amazon_estimator.RecordSet]) - A list of
                     :class:`sagemaker.amazon.amazon_estimator.RecordSet` objects,
                     where each instance is a different channel of training data.
-            hyperparameters (dict, optional): Specify the hyper parameters for the training. (Default: None)
+            hyperparameters (dict, optional): Specify the hyperparameters that are set before the model begins training. If hyperparameters provided are also specified in the estimator, the provided value will used. (Default: Hyperparameters specified in the estimator will be used for training.)
             mini_batch_size (int): Specify this argument only when estimator is a built-in estimator of an Amazon algorithm. For other estimators, batch size should be specified in the estimator.
             experiment_config (dict, optional): Specify the experiment config for the training. (Default: None)
             wait_for_completion (bool, optional): Boolean value set to `True` if the Task state should wait for the training job to complete before proceeding to the next step in the workflow. Set to `False` if the Task state should submit the training job and proceed to the next step. (default: True)
@@ -104,11 +108,9 @@ def __init__(self, state_id, estimator, job_name, data=None, hyperparameters=Non
             parameters['TrainingJobName'] = job_name
 
         if hyperparameters is not None:
-            merged_hyperparameters = {}
             if estimator.hyperparameters() is not None:
-                merged_hyperparameters.update(estimator.hyperparameters())
-            merged_hyperparameters.update(hyperparameters)
-            parameters['HyperParameters'] = merged_hyperparameters
+                hyperparameters = self.__merge_hyperparameters(hyperparameters, estimator.hyperparameters())
+            parameters['HyperParameters'] = hyperparameters
 
         if experiment_config is not None:
             parameters['ExperimentConfig'] = experiment_config
@@ -139,6 +141,26 @@ def get_expected_model(self, model_name=None):
         model.model_data = self.output()["ModelArtifacts"]["S3ModelArtifacts"]
         return model
 
+    """
+    Merges the hyperparameters supplied in the TrainingStep constructor with the hyperparameters
+    specified in the estimator. If there are duplicate entries, the value provided in the constructor
+    will be used.
+    """
+
+    def __merge_hyperparameters(self, training_step_hyperparameters, estimator_hyperparameters):
+        """
+        Args:
+            training_step_hyperparameters (dict): Hyperparameters supplied in the training step constructor
+            estimator_hyperparameters (dict): Hyperparameters specified in the estimator
+        """
+        merged_hyperparameters = estimator_hyperparameters.copy()
+        for key, value in training_step_hyperparameters.items():
+            if key in merged_hyperparameters:
+                logger.info(
+                    f"hyperparameter property: <{key}> with value: <{merged_hyperparameters[key]}> provided in the"
+                    f" estimator will be overwritten with value provided in constructor: <{value}>")
+            merged_hyperparameters[key] = value
+        return merged_hyperparameters
 
 class TransformStep(Task):
 

tests/unit/test_sagemaker_steps.py

@@ -484,7 +484,7 @@ def test_training_step_creation_with_framework(tensorflow_estimator):
 
 @patch('botocore.client.BaseClient._make_api_call', new=mock_boto_api_call)
 @patch.object(boto3.session.Session, 'region_name', 'us-east-1')
-def training_step_merges_hyperparameters_from_constructor_and_estimator(tensorflow_estimator):
+def test_training_step_merges_hyperparameters_from_constructor_and_estimator(tensorflow_estimator):
     step = TrainingStep('Training',
         estimator=tensorflow_estimator,
         data={'train': 's3://sagemaker/train'},
@@ -549,7 +549,7 @@ def training_step_merges_hyperparameters_from_constructor_and_estimator(tensorfl
 
 @patch('botocore.client.BaseClient._make_api_call', new=mock_boto_api_call)
 @patch.object(boto3.session.Session, 'region_name', 'us-east-1')
-def training_step_uses_constructor_hyperparameters_when_duplicates_supplied_in_estimator(tensorflow_estimator):
+def test_training_step_uses_constructor_hyperparameters_when_duplicates_supplied_in_estimator(tensorflow_estimator):
     step = TrainingStep('Training',
         estimator=tensorflow_estimator,
         data={'train': 's3://sagemaker/train'},