Skip to content

fix: emit HyperparameterTuner tags to model and endpoint #803

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

Merged
merged 2 commits into from
May 21, 2019
Merged

fix: emit HyperparameterTuner tags to model and endpoint #803

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
81fb7b1
feature: emit training jobs tags to estimator
imujjwal96 May 20, 2019
b0192b7
Merge branch 'master' into pulkit
jesterhazy May 21, 2019
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
2 changes: 2 additions & 0 deletions src/sagemaker/estimator.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -320,6 +320,8 @@ def attach(cls, training_job_name, sagemaker_session=None, model_channel_name='m

job_details = sagemaker_session.sagemaker_client.describe_training_job(TrainingJobName=training_job_name)
init_params = cls._prepare_init_params_from_job_description(job_details, model_channel_name)
tags = sagemaker_session.sagemaker_client.list_tags(ResourceArn=job_details['TrainingJobArn'])['Tags']
init_params.update(tags=tags)

estimator = cls(sagemaker_session=sagemaker_session, **init_params)
estimator.latest_training_job = _TrainingJob(sagemaker_session=sagemaker_session,
Expand Down
9 changes: 9 additions & 0 deletions tests/unit/test_chainer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,10 @@
{'ModelName': 'model-2'}]
}

LIST_TAGS_RESULT = {
'Tags': [{'Key': 'TagtestKey', 'Value': 'TagtestValue'}]
}


@pytest.fixture()
def sagemaker_session():
Expand All @@ -65,6 +69,7 @@ def sagemaker_session():
session.sagemaker_client.describe_training_job = Mock(return_value=describe)
session.sagemaker_client.describe_endpoint = Mock(return_value=ENDPOINT_DESC)
session.sagemaker_client.describe_endpoint_config = Mock(return_value=ENDPOINT_CONFIG_DESC)
session.sagemaker_client.list_tags = Mock(return_value=LIST_TAGS_RESULT)
session.default_bucket = Mock(name='default_bucket', return_value=BUCKET_NAME)
session.expand_role = Mock(name="expand_role", return_value=ROLE)
return session
Expand Down Expand Up @@ -222,6 +227,7 @@ def test_attach_with_additional_hyperparameters(sagemaker_session, chainer_versi
'StoppingCondition': {'MaxRuntimeInSeconds': 24 * 60 * 60},
'TrainingJobName': 'neo',
'TrainingJobStatus': 'Completed',
'TrainingJobArn': 'arn:aws:sagemaker:us-west-2:336:training-job/neo',
'OutputDataConfig': {'KmsKeyId': '',
'S3OutputPath': 's3://place/output/neo'},
'TrainingJobOutput': {'S3TrainingJobOutput': 's3://here/output.tar.gz'}}
Expand Down Expand Up @@ -404,6 +410,7 @@ def test_attach(sagemaker_session, chainer_version):
'StoppingCondition': {'MaxRuntimeInSeconds': 24 * 60 * 60},
'TrainingJobName': 'neo',
'TrainingJobStatus': 'Completed',
'TrainingJobArn': 'arn:aws:sagemaker:us-west-2:336:training-job/neo',
'OutputDataConfig': {'KmsKeyId': '',
'S3OutputPath': 's3://place/output/neo'},
'TrainingJobOutput': {'S3TrainingJobOutput': 's3://here/output.tar.gz'}}
Expand Down Expand Up @@ -446,6 +453,7 @@ def test_attach_wrong_framework(sagemaker_session):
'StoppingCondition': {'MaxRuntimeInSeconds': 24 * 60 * 60},
'TrainingJobName': 'neo',
'TrainingJobStatus': 'Completed',
'TrainingJobArn': 'arn:aws:sagemaker:us-west-2:336:training-job/neo',
'OutputDataConfig': {'KmsKeyId': '',
'S3OutputPath': 's3://place/output/neo'},
'TrainingJobOutput': {'S3TrainingJobOutput': 's3://here/output.tar.gz'}}
Expand Down Expand Up @@ -478,6 +486,7 @@ def test_attach_custom_image(sagemaker_session):
'StoppingCondition': {'MaxRuntimeInSeconds': 24 * 60 * 60},
'TrainingJobName': 'neo',
'TrainingJobStatus': 'Completed',
'TrainingJobArn': 'arn:aws:sagemaker:us-west-2:336:training-job/neo',
'OutputDataConfig': {'KmsKeyId': '',
'S3OutputPath': 's3://place/output/neo'},
'TrainingJobOutput': {'S3TrainingJobOutput': 's3://here/output.tar.gz'}}
Expand Down
7 changes: 7 additions & 0 deletions tests/unit/test_estimator.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -80,6 +80,7 @@
},
'TrainingJobName': 'neo',
'TrainingJobStatus': 'Completed',
'TrainingJobArn': 'arn:aws:sagemaker:us-west-2:336:training-job/neo',
'OutputDataConfig': {
'KmsKeyId': '',
'S3OutputPath': 's3://place/output/neo'
Expand Down Expand Up @@ -111,6 +112,10 @@
{'ModelName': 'model-2'}]
}

