feat: Support placeholders for TuningStep parameters

Author: ca-nguyen

src/stepfunctions/steps/sagemaker.py

@@ -444,7 +444,10 @@ def __init__(self, state_id, tuner, job_name, data, wait_for_completion=True, ta
                     :class:`sagemaker.amazon.amazon_estimator.RecordSet` objects,
                     where each instance is a different channel of training data.
             wait_for_completion(bool, optional): Boolean value set to `True` if the Task state should wait for the tuning job to complete before proceeding to the next step in the workflow. Set to `False` if the Task state should submit the tuning job and proceed to the next step. (default: True)
-            tags (list[dict], optional): `List to tags <https://docs.aws.amazon.com/sagemaker/latest/dg/API_Tag.html>`_ to associate with the resource.
+            tags (list[dict] or Placeholder, optional): `List to tags <https://docs.aws.amazon.com/sagemaker/latest/dg/API_Tag.html>`_ to associate with the resource.
+            parameters(dict, optional): The value of this field is merged with other arguments to become the request payload for SageMaker `CreateHyperParameterTuningJob<https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_CreateHyperParameterTuningJob.html>`_.
+                You can use `parameters` to override the value provided by other arguments and specify any field's value dynamically using `Placeholders<https://aws-step-functions-data-science-sdk.readthedocs.io/en/stable/placeholders.html?highlight=placeholder#stepfunctions.inputs.Placeholder>`_.
+
         """
         if wait_for_completion:
             """
@@ -462,19 +465,22 @@ def __init__(self, state_id, tuner, job_name, data, wait_for_completion=True, ta
             kwargs[Field.Resource.value] = get_service_integration_arn(SAGEMAKER_SERVICE_NAME,
                                                                        SageMakerApi.CreateHyperParameterTuningJob)
 
-        parameters = tuning_config(tuner=tuner, inputs=data, job_name=job_name).copy()
+        tuning_parameters = tuning_config(tuner=tuner, inputs=data, job_name=job_name).copy()
 
         if job_name is not None:
-            parameters['HyperParameterTuningJobName'] = job_name
+            tuning_parameters['HyperParameterTuningJobName'] = job_name
 
-        if 'S3Operations' in parameters:
-            del parameters['S3Operations']
+        if 'S3Operations' in tuning_parameters:
+            del tuning_parameters['S3Operations']
 
         if tags:
-            parameters['Tags'] = tags_dict_to_kv_list(tags)
+            tuning_parameters['Tags'] = tags if isinstance(tags, Placeholder) else tags_dict_to_kv_list(tags)
 
-        kwargs[Field.Parameters.value] = parameters
+        if Field.Parameters.value in kwargs and isinstance(kwargs[Field.Parameters.value], dict):
+            # Update tuning parameters with input parameters
+            merge_dicts(tuning_parameters, kwargs[Field.Parameters.value])
 
+        kwargs[Field.Parameters.value] = tuning_parameters
         super(TuningStep, self).__init__(state_id, **kwargs)
 
 

tests/integ/test_sagemaker_steps.py

@@ -257,6 +257,7 @@ def test_create_endpoint_step(trained_estimator, record_set_fixture, sfn_client,
         delete_sagemaker_model(model.name, sagemaker_session)
         # End of Cleanup
 
+
 def test_tuning_step(sfn_client, record_set_for_hyperparameter_tuning, sagemaker_role_arn, sfn_role_arn):
     job_name = generate_job_name()
 
@@ -308,6 +309,123 @@ def test_tuning_step(sfn_client, record_set_for_hyperparameter_tuning, sagemaker
         state_machine_delete_wait(sfn_client, workflow.state_machine_arn)
         # End of Cleanup
 
+
+def test_tuning_step_with_placeholders(sfn_client, record_set_for_hyperparameter_tuning, sagemaker_role_arn, sfn_role_arn):
+    kmeans = KMeans(
+        role=sagemaker_role_arn,
+        instance_count=1,
+        instance_type=INSTANCE_TYPE,
+        k=10
+    )
+
+    hyperparameter_ranges = {
+        "extra_center_factor": IntegerParameter(4, 10),
+        "mini_batch_size": IntegerParameter(10, 100),
+        "epochs": IntegerParameter(1, 2),
+        "init_method": CategoricalParameter(["kmeans++", "random"]),
+    }
+
+    tuner = HyperparameterTuner(
+        estimator=kmeans,
+        objective_metric_name="test:msd",
+        hyperparameter_ranges=hyperparameter_ranges,
+        objective_type="Maximize",
+        max_jobs=2,
+        max_parallel_jobs=1,
+    )
+
+    execution_input = ExecutionInput(schema={
+        'job_name': str,
+        'data_input': str,
+        'objective_metric_name': str,
+        'objective_type': str,
+        'max_jobs': int,
+        'max_parallel_jobs': int,
+        'early_stopping_type': str,
+        'strategy': str,
+    })
+
+    parameters = {
+        'HyperParameterTuningJobConfig': {
+            'HyperParameterTuningJobObjective': {
+                'MetricName': execution_input['objective_metric_name'],
+                'Type': execution_input['objective_type']
+            },
+            'ResourceLimits': {'MaxNumberOfTrainingJobs': execution_input['max_jobs'],
+                               'MaxParallelTrainingJobs': execution_input['max_parallel_jobs']},
+            'Strategy': execution_input['strategy'],
+            'TrainingJobEarlyStoppingType': execution_input['early_stopping_type']
+        },
+        'TrainingJobDefinition': {
+            'AlgorithmSpecification': {
+                'TrainingInputMode': 'File'
+            },
+            'InputDataConfig': execution_input['data_input']
+        }
+    }
+
+    # Build workflow definition
+    tuning_step = TuningStep('Tuning', tuner=tuner, job_name=execution_input['job_name'],
+                             data=record_set_for_hyperparameter_tuning, parameters=parameters)
+    tuning_step.add_retry(SAGEMAKER_RETRY_STRATEGY)
+    workflow_graph = Chain([tuning_step])
+
+    with timeout(minutes=DEFAULT_TIMEOUT_MINUTES):
+        # Create workflow and check definition
+        workflow = create_workflow_and_check_definition(
+            workflow_graph=workflow_graph,
+            workflow_name=unique_name_from_base("integ-test-tuning-step-workflow"),
+            sfn_client=sfn_client,
+            sfn_role_arn=sfn_role_arn
+        )
+
+        job_name = generate_job_name()
+        data_input = [
+            {
+                "DataSource": {
+                    "S3DataSource": {
+                        "S3DataType": "ManifestFile",
+                        "S3Uri": "s3://sagemaker-us-east-1-585192044892/sagemaker-record-sets/PCA-2021-10-19-00-19-10-799/.amazon.manifest",
+                        "S3DataDistributionType": "ShardedByS3Key"
+                    }
+                },
+                "ChannelName": "train"
+            },
+            {
+                "DataSource": {
+                    "S3DataSource": {
+                        "S3DataType": "ManifestFile",
+                        "S3Uri": "s3://sagemaker-us-east-1-585192044892/sagemaker-record-sets/PCA-2021-10-19-00-19-15-087/.amazon.manifest",
+                        "S3DataDistributionType": "ShardedByS3Key"
+                    }
+                },
+                "ChannelName": "test"
+            }
+        ]
+
+        inputs = {
+            'job_name': job_name,
+            'data_input': data_input,
+            'objective_metric_name': 'test:msd',
+            'objective_type': 'Minimize',
+            'max_jobs': 2,
+            'max_parallel_jobs': 2,
+            'early_stopping_type': 'Off',
+            'strategy': 'Bayesian',
+        }
+
+        # Execute workflow
+        execution = workflow.execute(inputs=inputs)
+        execution_output = execution.get_output(wait=True)
+
+        # Check workflow output
+        assert execution_output.get("HyperParameterTuningJobStatus") == "Completed"
+
+        # Cleanup
+        state_machine_delete_wait(sfn_client, workflow.state_machine_arn)
+        # End of Cleanup
+
+
 def test_processing_step(sklearn_processor_fixture, sagemaker_session, sfn_client, sfn_role_arn):
     region = boto3.session.Session().region_name
     input_data = 's3://sagemaker-sample-data-{}/processing/census/census-income.csv'.format(region)