Add integration test for transformer tags

Author: imujjwal96

tests/integ/test_transformer.py

@@ -148,6 +148,40 @@ 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):
+    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'}]
+
+    mx = MXNet(entry_point=script_path, role='SageMakerRole', train_instance_count=1,
+               train_instance_type='ml.c4.xlarge', sagemaker_session=sagemaker_session,
+               framework_version=mxnet_full_version)
+
+    train_input = mx.sagemaker_session.upload_data(path=os.path.join(data_path, 'train'),
+                                                   key_prefix='integ-test-data/mxnet_mnist/train')
+    test_input = mx.sagemaker_session.upload_data(path=os.path.join(data_path, 'test'),
+                                                  key_prefix='integ-test-data/mxnet_mnist/test')
+    job_name = unique_name_from_base('test-mxnet-transform')
+
+    with timeout(minutes=TRAINING_DEFAULT_TIMEOUT_MINUTES):
+        mx.fit({'train': train_input, 'test': test_input}, job_name=job_name)
+
+    transform_input_path = os.path.join(data_path, 'transform', 'data.csv')
+    transform_input_key_prefix = 'integ-test-data/mxnet_mnist/transform'
+    transform_input = mx.sagemaker_session.upload_data(path=transform_input_path,
+                                                       key_prefix=transform_input_key_prefix)
+
+    transformer = mx.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')