LIST_TAGS_RESULT = {
'Tags': [{'Key': 'TagtestKey', 'Value': 'TagtestValue'}]
}


class DummyFramework(Framework):
__framework_name__ = 'dummy'
Expand Down Expand Up @@ -157,6 +162,7 @@ def sagemaker_session():
return_value=DESCRIBE_TRAINING_JOB_RESULT)
sms.sagemaker_client.describe_endpoint = Mock(return_value=ENDPOINT_DESC)
sms.sagemaker_client.describe_endpoint_config = Mock(return_value=ENDPOINT_CONFIG_DESC)
sms.sagemaker_client.list_tags = Mock(return_value=LIST_TAGS_RESULT)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

the mock call is set up here, but there's no test that shows estimator uses the result of the call. adding an assert for this to test_attach_framework would cover this.

return sms


Expand Down Expand Up @@ -460,6 +466,7 @@ def test_attach_framework(sagemaker_session):
assert framework_estimator.subnets == ['foo']
assert framework_estimator.security_group_ids == ['bar']
assert framework_estimator.encrypt_inter_container_traffic is False
assert framework_estimator.tags == LIST_TAGS_RESULT['Tags']


def test_attach_without_hyperparameters(sagemaker_session):
Expand Down
10 changes: 10 additions & 0 deletions tests/unit/test_mxnet.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,10 @@
{'ModelName': 'model-2'}]
}

LIST_TAGS_RESULT = {
'Tags': [{'Key': 'TagtestKey', 'Value': 'TagtestValue'}]
}


@pytest.fixture()
def sagemaker_session():
Expand All @@ -66,6 +70,7 @@ def sagemaker_session():
session.sagemaker_client.describe_training_job = Mock(return_value=describe)
session.sagemaker_client.describe_endpoint = Mock(return_value=ENDPOINT_DESC)
session.sagemaker_client.describe_endpoint_config = Mock(return_value=ENDPOINT_CONFIG_DESC)
session.sagemaker_client.list_tags = Mock(return_value=LIST_TAGS_RESULT)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

again, there is no test that the tags actually get added to the estimator. adding an assert in test_attach would cover this. since we are testing parent class behavior, we don't need to add this for every framework, but it would be nice to test it in one of them.

