diff --git a/src/sagemaker/estimator.py b/src/sagemaker/estimator.py index 0e56d00fb6..116beac965 100644 --- a/src/sagemaker/estimator.py +++ b/src/sagemaker/estimator.py @@ -350,6 +350,11 @@ 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. 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. For more, see the implementation docs. @@ -374,7 +379,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..14adbfa847 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,11 @@ 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. 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) 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..acc5173cc8 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,11 @@ 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. 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: str: Name of the Amazon SageMaker ``Model`` created. @@ -583,12 +589,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 +712,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 +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. 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: @@ -724,10 +739,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 +1401,18 @@ 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 +def _create_model_request(name, role, container_def=None, tags=None): # pylint: disable=redefined-outer-name + request = {'ModelName': name, 'ExecutionRoleArn': role} + if isinstance(container_def, list): - return { - 'ModelName': name, - 'Containers': container_def, - 'ExecutionRoleArn': role - } + 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..ff4e7da30f 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,36 @@ 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..3a318e63a2 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,19 @@ 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 + tags = [{'Value': 'TagtestValue', 'Key': 'TagtestKey'}] + sagemaker_session.sagemaker_client.create_model.assert_called_with(ExecutionRoleArn=EXPANDED_ROLE, + ModelName=MODEL_NAME, + PrimaryContainer=PRIMARY_CONTAINER, + Tags=tags) + + @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 +916,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'})