fix: fixing propagation of tags to SageMaker endpoint

src/sagemaker/local/entities.py

@@ -327,15 +327,17 @@ def describe(self):
 
 class _LocalEndpointConfig(object):
 
- def __init__(self, config_name, production_variants):
+ def __init__(self, config_name, production_variants, tags=None):
  self.name = config_name
  self.production_variants = production_variants
+ self.tags = tags
  self.creation_time = datetime.datetime.now()
 
  def describe(self):
  response = {
  'EndpointConfigName': self.name,
  'EndpointConfigArn': _UNUSED_ARN,
+ 'Tags': self.tags,
  'CreationTime': self.creation_time,
  'ProductionVariants': self.production_variants
  }
@@ -348,7 +350,7 @@ class _LocalEndpoint(object):
  _IN_SERVICE = 'InService'
  _FAILED = 'Failed'
 
- def __init__(self, endpoint_name, endpoint_config_name, local_session=None):
+ def __init__(self, endpoint_name, endpoint_config_name, tags=None, local_session=None):
  # runtime import since there is a cyclic dependency between entities and local_session
  from sagemaker.local import LocalSession
  self.local_session = local_session or LocalSession()
@@ -357,6 +359,7 @@ def __init__(self, endpoint_name, endpoint_config_name, local_session=None):
  self.name = endpoint_name
  self.endpoint_config = local_client.describe_endpoint_config(endpoint_config_name)
  self.production_variant = self.endpoint_config['ProductionVariants'][0]
+ self.tags = tags
 
  model_name = self.production_variant['ModelName']
  self.primary_container = local_client.describe_model(model_name)['PrimaryContainer']
@@ -392,6 +395,7 @@ def describe(self):
  'EndpointConfigName': self.endpoint_config['EndpointConfigName'],
  'CreationTime': self.create_time,
  'ProductionVariants': self.endpoint_config['ProductionVariants'],
+ 'Tags': self.tags,
  'EndpointName': self.name,
  'EndpointArn': _UNUSED_ARN,
  'EndpointStatus': self.state

src/sagemaker/local/local_session.py

@@ -127,9 +127,9 @@ def describe_endpoint_config(self, EndpointConfigName):
  'Code': 'ValidationException', 'Message': 'Could not find local endpoint config'}}
  raise ClientError(error_response, 'describe_endpoint_config')
 
- def create_endpoint_config(self, EndpointConfigName, ProductionVariants):
+ def create_endpoint_config(self, EndpointConfigName, ProductionVariants, Tags=None):
  LocalSagemakerClient._endpoint_configs[EndpointConfigName] = _LocalEndpointConfig(
- EndpointConfigName, ProductionVariants)
+ EndpointConfigName, ProductionVariants, Tags)
 
  def describe_endpoint(self, EndpointName):
  if EndpointName not in LocalSagemakerClient._endpoints:
@@ -138,8 +138,8 @@ def describe_endpoint(self, EndpointName):
  else:
  return LocalSagemakerClient._endpoints[EndpointName].describe()
 
- def create_endpoint(self, EndpointName, EndpointConfigName):
- endpoint = _LocalEndpoint(EndpointName, EndpointConfigName, self.sagemaker_session)
+ def create_endpoint(self, EndpointName, EndpointConfigName, Tags=None):
+ endpoint = _LocalEndpoint(EndpointName, EndpointConfigName, Tags, self.sagemaker_session)
  LocalSagemakerClient._endpoints[EndpointName] = endpoint
  endpoint.serve()
 

src/sagemaker/model.py

@@ -271,7 +271,8 @@ def deploy(self, initial_instance_count, instance_type, accelerator_type=None, e
  model_name=self.name,
  initial_instance_count=initial_instance_count,
  instance_type=instance_type,
- accelerator_type=accelerator_type)
+ accelerator_type=accelerator_type,
+ tags=tags)
  self.sagemaker_session.update_endpoint(self.endpoint_name, endpoint_config_name)
  else:
  self.sagemaker_session.endpoint_from_production_variants(self.endpoint_name, [production_variant], tags)

src/sagemaker/session.py

@@ -749,7 +749,7 @@ def create_endpoint_config(self, name, model_name, initial_instance_count, insta
  )
  return name
 
- def create_endpoint(self, endpoint_name, config_name, wait=True):
+ def create_endpoint(self, endpoint_name, config_name, tags=None, wait=True):
  """Create an Amazon SageMaker ``Endpoint`` according to the endpoint configuration specified in the request.
 
  Once the ``Endpoint`` is created, client applications can send requests to obtain inferences.
@@ -764,7 +764,10 @@ def create_endpoint(self, endpoint_name, config_name, wait=True):
  str: Name of the Amazon SageMaker ``Endpoint`` created.
  """
  LOGGER.info('Creating endpoint with name {}'.format(endpoint_name))
- self.sagemaker_client.create_endpoint(EndpointName=endpoint_name, EndpointConfigName=config_name)
+
+ tags = tags or []
+
+ self.sagemaker_client.create_endpoint(EndpointName=endpoint_name, EndpointConfigName=config_name, Tags=tags)
  if wait:
  self.wait_for_endpoint(endpoint_name)
  return endpoint_name