session.wait_for_compilation_job = Mock(return_value=describe_compilation)
session.default_bucket = Mock(name='default_bucket', return_value=BUCKET_NAME)
session.expand_role = Mock(name="expand_role", return_value=ROLE)
Expand Down Expand Up @@ -346,6 +351,7 @@ def test_attach(sagemaker_session, mxnet_version):
'StoppingCondition': {'MaxRuntimeInSeconds': 24 * 60 * 60},
'TrainingJobName': 'neo',
'TrainingJobStatus': 'Completed',
'TrainingJobArn': 'arn:aws:sagemaker:us-west-2:336:training-job/neo',
'OutputDataConfig': {
'KmsKeyId': '',
'S3OutputPath': 's3://place/output/neo'
Expand All @@ -369,6 +375,7 @@ def test_attach(sagemaker_session, mxnet_version):
assert estimator.hyperparameters()['training_steps'] == '100'
assert estimator.source_dir == 's3://some/sourcedir.tar.gz'
assert estimator.entry_point == 'iris-dnn-classifier.py'
assert estimator.tags == LIST_TAGS_RESULT['Tags']


def test_attach_old_container(sagemaker_session):
Expand All @@ -392,6 +399,7 @@ def test_attach_old_container(sagemaker_session):
'StoppingCondition': {'MaxRuntimeInSeconds': 24 * 60 * 60},
'TrainingJobName': 'neo',
'TrainingJobStatus': 'Completed',
'TrainingJobArn': 'arn:aws:sagemaker:us-west-2:336:training-job/neo',
'OutputDataConfig': {'KmsKeyId': '', 'S3OutputPath': 's3://place/output/neo'},
'TrainingJobOutput': {'S3TrainingJobOutput': 's3://here/output.tar.gz'}}
sagemaker_session.sagemaker_client.describe_training_job = Mock(name='describe_training_job',
Expand Down Expand Up @@ -434,6 +442,7 @@ def test_attach_wrong_framework(sagemaker_session):
'StoppingCondition': {'MaxRuntimeInSeconds': 24 * 60 * 60},
'TrainingJobName': 'neo',
'TrainingJobStatus': 'Completed',
'TrainingJobArn': 'arn:aws:sagemaker:us-west-2:336:training-job/neo',
'OutputDataConfig': {'KmsKeyId': '', 'S3OutputPath': 's3://place/output/neo'},
'TrainingJobOutput': {'S3TrainingJobOutput': 's3://here/output.tar.gz'}}
sagemaker_session.sagemaker_client.describe_training_job = Mock(name='describe_training_job', return_value=rjd)
Expand Down Expand Up @@ -465,6 +474,7 @@ def test_attach_custom_image(sagemaker_session):
'StoppingCondition': {'MaxRuntimeInSeconds': 24 * 60 * 60},
'TrainingJobName': 'neo',
'TrainingJobStatus': 'Completed',
'TrainingJobArn': 'arn:aws:sagemaker:us-west-2:336:training-job/neo',
'OutputDataConfig': {'KmsKeyId': '', 'S3OutputPath': 's3://place/output/neo'},
'TrainingJobOutput': {'S3TrainingJobOutput': 's3://here/output.tar.gz'}}
sagemaker_session.sagemaker_client.describe_training_job = Mock(name='describe_training_job',
Expand Down
8 changes: 8 additions & 0 deletions tests/unit/test_pytorch.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,10 @@
{'ModelName': 'model-2'}]
}

LIST_TAGS_RESULT = {
'Tags': [{'Key': 'TagtestKey', 'Value': 'TagtestValue'}]
}


