change: allow specifying model name when creating a Transformer from an Estimator (#1398)

Author: laurenyu

src/sagemaker/estimator.py

@@ -826,6 +826,7 @@ def transformer(
         volume_kms_key=None,
         vpc_config_override=vpc_utils.VPC_CONFIG_DEFAULT,
         enable_network_isolation=None,
+        model_name=None,
     ):
         """Return a ``Transformer`` that uses a SageMaker Model based on the
         training job. It reuses the SageMaker Session and base job name used by
@@ -876,6 +877,8 @@ def transformer(
                 user entry script for inference. Also known as Internet-free mode.
                 If not specified, this setting is taken from the estimator's
                 current configuration.
+            model_name (str): Name to use for creating an Amazon SageMaker
+                model. If not specified, the name of the training job is used.
         """
         tags = tags or self.tags
 
@@ -884,9 +887,9 @@ def transformer(
                 "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
+            model_name = model_name or self._current_job_name
         else:
-            model_name = self.latest_training_job.name
+            model_name = model_name or self.latest_training_job.name
             if enable_network_isolation is None:
                 enable_network_isolation = self.enable_network_isolation()
 
@@ -1897,6 +1900,7 @@ def transformer(
         entry_point=None,
         vpc_config_override=vpc_utils.VPC_CONFIG_DEFAULT,
         enable_network_isolation=None,
+        model_name=None,
     ):
         """Return a ``Transformer`` that uses a SageMaker Model based on the
         training job. It reuses the SageMaker Session and base job name used by
@@ -1953,6 +1957,8 @@ def transformer(
                 user entry script for inference. Also known as Internet-free mode.
                 If not specified, this setting is taken from the estimator's
                 current configuration.
+            model_name (str): Name to use for creating an Amazon SageMaker
+                model. If not specified, the name of the training job is used.
 
         Returns:
             sagemaker.transformer.Transformer: a ``Transformer`` object that can be used to start a
@@ -1972,6 +1978,7 @@ def transformer(
                 vpc_config_override=vpc_config_override,
                 model_kms_key=self.output_kms_key,
                 enable_network_isolation=enable_network_isolation,
+                name=model_name,
             )
             model._create_sagemaker_model(instance_type, tags=tags)
 
@@ -1984,7 +1991,7 @@ def transformer(
                 "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
+            model_name = model_name or self._current_job_name
             transform_env = env or {}
 
         return Transformer(

src/sagemaker/tensorflow/estimator.py

@@ -780,6 +780,7 @@ def transformer(
         entry_point=None,
         vpc_config_override=VPC_CONFIG_DEFAULT,
         enable_network_isolation=None,
+        model_name=None,
     ):
         """Return a ``Transformer`` that uses a SageMaker Model based on the training job. It
         reuses the SageMaker Session and base job name used by the Estimator.
@@ -837,6 +838,8 @@ def transformer(
                 user entry script for inference. Also known as Internet-free mode.
                 If not specified, this setting is taken from the estimator's
                 current configuration.
+            model_name (str): Name to use for creating an Amazon SageMaker
+                model. If not specified, the name of the training job is used.
         """
         role = role or self.role
 
@@ -846,7 +849,7 @@ def transformer(
                 "this estimator is only used for building workflow config"
             )
             return Transformer(
-                self._current_job_name,
+                model_name or self._current_job_name,
                 instance_count,
                 instance_type,
                 strategy=strategy,
@@ -873,6 +876,7 @@ def transformer(
             endpoint_type=endpoint_type,
             entry_point=entry_point,
             enable_network_isolation=enable_network_isolation,
+            name=model_name,
         )
 
         return model.transformer(

tests/unit/test_estimator.py

@@ -1374,6 +1374,7 @@ def test_framework_transformer_creation_with_optional_params(name_from_image, sa
     env = {"FOO": "BAR"}
     new_role = "dummy-model-role"
     new_vpc_config = {"Subnets": ["x"], "SecurityGroupIds": ["y"]}
+    model_name = "model-name"
 
     transformer = fw.transformer(
         INSTANCE_COUNT,
@@ -1392,10 +1393,11 @@ def test_framework_transformer_creation_with_optional_params(name_from_image, sa
         model_server_workers=1,
         vpc_config_override=new_vpc_config,
         enable_network_isolation=True,
+        model_name=model_name,
     )
 
     sagemaker_session.create_model.assert_called_with(
-        MODEL_IMAGE,
+        model_name,
         new_role,
         MODEL_CONTAINER_DEF,
         vpc_config=new_vpc_config,
@@ -1413,6 +1415,7 @@ def test_framework_transformer_creation_with_optional_params(name_from_image, sa
     assert transformer.base_transform_job_name == base_name
     assert transformer.tags == TAGS
     assert transformer.volume_kms_key == kms_key
+    assert transformer.model_name == model_name
 
 
 def test_ensure_latest_training_job(sagemaker_session):
@@ -1492,6 +1495,7 @@ def test_estimator_transformer_creation_with_optional_params(create_model, sagem
     max_payload = 6
     env = {"FOO": "BAR"}
     new_vpc_config = {"Subnets": ["x"], "SecurityGroupIds": ["y"]}
+    model_name = "model-name"
 
     transformer = estimator.transformer(
         INSTANCE_COUNT,
@@ -1508,6 +1512,7 @@ def test_estimator_transformer_creation_with_optional_params(create_model, sagem
         role=ROLE,
         vpc_config_override=new_vpc_config,
         enable_network_isolation=True,
+        model_name=model_name,
     )
 
     create_model.assert_called_with(
@@ -1524,6 +1529,7 @@ def test_estimator_transformer_creation_with_optional_params(create_model, sagem
     assert transformer.env == env
     assert transformer.base_transform_job_name == base_name
     assert transformer.tags == TAGS
+    assert transformer.model_name == model_name
 
 
 # _TrainingJob 'utils'

tests/unit/test_tf_estimator.py

@@ -364,6 +364,7 @@ def test_transformer_creation_with_optional_args(create_model, sagemaker_session
     new_role = "role"
     model_server_workers = 2
     vpc_config = {"Subnets": ["1234"], "SecurityGroupIds": ["5678"]}
+    model_name = "model-name"
 
     tf.transformer(
         INSTANCE_COUNT,
@@ -384,6 +385,7 @@ def test_transformer_creation_with_optional_args(create_model, sagemaker_session
         entry_point=SERVING_SCRIPT_FILE,
         vpc_config_override=vpc_config,
         enable_network_isolation=True,
+        model_name=model_name,
     )
 
     create_model.assert_called_with(
@@ -393,6 +395,7 @@ def test_transformer_creation_with_optional_args(create_model, sagemaker_session
         endpoint_type="tensorflow-serving",
         entry_point=SERVING_SCRIPT_FILE,
         enable_network_isolation=True,
+        name=model_name,
     )
     model.transformer.assert_called_with(
         INSTANCE_COUNT,
@@ -432,6 +435,7 @@ def test_transformer_creation_without_optional_args(create_model, sagemaker_sess
         vpc_config_override="VPC_CONFIG_DEFAULT",
         entry_point=None,
         enable_network_isolation=False,
+        name=None,
     )
     model.transformer.assert_called_with(
         INSTANCE_COUNT,