Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: Support placeholders for TuningStep #173

Merged
merged 6 commits into from
Nov 3, 2021
Merged

feat: Support placeholders for TuningStep #173

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
39640f8
feat: Support placeholders for TuningStep parameters
ca-nguyen Oct 19, 2021
3b38a6a
Merge branch 'main' into support-placeholders-for-tuning-step
ca-nguyen Oct 21, 2021
8dd38ec
Merge branch 'main' into support-placeholders-for-tuning-step
ca-nguyen Oct 27, 2021
260b78d
Renamed unit test and removed unnecessary comments
ca-nguyen Oct 27, 2021
0ce5235
Update src/stepfunctions/steps/sagemaker.py
ca-nguyen Oct 27, 2021
eaf25dc
Merge branch 'main' into support-placeholders-for-tuning-step
ca-nguyen Oct 28, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 13 additions & 7 deletions src/stepfunctions/steps/sagemaker.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
"""
Expand All @@ -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)


Expand Down
101 changes: 94 additions & 7 deletions tests/integ/test_sagemaker_steps.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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):
Expand Down Expand Up @@ -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):
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand Down
Loading