@pytest.fixture(name='sagemaker_session')
def fixture_sagemaker_session():
Expand All @@ -63,6 +67,7 @@ def fixture_sagemaker_session():
session.sagemaker_client.describe_training_job = Mock(return_value=describe)
session.sagemaker_client.describe_endpoint = Mock(return_value=ENDPOINT_DESC)
session.sagemaker_client.describe_endpoint_config = Mock(return_value=ENDPOINT_CONFIG_DESC)
session.sagemaker_client.list_tags = Mock(return_value=LIST_TAGS_RESULT)
session.default_bucket = Mock(name='default_bucket', return_value=BUCKET_NAME)
session.expand_role = Mock(name="expand_role", return_value=ROLE)
return session
Expand Down Expand Up @@ -306,6 +311,7 @@ def test_attach(sagemaker_session, pytorch_version):
'StoppingCondition': {'MaxRuntimeInSeconds': 24 * 60 * 60},
'TrainingJobName': 'neo',
'TrainingJobStatus': 'Completed',
'TrainingJobArn': 'arn:aws:sagemaker:us-west-2:336:training-job/neo',
'OutputDataConfig': {'KmsKeyId': '',
'S3OutputPath': 's3://place/output/neo'},
'TrainingJobOutput': {'S3TrainingJobOutput': 's3://here/output.tar.gz'}}
Expand Down Expand Up @@ -348,6 +354,7 @@ def test_attach_wrong_framework(sagemaker_session):
'StoppingCondition': {'MaxRuntimeInSeconds': 24 * 60 * 60},
'TrainingJobName': 'neo',
'TrainingJobStatus': 'Completed',
'TrainingJobArn': 'arn:aws:sagemaker:us-west-2:336:training-job/neo',
'OutputDataConfig': {'KmsKeyId': '',
'S3OutputPath': 's3://place/output/neo'},
'TrainingJobOutput': {'S3TrainingJobOutput': 's3://here/output.tar.gz'}}
Expand Down Expand Up @@ -380,6 +387,7 @@ def test_attach_custom_image(sagemaker_session):
'StoppingCondition': {'MaxRuntimeInSeconds': 24 * 60 * 60},
'TrainingJobName': 'neo',
'TrainingJobStatus': 'Completed',
'TrainingJobArn': 'arn:aws:sagemaker:us-west-2:336:training-job/neo',
'OutputDataConfig': {'KmsKeyId': '',
'S3OutputPath': 's3://place/output/neo'},
'TrainingJobOutput': {'S3TrainingJobOutput': 's3://here/output.tar.gz'}}
Expand Down
8 changes: 8 additions & 0 deletions tests/unit/test_rl.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,10 @@
{'ModelName': 'model-2'}]
}

LIST_TAGS_RESULT = {
'Tags': [{'Key': 'TagtestKey', 'Value': 'TagtestValue'}]
}


@pytest.fixture(name='sagemaker_session')
def fixture_sagemaker_session():
Expand All @@ -60,6 +64,7 @@ def fixture_sagemaker_session():
session.sagemaker_client.describe_training_job = Mock(return_value=describe)
session.sagemaker_client.describe_endpoint = Mock(return_value=ENDPOINT_DESC)
session.sagemaker_client.describe_endpoint_config = Mock(return_value=ENDPOINT_CONFIG_DESC)
session.sagemaker_client.list_tags = Mock(return_value=LIST_TAGS_RESULT)
session.default_bucket = Mock(name='default_bucket', return_value=BUCKET_NAME)
session.expand_role = Mock(name="expand_role", return_value=ROLE)
return session
Expand Down Expand Up @@ -360,6 +365,7 @@ def test_attach(sagemaker_session, rl_coach_mxnet_version):
'StoppingCondition': {'MaxRuntimeInSeconds': 24 * 60 * 60},
'TrainingJobName': 'neo',
'TrainingJobStatus': 'Completed',
'TrainingJobArn': 'arn:aws:sagemaker:us-west-2:336:training-job/neo',
'OutputDataConfig': {'KmsKeyId': '',
'S3OutputPath': 's3://place/output/neo'},
'TrainingJobOutput': {
Expand Down Expand Up @@ -406,6 +412,7 @@ def test_attach_wrong_framework(sagemaker_session):
'StoppingCondition': {'MaxRuntimeInSeconds': 24 * 60 * 60},
'TrainingJobName': 'neo',
'TrainingJobStatus': 'Completed',
'TrainingJobArn': 'arn:aws:sagemaker:us-west-2:336:training-job/neo',
'OutputDataConfig': {'KmsKeyId': '',
'S3OutputPath': 's3://place/output/neo'},
'TrainingJobOutput': {'S3TrainingJobOutput': 's3://here/output.tar.gz'}}
Expand Down Expand Up @@ -438,6 +445,7 @@ def test_attach_custom_image(sagemaker_session):
'StoppingCondition': {'MaxRuntimeInSeconds': 24 * 60 * 60},
'TrainingJobName': 'neo',
'TrainingJobStatus': 'Completed',
'TrainingJobArn': 'arn:aws:sagemaker:us-west-2:336:training-job/neo',
'OutputDataConfig':
{'KmsKeyId': '',
'S3OutputPath': 's3://place/output/neo'},
Expand Down
8 changes: 8 additions & 0 deletions tests/unit/test_sklearn.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,10 @@
{'ModelName': 'model-2'}]
}