@@ -1052,7 +1055,7 @@ def endpoint_from_production_variants(self, name, production_variants, tags=None
  config_options['Tags'] = tags
 
  self.sagemaker_client.create_endpoint_config(**config_options)
- return self.create_endpoint(endpoint_name=name, config_name=name, wait=wait)
+ return self.create_endpoint(endpoint_name=name, config_name=name, tags=tags, wait=wait)
 
  def expand_role(self, role):
  """Expand an IAM role name into an ARN.

tests/integ/test_mxnet_train.py

@@ -78,6 +78,37 @@ def test_deploy_model(mxnet_training_job, sagemaker_session, mxnet_full_version)
  assert 'Could not find model' in str(exception.value)
 
 
+def test_deploy_model_with_tags(mxnet_training_job, sagemaker_session, mxnet_full_version):
+ endpoint_name = 'test-mxnet-deploy-model-{}'.format(sagemaker_timestamp())
+
+ with timeout_and_delete_endpoint_by_name(endpoint_name, sagemaker_session):
+ desc = sagemaker_session.sagemaker_client.describe_training_job(TrainingJobName=mxnet_training_job)
+ model_data = desc['ModelArtifacts']['S3ModelArtifacts']
+ script_path = os.path.join(DATA_DIR, 'mxnet_mnist', 'mnist.py')
+ model = MXNetModel(model_data, 'SageMakerRole', entry_point=script_path,
+ py_version=PYTHON_VERSION, sagemaker_session=sagemaker_session,
+ framework_version=mxnet_full_version)
+ tags = [{'Key': 'TagtestKey', 'Value': 'TagtestValue'}]
+ model.deploy(1, 'ml.m4.xlarge', endpoint_name=endpoint_name, tags=tags)
+
+ returned_model = sagemaker_session.describe_model(EndpointName=model.name)
+ returned_model_tags = sagemaker_session.list_tags(ResourceArn=returned_model['ModelArn'])['Tags']
+
+ endpoint = sagemaker_session.describe_endpoint(EndpointName=endpoint_name)
+ endpoint_tags = sagemaker_session.list_tags(ResourceArn=endpoint['EndpointArn'])['Tags']
+
+ endpoint_config = sagemaker_session.describe_endpoint_config(EndpointConfigName=endpoint['EndpointConfigName'])
+ endpoint_config_tags = sagemaker_session.list_tags(ResourceArn=endpoint_config['EndpointConfigArn'])['Tags']
+
+ production_variants = endpoint_config['ProductionVariants']
+
+ assert returned_model_tags == tags
+ assert endpoint_config_tags == tags
+ assert endpoint_tags == tags
+ assert production_variants[0]['InstanceType'] == 'ml.m4.xlarge'
+ assert production_variants[0]['InitialInstanceCount'] == 1
+
+
 def test_deploy_model_with_update_endpoint(mxnet_training_job, sagemaker_session, mxnet_full_version):
  endpoint_name = 'test-mxnet-deploy-model-{}'.format(sagemaker_timestamp())
 

tests/unit/test_create_deploy_entities.py

@@ -96,7 +96,7 @@ def test_create_endpoint_no_wait(sagemaker_session):
 
  assert returned_name == ENDPOINT_NAME
  sagemaker_session.sagemaker_client.create_endpoint.assert_called_once_with(
- EndpointName=ENDPOINT_NAME, EndpointConfigName=ENDPOINT_CONFIG_NAME)
+ EndpointName=ENDPOINT_NAME, EndpointConfigName=ENDPOINT_CONFIG_NAME, Tags=[])
 
 
 def test_create_endpoint_wait(sagemaker_session):
@@ -105,5 +105,5 @@ def test_create_endpoint_wait(sagemaker_session):
 
  assert returned_name == ENDPOINT_NAME
  sagemaker_session.sagemaker_client.create_endpoint.assert_called_once_with(
- EndpointName=ENDPOINT_NAME, EndpointConfigName=ENDPOINT_CONFIG_NAME)
+ EndpointName=ENDPOINT_NAME, EndpointConfigName=ENDPOINT_CONFIG_NAME, Tags=[])
  sagemaker_session.wait_for_endpoint.assert_called_once_with(ENDPOINT_NAME)

tests/unit/test_model.py

@@ -245,7 +245,8 @@ def test_deploy_update_endpoint(sagemaker_session, tmpdir):
  model_name=model.name,
  initial_instance_count=INSTANCE_COUNT,
  instance_type=INSTANCE_TYPE,
- accelerator_type=ACCELERATOR_TYPE
+ accelerator_type=ACCELERATOR_TYPE,
+ tags=None
  )
  config_name = sagemaker_session.create_endpoint_config(
  name=model.name,

tests/unit/test_session.py

@@ -910,7 +910,8 @@ def test_endpoint_from_production_variants(sagemaker_session):
  ims.sagemaker_client.describe_endpoint_config = Mock(side_effect=ex)
  sagemaker_session.endpoint_from_production_variants('some-endpoint', pvs)
  sagemaker_session.sagemaker_client.create_endpoint.assert_called_with(EndpointConfigName='some-endpoint',
- EndpointName='some-endpoint')
+ EndpointName='some-endpoint',
+ Tags=[])
  sagemaker_session.sagemaker_client.create_endpoint_config.assert_called_with(
  EndpointConfigName='some-endpoint',
  ProductionVariants=pvs)
@@ -936,7 +937,8 @@ def test_endpoint_from_production_variants_with_tags(sagemaker_session):
  tags = [{'ModelName': 'TestModel'}]
  sagemaker_session.endpoint_from_production_variants('some-endpoint', pvs, tags)
  sagemaker_session.sagemaker_client.create_endpoint.assert_called_with(EndpointConfigName='some-endpoint',
- EndpointName='some-endpoint')
+ EndpointName='some-endpoint',
+ Tags=tags)
  sagemaker_session.sagemaker_client.create_endpoint_config.assert_called_with(
  EndpointConfigName='some-endpoint',
  ProductionVariants=pvs,
@@ -953,7 +955,8 @@ def test_endpoint_from_production_variants_with_accelerator_type(sagemaker_sessi
  tags = [{'ModelName': 'TestModel'}]
  sagemaker_session.endpoint_from_production_variants('some-endpoint', pvs, tags)
  sagemaker_session.sagemaker_client.create_endpoint.assert_called_with(EndpointConfigName='some-endpoint',
- EndpointName='some-endpoint')
+ EndpointName='some-endpoint',
+ Tags=tags)
  sagemaker_session.sagemaker_client.create_endpoint_config.assert_called_with(
  EndpointConfigName='some-endpoint',
  ProductionVariants=pvs,

src/sagemaker/local/entities.py

@@ -327,15 +327,17 @@ def describe(self):
 
 class _LocalEndpointConfig(object):
 
- def __init__(self, config_name, production_variants):
+ def __init__(self, config_name, production_variants, tags=None):
  self.name = config_name
  self.production_variants = production_variants
+ self.tags = tags
  self.creation_time = datetime.datetime.now()
 
  def describe(self):
  response = {
  'EndpointConfigName': self.name,
  'EndpointConfigArn': _UNUSED_ARN,
+ 'Tags': self.tags,
  'CreationTime': self.creation_time,
  'ProductionVariants': self.production_variants
  }
@@ -348,7 +350,7 @@ class _LocalEndpoint(object):
  _IN_SERVICE = 'InService'
  _FAILED = 'Failed'
 
- def __init__(self, endpoint_name, endpoint_config_name, local_session=None):
+ def __init__(self, endpoint_name, endpoint_config_name, tags=None, local_session=None):
  # runtime import since there is a cyclic dependency between entities and local_session
  from sagemaker.local import LocalSession
  self.local_session = local_session or LocalSession()
@@ -357,6 +359,7 @@ def __init__(self, endpoint_name, endpoint_config_name, local_session=None):
  self.name = endpoint_name
  self.endpoint_config = local_client.describe_endpoint_config(endpoint_config_name)
  self.production_variant = self.endpoint_config['ProductionVariants'][0]
+ self.tags = tags
 
  model_name = self.production_variant['ModelName']
  self.primary_container = local_client.describe_model(model_name)['PrimaryContainer']
@@ -392,6 +395,7 @@ def describe(self):
  'EndpointConfigName': self.endpoint_config['EndpointConfigName'],
  'CreationTime': self.create_time,
  'ProductionVariants': self.endpoint_config['ProductionVariants'],
+ 'Tags': self.tags,
  'EndpointName': self.name,
  'EndpointArn': _UNUSED_ARN,
  'EndpointStatus': self.state

src/sagemaker/local/local_session.py

@@ -127,9 +127,9 @@ def describe_endpoint_config(self, EndpointConfigName):
  'Code': 'ValidationException', 'Message': 'Could not find local endpoint config'}}
  raise ClientError(error_response, 'describe_endpoint_config')
 
- def create_endpoint_config(self, EndpointConfigName, ProductionVariants):
+ def create_endpoint_config(self, EndpointConfigName, ProductionVariants, Tags=None):
  LocalSagemakerClient._endpoint_configs[EndpointConfigName] = _LocalEndpointConfig(
- EndpointConfigName, ProductionVariants)
+ EndpointConfigName, ProductionVariants, Tags)
 
  def describe_endpoint(self, EndpointName):
  if EndpointName not in LocalSagemakerClient._endpoints:
@@ -138,8 +138,8 @@ def describe_endpoint(self, EndpointName):
  else:
  return LocalSagemakerClient._endpoints[EndpointName].describe()
 
- def create_endpoint(self, EndpointName, EndpointConfigName):
- endpoint = _LocalEndpoint(EndpointName, EndpointConfigName, self.sagemaker_session)
+ def create_endpoint(self, EndpointName, EndpointConfigName, Tags=None):
+ endpoint = _LocalEndpoint(EndpointName, EndpointConfigName, Tags, self.sagemaker_session)
  LocalSagemakerClient._endpoints[EndpointName] = endpoint
  endpoint.serve()
 

src/sagemaker/model.py

@@ -271,7 +271,8 @@ def deploy(self, initial_instance_count, instance_type, accelerator_type=None, e
  model_name=self.name,
  initial_instance_count=initial_instance_count,
  instance_type=instance_type,
- accelerator_type=accelerator_type)
+ accelerator_type=accelerator_type,
+ tags=tags)
  self.sagemaker_session.update_endpoint(self.endpoint_name, endpoint_config_name)
  else:
  self.sagemaker_session.endpoint_from_production_variants(self.endpoint_name, [production_variant], tags)

src/sagemaker/session.py

@@ -749,7 +749,7 @@ def create_endpoint_config(self, name, model_name, initial_instance_count, insta
  )
  return name
 
- def create_endpoint(self, endpoint_name, config_name, wait=True):
+ def create_endpoint(self, endpoint_name, config_name, tags=None, wait=True):
  """Create an Amazon SageMaker ``Endpoint`` according to the endpoint configuration specified in the request.
 
  Once the ``Endpoint`` is created, client applications can send requests to obtain inferences.
@@ -764,7 +764,10 @@ def create_endpoint(self, endpoint_name, config_name, wait=True):
  str: Name of the Amazon SageMaker ``Endpoint`` created.
  """
  LOGGER.info('Creating endpoint with name {}'.format(endpoint_name))
- self.sagemaker_client.create_endpoint(EndpointName=endpoint_name, EndpointConfigName=config_name)
+
+ tags = tags or []
+
+ self.sagemaker_client.create_endpoint(EndpointName=endpoint_name, EndpointConfigName=config_name, Tags=tags)
  if wait:
  self.wait_for_endpoint(endpoint_name)
  return endpoint_name
@@ -1052,7 +1055,7 @@ def endpoint_from_production_variants(self, name, production_variants, tags=None
  config_options['Tags'] = tags
 
  self.sagemaker_client.create_endpoint_config(**config_options)
- return self.create_endpoint(endpoint_name=name, config_name=name, wait=wait)
+ return self.create_endpoint(endpoint_name=name, config_name=name, tags=tags, wait=wait)
 
  def expand_role(self, role):
  """Expand an IAM role name into an ARN.

tests/integ/test_mxnet_train.py

@@ -78,6 +78,37 @@ def test_deploy_model(mxnet_training_job, sagemaker_session, mxnet_full_version)
  assert 'Could not find model' in str(exception.value)
 
 
+def test_deploy_model_with_tags(mxnet_training_job, sagemaker_session, mxnet_full_version):
+ endpoint_name = 'test-mxnet-deploy-model-{}'.format(sagemaker_timestamp())
+
+ with timeout_and_delete_endpoint_by_name(endpoint_name, sagemaker_session):
+ desc = sagemaker_session.sagemaker_client.describe_training_job(TrainingJobName=mxnet_training_job)
+ model_data = desc['ModelArtifacts']['S3ModelArtifacts']
+ script_path = os.path.join(DATA_DIR, 'mxnet_mnist', 'mnist.py')
+ model = MXNetModel(model_data, 'SageMakerRole', entry_point=script_path,
+ py_version=PYTHON_VERSION, sagemaker_session=sagemaker_session,
+ framework_version=mxnet_full_version)
+ tags = [{'Key': 'TagtestKey', 'Value': 'TagtestValue'}]
+ model.deploy(1, 'ml.m4.xlarge', endpoint_name=endpoint_name, tags=tags)
+
+ returned_model = sagemaker_session.describe_model(EndpointName=model.name)
+ returned_model_tags = sagemaker_session.list_tags(ResourceArn=returned_model['ModelArn'])['Tags']
+
+ endpoint = sagemaker_session.describe_endpoint(EndpointName=endpoint_name)
+ endpoint_tags = sagemaker_session.list_tags(ResourceArn=endpoint['EndpointArn'])['Tags']
+
+ endpoint_config = sagemaker_session.describe_endpoint_config(EndpointConfigName=endpoint['EndpointConfigName'])
+ endpoint_config_tags = sagemaker_session.list_tags(ResourceArn=endpoint_config['EndpointConfigArn'])['Tags']
+
+ production_variants = endpoint_config['ProductionVariants']
+
+ assert returned_model_tags == tags
+ assert endpoint_config_tags == tags
+ assert endpoint_tags == tags
+ assert production_variants[0]['InstanceType'] == 'ml.m4.xlarge'
+ assert production_variants[0]['InitialInstanceCount'] == 1
+
+
 def test_deploy_model_with_update_endpoint(mxnet_training_job, sagemaker_session, mxnet_full_version):
  endpoint_name = 'test-mxnet-deploy-model-{}'.format(sagemaker_timestamp())
 

tests/unit/test_create_deploy_entities.py

@@ -96,7 +96,7 @@ def test_create_endpoint_no_wait(sagemaker_session):
 
  assert returned_name == ENDPOINT_NAME
  sagemaker_session.sagemaker_client.create_endpoint.assert_called_once_with(
- EndpointName=ENDPOINT_NAME, EndpointConfigName=ENDPOINT_CONFIG_NAME)
+ EndpointName=ENDPOINT_NAME, EndpointConfigName=ENDPOINT_CONFIG_NAME, Tags=[])
 
 
 def test_create_endpoint_wait(sagemaker_session):
@@ -105,5 +105,5 @@ def test_create_endpoint_wait(sagemaker_session):
 
  assert returned_name == ENDPOINT_NAME
  sagemaker_session.sagemaker_client.create_endpoint.assert_called_once_with(
- EndpointName=ENDPOINT_NAME, EndpointConfigName=ENDPOINT_CONFIG_NAME)
+ EndpointName=ENDPOINT_NAME, EndpointConfigName=ENDPOINT_CONFIG_NAME, Tags=[])
  sagemaker_session.wait_for_endpoint.assert_called_once_with(ENDPOINT_NAME)

tests/unit/test_model.py

@@ -245,7 +245,8 @@ def test_deploy_update_endpoint(sagemaker_session, tmpdir):
  model_name=model.name,
  initial_instance_count=INSTANCE_COUNT,
  instance_type=INSTANCE_TYPE,
- accelerator_type=ACCELERATOR_TYPE
+ accelerator_type=ACCELERATOR_TYPE,
+ tags=None
  )
  config_name = sagemaker_session.create_endpoint_config(
  name=model.name,

tests/unit/test_session.py

@@ -910,7 +910,8 @@ def test_endpoint_from_production_variants(sagemaker_session):
  ims.sagemaker_client.describe_endpoint_config = Mock(side_effect=ex)
  sagemaker_session.endpoint_from_production_variants('some-endpoint', pvs)
  sagemaker_session.sagemaker_client.create_endpoint.assert_called_with(EndpointConfigName='some-endpoint',
- EndpointName='some-endpoint')
+ EndpointName='some-endpoint',
+ Tags=[])
  sagemaker_session.sagemaker_client.create_endpoint_config.assert_called_with(
  EndpointConfigName='some-endpoint',
  ProductionVariants=pvs)
