change: refactor endpoint support for TF transformer

Author: caxiaohu

src/sagemaker/tensorflow/estimator.py

@@ -644,21 +644,27 @@ def transformer(
                 use the SageMaker Tensorflow Serving container.
         """
 
-        if endpoint_type == "tensorflow-serving":
-            self.script_mode = True
-        return super(TensorFlow, self).transformer(
+        role = role or self.role
+        model = self.create_model(
+            model_server_workers=model_server_workers,
+            role=role,
+            vpc_config_override=VPC_CONFIG_DEFAULT,
+            endpoint_type=endpoint_type,
+        )
+        return model.transformer(
             instance_count,
             instance_type,
-            strategy,
-            assemble_with,
-            output_path,
-            output_kms_key,
-            accept,
-            env,
-            max_concurrent_transforms,
-            max_payload,
-            tags,
-            role,
-            model_server_workers,
-            volume_kms_key,
+            strategy=strategy,
+            assemble_with=assemble_with,
+            output_path=output_path,
+            output_kms_key=output_kms_key,
+            accept=accept,
+            env=env,
+            max_concurrent_transforms=max_concurrent_transforms,
+            max_payload=max_payload,
+            tags=None,
+            volume_kms_key=volume_kms_key,
+
         )
+
+

tests/unit/test_tf_estimator.py

@@ -308,29 +308,26 @@ def test_create_model_with_optional_params(sagemaker_session):
     assert model.vpc_config == vpc_config
 
 
-@patch("sagemaker.tensorflow.estimator.TensorFlow._create_tfs_model")
-def test_transformer_creation_with_endpoint_type(create_tfs_model, sagemaker_session):
+@patch("sagemaker.tensorflow.estimator.TensorFlow.create_model")
+def test_transformer_creation_with_endpoint_type(create_model, sagemaker_session):
     tf = TensorFlow(
         entry_point=SCRIPT_PATH,
         role=ROLE,
         sagemaker_session=sagemaker_session,
         train_instance_count=INSTANCE_COUNT,
         train_instance_type=INSTANCE_TYPE,
     )
-
     tf.latest_training_job = _TrainingJob(sagemaker_session, JOB_NAME)
-    transformer = tf.transformer(INSTANCE_COUNT, INSTANCE_TYPE, endpoint_type="tensorflow-serving")
-    assert isinstance(transformer, Transformer)
-    assert transformer.sagemaker_session == sagemaker_session
-    assert transformer.instance_count == INSTANCE_COUNT
-    assert transformer.instance_type == INSTANCE_TYPE
-    assert tf.script_mode is True
-    assert tf._script_mode_enabled() is True
-    create_tfs_model.assert_called_once()
+
+    tf.transformer(INSTANCE_COUNT, INSTANCE_TYPE, model_server_workers=2, endpoint_type="tensorflow-serving")
+    create_model.assert_called_with(endpoint_type='tensorflow-serving',
+                                    model_server_workers=2,
+                                    role='Dummy',
+                                    vpc_config_override='VPC_CONFIG_DEFAULT')
 
 
-@patch("sagemaker.tensorflow.estimator.TensorFlow._create_default_model")
-def test_transformer_creation_without_endpoint_type(create_default_model, sagemaker_session):
+@patch("sagemaker.tensorflow.estimator.TensorFlow.create_model")
+def test_transformer_creation_without_endpoint_type(create_model, sagemaker_session):
 
     tf = TensorFlow(
         entry_point=SCRIPT_PATH,
@@ -341,14 +338,11 @@ def test_transformer_creation_without_endpoint_type(create_default_model, sagema
     )
 
     tf.latest_training_job = _TrainingJob(sagemaker_session, JOB_NAME)
-    transformer = tf.transformer(INSTANCE_COUNT, INSTANCE_TYPE)
-    assert isinstance(transformer, Transformer)
-    assert transformer.sagemaker_session == sagemaker_session
-    assert transformer.instance_count == INSTANCE_COUNT
-    assert transformer.instance_type == INSTANCE_TYPE
-    assert tf.script_mode is False
-    assert tf._script_mode_enabled() is False
-    create_default_model.assert_called_once()
+    transformer = tf.transformer(INSTANCE_COUNT, INSTANCE_TYPE, model_server_workers=4)
+    create_model.assert_called_with(endpoint_type=None,
+                                    model_server_workers=4,
+                                    role='Dummy',
+                                    vpc_config_override='VPC_CONFIG_DEFAULT')
 
 
 def test_create_model_with_custom_image(sagemaker_session):