LIST_TAGS_RESULT = {
'Tags': [{'Key': 'TagtestKey', 'Value': 'TagtestValue'}]
}


@pytest.fixture()
def sagemaker_session():
Expand All @@ -63,6 +67,7 @@ def sagemaker_session():
session.sagemaker_client.describe_training_job = Mock(return_value=describe)
session.sagemaker_client.describe_endpoint = Mock(return_value=ENDPOINT_DESC)
session.sagemaker_client.describe_endpoint_config = Mock(return_value=ENDPOINT_CONFIG_DESC)
session.sagemaker_client.list_tags = Mock(return_value=LIST_TAGS_RESULT)
session.default_bucket = Mock(name='default_bucket', return_value=BUCKET_NAME)
session.expand_role = Mock(name="expand_role", return_value=ROLE)
return session
Expand Down Expand Up @@ -311,6 +316,7 @@ def test_attach(sagemaker_session, sklearn_version):
'StoppingCondition': {'MaxRuntimeInSeconds': 24 * 60 * 60},
'TrainingJobName': 'neo',
'TrainingJobStatus': 'Completed',
'TrainingJobArn': 'arn:aws:sagemaker:us-west-2:336:training-job/neo',
'OutputDataConfig': {'KmsKeyId': '',
'S3OutputPath': 's3://place/output/neo'},
'TrainingJobOutput': {'S3TrainingJobOutput': 's3://here/output.tar.gz'}}
Expand Down Expand Up @@ -353,6 +359,7 @@ def test_attach_wrong_framework(sagemaker_session):
'StoppingCondition': {'MaxRuntimeInSeconds': 24 * 60 * 60},
'TrainingJobName': 'neo',
'TrainingJobStatus': 'Completed',
'TrainingJobArn': 'arn:aws:sagemaker:us-west-2:336:training-job/neo',
'OutputDataConfig': {'KmsKeyId': '',
'S3OutputPath': 's3://place/output/neo'},
'TrainingJobOutput': {'S3TrainingJobOutput': 's3://here/output.tar.gz'}}
Expand Down Expand Up @@ -385,6 +392,7 @@ def test_attach_custom_image(sagemaker_session):
'StoppingCondition': {'MaxRuntimeInSeconds': 24 * 60 * 60},
'TrainingJobName': 'neo',
'TrainingJobStatus': 'Completed',
'TrainingJobArn': 'arn:aws:sagemaker:us-west-2:336:training-job/neo',
'OutputDataConfig': {'KmsKeyId': '',
'S3OutputPath': 's3://place/output/neo'},
'TrainingJobOutput': {'S3TrainingJobOutput': 's3://here/output.tar.gz'}}
Expand Down
11 changes: 11 additions & 0 deletions tests/unit/test_tf_estimator.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,10 @@
{'ModelName': 'model-2'}]
}

LIST_TAGS_RESULT = {
'Tags': [{'Key': 'TagtestKey', 'Value': 'TagtestValue'}]
}


@pytest.fixture()
def sagemaker_session():
Expand All @@ -69,6 +73,7 @@ def sagemaker_session():
session.sagemaker_client.describe_training_job = Mock(return_value=describe)
session.sagemaker_client.describe_endpoint = Mock(return_value=ENDPOINT_DESC)
session.sagemaker_client.describe_endpoint_config = Mock(return_value=ENDPOINT_CONFIG_DESC)
session.sagemaker_client.list_tags = Mock(return_value=LIST_TAGS_RESULT)
return session