@@ -936,7 +937,8 @@ def test_endpoint_from_production_variants_with_tags(sagemaker_session):
  tags = [{'ModelName': 'TestModel'}]
  sagemaker_session.endpoint_from_production_variants('some-endpoint', pvs, tags)
  sagemaker_session.sagemaker_client.create_endpoint.assert_called_with(EndpointConfigName='some-endpoint',
- EndpointName='some-endpoint')
+ EndpointName='some-endpoint',
+ Tags=tags)
  sagemaker_session.sagemaker_client.create_endpoint_config.assert_called_with(
  EndpointConfigName='some-endpoint',
  ProductionVariants=pvs,
@@ -953,7 +955,8 @@ def test_endpoint_from_production_variants_with_accelerator_type(sagemaker_sessi
  tags = [{'ModelName': 'TestModel'}]
  sagemaker_session.endpoint_from_production_variants('some-endpoint', pvs, tags)
  sagemaker_session.sagemaker_client.create_endpoint.assert_called_with(EndpointConfigName='some-endpoint',
- EndpointName='some-endpoint')
+ EndpointName='some-endpoint',
+ Tags=tags)
  sagemaker_session.sagemaker_client.create_endpoint_config.assert_called_with(
  EndpointConfigName='some-endpoint',
  ProductionVariants=pvs,

src/sagemaker/local/entities.py

@@ -327,15 +327,17 @@ def describe(self):
 
 class _LocalEndpointConfig(object):
 
- def __init__(self, config_name, production_variants):
+ def __init__(self, config_name, production_variants, tags=None):
  self.name = config_name
  self.production_variants = production_variants
+ self.tags = tags
  self.creation_time = datetime.datetime.now()
 
  def describe(self):
  response = {
  'EndpointConfigName': self.name,
  'EndpointConfigArn': _UNUSED_ARN,
+ 'Tags': self.tags,
  'CreationTime': self.creation_time,
  'ProductionVariants': self.production_variants
  }
@@ -348,7 +350,7 @@ class _LocalEndpoint(object):
  _IN_SERVICE = 'InService'
  _FAILED = 'Failed'
 
- def __init__(self, endpoint_name, endpoint_config_name, local_session=None):
+ def __init__(self, endpoint_name, endpoint_config_name, tags=None, local_session=None):
  # runtime import since there is a cyclic dependency between entities and local_session
  from sagemaker.local import LocalSession
  self.local_session = local_session or LocalSession()
@@ -357,6 +359,7 @@ def __init__(self, endpoint_name, endpoint_config_name, local_session=None):
  self.name = endpoint_name
  self.endpoint_config = local_client.describe_endpoint_config(endpoint_config_name)
  self.production_variant = self.endpoint_config['ProductionVariants'][0]
+ self.tags = tags
 
  model_name = self.production_variant['ModelName']
  self.primary_container = local_client.describe_model(model_name)['PrimaryContainer']
@@ -392,6 +395,7 @@ def describe(self):
  'EndpointConfigName': self.endpoint_config['EndpointConfigName'],
  'CreationTime': self.create_time,
  'ProductionVariants': self.endpoint_config['ProductionVariants'],
+ 'Tags': self.tags,
  'EndpointName': self.name,
  'EndpointArn': _UNUSED_ARN,
  'EndpointStatus': self.state

src/sagemaker/local/local_session.py

@@ -127,9 +127,9 @@ def describe_endpoint_config(self, EndpointConfigName):
  'Code': 'ValidationException', 'Message': 'Could not find local endpoint config'}}
  raise ClientError(error_response, 'describe_endpoint_config')
 
- def create_endpoint_config(self, EndpointConfigName, ProductionVariants):
+ def create_endpoint_config(self, EndpointConfigName, ProductionVariants, Tags=None):
  LocalSagemakerClient._endpoint_configs[EndpointConfigName] = _LocalEndpointConfig(
- EndpointConfigName, ProductionVariants)
+ EndpointConfigName, ProductionVariants, Tags)
 
  def describe_endpoint(self, EndpointName):
  if EndpointName not in LocalSagemakerClient._endpoints:
@@ -138,8 +138,8 @@ def describe_endpoint(self, EndpointName):
  else:
  return LocalSagemakerClient._endpoints[EndpointName].describe()
 
- def create_endpoint(self, EndpointName, EndpointConfigName):
- endpoint = _LocalEndpoint(EndpointName, EndpointConfigName, self.sagemaker_session)
+ def create_endpoint(self, EndpointName, EndpointConfigName, Tags=None):
+ endpoint = _LocalEndpoint(EndpointName, EndpointConfigName, Tags, self.sagemaker_session)
  LocalSagemakerClient._endpoints[EndpointName] = endpoint
  endpoint.serve()
 

src/sagemaker/model.py

@@ -271,7 +271,8 @@ def deploy(self, initial_instance_count, instance_type, accelerator_type=None, e
  model_name=self.name,
  initial_instance_count=initial_instance_count,
  instance_type=instance_type,
- accelerator_type=accelerator_type)
+ accelerator_type=accelerator_type,
+ tags=tags)
  self.sagemaker_session.update_endpoint(self.endpoint_name, endpoint_config_name)
  else:
  self.sagemaker_session.endpoint_from_production_variants(self.endpoint_name, [production_variant], tags)

src/sagemaker/session.py

@@ -749,7 +749,7 @@ def create_endpoint_config(self, name, model_name, initial_instance_count, insta
  )
  return name
 
- def create_endpoint(self, endpoint_name, config_name, wait=True):
+ def create_endpoint(self, endpoint_name, config_name, tags=None, wait=True):
  """Create an Amazon SageMaker ``Endpoint`` according to the endpoint configuration specified in the request.
 
  Once the ``Endpoint`` is created, client applications can send requests to obtain inferences.
@@ -764,7 +764,10 @@ def create_endpoint(self, endpoint_name, config_name, wait=True):
  str: Name of the Amazon SageMaker ``Endpoint`` created.
  """
  LOGGER.info('Creating endpoint with name {}'.format(endpoint_name))
- self.sagemaker_client.create_endpoint(EndpointName=endpoint_name, EndpointConfigName=config_name)
+
+ tags = tags or []
+
+ self.sagemaker_client.create_endpoint(EndpointName=endpoint_name, EndpointConfigName=config_name, Tags=tags)
  if wait:
  self.wait_for_endpoint(endpoint_name)
  return endpoint_name
