fix: expose vpc_config_override in transformer() methods (#1042)

Author: laurenyu

src/sagemaker/estimator.py

@@ -47,7 +47,7 @@
 from sagemaker.session import Session
 from sagemaker.session import s3_input
 from sagemaker.transformer import Transformer
-from sagemaker.utils import base_name_from_image, name_from_base, name_from_image, get_config_value
+from sagemaker.utils import base_name_from_image, name_from_base, get_config_value
 from sagemaker import vpc_utils
 
 
@@ -667,6 +667,7 @@ def transformer(
         tags=None,
         role=None,
         volume_kms_key=None,
+        vpc_config_override=vpc_utils.VPC_CONFIG_DEFAULT,
     ):
         """Return a ``Transformer`` that uses a SageMaker Model based on the
         training job. It reuses the SageMaker Session and base job name used by
@@ -702,6 +703,11 @@ def transformer(
                 role from the Estimator will be used.
             volume_kms_key (str): Optional. KMS key ID for encrypting the volume
                 attached to the ML compute instance (default: None).
+            vpc_config_override (dict[str, list[str]]): Optional override for the
+                VpcConfig set on the model.
+                Default: use subnets and security groups from this Estimator.
+                * 'Subnets' (list[str]): List of subnet ids.
+                * 'SecurityGroupIds' (list[str]): List of security group ids.
         """
         tags = tags or self.tags
 
@@ -714,7 +720,7 @@ def transformer(
         else:
             model_name = self.latest_training_job.name
 
-        model = self.create_model()
+        model = self.create_model(vpc_config_override=vpc_config_override)
 
         # not all create_model() implementations have the same kwargs
         model.name = model_name
@@ -1635,6 +1641,7 @@ def transformer(
         model_server_workers=None,
         volume_kms_key=None,
         entry_point=None,
+        vpc_config_override=vpc_utils.VPC_CONFIG_DEFAULT,
     ):
         """Return a ``Transformer`` that uses a SageMaker Model based on the
         training job. It reuses the SageMaker Session and base job name used by
@@ -1676,25 +1683,29 @@ def transformer(
             entry_point (str): Path (absolute or relative) to the local Python source file which
                 should be executed as the entry point to training. If not specified, the training
                 entry point is used.
+            vpc_config_override (dict[str, list[str]]): Optional override for
+                the VpcConfig set on the model.
+                Default: use subnets and security groups from this Estimator.
+                * 'Subnets' (list[str]): List of subnet ids.
+                * 'SecurityGroupIds' (list[str]): List of security group ids.
 
         Returns:
             sagemaker.transformer.Transformer: a ``Transformer`` object that can be used to start a
                 SageMaker Batch Transform job.
         """
         role = role or self.role
+        tags = tags or self.tags
 
         if self.latest_training_job is not None:
             model = self.create_model(
-                role=role, model_server_workers=model_server_workers, entry_point=entry_point
+                role=role,
+                model_server_workers=model_server_workers,
+                entry_point=entry_point,
+                vpc_config_override=vpc_config_override,
             )
+            model._create_sagemaker_model(instance_type, tags=tags)
 
-            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
-            )
+            model_name = model.name
             transform_env = model.env.copy()
             if env is not None:
                 transform_env.update(env)
@@ -1706,7 +1717,6 @@ def transformer(
             model_name = self._current_job_name
             transform_env = env or {}
 
-        tags = tags or self.tags
         return Transformer(
             model_name,
             instance_count,

src/sagemaker/tensorflow/estimator.py

@@ -704,6 +704,7 @@ def transformer(
         volume_kms_key=None,
         endpoint_type=None,
         entry_point=None,
+        vpc_config_override=VPC_CONFIG_DEFAULT,
     ):
         """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.
@@ -746,13 +747,18 @@ def transformer(
                 should be executed as the entry point to training. If not specified and
                 ``endpoint_type`` is 'tensorflow-serving', no entry point is used. If
                 ``endpoint_type`` is also ``None``, then the training entry point is used.
+            vpc_config_override (dict[str, list[str]]): Optional override for
+                the VpcConfig set on the model.
+                Default: use subnets and security groups from this Estimator.
+                * 'Subnets' (list[str]): List of subnet ids.
+                * 'SecurityGroupIds' (list[str]): List of security group ids.
         """
 
         role = role or self.role
         model = self.create_model(
             model_server_workers=model_server_workers,
             role=role,
-            vpc_config_override=VPC_CONFIG_DEFAULT,
+            vpc_config_override=vpc_config_override,
             endpoint_type=endpoint_type,
             entry_point=entry_point,
         )

tests/unit/test_estimator.py

@@ -21,6 +21,7 @@
 import pytest
 from mock import ANY, MagicMock, Mock, patch
 
+from sagemaker import vpc_utils
 from sagemaker.amazon.amazon_estimator import registry
 from sagemaker.algorithm import AlgorithmEstimator
 from sagemaker.estimator import Estimator, EstimatorBase, Framework, _TrainingJob
@@ -112,9 +113,19 @@ class DummyFramework(Framework):
     def train_image(self):
         return IMAGE_NAME
 
-    def create_model(self, role=None, model_server_workers=None, entry_point=None):
+    def create_model(
+        self,
+        role=None,
+        model_server_workers=None,
+        entry_point=None,
+        vpc_config_override=vpc_utils.VPC_CONFIG_DEFAULT,
+    ):
         return DummyFrameworkModel(
-            self.sagemaker_session, vpc_config=self.get_vpc_config(), entry_point=entry_point
+            self.sagemaker_session,
+            vpc_config=self.get_vpc_config(vpc_config_override),
+            entry_point=entry_point,
+            enable_network_isolation=self.enable_network_isolation(),
+            role=role,
         )
 
     @classmethod
@@ -127,12 +138,11 @@ def _prepare_init_params_from_job_description(cls, job_details, model_channel_na
 
 
 class DummyFrameworkModel(FrameworkModel):
-    def __init__(self, sagemaker_session, entry_point=None, **kwargs):
+    def __init__(self, sagemaker_session, entry_point=None, role=ROLE, **kwargs):
         super(DummyFrameworkModel, self).__init__(
             MODEL_DATA,
             MODEL_IMAGE,
-            INSTANCE_TYPE,
-            ROLE,
+            role,
             entry_point or ENTRY_POINT,
             sagemaker_session=sagemaker_session,
             **kwargs
@@ -141,7 +151,7 @@ def __init__(self, sagemaker_session, entry_point=None, **kwargs):
     def create_predictor(self, endpoint_name):
         return None
 
-    def prepare_container_def(self, instance_type):
+    def prepare_container_def(self, instance_type, accelerator_type=None):
         return MODEL_CONTAINER_DEF
 
 
@@ -1280,22 +1290,30 @@ def test_init_with_source_dir_s3(strftime, sagemaker_session):
     assert fw._hyperparameters == expected_hyperparameters
 
 
-@patch("sagemaker.estimator.name_from_image", return_value=MODEL_IMAGE)
+@patch("sagemaker.model.utils.name_from_image", return_value=MODEL_IMAGE)
 def test_framework_transformer_creation(name_from_image, sagemaker_session):
+    vpc_config = {"Subnets": ["foo"], "SecurityGroupIds": ["bar"]}
     fw = DummyFramework(
         entry_point=SCRIPT_PATH,
         role=ROLE,
         train_instance_count=INSTANCE_COUNT,
         train_instance_type=INSTANCE_TYPE,
         sagemaker_session=sagemaker_session,
+        subnets=vpc_config["Subnets"],
+        security_group_ids=vpc_config["SecurityGroupIds"],
     )
     fw.latest_training_job = _TrainingJob(sagemaker_session, JOB_NAME)
 
     transformer = fw.transformer(INSTANCE_COUNT, INSTANCE_TYPE)
 
     name_from_image.assert_called_with(MODEL_IMAGE)
     sagemaker_session.create_model.assert_called_with(
-        MODEL_IMAGE, ROLE, MODEL_CONTAINER_DEF, None, tags=None
+        MODEL_IMAGE,
+        ROLE,
+        MODEL_CONTAINER_DEF,
+        tags=None,
+        vpc_config=vpc_config,
+        enable_network_isolation=False,
     )
 
     assert isinstance(transformer, Transformer)
@@ -1307,7 +1325,7 @@ def test_framework_transformer_creation(name_from_image, sagemaker_session):
     assert transformer.env == {}
 
 
-@patch("sagemaker.estimator.name_from_image", return_value=MODEL_IMAGE)
+@patch("sagemaker.model.utils.name_from_image", return_value=MODEL_IMAGE)
 def test_framework_transformer_creation_with_optional_params(name_from_image, sagemaker_session):
     base_name = "foo"
     vpc_config = {"Subnets": ["foo"], "SecurityGroupIds": ["bar"]}
@@ -1320,6 +1338,7 @@ def test_framework_transformer_creation_with_optional_params(name_from_image, sa
         base_job_name=base_name,
         subnets=vpc_config["Subnets"],
         security_group_ids=vpc_config["SecurityGroupIds"],
+        enable_network_isolation=True,
     )
     fw.latest_training_job = _TrainingJob(sagemaker_session, JOB_NAME)
 
@@ -1331,6 +1350,7 @@ def test_framework_transformer_creation_with_optional_params(name_from_image, sa
     max_payload = 6
     env = {"FOO": "BAR"}
     new_role = "dummy-model-role"
+    new_vpc_config = {"Subnets": ["x"], "SecurityGroupIds": ["y"]}
 
     transformer = fw.transformer(
         INSTANCE_COUNT,
@@ -1347,10 +1367,16 @@ def test_framework_transformer_creation_with_optional_params(name_from_image, sa
         env=env,
         role=new_role,
         model_server_workers=1,
+        vpc_config_override=new_vpc_config,
     )
 
     sagemaker_session.create_model.assert_called_with(
-        MODEL_IMAGE, new_role, MODEL_CONTAINER_DEF, vpc_config, tags=TAGS
+        MODEL_IMAGE,
+        new_role,
+        MODEL_CONTAINER_DEF,
+        vpc_config=new_vpc_config,
+        tags=TAGS,
+        enable_network_isolation=True,
     )
     assert transformer.strategy == strategy
     assert transformer.assemble_with == assemble_with