Expand Down Expand Up @@ -496,6 +501,7 @@ def test_attach(sagemaker_session, tf_version):
'StoppingCondition': {'MaxRuntimeInSeconds': 24 * 60 * 60},
'TrainingJobName': 'neo',
'TrainingJobStatus': 'Completed',
'TrainingJobArn': 'arn:aws:sagemaker:us-west-2:336:training-job/neo',
'OutputDataConfig': {'KmsKeyId': '', 'S3OutputPath': 's3://place/output/neo'},
'TrainingJobOutput': {'S3TrainingJobOutput': 's3://here/output.tar.gz'}}
sagemaker_session.sagemaker_client.describe_training_job = Mock(name='describe_training_job', return_value=rjd)
Expand Down Expand Up @@ -544,6 +550,7 @@ def test_attach_new_repo_name(sagemaker_session, tf_version):
'StoppingCondition': {'MaxRuntimeInSeconds': 24 * 60 * 60},
'TrainingJobName': 'neo',
'TrainingJobStatus': 'Completed',
'TrainingJobArn': 'arn:aws:sagemaker:us-west-2:336:training-job/neo',
'OutputDataConfig': {'KmsKeyId': '', 'S3OutputPath': 's3://place/output/neo'},
'TrainingJobOutput': {'S3TrainingJobOutput': 's3://here/output.tar.gz'}}
sagemaker_session.sagemaker_client.describe_training_job = Mock(name='describe_training_job', return_value=rjd)
Expand Down Expand Up @@ -593,6 +600,7 @@ def test_attach_old_container(sagemaker_session):
'StoppingCondition': {'MaxRuntimeInSeconds': 24 * 60 * 60},
'TrainingJobName': 'neo',
'TrainingJobStatus': 'Completed',
'TrainingJobArn': 'arn:aws:sagemaker:us-west-2:336:training-job/neo',
'OutputDataConfig': {'KmsKeyId': '', 'S3OutputPath': 's3://place/output/neo'},
'TrainingJobOutput': {'S3TrainingJobOutput': 's3://here/output.tar.gz'}}
sagemaker_session.sagemaker_client.describe_training_job = Mock(name='describe_training_job', return_value=rjd)
Expand Down Expand Up @@ -642,6 +650,7 @@ def test_attach_wrong_framework(sagemaker_session):
},
'TrainingJobName': 'neo',
'TrainingJobStatus': 'Completed',
'TrainingJobArn': 'arn:aws:sagemaker:us-west-2:336:training-job/neo',
'OutputDataConfig': {
'KmsKeyId': '',
'S3OutputPath': 's3://place/output/neo'
Expand Down Expand Up @@ -681,6 +690,7 @@ def test_attach_custom_image(sagemaker_session):
'StoppingCondition': {'MaxRuntimeInSeconds': 24 * 60 * 60},
'TrainingJobName': 'neo',
'TrainingJobStatus': 'Completed',
'TrainingJobArn': 'arn:aws:sagemaker:us-west-2:336:training-job/neo',
'OutputDataConfig': {'KmsKeyId': '', 'S3OutputPath': 's3://place/output/neo'},
'TrainingJobOutput': {'S3TrainingJobOutput': 's3://here/output.tar.gz'}}
sagemaker_session.sagemaker_client.describe_training_job = Mock(name='describe_training_job', return_value=rjd)
Expand Down Expand Up @@ -856,6 +866,7 @@ def test_tf_script_mode_attach(sagemaker_session, tf_version):
'StoppingCondition': {'MaxRuntimeInSeconds': 24 * 60 * 60},
'TrainingJobName': 'neo',
'TrainingJobStatus': 'Completed',
'TrainingJobArn': 'arn:aws:sagemaker:us-west-2:336:training-job/neo',
'OutputDataConfig': {'KmsKeyId': '', 'S3OutputPath': 's3://place/output/neo'},
'TrainingJobOutput': {'S3TrainingJobOutput': 's3://here/output.tar.gz'}}
sagemaker_session.sagemaker_client.describe_training_job = Mock(name='describe_training_job', return_value=rjd)
Expand Down
Loading