feat: Support placeholders for TuningStep (#173)

ca-nguyen · web-flow · commit 1ea834629914 · 2021-11-03T16:24:15.000-07:00
diff --git a/src/stepfunctions/steps/sagemaker.py b/src/stepfunctions/steps/sagemaker.py
@@ -465,7 +465,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 of tags <https://docs.aws.amazon.com/sagemaker/latest/dg/API_Tag.html>`_ to associate with the resource.
+            tags (list[dict] or Placeholder, optional): `List of 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:
             """
@@ -483,19 +486,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)
 
 
diff --git a/tests/integ/test_sagemaker_steps.py b/tests/integ/test_sagemaker_steps.py
@@ -104,7 +104,6 @@ def test_training_step(pca_estimator_fixture, record_set_fixture, sfn_client, sf
 
         # Cleanup
         state_machine_delete_wait(sfn_client, workflow.state_machine_arn)
-        # End of Cleanup
 
 
 def test_training_step_with_placeholders(pca_estimator_fixture, record_set_fixture, sfn_client, sfn_role_arn):
@@ -193,7 +192,7 @@ def test_model_step(trained_estimator, sfn_client, sagemaker_session, sfn_role_a
         state_machine_delete_wait(sfn_client, workflow.state_machine_arn)
         model_name = get_resource_name_from_arn(execution_output.get("ModelArn")).split("/")[1]
         delete_sagemaker_model(model_name, sagemaker_session)
-        # End of Cleanup
+
 
 
 def test_model_step_with_placeholders(trained_estimator, sfn_client, sagemaker_session, sfn_role_arn):
@@ -288,7 +287,6 @@ def test_transform_step(trained_estimator, sfn_client, sfn_role_arn):
 
         # Cleanup
         state_machine_delete_wait(sfn_client, workflow.state_machine_arn)
-        # End of Cleanup
 
 
 def test_transform_step_with_placeholder(trained_estimator, sfn_client, sfn_role_arn):
@@ -413,7 +411,7 @@ def test_endpoint_config_step(trained_estimator, sfn_client, sagemaker_session,
         state_machine_delete_wait(sfn_client, workflow.state_machine_arn)
         delete_sagemaker_endpoint_config(endpoint_config_name, sagemaker_session)
         delete_sagemaker_model(model.name, sagemaker_session)
-        # End of Cleanup
+
 
 def test_create_endpoint_step(trained_estimator, record_set_fixture, sfn_client, sagemaker_session, sfn_role_arn):
     # Setup: Create model and endpoint config for trained estimator in SageMaker
@@ -456,7 +454,7 @@ def test_create_endpoint_step(trained_estimator, record_set_fixture, sfn_client,
         delete_sagemaker_endpoint(endpoint_name, sagemaker_session)
         delete_sagemaker_endpoint_config(model.name, sagemaker_session)
         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()
@@ -507,7 +505,97 @@ def test_tuning_step(sfn_client, record_set_for_hyperparameter_tuning, sagemaker
 
         # Cleanup
         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,
+        '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'
+            }
+        }
+    }
+
+    # 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()
+
+        inputs = {
+            'job_name': job_name,
+            '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)
+
 
 def test_processing_step(sklearn_processor_fixture, sagemaker_session, sfn_client, sfn_role_arn):
     region = boto3.session.Session().region_name
@@ -561,7 +649,6 @@ def test_processing_step(sklearn_processor_fixture, sagemaker_session, sfn_clien
 
         # Cleanup
         state_machine_delete_wait(sfn_client, workflow.state_machine_arn)
-        # End of Cleanup
 
 
 def test_processing_step_with_placeholders(sklearn_processor_fixture, sagemaker_session, sfn_client, sfn_role_arn,
diff --git a/tests/unit/test_sagemaker_steps.py b/tests/unit/test_sagemaker_steps.py
