fix: preserve EnableNetworkIsolation setting in attach (#1063)

* fix: preserve EnableNetworkIsolation setting in attach

Author: icywang86rui

src/sagemaker/algorithm.py

@@ -53,6 +53,7 @@ def __init__(
         model_channel_name="model",
         metric_definitions=None,
         encrypt_inter_container_traffic=False,
+        **kwargs  # pylint: disable=W0613
     ):
         """Initialize an ``AlgorithmEstimator`` instance.
 
@@ -162,6 +163,8 @@ def __init__(
             model_channel_name:
             metric_definitions:
             encrypt_inter_container_traffic:
+            **kwargs: Additional kwargs. This is unused. It's only added for AlgorithmEstimator
+                to ignore the irrelevant arguments.
         """
         self.algorithm_arn = algorithm_arn
         super(AlgorithmEstimator, self).__init__(

src/sagemaker/amazon/amazon_estimator.py

@@ -59,6 +59,14 @@ def __init__(
                 default data location will be used.
             **kwargs:
         """
+
+        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
         )

src/sagemaker/estimator.py

@@ -620,6 +620,8 @@ class constructor
         init_params["base_job_name"] = job_details["TrainingJobName"]
         init_params["output_path"] = job_details["OutputDataConfig"]["S3OutputPath"]
         init_params["output_kms_key"] = job_details["OutputDataConfig"]["KmsKeyId"]
+        if "EnableNetworkIsolation" in job_details:
+            init_params["enable_network_isolation"] = job_details["EnableNetworkIsolation"]
 
         has_hps = "HyperParameters" in job_details
         init_params["hyperparameters"] = job_details["HyperParameters"] if has_hps else {}

tests/unit/test_estimator.py

@@ -79,6 +79,7 @@
     },
     "RoleArn": "arn:aws:iam::366:role/SageMakerRole",
     "ResourceConfig": {"VolumeSizeInGB": 30, "InstanceCount": 1, "InstanceType": "ml.c4.xlarge"},
+    "EnableNetworkIsolation": False,
     "StoppingCondition": {"MaxRuntimeInSeconds": 24 * 60 * 60},
     "TrainingJobName": "neo",
     "TrainingJobStatus": "Completed",
@@ -671,6 +672,7 @@ def test_enable_cloudwatch_metrics(sagemaker_session):
 def test_attach_framework(sagemaker_session):
     returned_job_description = RETURNED_JOB_DESCRIPTION.copy()
     returned_job_description["VpcConfig"] = {"Subnets": ["foo"], "SecurityGroupIds": ["bar"]}
+    returned_job_description["EnableNetworkIsolation"] = True
     sagemaker_session.sagemaker_client.describe_training_job = Mock(
         name="describe_training_job", return_value=returned_job_description
     )
@@ -694,6 +696,7 @@ def test_attach_framework(sagemaker_session):
     assert framework_estimator.security_group_ids == ["bar"]
     assert framework_estimator.encrypt_inter_container_traffic is False
     assert framework_estimator.tags == LIST_TAGS_RESULT["Tags"]
+    assert framework_estimator.enable_network_isolation() is True
 
 
 def test_attach_without_hyperparameters(sagemaker_session):

tox.ini

@@ -25,6 +25,7 @@ exclude =
 max-complexity = 10
 
 ignore =
+    C901,
     E203,  # whitespace before ':': Black disagrees with and explicitly violates this.
     FI10,
     FI12,