change: enable network isolation for amazon estimators (#1293)

Author: chenliu0831

src/sagemaker/amazon/amazon_estimator.py

@@ -44,7 +44,13 @@ class AmazonAlgorithmEstimatorBase(EstimatorBase):
     repo_version = None
 
     def __init__(
-        self, role, train_instance_count, train_instance_type, data_location=None, **kwargs
+        self,
+        role,
+        train_instance_count,
+        train_instance_type,
+        data_location=None,
+        enable_network_isolation=False,
+        **kwargs
     ):
         """Initialize an AmazonAlgorithmEstimatorBase.
 
@@ -63,6 +69,10 @@ def __init__(
                 "s3://example-bucket/some-key-prefix/". Objects will be saved in
                 a unique sub-directory of the specified location. If None, a
                 default data location will be used.
+            enable_network_isolation (bool): Specifies whether container will
+                run in network isolation mode. Network isolation mode restricts
+                the container access to outside networks (such as the internet).
+                Also known as internet-free mode (default: ``False``).
             **kwargs: Additional parameters passed to
                 :class:`~sagemaker.estimator.EstimatorBase`.
 
@@ -71,14 +81,6 @@ def __init__(
             You can find additional parameters for initializing this class at
             :class:`~sagemaker.estimator.EstimatorBase`.
         """
-
-        if "enable_network_isolation" in kwargs:
-            logger.debug(
-                "removing unused enable_network_isolation argument: %s",
-                str(kwargs["enable_network_isolation"]),
-            )
-            del kwargs["enable_network_isolation"]
-
         super(AmazonAlgorithmEstimatorBase, self).__init__(
             role, train_instance_count, train_instance_type, **kwargs
         )
@@ -87,6 +89,7 @@ def __init__(
             self.sagemaker_session.default_bucket()
         )
         self._data_location = data_location
+        self._enable_network_isolation = enable_network_isolation
 
     def train_image(self):
         """Placeholder docstring"""
@@ -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):