@@ -1052,7 +1055,7 @@ def endpoint_from_production_variants(self, name, production_variants, tags=None
  config_options['Tags'] = tags
 
  self.sagemaker_client.create_endpoint_config(**config_options)
- return self.create_endpoint(endpoint_name=name, config_name=name, wait=wait)
+ return self.create_endpoint(endpoint_name=name, config_name=name, tags=tags, wait=wait)
 
  def expand_role(self, role):
  """Expand an IAM role name into an ARN.

tests/integ/test_mxnet_train.py

@@ -78,6 +78,37 @@ def test_deploy_model(mxnet_training_job, sagemaker_session, mxnet_full_version)
  assert 'Could not find model' in str(exception.value)
 
 
+def test_deploy_model_with_tags(mxnet_training_job, sagemaker_session, mxnet_full_version):
+ endpoint_name = 'test-mxnet-deploy-model-{}'.format(sagemaker_timestamp())
+
+ with timeout_and_delete_endpoint_by_name(endpoint_name, sagemaker_session):
+ desc = sagemaker_session.sagemaker_client.describe_training_job(TrainingJobName=mxnet_training_job)
+ model_data = desc['ModelArtifacts']['S3ModelArtifacts']
+ script_path = os.path.join(DATA_DIR, 'mxnet_mnist', 'mnist.py')
+ model = MXNetModel(model_data, 'SageMakerRole', entry_point=script_path,
+ py_version=PYTHON_VERSION, sagemaker_session=sagemaker_session,
+ framework_version=mxnet_full_version)
+ tags = [{'Key': 'TagtestKey', 'Value': 'TagtestValue'}]
+ model.deploy(1, 'ml.m4.xlarge', endpoint_name=endpoint_name, tags=tags)
+
+ returned_model = sagemaker_session.describe_model(EndpointName=model.name)
+ returned_model_tags = sagemaker_session.list_tags(ResourceArn=returned_model['ModelArn'])['Tags']
+
+ endpoint = sagemaker_session.describe_endpoint(EndpointName=endpoint_name)
+ endpoint_tags = sagemaker_session.list_tags(ResourceArn=endpoint['EndpointArn'])['Tags']
+
+ endpoint_config = sagemaker_session.describe_endpoint_config(EndpointConfigName=endpoint['EndpointConfigName'])
+ endpoint_config_tags = sagemaker_session.list_tags(ResourceArn=endpoint_config['EndpointConfigArn'])['Tags']
+
+ production_variants = endpoint_config['ProductionVariants']
+
+ assert returned_model_tags == tags
+ assert endpoint_config_tags == tags
+ assert endpoint_tags == tags
+ assert production_variants[0]['InstanceType'] == 'ml.m4.xlarge'
+ assert production_variants[0]['InitialInstanceCount'] == 1
+
+
 def test_deploy_model_with_update_endpoint(mxnet_training_job, sagemaker_session, mxnet_full_version):
  endpoint_name = 'test-mxnet-deploy-model-{}'.format(sagemaker_timestamp())
 

tests/unit/test_create_deploy_entities.py

@@ -96,7 +96,7 @@ def test_create_endpoint_no_wait(sagemaker_session):
 
  assert returned_name == ENDPOINT_NAME
  sagemaker_session.sagemaker_client.create_endpoint.assert_called_once_with(
- EndpointName=ENDPOINT_NAME, EndpointConfigName=ENDPOINT_CONFIG_NAME)
+ EndpointName=ENDPOINT_NAME, EndpointConfigName=ENDPOINT_CONFIG_NAME, Tags=[])
 
 
 def test_create_endpoint_wait(sagemaker_session):
@@ -105,5 +105,5 @@ def test_create_endpoint_wait(sagemaker_session):
 
  assert returned_name == ENDPOINT_NAME
  sagemaker_session.sagemaker_client.create_endpoint.assert_called_once_with(
- EndpointName=ENDPOINT_NAME, EndpointConfigName=ENDPOINT_CONFIG_NAME)
+ EndpointName=ENDPOINT_NAME, EndpointConfigName=ENDPOINT_CONFIG_NAME, Tags=[])
  sagemaker_session.wait_for_endpoint.assert_called_once_with(ENDPOINT_NAME)

tests/unit/test_model.py

@@ -245,7 +245,8 @@ def test_deploy_update_endpoint(sagemaker_session, tmpdir):
  model_name=model.name,
  initial_instance_count=INSTANCE_COUNT,
  instance_type=INSTANCE_TYPE,
- accelerator_type=ACCELERATOR_TYPE
+ accelerator_type=ACCELERATOR_TYPE,
+ tags=None
  )
  config_name = sagemaker_session.create_endpoint_config(
  name=model.name,

tests/unit/test_session.py

@@ -910,7 +910,8 @@ def test_endpoint_from_production_variants(sagemaker_session):
  ims.sagemaker_client.describe_endpoint_config = Mock(side_effect=ex)
  sagemaker_session.endpoint_from_production_variants('some-endpoint', pvs)
  sagemaker_session.sagemaker_client.create_endpoint.assert_called_with(EndpointConfigName='some-endpoint',
- EndpointName='some-endpoint')
+ EndpointName='some-endpoint',
+ Tags=[])
  sagemaker_session.sagemaker_client.create_endpoint_config.assert_called_with(
  EndpointConfigName='some-endpoint',
  ProductionVariants=pvs)
@@ -936,7 +937,8 @@ def test_endpoint_from_production_variants_with_tags(sagemaker_session):
  tags = [{'ModelName': 'TestModel'}]
  sagemaker_session.endpoint_from_production_variants('some-endpoint', pvs, tags)
  sagemaker_session.sagemaker_client.create_endpoint.assert_called_with(EndpointConfigName='some-endpoint',
- EndpointName='some-endpoint')
+ EndpointName='some-endpoint',
+ Tags=tags)
  sagemaker_session.sagemaker_client.create_endpoint_config.assert_called_with(
  EndpointConfigName='some-endpoint',
  ProductionVariants=pvs,
@@ -953,7 +955,8 @@ def test_endpoint_from_production_variants_with_accelerator_type(sagemaker_sessi
  tags = [{'ModelName': 'TestModel'}]
  sagemaker_session.endpoint_from_production_variants('some-endpoint', pvs, tags)
  sagemaker_session.sagemaker_client.create_endpoint.assert_called_with(EndpointConfigName='some-endpoint',
- EndpointName='some-endpoint')
+ EndpointName='some-endpoint',
+ Tags=tags)
  sagemaker_session.sagemaker_client.create_endpoint_config.assert_called_with(
  EndpointConfigName='some-endpoint',
  ProductionVariants=pvs,

src/sagemaker/local/entities.py

@@ -327,15 +327,17 @@ def describe(self):
 
 class _LocalEndpointConfig(object):
 
- def __init__(self, config_name, production_variants):
+ def __init__(self, config_name, production_variants, tags=None):
  self.name = config_name
  self.production_variants = production_variants
+ self.tags = tags
  self.creation_time = datetime.datetime.now()
 
  def describe(self):
  response = {
  'EndpointConfigName': self.name,
  'EndpointConfigArn': _UNUSED_ARN,
+ 'Tags': self.tags,
  'CreationTime': self.creation_time,
  'ProductionVariants': self.production_variants
  }
@@ -348,7 +350,7 @@ class _LocalEndpoint(object):
  _IN_SERVICE = 'InService'
  _FAILED = 'Failed'
 
- def __init__(self, endpoint_name, endpoint_config_name, local_session=None):
+ def __init__(self, endpoint_name, endpoint_config_name, tags=None, local_session=None):
  # runtime import since there is a cyclic dependency between entities and local_session
  from sagemaker.local import LocalSession
  self.local_session = local_session or LocalSession()
@@ -357,6 +359,7 @@ def __init__(self, endpoint_name, endpoint_config_name, local_session=None):
  self.name = endpoint_name
  self.endpoint_config = local_client.describe_endpoint_config(endpoint_config_name)
  self.production_variant = self.endpoint_config['ProductionVariants'][0]
+ self.tags = tags
 
  model_name = self.production_variant['ModelName']
  self.primary_container = local_client.describe_model(model_name)['PrimaryContainer']
@@ -392,6 +395,7 @@ def describe(self):
  'EndpointConfigName': self.endpoint_config['EndpointConfigName'],
  'CreationTime': self.create_time,
  'ProductionVariants': self.endpoint_config['ProductionVariants'],
+ 'Tags': self.tags,
  'EndpointName': self.name,
  'EndpointArn': _UNUSED_ARN,
  'EndpointStatus': self.state

src/sagemaker/local/local_session.py

@@ -127,9 +127,9 @@ def describe_endpoint_config(self, EndpointConfigName):
  'Code': 'ValidationException', 'Message': 'Could not find local endpoint config'}}
  raise ClientError(error_response, 'describe_endpoint_config')
 
- def create_endpoint_config(self, EndpointConfigName, ProductionVariants):
+ def create_endpoint_config(self, EndpointConfigName, ProductionVariants, Tags=None):
  LocalSagemakerClient._endpoint_configs[EndpointConfigName] = _LocalEndpointConfig(
- EndpointConfigName, ProductionVariants)
+ EndpointConfigName, ProductionVariants, Tags)
 
  def describe_endpoint(self, EndpointName):
  if EndpointName not in LocalSagemakerClient._endpoints:
@@ -138,8 +138,8 @@ def describe_endpoint(self, EndpointName):
  else:
  return LocalSagemakerClient._endpoints[EndpointName].describe()
 
- def create_endpoint(self, EndpointName, EndpointConfigName):
- endpoint = _LocalEndpoint(EndpointName, EndpointConfigName, self.sagemaker_session)
+ def create_endpoint(self, EndpointName, EndpointConfigName, Tags=None):
+ endpoint = _LocalEndpoint(EndpointName, EndpointConfigName, Tags, self.sagemaker_session)
  LocalSagemakerClient._endpoints[EndpointName] = endpoint
  endpoint.serve()
 

src/sagemaker/model.py

@@ -271,7 +271,8 @@ def deploy(self, initial_instance_count, instance_type, accelerator_type=None, e
  model_name=self.name,
  initial_instance_count=initial_instance_count,
  instance_type=instance_type,
- accelerator_type=accelerator_type)
+ accelerator_type=accelerator_type,
+ tags=tags)
  self.sagemaker_session.update_endpoint(self.endpoint_name, endpoint_config_name)
  else:
  self.sagemaker_session.endpoint_from_production_variants(self.endpoint_name, [production_variant], tags)

src/sagemaker/session.py

@@ -749,7 +749,7 @@ def create_endpoint_config(self, name, model_name, initial_instance_count, insta
  )
  return name
 
- def create_endpoint(self, endpoint_name, config_name, wait=True):
+ def create_endpoint(self, endpoint_name, config_name, tags=None, wait=True):
  """Create an Amazon SageMaker ``Endpoint`` according to the endpoint configuration specified in the request.
 
  Once the ``Endpoint`` is created, client applications can send requests to obtain inferences.
@@ -764,7 +764,10 @@ def create_endpoint(self, endpoint_name, config_name, wait=True):
  str: Name of the Amazon SageMaker ``Endpoint`` created.
  """
  LOGGER.info('Creating endpoint with name {}'.format(endpoint_name))
- self.sagemaker_client.create_endpoint(EndpointName=endpoint_name, EndpointConfigName=config_name)
+
+ tags = tags or []
+
+ self.sagemaker_client.create_endpoint(EndpointName=endpoint_name, EndpointConfigName=config_name, Tags=tags)
  if wait:
  self.wait_for_endpoint(endpoint_name)
  return endpoint_name
@@ -1052,7 +1055,7 @@ def endpoint_from_production_variants(self, name, production_variants, tags=None
  config_options['Tags'] = tags
 
  self.sagemaker_client.create_endpoint_config(**config_options)
- return self.create_endpoint(endpoint_name=name, config_name=name, wait=wait)
+ return self.create_endpoint(endpoint_name=name, config_name=name, tags=tags, wait=wait)
 
  def expand_role(self, role):
  """Expand an IAM role name into an ARN.

tests/integ/test_mxnet_train.py

@@ -78,6 +78,37 @@ def test_deploy_model(mxnet_training_job, sagemaker_session, mxnet_full_version)
  assert 'Could not find model' in str(exception.value)
 
 
+def test_deploy_model_with_tags(mxnet_training_job, sagemaker_session, mxnet_full_version):
+ endpoint_name = 'test-mxnet-deploy-model-{}'.format(sagemaker_timestamp())
+
+ with timeout_and_delete_endpoint_by_name(endpoint_name, sagemaker_session):
+ desc = sagemaker_session.sagemaker_client.describe_training_job(TrainingJobName=mxnet_training_job)
+ model_data = desc['ModelArtifacts']['S3ModelArtifacts']
+ script_path = os.path.join(DATA_DIR, 'mxnet_mnist', 'mnist.py')
+ model = MXNetModel(model_data, 'SageMakerRole', entry_point=script_path,
+ py_version=PYTHON_VERSION, sagemaker_session=sagemaker_session,
+ framework_version=mxnet_full_version)
+ tags = [{'Key': 'TagtestKey', 'Value': 'TagtestValue'}]
+ model.deploy(1, 'ml.m4.xlarge', endpoint_name=endpoint_name, tags=tags)
+
+ returned_model = sagemaker_session.describe_model(EndpointName=model.name)
+ returned_model_tags = sagemaker_session.list_tags(ResourceArn=returned_model['ModelArn'])['Tags']
+
+ endpoint = sagemaker_session.describe_endpoint(EndpointName=endpoint_name)
+ endpoint_tags = sagemaker_session.list_tags(ResourceArn=endpoint['EndpointArn'])['Tags']
+
+ endpoint_config = sagemaker_session.describe_endpoint_config(EndpointConfigName=endpoint['EndpointConfigName'])
+ endpoint_config_tags = sagemaker_session.list_tags(ResourceArn=endpoint_config['EndpointConfigArn'])['Tags']
+
+ production_variants = endpoint_config['ProductionVariants']
+
+ assert returned_model_tags == tags
+ assert endpoint_config_tags == tags
+ assert endpoint_tags == tags
+ assert production_variants[0]['InstanceType'] == 'ml.m4.xlarge'
+ assert production_variants[0]['InitialInstanceCount'] == 1
+
+
 def test_deploy_model_with_update_endpoint(mxnet_training_job, sagemaker_session, mxnet_full_version):
  endpoint_name = 'test-mxnet-deploy-model-{}'.format(sagemaker_timestamp())
 

tests/unit/test_create_deploy_entities.py

@@ -96,7 +96,7 @@ def test_create_endpoint_no_wait(sagemaker_session):
 
  assert returned_name == ENDPOINT_NAME
  sagemaker_session.sagemaker_client.create_endpoint.assert_called_once_with(
- EndpointName=ENDPOINT_NAME, EndpointConfigName=ENDPOINT_CONFIG_NAME)
+ EndpointName=ENDPOINT_NAME, EndpointConfigName=ENDPOINT_CONFIG_NAME, Tags=[])
 
 
 def test_create_endpoint_wait(sagemaker_session):
@@ -105,5 +105,5 @@ def test_create_endpoint_wait(sagemaker_session):
 
  assert returned_name == ENDPOINT_NAME
  sagemaker_session.sagemaker_client.create_endpoint.assert_called_once_with(
- EndpointName=ENDPOINT_NAME, EndpointConfigName=ENDPOINT_CONFIG_NAME)
+ EndpointName=ENDPOINT_NAME, EndpointConfigName=ENDPOINT_CONFIG_NAME, Tags=[])
  sagemaker_session.wait_for_endpoint.assert_called_once_with(ENDPOINT_NAME)

tests/unit/test_model.py

@@ -245,7 +245,8 @@ def test_deploy_update_endpoint(sagemaker_session, tmpdir):
  model_name=model.name,
  initial_instance_count=INSTANCE_COUNT,
  instance_type=INSTANCE_TYPE,
- accelerator_type=ACCELERATOR_TYPE
+ accelerator_type=ACCELERATOR_TYPE,
+ tags=None
  )
  config_name = sagemaker_session.create_endpoint_config(
  name=model.name,

tests/unit/test_session.py

@@ -910,7 +910,8 @@ def test_endpoint_from_production_variants(sagemaker_session):
  ims.sagemaker_client.describe_endpoint_config = Mock(side_effect=ex)
  sagemaker_session.endpoint_from_production_variants('some-endpoint', pvs)
  sagemaker_session.sagemaker_client.create_endpoint.assert_called_with(EndpointConfigName='some-endpoint',
- EndpointName='some-endpoint')
+ EndpointName='some-endpoint',
+ Tags=[])
  sagemaker_session.sagemaker_client.create_endpoint_config.assert_called_with(
  EndpointConfigName='some-endpoint',
  ProductionVariants=pvs)
