change:enable network isolation for amazon estimators

src/sagemaker/amazon/amazon_estimator.py

@@ -74,10 +74,13 @@ def __init__(
 
         if "enable_network_isolation" in kwargs:
             logger.debug(
-                "removing unused enable_network_isolation argument: %s",
+                "removing unused enable_network_isolation argument in base class: %s",
                 str(kwargs["enable_network_isolation"]),
             )
+            self._enable_network_isolation = kwargs["enable_network_isolation"]
             del kwargs["enable_network_isolation"]
+        else:
+            self._enable_network_isolation = False
 
         super(AmazonAlgorithmEstimatorBase, self).__init__(
             role, train_instance_count, train_instance_type, **kwargs
@@ -98,6 +101,14 @@ def hyperparameters(self):
         """Placeholder docstring"""
         return hp.serialize_all(self)
 
+    def enable_network_isolation(self):
+        """If this Estimator can use network isolation when running.
+
+        Returns:
+            bool: Whether this Estimator can use network isolation or not.
+        """
+        return self._enable_network_isolation
+
     @property
     def data_location(self):
         """Placeholder docstring"""

tests/integ/test_pca.py

@@ -41,6 +41,7 @@ def test_pca(sagemaker_session, cpu_instance_type):
             train_instance_type=cpu_instance_type,
             num_components=48,
             sagemaker_session=sagemaker_session,
+            enable_network_isolation=True,
         )
 
         pca.algorithm_mode = "randomized"
@@ -50,7 +51,10 @@ def test_pca(sagemaker_session, cpu_instance_type):
 
     with timeout_and_delete_endpoint_by_name(job_name, sagemaker_session):
         pca_model = sagemaker.amazon.pca.PCAModel(
-            model_data=pca.model_data, role="SageMakerRole", sagemaker_session=sagemaker_session
+            model_data=pca.model_data,
+            role="SageMakerRole",
+            sagemaker_session=sagemaker_session,
+            enable_network_isolation=True,
         )
         predictor = pca_model.deploy(
             initial_instance_count=1, instance_type=cpu_instance_type, endpoint_name=job_name

tests/unit/test_amazon_estimator.py

@@ -93,6 +93,18 @@ def test_gov_ecr_uri():
 def test_init(sagemaker_session):
     pca = PCA(num_components=55, sagemaker_session=sagemaker_session, **COMMON_ARGS)
     assert pca.num_components == 55
+    assert pca.enable_network_isolation() is False
+
+
+def test_init_enable_network_isolation(sagemaker_session):
+    pca = PCA(
+        num_components=55,
+        sagemaker_session=sagemaker_session,
+        enable_network_isolation=True,
+        **COMMON_ARGS
+    )
+    assert pca.num_components == 55
+    assert pca.enable_network_isolation() is True
 
 
 def test_init_all_pca_hyperparameters(sagemaker_session):