fix: forward network_isolation parameter to Estimators when False (aws#4543)

A sagemaker config file would override this parameter if it was True in
the config but False in the Estimator parameters.

Closes aws#4542

Author: jmahlik

src/sagemaker/estimator.py

@@ -2434,6 +2434,7 @@ def start_new(cls, estimator, inputs, experiment_config):
         """
         train_args = cls._get_train_args(estimator, inputs, experiment_config)
 
+        logger.debug("Train args after processing defaults: %s", train_args)
         estimator.sagemaker_session.train(**train_args)
 
         return cls(estimator.sagemaker_session, estimator._current_job_name)
@@ -2499,7 +2500,13 @@ def _get_train_args(cls, estimator, inputs, experiment_config):
 
         # enable_network_isolation may be a pipeline variable place holder object
         # which is parsed in execution time
-        if estimator.enable_network_isolation():
+
+        # Should be defaulted to False
+        train_args["enable_network_isolation"] = False
+
+        # Only change it if it's explicitly passed so the sagemaker config
+        # doesn't override the kwarg.
+        if estimator.enable_network_isolation() is not None:
             train_args["enable_network_isolation"] = estimator.enable_network_isolation()
 
         if estimator.max_retry_attempts is not None:

tests/unit/sagemaker/huggingface/test_estimator.py

@@ -145,6 +145,7 @@ def _create_train_job(version, base_framework_version):
         "environment": None,
         "retry_strategy": None,
         "experiment_config": None,
+        "enable_network_isolation": False,
         "debugger_hook_config": {
             "CollectionConfigurations": [],
             "S3OutputPath": "s3://{}/".format(BUCKET_NAME),

tests/unit/sagemaker/tensorflow/test_estimator.py

@@ -141,6 +141,7 @@ def _create_train_job(tf_version, horovod=False, ps=False, py_version="py2", smd
         "vpc_config": None,
         "metric_definitions": None,
         "environment": None,
+        "enable_network_isolation": False,
         "experiment_config": None,
         "profiler_config": {
             "DisableProfiler": False,

tests/unit/sagemaker/training_compiler/test_huggingface_pytorch_compiler.py

@@ -143,6 +143,7 @@ def _create_train_job(
         "stop_condition": {"MaxRuntimeInSeconds": 24 * 60 * 60},
         "tags": None,
         "vpc_config": None,
+        "enable_network_isolation": False,
         "metric_definitions": None,
         "environment": None,
         "retry_strategy": None,

tests/unit/sagemaker/training_compiler/test_huggingface_tensorflow_compiler.py

@@ -144,6 +144,7 @@ def _create_train_job(
         "environment": None,
         "retry_strategy": None,
         "experiment_config": EXPERIMENT_CONFIG,
+        "enable_network_isolation": False,
         "debugger_hook_config": {
             "CollectionConfigurations": [],
             "S3OutputPath": "s3://{}/".format(BUCKET_NAME),

tests/unit/sagemaker/training_compiler/test_pytorch_compiler.py

@@ -143,6 +143,7 @@ def _create_train_job(
         "environment": None,
         "retry_strategy": None,
         "experiment_config": EXPERIMENT_CONFIG,
+        "enable_network_isolation": False,
         "debugger_hook_config": {
             "CollectionConfigurations": [],
             "S3OutputPath": "s3://{}/".format(BUCKET_NAME),

tests/unit/sagemaker/training_compiler/test_tensorflow_compiler.py

@@ -149,6 +149,7 @@ def _create_train_job(framework_version, instance_type, training_compiler_config
         "environment": None,
         "retry_strategy": None,
         "experiment_config": EXPERIMENT_CONFIG,
+        "enable_network_isolation": False,
         "debugger_hook_config": {
             "CollectionConfigurations": [],
             "S3OutputPath": "s3://{}/".format(BUCKET_NAME),

tests/unit/test_chainer.py

@@ -150,6 +150,7 @@ def _create_train_job(version, py_version):
         "tags": None,
         "vpc_config": None,
         "metric_definitions": None,
+        "enable_network_isolation": False,
         "environment": None,
         "experiment_config": None,
         "debugger_hook_config": {

tests/unit/test_estimator.py

@@ -1895,6 +1895,7 @@ def test_framework_with_spot_and_checkpoints(sagemaker_session):
         "encrypt_inter_container_traffic": True,
         "use_spot_instances": True,
         "checkpoint_s3_uri": "s3://mybucket/checkpoints/",
+        "enable_network_isolation": False,
         "checkpoint_local_path": "/tmp/checkpoints",
         "environment": None,
         "experiment_config": None,
@@ -3441,6 +3442,7 @@ def test_unsupported_type_in_dict():
     "vpc_config": None,
     "metric_definitions": None,
     "environment": None,
+    "enable_network_isolation": False,
     "experiment_config": None,
 }
 
@@ -3831,7 +3833,7 @@ def test_generic_to_fit_with_network_isolation(sagemaker_session):
 
     sagemaker_session.train.assert_called_once()
     args = sagemaker_session.train.call_args[1]
-    assert args["enable_network_isolation"]
+    assert args["enable_network_isolation"] is True
 
 
 def test_generic_to_fit_with_sagemaker_metrics_missing(sagemaker_session):

tests/unit/test_mxnet.py

@@ -167,6 +167,7 @@ def _get_train_args(job_name):
         "environment": None,
         "retry_strategy": None,
         "experiment_config": None,
+        "enable_network_isolation": False,
         "debugger_hook_config": {
             "CollectionConfigurations": [],
             "S3OutputPath": "s3://{}/".format(BUCKET_NAME),

tests/unit/test_pytorch.py

@@ -165,6 +165,7 @@ def _create_train_job(version, py_version):
         "environment": None,
         "retry_strategy": None,
         "experiment_config": None,
+        "enable_network_isolation": False,
         "debugger_hook_config": {
             "CollectionConfigurations": [],
             "S3OutputPath": "s3://{}/".format(BUCKET_NAME),

tests/unit/test_rl.py

@@ -155,6 +155,7 @@ def _create_train_job(toolkit, toolkit_version, framework):
         ],
         "environment": None,
         "experiment_config": None,
+        "enable_network_isolation": False,
         "debugger_hook_config": {
             "CollectionConfigurations": [],
             "S3OutputPath": "s3://{}/".format(BUCKET_NAME),

tests/unit/test_sklearn.py

@@ -142,6 +142,7 @@ def _create_train_job(version):
         "vpc_config": None,
         "environment": None,
         "experiment_config": None,
+        "enable_network_isolation": False,
         "debugger_hook_config": {
             "CollectionConfigurations": [],
             "S3OutputPath": "s3://{}/".format(BUCKET_NAME),

tests/unit/test_xgboost.py

@@ -155,6 +155,7 @@ def _create_train_job(version, instance_count=1, instance_type="ml.c4.4xlarge"):
         "vpc_config": None,
         "environment": None,
         "experiment_config": None,
+        "enable_network_isolation": False,
         "debugger_hook_config": {
             "CollectionConfigurations": [],
             "S3OutputPath": "s3://{}/".format(BUCKET_NAME),