Allow hyperparameters in Tensorflow estimator to be parameterized for a pipeline

Author: Payton Staub

src/sagemaker/estimator.py

@@ -2443,7 +2443,13 @@ def attach(cls, training_job_name, sagemaker_session=None, model_channel_name="m
     @staticmethod
     def _json_encode_hyperparameters(hyperparameters):
         """Placeholder docstring"""
-        return {str(k): json.dumps(v) for (k, v) in hyperparameters.items()}
+        current_hyperparameters = hyperparameters
+        if current_hyperparameters is not None:
+            hyperparameters = {
+                str(k): (v if isinstance(v, (Parameter, Expression, Properties)) else json.dumps(v))
+                for (k, v) in current_hyperparameters.items()
+            }
+        return hyperparameters
 
     @classmethod
     def _update_init_params(cls, hp, tf_arguments):

tests/unit/sagemaker/workflow/test_steps.py

@@ -15,6 +15,7 @@
 
 import pytest
 import sagemaker
+import os
 
 from mock import (
     Mock,
@@ -24,6 +25,7 @@
 
 from sagemaker.debugger import ProfilerConfig
 from sagemaker.estimator import Estimator
+from sagemaker.tensorflow import TensorFlow
 from sagemaker.inputs import TrainingInput, TransformInput, CreateModelInput
 from sagemaker.model import Model
 from sagemaker.processing import (
@@ -45,6 +47,10 @@
     CreateModelStep,
     CacheConfig,
 )
+from tests.unit import DATA_DIR
+
+SCRIPT_FILE = "dummy_script.py"
+SCRIPT_PATH = os.path.join(DATA_DIR, SCRIPT_FILE)
 
 REGION = "us-west-2"
 BUCKET = "my-bucket"
@@ -112,7 +118,7 @@ def test_custom_step():
     assert step.to_request() == {"Name": "MyStep", "Type": "Training", "Arguments": dict()}
 
 
-def test_training_step(sagemaker_session):
+def test_training_step_base_estimator(sagemaker_session):
     instance_type_parameter = ParameterString(name="InstanceType", default_value="c4.4xlarge")
     instance_count_parameter = ParameterInteger(name="InstanceCount", default_value=1)
     data_source_uri_parameter = ParameterString(
@@ -177,6 +183,100 @@ def test_training_step(sagemaker_session):
     assert step.properties.TrainingJobName.expr == {"Get": "Steps.MyTrainingStep.TrainingJobName"}
 
 
+def test_training_step_tensorflow(sagemaker_session):
+    instance_type_parameter = ParameterString(name="InstanceType", default_value="ml.p3.16xlarge")
+    instance_count_parameter = ParameterInteger(name="InstanceCount", default_value=1)
+    data_source_uri_parameter = ParameterString(
+        name="DataSourceS3Uri", default_value=f"s3://{BUCKET}/train_manifest"
+    )
+    training_epochs_parameter = ParameterInteger(name="TrainingEpochs", default_value=5)
+    training_batch_size_parameter = ParameterInteger(name="TrainingBatchSize", default_value=500)
+    estimator = TensorFlow(
+        entry_point=os.path.join(DATA_DIR, SCRIPT_FILE),
+        role=ROLE,
+        model_dir=False,
+        image_uri=IMAGE_URI,
+        source_dir="s3://mybucket/source",
+        framework_version="2.4.1",
+        py_version="py37",
+        instance_count=instance_count_parameter,
+        instance_type=instance_type_parameter,
+        sagemaker_session=sagemaker_session,
+        # subnets=subnets,
+        hyperparameters={
+            "batch-size": training_batch_size_parameter,
+            "epochs": training_epochs_parameter,
+        },
+        # security_group_ids=security_group_ids,
+        debugger_hook_config=False,
+        # Training using SMDataParallel Distributed Training Framework
+        distribution={"smdistributed": {"dataparallel": {"enabled": True}}},
+    )
+
+    inputs = TrainingInput(s3_data=data_source_uri_parameter)
+    cache_config = CacheConfig(enable_caching=True, expire_after="PT1H")
+    step = TrainingStep(
+        name="MyTrainingStep", estimator=estimator, inputs=inputs, cache_config=cache_config
+    )
+    step_request = step.to_request()
+    step_request['Arguments']['HyperParameters'].pop('sagemaker_job_name', None)
+    step_request['Arguments'].pop('ProfilerRuleConfigurations', None)
+    assert step_request == {
+        'Name':'MyTrainingStep',
+        'Type':'Training',
+        'Arguments':{
+            'AlgorithmSpecification':{
+                'TrainingInputMode':'File',
+                'TrainingImage':'fakeimage',
+                'EnableSageMakerMetricsTimeSeries':True
+            },
+            'OutputDataConfig':{
+                'S3OutputPath':'s3://my-bucket/'
+            },
+            'StoppingCondition':{
+                'MaxRuntimeInSeconds':86400
+            },
+            'ResourceConfig':{
+                'InstanceCount':instance_count_parameter,
+                'InstanceType':instance_type_parameter,
+                'VolumeSizeInGB':30
+            },
+            'RoleArn':'DummyRole',
+            'InputDataConfig':[
+                {
+                    'DataSource':{
+                    'S3DataSource':{
+                        'S3DataType':'S3Prefix',
+                        'S3Uri':data_source_uri_parameter,
+                        'S3DataDistributionType':'FullyReplicated'
+                    }
+                    },
+                    'ChannelName':'training'
+                }
+            ],
+            'HyperParameters':{
+                'batch-size':training_batch_size_parameter,
+                'epochs':training_epochs_parameter,
+                'sagemaker_submit_directory':'"s3://mybucket/source"',
+                'sagemaker_program':'"/Volumes/Unix/workplace/pstaub/sagemaker-python-sdk/tests/unit/../data/dummy_script.py"',
+                'sagemaker_container_log_level':'20',
+                'sagemaker_region':'"us-west-2"',
+                'sagemaker_distributed_dataparallel_enabled':'true',
+                'sagemaker_instance_type':instance_type_parameter,
+                'sagemaker_distributed_dataparallel_custom_mpi_options':'""'
+            },
+            'ProfilerConfig':{
+                'S3OutputPath':'s3://my-bucket/'
+            }
+        },
+        'CacheConfig':{
+            'Enabled':True,
+            'ExpireAfter':'PT1H'
+        }
+    }
+    assert step.properties.TrainingJobName.expr == {"Get": "Steps.MyTrainingStep.TrainingJobName"}
+
+
 def test_processing_step(sagemaker_session):
     processing_input_data_uri_parameter = ParameterString(
         name="ProcessingInputDataUri", default_value=f"s3://{BUCKET}/processing_manifest"