@@ -936,7 +937,8 @@ def test_endpoint_from_production_variants_with_tags(sagemaker_session):
  tags = [{'ModelName': 'TestModel'}]
  sagemaker_session.endpoint_from_production_variants('some-endpoint', pvs, tags)
  sagemaker_session.sagemaker_client.create_endpoint.assert_called_with(EndpointConfigName='some-endpoint',
- EndpointName='some-endpoint')
+ EndpointName='some-endpoint',
+ Tags=tags)
  sagemaker_session.sagemaker_client.create_endpoint_config.assert_called_with(
  EndpointConfigName='some-endpoint',
  ProductionVariants=pvs,
@@ -953,7 +955,8 @@ def test_endpoint_from_production_variants_with_accelerator_type(sagemaker_sessi
  tags = [{'ModelName': 'TestModel'}]
  sagemaker_session.endpoint_from_production_variants('some-endpoint', pvs, tags)
  sagemaker_session.sagemaker_client.create_endpoint.assert_called_with(EndpointConfigName='some-endpoint',
- EndpointName='some-endpoint')
+ EndpointName='some-endpoint',
+ Tags=tags)
  sagemaker_session.sagemaker_client.create_endpoint_config.assert_called_with(
  EndpointConfigName='some-endpoint',
  ProductionVariants=pvs,

src/sagemaker/local/entities.py

@@ -327,15 +327,17 @@ def describe(self):
 
 class _LocalEndpointConfig(object):
 
- def __init__(self, config_name, production_variants):
+ def __init__(self, config_name, production_variants, tags=None):
  self.name = config_name
  self.production_variants = production_variants
+ self.tags = tags
  self.creation_time = datetime.datetime.now()
 
  def describe(self):
  response = {
  'EndpointConfigName': self.name,
  'EndpointConfigArn': _UNUSED_ARN,
+ 'Tags': self.tags,
  'CreationTime': self.creation_time,
  'ProductionVariants': self.production_variants
  }
@@ -348,7 +350,7 @@ class _LocalEndpoint(object):
  _IN_SERVICE = 'InService'
  _FAILED = 'Failed'
 
- def __init__(self, endpoint_name, endpoint_config_name, local_session=None):
+ def __init__(self, endpoint_name, endpoint_config_name, tags=None, local_session=None):
  # runtime import since there is a cyclic dependency between entities and local_session
  from sagemaker.local import LocalSession
  self.local_session = local_session or LocalSession()
@@ -357,6 +359,7 @@ def __init__(self, endpoint_name, endpoint_config_name, local_session=None):
  self.name = endpoint_name
  self.endpoint_config = local_client.describe_endpoint_config(endpoint_config_name)
  self.production_variant = self.endpoint_config['ProductionVariants'][0]
+ self.tags = tags
 
  model_name = self.production_variant['ModelName']
  self.primary_container = local_client.describe_model(model_name)['PrimaryContainer']
@@ -392,6 +395,7 @@ def describe(self):
  'EndpointConfigName': self.endpoint_config['EndpointConfigName'],
  'CreationTime': self.create_time,
  'ProductionVariants': self.endpoint_config['ProductionVariants'],
+ 'Tags': self.tags,
  'EndpointName': self.name,
  'EndpointArn': _UNUSED_ARN,
  'EndpointStatus': self.state

src/sagemaker/local/local_session.py

@@ -127,9 +127,9 @@ def describe_endpoint_config(self, EndpointConfigName):
  'Code': 'ValidationException', 'Message': 'Could not find local endpoint config'}}
  raise ClientError(error_response, 'describe_endpoint_config')
 
- def create_endpoint_config(self, EndpointConfigName, ProductionVariants):
+ def create_endpoint_config(self, EndpointConfigName, ProductionVariants, Tags=None):
  LocalSagemakerClient._endpoint_configs[EndpointConfigName] = _LocalEndpointConfig(
- EndpointConfigName, ProductionVariants)
+ EndpointConfigName, ProductionVariants, Tags)
 
  def describe_endpoint(self, EndpointName):
  if EndpointName not in LocalSagemakerClient._endpoints:
@@ -138,8 +138,8 @@ def describe_endpoint(self, EndpointName):
  else:
  return LocalSagemakerClient._endpoints[EndpointName].describe()
 
- def create_endpoint(self, EndpointName, EndpointConfigName):
- endpoint = _LocalEndpoint(EndpointName, EndpointConfigName, self.sagemaker_session)
+ def create_endpoint(self, EndpointName, EndpointConfigName, Tags=None):
+ endpoint = _LocalEndpoint(EndpointName, EndpointConfigName, Tags, self.sagemaker_session)
  LocalSagemakerClient._endpoints[EndpointName] = endpoint
  endpoint.serve()
 

src/sagemaker/model.py

@@ -271,7 +271,8 @@ def deploy(self, initial_instance_count, instance_type, accelerator_type=None, e
  model_name=self.name,
  initial_instance_count=initial_instance_count,
  instance_type=instance_type,
- accelerator_type=accelerator_type)
+ accelerator_type=accelerator_type,
+ tags=tags)
  self.sagemaker_session.update_endpoint(self.endpoint_name, endpoint_config_name)
  else:
  self.sagemaker_session.endpoint_from_production_variants(self.endpoint_name, [production_variant], tags)

src/sagemaker/session.py

@@ -749,7 +749,7 @@ def create_endpoint_config(self, name, model_name, initial_instance_count, insta
  )
  return name
 
- def create_endpoint(self, endpoint_name, config_name, wait=True):
+ def create_endpoint(self, endpoint_name, config_name, tags=None, wait=True):
  """Create an Amazon SageMaker ``Endpoint`` according to the endpoint configuration specified in the request.
 
  Once the ``Endpoint`` is created, client applications can send requests to obtain inferences.
@@ -764,7 +764,10 @@ def create_endpoint(self, endpoint_name, config_name, wait=True):
  str: Name of the Amazon SageMaker ``Endpoint`` created.
  """
  LOGGER.info('Creating endpoint with name {}'.format(endpoint_name))
- self.sagemaker_client.create_endpoint(EndpointName=endpoint_name, EndpointConfigName=config_name)
+
+ tags = tags or []
+
+ self.sagemaker_client.create_endpoint(EndpointName=endpoint_name, EndpointConfigName=config_name, Tags=tags)
  if wait:
  self.wait_for_endpoint(endpoint_name)
  return endpoint_name
