Add integration test for estimator transformer

Author: imujjwal96

src/sagemaker/estimator.py

@@ -500,15 +500,16 @@ def transformer(self, instance_count, instance_type, strategy=None, assemble_wit
             volume_kms_key (str): Optional. KMS key ID for encrypting the volume attached to the ML
                 compute instance (default: None).
         """
+        tags = tags or self.tags
+
         if self.latest_training_job is not None:
-            model_name = self.sagemaker_session.create_model_from_job(self.latest_training_job.name, role=role)
+            model_name = self.sagemaker_session.create_model_from_job(self.latest_training_job.name, role=role,
+                                                                      tags=tags)
         else:
             logging.warning('No finished training job found associated with this estimator. Please make sure'
                             'this estimator is only used for building workflow config')
             model_name = self._current_job_name
 
-        tags = tags or self.tags
-
         return Transformer(model_name, instance_count, instance_type, strategy=strategy, assemble_with=assemble_with,
                            output_path=output_path, output_kms_key=output_kms_key, accept=accept,
                            max_concurrent_transforms=max_concurrent_transforms, max_payload=max_payload,
@@ -1061,6 +1062,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
+            tags = tags or self.tags
             self.sagemaker_session.create_model(model_name, role, container_def, vpc_config, tags=tags)
             transform_env = model.env.copy()
             if env is not None:

src/sagemaker/session.py

@@ -620,7 +620,8 @@ def create_model(self, name, role, container_defs, vpc_config=None,
         return name
 
     def create_model_from_job(self, training_job_name, name=None, role=None, primary_container_image=None,
-                              model_data_url=None, env=None, vpc_config_override=vpc_utils.VPC_CONFIG_DEFAULT):
+                              model_data_url=None, env=None, vpc_config_override=vpc_utils.VPC_CONFIG_DEFAULT,
+                              tags=None):
         """Create an Amazon SageMaker ``Model`` from a SageMaker Training Job.
 
         Args:
@@ -638,12 +639,13 @@ def create_model_from_job(self, training_job_name, name=None, role=None, primary
                 Default: use VpcConfig from training job.
                 * 'Subnets' (list[str]): List of subnet ids.
                 * 'SecurityGroupIds' (list[str]): List of security group ids.
+            tags(List[dict[str, str]]): Optional. The list of tags to add to the model. For more, see
+                https://docs.aws.amazon.com/sagemaker/latest/dg/API_Tag.html.
 
         Returns:
             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 {}

tests/integ/test_transformer.py

@@ -22,6 +22,7 @@
 from sagemaker import KMeans
 from sagemaker.mxnet import MXNet
 from sagemaker.transformer import Transformer
+from sagemaker.estimator import Estimator
 from sagemaker.utils import unique_name_from_base
 from tests.integ import DATA_DIR, TRAINING_DEFAULT_TIMEOUT_MINUTES, TRANSFORM_DEFAULT_TIMEOUT_MINUTES
 from tests.integ.kms_utils import get_or_create_kms_key
@@ -148,7 +149,7 @@ def test_transform_mxnet_vpc(sagemaker_session, mxnet_full_version):
         assert [security_group_id] == model_desc['VpcConfig']['SecurityGroupIds']
 
 
-def test_transform_mxnet_logs(sagemaker_session, mxnet_full_version):
+def test_transform_mxnet_tags(sagemaker_session, mxnet_full_version):
     data_path = os.path.join(DATA_DIR, 'mxnet_mnist')
     script_path = os.path.join(data_path, 'mnist.py')
     tags = [{'Key': 'some-tag', 'Value': 'value-for-tag'}]
@@ -182,6 +183,55 @@ def test_transform_mxnet_logs(sagemaker_session, mxnet_full_version):
         assert tags == model_tags
 
 
+def test_transform_byo_estimator(sagemaker_session):
+    data_path = os.path.join(DATA_DIR, 'one_p_mnist')
+    pickle_args = {} if sys.version_info.major == 2 else {'encoding': 'latin1'}
+    tags = [{'Key': 'some-tag', 'Value': 'value-for-tag'}]
+
+    # Load the data into memory as numpy arrays
+    train_set_path = os.path.join(data_path, 'mnist.pkl.gz')
+    with gzip.open(train_set_path, 'rb') as f:
+        train_set, _, _ = pickle.load(f, **pickle_args)
+
+    kmeans = KMeans(role='SageMakerRole', train_instance_count=1,
+                    train_instance_type='ml.c4.xlarge', k=10, sagemaker_session=sagemaker_session,
+                    output_path='s3://{}/'.format(sagemaker_session.default_bucket()))
+
+    # set kmeans specific hp
+    kmeans.init_method = 'random'
+    kmeans.max_iterators = 1
+    kmeans.tol = 1
+    kmeans.num_trials = 1
+    kmeans.local_init_method = 'kmeans++'
+    kmeans.half_life_time_size = 1
+    kmeans.epochs = 1
+
+    records = kmeans.record_set(train_set[0][:100])
+
+    job_name = unique_name_from_base('test-kmeans-attach')
+
+    with timeout(minutes=TRAINING_DEFAULT_TIMEOUT_MINUTES):
+        kmeans.fit(records, job_name=job_name)
+
+    transform_input_path = os.path.join(data_path, 'transform_input.csv')
+    transform_input_key_prefix = 'integ-test-data/one_p_mnist/transform'
+    transform_input = kmeans.sagemaker_session.upload_data(path=transform_input_path,
+                                                           key_prefix=transform_input_key_prefix)
+
+    estimator = Estimator.attach(training_job_name=job_name,
+                                 sagemaker_session=sagemaker_session)
+
+    transformer = estimator.transformer(1, 'ml.m4.xlarge', tags=tags)
+    transformer.transform(transform_input, content_type='text/csv')
+
+    with timeout_and_delete_model_with_transformer(transformer, sagemaker_session,
+                                                   minutes=TRANSFORM_DEFAULT_TIMEOUT_MINUTES):
+        transformer.wait()
+        model_desc = sagemaker_session.sagemaker_client.describe_model(ModelName=transformer.model_name)
+        model_tags = sagemaker_session.sagemaker_client.list_tags(ResourceArn=model_desc['ModelArn'])['Tags']
+        assert tags == model_tags
+
+
 def _create_transformer_and_transform_job(estimator, transform_input, volume_kms_key=None):
     transformer = estimator.transformer(1, 'ml.m4.xlarge', volume_kms_key=volume_kms_key)
     transformer.transform(transform_input, content_type='text/csv')

tests/unit/test_estimator.py

@@ -698,7 +698,7 @@ def test_estimator_transformer_creation(sagemaker_session):
 
     transformer = estimator.transformer(INSTANCE_COUNT, INSTANCE_TYPE)
 
-    sagemaker_session.create_model_from_job.assert_called_with(JOB_NAME, role=None)
+    sagemaker_session.create_model_from_job.assert_called_with(JOB_NAME, role=None, tags=None)
     assert isinstance(transformer, Transformer)
     assert transformer.sagemaker_session == sagemaker_session
     assert transformer.instance_count == INSTANCE_COUNT
@@ -728,7 +728,7 @@ def test_estimator_transformer_creation_with_optional_params(sagemaker_session):
                                         max_concurrent_transforms=max_concurrent_transforms, max_payload=max_payload,
                                         env=env, role=ROLE)
 
-    sagemaker_session.create_model_from_job.assert_called_with(JOB_NAME, role=ROLE)
+    sagemaker_session.create_model_from_job.assert_called_with(JOB_NAME, role=ROLE, tags=TAGS)
     assert transformer.strategy == strategy
     assert transformer.assemble_with == assemble_with
     assert transformer.output_path == OUTPUT_PATH

tests/unit/test_session.py

@@ -862,9 +862,20 @@ 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)
+
+
+def test_create_model_from_job_with_tags(sagemaker_session):
+    ims = sagemaker_session
+    ims.sagemaker_client.describe_training_job.return_value = COMPLETED_DESCRIBE_JOB_RESULT
+    ims.create_model_from_job(JOB_NAME, tags=TAGS)
+
     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,
@@ -876,7 +887,6 @@ def test_create_model_from_job(sagemaker_session):
 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'
@@ -885,7 +895,6 @@ 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
@@ -900,7 +909,6 @@ 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