feat: add network isolation support for PipelineModel (#1943)

* fix: pass enable_network_isolation to create_model

* add unit test for network isolation

Author: ajaykarpur

src/sagemaker/pipeline.py

@@ -26,7 +26,14 @@ class PipelineModel(object):
     """
 
     def __init__(
-        self, models, role, predictor_cls=None, name=None, vpc_config=None, sagemaker_session=None
+        self,
+        models,
+        role,
+        predictor_cls=None,
+        name=None,
+        vpc_config=None,
+        sagemaker_session=None,
+        enable_network_isolation=False,
     ):
         """Initialize a SageMaker `Model` instance.
 
@@ -57,13 +64,18 @@ def __init__(
                 object, used for SageMaker interactions (default: None). If not
                 specified, one is created using the default AWS configuration
                 chain.
+            enable_network_isolation (bool): Default False. if True, enables
+                network isolation in the endpoint, isolating the model
+                container. No inbound or outbound network calls can be made to
+                or from the model container.Boolean
         """
         self.models = models
         self.role = role
         self.predictor_cls = predictor_cls
         self.name = name
         self.vpc_config = vpc_config
         self.sagemaker_session = sagemaker_session
+        self.enable_network_isolation = enable_network_isolation
         self.endpoint_name = None
 
     def pipeline_container_def(self, instance_type):
@@ -157,7 +169,11 @@ def deploy(
 
         self.name = self.name or name_from_image(containers[0]["Image"])
         self.sagemaker_session.create_model(
-            self.name, self.role, containers, vpc_config=self.vpc_config
+            self.name,
+            self.role,
+            containers,
+            vpc_config=self.vpc_config,
+            enable_network_isolation=self.enable_network_isolation,
         )
 
         production_variant = sagemaker.production_variant(
@@ -214,7 +230,11 @@ def _create_sagemaker_pipeline_model(self, instance_type):
 
         self.name = self.name or name_from_image(containers[0]["Image"])
         self.sagemaker_session.create_model(
-            self.name, self.role, containers, vpc_config=self.vpc_config
+            self.name,
+            self.role,
+            containers,
+            vpc_config=self.vpc_config,
+            enable_network_isolation=self.enable_network_isolation,
         )
 
     def transformer(

tests/unit/test_pipeline_model.py

@@ -298,3 +298,43 @@ def test_delete_model(tfo, time, sagemaker_session):
 
     pipeline_model.delete_model()
     sagemaker_session.delete_model.assert_called_with(pipeline_model.name)
+
+
+@patch("tarfile.open")
+@patch("time.strftime", return_value=TIMESTAMP)
+def test_network_isolation(tfo, time, sagemaker_session):
+    framework_model = DummyFrameworkModel(sagemaker_session)
+    sparkml_model = SparkMLModel(
+        model_data=MODEL_DATA_2, role=ROLE, sagemaker_session=sagemaker_session
+    )
+    model = PipelineModel(
+        models=[framework_model, sparkml_model],
+        role=ROLE,
+        sagemaker_session=sagemaker_session,
+        enable_network_isolation=True,
+    )
+    model.deploy(instance_type=INSTANCE_TYPE, initial_instance_count=1)
+
+    sagemaker_session.create_model.assert_called_with(
+        model.name,
+        ROLE,
+        [
+            {
+                "Image": "mi-1",
+                "Environment": {
+                    "SAGEMAKER_PROGRAM": "blah.py",
+                    "SAGEMAKER_SUBMIT_DIRECTORY": "s3://mybucket/mi-1-2017-10-10-14-14-15/sourcedir.tar.gz",
+                    "SAGEMAKER_CONTAINER_LOG_LEVEL": "20",
+                    "SAGEMAKER_REGION": "us-west-2",
+                },
+                "ModelDataUrl": "s3://bucket/model_1.tar.gz",
+            },
+            {
+                "Image": "246618743249.dkr.ecr.us-west-2.amazonaws.com/sagemaker-sparkml-serving:2.2",
+                "Environment": {},
+                "ModelDataUrl": "s3://bucket/model_2.tar.gz",
+            },
+        ],
+        vpc_config=None,
+        enable_network_isolation=True,
+    )