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 1 commit
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 @@ -454,7 +454,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:
"""
Expand All @@ -472,19 +475,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
93 changes: 93 additions & 0 deletions tests/integ/test_sagemaker_steps.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -347,6 +347,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()

Expand Down Expand Up @@ -398,6 +399,98 @@ 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,
'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)
# 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)
Expand Down
235 changes: 233 additions & 2 deletions tests/unit/test_sagemaker_steps.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,12 +24,13 @@
from sagemaker.debugger import Rule, rule_configs, DebuggerHookConfig, CollectionConfig
from sagemaker.sklearn.processing import SKLearnProcessor
from sagemaker.processing import ProcessingInput, ProcessingOutput
from sagemaker.parameter import IntegerParameter, CategoricalParameter
from sagemaker.tuner import HyperparameterTuner

from unittest.mock import MagicMock, patch
from stepfunctions.inputs import ExecutionInput, StepInput
from stepfunctions.steps.fields import Field
from stepfunctions.steps.sagemaker import TrainingStep, TransformStep, ModelStep, EndpointStep, EndpointConfigStep,\
ProcessingStep
ProcessingStep, TuningStep
from stepfunctions.steps.sagemaker import tuning_config

from tests.unit.utils import mock_boto_api_call
Expand Down Expand Up @@ -1412,3 +1413,233 @@ def test_processing_step_creation_with_placeholders(sklearn_processor):
'Resource': 'arn:aws:states:::sagemaker:createProcessingJob.sync',
'End': True
}


@patch('botocore.client.BaseClient._make_api_call', new=mock_boto_api_call)
@patch.object(boto3.session.Session, 'region_name', 'us-east-1')
def test_tuning_step_creation_with_framework(tensorflow_estimator):
hyperparameter_ranges = {
"extra_center_factor": IntegerParameter(4, 10),
"epochs": IntegerParameter(1, 2),
"init_method": CategoricalParameter(["kmeans++", "random"]),
}

tuner = HyperparameterTuner(
estimator=tensorflow_estimator,
objective_metric_name="test:msd",
hyperparameter_ranges=hyperparameter_ranges,
objective_type="Minimize",
max_jobs=2,
max_parallel_jobs=2,
)

step = TuningStep('Tuning',
tuner=tuner,
data={'train': 's3://sagemaker/train'},
job_name='tensorflow-job',
tags=DEFAULT_TAGS
)

state_machine_definition = step.to_dict()
# The sagemaker_job_name is generated - expected name will be taken from the generated definition
generated_sagemaker_job_name = state_machine_definition['Parameters']['TrainingJobDefinition']\
['StaticHyperParameters']['sagemaker_job_name']
expected_definition = {
'Type': 'Task',
'Parameters': {
'HyperParameterTuningJobConfig': {
'HyperParameterTuningJobObjective': {
'MetricName': 'test:msd',
'Type': 'Minimize'
},
'ParameterRanges': {
'CategoricalParameterRanges': [
{
'Name': 'init_method',
'Values': ['"kmeans++"', '"random"']
}],
'ContinuousParameterRanges': [],
'IntegerParameterRanges': [
{
'MaxValue': '10',
'MinValue': '4',
'Name': 'extra_center_factor',
'ScalingType': 'Auto'
},
{
'MaxValue': '2',
'MinValue': '1',
'Name': 'epochs',
'ScalingType': 'Auto'
}
]
},
'ResourceLimits': {'MaxNumberOfTrainingJobs': 2,
'MaxParallelTrainingJobs': 2},
'Strategy': 'Bayesian',
'TrainingJobEarlyStoppingType': 'Off'
},
'HyperParameterTuningJobName': 'tensorflow-job',
'Tags': [{'Key': 'Purpose', 'Value': 'unittests'}],
'TrainingJobDefinition': {
'AlgorithmSpecification': {
'TrainingImage': '520713654638.dkr.ecr.us-east-1.amazonaws.com/sagemaker-tensorflow:1.13-gpu-py2',
'TrainingInputMode': 'File'
},
'InputDataConfig': [{'ChannelName': 'train',
'DataSource': {'S3DataSource': {
'S3DataDistributionType': 'FullyReplicated',
'S3DataType': 'S3Prefix',
'S3Uri': 's3://sagemaker/train'}}}],
'OutputDataConfig': {'S3OutputPath': 's3://sagemaker/models'},
'ResourceConfig': {'InstanceCount': 1,
'InstanceType': 'ml.p2.xlarge',
'VolumeSizeInGB': 30},
'RoleArn': 'execution-role',
'StaticHyperParameters': {
'checkpoint_path': '"s3://sagemaker/models/sagemaker-tensorflow/checkpoints"',
'evaluation_steps': '100',
'sagemaker_container_log_level': '20',
'sagemaker_estimator_class_name': '"TensorFlow"',
'sagemaker_estimator_module': '"sagemaker.tensorflow.estimator"',
'sagemaker_job_name': generated_sagemaker_job_name,
'sagemaker_program': '"tf_train.py"',
'sagemaker_region': '"us-east-1"',
'sagemaker_submit_directory': '"s3://sagemaker/source"',
'training_steps': '1000'},
'StoppingCondition': {'MaxRuntimeInSeconds': 86400}}},
'Resource': 'arn:aws:states:::sagemaker:createHyperParameterTuningJob.sync',
'End': True
}

assert state_machine_definition == expected_definition


@patch('botocore.client.BaseClient._make_api_call', new=mock_boto_api_call)
@patch.object(boto3.session.Session, 'region_name', 'us-east-1')
def test_tuning_step_creation_with_placeholders(tensorflow_estimator):
execution_input = ExecutionInput(schema={
'data_input': str,
'tags': list,
'objective_metric_name': str,
'hyperparameter_ranges': str,
'objective_type': str,
'max_jobs': int,
'max_parallel_jobs': int,
'early_stopping_type': str,
'strategy': str,
})

step_input = StepInput(schema={
'job_name': str
})

hyperparameter_ranges = {
"extra_center_factor": IntegerParameter(4, 10),
"epochs": IntegerParameter(1, 2),
"init_method": CategoricalParameter(["kmeans++", "random"]),
}

tuner = HyperparameterTuner(
estimator=tensorflow_estimator,
objective_metric_name="test:msd",
hyperparameter_ranges=hyperparameter_ranges,
objective_type="Minimize",
max_jobs=2,
max_parallel_jobs=2,
)

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'
},
'HyperParameterRanges': execution_input['hyperparameter_ranges'],
'InputDataConfig': execution_input['data_input']
}
}

step = TuningStep('Tuning',
tuner=tuner,
data={'train': 's3://sagemaker/train'},
job_name=step_input['job_name'],
tags=execution_input['tags'],
parameters=parameters
)

state_machine_definition = step.to_dict()
# The sagemaker_job_name is generated - expected name will be taken from the generated definition
generated_sagemaker_job_name = state_machine_definition['Parameters']['TrainingJobDefinition']['StaticHyperParameters']['sagemaker_job_name']
expected_parameters = {
'HyperParameterTuningJobConfig': {
'HyperParameterTuningJobObjective': {
'MetricName.$': "$$.Execution.Input['objective_metric_name']",
'Type.$': "$$.Execution.Input['objective_type']"
},
'ParameterRanges': {
'CategoricalParameterRanges': [
{
'Name': 'init_method',
'Values': ['"kmeans++"', '"random"']
}],
'ContinuousParameterRanges': [],
'IntegerParameterRanges': [
{
'MaxValue': '10',
'MinValue': '4',
'Name': 'extra_center_factor',
'ScalingType': 'Auto'
},
{
'MaxValue': '2',
'MinValue': '1',
'Name': 'epochs',
'ScalingType': 'Auto'
}
]
},
'ResourceLimits': {'MaxNumberOfTrainingJobs.$': "$$.Execution.Input['max_jobs']",
'MaxParallelTrainingJobs.$': "$$.Execution.Input['max_parallel_jobs']"},
'Strategy.$': "$$.Execution.Input['strategy']",
'TrainingJobEarlyStoppingType.$': "$$.Execution.Input['early_stopping_type']"
},
'HyperParameterTuningJobName.$': "$['job_name']",
'Tags.$': "$$.Execution.Input['tags']",
'TrainingJobDefinition': {
'AlgorithmSpecification': {
'TrainingImage': '520713654638.dkr.ecr.us-east-1.amazonaws.com/sagemaker-tensorflow:1.13-gpu-py2',
'TrainingInputMode': 'File'
},
'HyperParameterRanges.$': "$$.Execution.Input['hyperparameter_ranges']",
'InputDataConfig.$': "$$.Execution.Input['data_input']",
'OutputDataConfig': {'S3OutputPath': 's3://sagemaker/models'},
'ResourceConfig': {'InstanceCount': 1,
'InstanceType': 'ml.p2.xlarge',
'VolumeSizeInGB': 30},
'RoleArn': 'execution-role',
'StaticHyperParameters': {
'checkpoint_path': '"s3://sagemaker/models/sagemaker-tensorflow/checkpoints"',
'evaluation_steps': '100',
'sagemaker_container_log_level': '20',
'sagemaker_estimator_class_name': '"TensorFlow"',
'sagemaker_estimator_module': '"sagemaker.tensorflow.estimator"',
'sagemaker_job_name': generated_sagemaker_job_name,
'sagemaker_program': '"tf_train.py"',
'sagemaker_region': '"us-east-1"',
'sagemaker_submit_directory': '"s3://sagemaker/source"',
'training_steps': '1000'},
'StoppingCondition': {'MaxRuntimeInSeconds': 86400}
}
}

assert state_machine_definition['Parameters'] == expected_parameters