@@ -1052,7 +1055,7 @@ def endpoint_from_production_variants(self, name, production_variants, tags=None
  config_options['Tags'] = tags
 
  self.sagemaker_client.create_endpoint_config(**config_options)
- return self.create_endpoint(endpoint_name=name, config_name=name, wait=wait)
+ return self.create_endpoint(endpoint_name=name, config_name=name, tags=tags, wait=wait)
 
  def expand_role(self, role):
  """Expand an IAM role name into an ARN.

tests/integ/test_mxnet_train.py

@@ -78,6 +78,37 @@ def test_deploy_model(mxnet_training_job, sagemaker_session, mxnet_full_version)
  assert 'Could not find model' in str(exception.value)
 
 
+def test_deploy_model_with_tags(mxnet_training_job, sagemaker_session, mxnet_full_version):
+ endpoint_name = 'test-mxnet-deploy-model-{}'.format(sagemaker_timestamp())
+
+ with timeout_and_delete_endpoint_by_name(endpoint_name, sagemaker_session):
+ desc = sagemaker_session.sagemaker_client.describe_training_job(TrainingJobName=mxnet_training_job)
+ model_data = desc['ModelArtifacts']['S3ModelArtifacts']
+ script_path = os.path.join(DATA_DIR, 'mxnet_mnist', 'mnist.py')
+ model = MXNetModel(model_data, 'SageMakerRole', entry_point=script_path,
+ py_version=PYTHON_VERSION, sagemaker_session=sagemaker_session,
+ framework_version=mxnet_full_version)
+ tags = [{'Key': 'TagtestKey', 'Value': 'TagtestValue'}]
+ model.deploy(1, 'ml.m4.xlarge', endpoint_name=endpoint_name, tags=tags)
+
+ returned_model = sagemaker_session.describe_model(EndpointName=model.name)
+ returned_model_tags = sagemaker_session.list_tags(ResourceArn=returned_model['ModelArn'])['Tags']
+
+ endpoint = sagemaker_session.describe_endpoint(EndpointName=endpoint_name)
+ endpoint_tags = sagemaker_session.list_tags(ResourceArn=endpoint['EndpointArn'])['Tags']
+
+ endpoint_config = sagemaker_session.describe_endpoint_config(EndpointConfigName=endpoint['EndpointConfigName'])
+ endpoint_config_tags = sagemaker_session.list_tags(ResourceArn=endpoint_config['EndpointConfigArn'])['Tags']
+
+ production_variants = endpoint_config['ProductionVariants']
+
+ assert returned_model_tags == tags
+ assert endpoint_config_tags == tags
+ assert endpoint_tags == tags
+ assert production_variants[0]['InstanceType'] == 'ml.m4.xlarge'
+ assert production_variants[0]['InitialInstanceCount'] == 1
+
+
 def test_deploy_model_with_update_endpoint(mxnet_training_job, sagemaker_session, mxnet_full_version):
  endpoint_name = 'test-mxnet-deploy-model-{}'.format(sagemaker_timestamp())
 

tests/unit/test_create_deploy_entities.py

@@ -96,7 +96,7 @@ def test_create_endpoint_no_wait(sagemaker_session):
 
  assert returned_name == ENDPOINT_NAME
  sagemaker_session.sagemaker_client.create_endpoint.assert_called_once_with(
- EndpointName=ENDPOINT_NAME, EndpointConfigName=ENDPOINT_CONFIG_NAME)
+ EndpointName=ENDPOINT_NAME, EndpointConfigName=ENDPOINT_CONFIG_NAME, Tags=[])
 
 
 def test_create_endpoint_wait(sagemaker_session):
@@ -105,5 +105,5 @@ def test_create_endpoint_wait(sagemaker_session):
 
  assert returned_name == ENDPOINT_NAME
  sagemaker_session.sagemaker_client.create_endpoint.assert_called_once_with(
- EndpointName=ENDPOINT_NAME, EndpointConfigName=ENDPOINT_CONFIG_NAME)
+ EndpointName=ENDPOINT_NAME, EndpointConfigName=ENDPOINT_CONFIG_NAME, Tags=[])
  sagemaker_session.wait_for_endpoint.assert_called_once_with(ENDPOINT_NAME)

tests/unit/test_model.py

@@ -245,7 +245,8 @@ def test_deploy_update_endpoint(sagemaker_session, tmpdir):
  model_name=model.name,
  initial_instance_count=INSTANCE_COUNT,
  instance_type=INSTANCE_TYPE,
- accelerator_type=ACCELERATOR_TYPE
+ accelerator_type=ACCELERATOR_TYPE,
+ tags=None
  )
  config_name = sagemaker_session.create_endpoint_config(
  name=model.name,

tests/unit/test_session.py

@@ -910,7 +910,8 @@ def test_endpoint_from_production_variants(sagemaker_session):
  ims.sagemaker_client.describe_endpoint_config = Mock(side_effect=ex)
  sagemaker_session.endpoint_from_production_variants('some-endpoint', pvs)
  sagemaker_session.sagemaker_client.create_endpoint.assert_called_with(EndpointConfigName='some-endpoint',
- EndpointName='some-endpoint')
+ EndpointName='some-endpoint',
+ Tags=[])
  sagemaker_session.sagemaker_client.create_endpoint_config.assert_called_with(
  EndpointConfigName='some-endpoint',
  ProductionVariants=pvs)
@@ -936,7 +937,8 @@ def test_endpoint_from_production_variants_with_tags(sagemaker_session):
  tags = [{'ModelName': 'TestModel'}]
  sagemaker_session.endpoint_from_production_variants('some-endpoint', pvs, tags)
  sagemaker_session.sagemaker_client.create_endpoint.assert_called_with(EndpointConfigName='some-endpoint',
- EndpointName='some-endpoint')
+ EndpointName='some-endpoint',
+ Tags=tags)
  sagemaker_session.sagemaker_client.create_endpoint_config.assert_called_with(
  EndpointConfigName='some-endpoint',
  ProductionVariants=pvs,
@@ -953,7 +955,8 @@ def test_endpoint_from_production_variants_with_accelerator_type(sagemaker_sessi
  tags = [{'ModelName': 'TestModel'}]
  sagemaker_session.endpoint_from_production_variants('some-endpoint', pvs, tags)
  sagemaker_session.sagemaker_client.create_endpoint.assert_called_with(EndpointConfigName='some-endpoint',
- EndpointName='some-endpoint')
+ EndpointName='some-endpoint',
+ Tags=tags)
  sagemaker_session.sagemaker_client.create_endpoint_config.assert_called_with(
  EndpointConfigName='some-endpoint',
  ProductionVariants=pvs,

src/sagemaker/local/entities.py

@@ -327,15 +327,17 @@ def describe(self):
 
 class _LocalEndpointConfig(object):
 
- def __init__(self, config_name, production_variants):
+ def __init__(self, config_name, production_variants, tags=None):
  self.name = config_name
  self.production_variants = production_variants
+ self.tags = tags
  self.creation_time = datetime.datetime.now()
 
  def describe(self):
  response = {
  'EndpointConfigName': self.name,
  'EndpointConfigArn': _UNUSED_ARN,
+ 'Tags': self.tags,
  'CreationTime': self.creation_time,
  'ProductionVariants': self.production_variants
  }
@@ -348,7 +350,7 @@ class _LocalEndpoint(object):
  _IN_SERVICE = 'InService'
  _FAILED = 'Failed'
 
- def __init__(self, endpoint_name, endpoint_config_name, local_session=None):
+ def __init__(self, endpoint_name, endpoint_config_name, tags=None, local_session=None):
  # runtime import since there is a cyclic dependency between entities and local_session
  from sagemaker.local import LocalSession
  self.local_session = local_session or LocalSession()
@@ -357,6 +359,7 @@ def __init__(self, endpoint_name, endpoint_config_name, local_session=None):
  self.name = endpoint_name
  self.endpoint_config = local_client.describe_endpoint_config(endpoint_config_name)
  self.production_variant = self.endpoint_config['ProductionVariants'][0]
+ self.tags = tags
 
  model_name = self.production_variant['ModelName']
  self.primary_container = local_client.describe_model(model_name)['PrimaryContainer']
@@ -392,6 +395,7 @@ def describe(self):
  'EndpointConfigName': self.endpoint_config['EndpointConfigName'],
  'CreationTime': self.create_time,
  'ProductionVariants': self.endpoint_config['ProductionVariants'],
+ 'Tags': self.tags,
  'EndpointName': self.name,
  'EndpointArn': _UNUSED_ARN,
  'EndpointStatus': self.state

src/sagemaker/local/local_session.py

@@ -127,9 +127,9 @@ def describe_endpoint_config(self, EndpointConfigName):
  'Code': 'ValidationException', 'Message': 'Could not find local endpoint config'}}
  raise ClientError(error_response, 'describe_endpoint_config')
 
- def create_endpoint_config(self, EndpointConfigName, ProductionVariants):
+ def create_endpoint_config(self, EndpointConfigName, ProductionVariants, Tags=None):
  LocalSagemakerClient._endpoint_configs[EndpointConfigName] = _LocalEndpointConfig(
- EndpointConfigName, ProductionVariants)
+ EndpointConfigName, ProductionVariants, Tags)
 
  def describe_endpoint(self, EndpointName):
  if EndpointName not in LocalSagemakerClient._endpoints:
@@ -138,8 +138,8 @@ def describe_endpoint(self, EndpointName):
  else:
  return LocalSagemakerClient._endpoints[EndpointName].describe()
 
- def create_endpoint(self, EndpointName, EndpointConfigName):
- endpoint = _LocalEndpoint(EndpointName, EndpointConfigName, self.sagemaker_session)
+ def create_endpoint(self, EndpointName, EndpointConfigName, Tags=None):
+ endpoint = _LocalEndpoint(EndpointName, EndpointConfigName, Tags, self.sagemaker_session)
  LocalSagemakerClient._endpoints[EndpointName] = endpoint
  endpoint.serve()
 

src/sagemaker/model.py

@@ -271,7 +271,8 @@ def deploy(self, initial_instance_count, instance_type, accelerator_type=None, e
  model_name=self.name,
  initial_instance_count=initial_instance_count,
  instance_type=instance_type,
- accelerator_type=accelerator_type)
+ accelerator_type=accelerator_type,
+ tags=tags)
  self.sagemaker_session.update_endpoint(self.endpoint_name, endpoint_config_name)
  else:
  self.sagemaker_session.endpoint_from_production_variants(self.endpoint_name, [production_variant], tags)

src/sagemaker/session.py

@@ -749,7 +749,7 @@ def create_endpoint_config(self, name, model_name, initial_instance_count, insta
  )
  return name
 
- def create_endpoint(self, endpoint_name, config_name, wait=True):
+ def create_endpoint(self, endpoint_name, config_name, tags=None, wait=True):
  """Create an Amazon SageMaker ``Endpoint`` according to the endpoint configuration specified in the request.
 
  Once the ``Endpoint`` is created, client applications can send requests to obtain inferences.
@@ -764,7 +764,10 @@ def create_endpoint(self, endpoint_name, config_name, wait=True):
  str: Name of the Amazon SageMaker ``Endpoint`` created.
  """
  LOGGER.info('Creating endpoint with name {}'.format(endpoint_name))
- self.sagemaker_client.create_endpoint(EndpointName=endpoint_name, EndpointConfigName=config_name)
+
+ tags = tags or []
+
+ self.sagemaker_client.create_endpoint(EndpointName=endpoint_name, EndpointConfigName=config_name, Tags=tags)
  if wait:
  self.wait_for_endpoint(endpoint_name)
  return endpoint_name
@@ -1052,7 +1055,7 @@ def endpoint_from_production_variants(self, name, production_variants, tags=None
  config_options['Tags'] = tags
 
  self.sagemaker_client.create_endpoint_config(**config_options)
- return self.create_endpoint(endpoint_name=name, config_name=name, wait=wait)
+ return self.create_endpoint(endpoint_name=name, config_name=name, tags=tags, wait=wait)
 
  def expand_role(self, role):
  """Expand an IAM role name into an ARN.

tests/integ/test_mxnet_train.py

@@ -78,6 +78,37 @@ def test_deploy_model(mxnet_training_job, sagemaker_session, mxnet_full_version)
  assert 'Could not find model' in str(exception.value)
 
 
+def test_deploy_model_with_tags(mxnet_training_job, sagemaker_session, mxnet_full_version):
+ endpoint_name = 'test-mxnet-deploy-model-{}'.format(sagemaker_timestamp())
+
+ with timeout_and_delete_endpoint_by_name(endpoint_name, sagemaker_session):
+ desc = sagemaker_session.sagemaker_client.describe_training_job(TrainingJobName=mxnet_training_job)
+ model_data = desc['ModelArtifacts']['S3ModelArtifacts']
+ script_path = os.path.join(DATA_DIR, 'mxnet_mnist', 'mnist.py')
+ model = MXNetModel(model_data, 'SageMakerRole', entry_point=script_path,
+ py_version=PYTHON_VERSION, sagemaker_session=sagemaker_session,
+ framework_version=mxnet_full_version)
+ tags = [{'Key': 'TagtestKey', 'Value': 'TagtestValue'}]
+ model.deploy(1, 'ml.m4.xlarge', endpoint_name=endpoint_name, tags=tags)
+
+ returned_model = sagemaker_session.describe_model(EndpointName=model.name)
+ returned_model_tags = sagemaker_session.list_tags(ResourceArn=returned_model['ModelArn'])['Tags']
+
+ endpoint = sagemaker_session.describe_endpoint(EndpointName=endpoint_name)
+ endpoint_tags = sagemaker_session.list_tags(ResourceArn=endpoint['EndpointArn'])['Tags']
+
+ endpoint_config = sagemaker_session.describe_endpoint_config(EndpointConfigName=endpoint['EndpointConfigName'])
+ endpoint_config_tags = sagemaker_session.list_tags(ResourceArn=endpoint_config['EndpointConfigArn'])['Tags']
+
+ production_variants = endpoint_config['ProductionVariants']
+
+ assert returned_model_tags == tags
+ assert endpoint_config_tags == tags
+ assert endpoint_tags == tags
+ assert production_variants[0]['InstanceType'] == 'ml.m4.xlarge'
+ assert production_variants[0]['InitialInstanceCount'] == 1
+
+
 def test_deploy_model_with_update_endpoint(mxnet_training_job, sagemaker_session, mxnet_full_version):
  endpoint_name = 'test-mxnet-deploy-model-{}'.format(sagemaker_timestamp())
 

tests/unit/test_create_deploy_entities.py

@@ -96,7 +96,7 @@ def test_create_endpoint_no_wait(sagemaker_session):
 
  assert returned_name == ENDPOINT_NAME
  sagemaker_session.sagemaker_client.create_endpoint.assert_called_once_with(
- EndpointName=ENDPOINT_NAME, EndpointConfigName=ENDPOINT_CONFIG_NAME)
+ EndpointName=ENDPOINT_NAME, EndpointConfigName=ENDPOINT_CONFIG_NAME, Tags=[])
 
 
 def test_create_endpoint_wait(sagemaker_session):
@@ -105,5 +105,5 @@ def test_create_endpoint_wait(sagemaker_session):
 
  assert returned_name == ENDPOINT_NAME
  sagemaker_session.sagemaker_client.create_endpoint.assert_called_once_with(
- EndpointName=ENDPOINT_NAME, EndpointConfigName=ENDPOINT_CONFIG_NAME)
+ EndpointName=ENDPOINT_NAME, EndpointConfigName=ENDPOINT_CONFIG_NAME, Tags=[])
  sagemaker_session.wait_for_endpoint.assert_called_once_with(ENDPOINT_NAME)

tests/unit/test_model.py

@@ -245,7 +245,8 @@ def test_deploy_update_endpoint(sagemaker_session, tmpdir):
  model_name=model.name,
  initial_instance_count=INSTANCE_COUNT,
  instance_type=INSTANCE_TYPE,
- accelerator_type=ACCELERATOR_TYPE
+ accelerator_type=ACCELERATOR_TYPE,
+ tags=None
  )
  config_name = sagemaker_session.create_endpoint_config(
  name=model.name,

tests/unit/test_session.py

@@ -910,7 +910,8 @@ def test_endpoint_from_production_variants(sagemaker_session):
  ims.sagemaker_client.describe_endpoint_config = Mock(side_effect=ex)
  sagemaker_session.endpoint_from_production_variants('some-endpoint', pvs)
  sagemaker_session.sagemaker_client.create_endpoint.assert_called_with(EndpointConfigName='some-endpoint',
- EndpointName='some-endpoint')
+ EndpointName='some-endpoint',
+ Tags=[])
  sagemaker_session.sagemaker_client.create_endpoint_config.assert_called_with(
  EndpointConfigName='some-endpoint',
  ProductionVariants=pvs)
