fix: estimator hyperparameters in script mode (aws#3344)

Author: evakravi

src/sagemaker/estimator.py

@@ -2629,6 +2629,8 @@ def __init__(
             **kwargs,
         )
 
+        self.set_hyperparameters(**self._hyperparameters)
+
     def training_image_uri(self):
         """Returns the docker image to use for training.
 
@@ -2644,9 +2646,15 @@ def set_hyperparameters(self, **kwargs):
         training code on SageMaker. For convenience, this accepts other types
         for keys and values, but ``str()`` will be called to convert them before
         training.
+
+        If a source directory is specified, this method escapes the dict argument as JSON,
+        and updates the private hyperparameter attribute.
         """
-        for k, v in kwargs.items():
-            self._hyperparameters[k] = v
+        if self.source_dir:
+            self._hyperparameters.update(EstimatorBase._json_encode_hyperparameters(kwargs))
+        else:
+            for k, v in kwargs.items():
+                self._hyperparameters[k] = v
 
     def hyperparameters(self):
         """Returns the hyperparameters as a dictionary to use for training.

tests/unit/test_estimator.py

@@ -11,6 +11,7 @@
 # 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
+from copy import deepcopy
 
 import logging
 import json
@@ -3825,6 +3826,12 @@ def test_script_mode_estimator_same_calls_as_framework(
 
     model_uri = "s3://someprefix2/models/model.tar.gz"
     training_data_uri = "s3://bucket/mydata"
+    hyperparameters = {
+        "int_hyperparam": 1,
+        "string_hyperparam": "hello",
+        "stringified_numeric_hyperparam": "44",
+        "float_hyperparam": 1.234,
+    }
 
     generic_estimator = Estimator(
         entry_point=SCRIPT_PATH,
@@ -3838,6 +3845,7 @@ def test_script_mode_estimator_same_calls_as_framework(
         model_uri=model_uri,
         dependencies=[],
         debugger_hook_config={},
+        hyperparameters=deepcopy(hyperparameters),
     )
     generic_estimator.fit(training_data_uri)
 
@@ -3858,6 +3866,7 @@ def test_script_mode_estimator_same_calls_as_framework(
         model_uri=model_uri,
         dependencies=[],
         debugger_hook_config={},
+        hyperparameters=deepcopy(hyperparameters),
     )
     framework_estimator.fit(training_data_uri)
 
@@ -4394,3 +4403,51 @@ def test_insert_invalid_source_code_args():
     assert (
         "The entry_point should not be a pipeline variable " "when source_dir is a local path"
     ) in str(err.value)
+
+
+@patch("time.time", return_value=TIME)
+@patch("sagemaker.estimator.tar_and_upload_dir")
+@patch("sagemaker.model.Model._upload_code")
+def test_script_mode_estimator_escapes_hyperparameters_as_json(
+    patched_upload_code, patched_tar_and_upload_dir, sagemaker_session
+):
+    patched_tar_and_upload_dir.return_value = UploadedCode(
+        s3_prefix="s3://%s/%s" % ("bucket", "key"), script_name="script_name"
+    )
+    sagemaker_session.boto_region_name = REGION
+
+    instance_type = "ml.p2.xlarge"
+    instance_count = 1
+
+    training_data_uri = "s3://bucket/mydata"
+
+    jumpstart_source_dir = f"s3://{list(JUMPSTART_BUCKET_NAME_SET)[0]}/source_dirs/source.tar.gz"
+
+    hyperparameters = {
+        "int_hyperparam": 1,
+        "string_hyperparam": "hello",
+        "stringified_numeric_hyperparam": "44",
+        "float_hyperparam": 1.234,
+    }
+
+    generic_estimator = Estimator(
+        entry_point=SCRIPT_PATH,
+        role=ROLE,
+        region=REGION,
+        sagemaker_session=sagemaker_session,
+        instance_count=instance_count,
+        instance_type=instance_type,
+        source_dir=jumpstart_source_dir,
+        image_uri=IMAGE_URI,
+        model_uri=MODEL_DATA,
+        hyperparameters=hyperparameters,
+    )
+    generic_estimator.fit(training_data_uri)
+
+    formatted_hyperparams = EstimatorBase._json_encode_hyperparameters(hyperparameters)
+
+    assert (
+        set(formatted_hyperparams.items())
+        - set(sagemaker_session.train.call_args_list[0][1]["hyperparameters"].items())
+        == set()
+    )