feature: emit estimator transformer tags to model

Author: imujjwal96

src/sagemaker/estimator.py

@@ -1061,7 +1061,7 @@ def transformer(self, instance_count, instance_type, strategy=None, assemble_wit
             container_def = model.prepare_container_def(instance_type)
             model_name = model.name or name_from_image(container_def['Image'])
             vpc_config = model.vpc_config
-            self.sagemaker_session.create_model(model_name, role, container_def, vpc_config)
+            self.sagemaker_session.create_model(model_name, role, container_def, vpc_config, tags=tags)
             transform_env = model.env.copy()
             if env is not None:
                 transform_env.update(env)

src/sagemaker/session.py

@@ -643,6 +643,7 @@ def create_model_from_job(self, training_job_name, name=None, role=None, primary
             str: The name of the created ``Model``.
         """
         training_job = self.sagemaker_client.describe_training_job(TrainingJobName=training_job_name)
+        tags = self.sagemaker_client.list_tags(ResourceArn=training_job['TrainingJobArn'])['Tags']
         name = name or training_job_name
         role = role or training_job['RoleArn']
         env = env or {}
@@ -651,7 +652,7 @@ def create_model_from_job(self, training_job_name, name=None, role=None, primary
             model_data_url=model_data_url or training_job['ModelArtifacts']['S3ModelArtifacts'],
             env=env)
         vpc_config = _vpc_config_from_training_job(training_job, vpc_config_override)
-        return self.create_model(name, role, primary_container, vpc_config=vpc_config)
+        return self.create_model(name, role, primary_container, vpc_config=vpc_config, tags=tags)
 
     def create_model_package_from_algorithm(self, name, description, algorithm_arn, model_data):
         """Create a SageMaker Model Package from the results of training with an Algorithm Package

tests/unit/test_estimator.py

@@ -624,7 +624,7 @@ def test_framework_transformer_creation(name_from_image, sagemaker_session):
     transformer = fw.transformer(INSTANCE_COUNT, INSTANCE_TYPE)
 
     name_from_image.assert_called_with(MODEL_IMAGE)
-    sagemaker_session.create_model.assert_called_with(MODEL_IMAGE, ROLE, MODEL_CONTAINER_DEF, None)
+    sagemaker_session.create_model.assert_called_with(MODEL_IMAGE, ROLE, MODEL_CONTAINER_DEF, None, tags=None)
 
     assert isinstance(transformer, Transformer)
     assert transformer.sagemaker_session == sagemaker_session
@@ -659,7 +659,7 @@ def test_framework_transformer_creation_with_optional_params(name_from_image, sa
                                  max_concurrent_transforms=max_concurrent_transforms, max_payload=max_payload,
                                  volume_kms_key=kms_key, env=env, role=new_role, model_server_workers=1)
 
-    sagemaker_session.create_model.assert_called_with(MODEL_IMAGE, new_role, MODEL_CONTAINER_DEF, vpc_config)
+    sagemaker_session.create_model.assert_called_with(MODEL_IMAGE, new_role, MODEL_CONTAINER_DEF, vpc_config, tags=TAGS)
     assert transformer.strategy == strategy
     assert transformer.assemble_with == assemble_with
     assert transformer.output_path == OUTPUT_PATH

tests/unit/test_session.py

@@ -291,6 +291,7 @@ def test_s3_input_all_arguments():
 }
 
 COMPLETED_DESCRIBE_JOB_RESULT = dict(DEFAULT_EXPECTED_TRAIN_JOB_ARGS)
+COMPLETED_DESCRIBE_JOB_RESULT.update({'TrainingJobArn': 'arn:aws:sagemaker:us-west-2:336:training-job/' + JOB_NAME})
 COMPLETED_DESCRIBE_JOB_RESULT.update({'TrainingJobStatus': 'Completed'})
 COMPLETED_DESCRIBE_JOB_RESULT.update(
     {'ModelArtifacts': {
@@ -861,18 +862,21 @@ def test_create_model_failure(expand_container_def, sagemaker_session):
 def test_create_model_from_job(sagemaker_session):
     ims = sagemaker_session
     ims.sagemaker_client.describe_training_job.return_value = COMPLETED_DESCRIBE_JOB_RESULT
+    ims.sagemaker_client.list_tags.return_value = {'Tags': TAGS}
     ims.create_model_from_job(JOB_NAME)
 
     assert call(TrainingJobName=JOB_NAME) in ims.sagemaker_client.describe_training_job.call_args_list
     ims.sagemaker_client.create_model.assert_called_with(ExecutionRoleArn=EXPANDED_ROLE,
                                                          ModelName=JOB_NAME,
                                                          PrimaryContainer=PRIMARY_CONTAINER,
-                                                         VpcConfig=VPC_CONFIG)
+                                                         VpcConfig=VPC_CONFIG,
+                                                         Tags=TAGS)
 
 
 def test_create_model_from_job_with_image(sagemaker_session):
     ims = sagemaker_session
     ims.sagemaker_client.describe_training_job.return_value = COMPLETED_DESCRIBE_JOB_RESULT
+    ims.sagemaker_client.list_tags.return_value = {'Tags': TAGS}
     ims.create_model_from_job(JOB_NAME, primary_container_image='some-image')
     [create_model_call] = ims.sagemaker_client.create_model.call_args_list
     assert dict(create_model_call[1]['PrimaryContainer'])['Image'] == 'some-image'
@@ -881,6 +885,7 @@ def test_create_model_from_job_with_image(sagemaker_session):
 def test_create_model_from_job_with_container_def(sagemaker_session):
     ims = sagemaker_session
     ims.sagemaker_client.describe_training_job.return_value = COMPLETED_DESCRIBE_JOB_RESULT
+    ims.sagemaker_client.list_tags.return_value = {'Tags': TAGS}
     ims.create_model_from_job(JOB_NAME, primary_container_image='some-image', model_data_url='some-data',
                               env={'a': 'b'})
     [create_model_call] = ims.sagemaker_client.create_model.call_args_list
@@ -895,6 +900,7 @@ def test_create_model_from_job_with_vpc_config_override(sagemaker_session):
 
     ims = sagemaker_session
     ims.sagemaker_client.describe_training_job.return_value = COMPLETED_DESCRIBE_JOB_RESULT
+    ims.sagemaker_client.list_tags.return_value = {'Tags': TAGS}
     ims.create_model_from_job(JOB_NAME, vpc_config_override=vpc_config_override)
     assert ims.sagemaker_client.create_model.call_args[1]['VpcConfig'] == vpc_config_override