From c9d878a5f2bc8fbf7909f5e21eed0088c7293634 Mon Sep 17 00:00:00 2001 From: Marcio Dos Santos Date: Thu, 14 Mar 2019 12:47:26 -0700 Subject: [PATCH 1/2] Propagate Tags from estimator to model, endpoint, and endpoint config --- src/sagemaker/estimator.py | 5 ++- src/sagemaker/model.py | 4 +- src/sagemaker/session.py | 46 +++++++++++++++-------- tests/unit/test_create_deploy_entities.py | 4 +- tests/unit/test_estimator.py | 34 ++++++++++++++++- tests/unit/test_session.py | 25 +++++++++++- 6 files changed, 96 insertions(+), 22 deletions(-) diff --git a/src/sagemaker/estimator.py b/src/sagemaker/estimator.py index 0e56d00fb6..a4967f07e1 100644 --- a/src/sagemaker/estimator.py +++ b/src/sagemaker/estimator.py @@ -350,6 +350,8 @@ def deploy(self, initial_instance_count, instance_type, accelerator_type=None, e update_endpoint (bool): Flag to update the model in an existing Amazon SageMaker endpoint. If True, this will deploy a new EndpointConfig to an already existing endpoint and delete resources corresponding to the previous EndpointConfig. Default: False + tags(List[dict[str, str]]): Optional. The list of tags to attach to this specific endpoint. + **kwargs: Passed to invocation of ``create_model()``. Implementations may customize ``create_model()`` to accept ``**kwargs`` to customize model creation during deploy. For more, see the implementation docs. @@ -374,7 +376,8 @@ def deploy(self, initial_instance_count, instance_type, accelerator_type=None, e initial_instance_count=initial_instance_count, accelerator_type=accelerator_type, endpoint_name=endpoint_name, - update_endpoint=update_endpoint) + update_endpoint=update_endpoint, + tags=self.tags) @property def model_data(self): diff --git a/src/sagemaker/model.py b/src/sagemaker/model.py index 64076377ba..11964678fb 100644 --- a/src/sagemaker/model.py +++ b/src/sagemaker/model.py @@ -96,7 +96,7 @@ def enable_network_isolation(self): """ return False - def _create_sagemaker_model(self, instance_type, accelerator_type=None): + def _create_sagemaker_model(self, instance_type, accelerator_type=None, tags=None): """Create a SageMaker Model Entity Args: @@ -105,6 +105,8 @@ def _create_sagemaker_model(self, instance_type, accelerator_type=None): accelerator_type (str): Type of Elastic Inference accelerator to attach to an endpoint for model loading and inference, for example, 'ml.eia1.medium'. If not specified, no Elastic Inference accelerator will be attached to the endpoint. + tags(List[dict[str, str]]): Optional. The list of tags to add to the model. + """ container_def = self.prepare_container_def(instance_type, accelerator_type=accelerator_type) self.name = self.name or utils.name_from_image(container_def['Image']) diff --git a/src/sagemaker/session.py b/src/sagemaker/session.py index f46e413db1..802a166f9a 100644 --- a/src/sagemaker/session.py +++ b/src/sagemaker/session.py @@ -544,7 +544,8 @@ def transform(self, job_name, model_name, strategy, max_concurrent_transforms, m self.sagemaker_client.create_transform_job(**transform_request) def create_model(self, name, role, container_defs, vpc_config=None, - enable_network_isolation=False, primary_container=None): + enable_network_isolation=False, primary_container=None, + tags=None): """Create an Amazon SageMaker ``Model``. Specify the S3 location of the model artifacts and Docker image containing the inference code. Amazon SageMaker uses this information to deploy the @@ -570,6 +571,8 @@ def create_model(self, name, role, container_defs, vpc_config=None, You can also specify the return value of ``sagemaker.container_def()``, which is used to create more advanced container configurations, including model containers which need artifacts from S3. This field is deprecated, please use container_defs instead. + tags(List[dict[str, str]]): Optional. The list of tags to add to the model. + Returns: str: Name of the Amazon SageMaker ``Model`` created. @@ -583,12 +586,16 @@ def create_model(self, name, role, container_defs, vpc_config=None, container_defs = primary_container role = self.expand_role(role) - create_model_request = {} + if isinstance(container_defs, list): - create_model_request = _create_model_request(name=name, role=role, container_def=container_defs) + container_definition = container_defs else: - primary_container = _expand_container_def(container_defs) - create_model_request = _create_model_request(name=name, role=role, container_def=primary_container) + container_definition = _expand_container_def(container_defs) + + create_model_request = _create_model_request(name=name, + role=role, + container_def=container_definition, + tags=tags) if vpc_config: create_model_request['VpcConfig'] = vpc_config @@ -702,7 +709,8 @@ def wait_for_model_package(self, model_package_name, poll=5): model_package_name, status, reason)) return desc - def create_endpoint_config(self, name, model_name, initial_instance_count, instance_type, accelerator_type=None): + def create_endpoint_config(self, name, model_name, initial_instance_count, instance_type, + accelerator_type=None, tags=None): """Create an Amazon SageMaker endpoint configuration. The endpoint configuration identifies the Amazon SageMaker model (created using the @@ -717,6 +725,7 @@ def create_endpoint_config(self, name, model_name, initial_instance_count, insta instance_type (str): Type of EC2 instance to launch, for example, 'ml.c4.xlarge'. accelerator_type (str): Type of Elastic Inference accelerator to attach to the instance. For example, 'ml.eia1.medium'. For more information: https://docs.aws.amazon.com/sagemaker/latest/dg/ei.html + tags(List[dict[str, str]]): Optional. The list of tags to add to the endpoint config. Returns: @@ -724,10 +733,13 @@ def create_endpoint_config(self, name, model_name, initial_instance_count, insta """ LOGGER.info('Creating endpoint-config with name {}'.format(name)) + tags = tags or [] + self.sagemaker_client.create_endpoint_config( EndpointConfigName=name, ProductionVariants=[production_variant(model_name, instance_type, initial_instance_count, - accelerator_type=accelerator_type)] + accelerator_type=accelerator_type)], + Tags=tags ) return name @@ -1383,19 +1395,21 @@ def __init__(self, model_data, image, env=None): self.env = env -def _create_model_request(name, role, container_def=None): # pylint: disable=redefined-outer-name - if isinstance(container_def, list): - return { +def _create_model_request(name, role, container_def=None, tags=None): # pylint: disable=redefined-outer-name + request = { 'ModelName': name, - 'Containers': container_def, 'ExecutionRoleArn': role } + + if isinstance(container_def, list): + request['Containers'] = container_def else: - return { - 'ModelName': name, - 'PrimaryContainer': container_def, - 'ExecutionRoleArn': role - } + request['PrimaryContainer'] = container_def + + if tags: + request['Tags'] = tags + + return request def _deployment_entity_exists(describe_fn): diff --git a/tests/unit/test_create_deploy_entities.py b/tests/unit/test_create_deploy_entities.py index 283697f7b3..1b7e170317 100644 --- a/tests/unit/test_create_deploy_entities.py +++ b/tests/unit/test_create_deploy_entities.py @@ -70,7 +70,7 @@ def test_create_endpoint_config(sagemaker_session): 'InitialVariantWeight': 1, 'VariantName': 'AllTraffic'}] sagemaker_session.sagemaker_client.create_endpoint_config.assert_called_once_with( - EndpointConfigName=ENDPOINT_CONFIG_NAME, ProductionVariants=expected_pvs) + EndpointConfigName=ENDPOINT_CONFIG_NAME, ProductionVariants=expected_pvs, Tags=[]) def test_create_endpoint_config_with_accelerator(sagemaker_session): @@ -87,7 +87,7 @@ def test_create_endpoint_config_with_accelerator(sagemaker_session): 'VariantName': 'AllTraffic', 'AcceleratorType': ACCELERATOR_TYPE}] sagemaker_session.sagemaker_client.create_endpoint_config.assert_called_once_with( - EndpointConfigName=ENDPOINT_CONFIG_NAME, ProductionVariants=expected_pvs) + EndpointConfigName=ENDPOINT_CONFIG_NAME, ProductionVariants=expected_pvs, Tags=[]) def test_create_endpoint_no_wait(sagemaker_session): diff --git a/tests/unit/test_estimator.py b/tests/unit/test_estimator.py index 9fe493908a..308d08c1bc 100644 --- a/tests/unit/test_estimator.py +++ b/tests/unit/test_estimator.py @@ -18,7 +18,7 @@ from time import sleep import pytest -from mock import MagicMock, Mock, patch +from mock import ANY, MagicMock, Mock, patch from sagemaker.amazon.amazon_estimator import registry from sagemaker.algorithm import AlgorithmEstimator @@ -882,6 +882,38 @@ def test_unsupported_type_in_dict(): HP_TRAIN_CALL.update({'hyperparameters': STRINGIFIED_HYPERPARAMS}) +def test_fit_deploy_keep_tags(sagemaker_session): + tags = [{'Key': 'TagtestKey', 'Value': 'TagtestValue'}] + estimator = Estimator(IMAGE_NAME, + ROLE, + INSTANCE_COUNT, + INSTANCE_TYPE, + tags=tags, + sagemaker_session=sagemaker_session) + + estimator.fit() + + estimator.deploy(INSTANCE_COUNT, INSTANCE_TYPE) + + variant = [{ + 'InstanceType': 'c4.4xlarge', 'VariantName': 'AllTraffic', + 'ModelName': ANY, 'InitialVariantWeight': 1, + 'InitialInstanceCount': 1 + }] + + job_name = estimator._current_job_name + sagemaker_session.endpoint_from_production_variants.assert_called_with(job_name, + variant, + tags) + + sagemaker_session.create_model.assert_called_with( + ANY, + 'DummyRole', + {'ModelDataUrl': 's3://bucket/model.tar.gz', 'Environment': {}, 'Image': 'fakeimage'}, + enable_network_isolation=False, + vpc_config=None) + + def test_generic_to_fit_no_input(sagemaker_session): e = Estimator(IMAGE_NAME, ROLE, INSTANCE_COUNT, INSTANCE_TYPE, output_path=OUTPUT_PATH, sagemaker_session=sagemaker_session) diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py index 6c3d851e28..2b58282dc1 100644 --- a/tests/unit/test_session.py +++ b/tests/unit/test_session.py @@ -19,7 +19,7 @@ import pytest import six from botocore.exceptions import ClientError -from mock import MagicMock, Mock, patch, call, mock_open +from mock import ANY, MagicMock, Mock, patch, call, mock_open import sagemaker from sagemaker import s3_input, Session, get_execution_role @@ -758,6 +758,18 @@ def test_create_model(expand_container_def, sagemaker_session): PrimaryContainer=PRIMARY_CONTAINER) +@patch('sagemaker.session._expand_container_def', return_value=PRIMARY_CONTAINER) +def test_create_model_with_tags(expand_container_def, sagemaker_session): + tags = [{'Key': 'TagtestKey', 'Value': 'TagtestValue'}] + model = sagemaker_session.create_model(MODEL_NAME, ROLE, PRIMARY_CONTAINER, tags=tags) + + assert model == MODEL_NAME + sagemaker_session.sagemaker_client.create_model.assert_called_with(ExecutionRoleArn=EXPANDED_ROLE, + ModelName=MODEL_NAME, + PrimaryContainer=PRIMARY_CONTAINER, + Tags=[{'Value': 'TagtestValue', 'Key': 'TagtestKey'}]) + + @patch('sagemaker.session._expand_container_def', return_value=PRIMARY_CONTAINER) def test_create_model_with_primary_container(expand_container_def, sagemaker_session): model = sagemaker_session.create_model(MODEL_NAME, ROLE, container_defs=PRIMARY_CONTAINER) @@ -903,6 +915,17 @@ def test_endpoint_from_production_variants(sagemaker_session): ProductionVariants=pvs) +def test_create_endpoint_config_with_tags(sagemaker_session): + tags = [{'Key': 'TagtestKey', 'Value': 'TagtestValue'}] + + sagemaker_session.create_endpoint_config('endpoint-test', 'simple-model', 1, 'local', tags=tags) + + sagemaker_session.sagemaker_client.create_endpoint_config.assert_called_with( + EndpointConfigName='endpoint-test', + ProductionVariants= ANY, + Tags=tags) + + def test_endpoint_from_production_variants_with_tags(sagemaker_session): ims = sagemaker_session ims.sagemaker_client.describe_endpoint = Mock(return_value={'EndpointStatus': 'InService'}) From 2a9931927052a5267e65229c91e6fa6d6766bca5 Mon Sep 17 00:00:00 2001 From: Marcio Dos Santos Date: Thu, 14 Mar 2019 14:52:52 -0700 Subject: [PATCH 2/2] Fix flake8 --- src/sagemaker/estimator.py | 5 ++++- src/sagemaker/model.py | 5 ++++- src/sagemaker/session.py | 15 +++++++++------ tests/unit/test_estimator.py | 8 +++----- tests/unit/test_session.py | 5 +++-- 5 files changed, 23 insertions(+), 15 deletions(-) diff --git a/src/sagemaker/estimator.py b/src/sagemaker/estimator.py index a4967f07e1..116beac965 100644 --- a/src/sagemaker/estimator.py +++ b/src/sagemaker/estimator.py @@ -350,7 +350,10 @@ def deploy(self, initial_instance_count, instance_type, accelerator_type=None, e update_endpoint (bool): Flag to update the model in an existing Amazon SageMaker endpoint. If True, this will deploy a new EndpointConfig to an already existing endpoint and delete resources corresponding to the previous EndpointConfig. Default: False - tags(List[dict[str, str]]): Optional. The list of tags to attach to this specific endpoint. + tags(List[dict[str, str]]): Optional. The list of tags to attach to this specific endpoint. Example: + >>> tags = [{'Key': 'tagname', 'Value': 'tagvalue'}] + For more information about tags, see https://boto3.amazonaws.com/v1/documentation\ + /api/latest/reference/services/sagemaker.html#SageMaker.Client.add_tags **kwargs: Passed to invocation of ``create_model()``. Implementations may customize ``create_model()`` to accept ``**kwargs`` to customize model creation during deploy. diff --git a/src/sagemaker/model.py b/src/sagemaker/model.py index 11964678fb..14adbfa847 100644 --- a/src/sagemaker/model.py +++ b/src/sagemaker/model.py @@ -105,7 +105,10 @@ def _create_sagemaker_model(self, instance_type, accelerator_type=None, tags=Non accelerator_type (str): Type of Elastic Inference accelerator to attach to an endpoint for model loading and inference, for example, 'ml.eia1.medium'. If not specified, no Elastic Inference accelerator will be attached to the endpoint. - tags(List[dict[str, str]]): Optional. The list of tags to add to the model. + tags(List[dict[str, str]]): Optional. The list of tags to add to the model. Example: + >>> tags = [{'Key': 'tagname', 'Value': 'tagvalue'}] + For more information about tags, see https://boto3.amazonaws.com/v1/documentation\ + /api/latest/reference/services/sagemaker.html#SageMaker.Client.add_tags """ container_def = self.prepare_container_def(instance_type, accelerator_type=accelerator_type) diff --git a/src/sagemaker/session.py b/src/sagemaker/session.py index 802a166f9a..acc5173cc8 100644 --- a/src/sagemaker/session.py +++ b/src/sagemaker/session.py @@ -571,7 +571,10 @@ def create_model(self, name, role, container_defs, vpc_config=None, You can also specify the return value of ``sagemaker.container_def()``, which is used to create more advanced container configurations, including model containers which need artifacts from S3. This field is deprecated, please use container_defs instead. - tags(List[dict[str, str]]): Optional. The list of tags to add to the model. + tags(List[dict[str, str]]): Optional. The list of tags to add to the model. Example: + >>> tags = [{'Key': 'tagname', 'Value': 'tagvalue'}] + For more information about tags, see https://boto3.amazonaws.com/v1/documentation\ + /api/latest/reference/services/sagemaker.html#SageMaker.Client.add_tags Returns: @@ -725,7 +728,10 @@ def create_endpoint_config(self, name, model_name, initial_instance_count, insta instance_type (str): Type of EC2 instance to launch, for example, 'ml.c4.xlarge'. accelerator_type (str): Type of Elastic Inference accelerator to attach to the instance. For example, 'ml.eia1.medium'. For more information: https://docs.aws.amazon.com/sagemaker/latest/dg/ei.html - tags(List[dict[str, str]]): Optional. The list of tags to add to the endpoint config. + tags(List[dict[str, str]]): Optional. The list of tags to add to the endpoint config. Example: + >>> tags = [{'Key': 'tagname', 'Value': 'tagvalue'}] + For more information about tags, see https://boto3.amazonaws.com/v1/documentation\ + /api/latest/reference/services/sagemaker.html#SageMaker.Client.add_tags Returns: @@ -1396,10 +1402,7 @@ def __init__(self, model_data, image, env=None): def _create_model_request(name, role, container_def=None, tags=None): # pylint: disable=redefined-outer-name - request = { - 'ModelName': name, - 'ExecutionRoleArn': role - } + request = {'ModelName': name, 'ExecutionRoleArn': role} if isinstance(container_def, list): request['Containers'] = container_def diff --git a/tests/unit/test_estimator.py b/tests/unit/test_estimator.py index 308d08c1bc..ff4e7da30f 100644 --- a/tests/unit/test_estimator.py +++ b/tests/unit/test_estimator.py @@ -895,11 +895,9 @@ def test_fit_deploy_keep_tags(sagemaker_session): estimator.deploy(INSTANCE_COUNT, INSTANCE_TYPE) - variant = [{ - 'InstanceType': 'c4.4xlarge', 'VariantName': 'AllTraffic', - 'ModelName': ANY, 'InitialVariantWeight': 1, - 'InitialInstanceCount': 1 - }] + variant = [{'InstanceType': 'c4.4xlarge', 'VariantName': 'AllTraffic', + 'ModelName': ANY, 'InitialVariantWeight': 1, + 'InitialInstanceCount': 1}] job_name = estimator._current_job_name sagemaker_session.endpoint_from_production_variants.assert_called_with(job_name, diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py index 2b58282dc1..3a318e63a2 100644 --- a/tests/unit/test_session.py +++ b/tests/unit/test_session.py @@ -764,10 +764,11 @@ def test_create_model_with_tags(expand_container_def, sagemaker_session): model = sagemaker_session.create_model(MODEL_NAME, ROLE, PRIMARY_CONTAINER, tags=tags) assert model == MODEL_NAME + tags = [{'Value': 'TagtestValue', 'Key': 'TagtestKey'}] sagemaker_session.sagemaker_client.create_model.assert_called_with(ExecutionRoleArn=EXPANDED_ROLE, ModelName=MODEL_NAME, PrimaryContainer=PRIMARY_CONTAINER, - Tags=[{'Value': 'TagtestValue', 'Key': 'TagtestKey'}]) + Tags=tags) @patch('sagemaker.session._expand_container_def', return_value=PRIMARY_CONTAINER) @@ -922,7 +923,7 @@ def test_create_endpoint_config_with_tags(sagemaker_session): sagemaker_session.sagemaker_client.create_endpoint_config.assert_called_with( EndpointConfigName='endpoint-test', - ProductionVariants= ANY, + ProductionVariants=ANY, Tags=tags)