@@ -936,7 +937,8 @@ def test_endpoint_from_production_variants_with_tags(sagemaker_session):
  tags = [{'ModelName': 'TestModel'}]
  sagemaker_session.endpoint_from_production_variants('some-endpoint', pvs, tags)
  sagemaker_session.sagemaker_client.create_endpoint.assert_called_with(EndpointConfigName='some-endpoint',
- EndpointName='some-endpoint')
+ EndpointName='some-endpoint',
+ Tags=tags)
  sagemaker_session.sagemaker_client.create_endpoint_config.assert_called_with(
  EndpointConfigName='some-endpoint',
  ProductionVariants=pvs,
@@ -953,7 +955,8 @@ def test_endpoint_from_production_variants_with_accelerator_type(sagemaker_sessi
  tags = [{'ModelName': 'TestModel'}]
  sagemaker_session.endpoint_from_production_variants('some-endpoint', pvs, tags)
  sagemaker_session.sagemaker_client.create_endpoint.assert_called_with(EndpointConfigName='some-endpoint',
- EndpointName='some-endpoint')
+ EndpointName='some-endpoint',
+ Tags=tags)
  sagemaker_session.sagemaker_client.create_endpoint_config.assert_called_with(
  EndpointConfigName='some-endpoint',
  ProductionVariants=pvs,

src/sagemaker/local/entities.py

@@ -327,15 +327,17 @@ def describe(self):
 
 class _LocalEndpointConfig(object):
 
- def __init__(self, config_name, production_variants):
+ def __init__(self, config_name, production_variants, tags=None):
  self.name = config_name
  self.production_variants = production_variants
+ self.tags = tags
  self.creation_time = datetime.datetime.now()
 
  def describe(self):
  response = {
  'EndpointConfigName': self.name,
  'EndpointConfigArn': _UNUSED_ARN,
+ 'Tags': self.tags,
  'CreationTime': self.creation_time,
  'ProductionVariants': self.production_variants
  }
@@ -348,7 +350,7 @@ class _LocalEndpoint(object):
  _IN_SERVICE = 'InService'
  _FAILED = 'Failed'
 
- def __init__(self, endpoint_name, endpoint_config_name, local_session=None):
+ def __init__(self, endpoint_name, endpoint_config_name, tags=None, local_session=None):
  # runtime import since there is a cyclic dependency between entities and local_session
  from sagemaker.local import LocalSession
  self.local_session = local_session or LocalSession()
@@ -357,6 +359,7 @@ def __init__(self, endpoint_name, endpoint_config_name, local_session=None):
  self.name = endpoint_name
  self.endpoint_config = local_client.describe_endpoint_config(endpoint_config_name)
  self.production_variant = self.endpoint_config['ProductionVariants'][0]
+ self.tags = tags
 
  model_name = self.production_variant['ModelName']
  self.primary_container = local_client.describe_model(model_name)['PrimaryContainer']
@@ -392,6 +395,7 @@ def describe(self):
  'EndpointConfigName': self.endpoint_config['EndpointConfigName'],
  'CreationTime': self.create_time,
  'ProductionVariants': self.endpoint_config['ProductionVariants'],
+ 'Tags': self.tags,
  'EndpointName': self.name,
  'EndpointArn': _UNUSED_ARN,
  'EndpointStatus': self.state

src/sagemaker/local/local_session.py

@@ -127,9 +127,9 @@ def describe_endpoint_config(self, EndpointConfigName):
  'Code': 'ValidationException', 'Message': 'Could not find local endpoint config'}}
  raise ClientError(error_response, 'describe_endpoint_config')
 
- def create_endpoint_config(self, EndpointConfigName, ProductionVariants):
+ def create_endpoint_config(self, EndpointConfigName, ProductionVariants, Tags=None):
  LocalSagemakerClient._endpoint_configs[EndpointConfigName] = _LocalEndpointConfig(
- EndpointConfigName, ProductionVariants)
+ EndpointConfigName, ProductionVariants, Tags)
 
  def describe_endpoint(self, EndpointName):
  if EndpointName not in LocalSagemakerClient._endpoints:
@@ -138,8 +138,8 @@ def describe_endpoint(self, EndpointName):
  else:
  return LocalSagemakerClient._endpoints[EndpointName].describe()
 
- def create_endpoint(self, EndpointName, EndpointConfigName):
- endpoint = _LocalEndpoint(EndpointName, EndpointConfigName, self.sagemaker_session)
+ def create_endpoint(self, EndpointName, EndpointConfigName, Tags=None):
+ endpoint = _LocalEndpoint(EndpointName, EndpointConfigName, Tags, self.sagemaker_session)
  LocalSagemakerClient._endpoints[EndpointName] = endpoint
  endpoint.serve()
 

src/sagemaker/model.py

@@ -271,7 +271,8 @@ def deploy(self, initial_instance_count, instance_type, accelerator_type=None, e
  model_name=self.name,
  initial_instance_count=initial_instance_count,
  instance_type=instance_type,
- accelerator_type=accelerator_type)
+ accelerator_type=accelerator_type,
+ tags=tags)
  self.sagemaker_session.update_endpoint(self.endpoint_name, endpoint_config_name)
  else:
  self.sagemaker_session.endpoint_from_production_variants(self.endpoint_name, [production_variant], tags)

src/sagemaker/session.py

@@ -749,7 +749,7 @@ def create_endpoint_config(self, name, model_name, initial_instance_count, insta
  )
  return name
 
- def create_endpoint(self, endpoint_name, config_name, wait=True):
+ def create_endpoint(self, endpoint_name, config_name, tags=None, wait=True):
  """Create an Amazon SageMaker ``Endpoint`` according to the endpoint configuration specified in the request.
 
  Once the ``Endpoint`` is created, client applications can send requests to obtain inferences.
@@ -764,7 +764,10 @@ def create_endpoint(self, endpoint_name, config_name, wait=True):
  str: Name of the Amazon SageMaker ``Endpoint`` created.
  """
  LOGGER.info('Creating endpoint with name {}'.format(endpoint_name))
- self.sagemaker_client.create_endpoint(EndpointName=endpoint_name, EndpointConfigName=config_name)
+
+ tags = tags or []
+
+ self.sagemaker_client.create_endpoint(EndpointName=endpoint_name, EndpointConfigName=config_name, Tags=tags)
  if wait:
  self.wait_for_endpoint(endpoint_name)
  return endpoint_name
