change: use Estimator.create_model in Estimator.transformer (#1041)

Author: laurenyu

src/sagemaker/estimator.py

@@ -705,16 +705,23 @@ def transformer(
         """
         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, tags=tags
-            )
-        else:
+        if self.latest_training_job is None:
             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
+        else:
+            model_name = self.latest_training_job.name
+
+        model = self.create_model()
+
+        # not all create_model() implementations have the same kwargs
+        model.name = model_name
+        if role is not None:
+            model.role = role
+
+        model._create_sagemaker_model(instance_type, tags=tags)
 
         return Transformer(
             model_name,

tests/integ/test_transformer.py

@@ -279,14 +279,15 @@ def test_transform_byo_estimator(sagemaker_session, cpu_instance_type):
     with timeout(minutes=TRAINING_DEFAULT_TIMEOUT_MINUTES):
         kmeans.fit(records, job_name=job_name)
 
+    estimator = Estimator.attach(training_job_name=job_name, sagemaker_session=sagemaker_session)
+    estimator._enable_network_isolation = True
+
     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, cpu_instance_type, tags=tags)
     transformer.transform(transform_input, content_type="text/csv")
 
@@ -297,6 +298,8 @@ def test_transform_byo_estimator(sagemaker_session, cpu_instance_type):
         model_desc = sagemaker_session.sagemaker_client.describe_model(
             ModelName=transformer.model_name
         )
+        assert model_desc["EnableNetworkIsolation"]
+
         model_tags = sagemaker_session.sagemaker_client.list_tags(
             ResourceArn=model_desc["ModelArn"]
         )["Tags"]

tests/unit/test_estimator.py

@@ -1392,6 +1392,7 @@ def test_ensure_latest_training_job_failure(sagemaker_session):
     assert "Estimator is not associated with a training job" in str(e)
 
 
+@patch("sagemaker.estimator.Estimator.create_model", return_value=Mock())
 def test_estimator_transformer_creation(sagemaker_session):
     estimator = Estimator(
         image_name=IMAGE_NAME,
@@ -1401,11 +1402,9 @@ def test_estimator_transformer_creation(sagemaker_session):
         sagemaker_session=sagemaker_session,
     )
     estimator.latest_training_job = _TrainingJob(sagemaker_session, JOB_NAME)
-    sagemaker_session.create_model_from_job.return_value = JOB_NAME
 
     transformer = estimator.transformer(INSTANCE_COUNT, INSTANCE_TYPE)
 
-    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
@@ -1414,6 +1413,7 @@ def test_estimator_transformer_creation(sagemaker_session):
     assert transformer.tags is None
 
 
+@patch("sagemaker.estimator.Estimator.create_model", return_value=Mock())
 def test_estimator_transformer_creation_with_optional_params(sagemaker_session):
     base_name = "foo"
     estimator = Estimator(
@@ -1425,7 +1425,6 @@ def test_estimator_transformer_creation_with_optional_params(sagemaker_session):
         base_job_name=base_name,
     )
     estimator.latest_training_job = _TrainingJob(sagemaker_session, JOB_NAME)
-    sagemaker_session.create_model_from_job.return_value = JOB_NAME
 
     strategy = "MultiRecord"
     assemble_with = "Line"
@@ -1450,7 +1449,6 @@ def test_estimator_transformer_creation_with_optional_params(sagemaker_session):
         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