@@ -1052,7 +1055,7 @@ def endpoint_from_production_variants(self, name, production_variants, tags=None
  config_options['Tags'] = tags
 
  self.sagemaker_client.create_endpoint_config(**config_options)
- return self.create_endpoint(endpoint_name=name, config_name=name, wait=wait)
+ return self.create_endpoint(endpoint_name=name, config_name=name, tags=tags, wait=wait)
 
  def expand_role(self, role):
  """Expand an IAM role name into an ARN.

tests/integ/test_mxnet_train.py

@@ -78,6 +78,37 @@ def test_deploy_model(mxnet_training_job, sagemaker_session, mxnet_full_version)
  assert 'Could not find model' in str(exception.value)
 
 
+def test_deploy_model_with_tags(mxnet_training_job, sagemaker_session, mxnet_full_version):
+ endpoint_name = 'test-mxnet-deploy-model-{}'.format(sagemaker_timestamp())
+
+ with timeout_and_delete_endpoint_by_name(endpoint_name, sagemaker_session):
+ desc = sagemaker_session.sagemaker_client.describe_training_job(TrainingJobName=mxnet_training_job)
+ model_data = desc['ModelArtifacts']['S3ModelArtifacts']
+ script_path = os.path.join(DATA_DIR, 'mxnet_mnist', 'mnist.py')
+ model = MXNetModel(model_data, 'SageMakerRole', entry_point=script_path,
+ py_version=PYTHON_VERSION, sagemaker_session=sagemaker_session,
+ framework_version=mxnet_full_version)
+ tags = [{'Key': 'TagtestKey', 'Value': 'TagtestValue'}]
+ model.deploy(1, 'ml.m4.xlarge', endpoint_name=endpoint_name, tags=tags)
+
+ returned_model = sagemaker_session.describe_model(EndpointName=model.name)
+ returned_model_tags = sagemaker_session.list_tags(ResourceArn=returned_model['ModelArn'])['Tags']
+
+ endpoint = sagemaker_session.describe_endpoint(EndpointName=endpoint_name)
+ endpoint_tags = sagemaker_session.list_tags(ResourceArn=endpoint['EndpointArn'])['Tags']
+
+ endpoint_config = sagemaker_session.describe_endpoint_config(EndpointConfigName=endpoint['EndpointConfigName'])
+ endpoint_config_tags = sagemaker_session.list_tags(ResourceArn=endpoint_config['EndpointConfigArn'])['Tags']
+
+ production_variants = endpoint_config['ProductionVariants']
+
+ assert returned_model_tags == tags
+ assert endpoint_config_tags == tags
+ assert endpoint_tags == tags
+ assert production_variants[0]['InstanceType'] == 'ml.m4.xlarge'
+ assert production_variants[0]['InitialInstanceCount'] == 1
+
+
 def test_deploy_model_with_update_endpoint(mxnet_training_job, sagemaker_session, mxnet_full_version):
  endpoint_name = 'test-mxnet-deploy-model-{}'.format(sagemaker_timestamp())
 

tests/unit/test_create_deploy_entities.py

@@ -96,7 +96,7 @@ def test_create_endpoint_no_wait(sagemaker_session):
 
  assert returned_name == ENDPOINT_NAME
  sagemaker_session.sagemaker_client.create_endpoint.assert_called_once_with(
- EndpointName=ENDPOINT_NAME, EndpointConfigName=ENDPOINT_CONFIG_NAME)
+ EndpointName=ENDPOINT_NAME, EndpointConfigName=ENDPOINT_CONFIG_NAME, Tags=[])
 
 
 def test_create_endpoint_wait(sagemaker_session):
@@ -105,5 +105,5 @@ def test_create_endpoint_wait(sagemaker_session):
 
  assert returned_name == ENDPOINT_NAME
  sagemaker_session.sagemaker_client.create_endpoint.assert_called_once_with(
- EndpointName=ENDPOINT_NAME, EndpointConfigName=ENDPOINT_CONFIG_NAME)
+ EndpointName=ENDPOINT_NAME, EndpointConfigName=ENDPOINT_CONFIG_NAME, Tags=[])
  sagemaker_session.wait_for_endpoint.assert_called_once_with(ENDPOINT_NAME)

tests/unit/test_model.py

@@ -245,7 +245,8 @@ def test_deploy_update_endpoint(sagemaker_session, tmpdir):
  model_name=model.name,
  initial_instance_count=INSTANCE_COUNT,
  instance_type=INSTANCE_TYPE,
- accelerator_type=ACCELERATOR_TYPE
+ accelerator_type=ACCELERATOR_TYPE,
+ tags=None
  )
  config_name = sagemaker_session.create_endpoint_config(
  name=model.name,

tests/unit/test_session.py

@@ -910,7 +910,8 @@ def test_endpoint_from_production_variants(sagemaker_session):
  ims.sagemaker_client.describe_endpoint_config = Mock(side_effect=ex)
  sagemaker_session.endpoint_from_production_variants('some-endpoint', pvs)
  sagemaker_session.sagemaker_client.create_endpoint.assert_called_with(EndpointConfigName='some-endpoint',
- EndpointName='some-endpoint')
+ EndpointName='some-endpoint',
+ Tags=[])
  sagemaker_session.sagemaker_client.create_endpoint_config.assert_called_with(
  EndpointConfigName='some-endpoint',
  ProductionVariants=pvs)
@@ -936,7 +937,8 @@ def test_endpoint_from_production_variants_with_tags(sagemaker_session):
  tags = [{'ModelName': 'TestModel'}]
  sagemaker_session.endpoint_from_production_variants('some-endpoint', pvs, tags)
  sagemaker_session.sagemaker_client.create_endpoint.assert_called_with(EndpointConfigName='some-endpoint',
- EndpointName='some-endpoint')
+ EndpointName='some-endpoint',
+ Tags=tags)
  sagemaker_session.sagemaker_client.create_endpoint_config.assert_called_with(
  EndpointConfigName='some-endpoint',
  ProductionVariants=pvs,
@@ -953,7 +955,8 @@ def test_endpoint_from_production_variants_with_accelerator_type(sagemaker_sessi
  tags = [{'ModelName': 'TestModel'}]
  sagemaker_session.endpoint_from_production_variants('some-endpoint', pvs, tags)
  sagemaker_session.sagemaker_client.create_endpoint.assert_called_with(EndpointConfigName='some-endpoint',
- EndpointName='some-endpoint')
+ EndpointName='some-endpoint',
+ Tags=tags)
  sagemaker_session.sagemaker_client.create_endpoint_config.assert_called_with(
  EndpointConfigName='some-endpoint',
  ProductionVariants=pvs,

src/sagemaker/local/entities.py

@@ -327,15 +327,17 @@ def describe(self):
 
 class _LocalEndpointConfig(object):
 
- def __init__(self, config_name, production_variants):
+ def __init__(self, config_name, production_variants, tags=None):
  self.name = config_name
  self.production_variants = production_variants
+ self.tags = tags
  self.creation_time = datetime.datetime.now()
 
  def describe(self):
  response = {
  'EndpointConfigName': self.name,
  'EndpointConfigArn': _UNUSED_ARN,
+ 'Tags': self.tags,
  'CreationTime': self.creation_time,
  'ProductionVariants': self.production_variants
  }
@@ -348,7 +350,7 @@ class _LocalEndpoint(object):
  _IN_SERVICE = 'InService'
  _FAILED = 'Failed'
 
- def __init__(self, endpoint_name, endpoint_config_name, local_session=None):
+ def __init__(self, endpoint_name, endpoint_config_name, tags=None, local_session=None):
  # runtime import since there is a cyclic dependency between entities and local_session
  from sagemaker.local import LocalSession
  self.local_session = local_session or LocalSession()
@@ -357,6 +359,7 @@ def __init__(self, endpoint_name, endpoint_config_name, local_session=None):
  self.name = endpoint_name
  self.endpoint_config = local_client.describe_endpoint_config(endpoint_config_name)
  self.production_variant = self.endpoint_config['ProductionVariants'][0]
+ self.tags = tags
 
  model_name = self.production_variant['ModelName']
  self.primary_container = local_client.describe_model(model_name)['PrimaryContainer']
@@ -392,6 +395,7 @@ def describe(self):
  'EndpointConfigName': self.endpoint_config['EndpointConfigName'],
  'CreationTime': self.create_time,
  'ProductionVariants': self.endpoint_config['ProductionVariants'],
+ 'Tags': self.tags,
  'EndpointName': self.name,
  'EndpointArn': _UNUSED_ARN,
  'EndpointStatus': self.state

src/sagemaker/local/local_session.py

@@ -127,9 +127,9 @@ def describe_endpoint_config(self, EndpointConfigName):
  'Code': 'ValidationException', 'Message': 'Could not find local endpoint config'}}
  raise ClientError(error_response, 'describe_endpoint_config')
 
- def create_endpoint_config(self, EndpointConfigName, ProductionVariants):
+ def create_endpoint_config(self, EndpointConfigName, ProductionVariants, Tags=None):
  LocalSagemakerClient._endpoint_configs[EndpointConfigName] = _LocalEndpointConfig(
- EndpointConfigName, ProductionVariants)
+ EndpointConfigName, ProductionVariants, Tags)
 
  def describe_endpoint(self, EndpointName):
  if EndpointName not in LocalSagemakerClient._endpoints:
@@ -138,8 +138,8 @@ def describe_endpoint(self, EndpointName):
  else:
  return LocalSagemakerClient._endpoints[EndpointName].describe()
 
- def create_endpoint(self, EndpointName, EndpointConfigName):
- endpoint = _LocalEndpoint(EndpointName, EndpointConfigName, self.sagemaker_session)
+ def create_endpoint(self, EndpointName, EndpointConfigName, Tags=None):
+ endpoint = _LocalEndpoint(EndpointName, EndpointConfigName, Tags, self.sagemaker_session)
  LocalSagemakerClient._endpoints[EndpointName] = endpoint
  endpoint.serve()
 

src/sagemaker/model.py

@@ -271,7 +271,8 @@ def deploy(self, initial_instance_count, instance_type, accelerator_type=None, e
  model_name=self.name,
  initial_instance_count=initial_instance_count,
  instance_type=instance_type,
- accelerator_type=accelerator_type)
+ accelerator_type=accelerator_type,
+ tags=tags)
  self.sagemaker_session.update_endpoint(self.endpoint_name, endpoint_config_name)
  else:
  self.sagemaker_session.endpoint_from_production_variants(self.endpoint_name, [production_variant], tags)

src/sagemaker/session.py

@@ -749,7 +749,7 @@ def create_endpoint_config(self, name, model_name, initial_instance_count, insta
  )
  return name
 
- def create_endpoint(self, endpoint_name, config_name, wait=True):
+ def create_endpoint(self, endpoint_name, config_name, tags=None, wait=True):
  """Create an Amazon SageMaker ``Endpoint`` according to the endpoint configuration specified in the request.
 
  Once the ``Endpoint`` is created, client applications can send requests to obtain inferences.
@@ -764,7 +764,10 @@ def create_endpoint(self, endpoint_name, config_name, wait=True):
  str: Name of the Amazon SageMaker ``Endpoint`` created.
  """
  LOGGER.info('Creating endpoint with name {}'.format(endpoint_name))
- self.sagemaker_client.create_endpoint(EndpointName=endpoint_name, EndpointConfigName=config_name)
+
+ tags = tags or []
+
+ self.sagemaker_client.create_endpoint(EndpointName=endpoint_name, EndpointConfigName=config_name, Tags=tags)
  if wait:
  self.wait_for_endpoint(endpoint_name)
  return endpoint_name
@@ -1052,7 +1055,7 @@ def endpoint_from_production_variants(self, name, production_variants, tags=None
  config_options['Tags'] = tags
 
  self.sagemaker_client.create_endpoint_config(**config_options)
- return self.create_endpoint(endpoint_name=name, config_name=name, wait=wait)
+ return self.create_endpoint(endpoint_name=name, config_name=name, tags=tags, wait=wait)
 
  def expand_role(self, role):
  """Expand an IAM role name into an ARN.

tests/integ/test_mxnet_train.py

@@ -78,6 +78,37 @@ def test_deploy_model(mxnet_training_job, sagemaker_session, mxnet_full_version)
  assert 'Could not find model' in str(exception.value)
 
 
+def test_deploy_model_with_tags(mxnet_training_job, sagemaker_session, mxnet_full_version):
+ endpoint_name = 'test-mxnet-deploy-model-{}'.format(sagemaker_timestamp())
+
+ with timeout_and_delete_endpoint_by_name(endpoint_name, sagemaker_session):
+ desc = sagemaker_session.sagemaker_client.describe_training_job(TrainingJobName=mxnet_training_job)
+ model_data = desc['ModelArtifacts']['S3ModelArtifacts']
+ script_path = os.path.join(DATA_DIR, 'mxnet_mnist', 'mnist.py')
+ model = MXNetModel(model_data, 'SageMakerRole', entry_point=script_path,
+ py_version=PYTHON_VERSION, sagemaker_session=sagemaker_session,
+ framework_version=mxnet_full_version)
+ tags = [{'Key': 'TagtestKey', 'Value': 'TagtestValue'}]
+ model.deploy(1, 'ml.m4.xlarge', endpoint_name=endpoint_name, tags=tags)
+
+ returned_model = sagemaker_session.describe_model(EndpointName=model.name)
+ returned_model_tags = sagemaker_session.list_tags(ResourceArn=returned_model['ModelArn'])['Tags']
+
+ endpoint = sagemaker_session.describe_endpoint(EndpointName=endpoint_name)
+ endpoint_tags = sagemaker_session.list_tags(ResourceArn=endpoint['EndpointArn'])['Tags']
+
+ endpoint_config = sagemaker_session.describe_endpoint_config(EndpointConfigName=endpoint['EndpointConfigName'])
+ endpoint_config_tags = sagemaker_session.list_tags(ResourceArn=endpoint_config['EndpointConfigArn'])['Tags']
+
+ production_variants = endpoint_config['ProductionVariants']
+
+ assert returned_model_tags == tags
+ assert endpoint_config_tags == tags
+ assert endpoint_tags == tags
+ assert production_variants[0]['InstanceType'] == 'ml.m4.xlarge'
+ assert production_variants[0]['InitialInstanceCount'] == 1
+
+
 def test_deploy_model_with_update_endpoint(mxnet_training_job, sagemaker_session, mxnet_full_version):
  endpoint_name = 'test-mxnet-deploy-model-{}'.format(sagemaker_timestamp())
 

tests/unit/test_create_deploy_entities.py

@@ -96,7 +96,7 @@ def test_create_endpoint_no_wait(sagemaker_session):
 
  assert returned_name == ENDPOINT_NAME
  sagemaker_session.sagemaker_client.create_endpoint.assert_called_once_with(
- EndpointName=ENDPOINT_NAME, EndpointConfigName=ENDPOINT_CONFIG_NAME)
+ EndpointName=ENDPOINT_NAME, EndpointConfigName=ENDPOINT_CONFIG_NAME, Tags=[])
 
 
 def test_create_endpoint_wait(sagemaker_session):
@@ -105,5 +105,5 @@ def test_create_endpoint_wait(sagemaker_session):
 
  assert returned_name == ENDPOINT_NAME
  sagemaker_session.sagemaker_client.create_endpoint.assert_called_once_with(
- EndpointName=ENDPOINT_NAME, EndpointConfigName=ENDPOINT_CONFIG_NAME)
+ EndpointName=ENDPOINT_NAME, EndpointConfigName=ENDPOINT_CONFIG_NAME, Tags=[])
  sagemaker_session.wait_for_endpoint.assert_called_once_with(ENDPOINT_NAME)

tests/unit/test_model.py

@@ -245,7 +245,8 @@ def test_deploy_update_endpoint(sagemaker_session, tmpdir):
  model_name=model.name,
  initial_instance_count=INSTANCE_COUNT,
  instance_type=INSTANCE_TYPE,
- accelerator_type=ACCELERATOR_TYPE
+ accelerator_type=ACCELERATOR_TYPE,
+ tags=None
  )
  config_name = sagemaker_session.create_endpoint_config(
  name=model.name,

tests/unit/test_session.py

@@ -910,7 +910,8 @@ def test_endpoint_from_production_variants(sagemaker_session):
  ims.sagemaker_client.describe_endpoint_config = Mock(side_effect=ex)
  sagemaker_session.endpoint_from_production_variants('some-endpoint', pvs)
  sagemaker_session.sagemaker_client.create_endpoint.assert_called_with(EndpointConfigName='some-endpoint',
- EndpointName='some-endpoint')
+ EndpointName='some-endpoint',
+ Tags=[])
  sagemaker_session.sagemaker_client.create_endpoint_config.assert_called_with(
  EndpointConfigName='some-endpoint',
  ProductionVariants=pvs)
@@ -936,7 +937,8 @@ def test_endpoint_from_production_variants_with_tags(sagemaker_session):
  tags = [{'ModelName': 'TestModel'}]
  sagemaker_session.endpoint_from_production_variants('some-endpoint', pvs, tags)
  sagemaker_session.sagemaker_client.create_endpoint.assert_called_with(EndpointConfigName='some-endpoint',
- EndpointName='some-endpoint')
+ EndpointName='some-endpoint',
+ Tags=tags)
  sagemaker_session.sagemaker_client.create_endpoint_config.assert_called_with(
  EndpointConfigName='some-endpoint',
  ProductionVariants=pvs,
@@ -953,7 +955,8 @@ def test_endpoint_from_production_variants_with_accelerator_type(sagemaker_sessi
  tags = [{'ModelName': 'TestModel'}]
  sagemaker_session.endpoint_from_production_variants('some-endpoint', pvs, tags)
  sagemaker_session.sagemaker_client.create_endpoint.assert_called_with(EndpointConfigName='some-endpoint',
- EndpointName='some-endpoint')
+ EndpointName='some-endpoint',
+ Tags=tags)
  sagemaker_session.sagemaker_client.create_endpoint_config.assert_called_with(
  EndpointConfigName='some-endpoint',
  ProductionVariants=pvs,